fix(mars): enable torch bindings and CONTRIBUTING smoke

.gitignore

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

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&`)

src/CMakeLists.txt

@@ -889,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)
@@ -910,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(
@@ -945,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/torch/device_.h

@@ -50,6 +50,11 @@ struct TorchDeviceName<Device::Type::kHygon> {
   static constexpr std::string_view kValue{"cuda"};
 };
 
+template <>
+struct TorchDeviceName<Device::Type::kMars> {
+  static constexpr std::string_view kValue{"cuda"};
+};
+
 }  // namespace infini::ops::detail
 
 #endif

src/torch/mars/c10.h

@@ -0,0 +1,28 @@
+#ifndef INFINI_OPS_TORCH_MARS_C10_H_
+#define INFINI_OPS_TORCH_MARS_C10_H_
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime_api.h>
+
+#include "torch/c10.h"
+
+namespace infini::ops {
+
+template <>
+struct C10<Device::Type::kMars> {
+  static constexpr Device::Type kDeviceType = Device::Type::kMars;
+
+  using StreamGuard = c10::cuda::CUDAStreamGuard;
+
+  static c10::cuda::CUDAStream GetStreamFromExternal(void* stream,
+                                                     int device_index) {
+    return c10::cuda::getStreamFromExternal(
+        reinterpret_cast<cudaStream_t>(stream),
+        static_cast<c10::DeviceIndex>(device_index));
+  }
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_TORCH_MARS_C10_H_

src/torch/ops/add/add.cc

@@ -40,5 +40,6 @@ template class Operator<Add, Device::Type::kMetax, 1>;
 template class Operator<Add, Device::Type::kMoore, 1>;
 template class Operator<Add, Device::Type::kIluvatar, 1>;
 template class Operator<Add, Device::Type::kHygon, 1>;
+template class Operator<Add, Device::Type::kMars, 1>;
 
 }  // namespace infini::ops

src/torch/ops/gemm/gemm.cc

@@ -84,5 +84,6 @@ template class Operator<Gemm, Device::Type::kMetax, 2>;
 template class Operator<Gemm, Device::Type::kMoore, 2>;
 template class Operator<Gemm, Device::Type::kIluvatar, 2>;
 template class Operator<Gemm, Device::Type::kHygon, 2>;
+template class Operator<Gemm, Device::Type::kMars, 2>;
 
 }  // namespace infini::ops

src/torch/ops/reshape_and_cache/reshape_and_cache.cc

@@ -85,5 +85,6 @@ template class Operator<ReshapeAndCache, Device::Type::kMetax, 1>;
 template class Operator<ReshapeAndCache, Device::Type::kMoore, 1>;
 template class Operator<ReshapeAndCache, Device::Type::kIluvatar, 1>;
 template class Operator<ReshapeAndCache, Device::Type::kHygon, 1>;
+template class Operator<ReshapeAndCache, Device::Type::kMars, 1>;
 
 }  // namespace infini::ops

src/torch/ops/rotary_embedding/rotary_embedding.cc

@@ -83,5 +83,6 @@ template class Operator<RotaryEmbedding, Device::Type::kMetax, 1>;
 template class Operator<RotaryEmbedding, Device::Type::kMoore, 1>;
 template class Operator<RotaryEmbedding, Device::Type::kIluvatar, 1>;
 template class Operator<RotaryEmbedding, Device::Type::kHygon, 1>;
+template class Operator<RotaryEmbedding, Device::Type::kMars, 1>;
 
 }  // namespace infini::ops

src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc

@@ -78,5 +78,6 @@ template class Operator<ScaledDotProductAttention, Device::Type::kMetax, 1>;
 template class Operator<ScaledDotProductAttention, Device::Type::kMoore, 1>;
 template class Operator<ScaledDotProductAttention, Device::Type::kIluvatar, 1>;
 template class Operator<ScaledDotProductAttention, Device::Type::kHygon, 1>;
+template class Operator<ScaledDotProductAttention, Device::Type::kMars, 1>;
 
 }  // namespace infini::ops