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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/PatternMatch.h"
10 #include "mlir/IR/BlockAndValueMapping.h"
11 #include "mlir/IR/Operation.h"
12 #include "mlir/IR/Value.h"
13 
14 using namespace mlir;
15 
16 PatternBenefit::PatternBenefit(unsigned benefit) : representation(benefit) {
17   assert(representation == benefit && benefit != ImpossibleToMatchSentinel &&
18          "This pattern match benefit is too large to represent");
19 }
20 
21 unsigned short PatternBenefit::getBenefit() const {
22   assert(representation != ImpossibleToMatchSentinel &&
23          "Pattern doesn't match");
24   return representation;
25 }
26 
27 //===----------------------------------------------------------------------===//
28 // Pattern implementation
29 //===----------------------------------------------------------------------===//
30 
31 Pattern::Pattern(StringRef rootName, PatternBenefit benefit,
32                  MLIRContext *context)
33     : rootKind(OperationName(rootName, context)), benefit(benefit) {}
34 
35 // Out-of-line vtable anchor.
36 void Pattern::anchor() {}
37 
38 //===----------------------------------------------------------------------===//
39 // RewritePattern and PatternRewriter implementation
40 //===----------------------------------------------------------------------===//
41 
42 void RewritePattern::rewrite(Operation *op, PatternRewriter &rewriter) const {
43   llvm_unreachable("need to implement either matchAndRewrite or one of the "
44                    "rewrite functions!");
45 }
46 
47 LogicalResult RewritePattern::match(Operation *op) const {
48   llvm_unreachable("need to implement either match or matchAndRewrite!");
49 }
50 
51 /// Patterns must specify the root operation name they match against, and can
52 /// also specify the benefit of the pattern matching. They can also specify the
53 /// names of operations that may be generated during a successful rewrite.
54 RewritePattern::RewritePattern(StringRef rootName,
55                                ArrayRef<StringRef> generatedNames,
56                                PatternBenefit benefit, MLIRContext *context)
57     : Pattern(rootName, benefit, context) {
58   generatedOps.reserve(generatedNames.size());
59   std::transform(generatedNames.begin(), generatedNames.end(),
60                  std::back_inserter(generatedOps), [context](StringRef name) {
61                    return OperationName(name, context);
62                  });
63 }
64 
65 PatternRewriter::~PatternRewriter() {
66   // Out of line to provide a vtable anchor for the class.
67 }
68 
69 /// This method performs the final replacement for a pattern, where the
70 /// results of the operation are updated to use the specified list of SSA
71 /// values.
72 void PatternRewriter::replaceOp(Operation *op, ValueRange newValues) {
73   // Notify the rewriter subclass that we're about to replace this root.
74   notifyRootReplaced(op);
75 
76   assert(op->getNumResults() == newValues.size() &&
77          "incorrect # of replacement values");
78   op->replaceAllUsesWith(newValues);
79 
80   notifyOperationRemoved(op);
81   op->erase();
82 }
83 
84 /// This method erases an operation that is known to have no uses. The uses of
85 /// the given operation *must* be known to be dead.
86 void PatternRewriter::eraseOp(Operation *op) {
87   assert(op->use_empty() && "expected 'op' to have no uses");
88   notifyOperationRemoved(op);
89   op->erase();
90 }
91 
92 void PatternRewriter::eraseBlock(Block *block) {
93   for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
94     assert(op.use_empty() && "expected 'op' to have no uses");
95     eraseOp(&op);
96   }
97   block->erase();
98 }
99 
100 /// Merge the operations of block 'source' into the end of block 'dest'.
101 /// 'source's predecessors must be empty or only contain 'dest`.
102 /// 'argValues' is used to replace the block arguments of 'source' after
103 /// merging.
104 void PatternRewriter::mergeBlocks(Block *source, Block *dest,
105                                   ValueRange argValues) {
106   assert(llvm::all_of(source->getPredecessors(),
107                       [dest](Block *succ) { return succ == dest; }) &&
108          "expected 'source' to have no predecessors or only 'dest'");
109   assert(argValues.size() == source->getNumArguments() &&
110          "incorrect # of argument replacement values");
111 
112   // Replace all of the successor arguments with the provided values.
113   for (auto it : llvm::zip(source->getArguments(), argValues))
114     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
115 
116   // Splice the operations of the 'source' block into the 'dest' block and erase
117   // it.
118   dest->getOperations().splice(dest->end(), source->getOperations());
119   source->dropAllUses();
120   source->erase();
121 }
122 
123 /// Split the operations starting at "before" (inclusive) out of the given
124 /// block into a new block, and return it.
125 Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {
126   return block->splitBlock(before);
127 }
128 
129 /// 'op' and 'newOp' are known to have the same number of results, replace the
130 /// uses of op with uses of newOp
131 void PatternRewriter::replaceOpWithResultsOfAnotherOp(Operation *op,
132                                                       Operation *newOp) {
133   assert(op->getNumResults() == newOp->getNumResults() &&
134          "replacement op doesn't match results of original op");
135   if (op->getNumResults() == 1)
136     return replaceOp(op, newOp->getResult(0));
137   return replaceOp(op, newOp->getResults());
138 }
139 
140 /// Move the blocks that belong to "region" before the given position in
141 /// another region.  The two regions must be different.  The caller is in
142 /// charge to update create the operation transferring the control flow to the
143 /// region and pass it the correct block arguments.
144 void PatternRewriter::inlineRegionBefore(Region &region, Region &parent,
145                                          Region::iterator before) {
146   parent.getBlocks().splice(before, region.getBlocks());
147 }
148 void PatternRewriter::inlineRegionBefore(Region &region, Block *before) {
149   inlineRegionBefore(region, *before->getParent(), before->getIterator());
150 }
151 
152 /// Clone the blocks that belong to "region" before the given position in
153 /// another region "parent". The two regions must be different. The caller is
154 /// responsible for creating or updating the operation transferring flow of
155 /// control to the region and passing it the correct block arguments.
156 void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,
157                                         Region::iterator before,
158                                         BlockAndValueMapping &mapping) {
159   region.cloneInto(&parent, before, mapping);
160 }
161 void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,
162                                         Region::iterator before) {
163   BlockAndValueMapping mapping;
164   cloneRegionBefore(region, parent, before, mapping);
165 }
166 void PatternRewriter::cloneRegionBefore(Region &region, Block *before) {
167   cloneRegionBefore(region, *before->getParent(), before->getIterator());
168 }
169 
170 //===----------------------------------------------------------------------===//
171 // PatternMatcher implementation
172 //===----------------------------------------------------------------------===//
173 
174 RewritePatternMatcher::RewritePatternMatcher(
175     const OwningRewritePatternList &patterns) {
176   for (auto &pattern : patterns)
177     this->patterns.push_back(pattern.get());
178 
179   // Sort the patterns by benefit to simplify the matching logic.
180   std::stable_sort(this->patterns.begin(), this->patterns.end(),
181                    [](RewritePattern *l, RewritePattern *r) {
182                      return r->getBenefit() < l->getBenefit();
183                    });
184 }
185 
186 /// Try to match the given operation to a pattern and rewrite it.
187 bool RewritePatternMatcher::matchAndRewrite(Operation *op,
188                                             PatternRewriter &rewriter) {
189   for (auto *pattern : patterns) {
190     // Ignore patterns that are for the wrong root or are impossible to match.
191     if (pattern->getRootKind() != op->getName() ||
192         pattern->getBenefit().isImpossibleToMatch())
193       continue;
194 
195     // Try to match and rewrite this pattern. The patterns are sorted by
196     // benefit, so if we match we can immediately rewrite and return.
197     if (succeeded(pattern->matchAndRewrite(op, rewriter)))
198       return true;
199   }
200   return false;
201 }
202