feat: add graph runtime API

scripts/generate_public_headers.py

@@ -156,9 +156,31 @@ def _write_detail_headers(include_root, source_root, devices):
         _write_detail_header(include_root, source_root, relative_path)
 
 
-def _write_generated_header(include_root, devices):
+def _write_generated_header(include_root, source_root, devices):
     default_device = _default_device(devices)
     default_device_type = _DEVICE_TYPES[default_device]
+    public_runtime_functions = _public_runtime_functions_for_devices(
+        devices, source_root
+    )
+    has_graph_api = any(
+        function.name in {"StreamBeginCapture", "GraphLaunch"}
+        for function in public_runtime_functions
+    )
+    graph_declarations = (
+        """
+using Graph = void*;
+
+using GraphExec = void*;
+
+enum class StreamCaptureMode {
+  kStreamCaptureModeGlobal = 0,
+  kStreamCaptureModeThreadLocal = 1,
+  kStreamCaptureModeRelaxed = 2,
+};
+"""
+        if has_graph_api
+        else ""
+    )
     includes = [
         "#include <cstddef>",
         "#include <cstdint>",
@@ -177,7 +199,7 @@ def _write_generated_header(include_root, devices):
         includes.append(f"#include <infini/rt/{device}/runtime_.h>")
 
     runtime_declarations = "\n\n".join(
-        f"{function.signature()};" for function in _PUBLIC_RUNTIME_FUNCTIONS
+        f"{function.signature()};" for function in public_runtime_functions
     )
 
     path = include_root / "infini" / "rt" / "generated.h"
@@ -209,6 +231,7 @@ def _write_generated_header(include_root, devices):
 
 using Event = void*;
 
+{graph_declarations}
 using MemcpyKind = std::remove_cv_t<
     decltype(generated_detail::DefaultErrorRuntime::kMemcpyHostToHost)>;
 
@@ -366,6 +389,32 @@ def params_decl(self):
             _Param("Event", "end"),
         ),
     ),
+    _Function(
+        "Error",
+        "StreamBeginCapture",
+        (_Param("Stream", "stream"), _Param("StreamCaptureMode", "mode")),
+    ),
+    _Function(
+        "Error",
+        "StreamEndCapture",
+        (_Param("Stream", "stream"), _Param("Graph*", "graph")),
+    ),
+    _Function("Error", "GraphDestroy", (_Param("Graph", "graph"),)),
+    _Function(
+        "Error",
+        "GraphInstantiate",
+        (_Param("GraphExec*", "graph_exec"), _Param("Graph", "graph")),
+    ),
+    _Function(
+        "Error",
+        "GraphExecDestroy",
+        (_Param("GraphExec", "graph_exec"),),
+    ),
+    _Function(
+        "Error",
+        "GraphLaunch",
+        (_Param("GraphExec", "graph_exec"), _Param("Stream", "stream")),
+    ),
 )
 
 
@@ -395,6 +444,24 @@ def _runtime_arg(param, device):
         return (
             f"reinterpret_cast<typename Runtime<{device_type}>::Event*>({param.name})"
         )
+    if param.type == "Graph":
+        return f"reinterpret_cast<typename Runtime<{device_type}>::Graph>({param.name})"
+    if param.type == "Graph*":
+        return (
+            f"reinterpret_cast<typename Runtime<{device_type}>::Graph*>({param.name})"
+        )
+    if param.type == "GraphExec":
+        return (
+            f"reinterpret_cast<typename Runtime<{device_type}>::GraphExec>"
+            f"({param.name})"
+        )
+    if param.type == "GraphExec*":
+        return (
+            f"reinterpret_cast<typename Runtime<{device_type}>::GraphExec*>"
+            f"({param.name})"
+        )
+    if param.type == "StreamCaptureMode":
+        return f"RuntimeStreamCaptureMode<{device_type}>({param.name})"
 
     return param.name
 
@@ -452,8 +519,42 @@ def _devices_for_function(function, devices, source_root):
     )
 
 
+def _public_runtime_functions_for_devices(devices, source_root):
+    return tuple(
+        function
+        for function in _PUBLIC_RUNTIME_FUNCTIONS
+        if _devices_for_function(function, devices, source_root)
+    )
+
+
 def _write_runtime_dispatch(source_path, source_root, devices):
-    functions = _PUBLIC_RUNTIME_FUNCTIONS
+    functions = _public_runtime_functions_for_devices(devices, source_root)
+    stream_capture_mode_helper = (
+        """
+template <Device::Type device_type>
+auto RuntimeStreamCaptureMode(StreamCaptureMode mode) {
+  using DeviceRuntime = Runtime<device_type>;
+
+  switch (mode) {
+    case StreamCaptureMode::kStreamCaptureModeGlobal:
+      return DeviceRuntime::kStreamCaptureModeGlobal;
+    case StreamCaptureMode::kStreamCaptureModeThreadLocal:
+      return DeviceRuntime::kStreamCaptureModeThreadLocal;
+    case StreamCaptureMode::kStreamCaptureModeRelaxed:
+      return DeviceRuntime::kStreamCaptureModeRelaxed;
+  }
+
+  assert(false && "unsupported stream capture mode");
+  return DeviceRuntime::kStreamCaptureModeRelaxed;
+}
+"""
+        if any(
+            param.type == "StreamCaptureMode"
+            for function in functions
+            for param in function.params
+        )
+        else ""
+    )
     dispatch_functions = "\n".join(
         _write_runtime_dispatch_function(
             function,
@@ -535,6 +636,7 @@ def _write_runtime_dispatch(source_path, source_root, devices):
   return DeviceRuntime::kMemcpyHostToHost;
 }}
 
+{stream_capture_mode_helper}
 }}  // namespace
 
 {dispatch_functions}
@@ -566,7 +668,7 @@ def main():
         for wrapper_device, header_name, target in _DEVICE_HEADERS[device]:
             _write_wrapper(include_root, wrapper_device, header_name, target)
 
-    _write_generated_header(include_root, devices)
+    _write_generated_header(include_root, source_root, devices)
     _write_runtime_dispatch(pathlib.Path(args.source_output), source_root, devices)
 
 

src/CMakeLists.txt

@@ -108,7 +108,8 @@ if(WITH_ASCEND)
         "${ASCEND_HOME}/include/aclnnop")
     target_link_libraries(infinirt PUBLIC
         "${ASCEND_HOME}/lib64/libascendcl.so"
-        "${ASCEND_HAL_LIB}")
+        "${ASCEND_HAL_LIB}"
+        ${CMAKE_DL_LIBS})
 endif()
 
 target_include_directories(infinirt

src/native/ascend/runtime_.h

@@ -1,6 +1,8 @@
 #ifndef INFINI_RT_ASCEND_RUNTIME__H_
 #define INFINI_RT_ASCEND_RUNTIME__H_
 
+#include <dlfcn.h>
+
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
@@ -16,13 +18,19 @@ namespace infini::rt::runtime {
 
 template <>
 struct Runtime<Device::Type::kAscend>
-    : DeviceRuntime<Runtime<Device::Type::kAscend>> {
+    : GraphRuntime<Runtime<Device::Type::kAscend>> {
   using Error = aclError;
 
   using Stream = aclrtStream;
 
+  using Graph = void*;
+
+  using GraphExec = void*;
+
   using Event = void*;
 
+  using StreamCaptureMode = int;
+
   static constexpr Device::Type kDeviceType = Device::Type::kAscend;
 
   static constexpr Error kSuccess = ACL_SUCCESS;
@@ -41,7 +49,7 @@ struct Runtime<Device::Type::kAscend>
 
   static constexpr auto DeviceSynchronize = aclrtSynchronizeDevice;
 
-  static constexpr auto Malloc = [](void** ptr, size_t size) {
+  static constexpr auto Malloc = [](void** ptr, std::size_t size) {
     return aclrtMalloc(ptr, size, ACL_MEM_MALLOC_HUGE_FIRST);
   };
 
@@ -59,14 +67,14 @@ struct Runtime<Device::Type::kAscend>
 
   static Error MemGetInfo(std::size_t*, std::size_t*) { return Unsupported(); }
 
-  static constexpr auto Memcpy = [](void* dst, const void* src, size_t count,
-                                    aclrtMemcpyKind kind) {
+  static constexpr auto Memcpy = [](void* dst, const void* src,
+                                    std::size_t count, aclrtMemcpyKind kind) {
     return aclrtMemcpy(dst, count, src, count, kind);
   };
 
   static constexpr auto MemcpyAsync = [](void* dst, const void* src,
-                                         size_t count, aclrtMemcpyKind kind,
-                                         Stream stream) {
+                                         std::size_t count,
+                                         aclrtMemcpyKind kind, Stream stream) {
     return aclrtMemcpyAsync(dst, count, src, count, kind, stream);
   };
 
@@ -78,13 +86,14 @@ struct Runtime<Device::Type::kAscend>
 
   static constexpr auto kMemcpyDeviceToDevice = ACL_MEMCPY_DEVICE_TO_DEVICE;
 
-  static constexpr auto Memset = [](void* ptr, int value, size_t count) {
+  static constexpr auto Memset = [](void* ptr, int value, std::size_t count) {
     return aclrtMemset(ptr, count, value, count);
   };
 
-  static Error MemsetAsync(void*, int, std::size_t, Stream) {
-    return Unsupported();
-  }
+  static constexpr auto MemsetAsync = [](void* ptr, int value,
+                                         std::size_t count, Stream stream) {
+    return aclrtMemsetAsync(ptr, count, value, count, stream);
+  };
 
   static constexpr auto StreamCreate = aclrtCreateStream;
 
@@ -112,7 +121,126 @@ struct Runtime<Device::Type::kAscend>
 
   static Error EventElapsedTime(float*, Event, Event) { return Unsupported(); }
 
+  static constexpr StreamCaptureMode kStreamCaptureModeGlobal = 0;
+
+  static constexpr StreamCaptureMode kStreamCaptureModeThreadLocal = 1;
+
+  static constexpr StreamCaptureMode kStreamCaptureModeRelaxed = 2;
+
+  static Error StreamBeginCapture(Stream stream, StreamCaptureMode mode) {
+    const auto& api = GraphApi();
+    if (!api.Available()) {
+      return UnsupportedGraphApi();
+    }
+    return api.CaptureBegin(stream, mode);
+  }
+
+  static Error StreamEndCapture(Stream stream, Graph* graph) {
+    assert(graph != nullptr);
+    const auto& api = GraphApi();
+    if (!api.Available()) {
+      return UnsupportedGraphApi();
+    }
+    int capture_status = 0;
+    void* capturing_model_ri = nullptr;
+    const auto info_status =
+        api.CaptureGetInfo(stream, &capture_status, &capturing_model_ri);
+    if (info_status != ACL_SUCCESS) {
+      return info_status;
+    }
+    void* model_ri = nullptr;
+    const auto status = api.CaptureEnd(stream, &model_ri);
+    if (status != ACL_SUCCESS) {
+      return status;
+    }
+    *graph = model_ri;
+    return ACL_SUCCESS;
+  }
+
+  static Error GraphDestroy(Graph graph) {
+    const auto& api = GraphApi();
+    if (api.Destroy == nullptr || graph == nullptr) {
+      return ACL_SUCCESS;
+    }
+    return api.Destroy(graph);
+  }
+
+  static Error GraphInstantiate(GraphExec* graph_exec, Graph graph) {
+    assert(graph_exec != nullptr);
+    *graph_exec = graph;
+    return ACL_SUCCESS;
+  }
+
+  static Error GraphExecDestroy(GraphExec) { return ACL_SUCCESS; }
+
+  static Error GraphLaunch(GraphExec graph_exec, Stream stream) {
+    const auto& api = GraphApi();
+    if (!api.Available()) {
+      return UnsupportedGraphApi();
+    }
+    return api.ExecuteAsync(graph_exec, stream);
+  }
+
  private:
+  struct RiApi {
+    using CaptureBeginFn = Error (*)(Stream, int);
+    using CaptureGetInfoFn = Error (*)(Stream, int*, void**);
+    using CaptureEndFn = Error (*)(Stream, void**);
+    using RiExecuteAsyncFn = Error (*)(void*, Stream);
+    using RiDestroyFn = Error (*)(void*);
+
+    CaptureBeginFn CaptureBegin = nullptr;
+    CaptureGetInfoFn CaptureGetInfo = nullptr;
+    CaptureEndFn CaptureEnd = nullptr;
+    RiExecuteAsyncFn ExecuteAsync = nullptr;
+    RiDestroyFn Destroy = nullptr;
+
+    bool Available() const {
+      return CaptureBegin != nullptr && CaptureGetInfo != nullptr &&
+             CaptureEnd != nullptr && ExecuteAsync != nullptr;
+    }
+  };
+
+  static const RiApi& GraphApi() {
+    static const RiApi api = [] {
+      RiApi loaded{};
+      // Some CANN releases do not expose the aclmdlRI* graph symbols in
+      // headers or shared objects. Probe them at runtime so non-RI
+      // environments can still build and report graph capture as unsupported.
+      auto load_symbols = [](void* lib, RiApi* api) {
+        api->CaptureBegin = reinterpret_cast<RiApi::CaptureBeginFn>(
+            dlsym(lib, "aclmdlRICaptureBegin"));
+        api->CaptureGetInfo = reinterpret_cast<RiApi::CaptureGetInfoFn>(
+            dlsym(lib, "aclmdlRICaptureGetInfo"));
+        api->CaptureEnd = reinterpret_cast<RiApi::CaptureEndFn>(
+            dlsym(lib, "aclmdlRICaptureEnd"));
+        api->ExecuteAsync = reinterpret_cast<RiApi::RiExecuteAsyncFn>(
+            dlsym(lib, "aclmdlRIExecuteAsync"));
+        api->Destroy =
+            reinterpret_cast<RiApi::RiDestroyFn>(dlsym(lib, "aclmdlRIDestroy"));
+      };
+
+      load_symbols(RTLD_DEFAULT, &loaded);
+      if (loaded.Available()) {
+        return loaded;
+      }
+
+      void* lib = dlopen("libascendcl.so", RTLD_NOW | RTLD_NOLOAD);
+      if (lib == nullptr) {
+        lib = dlopen("libascendcl.so", RTLD_NOW);
+      }
+      if (lib == nullptr) {
+        return loaded;
+      }
+
+      load_symbols(lib, &loaded);
+      return loaded;
+    }();
+    return api;
+  }
+
+  static Error UnsupportedGraphApi() { return static_cast<Error>(1); }
+
   static Error Unsupported() { return static_cast<Error>(1); }
 };