From 9271fed062662c789e4d52be8edaf0278502500b Mon Sep 17 00:00:00 2001 From: yutianyu Date: Tue, 12 May 2026 11:30:19 +0800 Subject: [PATCH] feat: add reusable elementwise relu support --- README.md | 12 +- README.zh.md | 11 +- docs/how-to-add-an-operator.md | 34 +++- examples/06_relu.py | 30 +++ include/operator_runtime/detail/elementwise.h | 163 ++++++++++++++++ include/operator_runtime/operator_runtime.h | 1 + include/operator_runtime/ops/relu.h | 32 +++ include/operator_runtime/tensor_view.h | 13 +- ops/CMakeLists.txt | 4 + .../elementwise/metax/elementwise_metax.h | 184 ++++++++++++++++++ .../elementwise/nvidia/elementwise_nvidia.cuh | 184 ++++++++++++++++++ ops/common/metax/README.md | 12 -- ops/common/nvidia/elementwise.cuh | 45 ----- ops/common/nvidia/indexing.cuh | 12 -- ops/elementwise/relu/metax/relu_metax.maca | 84 ++++++++ ops/elementwise/relu/nvidia/relu_cuda.cu | 84 ++++++++ ops/elementwise/relu/relu_common.h | 76 ++++++++ python/operator_runtime/__init__.py | 2 + python/operator_runtime/_internal/__init__.py | 2 + python/operator_runtime/_internal/bindings.py | 49 +++++ python/operator_runtime/ops/__init__.py | 2 + python/operator_runtime/ops/relu.py | 80 ++++++++ tests/bench/relu.py | 45 +++++ tests/cases/relu.py | 35 ++++ tests/op_tests/test_relu.py | 115 +++++++++++ 25 files changed, 1226 insertions(+), 85 deletions(-) create mode 100644 examples/06_relu.py create mode 100644 include/operator_runtime/ops/relu.h create mode 100644 ops/common/elementwise/metax/elementwise_metax.h create mode 100644 ops/common/elementwise/nvidia/elementwise_nvidia.cuh delete mode 100644 ops/common/metax/README.md delete mode 100644 ops/common/nvidia/elementwise.cuh delete mode 100644 ops/common/nvidia/indexing.cuh create mode 100644 ops/elementwise/relu/metax/relu_metax.maca create mode 100644 ops/elementwise/relu/nvidia/relu_cuda.cu create mode 100644 ops/elementwise/relu/relu_common.h create mode 100644 python/operator_runtime/ops/relu.py create mode 100644 tests/bench/relu.py create mode 100644 tests/cases/relu.py create mode 100644 tests/op_tests/test_relu.py diff --git a/README.md b/README.md index e800c52..239ec18 100644 --- a/README.md +++ b/README.md @@ -3,7 +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. +1. Create a directory under `ops//` with backend implementations, or under + `ops/elementwise//` when the operator can 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. @@ -31,6 +32,7 @@ python/ | `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 | +| `relu` | runnable, elementwise fp16/fp32 with `negative_slope` | not implemented | not implemented | | `softmax` | runnable, row-wise fp32 | runnable when TileLang is installed | runnable, row-wise fp32 | ## Setup @@ -62,13 +64,17 @@ 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`. +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`. + +## 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. ## Production Mapping | Training concept | Production equivalent | | --- | --- | -| directory convention `ops//nvidia/*.cu` | build system auto-discovery / operator registry | +| directory convention `ops//nvidia/*.cu` and `ops/elementwise//nvidia/*.cu` | build system auto-discovery / operator registry | | C header `include/operator_runtime/ops/.h` | reviewed operator API contract | | descriptor lifecycle | create, workspace, execute, destroy | | `tests/cases/.py` | correctness, layout, and API contract coverage | diff --git a/README.zh.md b/README.zh.md index 3a86606..52a8bf2 100644 --- a/README.zh.md +++ b/README.zh.md @@ -2,7 +2,7 @@ 这个仓库是一个面向训练的 GPU 算子运行时。它体积很小，但工作流尽量贴近真实的算子库开发流程： -1. 在 `ops//` 下创建后端实现目录。 +1. 在 `ops//` 下创建后端实现目录；如果算子可以复用 elementwise 框架，则放到 `ops/elementwise//`。 2. 构建系统通过目录约定自动发现源码。 3. 实现后端专属的 descriptor 生命周期（create、workspace、execute、destroy）。 4. 提供 Python API，支持 out-of-place、out-variant 和 prepared 执行。 @@ -30,6 +30,7 @@ python/ | `copy` | 可运行 | 安装 TileLang 后可运行 | 可运行 | | `vector_add` | 可运行 | 安装 TileLang 后可运行 | 可运行 | | `reduce_sum` | 可运行，row-wise fp32 | 安装 TileLang 后可运行 | 可运行，row-wise fp32 | +| `relu` | 可运行，elementwise fp16/fp32，支持 `negative_slope` | 未实现 | 未实现 | | `softmax` | 可运行，row-wise fp32 | 安装 TileLang 后可运行 | 可运行，row-wise fp32 | ## 安装 @@ -61,13 +62,17 @@ python tests/run_ops.py --op all --backend nvidia --mode bench ./scripts/build_metax.sh test ``` -TileLang 后端需要安装 `tilelang` Python 包。MetaX 后端使用独立构建产物，构建时应关闭 NVIDIA 变体，并将后端源码放在 `ops/*/metax/*.maca`。 +TileLang 后端需要安装 `tilelang` Python 包。MetaX 后端使用独立构建产物，构建时应关闭 NVIDIA 变体，并将后端源码放在 `ops/*/metax/*.maca` 或 `ops/elementwise/*/metax/*.maca`。 + +## 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//nvidia/*.cu` 目录约定 | 构建系统自动发现 / 算子注册 | +| `ops//nvidia/*.cu` 和 `ops/elementwise//nvidia/*.cu` 目录约定 | 构建系统自动发现 / 算子注册 | | `include/operator_runtime/ops/.h` 头文件 | 经过评审的算子 API 契约 | | descriptor 生命周期 | create、workspace、execute、destroy | | `tests/cases/.py` | 正确性、布局和 API 契约覆盖 | diff --git a/docs/how-to-add-an-operator.md b/docs/how-to-add-an-operator.md index f608913..8ebc3ab 100644 --- a/docs/how-to-add-an-operator.md +++ b/docs/how-to-add-an-operator.md @@ -4,7 +4,7 @@ 在当前项目里，新增一个可运行算子的最小闭环包括四部分： -1. 在 `ops/<算子名>/` 下补后端实现。 +1. 在 `ops/<算子名>/` 下补后端实现；普通逐元素算子优先放到 `ops/elementwise/<算子名>/`。 2. 在 `python/operator_runtime/ops/` 下补 Python API。 3. 在 `tests/` 下补正确性测试和 benchmark 入口。 4. 重新构建并验证。 @@ -25,11 +25,18 @@ ## Step 2：创建算子目录 -先在 `ops/<算子名>/` 下建立对应目录。 +先选择算子目录： + +- 普通手写 kernel：放在 `ops/<算子名>/`。 +- 可复用公共 shape/stride/broadcast 逐元素框架的算子：放在 `ops/elementwise/<算子名>/`。 + +`relu` 是 elementwise 框架示例，目录为 `ops/elementwise/relu/nvidia/`，公共 NVIDIA launcher 位于 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`。 + +先在选定目录下建立对应后端目录。 当前建议至少补齐： -- `ops/<算子名>/nvidia/` +- `ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` - `python/operator_runtime/ops/<算子名>.py` - `tests/cases/<算子名>.py` - `tests/op_tests/test_<算子名>.py` @@ -40,9 +47,9 @@ ## Step 3：实现 NVIDIA 后端 -`ops/<算子名>/nvidia/` 这一层负责 C++/CUDA 实现。 +`ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` 这一层负责 C++/CUDA 实现。 -通常需要这几部分： +普通手写 kernel 通常需要这几部分： 1. kernel 文件 2. C API 头文件 @@ -50,6 +57,14 @@ 这里最关键的是保持现有命名约定一致，因为 Python 侧会按固定符号名去找函数。 +如果是 elementwise 算子，不需要每个算子重复写 shape/stride kernel。推荐复用 `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()`。 +4. execute 阶段只做 dtype dispatch，并调用 `oprt::elementwise::nvidia::launch(...)`。 +5. 算子自身只提供 device functor，例如 `relu` 的 `value > 0 ? value : value * negative_slope`。 + 一个 NVIDIA 算子需要完整提供四个生命周期接口： 1. create @@ -138,7 +153,7 @@ API contract 测试主要覆盖 shape 不匹配、dtype 不匹配、非 contiguo ## Step 9：重新配置和编译 -因为 `ops/CMakeLists.txt` 是通过 glob 自动发现 `ops/*/nvidia/*.cu`，所以新增 `.cu` 之后要重新配置。 +因为 `ops/CMakeLists.txt` 是通过 glob 自动发现 `ops/*/nvidia/*.cu` 和 `ops/elementwise/*/nvidia/*.cu`，所以新增 `.cu` 之后要重新配置。 顺序是： @@ -177,17 +192,18 @@ API contract 测试主要覆盖 shape 不匹配、dtype 不匹配、非 contiguo ## 现有模板怎么选 - `copy`：适合最简单的单输入单输出流程。 -- `vector_add`：适合标准 elementwise 双输入算子。 +- `vector_add`：适合标准 contiguous 双输入算子。 +- `relu`：适合复用 elementwise 框架的单输入逐元素算子，也展示了额外标量参数 `negative_slope` 的 C/Python 绑定方式。 - `reduce_sum`：适合带 reduce 维度的算子。 - `softmax`：适合带更明确 shape 约束和归一化逻辑的算子。 -如果新算子本质上是普通 elementwise，优先参考 `vector_add`。 +如果新算子本质上是普通 elementwise，优先参考 `relu` 和 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`；如果只想写最简单 contiguous kernel，再参考 `vector_add`。 ## 最终检查清单 提交前至少确认以下内容都已完成： -1. `ops/<算子名>/nvidia/` 已补齐实现。 +1. `ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` 已补齐实现。 2. `python/operator_runtime/ops/<算子名>.py` 已补齐。 3. 两个 `__init__.py` 已导出新接口。 4. `tests/cases/<算子名>.py` 已补数据。 diff --git a/examples/06_relu.py b/examples/06_relu.py new file mode 100644 index 0000000..ad0cb84 --- /dev/null +++ b/examples/06_relu.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "python")) +sys.path.insert(0, str(ROOT)) + +import torch + +from operator_runtime import relu + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--backend", default="nvidia", choices=["nvidia"]) + parser.add_argument("--negative-slope", type=float, default=0.0) + args = parser.parse_args() + + src = torch.randn((1024,), device="cuda", dtype=torch.float32) + out = relu(src, negative_slope=args.negative_slope, backend=args.backend) + expected = torch.where(src > 0, src, src * args.negative_slope) + torch.testing.assert_close(out, expected) + print(f"relu ok ({args.backend}, negative_slope={args.negative_slope})") + + +if __name__ == "__main__": + main() diff --git a/include/operator_runtime/detail/elementwise.h b/include/operator_runtime/detail/elementwise.h index 653787f..25d4382 100644 --- a/include/operator_runtime/detail/elementwise.h +++ b/include/operator_runtime/detail/elementwise.h @@ -1,11 +1,60 @@ #pragma once #include "operator_runtime/tensor_view.h" +#include "operator_runtime/detail/tensor_checks.h" #ifdef __cplusplus +#include +#include +#include +#include +#include + namespace oprt { +struct ElementwiseInfo { + struct InputInfo { + std::array shape{}; + std::array strides{}; + bool contiguous = false; + bool broadcasted = false; + }; + + oprt_dtype_t dtype = OPRT_DTYPE_F32; + int32_t ndim = 0; + int64_t elements = 0; + size_t input_count = 0; + std::array output_shape{}; + std::array output_strides{}; + bool output_contiguous = false; + std::vector inputs; + + size_t meta_bytes() const { + return sizeof(DeviceMeta) + input_count * sizeof(DeviceInputMeta); + } + + size_t workspace_bytes() const { + return input_count * sizeof(const void *) + meta_bytes(); + } + + struct DeviceInputMeta { + int64_t shape[OPRT_MAX_DIMS]; + int64_t strides[OPRT_MAX_DIMS]; + uint8_t contiguous; + uint8_t broadcasted; + }; + + struct DeviceMeta { + int32_t ndim; + uint32_t input_count; + int64_t elements; + int64_t output_shape[OPRT_MAX_DIMS]; + int64_t output_strides[OPRT_MAX_DIMS]; + uint8_t output_contiguous; + }; +}; + inline bool elementwise_fast_path(const oprt_tensor_view_t &out, const oprt_tensor_view_t &a) { return same_shape(out, a) && is_contiguous(out) && is_contiguous(a); @@ -18,6 +67,120 @@ inline bool elementwise_fast_path(const oprt_tensor_view_t &out, is_contiguous(out) && is_contiguous(a) && is_contiguous(b); } +inline bool elementwise_input_broadcastable_to(const oprt_tensor_view_t &out, + const oprt_tensor_view_t &in) { + if (out.ndim < 0 || out.ndim > OPRT_MAX_DIMS || in.ndim < 0 || in.ndim > OPRT_MAX_DIMS) { + return false; + } + if (in.ndim > out.ndim) { + return false; + } + int32_t dim_offset = out.ndim - in.ndim; + for (int32_t out_dim = 0; out_dim < out.ndim; ++out_dim) { + int32_t in_dim = out_dim - dim_offset; + if (in_dim < 0) { + continue; + } + int64_t in_extent = in.shape[in_dim]; + int64_t out_extent = out.shape[out_dim]; + if (in_extent != out_extent && in_extent != 1) { + return false; + } + } + return true; +} + +inline bool elementwise_has_broadcast(const oprt_tensor_view_t &out, + const oprt_tensor_view_t &in) { + if (has_broadcast_dim(in)) { + return true; + } + if (in.ndim != out.ndim) { + return true; + } + for (int32_t i = 0; i < out.ndim; ++i) { + if (in.shape[i] != out.shape[i]) { + return true; + } + } + return false; +} + +inline bool elementwise_same_shape_inputs(const oprt_tensor_view_t &out, + const std::vector &inputs) { + for (const auto *input : inputs) { + if (input == nullptr || !same_shape(out, *input)) { + return false; + } + } + return true; +} + +inline oprt_status_t create_elementwise_info(const oprt_tensor_view_t *out, + const std::vector &inputs, + ElementwiseInfo *info) { + if (info == nullptr || out == nullptr || inputs.empty()) { + return OPRT_ERR_INVALID_ARG; + } + + auto status = check_tensor(out); + if (status != OPRT_SUCCESS) { + return status; + } + if (has_broadcast_dim(*out)) { + return OPRT_ERR_INVALID_ARG; + } + + ElementwiseInfo next; + next.dtype = out->dtype; + next.ndim = out->ndim; + next.elements = numel(*out); + next.input_count = inputs.size(); + next.output_contiguous = is_contiguous(*out); + + for (int32_t i = 0; i < out->ndim; ++i) { + next.output_shape[i] = out->shape[i]; + next.output_strides[i] = out->strides[i]; + } + + next.inputs.reserve(inputs.size()); + for (const auto *input : inputs) { + status = check_tensor(input); + if (status != OPRT_SUCCESS) { + return status; + } + if (input->dtype != out->dtype) { + return OPRT_ERR_UNSUPPORTED_DTYPE; + } + if (!elementwise_input_broadcastable_to(*out, *input)) { + return OPRT_ERR_INVALID_ARG; + } + + ElementwiseInfo::InputInfo input_info; + input_info.contiguous = same_shape(*out, *input) && is_contiguous(*input); + input_info.broadcasted = elementwise_has_broadcast(*out, *input); + + int32_t dim_offset = out->ndim - input->ndim; + for (int32_t out_dim = 0; out_dim < out->ndim; ++out_dim) { + int32_t in_dim = out_dim - dim_offset; + if (in_dim < 0) { + input_info.shape[out_dim] = out->shape[out_dim]; + input_info.strides[out_dim] = 0; + input_info.broadcasted = true; + continue; + } + input_info.shape[out_dim] = out->shape[out_dim]; + input_info.strides[out_dim] = input->shape[in_dim] == 1 && out->shape[out_dim] != 1 + ? 0 + : input->strides[in_dim]; + } + next.inputs.push_back(input_info); + } + + *info = std::move(next); + return OPRT_SUCCESS; +} + } // namespace oprt #endif diff --git a/include/operator_runtime/operator_runtime.h b/include/operator_runtime/operator_runtime.h index 06b06e7..c2a687a 100644 --- a/include/operator_runtime/operator_runtime.h +++ b/include/operator_runtime/operator_runtime.h @@ -5,5 +5,6 @@ #include "operator_runtime/tensor_view.h" #include "operator_runtime/ops/copy.h" #include "operator_runtime/ops/reduce_sum.h" +#include "operator_runtime/ops/relu.h" #include "operator_runtime/ops/softmax.h" #include "operator_runtime/ops/vector_add.h" diff --git a/include/operator_runtime/ops/relu.h b/include/operator_runtime/ops/relu.h new file mode 100644 index 0000000..cfe9c08 --- /dev/null +++ b/include/operator_runtime/ops/relu.h @@ -0,0 +1,32 @@ +#pragma once + +#include "operator_runtime/api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +OPRT_EXPORT oprt_status_t oprt_create_relu_descriptor( + oprt_operator_descriptor_t *desc, + const oprt_tensor_view_t *out, + const oprt_tensor_view_t *in, + float negative_slope); + +OPRT_EXPORT oprt_status_t oprt_get_relu_workspace_size( + oprt_operator_descriptor_t desc, + size_t *size); + +OPRT_EXPORT oprt_status_t oprt_execute_relu( + 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_relu_descriptor( + oprt_operator_descriptor_t desc); + +#ifdef __cplusplus +} +#endif diff --git a/include/operator_runtime/tensor_view.h b/include/operator_runtime/tensor_view.h index 367f908..0977762 100644 --- a/include/operator_runtime/tensor_view.h +++ b/include/operator_runtime/tensor_view.h @@ -49,7 +49,18 @@ inline bool same_shape(const oprt_tensor_view_t &a, const oprt_tensor_view_t &b) return true; } +inline bool has_broadcast_dim(const oprt_tensor_view_t &view) { + if (view.ndim < 0 || view.ndim > OPRT_MAX_DIMS) { + return false; + } + for (int32_t i = 0; i < view.ndim; ++i) { + if (view.shape[i] > 1 && view.strides[i] == 0) { + return true; + } + } + return false; +} + } // namespace oprt #endif - diff --git a/ops/CMakeLists.txt b/ops/CMakeLists.txt index 80ca0eb..8fecdbb 100644 --- a/ops/CMakeLists.txt +++ b/ops/CMakeLists.txt @@ -1,6 +1,8 @@ # Convention: operator CUDA sources auto-discovered by directory layout: # ops//nvidia/*.cu -> compiled when CAMP_ENABLE_NVIDIA is ON +# ops/elementwise//nvidia/*.cu -> compiled when CAMP_ENABLE_NVIDIA is ON # ops//metax/*.maca -> compiled when CAMP_ENABLE_METAX is ON +# ops/elementwise//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 @@ -11,6 +13,7 @@ set(CAMP_NVIDIA_SOURCES "") if(CAMP_ENABLE_NVIDIA) file(GLOB_RECURSE CAMP_NVIDIA_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/*/nvidia/*.cu + ${CMAKE_CURRENT_SOURCE_DIR}/elementwise/*/nvidia/*.cu ) endif() @@ -18,6 +21,7 @@ set(CAMP_METAX_SOURCES "") if(CAMP_ENABLE_METAX) file(GLOB_RECURSE CAMP_METAX_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/*/metax/*.maca + ${CMAKE_CURRENT_SOURCE_DIR}/elementwise/*/metax/*.maca ) set_source_files_properties(${CAMP_METAX_SOURCES} PROPERTIES LANGUAGE MACA) endif() diff --git a/ops/common/elementwise/metax/elementwise_metax.h b/ops/common/elementwise/metax/elementwise_metax.h new file mode 100644 index 0000000..b95c3c4 --- /dev/null +++ b/ops/common/elementwise/metax/elementwise_metax.h @@ -0,0 +1,184 @@ +#pragma once + +#include "operator_runtime/detail/cuda_helpers.h" +#include "operator_runtime/detail/elementwise.h" + +#include +#include + +#include +#include +#include + +namespace oprt::elementwise::metax { + +template +__device__ inline T cast_scalar(float value) { + return static_cast(value); +} + +template <> +__device__ inline half cast_scalar(float value) { + return __float2half(value); +} + +__device__ inline int64_t index_to_offset(int64_t linear, + int32_t ndim, + const int64_t *shape, + const int64_t *strides) { + int64_t offset = 0; + for (int32_t dim = ndim - 1; dim >= 0; --dim) { + int64_t coord = linear % shape[dim]; + linear /= shape[dim]; + offset += coord * strides[dim]; + } + return offset; +} + +template +__device__ inline int64_t input_offset( + int64_t linear, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta) { + const auto &m = input_meta[Input]; + return m.contiguous != 0 + ? linear + : index_to_offset(linear, meta->ndim, m.shape, m.strides); +} + +template +__device__ inline T load_input( + int64_t linear, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta) { + const T *input = static_cast(inputs[Input]); + return input[input_offset(linear, meta, input_meta)]; +} + +template +__device__ inline T apply_at_index_impl( + int64_t linear, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta, + Op op, + std::index_sequence, + Args... args) { + static_assert(sizeof...(Is) > 0, "elementwise launch requires at least one input"); + return op(load_input(linear, inputs, meta, input_meta)..., args...); +} + +template +__device__ inline T apply_at_index( + int64_t linear, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta, + Op op, + Args... args) { + return apply_at_index_impl( + linear, + inputs, + meta, + input_meta, + op, + std::make_index_sequence{}, + args...); +} + +template +__global__ void elementwise_kernel( + T *__restrict__ out, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta, + Op op, + Args... args) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + const bool output_contiguous = meta->output_contiguous != 0; + + for (int64_t i = idx; i < meta->elements; i += stride) { + int64_t out_offset = output_contiguous + ? i + : index_to_offset(i, meta->ndim, meta->output_shape, meta->output_strides); + out[out_offset] = apply_at_index(i, inputs, meta, input_meta, op, args...); + } +} + +template +oprt_status_t launch( + const oprt::ElementwiseInfo &info, + void *workspace, + void *out, + const std::array &inputs, + oprt_stream_t stream, + Op op, + Args... args) { + if (info.input_count != N || (info.workspace_bytes() != 0 && workspace == nullptr)) { + return OPRT_ERR_INVALID_ARG; + } + if (info.elements == 0) { + return OPRT_SUCCESS; + } + + auto *workspace_bytes = static_cast(workspace); + auto **device_inputs = reinterpret_cast(workspace_bytes); + auto *device_meta = reinterpret_cast( + workspace_bytes + N * sizeof(const void *)); + auto *device_input_meta = reinterpret_cast(device_meta + 1); + + oprt::ElementwiseInfo::DeviceMeta host_meta{}; + host_meta.ndim = info.ndim; + host_meta.input_count = static_cast(N); + host_meta.elements = info.elements; + host_meta.output_contiguous = info.output_contiguous ? 1 : 0; + for (int32_t i = 0; i < info.ndim; ++i) { + host_meta.output_shape[i] = info.output_shape[i]; + host_meta.output_strides[i] = info.output_strides[i]; + } + + std::array host_input_meta{}; + for (size_t input = 0; input < N; ++input) { + const auto &src = info.inputs[input]; + auto &dst = host_input_meta[input]; + dst.contiguous = src.contiguous ? 1 : 0; + dst.broadcasted = src.broadcasted ? 1 : 0; + for (int32_t dim = 0; dim < info.ndim; ++dim) { + dst.shape[dim] = src.shape[dim]; + dst.strides[dim] = src.strides[dim]; + } + } + + cudaStream_t cuda_stream = oprt::as_cuda_stream(stream); + OPRT_CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(device_inputs, + inputs.data(), + N * sizeof(const void *), + cudaMemcpyHostToDevice, + cuda_stream)); + OPRT_CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(device_meta, + &host_meta, + sizeof(host_meta), + cudaMemcpyHostToDevice, + cuda_stream)); + OPRT_CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(device_input_meta, + host_input_meta.data(), + N * sizeof(oprt::ElementwiseInfo::DeviceInputMeta), + cudaMemcpyHostToDevice, + cuda_stream)); + + constexpr int threads = 256; + int blocks = oprt::blocks_for(info.elements, threads); + elementwise_kernel<<>>( + static_cast(out), + device_inputs, + device_meta, + device_input_meta, + op, + args...); + OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError()); + return OPRT_SUCCESS; +} + +} // namespace oprt::elementwise::metax diff --git a/ops/common/elementwise/nvidia/elementwise_nvidia.cuh b/ops/common/elementwise/nvidia/elementwise_nvidia.cuh new file mode 100644 index 0000000..8a43c61 --- /dev/null +++ b/ops/common/elementwise/nvidia/elementwise_nvidia.cuh @@ -0,0 +1,184 @@ +#pragma once + +#include "operator_runtime/detail/cuda_helpers.h" +#include "operator_runtime/detail/elementwise.h" + +#include +#include + +#include +#include +#include + +namespace oprt::elementwise::nvidia { + +template +__device__ inline T cast_scalar(float value) { + return static_cast(value); +} + +template <> +__device__ inline half cast_scalar(float value) { + return __float2half(value); +} + +__device__ inline int64_t index_to_offset(int64_t linear, + int32_t ndim, + const int64_t *shape, + const int64_t *strides) { + int64_t offset = 0; + for (int32_t dim = ndim - 1; dim >= 0; --dim) { + int64_t coord = linear % shape[dim]; + linear /= shape[dim]; + offset += coord * strides[dim]; + } + return offset; +} + +template +__device__ inline int64_t input_offset( + int64_t linear, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta) { + const auto &m = input_meta[Input]; + return m.contiguous != 0 + ? linear + : index_to_offset(linear, meta->ndim, m.shape, m.strides); +} + +template +__device__ inline T load_input( + int64_t linear, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta) { + const T *input = static_cast(inputs[Input]); + return input[input_offset(linear, meta, input_meta)]; +} + +template +__device__ inline T apply_at_index_impl( + int64_t linear, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta, + Op op, + std::index_sequence, + Args... args) { + static_assert(sizeof...(Is) > 0, "elementwise launch requires at least one input"); + return op(load_input(linear, inputs, meta, input_meta)..., args...); +} + +template +__device__ inline T apply_at_index( + int64_t linear, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta, + Op op, + Args... args) { + return apply_at_index_impl( + linear, + inputs, + meta, + input_meta, + op, + std::make_index_sequence{}, + args...); +} + +template +__global__ void elementwise_kernel( + T *__restrict__ out, + const void *const *__restrict__ inputs, + const oprt::ElementwiseInfo::DeviceMeta *__restrict__ meta, + const oprt::ElementwiseInfo::DeviceInputMeta *__restrict__ input_meta, + Op op, + Args... args) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + const bool output_contiguous = meta->output_contiguous != 0; + + for (int64_t i = idx; i < meta->elements; i += stride) { + int64_t out_offset = output_contiguous + ? i + : index_to_offset(i, meta->ndim, meta->output_shape, meta->output_strides); + out[out_offset] = apply_at_index(i, inputs, meta, input_meta, op, args...); + } +} + +template +oprt_status_t launch( + const oprt::ElementwiseInfo &info, + void *workspace, + void *out, + const std::array &inputs, + oprt_stream_t stream, + Op op, + Args... args) { + if (info.input_count != N || (info.workspace_bytes() != 0 && workspace == nullptr)) { + return OPRT_ERR_INVALID_ARG; + } + if (info.elements == 0) { + return OPRT_SUCCESS; + } + + auto *workspace_bytes = static_cast(workspace); + auto **device_inputs = reinterpret_cast(workspace_bytes); + auto *device_meta = reinterpret_cast( + workspace_bytes + N * sizeof(const void *)); + auto *device_input_meta = reinterpret_cast(device_meta + 1); + + oprt::ElementwiseInfo::DeviceMeta host_meta{}; + host_meta.ndim = info.ndim; + host_meta.input_count = static_cast(N); + host_meta.elements = info.elements; + host_meta.output_contiguous = info.output_contiguous ? 1 : 0; + for (int32_t i = 0; i < info.ndim; ++i) { + host_meta.output_shape[i] = info.output_shape[i]; + host_meta.output_strides[i] = info.output_strides[i]; + } + + std::array host_input_meta{}; + for (size_t input = 0; input < N; ++input) { + const auto &src = info.inputs[input]; + auto &dst = host_input_meta[input]; + dst.contiguous = src.contiguous ? 1 : 0; + dst.broadcasted = src.broadcasted ? 1 : 0; + for (int32_t dim = 0; dim < info.ndim; ++dim) { + dst.shape[dim] = src.shape[dim]; + dst.strides[dim] = src.strides[dim]; + } + } + + cudaStream_t cuda_stream = oprt::as_cuda_stream(stream); + OPRT_CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(device_inputs, + inputs.data(), + N * sizeof(const void *), + cudaMemcpyHostToDevice, + cuda_stream)); + OPRT_CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(device_meta, + &host_meta, + sizeof(host_meta), + cudaMemcpyHostToDevice, + cuda_stream)); + OPRT_CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(device_input_meta, + host_input_meta.data(), + N * sizeof(oprt::ElementwiseInfo::DeviceInputMeta), + cudaMemcpyHostToDevice, + cuda_stream)); + + constexpr int threads = 256; + int blocks = oprt::blocks_for(info.elements, threads); + elementwise_kernel<<>>( + static_cast(out), + device_inputs, + device_meta, + device_input_meta, + op, + args...); + OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError()); + return OPRT_SUCCESS; +} + +} // namespace oprt::elementwise::nvidia diff --git a/ops/common/metax/README.md b/ops/common/metax/README.md deleted file mode 100644 index b9794a9..0000000 --- a/ops/common/metax/README.md +++ /dev/null @@ -1,12 +0,0 @@ -# MetaX Backend - -The training runtime keeps MetaX interfaces aligned with the compiled C ABI -backends and builds MetaX as a separate variant through MACA/cu-bridge. - -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/common/nvidia/elementwise.cuh b/ops/common/nvidia/elementwise.cuh deleted file mode 100644 index 799a5d8..0000000 --- a/ops/common/nvidia/elementwise.cuh +++ /dev/null @@ -1,45 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace oprt::nvidia { - -template -__global__ void unary_contiguous_kernel(T *out, const T *in, int64_t n, UnaryOp op) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t stride = int64_t(blockDim.x) * gridDim.x; - for (int64_t i = idx; i < n; i += stride) { - out[i] = op(in[i]); - } -} - -template -__global__ void binary_contiguous_kernel(T *out, const T *a, const T *b, int64_t n, BinaryOp op) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t stride = int64_t(blockDim.x) * gridDim.x; - for (int64_t i = idx; i < n; i += stride) { - out[i] = op(a[i], b[i]); - } -} - -struct CopyOp { - template - __device__ T operator()(T value) const { - return value; - } -}; - -struct AddOp { - __device__ float operator()(float a, float b) const { - return a + b; - } - - __device__ half operator()(half a, half b) const { - return __hadd(a, b); - } -}; - -} // namespace oprt::nvidia - diff --git a/ops/common/nvidia/indexing.cuh b/ops/common/nvidia/indexing.cuh deleted file mode 100644 index 69c340a..0000000 --- a/ops/common/nvidia/indexing.cuh +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include - -namespace oprt::nvidia { - -__device__ inline int64_t contiguous_offset(int64_t linear) { - return linear; -} - -} // namespace oprt::nvidia - diff --git a/ops/elementwise/relu/metax/relu_metax.maca b/ops/elementwise/relu/metax/relu_metax.maca new file mode 100644 index 0000000..7f0636e --- /dev/null +++ b/ops/elementwise/relu/metax/relu_metax.maca @@ -0,0 +1,84 @@ +#include "operator_runtime/ops/relu.h" + +#include "ops/common/elementwise/metax/elementwise_metax.h" +#include "ops/elementwise/relu/relu_common.h" + +#include + +namespace { + +struct ReluOp { + template + __device__ T operator()(T value, float negative_slope) const { + T zero = oprt::elementwise::metax::cast_scalar(0.0f); + T slope = oprt::elementwise::metax::cast_scalar(negative_slope); + return value > zero ? value : value * slope; + } +}; + +template <> +__device__ half ReluOp::operator()(half value, float negative_slope) const { + half zero = __float2half(0.0f); + half slope = __float2half(negative_slope); + return __hgt(value, zero) ? value : __hmul(value, slope); +} + +template +oprt_status_t launch_relu(const oprt::elementwise::ReluDescriptor *desc, + void *workspace, + void *out, + const void *in, + oprt_stream_t stream) { + return oprt::elementwise::metax::launch( + desc->info, + workspace, + out, + std::array{in}, + stream, + ReluOp{}, + desc->negative_slope); +} + +} // namespace + +extern "C" OPRT_EXPORT oprt_status_t oprt_create_relu_descriptor( + oprt_operator_descriptor_t *desc, + const oprt_tensor_view_t *out, + const oprt_tensor_view_t *in, + float negative_slope) { + return oprt::elementwise::create_relu_descriptor_common(desc, out, in, negative_slope); +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_get_relu_workspace_size( + oprt_operator_descriptor_t desc, + size_t *size) { + return oprt::elementwise::get_relu_workspace_size_common(desc, size); +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_execute_relu( + oprt_operator_descriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *in, + oprt_stream_t stream) { + auto status = oprt::elementwise::validate_relu_execute_args(desc, workspace_size, out, in); + if (status != OPRT_SUCCESS) { + return status; + } + + auto *typed = static_cast(desc); + switch (typed->out_view.dtype) { + case OPRT_DTYPE_F16: + return launch_relu(typed, workspace, out, in, stream); + case OPRT_DTYPE_F32: + return launch_relu(typed, workspace, out, in, stream); + default: + return OPRT_ERR_UNSUPPORTED_DTYPE; + } +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_relu_descriptor( + oprt_operator_descriptor_t desc) { + return oprt::elementwise::destroy_relu_descriptor_common(desc); +} diff --git a/ops/elementwise/relu/nvidia/relu_cuda.cu b/ops/elementwise/relu/nvidia/relu_cuda.cu new file mode 100644 index 0000000..4633a32 --- /dev/null +++ b/ops/elementwise/relu/nvidia/relu_cuda.cu @@ -0,0 +1,84 @@ +#include "operator_runtime/ops/relu.h" + +#include "ops/common/elementwise/nvidia/elementwise_nvidia.cuh" +#include "ops/elementwise/relu/relu_common.h" + +#include + +namespace { + +struct ReluOp { + template + __device__ T operator()(T value, float negative_slope) const { + T zero = oprt::elementwise::nvidia::cast_scalar(0.0f); + T slope = oprt::elementwise::nvidia::cast_scalar(negative_slope); + return value > zero ? value : value * slope; + } +}; + +template <> +__device__ half ReluOp::operator()(half value, float negative_slope) const { + half zero = __float2half(0.0f); + half slope = __float2half(negative_slope); + return __hgt(value, zero) ? value : __hmul(value, slope); +} + +template +oprt_status_t launch_relu(const oprt::elementwise::ReluDescriptor *desc, + void *workspace, + void *out, + const void *in, + oprt_stream_t stream) { + return oprt::elementwise::nvidia::launch( + desc->info, + workspace, + out, + std::array{in}, + stream, + ReluOp{}, + desc->negative_slope); +} + +} // namespace + +extern "C" OPRT_EXPORT oprt_status_t oprt_create_relu_descriptor( + oprt_operator_descriptor_t *desc, + const oprt_tensor_view_t *out, + const oprt_tensor_view_t *in, + float negative_slope) { + return oprt::elementwise::create_relu_descriptor_common(desc, out, in, negative_slope); +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_get_relu_workspace_size( + oprt_operator_descriptor_t desc, + size_t *size) { + return oprt::elementwise::get_relu_workspace_size_common(desc, size); +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_execute_relu( + oprt_operator_descriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *in, + oprt_stream_t stream) { + auto status = oprt::elementwise::validate_relu_execute_args(desc, workspace_size, out, in); + if (status != OPRT_SUCCESS) { + return status; + } + + auto *typed = static_cast(desc); + switch (typed->out_view.dtype) { + case OPRT_DTYPE_F16: + return launch_relu(typed, workspace, out, in, stream); + case OPRT_DTYPE_F32: + return launch_relu(typed, workspace, out, in, stream); + default: + return OPRT_ERR_UNSUPPORTED_DTYPE; + } +} + +extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_relu_descriptor( + oprt_operator_descriptor_t desc) { + return oprt::elementwise::destroy_relu_descriptor_common(desc); +} diff --git a/ops/elementwise/relu/relu_common.h b/ops/elementwise/relu/relu_common.h new file mode 100644 index 0000000..4ede9b1 --- /dev/null +++ b/ops/elementwise/relu/relu_common.h @@ -0,0 +1,76 @@ +#pragma once + +#include "operator_runtime/descriptor.h" +#include "operator_runtime/detail/elementwise.h" + +namespace oprt::elementwise { + +struct ReluDescriptor final : oprt_operator_descriptor { + oprt_tensor_view_t out_view{}; + oprt_tensor_view_t in_view{}; + oprt::ElementwiseInfo info; + float negative_slope = 0.0f; + + const char *op_name() const override { + return "relu"; + } +}; + +inline oprt_status_t create_relu_descriptor_common( + oprt_operator_descriptor_t *desc, + const oprt_tensor_view_t *out, + const oprt_tensor_view_t *in, + float negative_slope) { + if (desc == nullptr) { + return OPRT_ERR_INVALID_ARG; + } + *desc = nullptr; + + if (out == nullptr || in == nullptr) { + return OPRT_ERR_INVALID_ARG; + } + + oprt::ElementwiseInfo info; + auto status = oprt::create_elementwise_info(out, {in}, &info); + if (status != OPRT_SUCCESS) { + 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; +} + +inline oprt_status_t validate_relu_execute_args(oprt_operator_descriptor_t desc, + size_t workspace_size, + void *out, + const void *in) { + 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; +} + +inline oprt_status_t destroy_relu_descriptor_common(oprt_operator_descriptor_t desc) { + delete desc; + return OPRT_SUCCESS; +} + +} // namespace oprt::elementwise diff --git a/python/operator_runtime/__init__.py b/python/operator_runtime/__init__.py index 5e54d1e..2d7d42f 100644 --- a/python/operator_runtime/__init__.py +++ b/python/operator_runtime/__init__.py @@ -3,6 +3,7 @@ from .ops import ( copy, copy_, prepare_copy, vector_add, vector_add_, prepare_vector_add, reduce_sum, reduce_sum_, prepare_reduce_sum, + relu, relu_, prepare_relu, softmax, softmax_, prepare_softmax, ) @@ -12,5 +13,6 @@ __all__ = [ "copy", "copy_", "prepare_copy", "vector_add", "vector_add_", "prepare_vector_add", "reduce_sum", "reduce_sum_", "prepare_reduce_sum", + "relu", "relu_", "prepare_relu", "softmax", "softmax_", "prepare_softmax", ] diff --git a/python/operator_runtime/_internal/__init__.py b/python/operator_runtime/_internal/__init__.py index a6311d1..17a4de7 100644 --- a/python/operator_runtime/_internal/__init__.py +++ b/python/operator_runtime/_internal/__init__.py @@ -6,6 +6,7 @@ from .bindings import ( OperatorRuntimeError, Status, bind_unary, + bind_relu, bind_binary, bind_reduce_like, check_status, @@ -24,6 +25,7 @@ __all__ = [ "OperatorRuntimeError", "Status", "bind_unary", + "bind_relu", "bind_binary", "bind_reduce_like", "check_status", diff --git a/python/operator_runtime/_internal/bindings.py b/python/operator_runtime/_internal/bindings.py index 23b6361..651ff50 100644 --- a/python/operator_runtime/_internal/bindings.py +++ b/python/operator_runtime/_internal/bindings.py @@ -84,6 +84,55 @@ def bind_unary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions return CFunctions(create, workspace, execute, destroy) +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) + + def bind_binary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions: lib = load_library() create_symbol = f"oprt_create_{name}_descriptor" diff --git a/python/operator_runtime/ops/__init__.py b/python/operator_runtime/ops/__init__.py index c1adfd5..b0147d5 100644 --- a/python/operator_runtime/ops/__init__.py +++ b/python/operator_runtime/ops/__init__.py @@ -1,11 +1,13 @@ from .copy import copy, copy_, prepare_copy from .vector_add import vector_add, vector_add_, prepare_vector_add from .reduce_sum import reduce_sum, reduce_sum_, prepare_reduce_sum +from .relu import relu, relu_, prepare_relu from .softmax import softmax, softmax_, prepare_softmax __all__ = [ "copy", "copy_", "prepare_copy", "vector_add", "vector_add_", "prepare_vector_add", "reduce_sum", "reduce_sum_", "prepare_reduce_sum", + "relu", "relu_", "prepare_relu", "softmax", "softmax_", "prepare_softmax", ] diff --git a/python/operator_runtime/ops/relu.py b/python/operator_runtime/ops/relu.py new file mode 100644 index 0000000..4ffef35 --- /dev/null +++ b/python/operator_runtime/ops/relu.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +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 + + +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") + + +def prepare_relu( + out: torch.Tensor, + src: torch.Tensor, + 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 build_prepared_op(funcs, create_args, (out, src), out) + + +def relu_( + out: torch.Tensor, + src: torch.Tensor, + negative_slope: float = 0.0, + backend: str | Backend = Backend.NVIDIA, +) -> torch.Tensor: + with prepare_relu(out, src, negative_slope, backend) as prepared: + prepared.run() + return out + + +def relu( + src: torch.Tensor, + negative_slope: float = 0.0, + backend: str | Backend = Backend.NVIDIA, +) -> torch.Tensor: + out = torch.empty_like(src) + return relu_(out, src, negative_slope, backend) diff --git a/tests/bench/relu.py b/tests/bench/relu.py new file mode 100644 index 0000000..81305d2 --- /dev/null +++ b/tests/bench/relu.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +PYTHON_DIR = ROOT / "python" +if str(PYTHON_DIR) not in sys.path: + sys.path.insert(0, str(PYTHON_DIR)) +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +import torch + +from operator_runtime_testing import cuda_time_ms, PerformanceResult +from tests.cases import relu as relu_cases + + +def bench_relu(backend: str) -> list[PerformanceResult]: + if backend not in ("nvidia", "metax"): + raise NotImplementedError("relu benchmark currently targets the C ABI elementwise framework") + + rows: list[PerformanceResult] = [] + for case in relu_cases.benchmark_cases(): + src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda") + out = torch.empty_like(src) + negative_slope = case["negative_slope"] + from operator_runtime import prepare_relu + with prepare_relu(out, src, negative_slope=negative_slope, backend=backend) as prepared: + runtime = cuda_time_ms(prepared.run_inputs, args=(src,)) + torch_ms = cuda_time_ms( + lambda src: torch.where(src > 0, src, src * negative_slope, out=out), + args=(src,), + ) + rows.append( + PerformanceResult( + "relu", + backend, + str(tuple(src.shape)), + str(src.dtype), + runtime, + torch_ms, + ) + ) + return rows diff --git a/tests/cases/relu.py b/tests/cases/relu.py new file mode 100644 index 0000000..eab1140 --- /dev/null +++ b/tests/cases/relu.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import torch + + +def correctness_cases(): + return [ + {"name": "relu_1k_fp32", "shape": (1024,), "dtype": torch.float32, "negative_slope": 0.0, "atol": 0, "rtol": 0}, + {"name": "relu_1k_fp16", "shape": (1024,), "dtype": torch.float16, "negative_slope": 0.0, "atol": 0, "rtol": 0}, + {"name": "leaky_relu_2d_fp32", "shape": (32, 64), "dtype": torch.float32, "negative_slope": 0.01, "atol": 1e-6, "rtol": 1e-6}, + {"name": "leaky_relu_2d_fp16", "shape": (32, 64), "dtype": torch.float16, "negative_slope": 0.01, "atol": 1e-3, "rtol": 1e-3}, + {"name": "non_contiguous_fp32", "shape": (16, 32), "dtype": torch.float32, "negative_slope": 0.01, "atol": 1e-6, "rtol": 1e-6}, + {"name": "broadcast_input_fp32", "src_shape": (64,), "out_shape": (32, 64), "dtype": torch.float32, "negative_slope": 0.01, "atol": 1e-6, "rtol": 1e-6}, + ] + + +def api_error_cases(): + return [ + {"name": "shape_mismatch", "shape": (16,), "out_shape": (8,), "dtype": torch.float32}, + {"name": "dtype_mismatch", "shape": (16,), "dtype": torch.float32, "out_dtype": torch.float16}, + {"name": "cpu_tensor", "shape": (16,), "dtype": torch.float32}, + {"name": "unsupported_dtype", "shape": (16,), "dtype": torch.float64}, + {"name": "broadcasted_output", "shape": (8, 16), "base_shape": (1, 16), "dtype": torch.float32}, + ] + + +def benchmark_cases(): + return [ + {"name": "contiguous_1001k_fp16", "shape": (1001 * 1024,), "dtype": torch.float16, "negative_slope": 0.0}, + {"name": "contiguous_4093k_fp16", "shape": (4093 * 1024,), "dtype": torch.float16, "negative_slope": 0.0}, + {"name": "contiguous_65521k_fp16", "shape": (65521 * 1024,), "dtype": torch.float16, "negative_slope": 0.0}, + {"name": "contiguous_1001k_fp32", "shape": (1001 * 1024,), "dtype": torch.float32, "negative_slope": 0.01}, + {"name": "contiguous_4093k_fp32", "shape": (4093 * 1024,), "dtype": torch.float32, "negative_slope": 0.01}, + {"name": "contiguous_65521k_fp32", "shape": (65521 * 1024,), "dtype": torch.float32, "negative_slope": 0.01}, + ] diff --git a/tests/op_tests/test_relu.py b/tests/op_tests/test_relu.py new file mode 100644 index 0000000..70e6eab --- /dev/null +++ b/tests/op_tests/test_relu.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import ctypes + +import pytest +import torch + +from operator_runtime import relu, relu_ +from operator_runtime.ops.relu import prepare_relu +from operator_runtime_testing import assert_close, require_cuda +from tests.cases import relu as relu_cases + + +def _reference(src: torch.Tensor, negative_slope: float) -> torch.Tensor: + return torch.where(src > 0, src, src * negative_slope) + + +@pytest.mark.parametrize("case", relu_cases.correctness_cases(), ids=lambda c: c["name"]) +def test_relu_correctness(case, backend): + require_cuda() + if backend not in ("nvidia", "metax"): + pytest.skip("relu currently targets the C ABI elementwise framework") + + if case["name"] == "broadcast_input_fp32": + src = torch.randn(case["src_shape"], dtype=case["dtype"], device="cuda") + out = torch.empty(case["out_shape"], dtype=case["dtype"], device="cuda") + relu_(out, src, negative_slope=case["negative_slope"], backend=backend) + ref_src = src.expand_as(out) + elif case["name"] == "non_contiguous_fp32": + src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda").t() + out = torch.empty_strided(src.shape, src.stride(), dtype=case["dtype"], device="cuda") + relu_(out, src, negative_slope=case["negative_slope"], backend=backend) + ref_src = src + else: + src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda") + out = relu(src, negative_slope=case["negative_slope"], backend=backend) + ref_src = src + + assert_close(out, _reference(ref_src, case["negative_slope"]), atol=case["atol"], rtol=case["rtol"]) + + +@pytest.mark.parametrize("case", relu_cases.api_error_cases(), ids=lambda c: c["name"]) +def test_relu_api_contract(case, backend): + require_cuda() + if backend not in ("nvidia", "metax"): + pytest.skip("relu currently targets the C ABI elementwise framework") + + if case["name"] == "shape_mismatch": + src = torch.randn(case["shape"], device="cuda", dtype=case["dtype"]) + out = torch.empty(case["out_shape"], device="cuda", dtype=case["dtype"]) + with pytest.raises(ValueError, match="broadcastable"): + relu_(out, src, backend=backend) + return + if case["name"] == "dtype_mismatch": + src = torch.randn(case["shape"], device="cuda", dtype=case["dtype"]) + out = torch.empty(case["shape"], device="cuda", dtype=case["out_dtype"]) + with pytest.raises(TypeError, match="matching dtypes"): + relu_(out, src, backend=backend) + return + if case["name"] == "cpu_tensor": + src = torch.randn(case["shape"], dtype=case["dtype"]) + out = torch.empty_like(src) + with pytest.raises(ValueError, match="CUDA tensors"): + relu_(out, src, backend=backend) + return + if case["name"] == "unsupported_dtype": + src = torch.randn(case["shape"], device="cuda", dtype=case["dtype"]) + with pytest.raises(TypeError, match="unsupported dtype"): + relu(src, backend=backend) + return + if case["name"] == "broadcasted_output": + base = torch.empty(case["base_shape"], device="cuda", dtype=case["dtype"]) + out = base.expand(case["shape"]) + src = torch.randn(case["shape"], device="cuda", dtype=case["dtype"]) + with pytest.raises(ValueError, match="writable output"): + relu_(out, src, backend=backend) + return + raise AssertionError(f"unhandled case: {case['name']}") + + +def test_prepared_relu_reuses_descriptor_with_new_inputs(backend): + require_cuda() + if backend not in ("nvidia", "metax"): + pytest.skip("descriptor lifecycle test targets C ABI backend") + + src = torch.randn((1024,), device="cuda", dtype=torch.float32) + out = torch.empty_like(src) + prepared = prepare_relu(out, src, negative_slope=0.01, backend=backend) + try: + prepared.run() + assert_close(out, _reference(src, 0.01), atol=1e-6, rtol=1e-6) + + next_src = torch.randn_like(src) + prepared.run_inputs(next_src) + assert_close(out, _reference(next_src, 0.01), atol=1e-6, rtol=1e-6) + finally: + prepared.destroy() + + +def test_prepared_relu_reports_insufficient_workspace(backend): + require_cuda() + if backend not in ("nvidia", "metax"): + pytest.skip("workspace contract test targets C ABI backend") + + src = torch.randn((16,), device="cuda", dtype=torch.float32) + out = torch.empty_like(src) + prepared = prepare_relu(out, src, backend=backend) + try: + original_workspace = prepared.workspace + prepared.workspace = torch.empty(0, dtype=torch.uint8, device="cuda") + with pytest.raises(RuntimeError, match="insufficient workspace"): + prepared.run() + prepared.workspace = original_workspace + finally: + prepared.destroy()