Home
last modified time | relevance | path

Searched refs:device_to_device_streams_ (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dlocal_device_state.cc58 device_to_device_streams_.reserve(kNumDeviceToDeviceStreams); in LocalDeviceState()
62 device_to_device_streams_.push_back(std::move(stream)); in LocalDeviceState()
143 (next_device_to_device_stream_ + 1) % device_to_device_streams_.size(); in GetDeviceToDeviceStream()
144 return device_to_device_streams_.at(i).get(); in GetDeviceToDeviceStream()
H A Dlocal_device_state.h179 std::vector<std::unique_ptr<se::Stream>> device_to_device_streams_; variable
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dxla_device_context.cc98 device_to_device_streams_(std::move(device_to_device_streams)), in XlaDeviceContext()
292 DCHECK_GT(device_to_device_streams_.size(), 0); in GetDeviceToDeviceStream()
295 next_stream_ = (next_stream_ + 1) % device_to_device_streams_.size(); in GetDeviceToDeviceStream()
H A Dxla_device_context.h81 return device_to_device_streams_.at(index).get(); in device_to_device_stream()
110 std::vector<std::shared_ptr<se::Stream>> device_to_device_streams_; variable
H A Dxla_device.cc237 device_to_device_streams_.resize(kNumDeviceToDeviceStreams); in XlaDevice()
338 for (std::shared_ptr<se::Stream>& stream : device_to_device_streams_) { in GetDeviceContextLocked()
344 device_to_device_streams = device_to_device_streams_; in GetDeviceContextLocked()
H A Dxla_device.h252 std::vector<std::shared_ptr<se::Stream>> device_to_device_streams_