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