feat: add driver API foundation and NVIDIA implementation

scripts/generate_public_headers.py

@@ -15,6 +15,7 @@
         ("nvidia", "data_type_.h", "native/cuda/nvidia/data_type_.h"),
         ("nvidia", "device_.h", "native/cuda/nvidia/device_.h"),
         ("nvidia", "runtime_.h", "native/cuda/nvidia/runtime_.h"),
+        ("nvidia", "driver_.h", "native/cuda/nvidia/driver_.h"),
     ),
     "iluvatar": (
         ("iluvatar", "data_type_.h", "native/cuda/iluvatar/data_type_.h"),
@@ -116,7 +117,7 @@ def _rewrite_detail_include(match):
 
 
 _DETAIL_INCLUDE_PATTERN = re.compile(
-    r'#include "((?:common|native)/[^"]+|data_type\.h|device\.h|dispatcher\.h|hash\.h|runtime\.h|tensor_view\.h)"'
+    r'#include "((?:common|native)/[^"]+|data_type\.h|device\.h|dispatcher\.h|driver\.h|hash\.h|runtime\.h|tensor_view\.h)"'
 )
 
 
@@ -148,6 +149,7 @@ def _write_detail_headers(include_root, source_root, devices):
         "data_type.h",
         "device.h",
         "dispatcher.h",
+        "driver.h",
         "hash.h",
         "runtime.h",
         "tensor_view.h",
@@ -200,6 +202,7 @@ def _write_generated_header(include_root, source_root, devices):
         "#include <type_traits>",
         f"#include {_detail_include('data_type.h')}",
         f"#include {_detail_include('device.h')}",
+        f"#include {_detail_include('driver.h')}",
         f"#include {_detail_include('hash.h')}",
         f"#include {_detail_include('runtime.h')}",
         f"#include {_detail_include('tensor_view.h')}",
@@ -211,6 +214,10 @@ def _write_generated_header(include_root, source_root, devices):
     for device in devices:
         includes.append(f"#include <infini/rt/{device}/runtime_.h>")
 
+    for device in devices:
+        if any(h == "driver_.h" for _, h, _ in _DEVICE_HEADERS[device]):
+            includes.append(f"#include <infini/rt/{device}/driver_.h>")
+
     runtime_declarations = "\n\n".join(
         f"{function.signature()};" for function in public_runtime_functions
     )

src/driver.h

@@ -0,0 +1,34 @@
+#ifndef INFINI_RT_DRIVER_H_
+#define INFINI_RT_DRIVER_H_
+
+#include <type_traits>
+
+#include "device.h"
+
+namespace infini::rt::driver {
+
+template <Device::Type device_type>
+struct Driver;
+
+template <typename Derived>
+struct DriverBase {
+  static constexpr bool Validate() {
+    static_assert(
+        std::is_same_v<std::remove_cv_t<decltype(Derived::kDeviceType)>,
+                       Device::Type>,
+        "`Driver` must define `static constexpr Device::Type kDeviceType`.");
+    static_assert(sizeof(typename Derived::Result) > 0,
+                  "`Driver` must define a `Result` type alias.");
+    static_assert(std::is_same_v<std::remove_cv_t<decltype(Derived::kSuccess)>,
+                                 typename Derived::Result>,
+                  "`Driver` must define `static constexpr Result kSuccess`.");
+    return true;
+  }
+};
+
+template <typename Derived>
+using DeviceDriver = DriverBase<Derived>;
+
+}  // namespace infini::rt::driver
+
+#endif

src/native/cuda/driver_.h

@@ -0,0 +1,47 @@
+#ifndef INFINI_RT_CUDA_DRIVER__H_
+#define INFINI_RT_CUDA_DRIVER__H_
+
+#include <type_traits>
+
+#include "driver.h"
+
+namespace infini::rt::driver {
+
+template <typename Derived>
+struct CudaDriver : DeviceDriver<Derived> {
+  static constexpr bool Validate() {
+    DeviceDriver<Derived>::Validate();
+    static_assert(sizeof(typename Derived::Function) > 0,
+                  "`Driver` must define a `Function` type alias.");
+    static_assert(sizeof(typename Derived::Module) > 0,
+                  "`Driver` must define a `Module` type alias.");
+    static_assert(sizeof(typename Derived::Stream) > 0,
+                  "`Driver` must define a `Stream` type alias.");
+    static_assert(std::is_invocable_v<decltype(Derived::ModuleLoadData),
+                                      typename Derived::Module*, const void*>,
+                  "`Driver::ModuleLoadData` must be callable with `(Module*, "
+                  "const void*)`.");
+    static_assert(std::is_invocable_v<decltype(Derived::ModuleGetFunction),
+                                      typename Derived::Function*,
+                                      typename Derived::Module, const char*>,
+                  "`Driver::ModuleGetFunction` must be callable with "
+                  "`(Function*, Module, const char*)`.");
+    static_assert(std::is_invocable_v<decltype(Derived::ModuleUnload),
+                                      typename Derived::Module>,
+                  "`Driver::ModuleUnload` must be callable with `(Module)`.");
+    static_assert(
+        std::is_invocable_v<decltype(Derived::LaunchKernel),
+                            typename Derived::Function, unsigned int,
+                            unsigned int, unsigned int, unsigned int,
+                            unsigned int, unsigned int, unsigned int,
+                            typename Derived::Stream, void**, void**>,
+        "`Driver::LaunchKernel` must be callable with "
+        "`(Function, unsigned int, unsigned int, unsigned int, unsigned int, "
+        "unsigned int, unsigned int, unsigned int, Stream, void**, void**)`.");
+    return true;
+  }
+};
+
+}  // namespace infini::rt::driver
+
+#endif

src/native/cuda/nvidia/driver_.h

@@ -0,0 +1,55 @@
+#ifndef INFINI_RT_NVIDIA_DRIVER__H_
+#define INFINI_RT_NVIDIA_DRIVER__H_
+
+// clang-format off
+#include <cuda.h>
+// clang-format on
+
+#include "native/cuda/driver_.h"
+#include "native/cuda/nvidia/device_.h"
+
+namespace infini::rt::driver {
+
+template <>
+struct Driver<Device::Type::kNvidia>
+    : CudaDriver<Driver<Device::Type::kNvidia>> {
+  using Result = CUresult;
+
+  using Module = CUmodule;
+
+  using Function = CUfunction;
+
+  using Stream = CUstream;
+
+  static constexpr Device::Type kDeviceType = Device::Type::kNvidia;
+
+  static constexpr Result kSuccess = CUDA_SUCCESS;
+
+  static constexpr auto ModuleLoadData = cuModuleLoadData;
+
+  static constexpr auto ModuleGetFunction = cuModuleGetFunction;
+
+  static constexpr auto ModuleUnload = cuModuleUnload;
+
+  static constexpr auto FuncGetAttribute = cuFuncGetAttribute;
+
+  static constexpr auto kFuncAttributeSharedSizeBytes =
+      CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES;
+
+  static constexpr auto FuncSetCacheConfig = cuFuncSetCacheConfig;
+
+  static constexpr auto kFuncCachePreferShared = CU_FUNC_CACHE_PREFER_SHARED;
+
+  static constexpr auto FuncSetAttribute = cuFuncSetAttribute;
+
+  static constexpr auto kFuncAttributeMaxDynamicSharedSizeBytes =
+      CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
+
+  static constexpr auto LaunchKernel = cuLaunchKernel;
+};
+
+static_assert(Driver<Device::Type::kNvidia>::Validate());
+
+}  // namespace infini::rt::driver
+
+#endif