
算子融合是提升推理性能最直接的手段之一。
两个相邻的算子,如果前一个的输出正好是后一个的输入,而且中间数据不卖给其他算子,就可以合成一个算子。合成之后,原来需要两次显存读写(写 HBM 再读 HBM)变成一个算子内部直接传,省掉一次显存搬运。对带宽敏感的模型(比如 Transformer),这个优化能带来 20-40% 的性能提升。
GE(Graph Engine)在 ATC 编译阶段做算子融合。
它的做法是先跑一遍图遍历,找出所有可以融合的算子对,然后按预定义的融合模板匹配。融合模板是硬编码在 GE 里的,不是动态搜索出来的。比如「Conv → BatchNorm → ReLU」是一个模板,「MatMul → Add → ReLU」是另一个模板。目前 GE 支持大约 30 种融合模板,覆盖了 CV 和 NLP 模型里最常见的算子组合。模板不匹配的算子对,即使理论上可以融合,GE 也不会去做。这是出于稳定性考虑------某些融合写法在特定输入 shape 下会有精度问题,硬编码模板等于只放开验证过的融合模式。
融合的顺序也有讲究。
GE 先做「纵向融合」(相邻算子合成),再做「横向融合」(多个相同算子合成一个)。纵向融合的收益通常更大,因为直接省掉了显存读写。横向融合的收益在于减少 kernel 启动次数------每个算子调用都要一次 kernel 启动,开销大约 1-2 微秒,批量合成之后只启动一次,省掉这部分开销。对层数多、单算子计算量小的模型(比如 MobileNet),横向融合的收益更明显。
融合之后,原来的两个(或更多)算子变成一个新算子,这个新算子的实现是 GE 预先写好的,不是动态生成的。比如 「Conv+BN+ReLU」融合之后,GE 调用的是一个叫 ConvBnRelu 的算子实现,这个实现在编译 CANN 的时候就已经编译好了,存在 OPP 目录里。ATC 编译时只是把这个算子选出来,替换掉原来的计算图节点。所以算子融合不会增加编译时间------匹配模板是 O(N) 的图遍历,30 层的模型也就几毫秒。
但不是所有能融合的算子都会融合。GE 有一个「融合收益评估」步骤:融合之后新算子的寄存器占用会不会超过 SRAM 上限?融合之后 cube 利用率会不会下降(有些算子组合起来反而不能并行)?这些问题 GE 会在融合之前评估一遍,收益为负的融合会被放弃。实际观测下来,ResNet50 有大约 70% 的算子对被融合,MobileNet 只有 40% 左右------后者的算子太多小算子,融合之后反而影响并行度。
算子融合对精度有没有影响?
理论上是没有的------融合只是把计算顺序换了,数学上是等价的。但实际实现里,融合算子的中间结果可能用不同的精度(比如 Conv 用 FP16 算,BN 也用 FP16,不融合的话 BN 可能会转成 FP32 算)。GE 的默认策略是「尽量用低精度」,所以融合后的模型精度可能跟原来略有差异(通常在 0.1% 以内)。如果对这个差异敏感,可以在 ATC 编译时加 --op_precision_mode=force_fp32,强制融合算子里的每个子算子都用 FP32,但性能会掉 15-25%。
最后说一下动态 shape 场景。
算子融合是针对静态计算图做的------编译时图结构已经确定了,才能做模板匹配。如果模型是动态 shape(比如输入图片的大小不固定),GE 会在编译时生成多个子图,每个子图对应一个常见的 shape 范围,每个子图内部独立做算子融合。运行时根据输入 shape 选对应的子图,所以动态 shape 模型的编译时间会比静态模型长 2-3 倍(要编译多个子图)。如果 shape 变化范围太大(比如从 224 到 1024 都有),GE 可能放弃融合某些算子,因为不同 shape 下的最优融合策略不一样。