| 114 | Name("_HostSend").Device(DEVICE_DEFAULT).HostMemory("tensor"), SendOp); |
| 115 | |
| 116 | RecvOp::RecvOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) { |
| 117 | string send_device; |
| 118 | OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device)); |
| 119 | string recv_device; |
| 120 | OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device)); |
| 121 | uint64 send_device_incarnation; |
| 122 | OP_REQUIRES_OK( |
| 123 | ctx, ctx->GetAttr("send_device_incarnation", |
| 124 | reinterpret_cast<int64*>(&send_device_incarnation))); |
| 125 | string tensor_name; |
| 126 | OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name)); |
| 127 | key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device, |
| 128 | send_device_incarnation, tensor_name); |
| 129 | // The vast majority of Recv nodes are outside any loop context, so |
| 130 | // proactively cache the rendezvous key for the top-level. |
| 131 | GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_); |
| 132 | OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_)); |
| 133 | if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) { |
| 134 | hostmem_sendrecv_ = false; |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | namespace { |
| 139 | Rendezvous::DoneCallback make_recv_callback(OpKernelContext* ctx, |
nothing calls this directly
no test coverage detected