fix: 更新 cuda driver，改正异常和警告

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

Cargo.toml

@@ -12,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 = "a815e41" }
-cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a815e41" }
-nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a815e41" }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a815e41" }
-search-corex-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "a815e41" }
+cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "db4692c" }
+cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "db4692c" }
+nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "db4692c" }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "db4692c" }
+search-corex-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "db4692c" }

operators/src/all_reduce/nccl.rs

@@ -1,6 +1,6 @@
 use super::{AllReduce, Args, ReduceOp, args::Meta};
 use crate::{
-    ByteOf, LaunchError, LaunchError, QueueAlloc,
+    ByteOf, LaunchError, QueueAlloc,
     cuda::{Gpu, NcclNode},
     rearrange,
 };
@@ -26,14 +26,6 @@ impl crate::Operator for Operator {
         }
     }
 
-    fn scheme(
-        &mut self,
-        _args: &Self::Args,
-        _max_workspace_size: usize,
-    ) -> Result<usize, LaunchError> {
-        Ok(0)
-    }
-
     fn launch<QA>(
         &self,
         args: &Self::Args,

operators/src/broadcast/nccl/mod.rs

@@ -1,6 +1,6 @@
 use super::{Args, Broadcast, args::Meta};
 use crate::{
-    ByteOf, LaunchError, LaunchError, QueueAlloc,
+    ByteOf, LaunchError, QueueAlloc,
     cuda::{Gpu, NcclNode},
     rearrange,
 };
@@ -26,14 +26,6 @@ impl crate::Operator for Operator {
         }
     }
 
-    fn scheme(
-        &mut self,
-        _args: &Self::Args,
-        _max_workspace_size: usize,
-    ) -> Result<usize, LaunchError> {
-        Ok(0)
-    }
-
     fn launch<QA>(
         &self,
         args: &Self::Args,

operators/src/rearrange/cuda/mod.rs

@@ -7,7 +7,6 @@ use crate::{
 use itertools::Itertools;
 use lru::LruCache;
 use std::cmp::max;
-use std::iter::repeat;
 use std::slice::{from_raw_parts, from_raw_parts_mut};
 use std::{
     ffi::CString,
@@ -71,21 +70,11 @@ impl ArrayStruct {
     }
 
     fn try_into_array<const N: usize>(self) -> Result<[ArrayType; N], LaunchError> {
-        if self.0.len() > N {
-            Err(rank_not_support("ArrayStruct::try_into_array"))
+        let ArrayStruct(vec) = self;
+        if vec.len() <= N {
+            Ok(std::array::from_fn(|i| vec.get(i).copied().unwrap_or(0)))
         } else {
-            let ArrayStruct(vec) = self;
-            if vec.len() == N {
-                Ok(vec.try_into().unwrap())
-            } else {
-                let vec_len = vec.len();
-                Ok(vec
-                    .into_iter()
-                    .chain(repeat(0).take(N - vec_len))
-                    .collect::<Vec<_>>()
-                    .try_into()
-                    .unwrap())
-            }
+            Err(rank_not_support("over length"))
         }
     }
 }
@@ -158,7 +147,7 @@ impl crate::Operator for Operator {
             .collect::<Vec<_>>();
 
         //分离维度，分成grid处理的维度和block处理的维度，与dst的维度相对应
-        let mut block_dim_choose = repeat(false).take(ndim).collect::<Vec<_>>();
+        let mut block_dim_choose = vec![false; ndim];
 
         //TODO 需要优化
         let max_block_size = 256;
@@ -596,9 +585,9 @@ mod test {
         let dt = ty::U64;
 
         let cpu_op = RefOp::new(&Cpu);
-        let gpu_op = Operator::new(&gpu);
+        let gpu_op = Operator::new(gpu);
 
-        let mut r_shape: [usize; N] = shape.clone();
+        let mut r_shape = shape;
         r_shape[0..TRANS_N].reverse();
 
         let trans_param: [usize; TRANS_N] =