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() argument46 device_to_device_stream_(device_to_device_stream) { in GPUDeviceContext()57 se::Stream* device_to_device_stream(int index) const { in device_to_device_stream() function
79 se::Stream* device_to_device_stream(int index) const { in device_to_device_stream() function
301 return device_to_device_stream(stream); in GetDeviceToDeviceStream()
172 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()