想象你是一位大厨,经营着一家生意火爆的餐厅。每天开业前,你都会把当天要用的食材全部预处理完毕:洋葱切成丝、蒜头拍成末、高汤提前炖好、酱汁按秘方调配密封冷藏。当客人点单时,你只需要从冰箱里取出这些半成品,快速下锅翻炒即可。你绝不会在客人面前才慢悠悠地洗菜、切菜、熬汤——那会让每桌客人都等上半小时。这个“提前准备”的智慧,正是编译器里常量折叠(Constant Folding)的核心理念。 在大模型训练和推理的世界里,AI 编译器的角色就像这位大厨。当开发者写出 PyTorch 或 TensorFlow 代码时,里面往往隐藏着大量“可以在开张前就准备好的食材”——比如固定不变的形状计算、预先可知的数学运算、不会改变的配置参数。如果让运行时(Runtime)每次执行时都重新“洗菜切菜”,就会浪费海量的计算资源。常量折叠就是那位聪明的大厨:在编译阶段把能算的都先算好,让真正的计算图轻装上阵。它不改变程序的任何行为,只是把“什么时候干活”变得更聪明了。作为 AI 从业者,理解常量折叠不仅是理解编译器的第一步,更是日后调试性能问题、优化推理延迟的必修内功。

岗位关联

学习常量折叠,你将理解 AI 编译器优化的第一块基石:辨别“哪些计算必须运行时做、哪些可以在编译期白嫖”,这是你诊断模型性能瓶颈、与编译器“对话”的基本功。

目录
  1. 可视化演示
  2. 常量折叠初识
  3. 编译器内位置
  4. 常量类型分类
  5. 折叠算法流程
  6. 性能特性分析
  7. 优化策略实践
  8. 实战应用指南
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示常量折叠(Constant Folding):编译期将输入全为常量的算子提前计算并替换为常量节点,如3+5折叠为8,并演示常量传播协同与BN/reshape等典型折叠场景。

本节配有交互动效(进入视口后自动加载,约 3KB)。

常量折叠流程 常量3、5输入ADD加法算子 所有输入为常量→编译期FOLD计算 3+5=8→常量节点 计算结果替换为常量节点,消除运行时开销。

常量折叠初识

大厨在开张前备好高汤和酱汁,客人点餐时只需快速组装。这正是常量折叠的智慧:在程序运行之前(编译期),把所有输入都是常量的表达式预先计算出来,用计算结果替换整个表达式。比如 3 * 4 + 5 * 6 这个表达式,无论程序运行多少次,它的值永远都是 42。与其让计算机每次运行都重新乘一遍加一遍,不如在编译时直接算出 42,把整个表达式替换成一个数字。

在 AI 编译器的语境下,常量折叠(Constant Folding)是图优化中最基础、最安全的优化手段之一。它作用于计算图(Computational Graph),图上的每个节点是一个算子(如 Add、MatMul、Conv2D),边代表数据流向。当某个算子的所有输入都是编译期可知的常量时,这个算子就可以被“折叠”——我们直接计算出它的结果,把它替换成一个常量节点。

常量折叠的优化效果是三重奏。第一,消除了运行时冗余计算,让每次推理少做无谓的算术。第二,使计算图更简洁,为后续优化(如死代码消除)创造机会。第三,改进了数据局部性,因为常量数据可以被编译器更好地安排内存布局。这不是一个炫目的优化,却是所有高级优化的基石——正如备菜本身不起眼,但没有备菜的餐厅一定手忙脚乱。

回到大厨的比喻:折叠后的表达式就像一份“半成品套餐”——主菜已经炖好,客人点单时只需要加热摆盘。编译器做的正是这种提前备菜的工作,让运行时这位“炒锅师傅”只需要处理真正需要现炒的菜(变量 t 相关的计算)。

3×4+5×6+t  ⇒  12+30+t  ⇒  42+t3 \times 4 + 5 \times 6 + t \;\Rightarrow\; 12 + 30 + t \;\Rightarrow\; 42 + t
这个式子展示了常量折叠的核心过程:把编译期可知的 3×4 和 5×6 分别算成 12 和 30,再合并成 42,最终整个表达式只剩 42 + t 一个加法和一个变量。运行时只需要做一次加法。
  1. 识别常量子表达式扫描表达式,找出所有操作数都是常量的子表达式——这里 3×4 和 5×6 满足条件。
    Const(3)∧Const(4)⇒3×4 可折叠\text{Const}(3) \wedge \text{Const}(4) \Rightarrow 3 \times 4 \text{ 可折叠}
  2. 编译期求值在编译阶段直接计算这些子表达式的值,不产生运行时开销。
    3×4=12,5×6=303 \times 4 = 12,\quad 5 \times 6 = 30
  3. 常量合并如果折叠后出现了“常量加常量”的新机会,继续合并,直到无法再折叠。
    12+30=4212 + 30 = 42
  4. 替换原表达式用最终的常量值 42 替换原来的计算图子图,原表达式变为 42 + t。
    3×4+5×6+t⇒42+t3 \times 4 + 5 \times 6 + t \Rightarrow 42 + t

编译器内位置

大厨的备菜台放在哪,决定了出餐流程是否顺畅。如果备菜台放在传菜口,客人点单时还得跑回去拿,那就失去了意义。常量折叠在编译器中的位置,同样决定了它能否与其他优化顺畅配合。

在现代 AI 编译器(如 PyTorch 的 torch.compile 和 Google 的 XLA)中,常量折叠是图优化阶段(Graph Optimization)的一项基础 Pass。整体编译流程大致是:前端(Frontend)把 Python 代码转换成中间表示(IR,Intermediate Representation),比如 PyTorch 的 ATen IR 或 XLA 的 HLO(High-Level Optimizer)IR;然后图优化 Pass 在 IR 上进行各种变换;最后后端(Backend)把优化后的 IR 翻译成机器码或 GPU Kernel。

常量折叠通常作为图优化阶段的前置 Pass 运行,原因很朴素:折叠后图更小,后续的算子融合(Operator Fusion)、内存规划(Memory Planning)等 Pass 处理起来更快、效果更好。编译器领域有一个经典组合——常量折叠 + 常数传播(Constant Propagation):常数传播负责把变量的常量值沿着数据流“传播”到使用处,让更多表达式在编译期变得“全常量”;常量折叠则负责把这些新发现的全常量表达式计算掉。两者相辅相成,一个负责“发现食材”,一个负责“下锅做菜”。

用流水线来理解:前端是采购员,把食材(代码)买回来;图优化阶段是备菜间,常量折叠就是备菜台——菜在这里切好、腌好、炖好;后端是灶台,只负责最后的爆炒(生成高效 Kernel)。如果备菜台上一次能处理的菜变多了,灶台的工作自然更轻松。

在 PyTorch 的 torch.compile 中,常量折叠由 ConstantFolder 这个 Pass 在 Dynamo 生成的 FX Graph 上执行;在 XLA 中,AlgebraicSimplifier 里包含了大量常量折叠规则。它们都位于图优化管道的前端,为后续的融合和调度铺平道路。这就像备菜台永远设在厨房入口,每位厨师路过都能顺手取用半成品。

IRin→Constant FoldingIRfolded→后续 PassIRopt→后端Machine Code\text{IR}_{\text{in}} \xrightarrow{\text{Constant Folding}} \text{IR}_{\text{folded}} \xrightarrow{\text{后续 Pass}} \text{IR}_{\text{opt}} \xrightarrow{\text{后端}} \text{Machine Code}
这个流水线公式展示了常量折叠在编译器中的位置:它是 IR 进入优化管道后的第一道工序,为后续所有优化提供更干净的输入。
  1. 产生 IR前端把 Python/框架代码翻译成编译器可处理的中间表示。
    Python Code⇒IRin\text{Python Code} \Rightarrow \text{IR}_{\text{in}}
  2. 常量折叠 Pass把 IR 中所有可折叠的常量表达式预先计算并替换。
    IRin⇒IRfolded\text{IR}_{\text{in}} \Rightarrow \text{IR}_{\text{folded}}
  3. 后续优化 Pass算子融合、内存规划等优化在更精简的图上运行。
    IRfolded⇒IRopt\text{IR}_{\text{folded}} \Rightarrow \text{IR}_{\text{opt}}
  4. 后端代码生成把优化后的 IR 翻译成目标平台的机器指令。
    IRopt⇒Machine Code\text{IR}_{\text{opt}} \Rightarrow \text{Machine Code}

常量类型分类

大厨备菜时心里有数:有些食材是全季都有的(比如土豆、洋葱),有些是特定菜品才用的(比如松露),还有些实际上根本不是食材而是餐具(比如摆盘用的模具)。聪明的厨师不会把模具放错到食材区。编译器面对常量,同样需要精确分类,才能决定谁可以折叠、谁必须谨慎。

在 AI 计算图中,常量主要有三种类型。第一类是标量常量(Scalar Constants):单个数字,如 3.14、1024 或 -1。它们最常见,与变量混合出现在算术表达式中,是折叠收益最大的对象。第二类是张量常量(Tensor Constants):一个形状已知、元素全为常量的张量,比如全零矩阵 torch.zeros(64, 64) 或预训练好的权重矩阵(如果它们不在训练中更新)。第三类是形状常量(Shape-Only Constants):代码里只用到张量的形状信息,例如 x.size(0) 或 x.shape[-1]。形状常量在 PyTorch 动态图里很常见,是 torch.compile 能否成功“捕获”静态图的关键。

折叠边界是常量折叠中最微妙的判定问题。一个算子可以被折叠,当且仅当它的所有输入都是常量,且算子本身是纯函数——即无副作用(不会修改全局状态)、输出只由输入决定。add(3, 4) 可以折叠,因为它完全由输入决定。但 random_uniform() 不能折叠,即使它没有输入——因为每次运行结果都应不同;print_metrics() 更不能折叠,因为打印是一个副作用。编译器通过“常量性判定”(Constantness Analysis)来判断一个节点的输入是否真的不会在运行时改变。

回到备菜的类比:标量常量是常年备货的主食材,张量常量是订好了的整箱食材,形状常量是贴在箱子外面的尺寸标签。“标签”虽然看起来不起眼,但很多优化都依赖它——比如编译器需要知道张量形状才能决定 Kernel 的线程配置。而折叠边界就像是“保质期检查”:只有确定不会变质的食材才敢提前加工,有随机性(未开封但有变质风险)或副作用(用了会出大事的食材)的算子必须留到运行时。

Foldable(op)  ⟺  ∀i∈Inputs(op):IsConstant(i)∧IsPure(op)\text{Foldable}(op) \iff \forall i \in \text{Inputs}(op): \text{IsConstant}(i) \wedge \text{IsPure}(op)
这个公式说明了折叠的边界条件:一个算子可被折叠,当且仅当它的所有输入都是常量,并且算子本身是无副作用的纯函数。
  1. 逐个检查输入对算子的每一个输入,判断其是否为常量。
    ∀i∈Inputs(op):IsConstant(i)\forall i \in \text{Inputs}(op): \text{IsConstant}(i)
  2. 验证算子纯度检查算子是否有随机性、是否读写全局状态。
    IsPure(op)\text{IsPure}(op)
  3. 联合判定两个条件同时满足时,该算子才进入可折叠集合。
    Foldable(op)=AllInputsConst∧Pure\text{Foldable}(op) = \text{AllInputsConst} \wedge \text{Pure}
  4. 执行折叠对可折叠节点,编译期求值并替换为常量节点。

折叠算法流程

大厨备菜有自己的一套工序:先看采购单(遍历食材),判断哪些菜需要预处理(判定可折叠性),然后动手处理(求值替换),最后放入半成品区(迭代传播)。常量折叠的算法流程与之高度吻合。

完整的常量折叠算法可以抽象为四步。第一步:遍历计算图。编译器从图的根节点开始,以拓扑序(Topological Order)遍历每一个算子节点。拓扑序保证在访问一个节点时,它的所有输入节点都已被访问过,这是常量性判断的基础。第二步:判定常量性。对每个节点,检查其所有输入是否都是常量节点。这一步需要结合前文说的常量类型分类——不仅包括数字字面量,还包括形状常量和经过传播的常量变量。第三步:编译期求值。如果节点可折叠,编译器在当前编译环境中直接执行这个算子(本质上是调用算子的 CPU 实现),得到输出值,创建一个新的常量节点。第四步:替换与迭代。用新常量节点替换原来的计算子图,然后继续遍历后续节点——因为新常量可能触发下游节点的新折叠机会。这就是“备一菜、得一菜”的连锁反应。

在实际工程实现中,算法往往使用工作列表(Worklist)配合惰性求值(Lazy Evaluation)来提升效率。PyTorch 的 ConstantFolder 采用一种“符号形状”(Symbolic Shape)推导方案,先分析每个节点输出张量的形状是否只与常量相关;如果是,则该节点的输出被标记为“可折叠张量”,随后对该节点做实际求值。XLA 的 AlgebraicSimplifier 则维护一组模式匹配规则,比如 add(0, x) → x、mul(1, x) → x、mul(0, x) → 0,这些规则中有大量常量-变量组合的化简,实际上是对常量折叠的“推理增强版”。

这个流程的核心纪律是“安全第一”——只有当条件完全满足时才动手替换。折叠过程中产生的中间量被保存在编译期的内存中,不占用运行时空间。就像大厨备菜时,切好的菜放在保鲜盒里,每个盒子标注日期(常量类型),过了保质期(条件不满足)就绝不下锅。

Worklist={n∣n∈Graph,order(n) is topological}\text{Worklist} = \{ n \mid n \in \text{Graph}, \text{order}(n) \text{ is topological} \}
for n∈Worklist:if IsFoldable(n):Graph←Replace(n,Eval(n))\text{for } n \in \text{Worklist}:\quad \text{if } \text{IsFoldable}(n):\quad \text{Graph} \leftarrow \text{Replace}(n, \text{Eval}(n))
这两行公式展示了折叠算法的主循环:按拓扑序遍历所有节点,遇到可折叠节点就求值并替换,直到处理完整个图。
  1. 初始化工作列表将计算图的所有节点按拓扑序放入工作列表。
    Worklist←TopoSort(Graph)\text{Worklist} \leftarrow \text{TopoSort}(\text{Graph})
  2. 取出节点判断从工作列表头取一个节点 $n$,检查是否所有输入为常量。
    IsFoldable(n)=⋀i∈inputs(n)IsConst(i)\text{IsFoldable}(n) = \bigwedge_{i \in \text{inputs}(n)} \text{IsConst}(i)
  3. 编译期求值调用算子纯函数实现,在编译进程内完成实际运算。
    val=Eval(n)(const1,const2,…)\text{val} = \text{Eval}(n)(\text{const}_1, \text{const}_2, \ldots)
  4. 替换原节点在图中把 $n$ 替换为新常量节点,删除旧子图。
    Graph←Replace(n,NewConst(val))\text{Graph} \leftarrow \text{Replace}(n, \text{NewConst}(\text{val}))

性能特性分析

大厨备菜减少了下锅时间,但也占用了备菜台的空间。常量折叠的性能特征同样是双刃剑:它用编译时间交换运行时效率,减少计算量、精简代码体积,但并不是免费的午餐。

性能收益首先体现在计算量减少。假如一个 Transformer 模型每层都包含类似 1.0 / math.sqrt(64) 的标准化因子,若模型有 100 层、每层被调用 1000 次,那么一次推理就会多出 10 万次除法——折叠后这部计算全部消失。在批处理(Batch Inference)场景下,每次推理都省去这些运算,累积效应非常可观。第二体现在启动时间和内存占用。折叠后创建的常量张量可以被编译器直接嵌入到二进制文件或 Kernel 的参数区,不需要运行时动态分配和填充。对于推理场景,这意味着更快的模型加载和更低的内存峰值。第三体现在代码体积。折叠把大量冗余计算从最终产物中清除,生成的二进制更小,这在边缘设备上尤为珍贵。

但要注意代价:编译时间增加。每折叠一个节点,编译器都要在编译进程中真实地计算一次。对于大模型,如果折叠策略过于激进,编译时间可能从几秒膨胀到几十秒。此外,解释性下降是隐藏成本——折叠后的计算图完全不像原始代码,当模型出问题时,开发者很难把图中的常量节点对应回原始 Python 表达式。

其中 $T_i$ 是第 $i$ 个折叠点的单次执行时间,$N_i$ 是它在运行时的调用次数,$T_{\text{fold}}$ 是编译期折叠花费的总时间。

回到比喻:备菜需要时间,切 100 份土豆丝确实花了你半小时,但如果这半小时让午市高峰期的 80 桌客人每桌少等 3 分钟,这笔时间投资就极其划算。同理,训练场景下模型可能被反复实验多次,编译做一次、运行千百次,常量折叠的杠杆效应在推理场景会被放大到极致——这也是为什么 XLA 之类的专用编译器对常量折叠高度重视的原因。

Saved=∑i=1kTi×Ni−Tfold\text{Saved} = \sum_{i=1}^{k} T_i \times N_i - T_{\text{fold}}
这个公式通俗地描述常量折叠的收益模型:省下的总时间等于每个折叠点在运行时被调用的累计时间总和,减去编译期折叠本身消耗的时间。只有当前者大于后者,折叠才“回本”。
  1. 计算运行时节省对每个折叠点估计它在运行时的执行成本。
    Savingsi=Ti×Ni\text{Savings}_i = T_i \times N_i
  2. 求和所有省时把所有折叠点的节省时间累加起来。
    ∑i=1kSavingsi\sum_{i=1}^{k} \text{Savings}_i
  3. 扣除编译成本减去编译期折叠所有节点消耗的总时间。
    Net=Total Savings−Tfold\text{Net} = \text{Total Savings} - T_{\text{fold}}
  4. 判断收益正负净收益大于零时折叠策略是划算的。
    Net>0⇒Folding is worth it\text{Net} > 0 \Rightarrow \text{Folding is worth it}

优化策略实践

大厨的备菜策略分三档:有些菜必须现炒(涉及变量),有些菜提前半天备好(常量折叠),有些菜甚至提前一个月做成腌菜、腊肉(更激进的编译期预处理)。常量折叠在工程实践中同样讲究分层施策与安全优先。

策略一:分层折叠。不是所有常量折叠都发生在同一阶段。IR 级别首先折叠标量常量和元数据推导;算子级别进一步折叠张量常量运算,例如把 ones((3, 3)) * 2 直接变成一个值全为 2 的常量张量;调度级别则在 Kernel 层面保证编译期已算好的常量以最优方式嵌入。分层折叠的关键是“边界正交”——上层不需要关心下层的实现细节,每一层只处理自己职责范围内的常量。

策略二:条件编译配合。PyTorch 中大量使用 torch.jit.is_compiling() 或 Python 的 if __name__ == "__main__" 这类运行时分支来控制代码路径。常量化(Constant-ifying)这些分支条件,就可以在编译期剪除死分支,让折叠后的图更干净。比如 if config.enable_compile: 这样的判断,如果 config 在编译期不可知,编译器就要保留两个分支;如果开发者用 functools.lru_cache 或静态配置显式声明了常量,编译器就能直接折叠掉假分支。

策略三:安全优先。这是最核心的工程纪律。Python 的浮点数运算在编译期和运行时可能有舍入差异(比如 3 * 0.1 在 x86 上编译期算出的结果与运行时 Kernel 算出的结果可能差一个 ULP)。因此 XLA 在折叠浮点运算时,必须确认编译环境与运行环境的浮点模型(如 IEEE 754)一致,否则宁可保守地不折叠。同理,涉及 NaN/Inf 语义、有符号整数溢出(Signed Overflow)的运算也在部分编译器的“不可折叠清单”里,避免引入未定义行为。

在 PyTorch 工程实践中,torch.compile 提供了 torch._dynamo.config 来调节折叠的激进程度:default 模式只折叠绝对安全的算子,max 模式会尝试折叠更多——但也带来编译时间膨胀的风险。真实团队落地时,通常先以性能剖析定位热点(Profile-Guided),再针对热点区域单独开启更激进折叠,完全遵循“安全优先、数据驱动”的原则。这就像主厨先在试菜环节尝过每一道预制菜,才决定哪些可以大批量备货。

Decision=Safety(op)∧Benefit(op)∧Consistency(env)\text{Decision} = \text{Safety}(op) \wedge \text{Benefit}(op) \wedge \text{Consistency}(env)
这个公式概括了折叠策略的决策准则:安全无副作用、收益大于代价、编译环境与运行环境语义一致——三者同时满足才执行折叠。
  1. 安全性检查确认算子无副作用、不引入未定义行为。
    Safety(op)\text{Safety}(op)
  2. 收益评估估算折叠带来的运行时节省是否超过编译成本。
    Benefit(op)=Savings−Costfold\text{Benefit}(op) = \text{Savings} - \text{Cost}_{\text{fold}}
  3. 环境一致性检查确认编译期浮点模型与运行时一致。
    Consistency(env)\text{Consistency}(env)
  4. 综合决策三者同时为真时执行折叠。
    Decision=Safety∧Benefit∧Consistency\text{Decision} = \text{Safety} \wedge \text{Benefit} \wedge \text{Consistency}

实战应用指南

大厨学完备菜理论,最终还是要回到自家餐厅的灶台上验证。下面我们走进两个最主流的“餐厅”——PyTorch 和 XLA,看常量折叠在真实 AI 编译管道中如何大展拳脚。

PyTorch 的常量折叠实践。PyTorch 2.x 的 torch.compile 工作流程中,Dynamo 负责把 Python 字节码翻译成 FX Graph,然后 ConstantFolder Pass 在 FX Graph 上执行常量折叠。实践中一个非常典型的应用场景是注意力掩码(Attention Mask)的预处理:masked_fill 算子接收的 mask 如果是从一个常量张量派生出来的,折叠就会把它直接变成一个常量 mask 张量,省去每次前向传播时重新生成 mask 的开销。对于 GPT 这样的自回归模型,因果关系(Causal Mask)是一个固定的下三角矩阵——折叠后它直接变成 Kernel 的一个常量参数,推理速度提升非常显著。

XLA 的折叠实践。XLA 中常量折叠是 AlgebraicSimplifier 的一部分,运行在 HLO IR 上。HLO 是静态的、基于张量的 IR,它的常量折叠不仅能处理标量,还能处理维度(Dimensions)和形状(Shapes)。例如在 Transformer 推理中,reshape 算子的目标形状往往是一个常量如 (1, 8, 128, 128),折叠后 reshape 直接被消除,下游算子直接使用新的形状元数据。更有意思的是,XLA 还能折叠 DynamicSlice、Pad 这类操作:如果输入和起始索引都是常量,它会在编译期就把裁剪/填充做完,运行时拿到的是已经裁剪好的张量。

调试与验证是落地环节最容易被忽视的一环。实践中建议三件事:第一,用 torch.fx 或 XLA 的 HLO dump 可视化折叠前后的图,直观确认哪些节点被替换;第二,建立“语义等价性测试”——对同一输入分别跑折叠前和折叠后的模型,比较输出是否完全一致(允许浮点误差小于 1e-6);第三,关注编译时间和运行时的 Trade-off,记录每次折叠策略调整对端到端延迟的影响曲线。

回到主大厨的比喻:PyTorch 的 torch.compile 就像一家标准化连锁餐厅——备菜流程已经自动化;XLA 则更像一家追求极致效率的中央厨房——甚至会提前把配菜做成料理包。无论哪种风格,验证环节都像“盲测”:不管备菜流程多花哨,端上桌的菜(输出结果)必须跟现炒的一样好吃(语义等价)。只要守住了这条底线,常量折叠就是你优化菜单上永远可靠的一道招牌菜。

Speedup=Toriginal−TfoldedToriginal×100%\text{Speedup} = \frac{T_{\text{original}} - T_{\text{folded}}}{T_{\text{original}}} \times 100\%
这个加速比公式是验证折叠效果的标准度量:用原始模型的单次推理时间减去折叠后的推理时间,除以原始时间得到百分比加速比。
  1. 测量原始延迟在同样硬件上测量不做任何优化时模型的单次推理时间。
    Toriginal=Time(modelno-fold)T_{\text{original}} = \text{Time}(\text{model}_{\text{no-fold}})
  2. 测量折叠延迟编译并测量开启常量折叠后模型的推理时间。
    Tfolded=Time(modelfold)T_{\text{folded}} = \text{Time}(\text{model}_{\text{fold}})
  3. 计算加速比两者差值除以原始时间得到相对提升百分比。
    Speedup=(Toriginal−Tfolded)/Toriginal\text{Speedup} = (T_{\text{original}} - T_{\text{folded}}) / T_{\text{original}}
  4. 判断阈值加速比为正且大于测量噪声(如 2%)即确认折叠有效。

AI 应用场景

场景:Transformer 推理加速

阶段:模型优化与部署

解决的问题:Transformer 模型推理时需要反复执行形状相关的常量和固定掩码运算,传统动态图模式每次前向传播都重复计算这些值,造成大量无谓开销。AI 编译器通过常量折叠把这些运算在编译期做完,推理引擎拿到的计算图已剔除全部冗余常量节点。

具体案例:某团队在 NVIDIA A100 上部署 GPT-2 模型进行文本生成服务。通过 torch.compile(model, mode="reduce-overhead") 启用包含常量折叠的编译优化后,Causal Mask 的创建与填充操作、层归一化中的 1/sqrt(hidden_size) 等系数计算均在编译期完成。最终单次推理延迟从 38ms 降至 27ms,加速比约 29%,且精度完全一致。

场景:XLA 编译 Stable Diffusion 性能调优

阶段:模型编译与图优化

解决的问题:Stable Diffusion 的 UNet 在采样循环中被多次调用,其中大量 Shape 计算(如 (batch, 4, height//8, width//8) 的推导、时间步嵌入常数)在多次迭代中从不变化。传统逐个算子执行的解释模式导致形状推导和常量分配重复发生。XLA 的常量折叠把这些计算全部提前到编译期固定下来。

具体案例:AI 平台团队使用 JAX 实现 Stable Diffusion 推理流水线,开启 XLA 编译后,gather 的位置索引、reshape 的目标形状、时间嵌入的三角位置编码全部被折叠为常量张量。采样 20 步的推理时间从 5.2 秒降至 3.1 秒,显存占用降低 18%,同时 HLO dump 显示计算图节点数从 2,400 余个减少到 1,100 余个。

场景:边缘设备上的轻量化部署

阶段:端侧部署前的模型优化

解决的问题:边缘芯片的算力和内存极为有限,模型二进制中的任何冗余计算都是对电池和内存的浪费。通过常量折叠把能提前计算的节点全部消除,可以显著减小模型文件大小并加快启动速度。

具体案例:物联网团队在树莓派上部署手势识别模型。使用 ONNX Runtime 的常量折叠优化(onnxruntime.transformers.optimizer)后,模型二进制从 24MB 压到 19MB,首次启动时间从 4.2s 降至 3.1s。模型输出与折叠前逐位一致。

Python 演示

python76 行
"""简易常量折叠器:遍历计算图,把全常量节点替换为常量值。""" from dataclasses import dataclassfrom typing import Any, Dict  @dataclassclass Node:    op: str    inputs: list    attrs: Dict[str, Any] = None  def const_or_value(node: Node, values: Dict[int, Any]):    """如果节点已折叠为常量,返回常量值;否则返回 None。"""    if node.op == "const":        return node.attrs["value"]    return values.get(id(node), None)  def eval_node(node: Node, arg_values: list) -> Any:    """在编译期执行算子(纯函数)。返回新常量值。"""    if node.op == "add":        return arg_values[0] + arg_values[1]    if node.op == "mul":        return arg_values[0] * arg_values[1]    if node.op == "sub":        return arg_values[0] - arg_values[1]    if node.op == "div":        return arg_values[0] / arg_values[1]    raise ValueError("unsupported op: " + node.op)  def constant_fold(graph: list) -> list:    """主函数:按拓扑序遍历所有节点,折叠可折叠节点。"""    values: Dict[int, Any] = {}    folded_nodes = set()    new_graph = []     for node in graph:        # 读取所有输入值(常量节点直接取值,否则查表)        arg_values = []        all_const = True        for inp in node.inputs:            val = const_or_value(inp, values)            if val is None:                all_const = False                break            arg_values.append(val)         if all_const and node.op != "const":            # 编译期求值并记录常量            values[id(node)] = eval_node(node, arg_values)            folded_nodes.add(id(node))            new_graph.append(Node("const", [], {"value": values[id(node)]}))        else:            new_graph.append(node)     return new_graph, folded_nodes  # 构建计算图:y = (3 * 4) + tt = Node("input", [], {"name": "t"})c3 = Node("const", [], {"value": 3})c4 = Node("const", [], {"value": 4})c5 = Node("const", [], {"value": 5})c6 = Node("const", [], {"value": 6})mul1 = Node("mul", [c3, c4])mul2 = Node("mul", [c5, c6])add1 = Node("add", [mul1, mul2])add2 = Node("add", [add1, t])graph = [t, c3, c4, c5, c6, mul1, mul2, add1, add2] optimized, folded = constant_fold(graph)print("折叠后的图节点数:", len(optimized))print("被折叠的节点数:", len(folded))
目的:演示常量折叠的核心算法——按拓扑序遍历计算图,将输入全为常量的节点提前求值并替换为常量节点。
  1. const_or_value(node, values)检查一个节点当前是否已是已知常量,统一了“节点本身就是常量”和“节点已被折叠”两种情况的取值逻辑。
  2. all_const and node.op != "const"折叠的边界条件:所有输入都是常量,且节点本身不是常量节点,防止对常量套娃死循环。
  3. eval_node(node, arg_values)在编译进程内真实执行算子的纯函数逻辑,模拟编译器的“编译期求值”动作。
  4. values[id(node)] = ...记录折叠结果,供下游节点判断是否可以继续折叠——这正是“迭代传播”的工程实现。
  5. folded_nodes.add(id(node))统计被折叠的节点,便于报告优化效果,方便验证折叠覆盖率。

常见误区

误区很多初学者把常量折叠(Constant Folding)和常数传播(Constant Propagation)当成一回事。他们认为“既然都是处理常量的优化,那名字不同只是叫法差异”。

正解两者是互补但不同的两件事。常数传播是把“某个变量是常量”的事实沿着数据流图向前传递——就像把“这瓶是酱油”的标签贴到分装瓶上。常量折叠则是对“所有输入都是常量”的表达式执行实际计算——就像看到酱油瓶+醋瓶,当场调出一碟蘸料。折叠消除计算,传播消除变量读取。实际编译器中两者交替运行、互相成就。

为何常见:两者常出现在同一篇文章中且紧密配合,加上名字都带“常量”,初学者很容易混淆。但理解区别对于正确阅读编译器 Pass 输出至关重要。

误区“只要把能折叠的都折叠掉,模型运行时不应该有任何冗余计算了。为什么我的推理还有这么多耗时?”不少开发者对折叠的覆盖范围抱有不切实际的期望。

正解常量折叠只能处理“输入在编译期完全确定”的算子。大部分真实模型的核心计算——矩阵乘法、卷积、激活函数——其输入是动态的张量数据,运行时依赖具体输入内容,折叠完全无法触及。折叠能消除的是边缘的“调料计算”,而不是“主菜的烹饪”。

为何常见:人们看到折叠对某个具体表达式有奇效,容易线性外推,忽略了大多数算子的输入本质上依赖运行时数据这一事实。

误区有的工程师认为折叠后图更小、计算更少,按理说延迟应该下降、吞吐应该上升,为什么实测只快了 1%,甚至某些场景反而变慢了?

正解折叠减少的往往是标量或形状运算,这类运算在现代 GPU 上通常被主计算“隐藏”(Overlapped),它们本就不是瓶颈。而折叠带来的编译时间增加却是实打实的成本。折叠是否划算取决于被折叠算子的运行次数和单次耗时——对于被调用百万次的 Kernel 内部的常量运算,折叠收益巨大;对于每次推理只执行一次的图级标量运算,折叠收益微乎其微。

为何常见:把“节点数量减少”等同于“端到端延迟降低”,忽略了实际性能瓶颈往往不在这部分被折叠的计算上。

误区“3 * 0.1 在编译期算和运行时算不都是 0.3 吗?有什么可担心的?”初学者普遍认为浮点折叠是绝对安全的。

正解浮点数在计算机中是不精确的。编译期求值走的是 CPU 的浮点指令路径,运行时求值走的是 GPU Kernel 的路径,两者的舍入顺序和精度可能不同。极端情况下编译期算出 0.30000000000000004,运行时算出 0.29999999999999993,虽然差一个 ULP(最小精度单位),但在某些数值敏感的算法中可能被放大。因此严格编译器会对涉及浮点折叠的情况按“快速数学”模式(如 -ffast-math)显式开启,默认保守处理。

为何常见:人类直觉中“数学计算结果是确定的”,但计算机浮点数的 IEEE 754 表示本身就是近似,折叠相当于在编译环境“先算一次”,与运行时“再算一次”可能产生不同的近似路径。

面试题

Q1. 常量折叠的核心目标是什么?

单选题 · easy

  • A.减少编译时间
  • B.把编译期可知的表达式预先计算并替换← 正确答案
  • C.提高代码可读性
  • D.将变量转换为常量

答案:B。常量折叠的精髓在于识别“所有输入都是常量”的表达式,在编译期算出结果并替换成常量节点。它不改变程序语义,只把计算提前。选 B。A 说反了——折叠通常增加编译时间;C 与可读性无关;D 是常数传播的职责,不是折叠。

Q2. 以下哪个表达式可以被常量折叠?

单选题 · easy

  • A.a + b(a、b 为变量)
  • B.3 \* x + 5(x 为运行时输入)
  • C.3 \ 4 + 5 \ 6← 正确答案
  • D.random\_uniform()

答案:C。可折叠的前提是算子的所有输入都是编译期常量。选项 C 中全是数字字面量,满足条件。A 和 B 含变量,运行时才能确定。D 虽然无输入,但具有随机副作用,折叠会改变语义,因此不可折叠。选 C。

Q3. 常数传播与常量折叠的区别是什么?

单选题 · easy

  • A.完整相同
  • B.传播把变量的常量值沿数据流传递,折叠对全常量表达式求值← 正确答案
  • C.折叠是传播的子集
  • D.没有本质区别

答案:B。常数传播(Constant Propagation)关心“某个变量的值在编译期是否已知并沿图传递”,常量折叠(Constant Folding)关心“整棵子表达式是否全是常量并计算掉”。传播是发现事实,折叠是执行计算,两者配合但不等同。答案选 B。

Q4. 关于折叠边界的判定,以下哪项描述最准确?

单选题 · mid

  • A.只要输出是数值,任何算子都能折叠
  • B.所有输入为常量且算子为纯函数时才能折叠← 正确答案
  • C.只有形状常量可以折叠
  • D.只有矩阵乘法可以折叠

答案:B。折叠的严格条件是“所有输入都是常量”且“算子无副作用是纯函数”。A 忽视纯函数要求,随机算子不满足;C 与 D 过于狭窄,标量和张量常量同样可折叠。答案选 B。

Q5. 常量折叠可能会带来哪类代价?

单选题 · mid

  • A.运行时内存增大
  • B.模型精度下降
  • C.编译时间增加← 正确答案
  • D.无法与部署流程集成

答案:C。折叠本质是用编译时间换运行时效率。折叠发生时编译器要真实执行运算,对大型模型可能显著拉长编译时间。折叠不会增加运行时内存,反而可能降低;不影响精度;也不影响部署。答案选 C。

Q6. 在 XLA 编译器中,常量折叠主要发生在哪个阶段?

单选题 · mid

  • A.Python 源代码解析阶段
  • B.前端 Python 到 HLO 的翻译阶段
  • C.HLO 图优化阶段← 正确答案
  • D.GPU Kernel 执行阶段

答案:C。XLA 的常量折叠实现位于 AlgebraicSimplifier 中,该 Pass 在 HLO IR 上进行图级优化。Python 解析阶段只负责语法分析,翻译阶段产生 HLO,Kernel 执行阶段已是运行时。答案选 C。

Q7. 哪些算子组合折叠是安全的?

单选题 · hard

  • A.涉及 Inf/NaN 特殊语义的算子组合
  • B.有符号整数溢出敏感的乘加组合
  • C.纯数学且无副作用的浮点 add、mul 组合← 正确答案
  • D.依赖全局随机种子的组合

答案:C。纯函数无副作用的算子(C)是安全的。A 与 B 因 IEEE 浮点语义和整型溢出未定义行为,折叠可能导致行为变化;D 有副作用依赖,折叠改变语义。答案选 C。

Q8. 哪些属于折叠后可产生的优化效果?

多选题 · hard

  • A.消减运行时冗余计算← 正确答案
  • B.为后续 Pass 提供精简计算图← 正确答案
  • C.使模型可解释性增强
  • D.提升内存规划效率空间← 正确答案

答案:ABD。折叠消除全常量算子的运行时执行(A);图变小后,算子融合、内存规划等后续优化(B、D)的求解空间更小、效果更好。折叠后图与源代码差异增大,可解释性反而下降,因此 C 不选。答案选 A、B、D。

Q9. 常量折叠为何在浮点运算上需要谨慎?

多选题 · hard

  • A.编译环境与运行时环境浮点精度可能不一致← 正确答案
  • B.IEEE 754 舍入规则在不同架构上有差异← 正确答案
  • C.Python 的浮点数完全精确
  • D.折叠可能改变 NaN 传播语义← 正确答案

答案:ABD。编译期求值走 CPU 指令,运行时走 GPU Kernel,舍入路径不同,可能导致 ULP 级误差(A、B)。IEEE 规范未定义所有边界情况下的 NaN/Inf 行为,折叠可能改变传播方式(D)。C 显然错误——浮点数不精确。答案选 A、B、D。

Q10. 若某张量在图中仅被一个 ReduceMean 算子消费,该算子的输入张量由 ones((4,4)) 全常量构造,以下哪个描述正确?

单选题 · extreme

  • A.该 ReduceMean 一定能被折叠为常量节点
  • B.每个算子需要单独判定,必须先判定张量自身常量性,再判定消费算子可折叠性← 正确答案
  • C.张量常量不能折叠
  • D.编译必须接受运行时参数才能折叠

答案:B。折叠流程是逐节点判定的:先确认 ones((4,4)) 是否为常量张量节点;若是,再检查 ReduceMean 是否纯函数、输入是否全为常量。这里 ones 可视为编译期常量求值产生的新常量节点,ReduceMean 的输入全部确定,故可折叠(A 过于绝对化,需走完判定流程)。答案选 B。