430 lines · cpp
1//===- PatternMatch.cpp - Base classes for pattern match ------------------===//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/IR/PatternMatch.h"10#include "mlir/IR/Iterators.h"11#include "mlir/IR/RegionKindInterface.h"12#include "llvm/ADT/SmallPtrSet.h"13 14using namespace mlir;15 16//===----------------------------------------------------------------------===//17// PatternBenefit18//===----------------------------------------------------------------------===//19 20PatternBenefit::PatternBenefit(unsigned benefit) : representation(benefit) {21 assert(representation == benefit && benefit != ImpossibleToMatchSentinel &&22 "This pattern match benefit is too large to represent");23}24 25unsigned short PatternBenefit::getBenefit() const {26 assert(!isImpossibleToMatch() && "Pattern doesn't match");27 return representation;28}29 30//===----------------------------------------------------------------------===//31// Pattern32//===----------------------------------------------------------------------===//33 34//===----------------------------------------------------------------------===//35// OperationName Root Constructors36//===----------------------------------------------------------------------===//37 38Pattern::Pattern(StringRef rootName, PatternBenefit benefit,39 MLIRContext *context, ArrayRef<StringRef> generatedNames)40 : Pattern(OperationName(rootName, context).getAsOpaquePointer(),41 RootKind::OperationName, generatedNames, benefit, context) {}42 43//===----------------------------------------------------------------------===//44// MatchAnyOpTypeTag Root Constructors45//===----------------------------------------------------------------------===//46 47Pattern::Pattern(MatchAnyOpTypeTag tag, PatternBenefit benefit,48 MLIRContext *context, ArrayRef<StringRef> generatedNames)49 : Pattern(nullptr, RootKind::Any, generatedNames, benefit, context) {}50 51//===----------------------------------------------------------------------===//52// MatchInterfaceOpTypeTag Root Constructors53//===----------------------------------------------------------------------===//54 55Pattern::Pattern(MatchInterfaceOpTypeTag tag, TypeID interfaceID,56 PatternBenefit benefit, MLIRContext *context,57 ArrayRef<StringRef> generatedNames)58 : Pattern(interfaceID.getAsOpaquePointer(), RootKind::InterfaceID,59 generatedNames, benefit, context) {}60 61//===----------------------------------------------------------------------===//62// MatchTraitOpTypeTag Root Constructors63//===----------------------------------------------------------------------===//64 65Pattern::Pattern(MatchTraitOpTypeTag tag, TypeID traitID,66 PatternBenefit benefit, MLIRContext *context,67 ArrayRef<StringRef> generatedNames)68 : Pattern(traitID.getAsOpaquePointer(), RootKind::TraitID, generatedNames,69 benefit, context) {}70 71//===----------------------------------------------------------------------===//72// General Constructors73//===----------------------------------------------------------------------===//74 75Pattern::Pattern(const void *rootValue, RootKind rootKind,76 ArrayRef<StringRef> generatedNames, PatternBenefit benefit,77 MLIRContext *context)78 : rootValue(rootValue), rootKind(rootKind), benefit(benefit),79 contextAndHasBoundedRecursion(context, false) {80 if (generatedNames.empty())81 return;82 generatedOps.reserve(generatedNames.size());83 llvm::append_range(generatedOps,84 llvm::map_range(generatedNames, [context](StringRef name) {85 return OperationName(name, context);86 }));87}88 89//===----------------------------------------------------------------------===//90// RewritePattern91//===----------------------------------------------------------------------===//92 93/// Out-of-line vtable anchor.94void RewritePattern::anchor() {}95 96//===----------------------------------------------------------------------===//97// RewriterBase98//===----------------------------------------------------------------------===//99 100bool RewriterBase::Listener::classof(const OpBuilder::Listener *base) {101 return base->getKind() == OpBuilder::ListenerBase::Kind::RewriterBaseListener;102}103 104RewriterBase::~RewriterBase() {105 // Out of line to provide a vtable anchor for the class.106}107 108void RewriterBase::replaceAllOpUsesWith(Operation *from, ValueRange to) {109 // Notify the listener that we're about to replace this op.110 if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))111 rewriteListener->notifyOperationReplaced(from, to);112 113 replaceAllUsesWith(from->getResults(), to);114}115 116void RewriterBase::replaceAllOpUsesWith(Operation *from, Operation *to) {117 // Notify the listener that we're about to replace this op.118 if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))119 rewriteListener->notifyOperationReplaced(from, to);120 121 replaceAllUsesWith(from->getResults(), to->getResults());122}123 124/// This method replaces the results of the operation with the specified list of125/// values. The number of provided values must match the number of results of126/// the operation. The replaced op is erased.127void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {128 assert(op->getNumResults() == newValues.size() &&129 "incorrect # of replacement values");130 131 // Replace all result uses. Also notifies the listener of modifications.132 replaceAllOpUsesWith(op, newValues);133 134 // Erase op and notify listener.135 eraseOp(op);136}137 138/// This method replaces the results of the operation with the specified new op139/// (replacement). The number of results of the two operations must match. The140/// replaced op is erased.141void RewriterBase::replaceOp(Operation *op, Operation *newOp) {142 assert(op && newOp && "expected non-null op");143 assert(op->getNumResults() == newOp->getNumResults() &&144 "ops have different number of results");145 146 // Replace all result uses. Also notifies the listener of modifications.147 replaceAllOpUsesWith(op, newOp->getResults());148 149 // Erase op and notify listener.150 eraseOp(op);151}152 153/// This method erases an operation that is known to have no uses. The uses of154/// the given operation *must* be known to be dead.155void RewriterBase::eraseOp(Operation *op) {156 assert(op->use_empty() && "expected 'op' to have no uses");157 auto *rewriteListener = dyn_cast_if_present<Listener>(listener);158 159 // If the current insertion point is before the erased operation, we adjust160 // the insertion point to be after the operation.161 if (getInsertionPoint() == op->getIterator())162 setInsertionPointAfter(op);163 164 // Fast path: If no listener is attached, the op can be dropped in one go.165 if (!rewriteListener) {166 op->erase();167 return;168 }169 170 // Helper function that erases a single op.171 auto eraseSingleOp = [&](Operation *op) {172#ifndef NDEBUG173 // All nested ops should have been erased already.174 assert(175 llvm::all_of(op->getRegions(), [&](Region &r) { return r.empty(); }) &&176 "expected empty regions");177 // All users should have been erased already if the op is in a region with178 // SSA dominance.179 if (!op->use_empty() && op->getParentOp())180 assert(mayBeGraphRegion(*op->getParentRegion()) &&181 "expected that op has no uses");182#endif // NDEBUG183 rewriteListener->notifyOperationErased(op);184 185 // Explicitly drop all uses in case the op is in a graph region.186 op->dropAllUses();187 op->erase();188 };189 190 // Nested ops must be erased one-by-one, so that listeners have a consistent191 // view of the IR every time a notification is triggered. Users must be192 // erased before definitions. I.e., post-order, reverse dominance.193 std::function<void(Operation *)> eraseTree = [&](Operation *op) {194 // Erase nested ops.195 for (Region &r : llvm::reverse(op->getRegions())) {196 // Erase all blocks in the right order. Successors should be erased197 // before predecessors because successor blocks may use values defined198 // in predecessor blocks. A post-order traversal of blocks within a199 // region visits successors before predecessors. Repeat the traversal200 // until the region is empty. (The block graph could be disconnected.)201 while (!r.empty()) {202 SmallVector<Block *> erasedBlocks;203 // Some blocks may have invalid successor, use a set including nullptr204 // to avoid null pointer.205 llvm::SmallPtrSet<Block *, 4> visited{nullptr};206 for (Block *b : llvm::post_order_ext(&r.front(), visited)) {207 // Visit ops in reverse order.208 for (Operation &op :209 llvm::make_early_inc_range(ReverseIterator::makeIterable(*b)))210 eraseTree(&op);211 // Do not erase the block immediately. This is not supprted by the212 // post_order iterator.213 erasedBlocks.push_back(b);214 }215 for (Block *b : erasedBlocks) {216 // Explicitly drop all uses in case there is a cycle in the block217 // graph.218 for (BlockArgument bbArg : b->getArguments())219 bbArg.dropAllUses();220 b->dropAllUses();221 eraseBlock(b);222 }223 }224 }225 // Then erase the enclosing op.226 eraseSingleOp(op);227 };228 229 eraseTree(op);230}231 232void RewriterBase::eraseBlock(Block *block) {233 assert(block->use_empty() && "expected 'block' to have no uses");234 235 for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {236 assert(op.use_empty() && "expected 'op' to have no uses");237 eraseOp(&op);238 }239 240 // Notify the listener that the block is about to be removed.241 if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))242 rewriteListener->notifyBlockErased(block);243 244 block->erase();245}246 247void RewriterBase::finalizeOpModification(Operation *op) {248 // Notify the listener that the operation was modified.249 if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))250 rewriteListener->notifyOperationModified(op);251}252 253void RewriterBase::replaceAllUsesExcept(254 Value from, Value to, const SmallPtrSetImpl<Operation *> &preservedUsers) {255 return replaceUsesWithIf(from, to, [&](OpOperand &use) {256 Operation *user = use.getOwner();257 return !preservedUsers.contains(user);258 });259}260 261void RewriterBase::replaceUsesWithIf(Value from, Value to,262 function_ref<bool(OpOperand &)> functor,263 bool *allUsesReplaced) {264 bool allReplaced = true;265 for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {266 bool replace = functor(operand);267 if (replace)268 modifyOpInPlace(operand.getOwner(), [&]() { operand.set(to); });269 allReplaced &= replace;270 }271 if (allUsesReplaced)272 *allUsesReplaced = allReplaced;273}274 275void RewriterBase::replaceUsesWithIf(ValueRange from, ValueRange to,276 function_ref<bool(OpOperand &)> functor,277 bool *allUsesReplaced) {278 assert(from.size() == to.size() && "incorrect number of replacements");279 bool allReplaced = true;280 for (auto it : llvm::zip_equal(from, to)) {281 bool r;282 replaceUsesWithIf(std::get<0>(it), std::get<1>(it), functor,283 /*allUsesReplaced=*/&r);284 allReplaced &= r;285 }286 if (allUsesReplaced)287 *allUsesReplaced = allReplaced;288}289 290void RewriterBase::inlineBlockBefore(Block *source, Block *dest,291 Block::iterator before,292 ValueRange argValues) {293 assert(argValues.size() == source->getNumArguments() &&294 "incorrect # of argument replacement values");295 296 // The source block will be deleted, so it should not have any users (i.e.,297 // there should be no predecessors).298 assert(source->hasNoPredecessors() &&299 "expected 'source' to have no predecessors");300 301 if (dest->end() != before) {302 // The source block will be inserted in the middle of the dest block, so303 // the source block should have no successors. Otherwise, the remainder of304 // the dest block would be unreachable.305 assert(source->hasNoSuccessors() &&306 "expected 'source' to have no successors");307 } else {308 // The source block will be inserted at the end of the dest block, so the309 // dest block should have no successors. Otherwise, the inserted operations310 // will be unreachable.311 assert(dest->hasNoSuccessors() && "expected 'dest' to have no successors");312 }313 314 // Replace all of the successor arguments with the provided values.315 for (auto it : llvm::zip(source->getArguments(), argValues))316 replaceAllUsesWith(std::get<0>(it), std::get<1>(it));317 318 // Move operations from the source block to the dest block and erase the319 // source block.320 if (!listener) {321 // Fast path: If no listener is attached, move all operations at once.322 dest->getOperations().splice(before, source->getOperations());323 } else {324 while (!source->empty())325 moveOpBefore(&source->front(), dest, before);326 }327 328 // If the current insertion point is within the source block, adjust the329 // insertion point to the destination block.330 if (getInsertionBlock() == source)331 setInsertionPoint(dest, getInsertionPoint());332 333 // Erase the source block.334 assert(source->empty() && "expected 'source' to be empty");335 eraseBlock(source);336}337 338void RewriterBase::inlineBlockBefore(Block *source, Operation *op,339 ValueRange argValues) {340 inlineBlockBefore(source, op->getBlock(), op->getIterator(), argValues);341}342 343void RewriterBase::mergeBlocks(Block *source, Block *dest,344 ValueRange argValues) {345 inlineBlockBefore(source, dest, dest->end(), argValues);346}347 348/// Split the operations starting at "before" (inclusive) out of the given349/// block into a new block, and return it.350Block *RewriterBase::splitBlock(Block *block, Block::iterator before) {351 // Fast path: If no listener is attached, split the block directly.352 if (!listener)353 return block->splitBlock(before);354 355 // `createBlock` sets the insertion point at the beginning of the new block.356 InsertionGuard g(*this);357 Block *newBlock =358 createBlock(block->getParent(), std::next(block->getIterator()));359 360 // If `before` points to end of the block, no ops should be moved.361 if (before == block->end())362 return newBlock;363 364 // Move ops one-by-one from the end of `block` to the beginning of `newBlock`.365 // Stop when the operation pointed to by `before` has been moved.366 while (before->getBlock() != newBlock)367 moveOpBefore(&block->back(), newBlock, newBlock->begin());368 369 return newBlock;370}371 372/// Move the blocks that belong to "region" before the given position in373/// another region. The two regions must be different. The caller is in374/// charge to update create the operation transferring the control flow to the375/// region and pass it the correct block arguments.376void RewriterBase::inlineRegionBefore(Region ®ion, Region &parent,377 Region::iterator before) {378 // Fast path: If no listener is attached, move all blocks at once.379 if (!listener) {380 parent.getBlocks().splice(before, region.getBlocks());381 return;382 }383 384 // Move blocks from the beginning of the region one-by-one.385 while (!region.empty())386 moveBlockBefore(®ion.front(), &parent, before);387}388void RewriterBase::inlineRegionBefore(Region ®ion, Block *before) {389 inlineRegionBefore(region, *before->getParent(), before->getIterator());390}391 392void RewriterBase::moveBlockBefore(Block *block, Block *anotherBlock) {393 moveBlockBefore(block, anotherBlock->getParent(),394 anotherBlock->getIterator());395}396 397void RewriterBase::moveBlockBefore(Block *block, Region *region,398 Region::iterator iterator) {399 Region *currentRegion = block->getParent();400 Region::iterator nextIterator = std::next(block->getIterator());401 block->moveBefore(region, iterator);402 if (listener)403 listener->notifyBlockInserted(block, /*previous=*/currentRegion,404 /*previousIt=*/nextIterator);405}406 407void RewriterBase::moveOpBefore(Operation *op, Operation *existingOp) {408 moveOpBefore(op, existingOp->getBlock(), existingOp->getIterator());409}410 411void RewriterBase::moveOpBefore(Operation *op, Block *block,412 Block::iterator iterator) {413 Block *currentBlock = op->getBlock();414 Block::iterator nextIterator = std::next(op->getIterator());415 op->moveBefore(block, iterator);416 if (listener)417 listener->notifyOperationInserted(418 op, /*previous=*/InsertPoint(currentBlock, nextIterator));419}420 421void RewriterBase::moveOpAfter(Operation *op, Operation *existingOp) {422 moveOpAfter(op, existingOp->getBlock(), existingOp->getIterator());423}424 425void RewriterBase::moveOpAfter(Operation *op, Block *block,426 Block::iterator iterator) {427 assert(iterator != block->end() && "cannot move after end of block");428 moveOpBefore(op, block, std::next(iterator));429}430