| 204 | } |
| 205 | |
| 206 | void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, |
| 207 | absl::string_view tensor_name, |
| 208 | Device* device, Tensor* cpu_tensor, |
| 209 | StatusCallback done) { |
| 210 | if (device_tensor->NumElements() == 0) { |
| 211 | VLOG(2) << "CopyDeviceTensorToCPU empty tensor"; |
| 212 | done(Status::OK()); |
| 213 | return; |
| 214 | } |
| 215 | VLOG(2) << "CopyDeviceTensorToCPU " |
| 216 | << reinterpret_cast<const void*>(device_tensor->tensor_data().data()) |
| 217 | << " " |
| 218 | << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data()) |
| 219 | << " " << device_tensor->NumElements() << " " |
| 220 | << cpu_tensor->shape().DebugString() << " " |
| 221 | << device_tensor->shape().DebugString(); |
| 222 | |
| 223 | std::shared_ptr<se::Stream> device_to_host_stream; |
| 224 | if (device_to_host_stream_) { |
| 225 | device_to_host_stream = device_to_host_stream_; |
| 226 | } else { |
| 227 | stream_executor::port::StatusOr<xla::StreamPool::Ptr> ptr_or_status = |
| 228 | client_->mutable_backend()->BorrowStream( |
| 229 | stream_->parent()->device_ordinal()); |
| 230 | if (!ptr_or_status.status().ok()) { |
| 231 | done(ptr_or_status.status()); |
| 232 | return; |
| 233 | } |
| 234 | device_to_host_stream = |
| 235 | std::shared_ptr<se::Stream>(std::move(ptr_or_status.ValueOrDie())); |
| 236 | } |
| 237 | |
| 238 | XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); |
| 239 | xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream.get()); |
| 240 | |
| 241 | // Transfer manager requires the shape of the shaped buffer to be the same as |
| 242 | // literal shape except for the layout. Set the literal to use xla_tensor's |
| 243 | // shape as it is derived from the cpu_tensor's shape using |
| 244 | // shape_representation_fn_. |
| 245 | xla::MutableBorrowingLiteral literal; |
| 246 | TF_CHECK_OK(HostTensorToMutableBorrowingLiteral( |
| 247 | xla::LayoutUtil::GetWithDefaultLayout( |
| 248 | xla_tensor->shaped_buffer().on_host_shape()), |
| 249 | cpu_tensor, &literal)); |
| 250 | |
| 251 | TensorReference ref(*device_tensor); |
| 252 | const bool device_allows_sync_on_completion = |
| 253 | device->AllowsSyncOnCompletion(); |
| 254 | // Explicitly capture device_to_host_stream to make sure the stream is alive |
| 255 | // before the transfer finishes. |
| 256 | transfer_manager_->TransferLiteralFromDevice( |
| 257 | device_to_host_stream.get(), xla_tensor->shaped_buffer(), literal, |
| 258 | [ref, xla_tensor, done, device_to_host_stream, |
| 259 | device_allows_sync_on_completion](xla::Status status) { |
| 260 | Status done_status = status; |
| 261 | VLOG(2) << "Transfer from device as literal: " |
| 262 | << xla_tensor->shaped_buffer().ToString(); |
| 263 | // For devices don't allow sync on completion, the device execution is |