Lecture 13 · ML 编译器（内核生成）：调度语言、多面体、自动调优与稀疏编译 · ML Compilers: Kernel Generation

对应 Maps to：Stanford CS349E / EE339 Efficient ML Inference at Scale（Autumn 2026），原课第 15 次课 · Thu Nov 12 · ML Compilers – Kernel Generation（Fredrik Kjolstad）。 依据说明：这次课没有指定阅读材料，课程网站也没有 slides。本讲以公开文献为主线：J. Ragan-Kelley et al., Halide（PLDI 2013）；T. Chen et al., TVM（OSDI 2018）与 L. Zheng et al., Ansor（OSDI 2020）；P. Tillet et al., Triton（MAPL 2019）；多面体编译的经典工作（Feautrier 1992、Pluto / Bondhugula et al. PLDI 2008）；F. Kjolstad et al., The Tensor Algebra Compiler（OOPSLA 2017）——TACO 是本课老师 Kjolstad 的工作。内容的取舍是我定的，不代表课堂上实际讲了什么。 配套代码 Code：code/lecture13/（numpy + 系统的 C 编译器；一个共享模块 + 5 个脚本，最长的约 10 秒） 先修：前置讲 P4（循环嵌套、依赖分析、循环变换、算法与调度分离、代价模型与自动调优）、第 10 讲（GPU 矩阵乘法内核的每一个旋钮）、第 12 讲（前端把模型变成一组内核）。 预计阅读时间：3 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 解释”算法与调度分离”，用 Halide 的调度原语描述一个流水线的几种实现，并说出生产者-消费者调度的三角取舍：并行性、局部性、重复计算。
2. 做边界推断：给定输出区域，算出每一级要算多大的区域；估算分块融合的重复计算量。
3. 用多面体模型的语言（迭代空间、依赖距离、仿射变换）判断一个变换是否合法、能否矩形分块、内层能否并行；说明倾斜（skewing）和波前（wavefront）的作用。
4. 说出 TVM、Triton 各自把调度的哪部分交给程序员、哪部分交给编译器；描述一个 GPU 矩阵乘法内核的调度空间，以及最优配置为什么随形状变化。
5. 比较自动调优的几种搜索策略，解释学习的代价模型的作用、局限和成本。
6. 说明稀疏张量编译的核心思想：每一维的存储格式决定生成什么样的循环；加法是并集合并、乘法是交集合并；稀疏什么时候才划算。

0.2 本讲地图 Map

§1  后端要做什么
§2  Halide：算法与调度分离；生产者在哪里算                       ← 代码 01
§3  多面体模型：合法的变换、倾斜、分块、波前                      ← 代码 02
§4  TVM：张量表达式、调度原语与 tensorize
§5  Triton 编译器替你做了什么
§6  GPU 矩阵乘法的调度空间                                       ← 代码 03
§7  自动调优：搜索与代价模型                                     ← 代码 04
§8  稀疏张量编译（TACO）                                        ← 代码 05
§9  现状（2025–2026）
§10–14  小结、常见误区、练习（含答案）、延伸阅读、下一讲预告

0.3 本讲的代码

sched_lab.py 有四个部分：

• A. 迷你 Halide：用运算符重载写 Func（算法），用一个字典写调度（root / inline / 分块 / 滑动窗口）；边界推断、运算量统计、生成 C 代码并用前置讲 P4 的 cbuild 编译、在你的 CPU 上计时。
• B. 二维多面体模型：依赖距离向量、幺模变换的合法性 / 可分块性 / 并行性；按变换后的顺序真的执行，和原顺序比较；用 P4 的 LRU cache 模型数缺失。
• C. 自动调优：一个 Triton 风格的 GPU 矩阵乘法配置空间（6000 个组合）；“硬件”是一个解析模型（H100 式参数 + 第 8、10 讲的资源账 + 若干假设），用来研究搜索，不是实测。
• D. 稀疏编译：CSR / DCSR 格式、按格式生成 SpMV 的代码、并集 / 交集合并的代码生成；C 实测稀疏和稠密的交叉点。

1. 后端要做什么 · What the backend does

第 12 讲的前端把模型变成了一组内核：每个内核是一个融合好的计算——逐元素运算、按行归约、带 epilogue 的矩阵乘法、注意力。每个内核在数学上是一个循环嵌套。后端要为它生成高效的机器码：

• 循环怎样分块，块多大（第 9 讲的复用、第 10 讲的寄存器 / 共享内存分块）；
• 循环的顺序，哪些循环并行、映射到线程块 / warp / 线程，哪些向量化；
• 数据经过哪些存储层次（共享内存？寄存器？），怎样流水；
• 多个阶段之间，中间结果是存起来、现算，还是一块一块地算。

有三种做法：

1. 调用库：cuBLAS、cuDNN 为常见的算子和形状准备了大量手写的内核，加上一套按形状选内核的启发式。最快，但只覆盖库里有的东西。
2. 模板：CUTLASS / CuTe 把 GEMM、卷积写成 C++ 模板，块大小、流水线级数、epilogue 是模板参数；融合的 epilogue 可以定制。
3. 生成：Halide、TVM、Triton、TorchInductor 从一个高层描述生成代码。它们共同的核心是前置讲 P4 §9 的算法与调度分离：算法说算什么，调度说怎么算；改调度不会改变结果，所以可以放心地搜索调度。

这一讲按”调度语言 → 变换的数学 → 搜索 → 稀疏”的顺序展开。

2. Halide：算法与调度分离 · Halide

代码：01_halide_schedules.py

2.1 算法

Halide（Ragan-Kelley et al. PLDI 2013）最初为图像处理设计。图像处理的流水线是很多个模板（stencil）运算串起来：每个输出像素依赖上一级的一小块邻域。经典的例子是两级的模糊：

blurx(x, y) = (in(x−1, y) + in(x, y) + in(x+1, y)) / 3
blury(x, y) = (blurx(x, y−1) + blurx(x, y) + blurx(x, y+1)) / 3

这就是 Halide 的算法：每个 Func 是定义在整个二维整数平面上的纯函数，只说”一个点怎么算”，不说循环、缓冲区、计算的顺序。论文的局部拉普拉斯滤波器（local Laplacian filter）有 99 级，Halide 的算法只有 62 行。

2.2 调度：生产者在哪里算

调度用一组原语描述怎样把算法变成循环：split、reorder、tile、vectorize、parallel、unroll（前置讲 P4 §8 的循环变换），以及 Halide 最独特的两个——compute_at 和 store_at：生产者（blurx）相对于消费者（blury）在哪一层循环里计算、在哪一层分配存储。几个典型的选择（代码 01 演示 2，4096×4096 的输出）：

边界推断（bounds inference）是让这一切成立的关键：要算 64×64 的输出，blurx 要多算上下各一行（64×66），输入还要再多左右各一列（66×66）。这些区域由编译器从算法里的偏移自动推出来；换一种调度，区域跟着变。分块多算的那一部分就是块边上的”光晕”（halo）：64×64 的块多算 $66/64 - 1 = 3\%$。

这张表就是 Halide 论文的核心论点：并行性、局部性、重复计算三者不可兼得。root 不重复、能并行，但局部性差；inline 局部性最好，但重复计算；滑动窗口不重复、局部性好，但串行；分块在三者之间折中，块的大小调节折中的位置。Halide 的实践是组合：先按条带切开（并行），每个条带内部再滑动（不重复）——用一点点重复计算换回并行性。

2.3 实测：最快的调度取决于机器

代码 01 演示 3 为每种调度生成 C 代码（单线程，-O3），在作者的机器（Apple M4 Pro）上一次运行的结果：

同一个算法、同样的结果，只换调度，最快和最慢差 3 倍多。有几处和”教科书直觉”不一样：

• 方块分块在这台机器上反而最慢（8192 时是 root 的 2.5 倍）。生成的代码被编译器向量化了，但每块只有几十个连续的数，一行一行地跳着读，对硬件预取和 TLB 都不友好；整行的条带保留了连续读，又把 blurx 留在 cache 里。我们没有进一步分析原因——这正说明调度的好坏很难靠直觉判断。
• 图像小的时候 root 很快：2048×2048 时 blurx 只有 16 MB，这台机器的 cache 和内存带宽都很大，“写出去再读回来”的代价不高；到 8192（blurx 268 MB）时，滑动窗口才明显胜出。
• 这个实验是单线程的，没有体现并行性——多核时滑动窗口的串行依赖就成了劣势。

Halide 论文的做法是自动调优：在调度空间里随机搜索 + 进化，找到的往往是”反直觉”的调度；论文报告用几小时写出的 Halide 程序，比专家花几周到几个月手工优化的 C、intrinsics、CUDA 实现快最多 5 倍。后来的 Halide autoscheduler（Adams et al. 2019）用学到的代价模型 + 树搜索自动生成调度。

2.4 深的流水线

代码 01 演示 4：4 级 3×3 方框模糊（每级 9 个点的平均）串起来，2048×2048：

全部内联时，第 4 级的每个点要第 3 级的 9 个点，每个又要第 2 级的 9 个……运算按 9 的幂增长：$9 + 81 + 729 + 6561 = 7380$。分块融合时，每一级的区域都比下一级大一圈，块越小”光晕”占比越大；块越大，中间存储越多（放不进 L1 / 共享内存时融合就没意义了）。在这台机器上 root 反而最快（中间结果 50 MB，写出去再读回来的代价不高）。

这和 ML 里的决定是同一类：把几个卷积层融合进一个内核（层越多、感受野越大、光晕越宽）、FlashAttention 宁可在反向重算 $S$（第 11 讲）、第 12 讲的融合 pass 选择”内联”一个逐元素运算还是把它写出来。

3. 多面体模型：合法的变换 · The polyhedral model

代码：02_polyhedral.py

3.1 迭代空间与依赖

Halide 的调度是”在结构化的流水线上做选择”；多面体模型处理更一般的问题：给定一个（仿射的）循环嵌套，所有合法的循环变换是什么，哪一个最好？

• 迭代空间：循环的每一次迭代是一个整数点，所有的点组成一个由线性不等式围成的集合（多面体里的整数点）。
• 依赖：迭代 $q$ 用到了迭代 $p$ 产生的值，就有一个依赖 $p \to q$；对规则的程序，依赖可以用距离向量 $d = q - p$ 描述（前置讲 P4 §7.2）。
• 变换：一个仿射映射 $\theta(p) = Tp$（加常数），新的执行顺序是 $\theta(p)$ 的字典序。交换、倾斜、反转都是特例；$T$ 是幺模矩阵（行列式 ±1）时，变换后的整数点和原来一一对应。

三条判断（代码 02 的 classify）：

3.2 一维 Jacobi：倾斜之后才能分块

$A[t+1][i] = (A[t][i-1] + A[t][i] + A[t][i+1]) / 3$，迭代点 $(t, i)$，依赖 $(1, -1)$、$(1, 0)$、$(1, 1)$：

原样的 $(1, -1)$ 有负分量，直接矩形分块非法：块的边上有依赖指向”右边还没执行的块”。倾斜（skewing）$i' = i + t$ 把所有依赖变成非负，就能分块了。枚举元素在 $[-2, 2]$ 的全部 2×2 幺模矩阵：合法的 26 个，可分块的 12 个。多面体编译器（Pluto）用整数线性规划在所有仿射变换里找”可分块、依赖距离最短”的那个，而不是暴力枚举。

代码 02 演示 2 真的按变换后的顺序执行（$N = 2000$、200 个时间步、块 50×50；每个时间步存一行，用 NaN 标记”还没算出来”）：

非法的顺序真的算错了。倾斜后分块的结果和原顺序逐位相同，cache 缺失从每点 0.250 降到 0.151——其中写新的一行总是缺失（每 8 个点 1 次 = 0.125），读的缺失从 0.125 降到约 0.026：原顺序每个时间步扫过整行（16 KB，比 cache 大），下一步再读时早已被挤出去；倾斜的块在一小块时空区域里连续推进几十个时间步。这就是时间分块（time tiling），多面体编译最典型的成功案例。

3.3 二维 Gauss–Seidel：波前并行

原地更新 $A[i][j] = (A[i-1][j] + A[i][j-1] + A[i+1][j] + A[i][j+1]) / 4$：流依赖和反依赖的距离都是 $(1, 0)$、$(0, 1)$。两层循环都带着依赖，原样执行时内层不能并行（“内层倒着执行”结果就变了，差 0.27）。波前变换 $(i + j,\ i)$：同一条反对角线上的点互不依赖。$N = 64$ 时外层 127 步，每步内层平均 32 个点可以并行（从 1 涨到 64 再降回 1），倒着执行结果完全一样。

3.4 矩阵乘法与多面体模型的边界

矩阵乘法只有一个依赖 $(0, 0, 1)$（同一个 $C[i][j]$ 沿 $k$ 累加）：任何排列都合法、任意分块都合法、$i$ 和 $j$ 可以任意并行。所以矩阵乘法的调度空间（§6）只有性能问题，没有正确性问题——如果允许重新结合加法（浮点结果略有不同），$k$ 也能切开并行（split-K；第 11 讲的 flash-decoding 就是在注意力的 KV 方向上做 split-K）。

多面体模型的工具链：ISL（整数集合库）、Pluto、PPCG（生成 GPU 代码）、Tensor Comprehensions、Tiramisu，以及 MLIR 的 affine 方言。它的边界也很清楚：要求循环边界和下标都是仿射的（循环变量的线性函数）——数据依赖的下标（稀疏、间接寻址、分页 KV 的块表）就超出了它的范围（§8）；整数线性规划在大的循环嵌套上很慢；它擅长找”合法的变换”，但选哪个变换最快还需要代价模型。现代的 ML 编译器大多只用它的一部分思想（依赖分析、仿射映射），而不是完整的多面体调度。

4. TVM：张量表达式、调度原语与 tensorize · TVM

TVM（Chen et al. OSDI 2018）把 Halide 的思路搬到深度学习的张量算子上：

A = te.placeholder((M, K)); B = te.placeholder((K, N))
k = te.reduce_axis((0, K))
C = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))   # 算法：张量表达式
s = te.create_schedule(C.op)                                              # 调度
io, jo, ii, ji = s[C].tile(C.op.axis[0], C.op.axis[1], 64, 64)            # 分块
s[C].bind(io, te.thread_axis("blockIdx.x"))                               # 映射到线程块 / 线程
AA = s.cache_read(A, "shared", [C]); s[AA].compute_at(s[C], ko)           # 经过共享内存，在 ko 这层循环里搬
s[C].tensorize(ii, mma_intrin)                                            # 把内层换成一条硬件指令（Tensor Core）

（示意，省略了部分步骤。）和 Halide 相比，TVM 的调度原语多了面向加速器的几个：bind（把循环映射到 GPU 的线程层次）、cache_read / cache_write（经过指定的存储层次）、tensorize（把一小块循环嵌套整体替换成一条硬件指令或一个微内核，如 Tensor Core 的 mma、第 9 讲 NKI 的矩阵引擎）、虚拟线程（隐藏延迟）。

调度空间太大，手写调度不现实，TVM 的回答是自动调优：

• AutoTVM：人写一个调度模板（哪些原语、参数待定），机器搜索参数；用 XGBoost 训练代价模型、模拟退火选候选、真机测量。
• Ansor（Zheng et al. OSDI 2020）：不再需要人写模板。先根据算子的结构自动生成”草图”（sketch，高层的分块结构），再随机填参数，用进化搜索 + 学到的代价模型精化；一个任务调度器在整个网络的多个子图之间分配调优时间。论文报告在 Intel CPU、ARM CPU、NVIDIA GPU 上比当时的最好结果快最多 3.8 倍、2.6 倍、1.7 倍。
• MetaSchedule（TVM 后来的版本）：把调度写成一个”带随机选择的程序”，统一了手写模板和自动生成。

5. Triton 编译器替你做了什么 · Inside the Triton compiler

第 10 讲 §10 说过，Triton 的程序员写”一个块做什么”，把块内的事交给编译器。具体是哪些事？Triton 把块级的程序降低成一系列 MLIR 方言（Triton IR → Triton GPU IR → LLVM IR），沿途决定：

• 布局（layout）：一个 BLOCK_M × BLOCK_K 的块怎样分布到线程和寄存器上——“阻塞布局”让相邻线程读相邻地址（合并访存，第 10 讲 §4），“MMA 布局”满足 Tensor Core 片段的要求（第 10 讲 §8.1），两者之间需要时插入布局转换（经过共享内存）。
• 共享内存：哪些块要放进共享内存、用什么 swizzle 避免 bank 冲突（第 10 讲 §8.2）、怎样复用共享内存（生存期分析，和第 12 讲 §7 的内存规划同一个问题）。
• 软件流水线：num_stages 级的异步拷贝（Ampere 的 cp.async、Hopper 的 TMA），把下一块的读和这一块的计算重叠（第 10 讲 §7、§8.4）。
• Tensor Core 指令：tl.dot 变成 mma / wgmma；num_warps 决定块内的 warp 怎样分工。

留给程序员（和自动调优）的是块的大小、warp 数、流水线级数、启动顺序——恰好是第 9 讲说的”最重要的那部分”（分块），也是 §6 的调度空间。TorchInductor 为大部分融合内核生成 Triton 代码，所以这一层的质量直接决定了 torch.compile 的性能。

6. GPU 矩阵乘法的调度空间 · The schedule space of a GPU GEMM

代码：03_gpu_schedule_space.py（时间来自解析模型，不是实测）

一个 Triton 风格的矩阵乘法有这些旋钮：BLOCK_M、BLOCK_N（16–256）、BLOCK_K（16–128）、num_warps（1–8）、num_stages（1–5）、GROUP_M（1、4、8），共 6000 个组合。其中只有 3744 个（62%）能启动：1200 个每线程寄存器超过 255（累加器太多）、900 个 warp 分到的块比一条 mma（16×8）还小、384 个共享内存超过 227 KB。这些不用测量就能判断，调优前先剪掉。

为了研究搜索，我们写了一个解析模型（sched_lab.gemm_time）当”硬件”：资源与占用率、每个 K 步的 Tensor Core / 共享内存 / L2 时间（第 10 讲 §8.3 的复用账）、Little 定律与流水线级数、波次与尾部、显存的波次模型、补齐的浪费。它给 4096³ 的最高约 900 TFLOPS，是 H100 峰值的九成，比真实的 cuBLAS 乐观——它的作用是给调度空间一个有合理结构的代价，不要拿它的绝对数字去预测真机。

地形（4096³）：最快的一组配置都是 256×128 的块、8 个 warp，BK 和 stages 合起来能藏住延迟即可。好的配置很稀少：3744 个里只有 15 个在最快的 5% 以内，2 倍以内的 336 个；随便选一个（中位数）要慢 5.2 倍。

在最优点附近一次改一个参数：块小一号慢 1.5–3 倍；BM 再小（16、32）或 warp 少于 8，就不能启动了（前者 warp 的块比 mma 小，后者累加器放不进寄存器）；stages 从 1 到 4 级，时间从 3.67 倍降到 1.00 倍，5 级不再变好。参数之间的相互作用很强，“一次调一个”的坐标下降容易卡住。

形状一变，最优配置就变：

前三种形状的最优恰好相同；把它用在 decode 形状上慢 2.9 倍（M = 16 要补齐到 256 行，而且只有 32 个块，大部分 SM 闲着）；反过来，decode 的最优用在其他形状上慢 6.6–14 倍。decode 的矩阵乘法受限于读权重（第 1 讲），要让所有 SM 都在读，所以用很小的块。这就是为什么 cuBLAS 为不同形状准备了很多内核和一套启发式，Triton 的 @triton.autotune 按 key（M、N、K）分别调优、缓存结果。

7. 自动调优：搜索与代价模型 · Autotuning

代码：04_autotuning.py

7.1 学习的代价模型

真实的调优每”测一次”要编译、运行一个内核（几秒）。代价模型用已经测过的配置训练一个回归模型，预测没测过的配置有多快，只把最有希望的拿去真测。代码 04 用最简单的模型——在 log(时间) 上做岭回归，特征是各参数的 log2 及两两乘积：

• 几十个测量就能训练出一个排序很准的模型：调度参数和时间之间的关系不光滑，但有很强的结构。
• 换形状：相近的形状（prefill 的 up 投影）上还能用，换到受限于带宽的 decode 形状就几乎没用了。Ansor / MetaSchedule 用”从程序里抽出来的特征”（循环长度、访存量、向量化宽度……）训练跨形状、跨算子的模型来缓解。
• 先剪枝再学习：如果把不能启动的配置（时间 = ∞，目标值记为”见过的最慢 × 2”）也混进训练集，同样的模型排序相关掉到 0.21（64 个）/ 0.36（256 个）。

7.2 搜索策略

三种策略，预算 16–256 次测量，20 次重复，测量带 2% 的对数正态噪声（找到的配置用无噪声的时间评价）：

到最优的 5% 以内，中位数需要：代价模型引导 24 次、进化搜索 48 次、随机搜索 196 次（20 次重复里只有 13 次在 256 次以内做到）。随机搜索的这个数可以手算：3744 个里有 15 个在 5% 以内，随机抽 256 个至少碰到一个的概率是 $1 - (1 - 15/3744)^{256} \approx 64\%$——和观察到的 13/20 一致。

7.3 调优的成本

一个推理服务：每层 4 种矩阵乘法 × decode 的 16 个批大小桶 + prefill 的 8 个长度桶 = 96 个形状。每次测量 3 秒：随机搜索到 5% 以内约 20.5 小时，代价模型引导约 3.8 小时。降低成本的办法：

1. 缓存：调好的结果按（算子、形状、GPU）存起来——Triton 的 autotune 缓存、TVM 的调优日志；
2. 按桶调优，实际形状向上补齐（第 12 讲 §8）；
3. 只在热点上调优：Ansor 的任务调度器按”这个子图占总时间多少、还有多大提升空间”分配预算；
4. 启发式和预先调好的表：cuBLASLt 内置按形状选内核的启发式；Triton 的 matmul 有一组默认配置。

自动调优还有一个常被忽略的问题：测量噪声。最优附近有好几个只差 1%–2% 的配置，2% 的噪声足以让调优器选错；真实的调优要预热、重复测量、取中位数。第 14 讲评测 LLM 写的内核时会遇到同样的问题。

8. 稀疏张量编译（TACO）· Sparse tensor compilation

代码：05_sparse_compilation.py

8.1 格式决定循环

前面所有的调度都假设张量是稠密的：循环可以随便分块、交换。稀疏张量只存非零元，循环必须”跟着数据结构走”。手写稀疏代码的困难在于组合爆炸：每种表达式 × 每种存储格式 × 每种循环顺序，都要一份手写的代码。

TACO（Kjolstad, Kamil, Chou, Lugato, Amarasinghe, OOPSLA 2017）的想法：张量的每一维各选一种格式——稠密（dense）或压缩（compressed），后来还有 singleton、hashed 等——编译器根据”表达式 + 每一维的格式”推导出遍历代码。

A =  5 0 0 1 0 0
     0 0 0 0 0 0
     0 2 0 0 0 3
     0 0 0 0 4 0
CSR  = (dense, compressed)：     pos = [0, 2, 2, 4, 5]，crd = [0, 3, 1, 5, 4]，vals = [5, 1, 2, 3, 4]
DCSR = (compressed, compressed)：crd0 = [0, 2, 3]（只记非空的行），pos1 = [0, 2, 4, 5]，crd1 = [0, 3, 1, 5, 4]

压缩的一维用 pos（第 $i$ 行的非零在 crd / vals 里的范围）和 crd（列号）描述。每个非零元要 8 字节（值 + 列号），比稠密的 4 字节多一倍——密度 50% 时 CSR 和稠密一样大。

$y = Ax$ 的代码由 $A$ 的格式决定（代码 05 演示 2 生成的）：

# CSR：第 0 维稠密 → 遍历所有行；第 1 维压缩 → 只遍历这一行的非零
for i in range(n):
    for p in range(pos[i], pos[i + 1]):
        j = crd[p]; acc += vals[p] * x[j]
# DCSR：第 0 维也压缩 → 只遍历非空的行
for p0 in range(pos0[0], pos0[1]):
    i = crd0[p0]
    for p1 in range(pos1[p0], pos1[p0 + 1]): ...

工作量：稠密 $nm$；CSR = 非零 + 行数；DCSR = 非零 + 非空行数。超稀疏的矩阵（2000 行里只有 2% 非空）上，CSR 光遍历空行就占了大半（2,760 vs 811 次迭代）。

8.2 合并：加法是并集，乘法是交集

两个稀疏操作数时，内层要同时走两个有序的列号序列（代码 05 演示 3 生成的 C = A + B）：

while pa < ea and pb < eb:                # 两边都没走完
    ja, jb = crd_A[pa], crd_B[pb]; j = min(ja, jb)
    if ja == j and jb == j:  输出 A + B；两边都前进
    elif ja == j:            输出 A（只有 A 有）；A 前进       ← 乘法时：不输出，只前进
    else:                    输出 B；B 前进                    ← 乘法时：同上
while pa < ea: 抄 A 剩下的                 ← 乘法时：不需要，直接结束
while pb < eb: 抄 B 剩下的

加法时”只有一边有”的元素也要输出（并集），乘法时只有”两边都有”才输出（交集）。TACO 用合并格（merge lattice）系统地推导这些分支：表达式里的 + 是并集、× 是交集，嵌套的表达式组合成一张格，每个格点对应生成代码里的一个 while 循环。它还利用格式做优化：$A \circ B$ 里如果 $B$ 是稠密的，就只遍历 $A$ 的非零、按下标直接取 $B$。结果的非零个数：密度 5% × 5% 时，加法约 9.7%（两者之和减去重叠），乘法只剩约 0.25%。结果还要”装配”成压缩格式，TACO 把计算和装配一起生成。

8.3 稀疏什么时候划算

代码 05 演示 4 在 CPU 上实测（$n = 4096$、FP32、单线程，两种代码都用 -O3 -ffast-math 编译；作者机器上一次运行）：

密度很低时 CSR 快得多；但每个非零元要多读一个列号、按列号随机访问 $x$（gather），内层循环长度不规则，编译器向量化不了，而稠密的代码连续读、向量化——在这台机器上，密度高于 10%–15% 稀疏就不划算了。

一个陷阱：只开 -O3 时，稠密代码里 acc += A[i][j]·x[j] 的浮点归约不允许重新结合，编译器不向量化，慢了 8 倍（8.2 ms vs 0.99 ms）。拿它当基准，会得出”密度 50% 的稀疏也更快”的错误结论——这个脚本的第一版就是这样。比较两种实现时，基准本身必须是好的。

2:4 结构化稀疏（第 3 讲）密度正好 50%：用 CSR 这种通用格式存，比稠密还慢。它的价值在于硬件：Ampere 起的稀疏 Tensor Core 用 2 位的元数据记录每 4 个里哪 2 个非零，在矩阵乘法单元里直接跳过，理论上快 2 倍。LLM 里的”稀疏”大多是结构化的：MoE 的专家路由（按专家分组的稠密矩阵乘法）、分块稀疏的注意力（FlashAttention 跳过整块）、分页 KV 的块表（第 7 讲）——都把不规则性挪到”块”的粒度上，块内仍然是稠密、规则的计算。稀疏编译的研究也在往这个方向走：带块结构的格式、为加速器生成稀疏代码（TACO 之后的 SAM、Stardust 等工作，Kjolstad 组）。

9. 现状（2025–2026）· State of the field

• MLIR 成为公共的基础设施：Triton、IREE、OpenXLA 的一部分、各家加速器的编译器都建在 MLIR 上；不同层次的 IR（张量、循环、硬件指令）用”方言”表示，降低（lowering）是一系列方言之间的转换。
• 布局成为一等公民：CUTLASS 的 CuTe 把”数据在线程、寄存器、共享内存之间怎样分布”写成可以组合、推导的代数对象（layout algebra）；Triton 的布局系统也在往这个方向发展（2025 年的”linear layouts”）。第 10 讲的 swizzle、mma 的片段布局都是布局。
• 瓦片级语言：ThunderKittens、TileLang、Triton 的 Gluon、CuTe DSL——介于”手写 CUDA”和”全自动编译器”之间，程序员控制分块和数据搬运，编译器处理线程级的细节。
• 自动调优无处不在：TorchInductor 的 max-autotune 为矩阵乘法在 Triton 模板和 cuBLAS 之间选择、为 Triton 内核调块大小；TVM 的 MetaSchedule；各家推理引擎离线调好按形状索引的配置表。
• LLM 进入循环：既然”提出一个候选 → 编译 → 测量 → 反馈”是一个闭环，候选可以由 LLM 生成，而且不限于固定参数空间——第 14 讲的 KernelBench 和 AlphaEvolve。

10. 小结 · Summary

1. 后端为每个融合好的内核生成代码：分块、顺序、并行映射、存储层次、流水线。库、模板、生成三种做法；生成的核心是算法与调度分离。
2. Halide：算法是定义在整个平面上的纯函数；调度用 compute_at / store_at 决定生产者在哪里算。root、inline、分块、滑动窗口在并行性、局部性、重复计算之间取舍；边界推断自动算出每一级的区域。实测的最优调度取决于机器（这台机器上方块分块反而最慢），所以要自动调优。
3. 多面体模型：迭代空间是整数点集，依赖是距离向量，变换是仿射映射。合法 ⟺ 依赖仍字典序为正；可分块 ⟺ 每个分量非负；内层并行 ⟺ 外层带走所有依赖。倾斜让 Jacobi 可以时间分块（读缺失少约 80%），波前让 Gauss–Seidel 内层可以并行。只适用于仿射的程序。
4. TVM 的调度原语加上 bind、cache_read、tensorize；AutoTVM（模板 + 模型）、Ansor（自动草图 + 进化 + 模型）。Triton 把布局、共享内存、流水线、Tensor Core 指令交给编译器，块大小等交给自动调优。
5. GPU 矩阵乘法的调度空间里，好的配置很稀少（4096³：3744 个能启动的里 15 个在 5% 以内），而且随形状变（decode 的细长矩阵乘法要小块）。
6. 自动调优：学习的代价模型几十个测量就能排得很准（Spearman 0.9+），引导搜索比随机搜索省一个数量级的测量；但换形状就不准了，而且要先剪掉不能启动的配置。调优的成本靠缓存、分桶、只调热点、启发式控制。
7. 稀疏编译（TACO）：每一维的格式决定生成的循环；加法并集合并、乘法交集合并。通用的稀疏格式只有在密度很低时才划算（这台机器上 < 10%–15%）；2:4 这样的结构化稀疏要靠硬件。

11. 常见误区 · Common pitfalls

1. “把中间结果融合进消费者总是更好”——inline 会重复计算（4 级模糊内联后每点 7380 次运算），分块会多算光晕、多占存储；在 cache 大的机器上 root 可能最快（§2.3、§2.4）。
2. “分块一定能提升局部性”——方块分块在我们的机器上比 root 慢 2.5 倍；块的形状要配合硬件预取和连续访问（§2.3）。
3. “只要能想到的循环变换都可以做”——变换必须保持依赖；Jacobi 直接分块就算错了，要先倾斜（§3.2）。
4. “多面体模型能处理任何循环”——它只处理仿射的下标和边界；稀疏、间接寻址、数据依赖的循环不在其中（§3.4）。
5. “一个形状上调好的配置可以到处用”——大 GEMM 的最优用在 decode 形状上慢约 3 倍，反过来慢 6.6–14 倍（§6）。
6. “代价模型可以替代测量”——它擅长排序和挑候选；换形状就不准，最终仍要真测（§7.1）。
7. “测一次就够了”——2% 的噪声足以在几个相近的配置之间选错（§7.3）。
8. “稀疏就会快”——CSR 每个非零元多一个列号、访问不规则；密度过一定阈值就比稠密慢；比较时基准必须是好的（§8.3）。
9. “50% 的稀疏能省一半时间”——在通用格式和通用硬件上不能；2:4 的 2 倍来自专门的硬件（§8.3）。

12. 练习 · Exercises

练习 1（边界推断） 三级流水线：f(x, y) = in(x−1, y) + in(x+2, y)；g(x, y) = f(x, y−1) + f(x, y+1)；输出 h(x, y) = g(x+1, y) · 2。要算 h 的 [0, 32) × [0, 32)，g、f、in 各要算哪个区域？

答案：g：x 偏移 +1 → [1, 33) × [0, 32)。f：y 偏移 −1、+1 → [1, 33) × [−1, 33)。in：x 偏移 −1、+2 → [0, 35) × [−1, 33)。

练习 2（分块的重复计算） 两级模糊（blurx 只在 y 方向有光晕），分块 T × T。blurx 每个输出点算几次？T = 8、64、256 呢？

答案：每块算 T × (T + 2) 个 blurx，比例 (T + 2)/T：T = 8 时 1.25，64 时 1.031，256 时 1.008。

练习 3（内联的爆炸） L 级 3×3 方框模糊（每级 9 次运算）全部内联，每个输出点多少次运算？验证 L = 4 时是 7380。

答案：第 L 级算 1 次（9 次运算），它需要第 L−1 级的 9 个点（每个 9 次运算），第 L−2 级的 81 个点……合计 $9(1 + 9 + 81 + \dots + 9^{L-1}) = 9(9^L - 1)/8$。L = 4：$9 \times 6560 / 8 = 7380$。

练习 4（合法性） 依赖距离 (1, −1) 和 (0, 1)。交换两层循环合法吗？能直接矩形分块吗？

答案：交换后 (−1, 1) 和 (1, 0)：(−1, 1) 字典序为负，非法。原样有负分量 −1，不能直接分块；倾斜 $i' = i + t$ 后变成 (1, 0)、(0, 1)，可以分块。

练习 5（倾斜因子） 依赖 (1, −2)、(1, 0)、(1, 2)（每个时间步看左右各两格）。最小的倾斜因子 $s$（$i' = i + s\,t$）是多少才能矩形分块？

答案：变换后 (1, s − 2)、(1, s)、(1, s + 2)，要 s − 2 ≥ 0，所以 s = 2。

练习 6（波前的并行度） N × N 的 Gauss–Seidel，用波前 (i + j, i)。外层多少步？平均每步多少个点可以并行？

答案：$i + j$ 从 2 到 2N，共 2N − 1 步；平均 $N^2 / (2N - 1)$——N = 64 时约 32（代码 02 演示 3）。

练习 7（共享内存的限制） BM = 256、BN = 128、BK = 64、stages = 4，FP16。共享内存要多少？BK = 128 呢？H100 每块最多 227 KB。

答案：$4 \times (256 + 128) \times 64 \times 2 = 196{,}608$ 字节 = 192 KiB，放得下；BK = 128 时 384 KiB，放不下（代码 03 演示 4 的”—”）。

练习 8（寄存器的限制） BM = 256、BN = 128，4 个 warp（排成 2×2）。每个 warp 算多大的块？每个线程要多少个 FP32 累加器？8 个 warp（4×2）呢？

答案：4 个 warp：每个 128 × 64，$128 \times 64 / 32 = 256$ 个累加器，超过 255 个寄存器，不能启动。8 个 warp：64 × 64，128 个累加器 + 约 40 个其他寄存器 = 168，可以。

练习 9（随机搜索的概率） 3744 个配置里有 15 个在最优的 5% 以内。随机（不重复）抽 256 个，至少抽中一个的概率约是多少？

答案：$1 - (1 - 15/3744)^{256} \approx 1 - e^{-1.03} \approx 64\%$。代码 04 的 20 次重复里有 13 次做到（65%）。

练习 10（CSR 的存储） 4096 × 4096 的 FP32 矩阵，列号用 4 字节整数。密度多少时 CSR 和稠密一样大？

答案：稠密 $4n^2$；CSR 约 $8 \cdot \text{nnz} + 4(n + 1)$。$8 \rho n^2 = 4 n^2$ → $\rho = 50\%$（忽略 pos）。

练习 11（合并） $C = A \circ B$（逐元素乘），A 是 CSR，B 是稠密的。生成的代码应该怎样遍历？工作量多少？

答案：只遍历 A 的非零：for i: for p in pos[i]..pos[i+1]: j = crd[p]; C[i, j] = vals[p] * B[i, j]。工作量 = A 的非零 + 行数，和 B 的大小无关——乘法是交集，稠密的那一边”总是有值”，不需要合并循环。

练习 12（调优的成本） 96 个形状，每个形状代价模型引导需要 48 次测量，每次 3 秒。总共多久？如果每个形状只调一次、结果缓存，下次部署同样的模型和 GPU 要多久？

答案：$96 \times 48 \times 3 = 13{,}824$ 秒 ≈ 3.8 小时；缓存命中时几乎为 0——这就是 Triton autotune 缓存、TVM 调优日志的意义。

13. 延伸阅读 · Further reading

主线

• J. Ragan-Kelley, C. Barnes, A. Adams, S. Paris, F. Durand, S. Amarasinghe, Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines, PLDI 2013.
• T. Chen et al., TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018. https://arxiv.org/abs/1802.04799
• L. Zheng et al., Ansor: Generating High-Performance Tensor Programs for Deep Learning, OSDI 2020. https://arxiv.org/abs/2006.06762
• P. Tillet, H. T. Kung, D. Cox, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019.
• F. Kjolstad, S. Kamil, S. Chou, D. Lugato, S. Amarasinghe, The Tensor Algebra Compiler, OOPSLA 2017；S. Chou, F. Kjolstad, S. Amarasinghe, Format Abstraction for Sparse Tensor Algebra Compilers, OOPSLA 2018.

更多

• U. Bondhugula et al., A Practical Automatic Polyhedral Parallelizer and Locality Optimizer（Pluto），PLDI 2008；P. Feautrier, Some efficient solutions to the affine scheduling problem, 1992；S. Verdoolaege, isl: An Integer Set Library for the Polyhedral Model, 2010。
• A. Adams et al., Learning to Optimize Halide with Tree Search and Random Programs, SIGGRAPH 2019（Halide autoscheduler）。
• J. Shao et al., Tensor Program Optimization with Probabilistic Programs（MetaSchedule），NeurIPS 2022。
• C. Lattner et al., MLIR, CGO 2021；NVIDIA CUTLASS / CuTe 文档（layout algebra）。

本仓库

• 前置讲 P4：循环变换、依赖分析、算法与调度分离、代价模型与自动调优（loop_ir.py）。
• 第 9 讲：Interstellar 用 Halide 的调度语言描述加速器——同一种语言既描述软件的调度，也描述硬件的数据流。
• 第 10 讲：手写的 GPU 矩阵乘法里每个”旋钮”的来历。

14. 下一讲预告 · Next lecture

这一讲的自动调优在一个固定的参数空间里搜索：块多大、几个 warp、几级流水线。如果候选不只是参数，而是整个程序呢？LLM 可以读一个 PyTorch 算子、写出一个 CUDA / Triton 内核；把”写 → 编译 → 测正确性 → 测速度 → 反馈”做成一个闭环，就成了一个能搜索程序空间的智能体。第 14 讲讲两件事：KernelBench——怎样公平地评测 LLM 写的内核（fast_p 指标、正确性检查、各种”作弊”方式，以及本讲 §7.3、§8.3 提到的测量陷阱）；AlphaEvolve——LLM 提出修改 + 自动评估器 + 进化搜索，在矩阵乘法算法和数据中心调度上找到了新东西。

→ 第 14 讲：智能体生成内核