Graph Optimization（图优化）

3 秒看懂

图优化 = 把 AI 模型的”计算蓝图”（计算图）自动重写，合并冗余步骤、删掉废操作、调整执行顺序，让同一模型跑得更快、占更少内存、耗更少算力——AI 编译器的核心引擎。

类比：就像一份菜谱的”精简版”——把”先切葱，再切葱”合并为”一次切葱”，把”加热→冷却→再加热”优化为直接用温的食材，步骤少了，出菜速度翻倍，厨房空间也省了。

3 分钟产业解释

什么是计算图

深度学习框架（PyTorch、TensorFlow）会把用户写的 Python 模型代码转换成一张有向无环图（DAG）：

• 节点：算子（Conv2d、MatMul、ReLU、Softmax…）
• 边：张量数据流（tensor）的依赖关系

这张图就是模型的”执行蓝图”。框架 runtime 按图逐节点执行。

为什么需要图优化

原始计算图是”教科书式”的：忠于用户代码，但不考虑硬件特性。常见问题：

谁在做图优化

• 推理侧：NVIDIA TensorRT、Intel OpenVINO、Qualcomm QNN、Apple Core ML、华为 MindSpore Lite
• 训练/通用编译：Google XLA（TPU 生态）、PyTorch 2.0 的 torch.compile（TorchInductor/TorchDynamo）、Apache TVM、微软 ONNX Runtime
• 学术/开源：MLIR（LLVM 子项目）提供统一 IR，成为新一代编译基础设施

产业定位：图优化是 AI 编译器的”中端（middle-end）“——前端把框架模型转为 IR，中端做图级优化，后端生成硬件指令。它是连接”模型”和”芯片”的关键中间层。

15 分钟专家深入

图优化的技术分层

图优化并非单一 pass，而是一个多层级的优化流水线：

┌─────────────────────────────────────────────────────────────┐
│                      图优化流水线                             │
├─────────────────────────────────────────────────────────────┤
│  Level 0: 通用图变换                                          │
│    - 常量折叠（Constant Folding）                              │
│    - 死代码消除（Dead Code Elimination）                       │
│    - 公共子表达式消除（CSE）                                    │
│    - 冗余转置消除（Transpose Elimination）                     │
├─────────────────────────────────────────────────────────────┤
│  Level 1: 数学等价变换                                         │
│    - 算子代数简化（如 Conv+BN 融合）                            │
│    - MatMul + Bias + Activation → Fused GEMM                 │
│    - Softmax 数值稳定化改写                                    │
├─────────────────────────────────────────────────────────────┤
│  Level 2: 硬件感知优化（需后端反馈）                             │
│    - 算子融合（Operator Fusion / Kernel Fusion）                │
│    - 内存规划（Memory Planning / Liveness Analysis）            │
│    - 数据布局优化（Layout Transform）                          │
│    - 并行策略（如 Tensor Parallelism 切分点选择）                │
├─────────────────────────────────────────────────────────────┤
│  Level 3: 代数化/Tile 级优化                                   │
│    - Loop Tiling / Blocking                                   │
│    - 向量化 / Tensorization                                    │
│    - Prefetch / Pipeline 调度                                  │
└─────────────────────────────────────────────────────────────┘

关键 Pass 详解

1. 算子融合（Operator Fusion）

核心思想：多个算子合为一个 kernel launch，减少 kernel 启动开销和中间张量的显存读写。

常见融合 pattern：

工程难点：

• 融合不是无限的——GPU 每个 kernel 的寄存器/共享内存有限，融合过度会导致 occupancy 下降
• 不同硬件（NVIDIA GPU / CPU / NPU）的融合规则不同，需要 cost model 或 profile 指导

2. 常量折叠（Constant Folding）

编译期将纯常量子图计算出结果，直接替换为常量张量。

# 优化前
x = tf.constant([1, 2, 3])
y = tf.constant([4, 5, 6])
z = x + y  # 运行时才计算

# 优化后（常量折叠）
z = tf.constant([5, 7, 9])  # 编译期直接算好

对量化模型特别重要：量化参数的 scale/zero_point 计算往往是纯常量图。

3. 死代码消除（DCE）

移除不影响最终输出的节点。常见于：

• 模型剪枝后残留的未连接算子
• 条件分支中不可达的子图
• 调试/日志算子在生产环境的移除

4. 内存规划（Memory Planning）

问题：计算图执行过程中，各张量的生命周期（live range）不同，如何分配显存使峰值最小？

方法：

• 图着色算法：对 interference graph 做着色，同色张量可共享内存
• Greedy by size：按张量大小贪心分配
• 更先进的方法：考虑碎片化、对齐约束、多 stream 并发

实际影响：对大模型推理尤其关键——相同硬件，优化前可能 OOM，优化后可跑。

5. 数据布局优化（Layout Optimization）

不同硬件对 tensor 内存布局的偏好不同：

• NVIDIA Tensor Core：对 NHWC、特定 tiling 格式有硬件加速
• Intel VNNI 指令：要求特定的数据排列
• NPU/TPU：通常有固定的 layout 要求

图优化器需要在图中插入最少的 layout transform 节点，同时满足各算子的 layout 约束——这是一个约束满足问题。

技术原理（最深）

计算图的 IR 表示

图优化操作的是 中间表示（Intermediate Representation, IR）。不同框架的 IR 各有特点：

MLIR 的特殊地位：它不是一个 IR，而是一个 IR 框架（IR Framework）——允许定义不同抽象层级的 dialect，用统一 pass infrastructure 做变换。当前业界趋势是向 MLIR 收敛。

Pass Infrastructure

图优化的执行依赖于 Pass Manager：

PassManager
  ├── GraphPass 1 (常量折叠)
  ├── GraphPass 2 (DCE)
  ├── GraphPass 3 (算子融合 - pattern matching)
  ├── GraphPass 4 (布局优化)
  ├── GraphPass 5 (内存规划)
  └── GraphPass 6 (目标代码生成 - delegation to backend)

每个 Pass：

• 输入：一个 IR（计算图）
• 输出：变换后的 IR
• 可配置：是否依赖后端、是否需要 profile 数据

Pass 排序问题：不同 Pass 之间可能有依赖关系（如 DCE 需要在融合之后再跑一次）。Pass Manager 需要处理这些依赖。

算子融合的形式化

算子融合的本质是 子图替换（Subgraph Replacement）：

Pattern:  [A] → [B] → [C]   (A、B、C 为算子类型)
Replace:  [FusedABC]         (如果 A,B,C 的 shape/类型满足条件)

实现方式：

1. 基于规则：人工定义可融合的 pattern 白名单（TensorRT 主要方式）
2. 基于搜索：TVM 的 AutoTVM/Ansor，搜索最优的算子组合和 tiling 策略
3. 基于 Cost Model：训练一个 ML 模型预测融合后的性能，指导是否融合

图示——垂直融合的内存效果：

未融合：
  Conv2D:  读 input[224×224×3] → 计算 → 写 output[112×112×64]
  BatchNorm: 读 output[112×112×64] → 计算 → 写 output'[112×112×64]
  ReLU:    读 output'[112×112×64] → 计算 → 写 output''[112×112×64]
  
  显存带宽: 3次读 + 3次写 (中间张量各读写一次)

融合后：
  Conv2D_BN_ReLU:  读 input → 计算(在寄存器/共享内存中完成BN和ReLU) → 写 output
  
  显存带宽: 1次读 + 1次写
  节省: ~67% 显存带宽 (此例)

内存规划算法详解

核心问题是 区间图着色（Interval Graph Coloring）：

张量生命周期（live range）：
  t0: ████████                    (张量 A)
  t1:     ████████████            (张量 B)
  t2:         ████████            (张量 C)
  t3:                 ████████████(张量 D)
  
  时刻: 0  1  2  3  4  5  6  7  8

Interference Graph: (生命周期重叠的张量连边)
  A-B, A-C, B-C, B-D, C-D

图着色（最少颜色 = 最少内存块）：
  A→色1, D→色1  (不重叠，可复用)
  B→色2
  C→色3

结果: 3块内存，峰值 = max(|A|+|B|+|C|, |B|+|D|, ...)

实际工程中还需考虑：张量对齐要求、多 stream 并发、动态 shape 的不确定性。

动态 Shape 的图优化挑战

许多模型有动态维度（如序列长度可变）：

• 静态图优化假设 shape 固定，可做更激进的优化
• 动态 shape 下，融合规则可能因 shape 变化而失效
• 解决方案：Shape-specialized compilation（为不同 shape 范围编译不同版本）、Trace-and-specialize

技术演进史

2015  TensorFlow 1.x 引入静态计算图（Graph Mode）
       └─ 开启"先定义图，再执行"的范式
       
2016  TensorRT 1.0 发布
       └─ 业界首个商业化图优化推理引擎（NVIDIA）
       
2017  TVM 论文发表（陈天奇等，华盛顿大学）
       └─ 提出端到端编译栈，引入 AutoTVM 自动调优
       
2018  ONNX 标准逐渐成型
       └─ 计算图的跨框架交换成为可能

2018  PyTorch 1.0 引入 TorchScript / JIT
       └─ 尝试在 eager mode 框架上建静态图优化能力
       
2019  MLIR 项目启动（Chris Lattner，LLVM 之父）
       └─ 统一编译基础设施的愿景
       
2022  XLA 在 TPU 上默认使用，但在 TensorFlow 2.9 中 CPU/GPU 尚未默认启用（需手动开启）
       └─ 端到端图优化 + 代码生成
       
2023  PyTorch 2.0 发布 torch.compile
       └─ TorchDynamo (图捕获) + TorchInductor (后端) 架构
       └─ 标志着 eager mode 框架全面拥抱图优化
       
2023  StableHLO 成为 MLIR 生态的稳定 HLO dialect
       └─ JAX/XLA 生态向 MLIR 收敛
       
2024  图优化与大模型训练/推理深度结合
       └─ MoE routing 的图级优化、投机解码的图变换、
          KV-cache 管理的图级推理等成为热点
       
2025  图优化走向"编译时-运行时协同"
       └─ 在线 profiling → 图重编译（如 PyTorch 的
          profile-guided optimization in Inductor）

关键转折点：

• 2022 年 torch.compile 是分水岭——证明了”易用性（eager mode）+ 高性能（图优化）“可以兼得
• MLIR 的崛起改变了图优化的基础设施层

技术路线对比

趋势判断：

• 推理侧：TensorRT 在 NVIDIA 生态仍占主导，但 TensorRT-LLM 等专用优化器在崛起
• 训练侧：PyTorch + torch.compile 逐步成为主流
• 编译基础设施：MLIR 方向是共识

上下游

上游（图优化的输入来源）

下游（图优化的输出去向）

产业链位置

┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐
│  模型训练 │ →  │ 模型导出  │ →  │ 图优化   │ →  │ 代码生成  │ → 硬件执行
│(PyTorch) │    │(ONNX等)  │    │(编译中端) │    │(编译后端) │
└──────────┘    └──────────┘    └──────────┘    └──────────┘
                                 ↑
                          我们讨论的是这里
                    "AI编译器的心脏"

关键指标

性能指标

编译指标

优化覆盖率指标

供需与市场数据

需求侧驱动力

1. 大模型推理成本：LLM 推理成本占 AI 部署总成本的大头（估算 >60%），图优化是降本核心手段
2. 端侧 AI：手机/IoT/车载芯片算力有限，图优化是”免费”获取性能的关键
3. 多硬件生态：云端 GPU/TPU、边缘 NPU、端侧 DSP，每个都需要图优化适配
4. 训练效率：大模型训练成本极高（单次 GPT-4 级训练估算数千万美元），训练侧图优化也有显著经济价值

供给侧格局

市场规模（估算）

• AI 编译器/优化工具作为独立市场尚无权威统计口径——它嵌入在推理框架、芯片 SDK、云服务中
• 间接参考：全球 AI 推理市场（包含硬件+软件）规模在数百亿美元量级 [多家行研估算]，图优化是推理软件栈的核心价值环节
• Modular（Chris Lattner 创办）2022 年融资估值超过 $1B [公开报道]，反映资本市场对 AI 编译基础设施的重视

代表公司与资本映射

资本映射逻辑：

• 纯做 AI 编译器的上市公司极少——它主要是巨头的基础设施投入
• 间接受益标的：NVIDIA（TensorRT 强化 CUDA 生态护城河）、AI 编译器创业公司（一级市场）
• 国内关注：华为昇腾生态的编译器成熟度、国产 GPU 厂商的编译器能力

投资逻辑

核心投资主题

1. “编译器即护城河”

   • NVIDIA 的 CUDA + TensorRT 生态是其 AI 芯片霸主地位的核心护城河
   • 图优化器的算子覆盖、融合质量、硬件适配是多年工程积累，难以快速复制

2. “编译器决定硬件可用性”

   • 同一款芯片，有无成熟的图优化工具链，性能可能差 2~10 倍 [估算]
   • 这直接决定芯片是否能进入实际部署

3. “AI 编译器创业窗口”

   • Modular 的 Mojo 语言试图从编程语言层面重新定义 AI 编译
   • 中国 AI 编译器人才稀缺，创业公司有机会

风险点

常见误读纠偏

误读 1：“图优化 = 量化”

纠偏：量化（INT8/FP16 量化）是图优化可应用的一个场景，但图优化的范围远大于量化。量化是数据类型变换，图优化是图结构变换（融合、消除、重排、内存规划等）。两者是正交维度，可叠加使用。

误读 2：“图优化能无限加速，XX 倍提升”

纠偏：图优化的加速受限于硬件 roofline 和原始图的优化空间：

• 如果模型已经是 compute-bound（计算瓶颈），减少显存带宽的融合收益有限
• 如果原始框架已经做了基本优化，进一步提升空间更小
• 实际加速倍数因模型、硬件、baseline 对比不同而差异巨大——宣传中的”10x”往往是以未优化 eager mode 为 baseline，参考价值有限

误读 3：“torch.compile 能替代 TensorRT”

纠偏：两者定位不同：

• torch.compile 是 PyTorch 训练+推理的通用编译优化，覆盖多种硬件
• TensorRT 是 NVIDIA GPU 推理专用的深度优化引擎，包含图优化 + 算子自动调优 + 精度校准
• 在 NVIDIA GPU 推理场景，TensorRT 的优化深度通常优于 torch.compile
• torch.compile 的优势在于易用性和训练支持

误读 4：“图优化是一次性编译，不需要运行时信息”

纠偏：现代图优化正走向 编译时-运行时协同：

• Profile-Guided Optimization (PGO)：先用运行时 profile 数据指导编译优化
• Speculation + Recompilation：如 PyTorch 2.0 会先假设静态 shape 编译，遇到新 shape 时重新编译
• Continuous Profiling：运行时持续监测，必要时触发重优化

学习路径

入门（~2 周）

1. 理解计算图：用 PyTorch 画一个简单模型的计算图（torch.jit.trace 或 torch.fx）
2. 跑一遍 ONNX：导出模型为 ONNX，用 Netron 可视化，感受”图长什么样”
3. 体验 torch.compile：对同一个模型，对比 eager mode vs torch.compile 的速度

进阶（~1-2 月）

4. 读 TVM 教程：TVM 官方教程 的 Relay pass 部分
5. 读 TensorRT 开发者指南：了解 fusion rules、layer fusion 的实操
6. 学 MLIR 基础：MLIR 官方教程——理解 dialect、pass、pattern rewrite

深入（~3-6 月）

7. 读论文：
   • “TVM: An Automated End-to-End Optimizing Compiler for Deep Learning” (OSDI 2018)
   • “Glow: Graph Lowering Compiler Techniques for Neural Networks” (arXiv 2018)
   • “MLIR: Scaling Compiler Infrastructure for Domain Specific Computation” (CGO 2021)
8. 读源码：PyTorch `to