算子融合策略

现代深度学习 (deep learning)加速器众所周知受限于带宽。尽管计算能力（FLOPs）呈指数级增长，内存带宽却相对滞后。因此，模型执行中的主要瓶颈通常不是芯片的计算速度，而是其在全局内存（HBM/DRAM）和片上寄存器之间移动张量的速度。算子融合是解决这种差异的最有效的图级优化。

通过识别计算图中可以将多个算子在单个核函数执行中完成的模式，融合减少了内存总流量。它避免了将中间结果写入全局内存，然后又立即读回以进行后续操作。这提升了核函数的算术强度，即浮点运算次数与访问字节数之比，从而使硬件能更接近其理论峰值计算性能。

垂直融合

垂直融合，也称为生产者-消费者融合，合并连续操作，其中一个节点的输出是下一个节点的输入。这是在 TVM 或 XLA 等标准编译器流程中普遍存在的优化方式。

考虑 ResNet 等架构中普遍存在的标准卷积层块：一个卷积运算，接着是偏置 (bias)相加和 ReLU 激活。没有融合时，执行流程需要三次独立的核函数启动：

1. Conv2D: 读取输入 + 权重 (weight) $\rightarrow$ 计算 $\rightarrow$ 将 Temp1 写入 HBM。
2. BiasAdd: 读取 Temp1 + 偏置 $\rightarrow$ 计算 $\rightarrow$ 将 Temp2 写入 HBM。
3. ReLU: 读取 Temp2 $\rightarrow$ 计算 $\rightarrow$ 将输出写入 HBM。

这个序列产生了显著的冗余。中间张量 Temp1 和 Temp2 是瞬时的；它们不被网络的任何其他部分需要。在融合策略中，编译器会生成一个单一的核函数。卷积输出保存在寄存器中或累积在本地缓存中（如共享内存），偏置被添加，并且 ReLU 非线性操作被应用，之后最终结果才写入全局内存。

离散执行与垂直融合之间的内存流量对比。融合核函数消除了中间往返全局内存的操作。

为实现这一点，编译器必须确保生产者节点（例如，卷积）没有其他消费者。如果 Temp1 被图的独立分支（例如，跳跃连接）使用，简单的融合意味着为第二个分支重新计算 Temp1，或者放弃融合以保存中间状态。大多数现代编译器使用成本模型来判定重新计算的成本是否低于存储中间张量的内存带宽成本。

水平融合

水平融合合并彼此独立但共享相同输入数据的算子。这种模式经常出现在 Transformers 中的多头注意力 (multi-head attention)机制 (attention mechanism)或 CNN 中的 Inception 模块中。

在 Transformer 中，单个输入嵌入 (embedding)通过三次独立的矩阵乘法（全连接层）被投影到查询 ($Q$)、键 ($K$) 和值 ($V$) 张量。如果分别执行，GPU 必须从内存中读取输入嵌入三次。

水平融合将这三个操作合并成一个单一的、更大的矩阵乘法。通过拼接 $Q, K, V$ 投影的权重 (weight)矩阵，编译器生成一个单一的核函数。该策略有两项益处：

1. 减少内存读取: 输入张量只读取一次，并用于所有三个投影。
2. 提升占用率: 较大的核函数通常比几个更小、更零散的核函数更有效地使 GPU 计算单元饱和。

融合规则和拓扑约束

并非所有图节点都能不加区分地融合。编译器必须根据内存访问模式对算子进行归类，以判定兼容性。我们通常将算子分为四类，以进行融合逻辑处理：

1. 单射（逐元素）: 一个输出元素对应一个输入元素的操作（例如，加法、ReLU、类型转换）。这些是最易于融合的。
2. 广播: 沿某个维度复制数据的操作（例如，将偏置 (bias)向量 (vector)添加到矩阵）。这些通常可以与单射操作融合。
3. 归约: 将张量映射到更小维度的操作（例如，求和、MaxPool）。将归约操作与单射算子融合是可行的（例如，BatchNorm 包含归约操作），但这需要复杂的循环调度来处理线程同步边界。
4. 不透明: 难以分析或需要全局同步的操作（例如，排序、非极大值抑制）。这些充当融合屏障。

融合算法通常遍历图（通常以后序方式）并根据这些类别贪婪地合并节点。一个普遍规则是，你可以将任意数量的单射操作融合到一个归约操作中，但没有分块优化（将在循环调度章节中讨论）时，你无法轻易将一个归约操作融合到另一个归约操作中。

算术强度对比。融合块大幅提升了计算与内存比，将工作负载从内存受限状态转向计算受限状态。

代码生成中的挑战

融合简化了图结构，但使代码生成更复杂。一个融合核函数需要更多寄存器来保存多个操作的活跃状态。如果寄存器压力超出每个线程的硬件限制，编译器通常会将数据溢出到本地内存（L1/L2 缓存）或限制活跃线程数量（占用率），这会降低性能。

此外，融合循环边界不匹配的算子需要谨慎的索引对齐 (alignment)。例如，将 $3 \times 3$ 卷积与 $2 \times 2$ 最大池化层融合是棘手的，因为它们的迭代空间不具备一对一映射关系。在这种情况下，编译器可能会采用“按需计算”调度原语，有效地将生产者循环嵌套在消费者循环内部，按需计算输入块。重新计算与内存存储之间的权衡是编译器工程中一个常见的话题。