diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 8efbdf7..cfe73dd 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -64,7 +64,7 @@ Please check all the platforms and/or backends this PR affects (i.e., code is to - [ ] OpenMPI - [ ] MPICH -- [ ] NCCL +- [ ] NCCL/RCCL - [ ] MCCL ## Performance Impact @@ -114,7 +114,7 @@ See `CONTRIBUTING.md` ยง Pull Requests for the official testing requirements and - [ ] OpenMPI - [ ] MPICH -- [ ] NCCL +- [ ] NCCL/RCCL - [ ] MCCL --- diff --git a/CMakeLists.txt b/CMakeLists.txt index 4277c1c..1b006ce 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -250,20 +250,40 @@ if(AUTO_DETECT_BACKENDS) endif() # Detect NCCL Dependencies - if(WITH_NVIDIA OR WITH_ILUVATAR) + if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON) + set(_NCCL_HINTS) + set(_NCCL_HEADER_NAMES nccl.h) + set(_NCCL_LIBRARY_NAMES nccl) + if (WITH_ILUVATAR) set(_NCCL_HINTS /usr/local/corex) message(STATUS "Iluvatar detected. Searching for NCCL in ${_NCCL_HINTS}") + elseif(WITH_HYGON) + foreach(_hygon_nccl_env DTKROOT DTK_ROOT ROCM_PATH) + if(DEFINED ENV{${_hygon_nccl_env}} AND NOT "$ENV{${_hygon_nccl_env}}" STREQUAL "") + list(APPEND _NCCL_HINTS "$ENV{${_hygon_nccl_env}}") + endif() + endforeach() + + if(DEFINED ENV{HIP_PATH} AND NOT "$ENV{HIP_PATH}" STREQUAL "") + get_filename_component(_HYGON_NCCL_ROOT_FROM_HIP "$ENV{HIP_PATH}/.." ABSOLUTE) + list(APPEND _NCCL_HINTS "${_HYGON_NCCL_ROOT_FROM_HIP}") + endif() + + list(APPEND _NCCL_HINTS /opt/dtk /opt/dtk/rccl) + set(_NCCL_HEADER_NAMES rccl/rccl.h rccl.h) + set(_NCCL_LIBRARY_NAMES rccl) + message(STATUS "Hygon detected. Searching for NCCL-compatible RCCL in ${_NCCL_HINTS}") endif() - find_path(AUTO_NCCL_INC NAMES nccl.h HINTS ${_NCCL_HINTS} PATH_SUFFIXES include QUIET) - find_library(AUTO_NCCL_LIB NAMES nccl HINTS ${_NCCL_HINTS} PATH_SUFFIXES lib lib64 QUIET) + find_path(AUTO_NCCL_INC NAMES ${_NCCL_HEADER_NAMES} HINTS ${_NCCL_HINTS} PATH_SUFFIXES include QUIET) + find_library(AUTO_NCCL_LIB NAMES ${_NCCL_LIBRARY_NAMES} HINTS ${_NCCL_HINTS} PATH_SUFFIXES lib lib64 rccl/lib QUIET) if(AUTO_NCCL_INC AND AUTO_NCCL_LIB) set(WITH_NCCL ON) message(STATUS "Auto-detected NCCL backend.") else() - message(STATUS "NCCL library/headers not found in standard paths.") + message(STATUS "NCCL-compatible library/headers not found in configured paths.") endif() else() message(STATUS "No suitable device environment, skipping NCCL detection.") @@ -470,12 +490,23 @@ if(WITH_OMPI OR WITH_MPICH) endif() if(WITH_NCCL) - if (NOT WITH_NVIDIA AND NOT WITH_ILUVATAR) - message(FATAL_ERROR "NCCL backend requires NVIDIA or Iluvatar GPU support. Please enable `WITH_NVIDIA` or `WITH_ILUVATAR`.") + if (NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON) + message(FATAL_ERROR "NCCL backend requires NVIDIA, Iluvatar, or Hygon GPU support. Please enable `WITH_NVIDIA`, `WITH_ILUVATAR`, or `WITH_HYGON`.") + endif() + + set(_NCCL_HEADER_NAMES nccl.h) + set(_NCCL_LIBRARY_NAMES nccl) + set(NCCL_COMPILE_DEFINITIONS) + + if(WITH_HYGON) + list(APPEND _NCCL_HINTS "${HYGON_DTK_ROOT}" "${HYGON_DTK_ROOT}/rccl") + set(_NCCL_HEADER_NAMES rccl/rccl.h rccl.h) + set(_NCCL_LIBRARY_NAMES rccl) + set(NCCL_COMPILE_DEFINITIONS INFINI_CCL_USE_RCCL) endif() - find_library(NCCL_LIB NAMES nccl HINTS ${_NCCL_HINTS} PATH_SUFFIXES lib lib64 REQUIRED) - find_path(NCCL_INC NAMES nccl.h HINTS ${_NCCL_HINTS} PATH_SUFFIXES include REQUIRED) + find_library(NCCL_LIB NAMES ${_NCCL_LIBRARY_NAMES} HINTS ${_NCCL_HINTS} PATH_SUFFIXES lib lib64 rccl/lib REQUIRED) + find_path(NCCL_INC NAMES ${_NCCL_HEADER_NAMES} HINTS ${_NCCL_HINTS} PATH_SUFFIXES include REQUIRED) include_directories(${NCCL_INC}) endif() diff --git a/README.md b/README.md index 071c38d..4a694d7 100644 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ cmake .. -DWITH_NVIDIA=ON -DWITH_OMPI=ON | **Backend (Communication) Options** ||| | `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_NCCL` | Enable NCCL/RCCL backend | `OFF` | | `WITH_MCCL` | Enable MCCL backend | `OFF` | | **Miscellaneous** ||| | `AUTO_DETECT_DEVICES` | Automatically detect available devices and enable corresponding support | `ON` | @@ -355,7 +355,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`.| +| **NCCL** | Partial | `WITH_NCCL=ON` | Requires NVIDIA or Iluvatar NCCL, or HYGON RCCL. Currently available when `WITH_NVIDIA=ON`, `WITH_ILUVATAR=ON`, or `WITH_HYGON=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 1f4b096..2a39876 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -61,7 +61,12 @@ foreach(source_file ${EXAMPLE_SOURCES}) endif() if(WITH_NCCL) + target_include_directories(${target_name} PRIVATE ${NCCL_INC}) target_link_libraries(${target_name} PRIVATE "${NCCL_LIB}") + + if(NCCL_COMPILE_DEFINITIONS) + target_compile_definitions(${target_name} PRIVATE ${NCCL_COMPILE_DEFINITIONS}) + endif() endif() if(WITH_MCCL) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 42c1171..84ffbcf 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -267,6 +267,10 @@ if(WITH_NCCL) target_sources(infiniccl PRIVATE ${NCCL_SRCS}) target_include_directories(infiniccl PRIVATE ${NCCL_INC}) target_link_libraries(infiniccl PRIVATE ${NCCL_LIB}) + + if(NCCL_COMPILE_DEFINITIONS) + target_compile_definitions(infiniccl PRIVATE ${NCCL_COMPILE_DEFINITIONS}) + endif() endif() # MCCL diff --git a/src/backend_device_map.h b/src/backend_device_map.h index 96766c0..7cf46b0 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 {}; + template <> struct IsSupportedCombination : std::true_type {}; diff --git a/src/backends/ccl/common/comm_instance.h b/src/backends/ccl/common/comm_instance.h index 2bd9293..bf74d18 100644 --- a/src/backends/ccl/common/comm_instance.h +++ b/src/backends/ccl/common/comm_instance.h @@ -20,6 +20,7 @@ struct CclCommInstance : public BackendCommInstance { (void)Api::Check(Api::CommDestroy(handle)); handle = Comm{}; } + return; } }; diff --git a/src/backends/ccl/nccl/api.h b/src/backends/ccl/nccl/api.h index e7b6119..dbbf24c 100644 --- a/src/backends/ccl/nccl/api.h +++ b/src/backends/ccl/nccl/api.h @@ -1,17 +1,19 @@ #ifndef INFINI_CCL_BACKENDS_CCL_NCCL_API_H_ #define INFINI_CCL_BACKENDS_CCL_NCCL_API_H_ -#include - #include #include "backends/ccl/common/api.h" +#include "backends/ccl/nccl/compat.h" #include "logging.h" #include "return_status_impl.h" #include "runtime.h" namespace infini::ccl { +template +struct NcclDataTypeTraits; + template struct NcclApi { static constexpr BackendType kBackendType = BackendType::kNccl; diff --git a/src/backends/ccl/nccl/checks.h b/src/backends/ccl/nccl/checks.h index defdf17..379bf6b 100644 --- a/src/backends/ccl/nccl/checks.h +++ b/src/backends/ccl/nccl/checks.h @@ -1,10 +1,9 @@ #ifndef INFINI_CCL_BACKENDS_CCL_NCCL_CHECKS_H_ #define INFINI_CCL_BACKENDS_CCL_NCCL_CHECKS_H_ -#include - #include +#include "backends/ccl/nccl/compat.h" #include "return_status_impl.h" #define INFINI_CHECK_NCCL(result) \ diff --git a/src/backends/ccl/nccl/compat.h b/src/backends/ccl/nccl/compat.h new file mode 100644 index 0000000..d8951ba --- /dev/null +++ b/src/backends/ccl/nccl/compat.h @@ -0,0 +1,10 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_COMPAT_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_COMPAT_H_ + +#if defined(INFINI_CCL_USE_RCCL) +#include +#else +#include +#endif + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_COMPAT_H_ diff --git a/src/backends/ccl/nccl/hygon/api.h b/src/backends/ccl/nccl/hygon/api.h new file mode 100644 index 0000000..a6092da --- /dev/null +++ b/src/backends/ccl/nccl/hygon/api.h @@ -0,0 +1,24 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_HYGON_API_H_ +#define INFINI_CCL_BACKENDS_CCL_NCCL_HYGON_API_H_ + +#include "backends/ccl/nccl/api.h" +#include "devices/hygon/runtime_.h" + +namespace infini::ccl { + +template <> +struct NcclDataTypeTraits { +#if defined(RCCL_BFLOAT16) && RCCL_BFLOAT16 + static constexpr ncclDataType_t kBFloat16 = ncclBfloat16; +#else + static constexpr ncclDataType_t kBFloat16 = ncclNumTypes; +#endif +}; + +template <> +struct CclApi + : NcclApi {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_NCCL_HYGON_API_H_ diff --git a/src/backends/ccl/nccl/iluvatar/api.h b/src/backends/ccl/nccl/iluvatar/api.h index 9cc84d6..8734174 100644 --- a/src/backends/ccl/nccl/iluvatar/api.h +++ b/src/backends/ccl/nccl/iluvatar/api.h @@ -6,6 +6,15 @@ namespace infini::ccl { +template <> +struct NcclDataTypeTraits { +#if defined(__CUDA_BF16_TYPES_EXIST__) + static constexpr ncclDataType_t kBFloat16 = ncclBfloat16; +#else + static constexpr ncclDataType_t kBFloat16 = ncclNumTypes; +#endif +}; + template <> struct CclApi : NcclApi {}; diff --git a/src/backends/ccl/nccl/nvidia/api.h b/src/backends/ccl/nccl/nvidia/api.h index 4c11990..7de7708 100644 --- a/src/backends/ccl/nccl/nvidia/api.h +++ b/src/backends/ccl/nccl/nvidia/api.h @@ -6,6 +6,15 @@ namespace infini::ccl { +template <> +struct NcclDataTypeTraits { +#if defined(__CUDA_BF16_TYPES_EXIST__) + static constexpr ncclDataType_t kBFloat16 = ncclBfloat16; +#else + static constexpr ncclDataType_t kBFloat16 = ncclNumTypes; +#endif +}; + template <> struct CclApi : NcclApi {}; diff --git a/src/backends/ccl/nccl/type_map.h b/src/backends/ccl/nccl/type_map.h index 8b1d786..deefd50 100644 --- a/src/backends/ccl/nccl/type_map.h +++ b/src/backends/ccl/nccl/type_map.h @@ -1,35 +1,33 @@ #ifndef INFINI_CCL_BACKENDS_CCL_NCCL_TYPE_MAP_H_ #define INFINI_CCL_BACKENDS_CCL_NCCL_TYPE_MAP_H_ -#include +#include #include "backends/ccl/common/api.h" +#include "backends/ccl/nccl/api.h" #include "comm_impl.h" #include "data_type_impl.h" #include "logging.h" namespace infini::ccl { -#if defined(__CUDA_BF16_TYPES_EXIST__) -constexpr ncclDataType_t kNcclBFloat16Val = ncclBfloat16; -#else -constexpr ncclDataType_t kNcclBFloat16Val = ncclNumTypes; -#endif - -static const ConstexprMap kNcclTypeMap{{{ - {DataType::kInt8, ncclInt8}, - {DataType::kInt16, ncclNumTypes}, - {DataType::kInt32, ncclInt32}, - {DataType::kInt64, ncclInt64}, - {DataType::kUInt8, ncclUint8}, - {DataType::kUInt16, ncclNumTypes}, - {DataType::kUInt32, ncclUint32}, - {DataType::kUInt64, ncclUint64}, - {DataType::kFloat32, ncclFloat32}, - {DataType::kFloat64, ncclFloat64}, - {DataType::kFloat16, ncclFloat16}, - {DataType::kBFloat16, kNcclBFloat16Val}, -}}}; +template +struct NcclDataTypeMap { + static constexpr ConstexprMap kMap{{{ + {DataType::kInt8, ncclInt8}, + {DataType::kInt16, ncclNumTypes}, + {DataType::kInt32, ncclInt32}, + {DataType::kInt64, ncclInt64}, + {DataType::kUInt8, ncclUint8}, + {DataType::kUInt16, ncclNumTypes}, + {DataType::kUInt32, ncclUint32}, + {DataType::kUInt64, ncclUint64}, + {DataType::kFloat32, ncclFloat32}, + {DataType::kFloat64, ncclFloat64}, + {DataType::kFloat16, ncclFloat16}, + {DataType::kBFloat16, NcclDataTypeTraits::kBFloat16}, + }}}; +}; static const ConstexprMap kNcclOpMap{{{ {ReductionOpType::kSum, ncclSum}, @@ -39,8 +37,9 @@ static const ConstexprMap kNcclOpMap{{{ {ReductionOpType::kAvg, ncclAvg}, }}}; +template inline ncclDataType_t DataTypeToNcclType(DataType dtype) { - auto nccl_dtype = kNcclTypeMap.at(dtype); + auto nccl_dtype = NcclDataTypeMap::kMap.at(dtype); if (nccl_dtype == ncclNumTypes) { // This means the requested data type is not supported by NCCL. @@ -63,7 +62,7 @@ struct CclTypeMap { static bool ToBackendDataType(DataType dtype, typename Api::DataType *backend_dtype) { - auto nccl_dtype = DataTypeToNcclType(dtype); + auto nccl_dtype = DataTypeToNcclType(dtype); if (nccl_dtype == ncclNumTypes) { return false; }