[L0] track ze queue usage before destroy

Signed-off-by: Neil R. Spruit <[email protected]>

Author: nrspruit

source/adapters/level_zero/queue.cpp

@@ -1314,10 +1314,14 @@ ur_result_t urQueueReleaseInternal(ur_queue_handle_t Queue) {
       for (auto &QueueGroup : QueueMap)
         for (auto &ZeQueue : QueueGroup.second.ZeQueues)
           if (ZeQueue) {
-            auto ZeResult = ZE_CALL_NOCHECK(zeCommandQueueDestroy, (ZeQueue));
-            // Gracefully handle the case that L0 was already unloaded.
-            if (ZeResult && ZeResult != ZE_RESULT_ERROR_UNINITIALIZED)
-              return ze2urResult(ZeResult);
+            {
+              std::scoped_lock<ur_shared_mutex> Lock(
+                  UrQueue->zeQueueInUseMutex);
+              auto ZeResult = ZE_CALL_NOCHECK(zeCommandQueueDestroy, (ZeQueue));
+              // Gracefully handle the case that L0 was already unloaded.
+              if (ZeResult && ZeResult != ZE_RESULT_ERROR_UNINITIALIZED)
+                return ze2urResult(ZeResult);
+            }
           }
   }
 
@@ -1392,7 +1396,6 @@ ur_result_t CleanupEventListFromResetCmdList(
 // runtime. Need to investigate whether relase can be done earlier, at sync
 // points such as this, to reduce total number of active Events.
 ur_result_t ur_queue_handle_t_::synchronize() {
-  // std::shared_lock<ur_shared_mutex> QueueLock(this->Mutex);
   this->Mutex.lock();
   if (!Healthy) {
     this->Mutex.unlock();
@@ -1406,7 +1409,10 @@ ur_result_t ur_queue_handle_t_::synchronize() {
 
     // wait for all commands previously submitted to this immediate command list
     Queue->Mutex.unlock();
-    ZE2UR_CALL(zeCommandListHostSynchronize, (ImmCmdList->first, UINT64_MAX));\
+    {
+      std::scoped_lock<ur_shared_mutex> Lock(Queue->zeQueueInUseMutex);
+      ZE2UR_CALL(zeCommandListHostSynchronize, (ImmCmdList->first, UINT64_MAX));
+    }
     Queue->Mutex.lock();
 
     // Cleanup all events from the synced command list.
@@ -1421,7 +1427,12 @@ ur_result_t ur_queue_handle_t_::synchronize() {
     // zero handle can have device scope, so we can't synchronize the last
     // event.
     if (isInOrderQueue() && !LastCommandEvent->IsDiscarded) {
-      ZE2UR_CALL(zeHostSynchronize, (LastCommandEvent->ZeEvent));
+      this->Mutex.unlock();
+      {
+        std::scoped_lock<ur_shared_mutex> Lock(this->zeQueueInUseMutex);
+        ZE2UR_CALL(zeHostSynchronize, (LastCommandEvent->ZeEvent));
+      }
+      this->Mutex.lock();
 
       // clean up all events known to have been completed as well,
       // so they can be reused later
@@ -1450,7 +1461,11 @@ ur_result_t ur_queue_handle_t_::synchronize() {
             for (auto &ZeQueue : QueueGroup.second.ZeQueues)
               if (ZeQueue) {
                 this->Mutex.unlock();
-                ZE2UR_CALL(zeHostSynchronize, (ZeQueue));
+                {
+                  std::scoped_lock<ur_shared_mutex> Lock(
+                      this->zeQueueInUseMutex);
+                  ZE2UR_CALL(zeHostSynchronize, (ZeQueue));
+                }
                 this->Mutex.lock();
               }
           }

source/adapters/level_zero/queue.hpp

@@ -237,6 +237,8 @@ struct ur_queue_handle_t_ : _ur_object {
   // Map of all command lists used in this queue.
   ur_command_list_map_t CommandListMap;
 
+  ur_shared_mutex zeQueueInUseMutex;
+
   // Helper data structure to hold all variables related to batching
   struct command_batch {
     // These two members are used to keep track of how often the
@@ -548,4 +550,4 @@ ur_result_t createEventAndAssociateQueue(
 // list.
 ur_result_t CleanupEventListFromResetCmdList(
     std::vector<ur_event_handle_t> &EventListToCleanup,
-    bool QueueLocked = false);
+    bool QueueLocked = false);