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 
12 using namespace mlir;
13 
14 //===----------------------------------------------------------------------===//
15 // PatternBenefit
16 //===----------------------------------------------------------------------===//
17 
18 PatternBenefit::PatternBenefit(unsigned benefit) : representation(benefit) {
19   assert(representation == benefit && benefit != ImpossibleToMatchSentinel &&
20          "This pattern match benefit is too large to represent");
21 }
22 
23 unsigned short PatternBenefit::getBenefit() const {
24   assert(!isImpossibleToMatch() && "Pattern doesn't match");
25   return representation;
26 }
27 
28 //===----------------------------------------------------------------------===//
29 // Pattern
30 //===----------------------------------------------------------------------===//
31 
32 Pattern::Pattern(StringRef rootName, PatternBenefit benefit,
33                  MLIRContext *context)
34     : rootKind(OperationName(rootName, context)), benefit(benefit) {}
35 Pattern::Pattern(PatternBenefit benefit, MatchAnyOpTypeTag tag)
36     : benefit(benefit) {}
37 Pattern::Pattern(StringRef rootName, ArrayRef<StringRef> generatedNames,
38                  PatternBenefit benefit, MLIRContext *context)
39     : Pattern(rootName, benefit, context) {
40   generatedOps.reserve(generatedNames.size());
41   std::transform(generatedNames.begin(), generatedNames.end(),
42                  std::back_inserter(generatedOps), [context](StringRef name) {
43                    return OperationName(name, context);
44                  });
45 }
46 Pattern::Pattern(ArrayRef<StringRef> generatedNames, PatternBenefit benefit,
47                  MLIRContext *context, MatchAnyOpTypeTag tag)
48     : Pattern(benefit, tag) {
49   generatedOps.reserve(generatedNames.size());
50   std::transform(generatedNames.begin(), generatedNames.end(),
51                  std::back_inserter(generatedOps), [context](StringRef name) {
52                    return OperationName(name, context);
53                  });
54 }
55 
56 //===----------------------------------------------------------------------===//
57 // RewritePattern
58 //===----------------------------------------------------------------------===//
59 
60 void RewritePattern::rewrite(Operation *op, PatternRewriter &rewriter) const {
61   llvm_unreachable("need to implement either matchAndRewrite or one of the "
62                    "rewrite functions!");
63 }
64 
65 LogicalResult RewritePattern::match(Operation *op) const {
66   llvm_unreachable("need to implement either match or matchAndRewrite!");
67 }
68 
69 /// Out-of-line vtable anchor.
70 void RewritePattern::anchor() {}
71 
72 //===----------------------------------------------------------------------===//
73 // PDLValue
74 //===----------------------------------------------------------------------===//
75 
76 void PDLValue::print(raw_ostream &os) {
77   if (!impl) {
78     os << "<Null-PDLValue>";
79     return;
80   }
81   if (Value val = impl.dyn_cast<Value>()) {
82     os << val;
83     return;
84   }
85   AttrOpTypeImplT aotImpl = impl.get<AttrOpTypeImplT>();
86   if (Attribute attr = aotImpl.dyn_cast<Attribute>())
87     os << attr;
88   else if (Operation *op = aotImpl.dyn_cast<Operation *>())
89     os << *op;
90   else
91     os << aotImpl.get<Type>();
92 }
93 
94 //===----------------------------------------------------------------------===//
95 // PDLPatternModule
96 //===----------------------------------------------------------------------===//
97 
98 void PDLPatternModule::mergeIn(PDLPatternModule &&other) {
99   // Ignore the other module if it has no patterns.
100   if (!other.pdlModule)
101     return;
102   // Steal the other state if we have no patterns.
103   if (!pdlModule) {
104     constraintFunctions = std::move(other.constraintFunctions);
105     rewriteFunctions = std::move(other.rewriteFunctions);
106     pdlModule = std::move(other.pdlModule);
107     return;
108   }
109   // Steal the functions of the other module.
110   for (auto &it : constraintFunctions)
111     registerConstraintFunction(it.first(), std::move(it.second));
112   for (auto &it : rewriteFunctions)
113     registerRewriteFunction(it.first(), std::move(it.second));
114 
115   // Merge the pattern operations from the other module into this one.
116   Block *block = pdlModule->getBody();
117   block->getTerminator()->erase();
118   block->getOperations().splice(block->end(),
119                                 other.pdlModule->getBody()->getOperations());
120 }
121 
122 //===----------------------------------------------------------------------===//
123 // Function Registry
124 
125 void PDLPatternModule::registerConstraintFunction(
126     StringRef name, PDLConstraintFunction constraintFn) {
127   auto it = constraintFunctions.try_emplace(name, std::move(constraintFn));
128   (void)it;
129   assert(it.second &&
130          "constraint with the given name has already been registered");
131 }
132 
133 void PDLPatternModule::registerRewriteFunction(StringRef name,
134                                                PDLRewriteFunction rewriteFn) {
135   auto it = rewriteFunctions.try_emplace(name, std::move(rewriteFn));
136   (void)it;
137   assert(it.second && "native rewrite function with the given name has "
138                       "already been registered");
139 }
140 
141 //===----------------------------------------------------------------------===//
142 // RewriterBase
143 //===----------------------------------------------------------------------===//
144 
145 RewriterBase::~RewriterBase() {
146   // Out of line to provide a vtable anchor for the class.
147 }
148 
149 /// This method replaces the uses of the results of `op` with the values in
150 /// `newValues` when the provided `functor` returns true for a specific use.
151 /// The number of values in `newValues` is required to match the number of
152 /// results of `op`.
153 void RewriterBase::replaceOpWithIf(
154     Operation *op, ValueRange newValues, bool *allUsesReplaced,
155     llvm::unique_function<bool(OpOperand &) const> functor) {
156   assert(op->getNumResults() == newValues.size() &&
157          "incorrect number of values to replace operation");
158 
159   // Notify the rewriter subclass that we're about to replace this root.
160   notifyRootReplaced(op);
161 
162   // Replace each use of the results when the functor is true.
163   bool replacedAllUses = true;
164   for (auto it : llvm::zip(op->getResults(), newValues)) {
165     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), functor);
166     replacedAllUses &= std::get<0>(it).use_empty();
167   }
168   if (allUsesReplaced)
169     *allUsesReplaced = replacedAllUses;
170 }
171 
172 /// This method replaces the uses of the results of `op` with the values in
173 /// `newValues` when a use is nested within the given `block`. The number of
174 /// values in `newValues` is required to match the number of results of `op`.
175 /// If all uses of this operation are replaced, the operation is erased.
176 void RewriterBase::replaceOpWithinBlock(Operation *op, ValueRange newValues,
177                                         Block *block, bool *allUsesReplaced) {
178   replaceOpWithIf(op, newValues, allUsesReplaced, [block](OpOperand &use) {
179     return block->getParentOp()->isProperAncestor(use.getOwner());
180   });
181 }
182 
183 /// This method replaces the results of the operation with the specified list of
184 /// values. The number of provided values must match the number of results of
185 /// the operation.
186 void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {
187   // Notify the rewriter subclass that we're about to replace this root.
188   notifyRootReplaced(op);
189 
190   assert(op->getNumResults() == newValues.size() &&
191          "incorrect # of replacement values");
192   op->replaceAllUsesWith(newValues);
193 
194   notifyOperationRemoved(op);
195   op->erase();
196 }
197 
198 /// This method erases an operation that is known to have no uses. The uses of
199 /// the given operation *must* be known to be dead.
200 void RewriterBase::eraseOp(Operation *op) {
201   assert(op->use_empty() && "expected 'op' to have no uses");
202   notifyOperationRemoved(op);
203   op->erase();
204 }
205 
206 void RewriterBase::eraseBlock(Block *block) {
207   for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
208     assert(op.use_empty() && "expected 'op' to have no uses");
209     eraseOp(&op);
210   }
211   block->erase();
212 }
213 
214 /// Merge the operations of block 'source' into the end of block 'dest'.
215 /// 'source's predecessors must be empty or only contain 'dest`.
216 /// 'argValues' is used to replace the block arguments of 'source' after
217 /// merging.
218 void RewriterBase::mergeBlocks(Block *source, Block *dest,
219                                ValueRange argValues) {
220   assert(llvm::all_of(source->getPredecessors(),
221                       [dest](Block *succ) { return succ == dest; }) &&
222          "expected 'source' to have no predecessors or only 'dest'");
223   assert(argValues.size() == source->getNumArguments() &&
224          "incorrect # of argument replacement values");
225 
226   // Replace all of the successor arguments with the provided values.
227   for (auto it : llvm::zip(source->getArguments(), argValues))
228     std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
229 
230   // Splice the operations of the 'source' block into the 'dest' block and erase
231   // it.
232   dest->getOperations().splice(dest->end(), source->getOperations());
233   source->dropAllUses();
234   source->erase();
235 }
236 
237 // Merge the operations of block 'source' before the operation 'op'. Source
238 // block should not have existing predecessors or successors.
239 void RewriterBase::mergeBlockBefore(Block *source, Operation *op,
240                                     ValueRange argValues) {
241   assert(source->hasNoPredecessors() &&
242          "expected 'source' to have no predecessors");
243   assert(source->hasNoSuccessors() &&
244          "expected 'source' to have no successors");
245 
246   // Split the block containing 'op' into two, one containing all operations
247   // before 'op' (prologue) and another (epilogue) containing 'op' and all
248   // operations after it.
249   Block *prologue = op->getBlock();
250   Block *epilogue = splitBlock(prologue, op->getIterator());
251 
252   // Merge the source block at the end of the prologue.
253   mergeBlocks(source, prologue, argValues);
254 
255   // Merge the epilogue at the end the prologue.
256   mergeBlocks(epilogue, prologue);
257 }
258 
259 /// Split the operations starting at "before" (inclusive) out of the given
260 /// block into a new block, and return it.
261 Block *RewriterBase::splitBlock(Block *block, Block::iterator before) {
262   return block->splitBlock(before);
263 }
264 
265 /// 'op' and 'newOp' are known to have the same number of results, replace the
266 /// uses of op with uses of newOp
267 void RewriterBase::replaceOpWithResultsOfAnotherOp(Operation *op,
268                                                    Operation *newOp) {
269   assert(op->getNumResults() == newOp->getNumResults() &&
270          "replacement op doesn't match results of original op");
271   if (op->getNumResults() == 1)
272     return replaceOp(op, newOp->getResult(0));
273   return replaceOp(op, newOp->getResults());
274 }
275 
276 /// Move the blocks that belong to "region" before the given position in
277 /// another region.  The two regions must be different.  The caller is in
278 /// charge to update create the operation transferring the control flow to the
279 /// region and pass it the correct block arguments.
280 void RewriterBase::inlineRegionBefore(Region &region, Region &parent,
281                                       Region::iterator before) {
282   parent.getBlocks().splice(before, region.getBlocks());
283 }
284 void RewriterBase::inlineRegionBefore(Region &region, Block *before) {
285   inlineRegionBefore(region, *before->getParent(), before->getIterator());
286 }
287 
288 /// Clone the blocks that belong to "region" before the given position in
289 /// another region "parent". The two regions must be different. The caller is
290 /// responsible for creating or updating the operation transferring flow of
291 /// control to the region and passing it the correct block arguments.
292 void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
293                                      Region::iterator before,
294                                      BlockAndValueMapping &mapping) {
295   region.cloneInto(&parent, before, mapping);
296 }
297 void RewriterBase::cloneRegionBefore(Region &region, Region &parent,
298                                      Region::iterator before) {
299   BlockAndValueMapping mapping;
300   cloneRegionBefore(region, parent, before, mapping);
301 }
302 void RewriterBase::cloneRegionBefore(Region &region, Block *before) {
303   cloneRegionBefore(region, *before->getParent(), before->getIterator());
304 }
305