feat: support NCCL on Iluvatar

CMakeLists.txt

@@ -210,9 +210,14 @@ if(AUTO_DETECT_BACKENDS)
     endif()
 
     # Detect NCCL Dependencies
-    if(WITH_NVIDIA)
-        find_path(AUTO_NCCL_INC NAMES nccl.h QUIET)
-        find_library(AUTO_NCCL_LIB NAMES nccl QUIET)
+    if(WITH_NVIDIA OR WITH_ILUVATAR)
+        if (WITH_ILUVATAR)
+            set(_NCCL_HINTS /usr/local/corex)
+            message(STATUS "Iluvatar detected. Searching for NCCL 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)
 
         if(AUTO_NCCL_INC AND AUTO_NCCL_LIB)
             set(WITH_NCCL ON)
@@ -334,12 +339,12 @@ if(WITH_OMPI OR WITH_MPICH)
 endif()
 
 if(WITH_NCCL)
-    if (NOT WITH_NVIDIA)
-        message(FATAL_ERROR "NCCL backend requires NVIDIA GPU support. Please enable WITH_NVIDIA.")
+    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`.")
     endif()
 
-    find_library(NCCL_LIB NAMES nccl REQUIRED)
-    find_path(NCCL_INC NAMES nccl.h REQUIRED)
+    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)
 
     include_directories(${NCCL_INC})
 endif()

README.md

@@ -352,7 +352,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 NCCL. Currently available only when `WITH_NVIDIA=ON`.|
+| **NCCL** | Partial | `WITH_NCCL=ON` | Requires NVIDIA or Iluvatar NCCL. Currently available when `WITH_NVIDIA=ON` or `WITH_ILUVATAR=ON`.|
 
 </details>
 

examples/CMakeLists.txt

@@ -53,6 +53,10 @@ foreach(source_file ${EXAMPLE_SOURCES})
         target_link_libraries(${target_name} PRIVATE MPI::MPI_CXX)
     endif()
 
+    if(WITH_NCCL)
+        target_link_libraries(${target_name} PRIVATE "${NCCL_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 

src/backend_device_map.h

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

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

@@ -0,0 +1,15 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_ILUVATAR_API_H_
+#define INFINI_CCL_BACKENDS_CCL_NCCL_ILUVATAR_API_H_
+
+#include "backends/ccl/nccl/api.h"
+#include "devices/iluvatar/runtime_.h"
+
+namespace infini::ccl {
+
+template <>
+struct CclApi<BackendType::kNccl, Device::Type::kIluvatar>
+    : NcclApi<Device::Type::kIluvatar> {};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_NCCL_ILUVATAR_API_H_