From 110f6ed7620535071d3d8298cab5011b6bc3031b Mon Sep 17 00:00:00 2001 From: Ceng23333 <441651826@qq.com> Date: Mon, 24 Aug 2026 10:03:52 +0800 Subject: [PATCH] fix(mars): enable torch bindings and CONTRIBUTING smoke Compile Mars torch sources with system g++ (USE_HPCC), map TorchDeviceName to cuda, and instantiate Mars torch backends. Co-authored-by: Cursor --- .gitignore | 5 ++++ scripts/generate_torch_ops.py | 2 ++ src/CMakeLists.txt | 15 ++++++++-- src/torch/device_.h | 5 ++++ src/torch/mars/c10.h | 28 +++++++++++++++++++ src/torch/ops/add/add.cc | 1 + src/torch/ops/gemm/gemm.cc | 1 + .../reshape_and_cache/reshape_and_cache.cc | 1 + .../ops/rotary_embedding/rotary_embedding.cc | 1 + .../scaled_dot_product_attention.cc | 1 + 10 files changed, 57 insertions(+), 3 deletions(-) create mode 100644 src/torch/mars/c10.h diff --git a/.gitignore b/.gitignore index c6501fb2d..6050e9104 100644 --- a/.gitignore +++ b/.gitignore @@ -262,3 +262,8 @@ __marimo__/ # Streamlit .streamlit/secrets.toml + +# htcc_wrapper sibling symlinks for device compilation +*.cc.cu +*.cpp.cu +*.cxx.cu diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index 482e0f5b3..5a6788cf6 100644 --- a/scripts/generate_torch_ops.py +++ b/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&`) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8f33d1dee..c39ad6bef 100644 --- a/src/CMakeLists.txt +++ b/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() diff --git a/src/torch/device_.h b/src/torch/device_.h index dfd3cae6b..3cec42dc1 100644 --- a/src/torch/device_.h +++ b/src/torch/device_.h @@ -50,6 +50,11 @@ struct TorchDeviceName { static constexpr std::string_view kValue{"cuda"}; }; +template <> +struct TorchDeviceName { + static constexpr std::string_view kValue{"cuda"}; +}; + } // namespace infini::ops::detail #endif diff --git a/src/torch/mars/c10.h b/src/torch/mars/c10.h new file mode 100644 index 000000000..f0ead156d --- /dev/null +++ b/src/torch/mars/c10.h @@ -0,0 +1,28 @@ +#ifndef INFINI_OPS_TORCH_MARS_C10_H_ +#define INFINI_OPS_TORCH_MARS_C10_H_ + +#include +#include +#include + +#include "torch/c10.h" + +namespace infini::ops { + +template <> +struct C10 { + 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(stream), + static_cast(device_index)); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_TORCH_MARS_C10_H_ diff --git a/src/torch/ops/add/add.cc b/src/torch/ops/add/add.cc index cc36f5bc0..07134d5f2 100644 --- a/src/torch/ops/add/add.cc +++ b/src/torch/ops/add/add.cc @@ -40,5 +40,6 @@ template class Operator; template class Operator; template class Operator; template class Operator; +template class Operator; } // namespace infini::ops diff --git a/src/torch/ops/gemm/gemm.cc b/src/torch/ops/gemm/gemm.cc index 01a4a2a8f..499dbedcf 100644 --- a/src/torch/ops/gemm/gemm.cc +++ b/src/torch/ops/gemm/gemm.cc @@ -84,5 +84,6 @@ template class Operator; template class Operator; template class Operator; template class Operator; +template class Operator; } // namespace infini::ops diff --git a/src/torch/ops/reshape_and_cache/reshape_and_cache.cc b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc index 99ca3c82a..60e76d7a7 100644 --- a/src/torch/ops/reshape_and_cache/reshape_and_cache.cc +++ b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc @@ -85,5 +85,6 @@ template class Operator; template class Operator; template class Operator; template class Operator; +template class Operator; } // namespace infini::ops diff --git a/src/torch/ops/rotary_embedding/rotary_embedding.cc b/src/torch/ops/rotary_embedding/rotary_embedding.cc index 06658ea3a..a74920438 100644 --- a/src/torch/ops/rotary_embedding/rotary_embedding.cc +++ b/src/torch/ops/rotary_embedding/rotary_embedding.cc @@ -83,5 +83,6 @@ template class Operator; template class Operator; template class Operator; template class Operator; +template class Operator; } // namespace infini::ops diff --git a/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc index 81489cf8c..db5cf10f4 100644 --- a/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc +++ b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc @@ -78,5 +78,6 @@ template class Operator; template class Operator; template class Operator; template class Operator; +template class Operator; } // namespace infini::ops