| 30 | EventPool::EventPool(bool allow_reuse) : allow_reuse_(allow_reuse) {} |
| 31 | |
| 32 | StatusOr<EventPool::Handle> EventPool::ThenAllocateAndRecordEvent( |
| 33 | se::Stream* stream) { |
| 34 | Handle event; |
| 35 | |
| 36 | if (allow_reuse_) { |
| 37 | event.pool_ = this; |
| 38 | absl::MutexLock lock(&mu_); |
| 39 | if (!free_events_.empty()) { |
| 40 | event.event_ = std::move(free_events_.top()); |
| 41 | free_events_.pop(); |
| 42 | } |
| 43 | } |
| 44 | if (!event.event_) { |
| 45 | event.event_ = absl::make_unique<se::Event>(stream->parent()); |
| 46 | TF_RET_CHECK(event.event_->Init()) << "Event initialization failed"; |
| 47 | } |
| 48 | stream->ThenRecordEvent(event.event_.get()); |
| 49 | return event; |
| 50 | } |
| 51 | |
| 52 | } // namespace xla |