You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
463 lines
19 KiB
463 lines
19 KiB
//===- PredicateTree.cpp - Predicate tree merging -------------------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "PredicateTree.h"
|
|
#include "mlir/Dialect/PDL/IR/PDL.h"
|
|
#include "mlir/Dialect/PDL/IR/PDLTypes.h"
|
|
#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/Interfaces/InferTypeOpInterface.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::pdl_to_pdl_interp;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Predicate List Building
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Compares the depths of two positions.
|
|
static bool comparePosDepth(Position *lhs, Position *rhs) {
|
|
return lhs->getIndex().size() < rhs->getIndex().size();
|
|
}
|
|
|
|
/// Collect the tree predicates anchored at the given value.
|
|
static void getTreePredicates(std::vector<PositionalPredicate> &predList,
|
|
Value val, PredicateBuilder &builder,
|
|
DenseMap<Value, Position *> &inputs,
|
|
Position *pos) {
|
|
// Make sure this input value is accessible to the rewrite.
|
|
auto it = inputs.try_emplace(val, pos);
|
|
|
|
// If this is an input value that has been visited in the tree, add a
|
|
// constraint to ensure that both instances refer to the same value.
|
|
if (!it.second &&
|
|
isa<pdl::AttributeOp, pdl::InputOp, pdl::TypeOp>(val.getDefiningOp())) {
|
|
auto minMaxPositions = std::minmax(pos, it.first->second, comparePosDepth);
|
|
predList.emplace_back(minMaxPositions.second,
|
|
builder.getEqualTo(minMaxPositions.first));
|
|
return;
|
|
}
|
|
|
|
// Check for a per-position predicate to apply.
|
|
switch (pos->getKind()) {
|
|
case Predicates::AttributePos: {
|
|
assert(val.getType().isa<pdl::AttributeType>() &&
|
|
"expected attribute type");
|
|
pdl::AttributeOp attr = cast<pdl::AttributeOp>(val.getDefiningOp());
|
|
predList.emplace_back(pos, builder.getIsNotNull());
|
|
|
|
// If the attribute has a type, add a type constraint.
|
|
if (Value type = attr.type()) {
|
|
getTreePredicates(predList, type, builder, inputs, builder.getType(pos));
|
|
|
|
// Check for a constant value of the attribute.
|
|
} else if (Optional<Attribute> value = attr.value()) {
|
|
predList.emplace_back(pos, builder.getAttributeConstraint(*value));
|
|
}
|
|
break;
|
|
}
|
|
case Predicates::OperandPos: {
|
|
assert(val.getType().isa<pdl::ValueType>() && "expected value type");
|
|
|
|
// Prevent traversal into a null value.
|
|
predList.emplace_back(pos, builder.getIsNotNull());
|
|
|
|
// If this is a typed input, add a type constraint.
|
|
if (auto in = val.getDefiningOp<pdl::InputOp>()) {
|
|
if (Value type = in.type()) {
|
|
getTreePredicates(predList, type, builder, inputs,
|
|
builder.getType(pos));
|
|
}
|
|
|
|
// Otherwise, recurse into the parent node.
|
|
} else if (auto parentOp = val.getDefiningOp<pdl::OperationOp>()) {
|
|
getTreePredicates(predList, parentOp.op(), builder, inputs,
|
|
builder.getParent(cast<OperandPosition>(pos)));
|
|
}
|
|
break;
|
|
}
|
|
case Predicates::OperationPos: {
|
|
assert(val.getType().isa<pdl::OperationType>() && "expected operation");
|
|
pdl::OperationOp op = cast<pdl::OperationOp>(val.getDefiningOp());
|
|
OperationPosition *opPos = cast<OperationPosition>(pos);
|
|
|
|
// Ensure getDefiningOp returns a non-null operation.
|
|
if (!opPos->isRoot())
|
|
predList.emplace_back(pos, builder.getIsNotNull());
|
|
|
|
// Check that this is the correct root operation.
|
|
if (Optional<StringRef> opName = op.name())
|
|
predList.emplace_back(pos, builder.getOperationName(*opName));
|
|
|
|
// Check that the operation has the proper number of operands and results.
|
|
OperandRange operands = op.operands();
|
|
ResultRange results = op.results();
|
|
predList.emplace_back(pos, builder.getOperandCount(operands.size()));
|
|
predList.emplace_back(pos, builder.getResultCount(results.size()));
|
|
|
|
// Recurse into any attributes, operands, or results.
|
|
for (auto it : llvm::zip(op.attributeNames(), op.attributes())) {
|
|
getTreePredicates(
|
|
predList, std::get<1>(it), builder, inputs,
|
|
builder.getAttribute(opPos,
|
|
std::get<0>(it).cast<StringAttr>().getValue()));
|
|
}
|
|
for (auto operandIt : llvm::enumerate(operands))
|
|
getTreePredicates(predList, operandIt.value(), builder, inputs,
|
|
builder.getOperand(opPos, operandIt.index()));
|
|
|
|
// Only recurse into results that are not referenced in the source tree.
|
|
for (auto resultIt : llvm::enumerate(results)) {
|
|
getTreePredicates(predList, resultIt.value(), builder, inputs,
|
|
builder.getResult(opPos, resultIt.index()));
|
|
}
|
|
break;
|
|
}
|
|
case Predicates::ResultPos: {
|
|
assert(val.getType().isa<pdl::ValueType>() && "expected value type");
|
|
pdl::OperationOp parentOp = cast<pdl::OperationOp>(val.getDefiningOp());
|
|
|
|
// Prevent traversing a null value.
|
|
predList.emplace_back(pos, builder.getIsNotNull());
|
|
|
|
// Traverse the type constraint.
|
|
unsigned resultNo = cast<ResultPosition>(pos)->getResultNumber();
|
|
getTreePredicates(predList, parentOp.types()[resultNo], builder, inputs,
|
|
builder.getType(pos));
|
|
break;
|
|
}
|
|
case Predicates::TypePos: {
|
|
assert(val.getType().isa<pdl::TypeType>() && "expected value type");
|
|
pdl::TypeOp typeOp = cast<pdl::TypeOp>(val.getDefiningOp());
|
|
|
|
// Check for a constraint on a constant type.
|
|
if (Optional<Type> type = typeOp.type())
|
|
predList.emplace_back(pos, builder.getTypeConstraint(*type));
|
|
break;
|
|
}
|
|
default:
|
|
llvm_unreachable("unknown position kind");
|
|
}
|
|
}
|
|
|
|
/// Collect all of the predicates related to constraints within the given
|
|
/// pattern operation.
|
|
static void collectConstraintPredicates(
|
|
pdl::PatternOp pattern, std::vector<PositionalPredicate> &predList,
|
|
PredicateBuilder &builder, DenseMap<Value, Position *> &inputs) {
|
|
for (auto op : pattern.body().getOps<pdl::ApplyConstraintOp>()) {
|
|
OperandRange arguments = op.args();
|
|
ArrayAttr parameters = op.constParamsAttr();
|
|
|
|
std::vector<Position *> allPositions;
|
|
allPositions.reserve(arguments.size());
|
|
for (Value arg : arguments)
|
|
allPositions.push_back(inputs.lookup(arg));
|
|
|
|
// Push the constraint to the furthest position.
|
|
Position *pos = *std::max_element(allPositions.begin(), allPositions.end(),
|
|
comparePosDepth);
|
|
PredicateBuilder::Predicate pred =
|
|
builder.getConstraint(op.name(), std::move(allPositions), parameters);
|
|
predList.emplace_back(pos, pred);
|
|
}
|
|
}
|
|
|
|
/// Given a pattern operation, build the set of matcher predicates necessary to
|
|
/// match this pattern.
|
|
static void buildPredicateList(pdl::PatternOp pattern,
|
|
PredicateBuilder &builder,
|
|
std::vector<PositionalPredicate> &predList,
|
|
DenseMap<Value, Position *> &valueToPosition) {
|
|
getTreePredicates(predList, pattern.getRewriter().root(), builder,
|
|
valueToPosition, builder.getRoot());
|
|
collectConstraintPredicates(pattern, predList, builder, valueToPosition);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Pattern Predicate Tree Merging
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
/// This class represents a specific predicate applied to a position, and
|
|
/// provides hashing and ordering operators. This class allows for computing a
|
|
/// frequence sum and ordering predicates based on a cost model.
|
|
struct OrderedPredicate {
|
|
OrderedPredicate(const std::pair<Position *, Qualifier *> &ip)
|
|
: position(ip.first), question(ip.second) {}
|
|
OrderedPredicate(const PositionalPredicate &ip)
|
|
: position(ip.position), question(ip.question) {}
|
|
|
|
/// The position this predicate is applied to.
|
|
Position *position;
|
|
|
|
/// The question that is applied by this predicate onto the position.
|
|
Qualifier *question;
|
|
|
|
/// The first and second order benefit sums.
|
|
/// The primary sum is the number of occurrences of this predicate among all
|
|
/// of the patterns.
|
|
unsigned primary = 0;
|
|
/// The secondary sum is a squared summation of the primary sum of all of the
|
|
/// predicates within each pattern that contains this predicate. This allows
|
|
/// for favoring predicates that are more commonly shared within a pattern, as
|
|
/// opposed to those shared across patterns.
|
|
unsigned secondary = 0;
|
|
|
|
/// A map between a pattern operation and the answer to the predicate question
|
|
/// within that pattern.
|
|
DenseMap<Operation *, Qualifier *> patternToAnswer;
|
|
|
|
/// Returns true if this predicate is ordered before `other`, based on the
|
|
/// cost model.
|
|
bool operator<(const OrderedPredicate &other) const {
|
|
// Sort by:
|
|
// * first and secondary order sums
|
|
// * lower depth
|
|
// * position dependency
|
|
// * predicate dependency.
|
|
auto *otherPos = other.position;
|
|
return std::make_tuple(other.primary, other.secondary,
|
|
otherPos->getIndex().size(), otherPos->getKind(),
|
|
other.question->getKind()) >
|
|
std::make_tuple(primary, secondary, position->getIndex().size(),
|
|
position->getKind(), question->getKind());
|
|
}
|
|
};
|
|
|
|
/// A DenseMapInfo for OrderedPredicate based solely on the position and
|
|
/// question.
|
|
struct OrderedPredicateDenseInfo {
|
|
using Base = DenseMapInfo<std::pair<Position *, Qualifier *>>;
|
|
|
|
static OrderedPredicate getEmptyKey() { return Base::getEmptyKey(); }
|
|
static OrderedPredicate getTombstoneKey() { return Base::getTombstoneKey(); }
|
|
static bool isEqual(const OrderedPredicate &lhs,
|
|
const OrderedPredicate &rhs) {
|
|
return lhs.position == rhs.position && lhs.question == rhs.question;
|
|
}
|
|
static unsigned getHashValue(const OrderedPredicate &p) {
|
|
return llvm::hash_combine(p.position, p.question);
|
|
}
|
|
};
|
|
|
|
/// This class wraps a set of ordered predicates that are used within a specific
|
|
/// pattern operation.
|
|
struct OrderedPredicateList {
|
|
OrderedPredicateList(pdl::PatternOp pattern) : pattern(pattern) {}
|
|
|
|
pdl::PatternOp pattern;
|
|
DenseSet<OrderedPredicate *> predicates;
|
|
};
|
|
} // end anonymous namespace
|
|
|
|
/// Returns true if the given matcher refers to the same predicate as the given
|
|
/// ordered predicate. This means that the position and questions of the two
|
|
/// match.
|
|
static bool isSamePredicate(MatcherNode *node, OrderedPredicate *predicate) {
|
|
return node->getPosition() == predicate->position &&
|
|
node->getQuestion() == predicate->question;
|
|
}
|
|
|
|
/// Get or insert a child matcher for the given parent switch node, given a
|
|
/// predicate and parent pattern.
|
|
std::unique_ptr<MatcherNode> &getOrCreateChild(SwitchNode *node,
|
|
OrderedPredicate *predicate,
|
|
pdl::PatternOp pattern) {
|
|
assert(isSamePredicate(node, predicate) &&
|
|
"expected matcher to equal the given predicate");
|
|
|
|
auto it = predicate->patternToAnswer.find(pattern);
|
|
assert(it != predicate->patternToAnswer.end() &&
|
|
"expected pattern to exist in predicate");
|
|
return node->getChildren().insert({it->second, nullptr}).first->second;
|
|
}
|
|
|
|
/// Build the matcher CFG by "pushing" patterns through by sorted predicate
|
|
/// order. A pattern will traverse as far as possible using common predicates
|
|
/// and then either diverge from the CFG or reach the end of a branch and start
|
|
/// creating new nodes.
|
|
static void propagatePattern(std::unique_ptr<MatcherNode> &node,
|
|
OrderedPredicateList &list,
|
|
std::vector<OrderedPredicate *>::iterator current,
|
|
std::vector<OrderedPredicate *>::iterator end) {
|
|
if (current == end) {
|
|
// We've hit the end of a pattern, so create a successful result node.
|
|
node = std::make_unique<SuccessNode>(list.pattern, std::move(node));
|
|
|
|
// If the pattern doesn't contain this predicate, ignore it.
|
|
} else if (list.predicates.find(*current) == list.predicates.end()) {
|
|
propagatePattern(node, list, std::next(current), end);
|
|
|
|
// If the current matcher node is invalid, create a new one for this
|
|
// position and continue propagation.
|
|
} else if (!node) {
|
|
// Create a new node at this position and continue
|
|
node = std::make_unique<SwitchNode>((*current)->position,
|
|
(*current)->question);
|
|
propagatePattern(
|
|
getOrCreateChild(cast<SwitchNode>(&*node), *current, list.pattern),
|
|
list, std::next(current), end);
|
|
|
|
// If the matcher has already been created, and it is for this predicate we
|
|
// continue propagation to the child.
|
|
} else if (isSamePredicate(node.get(), *current)) {
|
|
propagatePattern(
|
|
getOrCreateChild(cast<SwitchNode>(&*node), *current, list.pattern),
|
|
list, std::next(current), end);
|
|
|
|
// If the matcher doesn't match the current predicate, insert a branch as
|
|
// the common set of matchers has diverged.
|
|
} else {
|
|
propagatePattern(node->getFailureNode(), list, current, end);
|
|
}
|
|
}
|
|
|
|
/// Fold any switch nodes nested under `node` to boolean nodes when possible.
|
|
/// `node` is updated in-place if it is a switch.
|
|
static void foldSwitchToBool(std::unique_ptr<MatcherNode> &node) {
|
|
if (!node)
|
|
return;
|
|
|
|
if (SwitchNode *switchNode = dyn_cast<SwitchNode>(&*node)) {
|
|
SwitchNode::ChildMapT &children = switchNode->getChildren();
|
|
for (auto &it : children)
|
|
foldSwitchToBool(it.second);
|
|
|
|
// If the node only contains one child, collapse it into a boolean predicate
|
|
// node.
|
|
if (children.size() == 1) {
|
|
auto childIt = children.begin();
|
|
node = std::make_unique<BoolNode>(
|
|
node->getPosition(), node->getQuestion(), childIt->first,
|
|
std::move(childIt->second), std::move(node->getFailureNode()));
|
|
}
|
|
} else if (BoolNode *boolNode = dyn_cast<BoolNode>(&*node)) {
|
|
foldSwitchToBool(boolNode->getSuccessNode());
|
|
}
|
|
|
|
foldSwitchToBool(node->getFailureNode());
|
|
}
|
|
|
|
/// Insert an exit node at the end of the failure path of the `root`.
|
|
static void insertExitNode(std::unique_ptr<MatcherNode> *root) {
|
|
while (*root)
|
|
root = &(*root)->getFailureNode();
|
|
*root = std::make_unique<ExitNode>();
|
|
}
|
|
|
|
/// Given a module containing PDL pattern operations, generate a matcher tree
|
|
/// using the patterns within the given module and return the root matcher node.
|
|
std::unique_ptr<MatcherNode>
|
|
MatcherNode::generateMatcherTree(ModuleOp module, PredicateBuilder &builder,
|
|
DenseMap<Value, Position *> &valueToPosition) {
|
|
// Collect the set of predicates contained within the pattern operations of
|
|
// the module.
|
|
SmallVector<std::pair<pdl::PatternOp, std::vector<PositionalPredicate>>, 16>
|
|
patternsAndPredicates;
|
|
for (pdl::PatternOp pattern : module.getOps<pdl::PatternOp>()) {
|
|
std::vector<PositionalPredicate> predicateList;
|
|
buildPredicateList(pattern, builder, predicateList, valueToPosition);
|
|
patternsAndPredicates.emplace_back(pattern, std::move(predicateList));
|
|
}
|
|
|
|
// Associate a pattern result with each unique predicate.
|
|
DenseSet<OrderedPredicate, OrderedPredicateDenseInfo> uniqued;
|
|
for (auto &patternAndPredList : patternsAndPredicates) {
|
|
for (auto &predicate : patternAndPredList.second) {
|
|
auto it = uniqued.insert(predicate);
|
|
it.first->patternToAnswer.try_emplace(patternAndPredList.first,
|
|
predicate.answer);
|
|
}
|
|
}
|
|
|
|
// Associate each pattern to a set of its ordered predicates for later lookup.
|
|
std::vector<OrderedPredicateList> lists;
|
|
lists.reserve(patternsAndPredicates.size());
|
|
for (auto &patternAndPredList : patternsAndPredicates) {
|
|
OrderedPredicateList list(patternAndPredList.first);
|
|
for (auto &predicate : patternAndPredList.second) {
|
|
OrderedPredicate *orderedPredicate = &*uniqued.find(predicate);
|
|
list.predicates.insert(orderedPredicate);
|
|
|
|
// Increment the primary sum for each reference to a particular predicate.
|
|
++orderedPredicate->primary;
|
|
}
|
|
lists.push_back(std::move(list));
|
|
}
|
|
|
|
// For a particular pattern, get the total primary sum and add it to the
|
|
// secondary sum of each predicate. Square the primary sums to emphasize
|
|
// shared predicates within rather than across patterns.
|
|
for (auto &list : lists) {
|
|
unsigned total = 0;
|
|
for (auto *predicate : list.predicates)
|
|
total += predicate->primary * predicate->primary;
|
|
for (auto *predicate : list.predicates)
|
|
predicate->secondary += total;
|
|
}
|
|
|
|
// Sort the set of predicates now that the cost primary and secondary sums
|
|
// have been computed.
|
|
std::vector<OrderedPredicate *> ordered;
|
|
ordered.reserve(uniqued.size());
|
|
for (auto &ip : uniqued)
|
|
ordered.push_back(&ip);
|
|
std::stable_sort(
|
|
ordered.begin(), ordered.end(),
|
|
[](OrderedPredicate *lhs, OrderedPredicate *rhs) { return *lhs < *rhs; });
|
|
|
|
// Build the matchers for each of the pattern predicate lists.
|
|
std::unique_ptr<MatcherNode> root;
|
|
for (OrderedPredicateList &list : lists)
|
|
propagatePattern(root, list, ordered.begin(), ordered.end());
|
|
|
|
// Collapse the graph and insert the exit node.
|
|
foldSwitchToBool(root);
|
|
insertExitNode(&root);
|
|
return root;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// MatcherNode
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
MatcherNode::MatcherNode(TypeID matcherTypeID, Position *p, Qualifier *q,
|
|
std::unique_ptr<MatcherNode> failureNode)
|
|
: position(p), question(q), failureNode(std::move(failureNode)),
|
|
matcherTypeID(matcherTypeID) {}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// BoolNode
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
BoolNode::BoolNode(Position *position, Qualifier *question, Qualifier *answer,
|
|
std::unique_ptr<MatcherNode> successNode,
|
|
std::unique_ptr<MatcherNode> failureNode)
|
|
: MatcherNode(TypeID::get<BoolNode>(), position, question,
|
|
std::move(failureNode)),
|
|
answer(answer), successNode(std::move(successNode)) {}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// SuccessNode
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
SuccessNode::SuccessNode(pdl::PatternOp pattern,
|
|
std::unique_ptr<MatcherNode> failureNode)
|
|
: MatcherNode(TypeID::get<SuccessNode>(), /*position=*/nullptr,
|
|
/*question=*/nullptr, std::move(failureNode)),
|
|
pattern(pattern) {}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// SwitchNode
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
SwitchNode::SwitchNode(Position *position, Qualifier *question)
|
|
: MatcherNode(TypeID::get<SwitchNode>(), position, question) {}
|