4. 典型算子融合模式

各位同学，今天我们来聊聊算子融合的具体模式。说实话，这块内容是我在实际项目中踩坑最多、也收获最大的部分。你想想看，存算一体架构下，数据搬运的成本远高于计算本身。所以怎么把多个算子捏在一起，减少数据来回倒腾，就成了性能优化的关键。

我个人习惯把常见的融合模式分成三类：逐元素算子融合、归约算子融合、还有矩阵链融合。咱们一个一个来看。

4.1 逐元素算子融合：激活 + 归一化

先说说最简单的——逐元素算子融合。这类算子的特点是：每个输出元素只依赖于对应的输入元素，跟周围的数据没关系。典型的例子就是激活函数（ReLU、Sigmoid）和归一化操作（BatchNorm、LayerNorm）。

为什么要把它们融合？因为这两步在计算上是独立的，但数据流是连续的。如果不融合，就得把激活后的结果写回存储，再读出来做归一化。这一写一读，在存算一体架构下就是白花花的功耗和延迟。

举个例子，ReLU + BatchNorm 的融合。ReLU 就是 max(0, x)，BatchNorm 是 γ * (x - μ) / σ + β。融合后，我们可以直接计算：

// 融合前
y1 = relu(x)          // 激活
y2 = batchnorm(y1)    // 归一化

// 融合后
y = batchnorm_relu(x) // 一步到位

我在项目中遇到过一个问题：有些框架的 BatchNorm 实现里，均值和方差是统计出来的，不是固定的。这种情况下融合要小心，得把统计逻辑也考虑进去。我曾经因为没注意这个细节，导致训练时精度掉了 2 个点，排查了两天才找到原因。

4.2 归约算子融合：卷积 + 池化

接下来是归约算子融合。归约操作的特点是：多个输入元素经过某种运算得到一个输出元素。卷积和池化就是典型的归约算子。

卷积是加权求和，池化是取最大值或平均值。这两者经常在神经网络中前后出现。如果不融合，卷积的输出要写回存储，池化再读出来。你想想看，卷积的输出特征图通常很大，这一写一读的带宽开销相当可观。

融合的思路是：在卷积的计算过程中，直接对输出做池化。比如，卷积核滑动计算完一个窗口后，不急着写结果，而是等周围几个窗口都算完，直接做池化。

// 传统做法
conv_out = convolution(input, weight)  // 卷积
pool_out = max_pooling(conv_out, 2x2)  // 池化

// 融合做法
// 在卷积计算时，每算完 2x2 个输出位置，直接取最大值
for each 2x2 block in output:
    max_val = -inf
    for each position in block:
        val = convolution_at_position(input, weight, position)
        max_val = max(max_val, val)
    pool_out[block] = max_val

嗯，这里要注意：融合后，卷积的中间结果不需要全部存储，只需要保留当前窗口内的数据。这对存算一体架构的片上存储压力是很大的缓解。

4.3 矩阵链融合：全连接 + 全连接

最后说说矩阵链融合。这个模式说白了就是把多个全连接层合并成一个大的矩阵乘法。为什么能这么做？因为全连接层的本质就是矩阵乘法，多个全连接层串联，就是多个矩阵相乘。

假设有两个全连接层：第一层是 Y1 = X * W1 + b1，第二层是 Y2 = Y1 * W2 + b2。如果不融合，就得先算 Y1，存起来，再算 Y2。融合后，直接算 Y2 = X * (W1 * W2) + (b1 * W2 + b2)。

// 不融合
Y1 = matmul(X, W1) + b1
Y2 = matmul(Y1, W2) + b2

// 融合
W_combined = matmul(W1, W2)      // 权重融合
b_combined = matmul(b1, W2) + b2 // 偏置融合
Y = matmul(X, W_combined) + b_combined

这样做的好处很明显：原本需要两次矩阵乘法，现在只需要一次。而且中间结果 Y1 不用存储了。我个人习惯在模型部署阶段，把连续的几个全连接层都检查一遍，能融合的尽量融合。

不过，这里有个坑：如果全连接层之间有非线性激活函数，就不能直接融合了。因为激活函数破坏了线性关系。我遇到过有人强行融合带 ReLU 的全连接层，结果精度直接崩了。

这三种模式是算子融合的基础。实际项目中，往往需要组合使用。比如卷积+激活+池化，或者全连接+归一化+激活。嗯，这些组合模式我们后面会详细讲。

下面这张图展示了三种融合模式的数据流对比，可以帮你更直观地理解：

好了，这三种融合模式是算子融合的基石。理解它们之后，你就能看懂更复杂的融合策略了。记住一点：融合的本质是减少数据搬运，而不是减少计算量。在存算一体架构下，数据搬运的成本远高于计算，所以哪怕计算量不变，只要减少了数据读写，性能就能大幅提升。