kkio
/
Intro-ops
forked from ccf-ai-infra/Intro-ops
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add MACA-based MetaX backend support 

Co-authored-by: wawahejun <hejunlbbc@gmail.com>
This commit is contained in:
yutianyu 2026-05-06 04:08:49 +00:00
parent ec8ed33b27
commit 6e02c3eb55
32 changed files with 1005 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
CMakeLists.txt
View File

@ -10,14 +10,37 @@ set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
list(PREPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
if(CAMP_ENABLE_NVIDIA)
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
set(CMAKE_CUDA_ARCHITECTURES native)
if(CAMP_ENABLE_NVIDIA AND CAMP_ENABLE_METAX)
message(FATAL_ERROR "CAMP_ENABLE_NVIDIA and CAMP_ENABLE_METAX cannot both be ON; build separate variants")
endif()
if(CAMP_ENABLE_METAX)
find_package(MetaX REQUIRED)
set(ENV{CUCC_PATH} "${MetaX_CUBRIDGE_ROOT}")
set(ENV{CUDA_PATH} "${MetaX_CUBRIDGE_ROOT}")
list(APPEND CMAKE_MODULE_PATH "${MetaX_CMAKE_MODULE_DIR}")
set(CMAKE_MACA_SOURCE_FILE_EXTENSIONS maca;cu;cc;cpp;c)
endif()
if(CAMP_ENABLE_NVIDIA OR CAMP_ENABLE_METAX)
if(CAMP_ENABLE_METAX)
enable_language(MACA)
else()
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
set(CMAKE_CUDA_ARCHITECTURES native)
endif()
enable_language(CUDA)
include(cmake/cuda_helpers.cmake)
endif()
enable_language(CUDA)
include(cmake/cuda_helpers.cmake)
endif()
if(CAMP_ENABLE_METAX)
set(CAMP_METAX_COMPILE_OPTIONS
-D__MACA_NO_HALF_OPERATORS__
-D__FAST_BLOCK_RED__
)
endif()
add_subdirectory(ops)

15
README.md
View File

@ -28,10 +28,10 @@ python/
| Operator | NVIDIA C++ | TileLang | MetaX |
| --- | --- | --- | --- |
| `copy` | runnable | runnable when TileLang is installed | stub |
| `vector_add` | runnable | runnable when TileLang is installed | stub |
| `reduce_sum` | runnable, row-wise fp32 | runnable when TileLang is installed | stub |
| `softmax` | runnable, row-wise fp32 | runnable when TileLang is installed | stub |
| `copy` | runnable | runnable when TileLang is installed | runnable |
| `vector_add` | runnable | runnable when TileLang is installed | runnable |
| `reduce_sum` | runnable, row-wise fp32 | runnable when TileLang is installed | runnable, row-wise fp32 |
| `softmax` | runnable, row-wise fp32 | runnable when TileLang is installed | runnable, row-wise fp32 |
## Setup
@ -48,6 +48,10 @@ cmake .. -DCAMP_ENABLE_NVIDIA=ON -DCAMP_ENABLE_METAX=OFF
cmake --build . -j$(nproc)
```
```bash
./scripts/build_metax.sh build
```
## Validate
```bash
@ -55,9 +59,10 @@ python tests/run_ops.py --op copy --backend nvidia --mode all
CAMP_BUILD_DIR=build pytest tests/ -v --backend nvidia
pytest tests/ -v --backend tilelang
python tests/run_ops.py --op all --backend nvidia --mode bench
./scripts/build_metax.sh test
```
The TileLang backend requires the `tilelang` Python package.
The TileLang backend requires the `tilelang` Python package. The MetaX backend is built as a separate variant, should use `CAMP_ENABLE_NVIDIA=OFF`, and stores backend sources under `ops/*/metax/*.maca`.
## Production Mapping

15
README.zh.md
View File

@ -27,10 +27,10 @@ python/
| 算子 | NVIDIA C++ | TileLang | MetaX |
| --- | --- | --- | --- |
| `copy` | 可运行 | 安装 TileLang 后可运行 | stub |
| `vector_add` | 可运行 | 安装 TileLang 后可运行 | stub |
| `reduce_sum` | 可运行，row-wise fp32 | 安装 TileLang 后可运行 | stub |
| `softmax` | 可运行，row-wise fp32 | 安装 TileLang 后可运行 | stub |
| `copy` | 可运行 | 安装 TileLang 后可运行 | 可运行 |
| `vector_add` | 可运行 | 安装 TileLang 后可运行 | 可运行 |
| `reduce_sum` | 可运行，row-wise fp32 | 安装 TileLang 后可运行 | 可运行，row-wise fp32 |
| `softmax` | 可运行，row-wise fp32 | 安装 TileLang 后可运行 | 可运行，row-wise fp32 |
## 安装
@ -47,6 +47,10 @@ cmake .. -DCAMP_ENABLE_NVIDIA=ON -DCAMP_ENABLE_METAX=OFF
cmake --build . -j$(nproc)
```
```bash
./scripts/build_metax.sh build
```
## 验证
```bash
@ -54,9 +58,10 @@ python tests/run_ops.py --op copy --backend nvidia --mode all
CAMP_BUILD_DIR=build pytest tests/ -v --backend nvidia
pytest tests/ -v --backend tilelang
python tests/run_ops.py --op all --backend nvidia --mode bench
./scripts/build_metax.sh test
```
TileLang 后端需要安装 `tilelang` Python 包。
TileLang 后端需要安装 `tilelang` Python 包。MetaX 后端使用独立构建产物，构建时应关闭 NVIDIA 变体，并将后端源码放在 `ops/*/metax/*.maca`。
## 生产映射

16
cmake/CMakeMACAInformation.cmake Normal file
View File

@ -0,0 +1,16 @@
set(CMAKE_MACA_COMPILE_OBJECT
"<CMAKE_MACA_COMPILER> <DEFINES> <INCLUDES> <FLAGS> -o <OBJECT> -c <SOURCE>"
)
set(CMAKE_INCLUDE_FLAG_MACA "-I")
set(CMAKE_MACA_COMPILE_OPTIONS_PIC "-fPIC")
set(CMAKE_MACA_CREATE_SHARED_LIBRARY
"<CMAKE_MACA_COMPILER> <LANGUAGE_COMPILE_FLAGS> <LINK_FLAGS> -shared -o <TARGET> <OBJECTS> <LINK_LIBRARIES> --maca-path=$ENV{MACA_PATH}"
)
set(CMAKE_MACA_LINK_EXECUTABLE
"<CMAKE_MACA_COMPILER> <FLAGS> <LINK_FLAGS> -o <TARGET> <OBJECTS> <LINK_LIBRARIES> --maca-path=$ENV{MACA_PATH}"
)
set(CMAKE_MACA_INFORMATION_LOADED 1)

86
cmake/FindMetaX.cmake
View File

@ -1,6 +1,82 @@
# Placeholder for a future MetaX toolchain integration.
#
# A production implementation should detect MACA_PATH, htcc/mxcc, runtime
# libraries, include directories, and backend-specific source suffixes.
set(MetaX_FOUND OFF)
set(_MetaX_SEARCH_ROOTS "")
if(DEFINED ENV{MACA_PATH})
list(APPEND _MetaX_SEARCH_ROOTS "$ENV{MACA_PATH}")
endif()
list(APPEND _MetaX_SEARCH_ROOTS /opt/maca)
find_path(
MetaX_ROOT
NAMES include/mcr/mc_runtime.h
PATHS ${_MetaX_SEARCH_ROOTS}
NO_DEFAULT_PATH
)
find_program(
MetaX_MXCC_EXECUTABLE
NAMES mxcc
PATHS
${MetaX_ROOT}/mxgpu_llvm/bin
/opt/maca/mxgpu_llvm/bin
)
find_program(
MetaX_CUCC_EXECUTABLE
NAMES cucc
PATHS
${MetaX_ROOT}/tools/cu-bridge/bin
/opt/maca/tools/cu-bridge/bin
)
find_path(
MetaX_INCLUDE_DIR
NAMES mcr/mc_runtime.h
PATHS
${MetaX_ROOT}/include
/opt/maca/include
)
find_library(
MetaX_MCRUNTIME_LIBRARY
NAMES mcruntime libmcruntime.so
PATHS
${MetaX_ROOT}/lib
/opt/maca/lib
)
find_library(
MetaX_MXCRT_LIBRARY
NAMES mxc-runtime64 libmxc-runtime64.so
PATHS
${MetaX_ROOT}/lib
/opt/maca/lib
)
include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(
MetaX
REQUIRED_VARS
MetaX_ROOT
MetaX_MXCC_EXECUTABLE
MetaX_INCLUDE_DIR
MetaX_MCRUNTIME_LIBRARY
)
if(MetaX_FOUND)
set(MetaX_CUBRIDGE_ROOT "${MetaX_ROOT}/tools/cu-bridge")
set(MetaX_CMAKE_MODULE_DIR "${MetaX_CUBRIDGE_ROOT}/cmake_module/maca")
set(MetaX_INCLUDE_DIRS "${MetaX_INCLUDE_DIR}")
set(MetaX_LIBRARIES "${MetaX_MCRUNTIME_LIBRARY}")
if(MetaX_MXCRT_LIBRARY)
list(APPEND MetaX_LIBRARIES "${MetaX_MXCRT_LIBRARY}")
endif()
endif()
mark_as_advanced(
MetaX_ROOT
MetaX_MXCC_EXECUTABLE
MetaX_CUCC_EXECUTABLE
MetaX_INCLUDE_DIR
MetaX_MCRUNTIME_LIBRARY
MetaX_MXCRT_LIBRARY
)

2
examples/01_copy.py
View File

@ -15,7 +15,7 @@ from operator_runtime import copy
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang"])
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
src = torch.randn((1024,), device="cuda", dtype=torch.float32)

2
examples/02_vector_add.py
View File

@ -15,7 +15,7 @@ from operator_runtime import vector_add
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang"])
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
a = torch.randn((1024,), device="cuda", dtype=torch.float32)

2
examples/03_reduce_sum.py
View File

@ -15,7 +15,7 @@ from operator_runtime import reduce_sum
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang"])
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
src = torch.randn((32, 128), device="cuda", dtype=torch.float32)

2
examples/04_softmax.py
View File

@ -15,7 +15,7 @@ from operator_runtime import softmax
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang"])
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
src = torch.randn((32, 128), device="cuda", dtype=torch.float32)

18
ops/CMakeLists.txt
View File

@ -1,5 +1,6 @@
# Convention: operator CUDA sources auto-discovered by directory layout:
# ops/<operator>/nvidia/*.cu -> compiled when CAMP_ENABLE_NVIDIA is ON
# ops/<operator>/metax/*.maca -> compiled when CAMP_ENABLE_METAX is ON
# After adding a new operator, re-run cmake to pick up new files.
set(CAMP_COMMON_SOURCES
@ -13,7 +14,15 @@ if(CAMP_ENABLE_NVIDIA)
)
endif()
set(CAMP_OPERATOR_SOURCES ${CAMP_COMMON_SOURCES} ${CAMP_NVIDIA_SOURCES})
set(CAMP_METAX_SOURCES "")
if(CAMP_ENABLE_METAX)
file(GLOB_RECURSE CAMP_METAX_SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/*/metax/*.maca
)
set_source_files_properties(${CAMP_METAX_SOURCES} PROPERTIES LANGUAGE MACA)
endif()
set(CAMP_OPERATOR_SOURCES ${CAMP_COMMON_SOURCES} ${CAMP_NVIDIA_SOURCES} ${CAMP_METAX_SOURCES})
add_library(camp_ops SHARED ${CAMP_OPERATOR_SOURCES})
target_include_directories(camp_ops PUBLIC
@ -21,6 +30,13 @@ target_include_directories(camp_ops PUBLIC
"${CMAKE_CURRENT_SOURCE_DIR}/../include"
)
if(CAMP_ENABLE_METAX)
target_include_directories(camp_ops PRIVATE ${MetaX_INCLUDE_DIRS})
target_link_libraries(camp_ops PRIVATE ${MetaX_LIBRARIES})
target_compile_options(camp_ops PRIVATE ${CAMP_METAX_COMPILE_OPTIONS})
set_target_properties(camp_ops PROPERTIES LINKER_LANGUAGE MACA)
endif()
target_compile_definitions(camp_ops PRIVATE
$<$<BOOL:${CAMP_ENABLE_NVIDIA}>:CAMP_ENABLE_NVIDIA=1>
$<$<BOOL:${CAMP_ENABLE_METAX}>:CAMP_ENABLE_METAX=1>

19
ops/common/metax/README.md
View File

@ -1,13 +1,12 @@
# MetaX Backend Stub
# MetaX Backend
The training runtime keeps MetaX interfaces aligned with NVIDIA interfaces, but
does not build MetaX code by default.
The training runtime keeps MetaX interfaces aligned with the compiled C ABI
backends and builds MetaX as a separate variant through MACA/cu-bridge.
A production integration should add:
- MACA SDK discovery (`MACA_PATH`)
- htcc/mxcc compiler rules
- `.maca` source handling
- hcruntime/mcruntime link rules
- backend-specific CI on MetaX hardware
Current expectations:
- discover the MACA SDK from `MACA_PATH` or `/opt/maca`
- build with `cmake_maca` and `ninja_maca`
- compile `ops/*/metax/*.maca` into `libcamp_ops.so`
- load `_metax` ABI symbols from Python FFI
- validate on a MetaX-enabled PyTorch environment that exposes `cuda:0`

7
ops/copy/metax/README.md
View File

@ -1,5 +1,4 @@
# Copy MetaX Stub
The copy ABI mirrors the NVIDIA lifecycle. Add `.maca` sources and htcc/mxcc
rules when MetaX hardware is available.
# Copy MetaX Backend
The copy MetaX backend exports the standard descriptor lifecycle with `_metax`
symbols and is built from `ops/copy/metax/copy_metax.maca`.

27
ops/copy/metax/copy_metax.h
View File

@ -2,5 +2,30 @@
#include "operator_runtime/api.h"
// MetaX stub: keep the ABI contract visible while the real toolchain is absent.
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_copy_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *dst,
const oprt_tensor_view_t *src);
OPRT_EXPORT oprt_status_t oprt_get_copy_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_copy_metax(
oprt_operator_descriptor_t desc,
void *workspace,
size_t workspace_size,
void *dst,
const void *src,
oprt_stream_t stream);
OPRT_EXPORT oprt_status_t oprt_destroy_copy_descriptor_metax(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

119
ops/copy/metax/copy_metax.maca Normal file
View File

@ -0,0 +1,119 @@
#include "ops/copy/metax/copy_metax.h"
#include "operator_runtime/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/elementwise.h"
#include "operator_runtime/tensor_checks.h"
#include <cuda_fp16.h>
namespace {
struct CopyDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t dst_view;
oprt_tensor_view_t src_view;
int64_t elements = 0;
const char *op_name() const override {
return "copy";
}
};
template <typename T>
__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
for (int64_t i = idx; i < n; i += stride) {
dst[i] = src[i];
}
}
template <typename T>
oprt_status_t launch_copy(const CopyDescriptor *desc, void *dst, const void *src, oprt_stream_t stream) {
constexpr int threads = 256;
int blocks = oprt::blocks_for(desc->elements, threads);
cudaStream_t s = oprt::as_cuda_stream(stream);
copy_contiguous_kernel<T><<<blocks, threads, 0, s>>>(
static_cast<T *>(dst), static_cast<const T *>(src), desc->elements);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_copy_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *dst,
const oprt_tensor_view_t *src) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(dst);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(src);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*dst, *src);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_shape(*dst, *src);
if (status != OPRT_SUCCESS) {
return status;
}
if (!oprt::elementwise_fast_path(*dst, *src)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new CopyDescriptor();
typed->dst_view = *dst;
typed->src_view = *src;
typed->elements = oprt::numel(*dst);
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_copy_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size) {
if (desc == nullptr || size == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*size = desc->workspace_size;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_copy_metax(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *dst,
const void *src,
oprt_stream_t stream) {
if (desc == nullptr || dst == nullptr || src == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const CopyDescriptor *>(desc);
switch (typed->dst_view.dtype) {
case OPRT_DTYPE_F16:
return launch_copy<half>(typed, dst, src, stream);
case OPRT_DTYPE_F32:
return launch_copy<float>(typed, dst, src, stream);
default:
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_copy_descriptor_metax(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

7
ops/reduce_sum/metax/README.md
View File

@ -1,5 +1,4 @@
# Reduce Sum MetaX Stub
The reduce_sum ABI mirrors the NVIDIA lifecycle. Add `.maca` sources and
htcc/mxcc rules when MetaX hardware is available.
# Reduce Sum MetaX Backend
The reduce_sum MetaX backend exports the standard descriptor lifecycle with
`_metax` symbols and is built from `ops/reduce_sum/metax/reduce_sum_metax.maca`.

28
ops/reduce_sum/metax/reduce_sum_metax.h
View File

@ -2,5 +2,31 @@
#include "operator_runtime/api.h"
// MetaX stub: keep the ABI contract visible while the real toolchain is absent.
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_reduce_sum_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis);
OPRT_EXPORT oprt_status_t oprt_get_reduce_sum_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_reduce_sum_metax(
oprt_operator_descriptor_t desc,
void *workspace,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream);
OPRT_EXPORT oprt_status_t oprt_destroy_reduce_sum_descriptor_metax(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

132
ops/reduce_sum/metax/reduce_sum_metax.maca Normal file
View File

@ -0,0 +1,132 @@
#include "ops/reduce_sum/metax/reduce_sum_metax.h"
#include "operator_runtime/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/tensor_checks.h"
#include <cuda_runtime.h>
namespace {
struct ReduceSumDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t in_view;
int64_t rows = 0;
int64_t cols = 0;
int64_t axis = 1;
const char *op_name() const override {
return "reduce_sum";
}
};
__global__ void reduce_sum_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
int64_t row = blockIdx.x;
if (row >= rows) {
return;
}
const float *row_ptr = in + row * cols;
float partial = 0.0f;
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
partial += row_ptr[col];
}
extern __shared__ float shared[];
shared[threadIdx.x] = partial;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared[threadIdx.x] += shared[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
out[row] = shared[0];
}
}
bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
return in.dtype == OPRT_DTYPE_F32 &&
out.dtype == OPRT_DTYPE_F32 &&
in.ndim == 2 &&
out.ndim == 1 &&
axis == 1 &&
out.shape[0] == in.shape[0] &&
oprt::is_contiguous(in) &&
oprt::is_contiguous(out);
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_reduce_sum_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(in);
if (status != OPRT_SUCCESS) {
return status;
}
if (!is_rowwise_case(*out, *in, axis)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new ReduceSumDescriptor();
typed->out_view = *out;
typed->in_view = *in;
typed->rows = in->shape[0];
typed->cols = in->shape[1];
typed->axis = axis;
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_reduce_sum_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size) {
if (desc == nullptr || size == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*size = desc->workspace_size;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_reduce_sum_metax(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || in == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const ReduceSumDescriptor *>(desc);
cudaStream_t s = oprt::as_cuda_stream(stream);
constexpr int threads = 256;
reduce_sum_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), s>>>(
static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_reduce_sum_descriptor_metax(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

7
ops/softmax/metax/README.md
View File

@ -1,5 +1,4 @@
# Softmax MetaX Stub
The softmax ABI mirrors the NVIDIA lifecycle. Add `.maca` sources and htcc/mxcc
rules when MetaX hardware is available.
# Softmax MetaX Backend
The softmax MetaX backend exports the standard descriptor lifecycle with
`_metax` symbols and is built from `ops/softmax/metax/softmax_metax.maca`.

28
ops/softmax/metax/softmax_metax.h
View File

@ -2,5 +2,31 @@
#include "operator_runtime/api.h"
// MetaX stub: keep the ABI contract visible while the real toolchain is absent.
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_softmax_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis);
OPRT_EXPORT oprt_status_t oprt_get_softmax_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_softmax_metax(
oprt_operator_descriptor_t desc,
void *workspace,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream);
OPRT_EXPORT oprt_status_t oprt_destroy_softmax_descriptor_metax(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

155
ops/softmax/metax/softmax_metax.maca Normal file
View File

@ -0,0 +1,155 @@
#include "ops/softmax/metax/softmax_metax.h"
#include "operator_runtime/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/tensor_checks.h"
#include <cuda_runtime.h>
#include <float.h>
#include <math.h>
namespace {
struct SoftmaxDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t in_view;
int64_t rows = 0;
int64_t cols = 0;
int64_t axis = 1;
const char *op_name() const override {
return "softmax";
}
};
__global__ void softmax_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
int64_t row = blockIdx.x;
if (row >= rows) {
return;
}
const float *row_in = in + row * cols;
float *row_out = out + row * cols;
float local_max = -FLT_MAX;
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
local_max = fmaxf(local_max, row_in[col]);
}
extern __shared__ float shared[];
shared[threadIdx.x] = local_max;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared[threadIdx.x] = fmaxf(shared[threadIdx.x], shared[threadIdx.x + stride]);
}
__syncthreads();
}
float row_max = shared[0];
float local_sum = 0.0f;
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
float value = expf(row_in[col] - row_max);
row_out[col] = value;
local_sum += value;
}
shared[threadIdx.x] = local_sum;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared[threadIdx.x] += shared[threadIdx.x + stride];
}
__syncthreads();
}
float row_sum = shared[0];
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
row_out[col] /= row_sum;
}
}
bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
return in.dtype == OPRT_DTYPE_F32 &&
out.dtype == OPRT_DTYPE_F32 &&
in.ndim == 2 &&
out.ndim == 2 &&
axis == 1 &&
oprt::same_shape(out, in) &&
oprt::is_contiguous(in) &&
oprt::is_contiguous(out);
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_softmax_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(in);
if (status != OPRT_SUCCESS) {
return status;
}
if (!is_rowwise_case(*out, *in, axis)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new SoftmaxDescriptor();
typed->out_view = *out;
typed->in_view = *in;
typed->rows = in->shape[0];
typed->cols = in->shape[1];
typed->axis = axis;
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_softmax_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size) {
if (desc == nullptr || size == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*size = desc->workspace_size;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_softmax_metax(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || in == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const SoftmaxDescriptor *>(desc);
cudaStream_t s = oprt::as_cuda_stream(stream);
constexpr int threads = 256;
softmax_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), s>>>(
static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_softmax_descriptor_metax(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

7
ops/vector_add/metax/README.md
View File

@ -1,5 +1,4 @@
# Vector Add MetaX Stub
The vector_add ABI mirrors the NVIDIA lifecycle. Add `.maca` sources and
htcc/mxcc rules when MetaX hardware is available.
# Vector Add MetaX Backend
The vector_add MetaX backend exports the standard descriptor lifecycle with
`_metax` symbols and is built from `ops/vector_add/metax/vector_add_metax.maca`.

29
ops/vector_add/metax/vector_add_metax.h
View File

@ -2,5 +2,32 @@
#include "operator_runtime/api.h"
// MetaX stub: keep the ABI contract visible while the real toolchain is absent.
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_vector_add_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *a,
const oprt_tensor_view_t *b);
OPRT_EXPORT oprt_status_t oprt_get_vector_add_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_vector_add_metax(
oprt_operator_descriptor_t desc,
void *workspace,
size_t workspace_size,
void *out,
const void *a,
const void *b,
oprt_stream_t stream);
OPRT_EXPORT oprt_status_t oprt_destroy_vector_add_descriptor_metax(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

141
ops/vector_add/metax/vector_add_metax.maca Normal file
View File

@ -0,0 +1,141 @@
#include "ops/vector_add/metax/vector_add_metax.h"
#include "operator_runtime/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/elementwise.h"
#include "operator_runtime/tensor_checks.h"
#include <cuda_fp16.h>
namespace {
struct VectorAddDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t a_view;
oprt_tensor_view_t b_view;
int64_t elements = 0;
const char *op_name() const override {
return "vector_add";
}
};
template <typename T>
__device__ T add_values(T a, T b) {
return a + b;
}
template <>
__device__ inline half add_values<half>(half a, half b) {
return __hadd(a, b);
}
template <typename T>
__global__ void vector_add_contiguous_kernel(
T *__restrict__ out,
const T *__restrict__ a,
const T *__restrict__ b,
int64_t n) {
int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
for (int64_t i = idx; i < n; i += stride) {
out[i] = add_values(a[i], b[i]);
}
}
template <typename T>
oprt_status_t launch_vector_add(const VectorAddDescriptor *desc, void *out, const void *a, const void *b, oprt_stream_t stream) {
constexpr int threads = 256;
int blocks = oprt::blocks_for(desc->elements, threads);
cudaStream_t s = oprt::as_cuda_stream(stream);
vector_add_contiguous_kernel<T><<<blocks, threads, 0, s>>>(
static_cast<T *>(out), static_cast<const T *>(a), static_cast<const T *>(b), desc->elements);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_vector_add_descriptor_metax(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *a,
const oprt_tensor_view_t *b) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(a);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(b);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*out, *a);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*out, *b);
if (status != OPRT_SUCCESS) {
return status;
}
if (!oprt::elementwise_fast_path(*out, *a, *b)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new VectorAddDescriptor();
typed->out_view = *out;
typed->a_view = *a;
typed->b_view = *b;
typed->elements = oprt::numel(*out);
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_vector_add_workspace_size_metax(
oprt_operator_descriptor_t desc,
size_t *size) {
if (desc == nullptr || size == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*size = desc->workspace_size;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_vector_add_metax(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *a,
const void *b,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || a == nullptr || b == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const VectorAddDescriptor *>(desc);
switch (typed->out_view.dtype) {
case OPRT_DTYPE_F16:
return launch_vector_add<half>(typed, out, a, b, stream);
case OPRT_DTYPE_F32:
return launch_vector_add<float>(typed, out, a, b, stream);
default:
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_vector_add_descriptor_metax(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

93
python/operator_runtime/_internal/bindings.py
View File

@ -4,6 +4,8 @@ import ctypes
from dataclasses import dataclass
from typing import Callable
from operator_runtime.backend import Backend, normalize_backend
from .loader import load_library
from .tensor_view import TensorView
@ -33,17 +35,40 @@ class CFunctions:
destroy: Callable
def bind_unary(name: str) -> CFunctions:
def _backend_suffix(backend: str | Backend) -> str:
return normalize_backend(backend).value
def _missing_symbol_error(name: str, backend: str, symbol: str, exc: AttributeError) -> OperatorRuntimeError:
raise OperatorRuntimeError(
f"backend {backend} is unavailable in the loaded libcamp_ops.so; missing symbol {symbol}"
) from exc
def bind_unary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions:
lib = load_library()
create = getattr(lib, f"oprt_create_{name}_descriptor_nvidia")
suffix = _backend_suffix(backend)
create_symbol = f"oprt_create_{name}_descriptor_{suffix}"
try:
create = getattr(lib, create_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, create_symbol, exc)
create.argtypes = [ctypes.POINTER(Descriptor), ctypes.POINTER(TensorView), ctypes.POINTER(TensorView)]
create.restype = Status
workspace = getattr(lib, f"oprt_get_{name}_workspace_size_nvidia")
workspace_symbol = f"oprt_get_{name}_workspace_size_{suffix}"
try:
workspace = getattr(lib, workspace_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, workspace_symbol, exc)
workspace.argtypes = [Descriptor, ctypes.POINTER(ctypes.c_size_t)]
workspace.restype = Status
execute = getattr(lib, f"oprt_execute_{name}_nvidia")
execute_symbol = f"oprt_execute_{name}_{suffix}"
try:
execute = getattr(lib, execute_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, execute_symbol, exc)
execute.argtypes = [
Descriptor,
ctypes.c_void_p,
@ -54,15 +79,24 @@ def bind_unary(name: str) -> CFunctions:
]
execute.restype = Status
destroy = getattr(lib, f"oprt_destroy_{name}_descriptor_nvidia")
destroy_symbol = f"oprt_destroy_{name}_descriptor_{suffix}"
try:
destroy = getattr(lib, destroy_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, destroy_symbol, exc)
destroy.argtypes = [Descriptor]
destroy.restype = Status
return CFunctions(create, workspace, execute, destroy)
def bind_binary(name: str) -> CFunctions:
def bind_binary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions:
lib = load_library()
create = getattr(lib, f"oprt_create_{name}_descriptor_nvidia")
suffix = _backend_suffix(backend)
create_symbol = f"oprt_create_{name}_descriptor_{suffix}"
try:
create = getattr(lib, create_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, create_symbol, exc)
create.argtypes = [
ctypes.POINTER(Descriptor),
ctypes.POINTER(TensorView),
@ -71,11 +105,19 @@ def bind_binary(name: str) -> CFunctions:
]
create.restype = Status
workspace = getattr(lib, f"oprt_get_{name}_workspace_size_nvidia")
workspace_symbol = f"oprt_get_{name}_workspace_size_{suffix}"
try:
workspace = getattr(lib, workspace_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, workspace_symbol, exc)
workspace.argtypes = [Descriptor, ctypes.POINTER(ctypes.c_size_t)]
workspace.restype = Status
execute = getattr(lib, f"oprt_execute_{name}_nvidia")
execute_symbol = f"oprt_execute_{name}_{suffix}"
try:
execute = getattr(lib, execute_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, execute_symbol, exc)
execute.argtypes = [
Descriptor,
ctypes.c_void_p,
@ -87,15 +129,24 @@ def bind_binary(name: str) -> CFunctions:
]
execute.restype = Status
destroy = getattr(lib, f"oprt_destroy_{name}_descriptor_nvidia")
destroy_symbol = f"oprt_destroy_{name}_descriptor_{suffix}"
try:
destroy = getattr(lib, destroy_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, destroy_symbol, exc)
destroy.argtypes = [Descriptor]
destroy.restype = Status
return CFunctions(create, workspace, execute, destroy)
def bind_reduce_like(name: str) -> CFunctions:
def bind_reduce_like(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions:
lib = load_library()
create = getattr(lib, f"oprt_create_{name}_descriptor_nvidia")
suffix = _backend_suffix(backend)
create_symbol = f"oprt_create_{name}_descriptor_{suffix}"
try:
create = getattr(lib, create_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, create_symbol, exc)
create.argtypes = [
ctypes.POINTER(Descriptor),
ctypes.POINTER(TensorView),
@ -104,11 +155,19 @@ def bind_reduce_like(name: str) -> CFunctions:
]
create.restype = Status
workspace = getattr(lib, f"oprt_get_{name}_workspace_size_nvidia")
workspace_symbol = f"oprt_get_{name}_workspace_size_{suffix}"
try:
workspace = getattr(lib, workspace_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, workspace_symbol, exc)
workspace.argtypes = [Descriptor, ctypes.POINTER(ctypes.c_size_t)]
workspace.restype = Status
execute = getattr(lib, f"oprt_execute_{name}_nvidia")
execute_symbol = f"oprt_execute_{name}_{suffix}"
try:
execute = getattr(lib, execute_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, execute_symbol, exc)
execute.argtypes = [
Descriptor,
ctypes.c_void_p,
@ -119,7 +178,11 @@ def bind_reduce_like(name: str) -> CFunctions:
]
execute.restype = Status
destroy = getattr(lib, f"oprt_destroy_{name}_descriptor_nvidia")
destroy_symbol = f"oprt_destroy_{name}_descriptor_{suffix}"
try:
destroy = getattr(lib, destroy_symbol)
except AttributeError as exc:
_missing_symbol_error(name, suffix, destroy_symbol, exc)
destroy.argtypes = [Descriptor]
destroy.restype = Status
return CFunctions(create, workspace, execute, destroy)

2
python/operator_runtime/_internal/loader.py
View File

@ -17,6 +17,8 @@ def _candidate_library_paths() -> list[Path]:
[
repo_root / "build" / "libcamp_ops.so",
repo_root / "build" / "ops" / "libcamp_ops.so",
repo_root / "build-metax" / "libcamp_ops.so",
repo_root / "build-metax" / "ops" / "libcamp_ops.so",
]
)
return candidates

4
python/operator_runtime/ops/copy.py
View File

@ -27,10 +27,10 @@ def prepare_copy(out: torch.Tensor, src: torch.Tensor, backend: str | Backend =
from ops.copy.tilelang.copy_tl import prepare_copy_tl
return prepare_copy_tl(out, src)
if backend is not Backend.NVIDIA:
if backend not in (Backend.NVIDIA, Backend.METAX):
raise NotImplementedError(f"backend {backend.value} is not runnable")
funcs = bind_unary("copy")
funcs = bind_unary("copy", backend)
out_view = tensor_view(out)
src_view = tensor_view(src)
create_args = (ctypes.byref(out_view), ctypes.byref(src_view))

4
python/operator_runtime/ops/reduce_sum.py
View File

@ -34,10 +34,10 @@ def prepare_reduce_sum(
from ops.reduce_sum.tilelang.reduce_sum_tl import prepare_reduce_sum_tl
return prepare_reduce_sum_tl(out, src, dim=dim)
if backend is not Backend.NVIDIA:
if backend not in (Backend.NVIDIA, Backend.METAX):
raise NotImplementedError(f"backend {backend.value} is not runnable")
funcs = bind_reduce_like("reduce_sum")
funcs = bind_reduce_like("reduce_sum", backend)
out_view = tensor_view(out)
src_view = tensor_view(src)
create_args = (ctypes.byref(out_view), ctypes.byref(src_view), ctypes.c_int64(dim))

4
python/operator_runtime/ops/softmax.py
View File

@ -34,10 +34,10 @@ def prepare_softmax(
from ops.softmax.tilelang.softmax_tl import prepare_softmax_tl
return prepare_softmax_tl(out, src, dim=dim)
if backend is not Backend.NVIDIA:
if backend not in (Backend.NVIDIA, Backend.METAX):
raise NotImplementedError(f"backend {backend.value} is not runnable")
funcs = bind_reduce_like("softmax")
funcs = bind_reduce_like("softmax", backend)
out_view = tensor_view(out)
src_view = tensor_view(src)
create_args = (ctypes.byref(out_view), ctypes.byref(src_view), ctypes.c_int64(dim))

4
python/operator_runtime/ops/vector_add.py
View File

@ -32,10 +32,10 @@ def prepare_vector_add(
from ops.vector_add.tilelang.vector_add_tl import prepare_vector_add_tl
return prepare_vector_add_tl(out, a, b)
if backend is not Backend.NVIDIA:
if backend not in (Backend.NVIDIA, Backend.METAX):
raise NotImplementedError(f"backend {backend.value} is not runnable")
funcs = bind_binary("vector_add")
funcs = bind_binary("vector_add", backend)
out_view = tensor_view(out)
a_view = tensor_view(a)
b_view = tensor_view(b)

72
scripts/build_metax.sh Executable file
View File

@ -0,0 +1,72 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
BUILD_DIR="${BUILD_DIR:-${ROOT}/build-metax}"
MACA_PATH="${MACA_PATH:-/opt/maca}"
CUCC_PATH="${CUCC_PATH:-${MACA_PATH}/tools/cu-bridge}"
CUDA_PATH="${CUDA_PATH:-${CUCC_PATH}}"
PYTHON_BIN="${PYTHON_BIN:-python3}"
MODE="${1:-build}"
export MACA_PATH
export CUCC_PATH
export CUDA_PATH
export LD_LIBRARY_PATH="${MACA_PATH}/lib:${LD_LIBRARY_PATH:-}"
configure() {
mkdir -p "${BUILD_DIR}"
"${PYTHON_BIN}" -m cmake -G Ninja "${ROOT}" \
-B "${BUILD_DIR}" \
-DCAMP_ENABLE_NVIDIA=OFF \
-DCAMP_ENABLE_METAX=ON
}
build() {
"${PYTHON_BIN}" -m cmake --build "${BUILD_DIR}" -- -v
}
test_pytest() {
env CAMP_BUILD_DIR="${BUILD_DIR}" "${PYTHON_BIN}" -m pytest "${ROOT}/tests" -v --backend metax
}
test_run_ops() {
env CAMP_BUILD_DIR="${BUILD_DIR}" "${PYTHON_BIN}" "${ROOT}/tests/run_ops.py" --op all --backend metax --mode all
}
test_examples() {
env CAMP_BUILD_DIR="${BUILD_DIR}" "${PYTHON_BIN}" "${ROOT}/examples/01_copy.py" --backend metax
env CAMP_BUILD_DIR="${BUILD_DIR}" "${PYTHON_BIN}" "${ROOT}/examples/02_vector_add.py" --backend metax
env CAMP_BUILD_DIR="${BUILD_DIR}" "${PYTHON_BIN}" "${ROOT}/examples/03_reduce_sum.py" --backend metax
env CAMP_BUILD_DIR="${BUILD_DIR}" "${PYTHON_BIN}" "${ROOT}/examples/04_softmax.py" --backend metax
}
case "${MODE}" in
configure)
configure
;;
build)
configure
build
;;
test)
test_pytest
test_run_ops
test_examples
;;
all)
configure
build
test_pytest
test_run_ops
test_examples
;;
clean)
rm -rf "${BUILD_DIR}"
;;
*)
echo "usage: $0 {configure|build|test|all|clean}" >&2
exit 2
;;
esac

2
tests/op_tests/test_vector_add.py
View File

@ -45,7 +45,7 @@ def test_vector_add_api_contract(case, backend):
def test_prepared_vector_add_reuses_descriptor(backend):
require_cuda()
if backend != "nvidia":
if backend == "tilelang":
pytest.skip("descriptor lifecycle test targets C ABI backend")
a = torch.randn((1024,), device="cuda")
b = torch.randn((1024,), device="cuda")

2
tests/run_ops.py
View File

@ -110,7 +110,7 @@ def main() -> int:
ops = _discover_ops()
parser = argparse.ArgumentParser()
parser.add_argument("--op", choices=[*ops, "all"], default="all")
parser.add_argument("--backend", choices=["nvidia", "tilelang"], default="nvidia")
parser.add_argument("--backend", choices=["nvidia", "tilelang", "metax"], default="nvidia")
parser.add_argument("--mode", choices=["test", "bench", "all"], default="all")
args = parser.parse_args()