已合并
【PR】: fix reduce layer norm #1891
【PR】: fix reduce layer norm #1891
已合并
czways创建于 9 天前
czways
czways成员
9 天前

Pull Request

Inductor LayerNorm 双链多引用断图方向不一致导致 SortSubGraphsByDependency 依赖环编译失败

描述

请清晰准确地描述本次 Pull Request 的意图和变更内容。

变更类型

请选择本次引入的变更类型:

关联的Issue

Closes #293

如何测试

描述测试此变更的步骤和前提条件:
1.
2.

核对清单

其他信息

在此添加任何其他关于本次 PR 的说明。

PR #1891 修改说明:fix reduce layer norm

PR 链接:https://gitcode.com/cann/graph-autofusion/pull/1891
关联 Issue:https://gitcode.com/cann/graph-autofusion/issues/293Closes #293
分支:czways/developcann/graph-autofusion@develop
当前版本:head f0329e47(2026-09-01 更新,已移除 tiling 相关修改,ATT GetCacheReuseInfo 修复由他人另行交付)
变更规模:3 个文件,+209 / -16
完整根因推演与复盘:见 autofuse/bugshot/tmp/root_cause_derivation.md


1. 问题现象

Inductor 场景执行 LayerNorm 融合用例 tc_af_inductor_layernorm_backend_0193(输入 shape=(128,6,23),模型为两条 LayerNorm 路径 + 三路输出:sum(layernorm(x.transpose(0,1)))mean(layernorm(x))abs(layernorm(x)))时,Autofuse Optimize 阶段在 SortSubGraphsByDependencysorted[4] != total[7],拓扑排序卡死,整体编译以 Optimize fail ret 1343225857 中断。

7 个子图中只排出 4 个,子图 1 与子图 3 互相等待——本 PR 修复该依赖环的根因。

历史备注:根因修复后曾暴露下游 ATT tiling cache 的 GroupLevelCache body/tail 编译失败(f8cf01e 休眠缺陷)。该二级问题不在本 PR 修复范围内,已由他人另行修复交付;本 PR 现只包含 Optimize 断图根因修复。

2. 根因分析

2.1 图结构前提

Inductor 将整个模型融合为单张 HintGraph,含两条几乎镜像的 LayerNorm 计算链(A 链 = transpose 后 LN,B 链 = 原始输入 LN),两链之间无直接数据边,但共享同一批输入节点:broadcast2(÷23 常量)、broadcast8(beta 权重)、broadcast9(gamma 权重)、scalar1(eps)。

2.2 断图机制

GenerateGeneralCase 依次执行三个断图 pass(ReducePartitionPostFusion → PartitionNorm → ReducePartitionMultipleCitations)。断开一条边的标准产物是 4 个新节点:

producer → Store → Ws_pre    ~~(同名,之间无边)~~    Ws_post → Load → consumer

写侧句柄与读侧句柄是两个不同节点、同一个名字,运行时靠同名配对传递内存。剪哪条边就是在决定"谁供谁"——pre 落在生产侧子图,post 落在消费侧子图,形成一条"生产方必须先执行"的跨子图依赖。

2.3 故障机制:断图方向不一致制造真实依赖环

ReducePartitionMultipleCitations(修复前)对每个多输出共享节点独立决策:"保留第一个遍历到的 consumer 方向,断开其余"。由于各节点 consumer 遍历顺序不同,方向互相矛盾:

共享节点 保留的 consumer 剪断侧 产生的跨链依赖
broadcast2(÷23) A 链 truediv B 链 truediv1 A 供 B(子图1→子图3)
broadcast8(beta) B 链 mul3 A 链 mul5 B 供 A(子图3→子图1)
broadcast9(gamma) B 链 add2 A 链 add3 B 供 A(子图3→子图1)

子图依赖图:

0 ──→ 4 ──→ 1 ⇄ 3 ──→ 6        ← 1⇄3 双向依赖即环
2 ──→ 5 ──↗

原图是无环 DAG;环完全由断图方向不一致制造,且每条依赖边都对应真实 workspace 数据流(非名称匹配假依赖),无法在排序层补救。

2.4 已证伪的备选方案(避免重走弯路)

方案 结论
修改 Workspace 节点名消除同名 证伪:同名是 buffer 复用的有意设计,改名破坏 allocator 映射
忽略 Workspace 依赖 / 在 HasDataDependency 加特判 证伪:本 case 依赖边全部真实,特判可致 consumer 先于 producer 执行
SCC 缩点合并环上子图 回退:可让排序通过,但改变子图融合粒度、跳过 R-core 模板,治标不治本
仅把 GetOutNodes() 换成 GetOutDataNodes() 证伪:断图后 synthetic Load 节点(如 mean_Load)确实拥有多个数据输出,API 切换挡不住,必须按节点类型过滤
每节点独立取最小 Reduce ID 证伪:局部最优不等于全局一致,各节点方向仍可互相打架

3. 修复思路与实现

3.1 共享引用链统一 anchor(断图方向全局裁决)

ReducePartitionMultipleCitationsreduce_schedule_case_generator.cpp/.h)重写为两阶段,并拆分为三个职责单一的函数:

ReducePartitionMultipleCitations        主流程(单 Reduce 早退)
├─ CollectCitationGroups                阶段一:图改写前收集 citation
├─ BuildCitationGroupAnchors            阶段二 a:Union-Find 合并 + 选 anchor
└─ PartitionCitationGroups              阶段二 b:统一方向断图 + 去重

阶段一:收集 citation(先于任何本 pass 内的断边,避免被自己改写过的图污染)

  1. 多输出起点过滤:!IsLoad && !IsStore && !IsOps<Workspace> && GetOutDataNodes().size() > 1——显式排除 synthetic 节点(断图产物 Store/Load/Ws 虽确实有多个数据输出,但不是用户图的真实共享节点);
  2. FindOutputReduce 仅沿数据边GetOutNodes()GetOutDataNodes())递归探测每个输出最终流到的 Reduce,形成 citation 三元组 (source, citation, reduce);不处理控制边是正确选择——下游依赖矩阵(HasDataDependency)本身只按数据边推导,方向决策必须与被保护的约束同构;
  3. citation 按 (reduce ID, citation ID) 升序排序——保证后续决策确定,不依赖容器遍历顺序(对应编码红线"图改写结果必须确定")。

阶段二 a:共享链合并与 anchor 选取

  1. Union-Find:任意两个 citation 组若共享同一 Reduce,合并为同一条共享引用链。本图 7 个共享节点(scalar1、b2、b8、b9、rsqrt、rsqrt1、add3)因都同时流到 sum2/sum3 被合并成一条链;
  2. 每链取最小 Reduce ID 为唯一 anchor(本图 sum2=41 < sum3=51 → anchor=sum2)。取最小 ID 非性能偏好,是要一个确定、可复现的裁决规则。

阶段二 b:统一方向断图

  • 规则全链唯一:citation.reduce == anchor 的路径保留原始数据边;其余路径统一调用 PartitionByNode 断开。所有断边方向一致把非 anchor 侧(sum3 侧)隔离为纯消费侧,依赖图在构造上不可能成环;
  • (source, reduce) 对去重PartitionByNode(src, dst) 只剪断 src→dst 单边(实现含 peer == dst_node 守卫,不改道 src 其余出边)。修复前的旧实现通过 reduce_nodes 集合隐式按 (source, reduce) 去重——首个 Reduce 方向保留直连,每个新出现的 Reduce 只切一次;按 (source, citation) 边去重会切掉同一 source 流向同一非 anchor Reduce 的全部 citation 边,切分边集合较旧实现扩大、属行为变更,故对齐旧语义(每个非 anchor Reduce 只切一条 citation 边,其余保持直连,行为保守且不引入额外跨分量依赖)。

3.2 与旧实现的行为对比

维度 旧实现 新实现
多输出起点判定 GetOutNodes().size() > 1(含控制边、含 synthetic 节点) 仅数据边出度 > 1,排除 Load/Store/Workspace
下游 Reduce 探测 沿控制边也可达 仅沿数据边
断边方向 每节点独立,由遍历顺序决定(未定义,可成环) 共享链合并 + 最小 Reduce ID anchor 全局裁决(构造性无环)
去重语义 (source, reduce) 隐式去重 (source, reduce) 显式去重(对齐)
断边原语 PartitionByNode 族 未改(Workspace 命名/buffer 复用/codegen ABI 保持)

3.3 真实图验证(2026-08-31 bugshot 回归)

修复后 [TEMP_REDUCE] 插桩日志(验证时临时添加,PR 已清理)逐条确认:

  • 9 个多输出节点中无任何 Load/Store/Workspace 起点——synthetic 过滤生效;
  • 7 个 citation 组 Union-Find 合并为一条链,anchor 统一为 sum2(41);
  • 全部 6 条断边(b2→truedivb9→add3b8→mul5rsqrt→broadcast11add3→sum3scalar1→add方向一致断向 sum3 侧——对照修复前 b2→truediv1 断向 sum2 侧的方向反转,反向依赖对消除;
  • 原始 1 ⇄ 3 环消失,SortSubGraphsByDependency 通过。

4. 变更文件清单

文件 变更
autofuse/optimize/task_generator/reduce_schedule_case_generator.cpp ReducePartitionMultipleCitations 重写(拆分三函数);FindOutputReduce 数据边化;新增 <map>/<numeric>/<set>
autofuse/optimize/task_generator/reduce_schedule_case_generator.h 新增 Citation/CitationGroups 类型与三个函数声明(均在 private: 段,无 ABI 影响)
autofuse/tests/ut/optimize/task_generator/test_reduce_schedule_case_generator.cpp 新增 2 条 UT(见下)

注:本 PR 早期版本曾包含 ATT GetCacheReuseInfo 统一过滤(tiling_code_generator.cpp + att_generator_unittest.cpp 2 条 UT),因该问题已由他人修复,为避免重复交付已移除(2026-09-01,head 70cc7e5df0329e47)。

5. 测试设计

用例 覆盖点
TestReduce_Multi_Cita_Multi_Out_NoDependencyCycle 双 Reduce 交叉引用链(shared0/shared1 两级多输出、sum1/sum0 双 Reduce)——真正进入新算法路径(单 Reduce 图在 IsOnlyHasOneOrLessReduce 早退);断言调度成功 + kCommon 模板 + 两组 workspace 断点对恰各一对
TestReduce_Multi_Cita_SameReduce_NoDuplicatePartition 同一 source 多条 citation 流向同一 Reduce 的场景,断言 workspace 断点对数量符合预期。:按图构造 sum0 先创建(ID 更小)→ anchor=sum0,实际仅切 branch2(sum1 侧唯一 citation)一条边,(source, reduce) 与 (source, citation) 两种去重键结果相同——去重路径本身未被该用例区分,属残留覆盖缺口;建议补"非 anchor 侧 ≥2 条 citation"构造

已执行的验证

验证项 结果
optimize UT(ReduceScheduleCaseGeneratorTest 全套含新增 2 条) ✅ 通过
optimize ST(多 citation 多输出归一化拓扑强化用例) ✅ 通过
真实 bugshot:Optimize 拓扑排序(原故障点) ✅ 通过(0831 回归,环消除)
真实 bugshot:host C++ 编译(asc_tiling.py ⏳ 依赖 ATT 侧修复(他人交付)验证
真实 bugshot:设备端三路输出数值精度 ⏳ 待执行

6. 性能影响

  • 编译期:一次 citation 收集 + O(组数²×组内 citation) 的共享 Reduce 探测(组数与本图规模为个位数)+ 集合判重,均为编译期常数开销;FindOutputReduce 无 memo,重汇聚宽图存在重复子遍历(可观测但非运行时问题)
  • 运行时:零新增分支;断边原语未改,生成的 kernel 代码结构不变
  • 运行时性能变化点:切分方向确定化会改变部分存量图(多输出节点流向多 Reduce 的形态)的分图结构,融合粒度/性能特征可能与历史版本不同(方向为切分面收窄、更保守)——非 LayerNorm 场景的性能基线对比是合入前建议项

7. 遗留与后续

  1. ATT GetCacheReuseInfo 修复由他人交付后,在真实 bugshot 环境联合重跑:确认 asc_tiling.py host C++ 编译通过 → 设备端三路输出精度校验通过(E2E 闭环);
  2. E2E 全绿后关闭 Issue #293 并归档 bugshot 分析文档;
  3. 测试覆盖补强:新增"同一 source ≥2 条 citation 流向同一非 anchor Reduce"用例,真正触达 (source, reduce) 去重与并查集合并路径;
  4. (可选)FindOutputReduce 增加节点级 memo(node → 首个可达 reduce)与 visited 保护:重汇聚图上避免重复子遍历,总复杂度降为 O(V+E),同时消除理论上的递归栈溢出风险(阶段一无图改写,memo 安全);
  5. (既有问题,非本 PR 引入,建议另立 issue)PartitionByNode 对多输出口 source 会为每个 out_anchor 建 Store/Ws 写链——dst 未连接的端口产生无读者的悬空写链;本 PR 场景 source 均单口未触发。
likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 czways 的贡献)
czwaysczways成员
9 天前 添加了label:enhancement
czwaysczways成员
9 天前 创建了 pull request,commit 3a98fb28
atomgit-bot
atomgit-bot
9 天前 评论:

变更摘要

本 PR 旨在修复 reduce 算子调度用例生成中的切分逻辑(标题 fix reduce layer norm),主要重构了 ReducePartitionCaseGenerator::ReducePartitionMultipleCitations:不再简单收集多输出节点后逐个切分,而是按"引用组"(source/citation/reduce 三元组)组织多输出节点,仅沿数据边查找 reduce,并借助并查集合并共享 reduce 的组、以组内 reduce id 最小的锚点为基准去重切分,避免重复切分与潜在依赖环;同时为 PartitionByNode 增加边替换跟踪日志,并新增对应单元测试。

主要改动

  • 重构多引用切分逻辑ReducePartitionMultipleCitations 由收集 multi_output_nodes 改为构造 Citation 结构(source/citation/reduce)组成的 citation_groups,按 reduce id、citation id 排序后分组,并只统计 GetOutDataNodes().size() > 1 且非 Load/Store/Workspace 的节点。
  • 并查集合并与锚点去重:引入 parent/find_root 并查集,合并共享同一 reduce 的引用组,每组以 reduce id 最小的节点作为 group_anchors 锚点,仅对非锚点且未被 partitioned_edges 记录过的边调用 PartitionByNode,防止重复切分。
  • 改为仅沿数据边查找 reduceFindOutputReduce 及遍历逻辑由 GetOutNodes() 改为 GetOutDataNodes(),reduce 查找与引用方向均只沿数据边(data edges)进行。
  • 切边替换跟踪与日志PartitionByNode 中新增 replaced_edge 标记,实际移除 src->dst 边时置位并打印 [TEMP_REDUCE] 日志,若未找到直接边替换则输出告警日志 GELOGW
  • 新增单元测试test_reduce_schedule_case_generator.cpp 增加 TestReduce_Multi_Cita_Multi_Out_NoDependencyCycle,使用 AscGraphBuilder 构造 SumAbs/Tanh/Add/Relu 多引用的图,验证 GeneratorTask 返回成功且任务列表非空。
likedislike
不准确?
atomgit-bot
atomgit-bot
9 天前 评论:

代码审查

审查结论

变更文件审查确认

  • autofuse/optimize/task_generator/reduce_schedule_case_generator.cpp — 已审查。核心变更:重写 ReducePartitionMultipleCitations(citation 收集 + union-find 分组 + anchor 选取 + 切分去重)、FindOutputReduce 改为仅沿数据边递归、PartitionByNode 增加 replaced_edge 跟踪与日志。发现 2 个问题(见下)。
  • autofuse/tests/ut/optimize/task_generator/test_reduce_schedule_case_generator.cpp — 已审查。新增 TestReduce_Multi_Cita_Multi_Out_NoDependencyCycle,其构建图仅含一个 reduce,ReducePartitionMultipleCitations 在入口处即提前返回,新增测试未执行任何被重写的逻辑。发现 1 个问题(见下)。

Findings 统计

优先级 数量 摘要
P2 1 新增测试未覆盖重写后的切分逻辑(测试图仅 1 个 reduce,函数在 IsOnlyHasOneOrLessReduce 处提前返回,407-500 行新代码一行未执行)
P3 1 生产 pass 中遗留约 10 处 [TEMP_REDUCE] 临时 GELOGI/GELOGW 调试日志,线上 INFO/WARN 刷屏
P3 1 新 anchor 分组语义与旧实现不一致(多源共享 reduce 时切分边集合改变、单引用源节点可能被切、重复调用 PartitionByNode 会向已改写锚点重复附加 store/workspace/load 链),且无测试兜底(低置信度)

总体风险评估

本 PR 对 ReducePartitionMultipleCitations 的切分算法进行了完整重写,算法内部(union-find 合并、anchor 选取、去重)逻辑自洽,未发现必然崩溃或数据损坏类缺陷;但存在两方面风险:其一,唯一新增测试因图结构只有一个 reduce 而完全未触达新逻辑,重写的复杂分支逻辑处于零回归保护状态;其二,新旧实现切分边集合在跨源共享 reduce 场景下明显不同,属行为变更,在缺少测试与明确语义说明的情况下存在功能回归隐患。此外遗留的临时调试日志应在合并前清理。整体风险中等,建议补齐能真正执行新逻辑的测试用例后再合入。

类型 数量
🔴 阻塞 0
🟡 建议 1

💬 仅评论

likedislike
不准确?
CANN-robotCANN-robot成员
9 天前 添加了label:stat/needs-squash
此处折叠了54条消息 查看更多
yangyongqiang
yangyongqiang成员
7 天前 评论:

/approve

likedislike
CANN-robotCANN-robot成员
7 天前 添加了label:approved
CANN-robot
CANN-robot成员
7 天前 评论:

The MR can not be merged, because of CodeReview discussion not resolved

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike
yangyongqiangyangyongqiang成员
7 天前 解决了最后一个问题
CANN-robotCANN-robot成员
7 天前 合入了pull request