blob: 4382753458644bd032371730c57a6847f3c72ce1 [file] [log] [blame]
//===- RootOrdering.cpp - Optimal root ordering ---------------------------===//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// An implementation of Edmonds' optimal branching algorithm. This is a
// directed analogue of the minimum spanning tree problem for a given root.
#include "RootOrdering.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallVector.h"
#include <queue>
#include <utility>
using namespace mlir;
using namespace mlir::pdl_to_pdl_interp;
/// Returns the cycle implied by the specified parent relation, starting at the
/// given node.
static SmallVector<Value> getCycle(const DenseMap<Value, Value> &parents,
Value rep) {
SmallVector<Value> cycle;
Value node = rep;
do {
node = parents.lookup(node);
assert(node && "got an empty value in the cycle");
} while (node != rep);
return cycle;
/// Contracts the specified cycle in the given graph in-place.
/// The parentsCost map specifies, for each node in the cycle, the lowest cost
/// 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,
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;
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);
for (const auto &inner : outer->second) {
Value source = inner.first;
// Ignore edges within the cycle.
if (cycleSet.contains(source))
// Edge incoming to the cycle.
std::pair<unsigned, unsigned> cost = inner.second.cost;
assert(parentCost <= cost.first && "invalid parent cost");
// 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) {
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;
// Erase the node in the cycle.
} else {
// Target not in cycle => edges going away from or unrelated to the cycle.
DenseMap<Value, RootOrderingCost> &costs = outer->second;
Value bestSource;
std::pair<unsigned, unsigned> bestCost;
auto inner = costs.begin(), inner_e = costs.end();
while (inner != inner_e) {
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;
} else {
// There were going-away edges, contract them.
if (bestSource) {
costs[rep].cost = bestCost;
actualSource[target] = bestSource;
// Store the edges to the representative.
graph[rep] = std::move(repCosts);
OptimalBranching::OptimalBranching(RootOrderingGraph graph, Value root)
: graph(std::move(graph)), root(root) {}
unsigned OptimalBranching::solve() {
// Initialize the parents and total cost.
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;
// 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
// 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.
do {
auto it = graph.find(node);
assert(it != graph.end() && "the graph is not strongly connected");
Value &bestSource = parents[node];
unsigned &bestCost = parentCosts[node];
for (const auto &inner : it->second) {
const RootOrderingCost &cost = inner.second;
if (!bestSource /* initial */ || bestCost > cost.cost.first) {
bestSource = inner.first;
bestCost = cost.cost.first;
assert(bestSource && "the graph is not strongly connected");
node = bestSource;
totalCost += bestCost;
} while (!parents.count(node));
// If we reached a non-root node, we have a cycle.
if (parentCosts.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);
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]);
if (cycle[i] == entry)
parents[cycle[i]] = parent; // break the cycle
parents[cycle[i]] = cycle[i + 1];
// `parents` has a complete solution.
return totalCost;
OptimalBranching::preOrderTraversal(ArrayRef<Value> nodes) const {
// Invert the parent mapping.
DenseMap<Value, std::vector<Value>> children;
for (Value node : nodes) {
if (node != root) {
Value parent = parents.lookup(node);
assert(parent && "invalid parent");
// The result which simultaneously acts as a queue.
EdgeList result;
result.emplace_back(root, Value());
// Perform a BFS, pushing into the queue.
for (size_t i = 0; i < result.size(); ++i) {
Value node = result[i].first;
for (Value child : children[node])
result.emplace_back(child, node);
return result;