| 990 | } |
| 991 | |
| 992 | Status Service::TransferToServer(const TransferToServerRequest* arg, |
| 993 | TransferToServerResponse* result) { |
| 994 | TF_ASSIGN_OR_RETURN(Literal literal, |
| 995 | Literal::CreateFromProto(arg->literal())); |
| 996 | const Shape& shape = literal.shape(); |
| 997 | |
| 998 | std::vector<se::StreamExecutor*> replicas; |
| 999 | if (arg->has_device_handle()) { |
| 1000 | TF_ASSIGN_OR_RETURN(replicas, |
| 1001 | Replicas(*execute_backend_, arg->device_handle())); |
| 1002 | } else { |
| 1003 | TF_ASSIGN_OR_RETURN( |
| 1004 | replicas, Replicas(*execute_backend_, SingleComputationDeviceHandle())); |
| 1005 | } |
| 1006 | |
| 1007 | // Allocate memory in each replica and transfer the data to all replicas. |
| 1008 | std::vector<ScopedShapedBuffer> replicated_buffers; |
| 1009 | for (se::StreamExecutor* executor : replicas) { |
| 1010 | TF_ASSIGN_OR_RETURN( |
| 1011 | ScopedShapedBuffer shaped_buffer, |
| 1012 | execute_backend_->transfer_manager()->AllocateScopedShapedBuffer( |
| 1013 | shape, execute_backend_->memory_allocator(), |
| 1014 | executor->device_ordinal())); |
| 1015 | TF_ASSIGN_OR_RETURN(auto stream, execute_backend_->BorrowStream(executor)); |
| 1016 | TF_RETURN_IF_ERROR( |
| 1017 | execute_backend_->transfer_manager()->TransferLiteralToDevice( |
| 1018 | stream.get(), literal, shaped_buffer)); |
| 1019 | replicated_buffers.emplace_back(std::move(shaped_buffer)); |
| 1020 | } |
| 1021 | TF_ASSIGN_OR_RETURN(*result->mutable_data(), |
| 1022 | allocation_tracker_.RegisterReplicatedBuffers( |
| 1023 | std::move(replicated_buffers), |
| 1024 | StrCat("TransferToServer literal of shape ", |
| 1025 | ShapeUtil::HumanString(shape)))); |
| 1026 | |
| 1027 | return Status::OK(); |
| 1028 | } |
| 1029 | |
| 1030 | Status Service::TransferToInfeed(const TransferToInfeedRequest* arg, |
| 1031 | TransferToInfeedResponse* result) { |
no test coverage detected