| 415 | } |
| 416 | |
| 417 | StatusOr<TransferToServerResponse> LocalClient::TransferToLocalServer( |
| 418 | const ::xla::BorrowingLiteral& literal, int device_ordinal) { |
| 419 | const ::xla::Shape& shape = literal.shape(); |
| 420 | |
| 421 | TF_ASSIGN_OR_RETURN(::xla::ScopedShapedBuffer shaped_buffer, |
| 422 | backend().transfer_manager()->AllocateScopedShapedBuffer( |
| 423 | shape, backend().memory_allocator(), device_ordinal)); |
| 424 | TF_ASSIGN_OR_RETURN(auto stream, |
| 425 | mutable_backend()->BorrowStream(device_ordinal)); |
| 426 | TF_RETURN_IF_ERROR(backend().transfer_manager()->TransferLiteralToDevice( |
| 427 | stream.get(), literal, shaped_buffer)); |
| 428 | std::vector<::xla::ScopedShapedBuffer> replicated_buffer; |
| 429 | replicated_buffer.emplace_back(std::move(shaped_buffer)); |
| 430 | ::xla::TransferToServerResponse result; |
| 431 | TF_ASSIGN_OR_RETURN(*result.mutable_data(), |
| 432 | local_service_->RegisterReplicatedBuffers( |
| 433 | std::move(replicated_buffer), |
| 434 | absl::StrCat("TransferToServer literal of shape ", |
| 435 | ::xla::ShapeUtil::HumanString(shape)))); |
| 436 | |
| 437 | return result; |
| 438 | } |
| 439 | |
| 440 | } // namespace xla |
nothing calls this directly
no test coverage detected