AI 编译图优化：常量折叠不是玩具优化

AI 编译图优化：常量折叠不是玩具优化

AI 编译器里，常量折叠常被当成入门优化：能在编译期算出的东西就提前算掉。听起来简单，甚至有点像玩具。但在推理图里，常量折叠能减少运行时算子、降低内存占用、简化后续融合路径。尤其是部署端固定输入形状、固定权重、固定归一化参数时，它很实用。

常量折叠的价值，不在于少算一个小表达式，而在于让整张图更安静。

一、先识别常量子图

flowchart TD
  A[Const A] --> C[Add]
  B[Const B] --> C
  C --> D[Const Result]
  E[Input] --> F[Runtime Op]

常量子图的输入全部来自权重、固定参数或编译期已知值。比如固定 mean/std 的归一化参数、reshape 目标 shape、位置编码的一部分、mask 模板，都可能提前计算。

识别时要注意边界。只要子图依赖运行时输入，就不能折叠。编译器需要清楚区分 weight、attribute、shape constant 和真正 runtime tensor。

二、折叠结果要替换图节点

enum NodeKind {
    Const(Tensor),
    Op { op: String, inputs: Vec<NodeId> },
}

常量折叠不是把值算出来就结束，还要把原来的 op 节点替换成 Const 节点，并更新下游引用。被替换掉的中间节点如果不再被使用，可以从图里删除。这样后续 pass 才能看到更简单的图。

图优化 pass 之间有顺序关系。常量折叠后，某些算子融合机会才会出现；算子融合后，又可能产生新的常量折叠机会。因此优化通常要迭代到稳定状态，而不是只跑一次。

三、精度和 dtype 要一致

fold_rules:
  preserve_dtype
  preserve_shape
  check_overflow

常量折叠必须保持运行时语义。编译期用 f64 算，运行时本来用 fp16，结果可能不一致。量化模型里更要小心 scale、zero_point 和溢出。折叠不能为了方便改变 dtype。

测试要覆盖边界值。NaN、Inf、整型溢出、广播规则、空 shape，这些都可能让常量折叠出错。编译器优化看似数学正确，工程上必须尊重框架语义。

精度一致性的难点在量化路径。一个典型坑是：推理图用 INT8 量化，zero_point 为 128，scale 为 0.02，但常量折叠在编译期用 float 直接计算了 Add(Const_a, Const_b) 的结果，然后尝试以 INT8 形式回存。此时如果不重新量化——即根据原始 scale 和 zero_point 反算 INT8 表示——存入图的常量就会产生浮点到整型的截断误差，在 32 层网络里逐层放大后推理结果完全跑偏。正确的做法是：折叠后的常量必须沿用原始张量的量化参数，必要时做 requantize 而非简单 round，并在折叠前后对数值做逐元素容忍度检查。而不同框架——TensorFlow Lite、ONNX Runtime、TVM——在 requantize 的 rounding mode 和 saturation 处理上各有差异，跨框架移植优化 pass 时必须逐一确认语义对齐。

四、收益要可观测

{
  "folded_nodes": 18,
  "removed_ops": 11,
  "const_bytes_added": 4096,
  "runtime_ms_delta": "-3.2"
}

常量折叠也有成本。有时提前计算会增加常量体积，导致模型文件变大。优化报告要记录折叠节点数、删除算子数、常量体积变化和运行时收益。没有数据，优化 pass 只是自我感觉良好。

部署端还要看启动时间。模型加载时少算一点，推理时少算一点，文件体积大一点，这三者要平衡。编译优化最终服务的是部署目标。

常量折叠还要保留调试信息。图节点被替换后，如果推理结果异常，工程师需要知道原始子图是什么、折叠后的 tensor 来自哪些节点。优化 pass 如果只留下最终图，会让调试变得很痛苦。生产编译器必须能解释自己做过什么。

可以为每个 pass 输出差异报告：删除了哪些节点，新增了哪些常量，哪些节点因为动态输入无法折叠。这样优化结果可审计，也方便在某个 pass 引入 bug 时快速关闭。

五、总结

AI 编译图里的常量折叠不是玩具优化。它能简化图结构、减少运行时算子、打开后续融合机会，但必须保持 dtype、shape 和框架语义一致。

优秀的编译优化不是把图改得花哨，而是让运行时少做那些早就可以确定的事。