diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 6b94aa0..81fc28e 100644 --- a/.github/pull_request_template.md +++ b/.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 --- diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c52d13..d2bf727 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/README.md b/README.md index ce4cc60..4cca902 100644 --- a/README.md +++ b/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`.| diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 9c052c7..24495ea 100644 --- a/examples/CMakeLists.txt +++ b/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 diff --git a/scripts/gen_bridge.py b/scripts/gen_bridge.py index 9dc3cd8..211080b 100644 --- a/scripts/gen_bridge.py +++ b/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}"') diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index a17206f..06df487 100644 --- a/src/CMakeLists.txt +++ b/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 --- # ========================================================= diff --git a/src/backend.h b/src/backend.h index fbcbea0..0b02662 100644 --- a/src/backend.h +++ b/src/backend.h @@ -58,6 +58,11 @@ struct BackendPriority { static constexpr int value = 10; }; +template <> +struct BackendPriority { + static constexpr int value = 10; +}; + } // namespace infini::ccl #endif // INFINI_CCL_BACKEND_H_ diff --git a/src/backend_device_map.h b/src/backend_device_map.h index f5d8419..a5fa621 100644 --- a/src/backend_device_map.h +++ b/src/backend_device_map.h @@ -21,6 +21,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/api.h b/src/backends/ccl/mccl/api.h new file mode 100644 index 0000000..06398d2 --- /dev/null +++ b/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 + +#include + +#include "backends/ccl/common/api.h" +#include "logging.h" +#include "return_status_impl.h" +#include "runtime.h" + +namespace infini::ccl { + +template +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::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_ diff --git a/src/backends/ccl/mccl/checks.h b/src/backends/ccl/mccl/checks.h new file mode 100644 index 0000000..c9f109d --- /dev/null +++ b/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 + +#include + +#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_ diff --git a/src/backends/ccl/mccl/impl/all_reduce.h b/src/backends/ccl/mccl/impl/all_reduce.h new file mode 100644 index 0000000..826dd89 --- /dev/null +++ b/src/backends/ccl/mccl/impl/all_reduce.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_REDUCE_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_REDUCE_H_ + +#include "backends/ccl/common/impl/all_reduce.h" + +namespace infini::ccl { + +template +class AllReduceImpl + : public CclAllReduceImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_ALL_REDUCE_H_ diff --git a/src/backends/ccl/mccl/impl/comm_destroy.h b/src/backends/ccl/mccl/impl/comm_destroy.h new file mode 100644 index 0000000..ff616b8 --- /dev/null +++ b/src/backends/ccl/mccl/impl/comm_destroy.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_COMM_DESTROY_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_COMM_DESTROY_H_ + +#include "backends/ccl/common/impl/comm_destroy.h" + +namespace infini::ccl { + +template +class CommDestroyImpl + : public CclCommDestroyImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_COMM_DESTROY_H_ diff --git a/src/backends/ccl/mccl/impl/comm_init_rank.h b/src/backends/ccl/mccl/impl/comm_init_rank.h new file mode 100644 index 0000000..a42fdd6 --- /dev/null +++ b/src/backends/ccl/mccl/impl/comm_init_rank.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_COMM_INIT_RANK_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_COMM_INIT_RANK_H_ + +#include "backends/ccl/common/impl/comm_init_rank.h" + +namespace infini::ccl { + +template +class CommInitRankImpl + : public CclCommInitRankImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_COMM_INIT_RANK_H_ diff --git a/src/backends/ccl/mccl/impl/get_unique_id.h b/src/backends/ccl/mccl/impl/get_unique_id.h new file mode 100644 index 0000000..c68b0ea --- /dev/null +++ b/src/backends/ccl/mccl/impl/get_unique_id.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_GET_UNIQUE_ID_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_GET_UNIQUE_ID_H_ + +#include "backends/ccl/common/impl/get_unique_id.h" + +namespace infini::ccl { + +template +class GetUniqueIdImpl + : public CclGetUniqueIdImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_IMPL_GET_UNIQUE_ID_H_ diff --git a/src/backends/ccl/mccl/metax/api.h b/src/backends/ccl/mccl/metax/api.h new file mode 100644 index 0000000..0817a01 --- /dev/null +++ b/src/backends/ccl/mccl/metax/api.h @@ -0,0 +1,15 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_METAX_API_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_METAX_API_H_ + +#include "backends/ccl/mccl/api.h" +#include "devices/metax/runtime_.h" + +namespace infini::ccl { + +template <> +struct CclApi + : McclApi {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_METAX_API_H_ diff --git a/src/backends/ccl/mccl/type_map.h b/src/backends/ccl/mccl/type_map.h new file mode 100644 index 0000000..d3e2363 --- /dev/null +++ b/src/backends/ccl/mccl/type_map.h @@ -0,0 +1,77 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_MCCL_TYPE_MAP_H_ +#define INFINI_CCL_BACKENDS_CCL_MCCL_TYPE_MAP_H_ + +#include + +#include + +#include "backends/ccl/common/api.h" +#include "comm_impl.h" +#include "data_type_impl.h" +#include "logging.h" + +namespace infini::ccl { + +static const ConstexprMap kMcclTypeMap{{{ + {DataType::kInt8, mcclInt8}, + {DataType::kInt16, mcclNumTypes}, + {DataType::kInt32, mcclInt32}, + {DataType::kInt64, mcclInt64}, + {DataType::kUInt8, mcclUint8}, + {DataType::kUInt16, mcclNumTypes}, + {DataType::kUInt32, mcclUint32}, + {DataType::kUInt64, mcclUint64}, + {DataType::kFloat32, mcclFloat32}, + {DataType::kFloat64, mcclFloat64}, + {DataType::kFloat16, mcclFloat16}, + {DataType::kBFloat16, mcclBfloat16}, +}}}; + +static const ConstexprMap kMcclOpMap{{{ + {ReductionOpType::kSum, mcclSum}, + {ReductionOpType::kProd, mcclProd}, + {ReductionOpType::kMax, mcclMax}, + {ReductionOpType::kMin, mcclMin}, + {ReductionOpType::kAvg, mcclAvg}, +}}}; + +inline mcclDataType_t DataTypeToMcclType(DataType dtype) { + auto mccl_dtype = kMcclTypeMap.at(dtype); + + if (mccl_dtype == mcclNumTypes) { + LOG(("DataType '" + std::string(kDataTypeToDesc.at(dtype)) + + "' is not supported by the MCCL backend") + .c_str()); + } + + return mccl_dtype; +} + +inline mcclRedOp_t RedOpToMcclOp(ReductionOpType red_op) { + return kMcclOpMap.at(red_op); +} + +template +struct CclTypeMap { + using Api = CclApi; + + static bool ToBackendDataType(DataType dtype, + typename Api::DataType *backend_dtype) { + auto mccl_dtype = DataTypeToMcclType(dtype); + if (mccl_dtype == mcclNumTypes) { + return false; + } + *backend_dtype = mccl_dtype; + return true; + } + + static bool ToBackendRedOp(ReductionOpType red_op, + typename Api::RedOp *backend_op) { + *backend_op = RedOpToMcclOp(red_op); + return true; + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_MCCL_TYPE_MAP_H_