[UR][L0v2] Add initial record & replay implementation

Author: KFilipek

unified-runtime/source/adapters/level_zero/CMakeLists.txt

@@ -149,7 +149,6 @@ if(UR_BUILD_ADAPTER_L0_V2)
         ${CMAKE_CURRENT_SOURCE_DIR}/physical_mem.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/program.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/kernel_helpers.cpp
-        ${CMAKE_CURRENT_SOURCE_DIR}/graph.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/memory_helpers.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/mutable_helpers.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/usm_p2p.cpp
@@ -169,6 +168,7 @@ if(UR_BUILD_ADAPTER_L0_V2)
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event_provider_normal.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event_provider.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/v2/graph.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/kernel.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/memory.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/lockable.hpp

unified-runtime/source/adapters/level_zero/platform.cpp

@@ -526,6 +526,42 @@ ur_result_t ur_platform_handle_t_::initialize() {
   ZeMemGetPitchFor2dImageExt.Supported =
       ZeMemGetPitchFor2dImageExt.zeMemGetPitchFor2dImage != nullptr;
 
+  // Populate Graph Extension structure.
+  std::unordered_map<std::string, void **> ZeGraphFuncNameToAddrMap = {
+      {"zeGraphCreateExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeGraphCreateExp)},
+      {"zeCommandListBeginGraphCaptureExp",
+       reinterpret_cast<void **>(
+           &ZeGraphExt.zeCommandListBeginGraphCaptureExp)},
+      {"zeCommandListBeginCaptureIntoGraphExp",
+       reinterpret_cast<void **>(
+           &ZeGraphExt.zeCommandListBeginCaptureIntoGraphExp)},
+      {"zeCommandListEndGraphCaptureExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeCommandListEndGraphCaptureExp)},
+      {"zeCommandListInstantiateGraphExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeCommandListInstantiateGraphExp)},
+      {"zeCommandListAppendGraphExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeCommandListAppendGraphExp)},
+      {"zeGraphDestroyExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeGraphDestroyExp)},
+      {"zeExecutableGraphDestroyExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeExecutableGraphDestroyExp)},
+      {"zeCommandListIsGraphCaptureEnabledExp",
+       reinterpret_cast<void **>(
+           &ZeGraphExt.zeCommandListIsGraphCaptureEnabledExp)},
+      {"zeGraphIsEmptyExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeGraphIsEmptyExp)},
+      {"zeGraphDumpContentsExp",
+       reinterpret_cast<void **>(&ZeGraphExt.zeGraphDumpContentsExp)},
+  };
+
+  ZeGraphExt.Supported = true;
+  for (auto &[funcName, funcAddr] : ZeGraphFuncNameToAddrMap) {
+    ZE_CALL_NOCHECK(zeDriverGetExtensionFunctionAddress,
+                    (ZeDriver, funcName.c_str(), funcAddr));
+    ZeGraphExt.Supported &= (*funcAddr != nullptr);
+  }
+
   if (this->isDriverVersionNewerOrSimilar(1, 14, 36035)) {
     ZeCommandListAppendLaunchKernelWithArgumentsExt.Supported = true;
   } else {

unified-runtime/source/adapters/level_zero/platform.hpp

@@ -10,6 +10,7 @@
 #pragma once
 
 #include "common.hpp"
+#include "level_zero/driver_experimental/zex_api.h"
 #include "ur_api.h"
 #include "ze_api.h"
 #include "ze_ddi.h"
@@ -168,4 +169,34 @@ struct ur_platform_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter>,
     bool DriverSupportsCooperativeKernelLaunchWithArgs = false;
     bool DisableZeLaunchKernelWithArgs = false;
   } ZeCommandListAppendLaunchKernelWithArgumentsExt;
+
+  struct ZeGraphExtension {
+    bool Supported = false;
+    ze_result_t (*zeGraphCreateExp)(ze_context_handle_t hContext,
+                                    ze_graph_handle_t *phGraph, void *pNext);
+    ze_result_t (*zeCommandListBeginGraphCaptureExp)(
+        ze_command_list_handle_t hCommandList, void *pNext);
+    ze_result_t (*zeCommandListBeginCaptureIntoGraphExp)(
+        ze_command_list_handle_t hCommandList, ze_graph_handle_t hGraph,
+        void *pNext);
+    ze_result_t (*zeCommandListEndGraphCaptureExp)(
+        ze_command_list_handle_t hCommandList, ze_graph_handle_t *phGraph,
+        void *pNext);
+    ze_result_t (*zeCommandListInstantiateGraphExp)(
+        ze_graph_handle_t hGraph,
+        ze_executable_graph_handle_t *phExecutableGraph, void *pNext);
+    ze_result_t (*zeCommandListAppendGraphExp)(
+        ze_command_list_handle_t hCommandList,
+        ze_executable_graph_handle_t hGraph, void *pNext,
+        ze_event_handle_t hSignalEvent, uint32_t numWaitEvents,
+        ze_event_handle_t *phWaitEvents);
+    ze_result_t (*zeGraphDestroyExp)(ze_graph_handle_t hGraph);
+    ze_result_t (*zeExecutableGraphDestroyExp)(
+        ze_executable_graph_handle_t hGraph);
+    ze_result_t (*zeCommandListIsGraphCaptureEnabledExp)(
+        ze_command_list_handle_t hCommandList);
+    ze_result_t (*zeGraphIsEmptyExp)(ze_graph_handle_t hGraph);
+    ze_result_t (*zeGraphDumpContentsExp)(ze_graph_handle_t hGraph,
+                                          const char *filePath, void *pNext);
+  } ZeGraphExt;
 };

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

@@ -16,6 +16,7 @@
 #include "command_buffer.hpp"
 #include "common.hpp"
 #include "context.hpp"
+#include "graph.hpp"
 #include "kernel.hpp"
 #include "memory.hpp"
 
@@ -1302,3 +1303,78 @@ ur_result_t ur_command_list_manager::appendKernelLaunchWithArgsExp(
 
   return UR_RESULT_SUCCESS;
 }
+
+ur_result_t ur_command_list_manager::beginGraphCapture() {
+  if (!checkGraphExtensionSupport(hContext.get())) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  ZE2UR_CALL(hContext.get()
+                 ->getPlatform()
+                 ->ZeGraphExt.zeCommandListBeginGraphCaptureExp,
+             (getZeCommandList(), nullptr));
+  graphCapture.enableCapture();
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t
+ur_command_list_manager::beginCaptureIntoGraph(ur_exp_graph_handle_t hGraph) {
+  if (!checkGraphExtensionSupport(hContext.get())) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  ZE2UR_CALL(hContext.get()
+                 ->getPlatform()
+                 ->ZeGraphExt.zeCommandListBeginCaptureIntoGraphExp,
+             (getZeCommandList(), hGraph->getZeHandle(), nullptr));
+  graphCapture.enableCapture(hGraph);
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t
+ur_command_list_manager::endGraphCapture(ur_exp_graph_handle_t *phGraph) {
+  if (!checkGraphExtensionSupport(hContext.get())) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  ze_graph_handle_t zeGraph = nullptr;
+  ZE2UR_CALL(
+      hContext.get()->getPlatform()->ZeGraphExt.zeCommandListEndGraphCaptureExp,
+      (getZeCommandList(), &zeGraph, nullptr));
+  auto graph = graphCapture.getGraph();
+  graphCapture.disableCapture();
+
+  *phGraph =
+      graph ? graph : new ur_exp_graph_handle_t_(hContext.get(), zeGraph);
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t
+ur_command_list_manager::appendGraph(ur_exp_executable_graph_handle_t hGraph,
+                                     wait_list_view &waitListView,
+                                     ur_event_handle_t hEvent) {
+  if (!checkGraphExtensionSupport(hContext.get())) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  auto zeSignalEvent = getSignalEvent(hEvent, UR_COMMAND_ENQUEUE_GRAPH_EXP);
+  ZE2UR_CALL(
+      hContext.get()->getPlatform()->ZeGraphExt.zeCommandListAppendGraphExp,
+      (getZeCommandList(), hGraph->getZeHandle(), nullptr, zeSignalEvent,
+       waitListView.num, waitListView.handles));
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t ur_command_list_manager::isGraphCaptureActive(bool *pResult) {
+  if (!checkGraphExtensionSupport(hContext.get())) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  *pResult = graphCapture.isActive();
+
+  return UR_RESULT_SUCCESS;
+}

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

@@ -45,6 +45,25 @@ 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 context,
                           ur_device_handle_t device,
@@ -227,6 +246,15 @@ struct ur_command_list_manager {
       const ur_kernel_launch_ext_properties_t *launchPropList,
       wait_list_view &waitListView, ur_event_handle_t phEvent);
 
+  /************ Graph methods *************/
+  ur_result_t appendGraph(ur_exp_executable_graph_handle_t hGraph,
+                          wait_list_view &waitListView,
+                          ur_event_handle_t hEvent);
+  ur_result_t beginGraphCapture();
+  ur_result_t beginCaptureIntoGraph(ur_exp_graph_handle_t hGraph);
+  ur_result_t endGraphCapture(ur_exp_graph_handle_t *phGraph);
+  ur_result_t isGraphCaptureActive(bool *pResult);
+
   v2::raii::command_list_unique_handle &&releaseCommandList();
 
   void replaceCommandList(v2::raii::command_list_unique_handle &&cmdlist);
@@ -301,4 +329,6 @@ struct ur_command_list_manager {
   std::vector<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;
 };

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

@@ -0,0 +1,131 @@
+//===--------- graph.cpp - Level Zero Adapter -----------------------------===//
+//
+// Copyright (C) 2025 Intel Corporation
+//
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+// Exceptions. See LICENSE.TXT
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "graph.hpp"
+#include "../ur_interface_loader.hpp"
+#include "common.hpp"
+#include "context.hpp"
+#include "level_zero/driver_experimental/zex_api.h"
+
+ur_exp_graph_handle_t_::ur_exp_graph_handle_t_(ur_context_handle_t hContext)
+    : hContext(hContext) {
+  ZE2UR_CALL_THROWS(hContext->getPlatform()->ZeGraphExt.zeGraphCreateExp,
+                    (hContext->getZeHandle(), &zeGraph, nullptr));
+}
+
+ur_exp_graph_handle_t_::ur_exp_graph_handle_t_(ur_context_handle_t hContext,
+                                               ze_graph_handle_t zeGraph)
+    : hContext(hContext), zeGraph(zeGraph) {}
+
+ur_exp_graph_handle_t_::~ur_exp_graph_handle_t_() {
+  if (zeGraph) {
+    ze_result_t ZeResult = ZE_CALL_NOCHECK(
+        hContext->getPlatform()->ZeGraphExt.zeGraphDestroyExp, (zeGraph));
+    if (ZeResult != ZE_RESULT_SUCCESS) {
+      UR_LOG(WARN, "Failed to destroy graph handle: {}", ZeResult);
+    }
+  }
+}
+
+ur_exp_executable_graph_handle_t_::ur_exp_executable_graph_handle_t_(
+    ur_context_handle_t hContext, ur_exp_graph_handle_t hGraph)
+    : hContext(hContext) {
+  ZE2UR_CALL_THROWS(
+      hContext->getPlatform()->ZeGraphExt.zeCommandListInstantiateGraphExp,
+      (hGraph->getZeHandle(), &zeExGraph, nullptr));
+}
+
+ur_exp_executable_graph_handle_t_::~ur_exp_executable_graph_handle_t_() {
+  if (zeExGraph) {
+    ze_result_t ZeResult = ZE_CALL_NOCHECK(
+        hContext->getPlatform()->ZeGraphExt.zeExecutableGraphDestroyExp,
+        (zeExGraph));
+    if (ZeResult != ZE_RESULT_SUCCESS) {
+      UR_LOG(WARN, "Failed to destroy executable graph handle: {}", ZeResult);
+    }
+  }
+}
+
+namespace ur::level_zero {
+
+ur_result_t urGraphCreateExp(ur_context_handle_t hContext,
+                             ur_exp_graph_handle_t *phGraph) try {
+  if (!checkGraphExtensionSupport(hContext)) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  *phGraph = new ur_exp_graph_handle_t_(hContext);
+  return UR_RESULT_SUCCESS;
+} catch (...) {
+  return exceptionToResult(std::current_exception());
+}
+
+ur_result_t urGraphDestroyExp(ur_exp_graph_handle_t hGraph) try {
+  ur_context_handle_t hContext = hGraph->getContext();
+  if (!checkGraphExtensionSupport(hContext)) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  delete hGraph;
+  return UR_RESULT_SUCCESS;
+} catch (...) {
+  return exceptionToResult(std::current_exception());
+}
+
+ur_result_t urGraphInstantiateGraphExp(
+    ur_exp_graph_handle_t hGraph,
+    ur_exp_executable_graph_handle_t *phExecutableGraph) try {
+  ur_context_handle_t hContext = hGraph->getContext();
+  if (!checkGraphExtensionSupport(hContext)) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  *phExecutableGraph = new ur_exp_executable_graph_handle_t_(hContext, hGraph);
+  return UR_RESULT_SUCCESS;
+} catch (...) {
+  return exceptionToResult(std::current_exception());
+}
+
+ur_result_t urGraphExecutableGraphDestroyExp(
+    ur_exp_executable_graph_handle_t hExecutableGraph) try {
+  ur_context_handle_t hContext = hExecutableGraph->getContext();
+  if (!checkGraphExtensionSupport(hContext)) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  delete hExecutableGraph;
+  return UR_RESULT_SUCCESS;
+} catch (...) {
+  return exceptionToResult(std::current_exception());
+}
+
+ur_result_t urGraphIsEmptyExp(ur_exp_graph_handle_t hGraph, bool *pIsEmpty) {
+  ur_context_handle_t hContext = hGraph->getContext();
+  if (!checkGraphExtensionSupport(hContext)) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  ze_result_t zeResult =
+      ZE_CALL_NOCHECK(hContext->getPlatform()->ZeGraphExt.zeGraphIsEmptyExp,
+                      (hGraph->getZeHandle()));
+  if (zeResult == ZE_RESULT_ERROR_INVALID_GRAPH) {
+    return UR_RESULT_ERROR_INVALID_GRAPH;
+  }
+
+  *pIsEmpty = (zeResult == ZE_RESULT_QUERY_TRUE);
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t urGraphDumpContentsExp(ur_exp_graph_handle_t, const char *) {
+  UR_LOG(ERR, "{} function not implemented!", __FUNCTION__);
+  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+}
+
+} // namespace ur::level_zero