Searched refs:host_to_device_stream_ (Results 1 – 10 of 10) sorted by relevance
96 host_to_device_stream_(std::move(host_to_device_stream)), in XlaDeviceContext()103 CHECK(host_to_device_stream_ != nullptr); in XlaDeviceContext()168 host_to_device_stream_->ThenWaitFor(stream_.get()); in CopyCPUTensorToDevice()172 host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer())); in CopyCPUTensorToDevice()177 host_to_device_stream_->ThenRecordEvent(event.get()); in CopyCPUTensorToDevice()179 host_to_device_stream_.get()); in CopyCPUTensorToDevice()199 host_to_device_stream_->ThenDoHostCallback([ref]() { ref.Unref(); }); in CopyCPUTensorToDevice()202 host_to_device_stream_->ThenDoHostCallback([ref, done]() { in CopyCPUTensorToDevice()
78 return host_to_device_stream_.get(); in host_to_device_stream()96 bool UseMultipleStreams() const { return stream_ != host_to_device_stream_; } in UseMultipleStreams()103 std::shared_ptr<se::Stream> host_to_device_stream_; variable
249 std::shared_ptr<se::Stream> host_to_device_stream_ TF_GUARDED_BY(mu_);
336 &host_to_device_stream_, in GetDeviceContextLocked()343 host_to_device_stream = host_to_device_stream_; in GetDeviceContextLocked()
38 host_to_device_stream_(host_to_device_stream), in PluggableDeviceContext()45 se::Stream* host_to_device_stream() const { return host_to_device_stream_; } in host_to_device_stream()78 se::Stream* host_to_device_stream_; variable
45 host_to_device_stream_(host_to_device_stream), in GPUDeviceContext()57 se::Stream* host_to_device_stream() const { return host_to_device_stream_; } in host_to_device_stream()95 se::Stream* host_to_device_stream_; variable
66 host_to_device_stream_ = stream; in set_host_to_device_stream()71 return host_to_device_stream_; in host_to_device_stream()
167 stream_executor::Stream* host_to_device_stream_ = nullptr; variable
109 return host_to_device_stream_.get(); in host_to_device_stream()177 std::unique_ptr<se::Stream> host_to_device_stream_; variable
45 host_to_device_stream_ = std::make_unique<se::Stream>(executor); in LocalDeviceState()47 host_to_device_stream_->Init(); in LocalDeviceState()