feat(mars): add native Mars kernels and htcc wrapper

.gitignore

@@ -262,3 +262,8 @@ __marimo__/
 
 # Streamlit
 .streamlit/secrets.toml
+
+# htcc_wrapper sibling symlinks for device compilation
+*.cc.cu
+*.cpp.cu
+*.cxx.cu

CMakeLists.txt

@@ -17,6 +17,7 @@ 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_METAX "Enable MetaX backend" OFF)
+option(WITH_MARS "Enable Mars backend" OFF)
 option(WITH_CAMBRICON "Enable Cambricon backend" OFF)
 option(WITH_MOORE "Enable Moore backend" OFF)
 option(WITH_ASCEND "Enable Ascend backend" OFF)
@@ -114,7 +115,10 @@ if(AUTO_DETECT_DEVICES)
         message(STATUS "Auto-detected Hygon environment.")
     endif()
 
-    if(DEFINED ENV{MACA_PATH})
+    if(DEFINED ENV{HPCC_PATH})
+        set(WITH_MARS ON)
+        message(STATUS "Auto-detected Mars environment from HPCC_PATH")
+    elseif(DEFINED ENV{MACA_PATH})
         set(WITH_METAX ON)
         message(STATUS "Auto-detected MetaX environment from MACA_PATH")
     else()
@@ -156,7 +160,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_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND)
             if(${_gpu_backend})
                 set(_non_nvidia_gpu_detected TRUE)
             endif()
@@ -328,14 +332,14 @@ endif()
 
 # 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_METAX WITH_MARS 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_METAX`, `WITH_MARS`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.")
 endif()
 
 if(WITH_NINETOOTHED AND NOT WITH_NVIDIA)
@@ -472,6 +476,30 @@ if(WITH_METAX)
     find_library(MACA_BLAS_LIB    NAMES mcblas    HINTS "${MACA_PATH}/lib" REQUIRED)
 endif()
 
+if(WITH_MARS)
+    add_compile_definitions(WITH_MARS=1)
+
+    # Normally can be found at: `/opt/hpcc/`.
+    set(HPCC_PATH $ENV{HPCC_PATH})
+    if(NOT HPCC_PATH)
+        set(HPCC_PATH "/opt/hpcc")
+    endif()
+    set(CMAKE_C_COMPILER   ${CMAKE_CURRENT_SOURCE_DIR}/scripts/htcc_wrapper.sh)
+    set(CMAKE_CXX_COMPILER ${CMAKE_CURRENT_SOURCE_DIR}/scripts/htcc_wrapper.sh)
+
+    include_directories("${HPCC_PATH}/include")
+    include_directories("${HPCC_PATH}/tools/cu-bridge/include")
+    link_directories("${HPCC_PATH}/lib")
+    link_directories("${HPCC_PATH}/htgpu_llvm/lib")
+
+    # Libraries: hcruntime / hcdnn / hcblas. HPCC is the toolkit; Mars is the device.
+    find_library(HPCC_RUNTIME_LIB NAMES hcruntime HINTS "${HPCC_PATH}/lib" REQUIRED)
+    find_library(HPCC_DNN_LIB     NAMES hcdnn     HINTS "${HPCC_PATH}/lib" REQUIRED)
+    find_library(HPCC_BLAS_LIB    NAMES hcblas    HINTS "${HPCC_PATH}/lib" REQUIRED)
+    # htcc emits LLVM OpenMP (__kmpc_*), not GNU libgomp.
+    find_library(HPCC_OMP_LIB NAMES omp iomp5 HINTS "${HPCC_PATH}/htgpu_llvm/lib" REQUIRED)
+endif()
+
 if(WITH_MOORE)
     add_compile_definitions(WITH_MOORE=1)
 
@@ -539,11 +567,11 @@ 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_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND)
     add_compile_definitions(WITH_CPU=1)
 endif()
 
-if(WITH_TORCH OR WITH_METAX OR WITH_MOORE)
+if(WITH_TORCH OR WITH_METAX OR WITH_MARS OR WITH_MOORE)
     set(PYBIND11_ENABLE_EXTRAS OFF)
 endif()
 

examples/runtime_api.h

@@ -14,6 +14,9 @@
 #elif WITH_METAX
 #include "native/cuda/metax/ops/gemm/mcblas.h"
 #include "native/cuda/metax/runtime_.h"
+#elif WITH_MARS
+#include "native/cuda/mars/ops/gemm/hcblas.h"
+#include "native/cuda/mars/runtime_.h"
 #elif WITH_CAMBRICON
 #include "native/cambricon/ops/gemm/cnblas.h"
 #include "native/cambricon/runtime_.h"
@@ -39,6 +42,8 @@ using DefaultRuntimeUtils = Runtime<Device::Type::kNvidia>;
 using DefaultRuntimeUtils = Runtime<Device::Type::kIluvatar>;
 #elif WITH_METAX
 using DefaultRuntimeUtils = Runtime<Device::Type::kMetax>;
+#elif WITH_MARS
+using DefaultRuntimeUtils = Runtime<Device::Type::kMars>;
 #elif WITH_CAMBRICON
 using DefaultRuntimeUtils = Runtime<Device::Type::kCambricon>;
 #elif WITH_MOORE

scripts/generate_torch_ops.py

@@ -75,6 +75,7 @@
     "kMoore",
     "kIluvatar",
     "kHygon",
+    "kMars",
 )
 
 _C10_DEVICE_TYPES = (
@@ -83,6 +84,7 @@
     "kMetax",
     "kMoore",
     "kIluvatar",
+    "kMars",
 )
 
 # YAML scalar-type tokens → C++ types. Reference types (e.g. `const Scalar&`)

scripts/generate_wrappers.py

@@ -1721,8 +1721,10 @@ def _device_marker_headers(devices):
         "cambricon": "infini/rt/cambricon/device_.h",
         "ascend": "infini/rt/ascend/device_.h",
         "metax": "infini/rt/metax/device_.h",
+        "mars": "infini/rt/mars/device_.h",
         "moore": "infini/rt/moore/device_.h",
         "iluvatar": "infini/rt/iluvatar/device_.h",
+        "hygon": "infini/rt/hygon/device_.h",
     }
 
     return [paths[device] for device in devices if device in paths]
@@ -2108,7 +2110,7 @@ def _dispatch_gen_batch_size():
         nargs="+",
         default="cpu",
         type=str,
-        help="Devices to use. Please pick from `cpu`, `nvidia`, `cambricon`, `ascend`, `metax`, `moore`, `iluvatar`, and `hygon`. (default: `cpu`)",
+        help="Devices to use. Please pick from `cpu`, `nvidia`, `cambricon`, `ascend`, `metax`, `mars`, `moore`, `iluvatar`, and `hygon`. (default: `cpu`)",
     )
 
     parser.add_argument(

scripts/htcc_wrapper.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# HPCC analog of mxcc_wrapper.sh.
+# htcc only enables device compilation for `.cu`. Rewrite compile inputs.
+# Place the `.cu` symlink next to the source so same-directory `#include`
+# resolution (e.g. bindings `mul.h` vs C-API `generated/include/mul.h`) still
+# works — unlike rewriting into `/tmp`.
+ARGS=()
+skip_next=0
+has_compile=0
+for arg in "$@"; do
+    if [ $skip_next -eq 1 ]; then
+        skip_next=0
+        ARGS+=("$arg")
+        continue
+    fi
+    case "$arg" in
+        -pthread)
+            ;;
+        -B)
+            skip_next=1
+            ;;
+        -B*)
+            ;;
+        -c)
+            has_compile=1
+            ARGS+=("$arg")
+            ;;
+        *)
+            ARGS+=("$arg")
+            ;;
+    esac
+done
+
+if [ "$has_compile" -eq 1 ]; then
+    rewritten=()
+    for arg in "${ARGS[@]}"; do
+        case "$arg" in
+            *.cc|*.cpp|*.cxx)
+                if [ -f "$arg" ]; then
+                    abs=$(readlink -f "$arg")
+                    cu="${abs}.cu"
+                    ln -sfn "$abs" "$cu"
+                    rewritten+=("$cu")
+                else
+                    rewritten+=("$arg")
+                fi
+                ;;
+            *)
+                rewritten+=("$arg")
+                ;;
+        esac
+    done
+    ARGS=("${rewritten[@]}")
+fi
+
+HPCC_PATH="${HPCC_PATH:-/opt/hpcc}"
+exec "${HPCC_PATH}/htgpu_llvm/bin/htcc" "${ARGS[@]}"

src/CMakeLists.txt

@@ -291,6 +291,38 @@ if(WITH_METAX)
     list(APPEND DEVICE_LIST "metax")
 endif()
 
+if(WITH_MARS)
+    set(MARS_PATTERNS
+        "native/cuda/*.cc"
+        "native/cuda/*.cpp"
+        "native/cuda/mars/*.cc"
+        "native/cuda/mars/*.cu"
+    )
+
+    file(GLOB_RECURSE MARS_SOURCES CONFIGURE_DEPENDS ${MARS_PATTERNS})
+
+    set_source_files_properties(${MARS_SOURCES} PROPERTIES LANGUAGE CXX)
+
+    target_compile_definitions(infiniops PRIVATE WITH_MARS=1)
+    target_sources(infiniops PRIVATE ${MARS_SOURCES})
+
+    target_include_directories(infiniops PUBLIC
+        "${HPCC_PATH}/include"
+        "${HPCC_PATH}/tools/cu-bridge/include")
+    target_link_libraries(infiniops PUBLIC
+        ${HPCC_RUNTIME_LIB}
+        ${HPCC_DNN_LIB}
+        ${HPCC_BLAS_LIB}
+        ${HPCC_OMP_LIB}
+    )
+    set_property(TARGET infiniops APPEND PROPERTY
+        INSTALL_RPATH "${HPCC_PATH}/htgpu_llvm/lib")
+    set_property(TARGET infiniops APPEND PROPERTY
+        BUILD_RPATH "${HPCC_PATH}/htgpu_llvm/lib")
+
+    list(APPEND DEVICE_LIST "mars")
+endif()
+
 if(WITH_MOORE)
     set(MOORE_PATTERNS
         "native/cuda/*.cc"
@@ -857,9 +889,10 @@ if(TORCH_SOURCES)
             torch_compile=${INFINI_OPS_TORCH_COMPILE_JOBS})
     endif()
 
-    if(WITH_METAX OR WITH_MOORE)
-        # Vendor compilers (`mxcc`/`mcc`) cannot compile vendor-forked `torch`
-        # headers. Compile `torch` sources with the system C++ compiler instead.
+    if(WITH_METAX OR WITH_MOORE OR WITH_MARS)
+        # Vendor compilers (`mxcc`/`mcc`/`htcc`) cannot compile vendor-forked
+        # `torch` headers. Compile `torch` sources with the system C++ compiler
+        # instead.
         find_program(SYSTEM_CXX NAMES g++ c++)
 
         if(NOT SYSTEM_CXX)
@@ -878,6 +911,11 @@ if(TORCH_SOURCES)
                 "-I${MACA_PATH}/include/mcr"
                 "-I${MACA_PATH}/tools/cu-bridge/include")
         endif()
+        if(WITH_MARS)
+            list(APPEND _torch_vendor_include_flags
+                "-I${HPCC_PATH}/include"
+                "-I${HPCC_PATH}/tools/cu-bridge/include")
+        endif()
         if(WITH_MOORE)
             list(APPEND _torch_vendor_include_flags "-I${MUSA_ROOT}/include")
             execute_process(
@@ -913,6 +951,9 @@ if(TORCH_SOURCES)
         if(WITH_METAX)
             list(APPEND _torch_extra_flags "-DUSE_MACA=1" "-DWITH_METAX=1")
         endif()
+        if(WITH_MARS)
+            list(APPEND _torch_extra_flags "-DUSE_HPCC=1" "-DWITH_MARS=1")
+        endif()
         if(WITH_MOORE)
             list(APPEND _torch_extra_flags "-DWITH_MOORE=1")
         endif()

src/device.h

@@ -15,7 +15,7 @@ using DeviceEnabled = infini::rt::DeviceEnabled<device_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::kMars>;
 
 template <typename>
 struct ActiveDevicesImpl {

src/native/cuda/mars/blas.h

@@ -0,0 +1,51 @@
+#ifndef INFINI_OPS_MARS_BLAS_H_
+#define INFINI_OPS_MARS_BLAS_H_
+
+#include <utility>
+
+// clang-format off
+#include <hcblas/hcblas.h>
+// clang-format on
+
+#include "data_type.h"
+#include "native/cuda/blas.h"
+#include "native/cuda/mars/blas_utils.h"
+#include "native/cuda/mars/runtime_.h"
+
+namespace infini::ops {
+
+template <>
+struct Blas<Device::Type::kMars> : public Runtime<Device::Type::kMars> {
+  using BlasHandle = hcblasHandle_t;
+
+  static constexpr auto BLAS_OP_N = HCBLAS_OP_N;
+
+  static constexpr auto BLAS_OP_T = HCBLAS_OP_T;
+
+  static constexpr auto R_16F = HPCC_R_16F;
+
+  static constexpr auto R_16BF = HPCC_R_16BF;
+
+  static constexpr auto R_32F = HPCC_R_32F;
+
+  static constexpr auto BLAS_COMPUTE_32F = HCBLAS_COMPUTE_32F;
+
+  static constexpr auto BLAS_COMPUTE_32F_FAST_TF32 =
+      HCBLAS_COMPUTE_32F_FAST_TF32;
+
+  static constexpr auto BLAS_GEMM_DEFAULT = HCBLAS_GEMM_DEFAULT;
+
+  static constexpr auto BlasCreate = hcblasCreate;
+
+  static constexpr auto BlasSetStream = hcblasSetStream;
+
+  static constexpr auto BlasDestroy = hcblasDestroy;
+
+  static constexpr auto BlasGemmStridedBatchedEx = [](auto&&... args) {
+    return hcblasGemmStridedBatchedEx(std::forward<decltype(args)>(args)...);
+  };
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/mars/blas_utils.h

@@ -0,0 +1,30 @@
+#ifndef INFINI_OPS_MARS_BLAS_UTILS_H_
+#define INFINI_OPS_MARS_BLAS_UTILS_H_
+
+// clang-format off
+#include <hcblas/hcblas.h>
+// clang-format on
+
+#include "data_type.h"
+#include "native/cuda/blas_utils.h"
+
+namespace infini::ops {
+
+template <>
+struct BlasUtils<Device::Type::kMars> {
+  static auto GetDataType(DataType dtype) {
+    if (dtype == DataType::kFloat16) return HPCC_R_16F;
+    if (dtype == DataType::kBFloat16) return HPCC_R_16BF;
+    return HPCC_R_32F;
+  }
+
+  static auto GetComputeType(DataType dtype) {
+    if (dtype == DataType::kFloat16 || dtype == DataType::kBFloat16)
+      return HCBLAS_COMPUTE_32F;
+    return HCBLAS_COMPUTE_32F_FAST_TF32;
+  }
+};
+
+}  // namespace infini::ops
+
+#endif