| 263 | REGISTER_KERNEL_BUILDER(Name("_RefSend").Device(DEVICE_CPU), RefSendOp); |
| 264 | |
| 265 | RefRecvOp::RefRecvOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) { |
| 266 | string send_device; |
| 267 | OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device)); |
| 268 | string recv_device; |
| 269 | OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device)); |
| 270 | if (send_device != recv_device) { |
| 271 | LOG(FATAL) << "RefRecvOp require send_device equal recv_device, currently send_device" |
| 272 | << send_device << ", recv_device " << recv_device; |
| 273 | } |
| 274 | uint64 send_device_incarnation; |
| 275 | OP_REQUIRES_OK( |
| 276 | ctx, ctx->GetAttr("send_device_incarnation", |
| 277 | reinterpret_cast<int64*>(&send_device_incarnation))); |
| 278 | string tensor_name; |
| 279 | OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name)); |
| 280 | key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device, |
| 281 | send_device_incarnation, tensor_name); |
| 282 | // The vast majority of Recv nodes are outside any loop context, so |
| 283 | // proactively cache the rendezvous key for the top-level. |
| 284 | GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_); |
| 285 | OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_)); |
| 286 | if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) { |
| 287 | hostmem_sendrecv_ = false; |
| 288 | } |
| 289 | } |
| 290 | |
| 291 | void RefRecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { |
| 292 | OP_REQUIRES( |
nothing calls this directly
no test coverage detected