第 10 讲 · 高级 ML 编译：计算图级优化 · Advanced ML Compilation: Graph-Level Optimizations (TASO, PET)

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 10 · Wed 10/1 · Advanced ML Compilation: Graph-Level Optimizations。 必读论文：Z. Jia, O. Padon, J. Thomas, T. Warszawski, M. Zaharia, A. Aiken, TASO: Optimizing Deep Learning Computation with Automatic Generation of Graph Substitutions（SOSP 2019）；H. Wang, J. Zhai, M. Gao, Z. Ma, S. Tang, L. Zheng, Y. Li, K. Rong, Y. Chen, Z. Jia, PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections（OSDI 2021）。选读：Y. Yang, P. M. Phothilimthana, Y. R. Wang, M. Willsey, S. Roy, J. Pienaar, Equality Saturation for Tensor Graph Superoptimization（MLSys 2021，TENSAT）。 依据说明：本讲以原课 slides 10-graph-level-optimizations.pdf 的主线组织：

• 计算图级优化的例子：Conv + BN 融合、ResNet 块的一串改写；
• 手写规则不可行的原因；
• TASO：生成器 → 验证器 → 代价驱动的回溯搜索，以及”为什么是超优化器”；
• PET：完全等价与部分等价、变换生成、多线性张量程序、两条定理、修正内核、程序优化器；
• 从等价到不等价的优化谱系。

论文里的数字只做转述。graphsub.py（迷你 TASO：枚举、指纹、剪枝、随机测试验证、pair move、回溯搜索）、PET 例子的实现、“弱算力设备”以及全部实验，都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture10/（3 个脚本 + 共享模块 graphsub.py；只需 numpy；全部在 15 秒以内） 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标

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

1. 说清楚什么是计算图替换（graph substitution），为什么手写替换规则在”算子 × 网络结构 × 硬件”的组合下不可行。
2. 描述 TASO 的三个组件，并能用代码的数字说明每一步做了什么：
   • 生成器：枚举小图 → 指纹 → 候选替换 → 剪枝；
   • 验证器：用算子性质做形式化验证；
   • 优化器：代价驱动的回溯搜索，带参数 α。
3. 解释为什么好的图优化常常要”先变差、再变好”，以及 α 怎样在搜索时间和结果之间折中。
4. 区分优化器（optimizer）和超优化器（superoptimizer）。
5. 解释 PET 的部分等价变换和自动修正：
   • 什么是多线性张量程序（MLTP）；
   • 两条定理怎样把”检查所有位置 × 所有输入”变成”每个盒子几个位置 × 几次随机测试”；
   • 修正内核怎样恢复完全等价。
6. 把 TASO、PET 和量化、剪枝等”不等价但有精度保证”的优化放在同一个谱系里比较。

0.2 先修与本讲范围

0.3 本讲地图

§1  计算图级优化：从手写规则到自动生成
§2  TASO 概览
§3  替换生成器：枚举、指纹、剪枝                         ← 代码 01
§4  替换验证器：算子性质 + 定理证明器
§5  代价驱动的回溯搜索：为什么要允许暂时变差              ← 代码 02
§6  PET：部分等价的变换                                   ← 代码 03（演示 1）
§7  PET 的验证与修正：两条定理、盒子、修正内核             ← 代码 03（演示 2–4）
§8  从等价到不等价
§9  论文精读：TASO；PET；（选读 TENSAT）
§10 小结  §11 常见误区  §12 练习  §13 延伸阅读  §14 下一讲预告

0.4 关于 graphsub.py

代码用嵌套元组表示计算图里的张量表达式（“项”）：("matmul", None, A, B) 就是 A @ B。结构相同的子表达式自动被视为同一个节点（这等于做了公共子表达式消除），多输出的图就是一个”输出元组”。这种表示天然支持搜索需要的三种操作：比较两张图、哈希、模式匹配。

代价模型沿用第 1 讲的 roofline：每个算子按 max(FLOP / 峰值, 字节 / 带宽) 估算，再加 4 µs 的内核启动开销。只依赖权重的子图在编译时就算好（常量折叠），代价为 0。TASO 的代价模型则是在目标硬件上实测每个算子一次，再把图里所有算子的时间加起来。两者思路相同，只是数字来源不同。

1. 计算图级优化：从手写规则到自动生成 · Graph-level optimizations

1.1 什么是图替换

一个 DNN 是一张计算图：节点是张量代数算子（卷积、矩阵乘法……），边是张量。图替换（graph substitution）把图里匹配某个模式的子图，换成一个功能等价、但更快的子图。

最经典的例子是 slides 里的 Conv + BN 融合（第 2 讲做过）。推理时 BatchNorm 的均值、方差、缩放、偏移都是常数，所以

$\text{BN}(\text{Conv}(X, W)) = \text{Conv}\Big(X, \tfrac{\gamma}{\sigma} W\Big) + \Big(\beta - \tfrac{\gamma \mu}{\sigma}\Big),$

BN 被折进卷积核和偏置里，少了一个内核和一次读写。

1.2 一串改写：先变差，再变好

slides 第 9–11 页用一个 ResNet 式的块演示了一串改写：

• 左支：conv3x3 → relu → conv3x3；
• 右支：conv1x1 → relu；
• 两支相加，再 relu。

改写过程是：

1. 放大卷积核（enlarge）：把 1×1 卷积核四周补 0，变成 3×3。单看这一步是变慢的，计算量变成 9 倍。
2. 合并卷积：现在两个读同一个输入的卷积核大小相同了，可以沿输出通道拼成一个卷积，算完再拆开（split）。
3. 融合卷积和加法、融合卷积和 relu。

slides 指出：最终的图在 V100 上快 30%，在 K80 上却慢 10%。好不好取决于硬件。

1.3 为什么手写规则不可行

slides 把挑战概括成三个维度的乘积：

• 算子：新算子不断出现（分组卷积、深度可分离卷积、各种注意力）；
• 网络结构：NAS 搜出来的网络、Transformer 的各种变体；
• 硬件后端：每一代 GPU、每一种加速器的最优选择都不同。

为每一种组合手写规则不可能。TensorFlow 的手写图优化约 5.3 万行代码，而且手写的规则可能有微妙的错误（第 2 讲演示过：训练模式的 BN 不能折叠）。

2. TASO 概览 · Tensor Algebra SuperOptimizer

TASO 的关键想法：用自动生成、自动验证的替换，取代人工设计的图优化。 slides 列出三个好处：

• 工程量少：手写图优化 5.3 万行 → TASO 只需为每个算子写参考实现和数学性质，约 1,400 行。
• 性能更好：论文报告在 cuDNN 后端上比现有框架快 1.3–2.8 倍，在 TVM 后端上快 1.1–1.8 倍。
• 正确性更强：所有生成的替换都经过形式化验证。

工作流程分三步：

算子规格 ──→ 替换生成器 ──候选替换──→ 替换验证器 ──验证过的替换──→ 图优化器 ──→ 优化后的图
（参考实现 + 性质）  （枚举 + 指纹）            （算子性质 + Z3）            （代价驱动的回溯搜索）

前两步对每组算子离线做一次，生成的替换可以用在任何网络上；第三步对每个网络在线做。

3. 替换生成器：枚举、指纹、剪枝 · Substitution generator

3.1 枚举所有小图（代码 01 演示 1）

给定硬件后端支持的一组算子，枚举所有不超过 k 个算子的计算图：每次加一个算子，它的输入可以是初始的输入张量，也可以是前面算子的输出。同时排除”两个算子在同样的输入上做同样的事”这种重复计算。TASO 枚举到 4 个算子，slides 说有约 6,600 万个图。

代码只枚举到 3 个算子、单输出的图，分两个”世界”：

卷积世界加了”种类”规则：卷积的输入必须是激活、卷积核必须是权重。没有这条规则，生成器会枚举出”把卷积核当图片来卷积”这种无意义的图，三个算子的图会多出 9 倍。

3.2 指纹：把”两两比较”变成”分组”（演示 2）

两两比较 4 万个图要比约 8.6 亿对。TASO 的做法是给每个图算一个指纹：在一组固定的随机输入上求值，对输出取哈希（多输出的图用对输出顺序不敏感的对称哈希）。指纹不同的两个图一定不等价，所以只需比较指纹相同的图——同一组里任意两个图构成一个候选替换。

slides 说 TASO 就这样从枚举的图里得到了 28,744 条候选替换。

为什么用整数：浮点数的结合律不严格成立，(A @ B) @ C 和 A @ (B @ C) 会差最后几位。指纹用整数算就完全相同。

整数的范围要够大：如果只取 [-2, 2]，relu 之后大部分元素是 0，不等价的图很容易”撞车”。例如 relu(C) 和 A ⊙ relu(C)：

• 每个元素上两者相等的概率是 P(C ≤ 0) + P(C > 0)·P(A = 1) = 3/5 + 2/5·1/5 = 0.68；
• 16 个元素都相等的概率 0.68¹⁶ ≈ 0.21%；代码实测 2 万组里 51 组（0.26%）；
• 取 [-50, 50] 时降到 2e-5。

3.3 剪枝：去掉冗余的替换（演示 3）

很多候选替换是冗余的，即可以由别的替换推出来。TASO 用两种剪枝：

1. 输入重命名：(A + B) ⇔ (B + A) 和 (A + C) ⇔ (C + A) 只是名字不同，留一条。
2. 公共子图：
   • 如果两边有一个公共的子图，把它换成新的输入变量后替换仍然成立，就保留那条更一般的。例如 ((A @ B) + C) ⇔ (C + (A @ B)) 的公共部分是 A @ B，换成 D 就是交换律。
   • 如果公共部分包括了所有输出（两边的根是同一个算子、对应的子图等价），这条替换可以由子图上的替换推出。例如 relu(A + B) ⇔ relu(B + A)。

我们剩下的比 TASO 多。我们的公共子图剪枝只把一个公共子图换成新变量，TASO 的实现更彻底；剩下的很多是交换律、结合律的各种组合。

论文强调，剪枝不会丢掉任何优化机会：如果图 G 能通过一串替换变成 G’，剪枝之后仍然能（可能换一串替换）。冗余的替换不影响正确性，只会让搜索变慢。

3.4 机器找到了什么（演示 5）

验证通过的替换里能找到所有这些”教科书上的”改写，没有一条是人写进去的：

生成器只知道每个算子怎么算，不知道任何代数规律。

TASO 论文 §7.3 展示了几条它找到、而当时任何框架都没有的替换。例如借助一个常量卷积核，把两个平均池化加一次逐元素加法换成一个深度卷积（depth-wise convolution），提高了算子粒度，减少了内核启动。

4. 替换验证器：算子性质 + 定理证明器 · Substitution verifier

指纹相同只说明”在一组输入上相同”。TASO 分两步把关：

1. 随机测试：同一组里的图再在几组 [-1, 1] 的浮点随机输入上比较，差异超过 1e-5 的淘汰。论文说在这个阈值下，没有发现和整数测试不一致的情况。
2. 形式化验证：每个算子的数学性质写成一阶逻辑。例如”卷积对卷积核是双线性的""卷积对拼接满足分配律”，全部 43 条见论文表 2。然后用定理证明器 Z3 证明：每条候选替换的两边，对所有形状、所有输入都相等。

验证的代价：生成 743 条替换约 5 分钟，验证约 10 分钟。

这种方法还有两个附带的好处：

• 生成器产生了证不出来的替换，说明要么生成器有 bug，要么缺了算子性质——TASO 的作者就这样发现了生成器里的几个 bug；
• 性质本身也要检查：TASO 用 Z3 检查性质之间是否一致、有没有冗余，并在小尺寸上对照参考实现验证每条性质。

支持一个新算子，通常只需要专家写几条性质，花几个小时。

代码里的验证：我们没有实现定理证明器，只做了随机测试（8 组整数 + 4 组浮点）。演示 4 里被淘汰的替换很有教育意义：

• 矩阵世界淘汰 26 条。例如 A + A⊙A ⇔ relu(A + A⊙A)：x + x² = x(x+1) 对任何整数都 ≥ 0，所以在整数输入上 relu 不起作用；但对实数 x = −0.5 它是 −0.25。浮点测试一下就把它淘汰了。
• 卷积世界淘汰 43 条。例如指纹输入的 1×1 卷积核 K = [[29, 38], [−42, −45]]，每个输出通道的两个权重碰巧同号，于是 relu 放在卷积前后结果一样。张量越小，这种巧合越容易发生。

随机测试是概率性的，不是证明。第 11 讲的 Mirage 和 §7 的 PET 会给随机测试一个错误概率的上界。

5. 代价驱动的回溯搜索 · Cost-based backtracking search

5.1 算法

有了验证过的替换，怎样用在真实的网络上？最直接的做法是贪心：每次套用一条能让图变快的替换，直到没有可用的为止。

但 §1.2 的例子说明，好的优化常常要先变差、再变好：放大卷积核这一步单看是变慢的，贪心永远不会走。

TASO 采用 MetaFlow 的代价驱动回溯搜索（论文 Algorithm 2）：

P ← {原图}（按代价排序的优先队列）
while P 不空（且没超过预算）:
    G ← P 里代价最小的图
    for 每条替换 s，每个能匹配的位置:
        G' ← 在 G 上应用 s
        if Cost(G') < Cost(G_opt): G_opt ← G'
        if Cost(G') < α × Cost(G_opt): 把 G' 放进 P
return G_opt

• α：允许”暂时变差”的程度。α = 1 只接受严格变好的图，等于贪心；α 越大，探索的图越多。论文取 α = 1.05，优化一个网络不到 10 分钟。
• 代价：每个算子在目标硬件上实测一次，图的代价 = 所有算子之和。DNN 算子是没有分支的稠密线性代数，同样的形状和布局下耗时非常稳定，所以这个简单的模型足够准。
• 数据布局：TASO 还同时搜索每个张量的数据布局（行优先 / 列优先），因为不同算子偏好的布局不同。一条替换可能对应几种布局，都要试。代码没有实现这一点。
• 环：应用一条合法的替换可能在图里产生环，TASO 要检查。代码的”项”表示是树（带共享），天然没有环。

代码实现时有两个细节：

• 按 Algorithm 2 的字面写法，新的最好图会先成为 G_opt，再和 α × Cost(G_opt) 比较，于是它自己永远进不了队列，α = 1 时贪心只走一步。代码改为和更新前的最好代价比较，α = 1 才是真正的贪心下降。
• 多输出的合并：TASO 枚举多输出的图；代码只枚举单输出的图，另加一个通用的 pair move：把两个读同一个输入的同类节点 Y1、Y2 换成 split(concat(Y1, Y2))。这一步单看是变差的（多了一个拼接），但它让生成器找到的拼接合并规则 [A@B | A@C] → A@[B | C] 能派上用场。每次应用后都和原图做随机测试。

5.2 ResNet 式小块（代码 02 演示 1、2）

slides 的块，N = 32、C = 64、56×56、BF16，在 H100 上：

α = 1.5 找到的改写路径：

第 4 步是生成器找到的一条 3 个算子的规则，一步完成了”放大 1×1 卷积核 + 沿输出通道拼成一个 128 通道的卷积”。拼好的卷积核在编译时算好，运行时只剩一个读一次 X 的卷积。这正是 slides 里”放大卷积（变慢）→ 合并卷积”那条路径，只是机器自己找到的。

关键在第 3 步：代价从 62.2 升到 81.6 µs，是之前最好的 1.31 倍。α = 1.05 不允许这一步，所以找不到；演示 5 里 α = 1.3 也恰好不够。slides 还有一步”卷积 + 加法融合”，TASO 的卷积算子支持，我们的算子集里没有。

换一种硬件：一个假想的弱算力设备（4 TFLOP/s、100 GB/s、启动 10 µs）。

在弱设备上，把 1×1 放大成 3×3 意味着这一支的计算量变成 9 倍，而它是计算受限的。所以 H100 上最好的图拿到弱设备上反而慢 1.34 倍，弱设备上的搜索也不会选它。

这就是 slides 那句”V100 快 30%、K80 慢 10%“的道理：图优化的好坏取决于硬件，所以 TASO 的代价模型在目标硬件上实测每个算子，而不是用一套固定的规则。

5.3 注意力的 Q、K、V 投影（演示 3）

解码时（batch 16，4096 → 4096），三个投影 Q = X·Wq、K = X·Wk、V = X·Wv 是三个独立的矩阵乘法。在 H100、α = 1.5 下，搜索找到了这样一条路径：

最终三个投影合成一个矩阵乘法 X @ [Wk | Wq | Wv]，权重在编译时拼好；省掉了两次内核启动，X 只读一次。BERT 这类模型上，TASO 的收益很大一部分就来自这种合并。

在 H100 上，pair move 让代价最多变成之前最好的 1.11 倍，α = 1.05 找不到。在弱设备上，矩阵乘法本身慢得多，拼接只让代价变成 1.017 倍，α = 1.05 就够了。需要多大的 α，也取决于硬件。

5.4 LoRA：把适配器合并回权重（演示 4）

第 18 讲的 LoRA 在推理时计算 Y = X·W + (X·A)·B，其中 A 是 4096×16，B 是 16×4096。搜索的结果：

最好的图 X @ (W + A @ B) 里，W + A @ B 只依赖权重，编译时算好，运行时只剩一个矩阵乘法。这正是部署 LoRA 时”把适配器合并回权重”的标准做法，机器用两步自己推了出来：

1. 结合律：(X·A)·B → X·(A·B)。代价先变成原图的 1.27 倍（H100）或 1.92 倍（弱设备），因为两个很瘦的矩阵乘法被换成了一个满尺寸的矩阵乘法。
2. 分配律：X·W + X·(A·B) → X·(W + A·B)。代价一下子降到 0.54 倍（H100）。

所以在 H100 上要 α > 1.27，在弱设备上要 α > 1.92，才能走到这一步。

这个例子说明 α = 1.05 并不总是够用：有些重要的优化要求很大的 α，而 α 越大，搜索越贵。演示 5 里，ResNet 例子扩展的图数随 α 增长：α = 1.0 / 1.3 / 1.5 时分别是 3 / 33 / 150 张。

选读的 TENSAT 换了一种思路：等式饱和（equality saturation）用 e-graph 一次性记下所有等价的图，不按代价剪枝，最后再挑最好的。论文摘要报告，它找到的图比当时最好的方法快最多 16%，优化时间平均少 48 倍。

5.5 为什么 TASO 是”超优化器”

slides 的问题：优化器和超优化器有什么区别？

• 优化器（optimizer）：从输入程序出发，贪心地套用一系列已知的优化，逐步改进程序。传统编译器的 pass 就是这样：每个 pass 只做”肯定变好”的事。
• 超优化器（superoptimizer）：目标是为输入程序自动找到最优的程序。它在一个很大的等价程序空间里搜索，用代价模型比较，可以经过”更差”的中间程序。

TASO 自动生成替换、在等价图的空间里做允许变差的搜索，所以是超优化器。这个词最早来自 1987 年 Massalin 的工作：为一小段机器指令穷举找最短的等价指令序列。

6. PET：部分等价的变换 · Partially equivalent transformations

6.1 动机

TASO 只用完全等价的变换：任何输入、任何输出位置都相同。slides 第 32 页把变换分成两类：

PET 的问题是：能不能用部分等价的变换提升性能，同时仍然保证等价？ 它的回答是：部分等价的变换 + 自动生成的修正内核。

6.2 动机例子（代码 03 演示 1）

slides 和 PET 论文图 1 的例子：

1. 把卷积的输入沿宽度切成两半，用 reshape 和 transpose 放进 batch 维；
2. 做卷积；
3. 再用 reshape 和 transpose 放回去。

每一半单独补 0，所以接缝两侧的输出少了邻居的贡献。代码在 N = 1、C = 2、O = 2、H = 6、W = 8 上逐点比较，只有接缝两侧的第 3、4 列不一样，其余 36/48 个位置完全相同。

这样变换的好处是：对 ResNet 后面几层那种空间很小、通道很多的卷积，把空间维挪到 batch 维，可以换来更好的并行度或更适合库函数的形状。slides 说”变换 + 修正”让 ResNet-18 快了 1.2 倍。我们的 CPU 模拟不评估速度。

6.3 PET 的结构

输入程序 ──→ 变换生成器 ──变换后的程序──→ 变换修正器 ──修正后的程序──→ 程序优化器 ──→ 优化后的程序
             （沿用 TASO 的超优化）          （检查哪里不等价、生成修正内核）  （束搜索 + 融合、常量折叠等）

和 TASO 的对应关系：

• 变换生成器 ↔ 替换生成器：同样枚举小程序，但只要求”输入、输出的形状相同”。生成的程序可能完全等价，也可能部分等价。
• 变换修正器 ↔ 替换验证器：TASO 证明等价，PET 找出不等价的区域并修正。

PET 要解决两个难题：

1. 怎样生成部分等价的变换？——沿用 TASO 的超优化。
2. 怎样修正它们？——利用 DNN 计算的多线性。

7. PET 的验证与修正 · Verification and correction

7.1 难点

要知道变换后的程序 g 和原程序 f 哪里不等价，最朴素的做法是：对每一个输出位置，检查所有可能的输入。这需要 O(m × n) 次检查，m 是输出位置数，n 是输入的取值组合数——两者都大到无法枚举。

7.2 多线性张量程序（MLTP）

一个算子如果对每个输入都是线性的，就叫多线性的。例如矩阵乘法、卷积：固定一个输入，对另一个是线性的。只由多线性算子组成的程序叫多线性张量程序（MLTP）。

DNN 计算 = MLTP + 非线性激活，而计算量的大头在 MLTP 部分。PET 只对 MLTP 做部分等价的变换。

7.3 洞察 1：不需要检查所有输出位置（定理 1）

对一个输出位置，看它的”求和区间”：卷积核的哪些偏移量落在有效输入上（而不是落在填充的 0 上）。把求和区间相同的位置归成一个盒子（box）。

以 3×3 卷积、补 1 圈 0 为例，每个空间维分成 3 段：第一个位置、中间、最后一个位置，二维上共 9 个盒子（PET 论文图 4）。

定理 1：对两个 MLTP f 和 g，如果在一个盒子里的 m + 1 个特定位置（起点 v₀ 和沿每一维走一步的 v₀ + eᵢ，m 是输出的维数）上相等，那么在整个盒子上都相等。

直觉是：同一个盒子里，输出只是”同样的求和式整体平移”，多线性程序在平移下的行为由这 m + 1 个点决定。于是只需检查 O(1) 个位置/盒子，盒子的个数不随张量尺寸增长。

7.4 洞察 2：不需要考虑所有输入（定理 2）

定理 2：设 f、g 在位置 v 上不等价，n 是输入张量的个数。从有限域 F（整数模素数 p）里均匀随机取一组输入，f 和 g 在 v 上碰巧相等的概率至多是 n/p。

这是 Schwartz–Zippel 引理的推论：

• MLTP 的输出是输入的多项式，对 n 个输入各是一次的，总次数 n；
• 两个不同的多项式之差是一个非零多项式，它在随机一点上为 0 的概率不超过 次数 / |F|。

用 t 组随机输入测试，非等价的程序全部通过的概率 ≤ (n/p)^t。PET 取 p = 2³¹ − 1。在有限域里算有两个好处：没有浮点舍入误差，也不会整数溢出。

两条定理合起来（PET 论文表 2）：

7.5 代码里的验证（演示 2、4）

代码用的是 N = 2、C = 3、O = 4、H = 6、W = 8 的卷积：

• 原程序沿宽度有 3 个盒子 [0]、[1–6]、[7]；变换后有 6 个 [0]、[1–2]、[3]、[4]、[5–6]、[7]；
• 两者相交得到 6 段，乘以高度方向的 3 段，共 18 个区域。

对每个区域，在定理 1 选出的位置上做 3 次模 p 的随机测试（p = 1,000,003；比 PET 的 2³¹ − 1 小，是为了让 NumPy 的 int64 不溢出）：

• 判出 6 个不等价的区域，都在第 3、4 列（每个区域是一整条 N × O × 高度段），和逐点比较一致。
• 只检查了 66 个位置（全部输出有 384 个），每个位置 3 次测试。
• 非等价程序漏网的概率 ≤ (2/p)³ ≈ 8e-18。

出错概率的上界真的成立吗？ 用很小的素数把概率放大到看得见：在接缝处的一个位置上随机测试 3000 次，

都在上界以下。p = 2³¹ − 1 时，一次测试出错的概率 ≤ 9.3e-10。PET 论文明确说，这是用一个可控的、极小的不可靠概率换取验证速度。

一个”几乎对”的错误变换（卷积核被左右翻转）在 18/18 个区域都被判为不等价。需要修正的区域覆盖了整个输出，修正就等于重算一遍原程序，优化器不会选它。

7.6 修正内核（演示 3）

PET 的修正分两步：

1. 用原程序只重算不等价的盒子。盒子是多维空间里规则的长方体，可以看成原张量的子张量，所以修正内核就是”在子张量上运行原来的算子”。
2. 把修正内核和其他算子融合（例如融合进后面那个 reshape），减少额外的内核启动。

代码里的修正内核取输入的第 2–5 列做一次卷积，保留中间两列，写回输出的第 3、4 列。修正后和原程序完全相同（差为 0）。

额外的计算量 = 要修正的列数 / 宽度：宽度 8 时 25%，56 时 3.6%，112 时 1.8%。slides 说修正（融合之后）带来的开销不到 1%。

7.7 程序优化器

PET 把大程序切成小的子程序，对每个子程序生成变换并修正，然后：

• 用束搜索（beam search）在”变换哪些算子、变换几轮”的空间里挑最好的组合；
• 拼回去之后再做跨子程序的后优化：消除冗余、算子融合、常量折叠。

slides 说优化一个 DNN 不到 30 分钟。

论文的结果（V100）：

• ResNet-18 这种常用网络快 1.2 倍；
• CSRNet、BERT 等新模型最多快 2.5 倍。

8. 从等价到不等价 · From equivalent to non-equivalent optimizations

slides 第 56 页把这一讲放进一个谱系：

TASO                     PET                            剪枝、量化、蒸馏……
完全等价的优化     →     部分等价的优化 + 自动修正   →    不等价、但有精度保证的优化
（功能完全不变）          （修正后功能完全不变）            （功能改变，用精度损失换性能）

越往右，可用的优化越多、潜在的加速越大，但”正确”的含义越弱：

• TASO 保证每个输出都相同；
• PET 允许中间结果不同，但修正后最终输出相同；
• 量化（第 18 讲、A2）、剪枝则让输出改变，只保证精度损失在可接受的范围内。

slides 最后留了三个讨论题：

1. PET 在生成变换上和 TASO 有什么不同？——PET 只要求形状相同，接受部分等价的结果；TASO 只接受完全等价的。
2. PET 在验证 / 修正上和 TASO 有什么不同？——TASO 用算子性质 + Z3 证明对所有输入成立；PET 用两条定理 + 有限域随机测试找出不等价的区域，然后修正。
3. 怎样把图优化和内核优化结合起来？——这正是第 11 讲 Mirage 的主题：FlashAttention 这样的优化无法表示成计算图的改写。

9. 论文精读 · Paper reading

9.1 TASO: Optimizing Deep Learning Computation with Automatic Generation of Graph Substitutions（Jia et al., SOSP 2019）

1. 问题：现有框架（TensorFlow、TensorRT、TVM）的图优化靠专家手写的替换规则：工程量大（TensorFlow 约 5.3 万行）、容易出错、覆盖不全。新算子和新网络结构上尤其明显。

2. 想法：自动生成、自动验证替换，再用代价驱动的搜索应用它们。

• 生成器：枚举不超过 4 个算子的图，用整数输入上的指纹分组得到候选，再用浮点随机测试过滤。
• 剪枝：输入重命名、公共子图，28,744 → 743。
• 验证器：43 条用一阶逻辑写的算子性质 + Z3。
• 优化器：沿用 MetaFlow 的代价驱动回溯搜索（α = 1.05），同时搜索数据布局。

3. 验证：

• 生成 743 条替换约 5 分钟，验证不到 10 分钟。
• 五个网络（ResNet-50、ResNeXt-50、NasNet-A、NasRNN、BERT）上：ResNet-50 与现有框架持平，说明 TASO 能自动发现专家手写的替换；其余四个在 cuDNN 后端上快 1.3–2.8 倍，在 TVM 后端上快 1.1–1.8 倍。
• 加速来自新算子上的新替换，以及替换与布局的联合优化。

4. 从系统的角度看：TASO 把”写优化规则”变成了”写算子的规格”——这是本课程暗线”把专家技巧变成搜索空间 + 代价模型 + 正确性保证”的典型。它的三个组件分工清楚：

• 生成器只需要会算（参考实现）；
• 验证器只需要知道性质；
• 优化器只需要会比较代价。

新增一个算子只要几个小时。

5. 局限：

• 只考虑完全等价的变换（PET 要放宽这一点）；
• 枚举规模随算子数指数增长，4 个算子已经要存约 6,600 万个指纹；
• 回溯搜索受 α 限制，需要大幅变差的优化路径找不到（§5.4，TENSAT 要解决这一点）；
• 只改变图结构，不改变内核内部（第 11 讲）。

9.2 PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections（Wang et al., OSDI 2021）

1. 问题：只用完全等价的变换会错过很多优化机会；部分等价的变换性能更好，却会改变结果、损失精度。

2. 想法：用 TASO 式的超优化生成变换（包括部分等价的），再自动修正：

• 定理 1：多线性程序在同一个盒子里只需检查 m + 1 个位置；
• 定理 2：在有限域上随机测试，出错概率 ≤ (n/p)^t。

这样能快速找出不等价的区域，再用原程序只重算这些区域（修正内核），并把修正内核和相邻算子融合。程序优化器用束搜索组合完全等价和部分等价的变换。

3. 验证：

• 修正开销不到 1%；
• 优化一个网络不到 30 分钟；
• 端到端（V100）：ResNet-18 快 1.2 倍，CSRNet、BERT 等新模型最多快 2.5 倍。

4. 从系统的角度看：PET 把”等价”从每个中间结果放宽到最终输出，打开了一个更大的搜索空间。它的两条定理是”用数学结构换验证效率”的好例子：

• 利用多线性，把无限的检查变成有限的；
• 利用 Schwartz–Zippel，把证明变成随机测试，并给出可控的错误概率。

第 11 讲的 Mirage 会再次用到第二个思想。

5. 局限：

• 只适用于多线性部分，非线性激活要切开处理；
• 随机测试有一个（极小的）出错概率，不是证明；
• 修正内核增加了程序的复杂度；
• 变换的收益依赖底层内核库对特定形状的性能。

9.3 两篇论文之间的联系

PET 直接建立在 TASO 之上：变换生成器沿用 TASO 的超优化，两者的区别在于”接受什么样的变换”和”怎样保证正确”。

两者都把”写规则”变成了”生成 + 验证 + 搜索”，也都只改变计算图的结构。第 11 讲的 Mirage 把同样的思路推广到”图 + 内核 + 线程”三个层次。

10. 小结 · Summary

1. 图替换把匹配的子图换成功能等价、更快的子图。手写规则在”算子 × 结构 × 硬件”的组合下不可行（TensorFlow 约 5.3 万行）。

2. TASO 生成器：枚举小图 → 整数输入上的指纹分组 → 候选替换 → 剪枝。代码在矩阵世界从 92,869 条候选剪到 6,356 条，自动发现了结合律、分配律、拼接合并、放大卷积核等全部 10 条教科书改写。

3. 验证：整数指纹会”撞车”（x + x² 在整数上恒 ≥ 0），所以要浮点测试；TASO 最终靠 43 条算子性质 + Z3 形式化证明。

4. 代价驱动的回溯搜索：α 允许暂时变差。代码里三个例子都要先变差：

   • ResNet 块：中间步 1.31 倍，最后 85.6 → 57.8 µs；
   • Q/K/V 合并：1.11 倍，42.3 → 34.2 µs；
   • LoRA 合并：1.27 倍，39.7 → 21.4 µs。

   需要多大的 α、最优的图是什么，都取决于硬件（H100 的最优图在弱设备上慢 1.34 倍）。

5. PET 允许部分等价的变换，再修正：

   • 定理 1：同一个盒子只需检查 m + 1 个位置；
   • 定理 2：有限域随机测试，出错 ≤ (n/p)^t；
   • 修正内核只重算不等价的盒子，修正后与原程序完全相同。

   代码里 384 个位置只需检查 66 个。

6. 从等价到不等价的谱系：TASO（完全等价）→ PET（部分等价 + 修正）→ 量化、剪枝（有精度保证的不等价）。

11. 常见误区 · Common misconceptions

• “图优化就是把能融合的算子融合起来。” 融合只是其中一类。TASO 找到的替换还包括代数变换（结合律、分配律）、拼接合并、放大卷积核、改变数据布局……而且它们常常要组合起来用。
• “一条替换要么总是好，要么总是坏。” 放大卷积核在 H100 上（访存受限）几乎不增加代价，在弱算力设备上（计算受限）让这一支的计算量变成 9 倍。好坏取决于硬件，所以要按目标硬件的代价搜索。
• “贪心地应用让图变快的替换就能找到最好的图。” 本讲的三个例子都要先经过一个更慢的图。贪心（α = 1）在 ResNet 块上停在 62.2 µs，α = 1.5 能到 57.8 µs。
• “指纹相同就说明两个图等价。” 指纹只是在一组输入上相同。整数输入上 x + x² 恒 ≥ 0，relu 不起作用，但在实数上并非如此。所以 TASO 还要做浮点测试和形式化验证。
• “随机测试不能保证正确。” 不能证明，但能给出出错概率的上界。在有限域上，非等价的多线性程序通过 t 次测试的概率 ≤ (n/p)^t，p = 2³¹ − 1 时小得可以忽略。前提是程序真的是多项式（多线性），并且在有限域上算，而不是用浮点数。
• “部分等价的变换会损失精度。” 加上修正之后，最终输出和原程序完全相同。代码里修正后的差为 0。PET 的”部分等价”只是中间步骤，不是最终结果。
• “α 越大越好。” α 越大，找到的图可能越好，但搜索的图数急剧增加（ResNet 例子：α = 1.0 / 1.3 / 1.5 时扩展 3 / 33 / 150 张图）；α 太大时，没有节点数上限的搜索会被不断变大的图拖垮。

12. 练习 · Exercises（含参考答案）

练习 1　矩阵世界有 3 个 4×4 的输入、2 个一元算子（ᵀ、relu）、4 个二元算子（@、+、⊙、按列拼接）。验证”恰好 1 个算子的图有 42 个”。

答案：

• 一元：2 个算子 × 3 个输入 = 6。
• 二元：4 个算子 × 3 × 3 个有序输入对 = 36。所有组合的形状都合法：4×4 和 4×4 可以相乘、相加，也可以拼成 4×8。
• 共 42 个。

练习 2　（指纹撞车）整数取 [-50, 50] 时，relu(C) 和 A ⊙ relu(C) 在 4×4 输入上指纹相同的概率是多少？为什么 TASO 不担心这种概率？

答案：

• 每个元素上两者相等的概率 = P(C ≤ 0) + P(C > 0)·P(A = 1) = 51/101 + 50/101 · 1/101 ≈ 0.51，16 个元素都相等的概率 ≈ 0.51¹⁶ ≈ 2e-5（代码算出 2.1e-5）。
• 不担心的原因：指纹只用来分组，组内的候选还要经过浮点测试和形式化验证，撞车只会多产生几个候选，不会让错误的替换通过。

练习 3　写出 relu((A + B)) ⇔ relu((B + A)) 为什么被”公共子图”规则剪掉，以及剪掉它为什么不会丢掉优化机会。

答案：

• 两边的根都是 relu，子图 (A + B) 和 (B + A) 等价，所以这条替换可以由 A + B ⇔ B + A（作用在 relu 的输入上）推出来。
• 优化器能把替换用在图的任何位置（包括 relu 的输入），所以有了交换律就不需要这一条。

练习 4　在 H100 上，Q/K/V 合并路径里 pair move 那一步的代价是之前最好的 1.11 倍。如果 α = 1.10，搜索能找到最终的合并吗？如果把内核启动开销从 4 µs 降到 1 µs（例如用 CUDA Graph），需要的 α 会变大还是变小？

答案：

• 不能：1.11 > 1.10，那一步的图进不了队列。
• 启动开销降低后，拼接内核本身更便宜，pair move 让代价变差的比例更小，需要的 α 变小。但合并带来的收益（省两次启动）也变小了。硬件和运行时的参数会改变搜索的地形。

练习 5　LoRA 合并的第一步把 (X·A)·B 换成 X·(A·B)。设 X 是 512×4096，A 是 4096×r，B 是 r×4096，只算 FLOP：r = 16 时这一步的计算量变成多少倍？r 多大时这一步不会变差？

答案：

• 原来是两个瘦矩阵乘法：2·512·4096·r + 2·512·r·4096 = 4·512·4096·r。换成 X·(A·B)：A·B 是常量（编译时算），运行时是 512×4096×4096 的矩阵乘法，2·512·4096·4096。
• 比例 = 4096 / (2r)。r = 16 时是 128 倍（整个图里还有 X·W，所以图的代价只变成 1.27 倍）。
• r ≥ 2048 时这一步才不变差——但那时 LoRA 也就不”低秩”了。合并的收益在第二步（和 X·W 合并）。

练习 6　PET 的例子里，如果把宽度切成 4 段而不是 2 段，有多少列需要修正？修正的额外计算量占多少（宽度 W）？

答案：

• 3 个接缝，每个接缝两侧各 1 列，共 6 列；额外计算量 ≈ 6/W。
• 切得越碎，并行度越高，但修正越多——这是 PET 优化器要权衡的。

练习 7　PET 的定理 2 说非等价程序在一个随机点上相等的概率 ≤ n/p。为什么必须在有限域（整数模 p）上测试，而不是用 [-1, 1] 的随机浮点数？

答案：

• 浮点运算有舍入误差，两个等价的程序（加法顺序不同）也可能不完全相等，只能用阈值比较；而阈值又可能放过”差得很小”的非等价程序。
• Schwartz–Zippel 引理需要在一个域上取样：概率上界是”次数 / 域的大小”，整数模素数 p 是一个大小为 p 的有限域。
• 在有限域里算既精确又不会溢出，结论才有严格的概率保证。

练习 8　（开放）把 TASO 的搜索换成”等式饱和”（TENSAT）：先把所有能推出来的等价图记在一个 e-graph 里，再挑代价最低的。它和 α 回溯搜索相比，优缺点各是什么？

参考思路：

• 优点：不按代价剪枝，“先变差很多、再变好”的路径（LoRA 合并）不会被错过；e-graph 共享子结构，记下大量等价图的开销远小于一张张枚举。
• 缺点：e-graph 可能爆炸，要限制轮数或大小；从 e-graph 里”挑最好的图”本身是一个难的组合优化问题（TENSAT 用整数线性规划），而且多输出、共享子图让代价不再是简单的求和。
• TENSAT 报告：结果比当时最好的方法快最多 16%，优化时间平均少 48 倍。

13. 延伸阅读 · Further reading

• Z. Jia, J. Thomas, T. Warszawski, M. Gao, M. Zaharia, A. Aiken, Optimizing DNN Computation with Relaxed Graph Substitutions（SysML 2019，MetaFlow）——TASO 沿用的代价驱动回溯搜索。
• Y. Yang et al., Equality Saturation for Tensor Graph Superoptimization（MLSys 2021，TENSAT）——本讲选读。
• M. Willsey et al., egg: Fast and Extensible Equality Saturation（POPL 2021）——e-graph 库。
• H. Massalin, Superoptimizer: A Look at the Smallest Program（ASPLOS 1987）——“超优化”一词的来源。
• E. Schkufza, R. Sharma, A. Aiken, Stochastic Superoptimization（ASPLOS 2013）——用随机搜索（MCMC）做超优化。
• J. T. Schwartz（1980）、R. Zippel（1979）——Schwartz–Zippel 引理，PET 和 Mirage 的概率验证的数学基础。

14. 下一讲预告 · Next lecture

slides 第 11 讲开头的问题：FlashAttention 能不能表示成一次计算图优化？ 能不能把它写成 matmul、exp、add、mul、div 的某种组合？

不能。计算图里的每个算子都从显存读输入、往显存写输出，而 FlashAttention 的关键是分块、在线 softmax、把中间结果留在共享内存里——这需要代数变换 + 定制内核 + 调度变换同时进行。

第 11 讲《ML 超优化：Mirage》介绍一种同时在三个层次上搜索的超优化器：

• 内核图：哪些算子放进同一个内核；
• 线程块图：一个线程块怎样分块、怎样在共享内存里累加；
• 线程图：寄存器上的逐元素运算。

它用”抽象表达式”剪枝巨大的搜索空间，用有限域上的随机测试（本讲 PET 的定理 2 的推广）验证等价。我们会为 RMSNorm + MatMul 搜索融合方案，看机器怎样自己发现”把 RMSNorm 的分母挪到矩阵乘法之后”这个技巧。