| 449 | } |
| 450 | |
| 451 | StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::CopyToDevice( |
| 452 | std::shared_ptr<Device> dst_device) { |
| 453 | tensorflow::profiler::TraceMe traceme("PyLocalBuffer::CopyToDevice"); |
| 454 | std::shared_ptr<SharedDeviceBuffer> src_device_buffer = DeviceBuffer(); |
| 455 | TF_ASSIGN_OR_RETURN(LocalDeviceState * dst_local_device, |
| 456 | dst_device->GetLocalDeviceState()); |
| 457 | |
| 458 | if (dst_device.get() == device_.get()) { |
| 459 | return absl::make_unique<PyLocalBuffer>( |
| 460 | on_host_shape_, on_device_shape_, src_device_buffer, client_, device_); |
| 461 | } |
| 462 | LocalDeviceState* transfer_local_device = |
| 463 | client_->EnqueueD2DTransfersOnSrcStream() ? device_->local_device_state() |
| 464 | : dst_local_device; |
| 465 | |
| 466 | se::Stream* transfer_stream = |
| 467 | transfer_local_device->GetDeviceToDeviceStream(); |
| 468 | |
| 469 | TransferManager* transfer_manager = |
| 470 | client_->client()->backend().transfer_manager(); |
| 471 | TF_ASSIGN_OR_RETURN(ScopedShapedBuffer dst_buffer, |
| 472 | transfer_manager->AllocateScopedShapedBuffer( |
| 473 | on_host_shape_, client_->allocator(), |
| 474 | dst_local_device->device_ordinal())); |
| 475 | if (!transfer_manager->CanShapedBufferBeAccessedNow( |
| 476 | dst_local_device->compute_stream()->parent(), dst_buffer)) { |
| 477 | transfer_stream->ThenWaitFor(dst_local_device->compute_stream()); |
| 478 | } |
| 479 | TF_ASSIGN_OR_RETURN(ShapedBuffer src_buffer, AsShapedBuffer()); |
| 480 | |
| 481 | WaitForBufferDefinitionEventsOnStream(*src_device_buffer, transfer_stream); |
| 482 | |
| 483 | // Copy the leaf buffers. |
| 484 | for (const auto& leaf : src_buffer.buffers().leaves()) { |
| 485 | const ShapeIndex& index = leaf.first; |
| 486 | const se::DeviceMemoryBase& input_buffer = leaf.second; |
| 487 | const se::DeviceMemoryBase& output_buffer = dst_buffer.buffer(index); |
| 488 | TF_RET_CHECK(input_buffer.size() == output_buffer.size()) |
| 489 | << "input: " << input_buffer.size() |
| 490 | << " output: " << output_buffer.size(); |
| 491 | TF_RETURN_IF_ERROR(transfer_local_device->ThenMemcpyDeviceToDevice( |
| 492 | transfer_stream, dst_local_device->compute_stream(), input_buffer, |
| 493 | output_buffer)); |
| 494 | } |
| 495 | |
| 496 | // We hold on to the `src_device_buffer` until the transfer is finished. |
| 497 | transfer_local_device->ThenRelease(transfer_stream, |
| 498 | std::move(src_device_buffer)); |
| 499 | |
| 500 | // Write new tuple buffers. The destination buffers have different addresses, |
| 501 | // so we must construct tuple buffers from scratch instead of copying them. |
| 502 | if (dst_buffer.on_device_shape().IsTuple()) { |
| 503 | TF_RETURN_IF_ERROR(transfer_manager->WriteTupleIndexTablesAsync( |
| 504 | dst_local_device->host_to_device_stream(), dst_buffer)); |
| 505 | |
| 506 | // We need a single definition event, so make the device to device stream |
| 507 | // wait for the stream that wrote the tuple index tables on the destination |
| 508 | // device. |
no test coverage detected