GE 怎么做算子融合

算子融合是提升推理性能最直接的手段之一。

两个相邻的算子,如果前一个的输出正好是后一个的输入,而且中间数据不卖给其他算子,就可以合成一个算子。合成之后,原来需要两次显存读写(写 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 下的最优融合策略不一样。

相关推荐
W658034194 小时前
Meta Muse Glimmer 30B开源深度拆解:Apache 2.0+投机解码,24GB显卡跑本地Agent
ai·开源·大模型·apache·deepseek
kruptos10 小时前
分布式系统怎么“选主“?Raft 共识一次讲清
开发语言·分布式·php·共识算法
迅翼SwiftWing10 小时前
编队系列(总)|固定翼编队如何实现稳定飞行?通信架构、协同算法与工程约束深度解析
算法·架构
utmhikari10 小时前
【架构艺术】端到端监控告警和预案治理
架构·监控·告警·稳定性·sre·预案
今天AI了吗11 小时前
DeepSeek Harness 深度解析:从评测架构到实战落地
java·网络·数据库·人工智能·架构·java-ee
许彰午11 小时前
27-AuthService登录链路
java·低代码·架构
a11177611 小时前
基于 ROS 2 的 Franka Panda 机械臂视觉分拣系统
人工智能·开源
小智老师PMP11 小时前
2026深度解析|PMP第八版与NPDP核心侧重点本质区别(管理类证书怎么选)
开发语言·分布式·算法·职场和发展·产品经理
IT小白杨12 小时前
Facebook广告账户环境隔离实战:从指纹自洽到代理固定的完整工程方案
大数据·经验分享·架构·facebook·指纹浏览器
samble12 小时前
从零搭建工业控制系统(二十二):线程安全与并发控制
c#·wpf·并发·mvvm·线程安全·工业控制