| 274 | } |
| 275 | |
| 276 | XrtTensorHandle XrtTfContext::SendTensor( |
| 277 | std::unique_ptr<TensorProto> tensor_proto, int device_id, |
| 278 | bool host_memory) { |
| 279 | DataType dtype = tensor_proto->dtype(); |
| 280 | bool transfer_via_cpu_device = host_memory && device_id != cpu_device_id_; |
| 281 | int rpc_device_id = transfer_via_cpu_device ? cpu_device_id_ : device_id; |
| 282 | OperationId op_id; |
| 283 | { |
| 284 | absl::MutexLock lock(&mu_); |
| 285 | Operation* op = AddOperation(); |
| 286 | op_id = op->id; |
| 287 | } |
| 288 | |
| 289 | eager::EnqueueRequest request; |
| 290 | request.set_context_id(context_id_); |
| 291 | auto* send_tensor = request.add_queue()->mutable_send_tensor(); |
| 292 | send_tensor->set_op_id(op_id); |
| 293 | send_tensor->mutable_tensors()->AddAllocated(tensor_proto.release()); |
| 294 | send_tensor->set_device_name(devices_.at(rpc_device_id).name()); |
| 295 | auto response = std::make_shared<eager::EnqueueResponse>(); |
| 296 | auto context_ptr = shared_from_this(); |
| 297 | absl::Notification done; |
| 298 | eager_client_->EnqueueAsync( |
| 299 | &request, response.get(), |
| 300 | [context_ptr, op_id, response, &done](Status status) { |
| 301 | absl::MutexLock lock(&context_ptr->mu_); |
| 302 | if (!status.ok()) { |
| 303 | context_ptr->ReportError({op_id}, status); |
| 304 | } else { |
| 305 | context_ptr->DeleteOperation(op_id); |
| 306 | } |
| 307 | done.Notify(); |
| 308 | }); |
| 309 | XrtTensorHandle handle(context_ptr, rpc_device_id, TensorId{op_id, 0}); |
| 310 | |
| 311 | // TODO(phawkins): we block here to avoid a race. We must not |
| 312 | // enqueue any dependent operations until the SendTensor has been |
| 313 | // acknowledged. |
| 314 | done.WaitForNotification(); |
| 315 | |
| 316 | // TODO(phawkins): EagerService.SendTensor could use a host_memory option. |
| 317 | if (!transfer_via_cpu_device) { |
| 318 | return handle; |
| 319 | } |
| 320 | std::string wire_id = XrtGetUniqueWireID(); |
| 321 | EnqueueSend(this, handle, dtype, device_id, wire_id, /*host_memory=*/false); |
| 322 | return EnqueueRecv(this, dtype, rpc_device_id, device_id, wire_id, |
| 323 | /*host_memory=*/true); |
| 324 | } |
| 325 | |
| 326 | // This gets a unique wire ID. We add a random identifier so that if the |
| 327 | // worker has other clients that it is servicing, we don't have any collision. |