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

@@ -74,8 +74,7 @@ ur_exp_command_buffer_handle_t_::ur_exp_command_buffer_handle_t_(
       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) {
@@ -151,13 +150,25 @@ ur_event_handle_t ur_exp_command_buffer_handle_t_::getExecutionEventUnlocked() {
 
 ur_result_t ur_exp_command_buffer_handle_t_::registerExecutionEventUnlocked(
     ur_event_handle_t nextExecutionEvent) {
+  assert(nextExecutionEvent);
+  assert(nextExecutionEvent->getQueue());
+
+  auto nextExecutionQueue = nextExecutionEvent->getURQueueHandle();
+  auto currentExecutionQueue =
+      currentExecution ? currentExecution->getURQueueHandle() : nullptr;
+
   if (currentExecution) {
     UR_CALL(currentExecution->release());
-    currentExecution = nullptr;
   }
-  if (nextExecutionEvent) {
-    currentExecution = nextExecutionEvent;
+
+  currentExecution = nextExecutionEvent;
+
+  if (currentExecutionQueue != nextExecutionQueue) {
+    if (currentExecutionQueue)
+      UR_CALL(currentExecutionQueue->queueRelease());
+    UR_CALL(nextExecutionQueue->queueRetain());
   }
+
   return UR_RESULT_SUCCESS;
 }
 

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;