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

feat: add reusable elementwise relu support

This commit is contained in:
yutianyu 2026-05-12 11:30:19 +08:00
parent 757dbbec74
commit 9271fed062
25 changed files with 1226 additions and 85 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@ -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/<op>/` with backend implementations.
1. Create a directory under `ops/<op>/` with backend implementations, or under
`ops/elementwise/<op>/` 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/<op>/<backend>/`. 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/<op>/nvidia/*.cu` | build system auto-discovery / operator registry |
| directory convention `ops/<op>/nvidia/*.cu` and `ops/elementwise/<op>/nvidia/*.cu` | build system auto-discovery / operator registry |
| C header `include/operator_runtime/ops/<op>.h` | reviewed operator API contract |
| descriptor lifecycle | create, workspace, execute, destroy |
| `tests/cases/<op>.py` | correctness, layout, and API contract coverage |

11
README.zh.md
View File

@ -2,7 +2,7 @@
这个仓库是一个面向训练的 GPU 算子运行时。它体积很小，但工作流尽量贴近真实的算子库开发流程：
1. 在 `ops/<op>/` 下创建后端实现目录。
1. 在 `ops/<op>/` 下创建后端实现目录；如果算子可以复用 elementwise 框架，则放到 `ops/elementwise/<op>/`。
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/<op>/<backend>/`。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/<op>/nvidia/*.cu` 目录约定 | 构建系统自动发现 / 算子注册 |
| `ops/<op>/nvidia/*.cu` 和 `ops/elementwise/<op>/nvidia/*.cu` 目录约定 | 构建系统自动发现 / 算子注册 |
| `include/operator_runtime/ops/<op>.h` 头文件 | 经过评审的算子 API 契约 |
| descriptor 生命周期 | create、workspace、execute、destroy |
| `tests/cases/<op>.py` | 正确性、布局和 API 契约覆盖 |

34
docs/how-to-add-an-operator.md
View File

@ -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<T, N>(...)`。
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` 已补数据。

30
examples/06_relu.py Normal file
View File

@ -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()

163
include/operator_runtime/detail/elementwise.h
View File

@ -1,11 +1,60 @@
#pragma once
#include "operator_runtime/tensor_view.h"
#include "operator_runtime/detail/tensor_checks.h"
#ifdef __cplusplus
#include <array>
#include <cstddef>
#include <cstdint>
#include <utility>
#include <vector>
namespace oprt {
struct ElementwiseInfo {
struct InputInfo {
std::array<int64_t, OPRT_MAX_DIMS> shape{};
std::array<int64_t, OPRT_MAX_DIMS> 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<int64_t, OPRT_MAX_DIMS> output_shape{};
std::array<int64_t, OPRT_MAX_DIMS> output_strides{};
bool output_contiguous = false;
std::vector<InputInfo> 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<const oprt_tensor_view_t *> &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<const oprt_tensor_view_t *> &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

1
include/operator_runtime/operator_runtime.h
View File

@ -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"

32
include/operator_runtime/ops/relu.h Normal file
View File

@ -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

13
include/operator_runtime/tensor_view.h
View File

@ -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

4
ops/CMakeLists.txt
View File

@ -1,6 +1,8 @@
# Convention: operator CUDA sources auto-discovered by directory layout:
# ops/<operator>/nvidia/*.cu -> compiled when CAMP_ENABLE_NVIDIA is ON
# ops/elementwise/<operator>/nvidia/*.cu -> compiled when CAMP_ENABLE_NVIDIA is ON
# ops/<operator>/metax/*.maca -> compiled when CAMP_ENABLE_METAX is ON
# ops/elementwise/<operator>/metax/*.maca -> compiled when CAMP_ENABLE_METAX is ON
# After adding a new operator, re-run cmake to pick up new files.
set(CAMP_COMMON_SOURCES
@ -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()

184
ops/common/elementwise/metax/elementwise_metax.h Normal file
View File

@ -0,0 +1,184 @@
#pragma once
#include "operator_runtime/detail/cuda_helpers.h"
#include "operator_runtime/detail/elementwise.h"
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <array>
#include <cstdint>
#include <utility>
namespace oprt::elementwise::metax {
template <typename T>
__device__ inline T cast_scalar(float value) {
return static_cast<T>(value);
}
template <>
__device__ inline half cast_scalar<half>(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 <size_t Input>
__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 <typename T, size_t Input>
__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<const T *>(inputs[Input]);
return input[input_offset<Input>(linear, meta, input_meta)];
}
template <typename T, typename Op, size_t... Is, typename... Args>
__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<Is...>,
Args... args) {
static_assert(sizeof...(Is) > 0, "elementwise launch requires at least one input");
return op(load_input<T, Is>(linear, inputs, meta, input_meta)..., args...);
}
template <typename T, size_t N, typename Op, typename... Args>
__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<T>(
linear,
inputs,
meta,
input_meta,
op,
std::make_index_sequence<N>{},
args...);
}
template <typename T, size_t N, typename Op, typename... Args>
__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<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
int64_t stride = static_cast<int64_t>(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<T, N>(i, inputs, meta, input_meta, op, args...);
}
}
template <typename T, size_t N, typename Op, typename... Args>
oprt_status_t launch(
const oprt::ElementwiseInfo &info,
void *workspace,
void *out,
const std::array<const void *, N> &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<uint8_t *>(workspace);
auto **device_inputs = reinterpret_cast<const void **>(workspace_bytes);
auto *device_meta = reinterpret_cast<oprt::ElementwiseInfo::DeviceMeta *>(
workspace_bytes + N * sizeof(const void *));
auto *device_input_meta = reinterpret_cast<oprt::ElementwiseInfo::DeviceInputMeta *>(device_meta + 1);
oprt::ElementwiseInfo::DeviceMeta host_meta{};
host_meta.ndim = info.ndim;
host_meta.input_count = static_cast<uint32_t>(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<oprt::ElementwiseInfo::DeviceInputMeta, N> 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<T, N, Op, Args...><<<blocks, threads, 0, cuda_stream>>>(
static_cast<T *>(out),
device_inputs,
device_meta,
device_input_meta,
op,
args...);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace oprt::elementwise::metax

184
ops/common/elementwise/nvidia/elementwise_nvidia.cuh Normal file
View File

@ -0,0 +1,184 @@
#pragma once
#include "operator_runtime/detail/cuda_helpers.h"
#include "operator_runtime/detail/elementwise.h"
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <array>
#include <cstdint>
#include <utility>
namespace oprt::elementwise::nvidia {
template <typename T>
__device__ inline T cast_scalar(float value) {
return static_cast<T>(value);
}
template <>
__device__ inline half cast_scalar<half>(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 <size_t Input>
__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 <typename T, size_t Input>
__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<const T *>(inputs[Input]);
return input[input_offset<Input>(linear, meta, input_meta)];
}
template <typename T, typename Op, size_t... Is, typename... Args>
__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<Is...>,
Args... args) {
static_assert(sizeof...(Is) > 0, "elementwise launch requires at least one input");
return op(load_input<T, Is>(linear, inputs, meta, input_meta)..., args...);
}
template <typename T, size_t N, typename Op, typename... Args>
__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<T>(
linear,
inputs,
meta,
input_meta,
op,
std::make_index_sequence<N>{},
args...);
}
template <typename T, size_t N, typename Op, typename... Args>
__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<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
int64_t stride = static_cast<int64_t>(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<T, N>(i, inputs, meta, input_meta, op, args...);
}
}
template <typename T, size_t N, typename Op, typename... Args>
oprt_status_t launch(
const oprt::ElementwiseInfo &info,
void *workspace,
void *out,
const std::array<const void *, N> &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<uint8_t *>(workspace);
auto **device_inputs = reinterpret_cast<const void **>(workspace_bytes);
auto *device_meta = reinterpret_cast<oprt::ElementwiseInfo::DeviceMeta *>(
workspace_bytes + N * sizeof(const void *));
auto *device_input_meta = reinterpret_cast<oprt::ElementwiseInfo::DeviceInputMeta *>(device_meta + 1);
oprt::ElementwiseInfo::DeviceMeta host_meta{};
host_meta.ndim = info.ndim;
host_meta.input_count = static_cast<uint32_t>(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<oprt::ElementwiseInfo::DeviceInputMeta, N> 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<T, N, Op, Args...><<<blocks, threads, 0, cuda_stream>>>(
static_cast<T *>(out),
device_inputs,
device_meta,
device_input_meta,
op,
args...);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace oprt::elementwise::nvidia

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

@ -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`

45
ops/common/nvidia/elementwise.cuh
View File

@ -1,45 +0,0 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <stdint.h>
namespace oprt::nvidia {
template <typename T, typename UnaryOp>
__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 <typename T, typename BinaryOp>
__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 <typename T>
__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

12
ops/common/nvidia/indexing.cuh
View File

@ -1,12 +0,0 @@
#pragma once
#include <stdint.h>
namespace oprt::nvidia {
__device__ inline int64_t contiguous_offset(int64_t linear) {
return linear;
}
} // namespace oprt::nvidia

84
ops/elementwise/relu/metax/relu_metax.maca Normal file
View File

@ -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 <array>
namespace {
struct ReluOp {
template <typename T>
__device__ T operator()(T value, float negative_slope) const {
T zero = oprt::elementwise::metax::cast_scalar<T>(0.0f);
T slope = oprt::elementwise::metax::cast_scalar<T>(negative_slope);
return value > zero ? value : value * slope;
}
};
template <>
__device__ half ReluOp::operator()<half>(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 <typename T>
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<T, 1>(
desc->info,
workspace,
out,
std::array<const void *, 1>{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<const oprt::elementwise::ReluDescriptor *>(desc);
switch (typed->out_view.dtype) {
case OPRT_DTYPE_F16:
return launch_relu<half>(typed, workspace, out, in, stream);
case OPRT_DTYPE_F32:
return launch_relu<float>(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);
}

84
ops/elementwise/relu/nvidia/relu_cuda.cu Normal file
View File

@ -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 <array>
namespace {
struct ReluOp {
template <typename T>
__device__ T operator()(T value, float negative_slope) const {
T zero = oprt::elementwise::nvidia::cast_scalar<T>(0.0f);
T slope = oprt::elementwise::nvidia::cast_scalar<T>(negative_slope);
return value > zero ? value : value * slope;
}
};
template <>
__device__ half ReluOp::operator()<half>(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 <typename T>
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<T, 1>(
desc->info,
workspace,
out,
std::array<const void *, 1>{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<const oprt::elementwise::ReluDescriptor *>(desc);
switch (typed->out_view.dtype) {
case OPRT_DTYPE_F16:
return launch_relu<half>(typed, workspace, out, in, stream);
case OPRT_DTYPE_F32:
return launch_relu<float>(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);
}

76
ops/elementwise/relu/relu_common.h Normal file
View File

@ -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

2
python/operator_runtime/__init__.py
View File

@ -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",
]

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

@ -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",

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

@ -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"

2
python/operator_runtime/ops/__init__.py
View File

@ -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",
]

80
python/operator_runtime/ops/relu.py Normal file
View File

@ -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)

45
tests/bench/relu.py Normal file
View File

@ -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

35
tests/cases/relu.py Normal file
View File

@ -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},
]

115
tests/op_tests/test_relu.py Normal file
View File

@ -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()