已合并
InplaceAddLayerNormFusionPass 融合规则迁移 #9018
InplaceAddLayerNormFusionPass 融合规则迁移 #9018
已合并
rk创建于 15 天前
rk成员
15 天前

描述

本 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,继承 FusionBasePassRun 入口先做全局守卫(场景、平台,只判一次),再扫图取所有 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;
  • 属性透传:epsilonadditional_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成员
15 天前 创建了 pull request,commit 501f3aee
atomgit-bot
atomgit-bot
15 天前 评论:

变更摘要

该 PR 将 AddLayerNormInplaceAddLayerNorm 的融合规则从 FE 内置实现迁移到 op_graph,新增 graph base 路线的融合 Pass ZInplaceAddLayerNormFusionPass(继承 FusionBasePass,经 REG_FUSION_PASS 注册在 kAfterBuiltinFusionPass 阶段)。实现刻意只使用原生 Graph/GNode 接口改图,不引入 PatternFusionPasses::*SubgraphRewriterInferShapeUtil 等符号,以满足「高版本编译、最低 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
15 天前 评论:

代码审查

✅ 未发现问题

likedislike
CANN-robotCANN-robot成员
15 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
15 天前 评论:

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成员
10 天前 添加了label:approved
TangPC
TangPC成员
10 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
10 天前 添加了label:lgtm
CANN-robotCANN-robot成员
10 天前 关闭了关联的issue
CANN-robotCANN-robot成员
10 天前 合入了pull request