Searched refs:send_device_to_host_stream (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/gpu/ |
| H A D | gpu_util.cc | 127 auto send_device_to_host_stream = in SetProtoFromGPU() local
130 if (send_device_to_host_stream == nullptr) { in SetProtoFromGPU()
135 send_device_to_host_stream->ThenWaitFor(send_stream); in SetProtoFromGPU()
161 send_device_to_host_stream->ThenMemcpy(buf, gpu_src_ptr, total_bytes); in SetProtoFromGPU()
166 send_device_to_host_stream, [send_device_to_host_stream, done, proto, buf, in SetProtoFromGPU()
168 if (!send_device_to_host_stream->ok()) { in SetProtoFromGPU()
280 auto send_device_to_host_stream = in CopyGPUTensorToCPU() local
283 if (send_device_to_host_stream == nullptr) { in CopyGPUTensorToCPU()
288 send_device_to_host_stream->ThenWaitFor(send_stream); in CopyGPUTensorToCPU()
295 send_device_to_host_stream->ThenMemcpy(dst_ptr, gpu_src_ptr, total_bytes); in CopyGPUTensorToCPU()
[all …]
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/pluggable_device/ |
| H A D | pluggable_device_util.cc | 187 auto send_device_to_host_stream = in CopyPluggableDeviceTensorToCPU() local
190 if (send_device_to_host_stream == nullptr) { in CopyPluggableDeviceTensorToCPU()
196 send_device_to_host_stream->ThenWaitFor(send_stream); in CopyPluggableDeviceTensorToCPU()
203 send_device_to_host_stream->ThenMemcpy(dst_ptr, device_src_ptr, in CopyPluggableDeviceTensorToCPU()
210 send_device_to_host_stream, in CopyPluggableDeviceTensorToCPU()
211 [send_device_to_host_stream, done, input_ref]() { in CopyPluggableDeviceTensorToCPU()
212 if (!send_device_to_host_stream->ok()) { in CopyPluggableDeviceTensorToCPU()
|