3. 常见融合模式（二）：Element-wise算子融合、Add+ReLU融合、Mul+Add融合

好，咱们接着聊算子融合。上一节我们讲了卷积+BN+ReLU这种大块头的融合，今天来点更细碎的——Element-wise（逐元素）算子的融合。你别看它单个算子计算量不大，但在实际模型中，这类算子数量极多，尤其是各种 shortcut 连接、残差结构里，动不动就几十个 Add、Mul 堆在一起。如果不做融合，访存开销能把推理速度拖垮。

我个人习惯把 Element-wise 融合称为“零成本优化”。为什么？因为这类融合几乎不改变计算量，纯粹是减少访存和 kernel launch 的开销。说白了，就是让数据在寄存器里多待一会儿，别来回往显存里倒腾。

3.1 Element-wise 算子的特点与融合动机

Element-wise 算子，顾名思义，就是输出张量的每个元素只依赖于输入张量对应位置的元素。典型的包括：

• Add：逐元素相加
• Mul：逐元素相乘
• ReLU / Sigmoid / Tanh：激活函数
• Sub / Div：减法、除法
• Clamp / Clip：截断操作

这类算子的计算密度极低——计算量是 O(n)，访存量也是 O(n)，属于典型的 memory-bound 操作。你想想看，每个元素就读一次、写一次，计算就一个加法或乘法，计算单元大部分时间都在等数据。

我在项目中遇到过这样一个场景：一个 ResNet-50 的推理图，光 Add 节点就有 16 个，每个 Add 前后还跟着 ReLU。如果不做融合，这 16 个 Add+ReLU 就要 launch 32 个 kernel。融合后，只需要 16 个 kernel，甚至可以把连续的 Add+ReLU+Add 合并成一个。实测下来，端到端推理速度提升了 12%。

3.2 Add+ReLU 融合：最经典的组合

Add+ReLU 融合，可以说是 Element-wise 融合里最基础、最常用的模式。几乎每个残差网络的 shortcut 路径上都有它。

融合的原理很简单：把 Add 的计算结果直接喂给 ReLU，中间不写回显存。用伪代码表示就是：

// 未融合：两个 kernel
float* add_out = allocate_buffer(N);
for (int i = 0; i < N; i++) {
    add_out[i] = a[i] + b[i];
}
for (int i = 0; i < N; i++) {
    c[i] = max(0, add_out[i]);
}

// 融合后：一个 kernel
for (int i = 0; i < N; i++) {
    float tmp = a[i] + b[i];
    c[i] = tmp > 0 ? tmp : 0;
}

你看，融合后的代码少了中间数组 add_out 的分配和写回。在 NPU 上，这意味着少了一次全局内存的写入和一次读取。对于大张量（比如 224x224x64），这能省下几十 MB 的访存带宽。

3.3 Mul+Add 融合：FMA 指令的天然搭档

Mul+Add 融合，其实就是乘加操作。在数学上就是 c = a * b + d。这个模式在深度学习里太常见了：

• BatchNorm 的推理阶段：y = (x - mean) / sqrt(var + eps) * gamma + beta，本质上就是 Mul+Add
• LayerNorm / RMSNorm：同样包含乘加操作
• 权重缩放 + 偏置：量化推理中经常出现

为什么 Mul+Add 融合特别重要？因为几乎所有现代处理器（CPU、GPU、NPU）都原生支持 FMA（Fused Multiply-Add）指令。一条指令完成乘法和加法，延迟和单条乘法差不多。

// 未融合：两个 kernel
float* mul_out = allocate_buffer(N);
for (int i = 0; i < N; i++) {
    mul_out[i] = a[i] * b[i];
}
for (int i = 0; i < N; i++) {
    c[i] = mul_out[i] + d[i];
}

// 融合后：一个 kernel，使用 FMA
for (int i = 0; i < N; i++) {
    c[i] = fma(a[i], b[i], d[i]);  // 一条指令完成 a*b + d
}

我曾经在一个量化推理项目里踩过这个坑。当时把 BatchNorm 的 Mul+Add 融合后，精度掉了 0.3%。排查了半天，发现是 FMA 的舍入方式和原始实现不一致。后来改用双精度中间变量才解决。嗯，这里要注意，如果你的模型对精度极其敏感（比如医疗影像），最好在融合后做一次精度验证。

3.4 更复杂的 Element-wise Chain 融合

实际模型中，Element-wise 操作往往不是孤立的。它们会形成一条链，比如：

// 一个典型的 chain：Add -> ReLU -> Mul -> Add
t1 = a + b
t2 = relu(t1)
t3 = t2 * c
t4 = t3 + d

这种 chain 可以一口气融合成一个 kernel：

for (int i = 0; i < N; i++) {
    float tmp = a[i] + b[i];
    tmp = tmp > 0 ? tmp : 0;  // ReLU
    tmp = tmp * c[i];
    out[i] = tmp + d[i];
}

融合的关键在于：所有操作都是逐元素的，且数据依赖是线性的。只要满足这两个条件，理论上可以无限融合下去。当然，实际中要考虑寄存器压力——如果 chain 太长，中间变量太多，寄存器放不下，编译器会自动 spill 到栈上，反而变慢。

3.5 融合策略的图视角

从计算图的角度看，Element-wise 融合其实是在做子图匹配。编译器会遍历计算图，找到连续的 Element-wise 节点，然后合并成一个。

下面这张图展示了典型的融合流程：

编译器在做这种融合时，通常会遵循以下规则：

1. 拓扑顺序：从输入到输出，按拓扑序扫描节点
2. 类型匹配：只融合 Element-wise 类型的算子
3. 数据依赖：确保融合后的节点只有一个输出（或者输出被后续 Element-wise 节点独占）
4. 资源预算：检查寄存器数量，避免溢出

3.6 实际部署中的注意事项

最后，聊几个我在实际部署中踩过的坑：

问题	现象	解决方案
融合后精度下降	输出与原始模型不一致	检查 FMA 舍入模式；使用 FP32 中间变量
寄存器溢出	融合后性能反而下降	限制 chain 长度；使用局部内存（shared memory）
分支发散	ReLU 等条件操作导致 warp 内线程分支	使用 predicated 指令代替分支；确保数据分布均匀
融合后无法向量化	编译器无法生成 SIMD 指令	手动展开循环；使用内置向量类型（如 float4）

我曾经在一个项目里，把 Add+ReLU+Mul 融合后，性能反而慢了 5%。排查后发现是寄存器压力太大，编译器把中间变量 spill 到了 local memory。后来我把 chain 拆成两个（Add+ReLU 和 Mul），性能就正常了。所以，融合不是越多越好，要适可而止。

好了，Element-wise 融合就聊到这儿。记住一句话：能融合就融合，但别硬融。下一节我们讲更复杂的融合模式——涉及到 reshape 和 transpose 的那种，到时候再细聊。