算子融合基础：从理论到实战

大家好，我是你们的老朋友。今天咱们聊聊算子融合——这个听起来有点学术，但实际上每天都在帮你省算力的技术。

我记得刚入行那会儿，写模型就像搭积木，一个卷积接一个ReLU，中间还夹着BatchNorm。那时候GPU算力金贵，每毫秒都得精打细算。后来我发现，把这些小算子「粘」在一起，居然能省下30%以上的训练时间。嗯，这就是算子融合的魔力。

什么是算子融合？

说白了，就是把多个连续的计算步骤合并成一个。

举个例子。你写了个简单的网络层：

# 常见的三段式写法
x = torch.conv2d(input, weight, bias)
x = torch.batch_norm(x, mean, var)
x = torch.relu(x)

这三个操作，在底层是三个独立的kernel launch。每个kernel都要从显存读数据、算完、写回去。你想想看，中间那两次读写完全是浪费的——数据明明可以留在寄存器里直接算下一步。

算子融合就是干这个的：把conv、bn、relu合并成一个kernel，一次读取、一次计算、一次写回。

核心思想：减少中间数据的显存读写，让计算尽量在片上完成。

为什么需要算子融合？

我直接说结论：现代深度学习训练的瓶颈，往往不在计算，而在访存。

你看这张图就明白了：

我在项目中遇到过最夸张的情况：一个ResNet-50，光算子启动开销就占了总时间的15%。你想想，GPU每秒钟能跑几万亿次运算，结果有15%的时间在「准备干活」而不是「干活」——这谁能忍？

算子融合的好处，我总结成三点：

减少访存带宽压力——中间结果不用写回显存，直接在寄存器或L1 cache里传递
降低kernel launch开销——本来要启动3次，现在1次搞定
提升计算密度——每个线程做更多有用功，而不是等数据

实战小贴士：我习惯在模型训练前先用profiler跑一遍，看看哪些小算子频繁出现。如果发现连续3个以上element-wise操作，基本都可以考虑融合。

算子融合的数学原理

你可能会问：把几个算子合并，数学上真的等价吗？

答案是：对于线性或分段线性操作，完全等价。

咱们拿最经典的Conv+BN+ReLU来推导一下。

假设卷积的输出是：

z = W * x + b

BatchNorm的公式是：

y = γ * (z - μ) / σ + β

ReLU就是：

out = max(0, y)

把前两步合并：

y = (γ/σ) * (W * x + b - μ) + β
  = (γ/σ) * W * x + (γ/σ) * (b - μ) + β

看到了吗？这本质上还是一个线性变换！

令：

W' = (γ/σ) * W
b' = (γ/σ) * (b - μ) + β

那么整个Conv+BN就等价于：

y = W' * x + b'

再加上ReLU，就是：

out = max(0, W' * x + b')

完美融合，数学上分毫不差。

注意：这种融合只在推理阶段完全等价。训练时BN的均值和方差是动态计算的，不能简单合并。不过训练时我们可以做另一种融合——把前向计算中的连续操作合并，不影响反向传播。

计算图优化：编译器视角

算子融合在深度学习框架里，本质上是计算图优化的一环。

我习惯把计算图优化分成三个层次：

优化层次	典型操作	收益
代数化简	常量折叠、死代码消除	5-10%
算子融合	Conv+BN+ReLU、Add+ReLU	20-40%
内存规划	in-place操作、内存复用	10-20%

你看，算子融合是收益最高的那个。这也是为什么PyTorch的TorchScript、TensorFlow的XLA、TVM这些编译器，都把算子融合作为核心优化pass。

具体到实现，框架一般会做两件事：

模式匹配——在计算图上扫描，找到可以融合的算子序列（比如Conv→BN→ReLU这种固定模式）
子图替换——把匹配到的子图替换成一个融合后的算子节点

举个例子，在TVM中，你可以手动指定融合规则：

# TVM中的算子融合示例（简化）
@register_compute("conv2d_bn_relu")
def fused_conv_bn_relu(data, weight, gamma, beta, mean, var):
    # 卷积
    conv = topi.nn.conv2d(data, weight, strides=1, padding=1)
    # BN
    bn = topi.nn.batch_norm(conv, gamma, beta, mean, var)
    # ReLU
    out = topi.nn.relu(bn)
    return out

当然，实际框架的融合策略比这复杂得多。我记得有一次调试一个BERT模型，发现LayerNorm和后面的Add操作死活融合不了——因为LayerNorm的输出被两个不同的节点引用了。这种「多消费者」情况，融合反而会破坏计算图结构。

避坑指南：我曾经在融合Add+ReLU时踩过坑——如果Add的某个输入来自模型外部（比如残差连接），融合后梯度计算会出错。后来我养成了习惯：融合前一定检查所有输入的依赖关系。

小结

算子融合不是什么高深莫测的技术。说白了，就是让数据少跑路、让计算更集中。你想想，GPU的显存带宽就那么宽，每省一次读写，就是实打实的性能提升。

我个人建议，刚开始接触算子融合的同学，可以从最简单的Conv+BN+ReLU入手。先用profiler看看融合前后的性能差异，感受一下「数据局部性」带来的威力。等你真正理解了背后的访存模型，再去研究更复杂的融合策略——比如跨层融合、动态shape融合这些进阶话题。

嗯，今天就聊到这儿。记住一句话：好的算子融合，能让你的GPU「忙而不乱」。

📚 课程目录

算子融合基础：从理论到实战

什么是算子融合？

为什么需要算子融合？

算子融合的数学原理

计算图优化：编译器视角

小结