已合并
【PR】: fix reduce layer norm #1891
czways创建于 14 天前
【PR】: fix reduce layer norm #1891
已合并
共 3 个文件变更+209-16
| @@ -9,6 +9,9 @@ | |||
| 9 | */ | 9 | */ |
| 10 | 10 | ||
| 11 | 11 | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 12 | 15 | ||
| 13 | 16 | ||
| 14 | 17 | ||
| @@ -401,25 +404,109 @@ Status ReducePartitionCaseGenerator::ReducePartitionMultipleCitations(ascir::Imp | |||
| 401 | if (IsOnlyHasOneOrLessReduce(impl_graph)) { | 404 | if (IsOnlyHasOneOrLessReduce(impl_graph)) { |
| 402 | return ge::GRAPH_SUCCESS; | 405 | return ge::GRAPH_SUCCESS; |
| 403 | } | 406 | } |
| 404 | - std::vector<af::AscNodePtr> multi_output_nodes; | 407 | + CitationGroups citation_groups; |
| 408 | + GE_CHK_STATUS_RET(CollectCitationGroups(impl_graph, citation_groups)); | ||
| 409 | + std::vector<size_t> parent(citation_groups.size()); | ||
| 410 | + std::map<size_t, af::AscNodePtr> group_anchors; | ||
| 411 | + BuildCitationGroupAnchors(citation_groups, parent, group_anchors); | ||
| 412 | + return PartitionCitationGroups(impl_graph, citation_groups, parent, group_anchors); | ||
| 413 | +} | ||
| 414 | + | ||
| 415 | +Status ReducePartitionCaseGenerator::CollectCitationGroups(ascir::ImplGraph &impl_graph, | ||
| 416 | + CitationGroups &citation_groups) { | ||
| 405 | for (auto node : impl_graph.GetAllNodes()) { | 417 | for (auto node : impl_graph.GetAllNodes()) { |
| 406 | - if (node->GetOutNodes().size() > 1UL) { | 418 | + if (!ScheduleUtils::IsLoad(node) && !ScheduleUtils::IsStore(node) && |
| 407 | - multi_output_nodes.emplace_back(node); | 419 | + !af::ops::IsOps<af::ascir_op::Workspace>(node) && node->GetOutDataNodes().size() > 1UL) { |
| 420 | + std::vector<Citation> citations; | ||
| 421 | + for (const auto &output_node : node->GetOutDataNodes()) { | ||
| 422 | + auto citation = std::dynamic_pointer_cast<af::AscNode>(output_node); | ||
| 423 | + GE_CHECK_NOTNULL(citation); | ||
| 424 | + af::AscNodePtr reduce; | ||
| 425 | + if (FindOutputReduce(citation, reduce)) { | ||
| 426 | + citations.push_back({node, citation, reduce}); | ||
| 427 | + } | ||
| 428 | + } | ||
| 429 | + if (!citations.empty()) { | ||
| 430 | + std::sort(citations.begin(), citations.end(), [](const Citation &lhs, const Citation &rhs) { | ||
| 431 | + if (lhs.reduce->GetOpDescBarePtr()->GetId() != rhs.reduce->GetOpDescBarePtr()->GetId()) { | ||
| 432 | + return lhs.reduce->GetOpDescBarePtr()->GetId() < rhs.reduce->GetOpDescBarePtr()->GetId(); | ||
| 433 | + } | ||
| 434 | + return lhs.citation->GetOpDescBarePtr()->GetId() < rhs.citation->GetOpDescBarePtr()->GetId(); | ||
| 435 | + }); | ||
| 436 | + citation_groups.emplace_back(std::move(citations)); | ||
| 437 | + } | ||
| 408 | } | 438 | } |
| 409 | } | 439 | } |
| 410 | - std::sort(multi_output_nodes.begin(), multi_output_nodes.end(), [](const af::AscNodePtr &lhs, af::AscNodePtr &rhs) { | 440 | + return ge::GRAPH_SUCCESS; |
| 411 | - return lhs->GetOpDescBarePtr()->GetId() > rhs->GetOpDescBarePtr()->GetId(); | 441 | +} |
| 412 | - }); | 442 | + |
| 413 | - for (auto node : multi_output_nodes) { | 443 | +void ReducePartitionCaseGenerator::BuildCitationGroupAnchors(const CitationGroups &citation_groups, |
| 414 | - std::set<af::AscNodePtr> reduce_nodes; | 444 | + std::vector<size_t> &parent, |
| 415 | - for (const auto &output_node : node->GetOutNodes()) { | 445 | + std::map<size_t, af::AscNodePtr> &group_anchors) { |
| 416 | - af::AscNodePtr out_asc_node = std::dynamic_pointer_cast<af::AscNode>(output_node); | 446 | + std::iota(parent.begin(), parent.end(), 0UL); |
| 417 | - if (af::AscNodePtr reduce_node = nullptr; FindOutputReduce(out_asc_node, reduce_node)) { | 447 | + auto find_root = [&parent](size_t index) { |
| 418 | - if (!reduce_nodes.empty() && reduce_nodes.find(reduce_node) == reduce_nodes.end()) { | 448 | + while (parent[index] != index) { |
| 419 | - PartitionByNode(node, out_asc_node, impl_graph); | 449 | + parent[index] = parent[parent[index]]; |
| 450 | + index = parent[index]; | ||
| 451 | + } | ||
| 452 | + return index; | ||
| 453 | + }; | ||
| 454 | + for (size_t i = 0UL; i < citation_groups.size(); ++i) { | ||
| 455 | + for (size_t j = i + 1UL; j < citation_groups.size(); ++j) { | ||
| 456 | + bool shared_reduce = false; | ||
| 457 | + for (const auto &lhs : citation_groups[i]) { | ||
| 458 | + for (const auto &rhs : citation_groups[j]) { | ||
| 459 | + if (lhs.reduce == rhs.reduce) { | ||
| 460 | + shared_reduce = true; | ||
| 461 | + break; | ||
| 462 | + } | ||
| 463 | + } | ||
| 464 | + if (shared_reduce) { | ||
| 465 | + break; | ||
| 420 | } | 466 | } |
| 421 | - reduce_nodes.emplace(reduce_node); | ||
| 422 | } | 467 | } |
| 468 | + if (shared_reduce) { | ||
| 469 | + parent[find_root(j)] = find_root(i); | ||
| 470 | + } | ||
| 471 | + } | ||
| 472 | + } | ||
| 473 | + | ||
| 474 | + for (size_t i = 0UL; i < citation_groups.size(); ++i) { | ||
| 475 | + const auto root = find_root(i); | ||
| 476 | + for (const auto &citation : citation_groups[i]) { | ||
| 477 | + auto anchor = group_anchors.find(root); | ||
| 478 | + if (anchor == group_anchors.end() || | ||
| 479 | + citation.reduce->GetOpDescBarePtr()->GetId() < anchor->second->GetOpDescBarePtr()->GetId()) { | ||
| 480 | + group_anchors[root] = citation.reduce; | ||
| 481 | + } | ||
| 482 | + } | ||
| 483 | + } | ||
| 484 | +} | ||
| 485 | + | ||
| 486 | +Status ReducePartitionCaseGenerator::PartitionCitationGroups(ascir::ImplGraph &impl_graph, | ||
| 487 | + const CitationGroups &citation_groups, | ||
| 488 | + const std::vector<size_t> &parent, | ||
| 489 | + const std::map<size_t, af::AscNodePtr> &group_anchors) { | ||
| 490 | + auto find_root = [&parent](size_t index) { | ||
| 491 | + while (parent[index] != index) { | ||
| 492 | + index = parent[index]; | ||
| 493 | + } | ||
| 494 | + return index; | ||
| 495 | + }; | ||
| 496 | + // A source can have multiple citations that eventually reach the same non-anchor reduce. | ||
| 497 | + // Partitioning each citation creates duplicate workspace/load chains for one reduce. | ||
| 498 | + std::set<std::pair<const af::Node *, const af::Node *>> partitioned_source_reduces; | ||
| 499 | + for (size_t i = 0UL; i < citation_groups.size(); ++i) { | ||
| 500 | + const auto root = find_root(i); | ||
| 501 | + const auto anchor = group_anchors.at(root); | ||
| 502 | + for (const auto &citation : citation_groups[i]) { | ||
| 503 | + if (citation.reduce == anchor || | ||
| 504 | + !partitioned_source_reduces.emplace(citation.source.get(), citation.reduce.get()).second) { | ||
| 505 | + continue; | ||
| 506 | + } | ||
| 507 | + auto source = citation.source; | ||
| 508 | + auto citation_node = citation.citation; | ||
| 509 | + GE_CHK_STATUS_RET(PartitionByNode(source, citation_node, impl_graph)); | ||
| 423 | } | 510 | } |
| 424 | } | 511 | } |
| 425 | return ge::GRAPH_SUCCESS; | 512 | return ge::GRAPH_SUCCESS; |
| @@ -431,10 +518,10 @@ bool ReducePartitionCaseGenerator::FindOutputReduce(const af::AscNodePtr &node, | |||
| 431 | return true; | 518 | return true; |
| 432 | } | 519 | } |
| 433 | bool output_has_reduce = false; | 520 | bool output_has_reduce = false; |
| 434 | - if (node->GetOutNodes().empty()) { | 521 | + if (node->GetOutDataNodes().empty()) { |
| 435 | return output_has_reduce; | 522 | return output_has_reduce; |
| 436 | } | 523 | } |
| 437 | - for (const auto &output_node : node->GetOutNodes()) { | 524 | + for (const auto &output_node : node->GetOutDataNodes()) { |
| 438 | auto output_asc_node = std::dynamic_pointer_cast<af::AscNode>(output_node); | 525 | auto output_asc_node = std::dynamic_pointer_cast<af::AscNode>(output_node); |
| 439 | output_has_reduce = output_has_reduce || FindOutputReduce(output_asc_node, reduce_node); | 526 | output_has_reduce = output_has_reduce || FindOutputReduce(output_asc_node, reduce_node); |
| 440 | } | 527 | } |
| @@ -15,6 +15,7 @@ | |||
| 15 | 15 | ||
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | + | ||
| 18 | 19 | ||
| 19 | namespace optimize { | 20 | namespace optimize { |
| 20 | 21 | ||
| @@ -34,6 +35,13 @@ class ReducePartitionCaseGenerator : public FusionCaseGenerator { | |||
| 34 | const std::vector<std::string> &score_functions); | 35 | const std::vector<std::string> &score_functions); |
| 35 | 36 | ||
| 36 | private: | 37 | private: |
| 38 | + struct Citation { | ||
| 39 | + af::AscNodePtr source; | ||
| 40 | + af::AscNodePtr citation; | ||
| 41 | + af::AscNodePtr reduce; | ||
| 42 | + }; | ||
| 43 | + using CitationGroups = std::vector<std::vector<Citation>>; | ||
| 44 | + | ||
| 37 | Status GeneratorGeneralTask(ascir::HintGraph &optimize_graph, std::vector<ScheduleTask> &tasks); | 45 | Status GeneratorGeneralTask(ascir::HintGraph &optimize_graph, std::vector<ScheduleTask> &tasks); |
| 38 | Status GeneratorAllLoadTask(ascir::HintGraph &optimize_graph, std::vector<ScheduleTask> &tasks); | 46 | Status GeneratorAllLoadTask(ascir::HintGraph &optimize_graph, std::vector<ScheduleTask> &tasks); |
| 39 | Status GeneratorRCoreTask(ascir::HintGraph &optimize_graph, std::vector<ScheduleTask> &tasks) const; | 47 | Status GeneratorRCoreTask(ascir::HintGraph &optimize_graph, std::vector<ScheduleTask> &tasks) const; |
| @@ -56,6 +64,12 @@ class ReducePartitionCaseGenerator : public FusionCaseGenerator { | |||
| 56 | static bool IsOnlyHasOneOrLessReduce(const ascir::ImplGraph &impl_graph); | 64 | static bool IsOnlyHasOneOrLessReduce(const ascir::ImplGraph &impl_graph); |
| 57 | static bool CanFullLoadReduceFuse(const ascir::ImplGraph &impl_graph); | 65 | static bool CanFullLoadReduceFuse(const ascir::ImplGraph &impl_graph); |
| 58 | Status ReducePartitionMultipleCitations(ascir::ImplGraph &impl_graph); | 66 | Status ReducePartitionMultipleCitations(ascir::ImplGraph &impl_graph); |
| 67 | + Status CollectCitationGroups(ascir::ImplGraph &impl_graph, CitationGroups &citation_groups); | ||
| 68 | + void BuildCitationGroupAnchors(const CitationGroups &citation_groups, std::vector<size_t> &parent, | ||
| 69 | + std::map<size_t, af::AscNodePtr> &group_anchors); | ||
| 70 | + Status PartitionCitationGroups(ascir::ImplGraph &impl_graph, const CitationGroups &citation_groups, | ||
| 71 | + const std::vector<size_t> &parent, | ||
| 72 | + const std::map<size_t, af::AscNodePtr> &group_anchors); | ||
| 59 | bool FindOutputReduce(const af::AscNodePtr &node, af::AscNodePtr &reduce_node); | 73 | bool FindOutputReduce(const af::AscNodePtr &node, af::AscNodePtr &reduce_node); |
| 60 | Status PartitionReduceNode(af::AscNodePtr &src_node, ascir::ImplGraph &impl_graph); | 74 | Status PartitionReduceNode(af::AscNodePtr &src_node, ascir::ImplGraph &impl_graph); |
| 61 | 75 | ||
| @@ -23,6 +23,8 @@ using namespace optimize; | |||
| 23 | using namespace ge; | 23 | using namespace ge; |
| 24 | using namespace af::ops; | 24 | using namespace af::ops; |
| 25 | using namespace af::ascir_op; | 25 | using namespace af::ascir_op; |
| 26 | +using af::testing::AscGraphBuilder; | ||
| 27 | +using af::testing::Sym; | ||
| 26 | 28 | ||
| 27 | class ReduceScheduleCaseGeneratorTest : public ::testing::Test { | 29 | class ReduceScheduleCaseGeneratorTest : public ::testing::Test { |
| 28 | protected: | 30 | protected: |
| @@ -552,6 +554,96 @@ TEST_F(ReduceScheduleCaseGeneratorTest, TestReduce_Multi_Cita_Store) { | |||
| 552 | EXPECT_FALSE(HasReduceTemplateType(tasks, ReduceTemplateType::kAllLoad)); | 554 | EXPECT_FALSE(HasReduceTemplateType(tasks, ReduceTemplateType::kAllLoad)); |
| 553 | } | 555 | } |
| 554 | 556 | ||
| 557 | +TEST_F(ReduceScheduleCaseGeneratorTest, TestReduce_Multi_Cita_Multi_Out_NoDependencyCycle) { | ||
| 558 | + auto graph = AscGraphBuilder("reduce_multi_citation_multi_out") | ||
| 559 | + .Loops({Sym(128), Sym(64)}) | ||
| 560 | + .Data("data", 0) | ||
| 561 | + .Load("load", "data") | ||
| 562 | + .Abs("shared0", "load") | ||
| 563 | + .Relu("branch00", "shared0") | ||
| 564 | + .Op<af::ascir_op::Tanh>("branch01", {"shared0"}) | ||
| 565 | + .Add("merge0", "branch00", "branch01") | ||
| 566 | + .Abs("shared1", "merge0") | ||
| 567 | + .Relu("branch10", "shared1") | ||
| 568 | + .Op<af::ascir_op::Tanh>("branch11", {"shared1"}) | ||
| 569 | + .Relu("branch12", "shared1") | ||
| 570 | + .Add("merge1", "branch10", "branch11") | ||
| 571 | + .Sum("sum1", "branch12", {0, 1}) | ||
| 572 | + .Add("merge_final", "merge0", "merge1") | ||
| 573 | + .Sum("sum0", "merge_final", {0, 1}) | ||
| 574 | + .Store("store0", "sum0") | ||
| 575 | + .Output("output0", "store0", 0) | ||
| 576 | + .Store("store1", "sum1") | ||
| 577 | + .Output("output1", "store1", 1) | ||
| 578 | + .Build(); | ||
| 579 | + std::vector<ScheduleTask> tasks; | ||
| 580 | + ReducePartitionCaseGenerator generator; | ||
| 581 | + OptimizerOptions options; | ||
| 582 | + | ||
| 583 | + EXPECT_EQ(generator.GeneratorTask(graph, tasks, options), af::SUCCESS); | ||
| 584 | + EXPECT_TRUE(HasReduceTemplateType(tasks, ReduceTemplateType::kCommon)); | ||
| 585 | + | ||
| 586 | + size_t shared0_workspace_count = 0UL; | ||
| 587 | + size_t shared1_workspace_count = 0UL; | ||
| 588 | + for (const auto &task : tasks) { | ||
| 589 | + for (const auto &grouped_graph : task.grouped_graphs) { | ||
| 590 | + for (const auto &node : grouped_graph.GetAllNodes()) { | ||
| 591 | + if (!af::ops::IsOps<af::ascir_op::Workspace>(node)) { | ||
| 592 | + continue; | ||
| 593 | + } | ||
| 594 | + if (node->GetName().find("shared0_to_branch") != std::string::npos) { | ||
| 595 | + ++shared0_workspace_count; | ||
| 596 | + } | ||
| 597 | + if (node->GetName().find("shared1_to_branch") != std::string::npos) { | ||
| 598 | + ++shared1_workspace_count; | ||
| 599 | + } | ||
| 600 | + } | ||
| 601 | + } | ||
| 602 | + } | ||
| 603 | + // Both citation groups share sum0 and must be merged; sum1 is the lower-id anchor. | ||
| 604 | + // Each source/reduce pair is partitioned once, producing one workspace pair. | ||
| 605 | + EXPECT_EQ(shared0_workspace_count, 2UL); | ||
| 606 | + EXPECT_EQ(shared1_workspace_count, 2UL); | ||
| 607 | +} | ||
| 608 | + | ||
| 609 | +TEST_F(ReduceScheduleCaseGeneratorTest, TestReduce_Multi_Cita_SameReduce_NoDuplicatePartition) { | ||
| 610 | + auto graph = AscGraphBuilder("reduce_multi_citation_same_reduce") | ||
| 611 | + .Loops({Sym(128), Sym(64)}) | ||
| 612 | + .Data("data", 0) | ||
| 613 | + .Load("load", "data") | ||
| 614 | + .Abs("shared", "load") | ||
| 615 | + .Relu("branch0", "shared") | ||
| 616 | + .Op<af::ascir_op::Tanh>("branch1", {"shared"}) | ||
| 617 | + .Add("merge", "branch0", "branch1") | ||
| 618 | + .Sum("sum0", "merge", {0, 1}) | ||
| 619 | + .Relu("branch2", "shared") | ||
| 620 | + .Sum("sum1", "branch2", {0, 1}) | ||
| 621 | + .Store("store0", "sum0") | ||
| 622 | + .Output("output0", "store0", 0) | ||
| 623 | + .Store("store1", "sum1") | ||
| 624 | + .Output("output1", "store1", 1) | ||
| 625 | + .Build(); | ||
| 626 | + std::vector<ScheduleTask> tasks; | ||
| 627 | + ReducePartitionCaseGenerator generator; | ||
| 628 | + OptimizerOptions options; | ||
| 629 | + | ||
| 630 | + EXPECT_EQ(generator.GeneratorTask(graph, tasks, options), af::SUCCESS); | ||
| 631 | + | ||
| 632 | + size_t shared_workspace_count = 0UL; | ||
| 633 | + for (const auto &task : tasks) { | ||
| 634 | + for (const auto &grouped_graph : task.grouped_graphs) { | ||
| 635 | + for (const auto &node : grouped_graph.GetAllNodes()) { | ||
| 636 | + if (af::ops::IsOps<af::ascir_op::Workspace>(node) && | ||
| 637 | + node->GetName().find("shared_to_branch") != std::string::npos) { | ||
| 638 | + ++shared_workspace_count; | ||
| 639 | + } | ||
| 640 | + } | ||
| 641 | + } | ||
| 642 | + } | ||
| 643 | + // One partition creates a workspace pair; duplicate citation partitioning would create two pairs. | ||
| 644 | + EXPECT_EQ(shared_workspace_count, 2UL); | ||
| 645 | +} | ||
 | |||
| 646 | + | ||
| 555 | void ConstructReduceWithScalarData(AscGraph &graph) { | 647 | void ConstructReduceWithScalarData(AscGraph &graph) { |
| 556 | auto s0 = graph.CreateSizeVar(128); | 648 | auto s0 = graph.CreateSizeVar(128); |
| 557 | auto s1 = graph.CreateSizeVar(64); | 649 | auto s1 = graph.CreateSizeVar(64); |
🟡 Medium Priority
changed line → 新测试
TestReduce_Multi_Cita_Multi_Out_NoDependencyCycle构建的图中只有Sum("sum", "load", {0,1})这一个 reduce 算子（ScheduleUtils::IsReduce仅当compute_type == kComputeReduce时成立，而 Data/Load/Abs/Tanh/Add/Relu/Store/Output 均不满足）。affected behavior →ReducePartitionMultipleCitations第 404-406 行在IsOnlyHasOneOrLessReduce(impl_graph)为真时直接返回，而该函数判断的正是 reduce 数量 ≤1；本 PR 重写的全部新逻辑（407-500 行：citation 收集、union-find 分组、anchor 选取、partitioned_edges 去重、PartitionByNode 调用）在本次测试中一行都不会执行。failure mode → 即使新算法完全损坏（分组错误、anchor 选取错误、越界等），该测试仍会通过，只验证了"单 reduce 多输出图不崩溃"；测试名 "NoDependencyCycle" 给人已覆盖新逻辑的错觉。suggested fix → 构造一个至少含两个 reduce、且存在多输出节点同时引用多个 reduce 的图（例如两个 Sum 分支共享一个多输出源节点），并断言切分后生成了 kCommon 任务且关键边被改写；同时可断言多源共享 reduce 场景的预期切分行为。建议：在测试图中增加第二个 reduce 算子并让同一多输出源节点同时引用多个不同 reduce，使 IsOnlyHasOneOrLessReduce 返回 false，从而真正执行重写后的分组/切分逻辑；并增加对切分结果的断言（如引用 kCommon 任务、切分后源节点输出锚点被改写为 workspace/load 链）。