| 603 | } |
| 604 | |
| 605 | StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalExecutable::ExecuteHelper( |
| 606 | absl::Span<PyLocalBuffer* const> argument_handles, int replica, |
| 607 | int partition, const RunId& run_id) { |
| 608 | const int device_id = (*device_assignment_)(replica, partition); |
| 609 | std::shared_ptr<Device> device = LookupDevice(*client_, device_id); |
| 610 | CHECK_EQ(device->host_id(), client_->host_id()); |
| 611 | int device_ordinal = device->local_device_state()->device_ordinal(); |
| 612 | tensorflow::profiler::TraceMe traceme("LocalExecutable::Execute"); |
| 613 | VLOG(3) << "Replica " << replica << ", partition " << partition |
| 614 | << " mapped to device ordinal for execution: " << device_ordinal; |
| 615 | |
| 616 | absl::flat_hash_set<BufferDefinitionEvent*> events; |
| 617 | std::vector<std::shared_ptr<SharedDeviceBuffer>> device_buffers; |
| 618 | std::vector<ShapedBuffer> argument_buffers; |
| 619 | std::vector<const ShapedBuffer*> argument_buffer_ptrs; |
| 620 | device_buffers.reserve(argument_handles.size() + 1); |
| 621 | argument_buffers.reserve(argument_handles.size()); |
| 622 | argument_buffer_ptrs.reserve(argument_handles.size()); |
| 623 | for (int i = 0; i < argument_handles.size(); ++i) { |
| 624 | PyLocalBuffer* handle = argument_handles[i]; |
| 625 | std::shared_ptr<SharedDeviceBuffer> device_buffer = handle->DeviceBuffer(); |
| 626 | if (!device_buffer) { |
| 627 | return InvalidArgument( |
| 628 | "Deleted buffer passed to Execute() as argument %d to replica %d", i, |
| 629 | replica); |
| 630 | } |
| 631 | if (handle->device().get() != device.get()) { |
| 632 | return InvalidArgument( |
| 633 | "Buffer passed to Execute() as argument %d to replica %d is on " |
| 634 | "device %s, but replica is assigned to device %s.", |
| 635 | i, replica, handle->device()->DebugString(), device->DebugString()); |
| 636 | } |
| 637 | TF_ASSIGN_OR_RETURN(ShapedBuffer shaped_buffer, handle->AsShapedBuffer()); |
| 638 | argument_buffers.push_back(std::move(shaped_buffer)); |
| 639 | argument_buffer_ptrs.push_back(&argument_buffers.back()); |
| 640 | GetDeviceBufferDefinitionEvents(*device_buffer, &events); |
| 641 | device_buffers.push_back(std::move(device_buffer)); |
| 642 | VLOG(4) << "Argument " << i |
| 643 | << " buffer: " << argument_buffers.back().ToString(); |
| 644 | } |
| 645 | |
| 646 | LocalDeviceState* device_state = &client_->device_state(device_ordinal); |
| 647 | |
| 648 | for (BufferDefinitionEvent* event : events) { |
| 649 | event->WaitForEventOnStream(device_state->compute_stream()); |
| 650 | } |
| 651 | |
| 652 | ExecutableRunOptions options; |
| 653 | options.set_stream(device_state->compute_stream()); |
| 654 | options.set_host_to_device_stream(device_state->host_to_device_stream()); |
| 655 | options.set_allocator(client_->allocator()); |
| 656 | options.set_intra_op_thread_pool( |
| 657 | client_->client()->backend().eigen_intra_op_thread_pool_device()); |
| 658 | options.set_device_assignment(device_assignment_.get()); |
| 659 | options.set_run_id(run_id); |
| 660 | options.set_rng_seed(device_state->GetNewPrngSeed()); |
| 661 | |
| 662 | // The choice of where we wait is arbitrary; the reason for the wait is pacing |
nothing calls this directly
no test coverage detected