build: 更新 edition、cuda driver 和依赖项版本

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = ["operators"]
-resolver = "2"
+resolver = "3"
+package.edition = "2024"
 
 [workspace.dependencies]
 clrt = { git = "https://github.com/InfiniTensor/clrt", rev = "984ac7a" }
@@ -11,8 +12,8 @@ infini-op = { git = "https://github.com/InfiniTensor/infini-toolkit", rev = "e83
 infini-ccl = { git = "https://github.com/InfiniTensor/infini-toolkit", rev = "e8362c3" }
 search-infini-tools = { git = "https://github.com/InfiniTensor/infini-toolkit", rev = "e8362c3" }
 
-cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "f3ffbcc" }
-cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "f3ffbcc" }
-nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "f3ffbcc" }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "f3ffbcc" }
-search-corex-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "f3ffbcc" }
+cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "1751f0a" }
+cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "1751f0a" }
+nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "1751f0a" }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "1751f0a" }
+search-corex-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "1751f0a" }

operators/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "operators"
 version = "0.0.0"
-edition = "2021"
+edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [features]
@@ -13,13 +13,13 @@ nvidia-gpu = ["cuda", "cublas", "nccl", "fslock", "libloading"]
 iluvatar-gpu = ["cuda", "cublas", "fslock", "libloading"]
 
 [dependencies]
-digit-layout = "0.2"
-ndarray-layout = "0.1"
+digit-layout = "0.3"
+ndarray-layout = "0.2"
 rayon = "1.10"
-lru = "0.12"
+lru = "0.14"
 num-traits = "0.2"
 itertools = "0.14"
-half = "2.4"
+half = "2.6"
 log = "0.4"
 
 gemm = { version = "0.18", optional = true }

operators/src/add/cuda/mod.rs

@@ -5,10 +5,11 @@ use crate::{
     utils::gcd,
     ByteOf, LaunchError, QueueAlloc, SchemeDiversity,
 };
+use cuda::params;
 use digit_layout::DigitLayout;
 use lru::LruCache;
 use std::{
-    ffi::{c_uint, CString},
+    ffi::c_uint,
     sync::{Arc, Mutex},
 };
 
@@ -70,18 +71,15 @@ impl crate::Operator for Operator {
             b_base,
             ..
         } = args;
-        let params = cuda::params![c_base, a_base, b_base];
 
         self.schemes
             .lock()
             .unwrap()
             .get_or_insert(dt, || compile(&self.handle, dt))
             .launch(
-                CString::new("add").unwrap(),
-                grid_dims as c_uint,
-                block_dims as c_uint,
-                params.as_ptr(),
-                0,
+                c"add",
+                (grid_dims as c_uint, block_dims as c_uint, 0),
+                &params![*c_base, *a_base, *b_base].to_ptrs(),
                 queue_alloc.queue(),
             );
         Ok(())

operators/src/add_rows/cuda/mod.rs

@@ -6,6 +6,7 @@ use crate::{
     utils::gcd,
     ByteOf, LaunchError, QueueAlloc, SchemeDiversity,
 };
+use cuda::params;
 use digit_layout::DigitLayout;
 use lru::LruCache;
 use std::{
@@ -77,7 +78,7 @@ impl crate::Operator for Operator {
         let &[bsi] = cast(&[bsi], unit_idx as usize).as_slice() else {
             todo!()
         };
-        let params = cuda::params![dst_base, src_base, idx_base, bsd, msd, kss, bsi];
+        let params = params![*dst_base, *src_base, *idx_base, bsd, msd, kss, bsi];
         let block = gcd(self.max_threads_block, n);
         let dimx = n.div_ceil(block);
         let key = SchemeKey { dt: dst_layout.dt };
@@ -89,10 +90,8 @@ impl crate::Operator for Operator {
             .clone();
         scheme.module.launch(
             &scheme.name,
-            (b as _, m as _, dimx as _),
-            block as u32,
-            params.as_ptr(),
-            0,
+            ((b as _, m as _, dimx as _), block as u32, 0),
+            &params.to_ptrs(),
             queue_alloc.queue(),
         );
         Ok(())

operators/src/common/mod.rs

@@ -15,7 +15,7 @@ pub use tensor::TensorLayout;
 pub use unsigned::Unsigned;
 pub use workspace::Workspace;
 
-pub(crate) use diversity::{SchemeCacheSize, SchemeDiversity};
+pub(crate) use diversity::SchemeDiversity;
 
 pub mod utils {
     use super::{rank_not_support, type_mismatch, LaunchError};

operators/src/fuesd_softmax/cuda/mod.rs

@@ -6,6 +6,7 @@ use crate::{
     cuda::{Gpu, Handle, ModuleBox},
     strides_not_support, type_not_support, ByteOf, LaunchError, QueueAlloc,
 };
+use cuda::params;
 use digit_layout::types::F16;
 use std::{
     collections::HashMap,
@@ -73,26 +74,22 @@ impl crate::Operator for Operator {
         let sh = (sh / unit) as i32;
         let ss = (ss / unit) as i32;
         let att_len = att_len as u32;
-        let params = cuda::params![att_base, 0i32, sh, ss, att_len];
+        let params = params![*att_base, 0i32, sh, ss, att_len];
 
         if att_len <= block_size {
             scheme.module.launch(
                 &scheme.padding,
-                grid_dims,
-                att_len,
-                params.as_ptr(),
-                0,
+                (grid_dims, att_len, 0),
+                &params.to_ptrs(),
                 queue.queue(),
             );
         } else {
             let num_items_thread = att_len.div_ceil(block_size);
             let smem = (num_items_thread * block_size) as usize;
             scheme.module.launch(
                 &scheme.folding,
-                grid_dims,
-                block_size,
-                params.as_ptr(),
-                smem * size_of::<c_float>(),
+                (grid_dims, block_size, smem * size_of::<c_float>()),
+                &params.to_ptrs(),
                 queue.queue(),
             );
         }

operators/src/gelu/cuda/mod.rs

@@ -5,6 +5,7 @@ use crate::{
     utils::gcd,
     ByteOf, LaunchError, QueueAlloc,
 };
+use cuda::params;
 use digit_layout::types::F16;
 use std::{
     ffi::{c_uint, CString},
@@ -60,15 +61,12 @@ impl crate::Operator for Operator {
             return Err(strides_not_support(""));
         };
 
-        let params = cuda::params![base];
         let block = gcd(self.max_threads_block, d);
 
         self.module.launch(
             CString::new(NAME).unwrap(),
-            (n * d).div_ceil(block) as c_uint,
-            block as u32,
-            params.as_ptr(),
-            0,
+            ((n * d).div_ceil(block) as c_uint, block as c_uint, 0),
+            &params![*base].to_ptrs(),
             queue_alloc.queue(),
         );
         Ok(())

operators/src/handle/cuda/alloc.rs

@@ -108,7 +108,7 @@ impl<'ctx> Alloc<DevMem<'ctx>> for Stream<'ctx> {
 
     #[inline]
     fn free(&self, mem: DevMem<'ctx>) {
-        mem.drop_on(self)
+        Stream::free(self, mem);
     }
 }
 

operators/src/handle/cuda/module.rs

@@ -39,19 +39,11 @@ impl ModuleBox {
     pub fn launch(
         &self,
         name: impl AsRef<CStr>,
-        grid_dims: impl Into<Dim3>,
-        block_dims: impl Into<Dim3>,
-        params: *const *const c_void,
-        shared_mem: usize,
+        attrs: (impl Into<Dim3>, impl Into<Dim3>, usize),
+        params: &[*const c_void],
         stream: &Stream,
     ) {
-        self.load(name, stream.ctx()).launch(
-            grid_dims,
-            block_dims,
-            params,
-            shared_mem,
-            Some(stream),
-        )
+        stream.launch(&self.load(name, stream.ctx()), attrs, params);
     }
 }
 

operators/src/layer_norm/cuda/mod.rs

@@ -4,10 +4,11 @@ use crate::{
     layer_norm::args::Meta,
     shape_not_support, strides_not_support, ByteOf, LaunchError, QueueAlloc, SchemeDiversity,
 };
+use cuda::params;
 use digit_layout::DigitLayout;
 use lru::LruCache;
 use std::{
-    ffi::CString,
+    ffi::{c_uint, CString},
     sync::{Arc, Mutex},
 };
 
@@ -82,17 +83,27 @@ impl crate::Operator for Operator {
 
         let nsy = (nsy / unit) as i32;
         let nsx = (nsx / unit) as i32;
-        let params = cuda::params![y_base, nsy, x_base, nsx, scale_base, bias_base, epsilon];
+        let params = params![
+            *y_base,
+            nsy,
+            *x_base,
+            nsx,
+            *scale_base,
+            *bias_base,
+            *epsilon
+        ];
 
         scheme.module.launch(
             &scheme.name,
-            n as u32,
-            match scheme.ty {
-                SchemeType::Padding => d,
-                SchemeType::Folding { block_size } => block_size,
-            } as u32,
-            params.as_ptr(),
-            0,
+            (
+                n as c_uint,
+                match scheme.ty {
+                    SchemeType::Padding => d,
+                    SchemeType::Folding { block_size } => block_size,
+                } as c_uint,
+                0,
+            ),
+            &params.to_ptrs(),
             queue_alloc.queue(),
         );