llvm-project/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
River Riddle 76f3c2f3f3 [mlir][Pattern] Add better support for using interfaces/traits to match root operations in rewrite patterns
To match an interface or trait, users currently have to use the `MatchAny` tag. This tag can be quite problematic for compile time for things like the canonicalizer, as the `MatchAny` patterns may get applied to  *every* operation. This revision adds better support by bucketing interface/trait patterns based on which registered operations have them registered. This means that moving forward we will only attempt to match these patterns to operations that have this interface registered. Two simplify defining patterns that match traits and interfaces, two new utility classes have been added: OpTraitRewritePattern and OpInterfaceRewritePattern.

Differential Revision: https://reviews.llvm.org/D98986
2021-03-23 14:05:33 -07:00

112 lines
4.1 KiB
C++

//===- FrozenRewritePatternSet.cpp - Frozen Pattern List -------*- 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
//
//===----------------------------------------------------------------------===//
#include "mlir/Rewrite/FrozenRewritePatternSet.h"
#include "ByteCode.h"
#include "mlir/Conversion/PDLToPDLInterp/PDLToPDLInterp.h"
#include "mlir/Dialect/PDL/IR/PDLOps.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
using namespace mlir;
static LogicalResult convertPDLToPDLInterp(ModuleOp pdlModule) {
// Skip the conversion if the module doesn't contain pdl.
if (llvm::empty(pdlModule.getOps<pdl::PatternOp>()))
return success();
// Simplify the provided PDL module. Note that we can't use the canonicalizer
// here because it would create a cyclic dependency.
auto simplifyFn = [](Operation *op) {
// TODO: Add folding here if ever necessary.
if (isOpTriviallyDead(op))
op->erase();
};
pdlModule.getBody()->walk(simplifyFn);
/// Lower the PDL pattern module to the interpreter dialect.
PassManager pdlPipeline(pdlModule.getContext());
#ifdef NDEBUG
// We don't want to incur the hit of running the verifier when in release
// mode.
pdlPipeline.enableVerifier(false);
#endif
pdlPipeline.addPass(createPDLToPDLInterpPass());
if (failed(pdlPipeline.run(pdlModule)))
return failure();
// Simplify again after running the lowering pipeline.
pdlModule.getBody()->walk(simplifyFn);
return success();
}
//===----------------------------------------------------------------------===//
// FrozenRewritePatternSet
//===----------------------------------------------------------------------===//
FrozenRewritePatternSet::FrozenRewritePatternSet()
: impl(std::make_shared<Impl>()) {}
FrozenRewritePatternSet::FrozenRewritePatternSet(RewritePatternSet &&patterns)
: impl(std::make_shared<Impl>()) {
// Functor used to walk all of the operations registered in the context. This
// is useful for patterns that get applied to multiple operations, such as
// interface and trait based patterns.
std::vector<AbstractOperation *> abstractOps;
auto addToOpsWhen = [&](std::unique_ptr<RewritePattern> &pattern,
function_ref<bool(AbstractOperation *)> callbackFn) {
if (abstractOps.empty())
abstractOps = pattern->getContext()->getRegisteredOperations();
for (AbstractOperation *absOp : abstractOps) {
if (callbackFn(absOp)) {
OperationName opName(absOp);
impl->nativeOpSpecificPatternMap[opName].push_back(pattern.get());
}
}
impl->nativeOpSpecificPatternList.push_back(std::move(pattern));
};
for (std::unique_ptr<RewritePattern> &pat : patterns.getNativePatterns()) {
if (Optional<OperationName> rootName = pat->getRootKind()) {
impl->nativeOpSpecificPatternMap[*rootName].push_back(pat.get());
impl->nativeOpSpecificPatternList.push_back(std::move(pat));
continue;
}
if (Optional<TypeID> interfaceID = pat->getRootInterfaceID()) {
addToOpsWhen(pat, [&](AbstractOperation *absOp) {
return absOp->hasInterface(*interfaceID);
});
continue;
}
if (Optional<TypeID> traitID = pat->getRootTraitID()) {
addToOpsWhen(pat, [&](AbstractOperation *absOp) {
return absOp->hasTrait(*traitID);
});
continue;
}
impl->nativeAnyOpPatterns.push_back(std::move(pat));
}
// Generate the bytecode for the PDL patterns if any were provided.
PDLPatternModule &pdlPatterns = patterns.getPDLPatterns();
ModuleOp pdlModule = pdlPatterns.getModule();
if (!pdlModule)
return;
if (failed(convertPDLToPDLInterp(pdlModule)))
llvm::report_fatal_error(
"failed to lower PDL pattern module to the PDL Interpreter");
// Generate the pdl bytecode.
impl->pdlByteCode = std::make_unique<detail::PDLByteCode>(
pdlModule, pdlPatterns.takeConstraintFunctions(),
pdlPatterns.takeRewriteFunctions());
}
FrozenRewritePatternSet::~FrozenRewritePatternSet() {}