[UR][L0v2] Fork-join pattern

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

@@ -1345,7 +1345,6 @@ ur_result_t ur_command_list_manager::beginGraphCapture() {
   ZE2UR_CALL(hContextInternal->getPlatform()
                  ->ZeGraphExt.zeCommandListBeginGraphCaptureExp,
              (getZeCommandList(), nullptr));
-  graphCapture.enableCapture();
 
   return UR_RESULT_SUCCESS;
 }
@@ -1360,7 +1359,6 @@ ur_command_list_manager::beginCaptureIntoGraph(ur_exp_graph_handle_t hGraph) {
   ZE2UR_CALL(hContextInternal->getPlatform()
                  ->ZeGraphExt.zeCommandListBeginCaptureIntoGraphExp,
              (getZeCommandList(), hGraph->getZeHandle(), nullptr));
-  graphCapture.enableCapture(hGraph);
 
   return UR_RESULT_SUCCESS;
 }
@@ -1375,18 +1373,17 @@ ur_command_list_manager::endGraphCapture(ur_exp_graph_handle_t *phGraph) {
   ze_graph_handle_t zeGraph = nullptr;
   ZE2UR_CALL(hContext.get()->getPlatform()->ZeGraphExt.endGraphCapture,
              (getZeCommandList(), nullptr, &zeGraph));
-  auto graph = graphCapture.getGraph();
-  graphCapture.disableCapture();
 
-  if (!graph) {
+  {
     std::scoped_lock<ur_shared_mutex> lock(hContextInternal->GraphMapMutex);
-    graph = hContextInternal->getGraphFromZeHandle(zeGraph);
-    if (!graph) {
-      graph = new ur_exp_graph_handle_t_(hContextInternal, zeGraph);
-      hContextInternal->registerGraph(zeGraph, graph);
+    auto hUrGraph = hContextInternal->getGraphFromZeHandle(zeGraph);
+    if (!hUrGraph) {
+      hUrGraph = new ur_exp_graph_handle_t_(hContextInternal, zeGraph);
+      hContextInternal->registerGraph(zeGraph, hUrGraph);
     }
+
+    *phGraph = hUrGraph;
   }
-  *phGraph = graph;
 
   return UR_RESULT_SUCCESS;
 }
@@ -1432,35 +1429,24 @@ ur_result_t ur_command_list_manager::getGraph(ur_exp_graph_handle_t *phGraph) {
     return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   }
 
-  auto hCachedGraph = graphCapture.getGraph();
-  if (hCachedGraph) {
-    *phGraph = hCachedGraph;
-    return UR_RESULT_SUCCESS;
-  }
-
-  // Fork-join and implicit capture scenarios
   ze_graph_handle_t hZeGraph = nullptr;
   ze_result_t ZeResult =
       ZE_CALL_NOCHECK(zeGetGraph, (getZeCommandList(), &hZeGraph));
-
   if (ZeResult != ZE_RESULT_SUCCESS || !hZeGraph) {
     *phGraph = nullptr;
     return UR_RESULT_ERROR_INVALID_OPERATION;
   }
 
-  ur_exp_graph_handle_t hUrGraph = nullptr;
   {
     std::scoped_lock<ur_shared_mutex> lock(hContextInternal->GraphMapMutex);
-    hUrGraph = hContextInternal->getGraphFromZeHandle(hZeGraph);
+    auto hUrGraph = hContextInternal->getGraphFromZeHandle(hZeGraph);
     if (!hUrGraph) {
       hUrGraph = new ur_exp_graph_handle_t_(hContextInternal, hZeGraph);
       hContextInternal->registerGraph(hZeGraph, hUrGraph);
-      if (graphCapture.isActive()) {
-        graphCapture.enableCapture(hUrGraph);
-      }
     }
+
+    *phGraph = hUrGraph;
   }
-  *phGraph = hUrGraph;
 
   return UR_RESULT_SUCCESS;
 }

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

@@ -46,25 +46,6 @@ struct wait_list_view {
   }
 };
 
-struct graph_capture_tracking_data_t {
-  void enableCapture(ur_exp_graph_handle_t graph) {
-    capturedGraph = graph;
-    enableCapture();
-  }
-  void enableCapture() { active = true; }
-  void disableCapture() {
-    active = false;
-    capturedGraph = nullptr;
-  }
-
-  ur_exp_graph_handle_t getGraph() const { return capturedGraph; }
-  bool isActive() const { return active; }
-
-private:
-  bool active = false;
-  ur_exp_graph_handle_t capturedGraph = nullptr;
-};
-
 struct ur_command_list_manager {
   ur_command_list_manager(ur_context_handle_t hContext,
                           ur_device_handle_t hDevice,
@@ -336,8 +317,6 @@ struct ur_command_list_manager {
   std::unordered_set<ur_kernel_handle_t> submittedKernels;
   v2::raii::command_list_unique_handle zeCommandList;
   std::vector<ze_event_handle_t> waitList;
-
-  graph_capture_tracking_data_t graphCapture;
 };
 
 } // namespace ur::level_zero::v2

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

@@ -171,6 +171,14 @@ ur_result_t ur_queue_immediate_out_of_order_t::enqueueEventsWaitWithBarrierExt(
                        : &ur_command_list_manager::appendEventsWait;
 
   auto commandListManagersLocked = commandListManagers.lock();
+  bool captureActive = false;
+  commandListManagersLocked[captureCmdListManagerIdx].queryGraphCaptureActive(
+      &captureActive);
+  if (captureActive) {
+    return std::invoke(
+        barrierFn, commandListManagersLocked[captureCmdListManagerIdx],
+        waitListView, createEventIfRequested(eventPool.get(), phEvent, this));
+  }
 
   // Enqueue wait for the user-provider events on the first command list.
   UR_CALL(commandListManagersLocked[0].appendEventsWait(waitListView,