已合并
feat: 支持Gather后置Norm复合区域识别与多阶段融合候选生成 #2064
feat: 支持Gather后置Norm复合区域识别与多阶段融合候选生成 #2064
已合并
朱珉创建于 9月11日
朱珉成员
9月11日

描述

在 IndirectLoad 候选生成链路中支持 Gather 后置 Norm 复合区域(多个 Reduce + Broadcast + Elementwise 组合)的识别、校验与调度接入,并打通稳定 Softmax 的专用/通用双路径分流。

问题场景

含 IndirectLoad 的图后接 Norm 类计算(Softmax / LayerNorm 等)时,现有 Gather+Reduce 候选链存在三类问题:

  1. Gather+Softmax 候选被误杀:SoftmaxApiCall::Generate 固定把最后一个向量化轴当作 R,且专用 Softmax 节点保持完整输入输出形状;普通 Reduce 的 stride 推导(BuildPostReduceLayout)对它推不出归约轴,导致所有 Gather+Softmax 候选在布局校验阶段被淘汰。
  2. 多 Reduce 图断言失败:CollectOutputBoundaries 遇到第二个 Reduce 时触发图级断言 analysis.post_reduce == nullptr,LayerNorm 均值/方差双分支形态直接编译失败。
  3. SIMT 静默降级:FindPostReduceChain 遇多个 Reduce 返回空,has_post_reduce=false 使 SIMT 错误降级为直接写 GM 输出,绕过 Norm 统计阶段。

问题原因

  • 现有 RewrittenGraphAnalysis::post_reduce 为单 Reduce 假设,是 Gather+Reduce 全链路的布局锚点,但无法表达多统计 Reduce 区域;
  • Softmax Pattern 替换依赖 ReducePartitionCaseGenerator 的 AllLoad 流程,而该 Generator 遇 IndirectLoad 提前返回,主图中稳定 Softmax Pattern 永远不会被替换为专用节点;
  • 归约轴推导与尾轴 API 约束(SoftmaxApiCall 尾轴 R)之间缺乏显式分流。

修改方案

  1. 新增 norm_utils:NormInfo/NormStage 元数据结构,基于节点 ExtAttr(与 TemplateAxes/LogicalView 同模式)随候选图复制传递,不侵入通用 ScheduleTask 结构。
  2. 抽取 softmax_pattern_fusion_utils:将 Softmax 匹配逻辑拆为 MatchStableStructure(结构+dtype+布局+消费者闭合)与 MatchStableDedicated(+尾轴 R 约束);原 Pass 改为薄壳调用,全图行为不变。
  3. 候选图副本内 Softmax 规范化:NormalizeDirectPostSoftmax 仅替换原始输入可回溯到 IndirectLoad 输出、且满足专用 API 尾轴约束的 Pattern;非尾轴 Pattern 保留原始 Max/Sum 图进入通用路径。
  4. 多 Reduce 复合区域:CollectOutputBoundaries 收集全部 Reduce(多 Reduce 时不停在首个),逐个校验同轴 Broadcast 配对(stride 差推导归约轴 vs 零 stride 推导恢复轴,轴 ID 集合相等);CollectGatherNormInfo 校验入口/最终输出 View 一致后写入多 stage NormInfo。
  5. SIMD 逐 Reduce 独立校验:ValidateSimdCompositeLayouts 对每个 Reduce 独立执行 ValidateSimdPostReduceLayout(G≥R+后缀连续),避免整链转换次数误判;SeedCompositeReduceInputViews 为全部 Reduce 种向量化视图。
  6. SIMT 串行多阶段:FindPostReduceChain 保留首个 Reduce 作为 local target 锚点并继续遍历下游 Broadcast/Elementwise/Reduce,后续统计阶段由普通调度链执行。
  7. 分组完整性:FinalizeGroupedGraphLoweringMetadata 校验 Norm 区域节点与 IndirectLoad 同组(SIMD/SIMT 图级断言;SK 的 workspace 分区拆分为合法形态仅记录日志)。

修改效果

  • 修复后三类形态进入正常调度链:Gather+Softmax(专用 API)、Gather+LayerNorm(多 Reduce 共享输入)、Gather+串行多阶段;
  • 全部失败路径改为候选级淘汰并记录日志(原为图级断言失败);
  • 多阶段区域复用既有节点级 Reduce/Broadcast/Elementwise Codegen,不新增专用 Norm 设备 API;
  • 原有单 Reduce、无 Reduce 场景行为完全不变;
  • 未侵入 ScheduleTask 等通用结构。

变更类型

关联的Issue

无(如需方案讨论 Issue 可后续补充链接)

如何测试

  1. 已执行:bash build.sh --pkg 增量构建,全部修改文件编译通过(0 生产编译错误);
  2. 已验证:构建失败项仅为既有 ST 环境问题(TestBackendIndirectLoadBroadcastE2e 的 (index < tensors_.size()) 断言,干净基线同样失败);
  3. 未执行:UT/ST 专项测试、NPU 设备精度与性能验证(待后续补充)。

核对清单

其他信息

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 朱珉 的贡献)
朱朱珉成员
9月11日 创建了 pull request,commit 920894e5
CANN-robotCANN-robot成员
9月11日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
9月11日 评论:

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
repo-cann/graph-autofusion ✅ wangxiaotian995, xchu42 (2/2) ✅ wangxiaotian995 (1/1)

💡 Tip:

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

CLA Signature Pass

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

likedislike
CANN-robot
CANN-robot成员
9月11日 评论:

CLA Signature Pass

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

likedislike
CANN-robotCANN-robot成员
9月11日 将niuyuhu,yangyongqiang0606,liyuewei,xchu42,zhang_shengjie,depeng1994,wqtshg_wt,peiyang,shengnan666,XuebinYang,sjtulxh,zhanj,zhujingjing,xuyafei,wangxiaotian995设为评审人
此处折叠了1049条消息 查看更多
wangxiaotian995成员
9月21日 评论:

/lgtm

likedislike
CANN-robotCANN-robot成员
9月21日 添加了label:lgtm
wangxiaotian995成员
9月21日 评论:

/approve

likedislike
CANN-robotCANN-robot成员
9月21日 添加了label:approved
CANN-robotCANN-robot成员
9月21日 合入了pull request