Retain queue when storing execution event in command buffer

Otherwise the event might be returned to the eventPool
that has already been destroyed.

Author: igchor

unified-runtime/source/adapters/level_zero/v2/command_buffer.cpp

@@ -15,6 +15,45 @@
 #include "logger/ur_logger.hpp"
 #include "queue_handle.hpp"
 
+ur_result_t ur_execution_event_handle_t::assign(ur_event_handle_t hNewEvent) {
+  assert(hNewEvent);
+  assert(hNewEvent->getURQueueHandle());
+
+  auto newQueue = hNewEvent->getURQueueHandle();
+  auto currentQueue = hEvent ? hEvent->getURQueueHandle() : nullptr;
+
+  if (hEvent) {
+    UR_CALL(hEvent->release());
+  }
+
+  hEvent = hNewEvent;
+
+  if (newQueue != currentQueue) {
+    if (currentQueue)
+      UR_CALL(currentQueue->queueRelease());
+    UR_CALL(newQueue->queueRetain());
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_event_handle_t ur_execution_event_handle_t::get() { return hEvent; }
+
+ur_result_t ur_execution_event_handle_t::release() {
+  if (hEvent) {
+    assert(hEvent->getURQueueHandle());
+
+    auto hQueue = hEvent->getURQueueHandle();
+    UR_CALL_NOCHECK(hEvent->release());
+    UR_CALL_NOCHECK(hQueue->queueRelease());
+
+    hEvent = nullptr;
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+ur_execution_event_handle_t::~ur_execution_event_handle_t() { release(); }
+
 namespace {
 
 ur_result_t getZeKernelWrapped(ur_kernel_handle_t kernel,
@@ -69,13 +108,12 @@ ur_exp_command_buffer_handle_t_::ur_exp_command_buffer_handle_t_(
     : eventPool(context->getEventPoolCache(PoolCacheType::Regular)
                     .borrow(device->Id.value(),
                             isInOrder ? v2::EVENT_FLAGS_COUNTER : 0)),
-      context(context), device(device),
+      context(context), device(device), currentExecution(nullptr),
       isUpdatable(desc ? desc->isUpdatable : false),
       isInOrder(desc ? desc->isInOrder : false),
       commandListManager(
           context, device,
-          std::forward<v2::raii::command_list_unique_handle>(commandList))
-{}
+          std::forward<v2::raii::command_list_unique_handle>(commandList)) {}
 
 ur_exp_command_buffer_sync_point_t
 ur_exp_command_buffer_handle_t_::getSyncPoint(ur_event_handle_t event) {
@@ -146,25 +184,16 @@ ur_result_t ur_exp_command_buffer_handle_t_::finalizeCommandBuffer() {
   return UR_RESULT_SUCCESS;
 }
 ur_event_handle_t ur_exp_command_buffer_handle_t_::getExecutionEventUnlocked() {
-  return currentExecution;
+  return currentExecution.get();
 }
 
 ur_result_t ur_exp_command_buffer_handle_t_::registerExecutionEventUnlocked(
     ur_event_handle_t nextExecutionEvent) {
-  if (currentExecution) {
-    UR_CALL(currentExecution->release());
-    currentExecution = nullptr;
-  }
-  if (nextExecutionEvent) {
-    currentExecution = nextExecutionEvent;
-  }
+  UR_CALL(currentExecution.assign(nextExecutionEvent));
   return UR_RESULT_SUCCESS;
 }
 
 ur_exp_command_buffer_handle_t_::~ur_exp_command_buffer_handle_t_() {
-  if (currentExecution) {
-    currentExecution->release();
-  }
   for (auto &event : syncPoints) {
     event->release();
   }
@@ -181,14 +210,13 @@ ur_result_t ur_exp_command_buffer_handle_t_::applyUpdateCommands(
       this, device, context->getPlatform()->ZeDriverGlobalOffsetExtensionFound,
       numUpdateCommands, updateCommands));
 
-  if (currentExecution) {
+  if (currentExecution.get()) {
     // TODO: Move synchronization to command buffer enqueue
     // it would require to remember the update commands and perform update
     // before appending to the queue
     ZE2UR_CALL(zeEventHostSynchronize,
-               (currentExecution->getZeEvent(), UINT64_MAX));
-    currentExecution->release();
-    currentExecution = nullptr;
+               (currentExecution.get()->getZeEvent(), UINT64_MAX));
+    UR_CALL(currentExecution.release());
   }
 
   device_ptr_storage_t zeHandles;

unified-runtime/source/adapters/level_zero/v2/command_buffer.hpp

@@ -18,8 +18,26 @@
 #include "queue_api.hpp"
 #include <unordered_set>
 #include <ze_api.h>
+
 struct kernel_command_handle;
 
+struct ur_execution_event_handle_t {
+  ur_execution_event_handle_t(ur_event_handle_t event) : hEvent(event) {}
+
+  ur_execution_event_handle_t(const ur_execution_event_handle_t &) = delete;
+  ur_execution_event_handle_t &
+  operator=(const ur_execution_event_handle_t &) = delete;
+
+  ur_result_t assign(ur_event_handle_t hNewEvent);
+  ur_event_handle_t get();
+  ur_result_t release();
+
+  ~ur_execution_event_handle_t();
+
+private:
+  ur_event_handle_t hEvent;
+};
+
 struct ur_exp_command_buffer_handle_t_ : public ur_object {
   ur_exp_command_buffer_handle_t_(
       ur_context_handle_t context, ur_device_handle_t device,
@@ -72,7 +90,7 @@ struct ur_exp_command_buffer_handle_t_ : public ur_object {
   // Indicates if command-buffer was finalized.
   bool isFinalized = false;
 
-  ur_event_handle_t currentExecution = nullptr;
+  ur_execution_event_handle_t currentExecution;
 
 public:
   // Indicates if command-buffer commands can be updated after it is closed.

unified-runtime/source/adapters/level_zero/v2/event.cpp

@@ -192,6 +192,12 @@ ur_event_handle_t_::getEventEndTimestampAndHandle() {
 
 ur_queue_t_ *ur_event_handle_t_::getQueue() const { return hQueue; }
 
+ur_queue_handle_t ur_event_handle_t_::getURQueueHandle() const {
+  auto urHandle = reinterpret_cast<uintptr_t>(getQueue()) -
+                  ur_queue_handle_t_::queue_offset;
+  return reinterpret_cast<ur_queue_handle_t>(urHandle);
+}
+
 ur_context_handle_t ur_event_handle_t_::getContext() const { return hContext; }
 
 ur_command_t ur_event_handle_t_::getCommandType() const { return commandType; }
@@ -261,9 +267,7 @@ 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(hEvent->getURQueueHandle());
   }
   case UR_EVENT_INFO_CONTEXT: {
     return returnValue(hEvent->getContext());

unified-runtime/source/adapters/level_zero/v2/event.hpp

@@ -91,6 +91,9 @@ struct ur_event_handle_t_ : ur_object {
   // Queue associated with this event. Can be nullptr (for native events)
   ur_queue_t_ *getQueue() const;
 
+  // UR handle associated with this event's queue.
+  ur_queue_handle_t getURQueueHandle() const;
+
   // Context associated with this event
   ur_context_handle_t getContext() const;