[mlir] GreedyPatternRewriter: fix counting of iterations

[springerm]

Summary

The GreedyPatternRewriteDriver did previously not count the first iterations. I.e., when setting config.maxIterations = 1, two iterations were performed. In pratice, this number is not really important; we usually just a limit in some reasonable order of magnitude. However, this fix allows us to write better convergence/worklist tests with carefully crafted test patterns to purposely trigger edge cases in the driver.

Similarly, the first rewrite was previously not counted towards config.maxNumRewrites.

For consistency, OpPatternRewriteDriver now uses config.maxNumRewrites instead of config.maxIterations; this driver does not have "iterations", it consists of a single loop (corresponding to the inner loop in the GreedyPatternRewriteDriver).

[mehdi_amini]

LG, but I wonder if it wasn't intentional: that is it controlled how much the driver would "loop around" in a do {} while() fashion. The idea being that passing 0 wouldn't make sense otherwise.

[springerm]

In D141365#4039227, @mehdi_amini wrote:

LG, but I wonder if it wasn't intentional: that is it controlled how much the driver would "loop around" in a do {} while() fashion. The idea being that passing 0 wouldn't make sense otherwise.

0 would make it a "no-op" indeed.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

protected:
   /// efficiently remove operations from the worklist when they are erased, even
   /// if they aren't the root of a pattern.
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
-private:
+protected:
   /// Configuration information for how to simplify.
   GreedyRewriteConfig config;
 
+private:
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
auto insertKnownConstant = [&](Operation *op) {
     Attribute constValue;
     if (matchPattern(op, m_Constant(&constValue)))
       if (!folder.insertKnownConstant(op, constValue))
         return true;
     return false;
   };
 
   bool changed = false;
-  unsigned iteration = 0;
+  int64_t iteration = 0;
   do {
+    // Check if the iteration limit was reached.
+    if (iteration++ >= config.maxIterations &&
+        config.maxIterations != GreedyRewriteConfig::kNoLimit)
+      break;
+
     worklist.clear();
     worklistMap.clear();
 
     if (!config.useTopDownTraversal) {
       // Add operations to the worklist in postorder.
       for (auto &region : regions) {
         region.walk([&](Operation *op) {
           if (!insertKnownConstant(op))
if (!config.useTopDownTraversal) {
         worklistMap[worklist[i]] = i;
     }
 
     // These are scratch vectors used in the folding loop below.
     SmallVector<Value, 8> originalOperands, resultValues;
 
     changed = false;
     int64_t numRewrites = 0;
-    while (!worklist.empty()) {
+    while (!worklist.empty() &&
+           (numRewrites < config.maxNumRewrites ||
+            config.maxNumRewrites == GreedyRewriteConfig::kNoLimit)) {
       auto *op = popFromWorklist();
 
       // Nulls get added to the worklist when operations are removed, ignore
       // them.
       if (op == nullptr)
         continue;
 
       LLVM_DEBUG({
#ifndef NDEBUG
           matcher.matchAndRewrite(op, *this, canApply, onFailure, onSuccess);
       if (succeeded(matchResult))
         LLVM_DEBUG(logResultWithLine("success", "pattern matched"));
       else
         LLVM_DEBUG(logResultWithLine("failure", "pattern failed to match"));
 #else
       LogicalResult matchResult = matcher.matchAndRewrite(op, *this);
 #endif
+
       if (succeeded(matchResult)) {
         changed = true;
-        if (numRewrites++ >= config.maxNumRewrites &&
-            config.maxNumRewrites != GreedyRewriteConfig::kNoLimit)
-          break;
+        ++numRewrites;
       }
     }
 
     // After applying patterns, make sure that the CFG of each of the regions
     // is kept up to date.
     if (config.enableRegionSimplification)
       changed |= succeeded(simplifyRegions(*this, regions));
-  } while (changed && (iteration++ < config.maxIterations ||
-                       config.maxIterations == GreedyRewriteConfig::kNoLimit));
+  } while (changed);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
 void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
   // Check to see if the worklist already contains this op.
   if (worklistMap.count(op))
mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
   (void)regionIsIsolated;
   assert(llvm::all_of(regions, regionIsIsolated) &&
          "patterns can only be applied to operations IsolatedFromAbove");
 
   // Start the pattern driver.
   GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns, config);
   bool converged = driver.simplify(regions);
   LLVM_DEBUG(if (!converged) {
-    llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
+    llvm::dbgs() << "The pattern rewrite did not converge after scanning "
                  << config.maxIterations << " times\n";
   });
   return success(converged);
 }
 
 //===----------------------------------------------------------------------===//
 // OpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This is a simple driver for the PatternMatcher to apply patterns and perform
 /// folding on a single op. It repeatedly applies locally optimal patterns.
 class OpPatternRewriteDriver : public PatternRewriter {
 public:
   explicit OpPatternRewriteDriver(MLIRContext *ctx,
                                   const FrozenRewritePatternSet &patterns)
       : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
     // Apply a simple cost model based solely on pattern benefit.
     matcher.applyDefaultCostModel();
   }
 
-  LogicalResult simplifyLocally(Operation *op, int maxIterations, bool &erased);
+  LogicalResult simplifyLocally(Operation *op, int64_t maxNumRewrites,
+                                bool &erased);
 
   // These are hooks implemented for PatternRewriter.
 protected:
   /// If an operation is about to be removed, mark it so that we can let clients
   /// know.
   void notifyOperationRemoved(Operation *op) override {
     opErasedViaPatternRewrites = true;
   }
private:
   bool opErasedViaPatternRewrites = false;
 };
 
 } // namespace
 
 /// Performs the rewrites and folding only on `op`. The simplification
 /// converges if the op is erased as a result of being folded, replaced, or
 /// becoming dead, or no more changes happen in an iteration. Returns success if
-/// the rewrite converges in `maxIterations`. `erased` is set to true if `op`
+/// the rewrite converges in `maxNumRewrites`. `erased` is set to true if `op`
 /// gets erased.
 LogicalResult OpPatternRewriteDriver::simplifyLocally(Operation *op,
-                                                      int maxIterations,
+                                                      int64_t maxNumRewrites,
                                                       bool &erased) {
   bool changed = false;
   erased = false;
   opErasedViaPatternRewrites = false;
-  int iterations = 0;
-  // Iterate until convergence or until maxIterations. Deletion of the op as
+  int64_t numRewrites = 0;
+  // Iterate until convergence or until maxNumRewrites. Deletion of the op as
   // a result of being dead or folded is convergence.
   do {
+    if (numRewrites >= maxNumRewrites &&
+        maxNumRewrites != GreedyRewriteConfig::kNoLimit)
+      break;
+
     changed = false;
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
       op->erase();
       erased = true;
       return success();
     }
 
     // Try to fold this op.
     bool inPlaceUpdate;
     if (succeeded(folder.tryToFold(op, /*processGeneratedConstants=*/nullptr,
                                    /*preReplaceAction=*/nullptr,
                                    &inPlaceUpdate))) {
       changed = true;
       if (!inPlaceUpdate) {
         erased = true;
         return success();
       }
     }
 
     // Try to match one of the patterns. The rewriter is automatically
     // notified of any necessary changes, so there is nothing else to do here.
-    changed |= succeeded(matcher.matchAndRewrite(op, *this));
+    if (succeeded(matcher.matchAndRewrite(op, *this))) {
+      changed = true;
+      ++numRewrites;
+    }
     if ((erased = opErasedViaPatternRewrites))
       return success();
-  } while (changed && (++iterations < maxIterations ||
-                       maxIterations == GreedyRewriteConfig::kNoLimit));
+  } while (changed);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return failure(changed);
 }
 
 //===----------------------------------------------------------------------===//
 // MultiOpPatternRewriteDriver
 //===----------------------------------------------------------------------===//
bool MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops) {
   bool changed = false;
   worklist.clear();
   worklistMap.clear();
   for (Operation *op : ops)
     addToWorklist(op);
 
   // These are scratch vectors used in the folding loop below.
   SmallVector<Value, 8> originalOperands, resultValues;
-  while (!worklist.empty()) {
+  int64_t numRewrites = 0;
+  while (!worklist.empty() &&
+         (numRewrites < config.maxNumRewrites ||
+          config.maxNumRewrites == GreedyRewriteConfig::kNoLimit)) {
     Operation *op = popFromWorklist();
 
     // Nulls get added to the worklist when operations are removed, ignore
     // them.
     if (op == nullptr)
       continue;
 
     assert((!strictMode || strictModeFilteredOps.contains(op)) &&
if (succeeded(folder.tryToFold(op, processGeneratedConstants,
         // Op has been erased.
         continue;
       }
     }
 
     // Try to match one of the patterns. The rewriter is automatically
     // notified of any necessary changes, so there is nothing else to do
     // here.
-    changed |= succeeded(matcher.matchAndRewrite(op, *this));
+    if (succeeded(matcher.matchAndRewrite(op, *this))) {
+      changed = true;
+      ++numRewrites;
+    }
   }
 
   return changed;
 }
 
 /// Rewrites only `op` using the supplied canonicalization patterns and
 /// folding. `erased` is set to true if the op is erased as a result of being
 /// folded, replaced, or dead.
 LogicalResult mlir::applyOpPatternsAndFold(
     Operation *op, const FrozenRewritePatternSet &patterns, bool *erased) {
   // Start the pattern driver.
   GreedyRewriteConfig config;
   OpPatternRewriteDriver driver(op->getContext(), patterns);
   bool opErased;
   LogicalResult converged =
-      driver.simplifyLocally(op, config.maxIterations, opErased);
+      driver.simplifyLocally(op, config.maxNumRewrites, opErased);
   if (erased)
     *erased = opErased;
   LLVM_DEBUG(if (failed(converged)) {
-    llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
-                 << config.maxIterations << " times";
+    llvm::dbgs() << "The pattern rewrite did not converge after "
+                 << config.maxNumRewrites << " rewrites";
   });
   return converged;
 }
 
 bool mlir::applyOpPatternsAndFold(ArrayRef<Operation *> ops,
                                   const FrozenRewritePatternSet &patterns,
                                   bool strict) {
   if (ops.empty())
     return false;
 
   // Start the pattern driver.
   MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
                                      strict);
   return driver.simplifyLocally(ops);
 }