跳到主要内容
TileLang 系统教程

tilelang-tutorial.huandu.me

第 07 章 快速上手:安装与第一个 kernel

/code/07-quickstart/

01-first-kernel/README.md2022 字节
# 01-first-kernel · 一维向量加法

## 目的

本章「第一个 kernel」的完整、可直接复制的版本。一个文件里包含两种写法:

* **版本 A(`vector_add_direct`)**:最朴素的写法。只有 `T.Kernel` + `T.Parallel`,
  直接在全局内存上读写。对应官方教程 `docs/deeplearning_operators/elementwise.md`
  的第一段示例。
* **版本 B(`vector_add_buffered`)**:先 `T.alloc_fragment` 分配片上 buffer,用
  `T.copy` 把 tile 搬进来,算完再搬出去。对应官方
  `examples/elementwise/example_elementwise_add.py` 的结构。

两段代码解决的问题完全相同(`C = A + B`),差别只在「数据先不先落到片上 buffer」。
对照读它们,是理解 `T.copy` 与 `T.Parallel` 分工的最快方式。

## 运行前提

| 项 | 要求 |
| --- | --- |
| Python | >= 3.10 |
| tilelang | 已安装(wheel 或源码构建,见本章 3.2 / 3.3 节) |
| 加速器 | 一块 NVIDIA GPU;脚本用 `torch.cuda` 造数据 |
| torch | 已安装 CUDA 版 |

## 运行命令

```bash
python vecadd_tilelang.py                  # 两种写法都跑
python vecadd_tilelang.py --variant a      # 只跑版本 A
python vecadd_tilelang.py --n 4096 --block-n 128 --threads 128
```

## 预期输出

```text
[variant a] correctness OK (n=1048576, block_N=256, threads=256)
----- generated CUDA (variant a) -----
extern "C" __global__ void __launch_bounds__(256) main_kernel(...) {
  ...   # 预期看到 128-bit 向量化访存(uint4)与半精度成对加法
}
[variant b] correctness OK (n=1048576, block_N=256, threads=256)
```

**正确性判据**:`torch.testing.assert_close(c, a + b, rtol=1e-2, atol=1e-2)` 通过。
**性能口径**:本脚本不测性能。要测吞吐用 `kernel.get_profiler().do_bench()`
(见本章 6.3 节),并自行固定 shape / dtype / backend / warmup / repeat。

## 验证级别

**仅静态核对。**

作者本机是 macOS(Apple Silicon),既没有 NVIDIA GPU 也没有安装 tilelang,
**没有运行过这个脚本**。所谓「静态核对」指的是:

1. 脚本可以 `python3 -m py_compile` 通过(语法层面无误);
2. 用到的每个 API 名与签名都在 `_vendor/tilelang-v0.1.15/` 里逐个核对过,
   见本章 `references.md` 的 S 系列条目。

## 在真实机器上的验证步骤

1. 按本章 3.2 或 3.3 节安装 tilelang;`python -c "import tilelang; print(tilelang.__version__)"`
   能打印版本号。
2. `python -c "import torch; print(torch.cuda.is_available())"` 打印 `True`。
3. 运行 `python vecadd_tilelang.py`,确认两条 `correctness OK`。
4. 把打印出的 CUDA 源码贴到官方 `elementwise.md` 的生成结果旁边逐行比对
   (见本章第 5 节)。若看到 `uint4` 或 `nv_bfloat162`,说明命中了向量化。
5. 想确认缓存生效:连续运行两次,第二次不应重新走 nvcc(可用
   `TILELANG_VERBOSE=1 python vecadd_tilelang.py` 观察日志)。

← 浏览本章全部代码