| 21 | namespace xla { |
| 22 | |
| 23 | StreamPool::Ptr StreamPool::BorrowStream(se::StreamExecutor* executor) { |
| 24 | std::unique_ptr<se::Stream> stream; |
| 25 | { |
| 26 | tensorflow::mutex_lock lock(mu_); |
| 27 | if (!streams_.empty()) { |
| 28 | // Re-use an existing stream from the pool. |
| 29 | stream = std::move(streams_.back()); |
| 30 | streams_.pop_back(); |
| 31 | if (stream->ok()) { |
| 32 | VLOG(1) << stream->DebugStreamPointers() |
| 33 | << " StreamPool reusing existing stream"; |
| 34 | } else { |
| 35 | VLOG(1) << stream->DebugStreamPointers() |
| 36 | << " stream was not ok, StreamPool deleting"; |
| 37 | stream = nullptr; |
| 38 | } |
| 39 | } |
| 40 | } |
| 41 | |
| 42 | if (!stream) { |
| 43 | // Create a new stream. |
| 44 | stream = absl::make_unique<se::Stream>(executor); |
| 45 | stream->Init(); |
| 46 | VLOG(1) << stream->DebugStreamPointers() |
| 47 | << " StreamPool created new stream"; |
| 48 | } |
| 49 | |
| 50 | // Return the stream wrapped in Ptr, which has our special deleter semantics. |
| 51 | PtrDeleter deleter = {this}; |
| 52 | return Ptr(stream.release(), deleter); |
| 53 | } |
| 54 | |
| 55 | void StreamPool::ReturnStream(se::Stream* stream) { |
| 56 | if (stream->ok()) { |