常量折叠与传播

深度学习 (deep learning)模型中的并非所有操作都依赖于输入数据。许多计算涉及固定参数 (parameter)、超参数 (hyperparameter)或配置值，这些值在每次推理 (inference)运行时保持不变。当编译器遇到这些静态子图时，它提供了一个机会，可以在编译期间执行一次计算，而不是为每个输入样本重复执行，从而降低运行时负载。这种优化方法称为常量折叠。

常量折叠分析计算图，以找出所有输入都是静态已知值的操作。编译器评估这些操作，并将原始子图替换为包含预计算结果的单个节点。此过程与常量传播配合使用，新计算出的常量被传递到后续操作，可能由此引发更多折叠机会的连锁反应。

静态求值的机制

常量折叠的主要目标是缩小计算图。在典型的神经网络 (neural network)中，常量以多种形式出现：

• 模型权重 (weight)和偏置 (bias)： 这些在训练后固定。
• 超参数 (parameter) (hyperparameter)： 例如归一化 (normalization)epsilon、裁剪阈值或Leaky ReLU斜率等值。
• 形状算术： 用于确定reshape或slice操作的张量维度的计算。

考虑一个场景，模型包含一个源自两个常量的缩放因子。在Python中，这可能看起来像 x * (a / b)，其中 a 和 b 是在模型配置中定义的标量。如果没有优化，执行引擎在模型每次运行时都会执行 a / b 除法。常量折叠会检测到 a 和 b 是不可变的，在编译时计算商 c = a / b，然后重写图以执行 x * c。

以下图示说明了编译器如何识别并折叠静态子图。

一个图的转换，其中静态加法操作被预计算并替换为单个常量节点。

识别可折叠节点

为实现常量折叠，编译器必须遍历中间表示（IR），并判断哪些节点是静态求值的候选者。这通常通过数据流分析方法实现。

编译器遍历图，通常按拓扑顺序。对于每个节点，它检查其传入边（操作数）的属性。如果一个节点的所有操作数都是常量节点，则该节点被认为是“可折叠”的。

如果一个节点是可折叠的，编译器会调用解释器或轻量级执行引擎来计算结果。这个解释器的复杂程度各不相同。对于简单算术，编译器可以使用主机CPU的本地指令。对于复杂张量操作，编译器可能需要链接到算子的参考实现，以确保编译时结果与目标硬件的运行时行为一致。

传播与级联

这种优化的优势源于其递归特性。当一个节点被折叠后，它实际上成为图中的一个新常量。这个新常量可能正是允许下游节点也进行折叠的缺失部分。

设想一个表示以下计算的图： $y = x + \sqrt{100} + 5$

1. 初始状态： 图中包含 $x$ 的输入、常量 $100$ 和常量 $5$。
2. 第一次处理： 编译器看到 Sqrt 节点以 $100$（常量）作为输入。它计算 $\sqrt{100} = 10$。 Sqrt 节点被替换为包含 $10$ 的常量节点。
3. 传播： 下游的 Add 节点最初接收 Sqrt 的输出和常量 $5$。现在，它看到两个常量：$10$ 和 $5$。
4. 第二次处理： Add 节点现在可折叠。编译器计算 $10 + 5 = 15$。
5. 最终状态： 图被简化为 $y = x + 15$。

这种级联效应要求优化过程迭代运行或使用工作列表算法，直到图达到稳定状态，不再能进行任何简化。

批归一化 (normalization)中的优化

常量折叠的一个实用且影响大的应用出现在推理 (inference)期间的批归一化（BatchNorm）层中。在训练期间，BatchNorm 会减去批量均值并除以批量标准差。然而，在推理期间，这些统计数据是固定的移动平均值。

批归一化的标准推理公式如下：

$y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \gamma + \beta$

此处，$\mu$（均值）、$\sigma^2$（方差）、$\gamma$（缩放）和 $\beta$（偏移）都是学习参数 (parameter)，$\epsilon$ 是一个配置常量。

朴素执行会在运行时进行减法、加法、平方根、除法、乘法和另一个加法。然而，由于 $\mu, \sigma, \epsilon, \gamma, \beta$ 在推理时都是常量，编译器可以将它们折叠成两个不同的因子：一个乘法缩放因子 $w_{new}$ 和一个偏置 (bias)项 $b_{new}$。

$w_{new} = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}$ $b_{new} = \beta - \frac{\mu \gamma}{\sqrt{\sigma^2 + \epsilon}}$

运行时操作简化为线性变换 $y = w_{new}x + b_{new}$。如果批归一化层紧随一个卷积层之后，这些新常量通常可以直接折叠到卷积本身的权重 (weight)中，从而完全从运行时图中移除批归一化算子。

形状推断与静态切片

常量折叠不限于浮点张量数据。它大量用于与张量形状相关的整数算术。现代动态框架通常会生成复杂的子图，仅仅是为了计算 Reshape 或 Transpose 操作后张量的形状。

例如，将张量从 [Batch, Channels, Height, Width] 展平为 [Batch, Features] 涉及对维度大小的算术运算。如果输入图像大小是固定的（例如 $224 \times 224$），编译器可以预计算展平向量 (vector)的精确大小。这消除了推理 (inference)内核中的整数算术指令，并使内存分配器能够在模型运行前准确知道需要多少缓冲区空间。

权衡与局限

常量折叠通常有益，但也存在编译器工程师必须处理的极端情况。

代码大小膨胀： 如果常量计算生成的张量显著大于原始输入（例如，在小常量上使用 np.tile 或 torch.repeat），对其进行折叠可能会在编译后的二进制文件中存储一个巨大的数组。这会增加二进制文件大小和内存占用。如果结果常量超过某个大小阈值，编译器通常会使用启发式方法阻止折叠。

精度差异： 执行编译的机器可能与目标加速器有不同的浮点精度或舍入行为。在x86 CPU上预计算的值，与在嵌入 (embedding)式NPU或GPU上计算得到的结果，可能存在位级差异。对于数值敏感的模型，这种差异可能有效改变模型的准确性。

通过积极识别并处理静态图组件，常量折叠作为清理阶段，简化结构以暴露内存布局和循环调度方面的更多优化机会。