feat: support HYGON RCCL

.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
 
 ---

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()

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`.|
 
 </details>

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)

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

src/backend_device_map.h

@@ -21,6 +21,10 @@ template <>
 struct IsSupportedCombination<BackendType::kNccl, Device::Type::kIluvatar>
     : std::true_type {};
 
+template <>
+struct IsSupportedCombination<BackendType::kNccl, Device::Type::kHygon>
+    : std::true_type {};
+
 template <>
 struct IsSupportedCombination<BackendType::kMccl, Device::Type::kMetax>
     : std::true_type {};

src/backends/ccl/common/comm_instance.h

@@ -20,6 +20,7 @@ struct CclCommInstance : public BackendCommInstance {
       (void)Api::Check(Api::CommDestroy(handle));
       handle = Comm{};
     }
+    return;
   }
 };
 

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 <nccl.h>
-
 #include <cstddef>
 
 #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 <Device::Type device>
+struct NcclDataTypeTraits;
+
 template <Device::Type device>
 struct NcclApi {
   static constexpr BackendType kBackendType = BackendType::kNccl;

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 <nccl.h>
-
 #include <iostream>
 
+#include "backends/ccl/nccl/compat.h"
 #include "return_status_impl.h"
 
 #define INFINI_CHECK_NCCL(result) \

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 <rccl/rccl.h>
+#else
+#include <nccl.h>
+#endif
+
+#endif  // INFINI_CCL_BACKENDS_CCL_NCCL_COMPAT_H_

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<Device::Type::kHygon> {
+#if defined(RCCL_BFLOAT16) && RCCL_BFLOAT16
+  static constexpr ncclDataType_t kBFloat16 = ncclBfloat16;
+#else
+  static constexpr ncclDataType_t kBFloat16 = ncclNumTypes;
+#endif
+};
+
+template <>
+struct CclApi<BackendType::kNccl, Device::Type::kHygon>
+    : NcclApi<Device::Type::kHygon> {};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_NCCL_HYGON_API_H_

src/backends/ccl/nccl/iluvatar/api.h

@@ -6,6 +6,15 @@
 
 namespace infini::ccl {
 
+template <>
+struct NcclDataTypeTraits<Device::Type::kIluvatar> {
+#if defined(__CUDA_BF16_TYPES_EXIST__)
+  static constexpr ncclDataType_t kBFloat16 = ncclBfloat16;
+#else
+  static constexpr ncclDataType_t kBFloat16 = ncclNumTypes;
+#endif
+};
+
 template <>
 struct CclApi<BackendType::kNccl, Device::Type::kIluvatar>
     : NcclApi<Device::Type::kIluvatar> {};

src/backends/ccl/nccl/nvidia/api.h

@@ -6,6 +6,15 @@
 
 namespace infini::ccl {
 
+template <>
+struct NcclDataTypeTraits<Device::Type::kNvidia> {
+#if defined(__CUDA_BF16_TYPES_EXIST__)
+  static constexpr ncclDataType_t kBFloat16 = ncclBfloat16;
+#else
+  static constexpr ncclDataType_t kBFloat16 = ncclNumTypes;
+#endif
+};
+
 template <>
 struct CclApi<BackendType::kNccl, Device::Type::kNvidia>
     : NcclApi<Device::Type::kNvidia> {};

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 <nccl.h>
+#include <string>
 
 #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<DataType, ncclDataType_t, 12> 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 <Device::Type device>
+struct NcclDataTypeMap {
+  static constexpr ConstexprMap<DataType, ncclDataType_t, 12> 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<device>::kBFloat16},
+  }}};
+};
 
 static const ConstexprMap<ReductionOpType, ncclRedOp_t, 5> kNcclOpMap{{{
     {ReductionOpType::kSum, ncclSum},
@@ -39,8 +37,9 @@ static const ConstexprMap<ReductionOpType, ncclRedOp_t, 5> kNcclOpMap{{{
     {ReductionOpType::kAvg, ncclAvg},
 }}};
 
+template <Device::Type device>
 inline ncclDataType_t DataTypeToNcclType(DataType dtype) {
-  auto nccl_dtype = kNcclTypeMap.at(dtype);
+  auto nccl_dtype = NcclDataTypeMap<device>::kMap.at(dtype);
 
   if (nccl_dtype == ncclNumTypes) {
     // This means the requested data type is not supported by NCCL.
@@ -63,7 +62,7 @@ struct CclTypeMap<BackendType::kNccl, device> {
 
   static bool ToBackendDataType(DataType dtype,
                                 typename Api::DataType *backend_dtype) {
-    auto nccl_dtype = DataTypeToNcclType(dtype);
+    auto nccl_dtype = DataTypeToNcclType<device>(dtype);
     if (nccl_dtype == ncclNumTypes) {
       return false;
     }