XLA编译流程：算子融合

可视化演示

展示XLA算子融合：融合前后计算图对比（多内核HBM往返→单内核寄存器驻留）、逐元素归约生产者消费者三类融合与寄存器压力、内存读写与内核启动性能收益，交互切换核心维度。

算子融合概念

回到中央厨房。如果每切一刀都要把菜放回冰箱、下一步再拿出来，客人等到天黑也吃不上饭。神经网络里的“切菜”就是一个个算子：矩阵加法、乘法、激活函数、卷积……默认情况下，每个算子是一个独立的内核（kernel），各自从显存读入数据、计算、写回结果。下一个算子再读这个结果，又得重新搬运一趟。当算子数量成千上万时，GPU 大部分能量都耗在了“端盘子”上，而不是真正干活。

这种“盘子在灶台与冰箱间往返”的成本，在计算机里叫访存（内存访问）。一个由 N 个算子串联的计算链，分开执行时每个算子都要写一次中间结果、再被下一个算子读一次；而融合后，整个链路只读一次输入、写一次输出：

所以算子融合的本质，就是让中间结果“不落地”。总厨把切、焯、炒合成一锅出，省下的正是那一趟趟端盘子的时间。

XLA编译管线

中央厨房接单后有一套完整流程：先由服务员把客人需求写成工单（前端捕获），再由总厨把工单整理成标准菜谱（HLO 图），接着总厨逐段审视菜谱、决定哪些步骤可以合并（优化 Pass），最后把定稿的菜谱交给炉头执行（后端代码生成）。XLA 编译器就是这座厨房的大脑。

XLA 的完整管线大致是：前端把 PyTorch、TensorFlow 等框架的计算图翻译成 HLO（高级优化中间表示）；随后进入优化阶段，几十个 Pass 轮番对 HLO 图做变换，算子融合就发生在这里；优化后的 HLO 再次被翻译成 LLVM IR，最终由后端生成 GPU（PTX）或 TPU 的机器指令。

融合在图中就是一个“折叠”的操作：从一份 HLO 图里挑选 k 个节点，把它们替换成一个融合节点。

这条管线里，融合是唯一能“删边”的环节——总厨合并菜谱上的并列步骤，让整道菜的流程变短。

融合分类方式

厨房里的统筹有两种思路。第一种是横着批处理：十道菜都要用蒜末，那就一次剁一大碗，做菜时随手抓一把，不必每道菜重新剥蒜。第二种是纵着串流程：同一道菜，从洗、切到焯、炒一口气做完，中间不让半成品回冰箱。这两种思路，在编译器里分别叫横向融合和纵向融合。

• 横向融合：把同层、互不依赖的算子合并。比如多个形状相同的 elementwise 算子（加、乘、激活）并行执行，各自启动一个小 kernel 太浪费，合成一个大 kernel 后，启动开销从 N 次变 1 次。

• 纵向融合：把数据流上串联的算子合并。比如 Conv → BN → ReLU 是一条链，前一个算子的输出直接喂给后一个，合成一个核后中间结果不再访问显存。

• 输入融合：把多个上游算子的计算结果合并后一次性交付给下游消费者，常用于减少重复读取。

• 输出融合：把下游多个消费者的公共上游合并处理，同样是为了少搬几趟数据。

横向融合的特点是“并行批量”，纵向融合的特点是“串联到底”。实际模型中 Conv+BN+ReLU 是最经典的纵向融合案例：推理时 BN 的均值和方差已经固定，可以折叠成卷积的缩放系数，三段流水线瞬间变成一步到位。总厨心里那本账，就是看哪种统筹方式省得更多。

融合关键算法

中央厨房每天要面对成百上千道菜，总厨不可能把每两种食材的组合都试一遍。他会凭经验先找“相邻、且合并明显划算”的步骤下手：先并收益最大的，并完再看下一步。XLA 里的融合决策，用的正是这套思路——启发式搜索候选算子对，用代价模型评估收益，贪心迭代直到没有正收益为止。

具体的收益公式长这样：把两个算子 i 和 j 各自执行的代价相加，减去合并后一个核的代价，再归一化，得到融合带来的相对收益。

贪心策略会反复挑选 Gain 最大的相邻算子对执行融合，直到找不到正收益为止。同时，融合还要受峰值内存约束：如果融合核的中间结果需要在寄存器里同时存活，而数量超过了硬件寄存器上限，即使收益为正也会被否决。这就像总厨发现“这锅菜需要八个灶眼，但厨房只有四个”——再划算也得放弃。

性能优化特性

一锅出的好处不只是少洗几个碗。从执行的角度看，它至少带来三件好事：第一，访存次数骤降，半成品不再来回进出显存；第二，内核启动开销减少，GPU 不必为每个小算子付出调度和 launch 的固定成本；第三，融合核内部有机会做指令级并行编排，算子和算子之间的依赖变成寄存器间的直通数据，而不是等显存回包。

其中访存的收益最为关键。在访存受限（memory-bound）的场景下，算得快不快基本由“搬了多少数据”决定。假设 N 个算子的中间结果大小相近，都是 \bar{S}，带宽为 B，那么：

灶台火力（带宽）越紧张，一锅出的优势就越明显。大模型训练和推理往往就是这种“火力不足”的厨房——显存带宽供不应求，算子融合因此成为最立竿见影的优化手段。

实践与挑战

一锅出也有翻车的时候。锅太大，灶眼放不下；客人临时改分量，备菜方案全部作废。算子融合在工程实践中同样会遇到三类典型挑战。

第一是调试难度。融合核一旦生成，内部就是一个黑盒，原本可以逐个算子检查的中间结果全被藏起来了。工程师排查数值问题时，经常不得不先关闭融合再逐层对比。第二是编译耗时。启发式搜索和代价评估要遍历大量候选对，融合 pass 有时会让编译时间成倍增加，训练脚本里等编译的时间比训练还长。第三是动态形状。当输入张量的形状在运行时才能确定时（比如 NLP 里变长序列），编译器没法预编译最优的融合核，XLA 往往会回退到保守路径，放弃融合或触发重新编译。

实战中，融合的收益也很可观：ResNet 推理时把 Conv+BN+ReLU 融合，端到端时延常能下降 30% 以上；Transformer 解码阶段把 KV Cache 读取与注意力计算融合，能显著降低显存带宽压力。但这些数字都建立在“代价模型判断值得”的前提上——优秀的总厨既懂得合并，也懂得什么时候不该合。

AI 应用场景

Python 演示

import time  def busy_wait(ms):    end = time.perf_counter() + ms / 1000.0    while time.perf_counter() < end:        pass  def kernel(compute_ms, mem_mb, bandwidth_mb_per_ms):    busy_wait(compute_ms)    return compute_ms + mem_mb / bandwidth_mb_per_ms  def main():    # 四个算子的计算耗时（毫秒）    compute = [0.05, 0.08, 0.06, 0.30]    # 每个算子的中间结果大小（MB）    sizes = [1.0, 1.0, 1.0, 1.0]    bandwidth = 2.0  # MB/ms, 模拟较慢的显存带宽    n = len(compute)     sep_time = sum(kernel(compute[i], 2 * sizes[i], bandwidth) for i in range(n))    fused_time = kernel(sum(compute), sizes[0] + sizes[-1], bandwidth)     print("分开执行耗时: %.4f ms" % sep_time)    print("融合执行耗时: %.4f ms" % fused_time)    print("加速比: %.2f x" % (sep_time / fused_time))    print("访存对比: %d MB -> %d MB" % (2 * sum(sizes), sizes[0] + sizes[-1]))  if __name__ == "__main__":    main()

1. def kernel(compute_ms, mem_mb, bandwidth_mb_per_ms):一个 kernel 的耗时 = 计算时间 + 数据搬运时间
2. busy_wait(compute_ms)用空循环模拟 GPU 实际算数的时间
3. sum(kernel(compute[i], 2 * sizes[i], bandwidth) for i in range(n))分开执行时每个算子都写一次结果、下个算子再读一次，所以访存按 2 倍中间结果计算
4. kernel(sum(compute), sizes[0] + sizes[-1], bandwidth)融合核把计算量合并成一次执行，访存只剩输入与输出
5. sep_time / fused_time加速比展示访存受限场景下融合的收益

常见误区

面试题