| 199 | Name("_HostRecv").Device(DEVICE_DEFAULT).HostMemory("tensor"), RecvOp); |
| 200 | |
| 201 | RefSendOp::RefSendOp(OpKernelConstruction* ctx) : OpKernel(ctx) { |
| 202 | string send_device; |
| 203 | OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device)); |
| 204 | string recv_device; |
| 205 | OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device)); |
| 206 | if (send_device != recv_device) { |
| 207 | LOG(FATAL) << "RefSendOp require send_device equal recv_device, currently send_device" |
| 208 | << send_device << ", recv_device " << recv_device; |
| 209 | } |
| 210 | uint64 send_device_incarnation; |
| 211 | OP_REQUIRES_OK( |
| 212 | ctx, ctx->GetAttr("send_device_incarnation", |
| 213 | reinterpret_cast<int64*>(&send_device_incarnation))); |
| 214 | string tensor_name; |
| 215 | OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name)); |
| 216 | key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device, |
| 217 | send_device_incarnation, tensor_name); |
| 218 | // The vast majority of Send nodes are outside any loop context, so |
| 219 | // proactively cache the rendezvous key for the top-level. |
| 220 | GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_); |
| 221 | OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_)); |
| 222 | if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) { |
| 223 | hostmem_sendrecv_ = false; |
| 224 | } |
| 225 | } |
| 226 | |
| 227 | void RefSendOp::Compute(OpKernelContext* ctx) { |
| 228 | OP_REQUIRES( |
nothing calls this directly
no test coverage detected