4. 图重写入门

在深度学习编译器的前端与后端之间，计算图（Graph）是承载模型逻辑的核心中间表示。无论是 TensorFlow 的 GraphDef、PyTorch 的 TorchScript，还是 TVM 的 Relay/TIR，编译器都需要在图上执行一系列变换，以降低计算开销、减少内存占用、提升算子执行效率。这些变换统称为图重写（Graph Rewriting）。图重写是编译器优化中承上启下的关键环节：它不改变模型的数学语义，但通过结构调整和算子替换，使计算图更适合目标硬件执行。本章将系统讲解图重写的定义、常见优化模式（算子融合、常量折叠、死代码消除），并结合实际案例拆解其工程实现路径。

图重写核心优化模式总览 原始计算图 (IR / Graph) 算子融合 Conv+BN+ReLU 减少访存与内核启动 常量折叠 编译期预计算 消除运行时开销 死代码消除 删除无用节点 减少内存与计算 优化后计算图 更少节点、更低内存 更高执行效率 图重写 = 模式匹配 + 条件检查 + 变换应用 + 语义保持验证

4.1 图重写的核心原理

图重写本质上是一个基于模式匹配的树/图变换系统。编译器维护一组重写规则（Rewrite Rule），每条规则包含三个要素：源模式（LHS）、目标模式（RHS）和条件约束（Condition）。当计算图中的某个子图与源模式匹配，且满足条件约束时，编译器将该子图替换为目标模式。整个过程必须保证语义等价——即替换前后，对于任意合法输入，输出张量的数值保持一致（或仅在浮点舍入误差范围内有微小差异）。

从工程实现角度看，图重写通常分为四个阶段：

  1. 图遍历：采用拓扑序或后序遍历，访问图中的每个节点。
  2. 模式匹配：在访问节点时，尝试将当前子图与规则库中的 LHS 进行匹配。匹配过程需要考虑节点类型、边连接关系、属性值（如轴、步长、数据类型）等。
  3. 条件验证：即使结构匹配，还需验证附加条件。例如算子融合要求数据布局一致、常量折叠要求输入全部为常量、死代码消除要求节点没有外部副作用。
  4. 变换与重写：执行替换操作，更新图的拓扑结构，并重新建立数据依赖关系。重写后需要维护一个工作列表（Worklist），记录可能受影响的节点，以便后续迭代优化。

图重写系统通常采用贪心策略，即每次匹配成功后立即应用变换，直到没有新的规则可匹配为止。这种策略简单高效，但可能错过某些需要全局考虑的优化机会。因此，现代编译器（如 TVM 的 Relay）会结合pass 管理器，将多个重写规则组织成有序的 pass 序列，并在 pass 之间进行图规范化（如 DCE、常量折叠）以暴露更多优化机会。

关键理解：图重写不是简单的“字符串替换”，而是基于数据依赖关系的结构变换。每次重写都必须重新验证图的合法性（如类型正确、形状匹配、控制流完整），否则会导致编译错误或运行时崩溃。

4.2 常见优化模式之一：算子融合

算子融合是图重写中最重要、收益最明显的优化手段。其核心思想是：将多个连续的小算子合并为一个大的融合算子，从而减少内核启动次数、降低中间张量的内存读写开销、提升数据局部性。

以深度学习中最常见的 Conv2D + BatchNorm + ReLU 为例。在未融合的图中，推理时需要依次执行三个内核：卷积计算、BatchNorm 归一化、ReLU 激活。每次内核执行都需要从全局内存读取输入、写回输出。假设特征图尺寸为 [N, C, H, W]，则中间张量需要额外占用 N*C*H*W*4 字节的内存（FP32）。对于大模型，这会导致显著的内存带宽压力。

融合后的计算图将三个算子合并为一个 FusedConvBNReLU 算子。在推理阶段，BatchNorm 的参数可以预先折叠到卷积的权重和偏置中（因为推理时均值和方差是固定的），从而进一步减少计算量。融合后的内核在 GPU 上可以一次性完成卷积、缩放、偏移和激活操作，中间结果仅保存在寄存器或共享内存中，无需写回全局内存。

算子融合的判断条件通常包括：

  • 算子之间是否存在数据依赖（必须是直连的 producer-consumer 关系）。
  • 融合后是否改变数值语义（例如 ReLU 与 Conv 融合是安全的，但某些非线性变换可能不满足结合律）。
  • 目标硬件是否支持融合内核（例如 GPU 上需要手写 CUDA kernel 或使用 cuDNN 的融合 API）。
  • 数据布局是否一致（如 NCHW 与 NHWC 混用时不建议直接融合）。
融合类型示例主要收益典型硬件
逐元素融合Add + ReLU, Mul + Add减少内核启动，降低访存CPU/GPU
卷积融合Conv + BN + ReLU权重折叠，减少中间张量GPU/NPU
布局转换融合Transpose + Conv消除显式数据重排GPU
Reduce 融合Sum + Mean + Scale减少归约次数CPU/GPU

4.3 常见优化模式之二：常量折叠

常量折叠（Constant Folding）是指在编译期计算所有输入均为常量的子表达式，将结果直接替换为常量节点。例如，计算图中的 shape = [2, 3] 与 zeros = Constant(shape) 的拼接操作，如果拼接的输入全部是常量，编译器可以直接计算出结果张量，并在图中用一个新的常量节点替代整个子图。

常量折叠的核心价值在于：

  • 减少运行时计算开销，尤其是对于形状推导、索引计算、掩码生成等辅助操作。
  • 为后续优化（如死代码消除）创造机会——折叠后的常量节点如果未被使用，可以被 DCE 直接删除。
  • 降低内存占用，因为常量可以存储在只读内存区域，甚至直接嵌入到内核代码中。

实现常量折叠时，编译器需要维护一个常量传播机制。当某个节点的所有输入都是常量时，该节点被标记为“可折叠”。然后编译器调用该算子的 CPU 实现（或解释器）执行计算，得到输出常量。需要注意的是，并非所有算子都适合常量折叠：

  • 具有随机性的算子（如 dropout、随机正态分布）不能折叠，除非指定固定种子。
  • 涉及全局状态或文件 I/O 的算子不能折叠。
  • 浮点运算的折叠可能改变舍入行为，需要根据编译选项决定是否允许。
工程经验：在 TVM Relay 中，常量折叠 pass 通常与 SimplifyExpr 结合使用。例如 shape_of(Constant) 可以直接折叠为常量形状值，从而避免在运行时动态获取形状。

4.4 常见优化模式之三：死代码消除

死代码消除（Dead Code Elimination, DCE）用于删除计算图中对输出没有任何影响的节点。这些节点可能来自：模型定义时的冗余分支、调试用的打印节点、常量折叠后遗留的孤立子图、或者用户手动添加但未使用的中间变量。

DCE 的判断方法基于可达性分析：从图的输出节点（或带有副作用标记的节点，如 Print、Assert、Send）出发，反向遍历数据依赖边，标记所有可达节点。未被标记的节点即为死代码，可以安全删除。

DCE 的操作流程如下：

  1. 初始化：将所有输出节点和副作用节点加入工作列表。
  2. 遍历：从工作列表中取出节点，将其所有输入节点加入工作列表（如果尚未访问）。
  3. 标记：记录所有被访问的节点。
  4. 删除：遍历图中所有节点，删除未被标记的节点及其输出边。

DCE 的收益不仅在于减少计算量，更重要的是释放内存。在内存规划阶段，死节点对应的中间张量不会分配内存，从而降低峰值内存占用。此外，DCE 还能简化图结构，使后续的算子融合和内存复用更容易匹配。

注意：DCE 不能删除带有副作用的节点，即使其输出未被使用。例如，一个用于统计模型中间层输出的 Print 节点，虽然不影响最终预测结果，但用户期望看到打印信息，因此必须保留。编译器通常通过节点的 side_effect 属性来区分。

4.5 案例拆解：ResNet-50 推理图优化实战

下面通过一个 ResNet-50 推理场景，完整展示图重写如何协同工作。假设我们使用 TVM Relay 作为前端，目标硬件为 NVIDIA V100 GPU。

初始状态：从 PyTorch 导出的 ONNX 模型中，每个残差块包含 Conv2d → BatchNorm → ReLU 三个独立算子，且存在大量 Add（残差连接）和 Constant 节点（如 BatchNorm 的 gamma、beta、mean、var）。初始图共有约 350 个节点，其中约 120 个是常量节点。

优化目标：减少内核启动次数，降低中间张量内存占用，将推理延迟从 12ms 降低到 8ms 以下。

执行步骤：

  • Step 1：常量折叠。将 BatchNorm 的四个参数（gamma、beta、mean、var）与卷积权重进行数学合并，生成新的权重和偏置。这一步将 120 个常量节点缩减为 50 个，同时消除了 50 个乘法/加法节点。
  • Step 2：算子融合。将 Conv2d + BN + ReLU 融合为单个 FusedConvBNReLU 算子。对于残差连接中的 Add，如果其输入之一是 ReLU 的输出，则进一步融合为 FusedConvBNReLUAdd。融合后节点数从 350 降至 180。
  • Step 3：死代码消除。融合过程中产生的孤立常量节点（如旧的 BN 参数）以及未被使用的中间节点被 DCE 删除。最终节点数降至 150。
  • Step 4：内存规划。由于融合算子不再产生中间张量，内存分配器只需为每个残差块的输入和输出分配内存，峰值内存从 2.1GB 降至 1.4GB。

结果对比：

指标优化前优化后提升幅度
图节点数350150-57%
内核启动次数21095-55%
峰值内存2.1 GB1.4 GB-33%
推理延迟12.0 ms7.8 ms-35%

复盘经验：图重写的收益是叠加的。常量折叠为算子融合创造了更简洁的图结构，而算子融合产生的死节点又为 DCE 提供了素材。在实际工程中，建议按照“常量折叠 → 算子融合 → DCE → 内存规划”的顺序组织 pass，并在每个 pass 后重新运行 DCE 以清理中间产物。

4.6 常见误区与总结

图重写虽然基础，但在实际应用中存在几个常见误区：

  • 过度融合导致寄存器溢出：将过多算子融合为一个内核，可能导致寄存器使用超过硬件限制，反而降低性能。需要根据目标架构的寄存器文件大小设置融合上限。
  • 忽略浮点语义变化：某些融合（如将 a*b + c 融合为 FMA）会改变舍入行为。对于对数值精度敏感的训练场景，需要谨慎启用。
  • DCE 误删副作用节点：在分布式训练或模型调试中，某些节点虽然不参与前向计算，但承担梯度聚合或日志记录功能。删除它们会导致训练异常。
  • 常量折叠导致代码膨胀：如果常量过大（如大矩阵），将其嵌入内核代码会增加指令缓存压力。此时应保留为全局内存常量。

总结而言，图重写是深度学习编译器中成本最低、收益最稳定的优化手段。掌握算子融合、常量折叠和死代码消除三种基本模式，是理解后续内存规划、调度生成和自动调优的基础。在实际项目中，建议建立一套可复用的重写规则库，并配合图可视化工具（如 Netron）验证重写前后的结构差异，确保语义正确性。

← 上一章 📖 返回目录 下一章 →