Home
last modified time | relevance | path

Searched refs:src_device_buffer (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
H A Dtpu_client.cc408 std::shared_ptr<TpuSharedBuffer> src_device_buffer = DeviceBuffer(); in CopyToDevice() local
411 on_host_shape_, src_device_buffer, in CopyToDevice()
420 [driver, src_device_buffer](tpu_driver::BufferHandle* dst_handle) { in CopyToDevice()
422 for (auto& event : src_device_buffer->wait_for_use) { in CopyToDevice()
426 src_device_buffer->handle.get(), dst_handle, src_wait_for_use); in CopyToDevice()
433 src_device_buffer->wait_for_use.insert(src_device_buffer->wait_for_use.end(), in CopyToDevice()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dpjrt_stream_executor_client.cc1420 std::shared_ptr<TrackedDeviceBuffer> src_device_buffer) { in CopyToDeviceHelper() argument
1429 WaitForBufferDefinitionEventsOnStream(*src_device_buffer, transfer_stream); in CopyToDeviceHelper()
1466 std::move(src_device_buffer)); in CopyToDeviceHelper()
1517 ScopedHold src_device_buffer(this, ScopedHold::kUsage); in CopyToDevice() local
1526 AcquireHoldLocked(&src_device_buffer); in CopyToDevice()
1533 src_device_buffer.buffer()); in CopyToDevice()
1548 RecordUsage(std::move(src_device_buffer), device_->local_device_state(), in CopyToDevice()
H A Dtfrt_cpu_pjrt_client.cc1044 auto* src_device_buffer = AcquireUsage(usage_event); in CopyToDevice() local
1045 if (src_device_buffer == nullptr) { in CopyToDevice()
1050 int num_leaf_buffers = src_device_buffer->Buffers().size(); in CopyToDevice()
1059 auto src_buffer = src_device_buffer->Buffers()[i]; in CopyToDevice()
1070 const auto& src_definition_event = src_device_buffer->definition_event(); in CopyToDevice()
H A Dpjrt_stream_executor_client.h695 std::shared_ptr<TrackedDeviceBuffer> src_device_buffer);