| 53 | } |
| 54 | |
| 55 | SendOp::SendOp(OpKernelConstruction* ctx) : OpKernel(ctx) { |
| 56 | string send_device; |
| 57 | OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device", &send_device)); |
| 58 | string recv_device; |
| 59 | OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_device", &recv_device)); |
| 60 | uint64 send_device_incarnation; |
| 61 | OP_REQUIRES_OK( |
| 62 | ctx, ctx->GetAttr("send_device_incarnation", |
| 63 | reinterpret_cast<int64*>(&send_device_incarnation))); |
| 64 | string tensor_name; |
| 65 | OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name)); |
| 66 | key_prefix_ = GetRendezvousKeyPrefix(send_device, recv_device, |
| 67 | send_device_incarnation, tensor_name); |
| 68 | // The vast majority of Send nodes are outside any loop context, so |
| 69 | // proactively cache the rendezvous key for the top-level. |
| 70 | GetRendezvousKey(key_prefix_, {0, 0}, &parsed_key_.buf_); |
| 71 | OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key_.buf_, &parsed_key_)); |
| 72 | if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) { |
| 73 | hostmem_sendrecv_ = false; |
| 74 | } |
| 75 | } |
| 76 | |
| 77 | void SendOp::Compute(OpKernelContext* ctx) { |
| 78 | OP_REQUIRES( |
nothing calls this directly
no test coverage detected