1 //===- GreedyPatternRewriteDriver.cpp - A greedy rewriter -----------------===//
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 // This file implements mlir::applyPatternsGreedily.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/IR/PatternMatch.h"
14 #include "mlir/Interfaces/SideEffects.h"
15 #include "mlir/Transforms/FoldUtils.h"
16 #include "mlir/Transforms/RegionUtils.h"
17 #include "llvm/ADT/DenseMap.h"
18 #include "llvm/Support/CommandLine.h"
19 #include "llvm/Support/Debug.h"
20 #include "llvm/Support/raw_ostream.h"
21 
22 using namespace mlir;
23 
24 #define DEBUG_TYPE "pattern-matcher"
25 
26 /// The max number of iterations scanning for pattern match.
27 static unsigned maxPatternMatchIterations = 10;
28 
29 namespace {
30 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
31 /// applies the locally optimal patterns in a roughly "bottom up" way.
32 class GreedyPatternRewriteDriver : public PatternRewriter {
33 public:
34   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
35                                       const OwningRewritePatternList &patterns)
36       : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
37     worklist.reserve(64);
38   }
39 
40   /// Perform the rewrites while folding and erasing any dead ops. Return true
41   /// if the rewrite converges in `maxIterations`.
42   bool simplify(MutableArrayRef<Region> regions, int maxIterations);
43 
44   void addToWorklist(Operation *op) {
45     // Check to see if the worklist already contains this op.
46     if (worklistMap.count(op))
47       return;
48 
49     worklistMap[op] = worklist.size();
50     worklist.push_back(op);
51   }
52 
53   Operation *popFromWorklist() {
54     auto *op = worklist.back();
55     worklist.pop_back();
56 
57     // This operation is no longer in the worklist, keep worklistMap up to date.
58     if (op)
59       worklistMap.erase(op);
60     return op;
61   }
62 
63   /// If the specified operation is in the worklist, remove it.  If not, this is
64   /// a no-op.
65   void removeFromWorklist(Operation *op) {
66     auto it = worklistMap.find(op);
67     if (it != worklistMap.end()) {
68       assert(worklist[it->second] == op && "malformed worklist data structure");
69       worklist[it->second] = nullptr;
70       worklistMap.erase(it);
71     }
72   }
73 
74   // These are hooks implemented for PatternRewriter.
75 protected:
76   // Implement the hook for inserting operations, and make sure that newly
77   // inserted ops are added to the worklist for processing.
78   Operation *insert(Operation *op) override {
79     addToWorklist(op);
80     return OpBuilder::insert(op);
81   }
82 
83   // If an operation is about to be removed, make sure it is not in our
84   // worklist anymore because we'd get dangling references to it.
85   void notifyOperationRemoved(Operation *op) override {
86     addToWorklist(op->getOperands());
87     op->walk([this](Operation *operation) {
88       removeFromWorklist(operation);
89       folder.notifyRemoval(operation);
90     });
91   }
92 
93   // When the root of a pattern is about to be replaced, it can trigger
94   // simplifications to its users - make sure to add them to the worklist
95   // before the root is changed.
96   void notifyRootReplaced(Operation *op) override {
97     for (auto result : op->getResults())
98       for (auto *user : result.getUsers())
99         addToWorklist(user);
100   }
101 
102 private:
103   // Look over the provided operands for any defining operations that should
104   // be re-added to the worklist. This function should be called when an
105   // operation is modified or removed, as it may trigger further
106   // simplifications.
107   template <typename Operands>
108   void addToWorklist(Operands &&operands) {
109     for (Value operand : operands) {
110       // If the use count of this operand is now < 2, we re-add the defining
111       // operation to the worklist.
112       // TODO(riverriddle) This is based on the fact that zero use operations
113       // may be deleted, and that single use values often have more
114       // canonicalization opportunities.
115       if (!operand.use_empty() && !operand.hasOneUse())
116         continue;
117       if (auto *defInst = operand.getDefiningOp())
118         addToWorklist(defInst);
119     }
120   }
121 
122   /// The low-level pattern matcher.
123   RewritePatternMatcher matcher;
124 
125   /// The worklist for this transformation keeps track of the operations that
126   /// need to be revisited, plus their index in the worklist.  This allows us to
127   /// efficiently remove operations from the worklist when they are erased, even
128   /// if they aren't the root of a pattern.
129   std::vector<Operation *> worklist;
130   DenseMap<Operation *, unsigned> worklistMap;
131 
132   /// Non-pattern based folder for operations.
133   OperationFolder folder;
134 };
135 } // end anonymous namespace
136 
137 /// Performs the rewrites while folding and erasing any dead ops. Returns true
138 /// if the rewrite converges in `maxIterations`.
139 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions,
140                                           int maxIterations) {
141   // Add the given operation to the worklist.
142   auto collectOps = [this](Operation *op) { addToWorklist(op); };
143 
144   bool changed = false;
145   int i = 0;
146   do {
147     // Add all nested operations to the worklist.
148     for (auto &region : regions)
149       region.walk(collectOps);
150 
151     // These are scratch vectors used in the folding loop below.
152     SmallVector<Value, 8> originalOperands, resultValues;
153 
154     changed = false;
155     while (!worklist.empty()) {
156       auto *op = popFromWorklist();
157 
158       // Nulls get added to the worklist when operations are removed, ignore
159       // them.
160       if (op == nullptr)
161         continue;
162 
163       // If the operation is trivially dead - remove it.
164       if (isOpTriviallyDead(op)) {
165         notifyOperationRemoved(op);
166         op->erase();
167         changed = true;
168         continue;
169       }
170 
171       // Collects all the operands and result uses of the given `op` into work
172       // list. Also remove `op` and nested ops from worklist.
173       originalOperands.assign(op->operand_begin(), op->operand_end());
174       auto preReplaceAction = [&](Operation *op) {
175         // Add the operands to the worklist for visitation.
176         addToWorklist(originalOperands);
177 
178         // Add all the users of the result to the worklist so we make sure
179         // to revisit them.
180         for (auto result : op->getResults())
181           for (auto *userOp : result.getUsers())
182             addToWorklist(userOp);
183 
184         notifyOperationRemoved(op);
185       };
186 
187       // Try to fold this op.
188       bool inPlaceUpdate;
189       if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
190                                       &inPlaceUpdate)))) {
191         changed = true;
192         if (!inPlaceUpdate)
193           continue;
194       }
195 
196       // Make sure that any new operations are inserted at this point.
197       setInsertionPoint(op);
198 
199       // Try to match one of the patterns. The rewriter is automatically
200       // notified of any necessary changes, so there is nothing else to do here.
201       changed |= matcher.matchAndRewrite(op, *this);
202     }
203 
204     // After applying patterns, make sure that the CFG of each of the regions is
205     // kept up to date.
206     if (succeeded(simplifyRegions(regions))) {
207       folder.clear();
208       changed = true;
209     }
210   } while (changed && ++i < maxIterations);
211   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
212   return !changed;
213 }
214 
215 /// Rewrite the regions of the specified operation, which must be isolated from
216 /// above, by repeatedly applying the highest benefit patterns in a greedy
217 /// work-list driven manner. Return true if no more patterns can be matched in
218 /// the result operation regions.
219 /// Note: This does not apply patterns to the top-level operation itself.
220 ///
221 bool mlir::applyPatternsAndFoldGreedily(
222     Operation *op, const OwningRewritePatternList &patterns) {
223   return applyPatternsAndFoldGreedily(op->getRegions(), patterns);
224 }
225 
226 /// Rewrite the given regions, which must be isolated from above.
227 bool mlir::applyPatternsAndFoldGreedily(
228     MutableArrayRef<Region> regions, const OwningRewritePatternList &patterns) {
229   if (regions.empty())
230     return true;
231 
232   // The top-level operation must be known to be isolated from above to
233   // prevent performing canonicalizations on operations defined at or above
234   // the region containing 'op'.
235   auto regionIsIsolated = [](Region &region) {
236     return region.getParentOp()->isKnownIsolatedFromAbove();
237   };
238   (void)regionIsIsolated;
239   assert(llvm::all_of(regions, regionIsIsolated) &&
240          "patterns can only be applied to operations IsolatedFromAbove");
241 
242   // Start the pattern driver.
243   GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns);
244   bool converged = driver.simplify(regions, maxPatternMatchIterations);
245   LLVM_DEBUG(if (!converged) {
246     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
247                  << maxPatternMatchIterations << " times";
248   });
249   return converged;
250 }
251