231 lines · cpp
1//===- RootOrdering.cpp - Optimal root ordering ---------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// An implementation of Edmonds' optimal branching algorithm. This is a10// directed analogue of the minimum spanning tree problem for a given root.11//12//===----------------------------------------------------------------------===//13 14#include "RootOrdering.h"15 16#include "llvm/ADT/DenseMap.h"17#include "llvm/ADT/SmallVector.h"18#include <utility>19 20using namespace mlir;21using namespace mlir::pdl_to_pdl_interp;22 23/// Returns the cycle implied by the specified parent relation, starting at the24/// given node.25static SmallVector<Value> getCycle(const DenseMap<Value, Value> &parents,26 Value rep) {27 SmallVector<Value> cycle;28 Value node = rep;29 do {30 cycle.push_back(node);31 node = parents.lookup(node);32 assert(node && "got an empty value in the cycle");33 } while (node != rep);34 return cycle;35}36 37/// Contracts the specified cycle in the given graph in-place.38/// The parentsCost map specifies, for each node in the cycle, the lowest cost39/// among the edges entering that node. Then, the nodes in the cycle C are40/// replaced with a single node v_C (the first node in the cycle). All edges41/// (u, v) entering the cycle, v \in C, are replaced with a single edge42/// (u, v_C) with an appropriately chosen cost, and the selected node v is43/// marked in the output map actualTarget[u]. All edges (u, v) leaving the44/// cycle, u \in C, are replaced with a single edge (v_C, v), and the selected45/// node u is marked in the ouptut map actualSource[v].46static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,47 const DenseMap<Value, unsigned> &parentDepths,48 DenseMap<Value, Value> &actualSource,49 DenseMap<Value, Value> &actualTarget) {50 Value rep = cycle.front();51 DenseSet<Value> cycleSet(cycle.begin(), cycle.end());52 53 // Now, contract the cycle, marking the actual sources and targets.54 DenseMap<Value, RootOrderingEntry> repEntries;55 for (auto outer = graph.begin(), e = graph.end(); outer != e; ++outer) {56 Value target = outer->first;57 if (cycleSet.contains(target)) {58 // Target in the cycle => edges incoming to the cycle or within the cycle.59 unsigned parentDepth = parentDepths.lookup(target);60 for (const auto &inner : outer->second) {61 Value source = inner.first;62 // Ignore edges within the cycle.63 if (cycleSet.contains(source))64 continue;65 66 // Edge incoming to the cycle.67 std::pair<unsigned, unsigned> cost = inner.second.cost;68 assert(parentDepth <= cost.first && "invalid parent depth");69 70 // Subtract the cost of the parent within the cycle from the cost of71 // the edge incoming to the cycle. This update ensures that the cost72 // of the minimum-weight spanning arborescence of the entire graph is73 // the cost of arborescence for the contracted graph plus the cost of74 // the cycle, no matter which edge in the cycle we choose to drop.75 cost.first -= parentDepth;76 auto it = repEntries.find(source);77 if (it == repEntries.end() || it->second.cost > cost) {78 actualTarget[source] = target;79 // Do not bother populating the connector (the connector is only80 // relevant for the final traversal, not for the optimal branching).81 repEntries[source].cost = cost;82 }83 }84 // Erase the node in the cycle.85 graph.erase(outer);86 } else {87 // Target not in cycle => edges going away from or unrelated to the cycle.88 DenseMap<Value, RootOrderingEntry> &entries = outer->second;89 Value bestSource;90 std::pair<unsigned, unsigned> bestCost;91 auto inner = entries.begin(), innerE = entries.end();92 while (inner != innerE) {93 Value source = inner->first;94 if (cycleSet.contains(source)) {95 // Going-away edge => get its cost and erase it.96 if (!bestSource || bestCost > inner->second.cost) {97 bestSource = source;98 bestCost = inner->second.cost;99 }100 entries.erase(inner++);101 } else {102 ++inner;103 }104 }105 106 // There were going-away edges, contract them.107 if (bestSource) {108 entries[rep].cost = bestCost;109 actualSource[target] = bestSource;110 }111 }112 }113 114 // Store the edges to the representative.115 graph[rep] = std::move(repEntries);116}117 118OptimalBranching::OptimalBranching(RootOrderingGraph graph, Value root)119 : graph(std::move(graph)), root(root) {}120 121unsigned OptimalBranching::solve() {122 // Initialize the parents and total cost.123 parents.clear();124 parents[root] = Value();125 unsigned totalCost = 0;126 127 // A map that stores the cost of the optimal local choice for each node128 // in a directed cycle. This map is cleared every time we seed the search.129 DenseMap<Value, unsigned> parentDepths;130 parentDepths.reserve(graph.size());131 132 // Determine if the optimal local choice results in an acyclic graph. This is133 // done by computing the optimal local choice and traversing up the computed134 // parents. On success, `parents` will contain the parent of each node.135 for (const auto &outer : graph) {136 Value node = outer.first;137 if (parents.count(node)) // already visited138 continue;139 140 // Follow the trail of best sources until we reach an already visited node.141 // The code will assert if we cannot reach an already visited node, i.e.,142 // the graph is not strongly connected.143 parentDepths.clear();144 do {145 auto it = graph.find(node);146 assert(it != graph.end() && "the graph is not strongly connected");147 148 // Find the best local parent, taking into account both the depth and the149 // tie breaking rules.150 Value &bestSource = parents[node];151 std::pair<unsigned, unsigned> bestCost;152 for (const auto &inner : it->second) {153 const RootOrderingEntry &entry = inner.second;154 if (!bestSource /* initial */ || bestCost > entry.cost) {155 bestSource = inner.first;156 bestCost = entry.cost;157 }158 }159 assert(bestSource && "the graph is not strongly connected");160 parentDepths[node] = bestCost.first;161 node = bestSource;162 totalCost += bestCost.first;163 } while (!parents.count(node));164 165 // If we reached a non-root node, we have a cycle.166 if (parentDepths.count(node)) {167 // Determine the cycle starting at the representative node.168 SmallVector<Value> cycle = getCycle(parents, node);169 170 // The following maps disambiguate the source / target of the edges171 // going out of / into the cycle.172 DenseMap<Value, Value> actualSource, actualTarget;173 174 // Contract the cycle and recurse.175 contract(graph, cycle, parentDepths, actualSource, actualTarget);176 totalCost = solve();177 178 // Redirect the going-away edges.179 for (auto &p : parents)180 if (p.second == node)181 // The parent is the node representating the cycle; replace it182 // with the actual (best) source in the cycle.183 p.second = actualSource.lookup(p.first);184 185 // Redirect the unique incoming edge and copy the cycle.186 Value parent = parents.lookup(node);187 Value entry = actualTarget.lookup(parent);188 cycle.push_back(node); // complete the cycle189 for (size_t i = 0, e = cycle.size() - 1; i < e; ++i) {190 totalCost += parentDepths.lookup(cycle[i]);191 if (cycle[i] == entry)192 parents[cycle[i]] = parent; // break the cycle193 else194 parents[cycle[i]] = cycle[i + 1];195 }196 197 // `parents` has a complete solution.198 break;199 }200 }201 202 return totalCost;203}204 205OptimalBranching::EdgeList206OptimalBranching::preOrderTraversal(ArrayRef<Value> nodes) const {207 // Invert the parent mapping.208 DenseMap<Value, std::vector<Value>> children;209 for (Value node : nodes) {210 if (node != root) {211 Value parent = parents.lookup(node);212 assert(parent && "invalid parent");213 children[parent].push_back(node);214 }215 }216 217 // The result which simultaneously acts as a queue.218 EdgeList result;219 result.reserve(nodes.size());220 result.emplace_back(root, Value());221 222 // Perform a BFS, pushing into the queue.223 for (size_t i = 0; i < result.size(); ++i) {224 Value node = result[i].first;225 for (Value child : children[node])226 result.emplace_back(child, node);227 }228 229 return result;230}231