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

相关推荐
这个DBA有点耶13 分钟前
数据库迁移怎么做到“零翻车“:金仓 KDMS/KDTS/KFS 工具链实测
数据库·架构·dba
lucky_syq1 小时前
第5篇 · S1·下:前沿架构:MoE、Reasoning 模型、长上下文、多模态、SSM/Mamba 与模型谱系
人工智能·学习·架构
无凭1 小时前
DeerFlow 的可观测性(一):RunJournal 如何记录 Agent 运行过程
人工智能·开源
lucky_syq2 小时前
第3篇 · S1·上:什么是大语言模型 + Transformer 架构深讲
人工智能·语言模型·架构·transformer
SelectDB2 小时前
网易游戏 湖仓一体架构:Apache Doris / SelectDB 的技术能力与实践
开源
这个DBA有点耶3 小时前
迁移中的数据一致性挑战:如何确保百万行数据“搬得对”
mysql·架构·dba
zlinear数据采集卡4 小时前
数据采集卡从入门到精通(10):采样率与分辨率的核心关系——反比律、架构分布与过采样
arm开发·嵌入式硬件·算法·fpga开发·架构·开源
FserSuN4 小时前
回顾Spark的概念与应用
大数据·分布式·spark
m4Rk_5 小时前
【论文阅读】Agent 记忆机制(41):Agentic Plan Caching——将历史执行轨迹转化为可复用的计划记忆
论文阅读·人工智能·学习·开源·github
Hy行者勇哥5 小时前
软硬件一体化公司 企业标准体系架构
架构