老张是十里八乡最出名的掌勺师傅。明天他要在县城最大的酒店办一场三十桌的婚宴,菜谱摊开有一米长:从凉拌三丝到葱爆牛肉,两百多个步骤,每道工序都标着食材从哪来、往哪去。婚宴前夜,老张的徒弟小算盘不睡觉,捧着菜谱一页一页地“过堂”:他先把用量永远不变的复合调味料提前拌好——这一步叫常量折叠;再把一道从头到尾没人点的“冬瓜盅”整个划掉,连带着为它提前泡发的干贝也省了——这一步叫死代码消除;最后他发现“爆葱、滑肉、合炒”三件事其实可以一锅到底,省掉中间两次刷锅——这一步叫算子融合。 这个“小算盘”,就是 AI 编译器;那张菜谱,就是计算图。深度学习模型真正跑起来之前,编译器会对这张“菜谱”做一轮又一轮的化简与重排,让最终端上桌的每一道菜(每一次推理)又快又省。这篇文章带你走进这张菜谱的后厨,把计算图优化的全局流水线看个明白。
对 AI 从业者而言,掌握计算图优化是读懂推理引擎源码、定位性能瓶颈的底层能力。它能让你在部署模型时一眼看出哪里该融合、哪里该剪枝,而不是只会调参碰运气。
目录
可视化演示
本节配有交互动效(进入视口后自动加载,约 3KB)。
计算图基础:一张让后厨不乱的“主菜谱”
老张的菜谱为什么能同时管住三十桌宴席?因为它不写流水账,而是把每道菜拆成“一个动作 + 一份传递”:切好的葱丝放进“待用”的小碗,滑熟的牛肉片搁进“备用”的盘子,下一道工序只认碗和盘,不需要知道上一道是谁炒的。这张菜谱,就是后厨的“计算图”——两百个动作之间,靠一份份食材的流动彼此串联。
深度学习模型在编译器眼里也是这么一张图。把模型展开,每一个算子(卷积、矩阵乘、激活)是图上的一个节点,算子和算子之间流动的张量(可以理解成装在碗里的食材)是边。PyTorch 导出的 ONNX 图、TensorFlow 的 GraphDef,都是这张图的不同写法。
这张图通常是有向无环图(DAG):数据只能从上游流向下游,不存在循环。它至少带来两个好处:一是执行顺序可以由拓扑排序确定,调度器知道谁先谁后;二是优化器有了“定点下刀”的对象——所有优化都发生在图上,改节点、改边、改结构。
有了这张主菜谱,后厨才能井然有序;没有计算图,AI 编译器寸步难行。所以学习优化之前,先学会“读图”:看节点、看边、看数据怎么走。
- G 表示整张计算图它把模型的所有计算步骤装进一张有向无环图,编译器的所有优化都在它身上进行。
- V 是顶点集合每个顶点对应一个算子,例如卷积、加、乘,顶点是输入张量到输出张量的“工序”。
- E 是边集合每条边代表一份张量从上游算子流向下游算子,即数据依赖关系。
优化流水线:按“过堂”顺序走完的工序
小算盘“过堂”从来不是乱翻菜谱,而是有固定顺序:先核对调料配方,再划掉没人点的菜,最后合并重复工序。要是先合并工序再删菜,删掉的那道菜可能让前面的合并全白干。顺序,就是后厨的“流水线”。
编译器里每次对图做的一类修改,叫做一个 Pass(可以理解成菜谱上的一轮批注)。多个 Pass 按顺序执行,就组成优化流水线。为什么顺序重要?因为 Pass 会互相制造机会:常量折叠会产生“只剩单一输入”的节点,死代码消除正好把它们清理掉;融合决策又依赖前两步之后更干净的图结构。
工程实现上,Pass 管理器(Pass Manager)负责任务编排:按照预设顺序逐个执行,必要时重复执行直到图不再变化(达到不动点),再进入下一阶段。从图级到算子级,整个流水线像一条传送带,每个工位只处理自己负责的那类问题。
后厨的工序一旦乱了,宴席就砸了;优化流水线的顺序一旦错了,轻则白忙一场,重则让模型变慢。理解流水线,就是理解编译器“先干什么、后干什么”的全局节奏。
- G₀ 是原始图刚从模型导出、还未经过任何优化的计算图。
- Pᵢ 是第 i 个优化 Pass如常量折叠、死代码消除,每个 Pass 只负责一类改动。
- G_opt 是复合结果Pass 按固定顺序依次作用,后一个 Pass 的输入是前一个的输出。
常量折叠:提前拌好的“复合调料包”
菜谱上写着“盐 3 克、糖 2 克、味精 1 克”,而且这三样从开席到散席用量永远不变。小算盘直接在菜谱上改成一行字:“高级复合调料 6 克”。以后无论哪个帮厨来做,都不必临场各称各的——这就是折叠。
常量折叠(Constant Folding)说的是:当一个算子的全部输入在编译期就是常量时,编译器直接把结果算出来,替换为一个常量节点。比如图上有个 3.14 × 2 + 0 的子式,运行时要多绕两步,折叠后直接变成 6.28,运行时根本不用算。
折叠的收益不止省一点算术。更典型的是 shape 计算:张量的维度常常由常量表达式决定,折叠让编译器在编译期就确定形状,动态形状变成静态形状,内存规划和 Kernel 生成都能提前完成。这就像知道每桌坐十人,才好提前定三十桌的碗筷。
备菜时拌好的酱,下锅时一挤就完事。折叠省掉的不仅是称量时间,还有台面上的一堆瓶瓶罐罐——对应到图上,就是一批可以被删掉的中间节点。
- Op 是某个算子乘法、加法等,它的全部输入都是编译期已知的常量时,编译器就可以替运行时把结果算出来。
- cᵢ 是常量输入例如字面量 3.14、2.0,它们不再依赖任何前序节点。
- c_\* 是折叠产物算子被替换成一个新的常量节点,例如 3.14 × 2.0 → 6.28。
死代码消除:划掉“没人吃的冬瓜盅”
小算盘翻到菜谱第 47 页,发现一道“冬瓜盅”:从开席到散席,没有任何人点过它。他大笔一划,把这道菜连同为它提前泡发的干贝、占用的蒸笼一起撤掉。没人吃的东西,做得再精致也是白费火。
死代码消除(Dead Code Elimination,DCE)做的就是这件事:从图的输出根节点出发,沿着数据依赖反向标记所有“有人需要”的节点;凡是到不了任何输出的子图,都是死代码,直接剪掉。模型裁剪、条件分支消除留下的“孤儿”节点,都是 DCE 的常规猎物。
DCE 有一个和其他 Pass 不同的特点:它不做任何计算,也不改任何数值,纯粹是“丢垃圾”。但这一步极其重要——前一个 Pass 折叠、剪枝后留下的无用子图,如果不及时清理,会继续占用内存、拖慢后续 Pass 的扫描速度。
省下的食材是钱,省下的灶火是算力。DCE 的存在提醒我们:优化不仅要会“做”,还要会“扔”。扔掉注定没人吃的半成品,是后厨效率的底线。
- path 代表可达路径从节点 v 出发,能否沿着数据依赖边走到某个输出节点。
- output 是图的输出根模型的最终输出张量,是检查“有没有人需要我”的终点。
- V_keep 是保留集合只有能到达输出的节点才被保留,其余节点全部删除。
算子融合:用一锅炒代替“炒三次刷三次锅”
“葱爆牛肉”的菜谱分成三步:滑肉、爆葱、合炒,每步之间都要刷一次锅。小算盘跟老张商量:干脆一锅到底,滑完肉不倒掉底油,直接下葱爆香,再回肉合炒。省下的,是两次刷锅的时间。
算子融合(Operator Fusion)把多个相邻算子合并成一个可执行内核 Kernel(可以理解成一道一气呵成的菜)。分开执行时,每个中间结果都要写回内存、再被下一个算子读出来;融合后,这些中间值留在寄存器或共享内存里,省掉了一整轮“写出又读回”。
以 Conv+BN+ReLU 为例:三个算子分开跑,中间张量要在内存里“三进三出”;融合成一个 Kernel 后“一进一出”。在 Transformer 里,QKV 三个矩阵乘合并执行也是同一招。访存是当代芯片最贵的资源之一,融合因此成为收益最立竿见影的图优化。
刷一次锅的成本看似不高,可三十桌宴席、每桌十二道菜算下来,省下的就是上百次刷锅。对应到模型上,就是成千上万次省掉的内存往返——融合,是“少折腾中间人”的艺术。
- N 是张量元素个数比如一张 100×100 的特征图有 10,000 个元素。
- s 是单元素字节数float32 时 s=4,元素越多、单元素越宽,内存访问成本越高。
- 3 次访存 vs 1 次访存分开执行要“写中间结果→读中间结果→写最终结果”共 3 次访问,融合后只需读一次输入、写一次输出。
优化层级:总管管菜单,大厨管火候
小算盘管“今天做什么菜、菜和菜怎么搭配”;老张管“每道菜用几成火候、什么时候下锅”;帮厨管“刀怎么落、火苗怎么调”。三个层级各管各的事,又彼此咬合:总管的搭配,决定了大厨的发挥空间。
编译器优化也是分层的。图级优化(Graph-level)改变图的结构:折叠、消死、决定哪些算子要融合;算子级优化(Operator-level)针对单个 Kernel 的实现:分块(tiling)、向量化、循环重排;更往下还有指令级优化,直接贴着硬件走。
层级之间是协作关系:图级优化先把结构理顺,等于给算子级优化创造了机会;算子级优化把机会变成实实在在的加速。比如图级把 Conv+BN+ReLU 融合成一个节点,算子级才能针对这个融合 Kernel 做寄存器级的数据排布。
端什么菜是图级的事,怎么炒是算子级的事。哪怕切菜刀工再好,菜谱上多了一道没人吃的菜,也是白忙。理解优化层级,就是理解“谁在什么位置用力”。
- T_total 是单次推理总耗时端到端性能的度量,通常以时延或吞吐为准。
- T_memory 是访存耗时数据在内存与计算单元之间搬运的时间,算子融合等优化主要压缩这一项。
- T_compute 是计算耗时纯算术操作的时间,向量化、指令级优化主要压缩这一项。
代码生成:从菜谱到当天的“动线表”
菜谱改完了,还不能直接开席。小算盘把它翻译成当天的动线表:几点泡发干贝、几号灶炒哪道菜、传菜员几点到位、哪桌先上凉菜。这份精确到分钟的动线表,才是宴席真正执行的东西。
图优化的终点是代码生成。编译器把优化后的图翻译成能在目标设备上运行的内核代码(Kernel),并配置运行时的调度策略:算子在哪个设备执行、按照什么顺序启动。部署形态上,可以预编译(AOT,编好再上设备),也可以即时编译(JIT,到了设备上再编)。
图中的每个算子最终都对应一段可执行代码,融合后的 Kernel 是一整段;运行时拿到输入张量,按调度表逐个调用 Kernel,完成一次推理。图优化做得好不好,最终都会反映在“动线表”的执行时间上。
从菜谱到宴席,中间隔着的正是这份动线表。优化是“改菜谱”,代码生成是“排当天动线”,两者合起来,才是一次完整、可交付的推理部署。
- codegen 是代码生成函数把优化后的图翻译成目标设备上的内核代码,例如 CUDA kernel 或 CPU 汇编。
- Kernel 是可执行程序一个融合后的算子内核,输入张量 x 直接输出结果 y。
- 一次编译、多次执行Kernel 被预编译好,推理时只做函数调用,因此图优化的收益会被反复兑现。
AI 应用场景
场景 1:端侧部署前的“一键瘦身”
阶段:模型编译与推理优化(图级优化)
解决的问题:直接把训练好的模型放进手机或摄像头,几十个算子逐个启动,每次启动都有固定开销,中间张量在内存里写进写出,功耗高、延迟大。优化器在编译期依次执行常量折叠、死代码消除、算子融合,生成一个“轻装版”模型,运行时不感知优化过程。这就像小算盘在开席前把菜谱从头到尾删减合并,后厨只管照单执行。
具体案例:某安防摄像头的检测模型,导出为 ONNX 后经 TVM 编译优化:Conv+BN+ReLU 融合为一,BN 的均值方差折进卷积权重(本质是常量折叠),未使用的检测分支被 DCE 剪掉。单帧时延从 45ms 降到 28ms,功耗下降约 30%。
场景 2:LLM 推理中的 QKV 融合
阶段:算子级优化(Kernel 融合)
解决的问题:Transformer 每一层要算 Q、K、V 三个矩阵乘,分开执行意味着三次 kernel 启动、三份中间结果;把它们合并成一次大矩阵乘,就像老张把“切葱、切姜、切蒜”合成“一把抓”,启动开销和访存都省了。这是算子融合在生成式 AI 里的典型形态。
具体案例:vLLM、TensorRT-LLM 等推理框架在构建 engine 时,自动把多头注意力中的 QKV 投影融合为单个 GEMM,再与后面的 attention score 计算做二级融合,显著提升首 token 时延与吞吐。
场景 3:BN 参数折进卷积权重
阶段:训练后导出优化(常量折叠)
解决的问题:训练时 BN 层依赖 batch 统计量,一旦进入推理,均值 μ、方差 σ²、缩放 γ、偏移 β 全变成常量。如果推理时还逐层算归一化,纯属浪费。导出工具把 BN 的数学表达式直接“折叠”进卷积权重 $W' = W \cdot \gamma / \sqrt{\sigma^2 + \varepsilon}$,模型里再也没有 BN 层。
具体案例:PyTorch 模型转 ONNX 时,部署工具链自动完成 BN folding,导出后的图少了一层算子,推理加速明显,且数值完全等价。
Python 演示
class OpNode: def __init__(self, name, op, value=None, inputs=None): self.name = name self.op = op self.value = value self.inputs = inputs if inputs else [] self.outputs = [] def add_output(self, node): self.outputs.append(node) def build_graph(): a = OpNode('a', 'const', value=3.14) b = OpNode('b', 'const', value=2.0) z = OpNode('z', 'const', value=0.0) c = OpNode('c', 'mul', inputs=[a, b]) d = OpNode('d', 'add', inputs=[c, z]) out = OpNode('out', 'relu', inputs=[d]) dead = OpNode('dead', 'tanh', inputs=[d]) for n in [c, d, out, dead]: for inp in n.inputs: inp.add_output(n) return [a, b, z, c, d, out, dead] def constant_fold(nodes): val = {} for n in nodes: if n.op == 'const': val[n.name] = n.value changed = True while changed: changed = False for n in nodes: if n.name in val or not n.inputs: continue if all(i.name in val for i in n.inputs): if n.op == 'mul': val[n.name] = val[n.inputs[0].name] * val[n.inputs[1].name] elif n.op == 'add': val[n.name] = val[n.inputs[0].name] + val[n.inputs[1].name] elif n.op == 'relu': val[n.name] = max(0.0, val[n.inputs[0].name]) changed = True return val def dead_code_eliminate(nodes, output_names): keep = set() stack = list(output_names) by_name = {n.name: n for n in nodes} while stack: name = stack.pop() if name in keep: continue keep.add(name) for inp in by_name[name].inputs: stack.append(inp.name) return [n for n in nodes if n.name in keep] def main(): nodes = build_graph() print('original nodes:', [n.name for n in nodes]) vals = constant_fold(nodes) print('folded values :', {k: round(v, 4) for k, v in sorted(vals.items())}) kept = dead_code_eliminate(nodes, ['out']) print('kept nodes :', [n.name for n in kept]) if __name__ == '__main__': main()
stack = list(output_names)死代码消除从输出根节点出发,把“有人需要”的节点逐个入栈待查。keep.add(name)标记当前节点为“必须保留”,类似小算盘在菜谱上打勾。for inp in by_name[name].inputs:顺着数据依赖边把每个前驱节点加入待遍历栈。if all(i.name in val for i in n.inputs):只有全部输入都是常量时,这个算子才可以被折叠。elif n.op == 'relu':每种算子对应一条折叠规则,规则越多能折叠的子式越多。while changed:迭代折叠直到不再产生新的常量结果,对应流水线中的“不动点”。
常见误区
误区“我用 PyTorch 训练完,模型直接 torch.save 拿去部署就行,优化是框架自动做的,我什么都不用管。”
正解训练框架的前向确实会构建并简化计算图,但它以“可自动求导”为首要目标,产出的图往往带有反向传播的残留与冗余。真正做“婚宴统筹”的是部署阶段的编译器:ONNX Runtime、TensorRT、TVM 拿到模型后都会重新跑一遍完整的优化流水线,不同引擎的性能差异也正来自流水线的强弱。正确姿势是:训练完先导出再交给推理引擎编译,而不是把训练产物直接端上设备。
为何常见:现代框架封装太好,“一键导出”让用户看不到编译这一步,容易误以为图是“天生最优”的。
误区“死代码消除不就是编译器把我不用的 Python 变量删掉吗?跟 C 语言里的未使用变量是一回事。”
正解计算图里的“死代码”不是语言层变量,而是子图:从任何输出根节点出发都不可达的算子及其边。删掉它靠的是可达性分析,而不是有没有变量名。而且图里很多死代码正是上一个 Pass(如条件分支消除)制造的“孤儿”,只有流水线顺序对了才会被顺手清掉。如果把 DCE 当成“删变量”,就会错过它真正的用武之地。
为何常见:“死代码”这个词来自传统编译器,大家容易带着编程语言的经验去理解,而忘了图优化面对的是数据流而非源码。
误区“把 3.14×2 提前算成 6.28,编译器偷偷改了数值,模型精度肯定受影响。”
正解折叠是精确的代数化简,3.14×2 与 6.28 在数学上完全等价,编译器只是把运行时要做的事搬到编译期,数值本身没有被改动。唯一需要小心的是浮点运算不满足结合律,(a+b)+c 与 a+(b+c) 可能差一个末位;只有当编译器开启 fast-math 之类激进重排时才可能引入微小舍入差异,工程上通常可接受。把“量化”的精度损失经验迁移到折叠上,是张冠李戴。
为何常见:大家把“量化”损失精度的经验误套到“常量折叠”上,其实折叠不动数值,只是换了个时间算。
误区“融合不就是把 Conv 和 ReLU 的代码写进同一个函数吗?我写代码时也做过这种合并。”
正解融合的核心是让中间结果不落回全局内存:新 Kernel 在寄存器或共享内存里完成上一步的产出、下一步的消费。如果只是源码级别拼函数、访存结构不变,中间张量照样写进写出,性能不仅不会提升,还可能因寄存器压力变大而变慢。所以融合是“换一口锅一锅炒”,不是“把两个菜谱抄在同一页上”。
为何常见:“合并代码”是日常编程的直觉,但 GPU/CPU 的性能瓶颈在于内存搬运,这一点不做底层优化的人很难感知。
误区“优化肯定是多多益善,把能开的融合、折叠全开,模型一定跑得最快。”
正解每个 Pass 都有成本与副作用:激进融合会放大寄存器压力,让单个 Kernel 变慢;过度折叠让编译时间暴涨;错误的 Pass 顺序可能让两个优化互相抵消。工业级编译器遵循“收益大于成本”原则,甚至会在自动调优时对多个策略做实测筛选。后厨讲究“够用就好”,把所有灶台同时开到最大火,反而会把厨房烧了。
为何常见:直观上“做更多事”等于“更好”,但优化是有代价的交易,决策要看端到端收益,而不是 Pass 数量。
面试题
Q1. 下列关于常量折叠的说法,正确的是?
- A.在运行时把已知常量代入计算
- B.在编译期直接算出常量表达式结果← 正确答案
- C.把常量替换成变量再参与计算
- D.只对整数常量执行合并计算
答案:B。常量折叠发生在编译期而非运行时,编译器把输入全是常量的算子提前算成常量节点,运行时直接取值。A 说反了时间点;C 方向完全反;D 忽略了浮点常量同样能折叠。就像老张提前拌好复合调料,下锅时直接挤,不必临时称重。
Q2. 死代码消除判断“死”的标准是?
- A.节点名称带有 dead 这样的关键字
- B.节点计算出来的结果数值为零
- C.从任何最终输出根节点均不可达← 正确答案
- D.节点在源代码中出现的顺序靠后
答案:C。死代码消除以“数据可达性”为判据:从输出根节点反向遍历,碰不到的就是死代码。A 靠名字判断不科学;B 结果为零不代表没被使用;D 出现顺序与是否被消费无关。就像婚宴菜谱里没人点的“冬瓜盅”,判断标准是谁需要它,而不是它写在菜谱第几页。
Q3. 算子融合最主要的性能收益来自?
- A.减少算子编译所需的全部时间
- B.减少中间张量的内存读写开销← 正确答案
- C.增加设备上的并行线程数量
- D.提升模型推理的数值稳定精度
答案:B。融合最大的收益是省掉中间张量的内存读写,把“三进三出”变为“一进一出”。A 编译时间通常不减反增;C 不改变线程数也能提速;D 数值精度与融合无关。好比“一锅炒”省的是中间刷锅的时间,而不是多请几个帮厨。
Q4. 以下哪个属于图级优化?
- A.循环展开与循环重组优化
- B.向量化与数据布局转换优化
- C.算子融合与死代码消除优化← 正确答案
- D.寄存器分配与指令选择优化
答案:C。算子融合和死代码消除都改变图的结构(合并节点、删除子图),属于图级优化。循环展开、向量化、寄存器分配分别作用于循环级、指令级与硬件资源,不改图结构。这类似于“今天上什么菜”由总管定,而“火候怎么调”是大厨的事。
Q5. 关于优化流水线的说法,正确的有?
- A.Pass 的执行顺序会影响最终性能← 正确答案
- B.各 Pass 之间完全独立互不干扰
- C.常量折叠后再做死代码消除← 正确答案
- D.优化器通常只运行一个 Pass
答案:AC。Pass 顺序影响结果,比如先折叠再消死,折叠产生的“孤儿”节点才能被顺手删掉;Pass 之间并不独立,后者常依赖前者的产物。C 描述的正是典型顺序;D 说反了——一个完整流水线往往有几十个 Pass。后厨“先备料、再删菜、最后合并工序”,顺序乱了就白干。
Q6. 常量折叠对计算图 shape 信息的作用是?
- A.让部分动态形状变静态可推断← 正确答案
- B.完全消除模型所有动态维度
- C.只影响数值而不影响形状信息
- D.使模型不再需要任何网络输入
答案:A。shape 常常由常量运算决定,折叠让编译器在编译期算出形状,动态维度就变成静态可推断的,进而支撑内存规划与 Kernel 生成。B 过度绝对,不是所有动态都来自常量;C 恰好说反;D 与输入要求无关。提前知道每桌坐十人,才能定三十桌的碗筷。
Q7. 下列哪些是算子融合的常见工程形态?
- A.Conv+BN+ReLU 融合成单个 Kernel← 正确答案
- B.Attention 的 QKV 矩乘合并执行← 正确答案
- C.融合后中间结果仍写回全局内存
- D.减小 Kernel 启动次数对性能无益
答案:AB。Conv+BN+ReLU 融合和 QKV 合并都是工业界常见的融合形态,都能减少 Kernel 数量与中间访存。C、D 把收益说反了:融合 Kernel 恰恰要避免中间结果写回全局内存,启动次数减少也是重要收益。两两合并如同把“滑肉—爆葱—合炒”一锅端。
Q8. 编译器处理"3.14×2+0"这类子图时,最合理的优化是?
- A.先融合再折叠成单个算子
- B.直接折叠为常量 6.28← 正确答案
- C.保留原式以确保数值精度
- D.转化为死代码再予以删除
答案:B。3.14×2+0 的输入全是常量,编译器直接折叠成 6.28,这是常量折叠的典型场景。A 的顺序不必要;C 与正确优化矛盾;D 没到“死”的程度,它是被消费的活节点。复合调料包提前拌好,下锅时直接挤。
Q9. 计算图 G=(V,E) 中 E 表示?
- A.边的集合代表张量数据流向← 正确答案
- B.节点的集合代表计算算子
- C.参数的集合代表模型权重
- D.输出的集合代表最终结果
答案:A。计算图记作 G=(V,E),V 是算子节点,E 是张量流动的边。选项 B 说的是 V,C、D 分别是参数和输出,均不是 E 的定义。边就像菜谱上“这碗高汤从哪来到哪去”的箭头,没有边后厨就乱了。
Q10. 代码生成阶段需要做哪些事情?
- A.将优化后的图翻译成 Kernel 代码← 正确答案
- B.为端侧设备选择 AOT 或 JIT 策略← 正确答案
- C.重新加入训练过程的数据增强
- D.配置运行时算子的调度与启动← 正确答案
答案:ABD。代码生成把优化后的图翻译成 Kernel 并配置运行时调度,AOT/JIT 是两种部署形态。C 训练阶段的数据增强与推理无关,绝不该出现在生成阶段。从菜谱到动线表,正是“翻译成可执行工序”这一步。