【AI编译器入门④】编译中端：图优化Pass全解析——DCE、常量折叠、算子融合与布局变换

【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）：

1. 第一个kernel：读x和b → 计算x+b → 写回中间结果tmp
2. 第二个kernel：读tmp → 计算relu → 写回y

总共：3次读（x, b, tmp），2次写（tmp, y）

融合（一个kernel）：

1. 一个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这种）。

思路

1. 用模式匹配，找到cast(cast(x, fp32), fp16)这种模式
2. 直接替换为x（因为转了又转回来了）
3. 配合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：

1. DCE（死代码消除）：删掉没用的算子，省显存省算力
2. 常量折叠：编译期把能算的提前算完
3. 算子融合：合并相邻小算子，减少访存和launch------性价比最高的优化之一
4. 布局转换：调整数据布局适配硬件------双刃剑，用不好反而慢
5. 其他Pass：SimplifyInference、ESE、CombineParallelDense、Legalize、PartitionGraph
6. 自定义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
• 更多内容持续更新中......