issue/1021 - feat: support bf16 in infiniccl on moore_gpu_arch mp_31 with mccl

Author: spike-zhu

src/infiniccl/moore/infiniccl_moore.cc

@@ -23,6 +23,12 @@ inline mcclDataType_t getMcclDtype(infiniDtype_t datatype) {
         return mcclFloat;
     case INFINI_DTYPE_F16:
         return mcclHalf;
+
+#if MARCH_TYPE == 310
+    case INFINI_DTYPE_BF16:
+        return mcclBfloat16;
+#endif
+
     default:
         std::abort();
         return mcclHalf;
@@ -83,9 +89,16 @@ infiniStatus_t allReduce(
     infinicclComm_t comm,
     infinirtStream_t stream) {
 
-    if (datatype != INFINI_DTYPE_F32 && datatype != INFINI_DTYPE_F16) {
-        return INFINI_STATUS_BAD_PARAM;
-    }
+#if MARCH_TYPE == 310
+    CHECK_DTYPE(datatype,
+                INFINI_DTYPE_F32,
+                INFINI_DTYPE_F16,
+                INFINI_DTYPE_BF16);
+#else
+    CHECK_DTYPE(datatype,
+                INFINI_DTYPE_F32,
+                INFINI_DTYPE_F16);
+#endif
 
     CHECK_MCCL(mcclAllReduce(sendbuf, recvbuf, count, getMcclDtype(datatype),
                              getMcclRedOp(op), getMcclComm(comm), getMusaStream(stream)));

xmake.lua

@@ -180,6 +180,12 @@ option("moore-gpu")
     set_description("Whether to compile implementations for Moore Threads GPU")
 option_end()
 
+option("moore-gpu-arch")
+    set_default("mp_31")
+    set_showmenu(true)
+    set_description("Set Moore GPU architecture (e.g. mp_31)")
+option_end()
+
 if has_config("moore-gpu") then
     add_defines("ENABLE_MOORE_API")
     includes("xmake/moore.lua")

xmake/moore.lua

@@ -16,7 +16,22 @@ rule("mu")
         local mcc = MUSA_ROOT .. "/bin/mcc"
         local includedirs = table.concat(target:get("includedirs"), " ")
 
-        local args = {"-c", sourcefile, "-o", objectfile, "-I" .. MUSA_ROOT .. "/include", "-O3", "-fPIC", "-Wall", "-std=c++17", "-pthread"}
+        local args = {
+            "-c", sourcefile,
+            "-o", objectfile,
+            "-I" .. MUSA_ROOT .. "/include",
+            "-O3",
+            "-fPIC",
+            "-Wall",
+            "-std=c++17",
+            "-pthread"
+        }
+        local moore_gpu_arch = get_config("moore-gpu-arch")
+
+        if moore_gpu_arch == "mp_31" then
+            table.insert(args, 1, "--cuda-gpu-arch=mp_31")
+        end
+
         for _, includedir in ipairs(target:get("includedirs")) do
             table.insert(args, "-I" .. includedir)
         end
@@ -76,6 +91,12 @@ target("infiniccl-moore")
     if has_config("ccl") then
         add_links("libmccl.so")
         add_files("../src/infiniccl/moore/*.cc")
+        
+        -- Moore GPU arch with mp_31 support mcclBfloat16 in MCCL
+        if get_config("moore-gpu-arch") == "mp_31" then
+            add_defines("MARCH_TYPE=310")
+            add_cxxflags("-Wno-unused-function")
+        end
     end
     set_languages("cxx17")