[mlir] add a fluent API to GreedyRewriterConfig

This is similar to other configuration objects used across MLIR.

Author: ftynse

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

@@ -50,24 +50,41 @@ class GreedyRewriteConfig {
   ///
   /// Note: Only applicable when simplifying entire regions.
   bool useTopDownTraversal = false;
+  GreedyRewriteConfig &setUseTopDownTraversal(bool use = true) {
+    useTopDownTraversal = use;
+    return *this;
+  }
 
   /// Perform control flow optimizations to the region tree after applying all
   /// patterns.
   ///
   /// Note: Only applicable when simplifying entire regions.
   GreedySimplifyRegionLevel enableRegionSimplification =
       GreedySimplifyRegionLevel::Aggressive;
+  GreedyRewriteConfig &
+  setEnableRegionSimplification(GreedySimplifyRegionLevel level) {
+    enableRegionSimplification = level;
+    return *this;
+  }
 
   /// This specifies the maximum number of times the rewriter will iterate
   /// between applying patterns and simplifying regions. Use `kNoLimit` to
   /// disable this iteration limit.
   ///
   /// Note: Only applicable when simplifying entire regions.
   int64_t maxIterations = 10;
+  GreedyRewriteConfig &setMaxIterations(int64_t iterations) {
+    maxIterations = iterations;
+    return *this;
+  }
 
   /// This specifies the maximum number of rewrites within an iteration. Use
   /// `kNoLimit` to disable this limit.
   int64_t maxNumRewrites = kNoLimit;
+  GreedyRewriteConfig &setMaxNumRewrites(int64_t limit) {
+    maxNumRewrites = limit;
+    return *this;
+  }
 
   static constexpr int64_t kNoLimit = -1;
 
@@ -76,6 +93,10 @@ class GreedyRewriteConfig {
   /// (or the specified region, depending on which greedy rewrite entry point
   /// is used) is used as a scope.
   Region *scope = nullptr;
+  GreedyRewriteConfig &setScope(Region *scope) {
+    this->scope = scope;
+    return *this;
+  }
 
   /// Strict mode can restrict the ops that are added to the worklist during
   /// the rewrite.
@@ -88,16 +109,32 @@ class GreedyRewriteConfig {
   ///   were on the worklist at the very beginning) enqueued. All other ops are
   ///   excluded.
   GreedyRewriteStrictness strictMode = GreedyRewriteStrictness::AnyOp;
+  GreedyRewriteConfig &setStrictMode(GreedyRewriteStrictness mode) {
+    strictMode = mode;
+    return *this;
+  }
 
   /// An optional listener that should be notified about IR modifications.
   RewriterBase::Listener *listener = nullptr;
+  GreedyRewriteConfig &setListener(RewriterBase::Listener *listener) {
+    this->listener = listener;
+    return *this;
+  }
 
   /// Whether this should fold while greedily rewriting.
   bool fold = true;
+  GreedyRewriteConfig &setFold(bool enable = true) {
+    fold = enable;
+    return *this;
+  }
 
   /// If set to "true", constants are CSE'd (even across multiple regions that
   /// are in a parent-ancestor relationship).
   bool cseConstants = true;
+  GreedyRewriteConfig &setCSEConstants(bool enable = true) {
+    cseConstants = enable;
+    return *this;
+  }
 };
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

@@ -127,12 +127,13 @@ SimplifyBoundedAffineOpsOp::apply(transform::TransformRewriter &rewriter,
   patterns.insert<SimplifyAffineMinMaxOp<AffineMinOp>,
                   SimplifyAffineMinMaxOp<AffineMaxOp>>(getContext(), cstr);
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-  GreedyRewriteConfig config;
-  config.listener =
-      static_cast<RewriterBase::Listener *>(rewriter.getListener());
-  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
   // Apply the simplification pattern to a fixpoint.
-  if (failed(applyOpPatternsGreedily(targets, frozenPatterns, config))) {
+  if (failed(applyOpPatternsGreedily(
+          targets, frozenPatterns,
+          GreedyRewriteConfig()
+              .setListener(
+                  static_cast<RewriterBase::Listener *>(rewriter.getListener()))
+              .setStrictMode(GreedyRewriteStrictness::ExistingAndNewOps)))) {
     auto diag = emitDefiniteFailure()
                 << "affine.min/max simplification did not converge";
     return diag;

mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp

@@ -237,7 +237,8 @@ void AffineDataCopyGeneration::runOnOperation() {
   AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
   AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-  GreedyRewriteConfig config;
-  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
-  (void)applyOpPatternsGreedily(copyOps, frozenPatterns, config);
+  (void)applyOpPatternsGreedily(
+      copyOps, frozenPatterns,
+      GreedyRewriteConfig().setStrictMode(
+          GreedyRewriteStrictness::ExistingAndNewOps));
 }

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

@@ -109,7 +109,8 @@ void SimplifyAffineStructures::runOnOperation() {
     if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
       opsToSimplify.push_back(op);
   });
-  GreedyRewriteConfig config;
-  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
-  (void)applyOpPatternsGreedily(opsToSimplify, frozenPatterns, config);
+  (void)applyOpPatternsGreedily(
+      opsToSimplify, frozenPatterns,
+      GreedyRewriteConfig().setStrictMode(
+          GreedyRewriteStrictness::ExistingAndNewOps));
 }

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

@@ -315,11 +315,12 @@ LogicalResult mlir::affine::affineForOpBodySkew(AffineForOp forOp,
         // Simplify/canonicalize the affine.for.
         RewritePatternSet patterns(res.getContext());
         AffineForOp::getCanonicalizationPatterns(patterns, res.getContext());
-        GreedyRewriteConfig config;
-        config.strictMode = GreedyRewriteStrictness::ExistingOps;
         bool erased;
-        (void)applyOpPatternsGreedily(res.getOperation(), std::move(patterns),
-                                      config, /*changed=*/nullptr, &erased);
+        (void)applyOpPatternsGreedily(
+            res.getOperation(), std::move(patterns),
+            GreedyRewriteConfig().setStrictMode(
+                GreedyRewriteStrictness::ExistingAndNewOps),
+            /*changed=*/nullptr, &erased);
         if (!erased && !prologue)
           prologue = res;
         if (!erased)

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

@@ -494,10 +494,9 @@ struct IntRangeOptimizationsPass final
     RewritePatternSet patterns(ctx);
     populateIntRangeOptimizationsPatterns(patterns, solver);
 
-    GreedyRewriteConfig config;
-    config.listener = &listener;
-
-    if (failed(applyPatternsGreedily(op, std::move(patterns), config)))
+    if (failed(applyPatternsGreedily(
+            op, std::move(patterns),
+            GreedyRewriteConfig().setListener(&listener))))
       signalPassFailure();
   }
 };
@@ -520,13 +519,12 @@ struct IntRangeNarrowingPass final
     RewritePatternSet patterns(ctx);
     populateIntRangeNarrowingPatterns(patterns, solver, bitwidthsSupported);
 
-    GreedyRewriteConfig config;
     // We specifically need bottom-up traversal as cmpi pattern needs range
     // data, attached to its original argument values.
-    config.useTopDownTraversal = false;
-    config.listener = &listener;
-
-    if (failed(applyPatternsGreedily(op, std::move(patterns), config)))
+    if (failed(applyPatternsGreedily(
+            op, std::move(patterns),
+            GreedyRewriteConfig().setUseTopDownTraversal(false).setListener(
+                &listener))))
       signalPassFailure();
   }
 };

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp

@@ -463,15 +463,15 @@ struct BufferDeallocationSimplificationPass
                  SplitDeallocWhenNotAliasingAnyOther,
                  RetainedMemrefAliasingAlwaysDeallocatedMemref>(&getContext(),
                                                                 analysis);
+
+    populateDeallocOpCanonicalizationPatterns(patterns, &getContext());
     // We don't want that the block structure changes invalidating the
-    // `BufferOriginAnalysis` so we apply the rewrites witha `Normal` level of
+    // `BufferOriginAnalysis` so we apply the rewrites with `Normal` level of
     // region simplification
-    GreedyRewriteConfig config;
-    config.enableRegionSimplification = GreedySimplifyRegionLevel::Normal;
-    populateDeallocOpCanonicalizationPatterns(patterns, &getContext());
-
-    if (failed(
-            applyPatternsGreedily(getOperation(), std::move(patterns), config)))
+    if (failed(applyPatternsGreedily(
+            getOperation(), std::move(patterns),
+            GreedyRewriteConfig().setEnableRegionSimplification(
+                GreedySimplifyRegionLevel::Normal))))
       signalPassFailure();
   }
 };

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

@@ -3587,9 +3587,9 @@ transform::VectorizeChildrenAndApplyPatternsOp::applyToOne(
   vector::populateVectorStepLoweringPatterns(patterns);
 
   TrackingListener listener(state, *this);
-  GreedyRewriteConfig config;
-  config.listener = &listener;
-  if (failed(applyPatternsGreedily(target, std::move(patterns), config)))
+  if (failed(
+          applyPatternsGreedily(target, std::move(patterns),
+                                GreedyRewriteConfig().setListener(&listener))))
     return emitDefaultDefiniteFailure(target);
 
   results.push_back(target);

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

@@ -2327,10 +2327,9 @@ struct LinalgElementwiseOpFusionPass
     // Add constant folding patterns.
     populateConstantFoldLinalgOperations(patterns, defaultControlFn);
 
-    // Use TopDownTraversal for compile time reasons
-    GreedyRewriteConfig grc;
-    grc.useTopDownTraversal = true;
-    (void)applyPatternsGreedily(op, std::move(patterns), grc);
+    // Use TopDownTraversal for compile time reasons.
+    (void)applyPatternsGreedily(op, std::move(patterns),
+                                GreedyRewriteConfig().setUseTopDownTraversal());
   }
 };
 

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -1438,10 +1438,10 @@ SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
   if (!patterns)
     return success();
 
-  GreedyRewriteConfig config;
-  config.listener = this;
-  config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
-  return applyOpPatternsGreedily(ops, patterns.value(), config);
+  return applyOpPatternsGreedily(
+      ops, patterns.value(),
+      GreedyRewriteConfig().setListener(this).setStrictMode(
+          GreedyRewriteStrictness::ExistingAndNewOps));
 }
 
 void SliceTrackingListener::notifyOperationInserted(