算子融合：垂直融合

可视化演示

展示垂直融合三维度：融合原理（未融合3个kernel间2次VRAM中间张量读写→融合后1个kernel寄存器传递零中间访问）、Conv-BN-ReLU经典模式（三算子串行依赖模式匹配融合为单个FusedConvBNReLU kernel）、融合收益（VRAM读写6次→2次计算强度FLOP/Byte从1提升至3倍带宽瓶颈消除），三模式切换理解编译器核心优化技术。

融合动机

前文说过，后厨里真正的炒菜时间只占一小段，绝大多数时间花在装盒、入库、取盒、解冻上。假设传送带（内存带宽）每秒只能送 100 份食材，而灶台每秒能炒 1000 份，那么整条流水线就被传送带卡死了。GPU 也是同样的道理：显存带宽远低于芯片的浮点计算速度。对加法、激活这类“每份食材只需轻轻一炒”的算子，计算本身眨眼就完成，反而是一直在等数据进出显存。

一个 kernel 在 GPU 上执行至少访问两次内存：读输入、写输出。当 add 和 relu 两个算子组成链条时，add 写出的中间张量先落回显存，relu 再把它读出来——这个中间张量就白白贡献了两次全量读写。更别提每次 kernel 启动还有启动开销：好比每道工序都要重新点火、重新叫帮工，单次不过几微秒，算子一多就积少成多。

对访存密集（memory-bound）的算子来说，总耗时近似等于访存时间：

所以后厨改革的第一动机不是换一口更快的锅，而是少跑腿。垂直融合解决的正是“跑腿”问题：让数据待在离计算最近的地方，别动不动就回冷库。

工作原理

改革方案其实很朴素：把切菜台和炒锅搬到一起，甚至让同一位师傅左手切、右手炒。切好的土豆丝不装盒，顺砧板直接滑进锅；炒完出锅，同一只手顺势完成下一步。从“三道工序、三次入库”变成“一条龙、零入库”。

编译器做垂直融合，就是把计算图上串行的多个 kernel 合并成一个 kernel。合并后的 kernel 在同一个循环体里依次执行两个算子，前一个算子的输出保存在寄存器变量中，后一个算子立刻消费——这正是数据局部性（data locality）的胜利，尤其是时间局部性：刚算出来的数据马上用掉，不让它“冷却”到显存里。

写成伪代码，区别一目了然：

注意融合后循环体里的 t 就是一个寄存器变量，relu 立刻消费它，完全没有经过显存。编译器还会把整个循环按 tile 分块，让每个数据块在片上处理完再换下一块，进一步守住局部性。融合的本质，就是把“仓库中转”改成“手递手”。

# 融合前：两次循环，中间张量两次进出显存for i in range(N): mid[i] = x[i] + y[i]for i in range(N): out[i] = relu(mid[i]) # 融合后：一次循环，中间值留在寄存器for i in range(N):    t = x[i] + y[i]    out[i] = relu(t)

合法条件

并不是所有工序都能并台。比如第二道工序其实在等酱料而不是土豆丝，你把两台并在一起也没用；又比如第一道切丝、第二道要块，规格对不上。还有的工序会“夹带私货”——炸鸡时顺便记一笔账、改动别的食材，那就不能随便和别人合并。

编译器判断能否融合，第一步是做数据依赖分析：确认算子 B 的输入恰好是算子 A 的输出，且 B 不依赖更早的其他算子；如果 A 的输出元素是由多个输入归约而来的（比如求和），B 还必须能逐元素对齐。第二步是检查张量形状与内存布局是否一致，或者能否由广播规则静态推导。第三步是确认算子没有副作用——不修改全局状态、不依赖随机种子这类不可控输入。

综合起来，合法融合需要满足四条：单一的生产者—消费者关系；形状与布局一致或可推导；算子语义确定且无副作用；融合后的计算结果与原图完全等价。编译器通常会先做拓扑排序得到依赖图，再用模式匹配挑选出可融合组（fusion group）。在厨房里，这就意味着要先看菜谱、对刀工规格、查额外动作，确认无误才把两台并到一起——合法性检查不是流程负担，而是防止“出锅的菜不是客人要的菜”。

寄存器分配

并台之后新的烦恼来了：切好的菜要先放在操作台上等油热，可操作台只有那么大——GPU 里一个线程的寄存器数量是固定的，通常只有几十个。同一时刻要用的半成品太多，台面放不下，只能把多出来的菜先放回门口货架，用的时候再取。这来回一趟虽比冷库近，但也要时间。厨师管这叫“临时挪货”，编译器管这叫溢出（spill）。

融合后的 kernel 在一个循环体里同时执行多个算子，多个中间值在同一时刻存活（live），寄存器压力（register pressure）随之上升。编译器要做寄存器分配，先通过生命周期分析确定每个中间值的存活区间，再尽量让存活区间不重叠的中间值复用同一个寄存器。比如 t = x[i] + y[i] 算出后，t 只在紧接着的 relu 里用到，存活区间极短，很快就能被下一个循环迭代的变量复用。

当同时存活的中间值超过可用寄存器数量时，编译器只能把部分变量溢出到栈上（本地内存）。这正是“融合链越长越好”的危险所在：链条越长，同时存活的中间值越多，溢出风险越大。因此，实际编译器在融合时往往会评估额外寄存器需求，必要时放弃融合或改用共享内存中转。好的后厨管理不是无限并台，而是在“少跑腿”和“台面够用”之间找到平衡。

性能效果

改革效果可以用后厨的一笔账算清：从前每道工序都要“从冷库取料、完工送回冷库”各一趟；并台之后，全流程只取一次料、上一次菜。省下的跑腿趟数，全是纯赚的时间。

设一条链上有 k 个逐元素算子，每个张量有 N 个元素，每个元素占 s 字节。融合前每个算子都要从显存读输入、把输出写回显存，一共 2k 次全量访问；融合后全局只读一次输入、写一次输出，只有 2 次访问。于是搬运总量从 2kNs 降到了 2Ns：

账算完，老板笑了：同一间厨房，食材搬运减半甚至减到四分之一，出菜自然更快。实测量也印证了这一点——对访存密集的网络（如残差网络、检测网络的前几层）做垂直融合，端到端提速 1.3 到 3 倍很常见。反过来，像大矩阵乘法这类计算密集层，计算时间已经占满，省掉的那点访存占比很小，融合收益自然有限。用行话说，垂直融合把计算访存比（每搬运一字节能做的浮点运算次数）抬高了 k 倍，把算子从“等数据”的带宽受限区推向“算得过来”的计算受限区。

编译器实现

不同连锁餐厅的“厨房管理软件”对并台的支持很不一样。TVM 像一家讲究手工的精品餐厅：总管可以手工指定并台方案，也可以让自动搜索算法在上千种并台方式里挑最优；XLA 更像标准化的中央厨房：预先把常见菜谱做成模板，流水线一匹配就直接套用，稳定高效。

具体来说，TVM 提供算子调度 API，通过 fuse 这类调度原语把多个 compute 合并到同一个循环里；图优化阶段还有 FuseOps 这样的 pass 做自动融合，再结合自动调优（如 Ansor）同时搜索融合决策与分块大小。XLA 则把计算图表示为 HLO，对逐元素、广播等常见模式做模板化融合（fusion instruction），并且支持多输出融合，把 fan-out 出去的多个消费者一起包进同一个融合组。

两者的差异很有意思：TVM 偏向探索式，对任意算子的表达能力更自由，但搜索空间大、调优成本高；XLA 偏向规则匹配，融合模式稳定、易验证，但灵活性略逊。它们共通的地方在于：都先在图上做合法性分析，再决定融合方案，最后交给后端处理寄存器分配和代码生成。就像同一道土豆丝，精品店手工改刀、中央厨房模板化批量切——殊途同归，目标都是让菜别回冷库。

AI 应用场景

Python 演示

import numpy as np  def run_unfused(x, y):    n = len(x)    load = 0    store = 0    mid = np.empty(n)    for i in range(n):                # kernel 1: add        load += 2                      # 读 x[i] 和 y[i]        mid[i] = x[i] + y[i]    store += n                         # 写中间张量 mid    out = np.empty(n)    for i in range(n):                 # kernel 2: relu        load += 1                      # 读 mid[i]        out[i] = mid[i] if mid[i] > 0 else 0    store += n                         # 写输出 out    return out, load, store  def run_fused(x, y):    n = len(x)    load = 0    store = 0    out = np.empty(n)    for i in range(n):                 # 单一 kernel: add + relu        load += 2                      # 只读 x[i] 和 y[i]        h = x[i] + y[i]                # 中间值留在寄存器        out[i] = h if h > 0 else 0     # 立刻消费, 不再读中间量    store += n                         # 只写输出 out    return out, load, store  def main():    rng = np.random.default_rng(0)    x = rng.standard_normal(1000)    y = rng.standard_normal(1000)     out_a, la, sa = run_unfused(x, y)    out_b, lb, sb = run_fused(x, y)     assert np.allclose(out_a, out_b)   # 两种方式结果完全一致    print('未融合: load=%d store=%d 合计=%d' % (la, sa, la + sa))    print('融合后: load=%d store=%d 合计=%d' % (lb, sb, lb + sb))    print('内存访问降低倍数: %.1fx' % ((la + sa) / (lb + sb)))  if __name__ == '__main__':    main()

1. mid = np.empty(n)未融合时需要显式创建中间张量，这就是“半成品要先装盒”。
2. h = x[i] + y[i]加法结果直接留在“寄存器”变量 h 里，不落显存。
3. out[i] = h if h > 0 else 0紧接着做 relu，中间值零往返。
4. assert np.allclose(out_a, out_b)验证融合前后数学结果完全一致，合法性第一关。
5. (la + sa) / (lb + sb)算出内存访问的降幅，理论账（2 倍）和代码账对上。

常见误区

面试题