feat(mars): add WITH_MARS platform and htcc wrapper

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_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"

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/pybind11_utils.h

@@ -203,6 +203,7 @@ inline std::optional<Device::Type> TryDeviceTypeFromString(
       {"cambricon", Device::Type::kCambricon},
       {"ascend", Device::Type::kAscend},
       {"metax", Device::Type::kMetax},
+      {"mars", Device::Type::kMars},
       {"moore", Device::Type::kMoore},
       {"iluvatar", Device::Type::kIluvatar},
       {"hygon", Device::Type::kHygon},