Searched refs:device_to_host_stream (Results 1 – 7 of 7) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/ |
| H A D | xla_device_context.cc | 90 std::shared_ptr<se::Stream> device_to_host_stream, in XlaDeviceContext() argument
97 device_to_host_stream_(std::move(device_to_host_stream)), in XlaDeviceContext()
226 std::shared_ptr<se::Stream> device_to_host_stream; in CopyDeviceTensorToCPU() local
228 device_to_host_stream = device_to_host_stream_; in CopyDeviceTensorToCPU()
237 device_to_host_stream = in CopyDeviceTensorToCPU()
242 xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream.get()); in CopyDeviceTensorToCPU()
260 device_to_host_stream.get(), xla_tensor->shaped_buffer(), literal, in CopyDeviceTensorToCPU()
261 [this, ref, xla_tensor, done, device_to_host_stream, in CopyDeviceTensorToCPU()
280 !device_to_host_stream->RefreshStatus().ok()) { in CopyDeviceTensorToCPU()
285 [device_to_host_stream] {}); in CopyDeviceTensorToCPU()
|
| H A D | xla_device.cc | 332 std::shared_ptr<se::Stream> device_to_host_stream; in GetDeviceContextLocked() local
348 device_to_host_stream = nullptr; in GetDeviceContextLocked()
351 device_to_host_stream = stream_; in GetDeviceContextLocked()
371 stream_, host_to_device_stream, device_to_host_stream, in GetDeviceContextLocked()
|
| H A D | xla_device_context.h | 59 std::shared_ptr<se::Stream> device_to_host_stream,
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/pluggable_device/ |
| H A D | pluggable_device_context.h | 34 se::Stream* device_to_host_stream, in PluggableDeviceContext() argument
39 device_to_host_stream_(device_to_host_stream), in PluggableDeviceContext()
46 se::Stream* device_to_host_stream() const { return device_to_host_stream_; } in device_to_host_stream() function
|
| H A D | pluggable_device_util.cc | 189 ->device_to_host_stream(); in CopyPluggableDeviceTensorToCPU()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/ |
| H A D | gpu_device_context.h | 37 se::Stream* device_to_host_stream, in GPUDeviceContext() argument
46 device_to_host_stream_(device_to_host_stream), in GPUDeviceContext()
58 se::Stream* device_to_host_stream() const { return device_to_host_stream_; } in device_to_host_stream() function
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/gpu/ |
| H A D | gpu_util.cc | 129 ->device_to_host_stream(); in SetProtoFromGPU()
282 ->device_to_host_stream(); in CopyGPUTensorToCPU()
|