diff --git a/CMakeLists.txt b/CMakeLists.txt index fc7c381..854f22f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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) diff --git a/README.md b/README.md index a8ec095..02a5a9a 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/README.zh.md b/README.zh.md index 3f01c93..d168e29 100644 --- a/README.zh.md +++ b/README.zh.md @@ -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`。 ## 生产映射 diff --git a/cmake/CMakeMACAInformation.cmake b/cmake/CMakeMACAInformation.cmake new file mode 100644 index 0000000..52a3728 --- /dev/null +++ b/cmake/CMakeMACAInformation.cmake @@ -0,0 +1,16 @@ +set(CMAKE_MACA_COMPILE_OBJECT + " -o -c " +) + +set(CMAKE_INCLUDE_FLAG_MACA "-I") +set(CMAKE_MACA_COMPILE_OPTIONS_PIC "-fPIC") + +set(CMAKE_MACA_CREATE_SHARED_LIBRARY + " -shared -o --maca-path=$ENV{MACA_PATH}" +) + +set(CMAKE_MACA_LINK_EXECUTABLE + " -o --maca-path=$ENV{MACA_PATH}" +) + +set(CMAKE_MACA_INFORMATION_LOADED 1) diff --git a/cmake/FindMetaX.cmake b/cmake/FindMetaX.cmake index 5b1709f..b4858e6 100644 --- a/cmake/FindMetaX.cmake +++ b/cmake/FindMetaX.cmake @@ -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 +) diff --git a/examples/01_copy.py b/examples/01_copy.py index 60b6d76..e0c28b7 100644 --- a/examples/01_copy.py +++ b/examples/01_copy.py @@ -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) diff --git a/examples/02_vector_add.py b/examples/02_vector_add.py index aa5cd9f..b088288 100644 --- a/examples/02_vector_add.py +++ b/examples/02_vector_add.py @@ -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) diff --git a/examples/03_reduce_sum.py b/examples/03_reduce_sum.py index 53b4fdf..099ca82 100644 --- a/examples/03_reduce_sum.py +++ b/examples/03_reduce_sum.py @@ -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) diff --git a/examples/04_softmax.py b/examples/04_softmax.py index 440e88e..6b26c8d 100644 --- a/examples/04_softmax.py +++ b/examples/04_softmax.py @@ -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) diff --git a/ops/CMakeLists.txt b/ops/CMakeLists.txt index 36e569e..25b8af1 100644 --- a/ops/CMakeLists.txt +++ b/ops/CMakeLists.txt @@ -1,5 +1,6 @@ # Convention: operator CUDA sources auto-discovered by directory layout: # ops//nvidia/*.cu -> compiled when CAMP_ENABLE_NVIDIA is ON +# ops//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 $<$:CAMP_ENABLE_NVIDIA=1> $<$:CAMP_ENABLE_METAX=1> diff --git a/ops/common/metax/README.md b/ops/common/metax/README.md index 19df444..b9794a9 100644 --- a/ops/common/metax/README.md +++ b/ops/common/metax/README.md @@ -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` diff --git a/ops/copy/metax/README.md b/ops/copy/metax/README.md index 5fffc4b..01778af 100644 --- a/ops/copy/metax/README.md +++ b/ops/copy/metax/README.md @@ -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`. diff --git a/ops/copy/metax/copy_metax.h b/ops/copy/metax/copy_metax.h index 83911f5..092ed3f 100644 --- a/ops/copy/metax/copy_metax.h +++ b/ops/copy/metax/copy_metax.h @@ -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 diff --git a/ops/copy/metax/copy_metax.maca b/ops/copy/metax/copy_metax.maca new file mode 100644 index 0000000..367cde8 --- /dev/null +++ b/ops/copy/metax/copy_metax.maca @@ -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 + +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 +__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t i = idx; i < n; i += stride) { + dst[i] = src[i]; + } +} + +template +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<<>>( + static_cast(dst), static_cast(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(desc); + switch (typed->dst_view.dtype) { + case OPRT_DTYPE_F16: + return launch_copy(typed, dst, src, stream); + case OPRT_DTYPE_F32: + return launch_copy(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; +} diff --git a/ops/reduce_sum/metax/README.md b/ops/reduce_sum/metax/README.md index 557cba3..44306c8 100644 --- a/ops/reduce_sum/metax/README.md +++ b/ops/reduce_sum/metax/README.md @@ -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`. diff --git a/ops/reduce_sum/metax/reduce_sum_metax.h b/ops/reduce_sum/metax/reduce_sum_metax.h index 83911f5..fc7288a 100644 --- a/ops/reduce_sum/metax/reduce_sum_metax.h +++ b/ops/reduce_sum/metax/reduce_sum_metax.h @@ -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 diff --git a/ops/reduce_sum/metax/reduce_sum_metax.maca b/ops/reduce_sum/metax/reduce_sum_metax.maca new file mode 100644 index 0000000..e7b88ab --- /dev/null +++ b/ops/reduce_sum/metax/reduce_sum_metax.maca @@ -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 + +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(desc); + cudaStream_t s = oprt::as_cuda_stream(stream); + constexpr int threads = 256; + reduce_sum_rowwise_kernel<<rows, threads, threads * sizeof(float), s>>>( + static_cast(out), static_cast(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; +} diff --git a/ops/softmax/metax/README.md b/ops/softmax/metax/README.md index 66dc7d3..3a7d7ac 100644 --- a/ops/softmax/metax/README.md +++ b/ops/softmax/metax/README.md @@ -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`. diff --git a/ops/softmax/metax/softmax_metax.h b/ops/softmax/metax/softmax_metax.h index 83911f5..103c8af 100644 --- a/ops/softmax/metax/softmax_metax.h +++ b/ops/softmax/metax/softmax_metax.h @@ -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 diff --git a/ops/softmax/metax/softmax_metax.maca b/ops/softmax/metax/softmax_metax.maca new file mode 100644 index 0000000..6c3c692 --- /dev/null +++ b/ops/softmax/metax/softmax_metax.maca @@ -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 +#include +#include + +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(desc); + cudaStream_t s = oprt::as_cuda_stream(stream); + constexpr int threads = 256; + softmax_rowwise_kernel<<rows, threads, threads * sizeof(float), s>>>( + static_cast(out), static_cast(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; +} diff --git a/ops/vector_add/metax/README.md b/ops/vector_add/metax/README.md index a6211d3..9359a2f 100644 --- a/ops/vector_add/metax/README.md +++ b/ops/vector_add/metax/README.md @@ -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`. diff --git a/ops/vector_add/metax/vector_add_metax.h b/ops/vector_add/metax/vector_add_metax.h index 83911f5..75a035d 100644 --- a/ops/vector_add/metax/vector_add_metax.h +++ b/ops/vector_add/metax/vector_add_metax.h @@ -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 diff --git a/ops/vector_add/metax/vector_add_metax.maca b/ops/vector_add/metax/vector_add_metax.maca new file mode 100644 index 0000000..5279742 --- /dev/null +++ b/ops/vector_add/metax/vector_add_metax.maca @@ -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 + +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 +__device__ T add_values(T a, T b) { + return a + b; +} + +template <> +__device__ inline half add_values(half a, half b) { + return __hadd(a, b); +} + +template +__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(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t i = idx; i < n; i += stride) { + out[i] = add_values(a[i], b[i]); + } +} + +template +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<<>>( + static_cast(out), static_cast(a), static_cast(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(desc); + switch (typed->out_view.dtype) { + case OPRT_DTYPE_F16: + return launch_vector_add(typed, out, a, b, stream); + case OPRT_DTYPE_F32: + return launch_vector_add(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; +} diff --git a/python/operator_runtime/_internal/bindings.py b/python/operator_runtime/_internal/bindings.py index a4aeaf9..906e2ac 100644 --- a/python/operator_runtime/_internal/bindings.py +++ b/python/operator_runtime/_internal/bindings.py @@ -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) diff --git a/python/operator_runtime/_internal/loader.py b/python/operator_runtime/_internal/loader.py index 2db7349..55c7419 100644 --- a/python/operator_runtime/_internal/loader.py +++ b/python/operator_runtime/_internal/loader.py @@ -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 diff --git a/python/operator_runtime/ops/copy.py b/python/operator_runtime/ops/copy.py index 49a3677..b265800 100644 --- a/python/operator_runtime/ops/copy.py +++ b/python/operator_runtime/ops/copy.py @@ -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)) diff --git a/python/operator_runtime/ops/reduce_sum.py b/python/operator_runtime/ops/reduce_sum.py index 597216c..000a266 100644 --- a/python/operator_runtime/ops/reduce_sum.py +++ b/python/operator_runtime/ops/reduce_sum.py @@ -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)) diff --git a/python/operator_runtime/ops/softmax.py b/python/operator_runtime/ops/softmax.py index 4382b61..431528d 100644 --- a/python/operator_runtime/ops/softmax.py +++ b/python/operator_runtime/ops/softmax.py @@ -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)) diff --git a/python/operator_runtime/ops/vector_add.py b/python/operator_runtime/ops/vector_add.py index 706ee4b..893466e 100644 --- a/python/operator_runtime/ops/vector_add.py +++ b/python/operator_runtime/ops/vector_add.py @@ -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) diff --git a/scripts/build_metax.sh b/scripts/build_metax.sh new file mode 100755 index 0000000..5aa3e90 --- /dev/null +++ b/scripts/build_metax.sh @@ -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 diff --git a/tests/op_tests/test_vector_add.py b/tests/op_tests/test_vector_add.py index 0cdf21d..c187b26 100644 --- a/tests/op_tests/test_vector_add.py +++ b/tests/op_tests/test_vector_add.py @@ -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") diff --git a/tests/run_ops.py b/tests/run_ops.py index 0439d50..76d73cb 100644 --- a/tests/run_ops.py +++ b/tests/run_ops.py @@ -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()