[mlir][scf] Simplify affine.min ops after loop peeling

[springerm]

Summary

Simplify affine.min ops, enabling various other canonicalizations inside the peeled loop body.

affine.min ops such as:

map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>
%r = affine.min #affine.min #map(%iv)[%step, %ub]

are rewritten them into (in the case the peeled loop):

%r = %step

To determine how an affine.min op should be rewritten and to prove its correctness, FlatAffineConstraints is utilized.

Depends On D107730

[springerm]

rebase

[nicolasvasilache]

Generally looks good!

mlir/include/mlir/Dialect/SCF/Transforms.h

78

Does the description need mention of loop peeling anymore?
It seems to me this is a general feature that can apply to any value iv that is known to be bound by ub and to increment by step.

I would recommend the following:

1. split this into: a. a helper that constructs the set and returns whether it is empty b. an IR rewriter that calls the helper
2. rename "insideLoop" into a boolean that convey lesser-equal / greater than.

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

175

I would rephrase this in a more mathematical form; I'd love to "see" the constraints system you are building rather than have it described in text.

My experience is that a formatting constraints sets by block matrices where the 0 and subidentities appear clearly makes it quite easy to follow; see e.e. page 4 of http://icps.u-strasbg.fr/~bastoul/research/papers/VBGC06-ICS.pdf which should give an intuitive understanding of the layout of the constraints (in a different use case).

195

braces for multiline if please.

196

I am not sure I understand this note, could you explain?
It seems you want to ignore failures here but I am not sure how this works.

224

I do not see the "for each" behavior here?

234

Can we make the creation of those dims ahead of time?
The addInequality would be a bit more annoying to write but it is usually more readable to just fix the dimension of the constraint set before introducing inequalities (when possible).

256

I'd restructure this block and the previous a bit if possible; it wasn't immediately clear to me why you needed the extra insert here? Then after digging a bit more I saw that it is probably getSliceBounds that drops the particular dimension by projecting.

261

This would benefit from seeing the whole constraint set in math form.

[springerm]

rebase

[springerm]

rebase + update

[springerm]

minor updates

[springerm]

update

[springerm]

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

196

Update: I no longer use addLowerOrUpperBound here and add an equality directly.

224

The "for each" is in addLowerOrUpperBound (it is not written as a for-each and if you're not familiar with this function, it may be hard to see). addLowerOrUpperBound adds an inequality for each result in mapOpMap.

[ftynse]

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

268–269

I don't see anything in getSliceBounds that can clear the minOpValueUb so the first condition is always false. The second condition is currently also always false, but there is a TODO in getSliceBound that would trigger it, consider adding some debug spew here for when this becomes the case.

279

Nit: also assert that eq[kDimMinOpUb] == 0? Looking at the current system of constraints, it should be the case, but better to future-proof this.

282

292

I'd consider just removing the last inequality at the end of the loop, this will spare a copy on each iteration.

309

Are we sure that the UB AffineExpr has exactly three inputs? There doesn't seem to be any filtering of the affine.min that gets processed, so the original map can seemingly be arbitrarily complex as long as the emptiness check is satisfied.

[springerm]

address comments

[springerm]

mlir/include/mlir/Dialect/SCF/Transforms.h

78

Refactored as follows: The logic for solving the constraint system etc. is in canonicalizeAffineMinOp.

This function could even be moved out of the SCF dialect into a different file.

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

234

I rewrote large parts of the commit. Dimensions are still added "in the middle", though. This is more convenient, because I can pass the FlatAffineConstraints to a builder function, that adds pattern-specific constraints. At that point, I do not want to bother the user of canonicalizeAffineMinOp with unexpected extra dims. In the constraint builder, only those dims exist in the constraint set, that were specified by the caller (dims parameter).

279

I no longer use addEquality here.

309

Correct, this is wrong here. Had the fix in a dependent commit that's not out for review yet, but moved it into this one.

[aartbik]

mlir/include/mlir/Dialect/SCF/Transforms.h

49

can you elaborate on the "is beneficial" (i.e. we generally distinguish between enabling transformations and transformations that "cleanup" after the fact, and this is clearly the latter)

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

107

document the result value (ie. returns failure when not rewritten)

204

typo: replace with that bound?

[springerm]

rebase

[springerm]

rebase

[springerm]

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

107

Extended the comment. Also added detailed description of what loops are rewritten in the comment of peelAndCanonicalizeForLoop.

261

Added to the beginning of the function.

[springerm]

address comments

[springerm]

rebase

[nicolasvasilache]

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

165

Bound -> Bind

169

assoicated - > associated

180

So I am looking at the implementation of alignAffineMapWithValues as it now also gets used from here and I find that it could be significantly improved IMO.

I would just make everything explicit and use:

/// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
/// results and return a new AffineMap with the new results and with inferred
/// number of dims and symbols.
AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map) const;

To get there, there would be a helper function that would take an operand and return the AffineExpr to which the operand is aligned in the constraint system.
If not there yet it would add it and return a new AffineSymbolExpr (I think from looking at your code).

Then you can easily for_each those into he DenseMap and call the replace func.
This should significantly improve composability.

Does this make sense ?

(Not for this revision but would be nice as a followup).

189

"replaced with the bound" -> "folded away" ?

227

dims -> operands ?

251

It seems quite convoluted to pass a lambda with implicit invariant that after "addDim in this function is called then the number of columns must match".

I would get rid of constraintsBuilder entirely and just lift l 239 - 252 into the caller.

262

Same comment as above, I think constructing an explicit DenseMap<AffineExpr, AffineExpr> based on Values and just using the appropriate sparse replacement function would be significantly more readable.

324

In the grander order of things where we call FM, I find this micro-optimization to be more confusing than anything else (especially in the context of new symValues that you also have to assert against).

I'd prefer to see a fresh copy of the FlatAffineValueConstraint, add a constraint, test for emptiness and let if RAII away.

327

(I think) the block of code below somewhat reproduces AffineMap::compressUnusedDims and Symbols?
But this is mostly hidden by the fact that there are empt values that may be introduced for Dim ?

In light of this last point, I think the convolution has reached a point where the separation between alignment and Value / Attribute worlds would already be beneficial in this revision.

381

This seems unnecessarily convoluted, at the very least I would:

// Add loop peeling invariant. This is the main piece of knowledge that
// enables AffineMinOp simplification.
if (insideLoop) {
  // ub - iv >= step (equiv.: -iv + ub - step + 0 >= 0)
  // Intuitively: Inside the peeled loop, every iteration is a "full"
  // iteration, i.e., step divides the iteration space `ub - lb` evenly.
  auto lambda = ...;
  return canonicalizeAffineMinOp(rewriter, minOp,
                             /*dims=*/ValueRange{iv, ub, step}, lambda);
}
// ub - iv < step (equiv.: iv + -ub + step - 1 >= 0)
// Intuitively: `iv` is the split bound here, i.e., the iteration variable
// value of the very last iteration (in the unpeeled loop). At that point,
// there are less than `step` elements remaining. (Otherwise, the peeled
// loop would run for at least one more iteration.)
auto lambda = ...;
return canonicalizeAffineMinOp(rewriter, minOp,
                           /*dims=*/ValueRange{iv, ub, step}, lambda);

However, in light of the other comments I would refactor much more deeply and just pass a FlatAffineConstraints to the canonicalize function.

mlir/test/Dialect/SCF/for-loop-peeling.mlir

4

Hmm seems like this should simplify to (s1 - s0) mod s2.
Likely some missing affineexpr simplification pattern, however it may just be adding more canonicalization patterns that will continue to fail in slightly more complex cases.
Maybe a lot of this should be using FlatAffineConstraints directly.

168

#map3 -> #[[MAP3]]

201

I find a little hard to keep track of things at such a distance, could you please reorder in interleaved form?

// Most common case: Rewrite min(%ub - %iv, %step) to %step.
//      CHECK:     memref.store %[[STEP]]
%m0 = affine.min #map0(%ub, %iv)[%step]
memref.store %m0, %d[%c0] : memref<?xindex>

 // Increase %ub - %iv a little bit, pattern should still apply.
...

// At the end, the checks within the scf.if

Maybe even embed the affine_map into the affine.min since it is not reused ?

[springerm]

address comments

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

165

I think this should actually be "bound" (to bound). As in upper/lower bound.

227

These are the SSA values associated to the dimensions in the constraints set. These dims can be referred to in the constraintsBuilder. Note: dims are not the operands of the AffineMinOp. Those would be minOp.operands().

E.g.:
The caller of canonicalizeAffineMinOp adds two inequalities via the constraintBuilder. Those inequalities use iv, ub, step dimensions. dims contains the SSA values of those dimensions. Note: It is important that dims and the (in)equalities creates by the constraintBuilder stay in sync.

Looking at this code again, I think it would be better to use FlatAffineValueConstraints here instead of FlatAffineConstraints. Then I don't have to maintain extra helper variables for dims/syms SSA values. With the recent refactorings, I can just call the base class functions which do not have any "affine dialect" assumptions. I can try splitting the FlatAffineValueConstraints class in a subsequent commit, so that the "affine dialect" functionality is properly separated (as we talked last time) and cannot be called by accident in here.

I'll refactor this as follows:

• canonicalizeAffineMinOp takes a FlatAffineValueConstraints as an argument.
• dims parameter is gone.

251

Replaced by FlatAffineValueConstraints. There is need to keep track of dim/sym SSA values anymore.

324

That's how I had it in an earlier revision. This was a suggestion by another reviewer to avoid the overhead of copying the FlatAffineConstraints. Either one is fine with me.

327

I think I should use canonicalizeMapAndOperands here. It requires only a small change to support "empty" Values.

381

Replaced with FlatAffineValueConstraints, which is created by the caller and passed to canonicalizeAffineMinOp.

mlir/include/mlir/Dialect/SCF/Transforms.h

 
 #ifndef MLIR_DIALECT_SCF_TRANSFORMS_H_
 #define MLIR_DIALECT_SCF_TRANSFORMS_H_
 
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 
+class AffineMinOp;
 class ConversionTarget;
 struct LogicalResult;
 class MLIRContext;
 class Region;
 class RewriterBase;
 class TypeConverter;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 class Operation;
+class Value;
 
 namespace scf {
 
 class IfOp;
 class ForOp;
 class ParallelOp;
 class ForOp;
 
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
 /// into one scf.parallel operations. Uses a naive aliasing and dependency
 /// analysis.
 void naivelyFuseParallelOps(Region &region);
 
 /// Rewrite a for loop with bounds/step that potentially do not divide evenly
 /// into a for loop where the step divides the iteration space evenly, followed
 /// by an scf.if for the last (partial) iteration (if any). This transformation
 /// is called "loop peeling".
 ///
-/// Other patterns can simplify/canonicalize operations in the body of the loop
-/// and the scf.if. This is beneficial for a wide range of transformations such
+/// This transformation is beneficial for a wide range of transformations such
 /// as vectorization or loop tiling.
 ///
 /// E.g., assuming a lower bound of 0 (for illustration purposes):
 /// ```
 /// scf.for %iv = %c0 to %ub step %c4 {
 ///   (loop body)
 /// }
 /// ```
 /// is rewritten into the following pseudo IR:
 /// ```
 /// %newUb = %ub - (%ub mod %c4)
 /// scf.for %iv = %c0 to %newUb step %c4 {
 ///   (loop body)
 /// }
 /// scf.if %newUb < %ub {
 ///   (loop body)
 /// }
 /// ```
 ///
-/// This function rewrites the given scf.for loop in-place and creates a new
-/// scf.if operation (returned via `ifOp`) for the last iteration.
-///
-/// TODO: Simplify affine.min ops inside the new loop/if statement.
-LogicalResult peelForLoop(RewriterBase &b, ForOp forOp, scf::IfOp &ifOp);
+/// After loop peeling, this function tries to simplify/canonicalize affine.min
+/// operations in the body of the loop and the scf.if, taking advantage of the
+/// fact that every iteration of the peeled loop is a "full" iteration. This
+/// canonicalization is expected to enable further canonicalization
+/// opportunities through other patterns.
+///
+/// Note: This function rewrites the given scf.for loop in-place and creates a
+/// new scf.if operation for the last iteration. It replaces all uses of the
+/// unpeeled loop with the results of the newly generated scf.if.
+LogicalResult peelAndCanonicalizeForLoop(RewriterBase &rewriter, ForOp forOp);
 
 /// Tile a parallel loop of the form
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4, %arg5)
 ///
 /// into
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4*tileSize[0],

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

 //===- LoopSpecialization.cpp - scf.parallel/SCR.for specialization -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Specializes parallel loops and for loops for easier unrolling and
 // vectorization.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/SCF/Passes.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/BlockAndValueMapping.h"
static void specializeForLoopForUnrolling(ForOp op) {
   ifOp.getThenBodyBuilder().clone(*op.getOperation(), map);
   ifOp.getElseBodyBuilder().clone(*op.getOperation());
   op.erase();
 }
 
 /// Rewrite a for loop with bounds/step that potentially do not divide evenly
 /// into a for loop where the step divides the iteration space evenly, followed
 /// by an scf.if for the last (partial) iteration (if any).
-LogicalResult mlir::scf::peelForLoop(RewriterBase &b, ForOp forOp,
-                                     scf::IfOp &ifOp) {
+///
+/// This function rewrites the given scf.for loop in-place and creates a new
+/// scf.if operation for the last iteration. It replaces all uses of the
+/// unpeeled loop with the results of the newly generated scf.if.
+///
+/// The newly generated scf.if operation is returned via `ifOp`. The boundary
+/// at which the loop is split (new upper bound) is returned via `splitBound`.
+static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp, scf::IfOp &ifOp,
+                                 Value &splitBound) {
   RewriterBase::InsertionGuard guard(b);
   auto lbInt = getConstantIntValue(forOp.lowerBound());
   auto ubInt = getConstantIntValue(forOp.upperBound());
   auto stepInt = getConstantIntValue(forOp.step());
 
   // No specialization necessary if step already divides upper bound evenly.
   if (lbInt && ubInt && stepInt && (*ubInt - *lbInt) % *stepInt == 0)
     return failure();
   // No specialization necessary if step size is 1.
   if (stepInt == static_cast<int64_t>(1))
     return failure();
 
   auto loc = forOp.getLoc();
   AffineExpr dim0, dim1, dim2;
   bindDims(b.getContext(), dim0, dim1, dim2);
   // New upper bound: %ub - (%ub - %lb) mod %step
   auto modMap = AffineMap::get(3, 0, {dim1 - ((dim1 - dim0) % dim2)});
   b.setInsertionPoint(forOp);
-  Value splitBound = b.createOrFold<AffineApplyOp>(
+  splitBound = b.createOrFold<AffineApplyOp>(
       loc, modMap,
       ValueRange{forOp.lowerBound(), forOp.upperBound(), forOp.step()});
 
   // Set new upper loop bound.
   Value previousUb = forOp.upperBound();
   b.updateRootInPlace(forOp,
                       [&]() { forOp.upperBoundMutable().assign(splitBound); });
   b.setInsertionPointAfter(forOp);
b.cloneRegionBefore(forOp.region(), ifOp.thenRegion(),
                       ifOp.thenRegion().begin(), bvm);
   // Build else case.
   if (!resultTypes.empty())
     ifOp.getElseBodyBuilder().create<scf::YieldOp>(loc, forOp->getResults());
 
   return success();
 }
 
+/// Try to simplify an affine.min operation after loop peeling. This function
+/// detects affine.min operations such as (ub is the previous upper bound of the
+/// unpeeled loop):
+/// ```
+/// #map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>
+/// %r = affine.min #affine.min #map(%iv)[%step, %ub]
+/// ```
+/// and rewrites them into (in the case the peeled loop):
+/// ```
+/// %r = %step
+/// ```
+/// affine.min operations inside the generated scf.if operation are rewritten in
+/// a similar way.
+///
+/// This function builds up a set of constraints, capable of proving that:
+/// * min(step, ub - iv) == step inside the peeled loop
+/// * min(step, ub - iv) == ub - iv inside the if-statement
+///
+/// Note: `ub` is the previous upper bound of the loop (before peeling).
+/// `insideLoop` must be true for affine.min ops inside the loop and false for
+/// affine.min ops inside the scf.for op.
+///
+/// 1. Set up a constraint system with dimensions for `iv`, `ub` (previous
+///    upper bound), `step`, `minOp` and `minOpUb` (upper bound of `minOp`).
+/// 2. a. If inside the main loop: Add inequality: ub - iv >= step
+/// 2. b. If inside the if-statement: Add inequality: ub - iv < step
+/// 3. Add each result of `minOp` as a dimension `r_i` to the constraint system.
+/// 4. Compute an upper bound of `minOp` and bind it to `minOpUb`.
+/// 5. For each result of `minOp`: Prove that r_i >= minOpUb. If this is the
+///    case, upper_bound(minOp) == lower_bound(minOp) and `minOp` can be
+///    replaced with the that bound.
+static LogicalResult rewritePeeledAffineOp(RewriterBase &rewriter,
+                                           AffineMinOp minOp, Value iv,
+                                           Value ub, Value step,
+                                           bool insideLoop) {
+  RewriterBase::InsertionGuard guard(rewriter);
+  AffineMap minOpMap = minOp.getAffineMap();
+
+  // Set up constraint system.
+  FlatAffineConstraints constraints;
+  static const unsigned kDimIv = 0;      // Iteration variable
+  static const unsigned kDimUb = 1;      // Upper loop bound before peeling
+  static const unsigned kDimStep = 2;    // Loop step size
+  static const unsigned kDimMinOp = 3;   // `minOp`
+  static const unsigned kDimMinOpUb = 4; // Upper bound of `minOp`
+  // Add an SSA value as a dimension to the constraint system. If the SSA value
+  // is a constant, a new equality is added, setting the dimension to the
+  // constant value.
+  auto addDim = [&](unsigned dimId, Value value = nullptr) {
+    constraints.addDimId(dimId, value);
+    if (value) {
+      return constraints.addLowerOrUpperBound(
+          dimId, rewriter.getDimIdentityMap(), ValueRange{value},
+          /*eq=*/true, /*lower=*/true, /*composeMapAndOperands=*/false);
+    }
+    return success();
+  };
+
+  if (failed(addDim(kDimIv, iv)) || failed(addDim(kDimUb, ub)) ||
+      failed(addDim(kDimStep, step)) || failed(addDim(kDimMinOp)) ||
+      failed(addDim(kDimMinOpUb)))
+    return failure();
+
+  // Add loop peeling invariant. This is the main piece of knowledge that
+  // enables AffineMinOp simplification.
+  if (insideLoop) {
+    // ub - iv >= step (equiv.: -iv + ub - step + 0*minOp + 0*minOpUb + 0 >= 0)
+    // Intuitively: Inside the peeled loop, every iteration is "full" iteration,
+    // i.e., step divides the iteration space `ub - lb` evenly.
+    constraints.addInequality({-1, 1, -1, 0, 0, 0});
+  } else {
+    // ub - iv < step (equiv.: iv + -ub + step + 0*minOp + 0*minOpUb - 1 >= 0)
+    // Intuitively: `iv` is the split bound here, i.e., the iteration variable
+    // value of the very last iteration (in the unpeeled loop). At that point,
+    // there are less than `step` elements remaining. (Otherwise, the peeled
+    // loop would run for at least one more iteration.)
+    constraints.addInequality({1, -1, 1, 0, 0, -1});
+  }
+
+  // Add an inequality for each result expr_i of minOpMap: minOp <= expr_i
+  if (failed(constraints.addLowerOrUpperBound(
+          kDimMinOp, minOpMap, minOp.operands(), /*eq=*/false, /*lower=*/false,
+          /*composeMapAndOperands=*/false)))
+    return failure();
+
+  // Add a helper dimension r_i for each result expr_i of minOpMap.
+  unsigned resultStartDim = constraints.getNumDimIds();
+  unsigned numResults = minOpMap.getNumResults();
+  for (unsigned i = 0; i < minOpMap.getNumResults(); ++i) {
+    constraints.addDimId(resultStartDim + i);
+    // Add an equality: r_i = expr_i
+    if (failed(constraints.addLowerOrUpperBound(
+            resultStartDim + i, minOpMap.getSubMap({i}), minOp.operands(),
+            /*eq=*/true, /*lower=*/true, /*composeMapAndOperands=*/false)))
+      return failure();
+  }
+
+  // Try to compute an upper bound for minOp, expressed in terms of the other
+  // dimensions iv, ub and step.
+  SmallVector<AffineMap> minOpValLb(1), minOpValUb(1);
+  constraints.getSliceBounds(kDimMinOp, 1, minOp.getContext(), &minOpValLb,
+                             &minOpValUb);
+  if (minOpValUb.size() != 1 || minOpValUb[0].getNumResults() != 1)
+    return failure(); // No upper bound found.
+
+  // Add an equality: kDimMinOpUb = minOpValUb[0]
+  std::vector<SmallVector<int64_t, 8>> flattened;
+  if (failed(getFlattenedAffineExprs(minOpValUb[0], &flattened)))
+    return failure();
+  assert(flattened.size() == 1);
+  SmallVector<int64_t, 8> &eq = flattened[0];
+  eq.insert(eq.begin() + kDimMinOp, 0); // Add back kDimMinOp dimension.
+  assert(eq.size() == constraints.getNumCols());
+  eq[kDimMinOpUb] = -1;
+  constraints.addEquality(eq);
+
+  // Prove that each result of minOpMap has a lower bound that is equal (or
+  // greater than) the upper bound of minOp (`kDimMinOpUb`). In that case,
+  // minOp can be replaced with the bound. I.e., prove that for each result
+  // expr_i (represented by dimension r_i):
+  //
+  // r_i >= minOpUb
+  //
+  // To prove this inequality, add its negation to the constraint set and prove
+  // that the constraint set is empty.
+  for (unsigned i = resultStartDim; i < resultStartDim + numResults; ++i) {
+    FlatAffineConstraints newConstr(constraints);
+    // Add inequality: r_i < minOpUb (equiv.: minOpUb - r_i - 1 >= 0)
+    SmallVector<int64_t> ineq(newConstr.getNumCols(), 0);
+    ineq[kDimMinOpUb] = 1;
+    ineq[i] = -1;
+    ineq[newConstr.getNumCols() - 1] = -1;
+    newConstr.addInequality(ineq);
+    // If the constraint set it empty:
+    // \not \exists r_i: r_i < minOpUb
+    // ==> \forall r_i: r_i >= minOpUb
+    if (!newConstr.isEmpty())
+      return failure();
+  }
+
+  // Lower and upper bound of minOp are equal. Replace minOp with its upper
+  // bound.
+  auto newMap = AffineMap::get(/*dimCount=*/3, /*symbolCount=*/0,
+                               minOpValUb[0].getResult(0));
+  rewriter.setInsertionPoint(minOp);
+  rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, newMap,
+                                             ValueRange{iv, ub, step});
+  return success();
+}
+
+LogicalResult mlir::scf::peelAndCanonicalizeForLoop(RewriterBase &rewriter,
+                                                    ForOp forOp) {
+  Value ub = forOp.upperBound();
+  scf::IfOp ifOp;
+  Value splitBound;
+  if (failed(peelForLoop(rewriter, forOp, ifOp, splitBound)))
+    return failure();
+
+  // Rewrite affine.min ops.
+  forOp.walk([&](AffineMinOp minOp) {
+    (void)rewritePeeledAffineOp(rewriter, minOp, forOp.getInductionVar(), ub,
+                                forOp.step(), /*insideLoop=*/true);
+  });
+  ifOp.walk([&](AffineMinOp minOp) {
+    (void)rewritePeeledAffineOp(rewriter, minOp, splitBound, ub, forOp.step(),
+                                /*insideLoop=*/false);
+  });
+
+  return success();
+}
+
 static constexpr char kPeeledLoopLabel[] = "__peeled_loop__";
 
 namespace {
 struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
   using OpRewritePattern<ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForOp forOp,
                                 PatternRewriter &rewriter) const override {
     if (forOp->hasAttr(kPeeledLoopLabel))
       return failure();
-
-    scf::IfOp ifOp;
-    if (failed(peelForLoop(rewriter, forOp, ifOp)))
+    if (failed(peelAndCanonicalizeForLoop(rewriter, forOp)))
       return failure();
     // Apply label, so that the same loop is not rewritten a second time.
     rewriter.updateRootInPlace(forOp, [&]() {
       forOp->setAttr(kPeeledLoopLabel, rewriter.getUnitAttr());
     });
-
     return success();
   }
 };
 } // namespace
 
 namespace {
 struct ParallelLoopSpecialization
     : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {

 struct ForLoopPeeling : public SCFForLoopPeelingBase<ForLoopPeeling> {
   void runOnFunction() override {
     FuncOp funcOp = getFunction();
     MLIRContext *ctx = funcOp.getContext();
     RewritePatternSet patterns(ctx);
     patterns.add<ForLoopPeelingPattern>(ctx);
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 
     // Drop the marker.
     funcOp.walk([](ForOp op) { op->removeAttr(kPeeledLoopLabel); });
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {
   return std::make_unique<ParallelLoopSpecialization>();
 }
 
 std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() {
   return std::make_unique<ForLoopSpecialization>();
 }
 
 std::unique_ptr<Pass> mlir::createForLoopPeelingPass() {
   return std::make_unique<ForLoopPeeling>();
 }

mlir/test/Dialect/SCF/for-loop-peeling.mlir

 // RUN: mlir-opt %s -for-loop-peeling -canonicalize -split-input-file | FileCheck %s
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - s0) mod s2)>
-//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>
-//  CHECK-DAG: #[[MAP2:.*]] = affine_map<()[s0, s1, s2] -> (s0, s2 - (s2 - (s2 - s1) mod s0))>
+//  CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2))>
 //      CHECK: func @fully_dynamic_bounds(
 // CHECK-SAME:     %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
 //      CHECK:   %[[C0_I32:.*]] = constant 0 : i32
 //      CHECK:   %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[UB]], %[[STEP]]]
 //      CHECK:   %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]]
 // CHECK-SAME:       step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
-//      CHECK:     %[[MINOP:.*]] = affine.min #[[MAP1]](%[[IV]])[%[[STEP]], %[[UB]]]
-//      CHECK:     %[[CAST:.*]] = index_cast %[[MINOP]] : index to i32
+//      CHECK:     %[[CAST:.*]] = index_cast %[[STEP]] : index to i32
 //      CHECK:     %[[ADD:.*]] = addi %[[ACC]], %[[CAST]] : i32
 //      CHECK:     scf.yield %[[ADD]]
 //      CHECK:   }
 //      CHECK:   %[[HAS_MORE:.*]] = cmpi slt, %[[NEW_UB]], %[[UB]]
 //      CHECK:   %[[RESULT:.*]] = scf.if %[[HAS_MORE]] -> (i32) {
-//      CHECK:     %[[REM:.*]] = affine.min #[[MAP2]]()[%[[STEP]], %[[LB]], %[[UB]]]
+//      CHECK:     %[[REM:.*]] = affine.apply #[[MAP1]]()[%[[LB]], %[[UB]], %[[STEP]]]
 //      CHECK:     %[[CAST2:.*]] = index_cast %[[REM]]
 //      CHECK:     %[[ADD2:.*]] = addi %[[LOOP]], %[[CAST2]]
 //      CHECK:     scf.yield %[[ADD2]]
 //      CHECK:   } else {
 //      CHECK:     scf.yield %[[LOOP]]
 //      CHECK:   }
 //      CHECK:   return %[[RESULT]]
 #map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
 func @fully_dynamic_bounds(%lb : index, %ub: index, %step: index) -> i32 {
   %c0 = constant 0 : i32
   %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0) -> i32 {
     %s = affine.min #map(%ub, %iv)[%step]
     %casted = index_cast %s : index to i32
     %0 = addi %arg, %casted : i32
     scf.yield %0 : i32
   }
   return %r : i32
 }
 
 // -----
 
-//  CHECK-DAG: #[[MAP:.*]] = affine_map<(d0) -> (4, -d0 + 17)>
 //      CHECK: func @fully_static_bounds(
 //  CHECK-DAG:   %[[C0_I32:.*]] = constant 0 : i32
 //  CHECK-DAG:   %[[C1_I32:.*]] = constant 1 : i32
+//  CHECK-DAG:   %[[C4_I32:.*]] = constant 4 : i32
 //  CHECK-DAG:   %[[C0:.*]] = constant 0 : index
 //  CHECK-DAG:   %[[C4:.*]] = constant 4 : index
 //  CHECK-DAG:   %[[C16:.*]] = constant 16 : index
 //      CHECK:   %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C16]]
 // CHECK-SAME:       step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
-//      CHECK:     %[[MINOP:.*]] = affine.min #[[MAP]](%[[IV]])
-//      CHECK:     %[[CAST:.*]] = index_cast %[[MINOP]] : index to i32
-//      CHECK:     %[[ADD:.*]] = addi %[[ACC]], %[[CAST]] : i32
+//      CHECK:     %[[ADD:.*]] = addi %[[ACC]], %[[C4_I32]] : i32
 //      CHECK:     scf.yield %[[ADD]]
 //      CHECK:   }
 //      CHECK:   %[[RESULT:.*]] = addi %[[LOOP]], %[[C1_I32]] : i32
 //      CHECK:   return %[[RESULT]]
 #map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
 func @fully_static_bounds() -> i32 {
   %c0_i32 = constant 0 : i32
   %lb = constant 0 : index
   %step = constant 4 : index
   %ub = constant 17 : index
   %r = scf.for %iv = %lb to %ub step %step
                iter_args(%arg = %c0_i32) -> i32 {
     %s = affine.min #map(%ub, %iv)[%step]
     %casted = index_cast %s : index to i32
     %0 = addi %arg, %casted : i32
     scf.yield %0 : i32
   }
   return %r : i32
 }
 
 // -----
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * 4)>
-//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
-//  CHECK-DAG: #[[MAP2:.*]] = affine_map<()[s0] -> (4, s0 mod 4)>
+//  CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 mod 4)>
 //      CHECK: func @dynamic_upper_bound(
 // CHECK-SAME:     %[[UB:.*]]: index
 //  CHECK-DAG:   %[[C0_I32:.*]] = constant 0 : i32
+//  CHECK-DAG:   %[[C4_I32:.*]] = constant 4 : i32
 //  CHECK-DAG:   %[[C0:.*]] = constant 0 : index
 //  CHECK-DAG:   %[[C4:.*]] = constant 4 : index
 //      CHECK:   %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[UB]]]
 //      CHECK:   %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[NEW_UB]]
 // CHECK-SAME:       step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
-//      CHECK:     %[[MINOP:.*]] = affine.min #[[MAP1]](%[[IV]])[%[[UB]]]
-//      CHECK:     %[[CAST:.*]] = index_cast %[[MINOP]] : index to i32
-//      CHECK:     %[[ADD:.*]] = addi %[[ACC]], %[[CAST]] : i32
+//      CHECK:     %[[ADD:.*]] = addi %[[ACC]], %[[C4_I32]] : i32
 //      CHECK:     scf.yield %[[ADD]]
 //      CHECK:   }
 //      CHECK:   %[[HAS_MORE:.*]] = cmpi slt, %[[NEW_UB]], %[[UB]]
 //      CHECK:   %[[RESULT:.*]] = scf.if %[[HAS_MORE]] -> (i32) {
-//      CHECK:     %[[REM:.*]] = affine.min #[[MAP2]]()[%[[UB]]]
+//      CHECK:     %[[REM:.*]] = affine.apply #[[MAP1]]()[%[[UB]]]
 //      CHECK:     %[[CAST2:.*]] = index_cast %[[REM]]
 //      CHECK:     %[[ADD2:.*]] = addi %[[LOOP]], %[[CAST2]]
 //      CHECK:     scf.yield %[[ADD2]]
 //      CHECK:   } else {
 //      CHECK:     scf.yield %[[LOOP]]
 //      CHECK:   }
 //      CHECK:   return %[[RESULT]]
 #map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
%r = scf.for %iv = %lb to %ub step %step
     scf.yield %0 : i32
   }
   return %r : i32
 }
 
 // -----
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * 4)>
-//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
-//  CHECK-DAG: #[[MAP2:.*]] = affine_map<()[s0] -> (4, s0 mod 4)>
+//  CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 mod 4)>
 //      CHECK: func @no_loop_results(
 // CHECK-SAME:     %[[UB:.*]]: index, %[[MEMREF:.*]]: memref<i32>
+//  CHECK-DAG:   %[[C4_I32:.*]] = constant 4 : i32
 //  CHECK-DAG:   %[[C0:.*]] = constant 0 : index
 //  CHECK-DAG:   %[[C4:.*]] = constant 4 : index
 //      CHECK:   %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[UB]]]
 //      CHECK:   scf.for %[[IV:.*]] = %[[C0]] to %[[NEW_UB]] step %[[C4]] {
-//      CHECK:     %[[MINOP:.*]] = affine.min #[[MAP1]](%[[IV]])[%[[UB]]]
 //      CHECK:     %[[LOAD:.*]] = memref.load %[[MEMREF]][]
-//      CHECK:     %[[CAST:.*]] = index_cast %[[MINOP]] : index to i32
-//      CHECK:     %[[ADD:.*]] = addi %[[LOAD]], %[[CAST]] : i32
+//      CHECK:     %[[ADD:.*]] = addi %[[LOAD]], %[[C4_I32]] : i32
 //      CHECK:     memref.store %[[ADD]], %[[MEMREF]]
 //      CHECK:   }
 //      CHECK:   %[[HAS_MORE:.*]] = cmpi slt, %[[NEW_UB]], %[[UB]]
 //      CHECK:   scf.if %[[HAS_MORE]] {
-//      CHECK:     %[[REM:.*]] = affine.min #[[MAP2]]()[%[[UB]]]
+//      CHECK:     %[[REM:.*]] = affine.apply #[[MAP1]]()[%[[UB]]]
 //      CHECK:     %[[LOAD2:.*]] = memref.load %[[MEMREF]][]
 //      CHECK:     %[[CAST2:.*]] = index_cast %[[REM]]
 //      CHECK:     %[[ADD2:.*]] = addi %[[LOAD2]], %[[CAST2]]
 //      CHECK:     memref.store %[[ADD2]], %[[MEMREF]]
 //      CHECK:   }
 //      CHECK:   return
 #map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
 func @no_loop_results(%ub : index, %d : memref<i32>) {
   %c0_i32 = constant 0 : i32
   %lb = constant 0 : index
   %step = constant 4 : index
   scf.for %iv = %lb to %ub step %step {
     %s = affine.min #map(%ub, %iv)[%step]
     %r = memref.load %d[] : memref<i32>
     %casted = index_cast %s : index to i32
     %0 = addi %r, %casted : i32
     memref.store %0, %d[] : memref<i32>
   }
   return
 }
+
+// -----
+
+// Test rewriting of affine.min ops. Make sure that more general cases than
+// the ones above are successfully rewritten. Also make sure that the pattern
+// does not rewrite affine.min ops that should not be rewritten.
+
+//  CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 1)>
+//  CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1 - 1)>
+//  CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0, s1, s2] -> (s0, -d0 + s1, s2)>
+//  CHECK-DAG: #[[MAP4:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2))>
+//  CHECK-DAG: #[[MAP5:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2) + 1)>
+//  CHECK-DAG: #[[MAP6:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2) - 1)>
+//  CHECK-DAG: #[[MAP7:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0, s2 - (s2 - (s2 - s1) mod s0), s3)>
+//      CHECK: func @test_affine_min_rewrite(
+// CHECK-SAME:     %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index,
+// CHECK-SAME:     %[[MEMREF:.*]]: memref<?xindex>, %[[SOME_VAL:.*]]: index
+//      CHECK:   scf.for %[[IV:.*]] = %[[LB]] to %{{.*}} step %[[STEP]] {
+//      CHECK:     %[[RES2:.*]] = affine.apply #[[MAP1]]()[%[[STEP]]]
+//      CHECK:     %[[RES3:.*]] = affine.min #[[MAP2]](%[[IV]])[%[[STEP]], %[[UB]]]
+//      CHECK:     %[[RES4:.*]] = affine.min #map3(%[[IV]])[%[[STEP]], %[[UB]], %[[SOME_VAL]]]
+//      CHECK:     memref.store %[[STEP]]
+//      CHECK:     memref.store %[[STEP]]
+//      CHECK:     memref.store %[[RES2]]
+//      CHECK:     memref.store %[[RES3]]
+//      CHECK:     memref.store %[[RES4]]
+//      CHECK:   }
+//      CHECK:   scf.if {{.*}} {
+//      CHECK:     %[[RES_IF_0:.*]] = affine.apply #[[MAP4]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_1:.*]] = affine.apply #[[MAP5]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_2:.*]] = affine.apply #[[MAP5]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_3:.*]] = affine.apply #[[MAP6]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_4:.*]] = affine.min #[[MAP7]]()[%[[STEP]], %[[LB]], %[[UB]], %[[SOME_VAL]]]
+//      CHECK:     memref.store %[[RES_IF_0]]
+//      CHECK:     memref.store %[[RES_IF_1]]
+//      CHECK:     memref.store %[[RES_IF_2]]
+//      CHECK:     memref.store %[[RES_IF_3]]
+//      CHECK:     memref.store %[[RES_IF_4]]
+#map0 = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
+#map1 = affine_map<(d0, d1)[s0] -> (d0 - d1 + 1, s0)>
+#map2 = affine_map<(d0, d1)[s0] -> (s0 + 1, d0 - d1 + 1)>
+#map3 = affine_map<(d0, d1)[s0] -> (s0, d0 - d1 - 1)>
+#map4 = affine_map<(d0, d1, d2)[s0] -> (s0, d0 - d1, d2)>
+func @test_affine_min_rewrite(%lb : index, %ub: index,
+                              %step: index, %d : memref<?xindex>,
+                              %some_val: index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  %c4 = constant 4 : index
+  scf.for %iv = %lb to %ub step %step {
+    // Most common case: Rewrite min(%ub - %iv, %step) to %step.
+    %m0 = affine.min #map0(%ub, %iv)[%step]
+    // Increase %ub - %iv a little bit, pattern should still apply.
+    %m1 = affine.min #map1(%ub, %iv)[%step]
+    // Rewrite min(%ub - %iv + 1, %step + 1) to %step + 1.
+    %m2 = affine.min #map2(%ub, %iv)[%step]
+    // min(%ub - %iv - 1, %step) cannot be simplified because %ub - %iv - 1
+    // can be smaller than %step. (Can be simplified in if-statement.)
+    %m3 = affine.min #map3(%ub, %iv)[%step]
+    // min(%ub - %iv, %step, %some_val) cannot be simplified because the range
+    // of %some_val is unknown.
+    %m4 = affine.min #map4(%ub, %iv, %some_val)[%step]
+    memref.store %m0, %d[%c0] : memref<?xindex>
+    memref.store %m1, %d[%c1] : memref<?xindex>
+    memref.store %m2, %d[%c2] : memref<?xindex>
+    memref.store %m3, %d[%c3] : memref<?xindex>
+    memref.store %m4, %d[%c4] : memref<?xindex>
+  }
+  return
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     srcs = glob([
         "lib/Dialect/SCF/Transforms/*.cpp",
         "lib/Dialect/SCF/Transforms/*.h",
     ]),
     hdrs = ["include/mlir/Dialect/SCF/Passes.h"],
     includes = ["include"],
     deps = [
         ":Affine",
+        ":Analysis",
         ":DialectUtils",
         ":IR",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SCFPassIncGen",
         ":StandardOps",
         ":Support",