tilelang: implement copy and vector_add kernels

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_<op>.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 |

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

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