算子融合：Elementwise融合

可视化演示

融合原理

老周在车间里反复看那三条流水线：去壳车间把虾仁装盒，冷藏车花二十分钟运到下个车间；第二个车间拆盒、去虾线、再装盒；第三个车间又拆盒、腌制、再装盒。每一次“装盒—运输—拆盒”都是对冷库的一次全局访问。老周想，如果三个车间变成一张台上连续完成三道工序，那么只需要开工前取一次原料、完工后存一次成品。这个“台面改造”就是算子融合的直觉起点。

在深度学习中，逐元素算子（Elementwise）指的是对张量中每个元素独立施加相同函数的运算，例如 add、mul、ReLU、tanh，它们的输入输出形状一致，且输出第 i 个元素只依赖输入第 i 个元素。这类算子天然适合融合。若不融合，每个算子都会启动一个独立 kernel，每个 kernel 都从全局内存读一次数据、计算、再写回一次。三个算子连用，数组就经历三次读取与三次写入；而全局内存带宽是 GPU 最宝贵的资源之一。

设数组长度为 $L$，算子个数为 $N$，每次全局访问的耗时系数为 $C_m$，每次 kernel 启动开销为 $C_l$，则未融合与融合的总时间分别是：

老周的改造方案在纸上算完，立刻让车间节省了两趟冷藏车的往返。这个从“三趟冷库”到“一趟冷库”的根本性变化，正是 Elementwise 融合最核心的动机：尽量避免中间结果落到全局内存。

匹配规则

不是所有工序都能放心地合并到同一张工作台。老周意识到：如果第二道工序“去虾线”需要把虾先冷藏十五分钟让肉质收紧，那它就不能与第一道工序无缝衔接。同样，编译器在决定两个逐元素算子能否融合时，也要做严格的“合法性检查”。判断一个算子是不是 Elementwise，核心看它是否满足“逐元素访问模式”：输出每个位置的值只依赖输入对应位置的值，不跨元素、不涉及相邻索引。卷积、池化这类涉及邻域窗口的算子就不满足条件。

编译器识别逐元素算子后，还会检查算子之间的数据依赖与形状是否匹配。若前一个算子输出形状与后一个算子输入形状完全一致，索引映射是一一对应的，就可以安全融合。若中间有出度大于一的“分叉”——比如一个结果同时被 exp 和 log 使用——则要视具体编译器策略决定是否融合，因为这会增加寄存器压力。融合条件可写成：

回到车间，老周在墙上贴了一张“可同台工序清单”：只有那些不依赖冷库暂存、不需要等待时间的工序才能合并。AI 编译器里的匹配规则，就是这张“可同台清单”的工程化表达。

广播合并

老周遇到了新情况：第三车间的“腌制入味”工序，需要按“筐”投放调料。每 100 只虾一个筐，但有的批次一筐只有 50 只。如何让两个数量不同的批次在同一条工作台上同时加工？这就像两个形状不同的张量要做逐元素运算：一个形状是 [100]，另一个是 [50]。直接逐元素对齐会失败，因为索引对应不上。

广播（Broadcast）机制解决了这个问题：从最右侧的维度开始依次对齐，凡是长度为 1 的维度，会像“复制”一样被拉伸到与另一个张量的对应维度相同。比如形状 [100, 1] 与 [1, 50] 做逐元素运算，前者在第二维复制 50 份，后者在第一维复制 100 份，最终得到 [100, 50]。融合时，编译器会把广播逻辑“折叠”进同一个循环体，通过控制索引映射避免真正在内存中复制数据。输出形状的推导规则是每个维度取参与张量的较大值：

老周最终定制了一种“可伸缩调料筐”：50 只的批次自动补 50 个空位，工作时长和 100 只的批次完全一致。广播合并就像这种可伸缩筐，让不同形状的张量也能在融合后的同一个循环体里顺利加工。

内核生成

老周把去壳、去虾线、腌制三道工序写在了一张工序卡上，老师傅拿到卡片后，从冷库取一次原料，在台前连续完成全部三道工序，最后把成品送进冷库。在 AI 编译器中，“这张工序卡”就是融合后生成的单个 kernel。编译器不再为每个算子分别生成独立的内核函数，而是把它们解析为同一个循环体中的多次连续操作。

这段循环合并代码的伪代码逻辑是：遍历输入数组的每个元素 $i$，依次执行第一个算子、第二个算子、第三个算子，把最终结果写入输出数组。中间值不离开寄存器，直接作为下一个算子的输入。生成的内核可以抽象成如下形式：

工序卡让老师傅的双手从未离开过台面，融合 kernel 则让数据从未离开过寄存器。这就是内核生成的本质：把多个 kernel 的“启动—读—写”循环，压缩成一个 kernel 的“启动—连续计算—写”流程。

内存优化

老周算过一笔账：改造前，每道工序之间都要把半成品送进冷库，再取出来。一次冷库往返等于一次写入加一次读取。三道工序就有两次中途往返，虽然冷库本身容量很大，但抽开冷藏门、存取箱子的时间成本极高。Elementwise 融合的内存优化哲学如出一辙：尽可能减少“全局内存往返”的次数。

在未融合的执行流程中，如果数组长度为 $L$、算子个数为 $N$，那么每次算子的输出都要写回全局内存，下一次算子再读回，共发生 $N-1$ 次“写—读”中转。每次中转涉及 $2L$ 次内存访问。融合之后，这部分开销被完全消除：

融合还顺带提升了缓存局部性：单个 kernel 顺序读入一块数据后，在寄存器内完成所有运算，不会像未融合那样反复访问同一块全局内存区域。冷库的大门开关次数少了，整条生产线自然快得多。

向量化执行

老师傅手上有一把“四联刀”：一次下去能同时处理四只虾的去壳。这就是 SIMD（单指令多数据）的通俗版——一条机器指令同时操作多个数据元素。融合之后的 kernel 循环体更加规整，编译器可以轻松地把循环向量化，用一条指令处理多个连续元素。

向量化的关键参数是向量宽度：一条 SIMD 指令同时处理的元素个数。它由硬件 SIMD 寄存器宽度除以单个元素的数据位宽决定。以常见的 128 位 SIMD 寄存器与 32 位浮点数为例，向量宽度就是 4。融合后的循环可以直接改写为“每 4 个元素一组”的批处理：

但向量化还有前提：数据地址必须对齐。如果数组起点没有按 16 字节对齐，SIMD 加载指令就无法一次取齐数据。因此融合 kernel 在生成时会尽量申请对齐的内存块，并为循环首尾添加“标量头尾处理”。如果说融合是把流水线合并成一条，那向量化就是给这条线配上满载的集装箱卡车。

工程实践

老周的流水线改造最终动工了。工程实践层面，Elementwise 融合已经是主流 AI 编译器与推理引擎的标配能力。TVM、XLA 以及 PyTorch 的 Inductor 后端，都有专门的融合 pass 负责识别可融合的逐元素算子序列并生成单一 kernel。深度学习模型的很多基础结构都会受益：残差连接中的 Add+ReLU、归一化中的 Mul+Add、激活前后的 Linear+ReLU 等等。

工程落地的效果评估需要回到时间模型。设单次内存访存时间为 $M$、单算子的计算时间为 $C$，未融合总时间为 $N(M+C)$，融合后为 $M+NC$。这是因为融合只保留了一次整体的读入与写回，但 $N$ 个算子依然要在每个元素上各算一次：

实际项目中，开发者往往通过 profiling 工具观察融合前后的内存吞吐与 kernel 启动次数来验证收益。老周的车间改造也并非万能——若原料本来就极少、冷库距离极近，融合的边际收益就很小。Elementwise 融合的适用场景应具备“长数组、多次逐元素运算、访存受限”这三个特征，这为工程取舍提供了简洁的判据。

AI 应用场景

Python 演示

import time  class Kernel:    """A tiny simulation of a GPU kernel on a production line."""     def __init__(self, launch_cost=50.0, mem_cost=2.0):        self.launch_cost = launch_cost        self.mem_cost = mem_cost     def run(self, reads, writes):        """Total time = global memory access time + one launch."""        return (reads + writes) * self.mem_cost + self.launch_cost  def main():    L = 1_000_000    kernel = Kernel()     # Unfused: 3 kernels, each reads L and writes L once.    unfused_time = 0.0    for _ in range(3):        unfused_time += kernel.run(L, L)     # Fused: 1 kernel still reads L and writes L once.    fused_time = kernel.run(L, L)     print("unfused time: {:.1f} us".format(unfused_time))    print("fused time:   {:.1f} us".format(fused_time))    print("speedup:      {:.2f}x".format(unfused_time / fused_time))  if __name__ == "__main__":    main()

1. class Kernel:定义一个模拟 GPU 内核的类，用 launch_cost 表示启动开销，mem_cost 表示每次全局内存访问的耗时。
2. def run(self, reads, writes):计算一个内核执行的总时间：访问次数乘 mem_cost，再加上一次启动开销。
3. unfused_time += kernel.run(L, L)未融合方案中，3 个算子每个都要完整读一遍、写一遍数组，所以循环 3 次累加时间。
4. fused_time = kernel.run(L, L)融合方案只需一个内核，读一次、写一次，启动开销也只付一次。
5. print("speedup: {:.2f}x".format(...))输出加速比。当启动开销占比越高，融合的收益越明显。

常见误区

面试题