| 674 | } |
| 675 | |
| 676 | Status EagerRemoteExecute(EagerOperation* op, TensorHandle** retvals, |
| 677 | int* num_retvals) { |
| 678 | EagerContext* ctx = op->EagerContext(); |
| 679 | |
| 680 | // TODO(fishx): Remove following code when lazy tensor copy is ready. |
| 681 | if (op->Device() == nullptr) { |
| 682 | tensorflow::Device* device = nullptr; |
| 683 | string device_name = |
| 684 | DeviceNameUtils::ParsedNameToString(op->GetDeviceName()); |
| 685 | TF_RETURN_IF_ERROR(ctx->FindDeviceByName(device_name, &device)); |
| 686 | op->SetDevice(device); |
| 687 | } |
| 688 | |
| 689 | eager::EagerClient* eager_client = nullptr; |
| 690 | uint64 context_id = ctx->GetContextId(); |
| 691 | TF_RETURN_IF_ERROR(ctx->GetClient(op->GetDeviceName(), &eager_client)); |
| 692 | |
| 693 | std::unique_ptr<eager::EnqueueRequest> request(new eager::EnqueueRequest); |
| 694 | request->set_context_id(context_id); |
| 695 | |
| 696 | eager::Operation* remote_op = request->add_queue()->mutable_operation(); |
| 697 | |
| 698 | { |
| 699 | profiler::TraceMe activity("CopyInputToExpectedDevice", |
| 700 | profiler::TraceMeLevel::kInfo); |
| 701 | for (int i = 0; i < op->Inputs().size(); i++) { |
| 702 | tensorflow::TensorHandle* input = op->Inputs()[i]; |
| 703 | tensorflow::Device* input_device = input->device(); |
| 704 | const string* input_device_name = &input->DeviceOrHostCPU(ctx)->name(); |
| 705 | if (op->Device() != input_device && |
| 706 | // If the expected and actual devices are on the same task, don't |
| 707 | // explicitly copy, and instead depend on the copy to happen locally |
| 708 | // when the op is executed on the device. |
| 709 | !ctx->OnSameTask(op->Device(), input_device)) { |
| 710 | tensorflow::Device* remote_cpu_device; |
| 711 | TF_RETURN_IF_ERROR( |
| 712 | ctx->CPUDeviceOnTask(op->Device(), &remote_cpu_device)); |
| 713 | // TODO(b/110044833): It's possible the same tensor gets copied to the |
| 714 | // remote device repeatedly. |
| 715 | // Always copy to the remote CPU so that the actual device can be |
| 716 | // correctly determined after the kernel is selected/instantiated, since |
| 717 | // the op might have its inputs on host memory. |
| 718 | TensorHandle* handle = nullptr; |
| 719 | TF_RETURN_IF_ERROR(MaybeCopyInputToExpectedDevice( |
| 720 | op, op->Device(), i, remote_cpu_device, &handle)); |
| 721 | op->UpdateInput(i, handle); |
| 722 | input = handle; |
| 723 | input_device = remote_cpu_device; |
| 724 | input_device_name = &remote_cpu_device->name(); |
| 725 | // Unref handle since it has a ref as an input now |
| 726 | handle->Unref(); |
| 727 | } |
| 728 | |
| 729 | TF_RETURN_IF_ERROR(ctx->RemoteMgr()->SerializeRemoteTensorHandle( |
| 730 | input, remote_op->add_inputs(), input_device, *input_device_name)); |
| 731 | } |
| 732 | } |
| 733 |
no test coverage detected