本文系统解析AI编译器中端四大关键图优化Pass：死代码消除（DCE）通过引用计数清理冗余算子；常量折叠在编译期预计算已知常量表达式；算子融合合并相邻小算子以减少kernel launch与显存访存；布局变换（Layout Transform）适配硬件最优数据排布。同时涵盖SimplifyInference、ESE、Legalize等重要Pass及自定义Pass实现方法，强调Pass顺序、收益边界与常见误区。
【AI编译器入门④】编译中端：图优化Pass全解析——DCE、常量折叠、算子融合与布局变换
本文是「AI编译器」系列的第四篇。如果说前端是"入口"、后端是"出口"，那中端就是AI编译器的"大脑"——它在不改变推理结果的前提下，把计算图变得更干净、更高效。本文带你认识最核心的几类图优化Pass：死代码消除、常量折叠、算子融合、布局转换……以及如何手写一个自定义Pass。
一、写在前面：中端是干什么的？
中端的核心原则只有一句话：
在不改变模型推理逻辑与计算精度的前提下，精简图结构，为后端生成更高质量的代码。
什么意思呢？模型从前端解析出来的计算图，通常"不干净"——有很多冗余的东西：
- 没用的死代码（算出来但没人用）
- 无意义的Cast/Reshape（转过去又转回来）
- 可以合并的小算子（Conv+BN+ReLU可以合成一个）
- 编译期就能算完的常量（还在运行时算）
中端的工作就是把这些"脏东西"清理掉，把能合并的合并，能优化的优化——让图变得更高效。
先上一张中端优化全景图：
二、TVM图优化Pass一览
TVM有几十上百个Pass，我们挑最核心、最常用的来讲：
| Pass名称 | 通俗解释 | 一句话原理 |
|---|---|---|
| DeadCodeElimination | 死代码消除 | 删掉没人用的节点 |
| FoldConstant | 常量折叠 | 编译期把能算的提前算完 |
| FuseOps | 算子融合 | 把相邻小算子合并成一个 |
| SimplifyInference | 推理化简 | 去掉BN/Dropout等训练冗余 |
| AlterOpLayout | 布局转换 | 调整数据布局适配硬件 |
| EliminateCommonSubexpr | 公共子表达式消除 | 重复计算的只算一次 |
| CombineParallelDense | 并行全连接合并 | 多个Dense合并成一个大矩阵乘 |
| FoldScaleAxis | 轴向缩放折叠 | 把scale系数吸进权重里 |
| Legalize | 算子合法化 | 把算子换成后端支持的形式 |
| PartitionGraph | 计算图划分 | 把子图切分到不同硬件 |
| AnnotateUsedMemory | 内存占用标注 | 估算最小显存需求 |
下面我们挑最重要的几个详细讲。
三、死代码消除（DCE）：删掉没用的
3.1 什么是死代码？
死代码 = 做了没用、不影响输出、浪费算力的代码。
大模型里常见的冗余算子：
- 成对Cast：
fp16 → fp32 → fp16—— 转来转去又回来了 - 无用Reshape/Squeeze：shape根本没变
- 无意义Transpose：permute之后又转回去
- 恒等变换Identity：
x = identity(x)—— 等于没做 - 未使用的输出节点：算出来但没人用
这些算子不会改变最终结果，但会：
- 增加kernel launch次数
- 增加显存读写（中间结果要写回显存再读回来）
- 浪费算力
3.2 DCE怎么工作的？
原理很简单——引用计数法：
1. 遍历整个图，统计每个节点被引用了几次
2. 找出引用计数为0的节点（没人用的）
3. 删掉这些节点
4. 删除后，这些节点的输入可能也变成死节点了 → 递归清理
5. 重复直到没有死代码
就像垃圾回收——没人引用的对象，就回收掉。
3.3 实际案例：词表permute冗余
某模型的词表权重有2GB，经过两个permute之后，中间存储膨胀到6GB。但实际上这两个permute完全可以省掉——因为最终用的时候又转回来了。
原始链路：
weights(2GB) → permute → permute_out(2GB) → permute → final(2GB)
↓
unused_node(2GB) ← 死分支！
优化后：
weights(2GB) → 直接参与计算
省了多少？直接省了4GB的中间显存！还省了两次permute的计算时间。
💡 DCE看起来简单，但收益可能很大——尤其是大模型里，冗余算子积少成多。
四、常量折叠（Constant Folding）：编译期能算的提前算
4.1 什么是常量折叠？
对权重、固定超参等编译期已知的数据，提前完成运算，把可静态求解的表达式直接变成常量。
举个例子：
x = constant([1, 2, 3])
y = x + 2
z = y * 3
编译前：有三个算子（const、add、multiply）
编译后：只有一个常量 constant([9, 12, 15])
因为x是常量，所以y和z都可以在编译期算出来——运行时直接用结果，不用再算了。
4.2 执行流程
1. 遍历图，识别所有输入都是常量的子图
2. 用TVM解释器执行这个子图，得到常量结果
3. 把原计算节点替换成常量节点
4. 触发DCE删除已被替换的节点
4.3 收益有多大？
- 减少运行时的计算量（虽然常量计算通常不多）
- 减少算子数量（图更干净了，后续Pass更容易优化）
- 为其他Pass创造机会（比如FuseOps更容易找到融合模式）
💡 常见的onnx-simplifier工具，做的主要就是常量折叠+形状推导+死代码消除——在ONNX层先做一轮化简。
五、算子融合（FuseOps）：性价比最高的优化
5.1 什么是算子融合？
把多个相邻的简单算子合并为一个复合算子，减少中间结果的显存读写和kernel launch次数。
5.2 为什么融合能提速？
我们来算一笔账。假设有两个element-wise算子：y = relu(x + b)
不融合（两个kernel）：
- 第一个kernel：读x和b → 计算x+b → 写回中间结果tmp
- 第二个kernel：读tmp → 计算relu → 写回y
总共：3次读（x, b, tmp），2次写（tmp, y）
融合（一个kernel）：
- 一个kernel：读x和b → 计算x+b → 计算relu → 写回y
总共：2次读（x, b），1次写（y）
省了一次读和一次写——对于memory-bound的element-wise算子，这就是30%-50%的性能提升！
还省了一次kernel launch——虽然单次launch开销不大，但算子多了，积少成多。
5.3 TVM的融合策略
TVM的算子融合基于算子类型：
| 算子类型 | 说明 | 能否融合 |
|---|---|---|
| ElemWise | 逐元素算子（relu、add、mul） | 可以互相融合 |
| Broadcast | 广播算子 | 可以融合到elemwise |
| Injective | 一对一映射（reshape、transpose） | 可以注入到其他算子 |
| CommReduce | 归约算子（sum、max） | 可以和输出的elemwise融合 |
| OutEwiseFusable | 可以和输出的elemwise融合 | 如Dense/Conv的输出加个bias+relu |
5.4 典型融合模式
Conv2D → BiasAdd → ReLU → fused_conv2d_bias_relu
MatMul → Add → Softmax → fused_matmul_add_softmax
Dense → BiasAdd → ReLU → fused_dense_bias_relu
Add → Gelu → Add → fused_add_gelu_add
5.5 融合一定好吗？
不一定。融合也有代价：
- 融合后算子变复杂，可能无法调用高度优化的库kernel（比如cuBLAS）
- 某些硬件对融合后shape不友好，反而更慢
- 融合粒度过大，可能导致寄存器压力过大
所以有DefuseOps（拆分解融合）这个Pass——某些场景下拆开反而更快。
🎯 经验：融合是推理优化中性价比最高的手段之一——尤其是小算子多的模型。但不是融合得越多越好，要以实际性能为准。
六、布局转换（Layout Transform）：让数据排列更适合硬件
6.1 什么是布局转换？
修改卷积、矩阵乘权重的数据存储布局（如NCHW ↔ NHWC），适配硬件的最优计算格式。
比如：
- CPU上NCHW可能更快
- GPU上Tensor Core偏好NHWC
- 某些NPU偏好NC1HWC0这样的5D格式
6.2 两种布局转换Pass
| Pass | 作用 |
|---|---|
| ConvertLayout | 强制把整个模型转换为指定布局 |
| AlterOpLayout | 根据后端偏好，选择性地调整算子布局 |
# 把conv2d转成NHWC布局
mod = relay.transform.ConvertLayout(
{"nn.conv2d": ["NHWC", "OHWI"]}
)(mod)
6.3 收益与风险
收益：
- 硬件计算单元对特定布局更友好（如Tensor Core偏好NHWC）
- 减少运行时transpose开销，提前在编译期完成权重变换
风险：
- 某些算子只支持特定布局，转换后可能引入额外transpose
- 需要与后端对齐，否则可能产生错误结果
💡 布局转换是把双刃剑——转好了性能翻倍，转不好反而引入大量额外transpose。一定要和后端配合好。
七、其他重要Pass简介
7.1 SimplifyInference（推理化简）
把BN、Dropout、IN、GN等训练时的算子，重写成推理形式。
比如BatchNorm在推理时，可以折叠成scale+bias直接加到Conv权重里——省了一个算子。
7.2 EliminateCommonSubexpr（公共子表达式消除）
如果同一个表达式被计算了多次（比如同一个tensor被多个算子使用，而且路径完全相同），就合并成只算一次，多处引用。
7.3 CombineParallelDense（并行全连接合并）
如果多个Dense算子的输入相同、超参一致、只是输出通道不同，可以把它们的权重拼起来，合并成一个大的矩阵乘——减少kernel launch次数，提高计算密度。
7.4 Legalize（算子合法化）
面向后端做等价替换——比如把某个算子换成后端支持的形式。NPU后端通常需要大量的Legalize操作。
7.5 PartitionGraph（计算图划分）
按标注把计算图切分子图，分别卸载到不同硬件/不同后端执行。比如TensorRT子图走TensorRT，剩下的走TVM Codegen。
八、自定义Pass实战：简化冗余Cast
说了这么多，我们来手写一个简单的Pass——消除成对的Cast（fp32→fp16→fp32这种）。
思路
- 用模式匹配，找到
cast(cast(x, fp32), fp16)这种模式 - 直接替换为x（因为转了又转回来了）
- 配合DCE，清理剩下的死代码
代码框架
from tvm.relax.dpl import pattern, rewrite
# 定义模式：cast(cast(x, fp32), fp16)
inner_cast = pattern.is_op("relax.cast")(pattern.wildcard())
outer_cast = pattern.is_op("relax.cast")(inner_cast)
class SimplifyRedundantCast(rewrite.PatternRewriter):
def rewrite(self, expr):
if match := self.match(outer_cast, expr):
# 匹配成功，直接返回最内层的输入
return match[inner_cast].args[0]
return expr
mod = SimplifyRedundantCast().rewrite_module(mod)
💡 真实的实现要处理更多细节（比如dtype检查、嵌套情况等），这里只是展示核心思想。
九、常见误区
误区一：“Pass越多越好”
Pass有副作用和顺序问题。先FuseOps再AlterOpLayout，和先AlterOpLayout再FuseOps，结果可能完全不同。Pass pipeline需要精心设计，不是堆得越多越好。
误区二：“DCE只删死节点”
DCE会触发连锁反应——删了一个节点，它的输入可能变成新的死节点，需要迭代执行。DCE通常要跑好几轮才能收敛。
误区三：“融合一定能提升性能”
融合后算子变复杂，后端可能无法生成高效kernel。某些场景下需要DefuseOps拆分，或者限制融合粒度。
十、总结一下
今天这篇文章，我们讲了编译中端的核心Pass：
- DCE（死代码消除）：删掉没用的算子，省显存省算力
- 常量折叠：编译期把能算的提前算完
- 算子融合：合并相邻小算子，减少访存和launch——性价比最高的优化之一
- 布局转换：调整数据布局适配硬件——双刃剑，用不好反而慢
- 其他Pass：SimplifyInference、ESE、CombineParallelDense、Legalize、PartitionGraph
- 自定义Pass：基于模式匹配，可以快速实现图改写
下一篇预告：我们进入编译后端——CUDA Codegen、CUTLASS、TensorRT，三种后端有什么区别？各有什么优缺点？怎么选？
📚 本系列导航：
- 第1篇：AI编译器基础：三段式Pipeline与全链路协同
- 第2篇：TVM架构拆解：Relay、Relax、TIR与Runtime
- 第3篇：编译前端：ONNX、TorchScript与动态shape处理
- 第4篇：编译中端：图优化Pass全解析（本文）
- 第5篇：编译后端：CUDA Codegen、CUTLASS与TensorRT
- 更多内容持续更新中……
扫一扫
814
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
