| 21 | namespace gpu { |
| 22 | |
| 23 | se::Stream* InfeedManager::GetStream(se::StreamExecutor* executor) { |
| 24 | tensorflow::mutex_lock l(host_to_device_stream_mu_); |
| 25 | if (host_to_device_executor_ == nullptr) { |
| 26 | host_to_device_executor_ = executor; |
| 27 | host_to_device_stream_ = absl::make_unique<se::Stream>(executor); |
| 28 | host_to_device_stream_->Init(); |
| 29 | } |
| 30 | |
| 31 | if (executor != host_to_device_executor_) { |
| 32 | // The requested executor must be the same as the one for which |
| 33 | // the stream is cached. |
| 34 | return nullptr; |
| 35 | } |
| 36 | |
| 37 | return host_to_device_stream_.get(); |
| 38 | } |
| 39 | |
| 40 | InfeedManager* GetOrCreateInfeedManager() { |
| 41 | static InfeedManager* manager = new InfeedManager; |
no test coverage detected