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
PatternBenefit(unsigned benefit)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
getBenefit() const21 unsigned short PatternBenefit::getBenefit() const {
22 assert(!isImpossibleToMatch() && "Pattern doesn't match");
23 return representation;
24 }
25
26 //===----------------------------------------------------------------------===//
27 // Pattern implementation
28 //===----------------------------------------------------------------------===//
29
Pattern(StringRef rootName,PatternBenefit benefit,MLIRContext * context)30 Pattern::Pattern(StringRef rootName, PatternBenefit benefit,
31 MLIRContext *context)
32 : rootKind(OperationName(rootName, context)), benefit(benefit) {}
Pattern(PatternBenefit benefit,MatchAnyOpTypeTag)33 Pattern::Pattern(PatternBenefit benefit, MatchAnyOpTypeTag)
34 : benefit(benefit) {}
35
36 // Out-of-line vtable anchor.
anchor()37 void Pattern::anchor() {}
38
39 //===----------------------------------------------------------------------===//
40 // RewritePattern and PatternRewriter implementation
41 //===----------------------------------------------------------------------===//
42
rewrite(Operation * op,PatternRewriter & rewriter) const43 void RewritePattern::rewrite(Operation *op, PatternRewriter &rewriter) const {
44 llvm_unreachable("need to implement either matchAndRewrite or one of the "
45 "rewrite functions!");
46 }
47
match(Operation * op) const48 LogicalResult RewritePattern::match(Operation *op) const {
49 llvm_unreachable("need to implement either match or matchAndRewrite!");
50 }
51
RewritePattern(StringRef rootName,ArrayRef<StringRef> generatedNames,PatternBenefit benefit,MLIRContext * context)52 RewritePattern::RewritePattern(StringRef rootName,
53 ArrayRef<StringRef> generatedNames,
54 PatternBenefit benefit, MLIRContext *context)
55 : Pattern(rootName, benefit, context) {
56 generatedOps.reserve(generatedNames.size());
57 std::transform(generatedNames.begin(), generatedNames.end(),
58 std::back_inserter(generatedOps), [context](StringRef name) {
59 return OperationName(name, context);
60 });
61 }
RewritePattern(ArrayRef<StringRef> generatedNames,PatternBenefit benefit,MLIRContext * context,MatchAnyOpTypeTag tag)62 RewritePattern::RewritePattern(ArrayRef<StringRef> generatedNames,
63 PatternBenefit benefit, MLIRContext *context,
64 MatchAnyOpTypeTag tag)
65 : Pattern(benefit, tag) {
66 generatedOps.reserve(generatedNames.size());
67 std::transform(generatedNames.begin(), generatedNames.end(),
68 std::back_inserter(generatedOps), [context](StringRef name) {
69 return OperationName(name, context);
70 });
71 }
72
~PatternRewriter()73 PatternRewriter::~PatternRewriter() {
74 // Out of line to provide a vtable anchor for the class.
75 }
76
77 /// This method performs the final replacement for a pattern, where the
78 /// results of the operation are updated to use the specified list of SSA
79 /// values.
replaceOp(Operation * op,ValueRange newValues)80 void PatternRewriter::replaceOp(Operation *op, ValueRange newValues) {
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
92 /// This method erases an operation that is known to have no uses. The uses of
93 /// the given operation *must* be known to be dead.
eraseOp(Operation * op)94 void PatternRewriter::eraseOp(Operation *op) {
95 assert(op->use_empty() && "expected 'op' to have no uses");
96 notifyOperationRemoved(op);
97 op->erase();
98 }
99
eraseBlock(Block * block)100 void PatternRewriter::eraseBlock(Block *block) {
101 for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
102 assert(op.use_empty() && "expected 'op' to have no uses");
103 eraseOp(&op);
104 }
105 block->erase();
106 }
107
108 /// Merge the operations of block 'source' into the end of block 'dest'.
109 /// 'source's predecessors must be empty or only contain 'dest`.
110 /// 'argValues' is used to replace the block arguments of 'source' after
111 /// merging.
mergeBlocks(Block * source,Block * dest,ValueRange argValues)112 void PatternRewriter::mergeBlocks(Block *source, Block *dest,
113 ValueRange argValues) {
114 assert(llvm::all_of(source->getPredecessors(),
115 [dest](Block *succ) { return succ == dest; }) &&
116 "expected 'source' to have no predecessors or only 'dest'");
117 assert(argValues.size() == source->getNumArguments() &&
118 "incorrect # of argument replacement values");
119
120 // Replace all of the successor arguments with the provided values.
121 for (auto it : llvm::zip(source->getArguments(), argValues))
122 std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
123
124 // Splice the operations of the 'source' block into the 'dest' block and erase
125 // it.
126 dest->getOperations().splice(dest->end(), source->getOperations());
127 source->dropAllUses();
128 source->erase();
129 }
130
131 /// Split the operations starting at "before" (inclusive) out of the given
132 /// block into a new block, and return it.
splitBlock(Block * block,Block::iterator before)133 Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {
134 return block->splitBlock(before);
135 }
136
137 /// 'op' and 'newOp' are known to have the same number of results, replace the
138 /// uses of op with uses of newOp
replaceOpWithResultsOfAnotherOp(Operation * op,Operation * newOp)139 void PatternRewriter::replaceOpWithResultsOfAnotherOp(Operation *op,
140 Operation *newOp) {
141 assert(op->getNumResults() == newOp->getNumResults() &&
142 "replacement op doesn't match results of original op");
143 if (op->getNumResults() == 1)
144 return replaceOp(op, newOp->getResult(0));
145 return replaceOp(op, newOp->getResults());
146 }
147
148 /// Move the blocks that belong to "region" before the given position in
149 /// another region. The two regions must be different. The caller is in
150 /// charge to update create the operation transferring the control flow to the
151 /// region and pass it the correct block arguments.
inlineRegionBefore(Region & region,Region & parent,Region::iterator before)152 void PatternRewriter::inlineRegionBefore(Region ®ion, Region &parent,
153 Region::iterator before) {
154 parent.getBlocks().splice(before, region.getBlocks());
155 }
inlineRegionBefore(Region & region,Block * before)156 void PatternRewriter::inlineRegionBefore(Region ®ion, Block *before) {
157 inlineRegionBefore(region, *before->getParent(), before->getIterator());
158 }
159
160 /// Clone the blocks that belong to "region" before the given position in
161 /// another region "parent". The two regions must be different. The caller is
162 /// responsible for creating or updating the operation transferring flow of
163 /// control to the region and passing it the correct block arguments.
cloneRegionBefore(Region & region,Region & parent,Region::iterator before,BlockAndValueMapping & mapping)164 void PatternRewriter::cloneRegionBefore(Region ®ion, Region &parent,
165 Region::iterator before,
166 BlockAndValueMapping &mapping) {
167 region.cloneInto(&parent, before, mapping);
168 }
cloneRegionBefore(Region & region,Region & parent,Region::iterator before)169 void PatternRewriter::cloneRegionBefore(Region ®ion, Region &parent,
170 Region::iterator before) {
171 BlockAndValueMapping mapping;
172 cloneRegionBefore(region, parent, before, mapping);
173 }
cloneRegionBefore(Region & region,Block * before)174 void PatternRewriter::cloneRegionBefore(Region ®ion, Block *before) {
175 cloneRegionBefore(region, *before->getParent(), before->getIterator());
176 }
177
178 //===----------------------------------------------------------------------===//
179 // PatternMatcher implementation
180 //===----------------------------------------------------------------------===//
181
applyCostModel(CostModel model)182 void PatternApplicator::applyCostModel(CostModel model) {
183 // Separate patterns by root kind to simplify lookup later on.
184 patterns.clear();
185 anyOpPatterns.clear();
186 for (const auto &pat : owningPatternList) {
187 // If the pattern is always impossible to match, just ignore it.
188 if (pat->getBenefit().isImpossibleToMatch())
189 continue;
190 if (Optional<OperationName> opName = pat->getRootKind())
191 patterns[*opName].push_back(pat.get());
192 else
193 anyOpPatterns.push_back(pat.get());
194 }
195
196 // Sort the patterns using the provided cost model.
197 llvm::SmallDenseMap<RewritePattern *, PatternBenefit> benefits;
198 auto cmp = [&benefits](RewritePattern *lhs, RewritePattern *rhs) {
199 return benefits[lhs] > benefits[rhs];
200 };
201 auto processPatternList = [&](SmallVectorImpl<RewritePattern *> &list) {
202 // Special case for one pattern in the list, which is the most common case.
203 if (list.size() == 1) {
204 if (model(*list.front()).isImpossibleToMatch())
205 list.clear();
206 return;
207 }
208
209 // Collect the dynamic benefits for the current pattern list.
210 benefits.clear();
211 for (RewritePattern *pat : list)
212 benefits.try_emplace(pat, model(*pat));
213
214 // Sort patterns with highest benefit first, and remove those that are
215 // impossible to match.
216 std::stable_sort(list.begin(), list.end(), cmp);
217 while (!list.empty() && benefits[list.back()].isImpossibleToMatch())
218 list.pop_back();
219 };
220 for (auto &it : patterns)
221 processPatternList(it.second);
222 processPatternList(anyOpPatterns);
223 }
224
walkAllPatterns(function_ref<void (const RewritePattern &)> walk)225 void PatternApplicator::walkAllPatterns(
226 function_ref<void(const RewritePattern &)> walk) {
227 for (auto &it : owningPatternList)
228 walk(*it);
229 }
230
matchAndRewrite(Operation * op,PatternRewriter & rewriter,function_ref<bool (const RewritePattern &)> canApply,function_ref<void (const RewritePattern &)> onFailure,function_ref<LogicalResult (const RewritePattern &)> onSuccess)231 LogicalResult PatternApplicator::matchAndRewrite(
232 Operation *op, PatternRewriter &rewriter,
233 function_ref<bool(const RewritePattern &)> canApply,
234 function_ref<void(const RewritePattern &)> onFailure,
235 function_ref<LogicalResult(const RewritePattern &)> onSuccess) {
236 // Check to see if there are patterns matching this specific operation type.
237 MutableArrayRef<RewritePattern *> opPatterns;
238 auto patternIt = patterns.find(op->getName());
239 if (patternIt != patterns.end())
240 opPatterns = patternIt->second;
241
242 // Process the patterns for that match the specific operation type, and any
243 // operation type in an interleaved fashion.
244 // FIXME: It'd be nice to just write an llvm::make_merge_range utility
245 // and pass in a comparison function. That would make this code trivial.
246 auto opIt = opPatterns.begin(), opE = opPatterns.end();
247 auto anyIt = anyOpPatterns.begin(), anyE = anyOpPatterns.end();
248 while (opIt != opE && anyIt != anyE) {
249 // Try to match the pattern providing the most benefit.
250 RewritePattern *pattern;
251 if ((*opIt)->getBenefit() >= (*anyIt)->getBenefit())
252 pattern = *(opIt++);
253 else
254 pattern = *(anyIt++);
255
256 // Otherwise, try to match the generic pattern.
257 if (succeeded(matchAndRewrite(op, *pattern, rewriter, canApply, onFailure,
258 onSuccess)))
259 return success();
260 }
261 // If we break from the loop, then only one of the ranges can still have
262 // elements. Loop over both without checking given that we don't need to
263 // interleave anymore.
264 for (RewritePattern *pattern : llvm::concat<RewritePattern *>(
265 llvm::make_range(opIt, opE), llvm::make_range(anyIt, anyE))) {
266 if (succeeded(matchAndRewrite(op, *pattern, rewriter, canApply, onFailure,
267 onSuccess)))
268 return success();
269 }
270 return failure();
271 }
272
matchAndRewrite(Operation * op,const RewritePattern & pattern,PatternRewriter & rewriter,function_ref<bool (const RewritePattern &)> canApply,function_ref<void (const RewritePattern &)> onFailure,function_ref<LogicalResult (const RewritePattern &)> onSuccess)273 LogicalResult PatternApplicator::matchAndRewrite(
274 Operation *op, const RewritePattern &pattern, PatternRewriter &rewriter,
275 function_ref<bool(const RewritePattern &)> canApply,
276 function_ref<void(const RewritePattern &)> onFailure,
277 function_ref<LogicalResult(const RewritePattern &)> onSuccess) {
278 // Check that the pattern can be applied.
279 if (canApply && !canApply(pattern))
280 return failure();
281
282 // Try to match and rewrite this pattern. The patterns are sorted by
283 // benefit, so if we match we can immediately rewrite.
284 rewriter.setInsertionPoint(op);
285 if (succeeded(pattern.matchAndRewrite(op, rewriter)))
286 return success(!onSuccess || succeeded(onSuccess(pattern)));
287
288 if (onFailure)
289 onFailure(pattern);
290 return failure();
291 }
292