From 5185908cae5734354fea4b0226642ae3562a29b8 Mon Sep 17 00:00:00 2001 From: Ren_HL <1401141353@qq.com> Date: Thu, 23 Jul 2026 13:48:10 +0800 Subject: [PATCH] feat: support MCCL on Moore Threads --- CMakeLists.txt | 26 ++++++++++++++++++++++++-- README.md | 2 +- examples/CMakeLists.txt | 3 +++ src/CMakeLists.txt | 4 ++++ src/backend_device_map.h | 4 ++++ src/backends/ccl/mccl/moore/api.h | 15 +++++++++++++++ src/backends/ccl/mccl/type_map.h | 8 +++++++- src/devices/moore/runtime_.h | 4 ++++ 8 files changed, 62 insertions(+), 4 deletions(-) create mode 100644 src/backends/ccl/mccl/moore/api.h diff --git a/CMakeLists.txt b/CMakeLists.txt index d2bf727..cec2262 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/README.md b/README.md index 4cca902..661d5cc 100644 --- a/README.md +++ b/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`.| diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 24495ea..2a31f97 100644 --- a/examples/CMakeLists.txt +++ b/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) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 06df487..9584156 100644 --- a/src/CMakeLists.txt +++ b/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() # ========================================================= diff --git a/src/backend_device_map.h b/src/backend_device_map.h index a5fa621..96766c0 100644 --- a/src/backend_device_map.h +++ b/src/backend_device_map.h @@ -25,6 +25,10 @@ template <> struct IsSupportedCombination : std::true_type {}; +template <> +struct IsSupportedCombination + : std::true_type {}; + }; // namespace infini::ccl #endif // INFINI_CCL_BACKEND_DEVICE_MAP_H_ diff --git a/src/backends/ccl/mccl/moore/api.h b/src/backends/ccl/mccl/moore/api.h new file mode 100644 index 0000000..198f961 --- /dev/null +++ b/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 + : McclApi {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_MOORE_API_H_ diff --git a/src/backends/ccl/mccl/type_map.h b/src/backends/ccl/mccl/type_map.h index d3e2363..250445b 100644 --- a/src/backends/ccl/mccl/type_map.h +++ b/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 kMcclTypeMap{{{ {DataType::kInt8, mcclInt8}, {DataType::kInt16, mcclNumTypes}, @@ -24,7 +30,7 @@ static const ConstexprMap kMcclTypeMap{{{ {DataType::kFloat32, mcclFloat32}, {DataType::kFloat64, mcclFloat64}, {DataType::kFloat16, mcclFloat16}, - {DataType::kBFloat16, mcclBfloat16}, + {DataType::kBFloat16, kMcclBFloat16Val}, }}}; static const ConstexprMap kMcclOpMap{{{ diff --git a/src/devices/moore/runtime_.h b/src/devices/moore/runtime_.h index 9da15aa..1d062f3 100644 --- a/src/devices/moore/runtime_.h +++ b/src/devices/moore/runtime_.h @@ -46,6 +46,10 @@ struct Runtime static constexpr auto Memset = musaMemset; + static constexpr auto GetDevice = [](auto &&...args) { + return musaGetDevice(std::forward(args)...); + }; + static constexpr auto SetDevice = musaSetDevice; static constexpr auto DeviceSynchronize = [](auto &&...args) {