| 1054 | } // namespace |
| 1055 | |
| 1056 | Status EagerCopyToDevice(TensorHandle* h, EagerContext* ctx, |
| 1057 | EagerExecutor* executor, Device* device, bool mirror, |
| 1058 | TensorHandle** result) { |
| 1059 | Device* send_device = h->DeviceOrHostCPU(ctx); |
| 1060 | |
| 1061 | bool sender_is_local = send_device->IsLocal(); |
| 1062 | |
| 1063 | bool recver_is_local = device->IsLocal(); |
| 1064 | |
| 1065 | if (sender_is_local && recver_is_local) { |
| 1066 | return LocalEagerCopyToDevice(h, ctx, executor, device, result); |
| 1067 | } else { |
| 1068 | #if defined(IS_MOBILE_PLATFORM) |
| 1069 | return errors::Unimplemented( |
| 1070 | "Eager's remote execution is not available on mobile devices."); |
| 1071 | #else // !IS_MOBILE_PLATFORM |
| 1072 | if (mirror) { |
| 1073 | if (h->HasRemoteMirror(device)) { |
| 1074 | h->Ref(); |
| 1075 | *result = h; |
| 1076 | return Status::OK(); |
| 1077 | } |
| 1078 | } |
| 1079 | uint64 recv_op_id = 0; |
| 1080 | if (recver_is_local) { |
| 1081 | TF_RETURN_IF_ERROR(TensorHandle::CreateAsyncLocalHandle( |
| 1082 | /* d= */ device, |
| 1083 | /* op_device= */ device, /*resource_device=*/nullptr, h->dtype, ctx, |
| 1084 | result)); |
| 1085 | } else { |
| 1086 | eager::EagerClient* eager_client; |
| 1087 | uint64 context_id = ctx->GetContextId(); |
| 1088 | TF_RETURN_IF_ERROR(ctx->GetClient(device, &eager_client)); |
| 1089 | recv_op_id = ctx->RemoteMgr()->NextOpId(); |
| 1090 | auto tensor_handle_data = |
| 1091 | absl::make_unique<UnshapedRemoteTensorHandleData>( |
| 1092 | recv_op_id, 0, eager_client, context_id, ctx); |
| 1093 | if (mirror) { |
| 1094 | TF_RETURN_IF_ERROR( |
| 1095 | h->AddUnshapedRemoteMirror(std::move(tensor_handle_data), device)); |
| 1096 | h->Ref(); |
| 1097 | *result = h; |
| 1098 | } else { |
| 1099 | TF_RETURN_IF_ERROR(TensorHandle::CreateUnshapedRemoteHandle( |
| 1100 | std::move(tensor_handle_data), h->dtype, device, ctx, result)); |
| 1101 | } |
| 1102 | } |
| 1103 | auto node = absl::make_unique<eager::RemoteCopyNode>( |
| 1104 | ctx, executor, h, result[0], device, recv_op_id); |
| 1105 | Status s = executor->Async() ? executor->Add(std::move(node)) : node->Run(); |
| 1106 | if (!s.ok()) { |
| 1107 | result[0]->Unref(); |
| 1108 | } |
| 1109 | return s; |
| 1110 | #endif // !IS_MOBILE_PLATFORM |
| 1111 | } |
| 1112 | } |
| 1113 | } // namespace tensorflow |
no test coverage detected