build: 适配沐曦

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

Author: YdrMaster

Cargo.toml

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

operators/Cargo.toml

@@ -5,12 +5,20 @@ edition.workspace = true
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [features]
-default = ["common-cpu", "opencl", "infini", "nvidia-gpu", "iluvatar-gpu"]
+default = [
+    "common-cpu",
+    "opencl",
+    "infini",
+    "nvidia-gpu",
+    "iluvatar-gpu",
+    "metax-gpu",
+]
 common-cpu = ["gemm"]
 opencl = ["clrt"]
 infini = ["infini-rt", "infini-op", "infini-ccl"]
 nvidia-gpu = ["cuda", "cublas", "nccl", "fslock", "libloading"]
 iluvatar-gpu = ["cuda", "cublas", "fslock", "libloading"]
+metax-gpu = ["cuda", "cublas", "nccl", "fslock", "libloading"]
 
 [dependencies]
 digit-layout = "0.3"
@@ -42,6 +50,7 @@ search-cl-tools.workspace = true
 search-infini-tools.workspace = true
 search-cuda-tools.workspace = true
 search-corex-tools.workspace = true
+search-maca-tools.workspace = true
 
 [dev-dependencies]
 gemm = "0.18"

operators/build.rs

@@ -4,6 +4,7 @@ fn main() {
     use search_corex_tools::find_corex;
     use search_cuda_tools::{find_cuda_root, find_nccl_root};
     use search_infini_tools::{find_infini_ccl, find_infini_op, find_infini_rt};
+    use search_maca_tools::find_maca_root;
 
     let cpu = Cfg::new("use_cpu");
     let cl = Cfg::new("use_cl");
@@ -12,6 +13,7 @@ fn main() {
     let nvidia = Cfg::new("use_nvidia");
     let nccl = Cfg::new("use_nccl");
     let iluvatar = Cfg::new("use_iluvatar");
+    let metax = Cfg::new("use_metax");
 
     if cfg!(feature = "common-cpu") {
         cpu.define()
@@ -27,20 +29,22 @@ fn main() {
         infini.define()
     }
 
-    // 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 {
+    let use_iluvatar = cfg!(feature = "iluvatar-gpu") && find_corex().is_some();
+    let use_metax = cfg!(feature = "metax-gpu") && find_maca_root().is_some();
+    if use_metax {
+        metax.define();
+        nccl.define()
+    } else if use_iluvatar {
+        iluvatar.define()
+    } else if use_nvidia {
         nvidia.define();
         if find_nccl_root().is_some() {
             nccl.define()
         }
-        cuda.define();
+    }
+
+    if use_nvidia || use_iluvatar || use_metax {
+        cuda.define()
     }
 }

operators/src/add/cuda/mod.rs

@@ -162,7 +162,7 @@ mod test {
         rand::rng().fill(&mut b[..]);
         let data_ans = gpu.apply(|ctx| {
             let stream = ctx.stream();
-            #[cfg(use_nvidia)]
+            #[cfg(not(use_iluvatar))]
             let rt = &stream;
             #[cfg(use_iluvatar)]
             let rt = ctx;

operators/src/all_reduce/nccl.rs

@@ -60,10 +60,10 @@ impl crate::Operator for Operator {
 fn convert_enum(op: ReduceOp) -> nccl::ReduceType {
     use nccl::ReduceType::*;
     match op {
-        ReduceOp::Sum => ncclSum,
-        ReduceOp::Prod => ncclProd,
-        ReduceOp::Min => ncclMin,
-        ReduceOp::Max => ncclMax,
-        ReduceOp::Mean => ncclAvg,
+        ReduceOp::Sum => hcclSum,
+        ReduceOp::Prod => hcclProd,
+        ReduceOp::Min => hcclMin,
+        ReduceOp::Max => hcclMax,
+        ReduceOp::Mean => hcclAvg,
     }
 }

operators/src/attention/cuda.rs

@@ -69,7 +69,7 @@ mod test {
 
         let o_ans = gpu.apply(|ctx| {
             let stream = ctx.stream();
-            #[cfg(use_nvidia)]
+            #[cfg(not(use_iluvatar))]
             let rt = &stream;
             #[cfg(use_iluvatar)]
             let rt = ctx;

operators/src/attention_kv_cached/cuda.rs

@@ -80,7 +80,7 @@ mod test {
 
         let o_ans = gpu.apply(|ctx| {
             let stream = ctx.stream();
-            #[cfg(use_nvidia)]
+            #[cfg(not(use_iluvatar))]
             let rt = &stream;
             #[cfg(use_iluvatar)]
             let rt = ctx;

operators/src/fuesd_softmax/cuda/fused_softmax.cuh

@@ -28,7 +28,7 @@ static __device__ void block_padding(
     unsigned int const tok_id,
     unsigned int const seq_len) {
 
-    auto att_idx = threadIdx.x, att_len = blockDim.x;
+    unsigned int att_idx = threadIdx.x, att_len = blockDim.x;
     auto thread_data = mask(tok_id, seq_len, att_idx, att_len)
                          ? float(att[att_idx])
                          : -__FLT_MAX__;

operators/src/handle/cuda/alloc.rs

@@ -0,0 +1,37 @@
+local MACA_ROOT = os.getenv("MACA_PATH") or os.getenv("MACA_HOME") or os.getenv("MACA_ROOT")
+add_includedirs(MACA_ROOT .. "/include")
+add_linkdirs(MACA_ROOT .. "/lib")
+add_links("hcdnn", "hcblas", "hcruntime")
+
+rule("maca")
+    set_extensions(".maca")
+
+    on_load(function (target)
+        target:add("includedirs", "include")
+    end)
+
+    on_build_file(function (target, sourcefile)
+        local objectfile = target:objectfile(sourcefile)
+        os.mkdir(path.directory(objectfile))
+        local htcc = path.join(MACA_ROOT, "htgpu_llvm/bin/htcc")
+        local includedirs = target:get("includedirs") .. " "
+
+        local args = { "-x", "hpcc", "-c", sourcefile, "-o", objectfile, "-I" .. MACA_ROOT .. "/include", "-O3", "-fPIC", "-Werror", "-std=c++17"}
+
+        for _, includedir in ipairs(target:get("includedirs")) do
+            table.insert(args, "-I" .. includedir)
+        end
+
+        os.execv(htcc, args)
+        table.insert(target:objectfiles(), objectfile)
+    end)
+rule_end()
+
+target("lib")
+    set_kind("shared")
+    set_optimize("aggressive")
+    set_languages("cxx17")
+    set_warnings("all", "error")
+    add_cxflags("-lstdc++", "-fPIC", "-Wno-defaulted-function-deleted", "-Wno-strict-aliasing")
+    add_files("src.maca", {rule = "maca"})
+target_end()