231 lines · cpp
1//===- PDLExtensionOps.cpp - PDL extension for the Transform dialect ------===//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#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.h"10#include "mlir/Dialect/PDL/IR/PDLOps.h"11#include "mlir/IR/Builders.h"12#include "mlir/Rewrite/FrozenRewritePatternSet.h"13#include "mlir/Rewrite/PatternApplicator.h"14#include "llvm/ADT/ScopeExit.h"15 16using namespace mlir;17 18MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::transform::PDLMatchHooks)19 20#define GET_OP_CLASSES21#include "mlir/Dialect/Transform/PDLExtension/PDLExtensionOps.cpp.inc"22 23//===----------------------------------------------------------------------===//24// PatternApplicatorExtension25//===----------------------------------------------------------------------===//26 27namespace {28/// A TransformState extension that keeps track of compiled PDL pattern sets.29/// This is intended to be used along the WithPDLPatterns op. The extension30/// can be constructed given an operation that has a SymbolTable trait and31/// contains pdl::PatternOp instances. The patterns are compiled lazily and one32/// by one when requested; this behavior is subject to change.33class PatternApplicatorExtension : public transform::TransformState::Extension {34public:35 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PatternApplicatorExtension)36 37 /// Creates the extension for patterns contained in `patternContainer`.38 explicit PatternApplicatorExtension(transform::TransformState &state,39 Operation *patternContainer)40 : Extension(state), patterns(patternContainer) {}41 42 /// Appends to `results` the operations contained in `root` that matched the43 /// PDL pattern with the given name. Note that `root` may or may not be the44 /// operation that contains PDL patterns. Reports an error if the pattern45 /// cannot be found. Note that when no operations are matched, this still46 /// succeeds as long as the pattern exists.47 LogicalResult findAllMatches(StringRef patternName, Operation *root,48 SmallVectorImpl<Operation *> &results);49 50private:51 /// Map from the pattern name to a singleton set of rewrite patterns that only52 /// contains the pattern with this name. Populated when the pattern is first53 /// requested.54 // TODO: reconsider the efficiency of this storage when more usage data is55 // available. Storing individual patterns in a set and triggering compilation56 // for each of them has overhead. So does compiling a large set of patterns57 // only to apply a handful of them.58 llvm::StringMap<FrozenRewritePatternSet> compiledPatterns;59 60 /// A symbol table operation containing the relevant PDL patterns.61 SymbolTable patterns;62};63 64LogicalResult PatternApplicatorExtension::findAllMatches(65 StringRef patternName, Operation *root,66 SmallVectorImpl<Operation *> &results) {67 auto it = compiledPatterns.find(patternName);68 if (it == compiledPatterns.end()) {69 auto patternOp = patterns.lookup<pdl::PatternOp>(patternName);70 if (!patternOp)71 return failure();72 73 // Copy the pattern operation into a new module that is compiled and74 // consumed by the PDL interpreter.75 OwningOpRef<ModuleOp> pdlModuleOp = ModuleOp::create(patternOp.getLoc());76 auto builder = OpBuilder::atBlockEnd(pdlModuleOp->getBody());77 builder.clone(*patternOp);78 PDLPatternModule patternModule(std::move(pdlModuleOp));79 80 // Merge in the hooks owned by the dialect. Make a copy as they may be81 // also used by the following operations.82 auto *dialect =83 root->getContext()->getLoadedDialect<transform::TransformDialect>();84 for (const auto &[name, constraintFn] :85 dialect->getExtraData<transform::PDLMatchHooks>()86 .getPDLConstraintHooks()) {87 patternModule.registerConstraintFunction(name, constraintFn);88 }89 90 // Register a noop rewriter because PDL requires patterns to end with some91 // rewrite call.92 patternModule.registerRewriteFunction(93 "transform.dialect", [](PatternRewriter &, Operation *) {});94 95 it = compiledPatterns96 .try_emplace(patternOp.getName(), std::move(patternModule))97 .first;98 }99 100 PatternApplicator applicator(it->second);101 // We want to discourage direct use of PatternRewriter in APIs but In this102 // very specific case, an IRRewriter is not enough.103 PatternRewriter rewriter(root->getContext());104 applicator.applyDefaultCostModel();105 root->walk([&](Operation *op) {106 if (succeeded(applicator.matchAndRewrite(op, rewriter)))107 results.push_back(op);108 });109 110 return success();111}112} // namespace113 114//===----------------------------------------------------------------------===//115// PDLMatchHooks116//===----------------------------------------------------------------------===//117 118void transform::PDLMatchHooks::mergeInPDLMatchHooks(119 llvm::StringMap<PDLConstraintFunction> &&constraintFns) {120 // Steal the constraint functions from the given map.121 for (auto &it : constraintFns)122 pdlMatchHooks.registerConstraintFunction(it.getKey(), std::move(it.second));123}124 125const llvm::StringMap<PDLConstraintFunction> &126transform::PDLMatchHooks::getPDLConstraintHooks() const {127 return pdlMatchHooks.getConstraintFunctions();128}129 130//===----------------------------------------------------------------------===//131// PDLMatchOp132//===----------------------------------------------------------------------===//133 134DiagnosedSilenceableFailure135transform::PDLMatchOp::apply(transform::TransformRewriter &rewriter,136 transform::TransformResults &results,137 transform::TransformState &state) {138 auto *extension = state.getExtension<PatternApplicatorExtension>();139 assert(extension &&140 "expected PatternApplicatorExtension to be attached by the parent op");141 SmallVector<Operation *> targets;142 for (Operation *root : state.getPayloadOps(getRoot())) {143 if (failed(extension->findAllMatches(144 getPatternName().getLeafReference().getValue(), root, targets))) {145 emitDefiniteFailure()146 << "could not find pattern '" << getPatternName() << "'";147 }148 }149 results.set(llvm::cast<OpResult>(getResult()), targets);150 return DiagnosedSilenceableFailure::success();151}152 153void transform::PDLMatchOp::getEffects(154 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {155 onlyReadsHandle(getRootMutable(), effects);156 producesHandle(getOperation()->getOpResults(), effects);157 onlyReadsPayload(effects);158}159 160//===----------------------------------------------------------------------===//161// WithPDLPatternsOp162//===----------------------------------------------------------------------===//163 164DiagnosedSilenceableFailure165transform::WithPDLPatternsOp::apply(transform::TransformRewriter &rewriter,166 transform::TransformResults &results,167 transform::TransformState &state) {168 TransformOpInterface transformOp = nullptr;169 for (Operation &nested : getBody().front()) {170 if (!isa<pdl::PatternOp>(nested)) {171 transformOp = cast<TransformOpInterface>(nested);172 break;173 }174 }175 176 state.addExtension<PatternApplicatorExtension>(getOperation());177 auto guard = llvm::make_scope_exit(178 [&]() { state.removeExtension<PatternApplicatorExtension>(); });179 180 auto scope = state.make_region_scope(getBody());181 if (failed(mapBlockArguments(state)))182 return DiagnosedSilenceableFailure::definiteFailure();183 return state.applyTransform(transformOp);184}185 186void transform::WithPDLPatternsOp::getEffects(187 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {188 getPotentialTopLevelEffects(effects);189}190 191LogicalResult transform::WithPDLPatternsOp::verify() {192 Block *body = getBodyBlock();193 Operation *topLevelOp = nullptr;194 for (Operation &op : body->getOperations()) {195 if (isa<pdl::PatternOp>(op))196 continue;197 198 if (op.hasTrait<::mlir::transform::PossibleTopLevelTransformOpTrait>()) {199 if (topLevelOp) {200 InFlightDiagnostic diag =201 emitOpError() << "expects only one non-pattern op in its body";202 diag.attachNote(topLevelOp->getLoc()) << "first non-pattern op";203 diag.attachNote(op.getLoc()) << "second non-pattern op";204 return diag;205 }206 topLevelOp = &op;207 continue;208 }209 210 InFlightDiagnostic diag =211 emitOpError()212 << "expects only pattern and top-level transform ops in its body";213 diag.attachNote(op.getLoc()) << "offending op";214 return diag;215 }216 217 if (auto parent = getOperation()->getParentOfType<WithPDLPatternsOp>()) {218 InFlightDiagnostic diag = emitOpError() << "cannot be nested";219 diag.attachNote(parent.getLoc()) << "parent operation";220 return diag;221 }222 223 if (!topLevelOp) {224 InFlightDiagnostic diag = emitOpError()225 << "expects at least one non-pattern op";226 return diag;227 }228 229 return success();230}231