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 /// Merge the operations of block 'source' into the end of block 'dest'.
93 /// 'source's predecessors must be empty or only contain 'dest`.
94 /// 'argValues' is used to replace the block arguments of 'source' after
95 /// merging.
96 void PatternRewriter::mergeBlocks(Block *source, Block *dest,
97                                   ValueRange argValues) {
98   assert(llvm::all_of(source->getPredecessors(),
99                       [dest](Block *succ) { return succ == dest; }) &&
100          "expected 'source' to have no predecessors or only 'dest'");
101   assert(argValues.size() == source->getNumArguments() &&
102          "incorrect # of argument replacement values");
103 
104   // Replace all of the successor arguments with the provided values.
105   for (auto it : llvm::zip(source->getArguments(), argValues))
106     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
107 
108   // Splice the operations of the 'source' block into the 'dest' block and erase
109   // it.
110   dest->getOperations().splice(dest->end(), source->getOperations());
111   source->dropAllUses();
112   source->erase();
113 }
114 
115 /// Split the operations starting at "before" (inclusive) out of the given
116 /// block into a new block, and return it.
117 Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {
118   return block->splitBlock(before);
119 }
120 
121 /// 'op' and 'newOp' are known to have the same number of results, replace the
122 /// uses of op with uses of newOp
123 void PatternRewriter::replaceOpWithResultsOfAnotherOp(Operation *op,
124                                                       Operation *newOp) {
125   assert(op->getNumResults() == newOp->getNumResults() &&
126          "replacement op doesn't match results of original op");
127   if (op->getNumResults() == 1)
128     return replaceOp(op, newOp->getResult(0));
129   return replaceOp(op, newOp->getResults());
130 }
131 
132 /// Move the blocks that belong to "region" before the given position in
133 /// another region.  The two regions must be different.  The caller is in
134 /// charge to update create the operation transferring the control flow to the
135 /// region and pass it the correct block arguments.
136 void PatternRewriter::inlineRegionBefore(Region &region, Region &parent,
137                                          Region::iterator before) {
138   parent.getBlocks().splice(before, region.getBlocks());
139 }
140 void PatternRewriter::inlineRegionBefore(Region &region, Block *before) {
141   inlineRegionBefore(region, *before->getParent(), before->getIterator());
142 }
143 
144 /// Clone the blocks that belong to "region" before the given position in
145 /// another region "parent". The two regions must be different. The caller is
146 /// responsible for creating or updating the operation transferring flow of
147 /// control to the region and passing it the correct block arguments.
148 void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,
149                                         Region::iterator before,
150                                         BlockAndValueMapping &mapping) {
151   region.cloneInto(&parent, before, mapping);
152 }
153 void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,
154                                         Region::iterator before) {
155   BlockAndValueMapping mapping;
156   cloneRegionBefore(region, parent, before, mapping);
157 }
158 void PatternRewriter::cloneRegionBefore(Region &region, Block *before) {
159   cloneRegionBefore(region, *before->getParent(), before->getIterator());
160 }
161 
162 //===----------------------------------------------------------------------===//
163 // PatternMatcher implementation
164 //===----------------------------------------------------------------------===//
165 
166 RewritePatternMatcher::RewritePatternMatcher(
167     const OwningRewritePatternList &patterns) {
168   for (auto &pattern : patterns)
169     this->patterns.push_back(pattern.get());
170 
171   // Sort the patterns by benefit to simplify the matching logic.
172   std::stable_sort(this->patterns.begin(), this->patterns.end(),
173                    [](RewritePattern *l, RewritePattern *r) {
174                      return r->getBenefit() < l->getBenefit();
175                    });
176 }
177 
178 /// Try to match the given operation to a pattern and rewrite it.
179 bool RewritePatternMatcher::matchAndRewrite(Operation *op,
180                                             PatternRewriter &rewriter) {
181   for (auto *pattern : patterns) {
182     // Ignore patterns that are for the wrong root or are impossible to match.
183     if (pattern->getRootKind() != op->getName() ||
184         pattern->getBenefit().isImpossibleToMatch())
185       continue;
186 
187     // Try to match and rewrite this pattern. The patterns are sorted by
188     // benefit, so if we match we can immediately rewrite and return.
189     if (succeeded(pattern->matchAndRewrite(op, rewriter)))
190       return true;
191   }
192   return false;
193 }
194