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