From 7815bd28d69f62b5f6bb17a17b79417ab705b795 Mon Sep 17 00:00:00 2001
From: Ringssss <28751159+Ringssss@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:32:35 +0800
Subject: [PATCH] feat(hygon): expose driver modules and device attributes

---
 scripts/generate_public_headers.py |  1 +
 src/native/cuda/hygon/driver_.h    | 54 ++++++++++++++++++++++++++++++
 src/native/cuda/hygon/runtime_.h   | 13 +++++++
 3 files changed, 68 insertions(+)
 create mode 100644 src/native/cuda/hygon/driver_.h

diff --git a/scripts/generate_public_headers.py b/scripts/generate_public_headers.py
index 22c49d8..682f180 100644
--- a/scripts/generate_public_headers.py
+++ b/scripts/generate_public_headers.py
@@ -26,6 +26,7 @@
         ("hygon", "data_type_.h", "native/cuda/hygon/data_type_.h"),
         ("hygon", "device_.h", "native/cuda/hygon/device_.h"),
         ("hygon", "runtime_.h", "native/cuda/hygon/runtime_.h"),
+        ("hygon", "driver_.h", "native/cuda/hygon/driver_.h"),
     ),
     "thead": (
         ("thead", "data_type_.h", "native/cuda/thead/data_type_.h"),
diff --git a/src/native/cuda/hygon/driver_.h b/src/native/cuda/hygon/driver_.h
new file mode 100644
index 0000000..7bc8f5f
--- /dev/null
+++ b/src/native/cuda/hygon/driver_.h
@@ -0,0 +1,54 @@
+#ifndef INFINI_RT_HYGON_DRIVER__H_
+#define INFINI_RT_HYGON_DRIVER__H_
+
+// clang-format off
+#include <cuda.h>
+// clang-format on
+
+#include "native/cuda/driver_.h"
+#include "native/cuda/hygon/device_.h"
+
+namespace infini::rt::driver {
+
+template <>
+struct Driver<Device::Type::kHygon> : CudaDriver<Driver<Device::Type::kHygon>> {
+  using Result = CUresult;
+
+  using Module = CUmodule;
+
+  using Function = CUfunction;
+
+  using Stream = CUstream;
+
+  static constexpr Device::Type kDeviceType = Device::Type::kHygon;
+
+  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::kHygon>::Validate());
+
+}  // namespace infini::rt::driver
+
+#endif
diff --git a/src/native/cuda/hygon/runtime_.h b/src/native/cuda/hygon/runtime_.h
index df88c44..b82f561 100644
--- a/src/native/cuda/hygon/runtime_.h
+++ b/src/native/cuda/hygon/runtime_.h
@@ -40,6 +40,19 @@ struct Runtime<Device::Type::kHygon>
 
   static constexpr auto GetDeviceCount = cudaGetDeviceCount;
 
+  static constexpr auto DeviceGetAttribute = cudaDeviceGetAttribute;
+
+  static constexpr auto kDevAttrWarpSize = cudaDevAttrWarpSize;
+
+  static constexpr auto kDevAttrComputeCapabilityMajor =
+      cudaDevAttrComputeCapabilityMajor;
+
+  static constexpr auto kDevAttrComputeCapabilityMinor =
+      cudaDevAttrComputeCapabilityMinor;
+
+  static constexpr auto kDevAttrMaxSharedMemoryPerBlockOptin =
+      cudaDevAttrMaxSharedMemoryPerBlockOptin;
+
   static constexpr auto DeviceSynchronize = cudaDeviceSynchronize;
 
   static constexpr auto Malloc = [](auto&&... args) {