| 131 | } |
| 132 | |
| 133 | Status SliceSendOp::SendString(OpKernelContext* ctx, |
| 134 | const FrameAndIter& frame_iter, |
| 135 | const Tensor& input_t) { |
| 136 | Rendezvous::Args args; |
| 137 | args.device_context = ctx->op_device_context(); |
| 138 | args.alloc_attrs = AllocatorAttributes(); |
| 139 | Rendezvous::ParsedKey parsed_key; |
| 140 | |
| 141 | // send elements bytes. |
| 142 | Tensor elements_bytes_t; |
| 143 | TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_UINT64, input_t.shape(), |
| 144 | &elements_bytes_t)); |
| 145 | int64 num_elements = input_t.NumElements(); |
| 146 | auto input_flat = input_t.flat<tstring>(); |
| 147 | auto elements_bytes_flat = elements_bytes_t.flat<uint64>(); |
| 148 | for (int64 i = 0; i < num_elements; i++) { |
| 149 | elements_bytes_flat(i) = input_flat(i).size(); |
| 150 | } |
| 151 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, |
| 152 | "_slice_transfer_elements_bytes", frame_iter, &parsed_key.buf_); |
| 153 | VLOG(2) << "SliceSend " << parsed_key.buf_; |
| 154 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 155 | TF_RETURN_IF_ERROR(ctx->rendezvous()->Send(parsed_key, args, elements_bytes_t, |
| 156 | ctx->is_input_dead())); |
| 157 | |
| 158 | // send data. |
| 159 | args.alloc_attrs = ctx->input_alloc_attr(0); |
| 160 | Tensor data_t; |
| 161 | for (int64 i = 0; i < num_elements; i++) { |
| 162 | const std::string& elem = input_flat(i); |
| 163 | if (elem.size() <= slice_size_) { |
| 164 | TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_STRING, TensorShape({}), |
| 165 | &data_t)); |
| 166 | data_t.scalar<tstring>()() = elem; |
| 167 | std::string tensor_name_suffix = \ |
| 168 | strings::StrCat("_slice_transfer_data_", std::to_string(i)); |
| 169 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, tensor_name_suffix, |
| 170 | frame_iter, &parsed_key.buf_); |
| 171 | VLOG(2) << "SliceSend " << parsed_key.buf_; |
| 172 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 173 | TF_RETURN_IF_ERROR( |
| 174 | ctx->rendezvous()->FlowControlSend(tensor_name_, parsed_key, args, |
| 175 | data_t, ctx->is_input_dead())); |
| 176 | } else { |
| 177 | TF_RETURN_IF_ERROR(SendStringSlice(ctx, frame_iter, elem, i)); |
| 178 | } |
| 179 | } |
| 180 | |
| 181 | return Status::OK(); |
| 182 | } |
| 183 | |
| 184 | Status SliceSendOp::SendStringSlice(OpKernelContext* ctx, |
| 185 | const FrameAndIter& frame_iter, |
nothing calls this directly
no test coverage detected