已合并
InplaceAddLayerNormFusionPass 融合规则迁移 #9018
InplaceAddLayerNormFusionPass 融合规则迁移 #9018
已合并
rk创建于 8月21日
rk成员
8月21日

描述

本 PR 对 InplaceAddLayerNormFusionPass 图融合 pass 进行了迁移重构,使用 graph_metadef 的原生 Graph/GNode 接口改图,将图中的 AddLayerNorm 算子替换为 InplaceAddLayerNorm 算子,使 x1/x2 输入被原地写回,省去额外输出拷贝,降低显存占用与访存开销。配套 16 个 UT 用例覆盖正向替换与各守卫负向拦截。

改动原因

  • 该融合规则原先存在于其他路径,本 PR 将其迁移至 graph base(op_graph)路径,统一使用原生 Graph/GNode 接口改图,便于维护与后续演进;
  • InplaceAddLayerNorm 支持原地写回 x1/x2,推理场景下可减少显存与带宽开销。

改动方法

  • inplace_add_layer_norm_fusion_pass.{h,cpp}:实现 ZInplaceAddLayerNormFusionPass,继承 FusionBasePass。Run 入口先做全局守卫(场景、平台,只判一次),再扫图取所有 AddLayerNorm 候选节点,逐节点过守卫链后执行单节点替换(建点→接输入→改接输出→搬控制边→断旧边删旧点),替换失败时整图回退;
  • 守卫链:
    • 场景守卫:ge.graphRunMode==1(训练)跳过;CANN < 9.0.0 运行期 GetOptionValue 不可达时静默不改;
    • 平台守卫:仅 Ascend910B/Ascend910_93/Ascend950 支持,并取 L2 容量用于 shape 收益判定;
    • 独占性守卫:x1/x2 必须被本节点独占消费,避免原地写回踩踏他人数据;
    • shape 守卫:静态 shape 按 x1 大小相对 L2 的上下界判定收益;动态 shape 仅当最后一维为 5120(reduce 轴)时放行;
    • dtype 守卫:x1/x2/gamma 必须同 dtype;
  • 属性透传:epsilon、additional_output 原样拷贝到新节点;
  • tests/ut/op_graph/:UT 覆盖正向(无 bias/有 bias/bf16/输入输出改接/下游改接/属性透传/控制边迁移/910B 平台/动态 reduce 轴)与负向守卫(不支持平台/输入被共享/shape 过小/shape 过大/动态非 reduce 轴/混合 dtype)。

涉及算子:inplace_add_layer_norm(op_graph 融合规则)。

关联的Issue

测试

  • UT:tests/ut/op_graph/test_inplace_add_layer_norm_fusion_pass.cpp 共 16 个用例,覆盖正向替换与全部守卫负向拦截,全部通过;

文档更新

无

类型标签

AI/Agent生成声明

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 rk 的贡献)
Rrk成员
8月21日 创建了 pull request,commit 501f3aee
atomgit-bot
atomgit-bot
8月21日 评论:

变更摘要

该 PR 将 AddLayerNorm → InplaceAddLayerNorm 的融合规则从 FE 内置实现迁移到 op_graph,新增 graph base 路线的融合 Pass ZInplaceAddLayerNormFusionPass(继承 FusionBasePass,经 REG_FUSION_PASS 注册在 kAfterBuiltinFusionPass 阶段)。实现刻意只使用原生 Graph/GNode 接口改图,不引入 PatternFusionPass、es::*、SubgraphRewriter、InferShapeUtil 等符号,以满足「高版本编译、最低 8.5.0 运行」的兼容性要求;替换为 1:1 同构换型,算子数、拓扑、输入输出个数均不变,收益来自 InplaceAddLayerNorm 的原地写回语义(输出 0/3 复用输入 0/1 显存),并配套守卫链、DFX 事件与单元测试。

主要改动

  • 新增融合 Pass 实现:inplace_add_layer_norm_fusion_pass.cpp 实现 ZInplaceAddLayerNormFusionPass::Run,扫描图中 AddLayerNorm 节点,经 OperatorFactory::CreateOperator 建新节点、透传 epsilon/additional_output 属性,并逐位搬移数据边与控制边后删除旧节点,替换失败时整图回滚并设置错误消息。
  • 运行期兼容与训练场景守卫:通过 dlsym(RTLD_DEFAULT, ...) 解析 9.0.0 起提供的 CustomPassContext::GetOptionValue(mangled 符号 _ZNK2ge17CustomPassContext14GetOptionValueERKNS_12AscendStringERS1_),读取 ge.graphRunMode 判定训练场景(G1);运行时低于 9.0.0 或符号不可达时整体让位给老规则,编译期也仅在 GE_COMPILER_VERSION_NUM >= 90000000 时注册该规则。
  • 多级融合守卫:IsSupportedPlatform(G2,平台白名单 Ascend910B/Ascend910_93/Ascend950,统一用 short_soc_version 比对)、IsInplaceSafe(G4/G5,x1/x2 须被本节点独占消费)、IsShapeSupport(G6,静态 shape 须落在 L2 容量区间,动态 shape 末维须为 5120)、IsSameInputDataType(G7,x1/x2/gamma 须同 dtype),并输出 JSON 格式 DFX 事件(guard_rejected/skip/replacement_failed 等)。
  • 新增单元测试:test_inplace_add_layer_norm_fusion_pass.cpp 覆盖有/无 bias、bf16 的正向融合,输入输出接线与属性迁移、控制边搬运到新节点,以及平台、输入共享、shape 区间、混合 dtype、动态 shape 归约轴等守卫拒绝分支。
likedislike
不准确?
atomgit-bot
atomgit-bot
8月21日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CANN-robotCANN-robot成员
8月21日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
8月21日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.


PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
norm ✅ 钱泽洪, 汤平川 (2/2) ✅ 钱泽洪, 汤平川 (2/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

wangpengbo26, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了70条消息 查看更多
CANN-robotCANN-robot成员
8月27日 添加了label:approved
TangPC
TangPC成员
8月27日 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
8月27日 添加了label:lgtm
CANN-robotCANN-robot成员
8月27日 关闭了关联的issue
CANN-robotCANN-robot成员
8月27日 合入了pull request