llvm / llvm-project / 9e8a71caf02a503006eb08b5cd291adeaa20a9c2 / . / mlir / lib / Conversion / PDLToPDLInterp / RootOrdering.h

//===- RootOrdering.h - Optimal root ordering ------------------*- C++ -*-===// | |

// | |

// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | |

// See https://llvm.org/LICENSE.txt for license information. | |

// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |

// | |

//===----------------------------------------------------------------------===// | |

// | |

// This file contains definition for a cost graph over candidate roots and | |

// an implementation of an algorithm to determine the optimal ordering over | |

// these roots. Each edge in this graph indicates that the target root can be | |

// connected (via a chain of positions) to the source root, and their cost | |

// indicates the estimated cost of such traversal. The optimal root ordering | |

// is then formulated as that of finding a spanning arborescence (i.e., a | |

// directed spanning tree) of minimal weight. | |

// | |

//===----------------------------------------------------------------------===// | |

#ifndef MLIR_LIB_CONVERSION_PDLTOPDLINTERP_ROOTORDERING_H_ | |

#define MLIR_LIB_CONVERSION_PDLTOPDLINTERP_ROOTORDERING_H_ | |

#include "mlir/IR/Value.h" | |

#include "llvm/ADT/DenseMap.h" | |

#include "llvm/ADT/SmallVector.h" | |

#include <functional> | |

#include <vector> | |

namespace mlir { | |

namespace pdl_to_pdl_interp { | |

/// The information associated with an edge in the cost graph. Each node in | |

/// the cost graph corresponds to a candidate root detected in the pdl.pattern, | |

/// and each edge in the cost graph corresponds to connecting the two candidate | |

/// roots via a chain of operations. The cost of an edge is the smallest number | |

/// of upward traversals required to go from the source to the target root, and | |

/// the connector is a `Value` in the intersection of the two subtrees rooted at | |

/// the source and target root that results in that smallest number of upward | |

/// traversals. Consider the following pattern with 3 roots op3, op4, and op5: | |

/// | |

/// argA ---> op1 ---> op2 ---> op3 ---> res3 | |

/// ^ ^ | |

/// | | | |

/// argB argC | |

/// | | | |

/// v v | |

/// res4 <--- op4 op5 ---> res5 | |

/// ^ ^ | |

/// | | | |

/// op6 op7 | |

/// | |

/// 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 { | |

/// 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. | |

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>>; | |

/// 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 | |

/// could be reduced to O(N^3) with a more efficient algorithm. However, note | |

/// that the underlying implementation is very efficient, and N in our | |

/// instances tends to be very small (<10). | |

class OptimalBranching { | |

public: | |

/// A list of edges (child, parent). | |

using EdgeList = std::vector<std::pair<Value, Value>>; | |

/// Constructs the solver for the given graph and root value. | |

OptimalBranching(RootOrderingGraph graph, Value root); | |

/// Runs the Edmonds' algorithm for the current `graph`, returning the total | |

/// cost of the minimum-weight spanning arborescence (sum of the edge costs). | |

/// This function first determines the optimal local choice of the parents | |

/// and stores this choice in the `parents` mapping. If this choice results | |

/// in an acyclic graph, the function returns immediately. Otherwise, it | |

/// takes an arbitrary cycle, contracts it, and recurses on the new graph | |

/// (which is guaranteed to have fewer nodes than we began with). After we | |

/// return from recursion, we redirect the edges to/from the contracted node, | |

/// so the `parents` map contains a valid solution for the current graph. | |

unsigned solve(); | |

/// Returns the computed parent map. This is the unique predecessor for each | |

/// node (root) in the optimal branching. | |

const DenseMap<Value, Value> &getRootOrderingParents() const { | |

return parents; | |

} | |

/// Returns the computed edges as visited in the preorder traversal. | |

/// The specified array determines the order for breaking any ties. | |

EdgeList preOrderTraversal(ArrayRef<Value> nodes) const; | |

private: | |

/// The graph whose optimal branching we wish to determine. | |

RootOrderingGraph graph; | |

/// The root of the optimal branching. | |

Value root; | |

/// The computed parent mapping. This is the unique predecessor for each node | |

/// in the optimal branching. The keys of this map correspond to the keys of | |

/// the outer map of the input graph, and each value is one of the keys of | |

/// the inner map for this node. Also used as an intermediate (possibly | |

/// cyclical) result in the optimal branching algorithm. | |

DenseMap<Value, Value> parents; | |

}; | |

} // end namespace pdl_to_pdl_interp | |

} // end namespace mlir | |

#endif // MLIR_CONVERSION_PDLTOPDLINTERP_ROOTORDERING_H_ |