feat: add T-Head PPU runtime backend

CMakeLists.txt

@@ -33,6 +33,7 @@ option(WITH_CPU "Enable CPU backend" OFF)
 option(WITH_NVIDIA "Enable CUDA backend" OFF)
 option(WITH_ILUVATAR "Enable Iluvatar GPU backend" OFF)
 option(WITH_HYGON "Enable Hygon GPU backend" OFF)
+option(WITH_THEAD "Enable T-Head PPU backend" OFF)
 option(WITH_METAX "Enable MetaX backend" OFF)
 option(WITH_CAMBRICON "Enable Cambricon backend" OFF)
 option(WITH_MOORE "Enable Moore backend" OFF)
@@ -118,7 +119,7 @@ if(AUTO_DETECT_DEVICES)
 
     if(WITH_NVIDIA)
         set(_non_nvidia_gpu_detected FALSE)
-        foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_METAX WITH_CAMBRICON WITH_MOORE WITH_ASCEND)
+        foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_CAMBRICON WITH_MOORE WITH_ASCEND)
             if(${_gpu_backend})
                 set(_non_nvidia_gpu_detected TRUE)
             endif()
@@ -135,14 +136,14 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src)
 
 # Only one CUDA-like GPU backend can be enabled at a time.
 set(_gpu_backend_count 0)
-foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_METAX WITH_MOORE WITH_ASCEND)
+foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MOORE WITH_ASCEND)
     if(${_gpu_backend})
         math(EXPR _gpu_backend_count "${_gpu_backend_count} + 1")
     endif()
 endforeach()
 
 if(_gpu_backend_count GREATER 1)
-    message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_METAX`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.")
+    message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_THEAD`, `WITH_METAX`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.")
 endif()
 
 if(WITH_NVIDIA)
@@ -217,6 +218,44 @@ if(WITH_HYGON)
     find_package(CUDAToolkit REQUIRED)
 endif()
 
+if(WITH_THEAD)
+    add_compile_definitions(WITH_THEAD=1)
+
+    set(THEAD_CUDA_ROOT "")
+    foreach(_thead_cuda_env THEAD_CUDA_ROOT PPU_CUDA_ROOT)
+        if(NOT THEAD_CUDA_ROOT AND DEFINED ENV{${_thead_cuda_env}} AND
+           NOT "$ENV{${_thead_cuda_env}}" STREQUAL "")
+            set(THEAD_CUDA_ROOT "$ENV{${_thead_cuda_env}}")
+        endif()
+    endforeach()
+    if(NOT THEAD_CUDA_ROOT AND DEFINED ENV{PPU_SDK_ROOT} AND
+       NOT "$ENV{PPU_SDK_ROOT}" STREQUAL "")
+        set(THEAD_CUDA_ROOT "$ENV{PPU_SDK_ROOT}/CUDA_SDK")
+    endif()
+    if(NOT THEAD_CUDA_ROOT AND EXISTS "/usr/local/PPU_SDK/CUDA_SDK")
+        set(THEAD_CUDA_ROOT "/usr/local/PPU_SDK/CUDA_SDK")
+    endif()
+
+    if(NOT THEAD_CUDA_ROOT OR NOT EXISTS "${THEAD_CUDA_ROOT}/bin/nvcc")
+        message(FATAL_ERROR
+            "`WITH_THEAD` is `ON` but the Thead PPU CUDA SDK was not found. "
+            "Set `THEAD_CUDA_ROOT`, `PPU_CUDA_ROOT`, or `PPU_SDK_ROOT`.")
+    endif()
+
+    set(CMAKE_CUDA_COMPILER "${THEAD_CUDA_ROOT}/bin/nvcc"
+        CACHE FILEPATH "Thead PPU CUDA compiler" FORCE)
+    set(CUDAToolkit_ROOT "${THEAD_CUDA_ROOT}"
+        CACHE PATH "Thead PPU CUDA toolkit root" FORCE)
+    if(NOT CMAKE_CUDA_ARCHITECTURES)
+        set(CMAKE_CUDA_ARCHITECTURES 80 CACHE STRING
+            "Thead PPU CUDA architecture" FORCE)
+    endif()
+
+    message(STATUS "Thead PPU: CUDA toolkit ${THEAD_CUDA_ROOT}")
+    enable_language(CUDA)
+    find_package(CUDAToolkit REQUIRED)
+endif()
+
 if(WITH_METAX)
     add_compile_definitions(WITH_METAX=1)
 
@@ -294,7 +333,7 @@ if(WITH_ASCEND)
 endif()
 
 # If all other platforms are not enabled, CPU is enabled by default.
-if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_METAX AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND)
+if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND)
     set(WITH_CPU ON)
     add_compile_definitions(WITH_CPU=1)
 endif()
@@ -312,6 +351,9 @@ endif()
 if(WITH_HYGON)
     list(APPEND INFINI_RT_PUBLIC_HEADER_DEVICES hygon)
 endif()
+if(WITH_THEAD)
+    list(APPEND INFINI_RT_PUBLIC_HEADER_DEVICES thead)
+endif()
 if(WITH_METAX)
     list(APPEND INFINI_RT_PUBLIC_HEADER_DEVICES metax)
 endif()
@@ -342,7 +384,7 @@ endif()
 add_subdirectory(src)
 
 set(INFINI_RT_PACKAGE_NEEDS_CUDATOOLKIT OFF)
-if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON)
+if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON OR WITH_THEAD)
     set(INFINI_RT_PACKAGE_NEEDS_CUDATOOLKIT ON)
 endif()
 

cmake/InfiniRTConfig.cmake.in

@@ -9,6 +9,7 @@ set(InfiniRT_WITH_CPU @WITH_CPU@)
 set(InfiniRT_WITH_NVIDIA @WITH_NVIDIA@)
 set(InfiniRT_WITH_ILUVATAR @WITH_ILUVATAR@)
 set(InfiniRT_WITH_HYGON @WITH_HYGON@)
+set(InfiniRT_WITH_THEAD @WITH_THEAD@)
 set(InfiniRT_WITH_METAX @WITH_METAX@)
 set(InfiniRT_WITH_MOORE @WITH_MOORE@)
 set(InfiniRT_WITH_CAMBRICON @WITH_CAMBRICON@)

scripts/compare_performance_results.py

@@ -8,7 +8,7 @@ def _load_results(path):
     if isinstance(data, dict) and "results" in data:
         data = data["results"]
     if not isinstance(data, list):
-        raise ValueError(f"{path} must contain a JSON array or an object with results")
+        raise TypeError(f"{path} must contain a JSON array or an object with results")
     return data
 
 

scripts/generate_public_headers.py

@@ -3,7 +3,6 @@
 import pathlib
 import re
 
-
 _DETAIL_PREFIX = "infini/rt/detail"
 
 _DEVICE_HEADERS = {
@@ -27,6 +26,11 @@
         ("hygon", "device_.h", "native/cuda/hygon/device_.h"),
         ("hygon", "runtime_.h", "native/cuda/hygon/runtime_.h"),
     ),
+    "thead": (
+        ("thead", "data_type_.h", "native/cuda/thead/data_type_.h"),
+        ("thead", "device_.h", "native/cuda/thead/device_.h"),
+        ("thead", "runtime_.h", "native/cuda/thead/runtime_.h"),
+    ),
     "metax": (
         ("metax", "data_type_.h", "native/cuda/metax/data_type_.h"),
         ("metax", "device_.h", "native/cuda/metax/device_.h"),
@@ -54,6 +58,7 @@
     "nvidia": "Device::Type::kNvidia",
     "iluvatar": "Device::Type::kIluvatar",
     "hygon": "Device::Type::kHygon",
+    "thead": "Device::Type::kThead",
     "metax": "Device::Type::kMetax",
     "moore": "Device::Type::kMoore",
     "cambricon": "Device::Type::kCambricon",
@@ -64,6 +69,7 @@
     "nvidia",
     "iluvatar",
     "hygon",
+    "thead",
     "metax",
     "moore",
     "cambricon",

scripts/run_performance_tests.py

@@ -6,7 +6,6 @@
 import subprocess
 import sys
 
-
 _BACKEND_OPTIONS = (
     ("WITH_NVIDIA", "nvidia"),
     ("WITH_ILUVATAR", "iluvatar"),
@@ -39,7 +38,7 @@ def _read_cmake_cache(build_dir):
         return values
 
     for line in cache_path.read_text(encoding="utf-8", errors="replace").splitlines():
-        if line.startswith("//") or line.startswith("#") or "=" not in line:
+        if line.startswith(("//", "#")) or "=" not in line:
             continue
         key_type, value = line.split("=", 1)
         key = key_type.split(":", 1)[0]
@@ -105,8 +104,7 @@ def _run_executable(path):
         [str(path)],
         cwd=path.parent,
         text=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
+        capture_output=True,
         check=False,
     )
     if completed.stderr:

src/CMakeLists.txt

@@ -63,6 +63,20 @@ if(WITH_HYGON)
     )
 endif()
 
+if(WITH_THEAD)
+    enable_language(CUDA)
+
+    target_compile_definitions(infinirt PUBLIC WITH_THEAD=1)
+
+    find_package(CUDAToolkit REQUIRED)
+    target_link_libraries(infinirt PUBLIC CUDA::cudart)
+
+    set_target_properties(infinirt PROPERTIES
+        CUDA_STANDARD 17
+        CUDA_STANDARD_REQUIRED ON
+    )
+endif()
+
 if(WITH_METAX)
     target_compile_definitions(infinirt PRIVATE WITH_METAX=1)
     _infinirt_use_backend_runtime(

src/device.h

@@ -21,6 +21,7 @@ class Device {
     kMoore = 5,
     kIluvatar = 6,
     kHygon = 7,
+    kThead = 8,
     kCount
   };
 
@@ -64,6 +65,7 @@ class Device {
           {Type::kMoore, "moore"},
           {Type::kIluvatar, "iluvatar"},
           {Type::kHygon, "hygon"},
+          {Type::kThead, "thead"},
       }}};
 
   static constexpr ConstexprMap<std::string_view, Device::Type,
@@ -77,6 +79,7 @@ class Device {
           {"moore", Type::kMoore},
           {"iluvatar", Type::kIluvatar},
           {"hygon", Type::kHygon},
+          {"thead", Type::kThead},
       }}};
 
   int index_{0};
@@ -91,7 +94,7 @@ struct DeviceEnabled : std::false_type {};
 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::kHygon>;
+         Device::Type::kIluvatar, Device::Type::kHygon, Device::Type::kThead>;
 
 // Deferred computation of active devices. The `Filter` and `FilterList`
 // evaluation are nested inside a class template so that `DeviceEnabled`

src/native/cuda/thead/data_type_.h

@@ -0,0 +1,30 @@
+#ifndef INFINI_RT_THEAD_DATA_TYPE__H_
+#define INFINI_RT_THEAD_DATA_TYPE__H_
+
+// clang-format off
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+// clang-format on
+
+#include "data_type.h"
+#include "native/cuda/thead/device_.h"
+
+namespace infini::rt {
+
+using cuda_bfloat16 = nv_bfloat16;
+
+using cuda_bfloat162 = nv_bfloat162;
+
+template <>
+struct TypeMap<Device::Type::kThead, DataType::kFloat16> {
+  using type = half;
+};
+
+template <>
+struct TypeMap<Device::Type::kThead, DataType::kBFloat16> {
+  using type = __nv_bfloat16;
+};
+
+}  // namespace infini::rt
+
+#endif

src/native/cuda/thead/device_.h

@@ -0,0 +1,13 @@
+#ifndef INFINI_RT_THEAD_DEVICE__H_
+#define INFINI_RT_THEAD_DEVICE__H_
+
+#include "device.h"
+
+namespace infini::rt {
+
+template <>
+struct DeviceEnabled<Device::Type::kThead> : std::true_type {};
+
+}  // namespace infini::rt
+
+#endif