已合并
feat: 支持Gather后置Norm复合区域识别与融合调度(含SIMT side-input视图适配) #2329
feat: 支持Gather后置Norm复合区域识别与融合调度(含SIMT side-input视图适配) #2329
已合并
朱珉创建于 6 天前
朱珉成员
6 天前

描述

在 IndirectLoad 候选生成链路中支持 Gather 后置 Norm 复合区域(Softmax / LayerNorm 等多个 Reduce + Broadcast + Elementwise 组合)的识别、校验与调度接入,打通 Softmax 专用/通用双路径分流;同时适配 SIMT body 内行级广播 side-input 的坐标重建。

问题场景

含 IndirectLoad 的图后接 Norm 类计算时,现有 Gather+Reduce 候选链存在四类问题:

  1. Gather+Softmax 候选被误杀:SoftmaxApiCall::Generate 固定把最后一个向量化轴当作 R,且专用 Softmax 节点保持完整输入输出形状;普通 Reduce 的 stride 推导对它推不出归约轴,导致候选在布局校验阶段被淘汰。
  2. 多 Reduce 图断言失败:CollectOutputBoundaries 遇到第二个 Reduce 时触发图级断言,LayerNorm 均值/方差双分支形态直接编译失败。
  3. SIMT body 内行级广播 side-input 编译失败:带 arange/matmul 的图,softmax 复合区域的行级广播 side-input Load(如 [8,2048,1]/strides=[2048,1,0],每行读一个向量沿尾轴广播)经调度期视图 split 改写后 rank 失配且尺寸符号化,坐标重建断言失败。
  4. IL 模板 merge 后 vectorized_axis 不一致:tensor merge 只折叠 axis/repeats/strides 不处理 vectorized_axis,中间窗口公共检查打误导性 ERROR。

问题原因

  • 现有 RewrittenGraphAnalysis::post_reduce 为单 Reduce 假设,无法表达多统计 Reduce 区域
  • Softmax Pattern 替换依赖 ReducePartitionCaseGenerator 的 AllLoad 流程,遇 IndirectLoad 提前返回
  • SIMT body 是静态成员函数(无 tiling data 参数),坐标重建要求编译期常量,视图 split 符号化后兜底失效
  • tensor merge 与 vectorized_axis remap 分两步执行,中间窗口视图自洽性缺失

修改方案

复合区域识别与调度:

  1. 新增 norm_utils:NormInfo/NormStage 元数据,基于节点 ExtAttr 随候选图复制传递
  2. 抽取 softmax_pattern_fusion_utils:MatchStableStructure/MatchStableDedicated 拆分匹配
  3. CollectOutputBoundaries 收集全部 Reduce,逐个校验同轴 Broadcast 配对
  4. SIMD 逐 Reduce 独立校验(ValidateSimdCompositeLayouts);SIMT 串行多阶段
  5. simt_boundary_sync:SIMT 边界视图 split 同步与 vectorized_axis 两级重映射
  6. SIMD padded 视图强制 strided 降级(按视图 strides 逐行写,行内 pad 由 ReduceInit 清中性值)

SIMT side-input 视图适配:
7. SimtLoadMetadata 增加行级广播标记:generator 在视图改写前按原始结构判定(尾轴 stride==0、其余轴 row-major 稠密递推——从后往前),写入节点 attr 穿透 VF 克隆
8. codegen 坐标重建兜底:rank mismatch 断言前按 output_index / 尾轴宽 去尾维寻址(读[行,列]值沿尾轴广播的正确语义)
9. IL 模板 tensor merge 后,vectorized_axis 中失效旧轴立即替换为 merged 轴(消除中间窗口不一致,RemapOutputVectorizedAxes 的 direct_ancestor 路径无缝衔接到 inner 轴)

修改效果

  • Gather+Softmax(专用 API)、Gather+LayerNorm(多 Reduce 共享输入)、Gather+串行多阶段均进入正常调度链
  • 行级广播 side-input(含带 arange/matmul 图)正确寻址
  • SIMD padded 视图降级 strided,消除 index 越界读
  • 全部失败路径候选级淘汰并记录日志(原为图级断言失败)
  • 原有单 Reduce、无 Reduce、SIMD 模板场景行为不变
  • 未侵入通用 ScheduleTask 结构

变更类型

关联的Issue

无

如何测试

  1. 已执行:codegen 全量 55/55、optimize_ut 734、optimize_st 213 全部通过
  2. 已验证:行级广播兜底(标记/递推/寻址)各环节形态断言通过
  3. 待执行:生产带 arange/matmul 图复测

核对清单

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 朱珉 的贡献)
朱朱珉成员
6 天前 创建了 pull request,commit c59cf0b3
atomgit-bot
atomgit-bot
6 天前 评论:

变更摘要

本 PR 恢复被回退的 Gather 后置 Norm 复合区域融合能力,并叠加行级广播与视图一致性修复。核心是在候选图生成阶段引入 NormInfo/NormStage 元数据与 norm_utils,把 Gather(IndirectLoad)后置的 Reduce→Broadcast→Elementwise 链识别为轴保持复合区域;同时新增 simt_boundary_sync 在调度全部 split 完成后统一对齐 SIMT 边界节点视图,并通过 SoftmaxPatternFusionUtils、SimtLoadMetadata 行级广播标记、padded 视图强制 strided 降级等机制修复生产暴露的 rank mismatch 与 vectorized_axis 不一致问题。

主要改动

  • 新增 Norm 复合区域元数据与工具:新增 norm_utils.h/.cpp,定义 NormInfo(kSoftmaxDedicated/kGenericComposite 两种 kind)、NormStage 及 SetNormInfo/TryGetNormInfo/HasNormInfo,将区域入口/出口、成员节点、各 Reduce→Broadcast 组合与保持轴集合附着在 IndirectLoad 节点上随图传递。
  • 新增 SIMT 边界视图同步:新增 simt_boundary_sync.h/.cpp 的 SyncSimtBoundaryViews,对 SIMT 角色节点按实际 (outer, inner) split 逐层同步 tensor view,并将失效旧轴经 merge 链重映射到 inner 轴;schedule.cpp 在 DoScheduler 末尾调用,且在模板 merge 后立即把失效 vectorized_axis 替换为 merged 轴。
  • 行级广播 side-input 兜底:SimtLoadMetadata 增加 original_view 与 is_row_broadcast 字段,schedule_result.h 新增 SetRowBroadcastLoad/IsRowBroadcastLoad 标记;reg_indirect_load_api_call.cpp 在 rank mismatch 断言前按 IsDenseEquivalentView、退化视图、稠密尾段取模与 output_index / 尾轴宽 去尾维寻址逐级兜底。
  • Softmax 模式识别抽取与复用:将原 softmax_pattern_fusion_pass.cpp 中的匹配/替换逻辑抽取为 softmax_pattern_fusion_utils.h/.cpp,新增 MatchStableStructure/MatchStableDedicated/MatchStable 与 NormalizeDirectPostSoftmax,支持沿多输入 Elementwise 链回溯至 IndirectLoad 输出并在候选图副本内替换为专用 Softmax 节点。
  • SIMD padded 视图强制 strided 降级:FinalizeLoweringMetadata 与 GenerateSimd 通过比较向量化跨度与逻辑元素积判定输出视图被 pad,强制 SimdFallback::kStrided,避免 dense facade 线性写出与 padded 布局冲突。

其他配套调整:autoschedule.cpp 新增 IsAxisStreamingOnAnyOutput 循环轴豁免与可求解 tile 预建 case,node_cache_marker.cpp 恢复参数 Load 缓存起点,codegen_kernel_loop.cpp 增加固定间接寻址参数判定,indirect_load_utils.cpp 调整 post-reduce 链查找与 tpipe tensor 收集豁免,并补充 test_norm_utils、test_simt_boundary_sync、test_softmax_pattern_fusion_utils、IL_USER_SOFTMAX 端到端用例及多个 UT/ST 断言更新。

likedislike
不准确?
atomgit-bot
atomgit-bot
6 天前 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CANN-robotCANN-robot成员
6 天前 添加了label:stat/needs-squash
CANN-robotCANN-robot成员
6 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
6 天前 评论:

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 ✅ yangyongqiang0606, xchu42, xuyafei (3/2) ✅ yangyongqiang0606 (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-robotCANN-robot成员
6 天前 将niuyuhu,yangyongqiang0606,liyuewei,xchu42,zhang_shengjie,depeng1994,wqtshg_wt,peiyang,shengnan666,XuebinYang,sjtulxh,zhanj,zhujingjing,xuyafei,wangxiaotian995设为评审人
此处折叠了5条事件消息 查看更多
CANN-robotCANN-robot成员
6 天前 删除了label:stat/needs-squash
CANN-robot
CANN-robot成员
6 天前 评论:

Notification

This pull request has been changed(code update) or closed, so removes the following label(s): stat/needs-squash.

likedislike
CANN-robotCANN-robot成员
6 天前 添加了label:cann-cla/yes
此处折叠了7条事件消息 查看更多
朱朱珉成员
6 天前 修改了pull request 的描述
朱珉成员
6 天前 评论:

compile

likedislike
朱朱珉成员
6 天前 预合并成功(commit_id: 88a2fd6f1a26671e41ba428fb28fc4b0bb121641)
CANN-robot
CANN-robot成员
6 天前 评论:
🚀 CI 流水线已启动
📋 执行详情: 点击查看流水线
likedislike
CANN-robotCANN-robot成员
6 天前 添加了label:ci-pipeline-running
CANN-robotCANN-robot成员
6 天前 删除了label:ci-pipeline-running
CANN-robotCANN-robot成员
6 天前 添加了label:ci-pipeline-passed
xchu42
xchu42成员
6 天前 评论:

/lgtm

likedislike
xuyafei成员
6 天前 评论:

/lgtm

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

/approve

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