forked from ccf-ai-infra/Intro-ops
274 lines
8.5 KiB
Plaintext
274 lines
8.5 KiB
Plaintext
{
|
||
|
||
"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
|
||
|
||
}
|
||
|