[Mlir-commits] [mlir] b7144ab - [NFC] Clean up a few methods within GreedyPatternRewriter

[River Riddle]

Author: River Riddle
Date: 2021-10-21T17:14:35Z
New Revision: b7144ab7651702b7171f53cd778780fc3a52863e

URL: https://github.com/llvm/llvm-project/commit/b7144ab7651702b7171f53cd778780fc3a52863e
DIFF: https://github.com/llvm/llvm-project/commit/b7144ab7651702b7171f53cd778780fc3a52863e.diff

LOG: [NFC] Clean up a few methods within GreedyPatternRewriter

Move a few methods out of line and clean up comments.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index ba792a79cc16c..d4ac64ca8fd01 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -35,89 +35,40 @@ class GreedyPatternRewriteDriver : public PatternRewriter {
 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const FrozenRewritePatternSet &patterns,
-                                      const GreedyRewriteConfig &config)
-      : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
-    worklist.reserve(64);
-
-    // Apply a simple cost model based solely on pattern benefit.
-    matcher.applyDefaultCostModel();
-  }
+                                      const GreedyRewriteConfig &config);
 
+  /// Simplify the operations within the given regions.
   bool simplify(MutableArrayRef<Region> regions);
 
-  void addToWorklist(Operation *op) {
-    // Check to see if the worklist already contains this op.
-    if (worklistMap.count(op))
-      return;
+  /// Add the given operation to the worklist.
+  void addToWorklist(Operation *op);
 
-    worklistMap[op] = worklist.size();
-    worklist.push_back(op);
-  }
+  /// Pop the next operation from the worklist.
+  Operation *popFromWorklist();
 
-  Operation *popFromWorklist() {
-    auto *op = worklist.back();
-    worklist.pop_back();
+  /// If the specified operation is in the worklist, remove it.
+  void removeFromWorklist(Operation *op);
 
-    // This operation is no longer in the worklist, keep worklistMap up to date.
-    if (op)
-      worklistMap.erase(op);
-    return op;
-  }
-
-  /// If the specified operation is in the worklist, remove it.  If not, this is
-  /// a no-op.
-  void removeFromWorklist(Operation *op) {
-    auto it = worklistMap.find(op);
-    if (it != worklistMap.end()) {
-      assert(worklist[it->second] == op && "malformed worklist data structure");
-      worklist[it->second] = nullptr;
-      worklistMap.erase(it);
-    }
-  }
-
-  // These are hooks implemented for PatternRewriter.
 protected:
   // Implement the hook for inserting operations, and make sure that newly
   // inserted ops are added to the worklist for processing.
-  void notifyOperationInserted(Operation *op) override { addToWorklist(op); }
+  void notifyOperationInserted(Operation *op) override;
 
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further
   // simplifications.
   template <typename Operands>
-  void addToWorklist(Operands &&operands) {
-    for (Value operand : operands) {
-      // If the use count of this operand is now < 2, we re-add the defining
-      // operation to the worklist.
-      // TODO: This is based on the fact that zero use operations
-      // may be deleted, and that single use values often have more
-      // canonicalization opportunities.
-      if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
-        continue;
-      if (auto *defOp = operand.getDefiningOp())
-        addToWorklist(defOp);
-    }
-  }
+  void addToWorklist(Operands &&operands);
 
   // If an operation is about to be removed, make sure it is not in our
   // worklist anymore because we'd get dangling references to it.
-  void notifyOperationRemoved(Operation *op) override {
-    addToWorklist(op->getOperands());
-    op->walk([this](Operation *operation) {
-      removeFromWorklist(operation);
-      folder.notifyRemoval(operation);
-    });
-  }
+  void notifyOperationRemoved(Operation *op) override;
 
   // When the root of a pattern is about to be replaced, it can trigger
   // simplifications to its users - make sure to add them to the worklist
   // before the root is changed.
-  void notifyRootReplaced(Operation *op) override {
-    for (auto result : op->getResults())
-      for (auto *user : result.getUsers())
-        addToWorklist(user);
-  }
+  void notifyRootReplaced(Operation *op) override;
 
   /// The low-level pattern applicator.
   PatternApplicator matcher;
@@ -138,8 +89,16 @@ class GreedyPatternRewriteDriver : public PatternRewriter {
 };
 } // end anonymous namespace
 
-/// Performs the rewrites while folding and erasing any dead ops. Returns true
-/// if the rewrite converges in `maxIterations`.
+GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
+    MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
+    const GreedyRewriteConfig &config)
+    : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
+  worklist.reserve(64);
+
+  // Apply a simple cost model based solely on pattern benefit.
+  matcher.applyDefaultCostModel();
+}
+
 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
   bool changed = false;
   unsigned iteration = 0;
@@ -230,6 +189,67 @@ bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
   return !changed;
 }
 
+void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
+  // Check to see if the worklist already contains this op.
+  if (worklistMap.count(op))
+    return;
+
+  worklistMap[op] = worklist.size();
+  worklist.push_back(op);
+}
+
+Operation *GreedyPatternRewriteDriver::popFromWorklist() {
+  auto *op = worklist.back();
+  worklist.pop_back();
+
+  // This operation is no longer in the worklist, keep worklistMap up to date.
+  if (op)
+    worklistMap.erase(op);
+  return op;
+}
+
+void GreedyPatternRewriteDriver::removeFromWorklist(Operation *op) {
+  auto it = worklistMap.find(op);
+  if (it != worklistMap.end()) {
+    assert(worklist[it->second] == op && "malformed worklist data structure");
+    worklist[it->second] = nullptr;
+    worklistMap.erase(it);
+  }
+}
+
+void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
+  addToWorklist(op);
+}
+
+template <typename Operands>
+void GreedyPatternRewriteDriver::addToWorklist(Operands &&operands) {
+  for (Value operand : operands) {
+    // If the use count of this operand is now < 2, we re-add the defining
+    // operation to the worklist.
+    // TODO: This is based on the fact that zero use operations
+    // may be deleted, and that single use values often have more
+    // canonicalization opportunities.
+    if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
+      continue;
+    if (auto *defOp = operand.getDefiningOp())
+      addToWorklist(defOp);
+  }
+}
+
+void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
+  addToWorklist(op->getOperands());
+  op->walk([this](Operation *operation) {
+    removeFromWorklist(operation);
+    folder.notifyRemoval(operation);
+  });
+}
+
+void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op) {
+  for (auto result : op->getResults())
+    for (auto *user : result.getUsers())
+      addToWorklist(user);
+}
+
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner. Return success if no more patterns can be matched