fix: 兼容天数

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 = "077b425" }
-cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "077b425" }
-nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "077b425" }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "077b425" }
-search-corex-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "077b425" }
+cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "9eb3b45" }
+cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "9eb3b45" }
+nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "9eb3b45" }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "9eb3b45" }
+search-corex-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "9eb3b45" }

operators/build.rs

@@ -26,18 +26,21 @@ fn main() {
     {
         infini.define()
     }
-    let use_nvidia = cfg!(feature = "nvidia-gpu") && find_cuda_root().is_some();
+
+    // iluvatar
     let use_iluvatar = cfg!(feature = "iluvatar-gpu") && find_corex().is_some();
+    if use_iluvatar {
+        iluvatar.define();
+        cuda.define();
+        return;
+    }
+
+    let use_nvidia = cfg!(feature = "nvidia-gpu") && find_cuda_root().is_some();
     if use_nvidia {
         nvidia.define();
         if find_nccl_root().is_some() {
             nccl.define()
         }
-    }
-    if use_iluvatar {
-        iluvatar.define()
-    }
-    if use_nvidia || use_iluvatar {
-        cuda.define()
+        cuda.define();
     }
 }

operators/src/handle/cuda/alloc.rs

@@ -99,7 +99,6 @@ impl<'ctx> Alloc<DevMem<'ctx>> for &'ctx CurrentCtx {
     fn free(&self, _mem: DevMem<'ctx>) {}
 }
 
-#[cfg(use_nvidia)]
 impl<'ctx> Alloc<DevMem<'ctx>> for Stream<'ctx> {
     #[inline]
     fn alloc(&self, size: usize) -> DevMem<'ctx> {
@@ -112,7 +111,6 @@ impl<'ctx> Alloc<DevMem<'ctx>> for Stream<'ctx> {
     }
 }
 
-#[cfg(use_nvidia)]
 impl<'ctx> QueueAlloc for Stream<'ctx> {
     type Hardware = Gpu;
     type DevMem = DevMem<'ctx>;