diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index cfe73dd..76c344c 100644 --- a/.github/pull_request_template.md +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 1b006ce..b9786f2 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/README.md b/README.md index 4a694d7..6a287e1 100644 --- a/README.md +++ b/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. | @@ -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`.| diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 2a39876..266a548 100644 --- a/examples/CMakeLists.txt +++ b/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 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 84ffbcf..5f12ee5 100644 --- a/src/CMakeLists.txt +++ b/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") diff --git a/src/backend_device_map.h b/src/backend_device_map.h index 7cf46b0..f80fd6d 100644 --- a/src/backend_device_map.h +++ b/src/backend_device_map.h @@ -25,6 +25,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/nccl/thead/api.h b/src/backends/ccl/nccl/thead/api.h new file mode 100644 index 0000000..14dbd2b --- /dev/null +++ b/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 { +#if defined(__CUDA_BF16_TYPES_EXIST__) + 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_THEAD_API_H_ diff --git a/src/device.h b/src/device.h index 4809420..6d9f3c3 100644 --- a/src/device.h +++ b/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; + 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 { static constexpr int value = 5; }; +template <> +struct DevicePriority { + static constexpr int value = 5; +}; + enum class MemorySpace : std::uint8_t { kHost = 0, kDevice = 1, kUnknown }; template diff --git a/src/devices/thead/caster_.cuh b/src/devices/thead/caster_.cuh new file mode 100644 index 0000000..1b262f7 --- /dev/null +++ b/src/devices/thead/caster_.cuh @@ -0,0 +1,64 @@ +#ifndef INFINI_CCL_DEVICES_THEAD_CASTER_CUH_ +#define INFINI_CCL_DEVICES_THEAD_CASTER_CUH_ + +#include "caster.h" +#include "data_type_.h" + +namespace infini::ccl { + +template <> +struct HardwareCastImpl { + __host__ __device__ static float Apply(half x) { return __half2float(x); } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static half Apply(float x) { return __float2half(x); } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static float Apply(__nv_bfloat16 x) { + return __bfloat162float(x); + } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static __nv_bfloat16 Apply(float x) { + return __float2bfloat16(x); + } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static __nv_bfloat16 Apply(int x) { + return __int2bfloat16_rn(x); + } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static half Apply(int x) { return __int2half_rn(x); } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static __nv_bfloat16 Apply(double x) { + return __double2bfloat16(x); + } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static half Apply(double x) { return __double2half(x); } +}; + +template <> +struct HardwareCastImpl { + __host__ __device__ static half Apply(__nv_bfloat16 x) { return __half(x); } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_DEVICES_THEAD_CASTER_CUH_ diff --git a/src/devices/thead/data_type_.h b/src/devices/thead/data_type_.h new file mode 100644 index 0000000..e9b825e --- /dev/null +++ b/src/devices/thead/data_type_.h @@ -0,0 +1,26 @@ +#ifndef INFINI_CCL_DEVICES_THEAD_DATA_TYPE_H_ +#define INFINI_CCL_DEVICES_THEAD_DATA_TYPE_H_ + +// clang-format off +#include +#include +// clang-format on + +#include "data_type_impl.h" +#include "devices/thead/device_.h" + +namespace infini::ccl { + +template <> +struct TypeMap { + using type = half; +}; + +template <> +struct TypeMap { + using type = __nv_bfloat16; +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_DEVICES_THEAD_DATA_TYPE_H_ diff --git a/src/devices/thead/device_.h b/src/devices/thead/device_.h new file mode 100644 index 0000000..24930fd --- /dev/null +++ b/src/devices/thead/device_.h @@ -0,0 +1,32 @@ +#ifndef INFINI_CCL_DEVICES_THEAD_DEVICE_H_ +#define INFINI_CCL_DEVICES_THEAD_DEVICE_H_ + +#include + +#include "device.h" +#include "devices/cuda/checks.h" + +namespace infini::ccl { + +template <> +struct DeviceEnabled : std::true_type {}; + +template <> +MemorySpace GetMemorySpace(const void* ptr) { + if (!ptr) { + return MemorySpace::kHost; + } + + cudaPointerAttributes attr; + INFINI_CHECK_CUDA(cudaPointerGetAttributes(&attr, ptr)); + + if (attr.type == cudaMemoryTypeDevice || attr.type == cudaMemoryTypeManaged) { + return MemorySpace::kDevice; + } + + return MemorySpace::kHost; +} + +} // namespace infini::ccl + +#endif // INFINI_CCL_DEVICES_THEAD_DEVICE_H_ diff --git a/src/devices/thead/runtime_.h b/src/devices/thead/runtime_.h new file mode 100644 index 0000000..cc4bc07 --- /dev/null +++ b/src/devices/thead/runtime_.h @@ -0,0 +1,60 @@ +#ifndef INFINI_CCL_DEVICES_THEAD_RUNTIME_H_ +#define INFINI_CCL_DEVICES_THEAD_RUNTIME_H_ + +#include + +// clang-format off +#include +// clang-format on + +#include "devices/cuda/runtime_.h" +#include "devices/thead/device_.h" +#include "logging.h" +#include "return_status_impl.h" + +namespace infini::ccl { + +template <> +struct Runtime + : CudaRuntime> { + using Stream = cudaStream_t; + + static constexpr Device::Type kDeviceType = Device::Type::kThead; + + static constexpr auto Check = + [](auto status, ReturnStatus err_code = ReturnStatus::kSystemError) { + if (status != cudaSuccess) { + LOG(cudaGetErrorString(static_cast(status))); + return err_code; + } + return ReturnStatus::kSuccess; + }; + + static constexpr auto Malloc = [](auto&&... args) { + return cudaMalloc(std::forward(args)...); + }; + + static constexpr auto Memcpy = cudaMemcpy; + + static constexpr auto Free = cudaFree; + + static constexpr auto MemcpyHostToDevice = cudaMemcpyHostToDevice; + + static constexpr auto MemcpyDeviceToHost = cudaMemcpyDeviceToHost; + + static constexpr auto Memset = cudaMemset; + + static constexpr auto GetDevice = cudaGetDevice; + + static constexpr auto SetDevice = cudaSetDevice; + + static constexpr auto DeviceSynchronize = cudaDeviceSynchronize; + + static constexpr auto StreamSynchronize = cudaStreamSynchronize; +}; + +static_assert(Runtime::Validate()); + +} // namespace infini::ccl + +#endif // INFINI_CCL_DEVICES_THEAD_RUNTIME_H_