Star705
263
代码
代码
Issues65
Pull Requests157
Actions
项目讨论
Wiki
分析
Star705
263
已合并
【fix】: 解决同级多for循环变量未定义问题 #2018
WangYanMale创建于 18 天前
【fix】: 解决同级多for循环变量未定义问题 #2018
已合并
WangYanMale创建于 18 天前
从WangYanMale/graph-autofusion_dwsdevelop合入到cann/graph-autofusiondevelop
共 4 个文件变更+779-12
Mautofuse/optimize/schedule_utils.cpp+159-12
@@ -46,21 +46,21 @@ bool IsMulConsumerStruct(const af::NodePtr &node) {
46 return false;46 return false;
47}47}
48 48 
49-Status FindNodeSequence(af::Node *start_node, std::unordered_set<af::Node *> &reduce_sequences) {49+Status CollectPostReduceNodes(af::Node *start_node, std::unordered_set<af::Node *> &post_reduce_nodes) {
50 GE_ASSERT_NOTNULL(start_node);50 GE_ASSERT_NOTNULL(start_node);
51- if (reduce_sequences.count(start_node) > 0UL) {51+ if (post_reduce_nodes.count(start_node) > 0UL) {
52 return af::SUCCESS;52 return af::SUCCESS;
53 }53 }
54 std::queue<af::Node *> node_queue;54 std::queue<af::Node *> node_queue;
55 node_queue.emplace(start_node);55 node_queue.emplace(start_node);
56- reduce_sequences.emplace(start_node);56+ post_reduce_nodes.emplace(start_node);
57 while (!node_queue.empty()) {57 while (!node_queue.empty()) {
58 auto node = node_queue.front();58 auto node = node_queue.front();
59 node_queue.pop();59 node_queue.pop();
60 for (auto &out_node : node->GetOutDataNodes()) {60 for (auto &out_node : node->GetOutDataNodes()) {
61 GE_ASSERT_NOTNULL(out_node);61 GE_ASSERT_NOTNULL(out_node);
62- if (reduce_sequences.count(out_node.get()) == 0UL) {62+ if (post_reduce_nodes.count(out_node.get()) == 0UL) {
63- reduce_sequences.emplace(out_node.get());63+ post_reduce_nodes.emplace(out_node.get());
64 node_queue.emplace(out_node.get());64 node_queue.emplace(out_node.get());
65 }65 }
66 }66 }
@@ -93,6 +93,66 @@ bool IsNeedFixTopo(const af::AscGraph &graph, bool use_rdfs_v2) {
93 GELOGD("Skip fix topo: no reduce multi-consumer found in graph[%s].", graph.GetName().c_str());93 GELOGD("Skip fix topo: no reduce multi-consumer found in graph[%s].", graph.GetName().c_str());
94 return false;94 return false;
95}95}
96+ 
97+// BuildLoopGroups的扩散种子：记录节点、loop_axis与发现该种子的轮次（即当时正在填充的分组编号）
98+struct LoopGroupSeed {
99+ af::AscNode *node;
100+ int64_t loop_axis;
101+ size_t round;
102+};
103+ 
104+// 按当前拓扑序返回图中第一个reduce节点，不存在时返回nullptr
105+af::AscNode *FindFirstReduce(const af::AscGraph &graph) {
106+ for (const auto &node : graph.GetAllNodes()) {
107+ if (ScheduleUtils::IsReduce(node)) {
108+ return node.get();
109+ }
110+ }
111+ return nullptr;
112+}
113+ 
114+// 将一组邻接节点编入group_id分组：loop_axis与组轴相同或为-1的节点入组并入fill_queue继续扩散
115+// （-1节点按当前组的loop_axis参与排序），不同轴有效节点作为新种子入seed_queue并记录发现轮次
116+template <typename NodeContainer>
117+Status SpreadLoopGroupNeighbors(const NodeContainer &neighbors, const int64_t group_axis, const size_t group_id,
118+ std::vector<LoopGroup> &loop_groups,
119+ std::unordered_map<af::Node *, size_t> &node_to_group,
120+ std::queue<af::AscNode *> &fill_queue, std::queue<LoopGroupSeed> &seed_queue) {
121+ for (const auto &next_node : neighbors) {
122+ auto next_asc_node = std::dynamic_pointer_cast<af::AscNode>(next_node);
123+ GE_ASSERT_NOTNULL(next_asc_node);
124+ if (node_to_group.count(next_asc_node.get()) > 0UL) {
125+ continue; // 已编组的节点不再处理
126+ }
127+ const int64_t next_axis = next_asc_node->attr.sched.loop_axis;
128+ if ((next_axis == group_axis) || (next_axis == af::kIdNone)) {
129+ node_to_group.emplace(next_asc_node.get(), group_id);
130+ loop_groups[group_id].nodes.emplace_back(next_asc_node.get());
131+ fill_queue.emplace(next_asc_node.get());
132+ } else {
133+ seed_queue.emplace(LoopGroupSeed{next_asc_node.get(), next_axis, group_id}); // 不同轴边界节点作为新组种子
134+ }
135+ }
136+ return af::SUCCESS;
137+}
138+ 
139+// 从种子沿输入/输出方向BFS填充group_id分组，直至无新的同轴或-1邻接节点可吸收
140+Status FillLoopGroupFromSeed(const LoopGroupSeed &seed, const size_t group_id, std::vector<LoopGroup> &loop_groups,
141+ std::unordered_map<af::Node *, size_t> &node_to_group,
142+ std::queue<LoopGroupSeed> &seed_queue) {
143+ std::queue<af::AscNode *> fill_queue{{seed.node}}; // 当前组内待向输入/输出方向扩散的节点
144+ while (!fill_queue.empty()) {
145+ af::AscNode *cur = fill_queue.front();
146+ fill_queue.pop();
147+ // 向输出方向扩散
148+ GE_ASSERT_SUCCESS(SpreadLoopGroupNeighbors(cur->GetOutDataNodes(), seed.loop_axis, group_id, loop_groups,
149+ node_to_group, fill_queue, seed_queue));
150+ // 向输入方向扩散
151+ GE_ASSERT_SUCCESS(SpreadLoopGroupNeighbors(cur->GetInDataNodes(), seed.loop_axis, group_id, loop_groups,
152+ node_to_group, fill_queue, seed_queue));
153+ }
154+ return af::SUCCESS;
155+}
96} // namespace156} // namespace
97 157 
98std::vector<af::AxisId> ScheduleUtils::CalcReduceAxes(const std::vector<af::Expression> &src_strides,158std::vector<af::AxisId> ScheduleUtils::CalcReduceAxes(const std::vector<af::Expression> &src_strides,
@@ -349,6 +409,73 @@ bool ScheduleUtils::IsTailAxisAlignedBy(const af::AscNodePtr &node, const uint32
349 return GetTailAxisDataSize(node, size) && size % align_bytes == 0;409 return GetTailAxisDataSize(node, size) && size % align_bytes == 0;
350}410}
351 411 
412+// 判断图是否满足按循环轴分组的前置条件，成立条件：
413+// 1. 所有reduce节点的loop_axis均已赋值（存在未赋值说明尚未经过AutoScheduler，直接不成立）；
414+// 2. 全图有效loop_axis种类数 > 1（-1不计入，单一循环轴无需分组）。
415+// 其中reduce节点数量 > 1 的检查不是必须条件，与条件2重复，保留只是为了控制影响范围。
416+bool ScheduleUtils::IsNeedLoopGrouping(const af::AscGraph &graph) {
417+ size_t reduce_cnt = 0U;
418+ std::unordered_set<int64_t> unique_loop_axes;
419+ for (const auto &node : graph.GetAllNodes()) {
420+ const int64_t loop_axis = node->attr.sched.loop_axis;
421+ if (loop_axis != af::kIdNone) {
422+ unique_loop_axes.insert(loop_axis);
423+ }
424+ if (!IsReduce(node)) {
425+ continue;
426+ }
427+ ++reduce_cnt;
428+ if (loop_axis == af::kIdNone) {
429+ // 存在loop_axis未赋值的reduce节点，说明尚未经过AutoScheduler，不进入该方案
430+ return false;
431+ }
432+ }
433+ // reduce节点数量>1，且全图有效loop_axis种类数>1（-1不计入），才需要按循环轴分组
434+ return (reduce_cnt > 1U) && (unique_loop_axes.size() > 1U);
435+}
436+ 
437+// 两级BFS编组：外层seed_queue按发现顺序处理边界种子，种子记录发现轮次（即当时正在填充的分组编号）
438+// 与loop_axis；未编组种子出队时，若同轮次中已创建相同loop_axis的分组则并入该分组（不新增编号），
439+// 否则开创一个新分组（编号即loop_groups下标）；内层FillLoopGroupFromSeed从种子沿输入/输出方向扩散，
440+// 吸收loop_axis与本组相同或为-1的邻接节点（-1节点并入当前组并继续扩散，不再阻断）。
441+// 每个节点至多编组一次、每条边至多访问两次，整体复杂度O(V+E)。
442+Status ScheduleUtils::BuildLoopGroups(const af::AscGraph &graph, std::vector<LoopGroup> &loop_groups,
443+ std::unordered_map<af::Node *, size_t> &node_to_group) {
444+ loop_groups.clear();
445+ node_to_group.clear();
446+ af::AscNode *first_reduce = FindFirstReduce(graph);
447+ GE_ASSERT_NOTNULL(first_reduce);
448+ GE_ASSERT_TRUE(first_reduce->attr.sched.loop_axis != af::kIdNone, "The loop_axis of reduce node[%s] is not assigned.",
449+ first_reduce->GetNamePtr());
450+ 
451+ // 同轮次同loop_axis的种子归入同一分组：(发现轮次, loop_axis) -> 分组编号
452+ std::map<std::pair<size_t, int64_t>, size_t> seed_key_to_group;
453+ std::queue<LoopGroupSeed> seed_queue{{LoopGroupSeed{first_reduce, first_reduce->attr.sched.loop_axis, 0U}}};
454+ while (!seed_queue.empty()) {
455+ const LoopGroupSeed seed = seed_queue.front();
456+ seed_queue.pop();
457+ if (node_to_group.count(seed.node) > 0UL) {
458+ continue; // 已编组的节点不再处理
459+ }
460+ size_t group_id = loop_groups.size();
461+ const auto seed_key = std::make_pair(seed.round, seed.loop_axis);
462+ const auto existing_group = seed_key_to_group.find(seed_key);
463+ if (existing_group == seed_key_to_group.cend()) {
464+ loop_groups.emplace_back(LoopGroup{seed.loop_axis, {seed.node}});
465+ seed_key_to_group.emplace(seed_key, group_id);
466+ } else {
467+ group_id = existing_group->second; // 同轮次同轴的种子并入已建分组，不新增编号
468+ loop_groups[group_id].nodes.emplace_back(seed.node);
469+ }
470+ node_to_group.emplace(seed.node, group_id);
471+ GE_ASSERT_SUCCESS(FillLoopGroupFromSeed(seed, group_id, loop_groups, node_to_group, seed_queue));
472+ GELOGD("Build loop group[%zu]: loop_axis[%ld], node cnt[%zu].", group_id, seed.loop_axis,
473+ loop_groups[group_id].nodes.size());
474+ }
475+ GELOGD("Build %zu loop groups in graph[%s].", loop_groups.size(), graph.GetName().c_str());
476+ return af::SUCCESS;
477+}
478+ 
352Status ScheduleUtils::TopologicalSorting(af::AscGraph &graph, bool use_rdfs_v2) {479Status ScheduleUtils::TopologicalSorting(af::AscGraph &graph, bool use_rdfs_v2) {
353 auto compute_graph = af::AscGraphUtils::GetComputeGraph(graph);480 auto compute_graph = af::AscGraphUtils::GetComputeGraph(graph);
354 GE_ASSERT_NOTNULL(compute_graph);481 GE_ASSERT_NOTNULL(compute_graph);
@@ -360,19 +487,39 @@ Status ScheduleUtils::TopologicalSorting(af::AscGraph &graph, bool use_rdfs_v2)
360 return af::SUCCESS;487 return af::SUCCESS;
361 }488 }
362 489 
490+ if (IsNeedLoopGrouping(graph)) {
491+ GELOGI("Graph [%s] will be sorted with loop group rule.", graph.GetName().c_str());
492+ std::vector<LoopGroup> loop_groups;
493+ std::unordered_map<af::Node *, size_t> node_to_group;
494+ GE_ASSERT_SUCCESS(BuildLoopGroups(graph, loop_groups, node_to_group));
495+ const auto func = [&node_to_group](const af::NodePtr &node1, const af::NodePtr &node2) -> bool {
496+ auto it1 = node_to_group.find(node1.get());
497+ auto it2 = node_to_group.find(node2.get());
498+ const bool both_grouped = (it1 != node_to_group.cend()) && (it2 != node_to_group.cend());
499+ // 组间按for编号从小到大排（未入组节点不参与组序，退化为按topo序比较）
500+ if (both_grouped && it1->second != it2->second) {
501+ return it1->second < it2->second;
502+ }
503+ return node1->GetOpDescBarePtr()->GetId() < node2->GetOpDescBarePtr()->GetId(); // 其余按topo序
504+ };
505+ compute_graph->TopologicalSorting(func);
506+ return af::SUCCESS;
507+ }
508+ 
363 GELOGI("Graph [%s] will be sorted with a specific rule.", graph.GetName().c_str());509 GELOGI("Graph [%s] will be sorted with a specific rule.", graph.GetName().c_str());
364- std::unordered_set<af::Node *> reduce_sequences;510+ std::unordered_set<af::Node *> post_reduce_nodes;
365 for (const auto &node : graph.GetAllNodes()) {511 for (const auto &node : graph.GetAllNodes()) {
366 if (IsReduce(node)) {512 if (IsReduce(node)) {
367- GE_ASSERT_SUCCESS(FindNodeSequence(node.get(), reduce_sequences));513+ GE_ASSERT_SUCCESS(CollectPostReduceNodes(node.get(), post_reduce_nodes));
514+ break;
368 }515 }
369 }516 }
370- const auto func = [&reduce_sequences](const af::NodePtr &node1, const af::NodePtr &node2) -> bool {517+ const auto func = [&post_reduce_nodes](const af::NodePtr &node1, const af::NodePtr &node2) -> bool {
371- bool is_node1_in_reduce_seq = reduce_sequences.find(node1.get()) != reduce_sequences.end();518+ bool is_node1_post_reduce = post_reduce_nodes.find(node1.get()) != post_reduce_nodes.end();
372- bool is_node2_in_reduce_seq = reduce_sequences.find(node2.get()) != reduce_sequences.end();519+ bool is_node2_post_reduce = post_reduce_nodes.find(node2.get()) != post_reduce_nodes.end();
373- if (is_node1_in_reduce_seq && !is_node2_in_reduce_seq) {520+ if (is_node1_post_reduce && !is_node2_post_reduce) {
374 return false;521 return false;
375- } else if (!is_node1_in_reduce_seq && is_node2_in_reduce_seq) {522+ } else if (!is_node1_post_reduce && is_node2_post_reduce) {
376 return true;523 return true;
377 } else {524 } else {
378 return node1->GetOpDescBarePtr()->GetId() < node2->GetOpDescBarePtr()->GetId();525 return node1->GetOpDescBarePtr()->GetId() < node2->GetOpDescBarePtr()->GetId();
Mautofuse/optimize/schedule_utils.h+31-0
@@ -11,6 +11,8 @@
11#ifndef OPTIMIZE_SCHEDULE_UTILS_H_11#ifndef OPTIMIZE_SCHEDULE_UTILS_H_
12#define OPTIMIZE_SCHEDULE_UTILS_H_12#define OPTIMIZE_SCHEDULE_UTILS_H_
13 13 
14+#include <unordered_map>
15+ 
14#include "ascendc_ir/ascendc_ir_core/ascendc_ir_def.h"16#include "ascendc_ir/ascendc_ir_core/ascendc_ir_def.h"
15#include "graph/symbolizer/symbolic_utils.h"17#include "graph/symbolizer/symbolic_utils.h"
16#include "asc_graph_utils.h"18#include "asc_graph_utils.h"
@@ -22,6 +24,12 @@
22#include "schedule_result.h"24#include "schedule_result.h"
23 25 
24namespace optimize {26namespace optimize {
27+// 按loop_axis连通区域划分出的for循环节点分组，loop_groups中的下标即for循环编号
28+struct LoopGroup {
29+ int64_t loop_axis = af::kIdNone; // 组内统一的循环轴ID
30+ std::vector<af::Node *> nodes; // 组内节点，首个节点为该组的扩散种子
31+};
32+ 
25class ScheduleUtils {33class ScheduleUtils {
26 public:34 public:
27 static af::ComputeType GetComputeType(const af::AscNodePtr &node) {35 static af::ComputeType GetComputeType(const af::AscNodePtr &node) {
@@ -31,6 +39,29 @@ class ScheduleUtils {
31 // 后端默认采用逆dfs的拓扑排序方式39 // 后端默认采用逆dfs的拓扑排序方式
32 static Status TopologicalSorting(af::AscGraph &graph, bool use_rdfs_v2 = false);40 static Status TopologicalSorting(af::AscGraph &graph, bool use_rdfs_v2 = false);
33 41 
42+ // 判断图是否满足按循环轴分组的前置条件：
43+ // 1. 所有reduce节点的loop_axis均已赋值（保证该处理在AutoScheduler之后才生效）；
44+ // 2. 全图有效loop_axis种类数 > 1（-1不计入，单一循环轴无需分组）。
45+ // reduce节点数量 > 1 的检查与条件2重复，非必须，仅用于控制影响范围。
46+ static bool IsNeedLoopGrouping(const af::AscGraph &graph);
47+ 
48+ // 按loop_axis连通区域将图中节点划分为多个for循环分组，供TopologicalSorting使用：
49+ // 同循环节点在拓扑排序中相邻排列，避免被其它循环轴的节点穿插而硬拆成多个循环。
50+ // 以当前拓扑序的第一个reduce节点为初始种子，从种子沿输入/输出方向BFS扩散，
51+ // loop_axis相同或为-1的连通节点归入同一分组（-1节点并入当前组并继续扩散）；
52+ // loop_axis不同的有效节点作为新分组种子继续扩散（种子记录发现轮次与loop_axis），
53+ // 同一轮扩散中发现的同轴种子归入同一分组，不同轮发现的同轴区域拆分为不同分组。
54+ // 分组编号从0开始，即loop_groups的下标。
55+ // 规则约束：如下
56+ // 1. 同一分组内有效节点的loop_axis必然相同，loop_axis为-1的节点按所在组的轴参与排序；
57+ // 同轮发现的同轴种子（即使互不连通）归入同一分组，不同轮发现的同轴区域拆分为不同编号的分组；
58+ // 2. loop_axis为-1的节点（如Data/Output等Buffer节点）并入当前分组并继续扩散，不作为新分组种子，
59+ // 无法从首个reduce扩散到达的节点不会出现在分组结果中；
60+ // 3. 调用方需保证存在loop_axis已赋值的reduce节点（建议先经IsNeedLoopGrouping检查）。
61+ // 输出：loop_groups存储各分组及其轴，node_to_group记录节点到分组编号的映射，未入组节点不在其中。
62+ static Status BuildLoopGroups(const af::AscGraph &graph, std::vector<LoopGroup> &loop_groups,
63+ std::unordered_map<af::Node *, size_t> &node_to_group);
64+ 
34 static bool IsElewise(const af::AscNodePtr &node) {65 static bool IsElewise(const af::AscNodePtr &node) {
35 return node->attr.api.compute_type == af::ComputeType::kComputeElewise;66 return node->attr.api.compute_type == af::ComputeType::kComputeElewise;
36 }67 }
Mautofuse/tests/st/optimize/test_reduce.cpp+34-0
@@ -355,6 +355,26 @@ AscGraph ConstructNormStruct4Elewise4ReduceMultipleCitationsMulOut(const std::st
355 .Build();355 .Build();
356}356}
357 357 
358+// 双reduce水平融合 + 多引用结构（load0同输入喂abs0/abs1两路，两路reduce规约轴相同）：
359+// 走全流程Optimize验证loop group排序路径（BufQueAllocator在AutoScheduler后以v2模式重排序）
360+AscGraph ConstructTwoReduceSameAxisHorizontalFusion(const std::string &name) {
361+ return AscGraphBuilder(name)
362+ .Loops({Sym(128), Sym(64)})
363+ .Data("data0", 0)
364+ .Load("load0", "data0")
365+ .Abs("abs0", "load0")
366+ .Sum("sum0", "abs0", {1})
367+ .Relu("relu0", "sum0")
368+ .Store("store0", "relu0")
369+ .Output("output0", "store0", 0, af::DT_FLOAT)
370+ .Abs("abs1", "load0")
371+ .Sum("sum1", "abs1", {1})
372+ .Relu("relu1", "sum1")
373+ .Store("store1", "relu1")
374+ .Output("output1", "store1", 1, af::DT_FLOAT)
375+ .Build();
376+}
377+ 
358namespace optimize {378namespace optimize {
359class OptimizerReduceSt : public ::testing::Test {379class OptimizerReduceSt : public ::testing::Test {
360 protected:380 protected:
@@ -574,4 +594,18 @@ TEST_F(OptimizerReduceSt, TestReduce_Three_Elewise_Store_Multi_Citation_Multi_Ou
574 ASSERT_EQ(fused_scheduled_result.node_idx_to_scheduled_results.size(), 1UL);594 ASSERT_EQ(fused_scheduled_result.node_idx_to_scheduled_results.size(), 1UL);
575 ASSERT_EQ(fused_scheduled_result.node_idx_to_scheduled_results[0UL].size(), 2UL);595 ASSERT_EQ(fused_scheduled_result.node_idx_to_scheduled_results[0UL].size(), 2UL);
576}596}
597+ 
598+// 双reduce水平融合（同输入不同输出、规约轴相同）场景走全流程Optimize，
599+// 验证loop group排序分支与调度/内存分配流程兼容
600+TEST_F(OptimizerReduceSt, TestReduce_LoopGroupSortTwoReduceSameAxis) {
601+ auto graph = ConstructTwoReduceSameAxisHorizontalFusion("reduce_loop_group_sort");
602+ ::ascir::FusedScheduledResult fused_scheduled_result;
603+ optimize::Optimizer optimizer(optimize::OptimizerOptions{});
604+ Status res = optimizer.Optimize(graph, fused_scheduled_result);
605+ EXPECT_EQ(res, af::SUCCESS);
606+ ASSERT_EQ(fused_scheduled_result.node_idx_to_scheduled_results.size(), 1UL);
607+ ASSERT_FALSE(fused_scheduled_result.node_idx_to_scheduled_results[0].empty());
608+ ASSERT_FALSE(fused_scheduled_result.node_idx_to_scheduled_results[0][0].schedule_groups.empty());
609+ ASSERT_FALSE(fused_scheduled_result.node_idx_to_scheduled_results[0][0].schedule_groups[0].impl_graphs.empty());
610+}
577} // namespace optimize611} // namespace optimize
Aautofuse/tests/ut/optimize/test_loop_group_ut.cpp+555-0
@@ -0,0 +1,555 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <gtest/gtest.h>
12+ 
13+#include "ascendc_ir.h"
14+#include "ascir_ops.h"
15+#include "graph/ascendc_ir/utils/asc_graph_utils.h"
16+#include "schedule_utils.h"
17+ 
18+namespace optimize {
19+namespace {
20+// 水平融合双reduce（真实业务场景）：load0同输入分叉两路，两路reduce规约轴相同、各自输出。
21+// loop_axis赋值模拟真实调度分布：reduce及其前置节点同轴(B)，reduce后继节点异轴(A)。
22+// data -> load -> {abs0(B) -> max1(B,reduce) -> relu0(A) -> store0 -> output0,
23+// abs1(B) -> max2(B,reduce) -> relu1(A) -> store1 -> output1}
24+af::AscGraph CreateHorizontalFusionTwoReduce(const char *name) {
25+ af::AscGraph graph(name);
26+ af::ascir_op::Data data("data", graph);
27+ data.attr.api.compute_type = af::ComputeType::kComputeInvalid;
28+ data.attr.api.type = af::ApiType::kAPITypeBuffer;
29+ 
30+ af::ascir_op::Load load("load");
31+ load.x = data.y;
32+ load.attr.api.compute_type = af::ComputeType::kComputeLoad;
33+ load.attr.api.type = af::ApiType::kAPITypeCompute;
34+ 
35+ af::ascir_op::Abs abs0("abs0");
36+ abs0.x = load.y;
37+ abs0.attr.api.compute_type = af::ComputeType::kComputeElewise;
38+ abs0.attr.api.type = af::ApiType::kAPITypeCompute;
39+ 
40+ af::ascir_op::Abs abs1("abs1");
41+ abs1.x = load.y;
42+ abs1.attr.api.compute_type = af::ComputeType::kComputeElewise;
43+ abs1.attr.api.type = af::ApiType::kAPITypeCompute;
44+ 
45+ af::ascir_op::Max max1("max1");
46+ max1.x = abs0.y;
47+ max1.attr.api.compute_type = af::ComputeType::kComputeReduce;
48+ max1.attr.api.type = af::ApiType::kAPITypeCompute;
49+ max1.y.dtype = ge::DT_FLOAT;
50+ 
51+ af::ascir_op::Max max2("max2");
52+ max2.x = abs1.y;
53+ max2.attr.api.compute_type = af::ComputeType::kComputeReduce;
54+ max2.attr.api.type = af::ApiType::kAPITypeCompute;
55+ max2.y.dtype = ge::DT_FLOAT;
56+ 
57+ af::ascir_op::Relu relu0("relu0");
58+ relu0.x = max1.y;
59+ relu0.attr.api.compute_type = af::ComputeType::kComputeElewise;
60+ relu0.attr.api.type = af::ApiType::kAPITypeCompute;
61+ 
62+ af::ascir_op::Relu relu1("relu1");
63+ relu1.x = max2.y;
64+ relu1.attr.api.compute_type = af::ComputeType::kComputeElewise;
65+ relu1.attr.api.type = af::ApiType::kAPITypeCompute;
66+ 
67+ af::ascir_op::Store store0("store0");
68+ store0.x = relu0.y;
69+ store0.attr.api.compute_type = af::ComputeType::kComputeStore;
70+ store0.attr.api.type = af::ApiType::kAPITypeCompute;
71+ store0.y.dtype = ge::DT_FLOAT;
72+ 
73+ af::ascir_op::Output output0("output0");
74+ output0.x = store0.y;
75+ output0.attr.api.compute_type = af::ComputeType::kComputeInvalid;
76+ output0.attr.api.type = af::ApiType::kAPITypeBuffer;
77+ output0.y.dtype = ge::DT_FLOAT;
78+ 
79+ af::ascir_op::Store store1("store1");
80+ store1.x = relu1.y;
81+ store1.attr.api.compute_type = af::ComputeType::kComputeStore;
82+ store1.attr.api.type = af::ApiType::kAPITypeCompute;
83+ store1.y.dtype = ge::DT_FLOAT;
84+ 
85+ af::ascir_op::Output output1("output1");
86+ output1.x = store1.y;
87+ output1.attr.api.compute_type = af::ComputeType::kComputeInvalid;
88+ output1.attr.api.type = af::ApiType::kAPITypeBuffer;
89+ output1.y.dtype = ge::DT_FLOAT;
90+ return graph;
91+}
92+ 
93+// 同轴区域被异轴分支隔断：max1(A)的输出分叉，relu0支与max1同轴连通，
94+// exp0(B)支下游回到A轴形成第二个reduce区域（第2轮发现，与组0同轴但不连通）：
95+// data -> load -> abs0(A) -> max1(A,reduce) -> {relu0(A) -> store0 -> output0,
96+// exp0(B) -> abs2(A) -> max2(A,reduce) -> store1 -> output1}
97+af::AscGraph CreateSameAxisDisconnectedBranch(const char *name) {
98+ af::AscGraph graph(name);
99+ af::ascir_op::Data data("data", graph);
100+ data.attr.api.compute_type = af::ComputeType::kComputeInvalid;
101+ data.attr.api.type = af::ApiType::kAPITypeBuffer;
102+ 
103+ af::ascir_op::Load load("load");
104+ load.x = data.y;
105+ load.attr.api.compute_type = af::ComputeType::kComputeLoad;
106+ load.attr.api.type = af::ApiType::kAPITypeCompute;
107+ 
108+ af::ascir_op::Abs abs0("abs0");
109+ abs0.x = load.y;
110+ abs0.attr.api.compute_type = af::ComputeType::kComputeElewise;
111+ abs0.attr.api.type = af::ApiType::kAPITypeCompute;
112+ 
113+ af::ascir_op::Max max1("max1");
114+ max1.x = abs0.y;
115+ max1.attr.api.compute_type = af::ComputeType::kComputeReduce;
116+ max1.attr.api.type = af::ApiType::kAPITypeCompute;
117+ max1.y.dtype = ge::DT_FLOAT;
118+ 
119+ af::ascir_op::Relu relu0("relu0");
120+ relu0.x = max1.y;
121+ relu0.attr.api.compute_type = af::ComputeType::kComputeElewise;
122+ relu0.attr.api.type = af::ApiType::kAPITypeCompute;
123+ 
124+ af::ascir_op::Store store0("store0");
125+ store0.x = relu0.y;
126+ store0.attr.api.compute_type = af::ComputeType::kComputeStore;
127+ store0.attr.api.type = af::ApiType::kAPITypeCompute;
128+ store0.y.dtype = ge::DT_FLOAT;
129+ 
130+ af::ascir_op::Output output0("output0");
131+ output0.x = store0.y;
132+ output0.attr.api.compute_type = af::ComputeType::kComputeInvalid;
133+ output0.attr.api.type = af::ApiType::kAPITypeBuffer;
134+ output0.y.dtype = ge::DT_FLOAT;
135+ 
136+ af::ascir_op::Exp exp0("exp0");
137+ exp0.x = max1.y;
138+ exp0.attr.api.compute_type = af::ComputeType::kComputeElewise;
139+ exp0.attr.api.type = af::ApiType::kAPITypeCompute;
140+ 
141+ af::ascir_op::Abs abs2("abs2");
142+ abs2.x = exp0.y;
143+ abs2.attr.api.compute_type = af::ComputeType::kComputeElewise;
144+ abs2.attr.api.type = af::ApiType::kAPITypeCompute;
145+ 
146+ af::ascir_op::Max max2("max2");
147+ max2.x = abs2.y;
148+ max2.attr.api.compute_type = af::ComputeType::kComputeReduce;
149+ max2.attr.api.type = af::ApiType::kAPITypeCompute;
150+ max2.y.dtype = ge::DT_FLOAT;
151+ 
152+ af::ascir_op::Store store1("store1");
153+ store1.x = max2.y;
154+ store1.attr.api.compute_type = af::ComputeType::kComputeStore;
155+ store1.attr.api.type = af::ApiType::kAPITypeCompute;
156+ store1.y.dtype = ge::DT_FLOAT;
157+ 
158+ af::ascir_op::Output output1("output1");
159+ output1.x = store1.y;
160+ output1.attr.api.compute_type = af::ComputeType::kComputeInvalid;
161+ output1.attr.api.type = af::ApiType::kAPITypeBuffer;
162+ output1.y.dtype = ge::DT_FLOAT;
163+ return graph;
164+}
165+ 
166+// 独立双reduce链：两路各自独立读输入，无多引用结构，IsNeedFixTopo不触发：
167+// {data0 -> load0 -> abs0(B) -> max1(B,reduce) -> relu0(A) -> store0 -> output0,
168+// data1 -> load1 -> abs1(B) -> max2(B,reduce) -> relu1(A) -> store1 -> output1}
169+af::AscGraph CreateIndependentTwoReduceChains(const char *name) {
170+ af::AscGraph graph(name);
171+ af::ascir_op::Data data0("data0", graph);
172+ data0.attr.api.compute_type = af::ComputeType::kComputeInvalid;
173+ data0.attr.api.type = af::ApiType::kAPITypeBuffer;
174+ 
175+ af::ascir_op::Load load0("load0");
176+ load0.x = data0.y;
177+ load0.attr.api.compute_type = af::ComputeType::kComputeLoad;
178+ load0.attr.api.type = af::ApiType::kAPITypeCompute;
179+ 
180+ af::ascir_op::Abs abs0("abs0");
181+ abs0.x = load0.y;
182+ abs0.attr.api.compute_type = af::ComputeType::kComputeElewise;
183+ abs0.attr.api.type = af::ApiType::kAPITypeCompute;
184+ 
185+ af::ascir_op::Max max1("max1");
186+ max1.x = abs0.y;
187+ max1.attr.api.compute_type = af::ComputeType::kComputeReduce;
188+ max1.attr.api.type = af::ApiType::kAPITypeCompute;
189+ max1.y.dtype = ge::DT_FLOAT;
190+ 
191+ af::ascir_op::Relu relu0("relu0");
192+ relu0.x = max1.y;
193+ relu0.attr.api.compute_type = af::ComputeType::kComputeElewise;
194+ relu0.attr.api.type = af::ApiType::kAPITypeCompute;
195+ 
196+ af::ascir_op::Store store0("store0");
197+ store0.x = relu0.y;
198+ store0.attr.api.compute_type = af::ComputeType::kComputeStore;
199+ store0.attr.api.type = af::ApiType::kAPITypeCompute;
200+ store0.y.dtype = ge::DT_FLOAT;
201+ 
202+ af::ascir_op::Output output0("output0");
203+ output0.x = store0.y;
204+ output0.attr.api.compute_type = af::ComputeType::kComputeInvalid;
205+ output0.attr.api.type = af::ApiType::kAPITypeBuffer;
206+ output0.y.dtype = ge::DT_FLOAT;
207+ 
208+ af::ascir_op::Data data1("data1", graph);
209+ data1.attr.api.compute_type = af::ComputeType::kComputeInvalid;
210+ data1.attr.api.type = af::ApiType::kAPITypeBuffer;
211+ 
212+ af::ascir_op::Load load1("load1");
213+ load1.x = data1.y;
214+ load1.attr.api.compute_type = af::ComputeType::kComputeLoad;
215+ load1.attr.api.type = af::ApiType::kAPITypeCompute;
216+ 
217+ af::ascir_op::Abs abs1("abs1");
218+ abs1.x = load1.y;
219+ abs1.attr.api.compute_type = af::ComputeType::kComputeElewise;
220+ abs1.attr.api.type = af::ApiType::kAPITypeCompute;
221+ 
222+ af::ascir_op::Max max2("max2");
223+ max2.x = abs1.y;
224+ max2.attr.api.compute_type = af::ComputeType::kComputeReduce;
225+ max2.attr.api.type = af::ApiType::kAPITypeCompute;
226+ max2.y.dtype = ge::DT_FLOAT;
227+ 
228+ af::ascir_op::Relu relu1("relu1");
229+ relu1.x = max2.y;
230+ relu1.attr.api.compute_type = af::ComputeType::kComputeElewise;
231+ relu1.attr.api.type = af::ApiType::kAPITypeCompute;
232+ 
233+ af::ascir_op::Store store1("store1");
234+ store1.x = relu1.y;
235+ store1.attr.api.compute_type = af::ComputeType::kComputeStore;
236+ store1.attr.api.type = af::ApiType::kAPITypeCompute;
237+ store1.y.dtype = ge::DT_FLOAT;
238+ 
239+ af::ascir_op::Output output1("output1");
240+ output1.x = store1.y;
241+ output1.attr.api.compute_type = af::ComputeType::kComputeInvalid;
242+ output1.attr.api.type = af::ApiType::kAPITypeBuffer;
243+ output1.y.dtype = ge::DT_FLOAT;
244+ return graph;
245+}
246+ 
247+// 同轮同轴双种子：max1(A)的一轮扩散同时发现exp0(B)/abs1(C)/exp1(B)三个不连通分支种子：
248+// data -> load -> abs0(A) -> max1(A,reduce) -> {exp0(B) -> store0 -> output0,
249+// abs1(C) -> store1 -> output1,
250+// exp1(B) -> store2 -> output2}
251+af::AscGraph CreateSameRoundSameAxisSeeds(const char *name) {
252+ af::AscGraph graph(name);
253+ af::ascir_op::Data data("data", graph);
254+ data.attr.api.compute_type = af::ComputeType::kComputeInvalid;
255+ data.attr.api.type = af::ApiType::kAPITypeBuffer;
256+ 
257+ af::ascir_op::Load load("load");
258+ load.x = data.y;
259+ load.attr.api.compute_type = af::ComputeType::kComputeLoad;
260+ load.attr.api.type = af::ApiType::kAPITypeCompute;
261+ 
262+ af::ascir_op::Abs abs0("abs0");
263+ abs0.x = load.y;
264+ abs0.attr.api.compute_type = af::ComputeType::kComputeElewise;
265+ abs0.attr.api.type = af::ApiType::kAPITypeCompute;
266+ 
267+ af::ascir_op::Max max1("max1");
268+ max1.x = abs0.y;
269+ max1.attr.api.compute_type = af::ComputeType::kComputeReduce;
270+ max1.attr.api.type = af::ApiType::kAPITypeCompute;
271+ max1.y.dtype = ge::DT_FLOAT;
272+ 
273+ af::ascir_op::Exp exp0("exp0");
274+ exp0.x = max1.y;
275+ exp0.attr.api.compute_type = af::ComputeType::kComputeElewise;
276+ exp0.attr.api.type = af::ApiType::kAPITypeCompute;
277+ 
278+ af::ascir_op::Abs abs1("abs1");
279+ abs1.x = max1.y;
280+ abs1.attr.api.compute_type = af::ComputeType::kComputeElewise;
281+ abs1.attr.api.type = af::ApiType::kAPITypeCompute;
282+ 
283+ af::ascir_op::Exp exp1("exp1");
284+ exp1.x = max1.y;
285+ exp1.attr.api.compute_type = af::ComputeType::kComputeElewise;
286+ exp1.attr.api.type = af::ApiType::kAPITypeCompute;
287+ 
288+ af::ascir_op::Store store0("store0");
289+ store0.x = exp0.y;
290+ store0.attr.api.compute_type = af::ComputeType::kComputeStore;
291+ store0.attr.api.type = af::ApiType::kAPITypeCompute;
292+ store0.y.dtype = ge::DT_FLOAT;
293+ 
294+ af::ascir_op::Output output0("output0");
295+ output0.x = store0.y;
296+ output0.attr.api.compute_type = af::ComputeType::kComputeInvalid;
297+ output0.attr.api.type = af::ApiType::kAPITypeBuffer;
298+ output0.y.dtype = ge::DT_FLOAT;
299+ 
300+ af::ascir_op::Store store1("store1");
301+ store1.x = abs1.y;
302+ store1.attr.api.compute_type = af::ComputeType::kComputeStore;
303+ store1.attr.api.type = af::ApiType::kAPITypeCompute;
304+ store1.y.dtype = ge::DT_FLOAT;
305+ 
306+ af::ascir_op::Output output1("output1");
307+ output1.x = store1.y;
308+ output1.attr.api.compute_type = af::ComputeType::kComputeInvalid;
309+ output1.attr.api.type = af::ApiType::kAPITypeBuffer;
310+ output1.y.dtype = ge::DT_FLOAT;
311+ 
312+ af::ascir_op::Store store2("store2");
313+ store2.x = exp1.y;
314+ store2.attr.api.compute_type = af::ComputeType::kComputeStore;
315+ store2.attr.api.type = af::ApiType::kAPITypeCompute;
316+ store2.y.dtype = ge::DT_FLOAT;
317+ 
318+ af::ascir_op::Output output2("output2");
319+ output2.x = store2.y;
320+ output2.attr.api.compute_type = af::ComputeType::kComputeInvalid;
321+ output2.attr.api.type = af::ApiType::kAPITypeBuffer;
322+ output2.y.dtype = ge::DT_FLOAT;
323+ return graph;
324+}
325+ 
326+struct LoopAxisIds {
327+ int64_t axis_a;
328+ int64_t axis_b;
329+};
330+ 
331+LoopAxisIds CreateLoopAxisIds(af::AscGraph &graph) {
332+ return {graph.CreateAxis("z_a", af::Symbol(4)).id, graph.CreateAxis("z_b", af::Symbol(8)).id};
333+}
334+ 
335+void SetNodeLoopAxis(af::AscGraph &graph, const char *node_name, const int64_t loop_axis) {
336+ auto node = graph.FindNode(node_name);
337+ ASSERT_NE(node, nullptr);
338+ node->attr.sched.loop_axis = loop_axis;
339+}
340+ 
341+int64_t GetNodeId(af::AscGraph &graph, const char *node_name) {
342+ auto node = graph.FindNode(node_name);
343+ EXPECT_NE(node, nullptr);
344+ return (node == nullptr) ? -1 : node->GetOpDescBarePtr()->GetId();
345+}
346+ 
347+af::Node *GetNodePtr(af::AscGraph &graph, const char *node_name) {
348+ auto node = graph.FindNode(node_name);
349+ EXPECT_NE(node, nullptr);
350+ return (node == nullptr) ? nullptr : node.get();
351+}
352+ 
353+// 为水平融合图赋loop_axis：reduce及其前置节点为axis_b，reduce后继节点为axis_a
354+void SetHorizontalFusionLoopAxis(af::AscGraph &graph, const LoopAxisIds &axes) {
355+ SetNodeLoopAxis(graph, "load", axes.axis_b);
356+ SetNodeLoopAxis(graph, "abs0", axes.axis_b);
357+ SetNodeLoopAxis(graph, "abs1", axes.axis_b);
358+ SetNodeLoopAxis(graph, "max1", axes.axis_b);
359+ SetNodeLoopAxis(graph, "max2", axes.axis_b);
360+ SetNodeLoopAxis(graph, "relu0", axes.axis_a);
361+ SetNodeLoopAxis(graph, "relu1", axes.axis_a);
362+}
363+ 
364+TEST(LoopGroupTest, IsNeedLoopGroupingNotNeedSingleReduce) {
365+ af::AscGraph graph("single_reduce");
366+ af::ascir_op::Data data("data", graph);
367+ data.attr.api.compute_type = af::ComputeType::kComputeInvalid;
368+ data.attr.api.type = af::ApiType::kAPITypeBuffer;
369+ 
370+ af::ascir_op::Load load("load");
371+ load.x = data.y;
372+ load.attr.api.compute_type = af::ComputeType::kComputeLoad;
373+ load.attr.api.type = af::ApiType::kAPITypeCompute;
374+ 
375+ af::ascir_op::Abs abs0("abs0");
376+ abs0.x = load.y;
377+ abs0.attr.api.compute_type = af::ComputeType::kComputeElewise;
378+ abs0.attr.api.type = af::ApiType::kAPITypeCompute;
379+ 
380+ af::ascir_op::Max max1("max1");
381+ max1.x = abs0.y;
382+ max1.attr.api.compute_type = af::ComputeType::kComputeReduce;
383+ max1.attr.api.type = af::ApiType::kAPITypeCompute;
384+ max1.y.dtype = ge::DT_FLOAT;
385+ 
386+ af::ascir_op::Store store("store");
387+ store.x = max1.y;
388+ store.attr.api.compute_type = af::ComputeType::kComputeStore;
389+ store.attr.api.type = af::ApiType::kAPITypeCompute;
390+ store.y.dtype = ge::DT_FLOAT;
391+ 
392+ const auto axes = CreateLoopAxisIds(graph);
393+ SetNodeLoopAxis(graph, "max1", axes.axis_a);
394+ EXPECT_FALSE(ScheduleUtils::IsNeedLoopGrouping(graph));
395+}
396+ 
397+TEST(LoopGroupTest, IsNeedLoopGroupingNotNeedUnassignedLoopAxis) {
398+ auto graph = CreateHorizontalFusionTwoReduce("two_reduce_one_unassigned");
399+ const auto axes = CreateLoopAxisIds(graph);
400+ // max2保持-1（未经过AutoScheduler），存在未赋值reduce节点，不应进入分组
401+ SetNodeLoopAxis(graph, "load", axes.axis_b);
402+ SetNodeLoopAxis(graph, "abs0", axes.axis_b);
403+ SetNodeLoopAxis(graph, "abs1", axes.axis_b);
404+ SetNodeLoopAxis(graph, "max1", axes.axis_b);
405+ SetNodeLoopAxis(graph, "relu0", axes.axis_a);
406+ SetNodeLoopAxis(graph, "relu1", axes.axis_a);
407+ EXPECT_FALSE(ScheduleUtils::IsNeedLoopGrouping(graph));
408+}
409+ 
410+TEST(LoopGroupTest, IsNeedLoopGroupingNeedMultipleUniqueAxis) {
411+ auto graph = CreateHorizontalFusionTwoReduce("two_reduce_valid");
412+ const auto axes = CreateLoopAxisIds(graph);
413+ // 双reduce均已赋值，且全图有效loop_axis种类数>1（reduce链B轴+后继A轴）
414+ SetHorizontalFusionLoopAxis(graph, axes);
415+ EXPECT_TRUE(ScheduleUtils::IsNeedLoopGrouping(graph));
416+}
417+ 
418+TEST(LoopGroupTest, IsNeedLoopGroupingNotNeedSingleUniqueAxis) {
419+ auto graph = CreateHorizontalFusionTwoReduce("two_reduce_single_axis");
420+ const auto axes = CreateLoopAxisIds(graph);
421+ // 全图仅一种有效loop_axis，无需分组
422+ SetHorizontalFusionLoopAxis(graph, axes);
423+ SetNodeLoopAxis(graph, "relu0", axes.axis_b);
424+ SetNodeLoopAxis(graph, "relu1", axes.axis_b);
425+ EXPECT_FALSE(ScheduleUtils::IsNeedLoopGrouping(graph));
426+}
427+ 
428+TEST(LoopGroupTest, BuildLoopGroupsHorizontalFusionReduceSameGroup) {
429+ auto graph = CreateHorizontalFusionTwoReduce("horizontal_fusion");
430+ const auto axes = CreateLoopAxisIds(graph);
431+ SetHorizontalFusionLoopAxis(graph, axes);
432+ 
433+ std::vector<LoopGroup> loop_groups;
434+ std::unordered_map<af::Node *, size_t> node_to_group;
435+ ASSERT_EQ(ScheduleUtils::BuildLoopGroups(graph, loop_groups, node_to_group), af::SUCCESS);
436+ 
437+ // 组0（B轴）：两路reduce经共同输入load连通吸收进同一分组，data/load等-1节点随组扩散并入
438+ ASSERT_EQ(loop_groups.size(), 2UL);
439+ EXPECT_EQ(loop_groups[0].loop_axis, axes.axis_b);
440+ EXPECT_EQ(loop_groups[0].nodes.size(), 6UL);
441+ EXPECT_EQ(loop_groups[0].nodes[0], GetNodePtr(graph, "max1")); // 首个reduce为种子
442+ const size_t reduce_group = node_to_group[GetNodePtr(graph, "max1")];
443+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "max2")], reduce_group); // 水平融合：双reduce同组
444+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "abs0")], reduce_group);
445+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "abs1")], reduce_group);
446+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "load")], reduce_group);
447+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "data")], reduce_group);
448+ 
449+ // 组1（A轴）：两路reduce后继同轮同轴种子合并为一组，不被拆成两个循环
450+ EXPECT_EQ(loop_groups[1].loop_axis, axes.axis_a);
451+ EXPECT_EQ(loop_groups[1].nodes.size(), 6UL);
452+ const size_t post_group = node_to_group[GetNodePtr(graph, "relu0")];
453+ EXPECT_NE(post_group, reduce_group);
454+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "relu1")], post_group);
455+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "store0")], post_group);
456+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "store1")], post_group);
457+}
458+ 
459+TEST(LoopGroupTest, BuildLoopGroupsSameAxisDisconnectedSplitGroups) {
460+ auto graph = CreateSameAxisDisconnectedBranch("same_axis_disconnected");
461+ const auto axes = CreateLoopAxisIds(graph);
462+ SetNodeLoopAxis(graph, "abs0", axes.axis_a);
463+ SetNodeLoopAxis(graph, "max1", axes.axis_a);
464+ SetNodeLoopAxis(graph, "relu0", axes.axis_a);
465+ SetNodeLoopAxis(graph, "exp0", axes.axis_b);
466+ SetNodeLoopAxis(graph, "abs2", axes.axis_a);
467+ SetNodeLoopAxis(graph, "max2", axes.axis_a);
468+ 
469+ std::vector<LoopGroup> loop_groups;
470+ std::unordered_map<af::Node *, size_t> node_to_group;
471+ ASSERT_EQ(ScheduleUtils::BuildLoopGroups(graph, loop_groups, node_to_group), af::SUCCESS);
472+ 
473+ // 同轴但不连通的区域拆分为不同编号：组0和组2同为axis_a，组1为axis_b
474+ // （abs2/max2在exp0分支下游第2轮发现，与组0同轴但不连通，不并入组0）
475+ ASSERT_EQ(loop_groups.size(), 3UL);
476+ EXPECT_EQ(loop_groups[0].loop_axis, axes.axis_a);
477+ EXPECT_EQ(loop_groups[1].loop_axis, axes.axis_b);
478+ EXPECT_EQ(loop_groups[2].loop_axis, axes.axis_a);
479+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "max1")], 0UL);
480+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "relu0")], 0UL);
481+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "exp0")], 1UL);
482+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "abs2")], 2UL);
483+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "max2")], 2UL);
484+}
485+ 
486+TEST(LoopGroupTest, BuildLoopGroupsSameRoundSameAxisSeedsMerged) {
487+ auto graph = CreateSameRoundSameAxisSeeds("same_round_same_axis_seeds");
488+ const auto axes = CreateLoopAxisIds(graph);
489+ const int64_t axis_c = graph.CreateAxis("z_c", af::Symbol(2)).id;
490+ SetNodeLoopAxis(graph, "abs0", axes.axis_a);
491+ SetNodeLoopAxis(graph, "max1", axes.axis_a);
492+ SetNodeLoopAxis(graph, "exp0", axes.axis_b);
493+ SetNodeLoopAxis(graph, "abs1", axis_c);
494+ SetNodeLoopAxis(graph, "exp1", axes.axis_b);
495+ 
496+ std::vector<LoopGroup> loop_groups;
497+ std::unordered_map<af::Node *, size_t> node_to_group;
498+ ASSERT_EQ(ScheduleUtils::BuildLoopGroups(graph, loop_groups, node_to_group), af::SUCCESS);
499+ 
500+ // 组0的一轮扩散发现exp0/abs1/exp1三个种子：同轮同轴的exp0与exp1归入同一分组，
501+ // 不因abs1(C)分组创建在中间而被拆成两个B分组
502+ ASSERT_EQ(loop_groups.size(), 3UL);
503+ const size_t b_group = node_to_group[GetNodePtr(graph, "exp0")];
504+ EXPECT_EQ(loop_groups[b_group].loop_axis, axes.axis_b);
505+ EXPECT_EQ(loop_groups[b_group].nodes.size(), 6UL);
506+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "exp1")], b_group);
507+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "store0")], b_group);
508+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "store2")], b_group);
509+ EXPECT_NE(node_to_group[GetNodePtr(graph, "abs1")], b_group);
510+ EXPECT_EQ(node_to_group[GetNodePtr(graph, "max1")], 0UL);
511+}
512+ 
513+TEST(LoopGroupTest, TopologicalSortingLoopGroupRuleApplied) {
514+ auto graph = CreateHorizontalFusionTwoReduce("loop_group_sort");
515+ const auto axes = CreateLoopAxisIds(graph);
516+ SetHorizontalFusionLoopAxis(graph, axes);
517+ // 前置条件：load多引用触发IsNeedFixTopo，双reduce有效且多轴触发IsNeedLoopGrouping
518+ ASSERT_TRUE(ScheduleUtils::IsNeedLoopGrouping(graph));
519+ 
520+ ASSERT_EQ(ScheduleUtils::TopologicalSorting(graph, true), af::SUCCESS);
521+ 
522+ // 组内保持依赖topo序
523+ EXPECT_TRUE(GetNodeId(graph, "load") < GetNodeId(graph, "abs0"));
524+ EXPECT_TRUE(GetNodeId(graph, "abs0") < GetNodeId(graph, "max1"));
525+ EXPECT_TRUE(GetNodeId(graph, "relu0") < GetNodeId(graph, "store0"));
526+ // 两路reduce同组且整体先于reduce后继组：同级for循环不互相穿插
527+ EXPECT_TRUE(GetNodeId(graph, "max1") < GetNodeId(graph, "relu0"));
528+ EXPECT_TRUE(GetNodeId(graph, "max2") < GetNodeId(graph, "relu0"));
529+ EXPECT_TRUE(GetNodeId(graph, "max2") < GetNodeId(graph, "relu1"));
530+}
531+ 
532+TEST(LoopGroupTest, TopologicalSortingLoopGroupRuleSkippedWithoutMulConsumer) {
533+ auto graph = CreateIndependentTwoReduceChains("loop_group_sort_skipped");
534+ const auto axes = CreateLoopAxisIds(graph);
535+ SetNodeLoopAxis(graph, "load0", axes.axis_b);
536+ SetNodeLoopAxis(graph, "abs0", axes.axis_b);
537+ SetNodeLoopAxis(graph, "max1", axes.axis_b);
538+ SetNodeLoopAxis(graph, "relu0", axes.axis_a);
539+ SetNodeLoopAxis(graph, "load1", axes.axis_b);
540+ SetNodeLoopAxis(graph, "abs1", axes.axis_b);
541+ SetNodeLoopAxis(graph, "max2", axes.axis_b);
542+ SetNodeLoopAxis(graph, "relu1", axes.axis_a);
543+ ASSERT_TRUE(ScheduleUtils::IsNeedLoopGrouping(graph));
544+ 
545+ // 两路独立无多引用：IsNeedFixTopo不触发，分组排序分支不生效，保持依赖拓扑序
546+ ASSERT_EQ(ScheduleUtils::TopologicalSorting(graph, true), af::SUCCESS);
547+ EXPECT_TRUE(GetNodeId(graph, "load0") < GetNodeId(graph, "abs0"));
548+ EXPECT_TRUE(GetNodeId(graph, "abs0") < GetNodeId(graph, "max1"));
549+ EXPECT_TRUE(GetNodeId(graph, "max1") < GetNodeId(graph, "relu0"));
550+ EXPECT_TRUE(GetNodeId(graph, "load1") < GetNodeId(graph, "abs1"));
551+ EXPECT_TRUE(GetNodeId(graph, "abs1") < GetNodeId(graph, "max2"));
552+ EXPECT_TRUE(GetNodeId(graph, "relu1") < GetNodeId(graph, "store1"));
553+}
554+} // namespace
555+} // namespace optimize