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

Compare commits

merge into: FrRay:main
Branches Tags
FrRay:main
FrRay:master
ccf-ai-infra:main
ccf-ai-infra:master
...
pull from: FrRay:master
Branches Tags
FrRay:main
FrRay:master
ccf-ai-infra:main
ccf-ai-infra:master

No commits in common. "main" and "master" have entirely different histories.
main ... master

135 changed files with 1 additions and 9561 deletions
Show Stats Expand all files Collapse all files

9
.cnb.yml
View File

@ -1,9 +0,0 @@
$:
vscode:
- docker:
image: docker.cnb.cool/yutianyu.yi/image/ncu
runner:
tags: cnb:arch:amd64:gpu
services:
- vscode
- docker

36
.github/ISSUE_TEMPLATE/bug_report.md vendored
View File

@ -1,36 +0,0 @@
---
name: Bug Report
about: 报告代码或文档的错误
title: "[Bug] "
labels: bug
assignees: ""
---
## 环境信息
- OS: [如 Ubuntu 22.04]
- CUDA 版本: [如 12.1]
- GPU 型号: [如 NVIDIA L40]
- intro-ops 分支/commit: [如 main@abc1234]
- Python 版本: [如 3.12]
## 描述
请清晰描述你遇到的问题。
## 复现步骤
1. 执行 `...`
2. 看到错误 `...`
## 期望行为
请描述你期望发生什么。
## 实际行为
请描述实际发生了什么（附上完整的错误信息或截图）。
## 附加信息
如相关：日志、截图、benchmark 数据。

23
.github/ISSUE_TEMPLATE/feature_request.md vendored
View File

@ -1,23 +0,0 @@
---
name: Feature Request
about: 提出功能增强或新算子建议
title: "[Feature] "
labels: enhancement
assignees: ""
---
## 使用场景
请描述你的需求背景：在什么情况下需要这个功能？
## 提议方案
请描述你期望的功能或实现思路。
## 替代方案
是否考虑过其他替代方案？如有请描述。
## 附加信息
如相关：参考实现、论文链接、其他项目的类似功能。

22
.github/ISSUE_TEMPLATE/question.md vendored
View File

@ -1,22 +0,0 @@
---
name: Question
about: 使用问题或技术讨论
title: "[Question] "
labels: question
assignees: ""
---
## 问题描述
请清晰描述你的问题。
## 已尝试的方法
- [ ] 已查阅 README
- [ ] 已查阅 FAQ
- [ ] 已搜索 Issues
- [ ] 已查阅 troubleshooting 指南
## 环境信息（如相关）
- OS / CUDA 版本 / GPU 型号

26
.github/PULL_REQUEST_TEMPLATE.md vendored
View File

@ -1,26 +0,0 @@
## 描述
请简要描述此 PR 做了什么。
## 类型
- [ ] Bug 修复
- [ ] 新功能
- [ ] 文档更新
- [ ] 测试补充
- [ ] 代码重构
- [ ] 性能优化
## 验证
- [ ] 相关测试通过：`pytest tests/op_tests/ -v --backend nvidia`
- [ ] 代码风格检查通过（clang-format / black + isort）
- [ ] 新增代码有适当的测试覆盖
## 关联 Issue
Closes #
## 截图 / Benchmark（如适用）
（附上相关的测试结果截图或性能数据对比）

14
.gitignore vendored
View File

@ -1,14 +0,0 @@
build/
build-*/
*.egg-info/
__pycache__/
.pytest_cache/
.mypy_cache/
.ruff_cache/
*.pyc
*.so
*.o
*.a
*.log
/generated/
/third_party/

114
CMakeLists.txt
View File

@ -1,114 +0,0 @@
cmake_minimum_required(VERSION 3.22)
project(operator_runtime_training LANGUAGES CXX)
option(CAMP_ENABLE_NVIDIA "Build NVIDIA CUDA backend" ON)
option(CAMP_ENABLE_TILELANG "Enable TileLang Python backend metadata" ON)
option(CAMP_ENABLE_METAX "Build MetaX backend" OFF)
set(CAMP_ENABLE_CUTE "AUTO" CACHE STRING "Enable optional CuTe/CUTLASS headers for NVIDIA custom operators: AUTO, ON, or OFF")
set_property(CACHE CAMP_ENABLE_CUTE PROPERTY STRINGS AUTO ON OFF)
set(CAMP_CUTLASS_ROOT "" CACHE PATH "Optional CUTLASS checkout root for CuTe-based custom NVIDIA operators")
set(CAMP_CUTE_INCLUDE_DIRS "" CACHE STRING "Optional semicolon-separated CuTe/CUTLASS include directories")
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CUDA_STANDARD 17)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
list(PREPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
if(CAMP_ENABLE_NVIDIA AND CAMP_ENABLE_METAX)
message(FATAL_ERROR "CAMP_ENABLE_NVIDIA and CAMP_ENABLE_METAX cannot both be ON; build separate variants")
endif()
string(TOUPPER "${CAMP_ENABLE_CUTE}" CAMP_ENABLE_CUTE_MODE)
if(NOT CAMP_ENABLE_CUTE_MODE MATCHES "^(AUTO|ON|OFF)$")
message(FATAL_ERROR "CAMP_ENABLE_CUTE must be AUTO, ON, or OFF")
endif()
function(camp_append_cute_root root)
if(NOT root)
return()
endif()
list(APPEND CAMP_CUTE_CANDIDATE_INCLUDE_DIRS
"${root}"
"${root}/include"
"${root}/tools/util/include"
)
set(CAMP_CUTE_CANDIDATE_INCLUDE_DIRS "${CAMP_CUTE_CANDIDATE_INCLUDE_DIRS}" PARENT_SCOPE)
endfunction()
set(CAMP_CUTE_TARGET_INCLUDE_DIRS "")
set(CAMP_USE_CUTE OFF)
if(NOT CAMP_ENABLE_CUTE_MODE STREQUAL "OFF")
if(NOT CAMP_ENABLE_NVIDIA)
if(CAMP_ENABLE_CUTE_MODE STREQUAL "ON")
message(FATAL_ERROR "CAMP_ENABLE_CUTE=ON requires CAMP_ENABLE_NVIDIA=ON")
endif()
else()
set(CAMP_CUTE_CANDIDATE_INCLUDE_DIRS "")
camp_append_cute_root("${CAMP_CUTLASS_ROOT}")
camp_append_cute_root("${CMAKE_CURRENT_SOURCE_DIR}/third_party/cutlass")
foreach(_camp_cute_env_var CUTLASS_ROOT CUTLASS_HOME CUTLASS_PATH)
if(DEFINED ENV{${_camp_cute_env_var}} AND NOT "$ENV{${_camp_cute_env_var}}" STREQUAL "")
camp_append_cute_root("$ENV{${_camp_cute_env_var}}")
endif()
endforeach()
if(CAMP_CUTE_INCLUDE_DIRS)
list(APPEND CAMP_CUTE_CANDIDATE_INCLUDE_DIRS ${CAMP_CUTE_INCLUDE_DIRS})
endif()
if(CAMP_CUTE_CANDIDATE_INCLUDE_DIRS)
list(REMOVE_DUPLICATES CAMP_CUTE_CANDIDATE_INCLUDE_DIRS)
endif()
foreach(_camp_cute_include_dir IN LISTS CAMP_CUTE_CANDIDATE_INCLUDE_DIRS)
if(EXISTS "${_camp_cute_include_dir}/cute/tensor.hpp")
list(APPEND CAMP_CUTE_TARGET_INCLUDE_DIRS "${_camp_cute_include_dir}")
endif()
endforeach()
if(CAMP_CUTE_TARGET_INCLUDE_DIRS)
list(REMOVE_DUPLICATES CAMP_CUTE_TARGET_INCLUDE_DIRS)
set(CAMP_USE_CUTE ON)
message(STATUS "CuTe/CUTLASS support enabled: ${CAMP_CUTE_TARGET_INCLUDE_DIRS}")
elseif(CAMP_ENABLE_CUTE_MODE STREQUAL "ON")
message(FATAL_ERROR "CAMP_ENABLE_CUTE=ON requires cute/tensor.hpp via CAMP_CUTLASS_ROOT, CAMP_CUTE_INCLUDE_DIRS, or CUTLASS_ROOT/CUTLASS_HOME/CUTLASS_PATH")
else()
message(STATUS "CuTe/CUTLASS support disabled: cute/tensor.hpp was not found")
endif()
endif()
endif()
if(CAMP_ENABLE_METAX)
find_package(MetaX REQUIRED)
set(ENV{CUCC_PATH} "${MetaX_CUBRIDGE_ROOT}")
set(ENV{CUDA_PATH} "${MetaX_CUBRIDGE_ROOT}")
list(APPEND CMAKE_MODULE_PATH "${MetaX_CMAKE_MODULE_DIR}")
set(CMAKE_MACA_SOURCE_FILE_EXTENSIONS maca;cu;cc;cpp;c)
endif()
if(CAMP_ENABLE_NVIDIA OR CAMP_ENABLE_METAX)
if(CAMP_ENABLE_METAX)
enable_language(MACA)
else()
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
set(CMAKE_CUDA_ARCHITECTURES native)
endif()
enable_language(CUDA)
include(cmake/cuda_helpers.cmake)
endif()
endif()
if(CAMP_ENABLE_METAX)
set(CAMP_METAX_COMPILE_OPTIONS
-D__MACA_NO_HALF_OPERATORS__
-D__FAST_BLOCK_RED__
)
endif()
add_subdirectory(ops)

76
CMakePresets.json
View File

@ -1,76 +0,0 @@
{
"version": 6,
"cmakeMinimumRequired": {
"major": 3,
"minor": 22,
"patch": 0
},
"configurePresets": [
{
"name": "base",
"hidden": true,
"generator": "Ninja",
"binaryDir": "${sourceDir}/build/${presetName}",
"cacheVariables": {
"CMAKE_CXX_STANDARD": "17",
"CMAKE_CUDA_STANDARD": "17"
}
},
{
"name": "nvidia-debug",
"inherits": "base",
"cacheVariables": {
"CMAKE_BUILD_TYPE": "Debug",
"CAMP_ENABLE_NVIDIA": "ON",
"CAMP_ENABLE_METAX": "OFF",
"CAMP_ENABLE_CUTE": "AUTO",
"CMAKE_CUDA_ARCHITECTURES": "native"
}
},
{
"name": "nvidia-release",
"inherits": "base",
"cacheVariables": {
"CMAKE_BUILD_TYPE": "Release",
"CAMP_ENABLE_NVIDIA": "ON",
"CAMP_ENABLE_METAX": "OFF",
"CAMP_ENABLE_CUTE": "AUTO",
"CMAKE_CUDA_ARCHITECTURES": "native"
}
},
{
"name": "nvidia-cute",
"inherits": "nvidia-release",
"cacheVariables": {
"CAMP_ENABLE_CUTE": "ON"
}
},
{
"name": "metax-release",
"inherits": "base",
"cacheVariables": {
"CMAKE_BUILD_TYPE": "Release",
"CAMP_ENABLE_NVIDIA": "OFF",
"CAMP_ENABLE_METAX": "ON"
}
}
],
"buildPresets": [
{
"name": "nvidia-debug",
"configurePreset": "nvidia-debug"
},
{
"name": "nvidia-release",
"configurePreset": "nvidia-release"
},
{
"name": "nvidia-cute",
"configurePreset": "nvidia-cute"
},
{
"name": "metax-release",
"configurePreset": "metax-release"
}
]
}

175
README.md
View File

@ -1,175 +1,2 @@
# Operator Runtime Training Camp
# Intro-ops
A nano GPU operator runtime for learning kernel development. The framework uses a backend-agnostic public C API with separate NVIDIA and MetaX build variants, while Python, tests, and benchmarks sit on top of the same runtime contract.
## Quick Start
```bash
# activate environment
conda activate py312
# build NVIDIA (auto-fetches CUTLASS on first run)
bash scripts/build_nvidia.sh build
# run all NVIDIA tests
bash scripts/build_nvidia.sh test
# build MetaX
bash scripts/build_metax.sh build
# run MetaX tests
bash scripts/build_metax.sh test
# clean build
bash scripts/build_nvidia.sh clean
bash scripts/build_metax.sh clean
```
## 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
bash scripts/build_metax.sh env # show current MACA build environment
bash scripts/build_metax.sh configure # cmake configure only
bash scripts/build_metax.sh build # configure + build
bash scripts/build_metax.sh test # run pytest + run_ops + examples
bash scripts/build_metax.sh all # build + test
bash scripts/build_metax.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
# MetaX
PYTHONPATH=python:. CAMP_BUILD_DIR=build-metax pytest tests/op_tests/test_copy.py -v --backend metax
PYTHONPATH=python:. CAMP_BUILD_DIR=build-metax python tests/run_ops.py --op all --backend metax --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
```
include/operator_runtime/
operator_runtime.h <-- umbrella public runtime header
ops/<op>.h <-- backend-agnostic public C API
detail/*.h <-- shared internal helper layer
ops/<op>/nvidia/
kernel.cuh <-- implement your kernel here
<op>_cuda.cu <-- unified public symbol implementation
ops/<op>/metax/
<op>_metax.maca <-- MetaX kernel + launch implementation
ops/elementwise/<op>/
nvidia/*.cu <-- elementwise NVIDIA implementations
metax/*.maca <-- elementwise MetaX implementations
python/operator_runtime/
ops/<op>.py <-- Python wrapper over shared C API / TileLang
tests/
op_tests/test_<op>.py <-- correctness tests
cases/<op>.py <-- test cases
bench/<op>.py <-- benchmarks
```
## Architecture
- Public operator headers live under `include/operator_runtime/ops/*.h` and expose backend-agnostic symbols such as `oprt_create_copy_descriptor`.
- Backend-private code lives under `ops/.../<backend>/`.
- Python bindings use the same public symbol names for both compiled backends and select the active backend via `oprt_set_backend(...)` plus separate build outputs.
## Operators
| 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) |
## Workflow
1. Read the kernel scaffold in `ops/<op>/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_<op>.py -v`
5. Benchmark: `PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <op> --backend nvidia --mode bench`
## Adding a New Operator
See [docs/how-to-add-an-operator.md](docs/how-to-add-an-operator.md) for the full guide. Minimal steps:
1. Create `ops/<op>/nvidia/` with kernel and cuda source
2. Add Python bindings in `python/operator_runtime/ops/<op>.py`
3. Add test cases, correctness tests, and benchmarks under `tests/`
4. Re-run `bash scripts/build_nvidia.sh configure` when adding new NVIDIA `.cu` sources
5. Build and verify
## MetaX Backend
The training runtime also supports a MetaX build variant through `bash scripts/build_metax.sh ...`.
```bash
# build MetaX variant
bash scripts/build_metax.sh build
# run MetaX correctness + benchmark flow
bash scripts/build_metax.sh test
```
## TileLang on MetaX
Stock pip `tilelang` is not sufficient for this machine. Use the source-built `/root/tilelang-metax` tree for TileLang-on-MetaX validation when needed.
```bash
# build /root/tilelang-metax separately with USE_MACA=ON
CAMP_USE_TILELANG_METAX=1 \
CAMP_TILELANG_SOURCE_ROOT=/root/tilelang-metax \
bash scripts/build_metax.sh test
```
This mode exports:
- `PYTHONPATH=/root/tilelang-metax`
- `LD_LIBRARY_PATH=/root/tilelang-metax/build/lib:/opt/maca/lib:...`
and validates `copy/vector_add/reduce_sum/softmax` with `backend=tilelang` on MetaX.
## Environment Variables
| 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) |
| `CAMP_USE_TILELANG_METAX` | `0` | Set `1` to run TileLang-on-MetaX validation in `build_metax.sh` |
| `CAMP_TILELANG_SOURCE_ROOT` | `/root/tilelang-metax` | Source-built TileLang tree used for MetaX TileLang validation |

175
README.zh.md
View File

@ -1,175 +0,0 @@
# Intro-ops 训练营
面向 kernel 开发的 nano 级 GPU 算子运行时。框架现在采用后端无关的公共 C API，并提供独立的 NVIDIA 与 MetaX 构建变体；Python、测试和 benchmark 都建立在同一套运行时契约之上。
## 快速开始
```bash
# 激活环境
conda activate py312
# 构建 NVIDIA（首次运行自动拉取 CUTLASS）
bash scripts/build_nvidia.sh build
# 运行全部 NVIDIA 测试
bash scripts/build_nvidia.sh test
# 构建 MetaX
bash scripts/build_metax.sh build
# 运行 MetaX 测试
bash scripts/build_metax.sh test
# 清理构建
bash scripts/build_nvidia.sh clean
bash scripts/build_metax.sh clean
```
## 常用命令
### 构建脚本模式
```bash
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 env # 显示当前 MACA 构建环境
bash scripts/build_metax.sh configure # 仅 cmake 配置
bash scripts/build_metax.sh build # 配置 + 编译
bash scripts/build_metax.sh test # 运行 pytest + run_ops + examples
bash scripts/build_metax.sh all # 编译 + 测试
bash scripts/build_metax.sh clean # 删除构建目录
```
### 单算子测试
```bash
# 单算子正确性测试
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
# MetaX
PYTHONPATH=python:. CAMP_BUILD_DIR=build-metax pytest tests/op_tests/test_copy.py -v --backend metax
PYTHONPATH=python:. CAMP_BUILD_DIR=build-metax python tests/run_ops.py --op all --backend metax --mode all
```
### 强制重新构建
```bash
# 清除 cmake 缓存并重新构建
CAMP_FORCE_RECONFIGURE=1 bash scripts/build_nvidia.sh build
# 指定 GPU 架构
CMAKE_CUDA_ARCHITECTURES=89 bash scripts/build_nvidia.sh build
```
## 项目结构
```
include/operator_runtime/
operator_runtime.h <-- 公共运行时总头
ops/<op>.h <-- 后端无关的公开 C API
detail/*.h <-- 共享内部 helper 分层
ops/<op>/nvidia/
kernel.cuh <-- 在这里实现你的 kernel
<op>_cuda.cu <-- 统一公共符号实现
ops/<op>/metax/
<op>_metax.maca <-- MetaX kernel + launch 实现
ops/elementwise/<op>/
nvidia/*.cu <-- elementwise NVIDIA 实现
metax/*.maca <-- elementwise MetaX 实现
python/operator_runtime/
ops/<op>.py <-- 基于共享 C API / TileLang 的 Python 包装
tests/
op_tests/test_<op>.py <-- 正确性测试
cases/<op>.py <-- 测试用例
bench/<op>.py <-- 性能基准
```
## 架构说明
- 公开算子头文件位于 `include/operator_runtime/ops/*.h`，暴露统一的后端无关符号，例如 `oprt_create_copy_descriptor`。
- 后端私有实现位于 `ops/.../<backend>/`。
- Python 绑定对 NVIDIA 和 MetaX 复用同一套公开符号，通过 `oprt_set_backend(...)` 和独立构建产物选择后端。
## 算子列表
| 算子 | 难度 | 核心概念 |
| --- | --- | --- |
| `copy` | 入门 | grid-stride loop、向量化内存访问 |
| `vector_add` | 入门 | 逐元素计算 |
| `reduce_sum` | 中等 | shared memory、warp 归约 |
| `softmax` | 进阶 | 多趟归约 + 归一化 |
| `relu` | 参考实现 | elementwise 框架（已实现） |
## 开发流程
1. 阅读 `ops/<op>/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_<op>.py -v`
5. Benchmark：`PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <op> --backend nvidia --mode bench`
## 新增算子
完整指南见 [docs/how-to-add-an-operator.md](docs/how-to-add-an-operator.md)。最小步骤：
1. 创建 `ops/<op>/nvidia/`，包含 kernel 和 cuda 源文件
2. 在 `python/operator_runtime/ops/<op>.py` 添加 Python 绑定
3. 在 `tests/` 下添加测试用例、正确性测试和 benchmark
4. 新增 NVIDIA `.cu` 文件后，重新运行 `bash scripts/build_nvidia.sh configure`
5. 构建并验证
## MetaX 后端
训练营同时支持 MetaX 构建变体，可通过 `bash scripts/build_metax.sh ...` 使用。
```bash
# 构建 MetaX 变体
bash scripts/build_metax.sh build
# 运行 MetaX 正确性与 benchmark 流程
bash scripts/build_metax.sh test
```
## MetaX 上的 TileLang
本机不能直接使用 pip 官方版 `tilelang`。只有在需要验证 MetaX 上的 TileLang 时，才使用源码构建的 `/root/tilelang-metax`。
```bash
# 先单独把 /root/tilelang-metax 以 USE_MACA=ON 构建完成
CAMP_USE_TILELANG_METAX=1 \
CAMP_TILELANG_SOURCE_ROOT=/root/tilelang-metax \
bash scripts/build_metax.sh test
```
该模式会导出：
- `PYTHONPATH=/root/tilelang-metax`
- `LD_LIBRARY_PATH=/root/tilelang-metax/build/lib:/opt/maca/lib:...`
并在 MetaX 上验证 `copy/vector_add/reduce_sum/softmax` 的 `backend=tilelang`。
## 环境变量
| 变量 | 默认值 | 说明 |
| --- | --- | --- |
| `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 路径（跳过自动拉取） |
| `CAMP_USE_TILELANG_METAX` | `0` | 设为 `1` 时在 `build_metax.sh` 中执行 MetaX 上的 TileLang 验证 |
| `CAMP_TILELANG_SOURCE_ROOT` | `/root/tilelang-metax` | 用于 MetaX TileLang 验证的源码版 TileLang 根目录 |

16
cmake/CMakeMACAInformation.cmake
View File

@ -1,16 +0,0 @@
set(CMAKE_MACA_COMPILE_OBJECT
"<CMAKE_MACA_COMPILER> <DEFINES> <INCLUDES> <FLAGS> -o <OBJECT> -c <SOURCE>"
)
set(CMAKE_INCLUDE_FLAG_MACA "-I")
set(CMAKE_MACA_COMPILE_OPTIONS_PIC "-fPIC")
set(CMAKE_MACA_CREATE_SHARED_LIBRARY
"<CMAKE_MACA_COMPILER> <LANGUAGE_COMPILE_FLAGS> <LINK_FLAGS> -shared -o <TARGET> <OBJECTS> <LINK_LIBRARIES> --maca-path=$ENV{MACA_PATH}"
)
set(CMAKE_MACA_LINK_EXECUTABLE
"<CMAKE_MACA_COMPILER> <FLAGS> <LINK_FLAGS> -o <TARGET> <OBJECTS> <LINK_LIBRARIES> --maca-path=$ENV{MACA_PATH}"
)
set(CMAKE_MACA_INFORMATION_LOADED 1)

82
cmake/FindMetaX.cmake
View File

@ -1,82 +0,0 @@
set(_MetaX_SEARCH_ROOTS "")
if(DEFINED ENV{MACA_PATH})
list(APPEND _MetaX_SEARCH_ROOTS "$ENV{MACA_PATH}")
endif()
list(APPEND _MetaX_SEARCH_ROOTS /opt/maca)
find_path(
MetaX_ROOT
NAMES include/mcr/mc_runtime.h
PATHS ${_MetaX_SEARCH_ROOTS}
NO_DEFAULT_PATH
)
find_program(
MetaX_MXCC_EXECUTABLE
NAMES mxcc
PATHS
${MetaX_ROOT}/mxgpu_llvm/bin
/opt/maca/mxgpu_llvm/bin
)
find_program(
MetaX_CUCC_EXECUTABLE
NAMES cucc
PATHS
${MetaX_ROOT}/tools/cu-bridge/bin
/opt/maca/tools/cu-bridge/bin
)
find_path(
MetaX_INCLUDE_DIR
NAMES mcr/mc_runtime.h
PATHS
${MetaX_ROOT}/include
/opt/maca/include
)
find_library(
MetaX_MCRUNTIME_LIBRARY
NAMES mcruntime libmcruntime.so
PATHS
${MetaX_ROOT}/lib
/opt/maca/lib
)
find_library(
MetaX_MXCRT_LIBRARY
NAMES mxc-runtime64 libmxc-runtime64.so
PATHS
${MetaX_ROOT}/lib
/opt/maca/lib
)
include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(
MetaX
REQUIRED_VARS
MetaX_ROOT
MetaX_MXCC_EXECUTABLE
MetaX_INCLUDE_DIR
MetaX_MCRUNTIME_LIBRARY
)
if(MetaX_FOUND)
set(MetaX_CUBRIDGE_ROOT "${MetaX_ROOT}/tools/cu-bridge")
set(MetaX_CMAKE_MODULE_DIR "${MetaX_CUBRIDGE_ROOT}/cmake_module/maca")
set(MetaX_INCLUDE_DIRS "${MetaX_INCLUDE_DIR}")
set(MetaX_LIBRARIES "${MetaX_MCRUNTIME_LIBRARY}")
if(MetaX_MXCRT_LIBRARY)
list(APPEND MetaX_LIBRARIES "${MetaX_MXCRT_LIBRARY}")
endif()
endif()
mark_as_advanced(
MetaX_ROOT
MetaX_MXCC_EXECUTABLE
MetaX_CUCC_EXECUTABLE
MetaX_INCLUDE_DIR
MetaX_MCRUNTIME_LIBRARY
MetaX_MXCRT_LIBRARY
)

10
cmake/cuda_helpers.cmake
View File

@ -1,10 +0,0 @@
if(NOT DEFINED CMAKE_CUDA_STANDARD)
set(CMAKE_CUDA_STANDARD 17)
endif()
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
set(CMAKE_CUDA_ARCHITECTURES native)
endif()
set(CMAKE_CUDA_STANDARD_REQUIRED ON)

18
course/README.md
View File

@ -1,18 +0,0 @@
# 课程资料
本目录用于存放训练营课件和学习记录。
## 目录用途
- 课件、讲义、slides
- 个人学习笔记
- 实验记录和性能分析报告
## 建议组织方式
```
course/
slides/ # 课件
notes/ # 学习笔记
reports/ # 实验报告
```

56
course/checklist.md
View File

@ -1,56 +0,0 @@
# 学习进度清单
## Phase 1: Kernel 编写
### copy
- [ ] 理解 grid-stride loop 原理：为什么可以处理任意大小的 tensor
- [ ] 理解 global memory 合并访问（coalesced access）
- [ ] NVIDIA `kernel.cuh` TODO 完成
- [ ] TileLang `kernel.py` TODO 完成
- [ ] 两种后端测试全部通过
- [ ] benchmark 跑通，带宽数据合理
### vector_add
- [ ] 理解逐元素并行的线程网格布局
- [ ] 理解 tile-level 并行（`T.Parallel` vs `T.Serial` 的区别）
- [ ] NVIDIA `kernel.cuh` TODO 完成
- [ ] TileLang `kernel.py` TODO 完成
- [ ] 两种后端测试全部通过
- [ ] benchmark 跑通
### reduce_sum（进阶）
- [ ] 理解 shared memory 树形归约原理
- [ ] 理解 `__syncthreads()` 的使用时机和条件分支限制
- [ ] 理解 `T.Serial` 在归约场景中的作用
- [ ] NVIDIA `kernel.cuh` TODO 完成
- [ ] TileLang `kernel.py` TODO 完成
- [ ] 两种后端测试全部通过
- [ ] 尝试优化 bank conflict
- [ ] benchmark 跑通
### softmax（挑战）
- [ ] 理解数值稳定性问题：为什么需要减 max
- [ ] 理解 online softmax 算法（一遍扫描 vs 三趟扫描）
- [ ] 理解 log-sum-exp 的滚动更新逻辑
- [ ] 理解 TileLang 中为什么用 `exp2` / `log2` 替代 `exp` / `log`
- [ ] NVIDIA `kernel.cuh` TODO 完成
- [ ] TileLang `kernel.py` TODO 完成
- [ ] 两种后端测试全部通过
- [ ] 性能与 PyTorch 参考实现对比
- [ ] 尝试 warp-level 优化
## Phase 2: 进阶优化
- [ ] 阅读 `docs/tilelang-vs-cuda.md`，理解两种后端的差异
- [ ] 尝试修改 `.cu` 文件调整 launch 参数（block size、grid size）
- [ ] 尝试调整 shared memory 大小看性能变化
- [ ] 学习 CUTLASS/CuTe 集成基础
- [ ] 尝试提交 benchmark 数据
## Phase 3: 社区贡献
- [ ] 阅读 `docs/how-to-submit-first-pr.md`
- [ ] 阅读 `docs/how-to-add-an-operator.md`
- [ ] 找一个 good-first-issue 练手
- [ ] 提交第一个 PR 并被合并
- [ ] 尝试贡献文档翻译或 FAQ 补充

47
course/scripts/00_script_framework.md
View File

@ -1,47 +0,0 @@
# 录播/直播课程脚本框架
## 目标
为 intro-ops 四个算子的教学视频提供统一脚本模板。每个视频 **15-20 分钟**，按固定结构组织。
## 统一结构
```
[5min] 概念讲解：这个算子在深度学习中的用途 + 算法推导
[10min] 代码实操：IDE 中边写边讲，展示关键决策点
[3min] 测试验证：跑测试、看 benchmark、分析性能
[2min] 常见错误演示：故意写错然后排查
```
## 讲师准备清单
### 课前
- [ ] 确认录制环境：IDE 字体够大、终端清晰、无弹窗干扰
- [ ] 打开对应的 kernel 骨架文件（`.cuh` 和 `.py`）
- [ ] 准备好测试命令（粘贴即用）
- [ ] 准备一个"故意写错"的版本用于错误演示环节
### 课中
- [ ] 语速适中（190-210 字/分钟）
- [ ] 写代码时边写边说——不要沉默
- [ ] 每个关键决策点停下来解释"为什么"
- [ ] 错误演示后立即展示正确写法做对比
### 课后
- [ ] 提供本视频的代码 diff（骨架 → 完整实现）
- [ ] 在视频描述中附相关文档链接
- [ ] 设置一个"课后挑战"（如"把带宽从 X 提到 Y"）
## 发布节奏建议
| 平台 | 频率 | 内容类型 |
|------|------|---------|
| B 站 | 每周 1 期 | 完整教学视频 |
| 公众号 | 每周 1 篇 | 视频的文字摘要 + 代码片段 |
| 知乎 | 每周 1 篇 | 从视频中提炼一个深度问题单独讨论 |
## 字幕与多语言
- 视频字幕优先中文
- 代码和术语保留英文
- 关键概念在字幕中标注英文对照（如 "grid-stride loop（网格跨步循环）"）

129
course/scripts/01_copy_script.md
View File

@ -1,129 +0,0 @@
# Copy 算子 — 教学视频脚本
> 时长：15-20 分钟 | 难度：入门
---
## [5min] 概念讲解
### 开场（30s）
"今天我们从 intro-ops 训练营最简单的算子开始——copy。GPU 上做 copy 和 CPU 上不同——你写不出 `memcpy` 那种一行代码，因为 GPU 有成百上千个线程同时在工作。怎么让每个线程知道它该搬哪些数据？这就是 grid-stride loop 要解决的问题。"
### 算子在深度学习中的用途（1min）
"Copy 看起来 trivial，但在深度学习框架中无处不在。Tensor 的 `clone()`、`contiguous()`、数据加载中的 host-to-device 拷贝——底层都是类似的 copy kernel。理解 copy 就是理解 GPU 内存模型的第一步。"
### 算法推导（3.5min）
**关键画面：** 展示 grid-stride-loop Mermaid 图（`docs/diagrams/grid-stride-loop.md`）
讲解要点：
1. GPU 的线程组织：grid → block → thread
2. 每个线程如何计算自己的全局索引：`idx = blockIdx.x * blockDim.x + threadIdx.x`
3. stride = grid 总线程数 = `gridDim.x * blockDim.x`
4. 为什么需要循环？因为线程数可能少于元素数
5. 循环条件 `i < N` 保证了越界安全
**过渡语：** "好，理论就这么多。接下来我们在 IDE 里把它写出来。"
---
## [10min] 代码实操
### CUDA Kernel（5min）
**打开文件：** `ops/copy/nvidia/kernel.cuh`
**边写边讲：**
```
模板参数 T：支持 float、half 等多种数据类型
__global__：告诉 nvcc 这是 GPU 函数
命名空间 oprt::copy::nvidia：遵循项目约定
```
**关键决策点：**
- "为什么用 `int64_t` 而不是 `int`？——tensor 可能有超过 2^31 个元素"
- "为什么 `idx` 和 `stride` 的计算不放在循环条件里？——编译器优化和可读性"
- "为什么循环变量叫 `i` 不是 `idx`？——`idx` 是起始位置，`i` 是当前处理位置"
**写完后展示完整代码。**
### TileLang Kernel（5min）
**打开文件：** `ops/copy/tilelang/kernel.py`
**边写边讲：**
```
@tilelang.jit：JIT 编译装饰器——和 nvcc 不同，这里编译发生在运行时
T.Parallel(N // BLOCK_N)：告诉编译器这 N/BLOCK_N 个 tile 可以并行
T.copy：把全局内存的一整块搬进 fragment
```
**对比时刻：** "看，TileLang 不需要你写 grid、block、thread——编译器帮你做。这 5 行 Python 生成的 CUDA code 比你手写的可能还高效。"
---
## [3min] 测试验证
### 跑测试
```bash
# NVIDIA
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
pytest tests/op_tests/test_copy.py -v --backend nvidia
```
**展示：** 终端中测试全部绿色的画面。
### 跑 benchmark
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
python tests/run_ops.py --op copy --backend nvidia --mode all
```
**解读 benchmark 输出：**
- "看这个 Bandwidth——如果你跑不到峰值的 80%+，说明合并访问有问题"
- "copy 是纯内存搬运，它的性能天花板就是 GPU 内存带宽"
---
## [2min] 常见错误演示
### 错误 1：忘记 grid-stride loop（30s）
```cuda
// 错误——只处理了前 blockDim.x 个元素
int idx = threadIdx.x;
dst[idx] = src[idx]; // 超过 blockDim.x 的部分没被处理！
```
"这就是为什么需要 grid-stride loop——如果不写循环，每个线程只处理一个元素，你搬不完整个 tensor。"
### 错误 2：循环写反条件（30s）
```cuda
// 错误——死循环
for (int i = idx; i < n; i++) { // 忘记 += stride！
```
"每次循环 i 都只加 1——第一个线程就把所有活干完了，其他线程白等，而且还死循环。"
### 错误 3：TileLang 里用 T.Serial（1min）
```python
# 错误——用 T.Serial 跑 copy
for i in T.Serial(N // BLOCK_N):
...
```
"T.Serial 告诉编译器'这些 tile 必须顺序执行'——你的并行性全丢了。copy 的每个 tile 完全独立，应该用 T.Parallel。"
---
## 课后挑战
"把 copy kernel 的带宽从 60% 跑到 90% 以上。提示：检查你的 grid/block 尺寸是否合理。下周 vector_add 见！"

112
course/scripts/02_vector_add_script.md
View File

@ -1,112 +0,0 @@
# Vector Add 算子 — 教学视频脚本
> 时长：15-20 分钟 | 难度：入门
---
## [5min] 概念讲解
### 开场（30s）
"上次我们做了 copy——纯内存搬运。今天加一点计算：vector_add，逐元素加法。结构上和 copy 几乎一样，但引入了两个输入和一个辅助函数。"
### 算子在深度学习中的用途（1min）
"逐元素操作（element-wise ops）在深度学习中极常见：ReLU、dropout、residual connection 的加法、batch norm 的 scale+shift。vector_add 是所有这些操作的原型。"
### 算法推导（3.5min）
**关键画面：** 展示线程网格布局图（`docs/diagrams/thread-grid-layout.md`）
讲解要点：
1. 和 copy 一样的 grid-stride loop 模式
2. 区别：每次循环处理 `a[i] + b[i]` 而非 `src[i]`
3. 因为是逐元素，线程间仍然零通信
4. half 精度需要特殊处理：`__hadd()` 而非 `+`
**过渡语：** "vector_add 的 CUDA kernel 骨架比 copy 多了一个辅助函数——我们来写。"
---
## [10min] 代码实操
### CUDA Kernel（5min）
**打开文件：** `ops/vector_add/nvidia/kernel.cuh`
**Step 1: `add_values<T>` 泛型版（1min）**
```cuda
template <typename T>
__device__ T add_values(T a, T b) {
return a + b; // 泛型版本直接用 +
}
```
"`__device__` 表示这个函数运行在 GPU 上，只能被 kernel 或其他 device 函数调用。为什么抽出一个函数？因为 half 精度需要特化。"
**Step 2: `add_values<half>` 特化版（1min）**
```cuda
template <>
__device__ inline half add_values<half>(half a, half b) {
return __hadd(a, b); // half 专用加法指令
}
```
**关键决策点：** "为什么 half 用 `__hadd` 而不是 `+`？CUDA 中 `half` 是存储类型——算术操作需要先转 float 或调用 intrinsics。`__hadd` 是硬件原生指令，更快。"
**Step 3: kernel 函数（3min）**
和 copy kernel 几乎一样，只有行不同：`out[i] = add_values(a[i], b[i]);`
"注意 `__restrict__` 关键字——告诉编译器 `out`、`a`、`b` 指向不重叠的内存区域，可以做更激进的优化。"
### TileLang Kernel（5min）
**打开文件：** `ops/vector_add/tilelang/kernel.py`
**边写边讲：**
- 外层 `T.Parallel` 分发 tile，内层 `T.Parallel` 分发元素
- "两层都是 `T.Parallel`——因为所有 (a[i], b[i]) 独立计算"
**对比：** "CUDA 里你手动写了三层抽象（thread → block → grid），TileLang 两层 Parallel 就完成了。编译器把外层的 tile Parallel 映射到 block，内层的 element Parallel 映射到 thread。"
---
## [3min] 测试验证
### 跑测试
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
pytest tests/op_tests/test_vector_add.py -v --backend nvidia
```
**重点展示：** "注意测试里包含了 half 精度——如果你 `add_values<half>` 写错了，这里会炸。"
### Benchmark 解读
"vector_add 的理论带宽利用率和 copy 接近——因为它也是 memory-bound 的。如果你的带宽比 copy 低很多，检查是否有非合并访问。"
---
## [2min] 常见错误演示
### 错误 1：half 用 `a + b`（30s）
"CUDA 里 `half + half` 会自动提升为 float 再截断——能编译通过但结果有精度损失。必须用 `__hadd`。"
### 错误 2：不特化 half 就直接在 kernel 里用 `+`（30s）
"`add_values<T>` 里泛型版用 `a + b`，half 特化版用 `__hadd`。如果没写特化版，half 走泛型 `a + b` 也能跑——但性能差。测试能通过，benchmark 会暴露。"
### 错误 3：TileLang 忘记内层 Parallel（1min）
"如果内层用 `T.Serial`——tile 内的 256 个元素就变顺序执行了，线程完全没利用上。结果速度和单线程差不多。"
---
## 课后挑战
"尝试给 vector_add 添加第三个输入：`c = a + b + c`（in-place add）。提示：看 `operator_runtime/ops/vector_add.py` 的 `vector_add_` 函数签名。"

125
course/scripts/03_reduce_sum_script.md
View File

@ -1,125 +0,0 @@
# Reduce Sum 算子 — 教学视频脚本
> 时长：15-20 分钟 | 难度：进阶
---
## [5min] 概念讲解
### 开场（30s）
"前两个算子 copy 和 vector_add，线程之间是零通信的——各自算各自的。今天 reduce_sum 完全不同——线程需要互相通信，通过 shared memory 把部分结果合并为最终结果。"
### 算子在深度学习中的用途（1min）
"归约（reduction）操作在深度学习中非常常见：LayerNorm 里的 mean/std、attention 里的 softmax 分母 sum、loss 函数的最终求和——都是归约。理解归约是 GPU 编程的第一个分水岭。"
### 算法推导（3.5min）
**关键画面：** 展示树形归约流程图（`docs/diagrams/tree-reduction.md`）
讲解要点：
1. 目标：每行 N 个元素求和为 1 个值
2. Step 1: 各线程先独立累加自己负责的列 → 存入 shared memory
3. Step 2: 树形归约——`stride = blockDim.x/2, /4, ..., 1`，每次配对相加
4. 为什么叫"树形"——每步参与线程减半，log₂(blockDim.x) 步完成
5. 关键：每一步之后必须 `__syncthreads()`
**可视化辅助：** 手动画出 8→4→2→1 的归约树
**过渡语：** "听起来简单，但写起来有两个坑——我们直接在 IDE 里看。"
---
## [10min] 代码实操
### CUDA Kernel（6min）
**打开文件：** `ops/reduce_sum/nvidia/kernel.cuh`
**Step 1: 线程各自累加（2min）**
"每个线程用 grid-stride loop 跨步累加自己负责的列——这个和 copy 一样。但这次结果不是写到全局内存，而是写到 shared memory。"
**Step 2: shared memory 写入 + 同步（1min）**
```cuda
smem[threadIdx.x] = sum;
__syncthreads(); // 关键！
```
**关键决策点：** "为什么这里一定要 `__syncthreads()`？因为 tree reduction 下一步要读 `smem[tid+s]`——那是别的线程写的。不同步的话可能读到旧数据。"
**Step 3: 树形归约（2min）**
```cuda
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
smem[threadIdx.x] += smem[threadIdx.x + s];
}
__syncthreads(); // 在 if 外面！
}
```
"注意 `__syncthreads()` 在 if **外面**——这是最常见的 bug。如果放进 if，只有一半线程执行同步，另一半跳过，整个 block 死锁。"
**Step 4: 输出（1min）**
"只有 thread 0 写回全局内存——因为 smem[0] 已经是整行的和。"
### TileLang Kernel（4min）
**打开文件：** `ops/reduce_sum/tilelang/kernel.py`
**边写边讲：**
- "外层 `T.Parallel` 分发到行——不同行之间独立"
- "内层 `T.Serial` 顺序遍历列分块——因为累加器有状态依赖"
- "`T.reduce_sum(tile)` 替代手写树形归约——编译器为你生成最优代码"
- "注意 `T.alloc_fragment` 分配累加器——这是寄存器级别的存储"
**对比：** "CUDA 版 20+ 行，TileLang 版 8 行。`T.reduce_sum` 内部帮你处理了同步、bank conflict、warp divergence——这些在 CUDA 里都是你要手写的。"
---
## [3min] 测试验证
### 跑测试
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
pytest tests/op_tests/test_reduce_sum.py -v --backend nvidia
```
**重点：** "注意测试里的容差——归约涉及很多加法，浮点累加误差比 copy 大。这是正常的，只要在容差内就行。"
### Benchmark 解读
"reduce_sum 的计算量和内存访问量之比（算术强度）比 copy 高——这意味着它从 memory-bound 向 compute-bound 靠近。优化方向也变了：从追求带宽利用率转向减少 bank conflict。"
---
## [2min] 常见错误演示
### 错误 1：`__syncthreads()` 在 if 内（30s）
**故意写错并运行** → 死锁/挂起
"这是最容易踩的坑。症状是程序 hang 住不动。记住：`__syncthreads()` 永远放在条件分支外面。"
### 错误 2：忘记 shared memory 初始化（30s）
"如果把部分和写进 shared memory 之前没有清零——你存的是上一次 kernel launch 的垃圾数据。用 `extern __shared__` 时尤其注意。"
### 错误 3：blockDim.x 不是 2 的幂（30s）
"树形归约假设 block size 是 2 的幂。如果用 300 个线程——stride 从 150 开始，配对就会错位。建议 block size = 128/256/512。"
### 错误 4：写了 shared memory 但忘记声明（30s）
"`extern __shared__ float smem[]` 在 kernel 参数里声明还不够——launch 时 `<<<grid, block, shared_mem_size>>>` 第三个参数必须传。报错 `uses too much shared data` 时检查这里。"
---
## 课后挑战
"消除 reduce_sum 的 bank conflict：给 shared memory 加 padding，对比优化前后的 bandwidth。目标：提升 30%+ 带宽利用率。"

114
course/scripts/04_softmax_script.md
View File

@ -1,114 +0,0 @@
# Softmax 算子 — 教学视频脚本
> 时长：15-20 分钟 | 难度：挑战
---
## [5min] 概念讲解
### 开场（30s）
"softmax 是四个算子中最难的。它结合了 copy 的逐元素、reduce_sum 的归约——还要解决数值溢出问题。今天我们不只写正确的 softmax，还要理解为什么'正确'不是理所当然的。"
### 算子在深度学习中的用途（1min）
"softmax 是 transformer 的核心——每个 attention block 都调它。分类任务的最后一层也是 softmax。一个效率高 10% 的 softmax kernel 可以提升整个推理管线的吞吐。"
### 算法推导（3.5min）
**关键画面：** 展示 softmax 流水线图（`docs/diagrams/softmax-pipeline.md`）
讲解要点：
1. **朴素公式的问题：** `exp(100)` 溢出 → NaN。演示一个小 demo：`torch.tensor([100.0, 200.0, 300.0])` 的朴素 exp 直接炸
2. **稳定的公式：**
- `softmax(x_i) = exp(x_i - max) / Σ exp(x_j - max)`
- 为什么等价？分子分母同除 `exp(max)`
- 为什么稳定？`x_i - max ≤ 0`，所以 `exp() ≤ 1`，永不溢出
3. **三趟扫描流程：**
- Pass 1: 求 `max_val`（reduce max）
- Pass 2: 写 `exp(x - max)` 到 out + 累加 sum（reduce sum）
- Pass 3: `out[i] /= sum`
4. **Online Softmax（TileLang 版）：**
- 用 log-sum-exp (lse) 滚动更新，两趟完成
- `exp2`/`log2` 比 `exp`/`log` 在硬件上更快
**过渡语：** "理论比代码复杂——但代码本身并不可怕。来写。"
---
## [10min] 代码实操
### CUDA Kernel（5min）
**打开文件：** `ops/softmax/nvidia/kernel.cuh`
**Step 1: Pass 1 — 求行最大值（1.5min）**
"和 reduce_sum 的归约一模一样——只是把 `+=` 改成 `max()`。注意初始值：max 初始化为 `in[row * cols]` 而不是 0——因为输入可能是全负数。"
**Step 2: Pass 2 — exp + 累加 sum（2min）**
"第二趟扫描做了两件事：计算 `exp(x - max)` 写到 out，同时累加 sum。为什么要写 out？因为第三趟需要这些中间值——不写的话第三趟还得重新从全局内存读输入再算 exp，浪费带宽。"
**Step 3: Pass 3 — 归一化（1.5min）**
"最后一行 `__syncthreads()` 确保 sum 已经归约完成。除了 thread 0 知道 sum 之外——其他线程不需要知道 sum 就能做除法吗？不对——sum 还在 smem[0] 里，每个线程需要读 smem[0] 来做除法。所以第三趟之前也要 sync。"
### TileLang Kernel（5min）
**打开文件：** `ops/softmax/tilelang/kernel.py`
**关键概念：**
- "`log2_e = 1.44269504`：`exp(x) = 2^(x * log2(e))` 的转换系数"
- "Pass 1 用 `T.Serial` 遍历列分块——和 reduce_sum 一样，lse 有状态依赖"
- "`T.reduce_max` + `T.reduce_sum`：编译器为你选择最优的归约策略"
- "lse 更新公式：`m_new = max(lse, max(tile)); lse = m_new + log2(exp2(lse - m_new) + sum(exp2(tile - m_new)))`"
**不展开完整公式推导——引导学员去看 [docs/diagrams/softmax-pipeline.md](../../docs/diagrams/softmax-pipeline.md) 中的 LSE 滚动更新详解。**
---
## [3min] 测试验证
### 跑测试
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \
pytest tests/op_tests/test_softmax.py -v --backend nvidia
```
**重点展示：** "测试里会验证两件事——正确性（每行和 ≈ 1.0）和数值稳定性（大数据不出 NaN）。如果你偷懒没写减 max，大数值用例会直接炸。"
### Benchmark 解读
"softmax 的算术强度比 reduce_sum 更高——exp/log 是计算密集型操作。优化方向从内存带宽转向计算吞吐。如果你的算力利用率低，可能是 exp 计算没被流水线化。"
---
## [2min] 常见错误演示
### 错误 1：忘记减 max（30s）
"最经典的 bug。输入 `[100, 200, 300]`——朴素 exp 输出全是 NaN。减了 max 后正常输出 `[0, 0, 1]`。"
### 错误 2：Pass 2 和 Pass 3 之间没同步（30s）
"Pass 2 的 sum 归约完成后，thread 0 有正确的 sum——但 thread 1 可能还在写 smem。不 sync 的话 thread 1 在 Pass 3 读到的是旧 sum。"
### 错误 3：max 初始化为 0 而非第一个元素（30s）
"如果输入全是负数——max=0 比真实值大。exp(x - 0) 没问题（仍然 ≤ 1），但 exp(x - real_max) 的精度更好。不影响正确性但影响精度。"
### 错误 4：TileLang 里 Pass 1 用 T.Parallel（30s）
"lse 有跨 tile 的状态依赖——必须 T.Serial。用 T.Parallel 会读到未初始化的 lse 值。"
---
## 课后挑战
"把 CUDA 的三趟扫描改成两趟 online softmax（像 TileLang 版一样）。分析性能提升幅度，写一段注释解释为什么快。"

111
docs/FAQ.md
View File

@ -1,111 +0,0 @@
# 常见问题解答 (FAQ)
> 最后更新：2026-06-05
## 环境配置
### Q: 需要什么硬件才能运行 intro-ops？
**A:** 需要 NVIDIA GPU（支持 CUDA）或沐曦 MetaX GPU。Intel 集显、AMD GPU（无 ROCm 适配）、Apple Silicon 均无法运行。
### Q: 没有 NVIDIA GPU 怎么办？
**A:** 可以考虑租用云端 GPU（AutoDL、恒源云等平台），租一张 T4 或 GTX 1060 即可跑通全部算子。
### Q: CUDA Toolkit 需要什么版本？
**A:** 建议 CUDA 11.8 或以上。CMake 用 `"native"` 架构参数可自动适配你的 GPU。
### Q: 必须用 conda 管理 Python 环境吗？
**A:** 不强制，但推荐。主要依赖是 PyTorch >= 2.0、pytest >= 7.0、cmake >= 3.22。TileLang 后端还需 `pip install tilelang`。
### Q: Windows 上能跑吗？
**A:** 项目设计为 Linux 环境（构建脚本为 bash，产物为 `.so`）。Windows 上建议用 WSL2。
---
## 编译构建
### Q: 构建报 `FindCUDA failed` 怎么办？
**A:** 检查 CUDA Toolkit 是否安装，`nvcc --version` 是否可运行，`CMAKE_CUDA_COMPILER` 是否正确。
### Q: 每次改完 kernel 都要重新构建吗？
**A:** NVIDIA 后端需要重新构建（`bash scripts/build_nvidia.sh build`）。TileLang 后端是 JIT 编译，改完 Python 代码直接跑测试即可。
### Q: `CMake Error: Unknown CUDA architecture` 怎么处理？
**A:** 修改 `CMakePresets.json`，将 `CMAKE_CUDA_ARCHITECTURES` 改为 `"native"` 或你的 GPU 对应架构编号。
### Q: 构建成功但 `import` 时报找不到 `.so` 文件？
**A:** 检查环境变量 `CAMP_BUILD_DIR` 是否指向正确的构建目录（如 `build-nvidia`）。
---
## Kernel 编写
### Q: 四个算子应该按什么顺序学习？
**A:** copy → vector_add → reduce_sum → softmax，难度递增。copy 最简单（纯内存搬运），softmax 需要理解数值稳定性和 online 算法。
### Q: grid-stride loop 为什么能处理任意大小的 tensor？
**A:** 每个线程在循环中处理多个元素（步长为 `gridDim.x * blockDim.x`），而不是只处理一个。这样无论总元素数是多少，只要循环条件 `idx < N` 就能覆盖。
### Q: `__syncthreads()` 什么情况下必须用？
**A:** 当 block 内线程通过 shared memory 交换数据时，写入 shared memory 后必须 `__syncthreads()` 确保所有线程都完成写入，之后才能安全读取。典型场景：reduce_sum 的树形归约每一步之间。
### Q: softmax 为什么要减最大值？
**A:** `exp(88.7) ≈ 1.6e38`，接近 FP32 上限 `3.4e38`。更大的输入值会导致 `exp()` 溢出为 inf。减最大值等价于分子分母同除 `exp(max)`，数学结果不变但数值稳定。
### Q: TileLang 和 CUDA kernel 需要功能完全一致吗？
**A:** 是的，两者应通过相同的测试用例。但实现方式不同——CUDA 控制线程级别，TileLang 控制 tile 级别。
---
## 测试与调试
### Q: 测试报 `tensor not close` 但肉眼看不出来？
**A:** 检查 `rtol`/`atol` 设置。FP16 建议 `rtol=1e-3`，FP32 建议 `rtol=1e-5`。如果使用 TileLang 后端，注意它用 `exp2`/`log2` 而非 `exp`/`log`，会产生细微差异。
### Q: 如何单独跑一个算子的测试？
**A:**
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_copy.py -v --backend nvidia
```
### Q: 如何同时跑正确性和 benchmark？
**A:**
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op copy --backend nvidia --mode all
```
### Q: 测试结果不稳定怎么办？
**A:** 少量浮点误差波动（尤其是 reduce_sum 这类归约算子）是正常的。如果同一输入每次运行结果差异超过 `rtol=1e-5`（FP32）或 `rtol=1e-3`（FP16），检查是否有未初始化的内存或越界访问。
---
## 贡献流程
### Q: 我可以贡献什么？
**A:** 从 good-first-issue 标签入手比较合适。常见新手任务：修复文档错别字、补充测试用例、翻译文档、添加代码注释。详见 [how-to-submit-first-pr.md](how-to-submit-first-pr.md)。
### Q: PR 需要什么条件才能被合并？
**A:** 需要通过 CI（所有测试通过）、至少一位 reviewer 审核、代码风格符合规范、包含必要的测试。
### Q: 如何添加一个新算子？
**A:** 参考 [how-to-add-an-operator.md](how-to-add-an-operator.md)，完整流程包括 kernel 实现、C API、Python 绑定、测试和 benchmark。

54
docs/diagrams/grid-stride-loop.md
View File

@ -1,54 +0,0 @@
# Grid-Stride Loop 图示
## 概念
Grid-stride loop 是 GPU kernel 处理任意大小 tensor 的核心技术。每个线程不只处理一个元素，而是在循环中负责多个元素（跨步为 grid 总线程数），从而覆盖任意 N。
## Mermaid 图示
```mermaid
graph TB
subgraph "Global Memory (N=10 elements)"
M0["[0]"] --- M1["[1]"] --- M2["[2]"] --- M3["[3]"] --- M4["[4]"]
M5["[5]"] --- M6["[6]"] --- M7["[7]"] --- M8["[8]"] --- M9["[9]"]
end
subgraph "Grid = 2 Blocks × 2 Threads = 4 Threads"
subgraph "Block 0"
T0["Thread 0<br/>idx=0"]
T1["Thread 1<br/>idx=1"]
end
subgraph "Block 1"
T2["Thread 2<br/>idx=2"]
T3["Thread 3<br/>idx=3"]
end
end
T0 -->|"i=0"| M0
T0 -->|"i=4"| M4
T0 -->|"i=8"| M8
T1 -->|"i=1"| M1
T1 -->|"i=5"| M5
T1 -->|"i=9"| M9
T2 -->|"i=2"| M2
T2 -->|"i=6"| M6
T3 -->|"i=3"| M3
T3 -->|"i=7"| M7
```
## 代码对应
```cuda
// stride = gridDim.x * blockDim.x = 2 * 2 = 4
for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; // idx = 0,1,2,3
i < N; // N = 10
i += gridDim.x * blockDim.x) { // i += 4
dst[i] = src[i];
}
```
## 要点
- 每个线程负责多个元素，间隔为 stride（本例中 stride=4）
- 无论 N 多大，只要 `i < N` 条件在，就不会越界
- grid/block 尺寸可自由调整，不影响正确性，只影响性能

66
docs/diagrams/softmax-pipeline.md
View File

@ -1,66 +0,0 @@
# Softmax 流水线对比
## 朴素三趟扫描 vs Online Softmax
### 朴素实现（三趟扫描）
```
Pass 1: 求 row_max → max_val
Pass 2: 求 exp(x-max) 和 → sum (写入中间结果到 out)
Pass 3: 除以 sum → 最终归一化结果
```
### Online Softmax（两趟扫描）
```
Pass 1: 滚动更新 lse (log-sum-exp) → 不写 out
Pass 2: 用最终 lse 归一化 → 写入最终结果
```
## 流程图对比
```mermaid
flowchart LR
subgraph "朴素三趟扫描"
direction TB
N1["输入 x[0..N-1]"]
N2["Pass 1: 遍历全行<br/>max = max(x)"]
N3["Pass 2: 遍历全行<br/>sum = Σ exp(x - max)<br/>out[i] = exp(x[i]-max)"]
N4["Pass 3: 遍历全行<br/>out[i] /= sum"]
N5["输出 out[0..N-1]"]
N1 --> N2 --> N3 --> N4 --> N5
end
subgraph "Online Softmax（两趟扫描）"
direction TB
O1["输入 x[0..N-1]"]
O2["Pass 1: 滚动 lse<br/>lse = -∞<br/>for chunk in x:<br/> m_new = max(lse, max(chunk))<br/> lse = m_new + log(exp(lse-m_new)·out_old + Σ exp(chunk-m_new))"]
O3["Pass 2: 归一化<br/>for chunk in x:<br/> out[i] = exp2(log2_e · (x[i] - lse))"]
O4["输出 out[0..N-1]"]
O1 --> O2 --> O3 --> O4
end
```
## LSE 滚动更新详解
```mermaid
flowchart TD
Init["lse = -∞"]
Chunk1["处理 chunk 1<br/>m_new = max(-∞, max(chunk1)) = max(chunk1)<br/>lse = m_new + log(0 + Σ exp(chunk1 - m_new))"]
Chunk2["处理 chunk 2<br/>m_new = max(lse, max(chunk2))<br/>如果 m_new > lse: 旧结果需要 rescale<br/>lse = m_new + log(exp(lse-m_new) + Σ exp(chunk2-m_new))"]
Final["最终 lse 包含全行的 log-sum-exp"]
Init --> Chunk1 --> Chunk2 --> Final
```
## 要点
| 方面 | 朴素 | Online |
|------|------|--------|
| 扫描次数 | 3 趟 | 2 趟 |
| 全局内存写 | 2 次 (中间结果 + 最终) | 1 次 (仅最终) |
| 数值稳定性 | 靠减 max 保证 | 靠滚动 lse 保证 |
| TileLang 中的实现 | N/A | 用 `exp2`/`log2` 替代 `exp`/`log` |
- TileLang 版本使用 `exp2`/`log2` 因为硬件对 2 的幂运算支持更好
- `log2_e = 1.44269504` 是转换系数：`exp(x) = 2^(x * log2(e))`

71
docs/diagrams/thread-grid-layout.md
View File

@ -1,71 +0,0 @@
# 线程网格布局 — vector_add
## 概念
vector_add 是最经典的逐元素并行模式。输入 a 和 b 是两个等长一维 tensor，输出 c[i] = a[i] + b[i]。每个线程处理一组独立元素，线程之间无需通信。
## Mermaid 图示
```mermaid
graph TB
subgraph "输入 A (N=8)"
A0["a[0]"] --- A1["a[1]"] --- A2["a[2]"] --- A3["a[3]"]
A4["a[4]"] --- A5["a[5]"] --- A6["a[6]"] --- A7["a[7]"]
end
subgraph "输入 B (N=8)"
B0["b[0]"] --- B1["b[1]"] --- B2["b[2]"] --- B3["b[3]"]
B4["b[4]"] --- B5["b[5]"] --- B6["b[6]"] --- B7["b[7]"]
end
subgraph "Grid (4 Threads, stride=4)"
T0["Thread 0<br/>i=0,4"]
T1["Thread 1<br/>i=1,5"]
T2["Thread 2<br/>i=2,6"]
T3["Thread 3<br/>i=3,7"]
end
subgraph "输出 C (N=8)"
C0["c[0]"] --- C1["c[1]"] --- C2["c[2]"] --- C3["c[3]"]
C4["c[4]"] --- C5["c[5]"] --- C6["c[6]"] --- C7["c[7]"]
end
A0 --> T0 --> C0
B0 --> T0
A4 --> T0 --> C4
B4 --> T0
A1 --> T1 --> C1
B1 --> T1
A5 --> T1 --> C5
B5 --> T1
A2 --> T2 --> C2
B2 --> T2
A6 --> T2 --> C6
B6 --> T2
A3 --> T3 --> C3
B3 --> T3
A7 --> T3 --> C7
B7 --> T3
```
## 代码对应
```cuda
template <typename T>
__global__ void vector_add_contiguous_kernel(
T *out, const T *a, const T *b, int64_t n) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
int64_t stride = gridDim.x * blockDim.x;
for (int64_t i = idx; i < n; i += stride) {
out[i] = a[i] + b[i];
}
}
```
## 要点
- 每个 `(a[i], b[i])` 对独立计算，线程间零通信
- 相邻线程访问相邻内存地址 → 合并访问（coalesced access）
- TileLang 版本用 `T.Parallel` 表达同样的并行语义

85
docs/diagrams/tree-reduction.md
View File

@ -1,85 +0,0 @@
# 树形归约 — reduce_sum
## 概念
reduce_sum 对矩阵每行做求和归约。一个 block 处理一行，block 内线程先将各自负责的列累加为部分和，再通过 shared memory 做树形归约，最终 thread 0 输出整行的和。
## 树形归约过程（Block Size = 8）
```mermaid
flowchart TD
subgraph "Step 0: 初始化 Shared Memory"
S0["smem[0]=3"] --- S1["smem[1]=1"] --- S2["smem[2]=7"] --- S3["smem[3]=0"]
S4["smem[4]=4"] --- S5["smem[5]=1"] --- S6["smem[6]=6"] --- S7["smem[7]=3"]
end
subgraph "Step 1: stride=4"
A0["smem[0] += smem[4] → 7"] -..- A1["smem[1] += smem[5] → 2"]
A2["smem[2] += smem[6] → 13"] -..- A3["smem[3] += smem[7] → 3"]
end
subgraph "Step 2: stride=2"
B0["smem[0] += smem[2] → 20"] -..- B1["smem[1] += smem[3] → 5"]
end
subgraph "Step 3: stride=1"
C0["smem[0] += smem[1] → 25"]
end
S0 --> A0
S4 --> A0
S1 --> A1
S5 --> A1
S2 --> A2
S6 --> A2
S3 --> A3
S7 --> A3
A0 --> B0
A2 --> B0
A1 --> B1
A3 --> B1
B0 --> C0
B1 --> C0
C0 --> Result["out[row] = smem[0] = 25"]
```
## 代码对应
```cuda
__global__ void reduce_sum_rowwise_kernel(
float *out, const float *in, int64_t rows, int64_t cols) {
extern __shared__ float smem[];
int row = blockIdx.x;
int tid = threadIdx.x;
// 1. 各线程累加自己负责的列
float sum = 0.0f;
for (int c = tid; c < cols; c += blockDim.x) {
sum += in[row * cols + c];
}
smem[tid] = sum;
__syncthreads();
// 2. 树形归约
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
smem[tid] += smem[tid + s];
}
__syncthreads();
}
// 3. 输出
if (tid == 0) {
out[row] = smem[0];
}
}
```
## 要点
- 每步 stride 减半，log₂(blockDim.x) 步完成
- 每步之间必须 `__syncthreads()` —— 确保所有线程的 shared memory 写入对下一步可见
- `__syncthreads()` **不能**放在 `if (tid < s)` 分支内，否则死锁
- 优化方向：bank conflict 消除（添加 padding）

106
docs/en/phase1-kernel-writing.md
View File

@ -1,106 +0,0 @@
# Phase 1 Training Objective: Writing Kernels
## Goal
This phase has one primary task:
**Write kernel files for the existing operators.**
All other code (descriptor lifecycle, Python bindings, tests, build system) is already implemented. Your only job is to fill in the computation logic.
Default requirements:
1. **Priority: get the kernel correct and running.**
2. You do NOT need to pursue peak performance initially.
3. Complete the TODOs in `kernel.cuh` / `kernel.py` first, then consider aggressive optimizations.
If you've completed the basic goals and want to push performance further, you can modify the wrapper layer (e.g., `.cu` files for the NVIDIA backend), adjust launch policy, switch between kernels, or add more aggressive specializations. These are advanced topics and not required for Phase 1.
---
## Files to Write
Each operator has two backends, corresponding to two kernel files:
| Operator | NVIDIA kernel | TileLang kernel |
|----------|--------------|-----------------|
| `copy` | `ops/copy/nvidia/kernel.cuh` | `ops/copy/tilelang/kernel.py` |
| `vector_add` | `ops/vector_add/nvidia/kernel.cuh` | `ops/vector_add/tilelang/kernel.py` |
| `reduce_sum` | `ops/reduce_sum/nvidia/kernel.cuh` | `ops/reduce_sum/tilelang/kernel.py` |
| `softmax` | `ops/softmax/nvidia/kernel.cuh` | `ops/softmax/tilelang/kernel.py` |
Recommended order (easy to hard): copy → vector_add → reduce_sum → softmax.
---
## NVIDIA Kernel (`.cuh` files)
### What to Write
A `__global__` function placed under the operator's namespace. The function handles only computation logic, taking raw pointers as input/output. It does NOT involve any descriptor or public API.
By default, you do NOT need to modify the `.cu` files in the same directory this phase. The repository already provides compatibility-focused launch boilerplate. This phase only requires you to make `kernel.cuh` correct and runnable.
If you've completed the basic goal and want to push performance, treat the `.cu` file as an advanced optimization layer: adjust thread counts there, switch kernel variants, or add specialized paths. These are outside the scope of Phase 1.
### Difficulty Progression
**copy / vector_add**: Element-wise operations using the grid-stride loop pattern. Each thread handles several independent elements with no inter-thread communication.
**reduce_sum**: Row-wise reduction. One block per row. Threads within the block first independently accumulate their assigned columns, then perform a tree reduction via shared memory. `__syncthreads()` is required for synchronization.
**softmax**: The basic version uses a three-pass flow: first pass finds the row max (numerical stability), second pass computes `exp(x - max)` sum and writes intermediate results to output, third pass divides by the sum. Each pass requires intra-block synchronization.
### Concepts to Understand
- **grid-stride loop**: Why this pattern handles arbitrarily-sized tensors
- **shared memory reduction**: What each step of tree reduction does, and why `__syncthreads()` is necessary
- **softmax minus max**: Why computing `exp(x)` directly causes problems, and why subtracting the row max does not change the result
---
## TileLang Kernel (`.py` files)
### What to Write
A Python function decorated with `@tilelang.jit` that describes tile-level computation using the TileLang DSL. TileLang compiles it into real CUDA kernels.
As with NVIDIA, you only need to complete the TODOs in `kernel.py` this phase — no need to modify outer adapter code.
### Difficulty Progression
**copy**: Use the built-in `T.copy` to move data between two tiles. No manual loop needed.
**vector_add**: Use a `T.Parallel` loop to compute element-wise within a tile. Understand the difference between `T.Parallel` and `T.Serial`.
**reduce_sum**: Block-wise accumulation required. The outer loop uses `T.Serial` to traverse column-direction blocks sequentially (because state accumulates), and the inner loop uses `T.reduce_sum` to reduce over a fragment.
**softmax**: Two-pass scan using the online softmax algorithm. The first pass maintains a rolling log-sum-exp state; the second pass normalizes with the final LSE. Uses `exp2` / `log2` instead of `exp` / `log`.
### Concepts to Understand
- **`T.Parallel` vs `T.Serial`**: When can loop iterations run in parallel, and when must they be sequential
- **`T.alloc_fragment`**: A tile-level local buffer, corresponding to registers or shared memory
- **`T.copy`**: Moves a chunk of global memory into a fragment — not element-by-element assignment
- **online softmax**: Why a single pass can produce the correct normalization, and the rolling update logic for log-sum-exp
- **`exp2` / `log2`**: Why TileLang uses these instead of `exp` / `log`
---
## Verification
After writing each kernel, verify with the corresponding test:
```bash
# NVIDIA kernel (requires rebuild)
bash scripts/build_nvidia.sh build
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<op>.py -v --backend nvidia
# TileLang kernel (no build needed)
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<op>.py -v --backend tilelang
# Both correctness + benchmark
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <op> --backend nvidia --mode all
```
All four operators passing both backends = Phase 1 complete.

60
docs/glossary.md
View File

@ -1,60 +0,0 @@
# 术语表 (Glossary)
## GPU 编程
| 英文 | 中文 | 说明 |
|------|------|------|
| operator / op | 算子 | GPU 上执行的计算单元，如 copy、softmax |
| kernel | 核函数 | 在 GPU 上并行执行的函数，由 host 端调用 |
| grid | 网格 | CUDA 中最高层级的线程组织单元，由多个 block 组成 |
| block | 线程块 | grid 的子单元，block 内线程可以通过 shared memory 通信 |
| thread | 线程 | GPU 上最小的执行单元 |
| warp | 线程束 | GPU 调度的最小执行单元（NVIDIA 上为 32 个线程） |
| grid-stride loop | 网格跨步循环 | 让 kernel 处理任意大小输入的技术：每个线程负责多个元素，跨度为 grid 总线程数 |
| shared memory | 共享内存 | block 内线程共享的片上内存（SMEM），速度远超全局内存 |
| global memory | 全局内存 | GPU 的片外显存（HBM/DRAM），容量大但延迟高 |
| register | 寄存器 | 每个线程私有的最快存储 |
| SM (Streaming Multiprocessor) | 流式多处理器 | GPU 上的计算核心单元，一个 GPU 包含多个 SM，每个 SM 可并行执行多个 warp |
| coalesced access | 合并访问 | 相邻线程访问相邻地址，GPU 可将多次访问合并为一次内存事务 |
| stride | 步长 | tensor 在某个维度上相邻元素之间的内存地址间隔（以元素数计） |
| contiguous | 连续 | 描述 tensor 在内存中是否紧密排列、无间隙，合并访问的前提条件 |
| bank conflict | 存储体冲突 | shared memory 中多个线程同时访问同一 bank 的不同地址导致串行化 |
| __syncthreads() | 线程同步 | block 内所有线程到达此处后才能继续，用于 shared memory 写入后同步 |
| reduction | 归约 | 将多个值合并为一个值的操作（求和、求最大等） |
| online softmax | 在线 softmax | 一遍扫描完成 softmax 的算法，无需单独求 max 和 sum |
| occupancy | 占用率 | 每个 SM 上活跃 warp 数与理论最大值的比值，影响隐藏延迟的能力 |
| nvcc | NVIDIA CUDA 编译器 | 将 .cu 文件编译为 GPU 可执行的二进制 |
| CUDA | — | NVIDIA 的 GPU 通用并行计算平台与编程模型 |
## 数值精度
| 英文 | 中文 | 说明 |
|------|------|------|
| FP32 / float32 | 单精度浮点 | 32 位浮点数，约 7 位有效数字 |
| FP16 / float16 / half | 半精度浮点 | 16 位浮点数，约 3 位有效数字，显存占用为 FP32 的一半 |
| inf / NaN | 无穷/非数 | 浮点溢出或无效运算的结果 |
| tolerance / rtol / atol | 容差 | 测试中允许的数值误差上限 |
## 框架术语
| 英文 | 中文 | 说明 |
|------|------|------|
| descriptor | 描述符 | 封装算子运行所需全部信息（输入输出 shape、dtype、参数）的对象 |
| tensor_view | 张量视图 | 对 tensor 的轻量级引用，包含数据指针、shape、stride、dtype |
| backend | 后端 | 算子的具体硬件实现（NVIDIA CUDA、MetaX MACA、TileLang） |
| workspace | 工作空间 | kernel 执行时需要的临时内存 |
| launch / dispatch | 启动/分发 | 将 kernel 函数提交到 GPU 执行队列 |
| FFI | 外部函数接口 | Python 调用 C 动态库的机制（本项目用 ctypes） |
## TileLang 术语
| 英文 | 中文 | 说明 |
|------|------|------|
| TileLang | — | Python DSL，用 tile 级别的语义描述计算，JIT 编译为 GPU kernel |
| fragment | 片段 | tile 在单个线程上的局部数据块，对应寄存器或 shared memory |
| T.alloc_fragment | 分配片段 | 创建 tile 级别的局部 buffer |
| T.copy | 拷贝 | 在 tile 和全局内存之间搬移数据 |
| T.Parallel | 并行循环 | 循环的每次迭代可并行执行（不依赖前一次迭代的结果） |
| T.Serial | 顺序循环 | 循环必须按序执行（后一次依赖前一次的结果） |
| T.reduce_sum | 归约求和 | 对 fragment 内的元素做求和规约 |
| JIT | 即时编译 | 运行时将 TileLang Python 函数编译为 GPU 机器码 |

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

@ -1,131 +0,0 @@
# 如何开发一个新算子
## 目标
新增一个算子的最小闭环包括：
1. 在 `ops/<算子名>/nvidia/` 下实现 kernel 和 descriptor 生命周期。
2. 在 `python/operator_runtime/ops/` 下补 Python API。
3. 在 `tests/` 下补正确性测试和 benchmark。
4. 重新构建并验证。
训练营支持两条并行路径：
- **自定义算子**：`ops/<op>/nvidia/`，适合 kernel 教学、特殊 layout、reduce/softmax 等非逐元素结构。
- **elementwise 复用**：`ops/elementwise/<op>/nvidia/`，适合 unary/binary/broadcast 逐元素算子。
## Step 1：明确算子接口
开始前确认：
1. 几个输入、几个输出
2. 输出 shape 是否和输入一致
3. 是否要求 dtype 一致
4. 是否只支持 contiguous tensor
5. 是否需要额外参数（`dim`、`scalar` 等）
6. 是否需要 workspace
## Step 2：创建算子目录
最小文件集：
```
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
```
### 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 后端
一个算子需要四个生命周期接口：
1. `oprt_create_<op>_descriptor` — 检查输入合法性，保存运行信息
2. `oprt_get_<op>_workspace_size` — 返回临时内存大小
3. `oprt_execute_<op>` — 按 dtype dispatch 并 launch kernel
4. `oprt_destroy_<op>_descriptor` — 释放 descriptor
参考 `copy` 的实现：kernel 放 `kernel.cuh`，生命周期放 `<op>_cuda.cu`。
如果是 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`，通常暴露三个接口：
- `<op>` — out-of-place，自动分配输出
- `<op>_` — out-variant，调用方提供输出 tensor
- `prepare_<op>` — 创建 descriptor，支持多次执行复用
新增后在 `ops/__init__.py` 和 `operator_runtime/__init__.py` 中导出。
## Step 5：补测试
`tests/cases/<算子名>.py` 组织测试数据，分三类：
- `correctness_cases()` — 正确性用例（shape、dtype、tolerance）
- `api_error_cases()` — 异常输入用例
- `benchmark_cases()` — 性能测试规模
`tests/op_tests/test_<算子名>.py` 负责：
- 正确性对比（vs PyTorch）
- API contract 检查（shape/dtype 不匹配、非 contiguous）
- prepared 执行复用检查
`tests/bench/<算子名>.py` 负责性能入口。
## Step 6：构建和验证
```bash
# 新增 .cu 后必须重新 configure
bash scripts/build_nvidia.sh configure
bash scripts/build_nvidia.sh build
# 单算子验证
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend nvidia
# 正确性 + benchmark
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <算子名> --backend nvidia --mode all
# 确认没有破坏已有算子
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op all --backend nvidia --mode test
```
## 现有模板参考
| 模板 | 适用场景 |
| --- | --- |
| `copy` | 最简单的单输入单输出自定义算子 |
| `vector_add` | 双输入 contiguous 自定义算子 |
| `relu` | elementwise 框架复用，含标量参数 |
| `reduce_sum` | 带 reduce 维度的算子 |
| `softmax` | 多步归约 + 归一化 |

154
docs/how-to-submit-first-pr.md
View File

@ -1,154 +0,0 @@
# 如何为 intro-ops 提交第一个 PR
## 第一步：环境准备
### Fork 仓库 + Clone
1. 打开 https://github.com/metax/intro-ops ，点击右上角 **Fork** 按钮
2. Clone 你 fork 的仓库：
```bash
git clone https://github.com/<你的用户名>/intro-ops.git
cd intro-ops
```
3. 添加上游仓库（保持同步）：
```bash
git remote add upstream https://github.com/metax/intro-ops.git
```
### 配置开发环境
```bash
# 创建 conda 环境
conda create -n intro-ops python=3.12 -y
conda activate intro-ops
# 安装 Python 依赖
pip install -r requirements.txt
```
### 构建并验证基线
```bash
# NVIDIA 后端构建
bash scripts/build_nvidia.sh configure
bash scripts/build_nvidia.sh build
# 跑一遍全部测试，确保基线通过
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op all --backend nvidia --mode test
```
如果基线测试全部通过，说明环境配置正确，可以开始贡献。
---
## 第二步：选择你的贡献
在仓库的 [Issues](https://github.com/metax/intro-ops/issues) 页面，按标签筛选适合你的任务：
### 推荐的新手任务（good-first-issue）
| 类型 | 预计耗时 | 示例 |
|------|---------|------|
| 修复文档错别字/格式问题 | 15-30 min | 修正 README 中的拼写错误 |
| 补充测试用例 | 1-2 h | 为某个算子增加 corner case 测试 |
| 添加代码注释 | 30 min | 给关键算法步骤补充注释 |
| 翻译文档段落 | 1-2 h | 将中文文档翻译为英文（或反过来） |
| 补充 FAQ 条目 | 30 min | 把你遇到并解决的问题写成 FAQ |
### 不知道怎么选？
- 在 Issue 下留言询问
- 或者先读一遍文档，过程中发现任何可改进的地方记录下来
- 首次贡献从最小的事情做起，熟悉流程比贡献大小更重要
---
## 第三步：开发和自测
### 创建分支
```bash
git checkout -b fix/your-description
# 分支命名建议：
# fix/xxx — 修复
# feat/xxx — 新功能
# docs/xxx — 文档
# test/xxx — 测试
```
### 修改代码并本地验证
- **文档类**：修改后确认 Markdown 渲染正常
- **代码类**：修改后跑相关测试确保不引入回归：
```bash
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_copy.py -v --backend nvidia
```
### 代码风格
- C++ 代码：用 `clang-format` 格式化
- Python 代码：用 `black` + `isort` 格式化
---
## 第四步：提交 PR
### Commit 规范
```bash
git add <修改的文件>
git commit -m "docs: fix typo in README"
# commit message 格式：<类型>: <简短描述>
# 类型：docs / fix / feat / test / refactor
```
### 推送并创建 PR
```bash
git push origin fix/your-description
```
然后在 GitHub 上打开你的仓库，点击 **Compare & pull request** 按钮。
### 填写 PR 描述
一个好的 PR 描述应包含：
1. **做了什么**（一句话概括）
2. **为什么做**（问题背景）
3. **如何验证**（跑过哪些测试、附上截图或测试输出）
4. **关联 Issue**（如有，用 `Closes #123` 关联）
---
## 第五步：响应 Review
1. Review 意见是改进代码的机会，不要抵触
2. 如有不理解处，在 PR 评论区直接提问
3. 修改后 push 到同一分支，PR 会自动更新
4. 所有 Review 意见解决后，reviewer 会合并
---
## 合并之后
- 你的名字将出现在仓库的贡献者列表中
- 可以把你解决问题的经历写成经验分享（blog 或 Discussions）
- 继续探索更多贡献类型，挑战更难的任务
---
## 常见问题
**Q: 我的 PR 迟迟没人 review 怎么办？**
A: 通常 48h 内会有人响应。如果超时，可以在 PR 评论区 @ 仓库维护者。
**Q: 我不确定我的修改是否有价值？**
A: 可以先开一个 Issue 描述你的想法，得到确认后再动手。
**Q: 合并后有冲突怎么办？**
A: 将上游 main 合并到你的分支：
```bash
git fetch upstream
git merge upstream/main
# 解决冲突后 git push
```

121
docs/learning-path.md
View File

@ -1,121 +0,0 @@
# 分层学习路径
intro-ops 的四个算子从 copy 到 softmax 难度渐进。本文提供三级学习路线，帮你根据自己的基础和时间选择合适的路径。
---
## 入门级（约 1-2 周）
**目标：** 跑通 copy + vector_add，建立 GPU 编程基本概念
**前置知识：**
- C++ 基础（指针、模板、循环）
- 基本 GPU 概念：什么是 grid / block / thread
**推荐资源：**
- 《CUDA C Programming Guide》第 1-3 章（NVIDIA 官方，免费）
- intro-ops 的 `docs/glossary.md` 术语表
**学习路径：**
1. 阅读 `docs/phase1-kernel-writing.md` 了解整体目标
2. 阅读 `docs/glossary.md` 掌握基础术语
3. **copy 算子**：
- 理解 grid-stride loop（为什么这样写能处理任意大小 tensor）
- 完成 `ops/copy/nvidia/kernel.cuh` 的 TODO
- 完成 `ops/copy/tilelang/kernel.py` 的 TODO
- 跑通测试：`pytest tests/op_tests/test_copy.py -v --backend nvidia`
4. **vector_add 算子**：
- 理解逐元素并行和 `T.Parallel`
- 完成 `ops/vector_add/nvidia/kernel.cuh` 的 TODO
- 完成 `ops/vector_add/tilelang/kernel.py` 的 TODO
- 跑通测试
**产出：** 两个算子两种后端全部通过测试
**如果你卡住了：** 查看 `docs/troubleshooting.md` 的"运行时错误"章节
---
## 进阶级（约 2-4 周）
**目标：** 完成 reduce_sum + softmax，理解 shared memory 编程
**前置知识：**
- 已完成入门级
- 了解 shared memory 概念
- 了解线程同步机制
**推荐资源：**
- 《Professional CUDA C Programming》shared memory 章节
- `docs/tilelang-vs-cuda.md` 对比教程
**学习路径：**
1. **reduce_sum 算子**：
- 理解 shared memory 树形归约原理
- 理解 `__syncthreads()` 的使用时机（为什么不能放在条件分支内）
- 完成 `ops/reduce_sum/nvidia/kernel.cuh` 的 TODO
- 完成 `ops/reduce_sum/tilelang/kernel.py` 的 TODO（理解 `T.Serial` 的作用）
- 跑通测试
- 尝试优化 bank conflict（加 padding）
2. **softmax 算子**：
- 理解数值稳定性：为什么要减 max
- 理解 online softmax 算法（一遍扫描 vs 三趟扫描）
- 完成 `ops/softmax/nvidia/kernel.cuh` 的 TODO
- 完成 `ops/softmax/tilelang/kernel.py` 的 TODO（理解 `exp2`/`log2` 的用法）
- 跑通测试
3. 提交 benchmark 数据，对比自己的实现与 PyTorch 的差异
**产出：** 全部四个算子两种后端通过测试 + benchmark 数据
---
## 专精级（约 5 周+）
**目标：** 性能调优 + 贡献新算子，从"会用"到"能写"
**前置知识：**
- 已完成进阶级
- 对 GPU 架构有一定了解（SM、warp、memory hierarchy）
**推荐资源：**
- CUTLASS 官方文档
- `docs/how-to-add-an-operator.md` 新增算子指南
**可选方向：**
### 方向 A：性能调优
- 分析 benchmark 数据，找到瓶颈（内存带宽 / 计算 / 延迟）
- 调整 launch 参数（block size、grid size）
- 尝试 warp-level shuffle 优化
- 将你的优化写成经验分享
### 方向 B：CUTLASS / CuTe 集成
- 将 CUTLASS 的 tiling / warp-level 优化应用到现有算子
- 理解 CuTe 的 layout 和 tensor 抽象
### 方向 C：贡献新算子
- 阅读 `docs/how-to-add-an-operator.md`
- 选一个感兴趣的算子（如 LayerNorm、GELU、attention）
- 按流程实现：kernel → C API → Python 绑定 → 测试 → benchmark
- 提交 PR
### 方向 D：社区贡献
- 阅读 `docs/how-to-submit-first-pr.md`
- 翻译文档（中→英或英→中）
- 补充 FAQ / troubleshooting 条目
- Review 其他人的 PR
**产出：** PR 被合并 / 新算子通过 review / 在社区分享经验
---
## 学习建议
- **不要一次追求完美**：先写对再优化，能跑通的 kernel 比跑不快的 kernel 好一万倍
- **善用 checklist**：`course/checklist.md` 帮你追踪进度
- **遇到问题先查 FAQ 和 troubleshooting**：大部分常见问题已经有答案
- **把你的经验写下来**：帮你解决问题的经历，也可能是别人的 FAQ 条目

106
docs/phase1-kernel-writing.md
View File

@ -1,106 +0,0 @@
# 第一阶段训练目标：写 Kernel
## 目标
本阶段主要需要做一件事：
**给现有算子写 kernel 文件。**
其余所有代码（descriptor 生命周期、Python 绑定、测试、构建）都已经写好，剩下需要补全计算逻辑本身。
默认要求是：
1. **优先把 kernel 写对、跑通。**
2. 不需要一开始就追求最优性能。
3. 先完成 `kernel.cuh` / `kernel.py` 里的 TODO，再考虑更激进的优化。
如果你已经完成基础目标，并且学有余力，希望进一步冲性能，那么可以继续修改对应算子的包装层文件（例如 NVIDIA 后端的 `.cu` 文件），调整 launch policy、切换不同 kernel、加入更激进的 specialization。这些都属于进阶优化内容，不属于第一阶段的必做要求。
---
## 需要写的文件
每个算子有两种后端，对应两个 kernel 文件：
| 算子 | NVIDIA kernel | TileLang kernel |
|------|--------------|-----------------|
| `copy` | `ops/copy/nvidia/kernel.cuh` | `ops/copy/tilelang/kernel.py` |
| `vector_add` | `ops/vector_add/nvidia/kernel.cuh` | `ops/vector_add/tilelang/kernel.py` |
| `reduce_sum` | `ops/reduce_sum/nvidia/kernel.cuh` | `ops/reduce_sum/tilelang/kernel.py` |
| `softmax` | `ops/softmax/nvidia/kernel.cuh` | `ops/softmax/tilelang/kernel.py` |
建议按从简到难的顺序：copy → vector_add → reduce_sum → softmax。
---
## NVIDIA Kernel（`.cuh` 文件）
### 写什么
一个 `__global__` 函数，放在对应算子的命名空间下。函数只负责计算逻辑，输入输出是裸指针，不涉及任何 descriptor 或 API。
本阶段默认**不需要修改**同目录下的 `.cu` 文件。仓库已经提供了兼容性优先的 launch 样板，这一阶段只要求把 `kernel.cuh` 写对、跑通。
如果你已经完成基础目标，并且想继续冲性能，可以把 `.cu` 当作进阶优化层：在那里调整线程数、切换不同 kernel、加入更激进的 specialization。但这些不属于第一阶段的必做内容。
### 四个算子的难度递进
**copy / vector_add**：逐元素操作，用 grid-stride loop 模式。每个线程负责若干个独立的元素，线程之间没有通信。
**reduce_sum**：行规约。每行一个 block，block 内线程先各自累加自己负责的列，再通过 shared memory 做树形归约，最终 `smem[0]` 是整行的结果。需要用 `__syncthreads()` 同步。
**softmax**：基础版本使用三段流程：第一遍求行最大值（数值稳定性），第二遍求 `exp(x - max)` 的和并把中间结果写到输出，第三遍再除以 sum。每一遍都需要 block 内同步。
### 需要理解的概念
- **grid-stride loop**：为什么这样写可以处理任意大小的 tensor
- **shared memory 规约**：树形归约的每一步在做什么，为什么需要 `__syncthreads()`
- **softmax 减 max**：为什么直接算 `exp(x)` 会出问题，减去行最大值为什么不改变结果
---
## TileLang Kernel（`.py` 文件）
### 写什么
一个用 `@tilelang.jit` 装饰的 Python 函数，用 TileLang DSL 描述 tile 粒度的计算逻辑。TileLang 会把它编译成真正的 CUDA kernel。
和 NVIDIA 一样，本阶段默认只需要填写 `kernel.py` 里的 TODO，不需要修改外层适配代码。
### 四个算子的难度递进
**copy**：直接用内置 `T.copy` 在两个 tile 之间搬数据，无需手写循环。
**vector_add**：用 `T.Parallel` 循环在 tile 内逐元素计算，理解 `T.Parallel` 与 `T.Serial` 的区别。
**reduce_sum**：需要分块累加。外层用 `T.Serial` 顺序遍历列方向的分块（因为要累积状态），内层用 `T.reduce_sum` 对 fragment 做规约。
**softmax**：两遍扫描，使用 online softmax 算法。第一遍滚动维护 log-sum-exp 状态，第二遍用最终的 lse 归一化。计算中用 `exp2` / `log2` 替代 `exp` / `log`。
### 需要理解的概念
- **`T.Parallel` vs `T.Serial`**：什么情况下循环内的迭代可以并行，什么情况下必须顺序
- **`T.alloc_fragment`**：tile 级别的局部 buffer，对应寄存器或 shared memory
- **`T.copy`**：把全局内存的一块搬到 fragment，不是逐元素赋值
- **online softmax**：为什么一遍扫描就能算出正确的归一化，log-sum-exp 的滚动更新逻辑
- **`exp2` / `log2`**：为什么 TileLang 里用这两个而不是 `exp` / `log`
---
## 验证方式
每写完一个 kernel，用对应的测试验证：
```bash
# NVIDIA kernel（需要先重新编译）
bash scripts/build_nvidia.sh build
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<算子名>.py -v --backend nvidia
# 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
```
四个算子两种后端全部通过，阶段一完成。

163
docs/tilelang-vs-cuda.md
View File

@ -1,163 +0,0 @@
# TileLang vs CUDA 对比教程
intro-ops 的特色之一是同时支持 CUDA 和 TileLang 两种后端。本文通过 copy 算子的两种实现，帮助理解"同一个算法，两种 DSL 的表达差异"。
## 概述对比
| 维度 | CUDA | TileLang |
|------|------|----------|
| 抽象层级 | 线程级（thread/warp/block） | tile 级（fragment/tile） |
| 编写难度 | 需关注同步、shared memory 细节 | 声明式，编译器处理调度 |
| 编译方式 | nvcc 编译 `.cu` 文件 | JIT 编译，无需构建 |
| 调优手段 | 手动调整 grid/block/shared memory | `T.Parallel` / `T.Serial` 语义 |
| 适用场景 | 极致性能调优 | 快速原型 + 跨硬件 |
| 学习曲线 | 陡峭 | 平缓 |
---
## 案例一：copy 算子
### CUDA 实现
```cuda
template <typename T>
__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
// 1. 计算全局线程索引
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
// 2. 计算网格总步长
int64_t stride = gridDim.x * blockDim.x;
// 3. grid-stride loop
for (int64_t i = idx; i < n; i += stride) {
dst[i] = src[i];
}
}
```
**关键细节：**
- 你需要手动管理线程索引（`blockIdx.x`、`blockDim.x`、`threadIdx.x`）
- 你需要手动设计 grid-stride loop 来处理任意大小的 tensor
- launch 时你需要自己指定 grid size 和 block size
- 你需要理解"一个线程负责哪些元素"的映射关系
### TileLang 实现
```python
@tilelang.jit
def copy_kernel(src, BLOCK_N: int, dtype):
N = T.const("N")
src: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
for i in T.Parallel(N // BLOCK_N):
tile = T.copy(src[i * BLOCK_N : (i + 1) * BLOCK_N])
out[i * BLOCK_N : (i + 1) * BLOCK_N] = tile
return out
```
**关键细节：**
- `T.Parallel` 声明一个 tile 级别的并行循环——编译器负责把 tile 分配到线程
- `T.copy` 加载一个 tile 到局部 fragment——不需要手写线程索引
- 没有显式的 grid/block/thread 概念
- 编译器自动决定最优的 launch 参数
### 核心差异：谁负责什么
| 职责 | CUDA | TileLang |
|------|------|----------|
| 线程→元素映射 | 你手动写 `idx = blockIdx.x * blockDim.x + threadIdx.x` | 编译器生成 |
| 循环边界检查 | 你手动写 `i < n` | 编译器生成 |
| shared memory 管理 | 你手动声明、加载、同步 | 编译器处理 |
| launch 参数 | 你手动指定 grid/block size | 编译器自动选择 |
| 内存合并访问 | 你需要确保访问模式正确 | 编译器优化 |
---
## 案例二：reduce_sum 算子
### CUDA 实现
树形归约需要关注 shared memory 管理和线程同步：
```cuda
template <typename T>
__global__ void reduce_sum_kernel(T *out, const T *in, int64_t N, int64_t C) {
extern __shared__ float smem[];
int64_t row_id = blockIdx.x;
int tid = threadIdx.x;
// 1. 每个线程累加自己负责的列
float sum = 0.0f;
for (int64_t c = tid; c < C; c += blockDim.x) {
sum += (float)in[row_id * C + c];
}
smem[tid] = sum;
__syncthreads();
// 2. 树形归约
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (tid < s) {
smem[tid] += smem[tid + s];
}
__syncthreads();
}
// 3. 输出结果
if (tid == 0) {
out[row_id] = (T)smem[0];
}
}
```
**陷阱点：**
- `__syncthreads()` 必须放在条件分支外
- shared memory 大小需要作为 launch 参数传入
- 树形归约的 stride 和边界条件容易写错
### TileLang 实现
```python
@tilelang.jit
def reduce_sum_kernel(inp, BLOCK_C: int, dtype):
N = T.const("N")
C = T.const("C")
inp: T.Tensor((N, C), dtype)
out = T.empty((N,), dtype)
for n in T.Parallel(N):
acc = T.alloc_fragment((1,), dtype)
acc[0] = 0.0
for c in T.Serial(C // BLOCK_C):
tile = T.copy(inp[n, c * BLOCK_C : (c + 1) * BLOCK_C])
acc[0] += T.reduce_sum(tile)
out[n] = acc[0]
return out
```
**关键差异：**
- shared memory 归约被 `T.reduce_sum` 替代——编译器负责生成高效的归约代码
- `T.Serial` 保证分块累加的顺序（因为累加有依赖）
- 不需要手动管理 shared memory 大小和同步
---
## 选择建议
### 什么时候用 CUDA？
- 需要极致性能，想手动控制每一个优化细节
- 算子涉及复杂的 shared memory / warp-level 操作
- 需要嵌入已有 CUDA 生态（CUTLASS、cuBLAS）
### 什么时候用 TileLang？
- 快速原型：写一个可工作的 kernel 远比性能重要
- 跨硬件：同一份代码可编译到 NVIDIA / MetaX / 甚至未来其他硬件
- 学习阶段：先理解算法逻辑，不必被线程细节分心
### intro-ops 的建议路径
1. **入门**：先写 TileLang 版本，快速跑通算法逻辑
2. **进阶**：再写 CUDA 版本，理解线程级别的执行细节
3. **专精**：对比两个版本，分析编译器生成的差异，手动调优 CUDA

163
docs/troubleshooting.md
View File

@ -1,163 +0,0 @@
# 常见错误与排错指南
## 编译错误
### nvcc 版本不匹配
**症状**：`CMake Error: nvcc not found` 或 `nvcc fatal: Unsupported gpu architecture`
**原因**：CUDA Toolkit 版本与 CMake 预设的架构参数不兼容。
**解决**：
1. 检查 nvcc 版本：`nvcc --version`
2. 在 `CMakePresets.json` 中将 `CMAKE_CUDA_ARCHITECTURES` 改为 `"native"`，或显式指定你的 GPU 架构编号
3. 确保 `nvcc` 在 PATH 中
### CUDA 架构参数错误
**症状**：`cudaErrorNoKernelImageForDevice` 或运行时 kernel launch 失败
**原因**：编译时指定的 GPU 架构（如 `89` 对应 L40）与运行时 GPU 不匹配。
**解决**：
- 将 `CMAKE_CUDA_ARCHITECTURES` 设为 `"native"` 让 CMake 自动检测
- 或根据 GPU 型号查表设置正确的架构编号
### CUTLASS 拉取失败
**症状**：`FetchContent failed to download CUTLASS`
**原因**：GitHub 网络不可达或代理问题。
**解决**：
1. 检查网络：`git ls-remote https://github.com/NVIDIA/cutlass.git`
2. 配置代理：`git config --global http.proxy http://your-proxy:port`
3. 或手动下载 CUTLASS v3.7.0 放到 `third_party/cutlass/`
---
## 运行时错误
### 越界访问（grid-stride loop 边界条件）
**症状**：随机数值错误、CUDA `illegal memory access` 错误
**原因**：`idx < N` 条件缺失或写错，导致线程访问超出 tensor 范围的内存。
**解决**：
```cuda
// 正确写法：每次循环都检查 idx < N
for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
idx < N;
idx += gridDim.x * blockDim.x) {
out[idx] = in[idx]; // 安全：idx 始终 < N
}
```
### `__syncthreads()` 在条件分支内
**症状**：kernel 在 reduce_sum 或 softmax 中挂起（hang），或结果错误
**原因**：`__syncthreads()` 放在 `if` 分支内。CUDA 要求 block 内所有线程都到达同一个 `__syncthreads()`，如果有线程走 else 分支跳过同步点，整个 block 就会死锁。
**解决**：
```cuda
// 错误
if (threadIdx.x < N) {
smem[tid] = val;
__syncthreads(); // 部分线程不执行，死锁！
}
// 正确
smem[tid] = (threadIdx.x < N) ? val : 0;
__syncthreads(); // 全部线程都到达
```
### shared memory 大小不足
**症状**：编译错误 `uses too much shared data`
**原因**：申请的 shared memory 超过 GPU 的物理限制（通常 48KB-164KB/block）。
**解决**：
- 检查 `extern __shared__` 声明的数组大小
- 减小 block size 或 shared memory 使用量
- 分多轮处理
---
## 数值错误
### softmax 未减 max（大数值溢出）
**症状**：softmax 输出全为 NaN 或 inf
**原因**：`exp(x)` 在 `x > 88` 时溢出为 inf，直接除 inf 得到 NaN。
**解决**：先减行最大值再做 exp：
```cuda
float max_val = row[0];
for (int j = 1; j < N; j++) {
max_val = fmaxf(max_val, row[j * stride]);
}
float sum = 0;
for (int j = 0; j < N; j++) {
sum += expf(row[j * stride] - max_val);
}
```
### 浮点精度差异（多后端对比）
**症状**：NVIDIA 和 TileLang 的输出在小数点后几位不一致，测试报 `not close`
**原因**：不同后端使用不同数学库（`exp2`/`log2` vs `exp`/`log`），或规约顺序不同导致浮点累加误差。
**解决**：
- 检查测试容差设置是否合理（FP16 容差应比 FP32 宽松）
- 确认两边的算法逻辑一致（如 online vs 三趟 softmax）
- 规约顺序差异导致的误差在合理范围内（`rtol=1e-3` for FP16）则正常
### 规约顺序影响结果
**症状**：reduce_sum 结果每次运行都略有不同
**原因**：浮点加法不满足结合律，tree reduction 的配对顺序影响结果。
**解决**：这是正常的浮点行为，只要误差在容差范围内即可接受。
---
## 性能问题
### bank conflict
**症状**：reduce_sum kernel 带宽利用率远低于预期（如 < 50%）
**原因**：shared memory 访问时多个线程命中同一 bank。
**解决**：
- 添加 padding 偏移访问地址
- 使用 `__shared__ float smem[BLOCK_SIZE + PADDING]` 并错开 bank
- 将 stride 从 1 改为 2 开始归约
### 线程利用率低
**症状**：benchmark 带宽利用率低，但代码逻辑正确
**原因**：grid/block 尺寸选择不当，GPU SM 上的线程不够填满所有 core。
**解决**：
- block size 建议 128/256/512（取 32 的倍数）
- grid size 建议 `(N + block_size - 1) / block_size` 或更大
- 用 `cudaOccupancyMaxPotentialBlockSize` 自动计算最佳配置
### 全局内存未合并访问
**症状**：copy 或 vector_add 带宽利用率远低于峰值
**原因**：线程访问的内存地址不连续，导致多次内存事务而非一次合并访问。
**解决**：
- 确保相邻线程访问相邻内存地址（lane 0→地址 0，lane 1→地址 1…）
- 检查 stride 是否为 1
- 检查数据类型是否与访问模式对齐

28
examples/01_copy.py
View File

@ -1,28 +0,0 @@
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 copy
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
src = torch.randn((1024,), device="cuda", dtype=torch.float32)
out = copy(src, backend=args.backend)
torch.testing.assert_close(out, src)
print(f"copy ok ({args.backend})")
if __name__ == "__main__":
main()

29
examples/02_vector_add.py
View File

@ -1,29 +0,0 @@
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 vector_add
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
a = torch.randn((1024,), device="cuda", dtype=torch.float32)
b = torch.randn_like(a)
out = vector_add(a, b, backend=args.backend)
torch.testing.assert_close(out, a + b)
print(f"vector_add ok ({args.backend})")
if __name__ == "__main__":
main()

28
examples/03_reduce_sum.py
View File

@ -1,28 +0,0 @@
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 reduce_sum
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
src = torch.randn((32, 128), device="cuda", dtype=torch.float32)
out = reduce_sum(src, dim=1, backend=args.backend)
torch.testing.assert_close(out, torch.sum(src, dim=1), atol=1e-5, rtol=1e-5)
print(f"reduce_sum ok ({args.backend})")
if __name__ == "__main__":
main()

28
examples/04_softmax.py
View File

@ -1,28 +0,0 @@
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 softmax
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default="nvidia", choices=["nvidia", "tilelang", "metax"])
args = parser.parse_args()
src = torch.randn((32, 128), device="cuda", dtype=torch.float32)
out = softmax(src, dim=1, backend=args.backend)
torch.testing.assert_close(out, torch.softmax(src, dim=1), atol=1e-5, rtol=1e-5)
print(f"softmax ok ({args.backend})")
if __name__ == "__main__":
main()

37
examples/05_tilelang_copy_modes.py
View File

@ -1,37 +0,0 @@
from __future__ import annotations
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 ops.common.tilelang.eager_copy import copy_eager, copy_eager_
from ops.common.tilelang.lazy_out_idx_copy import copy_lazy_out_idx, copy_lazy_out_idx_
def main() -> None:
src = torch.randn((1024,), device="cuda", dtype=torch.float32)
eager = copy_eager(src)
torch.testing.assert_close(eager, src)
eager_out = torch.empty_like(src)
copy_eager_(eager_out, src)
torch.testing.assert_close(eager_out, src)
lazy = copy_lazy_out_idx(src)
torch.testing.assert_close(lazy, src)
lazy_out = torch.empty_like(src)
copy_lazy_out_idx_(lazy_out, src)
torch.testing.assert_close(lazy_out, src)
print("tilelang eager and lazy out_idx copy ok")
if __name__ == "__main__":
main()

30
examples/06_relu.py
View File

@ -1,30 +0,0 @@
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()

56
include/operator_runtime/api.h
View File

@ -1,56 +0,0 @@
#pragma once
#include <stdint.h>
#include <stddef.h>
#ifdef _WIN32
#define OPRT_EXPORT __declspec(dllexport)
#else
#define OPRT_EXPORT __attribute__((visibility("default")))
#endif
#ifdef __cplusplus
extern "C" {
#endif
typedef enum {
OPRT_SUCCESS = 0,
OPRT_ERR_INVALID_ARG = 1,
OPRT_ERR_UNSUPPORTED_DTYPE = 2,
OPRT_ERR_RUNTIME = 3,
OPRT_ERR_INSUFFICIENT_WORKSPACE = 4,
OPRT_ERR_NOT_SUPPORTED = 5
} oprt_status_t;
typedef enum {
OPRT_DTYPE_F16 = 0,
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;
typedef struct oprt_operator_descriptor *oprt_operator_descriptor_t;
typedef struct {
void *data;
oprt_dtype_t dtype;
int32_t ndim;
int64_t shape[OPRT_MAX_DIMS];
int64_t strides[OPRT_MAX_DIMS];
} 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

16
include/operator_runtime/descriptor.h
View File

@ -1,16 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __cplusplus
#include <string>
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

29
include/operator_runtime/detail/cuda_helpers.h
View File

@ -1,29 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __CUDACC__
#include <cuda_runtime.h>
#define OPRT_CUDA_RETURN_IF_ERROR(expr) \
do { \
cudaError_t err__ = (expr); \
if (err__ != cudaSuccess) { \
return OPRT_ERR_RUNTIME; \
} \
} while (0)
namespace oprt {
inline cudaStream_t as_cuda_stream(oprt_stream_t stream) {
return reinterpret_cast<cudaStream_t>(stream);
}
inline int blocks_for(int64_t n, int threads) {
int64_t blocks = (n + threads - 1) / threads;
return static_cast<int>(blocks > 0 ? blocks : 1);
}
} // namespace oprt
#endif

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

@ -1,249 +0,0 @@
#pragma once
#include "operator_runtime/descriptor.h"
#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);
}
inline bool elementwise_fast_path(const oprt_tensor_view_t &out,
const oprt_tensor_view_t &a,
const oprt_tensor_view_t &b) {
return same_shape(out, a) && same_shape(out, b) &&
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;
}
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

13
include/operator_runtime/detail/operation.h
View File

@ -1,13 +0,0 @@
#pragma once
#include <string>
namespace oprt {
struct OperationSpec {
std::string name;
std::string backend;
std::string kind;
};
} // namespace oprt

39
include/operator_runtime/detail/tensor_checks.h
View File

@ -1,39 +0,0 @@
#pragma once
#include "operator_runtime/tensor_view.h"
#ifdef __cplusplus
namespace oprt {
inline oprt_status_t check_tensor(const oprt_tensor_view_t *view) {
if (view == nullptr || view->data == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (view->ndim < 0 || view->ndim > OPRT_MAX_DIMS) {
return OPRT_ERR_INVALID_ARG;
}
for (int32_t i = 0; i < view->ndim; ++i) {
if (view->shape[i] < 0) {
return OPRT_ERR_INVALID_ARG;
}
}
if (view->dtype != OPRT_DTYPE_F16 && view->dtype != OPRT_DTYPE_F32) {
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
return OPRT_SUCCESS;
}
inline oprt_status_t check_same_dtype(const oprt_tensor_view_t &a,
const oprt_tensor_view_t &b) {
return a.dtype == b.dtype ? OPRT_SUCCESS : OPRT_ERR_UNSUPPORTED_DTYPE;
}
inline oprt_status_t check_same_shape(const oprt_tensor_view_t &a,
const oprt_tensor_view_t &b) {
return same_shape(a, b) ? OPRT_SUCCESS : OPRT_ERR_INVALID_ARG;
}
} // namespace oprt
#endif

10
include/operator_runtime/operator_runtime.h
View File

@ -1,10 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#include "operator_runtime/descriptor.h"
#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"

31
include/operator_runtime/ops/copy.h
View File

@ -1,31 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_copy_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *dst,
const oprt_tensor_view_t *src);
OPRT_EXPORT oprt_status_t oprt_get_copy_workspace_size(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_copy(
oprt_operator_descriptor_t desc,
void *workspace,
size_t workspace_size,
void *dst,
const void *src,
oprt_stream_t stream);
OPRT_EXPORT oprt_status_t oprt_destroy_copy_descriptor(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

32
include/operator_runtime/ops/reduce_sum.h
View File

@ -1,32 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_reduce_sum_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis);
OPRT_EXPORT oprt_status_t oprt_get_reduce_sum_workspace_size(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_reduce_sum(
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_reduce_sum_descriptor(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

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

@ -1,32 +0,0 @@
#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

32
include/operator_runtime/ops/softmax.h
View File

@ -1,32 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_softmax_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis);
OPRT_EXPORT oprt_status_t oprt_get_softmax_workspace_size(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_softmax(
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_softmax_descriptor(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

33
include/operator_runtime/ops/vector_add.h
View File

@ -1,33 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __cplusplus
extern "C" {
#endif
OPRT_EXPORT oprt_status_t oprt_create_vector_add_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *a,
const oprt_tensor_view_t *b);
OPRT_EXPORT oprt_status_t oprt_get_vector_add_workspace_size(
oprt_operator_descriptor_t desc,
size_t *size);
OPRT_EXPORT oprt_status_t oprt_execute_vector_add(
oprt_operator_descriptor_t desc,
void *workspace,
size_t workspace_size,
void *out,
const void *a,
const void *b,
oprt_stream_t stream);
OPRT_EXPORT oprt_status_t oprt_destroy_vector_add_descriptor(
oprt_operator_descriptor_t desc);
#ifdef __cplusplus
}
#endif

66
include/operator_runtime/tensor_view.h
View File

@ -1,66 +0,0 @@
#pragma once
#include "operator_runtime/api.h"
#ifdef __cplusplus
#include <algorithm>
#include <cstdint>
namespace oprt {
inline int64_t numel(const oprt_tensor_view_t &view) {
if (view.ndim < 0 || view.ndim > OPRT_MAX_DIMS) {
return 0;
}
int64_t total = 1;
for (int32_t i = 0; i < view.ndim; ++i) {
total *= view.shape[i];
}
return total;
}
inline bool is_contiguous(const oprt_tensor_view_t &view) {
if (view.ndim < 0 || view.ndim > OPRT_MAX_DIMS) {
return false;
}
int64_t expected = 1;
for (int32_t i = view.ndim - 1; i >= 0; --i) {
if (view.shape[i] == 1) {
continue;
}
if (view.strides[i] != expected) {
return false;
}
expected *= view.shape[i];
}
return true;
}
inline bool same_shape(const oprt_tensor_view_t &a, const oprt_tensor_view_t &b) {
if (a.ndim != b.ndim) {
return false;
}
for (int32_t i = 0; i < a.ndim; ++i) {
if (a.shape[i] != b.shape[i]) {
return false;
}
}
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

164
notebooks/00_environment_setup.ipynb
View File

@ -1,164 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 00 — 环境配置与验证\n",
"\n",
"本 Notebook 帮助你验证 intro-ops 开发环境是否正确配置。\n",
"\n",
"## 检查清单\n",
"\n",
"1. GPU 是否可用\n",
"2. CUDA Toolkit 版本\n",
"3. PyTorch 是否正确安装\n",
"4. intro-ops 是否成功构建\n",
"5. 运行一个简单的基线测试"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: 检查 GPU 和 CUDA"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"print(f\"PyTorch 版本: {torch.__version__}\")\n",
"print(f\"CUDA 可用: {torch.cuda.is_available()}\")\n",
"\n",
"if torch.cuda.is_available():\n",
" print(f\"GPU 型号: {torch.cuda.get_device_name(0)}\")\n",
" print(f\"CUDA 版本: {torch.version.cuda}\")\n",
" print(f\"GPU 数量: {torch.cuda.device_count()}\")\n",
" print(f\"当前设备: {torch.cuda.current_device()}\")\n",
"else:\n",
" print(\"⚠️ 未检测到 CUDA GPU。intro-ops 需要 NVIDIA GPU 才能运行。\")\n",
" print(\" 请确认 CUDA Toolkit 和 NVIDIA 驱动已正确安装。\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: 检查 intro-ops 构建产物"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import sys\n",
"\n",
"# 检查 CAMP_BUILD_DIR 环境变量\n",
"build_dir = os.environ.get(\"CAMP_BUILD_DIR\", \"build-nvidia\")\n",
"lib_path = os.path.join(build_dir, \"libcamp_ops.so\")\n",
"\n",
"print(f\"构建目录: {build_dir}\")\n",
"print(f\"库文件路径: {lib_path}\")\n",
"print(f\"库文件存在: {os.path.exists(lib_path)}\")\n",
"\n",
"if not os.path.exists(lib_path):\n",
" print(\"\\n⚠️ libcamp_ops.so 未找到。请先构建项目：\")\n",
" print(\" bash scripts/build_nvidia.sh configure\")\n",
" print(\" bash scripts/build_nvidia.sh build\")\n",
" print(f\"\\n 然后设置环境变量: export CAMP_BUILD_DIR={build_dir}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: 验证 Python 依赖"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import importlib\n",
"\n",
"deps = [\"torch\", \"pytest\", \"cmake\", \"ninja\"]\n",
"for dep in deps:\n",
" try:\n",
" mod = importlib.import_module(dep)\n",
" version = getattr(mod, \"__version__\", \"unknown\")\n",
" print(f\"✓ {dep}: {version}\")\n",
" except ImportError:\n",
" print(f\"✗ {dep}: 未安装\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 4: 运行基线测试\n",
"\n",
"如果以上检查全部通过，运行一个简单的 copy 算子测试验证端到端流程。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import subprocess\n",
"import sys\n",
"\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_copy.py\", \"-v\",\n",
" \"--backend\", \"nvidia\",\n",
" \"-x\", # 遇到第一个失败就停止\n",
"], capture_output=True, text=True)\n",
"\n",
"print(result.stdout)\n",
"if result.returncode != 0:\n",
" print(result.stderr)\n",
" print(\"\\n⚠️ 基线测试失败。请检查上面的错误信息。\")\n",
"else:\n",
" print(\"\\n✓ 环境配置正确，可以开始学习了！\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 常见问题\n",
"\n",
"- **`torch.cuda.is_available()` 返回 False**：检查 NVIDIA 驱动和 CUDA Toolkit 安装\n",
"- **`libcamp_ops.so` 不存在**：运行 `bash scripts/build_nvidia.sh configure && bash scripts/build_nvidia.sh build`\n",
"- **ImportError**：运行 `pip install -r requirements.txt`\n",
"\n",
"更多帮助见 `docs/FAQ.md` 和 `docs/troubleshooting.md`。"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

245
notebooks/01_copy_kernel.ipynb
View File

@ -1,245 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 01 — Copy 算子\n",
"\n",
"## 学习目标\n",
"\n",
"1. 理解 **grid-stride loop**：为什么它可以处理任意大小的 tensor\n",
"2. 完成 NVIDIA CUDA kernel 的 TODO\n",
"3. 完成 TileLang kernel 的 TODO\n",
"4. 验证正确性 + 查看 benchmark 结果\n",
"\n",
"Copy 是最简单的算子——把数据从 `src` 搬到 `dst`，不涉及任何计算。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 概念导入：Grid-Stride Loop\n",
"\n",
"GPU 有成千上万个线程，但 tensor 大小可能更大。Grid-stride loop 让每个线程处理多个元素：\n",
"\n",
"```python\n",
"# 伪代码\n",
"for i in range(thread_idx, N, grid_total_threads):\n",
" dst[i] = src[i]\n",
"```\n",
"\n",
"- `grid_total_threads = gridDim.x * blockDim.x`（所有线程总数）\n",
"- 步长 = grid_total_threads，每个线程“跨步”处理\n",
"- 循环条件 `i < N` 保证不越界\n",
"\n",
"详细图示见 [docs/diagrams/grid-stride-loop.md](../docs/diagrams/grid-stride-loop.md)。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## PyTorch 参考实现\n",
"\n",
"先看看用 PyTorch 怎么做 copy——训练营的目标就是实现和它一样的功能："
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"# PyTorch 版 copy\n",
"def pytorch_copy(src: torch.Tensor) -> torch.Tensor:\n",
" return src.clone()\n",
"\n",
"# 测试\n",
"src = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
"expected = pytorch_copy(src)\n",
"print(f\"输入: {src[:5]}\")\n",
"print(f\"输出: {expected[:5]}\")\n",
"print(f\"一致: {torch.allclose(src, expected)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## NVIDIA CUDA Kernel\n",
"\n",
"打开 `ops/copy/nvidia/kernel.cuh`，你会看到这个骨架。请填写 `TODO` 部分："
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 以下代码来自 ops/copy/nvidia/kernel.cuh，仅供参考——请直接编辑该文件\n",
"\n",
"KERNEL_SKELETON = \"\"\"\n",
"template <typename T>\n",
"__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {\n",
" // TODO: implement a grid-stride loop copy kernel.\n",
" //\n",
" // Suggested steps:\n",
" // 1. Compute the global thread index.\n",
" // int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;\n",
" // 2. Compute the grid-wide stride.\n",
" // int64_t stride = gridDim.x * blockDim.x;\n",
" // 3. Loop over i = idx; i < n; i += stride.\n",
" // 4. Copy src[i] to dst[i].\n",
"}\n",
"\"\"\"\n",
"print(KERNEL_SKELETON)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 你的任务\n",
"\n",
"编辑 `ops/copy/nvidia/kernel.cuh`，完成 TODO 后，重新构建并运行下面的验证："
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import subprocess, sys\n",
"\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_copy.py\", \"-v\",\n",
" \"--backend\", \"nvidia\"\n",
"], capture_output=True, text=True)\n",
"\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ NVIDIA copy kernel 测试全部通过！\")\n",
"else:\n",
" print(result.stderr[-500:] if result.stderr else \"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TileLang Kernel\n",
"\n",
"打开 `ops/copy/tilelang/kernel.py`，填写 TODO："
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"TILELANG_SKELETON = \"\"\"\n",
"import tilelang\n",
"import tilelang.language as T\n",
"\n",
"@tilelang.jit\n",
"def copy_kernel(src, BLOCK_N: int, dtype):\n",
" N = T.const(\"N\")\n",
" src: T.Tensor((N,), dtype)\n",
" out = T.empty((N,), dtype)\n",
"\n",
" # TODO: implement a tile-wise copy kernel.\n",
" #\n",
" # Suggested steps:\n",
" # 1. Launch one TileLang kernel over the N // BLOCK_N tiles.\n",
" # for i in T.Parallel(N // BLOCK_N):\n",
" # 2. Use T.copy to move one tile from src to out.\n",
" # tile = T.copy(src[i * BLOCK_N : (i + 1) * BLOCK_N])\n",
" # out[i * BLOCK_N : (i + 1) * BLOCK_N] = tile\n",
"\n",
" return out\n",
"\"\"\"\n",
"print(TILELANG_SKELETON)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 TileLang 版本\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_copy.py\", \"-v\",\n",
" \"--backend\", \"tilelang\"\n",
"], capture_output=True, text=True)\n",
"\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ TileLang copy kernel 测试全部通过！\")\n",
"else:\n",
" print(result.stderr[-500:] if result.stderr else \"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 对比：你的实现 vs PyTorch"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from operator_runtime import copy\n",
"\n",
"src = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
"my_out = copy(src, backend=\"nvidia\")\n",
"torch_out = src.clone()\n",
"\n",
"print(f\"输入: {src[:5]}\")\n",
"print(f\"我的 kernel: {my_out[:5]}\")\n",
"print(f\"PyTorch: {torch_out[:5]}\")\n",
"print(f\"一致: {torch.allclose(my_out, torch_out)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 检查清单\n",
"\n",
"- [ ] 理解 grid-stride loop 原理\n",
"- [ ] 理解 global memory 合并访问（coalesced access）\n",
"- [ ] NVIDIA `kernel.cuh` TODO 完成\n",
"- [ ] TileLang `kernel.py` TODO 完成\n",
"- [ ] 两种后端测试全部通过\n",
"- [ ] benchmark 跑通，带宽数据合理"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

239
notebooks/02_vector_add_kernel.ipynb
View File

@ -1,239 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 02 — Vector Add 算子\n",
"\n",
"## 学习目标\n",
"\n",
"1. 理解逐元素并行：每个线程独立处理一对 `(a[i], b[i])`\n",
"2. 理解 tile-level 并行（`T.Parallel` vs `T.Serial` 的区别）\n",
"3. 完成 NVIDIA + TileLang 两个 kernel 的 TODO\n",
"\n",
"Vector add 和 copy 结构相似，但引入了**双输入**和**逐元素计算**。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 概念导入：逐元素并行\n",
"\n",
"```\n",
"a = [a0, a1, a2, ..., aN-1]\n",
"b = [b0, b1, b2, ..., bN-1]\n",
" ↓ 逐元素相加\n",
"c = [a0+b0, a1+b1, a2+b2, ..., aN-1+bN-1]\n",
"```\n",
"\n",
"每个 `c[i] = a[i] + b[i]` 完全独立——不需要线程间通信。\n",
"Grid-stride loop 同样适用。\n",
"\n",
"详细图示见 [docs/diagrams/thread-grid-layout.md](../docs/diagrams/thread-grid-layout.md)。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## PyTorch 参考实现"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"a = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
"b = torch.randn_like(a)\n",
"expected = a + b\n",
"\n",
"print(f\"a[:5]: {a[:5]}\")\n",
"print(f\"b[:5]: {b[:5]}\")\n",
"print(f\"a+b[:5]: {expected[:5]}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## NVIDIA CUDA Kernel\n",
"\n",
"打开 `ops/vector_add/nvidia/kernel.cuh`，你需要完成两部分：\n",
"\n",
"1. **`add_values<T>` 辅助函数**：返回 `a + b`（泛型版本和 half 特化版）\n",
"2. **`vector_add_contiguous_kernel`**：grid-stride loop 遍历，`out[i] = add_values(a[i], b[i])`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"KERNEL_SKELETON = \"\"\"\n",
"template <typename T>\n",
"__device__ T add_values(T a, T b) {\n",
" // TODO: return the elementwise sum for generic types.\n",
"}\n",
"\n",
"template <>\n",
"__device__ inline half add_values<half>(half a, half b) {\n",
" // TODO: return the half-precision elementwise sum.\n",
" // 提示：half 的加法用 __hadd(a, b)\n",
"}\n",
"\n",
"template <typename T>\n",
"__global__ void vector_add_contiguous_kernel(\n",
" T *out, const T *a, const T *b, int64_t n) {\n",
" // TODO: grid-stride loop, out[i] = add_values(a[i], b[i])\n",
"}\n",
"\"\"\"\n",
"print(KERNEL_SKELETON)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 NVIDIA 版本\n",
"import subprocess, sys\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_vector_add.py\", \"-v\",\n",
" \"--backend\", \"nvidia\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ NVIDIA vector_add kernel 测试全部通过！\")\n",
"else:\n",
" print(result.stderr[-500:] if result.stderr else \"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TileLang Kernel\n",
"\n",
"打开 `ops/vector_add/tilelang/kernel.py`，注意这里使用了 `T.Parallel`。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"TILELANG_SKELETON = \"\"\"\n",
"@tilelang.jit\n",
"def vector_add_kernel(a, b, BLOCK_N: int, dtype):\n",
" N = T.const(\"N\")\n",
" a: T.Tensor((N,), dtype)\n",
" b: T.Tensor((N,), dtype)\n",
" out = T.empty((N,), dtype)\n",
"\n",
" # TODO: implement a tile-wise vector add kernel.\n",
" #\n",
" # Suggested steps:\n",
" # 1. T.Parallel(N // BLOCK_N) to iterate over tiles\n",
" # 2. Compute base = tile_idx * BLOCK_N\n",
" # 3. Inner T.Parallel(BLOCK_N): out[base+i] = a[base+i] + b[base+i]\n",
"\n",
" return out\n",
"\"\"\"\n",
"print(TILELANG_SKELETON)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### `T.Parallel` vs `T.Serial`\n",
"\n",
"| | `T.Parallel` | `T.Serial` |\n",
"|--|-------------|------------|\n",
"| 含义 | 循环迭代可以并行 | 循环迭代必须顺序执行 |\n",
"| 何时用 | 迭代之间无数据依赖 | 迭代之间有数据依赖（如累加器） |\n",
"| vector_add | ✓ 每个 `(a[i], b[i])` 独立 | N/A |"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 TileLang 版本\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_vector_add.py\", \"-v\",\n",
" \"--backend\", \"tilelang\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ TileLang vector_add kernel 测试全部通过！\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 对比验证"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from operator_runtime import vector_add\n",
"\n",
"a = torch.randn(1024, device=\"cuda\", dtype=torch.float32)\n",
"b = torch.randn_like(a)\n",
"\n",
"my_out = vector_add(a, b, backend=\"nvidia\")\n",
"torch_out = a + b\n",
"\n",
"print(f\"我的 kernel: {my_out[:5]}\")\n",
"print(f\"PyTorch: {torch_out[:5]}\")\n",
"print(f\"一致: {torch.allclose(my_out, torch_out)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 检查清单\n",
"\n",
"- [ ] 理解逐元素并行的线程网格布局\n",
"- [ ] 理解 `T.Parallel` vs `T.Serial` 的区别\n",
"- [ ] NVIDIA `kernel.cuh` TODO 完成（add_values + kernel 函数）\n",
"- [ ] TileLang `kernel.py` TODO 完成\n",
"- [ ] 两种后端测试全部通过\n",
"- [ ] benchmark 跑通"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

273
notebooks/03_reduce_sum_kernel.ipynb
View File

@ -1,273 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 03 — Reduce Sum 算子\n",
"\n",
"## 学习目标\n",
"\n",
"1. 理解 shared memory 树形归约原理\n",
"2. 理解 `__syncthreads()` 的使用时机和条件分支限制\n",
"3. 理解 `T.Serial` 在归约场景中的作用\n",
"4. 完成 NVIDIA + TileLang 两个 kernel 的 TODO\n",
"\n",
"Reduce sum 是第一个需要**线程间通信**的算子——从 copy/vector_add 的“各自为战“跨越到“协同计算“。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 概念导入：树形归约\n",
"\n",
"目标：把一行 N 个元素求和为 1 个值。每个 block 处理一行：\n",
"\n",
"```\n",
"Step 0: [3, 1, 7, 0, 4, 1, 6, 3] ← 线程各自的部分和存入 shared memory\n",
"Step 1: [7, 2, 13, 3] ← stride=4, 相邻配对相加\n",
"Step 2: [20, 5] ← stride=2\n",
"Step 3: [25] ← stride=1, 最终结果\n",
"```\n",
"\n",
"每步后必须 `__syncthreads()` ——因为下一步要读上一步别人写的数据。\n",
"\n",
"详细图示见 [docs/diagrams/tree-reduction.md](../docs/diagrams/tree-reduction.md)。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## PyTorch 参考实现"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"src = torch.randn(32, 128, device=\"cuda\", dtype=torch.float32)\n",
"expected = torch.sum(src, dim=1) # 对每行求和\n",
"\n",
"print(f\"输入 shape: {src.shape}\")\n",
"print(f\"输出 shape: {expected.shape}\")\n",
"print(f\"第一行: src[0, :5] = {src[0, :5]}\")\n",
"print(f\"第一行和: {expected[0]:.6f}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## NVIDIA CUDA Kernel\n",
"\n",
"打开 `ops/reduce_sum/nvidia/kernel.cuh`。这是第一个使用 shared memory 的 kernel。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"KERNEL_SKELETON = \"\"\"\n",
"__global__ void reduce_sum_rowwise_kernel(\n",
" float *out, const float *in, int64_t rows, int64_t cols) {\n",
"\n",
" // TODO: implement a row-wise reduce_sum kernel with shared memory.\n",
" //\n",
" // Suggested steps:\n",
" // 1. Use one block per row.\n",
" // int row = blockIdx.x;\n",
" // 2. Each thread accumulates its columns:\n",
" // float sum = 0;\n",
" // for (int c = threadIdx.x; c < cols; c += blockDim.x)\n",
" // sum += in[row * cols + c];\n",
" // 3. Store partial sum to shared memory, then __syncthreads().\n",
" // 4. Tree reduction:\n",
" // for (int s = blockDim.x/2; s > 0; s >>= 1) {\n",
" // if (threadIdx.x < s) smem[tid] += smem[tid + s];\n",
" // __syncthreads();\n",
" // }\n",
" // 5. Thread 0 writes smem[0] to out[row].\n",
"}\n",
"\"\"\"\n",
"print(KERNEL_SKELETON)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### ⚠️ 关键陷阱：`__syncthreads()` 不能在条件分支内\n",
"\n",
"```cuda\n",
"// 错误——死锁！\n",
"if (tid < s) {\n",
" smem[tid] += smem[tid + s];\n",
" __syncthreads(); // 只有 tid < s 的线程执行同步\n",
"}\n",
"\n",
"// 正确——所有线程都到达同步点\n",
"if (tid < s) {\n",
" smem[tid] += smem[tid + s];\n",
"}\n",
"__syncthreads(); // 全部线程到位\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 NVIDIA 版本\n",
"import subprocess, sys\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_reduce_sum.py\", \"-v\",\n",
" \"--backend\", \"nvidia\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ NVIDIA reduce_sum kernel 测试全部通过！\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TileLang Kernel\n",
"\n",
"打开 `ops/reduce_sum/tilelang/kernel.py`。TileLang 版本用 `T.Serial` + `T.reduce_sum` 替代手动树形归约。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"TILELANG_SKELETON = \"\"\"\n",
"@tilelang.jit\n",
"def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):\n",
" N, M = T.const(\"N, M\")\n",
" dtype = T.float32\n",
" src: T.Tensor((N, M), dtype)\n",
" out = T.empty((N,), dtype)\n",
"\n",
" # TODO: implement a tiled row-wise reduce_sum kernel.\n",
" #\n",
" # Key insight: outer loop uses T.Serial (accumulate state),\n",
" # inner loop uses T.reduce_sum (parallel reduction within tile).\n",
"\n",
" return out\n",
"\"\"\"\n",
"print(TILELANG_SKELETON)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 为什么外层用 `T.Serial`？\n",
"\n",
"分块累加时，每个 chunk 的求和结果需要**累加到同一个累加器**。这是有状态依赖的——后一个 chunk 必须在前一个完成之后才能累加。所以列方向的循环是 `T.Serial`。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 TileLang 版本\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_reduce_sum.py\", \"-v\",\n",
" \"--backend\", \"tilelang\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ TileLang reduce_sum kernel 测试全部通过！\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 对比验证"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from operator_runtime import reduce_sum\n",
"\n",
"src = torch.randn(32, 128, device=\"cuda\", dtype=torch.float32)\n",
"my_out = reduce_sum(src, dim=1, backend=\"nvidia\")\n",
"torch_out = torch.sum(src, dim=1)\n",
"\n",
"print(f\"我的 kernel[:5]: {my_out[:5]}\")\n",
"print(f\"PyTorch[:5]: {torch_out[:5]}\")\n",
"print(f\"最大误差: {(my_out - torch_out).abs().max().item():.2e}\")\n",
"print(f\"一致: {torch.allclose(my_out, torch_out, atol=1e-5, rtol=1e-5)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 进阶：Bank Conflict 优化\n",
"\n",
"如果你已经跑通基础版本，可以尝试优化 bank conflict：\n",
"\n",
"```cuda\n",
"// 加 padding 错开 bank 访问\n",
"__shared__ float smem[BLOCK_SIZE + PADDING];\n",
"// 或者归约从 stride=2 开始而非 stride=1\n",
"```\n",
"\n",
"对比优化前后的 benchmark 带宽数据。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 检查清单\n",
"\n",
"- [ ] 理解 shared memory 树形归约原理\n",
"- [ ] 理解 `__syncthreads()` 的使用时机和条件分支限制\n",
"- [ ] 理解 `T.Serial` 在归约场景中的作用\n",
"- [ ] NVIDIA `kernel.cuh` TODO 完成\n",
"- [ ] TileLang `kernel.py` TODO 完成\n",
"- [ ] 两种后端测试全部通过\n",
"- [ ] 尝试优化 bank conflict\n",
"- [ ] benchmark 跑通"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

260
notebooks/04_softmax_kernel.ipynb
View File

@ -1,260 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 04 — Softmax 算子\n",
"\n",
"## 学习目标\n",
"\n",
"1. 理解数值稳定性：为什么需要减 max\n",
"2. 理解 online softmax 算法（一遍扫描 vs 三趟扫描）\n",
"3. 理解 TileLang 中为什么用 `exp2`/`log2` 替代 `exp`/`log`\n",
"4. 完成 NVIDIA + TileLang 两个 kernel 的 TODO\n",
"\n",
"Softmax 是四个算子中最复杂的——既需要线程间通信（同步归约），又有数值稳定性陷阱。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 概念导入：Softmax 的数值稳定性\n",
"\n",
"### 朴素公式\n",
"\n",
"$$\\text{softmax}(x_i) = \\frac{e^{x_i}}{\\sum_j e^{x_j}}$$\n",
"\n",
"### 问题\n",
"\n",
"$e^{88.7} \\approx 1.6 \\times 10^{38}$，接近 FP32 上限。如果 $x_i = 100$，则 $e^{100}$ 溢出为 inf，inf/inf = NaN。\n",
"\n",
"### 解决：减最大值\n",
"\n",
"$$\\text{softmax}(x_i) = \\frac{e^{x_i - \\max(x)}}{\\sum_j e^{x_j - \\max(x)}}$$\n",
"\n",
"分子分母同除 $e^{\\max(x)}$，数学结果不变，但 $e^{x_i - \\max(x)}$ 最大为 $e^0 = 1$，永不溢出。\n",
"\n",
"详细图示见 [docs/diagrams/softmax-pipeline.md](../docs/diagrams/softmax-pipeline.md)。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## PyTorch 参考实现"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"src = torch.randn(32, 128, device=\"cuda\", dtype=torch.float32)\n",
"expected = torch.softmax(src, dim=1)\n",
"\n",
"print(f\"输入 shape: {src.shape}\")\n",
"print(f\"输出 shape: {expected.shape}\")\n",
"print(f\"每行和: {expected.sum(dim=1)[:5]}\") # 应该全为 1.0\n",
"\n",
"# 演示溢出问题\n",
"big = torch.tensor([100.0, 200.0, 300.0], device=\"cuda\")\n",
"naive = torch.exp(big) / torch.exp(big).sum()\n",
"print(f\"\\n朴素 softmax([100, 200, 300]): {naive}\") # NaN!\n",
"\n",
"stable = torch.softmax(big, dim=0)\n",
"print(f\"稳定 softmax([100, 200, 300]): {stable}\") # [0, 0, 1]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## NVIDIA CUDA Kernel（三趟扫描）\n",
"\n",
"打开 `ops/softmax/nvidia/kernel.cuh`。基础版本使用三趟扫描：\n",
"\n",
"1. Pass 1: 求行最大值 `max_val`\n",
"2. Pass 2: 写 `exp(x - max)` 到 out，同时累加 sum\n",
"3. Pass 3: 除 sum 归一化"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"KERNEL_SKELETON = \"\"\"\n",
"__global__ void softmax_rowwise_kernel(\n",
" float *out, const float *in, int64_t rows, int64_t cols) {\n",
"\n",
" // TODO: implement a numerically stable row-wise softmax kernel.\n",
" //\n",
" // Suggested steps:\n",
" // 1. One block per row.\n",
" // 2. Pass 1: find row max via shared memory reduction.\n",
" // 3. Pass 2: exp(x - max) → out, accumulate sum via reduction.\n",
" // 4. Pass 3: out[i] /= sum.\n",
"}\n",
"\"\"\"\n",
"print(KERNEL_SKELETON)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 为什么是三趟？\n",
"\n",
"因为 softmax 有两个需要全局信息才能计算的步骤：\n",
"- 求 max 需要全局信息（所有元素的最大值）\n",
"- 求 sum 需要全局信息（所有 exp(x-max) 的和）\n",
"\n",
"这两个全局信息都通过 shared memory 归约获得，每一步都需要 block 内同步。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 NVIDIA 版本\n",
"import subprocess, sys\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_softmax.py\", \"-v\",\n",
" \"--backend\", \"nvidia\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ NVIDIA softmax kernel 测试全部通过！\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TileLang Kernel（Online Softmax）\n",
"\n",
"打开 `ops/softmax/tilelang/kernel.py`。TileLang 版本使用 **online softmax**——只需两趟扫描。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"TILELANG_SKELETON = \"\"\"\n",
"@tilelang.jit\n",
"def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):\n",
" log2_e = 1.44269504 # log2(e) for exp/log2 conversion\n",
" N, M = T.const(\"N, M\")\n",
" dtype = T.float32\n",
" src: T.Tensor((N, M), dtype)\n",
" out = T.empty((N, M), dtype)\n",
"\n",
" # TODO: implement a tiled row-wise softmax with online algorithm.\n",
" #\n",
" # Two-pass approach:\n",
" # Pass 1: scroll through column tiles, update running log-sum-exp.\n",
" # Pass 2: scroll again, normalize each tile with final lse.\n",
"\n",
" return out\n",
"\"\"\"\n",
"print(TILELANG_SKELETON)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 为什么用 `exp2`/`log2`？\n",
"\n",
"GPU 硬件对 2 的幂运算 (`2^x`) 有专门的快速指令，比自然指数 (`e^x`) 快。\n",
"\n",
"转换公式：$\\exp(x) = 2^{x \\cdot \\log_2(e)}$\n",
"\n",
"其中 $\\log_2(e) \\approx 1.44269504$（代码中的 `log2_e` 常量）。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 验证 TileLang 版本\n",
"result = subprocess.run([\n",
" sys.executable, \"-m\", \"pytest\",\n",
" \"tests/op_tests/test_softmax.py\", \"-v\",\n",
" \"--backend\", \"tilelang\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode == 0:\n",
" print(\"✓ TileLang softmax kernel 测试全部通过！\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 对比验证"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from operator_runtime import softmax\n",
"\n",
"src = torch.randn(32, 128, device=\"cuda\", dtype=torch.float32)\n",
"my_out = softmax(src, dim=1, backend=\"nvidia\")\n",
"torch_out = torch.softmax(src, dim=1)\n",
"\n",
"print(f\"我的 kernel 行和: {my_out.sum(dim=1)[:5]}\")\n",
"print(f\"PyTorch 行和: {torch_out.sum(dim=1)[:5]}\")\n",
"print(f\"最大误差: {(my_out - torch_out).abs().max().item():.2e}\")\n",
"print(f\"一致: {torch.allclose(my_out, torch_out, atol=1e-5, rtol=1e-5)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 检查清单\n",
"\n",
"- [ ] 理解数值稳定性问题：为什么需要减 max\n",
"- [ ] 理解 online softmax 算法（一遍扫描 vs 三趟扫描）\n",
"- [ ] 理解 log-sum-exp 的滚动更新逻辑\n",
"- [ ] 理解 TileLang 中为什么用 `exp2`/`log2` 替代 `exp`/`log`\n",
"- [ ] NVIDIA `kernel.cuh` TODO 完成\n",
"- [ ] TileLang `kernel.py` TODO 完成\n",
"- [ ] 两种后端测试全部通过\n",
"- [ ] 性能与 PyTorch 参考实现对比\n",
"- [ ] 尝试 warp-level 优化"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

213
notebooks/05_benchmark_and_profile.ipynb
View File

@ -1,213 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 05 — Benchmark 与 Profiling 入门\n",
"\n",
"## 学习目标\n",
"\n",
"1. 运行 benchmark 获取 kernel 性能数据\n",
"2. 理解关键指标：带宽利用率、延迟\n",
"3. 学会对比自己的实现与 PyTorch 的性能差异\n",
"4. 初识 profiling 工具（nsys / ncu）"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: 运行 Benchmark\n",
"\n",
"intro-ops 的 `tests/bench/` 目录下每个算子都有 benchmark 脚本。"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import subprocess, sys\n",
"\n",
"# 对 copy 算子同时跑正确性和 benchmark\n",
"result = subprocess.run([\n",
" sys.executable, \"tests/run_ops.py\",\n",
" \"--op\", \"copy\",\n",
" \"--backend\", \"nvidia\",\n",
" \"--mode\", \"all\"\n",
"], capture_output=True, text=True)\n",
"print(result.stdout)\n",
"if result.returncode != 0:\n",
" print(result.stderr[-500:] if result.stderr else \"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: 理解 Benchmark 指标\n",
"\n",
"### 带宽利用率 (Bandwidth Utilization)\n",
"\n",
"$$\\text{利用率} = \\frac{\\text{实际带宽}}{\\text{峰值带宽}} \\times 100\\%$$\n",
"\n",
"- **实际带宽** = 数据量 / 执行时间（单位 GB/s）\n",
"- **峰值带宽** = GPU 理论最大内存带宽（如 L40 = 864 GB/s）\n",
"- copy 是纯内存搬运，瓶颈在带宽 → 追求高带宽利用率\n",
"- reduce_sum/softmax 有计算量，瓶颈可能在计算 → 对比算术强度\n",
"\n",
"### 关键指标\n",
"\n",
"| 指标 | 含义 | 好/坏参考 |\n",
"|------|------|----------|\n",
"| Bandwidth (GB/s) | 实际内存带宽 | 越高越好 |\n",
"| Bandwidth Utilization | 带宽利用率 | >80% 优秀，<50% 需优化 |\n",
"| Latency (ms) | 单次调用耗时 | 越低越好 |\n",
"| GFLOPS | 计算吞吐 | 越高越好 |"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: 对比所有算子"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"from operator_runtime import copy, vector_add, reduce_sum, softmax\n",
"import time\n",
"\n",
"def bench_op(name, fn, *args, warmup=10, iters=100):\n",
" for _ in range(warmup):\n",
" fn(*args)\n",
" torch.cuda.synchronize()\n",
" start = time.perf_counter()\n",
" for _ in range(iters):\n",
" fn(*args)\n",
" torch.cuda.synchronize()\n",
" elapsed_ms = (time.perf_counter() - start) / iters * 1000\n",
" print(f\"{name:15s}: {elapsed_ms:8.3f} ms\")\n",
" return elapsed_ms\n",
"\n",
"# FP32 测试\n",
"print(\"=== FP32 Benchmarks ===\")\n",
"src = torch.randn(1024 * 1024, device=\"cuda\", dtype=torch.float32)\n",
"bench_op(\"copy\", lambda x: copy(x, backend=\"nvidia\"), src)\n",
"\n",
"a = torch.randn(1024 * 1024, device=\"cuda\", dtype=torch.float32)\n",
"b = torch.randn_like(a)\n",
"bench_op(\"vector_add\", lambda x, y: vector_add(x, y, backend=\"nvidia\"), a, b)\n",
"\n",
"mat = torch.randn(1024, 1024, device=\"cuda\", dtype=torch.float32)\n",
"bench_op(\"reduce_sum\", lambda x: reduce_sum(x, dim=1, backend=\"nvidia\"), mat)\n",
"\n",
"mat2 = torch.randn(1024, 1024, device=\"cuda\", dtype=torch.float32)\n",
"bench_op(\"softmax\", lambda x: softmax(x, dim=1, backend=\"nvidia\"), mat2)\n",
"\n",
"print(\"\\n=== PyTorch Baseline ===\")\n",
"bench_op(\"torch.clone\", lambda x: x.clone(), src)\n",
"bench_op(\"torch.add\", lambda x, y: x + y, a, b)\n",
"bench_op(\"torch.sum\", lambda x: torch.sum(x, dim=1), mat)\n",
"bench_op(\"torch.softmax\", lambda x: torch.softmax(x, dim=1), mat2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 4: Profiling 入门\n",
"\n",
"### NVIDIA Nsight Systems (nsys)\n",
"\n",
"```bash\n",
"# 对 copy 算子做 profiling\n",
"nsys profile --stats=true \\\n",
" python tests/bench/copy.py --backend nvidia\n",
"\n",
"# 输出关键指标：\n",
"# - CUDA Kernel 执行时间\n",
"# - 内存传输时间\n",
"# - Occupancy\n",
"```\n",
"\n",
"### NVIDIA Nsight Compute (ncu)\n",
"\n",
"```bash\n",
"# 详细 kernel 分析\n",
"ncu --set full \\\n",
" python tests/bench/copy.py --backend nvidia\n",
"\n",
"# 关注：\n",
"# - Memory Throughput\n",
"# - Compute Throughput\n",
"# - Occupancy\n",
"# - Shared Memory Configuration\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 5: 性能分析检查清单\n",
"\n",
"当你的 kernel 正确运行但性能不理想时，按此顺序排查：\n",
"\n",
"### copy / vector_add\n",
"\n",
"1. 带宽利用率低于 60%？→ 检查合并访问（coalesced access）\n",
"2. grid/block 尺寸是否合理？→ 尝试 block_size=256/512\n",
"3. stride 是否为 1？\n",
"\n",
"### reduce_sum\n",
"\n",
"1. 带宽利用率低于 50%？→ 可能有 bank conflict\n",
"2. 尝试加 padding：`__shared__ float smem[BLOCK_SIZE + 1]`\n",
"3. 尝试 warp-level shuffle 替代 shared memory 归约\n",
"\n",
"### softmax\n",
"\n",
"1. 三趟扫描的带宽开销 = N × 3 次全局内存读写 → 能否减为两趟？\n",
"2. 用 online softmax 减少一次全局内存写\n",
"3. 用 `__shfl_down_sync` 在 warp 内做归约，节省 shared memory 带宽"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 参考：L40 GPU 峰值数据\n",
"\n",
"| 指标 | L40 | A100 |\n",
"|------|-----|------|\n",
"| 峰值带宽 | 864 GB/s | 2039 GB/s |\n",
"| FP32 TFLOPS | 90.5 | 19.5 |\n",
"| SM 数量 | 142 | 108 |\n",
"| Shared Memory / SM | 128 KB | 164 KB |\n",
"\n",
"> 你 GPU 的峰值数据见 `tests/perf_profiles/local_gpu.yaml`，benchmark 会自动读取该文件做对比。"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

59
ops/CMakeLists.txt
View File

@ -1,59 +0,0 @@
# 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
${CMAKE_CURRENT_SOURCE_DIR}/common/status.cc
)
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()
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()
set(CAMP_OPERATOR_SOURCES ${CAMP_COMMON_SOURCES} ${CAMP_NVIDIA_SOURCES} ${CAMP_METAX_SOURCES})
add_library(camp_ops SHARED ${CAMP_OPERATOR_SOURCES})
target_include_directories(camp_ops
PUBLIC
"${CMAKE_CURRENT_SOURCE_DIR}/../include"
PRIVATE
"${CMAKE_CURRENT_SOURCE_DIR}/.."
)
if(CAMP_USE_CUTE)
target_include_directories(camp_ops PRIVATE ${CAMP_CUTE_TARGET_INCLUDE_DIRS})
target_compile_definitions(camp_ops PRIVATE CAMP_ENABLE_CUTE=1)
target_compile_options(camp_ops PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
endif()
if(CAMP_ENABLE_METAX)
target_include_directories(camp_ops PRIVATE ${MetaX_INCLUDE_DIRS})
target_link_libraries(camp_ops PRIVATE ${MetaX_LIBRARIES})
target_compile_options(camp_ops PRIVATE ${CAMP_METAX_COMPILE_OPTIONS})
set_target_properties(camp_ops PROPERTIES LINKER_LANGUAGE MACA)
endif()
target_compile_definitions(camp_ops PRIVATE
$<$<BOOL:${CAMP_ENABLE_NVIDIA}>:CAMP_ENABLE_NVIDIA=1>
$<$<BOOL:${CAMP_ENABLE_METAX}>:CAMP_ENABLE_METAX=1>
)
if(CAMP_ENABLE_NVIDIA)
set_target_properties(camp_ops PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
endif()

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

@ -1,184 +0,0 @@
#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
View File

@ -1,184 +0,0 @@
#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

82
ops/common/status.cc
View File

@ -1,82 +0,0 @@
#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:
return "success";
case OPRT_ERR_INVALID_ARG:
return "invalid argument";
case OPRT_ERR_UNSUPPORTED_DTYPE:
return "unsupported dtype";
case OPRT_ERR_RUNTIME:
return "runtime error";
case OPRT_ERR_INSUFFICIENT_WORKSPACE:
return "insufficient workspace";
case OPRT_ERR_NOT_SUPPORTED:
return "not supported";
default:
return "unknown 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;
}

86
ops/common/tilelang/eager_copy.py
View File

@ -1,86 +0,0 @@
from dataclasses import dataclass
from functools import lru_cache
import tilelang
import tilelang.language as T
import torch
def _tl_dtype(dtype: torch.dtype):
if dtype is torch.float16:
return T.float16
if dtype is torch.float32:
return T.float32
raise TypeError(f"unsupported TileLang dtype: {dtype}")
@tilelang.jit
def _copy_eager_kernel(src, BLOCK_N: int, dtype):
N = T.const("N")
src: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
T.copy(
src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
)
return out
def _block_n(n: int) -> int:
return 1024 if n % 1024 == 0 else n
@lru_cache(maxsize=32)
def _compiled_copy_eager(n: int, block_n: int, dtype: torch.dtype):
return _copy_eager_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
@dataclass
class EagerCopyPrepared:
out: torch.Tensor
src: torch.Tensor
kernel: object
def run(self) -> None:
self.out.copy_(self.kernel(self.src))
def destroy(self) -> None:
pass
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def prepare_copy_eager(out: torch.Tensor, src: torch.Tensor) -> EagerCopyPrepared:
if out.shape != src.shape:
raise ValueError("eager copy expects matching shapes")
if out.dtype != src.dtype:
raise TypeError("eager copy expects matching dtypes")
if not out.is_cuda or not src.is_cuda:
raise ValueError("eager copy expects CUDA tensors")
if not out.is_contiguous() or not src.is_contiguous():
raise ValueError("eager copy v1 supports contiguous tensors only")
n = src.numel()
block_n = _block_n(n)
if n % block_n != 0:
raise ValueError("eager copy v1 requires N % BLOCK_N == 0")
kernel = _compiled_copy_eager(n, block_n, src.dtype)
return EagerCopyPrepared(out, src, kernel)
def copy_eager_(out: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
with prepare_copy_eager(out, src) as prepared:
prepared.run()
return out
def copy_eager(src: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(src)
return copy_eager_(out, src)

91
ops/common/tilelang/lazy_out_idx_copy.py
View File

@ -1,91 +0,0 @@
from dataclasses import dataclass
from functools import lru_cache
import tilelang
import tilelang.language as T
import torch
def _tl_dtype_str(dtype: torch.dtype) -> str:
if dtype is torch.float16:
return "float16"
if dtype is torch.float32:
return "float32"
raise TypeError(f"unsupported TileLang dtype: {dtype}")
@tilelang.jit(out_idx=[1])
def _copy_lazy_out_idx_kernel(n: int, block_n: int, dtype: str):
@T.prim_func
def main(src: T.Tensor((n,), dtype), out: T.Tensor((n,), dtype)):
with T.Kernel(n // block_n, threads=256) as pid_n:
T.copy(
src[pid_n * block_n : (pid_n + 1) * block_n],
out[pid_n * block_n : (pid_n + 1) * block_n],
)
return main
def _block_n(n: int) -> int:
return 1024 if n % 1024 == 0 else n
@lru_cache(maxsize=32)
def _compiled_copy_lazy_out_idx(n: int, block_n: int, dtype: torch.dtype):
return _copy_lazy_out_idx_kernel(n, block_n, _tl_dtype_str(dtype))
@dataclass
class LazyOutIdxCopyPrepared:
out: torch.Tensor
src: torch.Tensor
kernel: object
def run(self) -> None:
# out_idx makes TileLang allocate and return the output tensor. The
# training runtime keeps an out-variant API, so this adapter copies the
# result into the caller-owned output buffer.
self.out.copy_(self.kernel(self.src))
def destroy(self) -> None:
pass
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def prepare_copy_lazy_out_idx(out: torch.Tensor, src: torch.Tensor) -> LazyOutIdxCopyPrepared:
if out.shape != src.shape:
raise ValueError("lazy out_idx copy expects matching shapes")
if out.dtype != src.dtype:
raise TypeError("lazy out_idx copy expects matching dtypes")
if not out.is_cuda or not src.is_cuda:
raise ValueError("lazy out_idx copy expects CUDA tensors")
if not out.is_contiguous() or not src.is_contiguous():
raise ValueError("lazy out_idx copy v1 supports contiguous tensors only")
n = src.numel()
block_n = _block_n(n)
if n % block_n != 0:
raise ValueError("lazy out_idx copy v1 requires N % BLOCK_N == 0")
kernel = _compiled_copy_lazy_out_idx(n, block_n, src.dtype)
return LazyOutIdxCopyPrepared(out, src, kernel)
def copy_lazy_out_idx_(out: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
with prepare_copy_lazy_out_idx(out, src) as prepared:
prepared.run()
return out
def copy_lazy_out_idx(src: torch.Tensor) -> torch.Tensor:
n = src.numel()
block_n = _block_n(n)
if n % block_n != 0:
raise ValueError("lazy out_idx copy v1 requires N % BLOCK_N == 0")
kernel = _compiled_copy_lazy_out_idx(n, block_n, src.dtype)
return kernel(src)

4
ops/copy/metax/README.md
View File

@ -1,4 +0,0 @@
# Copy MetaX Backend
The copy MetaX backend exports the standard descriptor lifecycle with `_metax`
symbols and is built from `ops/copy/metax/copy_metax.maca`.

119
ops/copy/metax/copy_metax.maca
View File

@ -1,119 +0,0 @@
#include "operator_runtime/ops/copy.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/elementwise.h"
#include "operator_runtime/detail/tensor_checks.h"
#include <cuda_fp16.h>
namespace {
struct CopyDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t dst_view;
oprt_tensor_view_t src_view;
int64_t elements = 0;
const char *op_name() const override {
return "copy";
}
};
template <typename T>
__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
for (int64_t i = idx; i < n; i += stride) {
dst[i] = src[i];
}
}
template <typename T>
oprt_status_t launch_copy(const CopyDescriptor *desc, void *dst, const void *src, oprt_stream_t stream) {
constexpr int threads = 256;
int blocks = oprt::blocks_for(desc->elements, threads);
cudaStream_t s = oprt::as_cuda_stream(stream);
copy_contiguous_kernel<T><<<blocks, threads, 0, s>>>(
static_cast<T *>(dst), static_cast<const T *>(src), desc->elements);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_copy_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *dst,
const oprt_tensor_view_t *src) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(dst);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(src);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*dst, *src);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_shape(*dst, *src);
if (status != OPRT_SUCCESS) {
return status;
}
if (!oprt::elementwise_fast_path(*dst, *src)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new CopyDescriptor();
typed->dst_view = *dst;
typed->src_view = *src;
typed->elements = oprt::numel(*dst);
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_copy_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_copy(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *dst,
const void *src,
oprt_stream_t stream) {
if (desc == nullptr || dst == nullptr || src == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const CopyDescriptor *>(desc);
switch (typed->dst_view.dtype) {
case OPRT_DTYPE_F16:
return launch_copy<half>(typed, dst, src, stream);
case OPRT_DTYPE_F32:
return launch_copy<float>(typed, dst, src, stream);
default:
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_copy_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

119
ops/copy/nvidia/copy_cuda.cu
View File

@ -1,119 +0,0 @@
#include "operator_runtime/ops/copy.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/elementwise.h"
#include "operator_runtime/detail/tensor_checks.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "ops/copy/nvidia/kernel.cuh"
#ifdef CAMP_ENABLE_CUTE
#include "ops/copy/nvidia/cute_copy.cuh"
#endif
#include <cuda_fp16.h>
namespace {
struct CopyDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t dst_view;
oprt_tensor_view_t src_view;
int64_t elements = 0;
const char *op_name() const override {
return "copy";
}
};
template <typename T>
oprt_status_t launch_copy(const CopyDescriptor *desc, void *dst, const void *src, oprt_stream_t stream) {
constexpr int threads = 256;
int blocks = oprt::blocks_for(desc->elements, threads);
cudaStream_t s = oprt::as_cuda_stream(stream);
#ifdef CAMP_ENABLE_CUTE
oprt::copy::nvidia::launch_copy_contiguous_cute_l40<T>(
static_cast<T *>(dst), static_cast<const T *>(src), desc->elements, s);
#else
oprt::copy::nvidia::copy_contiguous_kernel<T><<<blocks, threads, 0, s>>>(
static_cast<T *>(dst), static_cast<const T *>(src), desc->elements);
#endif
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_copy_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *dst,
const oprt_tensor_view_t *src) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(dst);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(src);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*dst, *src);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_shape(*dst, *src);
if (status != OPRT_SUCCESS) {
return status;
}
if (!oprt::elementwise_fast_path(*dst, *src)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new CopyDescriptor();
typed->dst_view = *dst;
typed->src_view = *src;
typed->elements = oprt::numel(*dst);
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_copy_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_copy(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *dst,
const void *src,
oprt_stream_t stream) {
if (desc == nullptr || dst == nullptr || src == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const CopyDescriptor *>(desc);
switch (typed->dst_view.dtype) {
case OPRT_DTYPE_F16:
return launch_copy<half>(typed, dst, src, stream);
case OPRT_DTYPE_F32:
return launch_copy<float>(typed, dst, src, stream);
default:
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_copy_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

79
ops/copy/nvidia/cute_copy.cuh
View File

@ -1,79 +0,0 @@
#pragma once
#include <cute/tensor.hpp>
#include <cuda_runtime.h>
#include <cstdint>
namespace oprt::copy::nvidia {
template <typename T, int Threads, int ElementsPerAccess>
__global__ void copy_contiguous_cute_l40_kernel(T *dst, const T *src, int64_t n) {
using namespace cute;
constexpr int elements_per_block = Threads * ElementsPerAccess;
Tensor src_tensor = make_tensor(make_gmem_ptr(src), make_layout(make_shape(n)));
Tensor dst_tensor = make_tensor(make_gmem_ptr(dst), make_layout(make_shape(n)));
auto block_shape = make_shape(Int<elements_per_block>{});
auto block_coord = make_coord(blockIdx.x);
Tensor coords = make_identity_tensor(shape(src_tensor));
Tensor predicates = cute::lazy::transform(coords, [&](auto coord) {
return elem_less(coord, shape(src_tensor));
});
Tensor src_tile = local_tile(src_tensor, block_shape, block_coord);
Tensor dst_tile = local_tile(dst_tensor, block_shape, block_coord);
Tensor pred_tile = local_tile(predicates, block_shape, block_coord);
Layout thread_layout = make_layout(make_shape(Int<Threads>{}));
Layout value_layout = make_layout(make_shape(Int<ElementsPerAccess>{}));
using AccessType = uint_byte_t<sizeof(T) * ElementsPerAccess>;
using CopyOp = UniversalCopy<AccessType>;
using Atom = Copy_Atom<CopyOp, T>;
TiledCopy tiled_copy = make_tiled_copy(Atom{}, thread_layout, value_layout);
ThrCopy thread_copy = tiled_copy.get_thread_slice(threadIdx.x);
Tensor thread_src = thread_copy.partition_S(src_tile);
Tensor thread_dst = thread_copy.partition_D(dst_tile);
Tensor thread_pred = thread_copy.partition_S(pred_tile);
Tensor fragment = make_fragment_like(thread_src);
copy_if(tiled_copy, thread_pred, thread_src, fragment);
copy_if(tiled_copy, thread_pred, fragment, thread_dst);
}
template <typename T>
inline void launch_copy_contiguous_cute_l40(T *dst, const T *src, int64_t n, cudaStream_t stream) {
constexpr int threads = 256;
constexpr int bytes_per_access = 16;
constexpr int elements_per_access = bytes_per_access / static_cast<int>(sizeof(T));
static_assert(elements_per_access > 0, "copy element type is larger than the vector access");
if (n <= 0) {
return;
}
auto src_addr = reinterpret_cast<std::uintptr_t>(src);
auto dst_addr = reinterpret_cast<std::uintptr_t>(dst);
bool aligned = (src_addr % bytes_per_access == 0) && (dst_addr % bytes_per_access == 0);
if (aligned) {
constexpr int elements_per_block = threads * elements_per_access;
int64_t blocks64 = (n + elements_per_block - 1) / elements_per_block;
int blocks = static_cast<int>(blocks64 > 0 ? blocks64 : 1);
copy_contiguous_cute_l40_kernel<T, threads, elements_per_access><<<blocks, threads, 0, stream>>>(
dst, src, n);
} else {
int64_t blocks64 = (n + threads - 1) / threads;
int blocks = static_cast<int>(blocks64 > 0 ? blocks64 : 1);
copy_contiguous_cute_l40_kernel<T, threads, 1><<<blocks, threads, 0, stream>>>(
dst, src, n);
}
}
} // namespace oprt::copy::nvidia

19
ops/copy/nvidia/kernel.cuh
View File

@ -1,19 +0,0 @@
#pragma once
#include <cuda_fp16.h>
#include <stdint.h>
namespace oprt::copy::nvidia {
template <typename T>
__global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
// TODO: implement a grid-stride loop copy kernel.
//
// Suggested steps:
// 1. Compute the global thread index.
// 2. Compute the grid-wide stride.
// 3. Loop over i = idx; i < n; i += stride.
// 4. Copy src[i] to dst[i].
}
} // namespace oprt::copy::nvidia

73
ops/copy/tilelang/copy_tl.py
View File

@ -1,73 +0,0 @@
from __future__ import annotations
from functools import lru_cache
from dataclasses import dataclass
import tilelang.language as T
import torch
from ops.copy.tilelang.kernel import copy_kernel
def _tl_dtype(dtype: torch.dtype):
if dtype is torch.float16:
return T.float16
if dtype is torch.float32:
return T.float32
raise TypeError(f"unsupported TileLang dtype: {dtype}")
def _block_n(n: int) -> int:
return 1024 if n % 1024 == 0 else n
@lru_cache(maxsize=32)
def _compiled_copy(n: int, block_n: int, dtype: torch.dtype):
return copy_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
@dataclass
class TileLangCopyPrepared:
out: torch.Tensor
src: torch.Tensor
kernel: object
def run_inputs(self, *inputs: torch.Tensor) -> torch.Tensor:
src = self.src if not inputs else inputs[0]
if len(inputs) not in (0, 1):
raise ValueError(f"expected 0 or 1 input tensors, got {len(inputs)}")
return self.kernel(src)
def run_kernel(self) -> torch.Tensor:
return self.run_inputs()
def run(self) -> None:
self.out.copy_(self.run_inputs())
def destroy(self) -> None:
pass
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def prepare_copy_tl(out: torch.Tensor, src: torch.Tensor) -> TileLangCopyPrepared:
n = src.numel()
block_n = _block_n(n)
if n % block_n != 0:
raise ValueError("TileLang copy v1 requires N % BLOCK_N == 0")
kernel = _compiled_copy(n, block_n, src.dtype)
return TileLangCopyPrepared(out, src, kernel)
def copy_tl_(out: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
with prepare_copy_tl(out, src) as prepared:
prepared.run()
return out
def copy_tl(src: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(src)
return copy_tl_(out, src)

19
ops/copy/tilelang/kernel.py
View File

@ -1,19 +0,0 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit
def copy_kernel(src, BLOCK_N: int, dtype):
N = T.const("N")
src: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
# TODO: implement a tile-wise copy kernel.
#
# Suggested steps:
# 1. Launch one TileLang kernel over the N // BLOCK_N tiles.
# 2. Use T.copy to move one tile from src to out.
return out

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

@ -1,84 +0,0 @@
#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
View File

@ -1,84 +0,0 @@
#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);
}

60
ops/elementwise/relu/relu_common.h
View File

@ -1,60 +0,0 @@
#pragma once
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/elementwise.h"
namespace oprt::elementwise {
struct ReluDescriptor final : oprt::ElementwiseDescriptorBase {
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;
}
auto *typed = new ReluDescriptor();
auto status = oprt::init_elementwise_descriptor(typed, out, {in});
if (status != OPRT_SUCCESS) {
delete typed;
return status;
}
typed->negative_slope = negative_slope;
*desc = typed;
return OPRT_SUCCESS;
}
inline oprt_status_t get_relu_workspace_size_common(oprt_operator_descriptor_t desc,
size_t *size) {
return oprt::get_elementwise_workspace_size(desc, size);
}
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;
}
return oprt::validate_elementwise_execute_args(desc, workspace_size, out, {in});
}
inline oprt_status_t destroy_relu_descriptor_common(oprt_operator_descriptor_t desc) {
return oprt::destroy_elementwise_descriptor(desc);
}
} // namespace oprt::elementwise

4
ops/reduce_sum/metax/README.md
View File

@ -1,4 +0,0 @@
# Reduce Sum MetaX Backend
The reduce_sum MetaX backend exports the standard descriptor lifecycle with
`_metax` symbols and is built from `ops/reduce_sum/metax/reduce_sum_metax.maca`.

132
ops/reduce_sum/metax/reduce_sum_metax.maca
View File

@ -1,132 +0,0 @@
#include "operator_runtime/ops/reduce_sum.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/tensor_checks.h"
#include <cuda_runtime.h>
namespace {
struct ReduceSumDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t in_view;
int64_t rows = 0;
int64_t cols = 0;
int64_t axis = 1;
const char *op_name() const override {
return "reduce_sum";
}
};
__global__ void reduce_sum_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
int64_t row = blockIdx.x;
if (row >= rows) {
return;
}
const float *row_ptr = in + row * cols;
float partial = 0.0f;
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
partial += row_ptr[col];
}
extern __shared__ float shared[];
shared[threadIdx.x] = partial;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared[threadIdx.x] += shared[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
out[row] = shared[0];
}
}
bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
return in.dtype == OPRT_DTYPE_F32 &&
out.dtype == OPRT_DTYPE_F32 &&
in.ndim == 2 &&
out.ndim == 1 &&
axis == 1 &&
out.shape[0] == in.shape[0] &&
oprt::is_contiguous(in) &&
oprt::is_contiguous(out);
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_reduce_sum_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(in);
if (status != OPRT_SUCCESS) {
return status;
}
if (!is_rowwise_case(*out, *in, axis)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new ReduceSumDescriptor();
typed->out_view = *out;
typed->in_view = *in;
typed->rows = in->shape[0];
typed->cols = in->shape[1];
typed->axis = axis;
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_reduce_sum_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_reduce_sum(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || in == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const ReduceSumDescriptor *>(desc);
cudaStream_t s = oprt::as_cuda_stream(stream);
constexpr int threads = 256;
reduce_sum_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), s>>>(
static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_reduce_sum_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

19
ops/reduce_sum/nvidia/kernel.cuh
View File

@ -1,19 +0,0 @@
#pragma once
#include <cuda_runtime.h>
#include <stdint.h>
namespace oprt::reduce_sum::nvidia {
__global__ void reduce_sum_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
// TODO: implement a row-wise reduce_sum kernel with shared memory.
//
// Suggested steps:
// 1. Use one block per row.
// 2. Let each thread accumulate a partial sum over the row.
// 3. Store the partial sums in shared memory.
// 4. Reduce shared memory with a tree reduction.
// 5. Let thread 0 write the final row sum to out[row].
}
} // namespace oprt::reduce_sum::nvidia

105
ops/reduce_sum/nvidia/reduce_sum_cuda.cu
View File

@ -1,105 +0,0 @@
#include "operator_runtime/ops/reduce_sum.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/tensor_checks.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "ops/reduce_sum/nvidia/kernel.cuh"
#include <cuda_runtime.h>
namespace {
struct ReduceSumDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t in_view;
int64_t rows = 0;
int64_t cols = 0;
int64_t axis = 1;
const char *op_name() const override {
return "reduce_sum";
}
};
bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
return in.dtype == OPRT_DTYPE_F32 &&
out.dtype == OPRT_DTYPE_F32 &&
in.ndim == 2 &&
out.ndim == 1 &&
axis == 1 &&
out.shape[0] == in.shape[0] &&
oprt::is_contiguous(in) &&
oprt::is_contiguous(out);
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_reduce_sum_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(in);
if (status != OPRT_SUCCESS) {
return status;
}
if (!is_rowwise_case(*out, *in, axis)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new ReduceSumDescriptor();
typed->out_view = *out;
typed->in_view = *in;
typed->rows = in->shape[0];
typed->cols = in->shape[1];
typed->axis = axis;
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_reduce_sum_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_reduce_sum(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || in == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const ReduceSumDescriptor *>(desc);
cudaStream_t s = oprt::as_cuda_stream(stream);
constexpr int threads = 256;
oprt::reduce_sum::nvidia::reduce_sum_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), s>>>(
static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_reduce_sum_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

29
ops/reduce_sum/tilelang/kernel.py
View File

@ -1,29 +0,0 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
},
)
def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):
N, M = T.const("N, M")
dtype = T.float32
src: T.Tensor((N, M), dtype)
out = T.empty((N,), dtype)
# TODO: implement a tiled row-wise reduce_sum kernel.
#
# Suggested steps:
# 1. Launch one TileLang kernel over row tiles.
# 2. Allocate fragments for the input tile and running output tile.
# 3. Clear the running output fragment before accumulation.
# 4. Iterate over column tiles with T.Serial(M // BLOCK_M).
# 5. T.copy each input tile into the fragment.
# 6. Call T.reduce_sum on the fragment and accumulate into the output fragment.
# 7. Copy the final output fragment back to global memory.
return out

70
ops/reduce_sum/tilelang/reduce_sum_tl.py
View File

@ -1,70 +0,0 @@
from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
import torch
from ops.reduce_sum.tilelang.kernel import reduce_sum_kernel
@lru_cache(maxsize=32)
def _compiled_reduce_sum(n: int, m: int, block_n: int, block_m: int):
return reduce_sum_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m)
def _blocks(n: int, m: int) -> tuple[int, int]:
block_n = 16 if n % 16 == 0 else 1
block_m = 128 if m % 128 == 0 else m
return block_n, block_m
@dataclass
class TileLangReduceSumPrepared:
out: torch.Tensor
src: torch.Tensor
dim: int = 1
kernel: object | None = None
def run_inputs(self, *inputs: torch.Tensor) -> torch.Tensor:
if len(inputs) not in (0, 1):
raise ValueError(f"expected 0 or 1 input tensors, got {len(inputs)}")
src = self.src if not inputs else inputs[0]
assert self.kernel is not None
return self.kernel(src)
def run_kernel(self) -> torch.Tensor:
return self.run_inputs()
def run(self) -> None:
self.out.copy_(self.run_inputs())
def destroy(self) -> None:
pass
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def prepare_reduce_sum_tl(out: torch.Tensor, src: torch.Tensor, dim: int = 1) -> TileLangReduceSumPrepared:
if dim != 1:
raise ValueError("TileLang reduce_sum v1 supports dim=1")
n, m = src.shape
block_n, block_m = _blocks(n, m)
if n % block_n != 0 or m % block_m != 0:
raise ValueError("TileLang reduce_sum v1 requires divisible block sizes")
kernel = _compiled_reduce_sum(n, m, block_n, block_m)
return TileLangReduceSumPrepared(out, src, dim, kernel)
def reduce_sum_tl_(out: torch.Tensor, src: torch.Tensor, dim: int = 1) -> torch.Tensor:
with prepare_reduce_sum_tl(out, src, dim) as prepared:
prepared.run()
return out
def reduce_sum_tl(src: torch.Tensor, dim: int = 1) -> torch.Tensor:
out = torch.empty((src.shape[0],), dtype=src.dtype, device=src.device)
return reduce_sum_tl_(out, src, dim)

4
ops/softmax/metax/README.md
View File

@ -1,4 +0,0 @@
# Softmax MetaX Backend
The softmax MetaX backend exports the standard descriptor lifecycle with
`_metax` symbols and is built from `ops/softmax/metax/softmax_metax.maca`.

155
ops/softmax/metax/softmax_metax.maca
View File

@ -1,155 +0,0 @@
#include "operator_runtime/ops/softmax.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/tensor_checks.h"
#include <cuda_runtime.h>
#include <float.h>
#include <math.h>
namespace {
struct SoftmaxDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t in_view;
int64_t rows = 0;
int64_t cols = 0;
int64_t axis = 1;
const char *op_name() const override {
return "softmax";
}
};
__global__ void softmax_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
int64_t row = blockIdx.x;
if (row >= rows) {
return;
}
const float *row_in = in + row * cols;
float *row_out = out + row * cols;
float local_max = -FLT_MAX;
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
local_max = fmaxf(local_max, row_in[col]);
}
extern __shared__ float shared[];
shared[threadIdx.x] = local_max;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared[threadIdx.x] = fmaxf(shared[threadIdx.x], shared[threadIdx.x + stride]);
}
__syncthreads();
}
float row_max = shared[0];
float local_sum = 0.0f;
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
float value = expf(row_in[col] - row_max);
row_out[col] = value;
local_sum += value;
}
shared[threadIdx.x] = local_sum;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared[threadIdx.x] += shared[threadIdx.x + stride];
}
__syncthreads();
}
float row_sum = shared[0];
for (int64_t col = threadIdx.x; col < cols; col += blockDim.x) {
row_out[col] /= row_sum;
}
}
bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
return in.dtype == OPRT_DTYPE_F32 &&
out.dtype == OPRT_DTYPE_F32 &&
in.ndim == 2 &&
out.ndim == 2 &&
axis == 1 &&
oprt::same_shape(out, in) &&
oprt::is_contiguous(in) &&
oprt::is_contiguous(out);
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_softmax_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(in);
if (status != OPRT_SUCCESS) {
return status;
}
if (!is_rowwise_case(*out, *in, axis)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new SoftmaxDescriptor();
typed->out_view = *out;
typed->in_view = *in;
typed->rows = in->shape[0];
typed->cols = in->shape[1];
typed->axis = axis;
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_softmax_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_softmax(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || in == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const SoftmaxDescriptor *>(desc);
cudaStream_t s = oprt::as_cuda_stream(stream);
constexpr int threads = 256;
softmax_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), s>>>(
static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_softmax_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

22
ops/softmax/nvidia/kernel.cuh
View File

@ -1,22 +0,0 @@
#pragma once
#include <cuda_runtime.h>
#include <float.h>
#include <math.h>
#include <stdint.h>
namespace oprt::softmax::nvidia {
__global__ void softmax_rowwise_kernel(float *out, const float *in, int64_t rows, int64_t cols) {
// TODO: implement a numerically stable row-wise softmax kernel.
//
// Suggested steps:
// 1. Use one block per row.
// 2. Reduce to find the row maximum.
// 3. Compute exp(x - row_max), write the temporary values to out,
// and accumulate their sum.
// 4. Reduce to get the row sum.
// 5. Normalize each output element by row_sum.
}
} // namespace oprt::softmax::nvidia

105
ops/softmax/nvidia/softmax_cuda.cu
View File

@ -1,105 +0,0 @@
#include "operator_runtime/ops/softmax.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/tensor_checks.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "ops/softmax/nvidia/kernel.cuh"
#include <cuda_runtime.h>
namespace {
struct SoftmaxDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t in_view;
int64_t rows = 0;
int64_t cols = 0;
int64_t axis = 1;
const char *op_name() const override {
return "softmax";
}
};
bool is_rowwise_case(const oprt_tensor_view_t &out, const oprt_tensor_view_t &in, int64_t axis) {
return in.dtype == OPRT_DTYPE_F32 &&
out.dtype == OPRT_DTYPE_F32 &&
in.ndim == 2 &&
out.ndim == 2 &&
axis == 1 &&
oprt::same_shape(out, in) &&
oprt::is_contiguous(in) &&
oprt::is_contiguous(out);
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_softmax_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *in,
int64_t axis) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(in);
if (status != OPRT_SUCCESS) {
return status;
}
if (!is_rowwise_case(*out, *in, axis)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new SoftmaxDescriptor();
typed->out_view = *out;
typed->in_view = *in;
typed->rows = in->shape[0];
typed->cols = in->shape[1];
typed->axis = axis;
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_softmax_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_softmax(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *in,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || in == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const SoftmaxDescriptor *>(desc);
cudaStream_t s = oprt::as_cuda_stream(stream);
constexpr int threads = 256;
oprt::softmax::nvidia::softmax_rowwise_kernel<<<typed->rows, threads, threads * sizeof(float), s>>>(
static_cast<float *>(out), static_cast<const float *>(in), typed->rows, typed->cols);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_softmax_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

36
ops/softmax/tilelang/kernel.py
View File

@ -1,36 +0,0 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
},
)
def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):
log2_e = 1.44269504
N, M = T.const("N, M")
dtype = T.float32
src: T.Tensor((N, M), dtype)
out = T.empty((N, M), dtype)
# TODO: implement a tiled row-wise softmax kernel.
#
# Suggested steps:
# 1. Launch one TileLang kernel over row tiles.
# 2. Allocate fragments for src, out, temporary exp values, row max, row sum, and lse.
# 3. Initialize the running log-sum-exp state.
# 4. In a first T.Serial loop over column tiles:
# - copy the input tile into a fragment
# - reduce to get the tile max
# - compute exp2-based temporary values
# - reduce to get the tile sum
# - update the running lse
# 5. In a second T.Serial loop over column tiles:
# - copy the input tile again
# - normalize with the final lse
# - copy the result tile to global memory
return out

70
ops/softmax/tilelang/softmax_tl.py
View File

@ -1,70 +0,0 @@
from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
import torch
from ops.softmax.tilelang.kernel import softmax_kernel
@lru_cache(maxsize=32)
def _compiled_softmax(n: int, m: int, block_n: int, block_m: int):
return softmax_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m)
def _blocks(n: int, m: int) -> tuple[int, int]:
block_n = 16 if n % 16 == 0 else 1
block_m = 256 if m % 256 == 0 else m
return block_n, block_m
@dataclass
class TileLangSoftmaxPrepared:
out: torch.Tensor
src: torch.Tensor
dim: int = 1
kernel: object | None = None
def run_inputs(self, *inputs: torch.Tensor) -> torch.Tensor:
if len(inputs) not in (0, 1):
raise ValueError(f"expected 0 or 1 input tensors, got {len(inputs)}")
src = self.src if not inputs else inputs[0]
assert self.kernel is not None
return self.kernel(src)
def run_kernel(self) -> torch.Tensor:
return self.run_inputs()
def run(self) -> None:
self.out.copy_(self.run_inputs())
def destroy(self) -> None:
pass
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def prepare_softmax_tl(out: torch.Tensor, src: torch.Tensor, dim: int = 1) -> TileLangSoftmaxPrepared:
if dim != 1:
raise ValueError("TileLang softmax v1 supports dim=1")
n, m = src.shape
block_n, block_m = _blocks(n, m)
if n % block_n != 0 or m % block_m != 0:
raise ValueError("TileLang softmax v1 requires divisible block sizes")
kernel = _compiled_softmax(n, m, block_n, block_m)
return TileLangSoftmaxPrepared(out, src, dim, kernel)
def softmax_tl_(out: torch.Tensor, src: torch.Tensor, dim: int = 1) -> torch.Tensor:
with prepare_softmax_tl(out, src, dim) as prepared:
prepared.run()
return out
def softmax_tl(src: torch.Tensor, dim: int = 1) -> torch.Tensor:
out = torch.empty_like(src)
return softmax_tl_(out, src, dim)

4
ops/vector_add/metax/README.md
View File

@ -1,4 +0,0 @@
# Vector Add MetaX Backend
The vector_add MetaX backend exports the standard descriptor lifecycle with
`_metax` symbols and is built from `ops/vector_add/metax/vector_add_metax.maca`.

141
ops/vector_add/metax/vector_add_metax.maca
View File

@ -1,141 +0,0 @@
#include "operator_runtime/ops/vector_add.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/elementwise.h"
#include "operator_runtime/detail/tensor_checks.h"
#include <cuda_fp16.h>
namespace {
struct VectorAddDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t a_view;
oprt_tensor_view_t b_view;
int64_t elements = 0;
const char *op_name() const override {
return "vector_add";
}
};
template <typename T>
__device__ T add_values(T a, T b) {
return a + b;
}
template <>
__device__ inline half add_values<half>(half a, half b) {
return __hadd(a, b);
}
template <typename T>
__global__ void vector_add_contiguous_kernel(
T *__restrict__ out,
const T *__restrict__ a,
const T *__restrict__ b,
int64_t n) {
int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
for (int64_t i = idx; i < n; i += stride) {
out[i] = add_values(a[i], b[i]);
}
}
template <typename T>
oprt_status_t launch_vector_add(const VectorAddDescriptor *desc, void *out, const void *a, const void *b, oprt_stream_t stream) {
constexpr int threads = 256;
int blocks = oprt::blocks_for(desc->elements, threads);
cudaStream_t s = oprt::as_cuda_stream(stream);
vector_add_contiguous_kernel<T><<<blocks, threads, 0, s>>>(
static_cast<T *>(out), static_cast<const T *>(a), static_cast<const T *>(b), desc->elements);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_vector_add_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *a,
const oprt_tensor_view_t *b) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(a);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(b);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*out, *a);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*out, *b);
if (status != OPRT_SUCCESS) {
return status;
}
if (!oprt::elementwise_fast_path(*out, *a, *b)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new VectorAddDescriptor();
typed->out_view = *out;
typed->a_view = *a;
typed->b_view = *b;
typed->elements = oprt::numel(*out);
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_vector_add_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_vector_add(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *a,
const void *b,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || a == nullptr || b == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const VectorAddDescriptor *>(desc);
switch (typed->out_view.dtype) {
case OPRT_DTYPE_F16:
return launch_vector_add<half>(typed, out, a, b, stream);
case OPRT_DTYPE_F32:
return launch_vector_add<float>(typed, out, a, b, stream);
default:
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_vector_add_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

34
ops/vector_add/nvidia/kernel.cuh
View File

@ -1,34 +0,0 @@
#pragma once
#include <cuda_fp16.h>
#include <stdint.h>
namespace oprt::vector_add::nvidia {
template <typename T>
__device__ T add_values(T a, T b) {
// TODO: return the elementwise sum for generic types.
return T{};
}
template <>
__device__ inline half add_values<half>(half a, half b) {
// TODO: return the half-precision elementwise sum.
return half{};
}
template <typename T>
__global__ void vector_add_contiguous_kernel(T * __restrict__ out,
const T * __restrict__ a,
const T * __restrict__ b,
int64_t n) {
// TODO: implement a grid-stride loop vector add kernel.
//
// Suggested steps:
// 1. Compute the global thread index.
// 2. Compute the grid-wide stride.
// 3. Loop over i = idx; i < n; i += stride.
// 4. Write out[i] = add_values(a[i], b[i]).
}
} // namespace oprt::vector_add::nvidia

119
ops/vector_add/nvidia/vector_add_cuda.cu
View File

@ -1,119 +0,0 @@
#include "operator_runtime/ops/vector_add.h"
#include "operator_runtime/descriptor.h"
#include "operator_runtime/detail/elementwise.h"
#include "operator_runtime/detail/tensor_checks.h"
#include "operator_runtime/detail/cuda_helpers.h"
#include "ops/vector_add/nvidia/kernel.cuh"
#include <cuda_fp16.h>
namespace {
struct VectorAddDescriptor final : oprt_operator_descriptor {
oprt_tensor_view_t out_view;
oprt_tensor_view_t a_view;
oprt_tensor_view_t b_view;
int64_t elements = 0;
const char *op_name() const override {
return "vector_add";
}
};
template <typename T>
oprt_status_t launch_vector_add(const VectorAddDescriptor *desc, void *out, const void *a, const void *b, oprt_stream_t stream) {
constexpr int threads = 256;
int blocks = oprt::blocks_for(desc->elements, threads);
cudaStream_t s = oprt::as_cuda_stream(stream);
oprt::vector_add::nvidia::vector_add_contiguous_kernel<T><<<blocks, threads, 0, s>>>(
static_cast<T *>(out), static_cast<const T *>(a), static_cast<const T *>(b), desc->elements);
OPRT_CUDA_RETURN_IF_ERROR(cudaGetLastError());
return OPRT_SUCCESS;
}
} // namespace
extern "C" OPRT_EXPORT oprt_status_t oprt_create_vector_add_descriptor(
oprt_operator_descriptor_t *desc,
const oprt_tensor_view_t *out,
const oprt_tensor_view_t *a,
const oprt_tensor_view_t *b) {
if (desc == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
*desc = nullptr;
auto status = oprt::check_tensor(out);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(a);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_tensor(b);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*out, *a);
if (status != OPRT_SUCCESS) {
return status;
}
status = oprt::check_same_dtype(*out, *b);
if (status != OPRT_SUCCESS) {
return status;
}
if (!oprt::elementwise_fast_path(*out, *a, *b)) {
return OPRT_ERR_NOT_SUPPORTED;
}
auto *typed = new VectorAddDescriptor();
typed->out_view = *out;
typed->a_view = *a;
typed->b_view = *b;
typed->elements = oprt::numel(*out);
typed->workspace_size = 0;
*desc = typed;
return OPRT_SUCCESS;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_get_vector_add_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;
}
extern "C" OPRT_EXPORT oprt_status_t oprt_execute_vector_add(
oprt_operator_descriptor_t desc,
void *,
size_t workspace_size,
void *out,
const void *a,
const void *b,
oprt_stream_t stream) {
if (desc == nullptr || out == nullptr || a == nullptr || b == nullptr) {
return OPRT_ERR_INVALID_ARG;
}
if (workspace_size < desc->workspace_size) {
return OPRT_ERR_INSUFFICIENT_WORKSPACE;
}
auto *typed = static_cast<const VectorAddDescriptor *>(desc);
switch (typed->out_view.dtype) {
case OPRT_DTYPE_F16:
return launch_vector_add<half>(typed, out, a, b, stream);
case OPRT_DTYPE_F32:
return launch_vector_add<float>(typed, out, a, b, stream);
default:
return OPRT_ERR_UNSUPPORTED_DTYPE;
}
}
extern "C" OPRT_EXPORT oprt_status_t oprt_destroy_vector_add_descriptor(
oprt_operator_descriptor_t desc) {
delete desc;
return OPRT_SUCCESS;
}

21
ops/vector_add/tilelang/kernel.py
View File

@ -1,21 +0,0 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit
def vector_add_kernel(a, b, BLOCK_N: int, dtype):
N = T.const("N")
a: T.Tensor((N,), dtype)
b: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
# TODO: implement a tile-wise vector add kernel.
#
# Suggested steps:
# 1. Launch one TileLang kernel over the N // BLOCK_N tiles.
# 2. Compute the tile base offset.
# 3. Use T.Parallel(BLOCK_N) to fill out[base + i] = a[base + i] + b[base + i].
return out

74
ops/vector_add/tilelang/vector_add_tl.py
View File

@ -1,74 +0,0 @@
from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
import tilelang.language as T
import torch
from ops.vector_add.tilelang.kernel import vector_add_kernel
def _tl_dtype(dtype: torch.dtype):
if dtype is torch.float16:
return T.float16
if dtype is torch.float32:
return T.float32
raise TypeError(f"unsupported TileLang dtype: {dtype}")
def _block_n(n: int) -> int:
return 1024 if n % 1024 == 0 else n
@lru_cache(maxsize=32)
def _compiled_vector_add(n: int, block_n: int, dtype: torch.dtype):
return vector_add_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
@dataclass
class TileLangVectorAddPrepared:
out: torch.Tensor
a: torch.Tensor
b: torch.Tensor
kernel: object
def run_inputs(self, *inputs: torch.Tensor) -> torch.Tensor:
if len(inputs) not in (0, 2):
raise ValueError(f"expected 0 or 2 input tensors, got {len(inputs)}")
a, b = (self.a, self.b) if not inputs else inputs
return self.kernel(a, b)
def run_kernel(self) -> torch.Tensor:
return self.run_inputs()
def run(self) -> None:
self.out.copy_(self.run_inputs())
def destroy(self) -> None:
pass
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def prepare_vector_add_tl(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> TileLangVectorAddPrepared:
n = out.numel()
block_n = _block_n(n)
if n % block_n != 0:
raise ValueError("TileLang vector_add v1 requires N % BLOCK_N == 0")
kernel = _compiled_vector_add(n, block_n, out.dtype)
return TileLangVectorAddPrepared(out, a, b, kernel)
def vector_add_tl_(out: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
with prepare_vector_add_tl(out, a, b) as prepared:
prepared.run()
return out
def vector_add_tl(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(a)
return vector_add_tl_(out, a, b)

18
python/operator_runtime/__init__.py
View File

@ -1,18 +0,0 @@
from .backend import Backend, normalize_backend
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,
)
__all__ = [
"Backend",
"normalize_backend",
"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",
]

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

@ -1,35 +0,0 @@
from .loader import load_library
from .tensor_view import TensorView, tensor_view, dtype_to_oprt, current_stream_ptr, OPRT_MAX_DIMS
from .bindings import (
CFunctions,
Descriptor,
OperatorRuntimeError,
Status,
bind_elementwise,
bind_unary,
bind_relu,
bind_binary,
bind_reduce_like,
check_status,
)
from .prepared import PreparedOp
__all__ = [
"load_library",
"TensorView",
"tensor_view",
"dtype_to_oprt",
"current_stream_ptr",
"OPRT_MAX_DIMS",
"CFunctions",
"Descriptor",
"OperatorRuntimeError",
"Status",
"bind_elementwise",
"bind_unary",
"bind_relu",
"bind_binary",
"bind_reduce_like",
"check_status",
"PreparedOp",
]

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

@ -1,140 +0,0 @@
from __future__ import annotations
import ctypes
from dataclasses import dataclass
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 _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
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
workspace: Callable
execute: Callable
destroy: Callable
def _missing_symbol_error(name: str, symbol: str, exc: AttributeError) -> OperatorRuntimeError:
raise OperatorRuntimeError(
f"missing symbol {symbol} for operator {name}"
) from exc
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), *create_argtypes]
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, *execute_argtypes, 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)
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:
return bind_elementwise("relu", 1, [ctypes.c_float], backend)
def bind_binary(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions:
return bind_elementwise(name, 2, backend=backend)
def bind_reduce_like(name: str, backend: str | Backend = Backend.NVIDIA) -> CFunctions:
return _bind_lifecycle(
name,
backend,
[
ctypes.POINTER(TensorView),
ctypes.POINTER(TensorView),
ctypes.c_int64,
],
[ctypes.c_void_p, ctypes.c_void_p],
)

72
python/operator_runtime/_internal/loader.py
View File

@ -1,72 +0,0 @@
from __future__ import annotations
import ctypes
import os
from functools import lru_cache
from pathlib import 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] = []
repo_root = Path(__file__).resolve().parents[3]
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=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(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))

65
python/operator_runtime/_internal/prepared.py
View File

@ -1,65 +0,0 @@
from __future__ import annotations
import ctypes
from dataclasses import dataclass
from typing import Any
import torch
from .bindings import CFunctions, Descriptor, check_status
@dataclass
class PreparedOp:
funcs: CFunctions
descriptor: Descriptor
workspace: torch.Tensor | None
runner_args: tuple[Any, ...]
stream_tensor: torch.Tensor | None = None
def run(self) -> None:
self.run_inputs()
def run_inputs(self, *runner_tensors: torch.Tensor) -> None:
if runner_tensors:
expected = len(self.runner_args) - 1
if len(runner_tensors) != expected:
raise ValueError(f"expected {expected} input tensors, got {len(runner_tensors)}")
runner_args = (
self.runner_args[0],
*tuple(ctypes.c_void_p(tensor.data_ptr()) for tensor in runner_tensors),
)
else:
runner_args = self.runner_args
stream = torch.cuda.current_stream(device=self.stream_tensor.device if self.stream_tensor is not None else None)
workspace_ptr = None if self.workspace is None else ctypes.c_void_p(self.workspace.data_ptr())
workspace_size = 0 if self.workspace is None else self.workspace.numel()
status = self.funcs.execute(
self.descriptor,
workspace_ptr,
workspace_size,
*runner_args,
ctypes.c_void_p(stream.cuda_stream),
)
check_status(status)
def run_kernel(self) -> None:
self.run()
def destroy(self) -> None:
if self.descriptor:
check_status(self.funcs.destroy(self.descriptor))
self.descriptor = Descriptor()
def __enter__(self) -> "PreparedOp":
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.destroy()
def __del__(self) -> None:
try:
self.destroy()
except Exception:
pass

54
python/operator_runtime/_internal/tensor_view.py
View File

@ -1,54 +0,0 @@
from __future__ import annotations
import ctypes
import torch
OPRT_MAX_DIMS = 8
OPRT_DTYPE_F16 = 0
OPRT_DTYPE_F32 = 1
class TensorView(ctypes.Structure):
_fields_ = [
("data", ctypes.c_void_p),
("dtype", ctypes.c_int),
("ndim", ctypes.c_int32),
("shape", ctypes.c_int64 * OPRT_MAX_DIMS),
("strides", ctypes.c_int64 * OPRT_MAX_DIMS),
]
def dtype_to_oprt(dtype: torch.dtype) -> int:
if dtype is torch.float16:
return OPRT_DTYPE_F16
if dtype is torch.float32:
return OPRT_DTYPE_F32
raise TypeError(f"unsupported dtype: {dtype}")
def tensor_view(tensor: torch.Tensor) -> TensorView:
if tensor.ndim > OPRT_MAX_DIMS:
raise ValueError(f"ndim {tensor.ndim} exceeds OPRT_MAX_DIMS={OPRT_MAX_DIMS}")
shape = (ctypes.c_int64 * OPRT_MAX_DIMS)()
strides = (ctypes.c_int64 * OPRT_MAX_DIMS)()
for i, dim in enumerate(tensor.shape):
shape[i] = dim
for i, stride in enumerate(tensor.stride()):
strides[i] = stride
return TensorView(
ctypes.c_void_p(tensor.data_ptr()),
dtype_to_oprt(tensor.dtype),
tensor.ndim,
shape,
strides,
)
def current_stream_ptr(tensor: torch.Tensor | None = None) -> ctypes.c_void_p:
if not torch.cuda.is_available():
return ctypes.c_void_p()
device = tensor.device if tensor is not None and tensor.is_cuda else None
stream = torch.cuda.current_stream(device=device)
return ctypes.c_void_p(stream.cuda_stream)

19
python/operator_runtime/backend.py
View File

@ -1,19 +0,0 @@
from __future__ import annotations
from enum import Enum
class Backend(str, Enum):
NVIDIA = "nvidia"
TILELANG = "tilelang"
METAX = "metax"
def normalize_backend(backend: str | Backend) -> Backend:
if isinstance(backend, Backend):
return backend
try:
return Backend(backend.lower())
except ValueError as exc:
raise ValueError(f"unknown backend: {backend}") from exc

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

@ -1,13 +0,0 @@
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",
]

Some files were not shown because too many files have changed in this diff Show More