[MLIR] Introduce a new rewrite driver to simplify supplied list of ops

[bondhugula]

Summary

Introduce a new rewrite driver (MultiOpPatternRewriteDriver) to rewrite
a supplied list of ops and other ops. Provide a knob to restrict
rewrites strictly to those ops or also to affected ops (but still not to
completely related ops).

Discourse post
https://llvm.discourse.group/t/introducing-a-new-rewrite-driver-to-work-on-a-specified-list-of-ops/3899

This rewrite driver is commonly needed to run any simplification and
cleanup at the end of a transforms pass or transforms utility in a way
that only simplifies relevant IR. This makes it easy to write test cases
while not performing unrelated whole IR simplification that may
invalidate other state at the caller.

The introduced utility provides more freedom to developers of transforms
and transform utilities to perform focussed and local simplification. In
several cases, it provides greater efficiency as well as more
simplification when compared to repeatedly calling
applyOpPatternsAndFold; in other cases, it avoids the need to
undesirably call applyPatternsAndFoldGreedily to do unrelated
simplification in a FuncOp.

Update a few transformations that were earlier using
applyOpPatternsAndFold (SimplifyAffineStructures,
affineDataCopyGenerate, a linalg transform).

TODO

• Doc will be updated close to completion of review of this revision.
• OpPatternRewriteDriver can be removed as it's a special case of MultiOpPatternRewriteDriver, i.e., both can be merged.

[bondhugula]

Minor tweaks. NFC.

[bondhugula]

NFC. Improve comment.

[bondhugula]

Improve comment. NFC.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

399–401

Is this exact? Wouldn't the simplification of the operands added to the wordlist recursively triggers adding their own operands after simplification?

436

Nit: this is misleading as op is a pointer, Operation * ?

444

Can you tell me a bit more about when is this happening?

452

It should be removed of workOps here I think, otherwise another pattern can create ops and reuse the pointer that is freed now and we'd think it is part of workOps.
(in general keeping dangling pointers in set can be annoying to debug because it leads to bugs that are not deterministic)

467

Nit: I'm unsure that the autos here improves readability

468

The API supports contains now I believe

475

Can you match the name of the API? I had to look to see that it is processGeneratedConstants which is more explicit about what will happen here: it isn't just any op that we expect here but only constants (which makes sense, we don't expect the folder to create any other kind of ops).

489

Again: shouldn't it be removed from the workOps?

Also does inPlaceUpdate really implies that the op is removed? I would think that folding some attributes in-place would also result in this code path.

491

Nit: spurious braces

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

489

@mehdi_amini It's being removed from workOps in notifyOperationRemoved :-) It has to be done there otherwise erasure happening in canonicalization patterns would miss updating workOps. Likewise for your comment above.

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

399–401

Yes, this is exact. When strict is on, we prevent "outside" ops from getting into the worklist whether operand defs or result users.

436

Thanks. I missed it during copy and paste.

444

This code is common to all the rewrite drivers. When eraseOp is called, the ensuing hooks mark the item in the worklist null.

452

It's being done in notifyOperationRemoved cleanly! :-)

467

Thanks. Removed auto. (This part was was copied over from the other drivers.)

468

Thanks.

475

Sure, thanks.

491

I let the brace because there is a comment. Per LLVM style, this is recommended. (https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements)

[bondhugula]

Address review comments.

[bondhugula]

count(..) > 0 -> contains

[herhut]

Cool, thanks for adding this. Please also wait for @mehdi_amini to agree before landing.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

81

nit: as well as

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

407

I find workOps not very descriptive. It is a working set, so I kinda get that name workOps but when reading the code further below I find it confusing with also a worklist.

Maybe strictModeSet or strictModeFilter would be more descriptive?

413

nit: as well as

422

nit: remove once

423

nit: is to keep

[mehdi_amini]

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

90

(Nit: do you know that SmallVector size is optional now?)

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

407

Good point on the naming, I would throw in the pot the suggestion for strictModeFilteredOps :)

452

Ah, sweet! :)

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

555–556

Comment can be updated I think.

[bondhugula]

Address review comments from @herhut and @mehdi_amini.

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

90

Yes, but a force of habit here! Dropped it.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

407

Makes sense. I was also thinking of strictModeFilteredOps after @herhut suggestion.

[bondhugula]

Thanks for the comments @mehdi_amini and @herhut. I'll have a follow-up to add the documentation for this.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

 /// by selecting the highest benefits patterns in a greedy manner. Returns
 /// success if no more patterns can be matched. `erased` is set to true if `op`
 /// was folded away or erased as a result of becoming dead. Note: This does not
 /// apply any patterns recursively to the regions of `op`.
 LogicalResult applyOpPatternsAndFold(Operation *op,
                                      const FrozenRewritePatternSet &patterns,
                                      bool *erased = nullptr);
 
+/// Applies the specified rewrite patterns on `ops` while also trying to fold
+/// these ops as well as any other ops that were in turn created due to such
+/// rewrites. Furthermore, any pre-existing ops in the IR outside of `ops`
+/// remain completely unmodified if `strict` is set to true. If `strict` is
+/// false, other operations that use results of rewritten ops or supply operands
+/// to such ops are in turn simplified; any other ops still remain unmodified
+/// (i.e., regardless of `strict`). Note that ops in `ops` could be erased as a
+/// result of folding, becoming dead, or via pattern rewrites. If more far
+/// reaching simplification is desired, applyPatternsAndFoldGreedily should be
+/// used. Returns true if at all any IR was rewritten.
+bool applyOpPatternsAndFold(ArrayRef<Operation *> ops,
+                            const FrozenRewritePatternSet &patterns,
+                            bool strict);
+
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

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

void AffineDataCopyGeneration::runOnFunction() {
 
   // Promoting single iteration loops could lead to simplification of
   // contained load's/store's, and the latter could anyway also be
   // canonicalized.
   RewritePatternSet patterns(&getContext());
   AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
   AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-  for (Operation *op : copyOps)
-    (void)applyOpPatternsAndFold(op, frozenPatterns);
+  (void)applyOpPatternsAndFold(copyOps, frozenPatterns, /*strict=*/true);
 }

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

 mlir::createSimplifyAffineStructuresPass() {
   return std::make_unique<SimplifyAffineStructures>();
 }
 
 void SimplifyAffineStructures::runOnFunction() {
   auto func = getFunction();
   simplifiedAttributes.clear();
   RewritePatternSet patterns(func.getContext());
+  AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
   AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
   AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
-  AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+
+  // The simplification of affine attributes will likely simplify the op. Try to
+  // fold/apply canonicalization patterns when we have affine dialect ops.
+  SmallVector<Operation *> opsToSimplify;
   func.walk([&](Operation *op) {
     for (auto attr : op->getAttrs()) {
       if (auto mapAttr = attr.second.dyn_cast<AffineMapAttr>())
         simplifyAndUpdateAttribute(op, attr.first, mapAttr);
       else if (auto setAttr = attr.second.dyn_cast<IntegerSetAttr>())
         simplifyAndUpdateAttribute(op, attr.first, setAttr);
     }
 
-    // The simplification of the attribute will likely simplify the op. Try to
-    // fold / apply canonicalization patterns when we have affine dialect ops.
     if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
-      (void)applyOpPatternsAndFold(op, frozenPatterns);
+      opsToSimplify.push_back(op);
   });
+  (void)applyOpPatternsAndFold(opsToSimplify, frozenPatterns, /*strict=*/true);
 }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

public:
   }
 
   // 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); }
 
+  // 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);
+    }
+  }
+
   // 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);
     });
   }
 
   // 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);
   }
 
-private:
-  // 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 *defInst = operand.getDefiningOp())
-        addToWorklist(defInst);
-    }
-  }
-
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// The worklist for this transformation keeps track of the operations that
   /// need to be revisited, plus their index in the worklist.  This allows us to
   /// 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:
   /// Configuration information for how to simplify.
   GreedyRewriteConfig config;
 };
 } // end anonymous namespace
 
 /// Performs the rewrites while folding and erasing any dead ops. Returns true
 /// if the rewrite converges in `maxIterations`.
 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {

 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();
   }
 
-  /// Performs the rewrites and folding only on `op`. The simplification
-  /// converges if the op is erased as a result of being folded, replaced, or
-  /// dead, or no more changes happen in an iteration. Returns success if the
-  /// rewrite converges in `maxIterations`. `erased` is set to true if `op` gets
-  /// erased.
   LogicalResult simplifyLocally(Operation *op, int maxIterations, 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:
   OperationFolder folder;
 
   /// Set to true if the operation has been erased via pattern rewrites.
   bool opErasedViaPatternRewrites = false;
 };
 
 } // anonymous 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`
+/// gets erased.
 LogicalResult OpPatternRewriteDriver::simplifyLocally(Operation *op,
                                                       int maxIterations,
                                                       bool &erased) {
   bool changed = false;
   erased = false;
   opErasedViaPatternRewrites = false;
-  int i = 0;
+  int iterations = 0;
   // Iterate until convergence or until maxIterations. Deletion of the op as
   // a result of being dead or folded is convergence.
   do {
     changed = false;
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
       op->erase();
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 && ++i < maxIterations);
+  } while (changed && ++iterations < maxIterations);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return failure(changed);
 }
 
+//===----------------------------------------------------------------------===//
+// MultiOpPatternRewriteDriver
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// This is a specialized GreedyPatternRewriteDriver to apply patterns and
+/// perform folding for a supplied set of ops. It repeatedly simplifies while
+/// restricting the rewrites to only the provided set of ops or optionally
+/// to those directly affected by it (result users or operand providers).
+class MultiOpPatternRewriteDriver : public GreedyPatternRewriteDriver {
+public:
+  explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
+                                       const FrozenRewritePatternSet &patterns,
+                                       bool strict)
+      : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
+        strictMode(strict) {}
+
+  bool simplifyLocally(ArrayRef<Operation *> op);
+
+private:
+  // 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. If `strict` is set to true, only ops in
+  // `strictModeFilteredOps` are considered.
+  template <typename Operands>
+  void addOperandsToWorklist(Operands &&operands) {
+    for (Value operand : operands) {
+      if (auto *defOp = operand.getDefiningOp()) {
+        if (!strictMode || strictModeFilteredOps.contains(defOp))
+          addToWorklist(defOp);
+      }
+    }
+  }
+
+  void notifyOperationRemoved(Operation *op) override {
+    GreedyPatternRewriteDriver::notifyOperationRemoved(op);
+    if (strictMode)
+      strictModeFilteredOps.erase(op);
+  }
+
+  /// If `strictMode` is true, any pre-existing ops outside of
+  /// `strictModeFilteredOps` remain completely untouched by the rewrite driver.
+  /// If `strictMode` is false, operations that use results of (or supply
+  /// operands to) any rewritten ops stemming from the simplification of the
+  /// provided ops are in turn simplified; any other ops still remain untouched
+  /// (i.e., regardless of `strictMode`).
+  bool strictMode = false;
+
+  /// The list of ops we are restricting our rewrites to if `strictMode` is on.
+  /// These include the supplied set of ops as well as new ops created while
+  /// rewriting those ops. This set is not maintained when strictMode is off.
+  llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
+};
+
+} // end anonymous namespace
+
+/// Performs the specified rewrites on `ops` while also trying to fold these ops
+/// as well as any other ops that were in turn created due to these rewrite
+/// patterns. Any pre-existing ops outside of `ops` remain completely
+/// unmodified if `strictMode` is true. If `strictMode` is false, other
+/// operations that use results of rewritten ops or supply operands to such ops
+/// are in turn simplified; any other ops still remain unmodified (i.e.,
+/// regardless of `strictMode`). Note that ops in `ops` could be erased as a
+/// result of folding, becoming dead, or via pattern rewrites. Returns true if
+/// at all any changes happened.
+// Unlike `OpPatternRewriteDriver::simplifyLocally` which works on a single op
+// or GreedyPatternRewriteDriver::simplify, this method just iterates until
+// the worklist is empty. As our objective is to keep simplification "local",
+// there is no strong rationale to re-add all operations into the worklist and
+// rerun until an iteration changes nothing. If more widereaching simplification
+// is desired, GreedyPatternRewriteDriver should be used.
+bool MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops) {
+  if (strictMode) {
+    strictModeFilteredOps.clear();
+    strictModeFilteredOps.insert(ops.begin(), ops.end());
+  }
+
+  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()) {
+    Operation *op = popFromWorklist();
+
+    // Nulls get added to the worklist when operations are removed, ignore
+    // them.
+    if (op == nullptr)
+      continue;
+
+    // If the operation is trivially dead - remove it.
+    if (isOpTriviallyDead(op)) {
+      notifyOperationRemoved(op);
+      op->erase();
+      changed = true;
+      continue;
+    }
+
+    // Collects all the operands and result uses of the given `op` into work
+    // list. Also remove `op` and nested ops from worklist.
+    originalOperands.assign(op->operand_begin(), op->operand_end());
+    auto preReplaceAction = [&](Operation *op) {
+      // Add the operands to the worklist for visitation.
+      addOperandsToWorklist(originalOperands);
+
+      // Add all the users of the result to the worklist so we make sure
+      // to revisit them.
+      for (Value result : op->getResults())
+        for (Operation *userOp : result.getUsers()) {
+          if (!strictMode || strictModeFilteredOps.contains(userOp))
+            addToWorklist(userOp);
+        }
+      notifyOperationRemoved(op);
+    };
+
+    // Add the given operation generated by the folder to the worklist.
+    auto processGeneratedConstants = [this](Operation *op) {
+      // Newly created ops are also simplified -- these are also "local".
+      addToWorklist(op);
+      // When strict mode is off, we don't need to maintain
+      // strictModeFilteredOps.
+      if (strictMode)
+        strictModeFilteredOps.insert(op);
+    };
+
+    // Try to fold this op.
+    bool inPlaceUpdate;
+    if (succeeded(folder.tryToFold(op, processGeneratedConstants,
+                                   preReplaceAction, &inPlaceUpdate))) {
+      changed = true;
+      if (!inPlaceUpdate) {
+        // 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));
+  }
+
+  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);
   if (erased)
     *erased = opErased;
   LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << config.maxIterations << " times";
   });
   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);
+}

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

func @simplify_set(%a : index, %b : index) {
   return
 }
 
 // -----
 
 // CHECK-DAG: -> (s0 * 2 + 1)
 
 // Test "op local" simplification on affine.apply. DCE on addi will not happen.
-func @affine.apply(%N : index) {
+func @affine.apply(%N : index) -> index {
   %v = affine.apply affine_map<(d0, d1) -> (d0 + d1 + 1)>(%N, %N)
-  addi %v, %v : index
+  %res = addi %v, %v : index
   // CHECK: affine.apply #map{{.*}}()[%arg0]
   // CHECK-NEXT: addi
-  return
+  return %res: index
 }
 
 // -----
 
 // CHECK-LABEL: func @simplify_zero_dim_map
 func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {
   %out = affine.load %in[] : memref<f32>
   return %out : f32

mlir/test/Dialect/Linalg/fold-affine-min-scf.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-affine-min-scf-canonicalization-patterns | FileCheck %s
 
 // CHECK-LABEL: scf_for
 func @scf_for(%A : memref<i64>, %step : index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c7 = constant 7 : index
   %c4 = constant 4 : index
   %c16 = constant 16 : index
   %c1024 = constant 1024 : index
 
+  // CHECK:      %[[C2:.*]] = constant 2 : i64
   //      CHECK: scf.for
-  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
-  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
-  // CHECK-NEXT:   memref.store %[[C2I64]], %{{.*}}[] : memref<i64>
+  // CHECK-NEXT:   memref.store %[[C2]], %{{.*}}[] : memref<i64>
   scf.for %i = %c0 to %c4 step %c2 {
     %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
 
   //      CHECK: scf.for
-  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
-  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
-  // CHECK-NEXT:   memref.store %[[C2I64]], %{{.*}}[] : memref<i64>
+  // CHECK-NEXT:   memref.store %[[C2]], %{{.*}}[] : memref<i64>
   scf.for %i = %c1 to %c7 step %c2 {
     %1 = affine.min affine_map<(d0)[s0] -> (s0 - d0, 2)> (%i)[%c7]
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
 
   // This should not canonicalize because: 4 - %i may take the value 1 < 2.
   //     CHECK:   scf.for

 // CHECK-LABEL: scf_parallel
 func @scf_parallel(%A : memref<i64>, %step : index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c7 = constant 7 : index
   %c4 = constant 4 : index
 
+  // CHECK:   %[[C2:.*]] = constant 2 : i64
   // CHECK: scf.parallel
-  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
-  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
-  // CHECK-NEXT:   memref.store %[[C2I64]], %{{.*}}[] : memref<i64>
+  // CHECK-NEXT:   memref.store %[[C2]], %{{.*}}[] : memref<i64>
   scf.parallel (%i) = (%c0) to (%c4) step (%c2) {
     %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
 
   // CHECK: scf.parallel
-  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
-  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
-  // CHECK-NEXT:   memref.store %[[C2I64]], %{{.*}}[] : memref<i64>
+  // CHECK-NEXT:   memref.store %[[C2]], %{{.*}}[] : memref<i64>
   scf.parallel (%i) = (%c1) to (%c7) step (%c2) {
     %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
 
   // This example should simplify but affine_map is currently missing
   // semi-affine canonicalizations.

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

void TestAffineDataCopy::runOnFunction() {
   for (Operation *op : copyOps) {
     patterns.clear();
     if (isa<AffineLoadOp>(op)) {
       AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
     } else {
       assert(isa<AffineStoreOp>(op) && "expected affine store op");
       AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
     }
-    (void)applyOpPatternsAndFold(op, std::move(patterns));
   }
+  (void)applyOpPatternsAndFold(copyOps, std::move(patterns), /*strict=*/true);
 }
 
 namespace mlir {
 void registerTestAffineDataCopyPass() {
   PassRegistration<TestAffineDataCopy>();
 }
 } // namespace mlir

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

static void applyExtractSliceOfPadTensorSwapPattern(FuncOp funcOp) {
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
   RewritePatternSet foldPattern(funcOp.getContext());
   foldPattern.add<AffineMinSCFCanonicalizationPattern>(funcOp.getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(foldPattern));
 
-  // Explicitly walk and apply the pattern locally to avoid more general folding
+  // Explicitly apply the pattern on affected ops to avoid more general folding
   // on the rest of the IR.
-  funcOp.walk([&frozenPatterns](AffineMinOp minOp) {
-    (void)applyOpPatternsAndFold(minOp, frozenPatterns);
-  });
+  SmallVector<Operation *, 4> minOps;
+  funcOp.walk([&](AffineMinOp minOp) { minOps.push_back(minOp); });
+  (void)applyOpPatternsAndFold(minOps, frozenPatterns, /*strict=*/false);
 }
 
 // For now, just assume it is the zero of type.
 // In the future, it should be the zero of type + op.
 static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
   auto t = getElementTypeOrSelf(op.get());
   return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t));
 }