feat: support MCCL on Moore Threads (#52)

Author: Ren-HL

CMakeLists.txt

@@ -247,6 +247,24 @@ if(AUTO_DETECT_BACKENDS)
         else()
             message(STATUS "MCCL library/headers not found in MetaX paths.")
         endif()
+    elseif(WITH_MOORE)
+        set(_MCCL_HINTS)
+        foreach(_musa_env MUSA_ROOT MUSA_HOME MUSA_PATH)
+            if(DEFINED ENV{${_musa_env}} AND NOT "$ENV{${_musa_env}}" STREQUAL "")
+                list(APPEND _MCCL_HINTS "$ENV{${_musa_env}}")
+            endif()
+        endforeach()
+        list(APPEND _MCCL_HINTS /usr/local/musa)
+
+        find_path(AUTO_MCCL_INC NAMES mccl.h HINTS ${_MCCL_HINTS} PATH_SUFFIXES include QUIET)
+        find_library(AUTO_MCCL_LIB NAMES mccl HINTS ${_MCCL_HINTS} PATH_SUFFIXES lib lib64 QUIET)
+
+        if(AUTO_MCCL_INC AND AUTO_MCCL_LIB)
+            set(WITH_MCCL ON)
+            message(STATUS "Auto-detected MCCL backend.")
+        else()
+            message(STATUS "MCCL library/headers not found in Moore paths.")
+        endif()
     else()
         message(STATUS "No suitable device environment, skipping MCCL detection.")
     endif()
@@ -372,8 +390,8 @@ if(WITH_NCCL)
 endif()
 
 if(WITH_MCCL)
-    if (NOT WITH_METAX)
-        message(FATAL_ERROR "MCCL backend requires MetaX GPU support. Please enable `WITH_METAX`.")
+    if (NOT WITH_METAX AND NOT WITH_MOORE)
+        message(FATAL_ERROR "MCCL backend requires MetaX or Moore GPU support. Please enable `WITH_METAX` or `WITH_MOORE`.")
     endif()
 
     set(_MCCL_HINTS)
@@ -385,6 +403,10 @@ if(WITH_MCCL)
     endif()
     list(APPEND _MCCL_HINTS /opt/maca)
 
+    if(WITH_MOORE)
+        set(_MCCL_HINTS "${MUSA_ROOT}" /usr/local/musa)
+    endif()
+
     find_library(MCCL_LIB NAMES mccl HINTS ${_MCCL_HINTS} PATH_SUFFIXES lib lib64 REQUIRED)
     find_path(MCCL_INC NAMES mccl.h HINTS ${_MCCL_HINTS} PATH_SUFFIXES include REQUIRED)
 

README.md

@@ -354,7 +354,7 @@ export LD_LIBRARY_PATH=${INFINI_INSTALL}/lib:$LD_LIBRARY_PATH
 | **OpenMPI** | Full | `WITH_OMPI=ON` | The default backend. Requires the OpenMPI development package.|
 | **MPICH** | Full | `WITH_MPICH=ON` | Requires the MPICH development package.|
 | **NCCL** | Partial | `WITH_NCCL=ON` | Requires NVIDIA or Iluvatar NCCL. Currently available when `WITH_NVIDIA=ON` or `WITH_ILUVATAR=ON`.|
-| **MCCL** | Partial | `WITH_MCCL=ON` | Requires MetaX's MCCL. Currently only available when `WITH_METAX=ON`.|
+| **MCCL** | Partial | `WITH_MCCL=ON` | Requires MetaX or Moore MCCL. Currently available when `WITH_METAX=ON` or `WITH_MOORE=ON`.|
 
 </details>
 

examples/CMakeLists.txt

@@ -43,6 +43,9 @@ foreach(source_file ${EXAMPLE_SOURCES})
     if(WITH_MOORE)
         target_link_libraries(${target_name} PRIVATE ${MUSART_LIB})
         target_compile_options(${target_name} PRIVATE "-x" "musa")
+        if(WITH_MCCL)
+            target_compile_definitions(${target_name} PRIVATE INFINI_CCL_MCCL_BFLOAT16_UNSUPPORTED)
+        endif()
     endif()
 
     if(WITH_CAMBRICON)

src/CMakeLists.txt

@@ -243,6 +243,10 @@ if(WITH_MCCL)
     target_sources(infiniccl PRIVATE ${MCCL_SRCS})
     target_include_directories(infiniccl PRIVATE ${MCCL_INC})
     target_link_libraries(infiniccl PRIVATE ${MCCL_LIB})
+
+    if(WITH_MOORE)
+        target_compile_definitions(infiniccl PRIVATE INFINI_CCL_MCCL_BFLOAT16_UNSUPPORTED)
+    endif()
 endif()
 
 # =========================================================

src/backend_device_map.h

@@ -25,6 +25,10 @@ template <>
 struct IsSupportedCombination<BackendType::kMccl, Device::Type::kMetax>
     : std::true_type {};
 
+template <>
+struct IsSupportedCombination<BackendType::kMccl, Device::Type::kMoore>
+    : std::true_type {};
+
 };  // namespace infini::ccl
 
 #endif  // INFINI_CCL_BACKEND_DEVICE_MAP_H_

src/backends/ccl/mccl/moore/api.h

@@ -0,0 +1,15 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_MOORE_API_H_
+#define INFINI_CCL_BACKENDS_CCL_MCCL_MOORE_API_H_
+
+#include "backends/ccl/mccl/api.h"
+#include "devices/moore/runtime_.h"
+
+namespace infini::ccl {
+
+template <>
+struct CclApi<BackendType::kMccl, Device::Type::kMoore>
+    : McclApi<Device::Type::kMoore> {};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_MCCL_MOORE_API_H_

src/backends/ccl/mccl/type_map.h

@@ -12,6 +12,12 @@
 
 namespace infini::ccl {
 
+#if defined(INFINI_CCL_MCCL_BFLOAT16_UNSUPPORTED)
+constexpr mcclDataType_t kMcclBFloat16Val = mcclNumTypes;
+#else
+constexpr mcclDataType_t kMcclBFloat16Val = mcclBfloat16;
+#endif
+
 static const ConstexprMap<DataType, mcclDataType_t, 12> kMcclTypeMap{{{
     {DataType::kInt8, mcclInt8},
     {DataType::kInt16, mcclNumTypes},
@@ -24,7 +30,7 @@ static const ConstexprMap<DataType, mcclDataType_t, 12> kMcclTypeMap{{{
     {DataType::kFloat32, mcclFloat32},
     {DataType::kFloat64, mcclFloat64},
     {DataType::kFloat16, mcclFloat16},
-    {DataType::kBFloat16, mcclBfloat16},
+    {DataType::kBFloat16, kMcclBFloat16Val},
 }}};
 
 static const ConstexprMap<ReductionOpType, mcclRedOp_t, 5> kMcclOpMap{{{

src/devices/moore/runtime_.h

@@ -46,6 +46,10 @@ struct Runtime<Device::Type::kMoore>
 
   static constexpr auto Memset = musaMemset;
 
+  static constexpr auto GetDevice = [](auto &&...args) {
+    return musaGetDevice(std::forward<decltype(args)>(args)...);
+  };
+
   static constexpr auto SetDevice = musaSetDevice;
 
   static constexpr auto DeviceSynchronize = [](auto &&...args) {