diff --git a/README.md b/README.md index 239ec18..bd53bda 100644 --- a/README.md +++ b/README.md @@ -3,8 +3,8 @@ This repository is a training-oriented GPU operator runtime. It is intentionally small, but its workflow mirrors production operator libraries: -1. Create a directory under `ops//` with backend implementations, or under - `ops/elementwise//` when the operator can reuse the elementwise framework. +1. Create a custom backend implementation under `ops//`, or place the + operator under `ops/elementwise//` when it should reuse the elementwise framework. 2. The build system auto-discovers sources by directory convention. 3. Implement a backend-specific descriptor lifecycle (create, workspace, execute, destroy). 4. Expose a Python API with out-of-place, out-variant, and prepared execution. @@ -64,11 +64,16 @@ 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 MetaX backend is built as a separate variant, should use `CAMP_ENABLE_NVIDIA=OFF`, and stores backend sources under `ops/*/metax/*.maca` or `ops/elementwise/*/metax/*.maca`. +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` or `ops/elementwise/*/metax/*.maca`. The C ABI exposes a unified backend-selection interface; each build artifact accepts only the backend compiled into it, and returns `not supported` for unavailable backends. ## Elementwise Framework -Elementwise operators that share the common shape/stride/broadcast execution model should live under `ops/elementwise///`. The shared NVIDIA launcher is in `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`; each operator only needs to provide its public C API, descriptor lifecycle, dtype dispatch, and a small device functor. `relu` is the teaching example: `negative_slope=0.0` behaves like standard ReLU, while non-zero values behave like leaky ReLU. +The training camp supports two operator paths: + +- Custom operators live under `ops///`. This path is useful for teaching hand-written kernels, fixed contiguous fast paths, custom workspace needs, or non-elementwise structures. `copy` and `vector_add` demonstrate this path. +- Reusable elementwise operators live under `ops/elementwise///`. This path is useful for ordinary elementwise operators that share the shape/stride/broadcast execution model. Later exercises can ask students to reimplement `copy` / `add`-style operators with this path. + +The shared NVIDIA launcher is in `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`; shared descriptor helpers live in `include/operator_runtime/detail/elementwise.h`. Each elementwise operator usually only needs to provide its public C API, small dtype dispatch, and a device functor. On the Python side, use `ElementwiseOpSpec` to describe input count, scalar parameters, and broadcast semantics. `relu` is the teaching example: `negative_slope=0.0` behaves like standard ReLU, while non-zero values behave like leaky ReLU. ## Production Mapping diff --git a/README.zh.md b/README.zh.md index 52a8bf2..df63d39 100644 --- a/README.zh.md +++ b/README.zh.md @@ -2,7 +2,7 @@ 这个仓库是一个面向训练的 GPU 算子运行时。它体积很小，但工作流尽量贴近真实的算子库开发流程： -1. 在 `ops//` 下创建后端实现目录；如果算子可以复用 elementwise 框架，则放到 `ops/elementwise//`。 +1. 在 `ops//` 下创建自定义后端实现目录；如果算子选择复用 elementwise 框架，则放到 `ops/elementwise//`。 2. 构建系统通过目录约定自动发现源码。 3. 实现后端专属的 descriptor 生命周期（create、workspace、execute、destroy）。 4. 提供 Python API，支持 out-of-place、out-variant 和 prepared 执行。 @@ -62,11 +62,16 @@ python tests/run_ops.py --op all --backend nvidia --mode bench ./scripts/build_metax.sh test ``` -TileLang 后端需要安装 `tilelang` Python 包。MetaX 后端使用独立构建产物，构建时应关闭 NVIDIA 变体，并将后端源码放在 `ops/*/metax/*.maca` 或 `ops/elementwise/*/metax/*.maca`。 +TileLang 后端需要安装 `tilelang` Python 包。MetaX 后端使用独立构建产物，构建时应关闭 NVIDIA 变体，并将后端源码放在 `ops/*/metax/*.maca` 或 `ops/elementwise/*/metax/*.maca`。C ABI 暴露统一的 backend 选择接口；当前构建产物只接受已编译进来的 backend，请求未启用 backend 会返回 `not supported`。 ## Elementwise 框架 -普通逐元素算子如果共享 shape、stride、broadcast 的执行模型，推荐放在 `ops/elementwise///`。NVIDIA 公共 launcher 位于 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`；每个具体算子只需要提供公开 C API、descriptor 生命周期、dtype dispatch 和一个小的 device functor。`relu` 是教学示例：`negative_slope=0.0` 时等价于标准 ReLU，非 0 时等价于 leaky ReLU。 +训练营支持两条路径： + +- 自定义算子路径：放在 `ops///`，适合教学手写 kernel、固定 contiguous fast path、特殊 workspace 或非逐元素结构。`copy` 和 `vector_add` 是这条路径的演示。 +- elementwise 复用路径：放在 `ops/elementwise///`，适合共享 shape、stride、broadcast 执行模型的普通逐元素算子。后续作业可以让学生用这条路径重新实现 `copy` / `add` 类算子。 + +NVIDIA 公共 launcher 位于 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`；公共 descriptor helper 位于 `include/operator_runtime/detail/elementwise.h`。每个 elementwise 算子通常只需要提供公开 C API、少量 dtype dispatch 和一个 device functor。Python 侧优先用 `ElementwiseOpSpec` 描述输入数、标量参数和 broadcast 语义。`relu` 是 elementwise 教学示例：`negative_slope=0.0` 时等价于标准 ReLU，非 0 时等价于 leaky ReLU。 ## 生产映射 diff --git a/docs/how-to-add-an-operator.md b/docs/how-to-add-an-operator.md index 8ebc3ab..6b74cc0 100644 --- a/docs/how-to-add-an-operator.md +++ b/docs/how-to-add-an-operator.md @@ -2,9 +2,9 @@ ## 目标 -在当前项目里，新增一个可运行算子的最小闭环包括四部分： +在当前项目里，新增一个可运行算子的最小闭环包括四部分。训练营同时保留两条路径：自定义算子路径用于教学手写 kernel 和特殊算子结构，elementwise 复用路径用于普通逐元素算子。 -1. 在 `ops/<算子名>/` 下补后端实现；普通逐元素算子优先放到 `ops/elementwise/<算子名>/`。 +1. 在 `ops/<算子名>/` 下补自定义后端实现，或在 `ops/elementwise/<算子名>/` 下补复用 elementwise 框架的实现。 2. 在 `python/operator_runtime/ops/` 下补 Python API。 3. 在 `tests/` 下补正确性测试和 benchmark 入口。 4. 重新构建并验证。 @@ -21,16 +21,16 @@ 5. 是否需要额外参数，例如 `dim`、`scalar`。 6. 是否需要 workspace。 -这一步的目的，是确定后面 C API 和 Python 绑定该按哪种模式实现。 +这一步的目的，是确定后面 C API 和 Python 绑定该按哪种模式实现。不要把“逐元素”简单等同于“必须用 elementwise 框架”：训练营会先用 `copy`、`vector_add` 展示自定义算子写法，再把复用 elementwise 框架作为后续作业和工程化路径。 ## Step 2：创建算子目录 先选择算子目录： -- 普通手写 kernel：放在 `ops/<算子名>/`。 -- 可复用公共 shape/stride/broadcast 逐元素框架的算子：放在 `ops/elementwise/<算子名>/`。 +- 自定义算子：放在 `ops/<算子名>/`。适合教学底层 kernel、固定 contiguous fast path、特殊 workspace、reduce/softmax 等非普通 elementwise 算子。 +- elementwise 复用算子：放在 `ops/elementwise/<算子名>/`。适合共享 shape、stride、broadcast 执行模型的 unary / binary / 多输入逐元素算子。 -`relu` 是 elementwise 框架示例，目录为 `ops/elementwise/relu/nvidia/`，公共 NVIDIA launcher 位于 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`。 +`copy` 和 `vector_add` 是自定义算子教学示例。`relu` 是 elementwise 框架示例，目录为 `ops/elementwise/relu/nvidia/`，公共 NVIDIA launcher 位于 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`。 先在选定目录下建立对应后端目录。 @@ -49,7 +49,7 @@ `ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` 这一层负责 C++/CUDA 实现。 -普通手写 kernel 通常需要这几部分： +自定义算子通常需要这几部分： 1. kernel 文件 2. C API 头文件 @@ -57,11 +57,11 @@ 这里最关键的是保持现有命名约定一致，因为 Python 侧会按固定符号名去找函数。 -如果是 elementwise 算子，不需要每个算子重复写 shape/stride kernel。推荐复用 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`： +如果是 elementwise 复用算子，不需要每个算子重复写 shape/stride kernel。推荐复用 `include/operator_runtime/detail/elementwise.h` 里的 descriptor helper 和 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`： -1. descriptor 中保存 `oprt::ElementwiseInfo`。 -2. create 阶段调用 `oprt::create_elementwise_info(out, {inputs...}, &info)`。 -3. workspace 返回 `info.workspace_bytes()`。 +1. descriptor 继承 `oprt::ElementwiseDescriptorBase`。 +2. create 阶段调用 `oprt::init_elementwise_descriptor(desc, out, {inputs...})`。 +3. workspace 复用 `oprt::get_elementwise_workspace_size`。 4. execute 阶段只做 dtype dispatch，并调用 `oprt::elementwise::nvidia::launch(...)`。 5. 算子自身只提供 device functor，例如 `relu` 的 `value > 0 ? value : value * negative_slope`。 @@ -95,8 +95,18 @@ Python 入口放在 `python/operator_runtime/ops/<算子名>.py`。 2. `<算子名>_`：接收调用方提供的输出 tensor，执行一次。 3. `<算子名>`：自动分配输出 tensor，再调用 `<算子名>_`。 -如果你的算子签名和现有 unary、binary、reduce-like 模式一致，就沿用现有 helper。 -如果签名比较特殊，例如带额外标量参数，就单独写一层绑定。 +如果你的算子是自定义算子，并且签名和现有 unary、binary、reduce-like 模式一致，就沿用现有 helper。 +如果你的算子是 elementwise 复用算子，优先在 Python 里定义 `ElementwiseOpSpec`，再调用 `prepare_elementwise_op`。这个 spec 描述算子名、输入数、标量参数和 broadcast 语义，避免为每个 elementwise 算子重复写 FFI 绑定和基础校验。 + +示例： + +```python +_MY_OP_SPEC = ElementwiseOpSpec( + name="my_op", + input_count=1, + scalar_argtypes=(ctypes.c_float,), +) +``` ## Step 5：导出到公共 API @@ -191,13 +201,13 @@ API contract 测试主要覆盖 shape 不匹配、dtype 不匹配、非 contiguo ## 现有模板怎么选 -- `copy`：适合最简单的单输入单输出流程。 -- `vector_add`：适合标准 contiguous 双输入算子。 +- `copy`：适合最简单的自定义单输入单输出流程。 +- `vector_add`：适合标准 contiguous 双输入自定义算子。 - `relu`：适合复用 elementwise 框架的单输入逐元素算子，也展示了额外标量参数 `negative_slope` 的 C/Python 绑定方式。 - `reduce_sum`：适合带 reduce 维度的算子。 - `softmax`：适合带更明确 shape 约束和归一化逻辑的算子。 -如果新算子本质上是普通 elementwise，优先参考 `relu` 和 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`；如果只想写最简单 contiguous kernel，再参考 `vector_add`。 +如果新算子本质上是普通 elementwise，优先参考 `relu`、`ElementwiseOpSpec` 和 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`；如果课程目标是练习手写 contiguous kernel，再参考 `copy` 或 `vector_add` 的自定义路径。 ## 最终检查清单 diff --git a/include/operator_runtime/api.h b/include/operator_runtime/api.h index 0e5b65c..9e54257 100644 --- a/include/operator_runtime/api.h +++ b/include/operator_runtime/api.h @@ -27,6 +27,12 @@ typedef enum { OPRT_DTYPE_F32 = 1 } oprt_dtype_t; +typedef enum { + OPRT_BACKEND_AUTO = 0, + OPRT_BACKEND_NVIDIA = 1, + OPRT_BACKEND_METAX = 2 +} oprt_backend_t; + #define OPRT_MAX_DIMS 8 typedef void *oprt_stream_t; @@ -41,8 +47,10 @@ typedef struct { } oprt_tensor_view_t; OPRT_EXPORT const char *oprt_status_string(oprt_status_t status); +OPRT_EXPORT const char *oprt_backend_string(oprt_backend_t backend); +OPRT_EXPORT oprt_status_t oprt_set_backend(oprt_backend_t backend); +OPRT_EXPORT oprt_status_t oprt_get_backend(oprt_backend_t *backend); #ifdef __cplusplus } #endif - diff --git a/include/operator_runtime/descriptor.h b/include/operator_runtime/descriptor.h index 5d98a06..3d05504 100644 --- a/include/operator_runtime/descriptor.h +++ b/include/operator_runtime/descriptor.h @@ -10,7 +10,7 @@ struct oprt_operator_descriptor { virtual ~oprt_operator_descriptor() = default; virtual const char *op_name() const = 0; size_t workspace_size = 0; + oprt_backend_t backend = OPRT_BACKEND_AUTO; }; #endif - diff --git a/include/operator_runtime/detail/elementwise.h b/include/operator_runtime/detail/elementwise.h index 25d4382..1b23a9e 100644 --- a/include/operator_runtime/detail/elementwise.h +++ b/include/operator_runtime/detail/elementwise.h @@ -1,5 +1,6 @@ #pragma once +#include "operator_runtime/descriptor.h" #include "operator_runtime/tensor_view.h" #include "operator_runtime/detail/tensor_checks.h" @@ -181,6 +182,68 @@ inline oprt_status_t create_elementwise_info(const oprt_tensor_view_t *out, return OPRT_SUCCESS; } +struct ElementwiseDescriptorBase : oprt_operator_descriptor { + oprt_tensor_view_t out_view{}; + std::vector input_views; + oprt::ElementwiseInfo info; +}; + +inline oprt_status_t init_elementwise_descriptor(ElementwiseDescriptorBase *desc, + const oprt_tensor_view_t *out, + const std::vector &inputs) { + if (desc == nullptr || out == nullptr || inputs.empty()) { + return OPRT_ERR_INVALID_ARG; + } + + oprt::ElementwiseInfo info; + auto status = oprt::create_elementwise_info(out, inputs, &info); + if (status != OPRT_SUCCESS) { + return status; + } + + desc->out_view = *out; + desc->input_views.clear(); + desc->input_views.reserve(inputs.size()); + for (const auto *input : inputs) { + desc->input_views.push_back(*input); + } + desc->info = std::move(info); + desc->workspace_size = desc->info.workspace_bytes(); + return OPRT_SUCCESS; +} + +inline oprt_status_t get_elementwise_workspace_size(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; +} + +inline oprt_status_t validate_elementwise_execute_args(oprt_operator_descriptor_t desc, + size_t workspace_size, + void *out, + const std::vector &inputs) { + if (desc == nullptr || out == nullptr || inputs.empty()) { + return OPRT_ERR_INVALID_ARG; + } + for (const auto *input : inputs) { + if (input == nullptr) { + return OPRT_ERR_INVALID_ARG; + } + } + if (workspace_size < desc->workspace_size) { + return OPRT_ERR_INSUFFICIENT_WORKSPACE; + } + return OPRT_SUCCESS; +} + +inline oprt_status_t destroy_elementwise_descriptor(oprt_operator_descriptor_t desc) { + delete desc; + return OPRT_SUCCESS; +} + } // namespace oprt #endif diff --git a/ops/common/status.cc b/ops/common/status.cc index 3c0d51d..54ea95d 100644 --- a/ops/common/status.cc +++ b/ops/common/status.cc @@ -1,5 +1,15 @@ #include "operator_runtime/api.h" +namespace { +#ifdef CAMP_ENABLE_NVIDIA +oprt_backend_t g_current_backend = OPRT_BACKEND_NVIDIA; +#elif defined(CAMP_ENABLE_METAX) +oprt_backend_t g_current_backend = OPRT_BACKEND_METAX; +#else +oprt_backend_t g_current_backend = OPRT_BACKEND_AUTO; +#endif +} + extern "C" OPRT_EXPORT const char *oprt_status_string(oprt_status_t status) { switch (status) { case OPRT_SUCCESS: @@ -19,3 +29,54 @@ extern "C" OPRT_EXPORT const char *oprt_status_string(oprt_status_t status) { } } +extern "C" OPRT_EXPORT const char *oprt_backend_string(oprt_backend_t backend) { + switch (backend) { + case OPRT_BACKEND_AUTO: + return "auto"; + case OPRT_BACKEND_NVIDIA: + return "nvidia"; + case OPRT_BACKEND_METAX: + return "metax"; + default: + return "unknown"; + } +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_set_backend(oprt_backend_t backend) { + switch (backend) { + case OPRT_BACKEND_AUTO: +#ifdef CAMP_ENABLE_NVIDIA + g_current_backend = OPRT_BACKEND_NVIDIA; + return OPRT_SUCCESS; +#elif defined(CAMP_ENABLE_METAX) + g_current_backend = OPRT_BACKEND_METAX; + return OPRT_SUCCESS; +#else + return OPRT_ERR_NOT_SUPPORTED; +#endif + case OPRT_BACKEND_NVIDIA: +#ifdef CAMP_ENABLE_NVIDIA + g_current_backend = backend; + return OPRT_SUCCESS; +#else + return OPRT_ERR_NOT_SUPPORTED; +#endif + case OPRT_BACKEND_METAX: +#ifdef CAMP_ENABLE_METAX + g_current_backend = backend; + return OPRT_SUCCESS; +#else + return OPRT_ERR_NOT_SUPPORTED; +#endif + default: + return OPRT_ERR_INVALID_ARG; + } +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_get_backend(oprt_backend_t *backend) { + if (backend == nullptr) { + return OPRT_ERR_INVALID_ARG; + } + *backend = g_current_backend; + return OPRT_SUCCESS; +} diff --git a/ops/elementwise/relu/relu_common.h b/ops/elementwise/relu/relu_common.h index 4ede9b1..6258889 100644 --- a/ops/elementwise/relu/relu_common.h +++ b/ops/elementwise/relu/relu_common.h @@ -5,10 +5,7 @@ namespace oprt::elementwise { -struct ReluDescriptor final : oprt_operator_descriptor { - oprt_tensor_view_t out_view{}; - oprt_tensor_view_t in_view{}; - oprt::ElementwiseInfo info; +struct ReluDescriptor final : oprt::ElementwiseDescriptorBase { float negative_slope = 0.0f; const char *op_name() const override { @@ -30,29 +27,20 @@ inline oprt_status_t create_relu_descriptor_common( return OPRT_ERR_INVALID_ARG; } - oprt::ElementwiseInfo info; - auto status = oprt::create_elementwise_info(out, {in}, &info); + auto *typed = new ReluDescriptor(); + auto status = oprt::init_elementwise_descriptor(typed, out, {in}); if (status != OPRT_SUCCESS) { + delete typed; return status; } - - auto *typed = new ReluDescriptor(); - typed->out_view = *out; - typed->in_view = *in; - typed->info = std::move(info); typed->negative_slope = negative_slope; - typed->workspace_size = typed->info.workspace_bytes(); *desc = typed; return OPRT_SUCCESS; } inline oprt_status_t get_relu_workspace_size_common(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; + return oprt::get_elementwise_workspace_size(desc, size); } inline oprt_status_t validate_relu_execute_args(oprt_operator_descriptor_t desc, @@ -62,15 +50,11 @@ inline oprt_status_t validate_relu_execute_args(oprt_operator_descriptor_t desc, if (desc == nullptr || out == nullptr || in == nullptr) { return OPRT_ERR_INVALID_ARG; } - if (workspace_size < desc->workspace_size) { - return OPRT_ERR_INSUFFICIENT_WORKSPACE; - } - return OPRT_SUCCESS; + return oprt::validate_elementwise_execute_args(desc, workspace_size, out, {in}); } inline oprt_status_t destroy_relu_descriptor_common(oprt_operator_descriptor_t desc) { - delete desc; - return OPRT_SUCCESS; + return oprt::destroy_elementwise_descriptor(desc); } } // namespace oprt::elementwise diff --git a/python/operator_runtime/_internal/__init__.py b/python/operator_runtime/_internal/__init__.py index 17a4de7..6202358 100644 --- a/python/operator_runtime/_internal/__init__.py +++ b/python/operator_runtime/_internal/__init__.py @@ -5,6 +5,7 @@ from .bindings import ( Descriptor, OperatorRuntimeError, Status, + bind_elementwise, bind_unary, bind_relu, bind_binary, @@ -24,6 +25,7 @@ __all__ = [ "Descriptor", "OperatorRuntimeError", "Status", + "bind_elementwise", "bind_unary", "bind_relu", "bind_binary", diff --git a/python/operator_runtime/_internal/bindings.py b/python/operator_runtime/_internal/bindings.py index 651ff50..d9d98dc 100644 --- a/python/operator_runtime/_internal/bindings.py +++ b/python/operator_runtime/_internal/bindings.py @@ -2,31 +2,50 @@ from __future__ import annotations import ctypes from dataclasses import dataclass -from typing import Callable +from typing import Callable, Sequence from operator_runtime.backend import Backend +from operator_runtime.backend import normalize_backend from .loader import load_library from .tensor_view import TensorView Status = ctypes.c_int Descriptor = ctypes.c_void_p +OPRT_BACKEND_NVIDIA = 1 +OPRT_BACKEND_METAX = 2 class OperatorRuntimeError(RuntimeError): pass -def check_status(status: int) -> None: +def _backend_code(backend: str | Backend) -> int: + normalized = normalize_backend(backend) + if normalized is Backend.NVIDIA: + return OPRT_BACKEND_NVIDIA + if normalized is Backend.METAX: + return OPRT_BACKEND_METAX + raise NotImplementedError(f"backend {normalized.value} is not runnable through the C ABI") + + +def check_status(status: int, lib=None) -> None: if status == 0: return - lib = load_library() + if lib is None: + lib = load_library() lib.oprt_status_string.argtypes = [ctypes.c_int] lib.oprt_status_string.restype = ctypes.c_char_p message = lib.oprt_status_string(status).decode("utf-8") raise OperatorRuntimeError(message) +def set_runtime_backend(lib, backend: str | Backend) -> None: + lib.oprt_set_backend.argtypes = [ctypes.c_int] + lib.oprt_set_backend.restype = Status + check_status(lib.oprt_set_backend(_backend_code(backend)), lib) + + @dataclass(frozen=True) class CFunctions: create: Callable @@ -41,14 +60,20 @@ def _missing_symbol_error(name: str, symbol: str, exc: AttributeError) -> Operat ) from exc -def bind_unary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions: - lib = load_library() +def _bind_lifecycle( + name: str, + backend: str | Backend, + create_argtypes: Sequence[object], + execute_argtypes: Sequence[object], +) -> CFunctions: + lib = load_library(backend) + set_runtime_backend(lib, backend) create_symbol = f"oprt_create_{name}_descriptor" try: create = getattr(lib, create_symbol) except AttributeError as exc: _missing_symbol_error(name, create_symbol, exc) - create.argtypes = [ctypes.POINTER(Descriptor), ctypes.POINTER(TensorView), ctypes.POINTER(TensorView)] + create.argtypes = [ctypes.POINTER(Descriptor), *create_argtypes] create.restype = Status workspace_symbol = f"oprt_get_{name}_workspace_size" @@ -64,14 +89,7 @@ def bind_unary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions execute = getattr(lib, execute_symbol) except AttributeError as exc: _missing_symbol_error(name, execute_symbol, exc) - execute.argtypes = [ - Descriptor, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_void_p, - ] + execute.argtypes = [Descriptor, ctypes.c_void_p, ctypes.c_size_t, *execute_argtypes, ctypes.c_void_p] execute.restype = Status destroy_symbol = f"oprt_destroy_{name}_descriptor" @@ -84,147 +102,39 @@ def bind_unary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions return CFunctions(create, workspace, execute, destroy) +def bind_elementwise( + name: str, + input_count: int, + scalar_argtypes: Sequence[object] = (), + backend: str | Backend = Backend.NVIDIA, +) -> CFunctions: + if input_count <= 0: + raise ValueError("elementwise operators require at least one input") + tensor_views = [ctypes.POINTER(TensorView)] * (input_count + 1) + data_ptrs = [ctypes.c_void_p] * (input_count + 1) + return _bind_lifecycle(name, backend, [*tensor_views, *scalar_argtypes], data_ptrs) + + +def bind_unary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions: + return bind_elementwise(name, 1, backend=backend) + + def bind_relu(backend: str | Backend = Backend.NVIDIA) -> CFunctions: - lib = load_library() - name = "relu" - create_symbol = "oprt_create_relu_descriptor" - try: - create = getattr(lib, create_symbol) - except AttributeError as exc: - _missing_symbol_error(name, create_symbol, exc) - create.argtypes = [ - ctypes.POINTER(Descriptor), - ctypes.POINTER(TensorView), - ctypes.POINTER(TensorView), - ctypes.c_float, - ] - create.restype = Status - - workspace_symbol = "oprt_get_relu_workspace_size" - try: - workspace = getattr(lib, workspace_symbol) - except AttributeError as exc: - _missing_symbol_error(name, workspace_symbol, exc) - workspace.argtypes = [Descriptor, ctypes.POINTER(ctypes.c_size_t)] - workspace.restype = Status - - execute_symbol = "oprt_execute_relu" - try: - execute = getattr(lib, execute_symbol) - except AttributeError as exc: - _missing_symbol_error(name, execute_symbol, exc) - execute.argtypes = [ - Descriptor, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_void_p, - ] - execute.restype = Status - - destroy_symbol = "oprt_destroy_relu_descriptor" - try: - destroy = getattr(lib, destroy_symbol) - except AttributeError as exc: - _missing_symbol_error(name, destroy_symbol, exc) - destroy.argtypes = [Descriptor] - destroy.restype = Status - return CFunctions(create, workspace, execute, destroy) + return bind_elementwise("relu", 1, [ctypes.c_float], backend) def bind_binary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions: - lib = load_library() - create_symbol = f"oprt_create_{name}_descriptor" - try: - create = getattr(lib, create_symbol) - except AttributeError as exc: - _missing_symbol_error(name, create_symbol, exc) - create.argtypes = [ - ctypes.POINTER(Descriptor), - ctypes.POINTER(TensorView), - ctypes.POINTER(TensorView), - ctypes.POINTER(TensorView), - ] - create.restype = Status - - workspace_symbol = f"oprt_get_{name}_workspace_size" - try: - workspace = getattr(lib, workspace_symbol) - except AttributeError as exc: - _missing_symbol_error(name, workspace_symbol, exc) - workspace.argtypes = [Descriptor, ctypes.POINTER(ctypes.c_size_t)] - workspace.restype = Status - - execute_symbol = f"oprt_execute_{name}" - try: - execute = getattr(lib, execute_symbol) - except AttributeError as exc: - _missing_symbol_error(name, execute_symbol, exc) - execute.argtypes = [ - Descriptor, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_void_p, - ] - execute.restype = Status - - destroy_symbol = f"oprt_destroy_{name}_descriptor" - try: - destroy = getattr(lib, destroy_symbol) - except AttributeError as exc: - _missing_symbol_error(name, destroy_symbol, exc) - destroy.argtypes = [Descriptor] - destroy.restype = Status - return CFunctions(create, workspace, execute, destroy) + return bind_elementwise(name, 2, backend=backend) def bind_reduce_like(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions: - lib = load_library() - create_symbol = f"oprt_create_{name}_descriptor" - try: - create = getattr(lib, create_symbol) - except AttributeError as exc: - _missing_symbol_error(name, create_symbol, exc) - create.argtypes = [ - ctypes.POINTER(Descriptor), - ctypes.POINTER(TensorView), - ctypes.POINTER(TensorView), - ctypes.c_int64, - ] - create.restype = Status - - workspace_symbol = f"oprt_get_{name}_workspace_size" - try: - workspace = getattr(lib, workspace_symbol) - except AttributeError as exc: - _missing_symbol_error(name, workspace_symbol, exc) - workspace.argtypes = [Descriptor, ctypes.POINTER(ctypes.c_size_t)] - workspace.restype = Status - - execute_symbol = f"oprt_execute_{name}" - try: - execute = getattr(lib, execute_symbol) - except AttributeError as exc: - _missing_symbol_error(name, execute_symbol, exc) - execute.argtypes = [ - Descriptor, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_void_p, - ] - execute.restype = Status - - destroy_symbol = f"oprt_destroy_{name}_descriptor" - try: - destroy = getattr(lib, destroy_symbol) - except AttributeError as exc: - _missing_symbol_error(name, destroy_symbol, exc) - destroy.argtypes = [Descriptor] - destroy.restype = Status - return CFunctions(create, workspace, execute, destroy) + return _bind_lifecycle( + name, + backend, + [ + ctypes.POINTER(TensorView), + ctypes.POINTER(TensorView), + ctypes.c_int64, + ], + [ctypes.c_void_p, ctypes.c_void_p], + ) diff --git a/python/operator_runtime/_internal/loader.py b/python/operator_runtime/_internal/loader.py index 55c7419..b19fec8 100644 --- a/python/operator_runtime/_internal/loader.py +++ b/python/operator_runtime/_internal/loader.py @@ -6,28 +6,67 @@ from functools import lru_cache from pathlib import Path -def _candidate_library_paths() -> list[Path]: +def _backend_key(backend: object | None) -> str | None: + if backend is None: + return None + from operator_runtime.backend import normalize_backend + + return normalize_backend(backend).value + + +def _append_build_dir(candidates: list[Path], build_dir: str | None) -> None: + if not build_dir: + return + root = Path(build_dir) + candidates.extend([root / "libcamp_ops.so", root / "ops" / "libcamp_ops.so"]) + + +def _candidate_library_paths(backend: str | None = None) -> list[Path]: candidates: list[Path] = [] - build_dir = os.environ.get("CAMP_BUILD_DIR") - if build_dir: - root = Path(build_dir) - candidates.extend([root / "libcamp_ops.so", root / "ops" / "libcamp_ops.so"]) repo_root = Path(__file__).resolve().parents[3] - candidates.extend( - [ - 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", - ] - ) + + if backend == "nvidia": + _append_build_dir(candidates, os.environ.get("CAMP_NVIDIA_BUILD_DIR")) + _append_build_dir(candidates, os.environ.get("CAMP_BUILD_DIR")) + candidates.extend( + [ + repo_root / "build" / "libcamp_ops.so", + repo_root / "build" / "ops" / "libcamp_ops.so", + ] + ) + elif backend == "metax": + _append_build_dir(candidates, os.environ.get("CAMP_METAX_BUILD_DIR")) + _append_build_dir(candidates, os.environ.get("CAMP_BUILD_DIR")) + candidates.extend( + [ + repo_root / "build-metax" / "libcamp_ops.so", + repo_root / "build-metax" / "ops" / "libcamp_ops.so", + ] + ) + else: + _append_build_dir(candidates, os.environ.get("CAMP_BUILD_DIR")) + _append_build_dir(candidates, os.environ.get("CAMP_NVIDIA_BUILD_DIR")) + _append_build_dir(candidates, os.environ.get("CAMP_METAX_BUILD_DIR")) + candidates.extend( + [ + 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 -@lru_cache(maxsize=1) -def load_library() -> ctypes.CDLL: - for path in _candidate_library_paths(): +@lru_cache(maxsize=None) +def _load_library(backend: str | None) -> ctypes.CDLL: + for path in _candidate_library_paths(backend): if path.exists(): return ctypes.CDLL(str(path)) - searched = ", ".join(str(path) for path in _candidate_library_paths()) - raise FileNotFoundError(f"libcamp_ops.so not found; searched: {searched}") + searched = ", ".join(str(path) for path in _candidate_library_paths(backend)) + label = "default" if backend is None else backend + raise FileNotFoundError(f"libcamp_ops.so for backend {label} not found; searched: {searched}") + + +def load_library(backend: object | None = None) -> ctypes.CDLL: + return _load_library(_backend_key(backend)) diff --git a/python/operator_runtime/ops/_common.py b/python/operator_runtime/ops/_common.py index b6a1231..c780284 100644 --- a/python/operator_runtime/ops/_common.py +++ b/python/operator_runtime/ops/_common.py @@ -1,10 +1,31 @@ from __future__ import annotations import ctypes +from dataclasses import dataclass +from typing import Callable, Sequence import torch -from operator_runtime._internal import CFunctions, Descriptor, PreparedOp, check_status +from operator_runtime.backend import Backend, normalize_backend +from operator_runtime._internal import CFunctions, Descriptor, PreparedOp, bind_elementwise, check_status, tensor_view + + +def _broadcastable_to(src_shape: tuple[int, ...], out_shape: tuple[int, ...]) -> bool: + if len(src_shape) > len(out_shape): + return False + offset = len(out_shape) - len(src_shape) + for out_dim, out_extent in enumerate(out_shape): + src_dim = out_dim - offset + if src_dim < 0: + continue + src_extent = src_shape[src_dim] + if src_extent != out_extent and src_extent != 1: + return False + return True + + +def _has_broadcast_dim(tensor: torch.Tensor) -> bool: + return any(size > 1 and stride == 0 for size, stride in zip(tensor.shape, tensor.stride())) def build_prepared_op( @@ -22,3 +43,57 @@ def build_prepared_op( workspace = torch.empty(workspace_size.value, dtype=torch.uint8, device=stream_tensor.device) runner_args = tuple(ctypes.c_void_p(tensor.data_ptr()) for tensor in tensors) return PreparedOp(funcs, desc, workspace, runner_args, stream_tensor) + + +@dataclass(frozen=True) +class ElementwiseOpSpec: + name: str + input_count: int + scalar_argtypes: tuple[object, ...] = () + allow_broadcast: bool = True + + +def check_elementwise_tensors(spec: ElementwiseOpSpec, out: torch.Tensor, inputs: Sequence[torch.Tensor]) -> None: + if len(inputs) != spec.input_count: + raise ValueError(f"{spec.name} expects {spec.input_count} inputs, got {len(inputs)}") + if not out.is_cuda or any(not tensor.is_cuda for tensor in inputs): + raise ValueError(f"{spec.name} expects CUDA tensors") + if _has_broadcast_dim(out): + raise ValueError(f"{spec.name} expects a writable output tensor") + for tensor in inputs: + if spec.allow_broadcast: + if not _broadcastable_to(tuple(tensor.shape), tuple(out.shape)): + raise ValueError(f"{spec.name} expects inputs to be broadcastable to out") + elif out.shape != tensor.shape: + raise ValueError(f"{spec.name} expects matching shapes") + if out.dtype != tensor.dtype: + raise TypeError(f"{spec.name} expects matching dtypes") + + +def prepare_elementwise_op( + spec: ElementwiseOpSpec, + out: torch.Tensor, + inputs: Sequence[torch.Tensor], + scalars: Sequence[object] = (), + backend: str | Backend = Backend.NVIDIA, + scalar_converters: Sequence[Callable[[object], object]] = (), +) -> PreparedOp: + backend = normalize_backend(backend) + if backend not in (Backend.NVIDIA, Backend.METAX): + raise NotImplementedError(f"backend {backend.value} is not runnable") + + check_elementwise_tensors(spec, out, inputs) + if len(scalars) != len(spec.scalar_argtypes): + raise ValueError(f"{spec.name} expects {len(spec.scalar_argtypes)} scalar args, got {len(scalars)}") + if scalar_converters and len(scalar_converters) != len(scalars): + raise ValueError("scalar_converters must match scalars") + + funcs = bind_elementwise(spec.name, spec.input_count, spec.scalar_argtypes, backend) + views = [tensor_view(out), *(tensor_view(tensor) for tensor in inputs)] + converters: Sequence[Callable[[object], object]] = scalar_converters or tuple(lambda value: value for _ in scalars) + converted_scalars = [ + argtype(converter(value)) + for value, argtype, converter in zip(scalars, spec.scalar_argtypes, converters) + ] + create_args = (*[ctypes.byref(view) for view in views], *converted_scalars) + return build_prepared_op(funcs, create_args, (out, *inputs), out) diff --git a/python/operator_runtime/ops/relu.py b/python/operator_runtime/ops/relu.py index 4ffef35..d2a6d55 100644 --- a/python/operator_runtime/ops/relu.py +++ b/python/operator_runtime/ops/relu.py @@ -4,38 +4,16 @@ import ctypes import torch -from operator_runtime.backend import Backend, normalize_backend -from operator_runtime._internal import PreparedOp, bind_relu, tensor_view -from operator_runtime.ops._common import build_prepared_op +from operator_runtime.backend import Backend +from operator_runtime._internal import PreparedOp +from operator_runtime.ops._common import ElementwiseOpSpec, prepare_elementwise_op -def _broadcastable_to(src_shape: tuple[int, ...], out_shape: tuple[int, ...]) -> bool: - if len(src_shape) > len(out_shape): - return False - offset = len(out_shape) - len(src_shape) - for out_dim, out_extent in enumerate(out_shape): - src_dim = out_dim - offset - if src_dim < 0: - continue - src_extent = src_shape[src_dim] - if src_extent != out_extent and src_extent != 1: - return False - return True - - -def _has_broadcast_dim(tensor: torch.Tensor) -> bool: - return any(size > 1 and stride == 0 for size, stride in zip(tensor.shape, tensor.stride())) - - -def _check(out: torch.Tensor, src: torch.Tensor) -> None: - if not out.is_cuda or not src.is_cuda: - raise ValueError("relu expects CUDA tensors") - if _has_broadcast_dim(out): - raise ValueError("relu expects a writable output tensor") - if not _broadcastable_to(tuple(src.shape), tuple(out.shape)): - raise ValueError("relu expects src to be broadcastable to out") - if out.dtype != src.dtype: - raise TypeError("relu expects matching dtypes") +_RELU_SPEC = ElementwiseOpSpec( + name="relu", + input_count=1, + scalar_argtypes=(ctypes.c_float,), +) def prepare_relu( @@ -44,20 +22,14 @@ def prepare_relu( negative_slope: float = 0.0, backend: str | Backend = Backend.NVIDIA, ) -> PreparedOp: - backend = normalize_backend(backend) - _check(out, src) - if backend not in (Backend.NVIDIA, Backend.METAX): - raise NotImplementedError(f"backend {backend.value} is not runnable") - - funcs = bind_relu(backend) - out_view = tensor_view(out) - src_view = tensor_view(src) - create_args = ( - ctypes.byref(out_view), - ctypes.byref(src_view), - ctypes.c_float(float(negative_slope)), + return prepare_elementwise_op( + _RELU_SPEC, + out, + (src,), + (negative_slope,), + backend, + (float,), ) - return build_prepared_op(funcs, create_args, (out, src), out) def relu_(