Skip to content

Multi path fusion plan #1017

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Melikano wants to merge 5 commits into
base: main
Choose a base branch
from
+298 −42
Open

Multi path fusion plan #1017

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8efa723
refactor(tile-fusion): add reusable whole-group cost evaluation
Jul 27, 2026
b228629
feat(tile-fusion): explore cost-ranked fusion group candidates
Jul 27, 2026
41b23f7
test(tile-fusion): cover cost-guided fusion group search
Jul 27, 2026
815dc43
fix: separate lit tests
Jul 27, 2026
b183af7
chore: add a comment describing the bounded search
Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
210 changes: 168 additions & 42 deletions lib/PTO/Transforms/TileFusion/PTOFusionPlan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include "llvm/ADT/StringSwitch.h"

#include <algorithm>
#include <optional>

namespace mlir {
namespace pto {
Expand Down Expand Up @@ -245,6 +246,31 @@ countConnectionsToGroup(const pto::FusionBlockAnalysis &blockAnalysis,
return connections;
}

static unsigned
countInternalEdges(const pto::FusionBlockAnalysis &blockAnalysis,
ArrayRef<const pto::FusionComputeNode *> group) {
DenseSet<unsigned> memberIds;
for (const pto::FusionComputeNode *member : group)
memberIds.insert(member->id);

return llvm::count_if(blockAnalysis.edges,
[&](const pto::FusionDFGEdge &edge) {
return memberIds.contains(edge.producerNode) &&
memberIds.contains(edge.consumerNode);
});
}

static unsigned
countInternalConnections(const pto::FusionBlockAnalysis &blockAnalysis,
ArrayRef<const pto::FusionComputeNode *> group) {
unsigned connections = 0;
for (auto [index, lhs] : llvm::enumerate(group))
for (const pto::FusionComputeNode *rhs : group.drop_front(index + 1))
if (nodesHaveDirectDataFlowConnection(blockAnalysis, *lhs, *rhs))
++connections;
return connections;
}

static GroupFootprint
computeGroupFootprint(ArrayRef<const pto::FusionComputeNode *> members) {
DenseSet<Value> producedTiles;
Expand Down Expand Up @@ -272,6 +298,22 @@ computeGroupFootprint(ArrayRef<const pto::FusionComputeNode *> members) {
return footprint;
}

static PlanningCost
computePlanningCost(ArrayRef<const pto::FusionComputeNode *> members,
unsigned connectionCount, int64_t loopMergeBenefit,
unsigned liveTileLimit, unsigned vfParameterLimit) {
GroupFootprint footprint = computeGroupFootprint(members);

PlanningCost cost;
cost.dependencyBenefit = 4 * static_cast<int64_t>(connectionCount);
cost.loopMergeBenefit = loopMergeBenefit;
cost.liveTilePenalty = std::max<int64_t>(
0, static_cast<int64_t>(footprint.liveTileCount) - liveTileLimit);
cost.vfParameterPenalty = std::max<int64_t>(
0, static_cast<int64_t>(footprint.vfParameterCount) - vfParameterLimit);
return cost;
}

class CostModel {
public:
virtual ~CostModel() = default;
Expand All @@ -280,6 +322,10 @@ class CostModel {
const pto::FusionComputeNode &candidate)
const = 0;

virtual PlanningCost
evaluateGroup(const PlanningContext &ctx,
ArrayRef<const pto::FusionComputeNode *> members) const = 0;

virtual PlanningDecision
evaluateAppend(const PlanningContext &ctx,
ArrayRef<const pto::FusionComputeNode *> currentGroup,
Expand All @@ -304,6 +350,14 @@ class ConservativeGreedyCostModel final : public CostModel {
return decision;
}

PlanningCost evaluateGroup(
const PlanningContext &ctx,
ArrayRef<const pto::FusionComputeNode *> members) const override {
return computePlanningCost(
members, countInternalEdges(ctx.blockAnalysis, members),
/*loopMergeBenefit=*/2, /*liveTileLimit=*/4, /*vfParameterLimit=*/6);
}

PlanningDecision
evaluateAppend(const PlanningContext &ctx,
ArrayRef<const pto::FusionComputeNode *> currentGroup,
Expand Down Expand Up @@ -331,16 +385,10 @@ class ConservativeGreedyCostModel final : public CostModel {
SmallVector<const pto::FusionComputeNode *, 8> proposedGroup(
currentGroup.begin(), currentGroup.end());
proposedGroup.push_back(&candidate);
GroupFootprint footprint = computeGroupFootprint(proposedGroup);

decision.cost.dependencyBenefit =
4 * static_cast<int64_t>(
countEdgesFromGroup(ctx.blockAnalysis, currentGroup, candidate));
decision.cost.loopMergeBenefit = 2;
decision.cost.liveTilePenalty =
std::max<int64_t>(0, static_cast<int64_t>(footprint.liveTileCount) - 4);
decision.cost.vfParameterPenalty = std::max<int64_t>(
0, static_cast<int64_t>(footprint.vfParameterCount) - 6);
decision.cost = computePlanningCost(
proposedGroup,
countEdgesFromGroup(ctx.blockAnalysis, currentGroup, candidate),
/*loopMergeBenefit=*/2, /*liveTileLimit=*/4, /*vfParameterLimit=*/6);
decision.accept = decision.cost.total() > 0;
return decision;
}
Expand All @@ -364,6 +412,15 @@ class ConservativeDAGGreedyCostModel final : public CostModel {
return decision;
}

PlanningCost evaluateGroup(
const PlanningContext &ctx,
ArrayRef<const pto::FusionComputeNode *> members) const override {
return computePlanningCost(
members, countInternalConnections(ctx.blockAnalysis, members),
/*loopMergeBenefit=*/4, /*liveTileLimit=*/10,
/*vfParameterLimit=*/12);
}

PlanningDecision
evaluateAppend(const PlanningContext &ctx,
ArrayRef<const pto::FusionComputeNode *> currentGroup,
Expand Down Expand Up @@ -393,14 +450,9 @@ class ConservativeDAGGreedyCostModel final : public CostModel {
SmallVector<const pto::FusionComputeNode *, 8> proposedGroup(
currentGroup.begin(), currentGroup.end());
proposedGroup.push_back(&candidate);
GroupFootprint footprint = computeGroupFootprint(proposedGroup);

decision.cost.dependencyBenefit = 4 * static_cast<int64_t>(connectionCount);
decision.cost.loopMergeBenefit = 4;
decision.cost.liveTilePenalty = std::max<int64_t>(
0, static_cast<int64_t>(footprint.liveTileCount) - 10);
decision.cost.vfParameterPenalty = std::max<int64_t>(
0, static_cast<int64_t>(footprint.vfParameterCount) - 12);
decision.cost = computePlanningCost(
proposedGroup, connectionCount, /*loopMergeBenefit=*/4,
/*liveTileLimit=*/10, /*vfParameterLimit=*/12);
decision.accept = decision.cost.total() > 0;
return decision;
}
Expand Down Expand Up @@ -463,6 +515,59 @@ class ConservativeGreedyStrategyEngine final : public StrategyEngine {
};

class ConservativeDAGGreedyStrategyEngine final : public StrategyEngine {
// Keep a bounded, deterministically ranked frontier to prevent exponential
// compile-time growth. This selects the best group among retained candidates;
// it does not guarantee the global optimum across every connected subset.
static constexpr unsigned kMaxCandidatesPerGroupSize = 64;

struct GroupCandidate {
SmallVector<const pto::FusionComputeNode *, 8> members;
DenseSet<unsigned> memberIds;
PlanningCost cost;
};

static bool haveSameMembers(const GroupCandidate &lhs,
const GroupCandidate &rhs) {
if (lhs.members.size() != rhs.members.size())
return false;
return llvm::equal(
lhs.members, rhs.members,
[](const pto::FusionComputeNode *lhsNode,
const pto::FusionComputeNode *rhsNode) {
return lhsNode->id == rhsNode->id;
});
}

static bool isBetterCandidate(const GroupCandidate &lhs,
const GroupCandidate &rhs) {
if (lhs.cost.total() != rhs.cost.total())
return lhs.cost.total() > rhs.cost.total();
if (lhs.members.size() != rhs.members.size())
return lhs.members.size() > rhs.members.size();

for (auto [lhsNode, rhsNode] : llvm::zip(lhs.members, rhs.members)) {
if (lhsNode->blockOrder != rhsNode->blockOrder)
return lhsNode->blockOrder < rhsNode->blockOrder;
if (lhsNode->id != rhsNode->id)
return lhsNode->id < rhsNode->id;
}
return false;
}

static void
addOrUpdateCandidate(SmallVectorImpl<GroupCandidate> &candidates,
GroupCandidate candidate) {
auto existing = llvm::find_if(candidates, [&](const GroupCandidate &other) {
return haveSameMembers(candidate, other);
});
if (existing == candidates.end()) {
candidates.push_back(std::move(candidate));
return;
}
if (isBetterCandidate(candidate, *existing))
*existing = std::move(candidate);
}

public:
SmallVector<PlannedFusionGroup, 8>
planBlock(const PlanningContext &ctx,
Expand All @@ -478,36 +583,57 @@ class ConservativeDAGGreedyStrategyEngine final : public StrategyEngine {
if (!seedDecision.accept)
continue;

SmallVector<const pto::FusionComputeNode *, 8> groupMembers;
DenseSet<unsigned> groupNodeIds;
groupMembers.push_back(&seed);
groupNodeIds.insert(seed.id);

bool changed = true;
while (changed) {
changed = false;
for (const pto::FusionComputeNode &candidate :
ctx.blockAnalysis.computeNodes) {
if (assignedNodes.contains(candidate.id) ||
groupNodeIds.contains(candidate.id))
continue;

PlanningDecision appendDecision =
costModel.evaluateAppend(ctx, groupMembers, candidate);
if (!appendDecision.accept)
continue;

groupMembers.push_back(&candidate);
groupNodeIds.insert(candidate.id);
changed = true;
GroupCandidate seedCandidate;
seedCandidate.members.push_back(&seed);
seedCandidate.memberIds.insert(seed.id);
seedCandidate.cost =
costModel.evaluateGroup(ctx, seedCandidate.members);

SmallVector<GroupCandidate, 64> frontier;
frontier.push_back(std::move(seedCandidate));
std::optional<GroupCandidate> bestCandidate;

while (!frontier.empty()) {
SmallVector<GroupCandidate, 64> nextFrontier;
for (const GroupCandidate &current : frontier) {
for (const pto::FusionComputeNode &candidate :
ctx.blockAnalysis.computeNodes) {
if (assignedNodes.contains(candidate.id) ||
current.memberIds.contains(candidate.id))
continue;

PlanningDecision appendDecision =
costModel.evaluateAppend(ctx, current.members, candidate);
if (!appendDecision.accept)
continue;

GroupCandidate successor = current;
successor.members.push_back(&candidate);
successor.members = buildStableInGroupOrder(successor.members);
successor.memberIds.insert(candidate.id);
successor.cost =
costModel.evaluateGroup(ctx, successor.members);
addOrUpdateCandidate(nextFrontier, std::move(successor));
}
}

llvm::stable_sort(nextFrontier, isBetterCandidate);
if (nextFrontier.size() > kMaxCandidatesPerGroupSize)
nextFrontier.resize(kMaxCandidatesPerGroupSize);

for (const GroupCandidate &candidate : nextFrontier)
if (!bestCandidate ||
isBetterCandidate(candidate, *bestCandidate))
bestCandidate = candidate;

frontier = std::move(nextFrontier);
}

if (groupMembers.size() < 2)
if (!bestCandidate || bestCandidate->members.size() < 2)
continue;

PlannedFusionGroup group;
group.members = buildStableInGroupOrder(groupMembers);
group.members = bestCandidate->members;
groups.push_back(group);
for (const pto::FusionComputeNode *member : group.members)
assignedNodes.insert(member->id);
Expand Down
75 changes: 75 additions & 0 deletions test/lit/tile_fusion/fusion_plan_cost_guided_search.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
// CANN Open Software License Agreement Version 2.0 (the "License").
// Please refer to the License for details. You may not use this file except in compliance with the License.
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
// See LICENSE in the root of the software repository for the full text of the License.

// Guards cost-guided FusionPlan search across competing block-local DFG paths.
//
// RUN: ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-fusion-plan -o /dev/null 2>&1 | FileCheck %s

module {
// The six-op prefix has live-tile/VF-parameter footprint 12. The first
// branch is immediately profitable, but raises the footprint enough to
// reject both later appends:
//
// cost(prefix + first) = 24
// cost(prefix + second + join) = 28
//
// Candidate-group search must therefore choose the later, higher-cost path.
// The independent two-op component must still form the next fusion group.
func.func @prefer_higher_group_cost(
%arg0: !pto.tile_buf<vec, 32x32xf32>,
%arg1: !pto.tile_buf<vec, 32x32xf32>,
%arg2: !pto.tile_buf<vec, 32x32xf32>,
%arg3: !pto.tile_buf<vec, 32x32xf32>,
%arg4: !pto.tile_buf<vec, 32x32xf32>,
%arg5: !pto.tile_buf<vec, 32x32xf32>,
%arg6: !pto.tile_buf<vec, 32x32xf32>,
%arg7: !pto.tile_buf<vec, 32x32xf32>,
%arg8: !pto.tile_buf<vec, 32x32xf32>) {
%tmp0 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%tmp1 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%tmp2 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%tmp3 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%tmp4 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%tmp5 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%first = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%second = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%join = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%independent0 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>
%independent1 = pto.alloc_tile : !pto.tile_buf<vec, 32x32xf32>

pto.tadd ins(%arg0, %arg1 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%tmp0 : !pto.tile_buf<vec, 32x32xf32>)
pto.tadd ins(%tmp0, %arg2 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%tmp1 : !pto.tile_buf<vec, 32x32xf32>)
pto.tadd ins(%tmp1, %arg3 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%tmp2 : !pto.tile_buf<vec, 32x32xf32>)
pto.tadd ins(%tmp2, %arg4 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%tmp3 : !pto.tile_buf<vec, 32x32xf32>)
pto.tadd ins(%tmp3, %arg5 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%tmp4 : !pto.tile_buf<vec, 32x32xf32>)
pto.texp ins(%tmp4 : !pto.tile_buf<vec, 32x32xf32>) outs(%tmp5 : !pto.tile_buf<vec, 32x32xf32>)
pto.texp ins(%tmp5 : !pto.tile_buf<vec, 32x32xf32>) outs(%first : !pto.tile_buf<vec, 32x32xf32>)
pto.tadd ins(%tmp5, %arg6 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%second : !pto.tile_buf<vec, 32x32xf32>)
pto.tadd ins(%second, %tmp0 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%join : !pto.tile_buf<vec, 32x32xf32>)

pto.tadd ins(%arg7, %arg8 : !pto.tile_buf<vec, 32x32xf32>, !pto.tile_buf<vec, 32x32xf32>) outs(%independent0 : !pto.tile_buf<vec, 32x32xf32>)
pto.texp ins(%independent0 : !pto.tile_buf<vec, 32x32xf32>) outs(%independent1 : !pto.tile_buf<vec, 32x32xf32>)
return
}

}

// CHECK: // -----// IR Dump After FusionPlan (pto-fusion-plan) //----- //
// CHECK-LABEL: func.func @prefer_higher_group_cost(
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 0 : i64}
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 1 : i64}
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 2 : i64}
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 3 : i64}
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 4 : i64}
// CHECK: pto.texp{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 5 : i64}
// CHECK: pto.texp
// CHECK-NOT: pto.fusion.group_id
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 6 : i64}
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 0 : i64, pto.fusion.order = 7 : i64}
// CHECK: pto.tadd{{.*}}{pto.fusion.group_id = 1 : i64, pto.fusion.order = 0 : i64}
// CHECK: pto.texp{{.*}}{pto.fusion.group_id = 1 : i64, pto.fusion.order = 1 : i64}
Loading
Loading