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

描述

在 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_utilsNormInfo/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
(感谢 朱珉 的贡献)
朱珉成员
12 天前 创建了 pull request,commit 920894e5
CANN-robotCANN-robot成员
12 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
12 天前 评论:

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成员
12 天前 评论:

CLA Signature Pass

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

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

/lgtm

likedislike
CANN-robotCANN-robot成员
1 天前 添加了label:lgtm
wangxiaotian995成员
1 天前 评论:

/approve

likedislike
CANN-robotCANN-robot成员
1 天前 添加了label:approved
CANN-robotCANN-robot成员
1 天前 合入了pull request