diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 326ede4..bb23ef0 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 @@ -110,6 +111,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 48181fc..f052c92 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") @@ -32,6 +33,39 @@ option(AUTO_DETECT_DEVICES "Automatically detect available devices" ON) option(AUTO_DETECT_BACKENDS "Automatically detect available backends" OFF) option(BUILD_EXAMPLES "Build internal examples" ON) +function(_collect_thead_hints out_var) + if(NOT DEFINED _THEAD_HINTS_MEMO) + set(_hints + /usr/local/PPU_SDK + /usr/local/PPU_SDK/CUDA_SDK + ) + + foreach(_env IN ITEMS PPU_SDK PPU_HOME PPU_PATH) + if(DEFINED ENV{${_env}} AND NOT "$ENV{${_env}}" STREQUAL "") + list(APPEND _hints + "$ENV{${_env}}" + "$ENV{${_env}}/CUDA_SDK" + ) + endif() + endforeach() + + foreach(_env IN ITEMS CUDA_SDK CUDA_HOME CUDA_PATH CUDA_TOOLKIT_ROOT) + if(DEFINED ENV{${_env}} AND NOT "$ENV{${_env}}" STREQUAL "") + list(APPEND _hints + "$ENV{${_env}}" + "$ENV{${_env}}/.." + ) + endif() + endforeach() + + list(REMOVE_DUPLICATES _hints) + set(_THEAD_HINTS_MEMO "${_hints}") + set(_THEAD_HINTS_MEMO "${_hints}" PARENT_SCOPE) + endif() + + set(${out_var} "${_THEAD_HINTS_MEMO}" PARENT_SCOPE) +endfunction() + # ========================================================= # --- AUTO-DETECTION: DEVICES --- # ========================================================= @@ -39,6 +73,7 @@ if(AUTO_DETECT_DEVICES) message(STATUS "Auto-detecting available devices...") # NVIDIA + set(NVIDIA_REQUESTED ${WITH_NVIDIA}) set(NVIDIA_FOUND FALSE) file(GLOB NVIDIA_DEV_FILES "/dev/nvidia0") @@ -60,13 +95,6 @@ if(AUTO_DETECT_DEVICES) endif() endif() - if(NVIDIA_FOUND) - set(WITH_NVIDIA ON) - message(STATUS "Auto-detected NVIDIA environment.") - else() - message(STATUS "NVIDIA environment not detected.") - endif() - # Iluvatar set(ILUVATAR_FOUND FALSE) @@ -216,6 +244,59 @@ if(AUTO_DETECT_DEVICES) else() message(STATUS "Hygon environment not detected.") endif() + + # T-Head + set(THEAD_FOUND FALSE) + + if(WITH_THEAD) + set(THEAD_FOUND TRUE) + else() + _collect_thead_hints(_THEAD_HINTS) + foreach(_thead_hint IN LISTS _THEAD_HINTS) + if(EXISTS "${_thead_hint}/CUDA_SDK" OR EXISTS "${_thead_hint}/bin/ppu-smi" OR EXISTS "${_thead_hint}/ppu-smi/bin/ppu-smi") + set(THEAD_FOUND TRUE) + endif() + endforeach() + + if(NOT THEAD_FOUND) + find_program(THEAD_SMI_PATH + NAMES ppu-smi nvidia-smi + HINTS /usr/local/PPU_SDK/bin /usr/local/PPU_SDK/ppu-smi/bin + ) + if(THEAD_SMI_PATH) + execute_process( + COMMAND ${THEAD_SMI_PATH} -L + RESULT_VARIABLE SMI_RESULT + OUTPUT_VARIABLE SMI_OUTPUT + ERROR_QUIET + ) + if(SMI_RESULT EQUAL 0 AND SMI_OUTPUT MATCHES "PPU") + set(THEAD_FOUND TRUE) + endif() + endif() + endif() + endif() + + if(THEAD_FOUND) + set(WITH_THEAD ON) + message(STATUS "Auto-detected T-Head environment.") + else() + message(STATUS "T-Head environment not detected.") + endif() + + if(THEAD_FOUND AND NVIDIA_FOUND AND NOT NVIDIA_REQUESTED) + set(NVIDIA_FOUND FALSE) + set(WITH_NVIDIA OFF CACHE BOOL "Enable NVIDIA GPU support" FORCE) + message(STATUS "T-Head environment detected; ignoring NVIDIA auto-detection.") + endif() + + if(NVIDIA_FOUND) + set(WITH_NVIDIA ON) + message(STATUS "Auto-detected NVIDIA environment.") + else() + message(STATUS "NVIDIA environment not detected.") + endif() + endif() # ========================================================= @@ -251,12 +332,21 @@ if(AUTO_DETECT_BACKENDS) endif() # Detect NCCL Dependencies - if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON) + if(WITH_NVIDIA OR WITH_THEAD 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_INCLUDE_SUFFIXES include) + set(_NCCL_LIBRARY_SUFFIXES lib lib64) + + if(WITH_THEAD) + _collect_thead_hints(_THEAD_HINTS) + set(_NCCL_HINTS ${_THEAD_HINTS}) + list(APPEND _NCCL_INCLUDE_SUFFIXES targets/x86_64-linux/include) + list(APPEND _NCCL_LIBRARY_SUFFIXES targets/x86_64-linux/lib) + list(REMOVE_DUPLICATES _NCCL_HINTS) + message(STATUS "T-Head detected. Searching for NCCL in ${_NCCL_HINTS}") + elseif(WITH_ILUVATAR) set(_NCCL_HINTS /usr/local/corex) message(STATUS "Iluvatar detected. Searching for NCCL in ${_NCCL_HINTS}") elseif(WITH_HYGON) @@ -274,11 +364,18 @@ if(AUTO_DETECT_BACKENDS) list(APPEND _NCCL_HINTS /opt/dtk /opt/dtk/rccl) set(_NCCL_HEADER_NAMES rccl/rccl.h rccl.h) set(_NCCL_LIBRARY_NAMES rccl) + list(APPEND _NCCL_LIBRARY_SUFFIXES rccl/lib) message(STATUS "Hygon detected. Searching for NCCL-compatible RCCL in ${_NCCL_HINTS}") endif() - 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) + find_path(AUTO_NCCL_INC NAMES ${_NCCL_HEADER_NAMES} + HINTS ${_NCCL_HINTS} + PATH_SUFFIXES ${_NCCL_INCLUDE_SUFFIXES} + QUIET) + find_library(AUTO_NCCL_LIB NAMES ${_NCCL_LIBRARY_NAMES} + HINTS ${_NCCL_HINTS} + PATH_SUFFIXES ${_NCCL_LIBRARY_SUFFIXES} + QUIET) if(AUTO_NCCL_INC AND AUTO_NCCL_LIB) set(WITH_NCCL ON) @@ -352,6 +449,30 @@ endif() # ========================================================= # --- DEPENDENCY DISCOVERY --- # ========================================================= +if(WITH_THEAD) + _collect_thead_hints(_THEAD_CUDA_HINTS) + list(APPEND _THEAD_CUDA_HINTS /usr/local/cuda) + list(REMOVE_DUPLICATES _THEAD_CUDA_HINTS) + + set(THEAD_CUDA_ROOT "") + foreach(_thead_cuda_hint ${_THEAD_CUDA_HINTS}) + if(NOT THEAD_CUDA_ROOT AND EXISTS "${_thead_cuda_hint}/include/cuda_runtime.h") + set(THEAD_CUDA_ROOT "${_thead_cuda_hint}") + endif() + endforeach() + + if(NOT THEAD_CUDA_ROOT) + message(FATAL_ERROR "`WITH_THEAD` is `ON` but the T-Head CUDA-compatible SDK was not found. Set `PPU_SDK`, `PPU_HOME`, `CUDA_SDK`, `CUDA_HOME`, or `CUDA_PATH`.") + endif() + + set(CUDAToolkit_ROOT "${THEAD_CUDA_ROOT}" CACHE PATH "T-Head CUDA-compatible toolkit root" FORCE) + list(APPEND _NCCL_HINTS ${_THEAD_CUDA_HINTS}) + list(REMOVE_DUPLICATES _NCCL_HINTS) + + enable_language(CUDA) + find_package(CUDAToolkit REQUIRED) +endif() + if(WITH_NVIDIA) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) @@ -504,23 +625,37 @@ 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_THEAD AND NOT WITH_ILUVATAR AND NOT WITH_HYGON) + message(FATAL_ERROR "NCCL backend requires NVIDIA, T-Head, Iluvatar, or Hygon GPU support. Please enable `WITH_NVIDIA`, `WITH_THEAD`, `WITH_ILUVATAR`, or `WITH_HYGON`.") endif() set(_NCCL_HEADER_NAMES nccl.h) set(_NCCL_LIBRARY_NAMES nccl) set(NCCL_COMPILE_DEFINITIONS) - - if(WITH_HYGON) + set(_NCCL_INCLUDE_SUFFIXES include) + set(_NCCL_LIBRARY_SUFFIXES lib lib64) + + if(WITH_THEAD) + set(_NCCL_HINTS ${_THEAD_CUDA_HINTS}) + list(APPEND _NCCL_INCLUDE_SUFFIXES targets/x86_64-linux/include) + list(APPEND _NCCL_LIBRARY_SUFFIXES targets/x86_64-linux/lib) + elseif(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) + list(APPEND _NCCL_LIBRARY_SUFFIXES rccl/lib) endif() - 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) + list(REMOVE_DUPLICATES _NCCL_HINTS) + find_library(NCCL_LIB NAMES ${_NCCL_LIBRARY_NAMES} + HINTS ${_NCCL_HINTS} + PATH_SUFFIXES ${_NCCL_LIBRARY_SUFFIXES} + REQUIRED) + find_path(NCCL_INC NAMES ${_NCCL_HEADER_NAMES} + HINTS ${_NCCL_HINTS} + PATH_SUFFIXES ${_NCCL_INCLUDE_SUFFIXES} + REQUIRED) include_directories(${NCCL_INC}) endif() diff --git a/README.md b/README.md index 0158be0..3ffeb67 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` | @@ -189,7 +190,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 `t-head`. | | `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. | @@ -346,6 +347,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** | Full | Requires T-Head SDK. | @@ -356,7 +358,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 NVIDIA, Iluvatar, or T-Head NCCL, or HYGON RCCL. 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`.| | **CNCL** | Partial | `WITH_CNCL=ON` | Requires Cambricon CNCL. Available only when `WITH_CAMBRICON=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/examples/ccl/all_gather.cc b/examples/ccl/all_gather.cc index c3a41b1..fbf51d1 100644 --- a/examples/ccl/all_gather.cc +++ b/examples/ccl/all_gather.cc @@ -67,6 +67,8 @@ void PrintResult(const std::vector& output, const char* mode, int size, std::cout << "[r" << source << ": " << output[offset] << "] "; } std::cout << std::endl; + + return; } void WorkerThread(ThreadArgs args) { diff --git a/scripts/gen_bridge.py b/scripts/gen_bridge.py index 59d4dd2..fc52e88 100644 --- a/scripts/gen_bridge.py +++ b/scripts/gen_bridge.py @@ -143,7 +143,8 @@ def generate(project_root, output_dir, devices, backends): provider_root = CCL_PROVIDER_BACKENDS.get(bb) if provider_root: for dev in devices: - provider_path = f"{provider_root}/{dev}/api.h" + provider_dir = dev + provider_path = f"{provider_root}/{provider_dir}/api.h" if os.path.exists(os.path.join(src_dir, provider_path)): manifest_lines.append(f'#include "{provider_path}"') diff --git a/scripts/icclrun_logic.py b/scripts/icclrun_logic.py index c0f0d63..065cdf6 100644 --- a/scripts/icclrun_logic.py +++ b/scripts/icclrun_logic.py @@ -75,10 +75,15 @@ def __init__(self, manual_config_path): def _is_local(self, ip): """Check if the IP/hostname refers to the local machine.""" try: - local_ips = {"127.0.0.1", "localhost"} + local_ips = {"127.0.0.1", "localhost", "::1", socket.gethostname()} + if ip in local_ips: + return True # Add hostname-resolved IPs. - local_ips.update(socket.gethostbyname_ex(socket.gethostname())[2]) + try: + local_ips.update(socket.gethostbyname_ex(socket.gethostname())[2]) + except Exception: + pass # Add all interface IPs via `hostname -I`. try: diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c191b68..600629f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -68,6 +68,29 @@ if(WITH_NVIDIA) ) endif() +# T-Head +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 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 8c65d82..d7e365c 100644 --- a/src/backend_device_map.h +++ b/src/backend_device_map.h @@ -17,6 +17,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..71bb3a9 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::kThead, Device::Type::kQy>; // 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..3026e0f --- /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_