From 47a85a2478ad97966a23c62c14f08351db47bfeb Mon Sep 17 00:00:00 2001 From: gongchensu Date: Tue, 14 Jul 2026 09:08:31 +0000 Subject: [PATCH 1/4] feat(metax): enable standalone InfiniRT graph bridge --- src/infinicore/context/context_impl.cc | 1 + xmake.lua | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index 29a470679..d6991ebdd 100644 --- a/src/infinicore/context/context_impl.cc +++ b/src/infinicore/context/context_impl.cc @@ -92,6 +92,7 @@ ContextImpl::ContextImpl() { initializeRuntime(); initializeRuntime(); initializeRuntime(); + initializeRuntime(); initializeRuntime(); if (current_runtime_ == nullptr && !runtime_table_[static_cast(Device::Type::CPU)].empty()) { diff --git a/xmake.lua b/xmake.lua index 9ec14e9ca..ffc44ed58 100644 --- a/xmake.lua +++ b/xmake.lua @@ -365,6 +365,13 @@ local function add_external_infinirt() end end end + if has_config("metax-gpu") then + local maca_root = os.getenv("MACA_PATH") or os.getenv("MACA_HOME") or os.getenv("MACA_ROOT") or "/opt/maca" + local maca_include = path.join(maca_root, "include") + if os.isdir(maca_include) then + add_includedirs(maca_include, { public = true }) + end + end add_links("infinirt") end From cff33e171c6595542a28f6e0775783e8cdd58b4d Mon Sep 17 00:00:00 2001 From: gongchensu Date: Tue, 21 Jul 2026 16:05:26 +0800 Subject: [PATCH 2/4] feat(mars): add separate HPCC device backend --- include/infinicore.h | 1 + include/infinicore/adaptor/aten_adaptor.hpp | 2 +- include/infinicore/device.hpp | 1 + python/infinicore/device.py | 1 + src/bridge/infini/rt.hpp | 4 + src/infiniccl-test/main.cpp | 3 +- src/infiniccl/infiniccl.cc | 10 ++ src/infiniccl/mars/infiniccl_mars.cc | 93 +++++++++++++ src/infiniccl/mars/infiniccl_mars.h | 12 ++ src/infinicore-test/main.cc | 3 + src/infinicore/context/context_impl.cc | 1 + src/infinicore/device.cc | 2 + src/infinicore/nn/rmsnorm.cc | 1 + src/infinicore/ops/infiniops_impl.hpp | 4 + src/infinicore/ops/rotg/rotg_infiniop.cc | 3 +- src/infinicore/ops/rotm/rotm_infiniop.cc | 3 +- src/infinicore/ops/rotmg/rotmg_infiniop.cc | 3 +- src/infinicore/pybind11/device.hpp | 1 + src/infiniop/devices/handle.cc | 9 ++ src/infiniop/devices/mars/mars_handle.h | 20 +++ src/infiniop/ops/acos/operator.cc | 14 +- .../ops/adaptive_avg_pool1d/operator.cc | 14 +- .../ops/adaptive_avg_pool3d/operator.cc | 14 +- .../ops/adaptive_max_pool1d/operator.cc | 14 +- src/infiniop/ops/add/operator.cc | 14 +- src/infiniop/ops/add_rms_norm/operator.cc | 14 +- src/infiniop/ops/addbmm/operator.cc | 14 +- src/infiniop/ops/addcmul/operator.cc | 14 +- src/infiniop/ops/addr/operator.cc | 14 +- src/infiniop/ops/affine_grid/operator.cc | 14 +- src/infiniop/ops/all/operator.cc | 14 +- src/infiniop/ops/argwhere/operator.cc | 14 +- src/infiniop/ops/asin/operator.cc | 14 +- src/infiniop/ops/asinh/operator.cc | 14 +- src/infiniop/ops/asum/operator.cc | 14 +- src/infiniop/ops/atanh/operator.cc | 14 +- src/infiniop/ops/avg_pool1d/operator.cc | 14 +- src/infiniop/ops/avg_pool3d/operator.cc | 14 +- src/infiniop/ops/axpy/operator.cc | 14 +- .../operator.cc | 14 +- .../ops/bitwise_right_shift/operator.cc | 14 +- src/infiniop/ops/blas_amax/operator.cc | 14 +- src/infiniop/ops/blas_amin/operator.cc | 14 +- src/infiniop/ops/blas_copy/operator.cc | 14 +- src/infiniop/ops/blas_dot/operator.cc | 14 +- src/infiniop/ops/block_diag/operator.cc | 14 +- src/infiniop/ops/broadcast_to/operator.cc | 14 +- src/infiniop/ops/causal_conv1d/operator.cc | 14 +- src/infiniop/ops/causal_softmax/operator.cc | 14 +- src/infiniop/ops/cdist/operator.cc | 14 +- .../ops/chunk_gated_delta_rule/operator.cc | 14 +- src/infiniop/ops/clip/operator.cc | 14 +- src/infiniop/ops/cross_entropy/operator.cc | 14 +- src/infiniop/ops/deepseek_moe/operator.cc | 17 ++- src/infiniop/ops/diff/operator.cc | 14 +- src/infiniop/ops/digamma/operator.cc | 14 +- src/infiniop/ops/dist/operator.cc | 14 +- src/infiniop/ops/dot/operator.cc | 14 +- src/infiniop/ops/embedding/operator.cc | 11 +- src/infiniop/ops/equal/operator.cc | 14 +- src/infiniop/ops/erf/operator.cc | 14 +- src/infiniop/ops/erfc/operator.cc | 14 +- src/infiniop/ops/erfinv/operator.cc | 14 +- src/infiniop/ops/flipud/operator.cc | 14 +- src/infiniop/ops/float_power/operator.cc | 14 +- src/infiniop/ops/floor/operator.cc | 14 +- src/infiniop/ops/floor_divide/operator.cc | 14 +- src/infiniop/ops/fmin/operator.cc | 14 +- src/infiniop/ops/fmod/operator.cc | 14 +- src/infiniop/ops/fused_ffn/operator.cc | 14 +- .../fused_gated_delta_net_gating/operator.cc | 14 +- .../ops/gaussian_nll_loss/operator.cc | 14 +- src/infiniop/ops/gelu/operator.cc | 14 +- src/infiniop/ops/gelutanh/operator.cc | 14 +- src/infiniop/ops/gemm/operator.cc | 14 +- src/infiniop/ops/hardswish/operator.cc | 14 +- src/infiniop/ops/hardtanh/operator.cc | 14 +- .../ops/hinge_embedding_loss/operator.cc | 14 +- src/infiniop/ops/histc/operator.cc | 14 +- src/infiniop/ops/huber_loss/operator.cc | 14 +- src/infiniop/ops/hypot/operator.cc | 14 +- src/infiniop/ops/index_add/operator.cc | 14 +- src/infiniop/ops/index_copy/operator.cc | 14 +- src/infiniop/ops/inner/operator.cc | 14 +- src/infiniop/ops/interpolate/operator.cc | 14 +- src/infiniop/ops/kron/operator.cc | 14 +- src/infiniop/ops/kthvalue/operator.cc | 14 +- src/infiniop/ops/kv_caching/operator.cc | 14 +- src/infiniop/ops/layer_norm/operator.cc | 14 +- src/infiniop/ops/ldexp/operator.cc | 14 +- src/infiniop/ops/lerp/operator.cc | 14 +- src/infiniop/ops/log10/operator.cc | 14 +- src/infiniop/ops/log1p/operator.cc | 14 +- src/infiniop/ops/log_softmax/operator.cc | 14 +- src/infiniop/ops/logaddexp/operator.cc | 14 +- src/infiniop/ops/logaddexp2/operator.cc | 14 +- src/infiniop/ops/logcumsumexp/operator.cc | 14 +- src/infiniop/ops/logdet/operator.cc | 14 +- src/infiniop/ops/masked_select/operator.cc | 14 +- src/infiniop/ops/matrix_power/operator.cc | 14 +- src/infiniop/ops/mul/operator.cc | 14 +- .../ops/multi_margin_loss/operator.cc | 14 +- src/infiniop/ops/nrm2/operator.cc | 14 +- src/infiniop/ops/ones/operator.cc | 14 +- src/infiniop/ops/pad/operator.cc | 14 +- src/infiniop/ops/paged_attention/operator.cc | 14 +- .../ops/paged_attention_prefill/operator.cc | 14 +- src/infiniop/ops/paged_caching/operator.cc | 14 +- src/infiniop/ops/pixel_shuffle/operator.cc | 14 +- src/infiniop/ops/prelu/operator.cc | 14 +- src/infiniop/ops/random_sample/operator.cc | 14 +- src/infiniop/ops/rearrange/operator.cc | 11 +- src/infiniop/ops/reciprocal/operator.cc | 14 +- .../recurrent_gated_delta_rule/operator.cc | 14 +- src/infiniop/ops/relu6/operator.cc | 14 +- src/infiniop/ops/rms_norm/operator.cc | 14 +- src/infiniop/ops/rope/operator.cc | 14 +- src/infiniop/ops/rot/operator.cc | 14 +- src/infiniop/ops/rotg/operator.cc | 14 +- src/infiniop/ops/rotm/operator.cc | 14 +- src/infiniop/ops/rotmg/operator.cc | 14 +- src/infiniop/ops/scal/operator.cc | 14 +- src/infiniop/ops/scatter/operator.cc | 14 +- src/infiniop/ops/selu/operator.cc | 14 +- src/infiniop/ops/silu/operator.cc | 14 +- src/infiniop/ops/sinh/operator.cc | 14 +- src/infiniop/ops/smooth_l1_loss/operator.cc | 14 +- src/infiniop/ops/softplus/operator.cc | 14 +- src/infiniop/ops/softsign/operator.cc | 14 +- src/infiniop/ops/sub/operator.cc | 14 +- src/infiniop/ops/sum/operator.cc | 14 +- src/infiniop/ops/swap/operator.cc | 14 +- src/infiniop/ops/swiglu/operator.cc | 14 +- src/infiniop/ops/take/operator.cc | 14 +- src/infiniop/ops/tan/operator.cc | 14 +- src/infiniop/ops/tanhshrink/operator.cc | 14 +- src/infiniop/ops/topk/operator.cc | 14 +- src/infiniop/ops/topkrouter/operator.cc | 14 +- src/infiniop/ops/topksoftmax/operator.cc | 14 +- .../ops/triplet_margin_loss/operator.cc | 14 +- .../operator.cc | 14 +- src/infiniop/ops/unfold/operator.cc | 14 +- .../ops/upsample_bilinear/operator.cc | 14 +- src/infiniop/ops/upsample_nearest/operator.cc | 14 +- src/infiniop/ops/vander/operator.cc | 14 +- src/infiniop/ops/var/operator.cc | 14 +- src/infiniop/ops/var_mean/operator.cc | 14 +- src/infiniop/ops/zeros/operator.cc | 14 +- test/infinicore/framework/config.py | 9 ++ test/infinicore/framework/devices.py | 3 + test/infiniop/libinfiniop/devices.py | 3 + test/infiniop/libinfiniop/utils.py | 10 ++ xmake.lua | 58 +++++--- xmake/mars.lua | 127 ++++++++++++++++++ xmake/metax.lua | 22 +-- 155 files changed, 2025 insertions(+), 173 deletions(-) create mode 100644 src/infiniccl/mars/infiniccl_mars.cc create mode 100644 src/infiniccl/mars/infiniccl_mars.h create mode 100644 src/infiniop/devices/mars/mars_handle.h create mode 100644 xmake/mars.lua diff --git a/include/infinicore.h b/include/infinicore.h index 849c62c14..d24344b8b 100644 --- a/include/infinicore.h +++ b/include/infinicore.h @@ -48,6 +48,7 @@ typedef enum { INFINI_DEVICE_HYGON = 8, INFINI_DEVICE_QY = 9, INFINI_DEVICE_ALI = 10, + INFINI_DEVICE_MARS = 11, INFINI_DEVICE_TYPE_COUNT } infiniDevice_t; diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index 74b053dc5..5b8b5b58f 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -41,7 +41,7 @@ inline at::ScalarType to_at_dtype(DataType dtype) { inline at::Device to_at_device(const Device &device) { // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). // Treat MetaX/QY devices as CUDA devices for ATen tensor interoperability. - if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) { + if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::MARS || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) { return at::Device(at::kCUDA, device.getIndex()); } else if (device.getType() == Device::Type::CPU) { return at::Device(at::kCPU); diff --git a/include/infinicore/device.hpp b/include/infinicore/device.hpp index 4bdab1088..4ef245928 100644 --- a/include/infinicore/device.hpp +++ b/include/infinicore/device.hpp @@ -23,6 +23,7 @@ class Device { HYGON = INFINI_DEVICE_HYGON, QY = INFINI_DEVICE_QY, ALI = INFINI_DEVICE_ALI, + MARS = INFINI_DEVICE_MARS, COUNT = INFINI_DEVICE_TYPE_COUNT, }; diff --git a/python/infinicore/device.py b/python/infinicore/device.py index 8858a1235..e88e85cbb 100644 --- a/python/infinicore/device.py +++ b/python/infinicore/device.py @@ -83,6 +83,7 @@ def _from_infinicore_device(infinicore_device: _infinicore.Device): _infinicore.Device.Type.HYGON: "cuda", _infinicore.Device.Type.QY: "cuda", _infinicore.Device.Type.ALI: "cuda", + _infinicore.Device.Type.MARS: "cuda", } diff --git a/src/bridge/infini/rt.hpp b/src/bridge/infini/rt.hpp index c5bec7ecd..6370b5073 100644 --- a/src/bridge/infini/rt.hpp +++ b/src/bridge/infini/rt.hpp @@ -33,6 +33,8 @@ inline constexpr ::infini::rt::Device::Type translate_to(infiniDevice_t device) return ::infini::rt::Device::Type::kIluvatar; case INFINI_DEVICE_HYGON: return ::infini::rt::Device::Type::kHygon; + case INFINI_DEVICE_MARS: + return ::infini::rt::Device::Type::kMars; default: return ::infini::rt::Device::Type::kCount; } @@ -56,6 +58,8 @@ inline constexpr infiniDevice_t translate_from(::infini::rt::Device::Type device return INFINI_DEVICE_ILUVATAR; case ::infini::rt::Device::Type::kHygon: return INFINI_DEVICE_HYGON; + case ::infini::rt::Device::Type::kMars: + return INFINI_DEVICE_MARS; default: return INFINI_DEVICE_TYPE_COUNT; } diff --git a/src/infiniccl-test/main.cpp b/src/infiniccl-test/main.cpp index 8126aa15a..ae042c945 100644 --- a/src/infiniccl-test/main.cpp +++ b/src/infiniccl-test/main.cpp @@ -12,7 +12,7 @@ void printUsage() { std::cout << "infiniccl-test --" << std::endl << std::endl; std::cout << " --" << std::endl; - std::cout << " Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl + std::cout << " Specify the device type --(nvidia|cambricon|ascend|metax|mars|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl << std::endl; std::cout << "The program will run tests on all visible devices of the specified device type." << " Use Environmental Variables such as CUDA_VSIBLE_DEVICES to limit visible device IDs."; @@ -41,6 +41,7 @@ ParsedArgs parseArgs(int argc, char *argv[]) { else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON) else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND) else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX) + else PARSE_DEVICE("--mars", INFINI_DEVICE_MARS) else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE) else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR) else PARSE_DEVICE("--qy", INFINI_DEVICE_QY) diff --git a/src/infiniccl/infiniccl.cc b/src/infiniccl/infiniccl.cc index 48dbea425..a5c54a242 100644 --- a/src/infiniccl/infiniccl.cc +++ b/src/infiniccl/infiniccl.cc @@ -4,6 +4,7 @@ #include "./cambricon/infiniccl_cambricon.h" #include "./cuda/infiniccl_cuda.h" #include "./kunlun/infiniccl_kunlun.h" +#include "./mars/infiniccl_mars.h" #include "./metax/infiniccl_metax.h" #include "./moore/infiniccl_moore.h" @@ -25,6 +26,7 @@ __INFINI_C infiniStatus_t infinicclCommInitAll( COMM_INIT_ALL(INFINI_DEVICE_ASCEND, ascend); COMM_INIT_ALL(INFINI_DEVICE_CAMBRICON, cambricon); COMM_INIT_ALL(INFINI_DEVICE_METAX, metax); + COMM_INIT_ALL(INFINI_DEVICE_MARS, mars); COMM_INIT_ALL(INFINI_DEVICE_MOORE, moore); COMM_INIT_ALL(INFINI_DEVICE_KUNLUN, kunlun); COMM_INIT_ALL(INFINI_DEVICE_ALI, cuda); @@ -52,6 +54,7 @@ __INFINI_C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) { COMM_DESTROY(INFINI_DEVICE_ASCEND, ascend); COMM_DESTROY(INFINI_DEVICE_CAMBRICON, cambricon); COMM_DESTROY(INFINI_DEVICE_METAX, metax); + COMM_DESTROY(INFINI_DEVICE_MARS, mars); COMM_DESTROY(INFINI_DEVICE_MOORE, moore); COMM_DESTROY(INFINI_DEVICE_KUNLUN, kunlun); COMM_DESTROY(INFINI_DEVICE_ALI, cuda); @@ -78,6 +81,7 @@ __INFINI_C infiniStatus_t infinicclGroupStart(infinicclComm_t comm) { GROUP_START(INFINI_DEVICE_ASCEND, ascend); GROUP_START(INFINI_DEVICE_CAMBRICON, cambricon); GROUP_START(INFINI_DEVICE_METAX, metax); + GROUP_START(INFINI_DEVICE_MARS, mars); GROUP_START(INFINI_DEVICE_MOORE, moore); GROUP_START(INFINI_DEVICE_KUNLUN, kunlun); GROUP_START(INFINI_DEVICE_ALI, cuda); @@ -105,6 +109,7 @@ __INFINI_C infiniStatus_t infinicclGroupEnd(infinicclComm_t comm) { GROUP_END(INFINI_DEVICE_ASCEND, ascend); GROUP_END(INFINI_DEVICE_CAMBRICON, cambricon); GROUP_END(INFINI_DEVICE_METAX, metax); + GROUP_END(INFINI_DEVICE_MARS, mars); GROUP_END(INFINI_DEVICE_MOORE, moore); GROUP_END(INFINI_DEVICE_KUNLUN, kunlun); GROUP_END(INFINI_DEVICE_ALI, cuda); @@ -140,6 +145,7 @@ __INFINI_C infiniStatus_t infinicclAllReduce( ALL_REDUCE(INFINI_DEVICE_ASCEND, ascend); ALL_REDUCE(INFINI_DEVICE_CAMBRICON, cambricon); ALL_REDUCE(INFINI_DEVICE_METAX, metax); + ALL_REDUCE(INFINI_DEVICE_MARS, mars); ALL_REDUCE(INFINI_DEVICE_MOORE, moore); ALL_REDUCE(INFINI_DEVICE_KUNLUN, kunlun); ALL_REDUCE(INFINI_DEVICE_ALI, cuda); @@ -175,6 +181,7 @@ __INFINI_C infiniStatus_t infinicclAllGather( ALL_GATHER(INFINI_DEVICE_ASCEND, ascend); ALL_GATHER(INFINI_DEVICE_CAMBRICON, cambricon); ALL_GATHER(INFINI_DEVICE_METAX, metax); + ALL_GATHER(INFINI_DEVICE_MARS, mars); ALL_GATHER(INFINI_DEVICE_MOORE, moore); ALL_GATHER(INFINI_DEVICE_KUNLUN, kunlun); ALL_GATHER(INFINI_DEVICE_ALI, cuda); @@ -211,6 +218,7 @@ __INFINI_C infiniStatus_t infinicclAllGatherV( ALL_GATHER_V(INFINI_DEVICE_ASCEND, ascend); ALL_GATHER_V(INFINI_DEVICE_CAMBRICON, cambricon); ALL_GATHER_V(INFINI_DEVICE_METAX, metax); + ALL_GATHER_V(INFINI_DEVICE_MARS, mars); ALL_GATHER_V(INFINI_DEVICE_MOORE, moore); ALL_GATHER_V(INFINI_DEVICE_KUNLUN, kunlun); ALL_GATHER_V(INFINI_DEVICE_ALI, cuda); @@ -247,6 +255,7 @@ __INFINI_C infiniStatus_t infinicclReduceScatter( REDUCE_SCATTER(INFINI_DEVICE_ASCEND, ascend); REDUCE_SCATTER(INFINI_DEVICE_CAMBRICON, cambricon); REDUCE_SCATTER(INFINI_DEVICE_METAX, metax); + REDUCE_SCATTER(INFINI_DEVICE_MARS, mars); REDUCE_SCATTER(INFINI_DEVICE_MOORE, moore); REDUCE_SCATTER(INFINI_DEVICE_KUNLUN, kunlun); REDUCE_SCATTER(INFINI_DEVICE_ALI, cuda); @@ -284,6 +293,7 @@ __INFINI_C infiniStatus_t infinicclReduceScatterV( REDUCE_SCATTER_V(INFINI_DEVICE_ASCEND, ascend); REDUCE_SCATTER_V(INFINI_DEVICE_CAMBRICON, cambricon); REDUCE_SCATTER_V(INFINI_DEVICE_METAX, metax); + REDUCE_SCATTER_V(INFINI_DEVICE_MARS, mars); REDUCE_SCATTER_V(INFINI_DEVICE_MOORE, moore); REDUCE_SCATTER_V(INFINI_DEVICE_KUNLUN, kunlun); REDUCE_SCATTER_V(INFINI_DEVICE_ALI, cuda); diff --git a/src/infiniccl/mars/infiniccl_mars.cc b/src/infiniccl/mars/infiniccl_mars.cc new file mode 100644 index 000000000..ba3ebec70 --- /dev/null +++ b/src/infiniccl/mars/infiniccl_mars.cc @@ -0,0 +1,93 @@ +#include "infiniccl_mars.h" + +#include "../../utils.h" + +#include +#include + +#include +#include + +#define CHECK_HCCL(API__) CHECK_INTERNAL(API__, hcclSuccess) + +namespace { + +hcStream_t getMarsStream(infinirtStream_t stream) { + return stream == nullptr ? nullptr : static_cast(stream); +} + +hcclDataType_t getHcclDtype(infiniDtype_t datatype) { + switch (datatype) { + case INFINI_DTYPE_F32: + return hcclFloat; + case INFINI_DTYPE_F16: + return hcclHalf; + case INFINI_DTYPE_BF16: + return hcclBfloat16; + default: + std::abort(); + } +} + +hcclRedOp_t getHcclRedOp(infinicclReduceOp_t op) { + switch (op) { + case INFINICCL_SUM: + return hcclSum; + case INFINICCL_PROD: + return hcclProd; + case INFINICCL_MAX: + return hcclMax; + case INFINICCL_MIN: + return hcclMin; + case INFINICCL_AVG: + return hcclAvg; + default: + std::abort(); + } +} + +hcclComm_t getHcclComm(infinicclComm_t comm) { + return static_cast(comm->comm); +} + +} // namespace + +namespace infiniccl::mars { + +infiniStatus_t commInitAll(infinicclComm_t *comms, int ndevice, const int *device_ids) { + std::vector hccl_comms(ndevice); + CHECK_HCCL(hcclCommInitAll(hccl_comms.data(), ndevice, device_ids)); + + for (int i = 0; i < ndevice; ++i) { + comms[i] = new InfinicclComm{INFINI_DEVICE_MARS, device_ids[i], hccl_comms[i]}; + } + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t commDestroy(infinicclComm_t comm) { + CHECK_HCCL(hcclCommDestroy(getHcclComm(comm))); + delete comm; + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t allReduce(void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t datatype, + infinicclReduceOp_t op, + infinicclComm_t comm, + infinirtStream_t stream) { + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); + CHECK_HCCL(hcclAllReduce(sendbuf, + recvbuf, + count, + getHcclDtype(datatype), + getHcclRedOp(op), + getHcclComm(comm), + getMarsStream(stream))); + return INFINI_STATUS_SUCCESS; +} + +INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES(); + +} // namespace infiniccl::mars diff --git a/src/infiniccl/mars/infiniccl_mars.h b/src/infiniccl/mars/infiniccl_mars.h new file mode 100644 index 000000000..0001f0d91 --- /dev/null +++ b/src/infiniccl/mars/infiniccl_mars.h @@ -0,0 +1,12 @@ +#ifndef INFINICCL_MARS_H_ +#define INFINICCL_MARS_H_ + +#include "../infiniccl_impl.h" + +#if defined(ENABLE_MARS_API) && defined(ENABLE_CCL) +INFINICCL_DEVICE_API_IMPL(mars) +#else +INFINICCL_DEVICE_API_NOOP(mars) +#endif + +#endif // INFINICCL_MARS_H_ diff --git a/src/infinicore-test/main.cc b/src/infinicore-test/main.cc index c2c32170d..2a240df62 100644 --- a/src/infinicore-test/main.cc +++ b/src/infinicore-test/main.cc @@ -41,6 +41,7 @@ void printUsage() { << " cambricon" << std::endl << " ascend" << std::endl << " metax" << std::endl + << " mars" << std::endl << " moore" << std::endl << " iluvatar" << std::endl << " qy" << std::endl @@ -79,6 +80,8 @@ ParsedArgs parseArgs(int argc, char *argv[]) { args.device_type = INFINI_DEVICE_ASCEND; } else if (arg == "--metax") { args.device_type = INFINI_DEVICE_METAX; + } else if (arg == "--mars") { + args.device_type = INFINI_DEVICE_MARS; } else if (arg == "--moore") { args.device_type = INFINI_DEVICE_MOORE; } else if (arg == "--iluvatar") { diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index d6991ebdd..93f24f863 100644 --- a/src/infinicore/context/context_impl.cc +++ b/src/infinicore/context/context_impl.cc @@ -93,6 +93,7 @@ ContextImpl::ContextImpl() { initializeRuntime(); initializeRuntime(); initializeRuntime(); + initializeRuntime(); initializeRuntime(); if (current_runtime_ == nullptr && !runtime_table_[static_cast(Device::Type::CPU)].empty()) { diff --git a/src/infinicore/device.cc b/src/infinicore/device.cc index 77c1f4986..2a978eff0 100644 --- a/src/infinicore/device.cc +++ b/src/infinicore/device.cc @@ -43,6 +43,8 @@ std::string Device::toString(const Type &type) { return "HYGON"; case Type::ALI: return "ALI"; + case Type::MARS: + return "MARS"; case Type::COUNT: return "COUNT"; default: diff --git a/src/infinicore/nn/rmsnorm.cc b/src/infinicore/nn/rmsnorm.cc index 567d88e49..b01b5fa19 100644 --- a/src/infinicore/nn/rmsnorm.cc +++ b/src/infinicore/nn/rmsnorm.cc @@ -30,6 +30,7 @@ void RMSNorm::forward_inplace(Tensor &x, Tensor &residual) const { || device_.getType() == Device::Type::NVIDIA || device_.getType() == Device::Type::ILUVATAR || device_.getType() == Device::Type::METAX + || device_.getType() == Device::Type::MARS || device_.getType() == Device::Type::MOORE || device_.getType() == Device::Type::ALI || device_.getType() == Device::Type::CAMBRICON diff --git a/src/infinicore/ops/infiniops_impl.hpp b/src/infinicore/ops/infiniops_impl.hpp index befb830cd..5c3d475f0 100644 --- a/src/infinicore/ops/infiniops_impl.hpp +++ b/src/infinicore/ops/infiniops_impl.hpp @@ -51,6 +51,8 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) { return infini::ops::Device{infini::ops::Device::Type::kNvidia, static_cast(device.getIndex())}; case Device::Type::METAX: return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast(device.getIndex())}; + case Device::Type::MARS: + return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast(device.getIndex())}; case Device::Type::MOORE: return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast(device.getIndex())}; case Device::Type::ILUVATAR: @@ -64,6 +66,7 @@ inline bool isSupportedDevice(Device::Type device_type) { switch (device_type) { case Device::Type::NVIDIA: case Device::Type::METAX: + case Device::Type::MARS: case Device::Type::MOORE: case Device::Type::ILUVATAR: return true; @@ -76,6 +79,7 @@ template void registerSupportedDevices(Dispatcher &dispatcher, Function function) { dispatcher.registerDevice(Device::Type::NVIDIA, function); dispatcher.registerDevice(Device::Type::METAX, function); + dispatcher.registerDevice(Device::Type::MARS, function); dispatcher.registerDevice(Device::Type::MOORE, function); dispatcher.registerDevice(Device::Type::ILUVATAR, function); } diff --git a/src/infinicore/ops/rotg/rotg_infiniop.cc b/src/infinicore/ops/rotg/rotg_infiniop.cc index d5e7a0e99..79d911bcf 100644 --- a/src/infinicore/ops/rotg/rotg_infiniop.cc +++ b/src/infinicore/ops/rotg/rotg_infiniop.cc @@ -47,7 +47,8 @@ void calculate(Tensor x, Tensor y, Tensor c, Tensor s) { static bool registered = []() { Rotg::dispatcher().registerDevice({Device::Type::CPU, Device::Type::CAMBRICON, - Device::Type::METAX}, + Device::Type::METAX, + Device::Type::MARS}, &calculate, false); return true; diff --git a/src/infinicore/ops/rotm/rotm_infiniop.cc b/src/infinicore/ops/rotm/rotm_infiniop.cc index 9df5c9ab1..ccf0ceb79 100644 --- a/src/infinicore/ops/rotm/rotm_infiniop.cc +++ b/src/infinicore/ops/rotm/rotm_infiniop.cc @@ -47,7 +47,8 @@ void calculate(Tensor x, Tensor y, Tensor param) { static bool registered = []() { Rotm::dispatcher().registerDevice({Device::Type::CPU, Device::Type::CAMBRICON, - Device::Type::METAX}, + Device::Type::METAX, + Device::Type::MARS}, &calculate, false); return true; diff --git a/src/infinicore/ops/rotmg/rotmg_infiniop.cc b/src/infinicore/ops/rotmg/rotmg_infiniop.cc index ade0bd7aa..360693ebe 100644 --- a/src/infinicore/ops/rotmg/rotmg_infiniop.cc +++ b/src/infinicore/ops/rotmg/rotmg_infiniop.cc @@ -47,7 +47,8 @@ void calculate(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param) { static bool registered = []() { Rotmg::dispatcher().registerDevice({Device::Type::CPU, Device::Type::CAMBRICON, - Device::Type::METAX}, + Device::Type::METAX, + Device::Type::MARS}, &calculate, false); return true; diff --git a/src/infinicore/pybind11/device.hpp b/src/infinicore/pybind11/device.hpp index 74cb0a23b..83a6f30a7 100644 --- a/src/infinicore/pybind11/device.hpp +++ b/src/infinicore/pybind11/device.hpp @@ -23,6 +23,7 @@ inline void bind(py::module &m) { .value("KUNLUN", Device::Type::KUNLUN) .value("HYGON", Device::Type::HYGON) .value("ALI", Device::Type::ALI) + .value("MARS", Device::Type::MARS) .value("COUNT", Device::Type::COUNT); device diff --git a/src/infiniop/devices/handle.cc b/src/infiniop/devices/handle.cc index 9057f7203..6b17fe6bd 100644 --- a/src/infiniop/devices/handle.cc +++ b/src/infiniop/devices/handle.cc @@ -23,6 +23,9 @@ #ifdef ENABLE_METAX_API #include "metax/metax_handle.h" #endif +#ifdef ENABLE_MARS_API +#include "mars/mars_handle.h" +#endif __INFINI_C infiniStatus_t infiniopSetRuntimeDevice(infiniDevice_t device, int device_id) { auto rt_device = infinicore::bridge::infini::rt::translate_to(device); @@ -83,6 +86,9 @@ __INFINI_C infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr) { #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, mars); +#endif #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, hygon); #endif @@ -132,6 +138,9 @@ __INFINI_C infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, mars); +#endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, hygon); #endif diff --git a/src/infiniop/devices/mars/mars_handle.h b/src/infiniop/devices/mars/mars_handle.h new file mode 100644 index 000000000..e0545caeb --- /dev/null +++ b/src/infiniop/devices/mars/mars_handle.h @@ -0,0 +1,20 @@ +#ifndef __INFINIOP_MARS_HANDLE_H__ +#define __INFINIOP_MARS_HANDLE_H__ + +#include "../metax/metax_handle.h" + +namespace device::mars { + +struct Handle final : public device::metax::Handle { + explicit Handle(int device_id) + : device::metax::Handle(INFINI_DEVICE_MARS, device_id) {} + + static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id) { + *handle_ptr = new Handle(device_id); + return INFINI_STATUS_SUCCESS; + } +}; + +} // namespace device::mars + +#endif // __INFINIOP_MARS_HANDLE_H__ diff --git a/src/infiniop/ops/acos/operator.cc b/src/infiniop/ops/acos/operator.cc index 16f2649ae..26bdf1138 100644 --- a/src/infiniop/ops/acos/operator.cc +++ b/src/infiniop/ops/acos/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/acos_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/acos_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateAcosDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetAcosWorkspaceSize(infiniopAcosDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif @@ -137,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopAcos( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif @@ -175,6 +184,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAcosDescriptor(infiniopAcosDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif diff --git a/src/infiniop/ops/adaptive_avg_pool1d/operator.cc b/src/infiniop/ops/adaptive_avg_pool1d/operator.cc index 3c92f6d96..a21cf75f5 100644 --- a/src/infiniop/ops/adaptive_avg_pool1d/operator.cc +++ b/src/infiniop/ops/adaptive_avg_pool1d/operator.cc @@ -12,7 +12,7 @@ #endif // [Metax Support] -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/adaptive_avg_pool1d_metax.h" #endif @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateAdaptiveAvgPool1dDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -95,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopGetAdaptiveAvgPool1dWorkspaceSize(infiniopAdap #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -140,6 +146,9 @@ __INFINI_C infiniStatus_t infiniopAdaptiveAvgPool1d( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -179,6 +188,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAdaptiveAvgPool1dDescriptor(infiniopAda #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/operator.cc b/src/infiniop/ops/adaptive_avg_pool3d/operator.cc index 0b2036a6a..2d191f948 100644 --- a/src/infiniop/ops/adaptive_avg_pool3d/operator.cc +++ b/src/infiniop/ops/adaptive_avg_pool3d/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/adaptive_avg_pool3d_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/adaptive_avg_pool3d_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateAdaptiveAvgPool3DDescriptor( #endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetAdaptiveAvgPool3DWorkspaceSize( #endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -135,6 +141,9 @@ __INFINI_C infiniStatus_t infiniopAdaptiveAvgPool3D( #endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -170,6 +179,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAdaptiveAvgPool3DDescriptor(infiniopAda #endif #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/adaptive_max_pool1d/operator.cc b/src/infiniop/ops/adaptive_max_pool1d/operator.cc index 52d9e805f..548592b08 100644 --- a/src/infiniop/ops/adaptive_max_pool1d/operator.cc +++ b/src/infiniop/ops/adaptive_max_pool1d/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/adaptive_max_pool1d_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/adaptive_max_pool1d_metax.cuh" #endif #ifdef ENABLE_MOORE_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreateAdaptiveMaxPool1dDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetAdaptiveMaxPool1dWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -114,6 +120,9 @@ __INFINI_C infiniStatus_t infiniopAdaptiveMaxPool1d( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -146,6 +155,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAdaptiveMaxPool1dDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/add/operator.cc b/src/infiniop/ops/add/operator.cc index 9546a7577..3918660a8 100644 --- a/src/infiniop/ops/add/operator.cc +++ b/src/infiniop/ops/add/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/add_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/add_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -64,6 +64,9 @@ __INFINI_C infiniStatus_t infiniopCreateAddDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -113,6 +116,9 @@ __INFINI_C infiniStatus_t infiniopGetAddWorkspaceSize(infiniopAddDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -208,6 +214,9 @@ __INFINI_C infiniStatus_t infiniopAdd( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -259,6 +268,9 @@ infiniopDestroyAddDescriptor(infiniopAddDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/add_rms_norm/operator.cc b/src/infiniop/ops/add_rms_norm/operator.cc index 7d502770d..1e6138f13 100644 --- a/src/infiniop/ops/add_rms_norm/operator.cc +++ b/src/infiniop/ops/add_rms_norm/operator.cc @@ -15,7 +15,7 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/add_rms_norm_bang.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/add_rms_norm_metax.cuh" #endif #ifdef ENABLE_MOORE_API @@ -67,6 +67,9 @@ __INFINI_C infiniStatus_t infiniopCreateAddRMSNormDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -112,6 +115,9 @@ __INFINI_C infiniStatus_t infiniopGetAddRMSNormWorkspaceSize(infiniopAddRMSNormD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -168,6 +174,9 @@ __INFINI_C infiniStatus_t infiniopAddRMSNorm( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -215,6 +224,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAddRMSNormDescriptor(infiniopAddRMSNorm #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DESTROY(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/addbmm/operator.cc b/src/infiniop/ops/addbmm/operator.cc index 53ca28e2e..d0f48055e 100644 --- a/src/infiniop/ops/addbmm/operator.cc +++ b/src/infiniop/ops/addbmm/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/addbmm_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/addbmm_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -65,6 +65,9 @@ __INFINI_C infiniStatus_t infiniopCreateAddbmmDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -104,6 +107,9 @@ __INFINI_C infiniStatus_t infiniopGetAddbmmWorkspaceSize(infiniopAddbmmDescripto #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -151,6 +157,9 @@ __INFINI_C infiniStatus_t infiniopAddbmm( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -190,6 +199,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAddbmmDescriptor(infiniopAddbmmDescript #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/addcmul/operator.cc b/src/infiniop/ops/addcmul/operator.cc index ca071adb5..eefee6f04 100644 --- a/src/infiniop/ops/addcmul/operator.cc +++ b/src/infiniop/ops/addcmul/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/addcmul_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/addcmul_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateAddcmulDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -87,6 +90,9 @@ __INFINI_C infiniStatus_t infiniopGetAddcmulWorkspaceSize(infiniopAddcmulDescrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -130,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopAddcmul( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -165,6 +174,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAddcmulDescriptor(infiniopAddcmulDescri #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/addr/operator.cc b/src/infiniop/ops/addr/operator.cc index 1b5101352..fce1b0dba 100644 --- a/src/infiniop/ops/addr/operator.cc +++ b/src/infiniop/ops/addr/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/addr_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/addr_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -59,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopCreateAddrDescriptor( #endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -99,6 +102,9 @@ __INFINI_C infiniStatus_t infiniopGetAddrWorkspaceSize( #endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -145,6 +151,9 @@ __INFINI_C infiniStatus_t infiniopAddr( #endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -183,6 +192,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAddrDescriptor(infiniopAddrDescriptor_t #endif #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/affine_grid/operator.cc b/src/infiniop/ops/affine_grid/operator.cc index d5212f798..079d2d207 100644 --- a/src/infiniop/ops/affine_grid/operator.cc +++ b/src/infiniop/ops/affine_grid/operator.cc @@ -12,7 +12,7 @@ #ifdef ENABLE_MOORE_API #include "moore/affine_grid_moore.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/affine_grid_metax.h" #endif @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateAffineGridDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ILUVATAR_API CREATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif @@ -85,6 +88,9 @@ __INFINI_C infiniStatus_t infiniopGetAffineGridWorkspaceSize(infiniopAffineGridD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_QY_API GET(INFINI_DEVICE_QY, nvidia); #endif @@ -126,6 +132,9 @@ __INFINI_C infiniStatus_t infiniopAffineGrid( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ILUVATAR_API CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif @@ -165,6 +174,9 @@ infiniopDestroyAffineGridDescriptor(infiniopAffineGridDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ILUVATAR_API DELETE(INFINI_DEVICE_ILUVATAR, nvidia); #endif diff --git a/src/infiniop/ops/all/operator.cc b/src/infiniop/ops/all/operator.cc index 8ffb34b63..753ce82dd 100644 --- a/src/infiniop/ops/all/operator.cc +++ b/src/infiniop/ops/all/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/all_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/all_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateAllDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetAllWorkspaceSize(infiniopAllDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -141,6 +147,9 @@ __INFINI_C infiniStatus_t infiniopAll( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -180,6 +189,9 @@ infiniopDestroyAllDescriptor(infiniopAllDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/argwhere/operator.cc b/src/infiniop/ops/argwhere/operator.cc index e1eaccfd3..c123f7427 100644 --- a/src/infiniop/ops/argwhere/operator.cc +++ b/src/infiniop/ops/argwhere/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/argwhere_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/argwhere_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreateArgwhereDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -83,6 +86,9 @@ __INFINI_C infiniStatus_t infiniopGetArgwhereWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -128,6 +134,9 @@ __INFINI_C infiniStatus_t infiniopArgwhere( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -167,6 +176,9 @@ infiniopDestroyArgwhereDescriptor(infiniopArgwhereDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/asin/operator.cc b/src/infiniop/ops/asin/operator.cc index e372286b3..46e61f110 100644 --- a/src/infiniop/ops/asin/operator.cc +++ b/src/infiniop/ops/asin/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/asin_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/asin_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateAsinDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetAsinWorkspaceSize(infiniopAsinDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopAsin( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyAsinDescriptor(infiniopAsinDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/asinh/operator.cc b/src/infiniop/ops/asinh/operator.cc index 021ddb6b0..8daac4831 100644 --- a/src/infiniop/ops/asinh/operator.cc +++ b/src/infiniop/ops/asinh/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/asinh_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/asinh_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -41,6 +41,9 @@ __INFINI_C infiniStatus_t infiniopCreateAsinhDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -73,6 +76,9 @@ __INFINI_C infiniStatus_t infiniopGetAsinhWorkspaceSize(infiniopAsinhDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -109,6 +115,9 @@ __INFINI_C infiniStatus_t infiniopAsinh(infiniopAsinhDescriptor_t desc, #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -140,6 +149,9 @@ __INFINI_C infiniStatus_t infiniopDestroyAsinhDescriptor(infiniopAsinhDescriptor #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/asum/operator.cc b/src/infiniop/ops/asum/operator.cc index 51769f0d3..2b2784379 100644 --- a/src/infiniop/ops/asum/operator.cc +++ b/src/infiniop/ops/asum/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/asum_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/asum_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -33,6 +33,9 @@ __INFINI_C infiniStatus_t infiniopCreateAsumDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -57,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopGetAsumWorkspaceSize(infiniopAsumDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -88,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopAsum( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -113,6 +122,9 @@ infiniopDestroyAsumDescriptor(infiniopAsumDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/atanh/operator.cc b/src/infiniop/ops/atanh/operator.cc index 0472d6d43..000fa4136 100644 --- a/src/infiniop/ops/atanh/operator.cc +++ b/src/infiniop/ops/atanh/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/atanh_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/atanh_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateAtanhDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -86,6 +89,9 @@ __INFINI_C infiniStatus_t infiniopGetAtanhWorkspaceSize(infiniopAtanhDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -128,6 +134,9 @@ __INFINI_C infiniStatus_t infiniopAtanh( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -167,6 +176,9 @@ infiniopDestroyAtanhDescriptor(infiniopAtanhDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/avg_pool1d/operator.cc b/src/infiniop/ops/avg_pool1d/operator.cc index 54ea110fd..40e1b1548 100644 --- a/src/infiniop/ops/avg_pool1d/operator.cc +++ b/src/infiniop/ops/avg_pool1d/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/avg_pool1d_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/avg_pool1d_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateAvgPool1dDescriptor( #endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -92,6 +95,9 @@ __INFINI_C infiniStatus_t infiniopGetAvgPool1dWorkspaceSize(infiniopAvgPool1dDes #endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -134,6 +140,9 @@ __INFINI_C infiniStatus_t infiniopAvgPool1d( #endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -171,6 +180,9 @@ infiniopDestroyAvgPool1dDescriptor(infiniopAvgPool1dDescriptor_t desc) { #endif #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/avg_pool3d/operator.cc b/src/infiniop/ops/avg_pool3d/operator.cc index df58e1282..79a4b0299 100644 --- a/src/infiniop/ops/avg_pool3d/operator.cc +++ b/src/infiniop/ops/avg_pool3d/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) #include "nvidia/avg_pool3d_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/avg_pool3d_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateAvgPool3dDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetAvgPool3dWorkspaceSize(infiniopAvgPool3dDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopAvgPool3d( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -151,6 +160,9 @@ infiniopDestroyAvgPool3dDescriptor(infiniopAvgPool3dDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/axpy/operator.cc b/src/infiniop/ops/axpy/operator.cc index cd57223e4..87c461837 100644 --- a/src/infiniop/ops/axpy/operator.cc +++ b/src/infiniop/ops/axpy/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/axpy_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/axpy_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -35,6 +35,9 @@ __INFINI_C infiniStatus_t infiniopCreateAxpyDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopGetAxpyWorkspaceSize(infiniopAxpyDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -91,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopAxpy( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -116,6 +125,9 @@ infiniopDestroyAxpyDescriptor(infiniopAxpyDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc b/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc index 06180aa2c..b42b798d5 100644 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc +++ b/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/binary_cross_entropy_with_logits_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/binary_cross_entropy_with_logits_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -54,6 +54,9 @@ __INFINI_C infiniStatus_t infiniopCreateBCEWithLogitsDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -94,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopGetBCEWithLogitsWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -140,6 +146,9 @@ __INFINI_C infiniStatus_t infiniopBCEWithLogits( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -178,6 +187,9 @@ __INFINI_C infiniStatus_t infiniopDestroyBCEWithLogitsDescriptor(infiniopBCEWith #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/bitwise_right_shift/operator.cc b/src/infiniop/ops/bitwise_right_shift/operator.cc index b1f3d81ed..9fd0c0eb6 100644 --- a/src/infiniop/ops/bitwise_right_shift/operator.cc +++ b/src/infiniop/ops/bitwise_right_shift/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/bitwise_right_shift_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/bitwise_right_shift_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreateBitwiseRightShiftDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -78,6 +81,9 @@ __INFINI_C infiniStatus_t infiniopGetBitwiseRightShiftWorkspaceSize(infiniopBitw #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -120,6 +126,9 @@ __INFINI_C infiniStatus_t infiniopBitwiseRightShift( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -156,6 +165,9 @@ infiniopDestroyBitwiseRightShiftDescriptor(infiniopBitwiseRightShiftDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/blas_amax/operator.cc b/src/infiniop/ops/blas_amax/operator.cc index c6b48eeb4..ec1fd95af 100644 --- a/src/infiniop/ops/blas_amax/operator.cc +++ b/src/infiniop/ops/blas_amax/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/blas_amax_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/blas_amax_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -33,6 +33,9 @@ __INFINI_C infiniStatus_t infiniopCreateBlasAmaxDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -57,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopGetBlasAmaxWorkspaceSize(infiniopBlasAmaxDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -87,6 +93,9 @@ __INFINI_C infiniStatus_t infiniopBlasAmax( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -111,6 +120,9 @@ __INFINI_C infiniStatus_t infiniopDestroyBlasAmaxDescriptor(infiniopBlasAmaxDesc #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/blas_amin/operator.cc b/src/infiniop/ops/blas_amin/operator.cc index 7f960b773..5961b4238 100644 --- a/src/infiniop/ops/blas_amin/operator.cc +++ b/src/infiniop/ops/blas_amin/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/blas_amin_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/blas_amin_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -33,6 +33,9 @@ __INFINI_C infiniStatus_t infiniopCreateBlasAminDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -57,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopGetBlasAminWorkspaceSize(infiniopBlasAminDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -87,6 +93,9 @@ __INFINI_C infiniStatus_t infiniopBlasAmin( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -111,6 +120,9 @@ __INFINI_C infiniStatus_t infiniopDestroyBlasAminDescriptor(infiniopBlasAminDesc #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/blas_copy/operator.cc b/src/infiniop/ops/blas_copy/operator.cc index 394adc665..978aaa23c 100644 --- a/src/infiniop/ops/blas_copy/operator.cc +++ b/src/infiniop/ops/blas_copy/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/blas_copy_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/blas_copy_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -32,6 +32,9 @@ __INFINI_C infiniStatus_t infiniopCreateBlasCopyDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -56,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopGetBlasCopyWorkspaceSize(infiniopBlasCopyDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -86,6 +92,9 @@ __INFINI_C infiniStatus_t infiniopBlasCopy( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -110,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopDestroyBlasCopyDescriptor(infiniopBlasCopyDesc #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/blas_dot/operator.cc b/src/infiniop/ops/blas_dot/operator.cc index e28943b8b..12b52d0be 100644 --- a/src/infiniop/ops/blas_dot/operator.cc +++ b/src/infiniop/ops/blas_dot/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/blas_dot_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/blas_dot_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -35,6 +35,9 @@ __INFINI_C infiniStatus_t infiniopCreateBlasDotDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopGetBlasDotWorkspaceSize(infiniopBlasDotDescrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -90,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopBlasDot( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -114,6 +123,9 @@ __INFINI_C infiniStatus_t infiniopDestroyBlasDotDescriptor(infiniopBlasDotDescri #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/block_diag/operator.cc b/src/infiniop/ops/block_diag/operator.cc index 31e6198f0..d00384a48 100644 --- a/src/infiniop/ops/block_diag/operator.cc +++ b/src/infiniop/ops/block_diag/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/block_diag_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/block_diag_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -45,6 +45,9 @@ __INFINI_C infiniStatus_t infiniopCreateBlockDiagDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -79,6 +82,9 @@ __INFINI_C infiniStatus_t infiniopGetBlockDiagWorkspaceSize(infiniopBlockDiagDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -120,6 +126,9 @@ __INFINI_C infiniStatus_t infiniopBlockDiag( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -156,6 +165,9 @@ infiniopDestroyBlockDiagDescriptor(infiniopBlockDiagDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/broadcast_to/operator.cc b/src/infiniop/ops/broadcast_to/operator.cc index ea861f970..d5ff54614 100644 --- a/src/infiniop/ops/broadcast_to/operator.cc +++ b/src/infiniop/ops/broadcast_to/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/broadcast_to_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/broadcast_to_metax.h" #endif @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateBroadcastToDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetBroadcastToWorkspaceSize(infiniopBroadcastT #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -141,6 +147,9 @@ __INFINI_C infiniStatus_t infiniopBroadcastTo( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -182,6 +191,9 @@ __INFINI_C infiniStatus_t infiniopDestroyBroadcastToDescriptor(infiniopBroadcast #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/causal_conv1d/operator.cc b/src/infiniop/ops/causal_conv1d/operator.cc index 8666d2e62..9abb67820 100644 --- a/src/infiniop/ops/causal_conv1d/operator.cc +++ b/src/infiniop/ops/causal_conv1d/operator.cc @@ -5,7 +5,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "nvidia/causal_conv1d_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/causal_conv1d_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreateCausalConv1dDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -70,6 +73,9 @@ __INFINI_C infiniStatus_t infiniopGetCausalConv1dWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -110,6 +116,9 @@ __INFINI_C infiniStatus_t infiniopCausalConv1d( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif @@ -136,6 +145,9 @@ __INFINI_C infiniStatus_t infiniopDestroyCausalConv1dDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/src/infiniop/ops/causal_softmax/operator.cc b/src/infiniop/ops/causal_softmax/operator.cc index f8d2bfba6..87a4d4d39 100644 --- a/src/infiniop/ops/causal_softmax/operator.cc +++ b/src/infiniop/ops/causal_softmax/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/causal_softmax_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/causal_softmax_metax.h" #endif #ifdef ENABLE_ASCEND_API @@ -63,6 +63,9 @@ __INFINI_C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ASCEND_API CREATE(INFINI_DEVICE_ASCEND, ascend) #endif @@ -107,6 +110,9 @@ __INFINI_C infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSo #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ASCEND_API GET(INFINI_DEVICE_ASCEND, ascend) #endif @@ -162,6 +168,9 @@ __INFINI_C infiniStatus_t infiniopCausalSoftmax( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ASCEND_API CALCULATE(INFINI_DEVICE_ASCEND, ascend) #endif @@ -209,6 +218,9 @@ __INFINI_C infiniStatus_t infiniopDestroyCausalSoftmaxDescriptor(infiniopCausalS #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ASCEND_API DESTROY(INFINI_DEVICE_ASCEND, ascend) #endif diff --git a/src/infiniop/ops/cdist/operator.cc b/src/infiniop/ops/cdist/operator.cc index 0ed6541b7..fb7a3ed2d 100644 --- a/src/infiniop/ops/cdist/operator.cc +++ b/src/infiniop/ops/cdist/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/cdist_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/cdist_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateCdistDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -92,6 +95,9 @@ __INFINI_C infiniStatus_t infiniopGetCdistWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -136,6 +142,9 @@ __INFINI_C infiniStatus_t infiniopCdist( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -174,6 +183,9 @@ __INFINI_C infiniStatus_t infiniopDestroyCdistDescriptor(infiniopCdistDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc index 19b5198bf..939fff404 100644 --- a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc @@ -7,7 +7,7 @@ #if defined(ENABLE_NVIDIA_API) #include "nvidia/chunk_gated_delta_rule_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/chunk_gated_delta_rule_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -50,6 +50,9 @@ __INFINI_C infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -76,6 +79,9 @@ __INFINI_C infiniStatus_t infiniopGetChunkGatedDeltaRuleWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -109,6 +115,9 @@ __INFINI_C infiniStatus_t infiniopChunkGatedDeltaRule( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif @@ -134,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopDestroyChunkGatedDeltaRuleDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/src/infiniop/ops/clip/operator.cc b/src/infiniop/ops/clip/operator.cc index c933cafab..9eeb06f80 100644 --- a/src/infiniop/ops/clip/operator.cc +++ b/src/infiniop/ops/clip/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/clip_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/clip_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -51,6 +51,9 @@ __INFINI_C infiniStatus_t infiniopCreateClipDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -91,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun) #endif @@ -139,6 +145,9 @@ __INFINI_C infiniStatus_t infiniopClip( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -181,6 +190,9 @@ infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/cross_entropy/operator.cc b/src/infiniop/ops/cross_entropy/operator.cc index 75f35fcb7..6392b8f0f 100644 --- a/src/infiniop/ops/cross_entropy/operator.cc +++ b/src/infiniop/ops/cross_entropy/operator.cc @@ -13,7 +13,7 @@ #ifdef ENABLE_MOORE_API #include "moore/cross_entropy_moore.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/cross_entropy_metax.h" #endif @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateCrossEntropyDescriptor( #endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) +#endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -88,6 +91,9 @@ __INFINI_C infiniStatus_t infiniopGetCrossEntropyWorkspaceSize( #endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) +#endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -130,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopCrossEntropy( #endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) +#endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -166,6 +175,9 @@ __INFINI_C infiniStatus_t infiniopDestroyCrossEntropyDescriptor( #endif #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) +#endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/deepseek_moe/operator.cc b/src/infiniop/ops/deepseek_moe/operator.cc index d1ddfb58a..b75f6be08 100644 --- a/src/infiniop/ops/deepseek_moe/operator.cc +++ b/src/infiniop/ops/deepseek_moe/operator.cc @@ -6,7 +6,7 @@ || defined(ENABLE_ALI_API) #include "nvidia/deepseek_moe_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/deepseek_moe_metax.cuh" #endif #ifdef ENABLE_MOORE_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateDeepseekMoeDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetDeepseekMoeWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -124,6 +130,9 @@ __INFINI_C infiniStatus_t infiniopDeepseekMoe( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -168,6 +177,9 @@ __INFINI_C infiniStatus_t infiniopDeepseekMoeWithDevicePtrs( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -201,6 +213,9 @@ __INFINI_C infiniStatus_t infiniopDestroyDeepseekMoeDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/diff/operator.cc b/src/infiniop/ops/diff/operator.cc index 3037dc3b5..911031874 100644 --- a/src/infiniop/ops/diff/operator.cc +++ b/src/infiniop/ops/diff/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/diff_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/diff_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -47,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateDiffDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -81,6 +84,9 @@ __INFINI_C infiniStatus_t infiniopGetDiffWorkspaceSize(infiniopDiffDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -122,6 +128,9 @@ __INFINI_C infiniStatus_t infiniopDiff( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -158,6 +167,9 @@ infiniopDestroyDiffDescriptor(infiniopDiffDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/digamma/operator.cc b/src/infiniop/ops/digamma/operator.cc index d6be706ca..e14c8632b 100644 --- a/src/infiniop/ops/digamma/operator.cc +++ b/src/infiniop/ops/digamma/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/digamma_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/digamma_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateDigammaDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetDigammaWorkspaceSize(infiniopDigammaDescrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopDigamma( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyDigammaDescriptor(infiniopDigammaDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/dist/operator.cc b/src/infiniop/ops/dist/operator.cc index 6aba56b88..d713b9bb2 100644 --- a/src/infiniop/ops/dist/operator.cc +++ b/src/infiniop/ops/dist/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) #include "nvidia/dist_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/dist_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -47,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateDistDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -78,6 +81,9 @@ __INFINI_C infiniStatus_t infiniopGetDistWorkspaceSize(infiniopDistDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -117,6 +123,9 @@ __INFINI_C infiniStatus_t infiniopDist( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -150,6 +159,9 @@ infiniopDestroyDistDescriptor(infiniopDistDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/dot/operator.cc b/src/infiniop/ops/dot/operator.cc index 10f9fda57..08f1bf657 100644 --- a/src/infiniop/ops/dot/operator.cc +++ b/src/infiniop/ops/dot/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) #include "nvidia/dot_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/dot_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -45,6 +45,9 @@ __INFINI_C infiniStatus_t infiniopCreateDotDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -76,6 +79,9 @@ __INFINI_C infiniStatus_t infiniopGetDotWorkspaceSize(infiniopDotDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -115,6 +121,9 @@ __INFINI_C infiniStatus_t infiniopDot( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -148,6 +157,9 @@ infiniopDestroyDotDescriptor(infiniopDotDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/embedding/operator.cc b/src/infiniop/ops/embedding/operator.cc index 5301e57a5..027579601 100644 --- a/src/infiniop/ops/embedding/operator.cc +++ b/src/infiniop/ops/embedding/operator.cc @@ -15,7 +15,7 @@ #ifdef ENABLE_ASCEND_API #include "ascend/embedding_ascend.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/embedding_metax.cuh" #endif #ifdef ENABLE_MOORE_API @@ -67,6 +67,9 @@ __INFINI_C infiniStatus_t infiniopCreateEmbeddingDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -122,6 +125,9 @@ __INFINI_C infiniStatus_t infiniopEmbedding( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -168,6 +174,9 @@ __INFINI_C infiniStatus_t infiniopDestroyEmbeddingDescriptor(infiniopEmbeddingDe #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/equal/operator.cc b/src/infiniop/ops/equal/operator.cc index 9c5866c1e..58e8b640d 100644 --- a/src/infiniop/ops/equal/operator.cc +++ b/src/infiniop/ops/equal/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/equal_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/equal_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -47,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateEqualDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -84,6 +87,9 @@ __INFINI_C infiniStatus_t infiniopGetEqualWorkspaceSize(infiniopEqualDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -129,6 +135,9 @@ __INFINI_C infiniStatus_t infiniopEqual( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -168,6 +177,9 @@ infiniopDestroyEqualDescriptor(infiniopEqualDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/erf/operator.cc b/src/infiniop/ops/erf/operator.cc index af1327c3b..9f36eefa0 100644 --- a/src/infiniop/ops/erf/operator.cc +++ b/src/infiniop/ops/erf/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/erf_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/erf_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C __export infiniStatus_t infiniopCreateErfDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C __export infiniStatus_t infiniopGetErfWorkspaceSize(infiniopErfDescri #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C __export infiniStatus_t infiniopErf( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyErfDescriptor(infiniopErfDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/erfc/operator.cc b/src/infiniop/ops/erfc/operator.cc index ce283df19..3f13c23ca 100644 --- a/src/infiniop/ops/erfc/operator.cc +++ b/src/infiniop/ops/erfc/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/erfc_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/erfc_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C __export infiniStatus_t infiniopCreateErfcDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C __export infiniStatus_t infiniopGetErfcWorkspaceSize(infiniopErfcDesc #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C __export infiniStatus_t infiniopErfc( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyErfcDescriptor(infiniopErfcDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/erfinv/operator.cc b/src/infiniop/ops/erfinv/operator.cc index dc711c1ef..64de72634 100644 --- a/src/infiniop/ops/erfinv/operator.cc +++ b/src/infiniop/ops/erfinv/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/erfinv_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/erfinv_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C __export infiniStatus_t infiniopCreateErfinvDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C __export infiniStatus_t infiniopGetErfinvWorkspaceSize(infiniopErfinv #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C __export infiniStatus_t infiniopErfinv( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyErfinvDescriptor(infiniopErfinvDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/flipud/operator.cc b/src/infiniop/ops/flipud/operator.cc index 69a785ec6..6462c9c6e 100644 --- a/src/infiniop/ops/flipud/operator.cc +++ b/src/infiniop/ops/flipud/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/flipud_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/flipud_metax.h" #endif @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateFlipudDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -91,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopGetFlipudWorkspaceSize(infiniopFlipudDescripto #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -135,6 +141,9 @@ __INFINI_C infiniStatus_t infiniopFlipud( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -173,6 +182,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFlipudDescriptor(infiniopFlipudDescript #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/float_power/operator.cc b/src/infiniop/ops/float_power/operator.cc index de5485ef7..33d5feb33 100644 --- a/src/infiniop/ops/float_power/operator.cc +++ b/src/infiniop/ops/float_power/operator.cc @@ -9,7 +9,7 @@ #include "nvidia/float_power_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/float_power_metax.h" #endif @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateFloatPowerDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -94,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopGetFloatPowerWorkspaceSize(infiniopFloatPowerD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -139,6 +145,9 @@ __INFINI_C infiniStatus_t infiniopFloatPower( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -177,6 +186,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFloatPowerDescriptor(infiniopFloatPower #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/floor/operator.cc b/src/infiniop/ops/floor/operator.cc index 20a94f25c..009bd3f8a 100644 --- a/src/infiniop/ops/floor/operator.cc +++ b/src/infiniop/ops/floor/operator.cc @@ -11,7 +11,7 @@ #include "nvidia/floor_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/floor_metax.h" #endif @@ -58,6 +58,9 @@ __INFINI_C infiniStatus_t infiniopCreateFloorDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif // ========================================== // 添加 MOORE 分支 // ========================================== @@ -105,6 +108,9 @@ __INFINI_C infiniStatus_t infiniopGetFloorWorkspaceSize(infiniopFloorDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif @@ -154,6 +160,9 @@ __INFINI_C infiniStatus_t infiniopFloor( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif @@ -195,6 +204,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFloorDescriptor(infiniopFloorDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif diff --git a/src/infiniop/ops/floor_divide/operator.cc b/src/infiniop/ops/floor_divide/operator.cc index d2644e944..7d8b63b3c 100644 --- a/src/infiniop/ops/floor_divide/operator.cc +++ b/src/infiniop/ops/floor_divide/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/floor_divide_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/floor_divide_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -48,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopCreateFloorDivideDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -85,6 +88,9 @@ __INFINI_C infiniStatus_t infiniopGetFloorDivideWorkspaceSize(infiniopFloorDivid #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -130,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopFloorDivide( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -169,6 +178,9 @@ infiniopDestroyFloorDivideDescriptor(infiniopFloorDivideDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/fmin/operator.cc b/src/infiniop/ops/fmin/operator.cc index 4b41c93ae..f8f941c57 100644 --- a/src/infiniop/ops/fmin/operator.cc +++ b/src/infiniop/ops/fmin/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/fmin_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/fmin_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -48,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopCreateFminDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -85,6 +88,9 @@ __INFINI_C infiniStatus_t infiniopGetFminWorkspaceSize(infiniopFminDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -130,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopFmin( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -169,6 +178,9 @@ infiniopDestroyFminDescriptor(infiniopFminDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/fmod/operator.cc b/src/infiniop/ops/fmod/operator.cc index 100257d38..45128b236 100644 --- a/src/infiniop/ops/fmod/operator.cc +++ b/src/infiniop/ops/fmod/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/fmod_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/fmod_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreateFmodDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetFmodWorkspaceSize(infiniopFmodDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -116,6 +122,9 @@ __INFINI_C infiniStatus_t infiniopFmod( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -148,6 +157,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFmodDescriptor(infiniopFmodDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/fused_ffn/operator.cc b/src/infiniop/ops/fused_ffn/operator.cc index 616a1957f..e36a9ba87 100644 --- a/src/infiniop/ops/fused_ffn/operator.cc +++ b/src/infiniop/ops/fused_ffn/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/fused_ffn_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/fused_ffn_metax.cuh" #endif @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateFusedFFNDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_QY_API CREATE(INFINI_DEVICE_QY, nvidia); #endif @@ -89,6 +92,9 @@ __INFINI_C infiniStatus_t infiniopGetFusedFFNWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_QY_API GET(INFINI_DEVICE_QY, nvidia); #endif @@ -138,6 +144,9 @@ __INFINI_C infiniStatus_t infiniopFusedFFN( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_QY_API CALCULATE(INFINI_DEVICE_QY, nvidia); #endif @@ -177,6 +186,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFusedFFNDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_QY_API DESTROY(INFINI_DEVICE_QY, nvidia); #endif diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc b/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc index 14296e846..7756343c6 100644 --- a/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc +++ b/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc @@ -5,7 +5,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/fused_gated_delta_net_gating_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/fused_gated_delta_net_gating_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -50,6 +50,9 @@ infiniopCreateFusedGatedDeltaNetGatingDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -89,6 +92,9 @@ infiniopGetFusedGatedDeltaNetGatingWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -136,6 +142,9 @@ infiniopFusedGatedDeltaNetGating( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -174,6 +183,9 @@ infiniopDestroyFusedGatedDeltaNetGatingDescriptor( #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/gaussian_nll_loss/operator.cc b/src/infiniop/ops/gaussian_nll_loss/operator.cc index 063319894..64fdfa5d4 100644 --- a/src/infiniop/ops/gaussian_nll_loss/operator.cc +++ b/src/infiniop/ops/gaussian_nll_loss/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/gaussian_nll_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/gaussian_nll_loss_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateGaussianNllLossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -87,6 +90,9 @@ __INFINI_C infiniStatus_t infiniopGetGaussianNllLossWorkspaceSize(infiniopGaussi #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -130,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopGaussianNllLoss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -166,6 +175,9 @@ infiniopDestroyGaussianNllLossDescriptor(infiniopGaussianNllLossDescriptor_t des #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/gelu/operator.cc b/src/infiniop/ops/gelu/operator.cc index 69b590d9a..eff60fa29 100644 --- a/src/infiniop/ops/gelu/operator.cc +++ b/src/infiniop/ops/gelu/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/gelu_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/gelu_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateGeluDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -95,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -146,6 +152,9 @@ __INFINI_C infiniStatus_t infiniopGelu( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -191,6 +200,9 @@ infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/gelutanh/operator.cc b/src/infiniop/ops/gelutanh/operator.cc index 91a213a24..d7f489cd4 100644 --- a/src/infiniop/ops/gelutanh/operator.cc +++ b/src/infiniop/ops/gelutanh/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/gelutanh_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/gelutanh_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -48,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopCreateGeluTanhDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -84,6 +87,9 @@ __INFINI_C infiniStatus_t infiniopGetGeluTanhWorkspaceSize(infiniopGeluTanhDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -126,6 +132,9 @@ __INFINI_C infiniStatus_t infiniopGeluTanh( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -162,6 +171,9 @@ __INFINI_C infiniStatus_t infiniopDestroyGeluTanhDescriptor(infiniopGeluTanhDesc #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/gemm/operator.cc b/src/infiniop/ops/gemm/operator.cc index 81d9cb066..3a6059280 100644 --- a/src/infiniop/ops/gemm/operator.cc +++ b/src/infiniop/ops/gemm/operator.cc @@ -14,7 +14,7 @@ #ifdef ENABLE_ASCEND_API #include "ascend/gemm_ascend.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/gemm_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -71,6 +71,9 @@ __INFINI_C infiniStatus_t infiniopCreateGemmDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -125,6 +128,9 @@ infiniopGetGemmWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -186,6 +192,9 @@ __INFINI_C infiniStatus_t infiniopGemm( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -237,6 +246,9 @@ infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/hardswish/operator.cc b/src/infiniop/ops/hardswish/operator.cc index 918a6daf4..e7ca7cf05 100644 --- a/src/infiniop/ops/hardswish/operator.cc +++ b/src/infiniop/ops/hardswish/operator.cc @@ -11,7 +11,7 @@ #ifdef ENABLE_MOORE_API #include "moore/hardswish_moore.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/hardswish_metax.h" #endif @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateHardSwishDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetHardSwishWorkspaceSize(infiniopHardSwishDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif @@ -122,6 +128,9 @@ __INFINI_C infiniStatus_t infiniopHardSwish( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif @@ -157,6 +166,9 @@ __INFINI_C infiniStatus_t infiniopDestroyHardSwishDescriptor(infiniopHardSwishDe #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif diff --git a/src/infiniop/ops/hardtanh/operator.cc b/src/infiniop/ops/hardtanh/operator.cc index daf7581bb..0e07cf693 100644 --- a/src/infiniop/ops/hardtanh/operator.cc +++ b/src/infiniop/ops/hardtanh/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/hardtanh_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/hardtanh_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -47,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateHardTanhDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -81,6 +84,9 @@ __INFINI_C infiniStatus_t infiniopGetHardTanhWorkspaceSize(infiniopHardTanhDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -122,6 +128,9 @@ __INFINI_C infiniStatus_t infiniopHardTanh( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -158,6 +167,9 @@ infiniopDestroyHardTanhDescriptor(infiniopHardTanhDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/hinge_embedding_loss/operator.cc b/src/infiniop/ops/hinge_embedding_loss/operator.cc index 6ad907a41..018b95ae4 100644 --- a/src/infiniop/ops/hinge_embedding_loss/operator.cc +++ b/src/infiniop/ops/hinge_embedding_loss/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/hinge_embedding_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/hinge_embedding_loss_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateHingeEmbeddingLossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -83,6 +86,9 @@ __INFINI_C infiniStatus_t infiniopGetHingeEmbeddingLossWorkspaceSize(infiniopHin #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -125,6 +131,9 @@ __INFINI_C infiniStatus_t infiniopHingeEmbeddingLoss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -161,6 +170,9 @@ infiniopDestroyHingeEmbeddingLossDescriptor(infiniopHingeEmbeddingLossDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/histc/operator.cc b/src/infiniop/ops/histc/operator.cc index ed4d07662..bdde5e778 100644 --- a/src/infiniop/ops/histc/operator.cc +++ b/src/infiniop/ops/histc/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) #include "nvidia/histc_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/histc_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateHistcDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetHistcWorkspaceSize(infiniopHistcDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopHistc( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -151,6 +160,9 @@ infiniopDestroyHistcDescriptor(infiniopHistcDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/huber_loss/operator.cc b/src/infiniop/ops/huber_loss/operator.cc index beb0237de..2700d7415 100644 --- a/src/infiniop/ops/huber_loss/operator.cc +++ b/src/infiniop/ops/huber_loss/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/huber_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/huber_loss_metax.h" #endif @@ -59,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopCreateHuberLossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -97,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetHuberLossWorkspaceSize(infiniopHuberLossDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -142,6 +148,9 @@ __INFINI_C infiniStatus_t infiniopHuberLoss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -177,6 +186,9 @@ __INFINI_C infiniStatus_t infiniopDestroyHuberLossDescriptor(infiniopHuberLossDe #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/hypot/operator.cc b/src/infiniop/ops/hypot/operator.cc index 53f82b8ef..484d726ed 100644 --- a/src/infiniop/ops/hypot/operator.cc +++ b/src/infiniop/ops/hypot/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/hypot_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/hypot_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateHypotDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -87,6 +90,9 @@ __INFINI_C infiniStatus_t infiniopGetHypotWorkspaceSize(infiniopHypotDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia); #endif @@ -128,6 +134,9 @@ __INFINI_C infiniStatus_t infiniopHypot( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif @@ -163,6 +172,9 @@ __INFINI_C infiniStatus_t infiniopDestroyHypotDescriptor(infiniopHypotDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); #endif diff --git a/src/infiniop/ops/index_add/operator.cc b/src/infiniop/ops/index_add/operator.cc index f655628fa..97bd78b42 100644 --- a/src/infiniop/ops/index_add/operator.cc +++ b/src/infiniop/ops/index_add/operator.cc @@ -9,7 +9,7 @@ #include "nvidia/index_add_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/index_add_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateIndexAddDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -92,6 +95,9 @@ __INFINI_C infiniStatus_t infiniopGetIndexAddWorkspaceSize(infiniopIndexAddDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif // 【关键修复】启用 Moore 分支 #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); @@ -137,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopIndexAdd( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif // 【关键修复】启用 Moore 分支 #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); @@ -174,6 +183,9 @@ __INFINI_C infiniStatus_t infiniopDestroyIndexAddDescriptor(infiniopIndexAddDesc #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/index_copy/operator.cc b/src/infiniop/ops/index_copy/operator.cc index 97d078da4..4d097656f 100644 --- a/src/infiniop/ops/index_copy/operator.cc +++ b/src/infiniop/ops/index_copy/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/index_copy_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/index_copy_metax.h" #endif @@ -59,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopCreateIndexCopyDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -97,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetIndexCopyWorkspaceSize(infiniopIndexCopyDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -143,6 +149,9 @@ __INFINI_C infiniStatus_t infiniopIndexCopy( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -181,6 +190,9 @@ __INFINI_C infiniStatus_t infiniopDestroyIndexCopyDescriptor(infiniopIndexCopyDe #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/inner/operator.cc b/src/infiniop/ops/inner/operator.cc index 54ed677c2..50d438d21 100644 --- a/src/infiniop/ops/inner/operator.cc +++ b/src/infiniop/ops/inner/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/inner_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/inner_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -47,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateInnerDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetInnerWorkspaceSize(infiniopInnerDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -113,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopInner( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif @@ -143,6 +152,9 @@ __INFINI_C infiniStatus_t infiniopDestroyInnerDescriptor(infiniopInnerDescriptor #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/src/infiniop/ops/interpolate/operator.cc b/src/infiniop/ops/interpolate/operator.cc index 0197f8cbe..876bffd5d 100644 --- a/src/infiniop/ops/interpolate/operator.cc +++ b/src/infiniop/ops/interpolate/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/interpolate_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/interpolate_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -51,6 +51,9 @@ __INFINI_C infiniStatus_t infiniopCreateInterpolateDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -85,6 +88,9 @@ __INFINI_C infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolat #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -126,6 +132,9 @@ __INFINI_C infiniStatus_t infiniopInterpolate( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -162,6 +171,9 @@ infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/kron/operator.cc b/src/infiniop/ops/kron/operator.cc index f803f434a..ae528c34f 100644 --- a/src/infiniop/ops/kron/operator.cc +++ b/src/infiniop/ops/kron/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/kron_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/kron_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -45,6 +45,9 @@ __INFINI_C infiniStatus_t infiniopCreateKronDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -79,6 +82,9 @@ __INFINI_C infiniStatus_t infiniopGetKronWorkspaceSize(infiniopKronDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -121,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopKron( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -157,6 +166,9 @@ infiniopDestroyKronDescriptor(infiniopKronDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/kthvalue/operator.cc b/src/infiniop/ops/kthvalue/operator.cc index 094d0a1e9..dfd17fb01 100644 --- a/src/infiniop/ops/kthvalue/operator.cc +++ b/src/infiniop/ops/kthvalue/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/kthvalue_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/kthvalue_metax.h" #endif @@ -61,6 +61,9 @@ __INFINI_C infiniStatus_t infiniopCreateKthvalueDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -99,6 +102,9 @@ __INFINI_C infiniStatus_t infiniopGetKthvalueWorkspaceSize(infiniopKthvalueDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -144,6 +150,9 @@ __INFINI_C infiniStatus_t infiniopKthvalue( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -182,6 +191,9 @@ __INFINI_C infiniStatus_t infiniopDestroyKthvalueDescriptor(infiniopKthvalueDesc #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/kv_caching/operator.cc b/src/infiniop/ops/kv_caching/operator.cc index e83bb1240..cf56b2906 100644 --- a/src/infiniop/ops/kv_caching/operator.cc +++ b/src/infiniop/ops/kv_caching/operator.cc @@ -5,7 +5,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/kv_caching_nvidia.cuh" #endif -#if defined(ENABLE_METAX_API) +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/kv_caching_metax.h" #endif @@ -49,6 +49,9 @@ __INFINI_C infiniStatus_t infiniopCreateKVCachingDescriptor( #if defined(ENABLE_METAX_API) CREATE(INFINI_DEVICE_METAX, metax); #endif +#if defined(ENABLE_MARS_API) + CREATE(INFINI_DEVICE_MARS, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -87,6 +90,9 @@ __INFINI_C infiniStatus_t infiniopGetKVCachingWorkspaceSize( #if defined(ENABLE_METAX_API) GET_SIZE(INFINI_DEVICE_METAX, metax); #endif +#if defined(ENABLE_MARS_API) + GET_SIZE(INFINI_DEVICE_MARS, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -131,6 +137,9 @@ __INFINI_C infiniStatus_t infiniopKVCaching( #if defined(ENABLE_METAX_API) CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#if defined(ENABLE_MARS_API) + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -167,6 +176,9 @@ __INFINI_C infiniStatus_t infiniopDestroyKVCachingDescriptor( #if defined(ENABLE_METAX_API) DELETE(INFINI_DEVICE_METAX, metax); #endif +#if defined(ENABLE_MARS_API) + DELETE(INFINI_DEVICE_MARS, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/layer_norm/operator.cc b/src/infiniop/ops/layer_norm/operator.cc index 3925e845a..bac74e300 100644 --- a/src/infiniop/ops/layer_norm/operator.cc +++ b/src/infiniop/ops/layer_norm/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/layer_norm_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/layer_norm_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -64,6 +64,9 @@ __INFINI_C infiniStatus_t infiniopCreateLayerNormDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -106,6 +109,9 @@ __INFINI_C infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -168,6 +174,9 @@ __INFINI_C infiniStatus_t infiniopLayerNorm( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -210,6 +219,9 @@ infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/ldexp/operator.cc b/src/infiniop/ops/ldexp/operator.cc index 32e6074b8..fd6d3e77f 100644 --- a/src/infiniop/ops/ldexp/operator.cc +++ b/src/infiniop/ops/ldexp/operator.cc @@ -11,7 +11,7 @@ #include "nvidia/ldexp_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/ldexp_metax.h" #endif @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateLdexpDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -94,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopGetLdexpWorkspaceSize(infiniopLdexpDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -139,6 +145,9 @@ __INFINI_C infiniStatus_t infiniopLdexp( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -177,6 +186,9 @@ __INFINI_C infiniStatus_t infiniopDestroyLdexpDescriptor(infiniopLdexpDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/lerp/operator.cc b/src/infiniop/ops/lerp/operator.cc index 8c8ed13b3..2f9b84098 100644 --- a/src/infiniop/ops/lerp/operator.cc +++ b/src/infiniop/ops/lerp/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/lerp_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/lerp_metax.h" #endif @@ -59,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopCreateLerpDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -100,6 +103,9 @@ __INFINI_C infiniStatus_t infiniopGetLerpWorkspaceSize(infiniopLerpDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -149,6 +155,9 @@ __INFINI_C infiniStatus_t infiniopLerp( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -190,6 +199,9 @@ __INFINI_C infiniStatus_t infiniopDestroyLerpDescriptor(infiniopLerpDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/log10/operator.cc b/src/infiniop/ops/log10/operator.cc index 03f52a8e2..00f3bcd3f 100644 --- a/src/infiniop/ops/log10/operator.cc +++ b/src/infiniop/ops/log10/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/log10_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/log10_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateLog10Descriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetLog10WorkspaceSize(infiniopLog10Descriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopLog10( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyLog10Descriptor(infiniopLog10Descriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/log1p/operator.cc b/src/infiniop/ops/log1p/operator.cc index 2e68ce7ad..afc663611 100644 --- a/src/infiniop/ops/log1p/operator.cc +++ b/src/infiniop/ops/log1p/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/log1p_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/log1p_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateLog1pDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetLog1pWorkspaceSize(infiniopLog1pDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopLog1p( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyLog1pDescriptor(infiniopLog1pDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/log_softmax/operator.cc b/src/infiniop/ops/log_softmax/operator.cc index 894ce34dd..919d845ac 100644 --- a/src/infiniop/ops/log_softmax/operator.cc +++ b/src/infiniop/ops/log_softmax/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/log_softmax_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/log_softmax_metax.h" #endif @@ -55,6 +55,9 @@ __INFINI_C infiniStatus_t infiniopCreateLogSoftmaxDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -137,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopLogSoftmax( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -175,6 +184,9 @@ __INFINI_C infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmax #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/logaddexp/operator.cc b/src/infiniop/ops/logaddexp/operator.cc index c7ce3bc9f..43627b0ab 100644 --- a/src/infiniop/ops/logaddexp/operator.cc +++ b/src/infiniop/ops/logaddexp/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/logaddexp_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/logaddexp_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateLogAddExpDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -91,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDes #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -136,6 +142,9 @@ __INFINI_C infiniStatus_t infiniopLogAddExp( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -174,6 +183,9 @@ __INFINI_C infiniStatus_t infiniopDestroyLogAddExpDescriptor(infiniopLogAddExpDe #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/logaddexp2/operator.cc b/src/infiniop/ops/logaddexp2/operator.cc index 6962c574d..b3d7dd1d3 100644 --- a/src/infiniop/ops/logaddexp2/operator.cc +++ b/src/infiniop/ops/logaddexp2/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/logaddexp2_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/logaddexp2_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateLogAddExp2Descriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -91,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2D #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -136,6 +142,9 @@ __INFINI_C infiniStatus_t infiniopLogAddExp2( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -174,6 +183,9 @@ __INFINI_C infiniStatus_t infiniopDestroyLogAddExp2Descriptor(infiniopLogAddExp2 #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/logcumsumexp/operator.cc b/src/infiniop/ops/logcumsumexp/operator.cc index 9a699e905..0ecd56fbd 100644 --- a/src/infiniop/ops/logcumsumexp/operator.cc +++ b/src/infiniop/ops/logcumsumexp/operator.cc @@ -11,7 +11,7 @@ #include "nvidia/logcumsumexp_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/logcumsumexp_metax.h" #endif @@ -60,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopCreateLogCumSumExpDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -98,6 +101,9 @@ __INFINI_C infiniStatus_t infiniopGetLogCumSumExpWorkspaceSize(infiniopLogCumSum #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -142,6 +148,9 @@ __INFINI_C infiniStatus_t infiniopLogCumSumExp( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -180,6 +189,9 @@ __INFINI_C infiniStatus_t infiniopDestroyLogCumSumExpDescriptor(infiniopLogCumSu #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/logdet/operator.cc b/src/infiniop/ops/logdet/operator.cc index ae7a5804d..730635927 100644 --- a/src/infiniop/ops/logdet/operator.cc +++ b/src/infiniop/ops/logdet/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/logdet_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/logdet_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateLogdetDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetLogdetWorkspaceSize(infiniopLogdetDescripto #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopLogdet( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyLogdetDescriptor(infiniopLogdetDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/masked_select/operator.cc b/src/infiniop/ops/masked_select/operator.cc index f379565c1..e1188aca1 100644 --- a/src/infiniop/ops/masked_select/operator.cc +++ b/src/infiniop/ops/masked_select/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/masked_select_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/masked_select_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -45,6 +45,9 @@ __INFINI_C infiniStatus_t infiniopCreateMaskedSelectDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -75,6 +78,9 @@ __INFINI_C infiniStatus_t infiniopGetMaskedSelectWorkspaceSize(infiniopMaskedSel #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -112,6 +118,9 @@ __INFINI_C infiniStatus_t infiniopMaskedSelect( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif @@ -142,6 +151,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMaskedSelectDescriptor(infiniopMaskedSe #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/src/infiniop/ops/matrix_power/operator.cc b/src/infiniop/ops/matrix_power/operator.cc index f2933dcb7..1cf993784 100644 --- a/src/infiniop/ops/matrix_power/operator.cc +++ b/src/infiniop/ops/matrix_power/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/matrix_power_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/matrix_power_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -45,6 +45,9 @@ __INFINI_C __export infiniStatus_t infiniopCreateMatrixPowerDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -79,6 +82,9 @@ __INFINI_C __export infiniStatus_t infiniopGetMatrixPowerWorkspaceSize(infiniopM #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -120,6 +126,9 @@ __INFINI_C __export infiniStatus_t infiniopMatrixPower( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -156,6 +165,9 @@ infiniopDestroyMatrixPowerDescriptor(infiniopMatrixPowerDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/mul/operator.cc b/src/infiniop/ops/mul/operator.cc index 33f4e0c88..c8022fecc 100644 --- a/src/infiniop/ops/mul/operator.cc +++ b/src/infiniop/ops/mul/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/mul_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/mul_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -54,6 +54,9 @@ __INFINI_C infiniStatus_t infiniopCreateMulDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -97,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -149,6 +155,9 @@ __INFINI_C infiniStatus_t infiniopMul( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -194,6 +203,9 @@ infiniopDestroyMulDescriptor(infiniopMulDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/multi_margin_loss/operator.cc b/src/infiniop/ops/multi_margin_loss/operator.cc index 3c443ee30..e594ada56 100644 --- a/src/infiniop/ops/multi_margin_loss/operator.cc +++ b/src/infiniop/ops/multi_margin_loss/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/multi_margin_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/multi_margin_loss_metax.h" #endif @@ -63,6 +63,9 @@ __INFINI_C infiniStatus_t infiniopCreateMultiMarginLossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -101,6 +104,9 @@ __INFINI_C infiniStatus_t infiniopGetMultiMarginLossWorkspaceSize(infiniopMultiM #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -147,6 +153,9 @@ __INFINI_C infiniStatus_t infiniopMultiMarginLoss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -181,6 +190,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMultiMarginLossDescriptor(infiniopMulti #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/nrm2/operator.cc b/src/infiniop/ops/nrm2/operator.cc index d1ecfa2bb..3afb695b1 100644 --- a/src/infiniop/ops/nrm2/operator.cc +++ b/src/infiniop/ops/nrm2/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/nrm2_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/nrm2_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -32,6 +32,9 @@ __INFINI_C infiniStatus_t infiniopCreateNrm2Descriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -56,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopGetNrm2WorkspaceSize(infiniopNrm2Descriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -86,6 +92,9 @@ __INFINI_C infiniStatus_t infiniopNrm2( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -110,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopDestroyNrm2Descriptor(infiniopNrm2Descriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/ones/operator.cc b/src/infiniop/ops/ones/operator.cc index 934f03dd8..801348bbb 100644 --- a/src/infiniop/ops/ones/operator.cc +++ b/src/infiniop/ops/ones/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/ones_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/ones_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateOnesDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -86,6 +89,9 @@ __INFINI_C infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -134,6 +140,9 @@ __INFINI_C infiniStatus_t infiniopOnes( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -176,6 +185,9 @@ infiniopDestroyOnesDescriptor(infiniopOnesDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/pad/operator.cc b/src/infiniop/ops/pad/operator.cc index 246531317..33cf392d4 100644 --- a/src/infiniop/ops/pad/operator.cc +++ b/src/infiniop/ops/pad/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/pad_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/pad_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -60,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopCreatePadDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -100,6 +103,9 @@ __INFINI_C infiniStatus_t infiniopGetPadWorkspaceSize(infiniopPadDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -147,6 +153,9 @@ __INFINI_C infiniStatus_t infiniopPad( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -189,6 +198,9 @@ infiniopDestroyPadDescriptor(infiniopPadDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/paged_attention/operator.cc b/src/infiniop/ops/paged_attention/operator.cc index 3c3cce9f9..c8ea18198 100644 --- a/src/infiniop/ops/paged_attention/operator.cc +++ b/src/infiniop/ops/paged_attention/operator.cc @@ -8,7 +8,7 @@ #ifdef ENABLE_MOORE_API #include "moore/paged_attention_moore.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/paged_attention_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreatePagedAttentionDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia) #endif @@ -85,6 +88,9 @@ __INFINI_C infiniStatus_t infiniopGetPagedAttentionWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia) #endif @@ -128,6 +134,9 @@ __INFINI_C infiniStatus_t infiniopPagedAttention( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia) #endif @@ -166,6 +175,9 @@ __INFINI_C infiniStatus_t infiniopDestroyPagedAttentionDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia) #endif diff --git a/src/infiniop/ops/paged_attention_prefill/operator.cc b/src/infiniop/ops/paged_attention_prefill/operator.cc index bd06d36ca..066d75bc7 100644 --- a/src/infiniop/ops/paged_attention_prefill/operator.cc +++ b/src/infiniop/ops/paged_attention_prefill/operator.cc @@ -5,7 +5,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/paged_attention_prefill_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/paged_attention_prefill_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -48,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia) #endif @@ -87,6 +90,9 @@ __INFINI_C infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia) #endif @@ -133,6 +139,9 @@ __INFINI_C infiniStatus_t infiniopPagedAttentionPrefill( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia) #endif @@ -171,6 +180,9 @@ __INFINI_C infiniStatus_t infiniopDestroyPagedAttentionPrefillDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia) #endif diff --git a/src/infiniop/ops/paged_caching/operator.cc b/src/infiniop/ops/paged_caching/operator.cc index a32b97669..9d5f87bac 100644 --- a/src/infiniop/ops/paged_caching/operator.cc +++ b/src/infiniop/ops/paged_caching/operator.cc @@ -5,7 +5,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_QY_API) #include "nvidia/paged_caching_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/paged_caching_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -41,6 +41,9 @@ __INFINI_C infiniStatus_t infiniopCreatePagedCachingDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia) #endif @@ -83,6 +86,9 @@ __INFINI_C infiniStatus_t infiniopGetPagedCachingWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia) #endif @@ -129,6 +135,9 @@ __INFINI_C infiniStatus_t infiniopPagedCaching( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia) #endif @@ -170,6 +179,9 @@ __INFINI_C infiniStatus_t infiniopDestroyPagedCachingDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia) #endif diff --git a/src/infiniop/ops/pixel_shuffle/operator.cc b/src/infiniop/ops/pixel_shuffle/operator.cc index 2eb8f89e7..21e36e9a9 100644 --- a/src/infiniop/ops/pixel_shuffle/operator.cc +++ b/src/infiniop/ops/pixel_shuffle/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/pixel_shuffle_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/pixel_shuffle_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -45,6 +45,9 @@ __INFINI_C __export infiniStatus_t infiniopCreatePixelShuffleDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -79,6 +82,9 @@ __INFINI_C __export infiniStatus_t infiniopGetPixelShuffleWorkspaceSize(infiniop #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -120,6 +126,9 @@ __INFINI_C __export infiniStatus_t infiniopPixelShuffle( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -156,6 +165,9 @@ infiniopDestroyPixelShuffleDescriptor(infiniopPixelShuffleDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/prelu/operator.cc b/src/infiniop/ops/prelu/operator.cc index 7b87d0020..f108511e6 100644 --- a/src/infiniop/ops/prelu/operator.cc +++ b/src/infiniop/ops/prelu/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/prelu_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/prelu_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -44,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopCreatePreluDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -78,6 +81,9 @@ __INFINI_C infiniStatus_t infiniopGetPreluWorkspaceSize(infiniopPreluDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -120,6 +126,9 @@ __INFINI_C infiniStatus_t infiniopPrelu( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -156,6 +165,9 @@ infiniopDestroyPreluDescriptor(infiniopPreluDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/random_sample/operator.cc b/src/infiniop/ops/random_sample/operator.cc index b09e56b2a..58628d0f4 100644 --- a/src/infiniop/ops/random_sample/operator.cc +++ b/src/infiniop/ops/random_sample/operator.cc @@ -11,7 +11,7 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/random_sample_bang.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/random_sample_metax.h" #endif #ifdef ENABLE_ASCEND_API @@ -65,6 +65,9 @@ infiniopCreateRandomSampleDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ASCEND_API CREATE(INFINI_DEVICE_ASCEND, ascend); #endif @@ -119,6 +122,9 @@ __INFINI_C infiniStatus_t infiniopGetRandomSampleWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ASCEND_API GET(INFINI_DEVICE_ASCEND, ascend); #endif @@ -183,6 +189,9 @@ __INFINI_C infiniStatus_t infiniopRandomSample( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ASCEND_API CALCULATE(INFINI_DEVICE_ASCEND, ascend); #endif @@ -234,6 +243,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRandomSampleDescriptor( #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ASCEND_API DELETE(INFINI_DEVICE_ASCEND, ascend); #endif diff --git a/src/infiniop/ops/rearrange/operator.cc b/src/infiniop/ops/rearrange/operator.cc index f1828c807..6e0242141 100644 --- a/src/infiniop/ops/rearrange/operator.cc +++ b/src/infiniop/ops/rearrange/operator.cc @@ -14,7 +14,7 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/rearrange_bang.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rearrange_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -67,6 +67,9 @@ __INFINI_C infiniStatus_t infiniopCreateRearrangeDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -120,6 +123,9 @@ __INFINI_C infiniStatus_t infiniopRearrange( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -171,6 +177,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRearrangeDescriptor( #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/reciprocal/operator.cc b/src/infiniop/ops/reciprocal/operator.cc index a4730a44f..db6e848df 100644 --- a/src/infiniop/ops/reciprocal/operator.cc +++ b/src/infiniop/ops/reciprocal/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/reciprocal_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/reciprocal_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateReciprocalDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -86,6 +89,9 @@ __INFINI_C infiniStatus_t infiniopGetReciprocalWorkspaceSize(infiniopReciprocalD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -130,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopReciprocal( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -169,6 +178,9 @@ infiniopDestroyReciprocalDescriptor(infiniopReciprocalDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc index fc272fdf9..aeae19e1c 100644 --- a/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc +++ b/src/infiniop/ops/recurrent_gated_delta_rule/operator.cc @@ -7,7 +7,7 @@ #if defined(ENABLE_NVIDIA_API) #include "nvidia/recurrent_gated_delta_rule_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/recurrent_gated_delta_rule_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -72,6 +75,9 @@ __INFINI_C infiniStatus_t infiniopGetRecurrentGatedDeltaRuleWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -106,6 +112,9 @@ __INFINI_C infiniStatus_t infiniopRecurrentGatedDeltaRule( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore) #endif @@ -131,6 +140,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRecurrentGatedDeltaRuleDescriptor( #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/src/infiniop/ops/relu6/operator.cc b/src/infiniop/ops/relu6/operator.cc index 0767e904e..209cad391 100644 --- a/src/infiniop/ops/relu6/operator.cc +++ b/src/infiniop/ops/relu6/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/relu6_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/relu6_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateRelu6Descriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetRelu6WorkspaceSize(infiniopRelu6Descriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopRelu6( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroyRelu6Descriptor(infiniopRelu6Descriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/rms_norm/operator.cc b/src/infiniop/ops/rms_norm/operator.cc index 63bf27eb6..2bd8ca2c0 100644 --- a/src/infiniop/ops/rms_norm/operator.cc +++ b/src/infiniop/ops/rms_norm/operator.cc @@ -14,7 +14,7 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/rms_norm_bang.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rms_norm_metax.cuh" #endif #ifdef ENABLE_MOORE_API @@ -73,6 +73,9 @@ __INFINI_C infiniStatus_t infiniopCreateRMSNormDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -121,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopGetRMSNormWorkspaceSize(infiniopRMSNormDescrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -170,6 +176,9 @@ __INFINI_C infiniStatus_t infiniopRMSNorm(infiniopRMSNormDescriptor_t desc, void #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -218,6 +227,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRMSNormDescriptor(infiniopRMSNormDescri #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/rope/operator.cc b/src/infiniop/ops/rope/operator.cc index ad3196092..9afb88e2d 100644 --- a/src/infiniop/ops/rope/operator.cc +++ b/src/infiniop/ops/rope/operator.cc @@ -14,7 +14,7 @@ #ifdef ENABLE_CAMBRICON_API #include "bang/rope_bang.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rope_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -71,6 +71,9 @@ __INFINI_C infiniStatus_t infiniopCreateRoPEDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_ASCEND_API CREATE(INFINI_DEVICE_ASCEND, ascend); #endif @@ -119,6 +122,9 @@ __INFINI_C infiniStatus_t infiniopGetRoPEWorkspaceSize(infiniopRoPEDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -176,6 +182,9 @@ __INFINI_C infiniStatus_t infiniopRoPE( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -225,6 +234,9 @@ infiniopDestroyRoPEDescriptor(infiniopRoPEDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/rot/operator.cc b/src/infiniop/ops/rot/operator.cc index 6c1345d7a..238d4d9ac 100644 --- a/src/infiniop/ops/rot/operator.cc +++ b/src/infiniop/ops/rot/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/rot_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rot_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -37,6 +37,9 @@ __INFINI_C infiniStatus_t infiniopCreateRotDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -61,6 +64,9 @@ __INFINI_C infiniStatus_t infiniopGetRotWorkspaceSize(infiniopRotDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -93,6 +99,9 @@ __INFINI_C infiniStatus_t infiniopRot( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -117,6 +126,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRotDescriptor(infiniopRotDescriptor_t d #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/rotg/operator.cc b/src/infiniop/ops/rotg/operator.cc index 0fa83d664..31a1e01e2 100644 --- a/src/infiniop/ops/rotg/operator.cc +++ b/src/infiniop/ops/rotg/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/rotg_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rotg_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -34,6 +34,9 @@ __INFINI_C infiniStatus_t infiniopCreateRotgDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -58,6 +61,9 @@ __INFINI_C infiniStatus_t infiniopGetRotgWorkspaceSize(infiniopRotgDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -90,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopRotg( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -114,6 +123,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRotgDescriptor(infiniopRotgDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/rotm/operator.cc b/src/infiniop/ops/rotm/operator.cc index fcfe2060c..35e2860e6 100644 --- a/src/infiniop/ops/rotm/operator.cc +++ b/src/infiniop/ops/rotm/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/rotm_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rotm_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -34,6 +34,9 @@ __INFINI_C infiniStatus_t infiniopCreateRotmDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopGetRotmWorkspaceSize(infiniopRotmDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -91,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopRotm( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -117,6 +126,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRotmDescriptor(infiniopRotmDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/rotmg/operator.cc b/src/infiniop/ops/rotmg/operator.cc index 98c25bd5b..54aa474cc 100644 --- a/src/infiniop/ops/rotmg/operator.cc +++ b/src/infiniop/ops/rotmg/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/rotmg_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/rotmg_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -35,6 +35,9 @@ __INFINI_C infiniStatus_t infiniopCreateRotmgDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopGetRotmgWorkspaceSize(infiniopRotmgDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -92,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopRotmg( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -116,6 +125,9 @@ __INFINI_C infiniStatus_t infiniopDestroyRotmgDescriptor(infiniopRotmgDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/scal/operator.cc b/src/infiniop/ops/scal/operator.cc index 33fa8618a..451cfb11d 100644 --- a/src/infiniop/ops/scal/operator.cc +++ b/src/infiniop/ops/scal/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/scal_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/scal_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -32,6 +32,9 @@ __INFINI_C infiniStatus_t infiniopCreateScalDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -56,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopGetScalWorkspaceSize(infiniopScalDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -86,6 +92,9 @@ __INFINI_C infiniStatus_t infiniopScal( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -110,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopDestroyScalDescriptor(infiniopScalDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/scatter/operator.cc b/src/infiniop/ops/scatter/operator.cc index 30fa1cf1f..23fb0b755 100644 --- a/src/infiniop/ops/scatter/operator.cc +++ b/src/infiniop/ops/scatter/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/scatter_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/scatter_metax.h" #endif @@ -61,6 +61,9 @@ __INFINI_C infiniStatus_t infiniopCreateScatterDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -99,6 +102,9 @@ __INFINI_C infiniStatus_t infiniopGetScatterWorkspaceSize(infiniopScatterDescrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -145,6 +151,9 @@ __INFINI_C infiniStatus_t infiniopScatter( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -183,6 +192,9 @@ __INFINI_C infiniStatus_t infiniopDestroyScatterDescriptor(infiniopScatterDescri #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/selu/operator.cc b/src/infiniop/ops/selu/operator.cc index 3558b04d5..b2f3ef6a3 100644 --- a/src/infiniop/ops/selu/operator.cc +++ b/src/infiniop/ops/selu/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/selu_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/selu_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateSeluDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetSeluWorkspaceSize(infiniopSeluDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopSelu( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroySeluDescriptor(infiniopSeluDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/silu/operator.cc b/src/infiniop/ops/silu/operator.cc index 602d5d178..f41164948 100644 --- a/src/infiniop/ops/silu/operator.cc +++ b/src/infiniop/ops/silu/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/silu_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/silu_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateSiluDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetSiluWorkspaceSize(infiniopSiluDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -125,6 +131,9 @@ __INFINI_C infiniStatus_t infiniopSilu( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -164,6 +173,9 @@ infiniopDestroySiluDescriptor(infiniopSiluDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/sinh/operator.cc b/src/infiniop/ops/sinh/operator.cc index 9df60758f..6b5889a97 100644 --- a/src/infiniop/ops/sinh/operator.cc +++ b/src/infiniop/ops/sinh/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/sinh_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/sinh_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -43,6 +43,9 @@ __INFINI_C infiniStatus_t infiniopCreateSinhDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -77,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetSinhWorkspaceSize(infiniopSinhDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax) #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -118,6 +124,9 @@ __INFINI_C infiniStatus_t infiniopSinh( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -154,6 +163,9 @@ infiniopDestroySinhDescriptor(infiniopSinhDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/smooth_l1_loss/operator.cc b/src/infiniop/ops/smooth_l1_loss/operator.cc index 75a66fdc5..10c98d058 100644 --- a/src/infiniop/ops/smooth_l1_loss/operator.cc +++ b/src/infiniop/ops/smooth_l1_loss/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/smooth_l1_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/smooth_l1_loss_metax.h" #endif @@ -59,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopCreateSmoothL1LossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -97,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetSmoothL1LossWorkspaceSize(infiniopSmoothL1L #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -142,6 +148,9 @@ __INFINI_C infiniStatus_t infiniopSmoothL1Loss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -180,6 +189,9 @@ __INFINI_C infiniStatus_t infiniopDestroySmoothL1LossDescriptor(infiniopSmoothL1 #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/softplus/operator.cc b/src/infiniop/ops/softplus/operator.cc index edaa5f546..e42230aa2 100644 --- a/src/infiniop/ops/softplus/operator.cc +++ b/src/infiniop/ops/softplus/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/softplus_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/softplus_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateSoftplusDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -95,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopGetSoftplusWorkspaceSize(infiniopSoftplusDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -145,6 +151,9 @@ __INFINI_C infiniStatus_t infiniopSoftplus( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -189,6 +198,9 @@ infiniopDestroySoftplusDescriptor(infiniopSoftplusDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/softsign/operator.cc b/src/infiniop/ops/softsign/operator.cc index e11de375b..8ae2cc74c 100644 --- a/src/infiniop/ops/softsign/operator.cc +++ b/src/infiniop/ops/softsign/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/softsign_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/softsign_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -65,6 +65,9 @@ __INFINI_C infiniStatus_t infiniopCreateSoftsignDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -106,6 +109,9 @@ __INFINI_C infiniStatus_t infiniopGetSoftsignWorkspaceSize(infiniopSoftsignDescr #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -155,6 +161,9 @@ __INFINI_C infiniStatus_t infiniopSoftsign( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -198,6 +207,9 @@ infiniopDestroySoftsignDescriptor(infiniopSoftsignDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/sub/operator.cc b/src/infiniop/ops/sub/operator.cc index 60857b7aa..6fd0ed3e0 100644 --- a/src/infiniop/ops/sub/operator.cc +++ b/src/infiniop/ops/sub/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/sub_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/sub_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -48,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopCreateSubDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -88,6 +91,9 @@ __INFINI_C infiniStatus_t infiniopGetSubWorkspaceSize(infiniopSubDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -137,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopSub( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -179,6 +188,9 @@ infiniopDestroySubDescriptor(infiniopSubDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DELETE(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/sum/operator.cc b/src/infiniop/ops/sum/operator.cc index 9eb9c25f4..cb68bb094 100644 --- a/src/infiniop/ops/sum/operator.cc +++ b/src/infiniop/ops/sum/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/sum_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/sum_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -56,6 +56,9 @@ __INFINI_C infiniStatus_t infiniopCreateSumDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetSumWorkspaceSize(infiniopSumDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -141,6 +147,9 @@ __INFINI_C infiniStatus_t infiniopSum( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -180,6 +189,9 @@ infiniopDestroySumDescriptor(infiniopSumDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/swap/operator.cc b/src/infiniop/ops/swap/operator.cc index 22d688021..2269390be 100644 --- a/src/infiniop/ops/swap/operator.cc +++ b/src/infiniop/ops/swap/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/swap_cpu.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/swap_metax.h" #endif #ifdef ENABLE_CAMBRICON_API @@ -32,6 +32,9 @@ __INFINI_C infiniStatus_t infiniopCreateSwapDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -56,6 +59,9 @@ __INFINI_C infiniStatus_t infiniopGetSwapWorkspaceSize(infiniopSwapDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -86,6 +92,9 @@ __INFINI_C infiniStatus_t infiniopSwap( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -110,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopDestroySwapDescriptor(infiniopSwapDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/swiglu/operator.cc b/src/infiniop/ops/swiglu/operator.cc index b43cd47dd..0f1cc2706 100644 --- a/src/infiniop/ops/swiglu/operator.cc +++ b/src/infiniop/ops/swiglu/operator.cc @@ -11,7 +11,7 @@ #ifdef ENABLE_KUNLUN_API #include "kunlun/swiglu_kunlun.h" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #if defined(ENABLE_NINETOOTHED) #include "ninetoothed/swiglu.h" #else @@ -87,6 +87,9 @@ __INFINI_C infiniStatus_t infiniopCreateSwiGLUDescriptor( CREATE_CUDA(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API + CREATE_CUDA(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -151,6 +154,9 @@ __INFINI_C infiniStatus_t infiniopGetSwiGLUWorkspaceSize(infiniopSwiGLUDescripto GET_CUDA(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API + GET_CUDA(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -221,6 +227,9 @@ __INFINI_C infiniStatus_t infiniopSwiGLU( CALCULATE_CUDA(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API + CALCULATE_CUDA(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif @@ -286,6 +295,9 @@ infiniopDestroySwiGLUDescriptor(infiniopSwiGLUDescriptor_t desc) { DELETE_CUDA(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API + DELETE_CUDA(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif diff --git a/src/infiniop/ops/take/operator.cc b/src/infiniop/ops/take/operator.cc index 90eb24660..15de317fa 100644 --- a/src/infiniop/ops/take/operator.cc +++ b/src/infiniop/ops/take/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/take_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/take_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateTakeDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -91,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopGetTakeWorkspaceSize(infiniopTakeDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -135,6 +141,9 @@ __INFINI_C infiniStatus_t infiniopTake( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -173,6 +182,9 @@ __INFINI_C infiniStatus_t infiniopDestroyTakeDescriptor(infiniopTakeDescriptor_t #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/tan/operator.cc b/src/infiniop/ops/tan/operator.cc index a0ac17f16..67ee8f280 100644 --- a/src/infiniop/ops/tan/operator.cc +++ b/src/infiniop/ops/tan/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/tan_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/tan_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateTanDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetTanWorkspaceSize(infiniopTanDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -121,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopTan( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -157,6 +166,9 @@ infiniopDestroyTanDescriptor(infiniopTanDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/tanhshrink/operator.cc b/src/infiniop/ops/tanhshrink/operator.cc index a153c4feb..8ec7a0507 100644 --- a/src/infiniop/ops/tanhshrink/operator.cc +++ b/src/infiniop/ops/tanhshrink/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/tanhshrink_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/tanhshrink_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateTanhshrinkDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +83,9 @@ __INFINI_C infiniStatus_t infiniopGetTanhshrinkWorkspaceSize(infiniopTanhshrinkD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -121,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopTanhshrink( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -157,6 +166,9 @@ infiniopDestroyTanhshrinkDescriptor(infiniopTanhshrinkDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/topk/operator.cc b/src/infiniop/ops/topk/operator.cc index cd7f17291..d37f207d2 100644 --- a/src/infiniop/ops/topk/operator.cc +++ b/src/infiniop/ops/topk/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/topk_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/topk_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -60,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopCreateTopKDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -97,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetTopKWorkspaceSize(infiniopTopKDescriptor_t #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -147,6 +153,9 @@ __INFINI_C infiniStatus_t infiniopTopK( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -186,6 +195,9 @@ infiniopDestroyTopKDescriptor(infiniopTopKDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/topkrouter/operator.cc b/src/infiniop/ops/topkrouter/operator.cc index bf9f2e586..3e137d7c2 100644 --- a/src/infiniop/ops/topkrouter/operator.cc +++ b/src/infiniop/ops/topkrouter/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/topkrouter_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/topkrouter_metax.h" #endif #ifdef ENABLE_KUNLUN_API @@ -36,6 +36,9 @@ __INFINI_C infiniStatus_t infiniopCreateTopkrouterDescriptor(infiniopHandle_t ha #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CREATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -71,6 +74,9 @@ __INFINI_C infiniStatus_t infiniopGetTopkrouterWorkspaceSize(infiniopTopkrouterD #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API GET(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -109,6 +115,9 @@ __INFINI_C infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); #endif @@ -144,6 +153,9 @@ __INFINI_C infiniStatus_t infiniopDestroyTopkrouterDescriptor(infiniopTopkrouter #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_KUNLUN_API DESTROY(INFINI_DEVICE_KUNLUN, kunlun); #endif diff --git a/src/infiniop/ops/topksoftmax/operator.cc b/src/infiniop/ops/topksoftmax/operator.cc index 56bab41bc..d30e0f941 100644 --- a/src/infiniop/ops/topksoftmax/operator.cc +++ b/src/infiniop/ops/topksoftmax/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/topksoftmax_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/topksoftmax_metax.cuh" #endif #ifdef ENABLE_MOORE_API @@ -37,6 +37,9 @@ __INFINI_C infiniStatus_t infiniopCreateTopksoftmaxDescriptor(infiniopHandle_t h #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -76,6 +79,9 @@ __INFINI_C infiniStatus_t infiniopGetTopksoftmaxWorkspaceSize(infiniopTopksoftma #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -120,6 +126,9 @@ __INFINI_C infiniStatus_t infiniopTopksoftmax(infiniopTopksoftmaxDescriptor_t de #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -159,6 +168,9 @@ __INFINI_C infiniStatus_t infiniopDestroyTopksoftmaxDescriptor(infiniopTopksoftm #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DESTROY(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/triplet_margin_loss/operator.cc b/src/infiniop/ops/triplet_margin_loss/operator.cc index 16ce25a83..69539a41d 100644 --- a/src/infiniop/ops/triplet_margin_loss/operator.cc +++ b/src/infiniop/ops/triplet_margin_loss/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/triplet_margin_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/triplet_margin_loss_metax.h" #endif @@ -67,6 +67,9 @@ __INFINI_C infiniStatus_t infiniopCreateTripletMarginLossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -105,6 +108,9 @@ __INFINI_C infiniStatus_t infiniopGetTripletMarginLossWorkspaceSize(infiniopTrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -151,6 +157,9 @@ __INFINI_C infiniStatus_t infiniopTripletMarginLoss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -189,6 +198,9 @@ __INFINI_C infiniStatus_t infiniopDestroyTripletMarginLossDescriptor(infiniopTri #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc b/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc index 22cbdb39d..507c883ce 100644 --- a/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/triplet_margin_with_distance_loss_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/triplet_margin_with_distance_loss_metax.h" #endif @@ -63,6 +63,9 @@ __INFINI_C infiniStatus_t infiniopCreateTripletMarginWithDistanceLossDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -103,6 +106,9 @@ __INFINI_C infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize( #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -149,6 +155,9 @@ __INFINI_C infiniStatus_t infiniopTripletMarginWithDistanceLoss( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -188,6 +197,9 @@ __INFINI_C infiniStatus_t infiniopDestroyTripletMarginWithDistanceLossDescriptor #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/unfold/operator.cc b/src/infiniop/ops/unfold/operator.cc index c78b46753..0d1ab0870 100644 --- a/src/infiniop/ops/unfold/operator.cc +++ b/src/infiniop/ops/unfold/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/unfold_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/unfold_metax.h" #endif @@ -62,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopCreateUnfoldDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -100,6 +103,9 @@ __INFINI_C infiniStatus_t infiniopGetUnfoldWorkspaceSize(infiniopUnfoldDescripto #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -144,6 +150,9 @@ __INFINI_C infiniStatus_t infiniopUnfold( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -182,6 +191,9 @@ __INFINI_C infiniStatus_t infiniopDestroyUnfoldDescriptor(infiniopUnfoldDescript #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/upsample_bilinear/operator.cc b/src/infiniop/ops/upsample_bilinear/operator.cc index eb03e6a21..ec64b91c8 100644 --- a/src/infiniop/ops/upsample_bilinear/operator.cc +++ b/src/infiniop/ops/upsample_bilinear/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/upsample_bilinear_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/upsample_bilinear_metax.h" #endif @@ -55,6 +55,9 @@ __INFINI_C infiniStatus_t infiniopCreateUpsampleBilinearDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -93,6 +96,9 @@ __INFINI_C infiniStatus_t infiniopGetUpsampleBilinearWorkspaceSize(infiniopUpsam #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -137,6 +143,9 @@ __INFINI_C infiniStatus_t infiniopUpsampleBilinear( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -175,6 +184,9 @@ __INFINI_C infiniStatus_t infiniopDestroyUpsampleBilinearDescriptor(infiniopUpsa #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/upsample_nearest/operator.cc b/src/infiniop/ops/upsample_nearest/operator.cc index bad09b892..98cb3ac50 100644 --- a/src/infiniop/ops/upsample_nearest/operator.cc +++ b/src/infiniop/ops/upsample_nearest/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/upsample_nearest_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/upsample_nearest_metax.h" #endif @@ -53,6 +53,9 @@ __INFINI_C infiniStatus_t infiniopCreateUpsampleNearestDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -91,6 +94,9 @@ __INFINI_C infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsamp #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -135,6 +141,9 @@ __INFINI_C infiniStatus_t infiniopUpsampleNearest( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -173,6 +182,9 @@ __INFINI_C infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsam #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/vander/operator.cc b/src/infiniop/ops/vander/operator.cc index 6c9b33deb..fcbb4f3fa 100644 --- a/src/infiniop/ops/vander/operator.cc +++ b/src/infiniop/ops/vander/operator.cc @@ -10,7 +10,7 @@ #include "nvidia/vander_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/vander_metax.h" #endif @@ -58,6 +58,9 @@ __INFINI_C infiniStatus_t infiniopCreateVanderDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -96,6 +99,9 @@ __INFINI_C infiniStatus_t infiniopGetVanderWorkspaceSize(infiniopVanderDescripto #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -140,6 +146,9 @@ __INFINI_C infiniStatus_t infiniopVander( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -178,6 +187,9 @@ __INFINI_C infiniStatus_t infiniopDestroyVanderDescriptor(infiniopVanderDescript #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/var/operator.cc b/src/infiniop/ops/var/operator.cc index 9bccebb39..91e45af13 100644 --- a/src/infiniop/ops/var/operator.cc +++ b/src/infiniop/ops/var/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/var_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/var_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -58,6 +58,9 @@ __INFINI_C infiniStatus_t infiniopCreateVarDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -95,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopGetVarWorkspaceSize(infiniopVarDescriptor_t de #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -144,6 +150,9 @@ __INFINI_C infiniStatus_t infiniopVar( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -183,6 +192,9 @@ infiniopDestroyVarDescriptor(infiniopVarDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/var_mean/operator.cc b/src/infiniop/ops/var_mean/operator.cc index 0fd3b7142..f26cd9b64 100644 --- a/src/infiniop/ops/var_mean/operator.cc +++ b/src/infiniop/ops/var_mean/operator.cc @@ -9,7 +9,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/var_mean_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/var_mean_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -60,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopCreateVarMeanDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -97,6 +100,9 @@ __INFINI_C infiniStatus_t infiniopGetVarMeanWorkspaceSize(infiniopVarMeanDescrip #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -147,6 +153,9 @@ __INFINI_C infiniStatus_t infiniopVarMean( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -186,6 +195,9 @@ infiniopDestroyVarMeanDescriptor(infiniopVarMeanDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/src/infiniop/ops/zeros/operator.cc b/src/infiniop/ops/zeros/operator.cc index 7f71e6f69..f831ff19b 100644 --- a/src/infiniop/ops/zeros/operator.cc +++ b/src/infiniop/ops/zeros/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/zeros_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #include "metax/zeros_metax.h" #endif #ifdef ENABLE_MOORE_API @@ -52,6 +52,9 @@ __INFINI_C infiniStatus_t infiniopCreateZerosDescriptor( #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CREATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -94,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopGetZerosWorkspaceSize(infiniopZerosDescriptor_ #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + GET(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -144,6 +150,9 @@ __INFINI_C infiniStatus_t infiniopZeros( #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -188,6 +197,9 @@ infiniopDestroyZerosDescriptor(infiniopZerosDescriptor_t desc) { #ifdef ENABLE_METAX_API DELETE(INFINI_DEVICE_METAX, metax); #endif +#ifdef ENABLE_MARS_API + DELETE(INFINI_DEVICE_MARS, metax); +#endif #ifdef ENABLE_MOORE_API DELETE(INFINI_DEVICE_MOORE, moore); #endif diff --git a/test/infinicore/framework/config.py b/test/infinicore/framework/config.py index 058fda90c..8e43e2423 100644 --- a/test/infinicore/framework/config.py +++ b/test/infinicore/framework/config.py @@ -20,6 +20,7 @@ def get_supported_hardware_platforms(): ("--ascend", "Huawei Ascend NPUs (requires torch_npu)"), ("--iluvatar", "Iluvatar GPUs"), ("--metax", "Metax GPUs"), + ("--mars", "Mars GPUs with HPCC support"), ("--moore", "Moore Threads GPUs (requires torch_musa)"), ("--kunlun", "Kunlun XPUs (requires torch_xmlir)"), ("--hygon", "Hygon DCUs"), @@ -195,6 +196,14 @@ def get_test_devices(args): except ImportError: print("Warning: Metax GPU support not available") + if args.mars: + try: + import torch + + devices_to_test.append(InfiniDeviceEnum.MARS) + except ImportError: + print("Warning: Mars GPU support not available") + if args.moore: try: import torch diff --git a/test/infinicore/framework/devices.py b/test/infinicore/framework/devices.py index 59b0ad558..69ca0e829 100644 --- a/test/infinicore/framework/devices.py +++ b/test/infinicore/framework/devices.py @@ -10,6 +10,7 @@ class InfiniDeviceEnum: HYGON = 8 QY = 9 ALI = 10 + MARS = 11 InfiniDeviceNames = { @@ -24,6 +25,7 @@ class InfiniDeviceEnum: InfiniDeviceEnum.KUNLUN: "Kunlun", InfiniDeviceEnum.HYGON: "Hygon", InfiniDeviceEnum.ALI: "Ali", + InfiniDeviceEnum.MARS: "Mars", } torch_device_map = { @@ -38,4 +40,5 @@ class InfiniDeviceEnum: InfiniDeviceEnum.HYGON: "cuda", InfiniDeviceEnum.QY: "cuda", InfiniDeviceEnum.ALI: "cuda", + InfiniDeviceEnum.MARS: "cuda", } diff --git a/test/infiniop/libinfiniop/devices.py b/test/infiniop/libinfiniop/devices.py index bb776e242..3f410dd52 100644 --- a/test/infiniop/libinfiniop/devices.py +++ b/test/infiniop/libinfiniop/devices.py @@ -10,6 +10,7 @@ class InfiniDeviceEnum: HYGON = 8 QY = 9 ALI = 10 + MARS = 11 InfiniDeviceNames = { @@ -24,6 +25,7 @@ class InfiniDeviceEnum: InfiniDeviceEnum.HYGON: "Hygon", InfiniDeviceEnum.QY: "QY", InfiniDeviceEnum.ALI: "Ali", + InfiniDeviceEnum.MARS: "Mars", } # Mapping that maps InfiniDeviceEnum to torch device string @@ -39,4 +41,5 @@ class InfiniDeviceEnum: InfiniDeviceEnum.HYGON: "cuda", InfiniDeviceEnum.QY: "cuda", InfiniDeviceEnum.ALI: "cuda", + InfiniDeviceEnum.MARS: "cuda", } diff --git a/test/infiniop/libinfiniop/utils.py b/test/infiniop/libinfiniop/utils.py index 690425b54..a064af435 100644 --- a/test/infiniop/libinfiniop/utils.py +++ b/test/infiniop/libinfiniop/utils.py @@ -449,6 +449,11 @@ def get_args(): action="store_true", help="Run METAX GPU test", ) + parser.add_argument( + "--mars", + action="store_true", + help="Run Mars GPU test", + ) parser.add_argument( "--moore", action="store_true", @@ -520,6 +525,7 @@ def filter_tensor_dtypes_by_device(device, tensor_dtypes): InfiniDeviceEnum.CPU, InfiniDeviceEnum.NVIDIA, InfiniDeviceEnum.METAX, + InfiniDeviceEnum.MARS, InfiniDeviceEnum.ASCEND, InfiniDeviceEnum.ILUVATAR, InfiniDeviceEnum.CAMBRICON, @@ -781,6 +787,10 @@ def get_test_devices(args): import torch devices_to_test.append(InfiniDeviceEnum.METAX) + if args.mars: + import torch + + devices_to_test.append(InfiniDeviceEnum.MARS) if args.moore: import torch import torch_musa diff --git a/xmake.lua b/xmake.lua index ffc44ed58..407ae8742 100644 --- a/xmake.lua +++ b/xmake.lua @@ -149,25 +149,26 @@ option("metax-gpu") set_description("Whether to compile implementations for MetaX GPU") option_end() -option("use-mc") +option("mars-gpu") set_default(false) set_showmenu(true) - set_description("Use MC version") + set_description("Whether to compile implementations for Mars GPU with HPCC") option_end() if has_config("metax-gpu") then - add_defines("ENABLE_METAX_API") - if has_config("use-mc") then - add_defines("ENABLE_METAX_MC_API") - -- MACA torch build expects USE_MACA for ATen headers (e.g. C10_WARP_SIZE). - add_defines("USE_MACA") - else - -- HPCC torch build expects this for ATen headers on hpcc. - add_defines("USE_HPCC") - end + add_defines("ENABLE_METAX_API", "ENABLE_METAX_MC_API", "USE_MACA") includes("xmake/metax.lua") end +if has_config("mars-gpu") then + add_defines("ENABLE_MARS_API", "USE_HPCC") + includes("xmake/mars.lua") +end + +if has_config("metax-gpu") and has_config("mars-gpu") then + raise("--metax-gpu and --mars-gpu are separate backends and cannot be enabled together") +end + -- 摩尔线程 option("moore-gpu") set_default(false) @@ -248,7 +249,7 @@ if has_config("aten") then add_defines("_GLIBCXX_USE_CXX11_ABI=0") end if get_config("flash-attn") and get_config("flash-attn") ~= "" - and (has_config("nv-gpu") or has_config("metax-gpu") or has_config("qy-gpu") or has_config("hygon-dcu")) then + and (has_config("nv-gpu") or has_config("metax-gpu") or has_config("mars-gpu") or has_config("qy-gpu") or has_config("hygon-dcu")) then add_defines("ENABLE_FLASH_ATTN") end end @@ -372,6 +373,13 @@ local function add_external_infinirt() add_includedirs(maca_include, { public = true }) end end + if has_config("mars-gpu") then + local hpcc_root = os.getenv("HPCC_PATH") or "/opt/hpcc" + local hpcc_include = path.join(hpcc_root, "include") + if os.isdir(hpcc_include) then + add_includedirs(hpcc_include, { public = true }) + end + end add_links("infinirt") end @@ -419,10 +427,11 @@ local function get_infiniops_backend_cmake_arg() end add_backend("nv-gpu", "-DWITH_NVIDIA=ON") add_backend("metax-gpu", "-DWITH_METAX=ON") + add_backend("mars-gpu", "-DWITH_MARS=ON") add_backend("iluvatar-gpu", "-DWITH_ILUVATAR=ON") add_backend("moore-gpu", "-DWITH_MOORE=ON") if #enabled == 0 then - raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --iluvatar-gpu, or --moore-gpu") + raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --mars-gpu, --iluvatar-gpu, or --moore-gpu") end if #enabled > 1 then raise("InfiniOps can build only one GPU backend at a time") @@ -569,6 +578,9 @@ target("infiniop") if has_config("metax-gpu") then add_deps("infiniop-metax") end + if has_config("mars-gpu") then + add_deps("infiniop-mars") + end if has_config("moore-gpu") then add_deps("infiniop-moore") end @@ -606,6 +618,9 @@ target("infiniccl") if has_config("metax-gpu") then add_deps("infiniccl-metax") end + if has_config("mars-gpu") then + add_deps("infiniccl-mars") + end if has_config("iluvatar-gpu") then add_deps("infiniccl-iluvatar") end @@ -660,7 +675,7 @@ target("infinicore_cpp_api") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") add_includedirs("include") - if has_config("metax-gpu") and has_config("use-mc") and has_config("aten") then + if has_config("metax-gpu") and has_config("aten") then local maca_root = os.getenv("MACA_PATH") or os.getenv("MACA_HOME") or os.getenv("MACA_ROOT") or "/opt/maca" add_includedirs(maca_root .. "/include") add_includedirs(maca_root .. "/tools/cu-bridge/include") @@ -719,6 +734,9 @@ target("infinicore_cpp_api") if has_config("metax-gpu") then add_deps("flash-attn-metax") end + if has_config("mars-gpu") then + add_deps("flash-attn-mars") + end if has_config("qy-gpu") then add_deps("flash-attn-qy") end @@ -740,14 +758,10 @@ target("infinicore_cpp_api") end before_build(function (target) - -- MetaX + flash-attn: `flash_attn_2_cuda` may use a different `mha_fwd_kvcache` ABI - -- depending on the underlying stack version. When building with MACA (`--use-mc=y`), - -- the version file is typically `/opt/maca/Version.txt` (HPCC uses `/opt/hpcc/Version.txt`). - if has_config("metax-gpu") and get_config("flash-attn") and get_config("flash-attn") ~= "" then - local version_txt = "/opt/hpcc/Version.txt" - if not os.isfile(version_txt) and has_config("use-mc") then - version_txt = "/opt/maca/Version.txt" - end + -- Mars HPCC flash-attn uses the HPCC major version to select its ABI. + if has_config("mars-gpu") and get_config("flash-attn") and get_config("flash-attn") ~= "" then + local hpcc_root = os.getenv("HPCC_PATH") or "/opt/hpcc" + local version_txt = path.join(hpcc_root, "Version.txt") if os.isfile(version_txt) then local content = os.iorunv("cat", {version_txt}) or "" content = content:trim() diff --git a/xmake/mars.lua b/xmake/mars.lua new file mode 100644 index 000000000..c85c3d1c9 --- /dev/null +++ b/xmake/mars.lua @@ -0,0 +1,127 @@ +local HPCC_ROOT = os.getenv("HPCC_PATH") or os.getenv("HPCC_HOME") or "/opt/hpcc" +local FLASH_ATTN_ROOT = get_config("flash-attn") + +local FLASH_ATTN_MARS_CUDA_SO_CONTAINER_DEFAULT = + "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so" + +local function mars_flash_attn_cuda_so_path() + local env_path = os.getenv("FLASH_ATTN_2_CUDA_SO") + if env_path and env_path ~= "" then + env_path = env_path:trim() + if os.isfile(env_path) then + return env_path + end + print(string.format("warning: mars+flash-attn: FLASH_ATTN_2_CUDA_SO is not a file: %s", env_path)) + end + + local container_path = os.getenv("FLASH_ATTN_MARS_CUDA_SO_CONTAINER") + if not container_path or container_path == "" then + container_path = FLASH_ATTN_MARS_CUDA_SO_CONTAINER_DEFAULT + end + return container_path +end + +target("infinicore_cpp_api") + if FLASH_ATTN_ROOT and FLASH_ATTN_ROOT ~= "" then + before_link(function (target) + local flash_so = mars_flash_attn_cuda_so_path() + local flash_dir = path.directory(flash_so) + local flash_name = path.filename(flash_so) + target:add( + "shflags", + "-Wl,--no-as-needed -L" .. flash_dir .. " -l:" .. flash_name .. " -Wl,-rpath," .. flash_dir, + {force = true} + ) + end) + end +target_end() + +add_includedirs(HPCC_ROOT .. "/include") +add_linkdirs(HPCC_ROOT .. "/lib") +add_links("hcdnn", "hcblas", "hcruntime") + +rule("hpcc") + set_extensions(".maca") + + on_load(function (target) + target:add("includedirs", "include") + end) + + on_buildcmd_file(function (target, batchcmds, sourcefile, opt) + local objectfile = target:objectfile(sourcefile) + local htcc = path.join(HPCC_ROOT, "htgpu_llvm/bin/htcc") + local args = { + "-x", "hpcc", "-c", sourcefile, "-o", objectfile, + "-I" .. HPCC_ROOT .. "/include", "-O3", "-fPIC", "-Werror", "-std=c++17" + } + + for _, includedir in ipairs(target:get("includedirs")) do + table.insert(args, "-I" .. includedir) + end + for _, define in ipairs(target:get("defines")) do + table.insert(args, "-D" .. define) + end + + table.insert(target:objectfiles(), objectfile) + batchcmds:mkdir(path.directory(objectfile)) + batchcmds:show_progress(opt.progress, "${color.build.object}compiling.hpcc %s", sourcefile) + batchcmds:vrunv(htcc, args) + batchcmds:add_depfiles(sourcefile) + batchcmds:set_depmtime(os.mtime(objectfile)) + batchcmds:set_depcache(target:dependfile(objectfile)) + end) +rule_end() + +target("infiniop-mars") + set_kind("static") + on_install(function (target) end) + set_languages("cxx17") + set_warnings("all", "error") + add_cxflags("-lstdc++", "-fPIC", "-Wno-defaulted-function-deleted", "-Wno-strict-aliasing", {force = true}) + add_cxxflags("-lstdc++", "-fPIC", "-Wno-defaulted-function-deleted", "-Wno-strict-aliasing", {force = true}) + add_files("../src/infiniop/devices/metax/*.cc") + add_files("../src/infiniop/ops/*/metax/*.cc") + add_files("../src/infiniop/ops/*/metax/*.maca", {rule = "hpcc"}) + + if has_config("ninetoothed") then + add_includedirs(HPCC_ROOT .. "/include/hcr") + add_files("../build/ninetoothed/*.c", "../build/ninetoothed/*.cpp", { + cxflags = { + "-include stdlib.h", + "-Wno-return-type", + "-Wno-implicit-function-declaration", + "-Wno-builtin-declaration-mismatch" + } + }) + end +target_end() + +target("flash-attn-mars") + set_kind("phony") + set_default(false) + + if FLASH_ATTN_ROOT and FLASH_ATTN_ROOT ~= "" then + before_build(function (target) + local torch_dir = os.iorunv("python", {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() + local python_include = os.iorunv("python", {"-c", "import sysconfig; print(sysconfig.get_paths()['include'])"}):trim() + local python_lib_dir = os.iorunv("python", {"-c", "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))"}):trim() + target:add("includedirs", torch_dir .. "/include", torch_dir .. "/include/torch/csrc/api/include", python_include, {public = false}) + target:add("linkdirs", torch_dir .. "/lib", python_lib_dir, {public = false}) + end) + end +target_end() + +target("infiniccl-mars") + set_kind("static") + on_install(function (target) end) + set_warnings("all", "error") + if not is_plat("windows") then + add_cxflags("-fPIC") + add_cxxflags("-fPIC") + end + if has_config("ccl") then + add_links("libhccl.so") + add_files("../src/infiniccl/mars/*.cc") + end + set_languages("cxx17") +target_end() diff --git a/xmake/metax.lua b/xmake/metax.lua index 3e700fff2..e308394a5 100644 --- a/xmake/metax.lua +++ b/xmake/metax.lua @@ -60,11 +60,7 @@ target_end() add_includedirs(MACA_ROOT .. "/include") add_linkdirs(MACA_ROOT .. "/lib") -if has_config("use-mc") then - add_links("mcdnn", "mcblas", "mcruntime") -else - add_links("hcdnn", "hcblas", "hcruntime") -end +add_links("mcdnn", "mcblas", "mcruntime") rule("maca") set_extensions(".maca") @@ -77,14 +73,8 @@ rule("maca") local objectfile = target:objectfile(sourcefile) os.mkdir(path.directory(objectfile)) local args - local htcc - if has_config("use-mc") then - htcc = path.join(MACA_ROOT, "mxgpu_llvm/bin/mxcc") - args = { "-x", "maca", "-c", sourcefile, "-o", objectfile, "-I" .. MACA_ROOT .. "/include", "-O3", "-fPIC", "-Werror", "-std=c++17"} - else - htcc = path.join(MACA_ROOT, "htgpu_llvm/bin/htcc") - args = { "-x", "hpcc", "-c", sourcefile, "-o", objectfile, "-I" .. MACA_ROOT .. "/include", "-O3", "-fPIC", "-Werror", "-std=c++17"} - end + local htcc = path.join(MACA_ROOT, "mxgpu_llvm/bin/mxcc") + local args = { "-x", "maca", "-c", sourcefile, "-o", objectfile, "-I" .. MACA_ROOT .. "/include", "-O3", "-fPIC", "-Werror", "-std=c++17"} local includedirs = table.concat(target:get("includedirs"), " ") for _, includedir in ipairs(target:get("includedirs")) do table.insert(args, "-I" .. includedir) @@ -165,11 +155,7 @@ target("infiniccl-metax") add_cxxflags("-fPIC") end if has_config("ccl") then - if has_config("use-mc") then - add_links("libmccl.so") - else - add_links("libhccl.so") - end + add_links("libmccl.so") add_files("../src/infiniccl/metax/*.cc") end set_languages("cxx17") From 5e034e4af525b0048d8b6e8b9c75228b27730665 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Thu, 23 Jul 2026 09:40:17 +0800 Subject: [PATCH 3/4] fix(mars): complete HPCC runtime integration Complete the Mars/HPCC split across build configuration, runtime loading, ATen dispatch, operator tests, and environment setup. Validate external InfiniRT prefixes, support lib64 installs, and reject standalone InfiniOps on Mars instead of mapping it back to MetaX. Resolve the active Mars flash_attn_2_cuda extension and detect its C++ ABI from exported symbols because HPCC and flash-attn are versioned independently. Pass the Mars-only attention-mask and auxiliary arguments for MHA, varlen MHA, and KV-cache while leaving the MetaX and NVIDIA signatures unchanged. Add Mars device and preload coverage, and update the build and test documentation for the separate backend. --- README.md | 18 +- include/infinicore/adaptor/aten_adaptor.hpp | 6 +- .../adaptor/flash_attention_adaptor.hpp | 31 +-- python/infinicore/_preload.py | 38 ++-- scripts/metax_env.py | 49 +++-- scripts/set_env.py | 35 +++- src/infinicore-test/README.md | 1 + src/infinicore/adaptor/aten_adaptor.cc | 2 +- src/infinicore/ops/infiniops_impl.hpp | 4 - .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 16 +- .../ops/multi_head_attention/mha_flashattn.cc | 24 ++- .../mha_varlen_flashattn.cc | 16 +- src/infiniop-test/src/main.cpp | 3 +- src/infiniop/ops/relu/operator.cc | 22 +- src/infiniop/ops/topk/cuda/kernel.cuh | 2 +- src/infiniop/ops/var/cuda/kernel.cuh | 2 +- src/infiniop/ops/var_mean/cuda/kernel.cuh | 2 +- src/infinirt-test/main.cc | 2 + test/infinicore/test.py | 25 ++- test/infinicore/test_preload.py | 44 ++++ test/infiniop/libinfiniop/liboperators.py | 25 ++- xmake.lua | 197 ++++++++++++++---- xmake/mars.lua | 22 +- xmake/metax.lua | 8 - 24 files changed, 423 insertions(+), 171 deletions(-) create mode 100644 test/infinicore/test_preload.py diff --git a/README.md b/README.md index 21d490224..173b3e82d 100644 --- a/README.md +++ b/README.md @@ -97,8 +97,8 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS] | `--nv-gpu=[y\|n]` | 是否编译英伟达 GPU 接口实现 | n | `--ascend-npu=[y\|n]` | 是否编译昇腾 NPU 接口实现 | n | `--cambricon-mlu=[y\|n]` | 是否编译寒武纪 MLU 接口实现 | n -| `--metax-gpu=[y\|n]` | 是否编译沐曦 GPU 接口实现 | n -| `--use-mc=[y\|n]` | 是否沐曦 GPU 接口实现使用maca SDK | n +| `--metax-gpu=[y\|n]` | 是否编译 MetaX/MACA GPU 接口实现 | n +| `--mars-gpu=[y\|n]` | 是否编译 Mars/HPCC GPU 接口实现 | n | `--moore-gpu=[y\|n]` | 是否编译摩尔线程 GPU 接口实现 | n | `--iluvatar-gpu=[y\|n]` | 是否编译天数 GPU 接口实现 | n | `--qy-gpu=[y\|n]` | 是否编译QY GPU 接口实现 | n @@ -108,6 +108,12 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS] | `--ninetoothed=[y\|n]` | 是否编译九齿实现 | n | `--ccl=[y\|n]` | 是否编译 InfiniCCL 通信库接口实现 | n | `--graph=[y\|n]` | 是否编译 cuda graph 接口实现 | n +| `--mars-flash-attn-abi=[detect\|standard\|extended]` | Mars flash-attn ABI；默认检查实际扩展库符号 | detect + +Mars/HPCC 当前不支持 standalone InfiniOps 后端；配置 Mars 时请使用 +`--mars-gpu=true --infiniops=false`。 +使用独立 InfiniRT 安装时应设置 `INFINI_RT_ROOT` 或 `--infinirt-root`；如果 +显式前缀中缺少 `include/infini/rt.h` 或 `lib{,64}/libinfinirt.so`，配置会直接失败。 ##### 手动安装底层库 @@ -268,9 +274,9 @@ pip install -e . ```bash # 测试单算子 -python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali] +python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali] # 测试全部算子 -python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --ali] +python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --ali] ``` 使用 -h 查看更多参数。 @@ -279,9 +285,9 @@ python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia ```shell # 测试单算子 -python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali] +python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali] # 测试全部算子 -python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali] +python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --mars | --moore | --kunlun | --Hygon | --ali] ``` #### 通信库（InfiniCCL）测试 diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index 5b8b5b58f..caac4bf7f 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -8,7 +8,7 @@ #if defined(ENABLE_HYGON_API) #include #include -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) #include #include #include @@ -40,7 +40,7 @@ inline at::ScalarType to_at_dtype(DataType dtype) { inline at::Device to_at_device(const Device &device) { // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). - // Treat MetaX/QY devices as CUDA devices for ATen tensor interoperability. + // Treat CUDA-compatible devices as CUDA devices for ATen interoperability. if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::MARS || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) { return at::Device(at::kCUDA, device.getIndex()); } else if (device.getType() == Device::Type::CPU) { @@ -60,7 +60,7 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t); #if defined(ENABLE_HYGON_API) c10::hip::HIPStream get_hip_stream(); -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStream get_cuda_stream(); #endif diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index c5bf14858..4f7e61ec9 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -2,10 +2,9 @@ #pragma once #include "aten_adaptor.hpp" -// NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension -// exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds -// where the namespace is empty. -#if !defined(ENABLE_METAX_API) +// NVIDIA flash-attn-nvidia.so uses namespace flash. MetaX and Mars extension +// modules export the same entry points at global scope. +#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API) namespace flash { #endif std::vector @@ -14,6 +13,10 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) std::optional &alibi_slopes_, // num_heads or batch_size x num_heads +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + // Mars extensions accept an attention mask before the scalar options. + std::optional &attn_mask_, +#endif const float p_dropout, const float softmax_scale, bool is_causal, @@ -22,10 +25,10 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num const float softcap, const bool return_softmax, std::optional gen_ -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + // Mars extensions append an auxiliary tensor after the generator. , - std::optional &flash_attn_mars_ext_ + std::optional &s_aux_ #endif ); @@ -51,10 +54,10 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_hea const float softcap, const bool return_softmax, std::optional gen_ -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + // Mars extensions append an auxiliary tensor. , - std::optional &flash_attn_mars_ext_ + std::optional &s_aux_ #endif ); @@ -126,14 +129,14 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size const float softcap, bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 int num_splits -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + // Mars extensions append an auxiliary tensor. , - std::optional &flash_attn_mars_ext_ + std::optional &s_aux_ #endif ); -#if !defined(ENABLE_METAX_API) +#if !defined(ENABLE_METAX_API) && !defined(ENABLE_MARS_API) } // namespace flash #endif #endif // ENABLE_FLASH_ATTN diff --git a/python/infinicore/_preload.py b/python/infinicore/_preload.py index 69121ff2f..7106885a3 100644 --- a/python/infinicore/_preload.py +++ b/python/infinicore/_preload.py @@ -12,8 +12,7 @@ def _candidate_prefixes(path: str) -> List[str]: """ - Return HPCC install prefixes to search for libs. - Prefer HPCC_PATH; if absent and explicitly opted-in, fall back to /opt/hpcc. + Return unique runtime install prefixes to search for libraries. """ prefixes: List[str] = [] if path: @@ -31,14 +30,15 @@ def _candidate_prefixes(path: str) -> List[str]: def _try_load(paths: Iterable[str], name: str) -> bool: """Try to load a shared library from given paths or system search path.""" for path in paths: - full = os.path.join(path, "lib", name) - if os.path.exists(full): - try: - ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL) - return True - except OSError: - # Try next candidate - continue + for subdir in ("lib", "lib64"): + full = os.path.join(path, subdir, name) + if os.path.exists(full): + try: + ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL) + return True + except OSError: + # Try next candidate + continue # Last resort: rely on loader search path try: ctypes.CDLL(name, mode=ctypes.RTLD_GLOBAL) @@ -54,7 +54,9 @@ def preload_hpcc() -> None: This mirrors the behavior of torch's HPCC build that loads libtorch_global_deps.so, but avoids introducing a hard torch dependency. All failures are swallowed. """ - hpcc_path = os.getenv("HPCC_PATH") + hpcc_path = os.getenv("HPCC_PATH") or os.getenv("HPCC_HOME") + if not hpcc_path and os.getenv("INFINICORE_PRELOAD_HPCC"): + hpcc_path = "/opt/hpcc" if not hpcc_path: return @@ -160,7 +162,7 @@ def _should_preload_device(device_type: str) -> bool: Check if preload is needed for a specific device type. """ device_env_map = { - "METAX": ["HPCC_PATH", "INFINICORE_PRELOAD_HPCC"], # HPCC/METAX + "MARS": ["HPCC_PATH", "HPCC_HOME", "INFINICORE_PRELOAD_HPCC"], "HYGON": ["DTK_ROOT", "INFINICORE_PRELOAD_TORCH_HIP"], # Add other device types here as needed: # "ASCEND": ["ASCEND_PATH"], @@ -183,9 +185,9 @@ def preload_device(device_type: str) -> None: Preload runtime libraries for a specific device type if needed. Args: - device_type: Device type name (e.g., "METAX", "ASCEND", etc.) + device_type: Device type name (e.g., "MARS" or "HYGON") """ - if device_type == "METAX": + if device_type == "MARS": preload_hpcc() elif device_type == "HYGON": preload_torch_hip() @@ -203,11 +205,11 @@ def preload() -> None: This function detects available device types and preloads their runtime libraries if the environment indicates they are needed. """ - # Device types that may require preload. Keep Hygon-only preloads gated by - # Hygon environment markers so other CUDA-compatible platforms do not load - # unrelated torch/flash-attn libraries during package import. + # Device types that may require preload. Gate each preload by its own + # environment markers so CUDA-compatible platforms do not load one another's + # runtime libraries during package import. device_types = [ - "METAX", # HPCC/METAX + "MARS", "HYGON", # Add other device types here as they are implemented: # "ASCEND", diff --git a/scripts/metax_env.py b/scripts/metax_env.py index df35c57a8..b55dc77d7 100644 --- a/scripts/metax_env.py +++ b/scripts/metax_env.py @@ -8,15 +8,12 @@ def _first_existing_dir(paths: list[str]) -> str: return "" -def _metax_toolkit_root(use_mc: bool) -> str: - """Return toolkit root for MetaX builds (MACA when use-mc; otherwise HPCC).""" - if use_mc: - for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"): - v = os.environ.get(key, "").strip() - if v: - return v - return _first_existing_dir(["/opt/maca"]) - return _first_existing_dir(["/opt/hpcc"]) +def _toolkit_root(env_names: tuple[str, ...], fallback: str) -> str: + for key in env_names: + value = os.environ.get(key, "").strip() + if value: + return value + return _first_existing_dir([fallback]) def _prepend_path_var(name: str, prefixes: list[str]) -> None: @@ -37,16 +34,14 @@ def set_env_for_metax_gpu( """ Prepend compiler include paths needed when building ATen-enabled C++ against torch headers. - This chooses paths based on xmake backend flags (e.g. --metax-gpu) and toolkit selection - (e.g. MetaX HPCC vs MACA when --use-mc=y). + MetaX always uses the MACA SDK. Mars/HPCC is configured separately. """ d = parse_xmake_cli_flag_values(flags) if not truthy_flag_value(d.get("aten", "n")): return if truthy_flag_value(d.get("metax-gpu", "n")): - use_mc = truthy_flag_value(d.get("use-mc", "n")) - root = _metax_toolkit_root(use_mc=use_mc) + root = _toolkit_root(("MACA_PATH", "MACA_HOME", "MACA_ROOT"), "/opt/maca") if not root: return dirs = [ @@ -64,4 +59,32 @@ def set_env_for_metax_gpu( ] for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"): _prepend_path_var(var, dirs) + + +def set_env_for_mars_gpu( + flags: str, + *, + parse_xmake_cli_flag_values, + truthy_flag_value, +) -> None: + """Prepend HPCC compatibility headers for ATen-enabled Mars builds.""" + d = parse_xmake_cli_flag_values(flags) + if not truthy_flag_value(d.get("aten", "n")): + return + if not truthy_flag_value(d.get("mars-gpu", "n")): + return + + root = _toolkit_root(("HPCC_PATH", "HPCC_HOME"), "/opt/hpcc") + if not root: return + dirs = [ + os.path.join(root, "tools", "cu-bridge", "include"), + os.path.join(root, "include", "hcr"), + os.path.join(root, "include", "common"), + os.path.join(root, "include", "hcsparse"), + os.path.join(root, "include", "hcblas"), + os.path.join(root, "include", "hcsolver"), + os.path.join(root, "include"), + ] + for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"): + _prepend_path_var(var, dirs) diff --git a/scripts/set_env.py b/scripts/set_env.py index f489c28ad..9477e7cf2 100644 --- a/scripts/set_env.py +++ b/scripts/set_env.py @@ -1,7 +1,7 @@ import os import platform -from metax_env import set_env_for_metax_gpu +from metax_env import set_env_for_mars_gpu, set_env_for_metax_gpu def _parse_xmake_cli_flag_values(flags: str): @@ -38,12 +38,16 @@ def set_env_by_config(flags: str) -> None: parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values, truthy_flag_value=_truthy_flag_value, ) - else: - pass + if _truthy_flag_value(d.get("mars-gpu", "n")): + set_env_for_mars_gpu( + flags, + parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values, + truthy_flag_value=_truthy_flag_value, + ) def set_env(): - if os.environ.get("INFINI_ROOT") == None: + if os.environ.get("INFINI_ROOT") is None: os.environ["INFINI_ROOT"] = os.path.expanduser("~/.infini") if platform.system() == "Windows": @@ -56,10 +60,23 @@ def set_env(): if new_path not in os.environ.get("PATH", ""): os.environ["PATH"] = f"{new_path}:{os.environ.get('PATH', '')}" - new_lib_path = os.path.expanduser(os.environ.get("INFINI_ROOT") + "/lib") - if new_lib_path not in os.environ.get("LD_LIBRARY_PATH", ""): - os.environ["LD_LIBRARY_PATH"] = ( - f"{new_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}" - ) + new_lib_paths = [] + infinirt_root = os.environ.get("INFINI_RT_ROOT") + if infinirt_root: + for subdir in ("lib", "lib64"): + candidate = os.path.join(infinirt_root, subdir) + if os.path.isdir(candidate): + new_lib_paths.append(candidate) + new_lib_paths.append(os.path.expanduser(os.environ["INFINI_ROOT"] + "/lib")) + + current_lib_paths = [ + path + for path in os.environ.get("LD_LIBRARY_PATH", "").split(":") + if path + ] + for new_lib_path in reversed(new_lib_paths): + if new_lib_path not in current_lib_paths: + current_lib_paths.insert(0, new_lib_path) + os.environ["LD_LIBRARY_PATH"] = ":".join(current_lib_paths) else: raise RuntimeError("Unsupported platform.") diff --git a/src/infinicore-test/README.md b/src/infinicore-test/README.md index a6349349d..8f96c5f1e 100644 --- a/src/infinicore-test/README.md +++ b/src/infinicore-test/README.md @@ -61,6 +61,7 @@ xmake build infinicore-test ./infinicore-test --cambricon ./infinicore-test --ascend ./infinicore-test --metax +./infinicore-test --mars ./infinicore-test --moore ./infinicore-test --iluvatar ./infinicore-test --qy diff --git a/src/infinicore/adaptor/aten_adaptor.cc b/src/infinicore/adaptor/aten_adaptor.cc index 73c01dfdb..8ab0603ab 100644 --- a/src/infinicore/adaptor/aten_adaptor.cc +++ b/src/infinicore/adaptor/aten_adaptor.cc @@ -37,7 +37,7 @@ c10::hip::HIPStream get_hip_stream() { return c10::hip::getStreamFromExternal( hipStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex()); } -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStream get_cuda_stream() { return c10::cuda::getStreamFromExternal( cudaStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex()); diff --git a/src/infinicore/ops/infiniops_impl.hpp b/src/infinicore/ops/infiniops_impl.hpp index 5c3d475f0..befb830cd 100644 --- a/src/infinicore/ops/infiniops_impl.hpp +++ b/src/infinicore/ops/infiniops_impl.hpp @@ -51,8 +51,6 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) { return infini::ops::Device{infini::ops::Device::Type::kNvidia, static_cast(device.getIndex())}; case Device::Type::METAX: return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast(device.getIndex())}; - case Device::Type::MARS: - return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast(device.getIndex())}; case Device::Type::MOORE: return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast(device.getIndex())}; case Device::Type::ILUVATAR: @@ -66,7 +64,6 @@ inline bool isSupportedDevice(Device::Type device_type) { switch (device_type) { case Device::Type::NVIDIA: case Device::Type::METAX: - case Device::Type::MARS: case Device::Type::MOORE: case Device::Type::ILUVATAR: return true; @@ -79,7 +76,6 @@ template void registerSupportedDevices(Dispatcher &dispatcher, Function function) { dispatcher.registerDevice(Device::Type::NVIDIA, function); dispatcher.registerDevice(Device::Type::METAX, function); - dispatcher.registerDevice(Device::Type::MARS, function); dispatcher.registerDevice(Device::Type::MOORE, function); dispatcher.registerDevice(Device::Type::ILUVATAR, function); } diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index 0167c17df..cedb0f45e 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -5,12 +5,12 @@ #include #ifdef ENABLE_FLASH_ATTN -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) #include #endif #endif -#if defined(ENABLE_METAX_API) +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #define INFINICORE_FLASH_OP(name) ::name #else #define INFINICORE_FLASH_OP(name) flash::name @@ -45,7 +45,7 @@ void *plan(Tensor out, void run(void *planned_meta) { #ifdef ENABLE_FLASH_ATTN -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); #endif auto *p = reinterpret_cast(planned_meta); @@ -55,7 +55,7 @@ void run(void *planned_meta) { Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work); auto q = infinicore::adaptor::to_aten_tensor(p->q); -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); #elif defined(ENABLE_QY_API) @@ -84,8 +84,8 @@ void run(void *planned_meta) { auto out = use_dynamic_out ? std::optional(std::nullopt) : std::optional(out_tensor); -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + std::optional s_aux = std::nullopt; #endif auto result = INFINICORE_FLASH_OP(mha_fwd_kvcache)( @@ -109,9 +109,9 @@ void run(void *planned_meta) { 0.0f, false, 0 -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT , - flash_attn_mars_ext + s_aux #endif ); diff --git a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc index 13c96b94d..46ede944d 100644 --- a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc @@ -5,7 +5,7 @@ #include #ifdef ENABLE_FLASH_ATTN -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) #include #endif #endif @@ -39,11 +39,10 @@ void *plan(Tensor out, namespace { -// Only support nv for now -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) -// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_fwd` at global scope (no namespace), +#if defined(ENABLE_FLASH_ATTN) && (defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API)) +// MetaX/Mars `flash_attn_2_cuda` exports `mha_fwd` at global scope, // while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_fwd`. -#if defined(ENABLE_METAX_API) +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #define INFINICORE_FLASH_OP(name) ::name #else #define INFINICORE_FLASH_OP(name) flash::name @@ -53,8 +52,7 @@ namespace { } // namespace void run(void *planned_meta) { -// Only support nv for now -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) +#if defined(ENABLE_FLASH_ATTN) && (defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API)) c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); auto *p = reinterpret_cast(planned_meta); @@ -71,8 +69,9 @@ void run(void *planned_meta) { auto scale = p->scale; auto is_causal = p->is_causal; -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + std::optional attn_mask = std::nullopt; + std::optional s_aux = std::nullopt; #endif INFINICORE_FLASH_OP(mha_fwd) @@ -82,6 +81,9 @@ void run(void *planned_meta) { v, out, alibi_slopes, +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + attn_mask, +#endif 0.0, scale, is_causal, @@ -90,9 +92,9 @@ void run(void *planned_meta) { 0.0, false, std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT , - flash_attn_mars_ext + s_aux #endif ); diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc index 4a233bf5b..5842bb487 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc @@ -3,7 +3,7 @@ #ifdef ENABLE_ATEN #include "infinicore/adaptor/aten_adaptor.hpp" #include -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) #include #endif #endif @@ -53,9 +53,9 @@ void *plan(Tensor out, namespace { #ifdef ENABLE_FLASH_ATTN -// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_varlen_fwd` at global scope (no namespace), +// MetaX/Mars `flash_attn_2_cuda` exports `mha_varlen_fwd` at global scope, // while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_varlen_fwd`. -#if defined(ENABLE_METAX_API) +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #define INFINICORE_FLASH_OP(name) ::name #else #define INFINICORE_FLASH_OP(name) flash::name @@ -69,7 +69,7 @@ void run(void *planned_meta) { (void)planned_meta; throw std::runtime_error("ATen is not enabled in this build"); #else -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); #endif auto *p = reinterpret_cast(planned_meta); @@ -138,8 +138,8 @@ void run(void *planned_meta) { auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; auto scale = p->scale; -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT + std::optional s_aux = std::nullopt; #endif INFINICORE_FLASH_OP(mha_varlen_fwd) @@ -165,9 +165,9 @@ void run(void *planned_meta) { 0.0, false, std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) +#ifdef INFINICORE_FLASH_ATTN_MARS_EXT , - flash_attn_mars_ext + s_aux #endif ); diff --git a/src/infiniop-test/src/main.cpp b/src/infiniop-test/src/main.cpp index 00e16cc99..5ea9c2570 100644 --- a/src/infiniop-test/src/main.cpp +++ b/src/infiniop-test/src/main.cpp @@ -22,7 +22,7 @@ void printUsage() { std::cout << " Path to the test gguf file" << std::endl << std::endl; std::cout << " --[:id]" << std::endl; - std::cout << " (Optional) Specify the device type --(cpu|nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali) and device ID (optional). CPU by default." << std::endl + std::cout << " (Optional) Specify the device type --(cpu|nvidia|cambricon|ascend|metax|mars|moore|iluvatar|qy|kunlun|hygon|ali) and device ID (optional). CPU by default." << std::endl << std::endl; std::cout << " --warmup " << std::endl; std::cout << " (Optional) Number of warmups to perform before timing. Default to 0." << std::endl @@ -75,6 +75,7 @@ ParsedArgs parseArgs(int argc, char *argv[]) { PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON) PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND) PARSE_DEVICE("--metax", INFINI_DEVICE_METAX) + PARSE_DEVICE("--mars", INFINI_DEVICE_MARS) PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE) PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR) PARSE_DEVICE("--qy", INFINI_DEVICE_QY) diff --git a/src/infiniop/ops/relu/operator.cc b/src/infiniop/ops/relu/operator.cc index 59b6f8ada..8d00233f7 100644 --- a/src/infiniop/ops/relu/operator.cc +++ b/src/infiniop/ops/relu/operator.cc @@ -8,7 +8,7 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) #include "nvidia/relu_nvidia.cuh" #endif -#ifdef ENABLE_METAX_API +#if defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) #ifdef ENABLE_NINETOOTHED #include "metax/relu_metax.h" #endif @@ -50,6 +50,11 @@ __INFINI_C infiniStatus_t infiniopCreateReluDescriptor( CREATE(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API +#ifdef ENABLE_NINETOOTHED + CREATE(INFINI_DEVICE_MARS, metax); +#endif +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -89,6 +94,11 @@ __INFINI_C infiniStatus_t infiniopGetReluWorkspaceSize(infiniopReluDescriptor_t GET(INFINI_DEVICE_METAX, metax) #endif #endif +#ifdef ENABLE_MARS_API +#ifdef ENABLE_NINETOOTHED + GET(INFINI_DEVICE_MARS, metax) +#endif +#endif #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia); #endif @@ -136,6 +146,11 @@ __INFINI_C infiniStatus_t infiniopRelu( CALCULATE(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API +#ifdef ENABLE_NINETOOTHED + CALCULATE(INFINI_DEVICE_MARS, metax); +#endif +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -177,6 +192,11 @@ infiniopDestroyReluDescriptor(infiniopReluDescriptor_t desc) { DELETE(INFINI_DEVICE_METAX, metax); #endif #endif +#ifdef ENABLE_MARS_API +#ifdef ENABLE_NINETOOTHED + DELETE(INFINI_DEVICE_MARS, metax); +#endif +#endif #ifdef ENABLE_ALI_API DELETE(INFINI_DEVICE_ALI, nvidia); #endif diff --git a/src/infiniop/ops/topk/cuda/kernel.cuh b/src/infiniop/ops/topk/cuda/kernel.cuh index 199ca8765..df9fade03 100644 --- a/src/infiniop/ops/topk/cuda/kernel.cuh +++ b/src/infiniop/ops/topk/cuda/kernel.cuh @@ -47,7 +47,7 @@ __device__ __forceinline__ float to_float(half v) { return __half2float(v) #if defined(ENABLE_MOORE_API) using bf16_t = __mt_bfloat16; -#elif defined(ENABLE_METAX_API) +#elif defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) using bf16_t = __hpcc_bfloat16; #else // CUDA / NVIDIA / ILUVATAR diff --git a/src/infiniop/ops/var/cuda/kernel.cuh b/src/infiniop/ops/var/cuda/kernel.cuh index a39cac751..dd6b1473e 100644 --- a/src/infiniop/ops/var/cuda/kernel.cuh +++ b/src/infiniop/ops/var/cuda/kernel.cuh @@ -22,7 +22,7 @@ __inline__ __device__ half Nan() { #if defined(ENABLE_MOORE_API) using bf16_t = __mt_bfloat16; -#elif defined(ENABLE_METAX_API) +#elif defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) using bf16_t = __hpcc_bfloat16; #else using bf16_t = __nv_bfloat16; diff --git a/src/infiniop/ops/var_mean/cuda/kernel.cuh b/src/infiniop/ops/var_mean/cuda/kernel.cuh index 945cabf16..0bd914b20 100644 --- a/src/infiniop/ops/var_mean/cuda/kernel.cuh +++ b/src/infiniop/ops/var_mean/cuda/kernel.cuh @@ -22,7 +22,7 @@ __inline__ __device__ half Nan() { #if defined(ENABLE_MOORE_API) using bf16_t = __mt_bfloat16; -#elif defined(ENABLE_METAX_API) +#elif defined(ENABLE_METAX_API) || defined(ENABLE_MARS_API) using bf16_t = __hpcc_bfloat16; #else using bf16_t = __nv_bfloat16; diff --git a/src/infinirt-test/main.cc b/src/infinirt-test/main.cc index 5c3e97db8..eebcc4b15 100644 --- a/src/infinirt-test/main.cc +++ b/src/infinirt-test/main.cc @@ -18,6 +18,7 @@ void printUsage() { << " cambricon" << std::endl << " ascend" << std::endl << " metax" << std::endl + << " mars" << std::endl << " moore" << std::endl << " iluvatar" << std::endl << " qy" << std::endl @@ -51,6 +52,7 @@ ParsedArgs parseArgs(int argc, char *argv[]) { else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON) else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND) else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX) + else PARSE_DEVICE("--mars", INFINI_DEVICE_MARS) else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE) else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR) else PARSE_DEVICE("--qy", INFINI_DEVICE_QY) diff --git a/test/infinicore/test.py b/test/infinicore/test.py index 72828b567..4d20ad0c5 100644 --- a/test/infinicore/test.py +++ b/test/infinicore/test.py @@ -1,9 +1,9 @@ -import torch -from infinicore.lib import _infinicore -from infinicore.utils import to_torch_dtype import numpy as np +import torch import infinicore +from infinicore.lib import _infinicore +from infinicore.utils import to_torch_dtype def test(): @@ -186,10 +186,11 @@ def func6_initialize_device_relationship(): _infinicore.Device.Type.KUNLUN, # 7 "cuda" _infinicore.Device.Type.HYGON, # 8 "cuda" _infinicore.Device.Type.ALI, # 10 "cuda" + _infinicore.Device.Type.MARS, # 11 "cuda" ] if True: print("\n ---------- 测试 CPU") - all_device_count = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0] + all_device_count = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] infinicore_2_python_dict, python_2_infinicore_dict = ( _initialize_device_relationship(all_device_types, all_device_count) ) @@ -197,15 +198,25 @@ def func6_initialize_device_relationship(): print("python_2_infinicore_dict: ", python_2_infinicore_dict) print("\n ---------- 测试 CPU+NVIDIA") - all_device_count = [1, 2, 0, 0, 0, 0, 0, 0, 0, 0] + all_device_count = [1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + infinicore_2_python_dict, python_2_infinicore_dict = ( + _initialize_device_relationship(all_device_types, all_device_count) + ) + print("infinicore_2_python_dict", infinicore_2_python_dict) + print("python_2_infinicore_dict: ", python_2_infinicore_dict) + + print("\n ---------- 测试 CPU+MARS") + all_device_count = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2] infinicore_2_python_dict, python_2_infinicore_dict = ( _initialize_device_relationship(all_device_types, all_device_count) ) + assert len(infinicore_2_python_dict[_infinicore.Device.Type.MARS]) == 2 + assert len(python_2_infinicore_dict["cuda"]) == 2 print("infinicore_2_python_dict", infinicore_2_python_dict) print("python_2_infinicore_dict: ", python_2_infinicore_dict) print("\n ---------- 测试 CPU+NVIDIA+HYGON") - all_device_count = [1, 2, 0, 0, 0, 0, 0, 0, 0, 2] + all_device_count = [1, 2, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0] infinicore_2_python_dict, python_2_infinicore_dict = ( _initialize_device_relationship(all_device_types, all_device_count) ) @@ -213,7 +224,7 @@ def func6_initialize_device_relationship(): print("python_2_infinicore_dict: ", python_2_infinicore_dict) print("\n ---------- 测试 CPU+NVIDIA+HYGON+CAMBRICON+ASCEND") - all_device_count = [1, 2, 2, 2, 0, 0, 0, 0, 0, 2] + all_device_count = [1, 2, 2, 2, 0, 0, 0, 0, 0, 2, 0, 0] infinicore_2_python_dict, python_2_infinicore_dict = ( _initialize_device_relationship(all_device_types, all_device_count) ) diff --git a/test/infinicore/test_preload.py b/test/infinicore/test_preload.py new file mode 100644 index 000000000..dde9e8254 --- /dev/null +++ b/test/infinicore/test_preload.py @@ -0,0 +1,44 @@ +import ctypes +import os +import unittest +from unittest import mock + +from infinicore import _preload + + +class DevicePreloadTest(unittest.TestCase): + def test_hpcc_environment_selects_only_mars(self): + with mock.patch.dict(os.environ, {"HPCC_PATH": "/opt/hpcc"}, clear=True): + self.assertTrue(_preload._should_preload_device("MARS")) + self.assertFalse(_preload._should_preload_device("METAX")) + + def test_maca_environment_does_not_select_mars(self): + with mock.patch.dict(os.environ, {"MACA_PATH": "/opt/maca"}, clear=True): + self.assertFalse(_preload._should_preload_device("MARS")) + self.assertFalse(_preload._should_preload_device("METAX")) + + def test_hpcc_home_is_a_supported_runtime_prefix(self): + with mock.patch.dict(os.environ, {"HPCC_HOME": "/custom/hpcc"}, clear=True): + with mock.patch.object(_preload, "_try_load") as try_load: + _preload.preload_hpcc() + + self.assertEqual(try_load.call_count, 4) + self.assertEqual(try_load.call_args_list[0].args[0], ["/custom/hpcc"]) + + def test_runtime_loader_checks_lib64(self): + def exists(path): + return path == "/custom/hpcc/lib64/libhcruntime.so" + + with mock.patch.object(_preload.os.path, "exists", side_effect=exists): + with mock.patch.object(_preload.ctypes, "CDLL") as cdll: + self.assertTrue( + _preload._try_load(["/custom/hpcc"], "libhcruntime.so") + ) + + cdll.assert_called_once_with( + "/custom/hpcc/lib64/libhcruntime.so", mode=ctypes.RTLD_GLOBAL + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/infiniop/libinfiniop/liboperators.py b/test/infiniop/libinfiniop/liboperators.py index 48ce04fe4..8452aca26 100644 --- a/test/infiniop/libinfiniop/liboperators.py +++ b/test/infiniop/libinfiniop/liboperators.py @@ -10,6 +10,7 @@ from .structs import * INFINI_ROOT = os.getenv("INFINI_ROOT") or str(Path.home() / ".infini") +INFINI_RT_ROOT = os.getenv("INFINI_RT_ROOT") class InfiniLib: @@ -28,11 +29,9 @@ def __getattr__(self, name): # Open operators library def open_lib(): - def find_library_in_ld_path(subdir, library_name): - ld_library_path = os.path.join(INFINI_ROOT, subdir) - paths = ld_library_path.split(os.pathsep) - for path in paths: - full_path = os.path.join(path, library_name) + def find_library(root, subdirs, library_name): + for subdir in subdirs: + full_path = os.path.join(root, subdir, library_name) if os.path.isfile(full_path): return full_path return None @@ -40,17 +39,23 @@ def find_library_in_ld_path(subdir, library_name): system_name = platform.system() # Load the library if system_name == "Windows": - libop_path = find_library_in_ld_path("bin", "infiniop.dll") - librt_path = find_library_in_ld_path("bin", "infinirt.dll") + libop_path = find_library(INFINI_ROOT, ("bin",), "infiniop.dll") + librt_root = INFINI_RT_ROOT or INFINI_ROOT + librt_path = find_library(librt_root, ("bin",), "infinirt.dll") elif system_name == "Linux": - libop_path = find_library_in_ld_path("lib", "libinfiniop.so") - librt_path = find_library_in_ld_path("lib", "libinfinirt.so") + libop_path = find_library(INFINI_ROOT, ("lib", "lib64"), "libinfiniop.so") + librt_root = INFINI_RT_ROOT or INFINI_ROOT + librt_path = find_library( + librt_root, ("lib", "lib64"), "libinfinirt.so" + ) + else: + raise RuntimeError(f"Unsupported platform: {system_name}") assert libop_path is not None, ( "Cannot find infiniop.dll or libinfiniop.so. Check if INFINI_ROOT is set correctly." ) assert librt_path is not None, ( - "Cannot find infinirt.dll or libinfinirt.so. Check if INFINI_ROOT is set correctly." + "Cannot find infinirt.dll or libinfinirt.so. Check INFINI_RT_ROOT or INFINI_ROOT." ) librt = ctypes.CDLL(librt_path) diff --git a/xmake.lua b/xmake.lua index 407ae8742..73075739d 100644 --- a/xmake.lua +++ b/xmake.lua @@ -165,10 +165,6 @@ if has_config("mars-gpu") then includes("xmake/mars.lua") end -if has_config("metax-gpu") and has_config("mars-gpu") then - raise("--metax-gpu and --mars-gpu are separate backends and cannot be enabled together") -end - -- 摩尔线程 option("moore-gpu") set_default(false) @@ -243,6 +239,13 @@ option("flash-attn") set_description("Path to flash-attention repo. If not set, flash-attention will not used.") option_end() +option("mars-flash-attn-abi") + set_default("detect") + set_values("detect", "standard", "extended") + set_showmenu(true) + set_description("Mars flash-attn ABI: detect from the library, standard, or extended") +option_end() + if has_config("aten") then add_defines("ENABLE_ATEN") if has_config("iluvatar-gpu") then @@ -334,11 +337,115 @@ local function find_infinirt_library(infinirt_root, xmake_os) infinirt_library = path.join(infinirt_root, "lib64", "libinfinirt.so") end if not xmake_os.isfile(infinirt_library) then - raise("InfiniRT library not found under: " .. infinirt_root) + xmake_os.raise("InfiniRT library not found under: " .. infinirt_root) end return infinirt_library end +local function validate_infinirt_root(xmake_os) + local infinirt_root = get_infinirt_root() + if not infinirt_root or infinirt_root == "" then + return + end + + local public_header = path.join(infinirt_root, "include", "infini", "rt.h") + if not xmake_os.isfile(public_header) then + xmake_os.raise("InfiniRT public header not found: " .. public_header) + end + find_infinirt_library(infinirt_root, xmake_os) +end + +local function find_mars_flash_attn_library(xmake_os) + local configured = os.getenv("FLASH_ATTN_2_CUDA_SO") + if configured and configured ~= "" then + configured = configured:trim() + if not xmake_os.isfile(configured) then + xmake_os.raise("Mars flash-attn library not found: " .. configured) + end + return configured + end + + local find_spec_script = + "import importlib.util; s = importlib.util.find_spec('flash_attn_2_cuda'); print(s.origin if s and s.origin else '')" + local discovered = xmake_os.iorunv("python", {"-c", find_spec_script}):trim() + if discovered ~= "" and xmake_os.isfile(discovered) then + return discovered + end + + local fallback = os.getenv("FLASH_ATTN_MARS_CUDA_SO_CONTAINER") + or "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so" + if not xmake_os.isfile(fallback) then + xmake_os.raise("Mars flash-attn library not found; set FLASH_ATTN_2_CUDA_SO") + end + return fallback +end + +-- HPCC and flash-attn are versioned independently, so inspect the extension's +-- exported C++ signatures instead of deriving its ABI from the SDK version. +local function mars_flash_attn_uses_extended_abi(xmake_os) + local abi = get_config("mars-flash-attn-abi") or "detect" + if abi == "extended" then + return true + elseif abi == "standard" then + return false + end + + local flash_attn_library = find_mars_flash_attn_library(xmake_os) + local symbols = xmake_os.iorunv("nm", {"-D", "-C", flash_attn_library}) + local mha_fwd_symbol = nil + local mha_varlen_fwd_symbol = nil + local mha_fwd_kvcache_symbol = nil + for line in symbols:gmatch("[^\r\n]+") do + if line:find(" mha_fwd(", 1, true) then + mha_fwd_symbol = line + elseif line:find(" mha_varlen_fwd(", 1, true) then + mha_varlen_fwd_symbol = line + elseif line:find(" mha_fwd_kvcache(", 1, true) then + mha_fwd_kvcache_symbol = line + end + end + + if mha_fwd_symbol and mha_varlen_fwd_symbol and mha_fwd_kvcache_symbol then + local has_attention_mask = mha_fwd_symbol:find( + "std::optional&, std::optional&, std::optional&, float, float", + 1, + true + ) ~= nil + local fwd_has_s_aux = mha_fwd_symbol:find( + "std::optional, std::optional&)", + 1, + true + ) ~= nil + local varlen_has_s_aux = mha_varlen_fwd_symbol:find( + "std::optional, std::optional&)", + 1, + true + ) ~= nil + local kvcache_has_s_aux = mha_fwd_kvcache_symbol:match( + ", int, std::optional&%)$" + ) ~= nil + if has_attention_mask + and fwd_has_s_aux + and varlen_has_s_aux + and kvcache_has_s_aux + then + return true + elseif not has_attention_mask + and mha_fwd_symbol:find("std::optional)", 1, true) + and mha_varlen_fwd_symbol:find("std::optional)", 1, true) + and mha_fwd_kvcache_symbol:match(", int%)$") + then + return false + end + end + + xmake_os.raise( + "Unable to detect the Mars flash-attn ABI from " + .. flash_attn_library + .. "; set --mars-flash-attn-abi=standard or extended" + ) +end + local function add_external_infinirt() local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") local infinirt_root = get_infinirt_root() @@ -374,7 +481,7 @@ local function add_external_infinirt() end end if has_config("mars-gpu") then - local hpcc_root = os.getenv("HPCC_PATH") or "/opt/hpcc" + local hpcc_root = os.getenv("HPCC_PATH") or os.getenv("HPCC_HOME") or "/opt/hpcc" local hpcc_include = path.join(hpcc_root, "include") if os.isdir(hpcc_include) then add_includedirs(hpcc_include, { public = true }) @@ -418,7 +525,12 @@ local function filter_infiniops_ops_for_backend(infiniops_ops) return table.concat(filtered, ",") end -local function get_infiniops_backend_cmake_arg() +local function get_infiniops_backend_cmake_arg(xmake_os) + if has_config("mars-gpu") then + xmake_os.raise( + "standalone InfiniOps does not expose a distinct Mars backend; configure with --infiniops=false" + ) + end local enabled = {} local function add_backend(config, cmake_arg) if has_config(config) then @@ -427,14 +539,15 @@ local function get_infiniops_backend_cmake_arg() end add_backend("nv-gpu", "-DWITH_NVIDIA=ON") add_backend("metax-gpu", "-DWITH_METAX=ON") - add_backend("mars-gpu", "-DWITH_MARS=ON") add_backend("iluvatar-gpu", "-DWITH_ILUVATAR=ON") add_backend("moore-gpu", "-DWITH_MOORE=ON") if #enabled == 0 then - raise("InfiniOps integration requires one of --nv-gpu, --metax-gpu, --mars-gpu, --iluvatar-gpu, or --moore-gpu") + xmake_os.raise( + "InfiniOps integration requires one of --nv-gpu, --metax-gpu, --iluvatar-gpu, or --moore-gpu" + ) end if #enabled > 1 then - raise("InfiniOps can build only one GPU backend at a time") + xmake_os.raise("InfiniOps can build only one GPU backend at a time") end return enabled[1] end @@ -446,12 +559,15 @@ local function build_infiniops_external(xmake_os) local infiniops_root = path.absolute(get_config("infiniops-root") or "submodules/InfiniOps", os.projectdir()) local infiniops_builddir = path.join(infiniops_root, "build") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") + if not xmake_os.isdir(infiniops_root) then + xmake_os.raise("InfiniOps root not found: " .. infiniops_root) + end local infinirt_root = get_infinirt_root() local cmake_config_args = { "-S", infiniops_root, "-B", infiniops_builddir, "-DWITH_CPU=ON", - get_infiniops_backend_cmake_arg(), + get_infiniops_backend_cmake_arg(xmake_os), "-DGENERATE_OPERATOR_CALL_INSTANTIATIONS=ON", "-DGENERATE_PYTHON_BINDINGS=OFF", "-DCMAKE_BUILD_TYPE=Release" @@ -500,6 +616,12 @@ end target("infini-utils") set_kind("static") on_install(function (target) end) + on_load(function (target) + if has_config("metax-gpu") and has_config("mars-gpu") then + os.raise("--metax-gpu and --mars-gpu are separate backends and cannot be enabled together") + end + validate_infinirt_root(os) + end) set_languages("cxx17") set_warnings("all", "error") @@ -671,6 +793,18 @@ target("infinicore_cpp_api") add_deps("infiniop", "infiniccl") set_languages("cxx17") set_symbols("visibility") + on_load(function (target) + if has_config("infiniops") then + build_infiniops_external(os) + end + if has_config("mars-gpu") + and get_config("flash-attn") + and get_config("flash-attn") ~= "" + and mars_flash_attn_uses_extended_abi(os) + then + target:add("defines", "INFINICORE_FLASH_ATTN_MARS_EXT=1") + end + end) local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") @@ -682,6 +816,17 @@ target("infinicore_cpp_api") for _, include_dir in ipairs(os.dirs(maca_root .. "/include/*")) do add_includedirs(include_dir) end + end + if has_config("mars-gpu") and has_config("aten") then + local hpcc_root = os.getenv("HPCC_PATH") or os.getenv("HPCC_HOME") or "/opt/hpcc" + add_includedirs(hpcc_root .. "/include") + add_includedirs(hpcc_root .. "/tools/cu-bridge/include") + for _, include_dir in ipairs(os.dirs(hpcc_root .. "/include/*")) do + add_includedirs(include_dir) + end + end + if has_config("aten") and (has_config("metax-gpu") or has_config("mars-gpu")) then + -- The CUDA-compatibility SDK headers may shadow the system limits header. add_cxflags("-include", "climits") add_cxxflags("-includeclimits", {force = true}) add_defines("CHAR_BIT=8", "INT_MIN=(-2147483647 - 1)", "INT_MAX=2147483647", "UINT_MAX=4294967295U") @@ -694,10 +839,6 @@ target("infinicore_cpp_api") end if has_config("infiniops") then local infiniops_root = path.absolute(get_config("infiniops-root") or "submodules/InfiniOps", os.projectdir()) - if not os.isdir(infiniops_root) then - raise("InfiniOps root not found: " .. infiniops_root) - end - get_infiniops_backend_cmake_arg() local infinirt_root = get_infinirt_root() if infinirt_root and infinirt_root ~= "" then add_includedirs(infinirt_root .. "/include") @@ -708,9 +849,6 @@ target("infinicore_cpp_api") add_defines("ENABLE_INFINIOPS_API") add_links("infiniops") add_rpathdirs(INFINI_ROOT .. "/lib") - on_load(function (target) - build_infiniops_external(os) - end) after_install(function (target) local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") local infinirt_root = get_infinirt_root() @@ -746,8 +884,7 @@ target("infinicore_cpp_api") end -- Flash pip `.so` link flags: `before_link` runs in an xmake sandbox that cannot see helpers - -- from other included scripts; MetaX and QY each register their own hook in `xmake/metax.lua` - -- and `xmake/qy.lua`. + -- from other included scripts; platform scripts register their own hooks. -- Moore mate: -- enable Python bridge macro for flash-attn Moore path @@ -758,26 +895,6 @@ target("infinicore_cpp_api") end before_build(function (target) - -- Mars HPCC flash-attn uses the HPCC major version to select its ABI. - if has_config("mars-gpu") and get_config("flash-attn") and get_config("flash-attn") ~= "" then - local hpcc_root = os.getenv("HPCC_PATH") or "/opt/hpcc" - local version_txt = path.join(hpcc_root, "Version.txt") - if os.isfile(version_txt) then - local content = os.iorunv("cat", {version_txt}) or "" - content = content:trim() - local major_str = content:match("Version:(%d+)") or content:match("^(%d+)") - if major_str and major_str ~= "" then - local major = tonumber(major_str) - if major then - local define = "INFINICORE_HPCC_VERSION_MAJOR=" .. tostring(major) - target:add("defines", define) - target:add("cxflags", "-D" .. define) - target:add("cxxflags", "-D" .. define) - end - end - end - end - if has_config("aten") then local outdata = os.iorunv("python", {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() local TORCH_DIR = outdata diff --git a/xmake/mars.lua b/xmake/mars.lua index c85c3d1c9..85c011505 100644 --- a/xmake/mars.lua +++ b/xmake/mars.lua @@ -4,27 +4,37 @@ local FLASH_ATTN_ROOT = get_config("flash-attn") local FLASH_ATTN_MARS_CUDA_SO_CONTAINER_DEFAULT = "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so" -local function mars_flash_attn_cuda_so_path() - local env_path = os.getenv("FLASH_ATTN_2_CUDA_SO") +local function mars_flash_attn_cuda_so_path(xmake_os) + local env_path = xmake_os.getenv("FLASH_ATTN_2_CUDA_SO") if env_path and env_path ~= "" then env_path = env_path:trim() - if os.isfile(env_path) then + if xmake_os.isfile(env_path) then return env_path end - print(string.format("warning: mars+flash-attn: FLASH_ATTN_2_CUDA_SO is not a file: %s", env_path)) + xmake_os.raise(string.format("mars+flash-attn: FLASH_ATTN_2_CUDA_SO is not a file: %s", env_path)) end - local container_path = os.getenv("FLASH_ATTN_MARS_CUDA_SO_CONTAINER") + local find_spec_script = + "import importlib.util; s = importlib.util.find_spec('flash_attn_2_cuda'); print(s.origin if s and s.origin else '')" + local discovered = xmake_os.iorunv("python", {"-c", find_spec_script}):trim() + if discovered ~= "" and xmake_os.isfile(discovered) then + return discovered + end + + local container_path = xmake_os.getenv("FLASH_ATTN_MARS_CUDA_SO_CONTAINER") if not container_path or container_path == "" then container_path = FLASH_ATTN_MARS_CUDA_SO_CONTAINER_DEFAULT end + if not xmake_os.isfile(container_path) then + xmake_os.raise(string.format("mars+flash-attn: library not found: %s", container_path)) + end return container_path end target("infinicore_cpp_api") if FLASH_ATTN_ROOT and FLASH_ATTN_ROOT ~= "" then before_link(function (target) - local flash_so = mars_flash_attn_cuda_so_path() + local flash_so = mars_flash_attn_cuda_so_path(os) local flash_dir = path.directory(flash_so) local flash_name = path.filename(flash_so) target:add( diff --git a/xmake/metax.lua b/xmake/metax.lua index e308394a5..6f0705900 100644 --- a/xmake/metax.lua +++ b/xmake/metax.lua @@ -2,13 +2,6 @@ local MACA_ROOT = os.getenv("MACA_PATH") or os.getenv("MACA_HOME") or os.getenv("MACA_ROOT") local FLASH_ATTN_ROOT = get_config("flash-attn") --- MetaX flash-attn (pip `flash_attn_2_cuda`) may append an extra trailing argument --- (`flash_attn_mars_ext_`) depending on the underlying HPCC/MetaX stack version. -do - -- Intentionally empty: HPCC version parsing is deferred to `before_build` - -- on `infinicore_cpp_api` where `os.iorunv` is available in this xmake sandbox. -end - -- Resolve MetaX flash-attn .so path (used only from this file: `before_link` sandbox cannot see globals from `xmake.lua`). local FLASH_ATTN_METAX_CUDA_SO_CONTAINER_DEFAULT = "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so" @@ -42,7 +35,6 @@ local function metax_flash_attn_cuda_so_path() end -- MetaX flash-attn link flags for pip `flash_attn_2_cuda`. --- Version/ABI macros are set in `xmake.lua` for `infinicore_cpp_api` so they apply to all sources. target("infinicore_cpp_api") if get_config("flash-attn") and get_config("flash-attn") ~= "" then before_link(function (target) From ee2f301b271b58e975c955d58e87876c20d18984 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Thu, 23 Jul 2026 07:00:31 +0000 Subject: [PATCH 4/4] fix(metax): detect flash-attn extension ABI MetaX flash-attn packages do not expose one fixed C++ signature. The MACA 3.2.1 extension appends s_aux and also passes an attention mask to mha_fwd, while the split backend assumed the standard signature and failed to import _infinicore with an undefined symbol. Reuse symbol-based ABI detection for MetaX and Mars, but map the result to independent MetaX ABI defines so Mars configuration remains isolated. Pass the optional attention-mask and auxiliary arguments only when the selected extension exports them. --- .../adaptor/flash_attention_adaptor.hpp | 16 ++-- .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 4 +- .../ops/multi_head_attention/mha_flashattn.cc | 8 +- .../mha_varlen_flashattn.cc | 4 +- xmake.lua | 75 ++++++++++++++----- 5 files changed, 72 insertions(+), 35 deletions(-) diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp index 4f7e61ec9..e20e1d9ef 100644 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ b/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -13,8 +13,8 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) std::optional &alibi_slopes_, // num_heads or batch_size x num_heads -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT - // Mars extensions accept an attention mask before the scalar options. +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) + // Extended device ABIs accept an attention mask before the scalar options. std::optional &attn_mask_, #endif const float p_dropout, @@ -25,8 +25,8 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num const float softcap, const bool return_softmax, std::optional gen_ -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT - // Mars extensions append an auxiliary tensor after the generator. +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) + // Some device extensions append an auxiliary tensor after the generator. , std::optional &s_aux_ #endif @@ -54,8 +54,8 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_hea const float softcap, const bool return_softmax, std::optional gen_ -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT - // Mars extensions append an auxiliary tensor. +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) + // Some device extensions append an auxiliary tensor. , std::optional &s_aux_ #endif @@ -129,8 +129,8 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size const float softcap, bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 int num_splits -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT - // Mars extensions append an auxiliary tensor. +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) + // Some device extensions append an auxiliary tensor. , std::optional &s_aux_ #endif diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index cedb0f45e..c1be0f0ff 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -84,7 +84,7 @@ void run(void *planned_meta) { auto out = use_dynamic_out ? std::optional(std::nullopt) : std::optional(out_tensor); -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) std::optional s_aux = std::nullopt; #endif @@ -109,7 +109,7 @@ void run(void *planned_meta) { 0.0f, false, 0 -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) , s_aux #endif diff --git a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc index 46ede944d..9f0d5c309 100644 --- a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc @@ -69,8 +69,10 @@ void run(void *planned_meta) { auto scale = p->scale; auto is_causal = p->is_causal; -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) std::optional attn_mask = std::nullopt; +#endif +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) std::optional s_aux = std::nullopt; #endif @@ -81,7 +83,7 @@ void run(void *planned_meta) { v, out, alibi_slopes, -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) attn_mask, #endif 0.0, @@ -92,7 +94,7 @@ void run(void *planned_meta) { 0.0, false, std::nullopt -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) , s_aux #endif diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc index 5842bb487..95e50918a 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc @@ -138,7 +138,7 @@ void run(void *planned_meta) { auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; auto scale = p->scale; -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) std::optional s_aux = std::nullopt; #endif @@ -165,7 +165,7 @@ void run(void *planned_meta) { 0.0, false, std::nullopt -#ifdef INFINICORE_FLASH_ATTN_MARS_EXT +#if defined(INFINICORE_FLASH_ATTN_MARS_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_EXT) || defined(INFINICORE_FLASH_ATTN_METAX_S_AUX) , s_aux #endif diff --git a/xmake.lua b/xmake.lua index 73075739d..df37ba316 100644 --- a/xmake.lua +++ b/xmake.lua @@ -355,12 +355,12 @@ local function validate_infinirt_root(xmake_os) find_infinirt_library(infinirt_root, xmake_os) end -local function find_mars_flash_attn_library(xmake_os) +local function find_flash_attn_library(xmake_os) local configured = os.getenv("FLASH_ATTN_2_CUDA_SO") if configured and configured ~= "" then configured = configured:trim() if not xmake_os.isfile(configured) then - xmake_os.raise("Mars flash-attn library not found: " .. configured) + xmake_os.raise("flash-attn library not found: " .. configured) end return configured end @@ -372,25 +372,23 @@ local function find_mars_flash_attn_library(xmake_os) return discovered end - local fallback = os.getenv("FLASH_ATTN_MARS_CUDA_SO_CONTAINER") - or "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so" + local fallback + if has_config("mars-gpu") then + fallback = os.getenv("FLASH_ATTN_MARS_CUDA_SO_CONTAINER") + or "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-aarch64-linux-gnu.so" + else + fallback = "/opt/conda/lib/python3.10/site-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so" + end if not xmake_os.isfile(fallback) then - xmake_os.raise("Mars flash-attn library not found; set FLASH_ATTN_2_CUDA_SO") + xmake_os.raise("flash-attn library not found; set FLASH_ATTN_2_CUDA_SO") end return fallback end --- HPCC and flash-attn are versioned independently, so inspect the extension's --- exported C++ signatures instead of deriving its ABI from the SDK version. -local function mars_flash_attn_uses_extended_abi(xmake_os) - local abi = get_config("mars-flash-attn-abi") or "detect" - if abi == "extended" then - return true - elseif abi == "standard" then - return false - end - - local flash_attn_library = find_mars_flash_attn_library(xmake_os) +-- Device SDK and flash-attn are versioned independently, so inspect the +-- extension's exported C++ signatures instead of deriving its ABI from the SDK. +local function detect_flash_attn_abi(xmake_os) + local flash_attn_library = find_flash_attn_library(xmake_os) local symbols = xmake_os.iorunv("nm", {"-D", "-C", flash_attn_library}) local mha_fwd_symbol = nil local mha_varlen_fwd_symbol = nil @@ -429,23 +427,49 @@ local function mars_flash_attn_uses_extended_abi(xmake_os) and varlen_has_s_aux and kvcache_has_s_aux then - return true + return "extended" + elseif not has_attention_mask + and fwd_has_s_aux + and varlen_has_s_aux + and kvcache_has_s_aux + then + return "s_aux" elseif not has_attention_mask and mha_fwd_symbol:find("std::optional)", 1, true) and mha_varlen_fwd_symbol:find("std::optional)", 1, true) and mha_fwd_kvcache_symbol:match(", int%)$") then - return false + return "standard" end end xmake_os.raise( - "Unable to detect the Mars flash-attn ABI from " + "Unable to detect the flash-attn ABI from " .. flash_attn_library - .. "; set --mars-flash-attn-abi=standard or extended" ) end +local function mars_flash_attn_uses_extended_abi(xmake_os) + local abi = get_config("mars-flash-attn-abi") or "detect" + if abi == "extended" then + return true + elseif abi == "standard" then + return false + end + + local detected = detect_flash_attn_abi(xmake_os) + if detected == "extended" then + return true + elseif detected == "standard" then + return false + end + xmake_os.raise("Mars flash-attn does not support detected ABI: " .. detected) +end + +local function metax_flash_attn_abi(xmake_os) + return detect_flash_attn_abi(xmake_os) +end + local function add_external_infinirt() local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") local infinirt_root = get_infinirt_root() @@ -804,6 +828,17 @@ target("infinicore_cpp_api") then target:add("defines", "INFINICORE_FLASH_ATTN_MARS_EXT=1") end + if has_config("metax-gpu") + and get_config("flash-attn") + and get_config("flash-attn") ~= "" + then + local abi = metax_flash_attn_abi(os) + if abi == "extended" then + target:add("defines", "INFINICORE_FLASH_ATTN_METAX_EXT=1") + elseif abi == "s_aux" then + target:add("defines", "INFINICORE_FLASH_ATTN_METAX_S_AUX=1") + end + end end) local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")