Searched refs:device_to_device_stream (Results 1 – 7 of 7) sorted by relevance
35 gtl::InlinedVector<se::Stream*, 4> device_to_device_stream) in PluggableDeviceContext() argument40 device_to_device_stream_(device_to_device_stream) {} in PluggableDeviceContext()47 se::Stream* device_to_device_stream(int index) const { in device_to_device_stream() function
126 ->device_to_device_stream(dev_to_dev_stream_index); in DeviceToDeviceCopy()
38 gtl::InlinedVector<se::Stream*, 4> device_to_device_stream, in GPUDeviceContext() argument47 device_to_device_stream_(device_to_device_stream), in GPUDeviceContext()59 se::Stream* device_to_device_stream(int index) const { in device_to_device_stream() function
80 se::Stream* device_to_device_stream(int index) const { in device_to_device_stream() function
296 return device_to_device_stream(stream); in GetDeviceToDeviceStream()
175 should_use_substream ? dst_xla_context->device_to_device_stream(0) in TpuDeviceToDeviceCopy()
202 ->device_to_device_stream(dev_to_dev_stream_index); in DeviceToDeviceCopy()