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,9 +1,12 @@
 #ifndef INFINI_RT_ASCEND_RUNTIME__H_
 #define INFINI_RT_ASCEND_RUNTIME__H_
 
+#include <dlfcn.h>
+
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <type_traits>
 
 // clang-format off
 #include "acl/acl.h"
@@ -21,8 +24,14 @@ struct Runtime<Device::Type::kAscend>
 
   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 +50,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 +68,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,7 +87,7 @@ 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);
   };
 
@@ -112,6 +121,151 @@ 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;
+
+  struct RIApi {
+    using CaptureBeginFn = aclError (*)(aclrtStream, int);
+    using CaptureGetInfoFn = aclError (*)(aclrtStream, int*, void**);
+    using CaptureEndFn = aclError (*)(aclrtStream, void**);
+    using RIExecuteAsyncFn = aclError (*)(void*, aclrtStream);
+    using RIDestroyFn = aclError (*)(void*);
+
+    CaptureBeginFn capture_begin = nullptr;
+    CaptureGetInfoFn capture_get_info = nullptr;
+    CaptureEndFn capture_end = nullptr;
+    RIExecuteAsyncFn execute_async = nullptr;
+    RIDestroyFn destroy = nullptr;
+
+    bool available() const {
+      return capture_begin != nullptr && capture_get_info != nullptr &&
+             capture_end != nullptr && execute_async != nullptr;
+    }
+  };
+
+  static const RIApi& GraphApi() {
+    static const RIApi api = [] {
+      RIApi loaded{};
+      auto load_symbols = [](void* lib, RIApi* api) {
+        api->capture_begin = reinterpret_cast<RIApi::CaptureBeginFn>(
+            dlsym(lib, "aclmdlRICaptureBegin"));
+        api->capture_get_info = reinterpret_cast<RIApi::CaptureGetInfoFn>(
+            dlsym(lib, "aclmdlRICaptureGetInfo"));
+        api->capture_end = reinterpret_cast<RIApi::CaptureEndFn>(
+            dlsym(lib, "aclmdlRICaptureEnd"));
+        api->execute_async = 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 aclError UnsupportedGraphApi() { return static_cast<aclError>(1); }
+
+  static aclError StreamBeginCapture(Stream stream, StreamCaptureMode mode) {
+    const auto& api = GraphApi();
+    if (!api.available()) {
+      return UnsupportedGraphApi();
+    }
+    return api.capture_begin(stream, mode);
+  }
+
+  static aclError 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.capture_get_info(stream, &capture_status, &capturing_model_ri);
+    if (info_status != ACL_SUCCESS) {
+      return info_status;
+    }
+    void* model_ri = nullptr;
+    const auto status = api.capture_end(stream, &model_ri);
+    if (status != ACL_SUCCESS) {
+      return status;
+    }
+    *graph = model_ri;
+    return ACL_SUCCESS;
+  }
+
+  static aclError GraphDestroy(Graph graph) {
+    const auto& api = GraphApi();
+    if (api.destroy == nullptr || graph == nullptr) {
+      return ACL_SUCCESS;
+    }
+    return api.destroy(graph);
+  }
+
+  static aclError GraphInstantiate(GraphExec* graph_exec, Graph graph) {
+    assert(graph_exec != nullptr);
+    *graph_exec = graph;
+    return ACL_SUCCESS;
+  }
+
+  static aclError GraphExecDestroy(GraphExec) { return ACL_SUCCESS; }
+
+  static aclError GraphLaunch(GraphExec graph_exec, Stream stream) {
+    const auto& api = GraphApi();
+    if (!api.available()) {
+      return UnsupportedGraphApi();
+    }
+    return api.execute_async(graph_exec, stream);
+  }
+
+  static constexpr bool Validate() {
+    DeviceRuntime<Runtime<Device::Type::kAscend>>::Validate();
+    static_assert(sizeof(Graph) > 0,
+                  "`Runtime` must define a `Graph` type alias.");
+    static_assert(sizeof(GraphExec) > 0,
+                  "`Runtime` must define a `GraphExec` type alias.");
+    static_assert(std::is_invocable_v<decltype(StreamBeginCapture), Stream,
+                                      StreamCaptureMode>,
+                  "`Runtime::StreamBeginCapture` must be callable with "
+                  "`(Stream, StreamCaptureMode)`.");
+    static_assert(
+        std::is_invocable_v<decltype(StreamEndCapture), Stream, Graph*>,
+        "`Runtime::StreamEndCapture` must be callable with "
+        "`(Stream, Graph*)`.");
+    static_assert(std::is_invocable_v<decltype(GraphDestroy), Graph>,
+                  "`Runtime::GraphDestroy` must be callable with `(Graph)`.");
+    static_assert(
+        std::is_invocable_v<decltype(GraphInstantiate), GraphExec*, Graph>,
+        "`Runtime::GraphInstantiate` must be callable with "
+        "`(GraphExec*, Graph)`.");
+    static_assert(
+        std::is_invocable_v<decltype(GraphExecDestroy), GraphExec>,
+        "`Runtime::GraphExecDestroy` must be callable with `(GraphExec)`.");
+    static_assert(
+        std::is_invocable_v<decltype(GraphLaunch), GraphExec, Stream>,
+        "`Runtime::GraphLaunch` must be callable with `(GraphExec, Stream)`.");
+    return true;
+  }
+
  private:
   static Error Unsupported() { return static_cast<Error>(1); }
 };