第 7 章：算子融合

算子融合是 AI 编译器优化的核心技术之一，通过将多个细粒度算子合并为单个粗粒度算子，显著减少内存访问开销和核函数启动开销。在自动驾驶和具身智能等对延迟敏感的场景中，算子融合可以将推理延迟降低 30-50%，同时减少内存带宽压力。本章将深入探讨算子融合的理论基础、实现策略和性能分析模型。

7.1 垂直融合与水平融合

算子融合可以从两个维度进行分类：垂直融合关注生产者-消费者关系的算子链，水平融合则并行合并独立的算子。理解这两种融合模式的特点和适用场景是设计高效融合策略的基础。

7.1.1 垂直融合（Producer-Consumer Fusion）

垂直融合将具有生产者-消费者关系的算子序列合并为单个算子。考虑计算图中的算子链：

$$ y = \phi_n(\phi_{n-1}(...\phi_2(\phi_1(x)))) $$

其中每个 $\phi_i$ 代表一个算子变换。未融合时，中间结果需要写入和读取全局内存：

算子链执行流程（未融合）：
    x → [全局内存] → φ₁ → [全局内存] → φ₂ → ... → φₙ → y

    内存访问次数：2n 次（n次读 + n次写）
    核函数启动：n 次

融合后的执行流程：

算子链执行流程（融合后）：
    x → [寄存器/共享内存] → φ₁∘φ₂∘...∘φₙ → y

    内存访问次数：2 次（1次读 + 1次写）
    核函数启动：1 次

融合收益分析：

设输入张量大小为 $N$，每个元素占用 $b$ 字节，内存带宽为 $B$，则：

• 未融合的内存传输时间：$T_{unfused} = \frac{2nNb}{B}$
• 融合后的内存传输时间：$T_{fused} = \frac{2Nb}{B}$
• 加速比：$S = \frac{T_{unfused}}{T_{fused}} = n$

7.1.2 水平融合（Parallel Fusion）

水平融合将多个独立的算子并行执行，共享输入数据的内存读取：

并行算子结构：
         ┌→ φ₁(x) → y₁
    x → ─┼→ φ₂(x) → y₂
         └→ φ₃(x) → y₃

这种模式特别适合于多头注意力机制、多尺度特征提取等场景。

融合条件判定：

两个算子 $\phi_1$ 和 $\phi_2$ 可以水平融合的必要条件：

1. 输入兼容性：$\text{input}(\phi_1) \cap \text{input}(\phi_2) \neq \emptyset$
2. 无数据依赖：$\nexists$ 路径 $\phi_1 \rightsquigarrow \phi_2$ 或 $\phi_2 \rightsquigarrow \phi_1$
3. 资源约束满足：$R(\phi_1) + R(\phi_2) \leq R_{max}$

其中 $R(\cdot)$ 表示资源需求（寄存器、共享内存等）。

7.1.3 混合融合策略

实际编译器常采用混合策略，同时进行垂直和水平融合：

混合融合示例：
    x → [Conv → BN → ReLU] ─┬→ [MaxPool] → y₁
                             └→ [AvgPool] → y₂

    垂直融合：Conv-BN-ReLU
    水平融合：MaxPool-AvgPool

融合决策的成本模型：

给定算子图 $\mathcal{G} = (V, E)$，融合决策可以形式化为划分问题：

$$ \min_{\mathcal{P}} \sum_{P \in \mathcal{P}} C(P) + \sum_{(P_i, P_j) \in E'} T(P_i, P_j) $$

其中：

• $\mathcal{P}$ 是节点集 $V$ 的一个划分
• $C(P)$ 是融合块 $P$ 的计算成本
• $T(P_i, P_j)$ 是块间数据传输成本
• $E'$ 是划分后的块间边集

7.1.4 自动驾驶场景的特殊考虑

在自动驾驶场景中，算子融合需要考虑实时性约束：

1. 确定性延迟：融合后的算子执行时间必须可预测
2. 优先级感知：关键路径（如障碍物检测）的融合优先级更高
3. 异构执行：某些算子可能需要在不同硬件上执行（GPU/DSP/NPU）

延迟约束下的融合：

设系统延迟上界为 $L_{max}$，则融合决策需满足：

$$ \max_{path \in \text{CriticalPaths}} \sum_{P \in path} L(P) \leq L_{max} $$

其中 $L(P)$ 是融合块 $P$ 的延迟。

7.2 融合规则与模式匹配

算子融合的自动化需要定义清晰的融合规则和高效的模式匹配算法。本节介绍如何形式化融合规则并实现高效的图模式匹配。

7.2.1 融合规则的形式化定义

融合规则可以表示为模式-动作对 $(Pattern, Action)$：

模式定义语言：

$$ Pattern ::= Op | Pattern \rightarrow Pattern | Pattern \parallel Pattern | Pattern^* $$

其中：

• $Op$ 表示基本算子
• $\rightarrow$ 表示顺序连接
• $\parallel$ 表示并行结构
• $*$ 表示重复模式

常见融合模式库：

1. Element-wise链融合： $$EW^+ ::= (Add | Mul | Activation)^+$$

2. 卷积-归一化-激活融合： $$CBA ::= Conv \rightarrow BatchNorm \rightarrow Activation$$

3. 注意力机制融合： $$Attention ::= (Q \parallel K \parallel V) \rightarrow MatMul \rightarrow Softmax \rightarrow MatMul$$

7.2.2 图匹配算法

模式匹配可以转化为子图同构问题。采用 VF2 算法的变种进行高效匹配：

算法框架：

算法：FusionPatternMatch
输入：计算图 G = (V, E)，模式集合 P
输出：匹配实例集合 M

1. 初始化 M = ∅
2. 对每个模式 p ∈ P：
   a. 构建模式的特征向量 f(p)
   b. 筛选候选起始节点 C = {v ∈ V | compatible(v, p.root)}
   c. 对每个 c ∈ C：

      - 执行递归匹配 match(c, p.root)
      - 若匹配成功，加入 M
3. 解决匹配冲突（见 7.2.3）
4. 返回 M

匹配复杂度优化：

1. 特征索引：预计算节点特征，使用哈希表加速查找
2. 剪枝策略：基于度数、类型等属性提前剪枝
3. 增量匹配：利用已匹配结果加速后续匹配

时间复杂度：$O(|V|^{|P|})$ 最坏情况，实际中通过剪枝可达到 $O(|V| \cdot |P|)$。

7.2.3 冲突检测与解决

多个融合模式可能产生冲突，需要设计冲突解决机制：

冲突类型：

1. 重叠冲突：两个模式共享节点
2. 资源冲突：融合后超出硬件资源限制
3. 依赖冲突：破坏原有数据依赖关系

冲突解决策略：

采用基于收益的贪心算法： $$ \text{score}(m) = \frac{\text{benefit}(m)}{\text{cost}(m)} \cdot \text{priority}(m) $$ 其中：

• $\text{benefit}(m)$ = 内存节省 + 延迟减少
• $\text{cost}(m)$ = 寄存器使用 + 编译复杂度
• $\text{priority}(m)$ = 路径关键度

7.2.4 动态模式学习

除了预定义模式，可以通过机器学习发现新的融合模式：

特征提取：

对于子图 $S = (V_S, E_S)$，提取特征向量： $$ \vec{f}_S = [\text{op_types}, \text{tensor_shapes}, \text{compute_intensity}, \text{memory_pattern}] $$ 融合收益预测：

训练回归模型 $g: \vec{f}_S \rightarrow \mathbb{R}$，预测融合收益： $$ \hat{y} = g(\vec{f}_S; \theta) = \sigma(W_2 \cdot \text{ReLU}(W_1 \cdot \vec{f}_S + b_1) + b_2) $$

7.3 内存带宽优化

内存带宽是现代 AI 加速器的主要瓶颈。算子融合通过减少内存访问次数，直接优化带宽利用率。

7.3.1 带宽瓶颈分析

Roofline 模型：

计算强度（Operational Intensity）定义为： $$ I = \frac{\text{FLOPs}}{\text{Bytes Accessed}} $$ 性能上界： $$ P = \min(P_{peak}, I \times B) $$ 其中 $P_{peak}$ 是峰值计算性能，$B$ 是内存带宽。

Roofline 模型示意图：
    性能
    ↑
    │     ╱─────── 计算受限区域 (P_peak)
    │    ╱ 
    │   ╱  带宽受限区域
    │  ╱   (斜率 = B)
    │ ╱
    └────────────────→ 计算强度
         I_c

临界点 $I_c = \frac{P_{peak}}{B}$ 决定了算子是计算受限还是带宽受限。

7.3.2 融合对带宽的影响

单算子带宽需求：

对于算子 $\phi$，带宽需求为： $$ BW(\phi) = \frac{M_{in}(\phi) + M_{out}(\phi)}{T_{exec}(\phi)} $$ 融合后的带宽优化：

融合 $n$ 个算子后： $$ BW_{fused} = \frac{M_{in}(\phi_1) + M_{out}(\phi_n) + \sum_{i=1}^{n-1} M_{intermediate}^{reg}(\phi_i)}{T_{exec}^{fused}} $$ 其中 $M_{intermediate}^{reg}$ 表示通过寄存器传递的中间数据（不占用全局内存带宽）。

带宽节省率： $$ \eta_{BW} = 1 - \frac{BW_{fused}}{\sum_{i=1}^n BW(\phi_i)} = 1 - \frac{M_{boundary}}{M_{total}} $$

7.3.3 缓存局部性优化

融合可以改善缓存局部性，减少缓存缺失：

时间局部性增强：

融合减少了数据重用的时间间隔： $$ \Delta t_{reuse} = t_{consume} - t_{produce} $$ 融合后 $\Delta t_{reuse} \approx 0$，数据在缓存中的概率提高。

空间局部性优化：

通过循环变换优化访问模式：

循环嵌套优化（伪代码）：
未融合：
  for i in [0, N):
    tmp[i] = φ₁(x[i])
  for i in [0, N):
    y[i] = φ₂(tmp[i])

融合后：
  for i in [0, N) step TILE_SIZE:
    for j in [i, min(i+TILE_SIZE, N)):
      y[j] = φ₂(φ₁(x[j]))  // 寄存器级数据传递

7.3.4 预取策略协同

算子融合需要与硬件预取机制协同：

预取距离计算： $$ d_{prefetch} = \lceil \frac{L_{mem}}{T_{compute}} \rceil $$ 其中 $L_{mem}$ 是内存延迟，$T_{compute}$ 是计算时间。

融合对预取的影响：

1. 融合增加了计算密度，可能需要调整预取距离
2. 减少了不规则访问，提高预取命中率
3. 需要考虑融合后的工作集大小对缓存的压力

7.4 融合的收益分析模型

准确的收益分析模型是制定融合决策的基础。本节建立综合考虑性能、功耗和资源的分析框架。

7.4.1 性能预测模型

静态分析模型：

融合后的执行时间预测： $$ T_{fused} = \max(T_{compute}, T_{memory}) + T_{overhead} $$ 其中：

• $T_{compute} = \frac{\sum_{i} \text{FLOPs}(\phi_i)}{P_{throughput}}$
• $T_{memory} = \frac{M_{in} + M_{out}}{B_{effective}}$
• $T_{overhead}$ = 核函数启动 + 同步开销

动态校正因子：

实际执行时间通常需要校正： $$ T_{actual} = T_{fused} \times (1 + \alpha_{cache} + \alpha_{conflict} + \alpha_{divergence}) $$ 校正因子通过 profiling 获得。

7.4.2 成本-收益分析

多目标优化框架：

融合决策需要平衡多个目标： $$ \text{Objective} = w_1 \cdot \Delta T + w_2 \cdot \Delta M + w_3 \cdot \Delta E - w_4 \cdot C_{compile} $$ 其中：

• $\Delta T$ = 延迟改善
• $\Delta M$ = 内存节省
• $\Delta E$ = 能耗降低
• $C_{compile}$ = 编译复杂度增加

帕累托最优解：

在多目标空间中寻找帕累托前沿： $$ \text{Pareto}(\mathcal{F}) = \{f \in \mathcal{F} | \nexists g \in \mathcal{F}, g \succ f\} $$

7.4.3 动态决策机制

运行时可能需要动态调整融合策略：

自适应阈值：

根据运行时统计信息调整融合阈值： $$ \theta_{t+1} = \theta_t + \eta \cdot (R_{observed} - R_{predicted}) $$ 其中 $R$ 表示收益，$\eta$ 是学习率。

在线学习框架：

使用强化学习优化融合决策：

• 状态空间：$s = [\text{graph_features}, \text{hardware_state}]$
• 动作空间：$a \in \{fuse, not_fuse\}$
• 奖励函数：$r = -T_{actual} - \lambda \cdot M_{used}$

7.4.4 案例分析：Transformer 模型

以 BERT 模型为例分析融合收益：

注意力块融合：

原始计算流程：
Q = Linear(X) → K = Linear(X) → V = Linear(X)
→ Scores = Q @ K^T → Attention = Softmax(Scores) @ V

融合方案：
[QKV生成融合] → [注意力计算融合]

性能分析：

• 内存访问减少：$3 \times d_{model} \times seq_len \times batch$ 字节
• 核函数调用：从 6 次减少到 2 次
• 实测加速比：1.3-1.5x（取决于序列长度）

具身智能场景优化：

在机器人控制中，融合可以显著降低感知-决策延迟：

感知管线融合：
[图像预处理] → [特征提取] → [目标检测] → [深度估计]
    ↓
[融合后的端到端感知模块]

延迟改善：50ms → 30ms（40% 降低）

本章小结

算子融合是 AI 编译器优化的核心技术，通过减少内存访问和核函数调用开销，显著提升推理性能。本章的关键要点包括：

1. 融合策略分类：垂直融合优化生产者-消费者链，水平融合并行执行独立算子，混合策略结合两者优势

2. 关键公式回顾： - 垂直融合加速比：$S = n$（n 为融合算子数） - 计算强度：$I = \frac{\text{FLOPs}}{\text{Bytes Accessed}}$ - 带宽节省率：$\eta_{BW} = 1 - \frac{M_{boundary}}{M_{total}}$ - 融合决策目标：$\min_{\mathcal{P}} \sum_{P \in \mathcal{P}} C(P) + \sum_{(P_i, P_j) \in E'} T(P_i, P_j)$

3. 模式匹配机制：通过形式化的模式定义语言和高效的子图匹配算法，自动识别可融合的算子组合

4. 内存带宽优化：融合直接减少全局内存访问，提高缓存局部性，是解决带宽瓶颈的有效手段

5. 收益分析框架：综合考虑性能、内存、功耗等多个维度，通过成本模型指导融合决策

6. 场景化优化：在自动驾驶和具身智能场景中，需要特别考虑实时性约束和确定性延迟要求

算子融合的成功实施需要深入理解硬件特性、准确的性能建模和灵活的策略调整。随着模型规模和硬件复杂度的增加，自动化和智能化的融合策略将变得越来越重要。

练习题

基础题

练习 7.1：给定算子链 Conv → BatchNorm → ReLU → MaxPool，其中 Conv 输出张量大小为 $256 \times 112 \times 112 \times 32$（NHWC 格式），数据类型为 FP16。假设内存带宽为 900 GB/s，计算：

1. 未融合时的总内存传输量
2. 完全融合后的内存传输量
3. 理论带宽节省率

练习 7.2：使用 Roofline 模型分析一个矩阵乘法算子 $C = A \times B$，其中 $A \in \mathbb{R}^{1024 \times 768}$，$B \in \mathbb{R}^{768 \times 512}$。硬件峰值性能为 312 TFLOPS（FP16），内存带宽为 2 TB/s。判断该算子是计算受限还是带宽受限？

Hint：计算 FLOPs 和内存访问量，求出计算强度。

练习 7.3：设计一个简单的融合规则，用形式化语言描述 "连续的逐元素操作可以融合" 这一模式，并给出三个符合该模式的具体算子链示例。

挑战题

练习 7.4：考虑一个简化的注意力机制计算： $$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$ 其中 $Q, K, V \in \mathbb{R}^{B \times L \times D}$，$B=32$（批大小），$L=512$（序列长度），$D=768$（隐藏维度）。设计一个两阶段融合方案，并分析每个阶段的内存访问模式和潜在的优化机会。

Hint：考虑哪些操作共享输入，哪些产生大的中间结果。

输入 X → [Q,K,V 并行生成] → [QK^T 计算 + scale] → [Softmax]

[Softmax 输出] + V → [矩阵乘法] → 输出

练习 7.5：在 NUMA 系统上进行算子融合时，需要考虑跨 NUMA 节点的数据访问开销。假设有两个 NUMA 节点，本地内存访问延迟为 100ns，远程访问延迟为 150ns。设计一个 NUMA 感知的融合成本模型。

练习 7.6：动态 shape 场景下，融合决策可能需要在运行时进行。设计一个基于历史执行信息的自适应融合策略，包括：

1. 需要收集哪些运行时信息
2. 如何根据这些信息调整融合决策
3. 如何处理 shape 变化导致的重编译开销

ProfileData = {
  shape_histogram: Map<Shape, Count>,
  fusion_perf: Map<(FusionPattern, Shape), Latency>,
  memory_pressure: TimeSeries<MemUsage>,
  recompile_cost: Map<Shape, CompileTime>
}

if predict_future_shape(history) not in cache:
  trigger_async_compilation()

练习 7.7：在自动驾驶场景中，某些融合可能会增加延迟抖动（jitter），影响实时性。设计一个考虑延迟分布（而非仅平均延迟）的融合决策框架，确保 99 分位延迟满足要求。

can_fuse = (T99_fused < T_max) AND 
           (jitter_fused < jitter_baseline * 1.2)

实时统计延迟分布
if percentile_99 > threshold:
  rollback to conservative fusion
  adjust fusion parameters

练习 7.8（开放性思考题）：未来的 AI 编译器可能需要支持投机执行，例如在大语言模型的投机解码中。讨论算子融合如何与投机执行机制协同工作，需要解决哪些新的挑战？

Layer 1: 确定性融合（always beneficial）
Layer 2: 投机友好融合（low rollback cost）
Layer 3: 条件融合（based on speculation success）

常见陷阱与错误

1. 过度融合：并非所有算子都适合融合，过度融合可能导致： - 寄存器溢出，性能反而下降 - 编译时间exponential增长 - 调试困难，错误定位复杂

2. 忽视硬件约束： - 不同硬件的融合收益差异很大 - GPU 的 shared memory 限制 - TPU 的 systolic array 对融合模式的要求 - 移动端的功耗约束

3. 静态分析的局限： - 运行时行为可能与静态预测不符 - 缓存效应难以准确建模 - 需要 profiling 数据校正

4. 依赖分析错误： - 忽视隐式依赖（如随机数生成器状态） - 错误处理 inplace 操作 - 别名分析不准确导致错误融合

5. 数值稳定性问题： - 融合可能改变计算顺序 - 浮点运算的非结合性 - 需要保证数值一致性

最佳实践检查清单

设计阶段

• [ ] 明确定义融合规则和优先级
• [ ] 建立准确的成本模型
• [ ] 考虑目标硬件特性
• [ ] 设计冲突解决机制
• [ ] 规划动态 shape 支持

实现阶段

• [ ] 实现高效的模式匹配算法
• [ ] 正确处理内存分配和生命周期
• [ ] 保证融合的正确性（依赖、数值）
• [ ] 添加融合的开关和配置选项
• [ ] 实现性能 profiling 接口

优化阶段

• [ ] 收集真实负载的 profiling 数据
• [ ] 基于数据调整融合策略
• [ ] 优化编译时间与运行时间的平衡
• [ ] 验证不同 batch size 下的表现
• [ ] 测试边界条件和异常情况

部署阶段

• [ ] 提供融合效果的可视化工具
• [ ] 记录融合决策日志便于调试
• [ ] 监控生产环境的融合效果
• [ ] 建立融合策略的 A/B 测试机制
• [ ] 准备回滚方案