Searched refs:host_to_device_stream_ (Results 1 – 9 of 9) sorted by relevance
80 host_to_device_stream_(std::move(host_to_device_stream)), in XlaDeviceContext()87 CHECK(host_to_device_stream_ != nullptr); in XlaDeviceContext()158 host_to_device_stream_->ThenWaitFor(stream_.get()); in CopyCPUTensorToDevice()162 host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer())); in CopyCPUTensorToDevice()167 host_to_device_stream_->ThenRecordEvent(event.get()); in CopyCPUTensorToDevice()169 host_to_device_stream_.get()); in CopyCPUTensorToDevice()189 host_to_device_stream_->ThenDoHostCallback([ref]() { ref.Unref(); }); in CopyCPUTensorToDevice()192 host_to_device_stream_->ThenDoHostCallback([ref, done]() { in CopyCPUTensorToDevice()
76 return host_to_device_stream_.get(); in host_to_device_stream()90 bool UseMultipleStreams() const { return stream_ != host_to_device_stream_; } in UseMultipleStreams()97 std::shared_ptr<se::Stream> host_to_device_stream_; variable
291 &host_to_device_stream_, in GetDeviceContextLocked()298 host_to_device_stream = host_to_device_stream_; in GetDeviceContextLocked()
220 std::shared_ptr<se::Stream> host_to_device_stream_ GUARDED_BY(mu_);
27 host_to_device_stream_ = absl::make_unique<se::Stream>(executor); in GetStream()28 host_to_device_stream_->Init(); in GetStream()37 return host_to_device_stream_.get(); in GetStream()
77 std::unique_ptr<se::Stream> host_to_device_stream_
38 host_to_device_stream_(host_to_device_stream), in GPUDeviceContext()45 se::Stream* host_to_device_stream() const { return host_to_device_stream_; } in host_to_device_stream()76 se::Stream* host_to_device_stream_; variable
50 host_to_device_stream_ = stream; in set_host_to_device_stream()55 return host_to_device_stream_; in host_to_device_stream()
91 stream_executor::Stream* host_to_device_stream_ = nullptr; variable