refactor(cust_raw): consolidate CUDA, cuDNN, OptiX bindgen and remove find_cuda_helper

1. Consolidation of bindgen related "*-sys" packages
  - Remove the common dependency of `find_cuda_helper`. Use the cargo
    metadata mechanism instead.
  - Merged all CUDA bindgen-generated code into the cust_raw crate for
    simplicity and maintainability.
  - Add CUDA Runtime API bindgen support.
2. cuDNN and OptiX Integration
  - Split cudnn into cudnn (high-level API) and cudnn-sys (low-level
    bindgens) for better abstraction.
  - Split optix into optix (high-level API) and optix-sys (low-level
    bindgens) for better abstraction.
3. CUDA 12+ Support
  - Updated cust to support CUDA versions >= 12.
  - Added compatibility for CUDA 12.3+ graph API changes:
    - Renamed cuGraphKernelNodeGetParams →
      cuGraphKernelNodeGetParams_v2.
    - Enabled conditional node support for CUDA >= 12.3.

Co-authored-by: Adam Basfop Cavendish <[email protected]>
Co-authored-by: Jorge Ortega <[email protected]>

Author: adamcavendish

Cargo.toml

@@ -11,11 +11,10 @@ members = [
   "examples/optix/*",
   "examples/cuda/cpu/*",
   "examples/cuda/gpu/*",
-
 ]
 
 exclude = [
-    "crates/optix/examples/common"
+    "crates/optix/examples/common",
 ]
 
 [profile.dev.package.rustc_codegen_nvvm]

crates/blastoff/Cargo.toml

@@ -7,8 +7,8 @@ repository = "https://github.com/Rust-GPU/Rust-CUDA"
 
 [dependencies]
 bitflags = "2.8"
-cublas_sys = { version = "0.1", path = "../cublas_sys" }
 cust = { version = "0.3", path = "../cust", features = ["impl_num_complex"] }
+cust_raw = { path = "../cust_raw", features = ["cublas"] }
 num-complex = "0.4.6"
 half = { version = "2.4.1", optional = true }
 

crates/blastoff/src/context.rs

@@ -1,15 +1,20 @@
-use crate::{error::*, sys};
-use cust::stream::Stream;
 use std::ffi::CString;
 use std::mem::{self, MaybeUninit};
 use std::os::raw::c_char;
 use std::ptr;
 
-type Result<T, E = Error> = std::result::Result<T, E>;
+use cust::stream::Stream;
+use cust_raw::cublas_sys;
+use cust_raw::driver_sys;
+
+use super::error::DropResult;
+use super::error::ToResult as _;
+
+type Result<T, E = super::error::Error> = std::result::Result<T, E>;
 
 bitflags::bitflags! {
     /// Configures precision levels for the math in cuBLAS.
-    #[derive(Default)]
+    #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
     pub struct MathMode: u32 {
         /// Highest performance mode which uses compute and intermediate storage precisions
         /// with at least the same number of mantissa and exponent bits as requested. Will
@@ -68,7 +73,7 @@ bitflags::bitflags! {
 /// - [Matrix Multiplication <span style="float:right;">`gemm`</span>](CublasContext::gemm)
 #[derive(Debug)]
 pub struct CublasContext {
-    pub(crate) raw: sys::v2::cublasHandle_t,
+    pub(crate) raw: cublas_sys::cublasHandle_t,
 }
 
 impl CublasContext {
@@ -87,10 +92,10 @@ impl CublasContext {
     pub fn new() -> Result<Self> {
         let mut raw = MaybeUninit::uninit();
         unsafe {
-            sys::v2::cublasCreate_v2(raw.as_mut_ptr()).to_result()?;
-            sys::v2::cublasSetPointerMode_v2(
+            cublas_sys::cublasCreate_v2(raw.as_mut_ptr()).to_result()?;
+            cublas_sys::cublasSetPointerMode_v2(
                 raw.assume_init(),
-                sys::v2::cublasPointerMode_t::CUBLAS_POINTER_MODE_DEVICE,
+                cublas_sys::cublasPointerMode_t::CUBLAS_POINTER_MODE_DEVICE,
             )
             .to_result()?;
             Ok(Self {
@@ -107,7 +112,7 @@ impl CublasContext {
 
         unsafe {
             let inner = mem::replace(&mut ctx.raw, ptr::null_mut());
-            match sys::v2::cublasDestroy_v2(inner).to_result() {
+            match cublas_sys::cublasDestroy_v2(inner).to_result() {
                 Ok(()) => {
                     mem::forget(ctx);
                     Ok(())
@@ -122,7 +127,7 @@ impl CublasContext {
         let mut raw = MaybeUninit::<u32>::uninit();
         unsafe {
             // getVersion can't fail
-            sys::v2::cublasGetVersion_v2(self.raw, raw.as_mut_ptr().cast())
+            cublas_sys::cublasGetVersion_v2(self.raw, raw.as_mut_ptr().cast())
                 .to_result()
                 .unwrap();
 
@@ -140,17 +145,17 @@ impl CublasContext {
     ) -> Result<T> {
         unsafe {
             // cudaStream_t is the same as CUstream
-            sys::v2::cublasSetStream_v2(
+            cublas_sys::cublasSetStream_v2(
                 self.raw,
-                mem::transmute::<*mut cust::sys::CUstream_st, *mut cublas_sys::v2::CUstream_st>(
+                mem::transmute::<*mut driver_sys::CUstream_st, *mut cublas_sys::CUstream_st>(
                     stream.as_inner(),
                 ),
             )
             .to_result()?;
             let res = func(self)?;
             // reset the stream back to NULL just in case someone calls with_stream, then drops the stream, and tries to
             // execute a raw sys function with the context's handle.
-            sys::v2::cublasSetStream_v2(self.raw, ptr::null_mut()).to_result()?;
+            cublas_sys::cublasSetStream_v2(self.raw, ptr::null_mut()).to_result()?;
             Ok(res)
         }
     }
@@ -180,12 +185,12 @@ impl CublasContext {
     /// ```
     pub fn set_atomics_mode(&self, allowed: bool) -> Result<()> {
         unsafe {
-            Ok(sys::v2::cublasSetAtomicsMode(
+            Ok(cublas_sys::cublasSetAtomicsMode(
                 self.raw,
                 if allowed {
-                    sys::v2::cublasAtomicsMode_t::CUBLAS_ATOMICS_ALLOWED
+                    cublas_sys::cublasAtomicsMode_t::CUBLAS_ATOMICS_ALLOWED
                 } else {
-                    sys::v2::cublasAtomicsMode_t::CUBLAS_ATOMICS_NOT_ALLOWED
+                    cublas_sys::cublasAtomicsMode_t::CUBLAS_ATOMICS_NOT_ALLOWED
                 },
             )
             .to_result()?)
@@ -210,10 +215,10 @@ impl CublasContext {
     pub fn get_atomics_mode(&self) -> Result<bool> {
         let mut mode = MaybeUninit::uninit();
         unsafe {
-            sys::v2::cublasGetAtomicsMode(self.raw, mode.as_mut_ptr()).to_result()?;
+            cublas_sys::cublasGetAtomicsMode(self.raw, mode.as_mut_ptr()).to_result()?;
             Ok(match mode.assume_init() {
-                sys::v2::cublasAtomicsMode_t::CUBLAS_ATOMICS_ALLOWED => true,
-                sys::v2::cublasAtomicsMode_t::CUBLAS_ATOMICS_NOT_ALLOWED => false,
+                cublas_sys::cublasAtomicsMode_t::CUBLAS_ATOMICS_ALLOWED => true,
+                cublas_sys::cublasAtomicsMode_t::CUBLAS_ATOMICS_NOT_ALLOWED => false,
             })
         }
     }
@@ -233,9 +238,9 @@ impl CublasContext {
     /// ```
     pub fn set_math_mode(&self, math_mode: MathMode) -> Result<()> {
         unsafe {
-            Ok(sys::v2::cublasSetMathMode(
+            Ok(cublas_sys::cublasSetMathMode(
                 self.raw,
-                mem::transmute::<u32, cublas_sys::v2::cublasMath_t>(math_mode.bits()),
+                mem::transmute::<u32, cublas_sys::cublasMath_t>(math_mode.bits()),
             )
             .to_result()?)
         }
@@ -258,7 +263,7 @@ impl CublasContext {
     pub fn get_math_mode(&self) -> Result<MathMode> {
         let mut mode = MaybeUninit::uninit();
         unsafe {
-            sys::v2::cublasGetMathMode(self.raw, mode.as_mut_ptr()).to_result()?;
+            cublas_sys::cublasGetMathMode(self.raw, mode.as_mut_ptr()).to_result()?;
             Ok(MathMode::from_bits(mode.assume_init() as u32)
                 .expect("Invalid MathMode from cuBLAS"))
         }
@@ -298,7 +303,7 @@ impl CublasContext {
             let path = log_file_name.map(|p| CString::new(p).expect("nul in log_file_name"));
             let path_ptr = path.map_or(ptr::null(), |s| s.as_ptr());
 
-            sys::v2::cublasLoggerConfigure(
+            cublas_sys::cublasLoggerConfigure(
                 enable as i32,
                 log_to_stdout as i32,
                 log_to_stderr as i32,
@@ -315,7 +320,7 @@ impl CublasContext {
     ///
     /// The callback must not panic and unwind.
     pub unsafe fn set_logger_callback(callback: Option<unsafe extern "C" fn(*const c_char)>) {
-        sys::v2::cublasSetLoggerCallback(callback)
+        cublas_sys::cublasSetLoggerCallback(callback)
             .to_result()
             .unwrap();
     }
@@ -324,7 +329,7 @@ impl CublasContext {
     pub fn get_logger_callback() -> Option<unsafe extern "C" fn(*const c_char)> {
         let mut cb = MaybeUninit::uninit();
         unsafe {
-            sys::v2::cublasGetLoggerCallback(cb.as_mut_ptr())
+            cublas_sys::cublasGetLoggerCallback(cb.as_mut_ptr())
                 .to_result()
                 .unwrap();
             cb.assume_init()
@@ -335,7 +340,7 @@ impl CublasContext {
 impl Drop for CublasContext {
     fn drop(&mut self) {
         unsafe {
-            sys::v2::cublasDestroy_v2(self.raw);
+            cublas_sys::cublasDestroy_v2(self.raw);
         }
     }
 }

crates/blastoff/src/error.rs

@@ -1,7 +1,8 @@
-use crate::sys;
-use cust::error::CudaError;
 use std::{ffi::CStr, fmt::Display};
 
+use cust::error::CudaError;
+use cust_raw::cublas_sys;
+
 /// Result that contains the un-dropped value on error.
 pub type DropResult<T> = std::result::Result<(), (CublasError, T)>;
 
@@ -24,7 +25,7 @@ impl std::error::Error for CublasError {}
 impl Display for CublasError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         unsafe {
-            let ptr = sys::v2::cublasGetStatusString(self.into_raw());
+            let ptr = cublas_sys::cublasGetStatusString(self.into_raw());
             let cow = CStr::from_ptr(ptr).to_string_lossy();
             f.write_str(cow.as_ref())
         }
@@ -35,39 +36,41 @@ pub trait ToResult {
     fn to_result(self) -> Result<(), CublasError>;
 }
 
-impl ToResult for sys::v2::cublasStatus_t {
+impl ToResult for cublas_sys::cublasStatus_t {
     fn to_result(self) -> Result<(), CublasError> {
+        use cust_raw::cublas_sys::cublasStatus_t::*;
         use CublasError::*;
 
         Err(match self {
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_SUCCESS => return Ok(()),
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_NOT_INITIALIZED => NotInitialized,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_ALLOC_FAILED => AllocFailed,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_INVALID_VALUE => InvalidValue,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_ARCH_MISMATCH => ArchMismatch,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_MAPPING_ERROR => MappingError,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_EXECUTION_FAILED => ExecutionFailed,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_INTERNAL_ERROR => InternalError,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_NOT_SUPPORTED => NotSupported,
-            sys::v2::cublasStatus_t::CUBLAS_STATUS_LICENSE_ERROR => LicenseError,
+            CUBLAS_STATUS_SUCCESS => return Ok(()),
+            CUBLAS_STATUS_NOT_INITIALIZED => NotInitialized,
+            CUBLAS_STATUS_ALLOC_FAILED => AllocFailed,
+            CUBLAS_STATUS_INVALID_VALUE => InvalidValue,
+            CUBLAS_STATUS_ARCH_MISMATCH => ArchMismatch,
+            CUBLAS_STATUS_MAPPING_ERROR => MappingError,
+            CUBLAS_STATUS_EXECUTION_FAILED => ExecutionFailed,
+            CUBLAS_STATUS_INTERNAL_ERROR => InternalError,
+            CUBLAS_STATUS_NOT_SUPPORTED => NotSupported,
+            CUBLAS_STATUS_LICENSE_ERROR => LicenseError,
         })
     }
 }
 
 impl CublasError {
-    pub fn into_raw(self) -> sys::v2::cublasStatus_t {
+    pub fn into_raw(self) -> cublas_sys::cublasStatus_t {
+        use cust_raw::cublas_sys::cublasStatus_t::*;
         use CublasError::*;
 
         match self {
-            NotInitialized => sys::v2::cublasStatus_t::CUBLAS_STATUS_NOT_INITIALIZED,
-            AllocFailed => sys::v2::cublasStatus_t::CUBLAS_STATUS_ALLOC_FAILED,
-            InvalidValue => sys::v2::cublasStatus_t::CUBLAS_STATUS_INVALID_VALUE,
-            ArchMismatch => sys::v2::cublasStatus_t::CUBLAS_STATUS_ARCH_MISMATCH,
-            MappingError => sys::v2::cublasStatus_t::CUBLAS_STATUS_MAPPING_ERROR,
-            ExecutionFailed => sys::v2::cublasStatus_t::CUBLAS_STATUS_EXECUTION_FAILED,
-            InternalError => sys::v2::cublasStatus_t::CUBLAS_STATUS_INTERNAL_ERROR,
-            NotSupported => sys::v2::cublasStatus_t::CUBLAS_STATUS_NOT_SUPPORTED,
-            LicenseError => sys::v2::cublasStatus_t::CUBLAS_STATUS_LICENSE_ERROR,
+            NotInitialized => CUBLAS_STATUS_NOT_INITIALIZED,
+            AllocFailed => CUBLAS_STATUS_ALLOC_FAILED,
+            InvalidValue => CUBLAS_STATUS_INVALID_VALUE,
+            ArchMismatch => CUBLAS_STATUS_ARCH_MISMATCH,
+            MappingError => CUBLAS_STATUS_MAPPING_ERROR,
+            ExecutionFailed => CUBLAS_STATUS_EXECUTION_FAILED,
+            InternalError => CUBLAS_STATUS_INTERNAL_ERROR,
+            NotSupported => CUBLAS_STATUS_NOT_SUPPORTED,
+            LicenseError => CUBLAS_STATUS_LICENSE_ERROR,
         }
     }
 }

crates/blastoff/src/lib.rs

@@ -10,7 +10,7 @@
 #![allow(clippy::too_many_arguments)]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 
-pub use cublas_sys as sys;
+pub use cust_raw::cublas_sys;
 use num_complex::{Complex32, Complex64};
 
 pub use context::*;
@@ -39,34 +39,34 @@ pub trait BlasDatatype: private::Sealed + cust::memory::DeviceCopy {
     /// The corresponding float type. For complex numbers this means their backing
     /// precision, and for floats it is just themselves.
     type FloatTy: Float;
-    fn to_raw(&self) -> sys::v2::cudaDataType;
+    fn to_raw(&self) -> cublas_sys::cudaDataType;
 }
 
 impl BlasDatatype for f32 {
     type FloatTy = f32;
-    fn to_raw(&self) -> sys::v2::cudaDataType {
-        sys::v2::cudaDataType::CUDA_R_32F
+    fn to_raw(&self) -> cublas_sys::cudaDataType {
+        cublas_sys::cudaDataType::CUDA_R_32F
     }
 }
 
 impl BlasDatatype for f64 {
     type FloatTy = f64;
-    fn to_raw(&self) -> sys::v2::cudaDataType {
-        sys::v2::cudaDataType::CUDA_R_64F
+    fn to_raw(&self) -> cublas_sys::cudaDataType {
+        cublas_sys::cudaDataType::CUDA_R_64F
     }
 }
 
 impl BlasDatatype for Complex32 {
     type FloatTy = f32;
-    fn to_raw(&self) -> sys::v2::cudaDataType {
-        sys::v2::cudaDataType::CUDA_C_32F
+    fn to_raw(&self) -> cublas_sys::cudaDataType {
+        cublas_sys::cudaDataType::CUDA_C_32F
     }
 }
 
 impl BlasDatatype for Complex64 {
     type FloatTy = f64;
-    fn to_raw(&self) -> sys::v2::cudaDataType {
-        sys::v2::cudaDataType::CUDA_C_64F
+    fn to_raw(&self) -> cublas_sys::cudaDataType {
+        cublas_sys::cudaDataType::CUDA_C_64F
     }
 }
 
@@ -106,11 +106,11 @@ pub enum MatrixOp {
 
 impl MatrixOp {
     /// Returns the corresponding `cublasOperation_t` for this operation.
-    pub fn to_raw(self) -> sys::v2::cublasOperation_t {
+    pub fn to_raw(self) -> cublas_sys::cublasOperation_t {
         match self {
-            MatrixOp::None => sys::v2::cublasOperation_t::CUBLAS_OP_N,
-            MatrixOp::Transpose => sys::v2::cublasOperation_t::CUBLAS_OP_T,
-            MatrixOp::ConjugateTranspose => sys::v2::cublasOperation_t::CUBLAS_OP_C,
+            MatrixOp::None => cublas_sys::cublasOperation_t::CUBLAS_OP_N,
+            MatrixOp::Transpose => cublas_sys::cublasOperation_t::CUBLAS_OP_T,
+            MatrixOp::ConjugateTranspose => cublas_sys::cublasOperation_t::CUBLAS_OP_C,
         }
     }
 }

crates/blastoff/src/raw/level1.rs

@@ -1,7 +1,10 @@
-use crate::{sys::v2::*, BlasDatatype};
-use num_complex::{Complex32, Complex64};
 use std::os::raw::c_int;
 
+use cust_raw::cublas_sys::*;
+use num_complex::{Complex32, Complex64};
+
+use crate::BlasDatatype;
+
 pub trait Level1: BlasDatatype {
     unsafe fn amax(
         handle: cublasHandle_t,

crates/blastoff/src/raw/level3.rs

@@ -1,7 +1,10 @@
-use crate::{sys::v2::*, GemmDatatype};
-use num_complex::{Complex32, Complex64};
 use std::os::raw::c_int;
 
+use cust_raw::cublas_sys::*;
+use num_complex::{Complex32, Complex64};
+
+use crate::GemmDatatype;
+
 pub trait GemmOps: GemmDatatype {
     unsafe fn gemm(
         handle: cublasHandle_t,