Delay queue destruction until all events are relased

igchor · igchor · commit b5f067ffb514 · 2025-06-17T18:29:37.000Z
diff --git a/unified-runtime/source/adapters/level_zero/v2/command_buffer.cpp b/unified-runtime/source/adapters/level_zero/v2/command_buffer.cpp
@@ -68,7 +68,12 @@ ur_exp_command_buffer_handle_t_::ur_exp_command_buffer_handle_t_(
     const ur_exp_command_buffer_desc_t *desc)
     : eventPool(context->getEventPoolCache(PoolCacheType::Regular)
                     .borrow(device->Id.value(),
-                            isInOrder ? v2::EVENT_FLAGS_COUNTER : 0)),
+                            isInOrder ? v2::EVENT_FLAGS_COUNTER : 0,
+                            [this] {
+                              if (this->RefCount.load() == 0) {
+                                delete this;
+                              }
+                            })),
       context(context), device(device),
       isUpdatable(desc ? desc->isUpdatable : false),
       isInOrder(desc ? desc->isInOrder : false),
@@ -226,6 +231,18 @@ ur_event_handle_t ur_exp_command_buffer_handle_t_::createEventIfRequested(
   return event;
 }
 
+ur_result_t ur_exp_command_buffer_handle_t_::release() {
+  // Command buffer can only be released if all events were returned to the
+  // event pool. If the event pool is not full, we delay queue destruction
+  // until all events are relased. Queue destruction will happen on
+  // the last eventPool::free.
+
+  if (eventPool->isFull()) {
+    delete this;
+  }
+  return UR_RESULT_SUCCESS;
+}
+
 namespace ur::level_zero {
 
 ur_result_t
@@ -276,7 +293,7 @@ urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t hCommandBuffer) try {
     ZE2UR_CALL(zeEventHostSynchronize,
                (executionEvent->getZeEvent(), UINT64_MAX));
   }
-  delete hCommandBuffer;
+  return hCommandBuffer->release();
   return UR_RESULT_SUCCESS;
 } catch (...) {
   return exceptionToResult(std::current_exception());
diff --git a/unified-runtime/source/adapters/level_zero/v2/command_buffer.hpp b/unified-runtime/source/adapters/level_zero/v2/command_buffer.hpp
@@ -53,6 +53,8 @@ struct ur_exp_command_buffer_handle_t_ : public ur_object {
   ur_event_handle_t
   createEventIfRequested(ur_exp_command_buffer_sync_point_t *retSyncPoint);
 
+  ur_result_t release();
+
 private:
   v2::raii::cache_borrowed_event_pool eventPool;
 
diff --git a/unified-runtime/source/adapters/level_zero/v2/event.cpp b/unified-runtime/source/adapters/level_zero/v2/event.cpp
@@ -261,9 +261,8 @@ ur_result_t urEventGetInfo(ur_event_handle_t hEvent, ur_event_info_t propName,
     return returnValue(hEvent->RefCount.load());
   }
   case UR_EVENT_INFO_COMMAND_QUEUE: {
-    auto urQueueHandle = reinterpret_cast<uintptr_t>(hEvent->getQueue()) -
-                         ur_queue_handle_t_::queue_offset;
-    return returnValue(urQueueHandle);
+    return returnValue(
+        ur_queue_handle_t_::queuePtrToHandle(hEvent->getQueue()));
   }
   case UR_EVENT_INFO_CONTEXT: {
     return returnValue(hEvent->getContext());
diff --git a/unified-runtime/source/adapters/level_zero/v2/event_pool.cpp b/unified-runtime/source/adapters/level_zero/v2/event_pool.cpp
@@ -54,6 +54,12 @@ void event_pool::free(ur_event_handle_t event) {
   // The event is still in the pool, so we need to increment the refcount
   assert(event->RefCount.load() == 0);
   event->RefCount.increment();
+
+  // All events are returned to the cache
+  if (events.size() == freelist.size()) {
+    lock.unlock();
+    cleanupCallback();
+  }
 }
 
 event_provider *event_pool::getProvider() const { return provider.get(); }
@@ -62,4 +68,9 @@ event_flags_t event_pool::getFlags() const {
   return getProvider()->eventFlags();
 }
 
+bool event_pool::isFull() const {
+  std::unique_lock<std::mutex> lock(*mutex);
+  return events.size() == freelist.size();
+}
+
 } // namespace v2
diff --git a/unified-runtime/source/adapters/level_zero/v2/event_pool.hpp b/unified-runtime/source/adapters/level_zero/v2/event_pool.hpp
@@ -42,6 +42,10 @@ class event_pool {
 
   ~event_pool() = default;
 
+  void setCleanupCallbackUnlocked(std::function<void(void)> cleanupCallback) {
+    this->cleanupCallback = cleanupCallback;
+  }
+
   // Allocate an event from the pool. Thread safe.
   ur_event_handle_t allocate();
 
@@ -51,6 +55,8 @@ class event_pool {
   event_provider *getProvider() const;
   event_flags_t getFlags() const;
 
+  bool isFull() const;
+
 private:
   ur_context_handle_t hContext;
   std::unique_ptr<event_provider> provider;
@@ -59,6 +65,8 @@ class event_pool {
   std::vector<ur_event_handle_t> freelist;
 
   std::unique_ptr<std::mutex> mutex;
+
+  std::function<void(void)> cleanupCallback;
 };
 
 // Only create an event when requested by the user.
diff --git a/unified-runtime/source/adapters/level_zero/v2/event_pool_cache.cpp b/unified-runtime/source/adapters/level_zero/v2/event_pool_cache.cpp
@@ -20,8 +20,9 @@ event_pool_cache::event_pool_cache(ur_context_handle_t hContext,
   pools.resize(max_devices * (1ULL << EVENT_FLAGS_USED_BITS));
 }
 
-raii::cache_borrowed_event_pool event_pool_cache::borrow(DeviceId id,
-                                                         event_flags_t flags) {
+raii::cache_borrowed_event_pool
+event_pool_cache::borrow(DeviceId id, event_flags_t flags,
+                         std::function<void(void)> cleanupCb) {
   std::unique_lock<ur_mutex> Lock(mutex);
 
   event_descriptor event_desc{id, flags};
@@ -39,6 +40,8 @@ raii::cache_borrowed_event_pool event_pool_cache::borrow(DeviceId id,
   auto pool = vec.back().release();
   vec.pop_back();
 
+  pool->setCleanupCallbackUnlocked(cleanupCb);
+
   return raii::cache_borrowed_event_pool(
       pool, [this, id, flags](event_pool *pool) {
         std::unique_lock<ur_mutex> Lock(mutex);
diff --git a/unified-runtime/source/adapters/level_zero/v2/event_pool_cache.hpp b/unified-runtime/source/adapters/level_zero/v2/event_pool_cache.hpp
@@ -38,7 +38,8 @@ class event_pool_cache {
   event_pool_cache(ur_context_handle_t hContext, size_t max_devices,
                    ProviderCreateFunc);
 
-  raii::cache_borrowed_event_pool borrow(DeviceId, event_flags_t flags);
+  raii::cache_borrowed_event_pool borrow(DeviceId, event_flags_t flags,
+                                         std::function<void(void)> cleanupCb);
 
 private:
   ur_context_handle_t hContext;
diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_handle.hpp b/unified-runtime/source/adapters/level_zero/v2/queue_handle.hpp
@@ -27,6 +27,13 @@ struct ur_queue_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter> {
   static constexpr uintptr_t queue_offset =
       sizeof(ur::handle_base<ur::level_zero::ddi_getter>);
 
+  template <typename Q> static ur_queue_handle_t queuePtrToHandle(Q *queue) {
+    if (!queue)
+      return nullptr;
+    return reinterpret_cast<ur_queue_handle_t>(
+        reinterpret_cast<uintptr_t>(queue) - queue_offset);
+  }
+
   template <typename T, class... Args>
   ur_queue_handle_t_(std::in_place_type_t<T>, Args &&...args)
       : ur::handle_base<ur::level_zero::ddi_getter>(),
@@ -61,8 +68,7 @@ struct ur_queue_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter> {
         [queueHandle = this](auto &q) {
           if (!q.RefCount.decrementAndTest())
             return UR_RESULT_SUCCESS;
-          delete queueHandle;
-          return UR_RESULT_SUCCESS;
+          return q.release();
         },
         queue_data);
   }
diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.cpp b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.cpp
@@ -12,6 +12,7 @@
 #include "command_buffer.hpp"
 #include "kernel.hpp"
 #include "memory.hpp"
+#include "queue_handle.hpp"
 #include "ur.hpp"
 
 #include "../common/latency_tracker.hpp"
@@ -29,7 +30,13 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t(
     event_flags_t eventFlags, ur_queue_flags_t flags)
     : hContext(hContext), hDevice(hDevice),
       eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate)
-                    .borrow(hDevice->Id.value(), eventFlags)),
+                    .borrow(hDevice->Id.value(), eventFlags,
+                            [this] {
+                              if (this->RefCount.load() == 0) {
+                                delete ur_queue_handle_t_::queuePtrToHandle(
+                                    this);
+                              }
+                            })),
       commandListManager(
           hContext, hDevice,
           hContext->getCommandListCache().getImmediateCommandList(
@@ -46,12 +53,30 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t(
     event_flags_t eventFlags, ur_queue_flags_t flags)
     : hContext(hContext), hDevice(hDevice),
       eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate)
-                    .borrow(hDevice->Id.value(), eventFlags)),
+                    .borrow(hDevice->Id.value(), eventFlags,
+                            [this] {
+                              if (this->RefCount.load() == 0) {
+                                delete ur_queue_handle_t_::queuePtrToHandle(
+                                    this);
+                              }
+                            })),
       commandListManager(hContext, hDevice, std::move(commandListHandle)),
       flags(flags) {
   ur::level_zero::urContextRetain(hContext);
 }
 
+ur_result_t ur_queue_immediate_in_order_t::release() {
+  // Command buffer can only be released if all events were returned to the
+  // event pool. If the event pool is not full, we delay queue destruction
+  // until all events are relased. Queue destruction will happen on
+  // the last eventPool::free.
+
+  if (eventPool->isFull())
+    delete ur_queue_handle_t_::queuePtrToHandle(this);
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t
 ur_queue_immediate_in_order_t::queueGetInfo(ur_queue_info_t propName,
                                             size_t propSize, void *pPropValue,
diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.hpp b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_in_order.hpp
@@ -48,6 +48,7 @@ struct ur_queue_immediate_in_order_t : ur_object, ur_queue_t_ {
 
   ~ur_queue_immediate_in_order_t();
 
+  ur_result_t release();
   ur_result_t queueGetInfo(ur_queue_info_t propName, size_t propSize,
                            void *pPropValue, size_t *pPropSizeRet) override;
   ur_result_t queueGetNativeHandle(ur_queue_native_desc_t *pDesc,
diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.cpp b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.cpp
@@ -10,6 +10,7 @@
 
 #include "queue_immediate_out_of_order.hpp"
 #include "../common/latency_tracker.hpp"
+#include "queue_handle.hpp"
 #include "ur.hpp"
 
 namespace v2 {
@@ -34,7 +35,13 @@ ur_queue_immediate_out_of_order_t::ur_queue_immediate_out_of_order_t(
     event_flags_t eventFlags, ur_queue_flags_t flags)
     : hContext(hContext), hDevice(hDevice),
       eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate)
-                    .borrow(hDevice->Id.value(), eventFlags)),
+                    .borrow(hDevice->Id.value(), eventFlags,
+                            [this] {
+                              if (this->RefCount.load() == 0) {
+                                delete ur_queue_handle_t_::queuePtrToHandle(
+                                    this);
+                              }
+                            })),
       commandListManagers(createCommandListManagers<numCommandLists>(
           hContext, hDevice, ordinal, priority, index)),
       flags(flags) {
@@ -45,6 +52,18 @@ ur_queue_immediate_out_of_order_t::ur_queue_immediate_out_of_order_t(
   ur::level_zero::urContextRetain(hContext);
 }
 
+ur_result_t ur_queue_immediate_out_of_order_t::release() {
+  // Command buffer can only be released if all events were returned to the
+  // event pool. If the event pool is not full, we delay queue destruction
+  // until all events are relased. Queue destruction will happen on
+  // the last eventPool::free.
+
+  if (eventPool->isFull())
+    delete ur_queue_handle_t_::queuePtrToHandle(this);
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t ur_queue_immediate_out_of_order_t::queueGetInfo(
     ur_queue_info_t propName, size_t propSize, void *pPropValue,
     size_t *pPropSizeRet) {
diff --git a/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.hpp b/unified-runtime/source/adapters/level_zero/v2/queue_immediate_out_of_order.hpp
@@ -59,6 +59,7 @@ struct ur_queue_immediate_out_of_order_t : ur_object, ur_queue_t_ {
 
   ~ur_queue_immediate_out_of_order_t();
 
+  ur_result_t release();
   ur_result_t queueGetInfo(ur_queue_info_t propName, size_t propSize,
                            void *pPropValue, size_t *pPropSizeRet) override;
   ur_result_t queueGetNativeHandle(ur_queue_native_desc_t *pDesc,
diff --git a/unified-runtime/test/adapters/level_zero/v2/event_pool_test.cpp b/unified-runtime/test/adapters/level_zero/v2/event_pool_test.cpp
@@ -168,9 +168,9 @@ UUR_DEVICE_TEST_SUITE_WITH_PARAM(EventPoolTest, testing::ValuesIn(test_cases),
                                  printParams<EventPoolTest>);
 
 TEST_P(EventPoolTest, InvalidDevice) {
-  auto pool = cache->borrow(MAX_DEVICES, getParam().flags);
+  auto pool = cache->borrow(MAX_DEVICES, getParam().flags, [] {});
   ASSERT_EQ(pool, nullptr);
-  pool = cache->borrow(MAX_DEVICES + 10, getParam().flags);
+  pool = cache->borrow(MAX_DEVICES + 10, getParam().flags, [] {});
   ASSERT_EQ(pool, nullptr);
 }
 
@@ -179,7 +179,7 @@ TEST_P(EventPoolTest, Basic) {
     ur_event_handle_t first;
     ze_event_handle_t zeFirst;
     {
-      auto pool = cache->borrow(device->Id.value(), getParam().flags);
+      auto pool = cache->borrow(device->Id.value(), getParam().flags, [] {});
 
       first = pool->allocate();
       first->setQueue(nullptr);
@@ -191,7 +191,7 @@ TEST_P(EventPoolTest, Basic) {
     ur_event_handle_t second;
     ze_event_handle_t zeSecond;
     {
-      auto pool = cache->borrow(device->Id.value(), getParam().flags);
+      auto pool = cache->borrow(device->Id.value(), getParam().flags, [] {});
 
       second = pool->allocate();
       second->setQueue(nullptr);
@@ -211,7 +211,7 @@ TEST_P(EventPoolTest, Threaded) {
   for (int iters = 0; iters < 3; ++iters) {
     for (int th = 0; th < 10; ++th) {
       threads.emplace_back([&] {
-        auto pool = cache->borrow(device->Id.value(), getParam().flags);
+        auto pool = cache->borrow(device->Id.value(), getParam().flags, [] {});
         std::vector<ur_event_handle_t> events;
         for (int i = 0; i < 100; ++i) {
           events.push_back(pool->allocate());
@@ -231,7 +231,7 @@ TEST_P(EventPoolTest, Threaded) {
 }
 
 TEST_P(EventPoolTest, ProviderNormalUseMostFreePool) {
-  auto pool = cache->borrow(device->Id.value(), getParam().flags);
+  auto pool = cache->borrow(device->Id.value(), getParam().flags, [] {});
   std::list<ur_event_handle_t> events;
   for (int i = 0; i < 128; ++i) {
     auto event = pool->allocate();

Original file line number	Diff line number	Diff line change
`@@ -261,9 +261,8 @@ ur_result_t urEventGetInfo(ur_event_handle_t hEvent, ur_event_info_t propName,`
`261`	`261`	`return returnValue(hEvent->RefCount.load());`
`262`	`262`	`}`
`263`	`263`	`case UR_EVENT_INFO_COMMAND_QUEUE: {`
`264`		`- auto urQueueHandle = reinterpret_cast<uintptr_t>(hEvent->getQueue()) -`
`265`		`- ur_queue_handle_t_::queue_offset;`
`266`		`- return returnValue(urQueueHandle);`
	`264`	`+ return returnValue(`
	`265`	`+ ur_queue_handle_t_::queuePtrToHandle(hEvent->getQueue()));`
`267`	`266`	`}`
`268`	`267`	`case UR_EVENT_INFO_CONTEXT: {`
`269`	`268`	`return returnValue(hEvent->getContext());`