[MLIR][PDL] Clear up the terminology in the root ordering graph.

[sfuniak]

Summary

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.

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

for (const auto &connectorRootsDepths : connectorsRootsDepths) {
     if (rootsDepths.size() == 1)
       continue;
 
     for (const RootDepth &p : rootsDepths) {
       for (const RootDepth &q : rootsDepths) {
         if (&p == &q)
           continue;
         // Insert or retrieve the property of edge from p to q.
-        RootOrderingCost &cost = graph[q.root][p.root];
-        if (!cost.connector /* new edge */ || cost.cost.first > q.depth) {
-          if (!cost.connector)
-            cost.cost.second = nextID++;
-          cost.cost.first = q.depth;
-          cost.connector = value;
+        RootOrderingEntry &entry = graph[q.root][p.root];
+        if (!entry.connector /* new edge */ || entry.cost.first > q.depth) {
+          if (!entry.connector)
+            entry.cost.second = nextID++;
+          entry.cost.first = q.depth;
+          entry.connector = value;
         }
       }
     }
   }
 
   assert((llvm::hasSingleElement(roots) || graph.size() == roots.size()) &&
          "the pattern contains a candidate root disconnected from the others");
 }
static Value buildPredicateList(pdl::PatternOp pattern,
   RootOrderingGraph graph;
   ParentMaps parentMaps;
   buildCostGraph(roots, graph, parentMaps);
   LLVM_DEBUG({
     llvm::dbgs() << "Graph:\n";
     for (auto &target : graph) {
       llvm::dbgs() << "  * " << target.first << "\n";
       for (auto &source : target.second) {
-        RootOrderingCost c = source.second;
-        llvm::dbgs() << "      <- " << source.first << ": " << c.cost.first
-                     << ":" << c.cost.second << " via " << c.connector.getLoc()
-                     << "\n";
+        RootOrderingEntry &entry = source.second;
+        llvm::dbgs() << "      <- " << source.first << ": " << entry.cost.first
+                     << ":" << entry.cost.second << " via "
+                     << entry.connector.getLoc() << "\n";
       }
     }
   });
 
   // Solve the optimal branching problem for each candidate root, or use the
   // provided one.
   Value bestRoot = pattern.getRewriter().root();
   OptimalBranching::EdgeList bestEdges;

mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.h

 /// The cost of the edge op3 -> op4 is 1 (the upward traversal argB -> op4),
 /// with argB being the connector `Value` and similarly for op3 -> op5 (cost 1,
 /// connector argC). The cost of the edge op4 -> op3 is 3 (upward traversals
 /// argB -> op1 -> op2 -> op3, connector argB), while the cost of edge op5 ->
 /// op3 is 2 (uwpard traversals argC -> op2 -> op3). There are no edges between
 /// op4 and op5 in the cost graph, because the subtrees rooted at these two
 /// roots do not intersect. It is easy to see that the optimal root for this
 /// pattern is op3, resulting in the spanning arborescence op3 -> {op4, op5}.
-struct RootOrderingCost {
+struct RootOrderingEntry {
   /// The depth of the connector `Value` w.r.t. the target root.
   ///
-  /// This is a pair where the first entry is the actual cost, and the second
-  /// entry is a priority for breaking ties (with 0 being the highest).
-  /// Typically, the priority is a unique edge ID.
+  /// This is a pair where the first value is the additive cost (the depth of
+  /// the connector), and the second value is a priority for breaking ties
+  /// (with 0 being the highest). Typically, the priority is a unique edge ID.
   std::pair<unsigned, unsigned> cost;
 
   /// The connector value in the intersection of the two subtrees rooted at
   /// the source and target root that results in that smallest depth w.r.t.
   /// the target root.
   Value connector;
 };
 
 /// A directed graph representing the cost of ordering the roots in the
 /// predicate tree. It is represented as an adjacency map, where the outer map
 /// is indexed by the target node, and the inner map is indexed by the source
 /// node. Each edge is associated with a cost and the underlying connector
 /// value.
-using RootOrderingGraph = DenseMap<Value, DenseMap<Value, RootOrderingCost>>;
+using RootOrderingGraph = DenseMap<Value, DenseMap<Value, RootOrderingEntry>>;
 
 /// The optimal branching algorithm solver. This solver accepts a graph and the
 /// root in its constructor, and is invoked via the solve() member function.
 /// This is a direct implementation of the Edmonds' algorithm, see
 /// https://en.wikipedia.org/wiki/Edmonds%27_algorithm. The worst-case
 /// computational complexity of this algorithm is O(N^3), for a single root.
 /// The PDL-to-PDLInterp lowering calls this N times (once for each candidate
 /// root), so the overall complexity root ordering is O(N^4). If needed, this

mlir/lib/Conversion/PDLToPDLInterp/RootOrdering.cpp

 /// among the edges entering that node. Then, the nodes in the cycle C are
 /// replaced with a single node v_C (the first node in the cycle). All edges
 /// (u, v) entering the cycle, v \in C, are replaced with a single edge
 /// (u, v_C) with an appropriately chosen cost, and the selected node v is
 /// marked in the output map actualTarget[u]. All edges (u, v) leaving the
 /// cycle, u \in C, are replaced with a single edge (v_C, v), and the selected
 /// node u is marked in the ouptut map actualSource[v].
 static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
-                     const DenseMap<Value, unsigned> &parentCosts,
+                     const DenseMap<Value, unsigned> &parentDepths,
                      DenseMap<Value, Value> &actualSource,
                      DenseMap<Value, Value> &actualTarget) {
   Value rep = cycle.front();
   DenseSet<Value> cycleSet(cycle.begin(), cycle.end());
 
   // Now, contract the cycle, marking the actual sources and targets.
-  DenseMap<Value, RootOrderingCost> repCosts;
+  DenseMap<Value, RootOrderingEntry> repEntries;
   for (auto outer = graph.begin(), e = graph.end(); outer != e; ++outer) {
     Value target = outer->first;
     if (cycleSet.contains(target)) {
       // Target in the cycle => edges incoming to the cycle or within the cycle.
-      unsigned parentCost = parentCosts.lookup(target);
+      unsigned parentDepth = parentDepths.lookup(target);
       for (const auto &inner : outer->second) {
         Value source = inner.first;
         // Ignore edges within the cycle.
         if (cycleSet.contains(source))
           continue;
 
         // Edge incoming to the cycle.
         std::pair<unsigned, unsigned> cost = inner.second.cost;
-        assert(parentCost <= cost.first && "invalid parent cost");
+        assert(parentDepth <= cost.first && "invalid parent depth");
 
         // Subtract the cost of the parent within the cycle from the cost of
         // the edge incoming to the cycle. This update ensures that the cost
         // of the minimum-weight spanning arborescence of the entire graph is
         // the cost of arborescence for the contracted graph plus the cost of
         // the cycle, no matter which edge in the cycle we choose to drop.
-        cost.first -= parentCost;
-        auto it = repCosts.find(source);
-        if (it == repCosts.end() || it->second.cost > cost) {
+        cost.first -= parentDepth;
+        auto it = repEntries.find(source);
+        if (it == repEntries.end() || it->second.cost > cost) {
           actualTarget[source] = target;
           // Do not bother populating the connector (the connector is only
           // relevant for the final traversal, not for the optimal branching).
-          repCosts[source].cost = cost;
+          repEntries[source].cost = cost;
         }
       }
       // Erase the node in the cycle.
       graph.erase(outer);
     } else {
       // Target not in cycle => edges going away from or unrelated to the cycle.
-      DenseMap<Value, RootOrderingCost> &costs = outer->second;
+      DenseMap<Value, RootOrderingEntry> &entries = outer->second;
       Value bestSource;
       std::pair<unsigned, unsigned> bestCost;
-      auto inner = costs.begin(), innerE = costs.end();
+      auto inner = entries.begin(), innerE = entries.end();
       while (inner != innerE) {
         Value source = inner->first;
         if (cycleSet.contains(source)) {
           // Going-away edge => get its cost and erase it.
           if (!bestSource || bestCost > inner->second.cost) {
             bestSource = source;
             bestCost = inner->second.cost;
           }
-          costs.erase(inner++);
+          entries.erase(inner++);
         } else {
           ++inner;
         }
       }
 
       // There were going-away edges, contract them.
       if (bestSource) {
-        costs[rep].cost = bestCost;
+        entries[rep].cost = bestCost;
         actualSource[target] = bestSource;
       }
     }
   }
 
   // Store the edges to the representative.
-  graph[rep] = std::move(repCosts);
+  graph[rep] = std::move(repEntries);
 }
 
 OptimalBranching::OptimalBranching(RootOrderingGraph graph, Value root)
     : graph(std::move(graph)), root(root) {}
 
 unsigned OptimalBranching::solve() {
   // Initialize the parents and total cost.
   parents.clear();
   parents[root] = Value();
   unsigned totalCost = 0;
 
   // A map that stores the cost of the optimal local choice for each node
   // in a directed cycle. This map is cleared every time we seed the search.
-  DenseMap<Value, unsigned> parentCosts;
-  parentCosts.reserve(graph.size());
+  DenseMap<Value, unsigned> parentDepths;
+  parentDepths.reserve(graph.size());
 
   // Determine if the optimal local choice results in an acyclic graph. This is
   // done by computing the optimal local choice and traversing up the computed
   // parents. On success, `parents` will contain the parent of each node.
   for (const auto &outer : graph) {
     Value node = outer.first;
     if (parents.count(node)) // already visited
       continue;
 
     // Follow the trail of best sources until we reach an already visited node.
     // The code will assert if we cannot reach an already visited node, i.e.,
     // the graph is not strongly connected.
-    parentCosts.clear();
+    parentDepths.clear();
     do {
       auto it = graph.find(node);
       assert(it != graph.end() && "the graph is not strongly connected");
 
+      // Find the best local parent, taking into account both the depth and the
+      // tie breaking rules.
       Value &bestSource = parents[node];
-      unsigned &bestCost = parentCosts[node];
+      std::pair<unsigned, unsigned> bestCost;
       for (const auto &inner : it->second) {
-        const RootOrderingCost &cost = inner.second;
-        if (!bestSource /* initial */ || bestCost > cost.cost.first) {
+        const RootOrderingEntry &entry = inner.second;
+        if (!bestSource /* initial */ || bestCost > entry.cost) {
           bestSource = inner.first;
-          bestCost = cost.cost.first;
+          bestCost = entry.cost;
         }
       }
       assert(bestSource && "the graph is not strongly connected");
+      parentDepths[node] = bestCost.first;
       node = bestSource;
-      totalCost += bestCost;
+      totalCost += bestCost.first;
     } while (!parents.count(node));
 
     // If we reached a non-root node, we have a cycle.
-    if (parentCosts.count(node)) {
+    if (parentDepths.count(node)) {
       // Determine the cycle starting at the representative node.
       SmallVector<Value> cycle = getCycle(parents, node);
 
       // The following maps disambiguate the source / target of the edges
       // going out of / into the cycle.
       DenseMap<Value, Value> actualSource, actualTarget;
 
       // Contract the cycle and recurse.
-      contract(graph, cycle, parentCosts, actualSource, actualTarget);
+      contract(graph, cycle, parentDepths, actualSource, actualTarget);
       totalCost = solve();
 
       // Redirect the going-away edges.
       for (auto &p : parents)
         if (p.second == node)
           // The parent is the node representating the cycle; replace it
           // with the actual (best) source in the cycle.
           p.second = actualSource.lookup(p.first);
 
       // Redirect the unique incoming edge and copy the cycle.
       Value parent = parents.lookup(node);
       Value entry = actualTarget.lookup(parent);
       cycle.push_back(node); // complete the cycle
       for (size_t i = 0, e = cycle.size() - 1; i < e; ++i) {
-        totalCost += parentCosts.lookup(cycle[i]);
+        totalCost += parentDepths.lookup(cycle[i]);
         if (cycle[i] == entry)
           parents[cycle[i]] = parent; // break the cycle
         else
           parents[cycle[i]] = cycle[i + 1];
       }
 
       // `parents` has a complete solution.
       break;