Star1148
739
代码
代码
Issues117
Pull Requests297
Actions
项目讨论
Wiki
分析
Star1148
739
已合并
fix(pass): Preserve assemble fanout during chain fusion #6065
zhangxiangjie创建于 9月8日
fix(pass): Preserve assemble fanout during chain fusion #6065
已合并
zhangxiangjie创建于 9月8日
从zhangxiangjie/pyptofix/a5-c2v2-accuracy合入到cann/pyptomaster
共 4 个文件变更+131-12
Mframework/src/passes/pass_utils/merge_view_assemble_utils.cpp+17-6
@@ -720,6 +720,7 @@ const MergeViewAssembleUtils::ConsumerCacheEntry& MergeViewAssembleUtils::BuildT
720 720 
721Status MergeViewAssembleUtils::BuildConsumerCache(Function& function)721Status MergeViewAssembleUtils::BuildConsumerCache(Function& function)
722{722{
723+ hasTokenDependencies_ = false;
723 auto operations = function.Operations(false);724 auto operations = function.Operations(false);
724 consumerCache_.reserve(operations.size());725 consumerCache_.reserve(operations.size());
725 tensorConsumerCache_.reserve(operations.size());726 tensorConsumerCache_.reserve(operations.size());
@@ -742,6 +743,7 @@ Status MergeViewAssembleUtils::BuildConsumerCache(Function& function)
742 recordTensor(outcast);743 recordTensor(outcast);
743 }744 }
744 for (auto& operation : operations) {745 for (auto& operation : operations) {
746+ hasTokenDependencies_ |= !operation.tokens_.empty() || !operation.result_token_.empty();
745 for (const auto& input : operation.GetIOperands()) {747 for (const auto& input : operation.GetIOperands()) {
746 recordTensor(input);748 recordTensor(input);
747 }749 }
@@ -914,8 +916,9 @@ bool MergeViewAssembleUtils::HasSplitVersionContribution(const LogicalTensorPtr&
914bool MergeViewAssembleUtils::BuildProducerGroupFusion(Function& function, const LogicalTensorPtr& middle,916bool MergeViewAssembleUtils::BuildProducerGroupFusion(Function& function, const LogicalTensorPtr& middle,
915 const ConsumerCacheEntry& consumers)917 const ConsumerCacheEntry& consumers)
916{918{
917- if (middle == nullptr || consumers.hasAssembleChainStopper || consumers.assembleConsumers.size() != 1 ||919+ if (middle == nullptr || (hasTokenDependencies_ && consumers.hasAssembleChainStopper) ||
918- !consumers.allProducersAreAssembleLike || IsFunctionBoundaryTensor(function, middle)) {920+ consumers.assembleConsumers.size() != 1 || !consumers.allProducersAreAssembleLike ||
921+ IsFunctionBoundaryTensor(function, middle)) {
919 return false;922 return false;
920 }923 }
921 auto* downstream = consumers.assembleConsumers.front();924 auto* downstream = consumers.assembleConsumers.front();
@@ -1129,7 +1132,9 @@ Status MergeViewAssembleUtils::AppendProducerGroupFusions(Function& function)
1129 }1132 }
1130 RewriteProducerGroupTokens(function, fusion, replacements);1133 RewriteProducerGroupTokens(function, fusion, replacements);
1131 for (auto* producer : fusion.producers) {1134 for (auto* producer : fusion.producers) {
1132- producer->SetAsDeleted();1135+ if (!GetConsumers(*producer).hasAssembleChainStopper) {
1136+ producer->SetAsDeleted();
1137+ }
1133 }1138 }
1134 fusion.downstream->SetAsDeleted();1139 fusion.downstream->SetAsDeleted();
1135 }1140 }
@@ -1396,7 +1401,7 @@ Status MergeViewAssembleUtils::MergeAssembleChain(Function& function, Operation&
1396 1401 
1397 // 2. 处理消费者1402 // 2. 处理消费者
1398 const auto& consumers = GetConsumers(operation);1403 const auto& consumers = GetConsumers(operation);
1399- bool chainEnd = consumers.assembleConsumers.empty() || consumers.hasAssembleChainStopper;1404+ bool chainEnd = consumers.assembleConsumers.empty() || (hasTokenDependencies_ && consumers.hasAssembleChainStopper);
1400 Status status = ProcessAssembleConsumers(function, consumers, chain, chainEnd, effectiveScopeId);1405 Status status = ProcessAssembleConsumers(function, consumers, chain, chainEnd, effectiveScopeId);
1401 if (status != SUCCESS) {1406 if (status != SUCCESS) {
1402 return status;1407 return status;
@@ -1424,7 +1429,7 @@ Status MergeViewAssembleUtils::ProcessAssembleConsumers(Function& function, cons
1424 std::vector<Operation*>& chain, bool& chainEnd,1429 std::vector<Operation*>& chain, bool& chainEnd,
1425 int effectiveScopeId)1430 int effectiveScopeId)
1426{1431{
1427- if (consumers.hasAssembleChainStopper || consumers.assembleConsumers.empty()) {1432+ if ((hasTokenDependencies_ && consumers.hasAssembleChainStopper) || consumers.assembleConsumers.empty()) {
1428 return SUCCESS;1433 return SUCCESS;
1429 }1434 }
1430 Operation* currentOp = chain.back();1435 Operation* currentOp = chain.back();
@@ -1515,8 +1520,14 @@ Status MergeViewAssembleUtils::ProcessAssembleChainEnd(Function& function, std::
1515 GetRmwModeAttrKey(rmwModeAttr), GetMergedAssembleOpcode(chain), tokenDependency,1520 GetRmwModeAttrKey(rmwModeAttr), GetMergedAssembleOpcode(chain), tokenDependency,
1516 atomicSemanticAttr.fromReduceAcc, atomicSemanticAttr.fromExplicitRmw);1521 atomicSemanticAttr.fromReduceAcc, atomicSemanticAttr.fromExplicitRmw);
1517 ClearLinearTokenDependency(function, chain, tokenDependency);1522 ClearLinearTokenDependency(function, chain, tokenDependency);
1523+ // Keep intermediate writes needed by side consumers, as in the legacy chain fusion.
1524+ const bool hasSideConsumer = std::any_of(chain.begin(), chain.end() - 1, [this](const Operation* op) {
1525+ return GetConsumers(*op).hasAssembleChainStopper;
1526+ });
1518 for (auto* op : chain) {1527 for (auto* op : chain) {
1519- op->SetAsDeleted();1528+ if (!hasSideConsumer || op == chain.back()) {
1529+ op->SetAsDeleted();
1530+ }
1520 }1531 }
1521 function.GetTensorMap().Erase(endTensor);1532 function.GetTensorMap().Erase(endTensor);
1522 1533 
Mframework/src/passes/pass_utils/merge_view_assemble_utils.h+1-0
@@ -213,6 +213,7 @@ public:
213 213 
214 // Cleanup methods214 // Cleanup methods
215 Status CleanUp(Function& function);215 Status CleanUp(Function& function);
216+ bool hasTokenDependencies_ = false;
216 std::unordered_set<int> visitedOp_;217 std::unordered_set<int> visitedOp_;
217 std::unordered_map<int, const ConsumerCacheEntry*> consumerCache_;218 std::unordered_map<int, const ConsumerCacheEntry*> consumerCache_;
218 std::unordered_map<int, ConsumerCacheEntry> tensorConsumerCache_;219 std::unordered_map<int, ConsumerCacheEntry> tensorConsumerCache_;
Mframework/tests/ut/passes/src/test_merge_view_assemble.cpp+100-0
@@ -1677,6 +1677,57 @@ TEST_F(MergeViewAssembleTest, NestedCompleteProducerGroupFusionsDoNotOverlap)
1677 EXPECT_NE(function->GetTensorMap().GetTensorByMagic(concatMiddle->GetMagic()), nullptr);1677 EXPECT_NE(function->GetTensorMap().GetTensorByMagic(concatMiddle->GetMagic()), nullptr);
1678}1678}
1679 1679 
1680+TEST_F(MergeViewAssembleTest, ProducerGroupWithIntermediateViewConsumerKeepsDataEdge)
1681+{
1682+ for (bool withToken : {false, true}) {
1683+ Program program;
1684+ auto function = std::make_unique<Function>(program, "producer_group_fanout", "producer_group_fanout", nullptr);
1685+ IRBuilder builder;
1686+ auto firstInput = builder.CreateTensorVar(*function, DataType::DT_FP32, {2, 4});
1687+ auto secondInput = builder.CreateTensorVar(*function, DataType::DT_FP32, {2, 4});
1688+ auto middle = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
1689+ auto viewOutput = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
1690+ auto output = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
1691+ function->inCasts_ = {firstInput, secondInput};
1692+ function->outCasts_ = {viewOutput, output};
1693+ 
1694+ auto& first = builder.CreateTensorOpStmt(*function, Opcode::OP_ASSEMBLE, {firstInput}, {middle});
1695+ first.SetOpAttribute(std::make_shared<AssembleOpAttribute>(std::vector<int64_t>{0, 0}));
1696+ auto& second = builder.CreateTensorOpStmt(*function, Opcode::OP_ASSEMBLE, {secondInput}, {middle});
1697+ second.SetOpAttribute(std::make_shared<AssembleOpAttribute>(std::vector<int64_t>{2, 0}));
1698+ auto& view = builder.CreateTensorOpStmt(*function, Opcode::OP_VIEW, {middle}, {viewOutput});
1699+ view.SetOpAttribute(std::make_shared<ViewOpAttribute>(std::vector<int64_t>{0, 0}));
1700+ auto& downstream = builder.CreateTensorOpStmt(*function, Opcode::OP_ASSEMBLE, {middle}, {output});
1701+ downstream.SetOpAttribute(std::make_shared<AssembleOpAttribute>(std::vector<int64_t>{0, 0}));
1702+ if (withToken) {
1703+ AddTokenEdge(*function, first, view);
1704+ }
1705+ function->BuildTensorMap();
1706+ 
1707+ MergeViewAssemble pass;
1708+ ASSERT_EQ(pass.RunOnFunction(*function), SUCCESS);
1709+ EXPECT_EQ(middle->GetProducers().size(), 2);
1710+ EXPECT_EQ(view.GetIOperands().front(), middle);
1711+ EXPECT_FALSE(first.IsDeleted());
1712+ EXPECT_FALSE(second.IsDeleted());
1713+ if (withToken) {
1714+ ASSERT_EQ(output->GetProducers().size(), 1);
1715+ EXPECT_EQ((*output->GetProducers().begin())->GetIOperands().front(), middle);
1716+ ASSERT_EQ(view.tokens_.size(), 1);
1717+ EXPECT_TRUE(function->GetVarDependency().HasProducer(view.tokens_.front(), ToStmtPtr(first)));
1718+ } else {
1719+ ASSERT_EQ(output->GetProducers().size(), 2);
1720+ for (auto* producer : output->GetProducers()) {
1721+ const auto& input = producer->GetIOperands().front();
1722+ EXPECT_TRUE(input == firstInput || input == secondInput);
1723+ auto attr = std::dynamic_pointer_cast<AssembleOpAttribute>(producer->GetOpAttribute());
1724+ ASSERT_NE(attr, nullptr);
1725+ EXPECT_EQ(attr->GetToOffset(), (std::vector<int64_t>{input == firstInput ? 0 : 2, 0}));
1726+ }
1727+ }
1728+ }
1729+}
1730+ 
1680TEST_F(MergeViewAssembleTest, ProducerGroupWithCoverageHoleDoesNotFuse)1731TEST_F(MergeViewAssembleTest, ProducerGroupWithCoverageHoleDoesNotFuse)
1681{1732{
1682 Program program;1733 Program program;
@@ -2059,6 +2110,55 @@ TEST_F(MergeViewAssembleTest, AssembleChainWithIntermediateViewConsumerKeepsData
2059 ASSERT_EQ(remainingView->GetIOperands().size(), 1);2110 ASSERT_EQ(remainingView->GetIOperands().size(), 1);
2060 EXPECT_EQ(remainingView->GetIOperands().front(), middle);2111 EXPECT_EQ(remainingView->GetIOperands().front(), middle);
2061 EXPECT_FALSE(middle->GetProducers().empty());2112 EXPECT_FALSE(middle->GetProducers().empty());
2113+ ASSERT_EQ(output->GetProducers().size(), 1);
2114+ EXPECT_EQ((*output->GetProducers().begin())->GetIOperands().front(), input);
2115+}
2116+ 
2117+TEST_F(MergeViewAssembleTest, AssembleChainWithIntermediateTokenConsumerKeepsDataEdge)
2118+{
2119+ Program program;
2120+ auto function = std::make_unique<Function>(program, "assemble_view_fanout", "assemble_view_fanout", nullptr);
2121+ IRBuilder builder;
2122+ auto input = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
2123+ auto middle = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
2124+ auto viewOutput = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
2125+ auto output = builder.CreateTensorVar(*function, DataType::DT_FP32, {4, 4});
2126+ function->inCasts_ = {input};
2127+ function->outCasts_ = {viewOutput, output};
2128+ 
2129+ auto& first = builder.CreateTensorOpStmt(*function, Opcode::OP_ASSEMBLE, {input}, {middle});
2130+ first.SetOpAttribute(std::make_shared<AssembleOpAttribute>(std::vector<int64_t>{0, 0}));
2131+ auto& view = builder.CreateTensorOpStmt(*function, Opcode::OP_VIEW, {middle}, {viewOutput});
2132+ view.SetOpAttribute(std::make_shared<ViewOpAttribute>(std::vector<int64_t>{0, 0}));
2133+ auto& second = builder.CreateTensorOpStmt(*function, Opcode::OP_ASSEMBLE, {middle}, {output});
2134+ second.SetOpAttribute(std::make_shared<AssembleOpAttribute>(std::vector<int64_t>{0, 0}));
2135+ auto token = AddTokenEdge(*function, first, view);
2136+ function->BuildTensorMap();
2137+ 
2138+ MergeViewAssemble pass;
2139+ ASSERT_EQ(pass.RunOnFunction(*function), SUCCESS);
2140+ 
2141+ size_t assembleCount = 0;
2142+ size_t viewCount = 0;
2143+ Operation* remainingView = nullptr;
2144+ for (auto& op : function->Operations(false)) {
2145+ if (op.GetOpcode() == Opcode::OP_ASSEMBLE) {
2146+ ++assembleCount;
2147+ } else if (op.GetOpcode() == Opcode::OP_VIEW) {
2148+ ++viewCount;
2149+ remainingView = &op;
2150+ }
2151+ }
2152+ EXPECT_EQ(assembleCount, 2);
2153+ EXPECT_EQ(viewCount, 1);
2154+ ASSERT_NE(remainingView, nullptr);
2155+ ASSERT_EQ(remainingView->GetIOperands().size(), 1);
2156+ EXPECT_EQ(remainingView->GetIOperands().front(), middle);
2157+ EXPECT_FALSE(middle->GetProducers().empty());
2158+ ASSERT_EQ(output->GetProducers().size(), 1);
2159+ EXPECT_EQ((*output->GetProducers().begin())->GetIOperands().front(), middle);
2160+ EXPECT_TRUE(function->GetVarDependency().HasProducer(token, ToStmtPtr(first)));
2161+ EXPECT_TRUE(function->GetVarDependency().HasConsumer(token, ToStmtPtr(view)));
2062}2162}
2063 2163 
2064TEST_F(MergeViewAssembleTest, ShmemSetPredicateConsumerStopsViewChainFusion)2164TEST_F(MergeViewAssembleTest, ShmemSetPredicateConsumerStopsViewChainFusion)
Mframework/tests/ut/passes/src/test_remove_redundant_op.cpp+13-6
@@ -112,7 +112,8 @@ TEST_F(RemoveRedundantOpTest, TestIntermediateOutcast)
112 Function* func = Program::GetInstance().GetFunctionByRawName("TENSOR_RemoveRedundantOpFunction");112 Function* func = Program::GetInstance().GetFunctionByRawName("TENSOR_RemoveRedundantOpFunction");
113 npu::tile_fwk::RemoveRedundantOp removeRedundantOp;113 npu::tile_fwk::RemoveRedundantOp removeRedundantOp;
114 auto oriOpList = func->Operations(true);114 auto oriOpList = func->Operations(true);
115- EXPECT_EQ(oriOpList.size(), 16) << "Before the Pass, there should be 16 operations";115+ // Assemble chain fusion preserves the four tiled VIEW consumers of the transpose result.
116+ EXPECT_EQ(oriOpList.size(), 19) << "Before the Pass, there should be 19 operations";
116 int ori_view_count = 0;117 int ori_view_count = 0;
117 int ori_assemble_count = 0;118 int ori_assemble_count = 0;
118 for (auto& op : oriOpList) {119 for (auto& op : oriOpList) {
@@ -122,16 +123,16 @@ TEST_F(RemoveRedundantOpTest, TestIntermediateOutcast)
122 ori_assemble_count += 1;123 ori_assemble_count += 1;
123 }124 }
124 }125 }
125- EXPECT_EQ(ori_view_count, 0) << "There should be no VIEW ops before RemoveRedundantOp";126+ EXPECT_EQ(ori_view_count, 4) << "There should be 4 VIEW ops before RemoveRedundantOp";
126- EXPECT_EQ(ori_assemble_count, 2) << "There should be 2 ASSEMBLE ops before RemoveRedundantOp";127+ EXPECT_EQ(ori_assemble_count, 1) << "There should be 1 ASSEMBLE op before RemoveRedundantOp";
127 EXPECT_EQ(removeRedundantOp.PreCheck(*func), SUCCESS);128 EXPECT_EQ(removeRedundantOp.PreCheck(*func), SUCCESS);
128 EXPECT_EQ(removeRedundantOp.RunOnFunction(*func), SUCCESS);129 EXPECT_EQ(removeRedundantOp.RunOnFunction(*func), SUCCESS);
129 EXPECT_EQ(removeRedundantOp.PostCheck(*func), SUCCESS);130 EXPECT_EQ(removeRedundantOp.PostCheck(*func), SUCCESS);
130 PrintGraphInfoRemoveRedundantOp(func);131 PrintGraphInfoRemoveRedundantOp(func);
131 // ================== Verify the effect of the Pass ==================132 // ================== Verify the effect of the Pass ==================
132 auto updated_operations = func->Operations(true);133 auto updated_operations = func->Operations(true);
133- int opSize = 15;134+ int opSize = 19;
134- EXPECT_EQ(updated_operations.size(), opSize) << "After the Pass, there should be 15 operations";135+ EXPECT_EQ(updated_operations.size(), opSize) << "After the Pass, there should be 19 operations";
135 EXPECT_EQ(updated_operations[0].GetOpcode(), Opcode::OP_SLICE) << "The first operation should be SILCE";136 EXPECT_EQ(updated_operations[0].GetOpcode(), Opcode::OP_SLICE) << "The first operation should be SILCE";
136 int view_count = 0;137 int view_count = 0;
137 int assemble_count = 0;138 int assemble_count = 0;
@@ -143,11 +144,17 @@ TEST_F(RemoveRedundantOpTest, TestIntermediateOutcast)
143 assemble_count += 1;144 assemble_count += 1;
144 }145 }
145 }146 }
146- EXPECT_EQ(view_count, 0) << "There should be no VIEW ops after RemoveRedundantOp";147+ EXPECT_EQ(view_count, 4) << "There should be 4 VIEW ops after RemoveRedundantOp";
147 EXPECT_EQ(assemble_count, 1) << "There should be 1 Assemble op after RemoveRedundantOp";148 EXPECT_EQ(assemble_count, 1) << "There should be 1 Assemble op after RemoveRedundantOp";
148 const auto& outcasts = func->GetOutcast();149 const auto& outcasts = func->GetOutcast();
149 ASSERT_GE(outcasts.size(), 2U);150 ASSERT_GE(outcasts.size(), 2U);
150 ASSERT_EQ(outcasts[0]->GetProducers().size(), 1U);151 ASSERT_EQ(outcasts[0]->GetProducers().size(), 1U);
152+ // Keep the outcast ASSEMBLE when its input also feeds internal consumers.
153+ const auto* outcastAssemble = *outcasts[0]->GetProducers().begin();
154+ EXPECT_EQ(outcastAssemble->GetOpcode(), Opcode::OP_ASSEMBLE);
155+ ASSERT_EQ(outcastAssemble->GetIOperands().size(), 1U);
156+ EXPECT_GT(outcastAssemble->GetIOperands().front()->GetConsumers().size(), 1U);
157+ EXPECT_TRUE(outcasts[0]->GetConsumers().empty());
151 const auto& validShape = outcasts[0]->GetDynValidShape();158 const auto& validShape = outcasts[0]->GetDynValidShape();
152 ASSERT_EQ(validShape.size(), resShape.size());159 ASSERT_EQ(validShape.size(), resShape.size());
153 EXPECT_EQ(validShape[0].Simplify().Concrete(), n);160 EXPECT_EQ(validShape[0].Simplify().Concrete(), n);