[UR][L0v2] Bound retained kernels with forced synchronization

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

@@ -114,8 +114,8 @@ void wait_list_view::addEvent(ur_event_handle_t phEvent) {
 
 ur_command_list_manager::ur_command_list_manager(
     ur_context_handle_t hContext, ur_device_handle_t hDevice,
-    v2::raii::command_list_unique_handle &&commandList)
-    : hContext(hContext), hDevice(hDevice),
+    v2::raii::command_list_unique_handle &&commandList, bool isImmediate)
+    : hContext(hContext), hDevice(hDevice), isImmediate(isImmediate),
       zeCommandList(std::move(commandList)) {}
 
 v2::raii::command_list_unique_handle &&
@@ -239,6 +239,23 @@ ur_result_t ur_command_list_manager::appendKernelLaunchLocked(
     ur_event_handle_t phEvent, bool cooperative, bool callWithArgs,
     void *pNext) {
 
+  constexpr size_t maxPendingKernels = 10000;
+  if (isImmediate && pendingKernels.size() >= maxPendingKernels) {
+    bool captureActive = false;
+    auto result = queryGraphCaptureActive(&captureActive);
+    if (result != UR_RESULT_SUCCESS &&
+        result != UR_RESULT_ERROR_UNSUPPORTED_FEATURE)
+      return result;
+    if (!captureActive) {
+      // Synchronize before launching so a failure cannot cause a submitted
+      // kernel to be retried. Captured kernels must remain available for
+      // replay.
+      ZE2UR_CALL(zeCommandListHostSynchronize,
+                 (getZeCommandList(), UINT64_MAX));
+      UR_CALL(releasePendingKernels());
+    }
+  }
+
   ze_group_count_t zeThreadGroupDimensions{1, 1, 1};
   uint32_t WG[3]{};
   UR_CALL(calculateKernelWorkDimensions(hZeKernel, hDevice.get(),
@@ -276,7 +293,19 @@ ur_result_t ur_command_list_manager::appendKernelLaunchLocked(
                 zeSignalEvent, waitListView.num, waitListView.handles));
   }
 
-  recordSubmittedKernel(hKernel);
+  bool captureActive = false;
+  if (isImmediate) {
+    // A wait-list dependency can make this launch join another list's capture.
+    // Query after submission, but never report a submitted launch as failed.
+    auto result = queryGraphCaptureActive(&captureActive);
+    if (result != UR_RESULT_SUCCESS &&
+        result != UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
+      UR_LOG(ERR, "Failed to query graph capture after kernel launch: {}",
+             result);
+      captureActive = true;
+    }
+  }
+  recordSubmittedKernel(hKernel, captureActive);
 
   postSubmit(hZeKernel, pGlobalWorkOffset);
 
@@ -1072,9 +1101,22 @@ ur_result_t ur_command_list_manager::appendNativeCommandExp(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
-void ur_command_list_manager::recordSubmittedKernel(
-    ur_kernel_handle_t hKernel) {
-  auto [_, inserted] = submittedKernels.insert(hKernel);
+ur_result_t ur_command_list_manager::releasePendingKernels() {
+  while (!pendingKernels.empty()) {
+    auto kernel = *pendingKernels.begin();
+    pendingKernels.erase(pendingKernels.begin());
+    UR_CALL(kernel->release());
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+void ur_command_list_manager::recordSubmittedKernel(ur_kernel_handle_t hKernel,
+                                                    bool captureActive) {
+  // Regular lists can be replayed, and captured kernels belong to a graph,
+  // not to the immediately executing command stream.
+  const bool canRetire = isImmediate && !captureActive;
+  auto &kernels = canRetire ? pendingKernels : submittedKernels;
+  auto [_, inserted] = kernels.insert(hKernel);
   if (inserted) {
     hKernel->RefCount.retain();
   }
@@ -1120,6 +1162,8 @@ ur_result_t ur_command_list_manager::appendEventsWaitWithBarrier(
 }
 
 ur_result_t ur_command_list_manager::releaseSubmittedKernels() {
+  UR_CALL(releasePendingKernels());
+
   // Free deferred kernels
   for (auto &hKernel : submittedKernels) {
     UR_CALL(hKernel->release());
@@ -1434,6 +1478,9 @@ ur_result_t ur_command_list_manager::queryGraphCaptureActive(bool *pResult) {
   ze_result_t ZeResult = ZeGraphExt.normalizeGraphQueryResult(ZE_CALL_NOCHECK(
       ZeGraphExt.zeCommandListIsGraphCaptureEnabledExp, (getZeCommandList())));
 
+  if (ZeResult != ZE_RESULT_QUERY_TRUE && ZeResult != ZE_RESULT_QUERY_FALSE)
+    return ze2urResult(ZeResult);
+
   *pResult = (ZeResult == ZE_RESULT_QUERY_TRUE);
 
   return UR_RESULT_SUCCESS;

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

@@ -49,7 +49,8 @@ struct wait_list_view {
 struct ur_command_list_manager {
   ur_command_list_manager(ur_context_handle_t hContext,
                           ur_device_handle_t hDevice,
-                          v2::raii::command_list_unique_handle &&commandList);
+                          v2::raii::command_list_unique_handle &&commandList,
+                          bool isImmediate = false);
   ur_command_list_manager(const ur_command_list_manager &src) = delete;
   ur_command_list_manager(ur_command_list_manager &&src) = default;
 
@@ -269,7 +270,8 @@ struct ur_command_list_manager {
       ur_event_handle_t phEvent, wait_list_view &waitListView,
       ur_command_t callerCommand);
 
-  void recordSubmittedKernel(ur_kernel_handle_t hKernel);
+  ur_result_t releasePendingKernels();
+  void recordSubmittedKernel(ur_kernel_handle_t hKernel, bool captureActive);
 
   ze_event_handle_t getSignalEvent(ur_event_handle_t hUserEvent,
                                    ur_command_t commandType);
@@ -315,6 +317,8 @@ struct ur_command_list_manager {
   v2::raii::ur_device_handle_t hDevice;
 
   std::unordered_set<ur_kernel_handle_t> submittedKernels;
+  bool isImmediate;
+  std::unordered_set<ur_kernel_handle_t> pendingKernels;
   v2::raii::command_list_unique_handle zeCommandList;
   std::vector<ze_event_handle_t> waitList;
 };

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

@@ -34,7 +34,8 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t(
               hDevice->ZeDevice,
               {true, ordinal, true /* always enable copy offload */},
               ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, priority, index,
-              priorityLevel)),
+              priorityLevel),
+          true),
       flags(flags),
       eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate)
                     .borrow(hDevice->Id.value(), eventFlags)) {}
@@ -44,7 +45,7 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t(
     raii::command_list_unique_handle commandListHandle,
     event_flags_t eventFlags, ur_queue_flags_t flags)
     : hContext(hContext), hDevice(hDevice),
-      commandListManager(hContext, hDevice, std::move(commandListHandle)),
+      commandListManager(hContext, hDevice, std::move(commandListHandle), true),
       flags(flags),
       eventPool(hContext->getEventPoolCache(PoolCacheType::Immediate)
                     .borrow(hDevice->Id.value(), eventFlags)) {}

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

@@ -25,8 +25,8 @@ std::array<ur_command_list_manager, N> createCommandListManagers(
         hContext->getCommandListCache().getImmediateCommandList(
             hDevice->ZeDevice,
             {true, ordinal, true /* always enable copy offload */},
-            ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, priority, index,
-            priorityLevel));
+            ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS, priority, index, priorityLevel),
+        true);
   });
 }