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 下的最优融合策略不一样。

相关推荐
高级c9 小时前
BLAS 高性能算子库与 GEMM 优化原理
架构·cann
SuniaWang10 小时前
《Agentx专栏》02-技术选型:预算有限时如何做出正确的技术决策
java·spring·架构·langchain·milvus·agenx·opl
runafterhit10 小时前
开源知识库GitHub使用经验总结
开源·github
AI小百科10 小时前
目前开源AI编辑器面临的主要挑战是什么
人工智能·开源·编辑器
yantuguiguziPGJ10 小时前
机器人机械结构开源周报
机器人·开源
Yuk丶10 小时前
厌倦了假AI对话?用本地大模型给UE注入真智能(已开源!)
c++·人工智能·开源·ue4·游戏程序·ue4客户端开发
星栈10 小时前
别让 API 跳去登录页:我在 Axum 里做了认证失败双通道
前端·后端·开源
a11177611 小时前
VR看房 网页(开源 threejs)html
前端·开源·html·vr
ting945200011 小时前
TestSprite 3.0 深度技术解析:端到端 AI 自动化测试架构、核心能力与底层实现原理
人工智能·架构