From 9afe211d658d883771b6ade9144f7a5d1823c8a3 Mon Sep 17 00:00:00 2001 From: yutianyu Date: Wed, 13 May 2026 16:08:21 +0800 Subject: [PATCH] docs: streamline README and guides Co-Authored-By: wawahejun --- README.md | 210 ++++++++++------------- README.zh.md | 169 +++++++++--------- course/README.md | 18 ++ docs/how-to-add-an-operator.md | 302 ++++++++++----------------------- docs/phase1-kernel-writing.md | 11 +- requirements.txt | 6 +- 6 files changed, 296 insertions(+), 420 deletions(-) create mode 100644 course/README.md diff --git a/README.md b/README.md index 6867ff9..efd9aa2 100644 --- a/README.md +++ b/README.md @@ -1,134 +1,112 @@ # Operator Runtime Training Camp -This repository is a training-oriented nano GPU operator runtime. It is -intentionally small, but its workflow mirrors production operator libraries: +A nano GPU operator runtime for learning kernel development. The framework handles descriptor lifecycle, Python FFI, tests, and benchmarks — students focus on writing kernels in `ops//nvidia/kernel.cuh`. -1. Choose one of two parallel operator paths: write a custom backend - implementation under `ops//`, or reuse the elementwise framework under - `ops/elementwise//`. -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. -5. Validate correctness against PyTorch and benchmark steady-state execution. +## Quick Start -The repository is a framework and exercise scaffold. The descriptor lifecycle, -Python API, tests, benchmark entry points, and example directories are wired up; -some phase-1 kernels are intentionally left as TODOs for students to complete. +```bash +# activate environment +conda activate py312 -## Python Layout +# build (auto-fetches CUTLASS on first run) +bash scripts/build_nvidia.sh build -```text -python/ - operator_runtime/ - backend.py - ops/ - _internal/ - operator_runtime_testing/ +# run all tests +bash scripts/build_nvidia.sh test + +# clean build +bash scripts/build_nvidia.sh clean ``` -- `operator_runtime.ops` contains public operator bindings. -- `operator_runtime._internal` contains private FFI/runtime plumbing. -- `operator_runtime_testing` contains test-only helpers such as assertions and benchmark utilities. +## Common Commands + +### Build Script Modes + +```bash +bash scripts/build_nvidia.sh env # show current build environment variables +bash scripts/build_nvidia.sh configure # cmake configure only +bash scripts/build_nvidia.sh build # configure + build +bash scripts/build_nvidia.sh test # run pytest + run_ops + examples +bash scripts/build_nvidia.sh all # build + test +bash scripts/build_nvidia.sh clean # remove build directory +``` + +### Single Operator Testing + +```bash +# correctness test for one operator +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_copy.py -v + +# run_ops supports --mode: test, bench, all +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode test +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode bench +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode all + +# all operators +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op all --backend nvidia --mode all +``` + +### Force Rebuild + +```bash +# clear cmake cache and rebuild +CAMP_FORCE_RECONFIGURE=1 bash scripts/build_nvidia.sh build + +# target specific GPU architecture +CMAKE_CUDA_ARCHITECTURES=89 bash scripts/build_nvidia.sh build +``` + +## Project Structure + +``` +ops//nvidia/ + kernel.cuh <-- implement your kernel here + _cuda.cu <-- descriptor lifecycle (provided) +include/operator_runtime/ + ops/.h <-- public C API (provided) +python/operator_runtime/ + ops/.py <-- Python bindings (provided) +tests/ + op_tests/test_.py <-- correctness tests + cases/.py <-- test cases + bench/.py <-- benchmarks +``` ## Operators -| Operator | NVIDIA C++ | TileLang | MetaX | -| --- | --- | --- | --- | -| `copy` | phase-1 TODO kernel scaffold | phase-1 TODO kernel scaffold when TileLang is installed | scaffolded backend | -| `vector_add` | phase-1 TODO kernel scaffold | phase-1 TODO kernel scaffold when TileLang is installed | scaffolded backend | -| `reduce_sum` | phase-1 TODO row-wise fp32 kernel scaffold | phase-1 TODO kernel scaffold when TileLang is installed | scaffolded row-wise fp32 backend | -| `relu` | runnable elementwise fp16/fp32 example with `negative_slope` | not implemented | runnable elementwise fp16/fp32 example | -| `softmax` | phase-1 TODO row-wise fp32 kernel scaffold | phase-1 TODO kernel scaffold when TileLang is installed | scaffolded row-wise fp32 backend | +| Operator | Difficulty | Key Concept | +| --- | --- | --- | +| `copy` | easy | grid-stride loop, vectorized memory access | +| `vector_add` | easy | elementwise computation | +| `reduce_sum` | medium | shared memory, warp reduction | +| `softmax` | hard | multi-pass reduction + normalization | +| `relu` | reference | elementwise framework (already implemented) | -For scaffolded phase-1 operators, a successful build only proves that the -framework wiring is present. Full correctness tests are expected to fail until -the TODO kernels are implemented. +## Workflow -## Setup +1. Read the kernel scaffold in `ops//nvidia/kernel.cuh` +2. Implement the TODO kernel +3. Build: `bash scripts/build_nvidia.sh build` +4. Test: `PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_.py -v` +5. Benchmark: `PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op --backend nvidia --mode bench` -```bash -pip install -r requirements.txt -``` +## Adding a New Operator -## Build +See [docs/how-to-add-an-operator.md](docs/how-to-add-an-operator.md) for the full guide. Minimal steps: -```bash -./scripts/build_nvidia.sh build -``` +1. Create `ops//nvidia/` with kernel and cuda source +2. Add Python bindings in `python/operator_runtime/ops/.py` +3. Add test cases, correctness tests, and benchmarks under `tests/` +4. Re-run `bash scripts/build_nvidia.sh configure` (CMake re-discovers new `.cu` files) +5. Build and verify -```bash -./scripts/build_metax.sh build -``` +## Environment Variables -`build_nvidia.sh` prefers `third_party/cutlass` when present, and otherwise -falls back to `CAMP_CUTLASS_ROOT` or `CUTLASS_ROOT`. -The default build directories are `build-nvidia` and `build-metax`. Use -`BUILD_DIR=...` to override them, and `CAMP_FORCE_RECONFIGURE=1` when you want a -script to clear an existing CMake cache before reconfiguring. - -```bash -BUILD_DIR=build ./scripts/build_nvidia.sh build -CAMP_FORCE_RECONFIGURE=1 ./scripts/build_nvidia.sh configure -cmake --preset nvidia-release -cmake --build --preset nvidia-release -``` - -The default build produces one backend variant of `libcamp_ops.so` for the local -environment. NVIDIA and MetaX are intentionally built as separate variants; the -C/Python API keeps the same backend-selection interface, and a library returns -`not supported` when asked to use a backend that was not compiled into it. - -Custom NVIDIA operators can be written as ordinary CUDA C++ kernels or as -CuTe/CUTLASS-style C++ implementations. CuTe/CUTLASS is not a separate backend; -it is an optional NVIDIA-backend implementation dependency. When -`CAMP_ENABLE_NVIDIA=ON`, the build auto-detects `cute/tensor.hpp` from -`CAMP_CUTLASS_ROOT`, `CAMP_CUTE_INCLUDE_DIRS`, `third_party/cutlass`, or the -`CUTLASS_ROOT` / `CUTLASS_HOME` / `CUTLASS_PATH` environment variables. If CuTe -is found, `CAMP_ENABLE_CUTE=1` is defined for `camp_ops`; if not found, ordinary -CUDA C++ operators still build. Set `CAMP_ENABLE_CUTE=ON` only when you want a -missing CuTe installation to be a configuration error: - -```bash -git clone https://github.com/NVIDIA/cutlass.git third_party/cutlass -cmake .. -DCAMP_ENABLE_NVIDIA=ON -``` - -```bash -cmake .. \ - -DCAMP_ENABLE_NVIDIA=ON \ - -DCAMP_ENABLE_CUTE=ON \ - -DCAMP_CUTLASS_ROOT=/path/to/cutlass -``` - -## Validate - -```bash -python tests/run_ops.py --op copy --backend nvidia --mode all -CAMP_BUILD_DIR=build-nvidia pytest tests/ -v --backend nvidia -pytest tests/ -v --backend tilelang -python tests/run_ops.py --op all --backend nvidia --mode bench -./scripts/build_metax.sh test -``` - -The TileLang backend requires the `tilelang` Python package. The 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`. - -## Operator Paths - -The training camp supports two operator paths: - -- Custom operators live under `ops///`. This path is useful for teaching kernel writing, fixed contiguous fast paths, special layout or workspace needs, CuTe/CUTLASS-style implementations, and non-elementwise structures. `copy` and `vector_add` demonstrate this path. It is a first-class path, not a temporary step before moving everything into the elementwise framework. -- Reusable elementwise operators live under `ops/elementwise///`. This path is useful for ordinary unary, binary, multi-input, broadcast, or stride-aware elementwise operators that share the same execution model. Later exercises can ask students to reimplement `copy` / `add`-style operators with this path as a comparison exercise. - -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 - -| Training concept | Production equivalent | -| --- | --- | -| 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 | -| `PerformanceResult` | profiler report row with latency, bytes, flops, bandwidth | -| eager TileLang kernel | puzzle-stage kernel using `T.empty(...)` return values | -| lazy TileLang `out_idx` template | TileOPs-style kernel factory and output-position contract | +| Variable | Default | Description | +| --- | --- | --- | +| `CMAKE_CUDA_ARCHITECTURES` | `native` | Target GPU arch (e.g. `89` for L40) | +| `CAMP_FORCE_RECONFIGURE` | `0` | Set `1` to clear CMake cache before configure | +| `CAMP_ENABLE_CUTE` | `AUTO` | CuTe/CUTLASS support: `AUTO`, `ON`, `OFF` | +| `BUILD_DIR` | `build-nvidia` | Build output directory | +| `CAMP_BUILD_DIR` | — | Tell Python where to find `libcamp_ops.so` | +| `CAMP_CUTLASS_ROOT` | — | Override CUTLASS path (skips auto-fetch) | diff --git a/README.zh.md b/README.zh.md index d4e7297..8771409 100644 --- a/README.zh.md +++ b/README.zh.md @@ -1,113 +1,112 @@ -# Intro-ops训练营 +# Intro-ops 训练营 -这个仓库是一个面向训练的 nano 级 GPU 算子运行时。它体积很小，但工作流尽量贴近真实的算子库开发流程： +面向 kernel 开发的 nano 级 GPU 算子运行时。框架已处理好 descriptor 生命周期、Python FFI、测试和 benchmark——学生只需在 `ops//nvidia/kernel.cuh` 中实现 kernel。 -1. 在两条并行路径中选择一条：在 `ops//` 下创建自定义后端实现，或在 `ops/elementwise//` 下复用 elementwise 框架。 -2. 构建系统通过目录约定自动发现源码。 -3. 实现后端专属的 descriptor 生命周期（create、workspace、execute、destroy）。 -4. 提供 Python API，支持 out-of-place、out-variant 和 prepared 执行。 -5. 通过 PyTorch 做正确性验证，并做稳定态性能测试。 - -这个仓库当前是框架和练习骨架：descriptor 生命周期、Python API、测试、benchmark 入口和示例目录已经接好；部分第一阶段 kernel 逻辑刻意保留 TODO，留给学生补全。 - -## Python 目录结构 - -```text -python/ - operator_runtime/ - backend.py - ops/ - _internal/ - operator_runtime_testing/ -``` - -- `operator_runtime.ops` 放公开算子绑定。 -- `operator_runtime._internal` 放私有 FFI 和运行时细节。 -- `operator_runtime_testing` 放断言、benchmark 等仅测试使用的工具。 - -## 算子 - -| 算子 | NVIDIA C++ | TileLang | MetaX | -| --- | --- | --- | --- | -| `copy` | 第一阶段 TODO kernel 骨架 | 安装 TileLang 后可用的第一阶段 TODO kernel 骨架 | 后端骨架已接好 | -| `vector_add` | 第一阶段 TODO kernel 骨架 | 安装 TileLang 后可用的第一阶段 TODO kernel 骨架 | 后端骨架已接好 | -| `reduce_sum` | 第一阶段 TODO row-wise fp32 kernel 骨架 | 安装 TileLang 后可用的第一阶段 TODO kernel 骨架 | row-wise fp32 后端骨架已接好 | -| `relu` | 可运行的 elementwise fp16/fp32 示例，支持 `negative_slope` | 未实现 | 可运行的 elementwise fp16/fp32 示例 | -| `softmax` | 第一阶段 TODO row-wise fp32 kernel 骨架 | 安装 TileLang 后可用的第一阶段 TODO kernel 骨架 | row-wise fp32 后端骨架已接好 | - -对第一阶段骨架算子来说，编译通过只说明框架链路已接好；在补全 TODO kernel 之前，完整 correctness 测试预期不会全部通过。 - -## 安装 +## 快速开始 ```bash -pip install -r requirements.txt +# 激活环境 +conda activate py312 + +# 构建（首次运行自动拉取 CUTLASS） +bash scripts/build_nvidia.sh build + +# 运行全部测试 +bash scripts/build_nvidia.sh test + +# 清理构建 +bash scripts/build_nvidia.sh clean ``` -## 构建 +## 常用命令 + +### 构建脚本模式 ```bash -./scripts/build_nvidia.sh build +bash scripts/build_nvidia.sh env # 显示当前构建环境变量 +bash scripts/build_nvidia.sh configure # 仅 cmake 配置 +bash scripts/build_nvidia.sh build # 配置 + 编译 +bash scripts/build_nvidia.sh test # 运行 pytest + run_ops + examples +bash scripts/build_nvidia.sh all # 编译 + 测试 +bash scripts/build_nvidia.sh clean # 删除构建目录 ``` +### 单算子测试 + ```bash -./scripts/build_metax.sh build +# 单算子正确性测试 +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_copy.py -v + +# run_ops 支持 --mode: test, bench, all +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode test +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode bench +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode all + +# 全部算子 +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op all --backend nvidia --mode all ``` -`build_nvidia.sh` 会优先使用仓库内的 `third_party/cutlass`；如果不存在，则回退到 `CAMP_CUTLASS_ROOT` 或 `CUTLASS_ROOT`。 -默认构建目录分别是 `build-nvidia` 和 `build-metax`。需要覆盖时传 `BUILD_DIR=...`；需要在重配前清理已有 CMake cache 时，传 `CAMP_FORCE_RECONFIGURE=1`。 +### 强制重新构建 ```bash -BUILD_DIR=build ./scripts/build_nvidia.sh build -CAMP_FORCE_RECONFIGURE=1 ./scripts/build_nvidia.sh configure -cmake --preset nvidia-release -cmake --build --preset nvidia-release +# 清除 cmake 缓存并重新构建 +CAMP_FORCE_RECONFIGURE=1 bash scripts/build_nvidia.sh build + +# 指定 GPU 架构 +CMAKE_CUDA_ARCHITECTURES=89 bash scripts/build_nvidia.sh build ``` -默认构建会为本机环境生成一个 backend 版本的 `libcamp_ops.so`。NVIDIA 和 MetaX 刻意作为独立变体构建；C/Python API 保持统一的 backend 选择接口，如果请求当前库未编译进来的 backend，会返回 `not supported`。 +## 项目结构 -自定义 NVIDIA 算子既可以写普通 CUDA C++ kernel，也可以写 CuTe/CUTLASS 风格 C++ 实现。CuTe/CUTLASS 不是单独 backend，而是 NVIDIA backend 内部的可选实现依赖。`CAMP_ENABLE_NVIDIA=ON` 时，构建会从 `CAMP_CUTLASS_ROOT`、`CAMP_CUTE_INCLUDE_DIRS`、`third_party/cutlass` 或 `CUTLASS_ROOT` / `CUTLASS_HOME` / `CUTLASS_PATH` 环境变量自动探测 `cute/tensor.hpp`。如果找到，会给 `camp_ops` 定义 `CAMP_ENABLE_CUTE=1`；如果没找到，普通 CUDA C++ 算子仍然正常构建。只有在希望找不到 CuTe 时直接配置失败，才显式传 `CAMP_ENABLE_CUTE=ON`： - -```bash -git clone https://github.com/NVIDIA/cutlass.git third_party/cutlass -cmake .. -DCAMP_ENABLE_NVIDIA=ON +``` +ops//nvidia/ + kernel.cuh <-- 在这里实现你的 kernel + _cuda.cu <-- descriptor 生命周期（已提供） +include/operator_runtime/ + ops/.h <-- 公开 C API（已提供） +python/operator_runtime/ + ops/.py <-- Python 绑定（已提供） +tests/ + op_tests/test_.py <-- 正确性测试 + cases/.py <-- 测试用例 + bench/.py <-- 性能基准 ``` -```bash -cmake .. \ - -DCAMP_ENABLE_NVIDIA=ON \ - -DCAMP_ENABLE_CUTE=ON \ - -DCAMP_CUTLASS_ROOT=/path/to/cutlass -``` +## 算子列表 -## 验证 +| 算子 | 难度 | 核心概念 | +| --- | --- | --- | +| `copy` | 入门 | grid-stride loop、向量化内存访问 | +| `vector_add` | 入门 | 逐元素计算 | +| `reduce_sum` | 中等 | shared memory、warp 归约 | +| `softmax` | 进阶 | 多趟归约 + 归一化 | +| `relu` | 参考实现 | elementwise 框架（已实现） | -```bash -python tests/run_ops.py --op copy --backend nvidia --mode all -CAMP_BUILD_DIR=build-nvidia pytest tests/ -v --backend nvidia -pytest tests/ -v --backend tilelang -python tests/run_ops.py --op all --backend nvidia --mode bench -./scripts/build_metax.sh test -``` +## 开发流程 -TileLang 后端需要安装 `tilelang` Python 包。MetaX 后端使用独立构建产物，构建时应关闭 NVIDIA 变体，并将后端源码放在 `ops/*/metax/*.maca` 或 `ops/elementwise/*/metax/*.maca`。 +1. 阅读 `ops//nvidia/kernel.cuh` 中的 kernel 骨架 +2. 实现 TODO kernel +3. 构建：`bash scripts/build_nvidia.sh build` +4. 测试：`PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_.py -v` +5. Benchmark：`PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op --backend nvidia --mode bench` -## 算子开发路径 +## 新增算子 -训练营支持两条路径： +完整指南见 [docs/how-to-add-an-operator.md](docs/how-to-add-an-operator.md)。最小步骤： -- 自定义算子路径：放在 `ops///`，适合教学 kernel 编写、固定 contiguous fast path、特殊 layout、特殊 workspace、CuTe/CUTLASS 风格实现或非逐元素结构。`copy` 和 `vector_add` 是这条路径的演示。这是一条正式路径，不是所有算子迁移到 elementwise 前的临时阶段。 -- elementwise 复用路径：放在 `ops/elementwise///`，适合共享 shape、stride、broadcast 执行模型的普通 unary、binary、多输入逐元素算子。后续作业可以让学生用这条路径重新实现 `copy` / `add` 类算子作为对照训练。 +1. 创建 `ops//nvidia/`，包含 kernel 和 cuda 源文件 +2. 在 `python/operator_runtime/ops/.py` 添加 Python 绑定 +3. 在 `tests/` 下添加测试用例、正确性测试和 benchmark +4. 重新运行 `bash scripts/build_nvidia.sh configure`（CMake 重新发现新 `.cu` 文件） +5. 构建并验证 -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。 +## 环境变量 -## 生产映射 - -| 训练概念 | 生产等价物 | -| --- | --- | -| `ops//nvidia/*.cu` 和 `ops/elementwise//nvidia/*.cu` 目录约定 | 构建系统自动发现 / 算子注册 | -| `include/operator_runtime/ops/.h` 头文件 | 经过评审的算子 API 契约 | -| descriptor 生命周期 | create、workspace、execute、destroy | -| `tests/cases/.py` | 正确性、布局和 API 契约覆盖 | -| `PerformanceResult` | 包含延迟、字节数、FLOPs、带宽的 profiler 报表行 | -| eager TileLang kernel | 使用 `T.empty(...)` 返回值的 puzzle-stage kernel | -| lazy TileLang `out_idx` template | TileOPs 风格的 kernel factory 和输出位置契约 | +| 变量 | 默认值 | 说明 | +| --- | --- | --- | +| `CMAKE_CUDA_ARCHITECTURES` | `native` | 目标 GPU 架构（如 L40 用 `89`） | +| `CAMP_FORCE_RECONFIGURE` | `0` | 设为 `1` 在配置前清除 CMake 缓存 | +| `CAMP_ENABLE_CUTE` | `AUTO` | CuTe/CUTLASS 支持：`AUTO`、`ON`、`OFF` | +| `BUILD_DIR` | `build-nvidia` | 构建输出目录 | +| `CAMP_BUILD_DIR` | — | 告诉 Python 在哪里找 `libcamp_ops.so` | +| `CAMP_CUTLASS_ROOT` | — | 覆盖 CUTLASS 路径（跳过自动拉取） | diff --git a/course/README.md b/course/README.md new file mode 100644 index 0000000..c4a705e --- /dev/null +++ b/course/README.md @@ -0,0 +1,18 @@ +# 课程资料 + +本目录用于存放训练营课件和学习记录。 + +## 目录用途 + +- 课件、讲义、slides +- 个人学习笔记 +- 实验记录和性能分析报告 + +## 建议组织方式 + +``` +course/ + slides/ # 课件 + notes/ # 学习笔记 + reports/ # 实验报告 +``` diff --git a/docs/how-to-add-an-operator.md b/docs/how-to-add-an-operator.md index badbc74..97eb728 100644 --- a/docs/how-to-add-an-operator.md +++ b/docs/how-to-add-an-operator.md @@ -2,256 +2,130 @@ ## 目标 -在当前项目里，新增一个算子的最小闭环包括四部分。训练营同时保留两条并行路径：自定义算子路径用于教学 kernel 编写、特殊 layout、特殊 workspace 和非普通逐元素结构，elementwise 复用路径用于普通 unary / binary / broadcast / stride 逐元素算子。 +新增一个算子的最小闭环包括： -1. 在 `ops/<算子名>/` 下补自定义后端实现，或在 `ops/elementwise/<算子名>/` 下补复用 elementwise 框架的实现。 +1. 在 `ops/<算子名>/nvidia/` 下实现 kernel 和 descriptor 生命周期。 2. 在 `python/operator_runtime/ops/` 下补 Python API。 -3. 在 `tests/` 下补正确性测试和 benchmark 入口。 +3. 在 `tests/` 下补正确性测试和 benchmark。 4. 重新构建并验证。 -当前仓库是训练营框架和练习骨架。`copy`、`vector_add`、`reduce_sum`、`softmax` 的第一阶段 kernel 文件仍保留 TODO，用来让学生补齐计算逻辑；`relu` 是已经接通的 elementwise 示例。 - +训练营支持两条并行路径： +- **自定义算子**：`ops//nvidia/`，适合 kernel 教学、特殊 layout、reduce/softmax 等非逐元素结构。 +- **elementwise 复用**：`ops/elementwise//nvidia/`，适合 unary/binary/broadcast 逐元素算子。 ## Step 1：明确算子接口 -开始前先确认： +开始前确认： -1. 这个算子有几个输入、几个输出。 -2. 输出 shape 是否和输入一致。 -3. 是否要求输入输出 dtype 一致。 -4. 是否只支持 contiguous tensor。 -5. 是否需要额外参数，例如 `dim`、`scalar`。 -6. 是否需要 workspace。 - -这一步的目的，是确定后面 C API 和 Python 绑定该按哪种模式实现。不要把“逐元素”简单等同于“必须用 elementwise 框架”：训练营会先用 `copy`、`vector_add` 展示自定义算子写法，也会保留 elementwise 复用路径作为普通逐元素算子的工程化选择。两条路径是并行选项，不是“先自定义、最后都迁移到 elementwise”的过渡关系。 +1. 几个输入、几个输出 +2. 输出 shape 是否和输入一致 +3. 是否要求 dtype 一致 +4. 是否只支持 contiguous tensor +5. 是否需要额外参数（`dim`、`scalar` 等） +6. 是否需要 workspace ## Step 2：创建算子目录 -先选择算子目录： +最小文件集： -- 自定义算子：放在 `ops/<算子名>/`。适合教学底层 kernel、固定 contiguous fast path、特殊 layout、特殊 workspace、reduce/softmax 等非普通 elementwise 算子，也允许写 CuTe/CUTLASS 风格 C++ 实现。 -- elementwise 复用算子：放在 `ops/elementwise/<算子名>/`。适合共享 shape、stride、broadcast 执行模型的 unary / binary / 多输入逐元素算子。 - -`copy` 和 `vector_add` 是自定义算子教学示例。`relu` 是 elementwise 框架示例，目录为 `ops/elementwise/relu/nvidia/`，公共 NVIDIA launcher 位于 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`。 - -先在选定目录下建立对应后端目录。 - -当前建议至少补齐： - -- `ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` -- `python/operator_runtime/ops/<算子名>.py` -- `tests/cases/<算子名>.py` -- `tests/op_tests/test_<算子名>.py` -- `tests/bench/<算子名>.py` - -如果后续要支持 TileLang 或 MetaX，再分别补 `tilelang/` 或 `metax/`。 -但对当前流程来说，NVIDIA 版本是最小必需项。 - -## Step 2.5：选择实现风格 - -自定义 NVIDIA 算子可以采用两种 C++ 实现风格： - -1. 原生 CUDA C++：直接在 `.cu` / `.cuh` 里写 `__global__` kernel。这是第一阶段默认训练方式，`copy`、`vector_add`、`reduce_sum`、`softmax` 都按这个方式组织。 -2. CuTe/CUTLASS 风格 C++：仍然放在 `ops/<算子名>/nvidia/`，仍然导出同一组 descriptor 生命周期符号，但内部可以使用 CuTe tensor、layout、copy atom、TiledMMA 等抽象。 - -CuTe/CUTLASS 不是单独 backend，而是 NVIDIA backend 内部的可选实现依赖。`CAMP_ENABLE_NVIDIA=ON` 时，构建会默认自动探测 `cute/tensor.hpp`，来源包括 `CAMP_CUTLASS_ROOT`、`CAMP_CUTE_INCLUDE_DIRS`、`third_party/cutlass` 或 `CUTLASS_ROOT` / `CUTLASS_HOME` / `CUTLASS_PATH` 环境变量。找到后目标会定义 `CAMP_ENABLE_CUTE=1` 并把 include 目录加到 `camp_ops`；找不到时继续构建普通 CUDA C++ 算子。 - -推荐约定是把 CUTLASS 源码 clone 到仓库根目录下的 `third_party/cutlass`，这样 `CAMP_ENABLE_NVIDIA=ON` 时就能自动探测，不需要每次手工传路径。 - -如果你希望“找不到 CuTe 就直接失败”，可以显式打开强制模式： - -```bash -cmake .. \ - -DCAMP_ENABLE_NVIDIA=ON \ - -DCAMP_ENABLE_CUTE=ON \ - -DCAMP_CUTLASS_ROOT=/path/to/cutlass +``` +ops/<算子名>/nvidia/ + kernel.cuh # kernel 实现 + <算子名>_cuda.cu # descriptor 生命周期 + launch +include/operator_runtime/ops/<算子名>.h # C API 头文件 +python/operator_runtime/ops/<算子名>.py # Python 绑定 +tests/cases/<算子名>.py # 测试数据 +tests/op_tests/test_<算子名>.py # 正确性测试 +tests/bench/<算子名>.py # benchmark ``` -也可以用 `"-DCAMP_CUTE_INCLUDE_DIRS=/path/a;/path/b"` 显式指定 include 目录。这个选项只提供框架扩展点，不要求第一阶段学生必须使用 CuTe。 +### include 目录结构与创建原则 + +``` +include/operator_runtime/ + api.h # 公共类型定义（status、dtype、tensor_view 等） + descriptor.h # descriptor 基类 + tensor_view.h # tensor view 结构体 + operator_runtime.h # 汇总头文件 + ops/ + <算子名>.h # 每个算子一个头文件，声明四个生命周期 C 函数 + detail/ + cuda_helpers.h # CUDA 工具函数（blocks_for、stream 转换等） + elementwise.h # elementwise descriptor helper + tensor_checks.h # tensor 校验工具 + operation.h # operation 基类 +``` + +创建原则： + +- `ops/<算子名>.h` 是算子的公开 C API 契约，只声明四个生命周期函数（create/workspace/execute/destroy），不暴露实现细节。 +- 所有函数使用 `extern "C"` + `OPRT_EXPORT`，保证 Python FFI 可以按符号名 dlsym。 +- 参数类型只使用 `api.h` 中定义的公共类型（`oprt_status_t`、`oprt_tensor_view_t`、`oprt_operator_descriptor_t`、`oprt_stream_t`）。 +- `detail/` 下放内部实现工具，不对外暴露，算子实现可以 include 但用户代码不应依赖。 +- 新增算子只需在 `ops/` 下加一个头文件，不需要修改其他头文件。 ## Step 3：实现 NVIDIA 后端 -`ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` 这一层负责 C++/CUDA 实现。 +一个算子需要四个生命周期接口： -自定义算子通常需要这几部分： +1. `oprt_create__descriptor` — 检查输入合法性，保存运行信息 +2. `oprt_get__workspace_size` — 返回临时内存大小 +3. `oprt_execute_` — 按 dtype dispatch 并 launch kernel +4. `oprt_destroy__descriptor` — 释放 descriptor -1. kernel 文件 -2. C API 头文件 -3. C API 实现文件 +参考 `copy` 的实现：kernel 放 `kernel.cuh`，生命周期放 `_cuda.cu`。 -这里最关键的是保持现有命名约定一致，因为 Python 侧会按固定符号名去找函数。 - -如果是 elementwise 复用算子，不需要每个算子重复写 shape/stride kernel。推荐复用 `include/operator_runtime/detail/elementwise.h` 里的 descriptor helper 和 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`： - -1. descriptor 继承 `oprt::ElementwiseDescriptorBase`。 -2. create 阶段调用 `oprt::init_elementwise_descriptor(desc, out, {inputs...})`。 -3. workspace 复用 `oprt::get_elementwise_workspace_size`。 -4. execute 阶段只做 dtype dispatch，并调用 `oprt::elementwise::nvidia::launch(...)`。 -5. 算子自身只提供 device functor，例如 `relu` 的 `value > 0 ? value : value * negative_slope`。 - -elementwise 路径推荐覆盖： - -1. 输出 shape 由输入 broadcast 得出，或由调用方提供的 out 固定。 -2. 每个输出元素可以独立计算。 -3. 输入之间只需要普通 shape/stride/broadcast 索引。 -4. 不需要跨元素同步、规约、排序、扫描、复杂临时 workspace 或特殊数据布局。 - -不满足这些条件时，优先走自定义算子路径。 - -一个 NVIDIA 算子需要完整提供四个生命周期接口： - -1. create -2. workspace -3. execute -4. destroy - -整体流程可以理解为： - -1. `create`：检查输入是否合法，并保存运行所需信息。 -2. `workspace`：返回执行需要的临时内存大小。 -3. `execute`：按 dtype 和参数启动 kernel。 -4. `destroy`：释放 descriptor。 +如果是 elementwise 算子，复用 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh` 和 `include/operator_runtime/detail/elementwise.h`，只需提供 device functor。参考 `relu`。 ## Step 4：补 Python 绑定 -Python 入口放在 `python/operator_runtime/ops/<算子名>.py`。 +Python 入口放在 `python/operator_runtime/ops/<算子名>.py`，通常暴露三个接口： -当前项目对外通常暴露三类接口： +- `` — out-of-place，自动分配输出 +- `_` — out-variant，调用方提供输出 tensor +- `prepare_` — 创建 descriptor，支持多次执行复用 -1. `prepare_<算子名>` -2. `<算子名>_` -3. `<算子名>` +新增后在 `ops/__init__.py` 和 `operator_runtime/__init__.py` 中导出。 -职责分工一般是： +## Step 5：补测试 -1. `prepare_<算子名>`：完成参数检查、绑定底层函数、创建 descriptor 和 workspace。 -2. `<算子名>_`：接收调用方提供的输出 tensor，执行一次。 -3. `<算子名>`：自动分配输出 tensor，再调用 `<算子名>_`。 +`tests/cases/<算子名>.py` 组织测试数据，分三类： -如果你的算子是自定义算子，并且签名和现有 unary、binary、reduce-like 模式一致，就沿用现有 helper。 -如果你的算子是 elementwise 复用算子，优先在 Python 里定义 `ElementwiseOpSpec`，再调用 `prepare_elementwise_op`。这个 spec 描述算子名、输入数、标量参数和 broadcast 语义，避免为每个 elementwise 算子重复写 FFI 绑定和基础校验。 +- `correctness_cases()` — 正确性用例（shape、dtype、tolerance） +- `api_error_cases()` — 异常输入用例 +- `benchmark_cases()` — 性能测试规模 -示例： +`tests/op_tests/test_<算子名>.py` 负责： -```python -_MY_OP_SPEC = ElementwiseOpSpec( - name="my_op", - input_count=1, - scalar_argtypes=(ctypes.c_float,), -) -``` - -## Step 5：导出到公共 API - -新增 Python 文件后，还要把公开接口补到： - -1. `python/operator_runtime/ops/__init__.py` -2. `python/operator_runtime/__init__.py` - -至少需要把这三个名字暴露出去： - -1. `<算子名>` -2. `<算子名>_` -3. `prepare_<算子名>` - -这样测试和用户代码才能直接导入。 - -## Step 6：补测试用例数据 - -`tests/cases/<算子名>.py` 只负责组织测试数据。 - -当前仓库一般分成三类： - -1. `correctness_cases()` -2. `api_error_cases()` -3. `benchmark_cases()` - -建议在这里把不同 shape、dtype、容差、错误场景、性能场景分开整理，避免把 case 直接写死在测试函数里。 - -## Step 7：补正确性测试 - -`tests/op_tests/test_<算子名>.py` 主要负责三件事： - -1. 正确性对比 -2. API contract 检查 -3. prepared 执行复用检查 - -正确性测试通常是拿 PyTorch 结果做对照。 -API contract 测试主要覆盖 shape 不匹配、dtype 不匹配、非 contiguous 等场景。 -如果底层走 descriptor lifecycle，建议补一个 prepared 多次执行的测试，确认 descriptor 可以复用。 - -## Step 8：补 benchmark +- 正确性对比（vs PyTorch） +- API contract 检查（shape/dtype 不匹配、非 contiguous） +- prepared 执行复用检查 `tests/bench/<算子名>.py` 负责性能入口。 -当前流程里，一般会： +## Step 6：构建和验证 -1. 从 `benchmark_cases()` 取输入规模。 -2. 构造 CUDA tensor。 -3. 测量自定义算子耗时。 -4. 测量对应 PyTorch 实现耗时。 -5. 汇总成性能结果。 +```bash +# 新增 .cu 后必须重新 configure +bash scripts/build_nvidia.sh configure +bash scripts/build_nvidia.sh build -这一步的目标不是做复杂分析，而是保证新算子已经接入仓库现有 benchmark 流程。 +# 单算子验证 +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend nvidia -## Step 9：重新配置和编译 +# 正确性 + benchmark +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <算子名> --backend nvidia --mode all -因为 `ops/CMakeLists.txt` 是通过 glob 自动发现 `ops/*/nvidia/*.cu` 和 `ops/elementwise/*/nvidia/*.cu`，所以新增 `.cu` 之后要重新配置。 +# 确认没有破坏已有算子 +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op all --backend nvidia --mode test +``` -顺序是： +## 现有模板参考 -1. 重新执行 `./scripts/build_nvidia.sh configure`，或手工进入 `build-nvidia/` 后重新执行 `cmake ..`。 -2. 再执行 `./scripts/build_nvidia.sh build`，或手工执行编译。 - -如果你只改了 Python，不新增 `.cu`，通常不需要重新配置。 -但只要新加了 NVIDIA 源文件，就必须重新跑一次 CMake。 - -当前构建策略是统一 API、单 backend 产物：NVIDIA 和 MetaX 不要求编进同一个 `libcamp_ops.so`。每个环境按本机硬件构建一个 backend 版本；Python / C ABI 保持统一 backend 参数；当前库只接受已编译进来的 backend，未启用 backend 返回 `not supported`。 - -## Step 10：验证 - -建议按下面顺序验证： - -1. 先跑单算子的 pytest。 -2. 再跑更大范围的算子测试。 -3. 最后跑 benchmark 或统一入口。 - -优先确认： - -1. 输出结果和 PyTorch 一致。 -2. 错误输入能稳定报错。 -3. descriptor/workspace 流程正常。 -4. 没有破坏已有算子。 - -## 推荐开发顺序 - -如果你要新增的是一个普通算子，建议按这个顺序推进： - -1. 先参考最接近的现有算子选模板。 -2. 先打通 NVIDIA 后端最小可运行版本。 -3. 再补 Python API。 -4. 再补 cases、pytest、benchmark。 -5. 最后补其他后端或做性能优化。 - -## 现有模板怎么选 - -- `copy`：适合最简单的自定义单输入单输出流程，当前第一阶段 kernel 逻辑保留 TODO。 -- `vector_add`：适合标准 contiguous 双输入自定义算子，当前第一阶段 kernel 逻辑保留 TODO。 -- `relu`：适合复用 elementwise 框架的单输入逐元素算子，也展示了额外标量参数 `negative_slope` 的 C/Python 绑定方式。 -- `reduce_sum`：适合带 reduce 维度的算子，当前第一阶段 kernel 逻辑保留 TODO。 -- `softmax`：适合带更明确 shape 约束和归一化逻辑的算子，当前第一阶段 kernel 逻辑保留 TODO。 - -如果新算子本质上是普通 elementwise，优先参考 `relu`、`ElementwiseOpSpec` 和 `ops/common/elementwise/nvidia/elementwise_nvidia.cuh`；如果课程目标是练习手写 contiguous kernel，再参考 `copy` 或 `vector_add` 的自定义路径。 - -## 最终检查清单 - -提交前至少确认以下内容都已完成： - -1. `ops/<算子名>/nvidia/` 或 `ops/elementwise/<算子名>/nvidia/` 已补齐实现。 -2. `python/operator_runtime/ops/<算子名>.py` 已补齐。 -3. 两个 `__init__.py` 已导出新接口。 -4. `tests/cases/<算子名>.py` 已补数据。 -5. `tests/op_tests/test_<算子名>.py` 已补测试。 -6. `tests/bench/<算子名>.py` 已补 benchmark。 -7. 新增 `.cu` 后已经重新执行过 `./scripts/build_nvidia.sh configure` 或等价的 `cmake ..`。 -8. 至少完成一次单算子验证。 +| 模板 | 适用场景 | +| --- | --- | +| `copy` | 最简单的单输入单输出自定义算子 | +| `vector_add` | 双输入 contiguous 自定义算子 | +| `relu` | elementwise 框架复用，含标量参数 | +| `reduce_sum` | 带 reduce 维度的算子 | +| `softmax` | 多步归约 + 归一化 | diff --git a/docs/phase1-kernel-writing.md b/docs/phase1-kernel-writing.md index ff56bc7..6135cd3 100644 --- a/docs/phase1-kernel-writing.md +++ b/docs/phase1-kernel-writing.md @@ -93,11 +93,14 @@ ```bash # NVIDIA kernel（需要先重新编译） -./scripts/build_nvidia.sh build -CAMP_BUILD_DIR=/workspace/build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend nvidia +bash scripts/build_nvidia.sh build +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend nvidia -# TileLang kernel -CAMP_BUILD_DIR=/workspace/build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend tilelang +# TileLang kernel（不需要编译） +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend tilelang + +# 同时跑正确性 + benchmark +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <算子名> --backend nvidia --mode all ``` 四个算子两种后端全部通过，阶段一完成。 diff --git a/requirements.txt b/requirements.txt index 1df26ab..5451817 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,8 +3,12 @@ # Runtime and benchmarking (required) torch>=2.0 -# TileLang backend (required) +# TileLang backend (optional) tilelang +# Build tools +cmake>=3.22 +ninja + # Testing framework pytest>=7.0