【AI编译器入门④】编译中端图优化Pass全解析——DCE、常量折叠、算子融合与布局变换本文是「AI编译器」系列的第四篇。如果说前端是入口、后端是出口那中端就是AI编译器的大脑——它在不改变推理结果的前提下把计算图变得更干净、更高效。本文带你认识最核心的几类图优化Pass死代码消除、常量折叠、算子融合、布局转换……以及如何手写一个自定义Pass。一、写在前面中端是干什么的中端的核心原则只有一句话在不改变模型推理逻辑与计算精度的前提下精简图结构为后端生成更高质量的代码。什么意思呢模型从前端解析出来的计算图通常不干净——有很多冗余的东西没用的死代码算出来但没人用无意义的Cast/Reshape转过去又转回来可以合并的小算子ConvBNReLU可以合成一个编译期就能算完的常量还在运行时算中端的工作就是把这些脏东西清理掉把能合并的合并能优化的优化——让图变得更高效。先上一张中端优化全景图二、TVM图优化Pass一览TVM有几十上百个Pass我们挑最核心、最常用的来讲Pass名称通俗解释一句话原理DeadCodeElimination死代码消除删掉没人用的节点FoldConstant常量折叠编译期把能算的提前算完FuseOps算子融合把相邻小算子合并成一个SimplifyInference推理化简去掉BN/Dropout等训练冗余AlterOpLayout布局转换调整数据布局适配硬件EliminateCommonSubexpr公共子表达式消除重复计算的只算一次CombineParallelDense并行全连接合并多个Dense合并成一个大矩阵乘FoldScaleAxis轴向缩放折叠把scale系数吸进权重里Legalize算子合法化把算子换成后端支持的形式PartitionGraph计算图划分把子图切分到不同硬件AnnotateUsedMemory内存占用标注估算最小显存需求下面我们挑最重要的几个详细讲。三、死代码消除DCE删掉没用的3.1 什么是死代码死代码 做了没用、不影响输出、浪费算力的代码。大模型里常见的冗余算子成对Castfp16 → fp32 → fp16—— 转来转去又回来了无用Reshape/Squeezeshape根本没变无意义Transposepermute之后又转回去恒等变换Identityx 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 什么是常量折叠对权重、固定超参等编译期已知的数据提前完成运算把可静态求解的表达式直接变成常量。举个例子xconstant([1,2,3])yx2zy*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 → 计算xb → 写回中间结果tmp第二个kernel读tmp → 计算relu → 写回y总共3次读x, b, tmp2次写tmp, y融合一个kernel一个kernel读x和b → 计算xb → 计算relu → 写回y总共2次读x, b1次写y省了一次读和一次写——对于memory-bound的element-wise算子这就是30%-50%的性能提升还省了一次kernel launch——虽然单次launch开销不大但算子多了积少成多。5.3 TVM的融合策略TVM的算子融合基于算子类型算子类型说明能否融合ElemWise逐元素算子relu、add、mul可以互相融合Broadcast广播算子可以融合到elemwiseInjective一对一映射reshape、transpose可以注入到其他算子CommReduce归约算子sum、max可以和输出的elemwise融合OutEwiseFusable可以和输出的elemwise融合如Dense/Conv的输出加个biasrelu5.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_add5.5 融合一定好吗不一定。融合也有代价融合后算子变复杂可能无法调用高度优化的库kernel比如cuBLAS某些硬件对融合后shape不友好反而更慢融合粒度过大可能导致寄存器压力过大所以有DefuseOps拆分解融合这个Pass——某些场景下拆开反而更快。经验融合是推理优化中性价比最高的手段之一——尤其是小算子多的模型。但不是融合得越多越好要以实际性能为准。六、布局转换Layout Transform让数据排列更适合硬件6.1 什么是布局转换修改卷积、矩阵乘权重的数据存储布局如NCHW ↔ NHWC适配硬件的最优计算格式。比如CPU上NCHW可能更快GPU上Tensor Core偏好NHWC某些NPU偏好NC1HWC0这样的5D格式6.2 两种布局转换PassPass作用ConvertLayout强制把整个模型转换为指定布局AlterOpLayout根据后端偏好选择性地调整算子布局# 把conv2d转成NHWC布局modrelay.transform.ConvertLayout({nn.conv2d:[NHWC,OHWI]})(mod)6.3 收益与风险收益硬件计算单元对特定布局更友好如Tensor Core偏好NHWC减少运行时transpose开销提前在编译期完成权重变换风险某些算子只支持特定布局转换后可能引入额外transpose需要与后端对齐否则可能产生错误结果布局转换是把双刃剑——转好了性能翻倍转不好反而引入大量额外transpose。一定要和后端配合好。七、其他重要Pass简介7.1 SimplifyInference推理化简把BN、Dropout、IN、GN等训练时的算子重写成推理形式。比如BatchNorm在推理时可以折叠成scalebias直接加到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——消除成对的Castfp32→fp16→fp32这种。思路用模式匹配找到cast(cast(x, fp32), fp16)这种模式直接替换为x因为转了又转回来了配合DCE清理剩下的死代码代码框架fromtvm.relax.dplimportpattern,rewrite# 定义模式cast(cast(x, fp32), fp16)inner_castpattern.is_op(relax.cast)(pattern.wildcard())outer_castpattern.is_op(relax.cast)(inner_cast)classSimplifyRedundantCast(rewrite.PatternRewriter):defrewrite(self,expr):ifmatch:self.match(outer_cast,expr):# 匹配成功直接返回最内层的输入returnmatch[inner_cast].args[0]returnexpr modSimplifyRedundantCast().rewrite_module(mod) 真实的实现要处理更多细节比如dtype检查、嵌套情况等这里只是展示核心思想。九、常见误区误区一“Pass越多越好”Pass有副作用和顺序问题。先FuseOps再AlterOpLayout和先AlterOpLayout再FuseOps结果可能完全不同。Pass pipeline需要精心设计不是堆得越多越好。误区二“DCE只删死节点”DCE会触发连锁反应——删了一个节点它的输入可能变成新的死节点需要迭代执行。DCE通常要跑好几轮才能收敛。误区三“融合一定能提升性能”融合后算子变复杂后端可能无法生成高效kernel。某些场景下需要DefuseOps拆分或者限制融合粒度。十、总结一下今天这篇文章我们讲了编译中端的核心PassDCE死代码消除删掉没用的算子省显存省算力常量折叠编译期把能算的提前算完算子融合合并相邻小算子减少访存和launch——性价比最高的优化之一布局转换调整数据布局适配硬件——双刃剑用不好反而慢其他PassSimplifyInference、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更多内容持续更新中……