Skip to content

Commit 9eb8e7b

Browse files
sfuniakMogball
authored andcommitted
[MLIR][PDL] Clear up the terminology in the root ordering graph.
Previously, we defined a struct named `RootOrderingCost`, which stored the cost (a pair consisting of the depth of the connector and a tie breaking ID), as well as the connector itself. This created some confusion, because we would sometimes write, e.g., `cost.cost.first` (the first `cost` referring to the struct, the second one referring to the `cost` field, and `first` referring to the depth). In order to address this confusion, here we rename `RootOrderingCost` to `RootOrderingEntry` (keeping the fields and their names as-is). This clarification exposed non-determinism in the optimal branching algorithm. When choosing the best local parent, we were previuosly only considering its depth (`cost.first`) and not the tie-breaking ID (`cost.second`). This led to non-deterministic choice of the parent when multiple potential parents had the same depth. The solution is to compare both the depth and the tie-breaking ID. Testing: Rely on existing unit tests. Non-detgerminism is hard to unit-test. Reviewed By: rriddle, Mogball Differential Revision: https://reviews.llvm.org/D116079
1 parent 42ac4f3 commit 9eb8e7b

File tree

3 files changed

+42
-39
lines changed

3 files changed

+42
-39
lines changed

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -487,12 +487,12 @@ static void buildCostGraph(ArrayRef<Value> roots, RootOrderingGraph &graph,
487487
if (&p == &q)
488488
continue;
489489
// Insert or retrieve the property of edge from p to q.
490-
RootOrderingCost &cost = graph[q.root][p.root];
491-
if (!cost.connector /* new edge */ || cost.cost.first > q.depth) {
492-
if (!cost.connector)
493-
cost.cost.second = nextID++;
494-
cost.cost.first = q.depth;
495-
cost.connector = value;
490+
RootOrderingEntry &entry = graph[q.root][p.root];
491+
if (!entry.connector /* new edge */ || entry.cost.first > q.depth) {
492+
if (!entry.connector)
493+
entry.cost.second = nextID++;
494+
entry.cost.first = q.depth;
495+
entry.connector = value;
496496
}
497497
}
498498
}
@@ -570,10 +570,10 @@ static Value buildPredicateList(pdl::PatternOp pattern,
570570
for (auto &target : graph) {
571571
llvm::dbgs() << " * " << target.first << "\n";
572572
for (auto &source : target.second) {
573-
RootOrderingCost c = source.second;
574-
llvm::dbgs() << " <- " << source.first << ": " << c.cost.first
575-
<< ":" << c.cost.second << " via " << c.connector.getLoc()
576-
<< "\n";
573+
RootOrderingEntry &entry = source.second;
574+
llvm::dbgs() << " <- " << source.first << ": " << entry.cost.first
575+
<< ":" << entry.cost.second << " via "
576+
<< entry.connector.getLoc() << "\n";
577577
}
578578
}
579579
});

mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -46,19 +46,19 @@ static SmallVector<Value> getCycle(const DenseMap<Value, Value> &parents,
4646
/// cycle, u \in C, are replaced with a single edge (v_C, v), and the selected
4747
/// node u is marked in the ouptut map actualSource[v].
4848
static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
49-
const DenseMap<Value, unsigned> &parentCosts,
49+
const DenseMap<Value, unsigned> &parentDepths,
5050
DenseMap<Value, Value> &actualSource,
5151
DenseMap<Value, Value> &actualTarget) {
5252
Value rep = cycle.front();
5353
DenseSet<Value> cycleSet(cycle.begin(), cycle.end());
5454

5555
// Now, contract the cycle, marking the actual sources and targets.
56-
DenseMap<Value, RootOrderingCost> repCosts;
56+
DenseMap<Value, RootOrderingEntry> repEntries;
5757
for (auto outer = graph.begin(), e = graph.end(); outer != e; ++outer) {
5858
Value target = outer->first;
5959
if (cycleSet.contains(target)) {
6060
// Target in the cycle => edges incoming to the cycle or within the cycle.
61-
unsigned parentCost = parentCosts.lookup(target);
61+
unsigned parentDepth = parentDepths.lookup(target);
6262
for (const auto &inner : outer->second) {
6363
Value source = inner.first;
6464
// Ignore edges within the cycle.
@@ -67,30 +67,30 @@ static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
6767

6868
// Edge incoming to the cycle.
6969
std::pair<unsigned, unsigned> cost = inner.second.cost;
70-
assert(parentCost <= cost.first && "invalid parent cost");
70+
assert(parentDepth <= cost.first && "invalid parent depth");
7171

7272
// Subtract the cost of the parent within the cycle from the cost of
7373
// the edge incoming to the cycle. This update ensures that the cost
7474
// of the minimum-weight spanning arborescence of the entire graph is
7575
// the cost of arborescence for the contracted graph plus the cost of
7676
// the cycle, no matter which edge in the cycle we choose to drop.
77-
cost.first -= parentCost;
78-
auto it = repCosts.find(source);
79-
if (it == repCosts.end() || it->second.cost > cost) {
77+
cost.first -= parentDepth;
78+
auto it = repEntries.find(source);
79+
if (it == repEntries.end() || it->second.cost > cost) {
8080
actualTarget[source] = target;
8181
// Do not bother populating the connector (the connector is only
8282
// relevant for the final traversal, not for the optimal branching).
83-
repCosts[source].cost = cost;
83+
repEntries[source].cost = cost;
8484
}
8585
}
8686
// Erase the node in the cycle.
8787
graph.erase(outer);
8888
} else {
8989
// Target not in cycle => edges going away from or unrelated to the cycle.
90-
DenseMap<Value, RootOrderingCost> &costs = outer->second;
90+
DenseMap<Value, RootOrderingEntry> &entries = outer->second;
9191
Value bestSource;
9292
std::pair<unsigned, unsigned> bestCost;
93-
auto inner = costs.begin(), innerE = costs.end();
93+
auto inner = entries.begin(), innerE = entries.end();
9494
while (inner != innerE) {
9595
Value source = inner->first;
9696
if (cycleSet.contains(source)) {
@@ -99,22 +99,22 @@ static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
9999
bestSource = source;
100100
bestCost = inner->second.cost;
101101
}
102-
costs.erase(inner++);
102+
entries.erase(inner++);
103103
} else {
104104
++inner;
105105
}
106106
}
107107

108108
// There were going-away edges, contract them.
109109
if (bestSource) {
110-
costs[rep].cost = bestCost;
110+
entries[rep].cost = bestCost;
111111
actualSource[target] = bestSource;
112112
}
113113
}
114114
}
115115

116116
// Store the edges to the representative.
117-
graph[rep] = std::move(repCosts);
117+
graph[rep] = std::move(repEntries);
118118
}
119119

120120
OptimalBranching::OptimalBranching(RootOrderingGraph graph, Value root)
@@ -128,8 +128,8 @@ unsigned OptimalBranching::solve() {
128128

129129
// A map that stores the cost of the optimal local choice for each node
130130
// in a directed cycle. This map is cleared every time we seed the search.
131-
DenseMap<Value, unsigned> parentCosts;
132-
parentCosts.reserve(graph.size());
131+
DenseMap<Value, unsigned> parentDepths;
132+
parentDepths.reserve(graph.size());
133133

134134
// Determine if the optimal local choice results in an acyclic graph. This is
135135
// done by computing the optimal local choice and traversing up the computed
@@ -142,27 +142,30 @@ unsigned OptimalBranching::solve() {
142142
// Follow the trail of best sources until we reach an already visited node.
143143
// The code will assert if we cannot reach an already visited node, i.e.,
144144
// the graph is not strongly connected.
145-
parentCosts.clear();
145+
parentDepths.clear();
146146
do {
147147
auto it = graph.find(node);
148148
assert(it != graph.end() && "the graph is not strongly connected");
149149

150+
// Find the best local parent, taking into account both the depth and the
151+
// tie breaking rules.
150152
Value &bestSource = parents[node];
151-
unsigned &bestCost = parentCosts[node];
153+
std::pair<unsigned, unsigned> bestCost;
152154
for (const auto &inner : it->second) {
153-
const RootOrderingCost &cost = inner.second;
154-
if (!bestSource /* initial */ || bestCost > cost.cost.first) {
155+
const RootOrderingEntry &entry = inner.second;
156+
if (!bestSource /* initial */ || bestCost > entry.cost) {
155157
bestSource = inner.first;
156-
bestCost = cost.cost.first;
158+
bestCost = entry.cost;
157159
}
158160
}
159161
assert(bestSource && "the graph is not strongly connected");
162+
parentDepths[node] = bestCost.first;
160163
node = bestSource;
161-
totalCost += bestCost;
164+
totalCost += bestCost.first;
162165
} while (!parents.count(node));
163166

164167
// If we reached a non-root node, we have a cycle.
165-
if (parentCosts.count(node)) {
168+
if (parentDepths.count(node)) {
166169
// Determine the cycle starting at the representative node.
167170
SmallVector<Value> cycle = getCycle(parents, node);
168171

@@ -171,7 +174,7 @@ unsigned OptimalBranching::solve() {
171174
DenseMap<Value, Value> actualSource, actualTarget;
172175

173176
// Contract the cycle and recurse.
174-
contract(graph, cycle, parentCosts, actualSource, actualTarget);
177+
contract(graph, cycle, parentDepths, actualSource, actualTarget);
175178
totalCost = solve();
176179

177180
// Redirect the going-away edges.
@@ -186,7 +189,7 @@ unsigned OptimalBranching::solve() {
186189
Value entry = actualTarget.lookup(parent);
187190
cycle.push_back(node); // complete the cycle
188191
for (size_t i = 0, e = cycle.size() - 1; i < e; ++i) {
189-
totalCost += parentCosts.lookup(cycle[i]);
192+
totalCost += parentDepths.lookup(cycle[i]);
190193
if (cycle[i] == entry)
191194
parents[cycle[i]] = parent; // break the cycle
192195
else

mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.h

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,12 @@ namespace pdl_to_pdl_interp {
5656
/// op4 and op5 in the cost graph, because the subtrees rooted at these two
5757
/// roots do not intersect. It is easy to see that the optimal root for this
5858
/// pattern is op3, resulting in the spanning arborescence op3 -> {op4, op5}.
59-
struct RootOrderingCost {
59+
struct RootOrderingEntry {
6060
/// The depth of the connector `Value` w.r.t. the target root.
6161
///
62-
/// This is a pair where the first entry is the actual cost, and the second
63-
/// entry is a priority for breaking ties (with 0 being the highest).
64-
/// Typically, the priority is a unique edge ID.
62+
/// This is a pair where the first value is the additive cost (the depth of
63+
/// the connector), and the second value is a priority for breaking ties
64+
/// (with 0 being the highest). Typically, the priority is a unique edge ID.
6565
std::pair<unsigned, unsigned> cost;
6666

6767
/// The connector value in the intersection of the two subtrees rooted at
@@ -75,7 +75,7 @@ struct RootOrderingCost {
7575
/// is indexed by the target node, and the inner map is indexed by the source
7676
/// node. Each edge is associated with a cost and the underlying connector
7777
/// value.
78-
using RootOrderingGraph = DenseMap<Value, DenseMap<Value, RootOrderingCost>>;
78+
using RootOrderingGraph = DenseMap<Value, DenseMap<Value, RootOrderingEntry>>;
7979

8080
/// The optimal branching algorithm solver. This solver accepts a graph and the
8181
/// root in its constructor, and is invoked via the solve() member function.

0 commit comments

Comments
 (0)