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