From 61ff779c4c065a47943ba1d2281b5b7d57a33196 Mon Sep 17 00:00:00 2001 From: Eugene Zhulenev Date: Thu, 31 Aug 2023 16:02:18 -0700 Subject: [PATCH] [cuda2] Fix event_pool reference counting --- experimental/cuda2/event_pool.c | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/experimental/cuda2/event_pool.c b/experimental/cuda2/event_pool.c index 11d7f78ff8c7..66caf6ca37f4 100644 --- a/experimental/cuda2/event_pool.c +++ b/experimental/cuda2/event_pool.c @@ -81,7 +81,6 @@ static inline iree_status_t iree_hal_cuda2_event_create( "cuEventCreate"); if (iree_status_is_ok(status)) { *out_event = event; - iree_hal_cuda2_event_pool_retain(pool); // +1 } else { iree_atomic_ref_count_dec(&event->ref_count); // -> 0 iree_hal_cuda2_event_destroy(event); @@ -104,7 +103,7 @@ void iree_hal_cuda2_event_release(iree_hal_cuda2_event_t* event) { iree_hal_cuda2_event_pool_t* pool = event->pool; // Release back to the pool if the reference count becomes 0. iree_hal_cuda2_event_pool_release_event(pool, 1, &event); - // Drop our reference to the pool itself. + // Drop our reference to the pool itself when we return event to it. iree_hal_cuda2_event_pool_release(pool); // -1 } } @@ -267,6 +266,12 @@ iree_status_t iree_hal_cuda2_event_pool_acquire( IREE_TRACE_ZONE_END(z1); } + // Retain a reference to a pool when we pass it to the caller. When caller + // returns event to the pool we'll release the reference. + for (unsigned i = 0; i < event_count; ++i) { + iree_hal_cuda2_event_pool_retain(out_events[i]->pool); // +1 + } + IREE_TRACE_ZONE_END(z0); return iree_ok_status(); }