Searched refs:transfer_stream (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_executable_interface.cc | 35 se::Stream* transfer_stream) { in PopulateResultTupleBuffers() argument
41 transfer_stream ? transfer_stream : stream, result)); in PopulateResultTupleBuffers()
42 if (transfer_stream && transfer_stream != stream) { in PopulateResultTupleBuffers()
43 stream->ThenWaitFor(transfer_stream); in PopulateResultTupleBuffers()
58 se::Stream* transfer_stream) { in AllocateOutputMemoryWithInputReuse() argument
164 PopulateResultTupleBuffers(result.Result(), stream, transfer_stream)); in AllocateOutputMemoryWithInputReuse()
|
| D | tpu_executable_interface.h | 70 se::Stream* transfer_stream = nullptr);
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tpu_client.cc | 54 Status ThenMemcpyDeviceToDevice(se::Stream* transfer_stream,
67 se::Stream* transfer_stream, se::Stream* dst_stream, in ThenMemcpyDeviceToDevice() argument
70 transfer_stream->implementation()); in ThenMemcpyDeviceToDevice()
|
| D | local_device_state.cc | 90 se::Stream* transfer_stream, se::Stream* dst_stream, in ThenMemcpyDeviceToDevice() argument
95 transfer_stream->ThenMemcpyD2D(&dst_buffer, src_buffer, dst_buffer.size()); in ThenMemcpyDeviceToDevice()
|
| D | pjrt_stream_executor_client.cc | 1295 LocalDeviceState* transfer_local_device, se::Stream* transfer_stream, in CopyToDeviceHelper() argument
1300 dst_device, dst_local_device, transfer_stream, in CopyToDeviceHelper()
1305 WaitForBufferDefinitionEventsOnStream(*src_device_buffer, transfer_stream); in CopyToDeviceHelper()
1323 transfer_stream, dst_local_device->compute_stream(), input_buffer, in CopyToDeviceHelper()
1331 transfer_stream)); in CopyToDeviceHelper()
1335 StallStreamOnError(transfer_local_device, transfer_stream); in CopyToDeviceHelper()
1343 ->ThenRelease(transfer_stream, src_device_buffer); in CopyToDeviceHelper()
1387 se::Stream* transfer_stream = in CopyToDevice() local
1405 dst_device, dst_local_device, transfer_local_device, transfer_stream, in CopyToDevice()
1424 transfer_local_device, event, transfer_stream, in CopyToDevice()
|
| D | local_device_state.h | 128 virtual Status ThenMemcpyDeviceToDevice(se::Stream* transfer_stream,
|
| D | pjrt_stream_executor_client.h | 584 se::Stream* transfer_stream,
|