diff --git a/course/notes/tilelang-basics.md b/course/notes/tilelang-basics.md new file mode 100644 index 0000000..3a17cec --- /dev/null +++ b/course/notes/tilelang-basics.md @@ -0,0 +1,124 @@ +# TileLang 基础知识点 + +## 1. 整体执行模型 + +TileLang kernel 的执行分为两个层级:**块(block)** 和 **线程(thread)**。 + +``` +T.Kernel(块数, threads=每块线程数) + │ + ├─ Block 0 (pid=0) ── N个线程 ── 处理自己的数据切片 + ├─ Block 1 (pid=1) ── N个线程 ── 处理自己的数据切片 + ├─ Block 2 (pid=2) ── N个线程 ── 处理自己的数据切片 + └─ ... +``` + +- **块之间完全独立**,各自算各自的,互不通信 +- **每个块执行同一份代码**,通过 `pid`(program ID)区分"我是哪个块" + +--- + +## 2. `T.Kernel` 的两个参数 + +```python +with T.Kernel(块数, threads=256) as pid_n: + ... +``` + +| 参数 | 含义 | 例子 | +|------|------|------| +| 第一个位置参数 | **块数**(block count) | `N // BLOCK_N` | +| `threads=` | **每块线程数**(硬件并行度) | `256` | +| `as pid_n` | 当前块的编号(0, 1, 2, ...) | — | + +--- + +## 3. `pid_n` 是什么 + +`pid_n` = **program ID**,当前线程块的编号。从 0 开始递增。 + +每个块用自己的 `pid_n` 算出该处理哪段数据: + +```python +base = pid_n * BLOCK_N # 块0→0, 块1→256, 块2→512, ... +``` + +--- + +## 4. `BLOCK_N` vs `threads` —— 数据 vs 硬件 + +| 参数 | 含义 | 类比 | +|------|------|------| +| `BLOCK_N` | **tile 大小**——每块处理多少元素 | 一份外卖要炒几个菜 | +| `threads=256` | **硬件线程数**——每块里多少工人 | 后厨里有几个厨师 | + +它们**不需要相等**。`BLOCK_N` 管数据分块,`threads` 管并行度。 + +当 `BLOCK_N > threads` 时,每个线程会自动多轮循环(grid-stride loop): + +``` +BLOCK_N=1024, threads=256 → 每个线程跑 1024/256 = 4 个元素 + +线程 0 → i=0, i=256, i=512, i=768 +线程 1 → i=1, i=257, i=513, i=769 +... +线程255 → i=255, i=511, i=767, i=1023 +``` + +--- + +## 5. `T.Parallel` vs `T.Serial` + +```python +for i in T.Parallel(BLOCK_N): # 并行——迭代间无依赖 + out[i] = a[i] + b[i] + +for i in T.Serial(M // BLOCK_M): # 串行——迭代间有依赖(如累加) + acc += T.reduce_sum(tile) +``` + +| | `T.Parallel` | `T.Serial` | +|---|---|---| +| 执行方式 | 所有迭代同时进行 | 必须一步步来 | +| 适用场景 | 逐元素操作(copy, vector_add) | 有累积状态的操作(reduce_sum, softmax) | +| 对应概念 | GPU 线程并行 | 顺序循环,依赖前一次结果 | + +--- + +## 6. `T.copy` 的正确用法 + +`T.copy` 搬的是**切片(tile)**,不是整个 tensor。 + +```python +# ❌ 错误:每个块都搬整个 tensor +T.copy(src, out) + +# ✅ 正确:每个块只搬自己的切片 +T.copy( + src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], + out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], +) +``` + +--- + +## 7. 常用验证命令 + +```bash +conda activate py312 + +# 单算子 TileLang 测试 +PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia \ + pytest tests/op_tests/test_.py -v --backend tilelang +``` + +--- + +## 8. 算子学习顺序 + +| 顺序 | 算子 | 关键词 | 新概念 | +|------|------|--------|--------| +| 1 | `copy` | `T.Kernel` + `T.copy` | 块、pid、tile 切片 | +| 2 | `vector_add` | `for i in T.Parallel(...)` | 手写并行循环 | +| 3 | `reduce_sum` | `T.Serial` + `T.reduce_sum` | 顺序循环、fragment、规约 | +| 4 | `softmax` | online softmax + `exp2`/`log2` | 两遍扫描、log-sum-exp | diff --git a/ops/copy/tilelang/kernel.py b/ops/copy/tilelang/kernel.py index 3f5e182..c1ec213 100644 --- a/ops/copy/tilelang/kernel.py +++ b/ops/copy/tilelang/kernel.py @@ -15,5 +15,10 @@ def copy_kernel(src, BLOCK_N: int, dtype): # 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. + 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 diff --git a/ops/vector_add/tilelang/kernel.py b/ops/vector_add/tilelang/kernel.py index 21fd3c7..77565c6 100644 --- a/ops/vector_add/tilelang/kernel.py +++ b/ops/vector_add/tilelang/kernel.py @@ -17,5 +17,9 @@ def vector_add_kernel(a, b, BLOCK_N: int, dtype): # 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]. + with T.Kernel(N // BLOCK_N, threads = 256) as pid_n: + base = pid_n * BLOCK_N + for i in T.Parallel(BLOCK_N): + out[base + i] = a[base + i] + b[base + i] return out