Home
last modified time | relevance | path

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

/external/ml_dtypes/ml_dtypes/tests/
Dintn_test.cc405 IntN* src_device_buffer = (IntN*)device.allocate(kNumElems * sizeof(IntN)); in TYPED_TEST() local
410 src_device_buffer, kNumElems); in TYPED_TEST()
423 device.memcpyHostToDevice(src_device_buffer, src_cpu.data(), in TYPED_TEST()
438 device.memcpyHostToDevice(src_device_buffer, src_cpu.data(), in TYPED_TEST()
441 device.memcpyDeviceToHost(src_cpu.data(), src_device_buffer, in TYPED_TEST()
451 device.deallocate(src_device_buffer); in TYPED_TEST()
Dfloat8_test.cc1094 Float8* src_device_buffer = in TYPED_TEST() local
1100 src_device_buffer, kNumElems); in TYPED_TEST()
1122 device.memcpyHostToDevice(src_device_buffer, src_cpu.data(), in TYPED_TEST()
1137 device.memcpyHostToDevice(src_device_buffer, src_cpu.data(), in TYPED_TEST()
1140 device.memcpyDeviceToHost(src_cpu.data(), src_device_buffer, in TYPED_TEST()
1150 device.deallocate(src_device_buffer); in TYPED_TEST()
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
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()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
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()
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()
Dpjrt_stream_executor_client.h695 std::shared_ptr<TrackedDeviceBuffer> src_device_buffer);