feat(thead): support T-Head PPU with NCCL

.github/pull_request_template.md

@@ -59,6 +59,7 @@ Please check all the platforms and/or backends this PR affects (i.e., code is to
 - [ ] Moore Threads GPU
 - [ ] Cambricon MLU
 - [ ] HYGON DCU
+- [ ] T-Head PPU
 
 ### Backend
 
@@ -109,6 +110,7 @@ See `CONTRIBUTING.md` § Pull Requests for the official testing requirements and
 - [ ] Moore Threads GPU
 - [ ] Cambricon MLU
 - [ ] HYGON DCU
+- [ ] T-Head PPU
 
 ### Test Involved Backend
 

CMakeLists.txt

@@ -13,6 +13,7 @@ option(WITH_METAX       "Enable MetaX GPU support"      OFF)
 option(WITH_MOORE       "Enable Moore GPU support"      OFF)
 option(WITH_CAMBRICON   "Enable Cambricon MLU support"  OFF)
 option(WITH_HYGON       "Enable Hygon DCU support"      OFF)
+option(WITH_THEAD       "Enable T-Head PPU support"     OFF)
 
 set(WITH_CPU ON CACHE INTERNAL "CPU backend is always enabled")
 
@@ -40,21 +41,28 @@ if(AUTO_DETECT_DEVICES)
     # NVIDIA
     set(NVIDIA_FOUND FALSE)
 
-    file(GLOB NVIDIA_DEV_FILES "/dev/nvidia0")
+    if(WITH_THEAD)
+        # Explicitly selecting T-Head means NVIDIA auto-detection
+        # must not take effect, because the T-Head SDK provides
+        # an `nvidia-smi` compatible command.
+        message(STATUS "Skipping NVIDIA auto-detection because `WITH_THEAD=ON`.")
+    else()
+        file(GLOB NVIDIA_DEV_FILES "/dev/nvidia0")
 
-    if(NVIDIA_DEV_FILES)
-        set(NVIDIA_FOUND TRUE)
-    elseif(NOT NVIDIA_FOUND)
-        find_program(NVIDIA_SMI_PATH nvidia-smi)
-        if(NVIDIA_SMI_PATH)
-            execute_process(
-                COMMAND ${NVIDIA_SMI_PATH} -L
-                RESULT_VARIABLE SMI_RESULT
-                OUTPUT_QUIET
-                ERROR_QUIET
-            )
-            if(SMI_RESULT EQUAL 0)
-                set(NVIDIA_FOUND TRUE)
+        if(NVIDIA_DEV_FILES)
+            set(NVIDIA_FOUND TRUE)
+        elseif(NOT NVIDIA_FOUND)
+            find_program(NVIDIA_SMI_PATH nvidia-smi)
+            if(NVIDIA_SMI_PATH)
+                execute_process(
+                    COMMAND ${NVIDIA_SMI_PATH} -L
+                    RESULT_VARIABLE SMI_RESULT
+                    OUTPUT_QUIET
+                    ERROR_QUIET
+                )
+                if(SMI_RESULT EQUAL 0)
+                    set(NVIDIA_FOUND TRUE)
+                endif()
             endif()
         endif()
     endif()
@@ -66,6 +74,36 @@ if(AUTO_DETECT_DEVICES)
         message(STATUS "NVIDIA environment not detected.")
     endif()
 
+    # T-Head PPU
+    set(THEAD_FOUND FALSE)
+
+    if(WITH_NVIDIA)
+        message(STATUS "Skipping T-Head auto-detection because `WITH_NVIDIA=ON`.")
+    else()
+        file(GLOB THEAD_DEV_FILES "/dev/alixpu_ppu*")
+
+        if(THEAD_DEV_FILES)
+            set(THEAD_FOUND TRUE)
+        elseif(DEFINED ENV{CUDA_SDK}
+               AND NOT "$ENV{CUDA_SDK}" STREQUAL ""
+               AND EXISTS "$ENV{CUDA_SDK}/bin/nvcc")
+            set(THEAD_FOUND TRUE)
+        elseif(DEFINED ENV{PPU_PATH}
+               AND NOT "$ENV{PPU_PATH}" STREQUAL ""
+               AND EXISTS "$ENV{PPU_PATH}/CUDA_SDK/bin/nvcc")
+            set(THEAD_FOUND TRUE)
+        elseif(EXISTS "/usr/local/PPU_SDK/CUDA_SDK/bin/nvcc")
+            set(THEAD_FOUND TRUE)
+        endif()
+    endif()
+
+    if(THEAD_FOUND)
+        set(WITH_THEAD ON)
+        message(STATUS "T-Head environment detected.")
+    else()
+        message(STATUS "T-Head environment not detected.")
+    endif()
+
     # Iluvatar
     set(ILUVATAR_FOUND FALSE)
 
@@ -217,6 +255,18 @@ if(AUTO_DETECT_DEVICES)
     endif()
 endif()
 
+set(_THEAD_SDK_ROOT "")
+
+if(WITH_THEAD)
+    if(DEFINED ENV{CUDA_SDK} AND NOT "$ENV{CUDA_SDK}" STREQUAL "")
+        set(_THEAD_SDK_ROOT "$ENV{CUDA_SDK}")
+    elseif(DEFINED ENV{PPU_PATH} AND NOT "$ENV{PPU_PATH}" STREQUAL "")
+        set(_THEAD_SDK_ROOT "$ENV{PPU_PATH}/CUDA_SDK")
+    else()
+        set(_THEAD_SDK_ROOT "/usr/local/PPU_SDK/CUDA_SDK")
+    endif()
+endif()
+
 # =========================================================
 # --- AUTO-DETECTION: BACKENDS ---
 # =========================================================
@@ -250,7 +300,7 @@ if(AUTO_DETECT_BACKENDS)
     endif()
 
     # Detect NCCL Dependencies
-    if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON)
+    if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON OR WITH_THEAD)
         set(_NCCL_HINTS)
         set(_NCCL_HEADER_NAMES nccl.h)
         set(_NCCL_LIBRARY_NAMES nccl)
@@ -274,6 +324,13 @@ if(AUTO_DETECT_BACKENDS)
             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}")
+        elseif(WITH_THEAD)
+            list(APPEND _NCCL_HINTS
+                "${_THEAD_SDK_ROOT}"
+                "${_THEAD_SDK_ROOT}/targets/${CMAKE_SYSTEM_PROCESSOR}-linux"
+            )
+            message(STATUS
+                "T-Head detected. Searching for NCCL in ${_NCCL_HINTS}")
         endif()
 
         find_path(AUTO_NCCL_INC NAMES ${_NCCL_HEADER_NAMES} HINTS ${_NCCL_HINTS} PATH_SUFFIXES include QUIET)
@@ -343,6 +400,18 @@ if(WITH_NVIDIA)
     find_package(CUDAToolkit REQUIRED)
 endif()
 
+if(WITH_THEAD)
+    if(NOT EXISTS "${_THEAD_SDK_ROOT}/bin/nvcc")
+        message(FATAL_ERROR
+            "T-Head PPU CUDA compiler was not found under "
+            "`${_THEAD_SDK_ROOT}/bin/nvcc`.")
+    endif()
+
+    set(CMAKE_CUDA_COMPILER "${_THEAD_SDK_ROOT}/bin/nvcc" CACHE FILEPATH
+        "T-Head PPU CUDA compiler" FORCE)
+    enable_language(CUDA)
+    find_package(CUDAToolkit REQUIRED)
+endif()
 if(WITH_ILUVATAR)
     find_program(ILUVATAR_CUDA_COMPILER NAMES clang++ HINTS /usr/local/corex/bin)
 
@@ -490,8 +559,8 @@ if(WITH_OMPI OR WITH_MPICH)
 endif()
 
 if(WITH_NCCL)
-    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`.")
+    if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD)
+        message(FATAL_ERROR "NCCL backend requires NVIDIA, Iluvatar, Hygon, or T-Head device support.")
     endif()
 
     set(_NCCL_HEADER_NAMES nccl.h)
@@ -503,6 +572,11 @@ if(WITH_NCCL)
         set(_NCCL_HEADER_NAMES rccl/rccl.h rccl.h)
         set(_NCCL_LIBRARY_NAMES rccl)
         set(NCCL_COMPILE_DEFINITIONS INFINI_CCL_USE_RCCL)
+    elseif(WITH_THEAD)
+        list(APPEND _NCCL_HINTS
+            "${_THEAD_SDK_ROOT}"
+            "${_THEAD_SDK_ROOT}/targets/${CMAKE_SYSTEM_PROCESSOR}-linux"
+        )
     endif()
 
     find_library(NCCL_LIB NAMES ${_NCCL_LIBRARY_NAMES} HINTS ${_NCCL_HINTS} PATH_SUFFIXES lib lib64 rccl/lib REQUIRED)

README.md

@@ -138,6 +138,7 @@ cmake .. -DWITH_NVIDIA=ON -DWITH_OMPI=ON
 | `WITH_MOORE`      | Enable Moore Threads GPU support  | `OFF` |
 | `WITH_CAMBRICON`  | Enable Cambricon MLU support  | `OFF` |
 | `WITH_HYGON`      | Enable HYGON DCU support  | `OFF` |
+| `WITH_THEAD`      | Enable T-Head PPU support | `OFF` |
 | `WITH_CPU`        | CPU support (always enabled) | `ON` (internal, not user‑settable) |
 | **Backend (Communication) Options** |||
 | `WITH_OMPI`   | Enable OpenMPI backend | `ON` if no backend specified, otherwise `OFF` |
@@ -188,7 +189,7 @@ After having a successful build and a complete `cluster.yaml`, we are ready for
 | `nodes[].ip` | Yes | Node | Node IP address or hostname. Use `localhost` or `127.0.0.1` for the local node. |
 | `nodes[].user` | No | Node | SSH user for this node. Overrides `common_user`. |
 | `nodes[].dir` | No | Node | Node-specific project source directory. Overrides `common_dir` for this node and is useful when the project path differs across hosts. |
-| `nodes[].type` | Yes | Node | Architecture/build label used in build, install, and wrapper paths. Common values include `cpu`, `nvidia`, `iluvatar`, `metax`, `moore`, `cambricon`, and `hygon`. |
+| `nodes[].type` | Yes | Node | Architecture/build label used in build, install, and wrapper paths. Common values include `cpu`, `nvidia`, `iluvatar`, `metax`, `moore`, `cambricon`, `hygon`, and `thead`. |
 | `nodes[].slots` | No | Node | Number of processes to launch on this node. This usually matches the number of devices assigned to the node. Defaults to `8`. |
 | `nodes[].cmake_flags` | No | Node | Node-specific CMake options used during `--build`, such as `-DUSE_CUDA=ON` or `-DUSE_MACA=ON`. Overrides global `cmake_flags` for this node. |
 | `nodes[].backend_env` | No | Node | Node-specific runtime environment variables, such as `CUDA_VISIBLE_DEVICES`, `UCX_TLS`, or `UCX_NET_DEVICES`. Overrides or prepends to global `backend_env` values for this node. |
@@ -345,6 +346,7 @@ export LD_LIBRARY_PATH=${INFINI_INSTALL}/lib:$LD_LIBRARY_PATH
 | **Moore Threads** | Full | Requires MUSA SDK and at least one of `MACA_ROOT` (default `/usr/local/musa`), `MACA_PATH`, and `MUSA_HOME` to be set. |
 | **Cambricon**  | Full | Requires CNToolKit and `NEUWARE_HOME` to be set. |
 | **HYGON**      | Full | Requires HYGON DTK and HYHAL. |
+| **T-Head PPU** | Full | Requires the T-Head PPU SDK. Set `CUDA_SDK` or `PPU_PATH` for a non-default installation. |
 
 </details>
 
@@ -355,7 +357,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, or HYGON RCCL. Currently available when `WITH_NVIDIA=ON`, `WITH_ILUVATAR=ON`, or `WITH_HYGON=ON`.|
+| **NCCL** | Partial | `WITH_NCCL=ON` | Requires an NCCL-compatible library for NVIDIA, Iluvatar, T-Head, or HYGON. Currently available when `WITH_NVIDIA=ON`, `WITH_ILUVATAR=ON`, `WITH_HYGON=ON`, or `WITH_THEAD=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

@@ -26,6 +26,10 @@ foreach(source_file ${EXAMPLE_SOURCES})
         target_link_libraries(${target_name} PRIVATE CUDA::cudart)
     endif()
 
+    if(WITH_THEAD)
+        target_link_libraries(${target_name} PRIVATE CUDA::cudart)
+    endif()
+
     if(WITH_ILUVATAR)
         set_source_files_properties(${source_file} PROPERTIES LANGUAGE CXX)
         set_target_properties(${target_name} PROPERTIES

src/CMakeLists.txt

@@ -68,6 +68,29 @@ if(WITH_NVIDIA)
     )
 endif()
 
+# T-Head PPU
+if(WITH_THEAD)
+    list(APPEND DEVICE_LIST "thead")
+
+    set(THEAD_PATTERNS
+        "devices/cuda/*.cc"
+        "devices/cuda/*.cpp"
+        "devices/cuda/*.cu"
+        "devices/thead/*.cc"
+        "devices/thead/*.cpp"
+        "devices/thead/*.cu"
+    )
+    file(GLOB_RECURSE THEAD_SOURCES ${THEAD_PATTERNS})
+
+    target_sources(infiniccl PRIVATE ${THEAD_SOURCES})
+    target_link_libraries(infiniccl PRIVATE CUDA::cudart CUDA::cuda_driver)
+
+    set_target_properties(infiniccl PROPERTIES
+        CUDA_STANDARD 17
+        CUDA_STANDARD_REQUIRED ON
+    )
+endif()
+
 # Iluvatar
 if(WITH_ILUVATAR)
     list(APPEND DEVICE_LIST "iluvatar")

src/backend_device_map.h

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

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

@@ -0,0 +1,24 @@
+#ifndef INFINI_CCL_BACKENDS_CCL_NCCL_THEAD_API_H_
+#define INFINI_CCL_BACKENDS_CCL_NCCL_THEAD_API_H_
+
+#include "backends/ccl/nccl/api.h"
+#include "devices/thead/runtime_.h"
+
+namespace infini::ccl {
+
+template <>
+struct NcclDataTypeTraits<Device::Type::kThead> {
+#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::kThead>
+    : NcclApi<Device::Type::kThead> {};
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_BACKENDS_CCL_NCCL_THEAD_API_H_

src/device.h

@@ -22,6 +22,7 @@ class Device {
     kKunlun = 7,
     kHygon = 8,
     kQy = 9,
+    kThead = 10,
     kCount
   };
 
@@ -67,6 +68,7 @@ class Device {
           {Type::kKunlun, "kunlun"},
           {Type::kHygon, "hygon"},
           {Type::kQy, "qy"},
+          {Type::kThead, "thead"},
       }}};
 
   static constexpr ConstexprMap<std::string_view, Device::Type,
@@ -82,6 +84,7 @@ class Device {
           {"kunlun", Type::kKunlun},
           {"hygon", Type::kHygon},
           {"qy", Type::kQy},
+          {"thead", Type::kThead},
       }}};
 
   int index_{0};
@@ -97,7 +100,7 @@ using AllDeviceTypes =
     List<Device::Type::kCpu, Device::Type::kNvidia, Device::Type::kCambricon,
          Device::Type::kAscend, Device::Type::kMetax, Device::Type::kMoore,
          Device::Type::kIluvatar, Device::Type::kKunlun, Device::Type::kHygon,
-         Device::Type::kQy>;
+         Device::Type::kQy, Device::Type::kThead>;
 
 // Deferred computation of active devices. The `Filter` and `FilterList`
 // evaluation are nested inside a class template so that `DeviceEnabled`
@@ -161,6 +164,11 @@ struct DevicePriority<Device::Type::kHygon> {
   static constexpr int value = 5;
 };
 
+template <>
+struct DevicePriority<Device::Type::kThead> {
+  static constexpr int value = 5;
+};
+
 enum class MemorySpace : std::uint8_t { kHost = 0, kDevice = 1, kUnknown };
 
 template <Device::Type kDev>