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> void addToWorklist(Operands &&operands) {
108     for (Value operand : operands) {
109       // If the use count of this operand is now < 2, we re-add the defining
110       // operation to the worklist.
111       // TODO(riverriddle) This is based on the fact that zero use operations
112       // may be deleted, and that single use values often have more
113       // canonicalization opportunities.
114       if (!operand.use_empty() && !operand.hasOneUse())
115         continue;
116       if (auto *defInst = operand.getDefiningOp())
117         addToWorklist(defInst);
118     }
119   }
120 
121   /// The low-level pattern matcher.
122   RewritePatternMatcher matcher;
123 
124   /// The worklist for this transformation keeps track of the operations that
125   /// need to be revisited, plus their index in the worklist.  This allows us to
126   /// efficiently remove operations from the worklist when they are erased, even
127   /// if they aren't the root of a pattern.
128   std::vector<Operation *> worklist;
129   DenseMap<Operation *, unsigned> worklistMap;
130 
131   /// Non-pattern based folder for operations.
132   OperationFolder folder;
133 };
134 } // end anonymous namespace
135 
136 /// Perform the rewrites while folding and erasing any dead ops.
137 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions,
138                                           int maxIterations) {
139   // Add the given operation to the worklist.
140   auto collectOps = [this](Operation *op) { addToWorklist(op); };
141 
142   bool changed = false;
143   int i = 0;
144   do {
145     // Add all nested operations to the worklist.
146     for (auto &region : regions)
147       region.walk(collectOps);
148 
149     // These are scratch vectors used in the folding loop below.
150     SmallVector<Value, 8> originalOperands, resultValues;
151 
152     changed = false;
153     while (!worklist.empty()) {
154       auto *op = popFromWorklist();
155 
156       // Nulls get added to the worklist when operations are removed, ignore
157       // them.
158       if (op == nullptr)
159         continue;
160 
161       // If the operation is trivially dead - remove it.
162       if (isOpTriviallyDead(op)) {
163         notifyOperationRemoved(op);
164         op->erase();
165         changed = true;
166         continue;
167       }
168 
169       // Collects all the operands and result uses of the given `op` into work
170       // list. Also remove `op` and nested ops from worklist.
171       originalOperands.assign(op->operand_begin(), op->operand_end());
172       auto preReplaceAction = [&](Operation *op) {
173         // Add the operands to the worklist for visitation.
174         addToWorklist(originalOperands);
175 
176         // Add all the users of the result to the worklist so we make sure
177         // to revisit them.
178         for (auto result : op->getResults())
179           for (auto *userOp : result.getUsers())
180             addToWorklist(userOp);
181 
182         notifyOperationRemoved(op);
183       };
184 
185       // Try to fold this op.
186       if (succeeded(folder.tryToFold(op, collectOps, preReplaceAction))) {
187         changed = true;
188         continue;
189       }
190 
191       // Make sure that any new operations are inserted at this point.
192       setInsertionPoint(op);
193 
194       // Try to match one of the patterns. The rewriter is automatically
195       // notified of any necessary changes, so there is nothing else to do here.
196       changed |= matcher.matchAndRewrite(op, *this);
197     }
198 
199     // After applying patterns, make sure that the CFG of each of the regions is
200     // kept up to date.
201     if (succeeded(simplifyRegions(regions))) {
202       folder.clear();
203       changed = true;
204     }
205   } while (changed && ++i < maxIterations);
206   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
207   return !changed;
208 }
209 
210 /// Rewrite the regions of the specified operation, which must be isolated from
211 /// above, by repeatedly applying the highest benefit patterns in a greedy
212 /// work-list driven manner. Return true if no more patterns can be matched in
213 /// the result operation regions.
214 /// Note: This does not apply patterns to the top-level operation itself.
215 ///
216 bool mlir::applyPatternsAndFoldGreedily(
217     Operation *op, const OwningRewritePatternList &patterns) {
218   return applyPatternsAndFoldGreedily(op->getRegions(), patterns);
219 }
220 
221 /// Rewrite the given regions, which must be isolated from above.
222 bool mlir::applyPatternsAndFoldGreedily(
223     MutableArrayRef<Region> regions, const OwningRewritePatternList &patterns) {
224   if (regions.empty())
225     return true;
226 
227   // The top-level operation must be known to be isolated from above to
228   // prevent performing canonicalizations on operations defined at or above
229   // the region containing 'op'.
230   auto regionIsIsolated = [](Region &region) {
231     return region.getParentOp()->isKnownIsolatedFromAbove();
232   };
233   (void)regionIsIsolated;
234   assert(llvm::all_of(regions, regionIsIsolated) &&
235          "patterns can only be applied to operations IsolatedFromAbove");
236 
237   // Start the pattern driver.
238   GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns);
239   bool converged = driver.simplify(regions, maxPatternMatchIterations);
240   LLVM_DEBUG(if (!converged) {
241     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
242                  << maxPatternMatchIterations << " times";
243   });
244   return converged;
245 }
246