feat: support MCCL on MetaX (#50)

* feat: support MCCL on MetaX

* docs: update `README.md` to document MetaX MCCL support

* docs: update the GitHub PR template to include MCCL as a valid backend

* fix: add the type mapping for `DataType::kBFloat16` for MCCL

Author: Ziminli

.github/pull_request_template.md

@@ -64,6 +64,7 @@ Please check all the platforms and/or backends this PR affects (i.e., code is to
 - [ ] OpenMPI
 - [ ] MPICH
 - [ ] NCCL
+- [ ] MCCL
 
 ## Performance Impact
 
@@ -112,6 +113,7 @@ See `CONTRIBUTING.md` § Pull Requests for the official testing requirements and
 - [ ] OpenMPI
 - [ ] MPICH
 - [ ] NCCL
+- [ ] MCCL
 
 ---
 

CMakeLists.txt

@@ -21,6 +21,7 @@ set(WITH_CPU ON CACHE INTERNAL "CPU backend is always enabled")
 option(WITH_OMPI "Enable OpenMPI backend" OFF)
 option(WITH_MPICH "Enable MPICH backend" OFF)
 option(WITH_NCCL "Enable NCCL backend"    OFF)
+option(WITH_MCCL "Enable MCCL backend"    OFF)
 
 # =========================================================
 # --- MISC. BUILD OPTIONS ---
@@ -228,10 +229,31 @@ if(AUTO_DETECT_BACKENDS)
     else()
         message(STATUS "No suitable device environment, skipping NCCL detection.")
     endif()
+
+    # Detect MCCL Dependencies
+    if(WITH_METAX)
+        set(_MCCL_HINTS)
+        if(DEFINED ENV{MACA_PATH})
+            list(APPEND _MCCL_HINTS "$ENV{MACA_PATH}")
+        endif()
+        list(APPEND _MCCL_HINTS /opt/maca)
+
+        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 MetaX paths.")
+        endif()
+    else()
+        message(STATUS "No suitable device environment, skipping MCCL detection.")
+    endif()
 endif()
 
 # Fallback: If no backends are enabled or auto-detected, fall back to OpenMPI as the default bootstrap profile.
-if(NOT WITH_OMPI AND NOT WITH_MPICH AND NOT WITH_NCCL)
+if(NOT WITH_OMPI AND NOT WITH_MPICH AND NOT WITH_NCCL AND NOT WITH_MCCL)
     set(WITH_OMPI ON)
     message(STATUS "No backend specified or detected. Defaulting to `WITH_OMPI=ON`")
 endif()
@@ -349,6 +371,26 @@ if(WITH_NCCL)
     include_directories(${NCCL_INC})
 endif()
 
+if(WITH_MCCL)
+    if (NOT WITH_METAX)
+        message(FATAL_ERROR "MCCL backend requires MetaX GPU support. Please enable `WITH_METAX`.")
+    endif()
+
+    set(_MCCL_HINTS)
+    if(MACA_PATH)
+        list(APPEND _MCCL_HINTS "${MACA_PATH}")
+    endif()
+    if(DEFINED ENV{MACA_PATH})
+        list(APPEND _MCCL_HINTS "$ENV{MACA_PATH}")
+    endif()
+    list(APPEND _MCCL_HINTS /opt/maca)
+
+    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)
+
+    include_directories(${MCCL_INC})
+endif()
+
 # Python is required for code generation.
 find_package(Python3 REQUIRED)
 

README.md

@@ -142,6 +142,7 @@ cmake .. -DWITH_NVIDIA=ON -DWITH_OMPI=ON
 | `WITH_OMPI`   | Enable OpenMPI backend | `ON` if no backend specified, otherwise `OFF` |
 | `WITH_MPICH`  | Enable MPICH backend | `OFF` |
 | `WITH_NCCL`   | Enable NCCL backend | `OFF` |
+| `WITH_MCCL`   | Enable MCCL backend | `OFF` |
 | **Miscellaneous** |||
 | `AUTO_DETECT_DEVICES` | Automatically detect available devices and enable corresponding support | `ON` |
 | `AUTO_DETECT_BACKENDS` | Automatically detect available communication backends and enable corresponding support | `OFF` |
@@ -353,6 +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`.|
 
 </details>
 

examples/CMakeLists.txt

@@ -57,6 +57,10 @@ foreach(source_file ${EXAMPLE_SOURCES})
         target_link_libraries(${target_name} PRIVATE "${NCCL_LIB}")
     endif()
 
+    if(WITH_MCCL)
+        target_link_libraries(${target_name} PRIVATE "${MCCL_LIB}")
+    endif()
+
     # Explicitly allow examples to "peek" into the internal `src` and binary dirs.
     # This is necessary because these were marked `PRIVATE` in the library's CMake.
     target_include_directories(${target_name} PRIVATE 

scripts/gen_bridge.py

@@ -31,10 +31,17 @@
     "ompi": ["backends/mpi/ompi/impl"],
     "mpich": ["backends/mpi/ompi/impl"],
     "nccl": ["backends/ccl/nccl/impl"],
+    "mccl": ["backends/ccl/mccl/impl"],
 }
 
 BACKEND_COMMON_HEADERS = {
     "nccl": ["backends/ccl/nccl/type_map.h"],
+    "mccl": ["backends/ccl/mccl/type_map.h"],
+}
+
+CCL_PROVIDER_BACKENDS = {
+    "nccl": "backends/ccl/nccl",
+    "mccl": "backends/ccl/mccl",
 }
 
 # =================================================================
@@ -130,9 +137,10 @@ def generate(project_root, output_dir, devices, backends):
 
         manifest_lines.append(f"\n// --- BACKEND: {bb.upper()} ---")
 
-        if bb == "nccl":
+        provider_root = CCL_PROVIDER_BACKENDS.get(bb)
+        if provider_root:
             for dev in devices:
-                provider_path = f"backends/ccl/nccl/{dev}/api.h"
+                provider_path = f"{provider_root}/{dev}/api.h"
                 if os.path.exists(os.path.join(src_dir, provider_path)):
                     manifest_lines.append(f'#include "{provider_path}"')
 

src/CMakeLists.txt

@@ -235,6 +235,16 @@ if(WITH_NCCL)
     target_link_libraries(infiniccl PRIVATE ${NCCL_LIB})
 endif()
 
+# MCCL
+if(WITH_MCCL)
+    list(APPEND BACKEND_LIST "mccl")
+    file(GLOB_RECURSE MCCL_SRCS "backends/ccl/mccl/*.cc" "backends/ccl/mccl/*.cpp" "backends/ccl/mccl/*.maca")
+
+    target_sources(infiniccl PRIVATE ${MCCL_SRCS})
+    target_include_directories(infiniccl PRIVATE ${MCCL_INC})
+    target_link_libraries(infiniccl PRIVATE ${MCCL_LIB})
+endif()
+
 # =========================================================
 # --- File Generation ---
 # =========================================================

src/backend.h

@@ -58,6 +58,11 @@ struct BackendPriority<BackendType::kNccl> {
   static constexpr int value = 10;
 };
 
+template <>
+struct BackendPriority<BackendType::kMccl> {
+  static constexpr int value = 10;
+};
+
 }  // namespace infini::ccl
 
 #endif  // INFINI_CCL_BACKEND_H_

src/backend_device_map.h

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

src/backends/ccl/mccl/api.h

@@ -0,0 +1,53 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_API_H_
+#define INFINI_CCL_BACKENDS_CCL_MCCL_API_H_
+
+#include <mccl.h>
+
+#include <cstddef>
+
+#include "backends/ccl/common/api.h"
+#include "logging.h"
+#include "return_status_impl.h"
+#include "runtime.h"
+
+namespace infini::ccl {
+
+template <Device::Type device>
+struct McclApi {
+  static constexpr BackendType kBackendType = BackendType::kMccl;
+  static constexpr Device::Type kDeviceType = device;
+
+  using Comm = mcclComm_t;
+  using UniqueId = mcclUniqueId;
+  using Result = mcclResult_t;
+  using DataType = mcclDataType_t;
+  using RedOp = mcclRedOp_t;
+  using Stream = typename Runtime<device>::Stream;
+
+  static ReturnStatus Check(Result result) {
+    if (result != mcclSuccess) {
+      LOG(mcclGetErrorString(result));
+      return ReturnStatus::kSystemError;
+    }
+    return ReturnStatus::kSuccess;
+  }
+
+  static Result GetUniqueId(UniqueId *id) { return mcclGetUniqueId(id); }
+
+  static Result CommInitRank(Comm *comm, int nranks, UniqueId id, int rank) {
+    return mcclCommInitRank(comm, nranks, id, rank);
+  }
+
+  static Result CommDestroy(Comm comm) { return mcclCommDestroy(comm); }
+
+  static Result AllReduce(const void *send_buff, void *recv_buff, size_t count,
+                          DataType data_type, RedOp op, Comm comm,
+                          Stream stream) {
+    return mcclAllReduce(send_buff, recv_buff, count, data_type, op, comm,
+                         stream);
+  }
+};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_MCCL_API_H_

src/backends/ccl/mccl/checks.h

@@ -0,0 +1,31 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_CHECKS_H_
+#define INFINI_CCL_BACKENDS_CCL_MCCL_CHECKS_H_
+
+#include <mccl.h>
+
+#include <iostream>
+
+#include "return_status_impl.h"
+
+#define INFINI_CHECK_MCCL(result) \
+  ::infini::ccl::detail::CheckMcclImpl((result), __FILE__, __LINE__)
+
+namespace infini::ccl {
+
+namespace detail {
+
+inline ReturnStatus CheckMcclImpl(mcclResult_t mccl_result, const char *file,
+                                  int line) {
+  if (mccl_result != mcclSuccess) {
+    std::cerr << "backend(mccl) MCCL error code: " << mccl_result << " at line "
+              << line << " in " << file << std::endl;
+    std::abort();
+  }
+  return ReturnStatus::kSuccess;
+}
+
+}  // namespace detail
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_MCCL_CHECKS_H_