[mlir] Limit number of pattern rewrites in CanonicalizerPass

[springerm]

Summary

The greedy pattern rewriter consists of two nested loops. config.maxIterations (which configurable on the CanonicalizerPass) controls the maximum number of iterations of the outer loops.

/// This specifies the maximum number of times the rewriter will iterate
/// between applying patterns and simplifying regions. Use `kNoLimit` to
/// disable this iteration limit.
int64_t maxIterations = 10;

This change adds config.maxNumRewrites which controls the maximum number of pattern rewrites within an iteration. (It effectively control the maximum number of iterations of the inner loop.)

This flag is meant for debugging and useful in cases where one or multiple faulty patterns can be applied indefinitely, resulting in an infinite loop. Such bugs are difficult to debug even with mlir-opt -debug because no output will printed to the terminal until the process terminates.

[mehdi_amini]

LGTM, but I don't quite get this part of the description:

Such bugs are difficult to debug even with mlir-opt -debug because no output will printed to the terminal until the process terminates.

I expect all the debug to show up as it goes, I see a bunch of logger.getOStream() << "\n"; that should trigger printing.
Should we add a call to logger.flush(); between each iteration?

[springerm]

In D140525#4013979, @mehdi_amini wrote:

LGTM, but I don't quite get this part of the description:

Such bugs are difficult to debug even with mlir-opt -debug because no output will printed to the terminal until the process terminates.

I expect all the debug to show up as it goes, I see a bunch of logger.getOStream() << "\n"; that should trigger printing.
Should we add a call to logger.flush(); between each iteration?

Hmm it could be that this is just an issue when running a test case in Google3. Will drop that part from the commit message.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

public:
   /// larger patterns when given an ambiguous pattern set.
   bool useTopDownTraversal = false;
 
   // Perform control flow optimizations to the region tree after applying all
   // patterns.
   bool enableRegionSimplification = true;
 
   /// This specifies the maximum number of times the rewriter will iterate
-  /// between applying patterns and simplifying regions. Use `kNoIterationLimit`
-  /// to disable this iteration limit.
+  /// between applying patterns and simplifying regions. Use `kNoLimit` to
+  /// disable this iteration limit.
   int64_t maxIterations = 10;
 
-  static constexpr int64_t kNoIterationLimit = -1;
+  /// This specifies the maximum number of rewrites within an iteration. Use
+  /// `kNoLimit` to disable this limit.
+  int64_t maxNumRewrites = kNoLimit;
+
+  static constexpr int64_t kNoLimit = -1;
 };
 
 //===----------------------------------------------------------------------===//
 // applyPatternsGreedily
 //===----------------------------------------------------------------------===//
 
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy

mlir/include/mlir/Transforms/Passes.td

def Canonicalizer : Pass<"canonicalize"> {
   }];
   let constructor = "mlir::createCanonicalizerPass()";
   let options = [
     Option<"topDownProcessingEnabled", "top-down", "bool",
            /*default=*/"true",
            "Seed the worklist in general top-down order">,
     Option<"enableRegionSimplification", "region-simplify", "bool",
            /*default=*/"true",
-           "Seed the worklist in general top-down order">,
+           "Perform control flow optimizations to the region tree">,
     Option<"maxIterations", "max-iterations", "int64_t",
            /*default=*/"10",
-           "Seed the worklist in general top-down order">
+           "Max. iterations between applying patterns / simplifying regions">,
+    Option<"maxNumRewrites", "max-num-rewrites", "int64_t", /*default=*/"-1",
+           "Max. number of pattern rewrites within an iteration">
   ] # RewritePassUtils.options;
 }
 
 def ControlFlowSink : Pass<"control-flow-sink"> {
   let summary = "Sink operations into conditional blocks";
   let description = [{
     This pass implements control-flow sink on operations that implement
     `RegionBranchOpInterface` by moving dominating operations whose only uses

mlir/lib/Transforms/Canonicalizer.cpp

 struct Canonicalizer : public impl::CanonicalizerBase<Canonicalizer> {
   Canonicalizer() = default;
   Canonicalizer(const GreedyRewriteConfig &config,
                 ArrayRef<std::string> disabledPatterns,
                 ArrayRef<std::string> enabledPatterns) {
     this->topDownProcessingEnabled = config.useTopDownTraversal;
     this->enableRegionSimplification = config.enableRegionSimplification;
     this->maxIterations = config.maxIterations;
+    this->maxNumRewrites = config.maxNumRewrites;
     this->disabledPatterns = disabledPatterns;
     this->enabledPatterns = enabledPatterns;
   }
 
   /// Initialize the canonicalizer by building the set of patterns used during
   /// execution.
   LogicalResult initialize(MLIRContext *context) override {
     RewritePatternSet owningPatterns(context);
     for (auto *dialect : context->getLoadedDialects())
       dialect->getCanonicalizationPatterns(owningPatterns);
     for (RegisteredOperationName op : context->getRegisteredOperations())
       op.getCanonicalizationPatterns(owningPatterns, context);
 
     patterns = FrozenRewritePatternSet(std::move(owningPatterns),
                                        disabledPatterns, enabledPatterns);
     return success();
   }
   void runOnOperation() override {
     GreedyRewriteConfig config;
     config.useTopDownTraversal = topDownProcessingEnabled;
     config.enableRegionSimplification = enableRegionSimplification;
     config.maxIterations = maxIterations;
+    config.maxNumRewrites = maxNumRewrites;
     (void)applyPatternsAndFoldGreedily(getOperation(), patterns, config);
   }
 
   FrozenRewritePatternSet patterns;
 };
 } // namespace
 
 /// Create a Canonicalizer pass.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

if (!config.useTopDownTraversal) {
       for (size_t i = 0, e = worklist.size(); i != e; ++i)
         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()) {
       auto *op = popFromWorklist();
 
       // Nulls get added to the worklist when operations are removed, ignore
       // them.
       if (op == nullptr)
         continue;
 
#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
-      changed |= succeeded(matchResult);
+      if (succeeded(matchResult)) {
+        changed = true;
+        if (numRewrites++ >= config.maxNumRewrites &&
+            config.maxNumRewrites != GreedyRewriteConfig::kNoLimit)
+          break;
+      }
     }
 
     // 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::kNoIterationLimit));
+  } while (changed && (iteration++ < config.maxIterations ||
+                       config.maxIterations == GreedyRewriteConfig::kNoLimit));
 
   // 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))
if (succeeded(folder.tryToFold(op, /*processGeneratedConstants=*/nullptr,
       }
     }
 
     // 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 ((erased = opErasedViaPatternRewrites))
       return success();
-  } while (changed &&
-           (++iterations < maxIterations ||
-            maxIterations == GreedyRewriteConfig::kNoIterationLimit));
+  } while (changed && (++iterations < maxIterations ||
+                       maxIterations == GreedyRewriteConfig::kNoLimit));
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return failure(changed);
 }
 
 //===----------------------------------------------------------------------===//
 // MultiOpPatternRewriteDriver
 //===----------------------------------------------------------------------===//