[MLIR] Introduce applyOpPatternsAndFold for op local rewrites

[bondhugula]

Summary

Introduce mlir::applyOpPatternsAndFold which applies patterns as well as
any folding only on a specified op (in contrast to
applyPatternsAndFoldGreedily which applies patterns only on the regions
of an op isolated from above). The caller is made aware of the op being
folded away or erased.

Depends on D77485.

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

251

This is wrong, you don't handle when the op is erased by a pattern. This really feels like should be its own driver, it doesn't really share anything.

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

251

Hmm... if the op was erased in a pattern by just erase(), there is no way to know if it was erased? If the op is erased via eraseOp or replaceOp, checking for null at the op's position via worklistMap will address that (we'll have to just keep the op in the worklist before matchAndRewrite). The other issue is that the erased op's operands will also get into the worklist and we don't want them to, but this can be addressed by clearing out the worklist. This could use a separate driver as you say, but it doesn't immediately address the first issue - how do we detect erasure if done via op.erase()?

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

251

how do we detect erasure if done via op.erase()?

We don't have to. It isn't valid to do IR mutations outside of the PatternRewriter.

I thought this was documented somewhere, but https://mlir.llvm.org/docs/QuickstartRewrites/ seems to be extremely out of date. I have an in-progress update to https://mlir.llvm.org/docs/Canonicalization/ I'll add that in as part of it.

[bondhugula]

Implementing this functionality has hit a dead-end; there's no easy way out.

[bondhugula]

Update to use worklist free driver. This isn't still well tested - so could
wait till it could be used at a few places.

[bondhugula]

Fix some gratuitous changes. This isn't still well-tested and could wait for
some of the pending revisions to get in.

[bondhugula]

rebase and permute commit stack

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

251

I don't recall that requirement ever being documented, although at one point some instances of op.erase() were replaced by rewriter.eraseOp(op). The document itself hasn't changed core content wise since Mar 2019. I couldn't find it in your recent update to the canonicalization doc. Is another one pending?

If that's the constraint on how one can erase, then I think I have it implemented correctly now. It needs better testing though.

[bondhugula]

Rebase.

[rriddle]

Can you add test coverage?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

251

Yes, there is another one pending.

254

Can you use a comment block to split these?

257

"bottom up" doesn't really make sense here.

266

Please use /// for the top-level comments.

330

while (!erased && changed && ...)

?

[bondhugula]

Address review comments.

[bondhugula]

Thanks for the review.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

254

Not sure I understood it right - but PTAL at what I did.

330

That won't be accurate since although you'll correctly exit the iteration and return, you will be returning !changed, which will be false instead of being true (i.e., the driver has converged).

[rriddle]

Can you add test coverage?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

330

Ah yeah, duh. Thanks.

[bondhugula]

In D77487#1975842, @rriddle wrote:

Can you add test coverage?

Sure; I'm trying to find a good set of places to call this from. Thanks for the review.

[bondhugula]

Add testing.

[bondhugula]

Drop unrelated changes.

[bondhugula]

Ping @rriddle - do you have more comments here? I added test cases and the output parameter on the method has been made optional. (A test case that exercises the op being erased is in the child revision. )

[nicolasvasilache]

Nice!

I expect this to be very useful for more easily expressing fused passes.

[bondhugula]

Comment updates - reset erased to false in driver's simplifyLocally.

[rriddle]

I have a growing concern that rebuilding the pattern list is going to become extremely costly for compile time. I don't have a concrete suggestion at the moment, but we should avoid rebuilding if possible.

That being said, I think this is a great start to doing local canonicalization. Thank you for pushing this!

mlir/include/mlir/IR/PatternMatch.h

464

I'm of the opinion that we should remove the return value from these methods. It is not used anywhere, and it isn't clear what the user would use it for in any of the existing usages.

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

233

This seems extremely costly. Can we create these once and just apply them if necessary?

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

98

Same here. Can we avoid rebuilding a pattern set for each operation?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

275

Please use /// here an below.

mlir/lib/Transforms/Utils/LoopUtils.cpp

318

nit: Use // here.

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

100

nit: isa<>() || isa<>()

[bondhugula]

NFC - clean up erasure logic

[bondhugula]

clean up and fix erasure further

[bondhugula]

Address review comments

mlir/include/mlir/IR/PatternMatch.h

464

The child revision has the use case where it's essential - let's discuss there if there are alternatives.

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

98

Sure - done.

mlir/include/mlir/IR/PatternMatch.h

 /// Note: These methods also perform folding and simple dead-code elimination
 ///       before attempting to match any of the provided patterns.
 ///
 bool applyPatternsAndFoldGreedily(Operation *op,
                                   const OwningRewritePatternList &patterns);
 /// Rewrite the given regions, which must be isolated from above.
 bool applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
                                   const OwningRewritePatternList &patterns);
+
+/// Applies the specified patterns on `op` alone while also trying to fold it,
+/// by selecting the highest benefits patterns in a greedy manner. Returns true
+/// 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`.
+bool applyOpPatternsAndFold(Operation *op,
+                            const OwningRewritePatternList &patterns,
+                            bool *erased = nullptr);
 } // end namespace mlir
 
 #endif // MLIR_PATTERN_MATCH_H

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

void AffineDataCopyGeneration::runOnFunction() {
   DenseSet<Operation *> copyNests;
 
   // Clear recorded copy nests.
   copyNests.clear();
 
   for (auto &block : f)
     runOnBlock(&block, copyNests);
 
-  // Promote any single iteration loops in the copy nests.
+  // Promote any single iteration loops in the copy nests and collect
+  // load/stores to simplify.
+  SmallVector<Operation *, 4> copyOps;
   for (auto nest : copyNests)
-    nest->walk([](AffineForOp forOp) { promoteIfSingleIteration(forOp); });
+    // With a post order walk, the erasure of loops does not affect
+    // continuation of the walk or the collection of load/store ops.
+    nest->walk([&](Operation *op) {
+      if (auto forOp = dyn_cast<AffineForOp>(op))
+        promoteIfSingleIteration(forOp);
+      else if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op))
+        copyOps.push_back(op);
+    });
 
   // Promoting single iteration loops could lead to simplification of
-  // load's/store's. We will run canonicalization patterns on load/stores.
-  // TODO: this whole function load/store canonicalization should be replaced by
-  // canonicalization that is limited to only the load/store ops
-  // introduced/touched by this pass (those inside 'copyNests'). This would be
-  // possible once the necessary support is available in the pattern rewriter.
-  if (!copyNests.empty()) {
-    OwningRewritePatternList patterns;
-    AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext());
-    AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext());
-    applyPatternsAndFoldGreedily(f, std::move(patterns));
-  }
+  // contained load's/store's, and the latter could anyway also be
+  // canonicalized.
+  for (auto op : copyOps)
+    applyOpPatternsAndFold(op, std::move(patterns));
 }

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

 //===- SimplifyAffineStructures.cpp ---------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This file implements a pass to simplify affine structures.
+// This file implements a pass to simplify affine structures in operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/Utils.h"
 
 #define DEBUG_TYPE "simplify-affine-structure"
 
 using namespace mlir;
 
 namespace {
 

 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createSimplifyAffineStructuresPass() {
   return std::make_unique<SimplifyAffineStructures>();
 }
 
 void SimplifyAffineStructures::runOnFunction() {
   auto func = getFunction();
   simplifiedAttributes.clear();
-  func.walk([&](Operation *opInst) {
-    for (auto attr : opInst->getAttrs()) {
+  OwningRewritePatternList patterns;
+  AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
+  AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
+  AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
+  func.walk([&](Operation *op) {
+    for (auto attr : op->getAttrs()) {
       if (auto mapAttr = attr.second.dyn_cast<AffineMapAttr>())
-        simplifyAndUpdateAttribute(opInst, attr.first, mapAttr);
+        simplifyAndUpdateAttribute(op, attr.first, mapAttr);
       else if (auto setAttr = attr.second.dyn_cast<IntegerSetAttr>())
-        simplifyAndUpdateAttribute(opInst, attr.first, setAttr);
+        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>(op) || isa<AffineIfOp>(op) || isa<AffineApplyOp>(op))
+      applyOpPatternsAndFold(op, patterns);
   });
 
   // Turn memrefs' non-identity layouts maps into ones with identity. Collect
   // alloc ops first and then process since normalizeMemRef replaces/erases ops
   // during memref rewriting.
   SmallVector<AllocOp, 4> allocOps;
   func.walk([&](AllocOp op) { allocOps.push_back(op); });
   for (auto allocOp : allocOps) {
     normalizeMemRef(allocOp);
   }
 }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 
 using namespace mlir;
 
 #define DEBUG_TYPE "pattern-matcher"
 
 /// The max number of iterations scanning for pattern match.
 static unsigned maxPatternMatchIterations = 10;
 
+//===----------------------------------------------------------------------===//
+// GreedyPatternRewriteDriver
+//===----------------------------------------------------------------------===//
+
 namespace {
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns in a roughly "bottom up" way.
 class GreedyPatternRewriteDriver : public PatternRewriter {
 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const OwningRewritePatternList &patterns)
       : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
     worklist.reserve(64);
   }
 
-  /// Perform the rewrites while folding and erasing any dead ops. Return true
-  /// if the rewrite converges in `maxIterations`.
   bool simplify(MutableArrayRef<Region> regions, int maxIterations);
 
   void addToWorklist(Operation *op) {
     // Check to see if the worklist already contains this op.
     if (worklistMap.count(op))
       return;
 
     worklistMap[op] = worklist.size();
bool mlir::applyPatternsAndFoldGreedily(
 
   // Start the pattern driver.
   GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns);
   bool converged = driver.simplify(regions, maxPatternMatchIterations);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << maxPatternMatchIterations << " times";
   });
   return 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 OwningRewritePatternList &patterns)
+      : PatternRewriter(ctx), matcher(patterns), folder(ctx) {}
+
+  bool simplifyLocally(Operation *op, int maxIterations, bool &erased);
+
+  /// No additional action needed other than inserting the op.
+  Operation *insert(Operation *op) override { return OpBuilder::insert(op); }
+
+  // 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;
+  }
+
+  // When a root is going to be replaced, its removal will be notified as well.
+  // So there is nothing to do here.
+  void notifyRootReplaced(Operation *op) override {}
+
+private:
+  /// The low-level pattern matcher.
+  RewritePatternMatcher matcher;
+
+  /// Non-pattern based folder for operations.
+  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 dead, or no
+/// more changes happen in an iteration. Returns true if the rewrite converges
+/// in `maxIterations`. `erased` is set to true if `op` gets erased.
+bool OpPatternRewriteDriver::simplifyLocally(Operation *op, int maxIterations,
+                                             bool &erased) {
+  bool changed = false;
+  erased = false;
+  opErasedViaPatternRewrites = false;
+  int i = 0;
+  // Iterate until convergence or until maxIterations. Deletion of the op as
+  // a result of being dead or folded is convergence.
+  do {
+    // If the operation is trivially dead - remove it.
+    if (isOpTriviallyDead(op)) {
+      op->erase();
+      erased = true;
+      return true;
+    }
+
+    // Try to fold this op.
+    bool inPlaceUpdate;
+    if (succeeded(folder.tryToFold(op, /*processGeneratedConstants=*/nullptr,
+                                   /*preReplaceAction=*/nullptr,
+                                   &inPlaceUpdate))) {
+      changed = true;
+      if (!inPlaceUpdate) {
+        erased = true;
+        return true;
+      }
+    }
+
+    // Make sure that any new operations are inserted at this point.
+    setInsertionPoint(op);
+
+    // 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 |= matcher.matchAndRewrite(op, *this);
+    if ((erased = opErasedViaPatternRewrites))
+      return true;
+  } while (changed && ++i < maxIterations);
+
+  // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
+  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.
+bool mlir::applyOpPatternsAndFold(Operation *op,
+                                  const OwningRewritePatternList &patterns,
+                                  bool *erased) {
+  // Start the pattern driver.
+  OpPatternRewriteDriver driver(op->getContext(), patterns);
+  bool opErased;
+  bool converged =
+      driver.simplifyLocally(op, maxPatternMatchIterations, opErased);
+  if (erased)
+    *erased = opErased;
+  LLVM_DEBUG(if (!converged) {
+    llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
+                 << maxPatternMatchIterations << " times";
+  });
+  return converged;
+}

mlir/lib/Transforms/Utils/LoopUtils.cpp

 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
if (!opGroupQueue.empty()) {
         opGroupQueue.clear();
         lbShift += tripCount * step;
       } else {
         res = generateShiftedLoop(b.getShiftedAffineMap(origLbMap, lbShift),
                                   b.getShiftedAffineMap(origLbMap, d),
                                   opGroupQueue, /*offset=*/0, forOp, b);
         lbShift = d * step;
       }
-      if (!prologue && res)
+
+      if (res) {
+        // Simplify/canonicalize the affine.for.
+        OwningRewritePatternList patterns;
+        AffineForOp::getCanonicalizationPatterns(patterns, res.getContext());
+        bool erased;
+        applyOpPatternsAndFold(res, std::move(patterns), &erased);
+
+        if (!erased && !prologue)
           prologue = res;
+        if (!erased)
           epilogue = res;
+      }
     } else {
       // Start of first interval.
       lbShift = d * step;
     }
     // Augment the list of operations that get into the current open interval.
     opGroupQueue.push_back({d, sortedOpGroups[d]});
   }
 
bool mlir::isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops,
   assert(loopPermMap.size() == loops.size());
   unsigned maxLoopDepth = loops.size();
   std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
   getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
   return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap);
 }
 
 /// Return true if `loops` is a perfect nest.
-static bool LLVM_ATTRIBUTE_UNUSED isPerfectlyNested(ArrayRef<AffineForOp> loops) {
+static bool LLVM_ATTRIBUTE_UNUSED
+isPerfectlyNested(ArrayRef<AffineForOp> loops) {
   auto outerLoop = loops.front();
   for (auto loop : loops.drop_front()) {
     auto parentForOp = dyn_cast<AffineForOp>(loop.getParentOp());
     // parentForOp's body should be just this loop and the terminator.
     if (parentForOp != outerLoop ||
         parentForOp.getBody()->getOperations().size() != 2)
       return false;
     outerLoop = loop;

mlir/test/Dialect/Affine/affine-data-copy.mlir

affine.for %ii = affine_map<(d0) -> (d0)>(%i) to min #map_ub(%i) {
       %5 = affine.load %A[%ii] : memref<4096xf32>
       %6 = mulf %5, %5 : f32
       affine.store %6, %A[%ii] : memref<4096xf32>
     }
   }
   return %A : memref<4096xf32>
 }
 // CHECK:      affine.for %[[IV1:.*]] = 0 to 4096 step 100
-// CHECK-NEXT:   %[[BUF:.*]] = alloc() : memref<100xf32>
+// CHECK:        %[[BUF:.*]] = alloc() : memref<100xf32>
 // CHECK-NEXT:   affine.for %[[IV2:.*]] = #[[MAP_IDENTITY]](%[[IV1]]) to min #[[MAP_MIN_UB1]](%[[IV1]]) {
 // CHECK-NEXT:     affine.load %{{.*}}[%[[IV2]]] : memref<4096xf32>
 // CHECK-NEXT:     affine.store %{{.*}}, %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
 // CHECK-NEXT:   }
 // CHECK-NEXT:   affine.for %[[IV2:.*]] = #[[MAP_IDENTITY]](%[[IV1]]) to min #[[MAP_MIN_UB2]](%[[IV1]]) {
 // CHECK-NEXT:     affine.load %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
 // CHECK-NEXT:     mulf
 // CHECK-NEXT:     affine.store %{{.*}}, %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
 // CHECK-NEXT:   }
-// CHECK-NEXT:   affine.for %[[IV2:.*]] = #[[MAP_IDENTITY]](%[[IV1]]) to min #[[MAP_MIN_UB1]](%[[IV1]]) {
+// CHECK:        affine.for %[[IV2:.*]] = #[[MAP_IDENTITY]](%[[IV1]]) to min #[[MAP_MIN_UB1]](%[[IV1]]) {
 // CHECK-NEXT:     affine.load %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>
 // CHECK-NEXT:     affine.store %{{.*}}, %{{.*}}[%[[IV2]]] : memref<4096xf32>
 // CHECK-NEXT:   }
 // CHECK-NEXT:   dealloc %[[BUF]] : memref<100xf32>
 // CHECK-NEXT: }
 
 // -----
 
 // Lower bound is a max; upper bound is a min. This pattern typically appears
 // with multi-level tiling when the tile sizes used don't divide loop trip
 // counts.
 
-#lb = affine_map<(d0, d1) -> (d0 * 512, d1 * 6)>
-#ub = affine_map<(d0, d1) -> (d0 * 512 + 512, d1 * 6 + 6)>
+#lb = affine_map<()[s0, s1] -> (s0 * 512, s1 * 6)>
+#ub = affine_map<()[s0, s1] -> (s0 * 512 + 512, s1 * 6 + 6)>
 
 // CHECK-DAG: #[[LB:.*]] = affine_map<()[s0, s1] -> (s0 * 512, s1 * 6)>
 // CHECK-DAG: #[[UB:.*]] = affine_map<()[s0, s1] -> (s0 * 512 + 512, s1 * 6 + 6)>
 
 // CHECK-LABEL: max_lower_bound(%{{.*}}: memref<2048x516xf64>,
 // CHECK-SAME: [[i:arg[0-9]+]]
 // CHECK-SAME: [[j:arg[0-9]+]]
 func @max_lower_bound(%M: memref<2048x516xf64>, %i : index, %j : index) {
   affine.for %ii = 0 to 2048 {
-    affine.for %jj = max #lb(%i, %j) to min #ub(%i, %j) {
+    affine.for %jj = max #lb()[%i, %j] to min #ub()[%i, %j] {
       affine.load %M[%ii, %jj] : memref<2048x516xf64>
     }
   }
   return
 }
 
 // CHECK:      %[[BUF=.*]] = alloc() : memref<2048x6xf64>
 // CHECK-NEXT: affine.for %[[ii:.*]] = 0 to 2048 {

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

 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -simplify-affine-structures | FileCheck %s
 
-// CHECK-DAG: #[[SET_EMPTY_2D:.*]] = affine_set<(d0, d1) : (1 == 0)>
+// CHECK-DAG: #[[SET_EMPTY:.*]] = affine_set<() : (1 == 0)>
 // CHECK-DAG: #[[SET_2D:.*]] = affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0)>
-// CHECK-DAG: #[[SET_EMPTY_2D_2S:.*]] = affine_set<(d0, d1)[s0, s1] : (1 == 0)>
-// CHECK-DAG: #[[SET_2D_2S:.*]] = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0, d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0, d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0, d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0)>
-// CHECK-DAG: #[[SET_EMPTY_1D:.*]] = affine_set<(d0) : (1 == 0)>
-// CHECK-DAG: #[[SET_EMPTY_1D_2S:.*]] = affine_set<(d0)[s0, s1] : (1 == 0)>
-// CHECK-DAG: #[[SET_EMPTY_3D:.*]] = affine_set<(d0, d1, d2) : (1 == 0)>
+// CHECK-DAG: #[[SET_7_11:.*]] = affine_set<(d0, d1) : (d0 * 7 + d1 * 5 + 88 == 0, d0 * 5 - d1 * 11 + 60 == 0, d0 * 11 + d1 * 7 - 24 == 0, d0 * 7 + d1 * 5 + 88 == 0)>
+
+// An external function that we will use in bodies to avoid DCE.
+func @external() -> ()
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set0() {
 func @test_gaussian_elimination_empty_set0() {
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: [[SET_EMPTY_2D]](%arg0, %arg1)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1) : (2 == 0)>(%arg0, %arg1) {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set1() {
 func @test_gaussian_elimination_empty_set1() {
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: [[SET_EMPTY_2D]](%arg0, %arg1)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1) : (1 >= 0, -1 >= 0)> (%arg0, %arg1) {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // CHECK-LABEL: func @test_gaussian_elimination_non_empty_set2() {
 func @test_gaussian_elimination_non_empty_set2() {
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
       // CHECK: #[[SET_2D]](%arg0, %arg1)
       affine.if affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0, d1 + 101 >= 0)>(%arg0, %arg1) {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set3() {
 func @test_gaussian_elimination_empty_set3() {
   %c7 = constant 7 : index
   %c11 = constant 11 : index
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: #[[SET_EMPTY_2D_2S]](%arg0, %arg1)[%c7, %c11]
+      // CHECK: #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1)[s0, s1] : (d0 - s0 == 0, d0 + s0 == 0, s0 - 1 == 0)>(%arg0, %arg1)[%c7, %c11] {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // Set for test case: test_gaussian_elimination_non_empty_set4
 #set_2d_non_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
                                        d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0,
                                        d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0,
                                        d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0)>
 
 // CHECK-LABEL: func @test_gaussian_elimination_non_empty_set4() {
 func @test_gaussian_elimination_non_empty_set4() {
   %c7 = constant 7 : index
   %c11 = constant 11 : index
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: #[[SET_2D_2S]](%arg0, %arg1)[%c7, %c11]
+      // CHECK: #[[SET_7_11]](%arg0, %arg1)
       affine.if #set_2d_non_empty(%arg0, %arg1)[%c7, %c11] {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // Add invalid constraints to previous non-empty set to make it empty.
-// Set for test case: test_gaussian_elimination_empty_set5
 #set_2d_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
                                        d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0,
                                        d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0,
                                        d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,
                                        d0 - 1 == 0, d0 + 2 == 0)>
 
 // CHECK-LABEL: func @test_gaussian_elimination_empty_set5() {
 func @test_gaussian_elimination_empty_set5() {
   %c7 = constant 7 : index
   %c11 = constant 11 : index
   affine.for %arg0 = 1 to 10 {
     affine.for %arg1 = 1 to 100 {
-      // CHECK: #[[SET_EMPTY_2D_2S]](%arg0, %arg1)[%c7, %c11]
+      // CHECK: #[[SET_EMPTY]]()
       affine.if #set_2d_empty(%arg0, %arg1)[%c7, %c11] {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // This is an artificially created system to exercise the worst case behavior of
 // FM elimination - as a safeguard against improperly constructed constraint
#set_fuzz_virus = affine_set<(d0, d1, d2, d3, d4, d5) : (
                             mod 2038 mod 2390 mod 2039 floordiv 55 >= 0
 )>
 
 // CHECK-LABEL: func @test_fuzz_explosion
 func @test_fuzz_explosion(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) {
   affine.for %arg4 = 1 to 10 {
     affine.for %arg5 = 1 to 100 {
       affine.if #set_fuzz_virus(%arg4, %arg5, %arg0, %arg1, %arg2, %arg3) {
+        call @external() : () -> ()
       }
     }
   }
   return
 }
 
 // CHECK-LABEL: func @test_empty_set(%arg0: index) {
 func @test_empty_set(%N : index) {
   affine.for %i = 0 to 10 {
     affine.for %j = 0 to 10 {
-      // CHECK: affine.if #[[SET_EMPTY_2D]](%arg1, %arg2)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1) : (d0 - d1 >= 0, d1 - d0 - 1 >= 0)>(%i, %j) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[SET_EMPTY_1D]](%arg1)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {
         "bar"() : () -> ()
       }
-      // CHECK: affine.if #[[SET_EMPTY_1D]](%arg1)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[SET_EMPTY_1D_2S]](%arg1)[%arg0, %arg0]
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0)[s0, s1] : (d0 >= 0, -d0 + s0 - 1 >= 0, -s0 >= 0)>(%i)[%N, %N] {
         "bar"() : () -> ()
       }
-      // CHECK: affine.if #[[SET_EMPTY_3D]](%arg1, %arg2, %arg0)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       // The set below implies d0 = d1; so d1 >= d0, but d0 >= d1 + 1.
       affine.if affine_set<(d0, d1, d2) : (d0 - d1 == 0, d2 - d0 >= 0, d0 - d1 - 1 >= 0)>(%i, %j, %N) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[SET_EMPTY_2D]](%arg1, %arg2)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       // The set below has rational solutions but no integer solutions; GCD test catches it.
       affine.if affine_set<(d0, d1) : (d0*2 -d1*2 - 1 == 0, d0 >= 0, -d0 + 100 >= 0, d1 >= 0, -d1 + 100 >= 0)>(%i, %j) {
         "foo"() : () -> ()
       }
-      // CHECK: affine.if #[[SET_EMPTY_2D]](%arg1, %arg2)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1) : (d1 == 0, d0 - 1 >= 0, - d0 - 1 >= 0)>(%i, %j) {
         "foo"() : () -> ()
       }
     }
   }
   // The tests below test GCDTightenInequalities().
   affine.for %k = 0 to 10 {
     affine.for %l = 0 to 10 {
       // Empty because no multiple of 8 lies between 4 and 7.
-      // CHECK: affine.if #[[SET_EMPTY_1D]](%arg1)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0) : (8*d0 - 4 >= 0, -8*d0 + 7 >= 0)>(%k) {
         "foo"() : () -> ()
       }
       // Same as above but with equalities and inequalities.
-      // CHECK: affine.if #[[SET_EMPTY_2D]](%arg1, %arg2)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1) : (d0 - 4*d1 == 0, 4*d1 - 5 >= 0, -4*d1 + 7 >= 0)>(%k, %l) {
         "foo"() : () -> ()
       }
       // Same as above but with a combination of multiple identifiers. 4*d0 +
       // 8*d1 here is a multiple of 4, and so can't lie between 9 and 11. GCD
       // tightening will tighten constraints to 4*d0 + 8*d1 >= 12 and 4*d0 +
       // 8*d1 <= 8; hence infeasible.
-      // CHECK: affine.if #[[SET_EMPTY_2D]](%arg1, %arg2)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1) : (4*d0 + 8*d1 - 9 >= 0, -4*d0 - 8*d1 + 11 >= 0)>(%k, %l) {
         "foo"() : () -> ()
       }
       // Same as above but with equalities added into the mix.
-      // CHECK: affine.if #[[SET_EMPTY_3D]](%arg1, %arg1, %arg2)
+      // CHECK: affine.if #[[SET_EMPTY]]()
       affine.if affine_set<(d0, d1, d2) : (d0 - 4*d2 == 0, d0 + 8*d1 - 9 >= 0, -d0 - 8*d1 + 11 >= 0)>(%k, %k, %l) {
         "foo"() : () -> ()
       }
     }
   }
 
   affine.for %m = 0 to 10 {
-    // CHECK: affine.if #[[SET_EMPTY_1D]](%arg{{[0-9]+}})
+    // CHECK: affine.if #[[SET_EMPTY]]()
     affine.if affine_set<(d0) : (d0 mod 2 - 3 == 0)> (%m) {
       "foo"() : () -> ()
     }
   }
 
   return
 }
 
 // -----
 
-// CHECK-DAG: #[[SET_2D:.*]] = affine_set<(d0, d1) : (d0 >= 0, -d0 + 50 >= 0)
-// CHECK-DAG: #[[SET_EMPTY:.*]] = affine_set<(d0, d1) : (1 == 0)
-// CHECK-DAG: #[[SET_UNIV:.*]] = affine_set<(d0, d1) : (0 == 0)
+// An external function that we will use in bodies to avoid DCE.
+func @external() -> ()
+
+// CHECK-DAG: #[[SET:.*]] = affine_set<()[s0] : (s0 >= 0, -s0 + 50 >= 0)
+// CHECK-DAG: #[[EMPTY_SET:.*]] = affine_set<() : (1 == 0)
+// CHECK-DAG: #[[UNIV_SET:.*]] = affine_set<() : (0 == 0)
 
 // CHECK-LABEL: func @simplify_set
 func @simplify_set(%a : index, %b : index) {
-  // CHECK: affine.if #[[SET_2D]]
+  // CHECK: affine.if #[[SET]]
   affine.if affine_set<(d0, d1) : (d0 - d1 + d1 + d0 >= 0, 2 >= 0, d0 >= 0, -d0 + 50 >= 0, -d0 + 100 >= 0)>(%a, %b) {
+    call @external() : () -> ()
   }
-  // CHECK: affine.if #[[SET_EMPTY]]
+  // CHECK: affine.if #[[EMPTY_SET]]
   affine.if affine_set<(d0, d1) : (d0 mod 2 - 1 == 0, d0 - 2 * (d0 floordiv 2) == 0)>(%a, %b) {
+    call @external() : () -> ()
   }
-  // CHECK: affine.if #[[SET_UNIV]]
+  // CHECK: affine.if #[[UNIV_SET]]
   affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%a, %b) {
-  }
-  return
+    call @external() : () -> ()
+  }
+  return
+
+// -----
+
+// CHECK-DAG: -> (s0 * 2 + 1)
+
+// Test "op local" simplification on affine.apply. DCE on addi will not happen.
+func @affine.apply(%N : index) {
+  %v = affine.apply affine_map<(d0, d1) -> (d0 + d1 + 1)>(%N, %N)
+  addi %v, %v : index
+  // CHECK: affine.apply #map{{.*}}()[%arg0]
+  // CHECK-NEXT: addi
+  return
 }

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

if (clMemRefFilter) {
     affineDataCopyGenerate(loopNest, copyOptions, load.getMemRef(), copyNests);
   } else if (clTestGenerateCopyForMemRegion) {
     CopyGenerateResult result;
     MemRefRegion region(loopNest.getLoc());
     region.compute(load, /*loopDepth=*/0);
     generateCopyForMemRegion(region, loopNest, copyOptions, result);
   }
 
-  // Promote any single iteration loops in the copy nests.
+  // Promote any single iteration loops in the copy nests and simplify
+  // load/stores.
+  SmallVector<Operation *, 4> copyOps;
   for (auto nest : copyNests)
-    nest->walk([](AffineForOp forOp) { promoteIfSingleIteration(forOp); });
-
-  // Promoting single iteration loops could lead to simplification
-  // of load's/store's. We will run the canonicalization patterns again.
+    // With a post order walk, the erasure of loops does not affect
+    // continuation of the walk or the collection of load/store ops.
+    nest->walk([&](Operation *op) {
+      if (auto forOp = dyn_cast<AffineForOp>(op))
+        promoteIfSingleIteration(forOp);
+      else if (auto loadOp = dyn_cast<AffineLoadOp>(op))
+        copyOps.push_back(loadOp);
+      else if (auto storeOp = dyn_cast<AffineStoreOp>(op))
+        copyOps.push_back(storeOp);
+    });
+
+  // Promoting single iteration loops could lead to simplification of
+  // generated load's/store's, and the latter could anyway also be
+  // canonicalized.
   OwningRewritePatternList patterns;
+  for (auto 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());
-  applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
+    }
+    applyOpPatternsAndFold(op, std::move(patterns));
+  }
 }
 
 namespace mlir {
 void registerTestAffineDataCopyPass() {
   PassRegistration<TestAffineDataCopy>(
       PASS_NAME, "Tests affine data copy utility functions.");
 }
 } // namespace mlir