[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

46–47

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/Analysis/AffineStructures.h

public:
 
   /// Add identifiers of the specified kind - specified positions are relative
   /// to the kind of identifier. The coefficient column corresponding to the
   /// added identifier is initialized to zero.
   void addDimId(unsigned pos);
   void addSymbolId(unsigned pos);
   void addLocalId(unsigned pos);
   virtual unsigned addId(IdKind kind, unsigned pos);
+  /// Add identifiers of the specified kind at the end of the table. Return the
+  /// position of the column. The coefficient column corresponding to the
+  /// added identifier is initialized to zero.
+  unsigned addDimId();
+  unsigned addSymbolId();
+  unsigned addLocalId();
 
   /// Composes an affine map whose dimensions and symbols match one to one with
   /// the dimensions and symbols of this FlatAffineConstraints. The results of
   /// the map `other` are added as the leading dimensions of this constraint
   /// system. Returns failure if `other` is a semi-affine map.
   LogicalResult composeMatchingMap(AffineMap other);
 
   /// Projects out (aka eliminates) `num` identifiers starting at position
public:
   /// added identifier is initialized to zero. `val` is the Value corresponding
   /// to the identifier that can optionally be provided.
   void addDimId(unsigned pos, Value val);
   using FlatAffineConstraints::addDimId;
   void addSymbolId(unsigned pos, Value val);
   using FlatAffineConstraints::addSymbolId;
   unsigned addId(IdKind kind, unsigned pos) override;
   unsigned addId(IdKind kind, unsigned pos, Value val);
+  /// Add identifiers of the specified kind at the end of the table. Return the
+  /// position of the column. The coefficient column corresponding to the
+  /// added identifier is initialized to zero. `val` is the Value corresponding
+  /// to the identifier that can optionally be provided.
+  unsigned addDimId(Value val);
+  unsigned addSymbolId(Value val);
 
   /// Add the specified values as a dim or symbol id depending on its nature, if
   /// it already doesn't exist in the system. `val` has to be either a terminal
   /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
   /// symbols or loop IVs. The identifier is added to the end of the existing
   /// dims or symbols. Additional information on the identifier is extracted
   /// from the IR and added to the constraint system.
   void addInductionVarOrTerminalSymbol(Value val);
public:
   inline void getAllValues(SmallVectorImpl<Value> *values) const {
     getValues(0, numIds, values);
   }
 
   inline ArrayRef<Optional<Value>> getMaybeValues() const {
     return {values.data(), values.size()};
   }
 
+  inline ArrayRef<Optional<Value>> getMaybeDimValues() const {
+    return {values.data(), getNumDimIds()};
+  }
+
+  inline ArrayRef<Optional<Value>> getMaybeSymbolValues() const {
+    return {values.data() + getNumDimIds(), getNumSymbolIds()};
+  }
+
+  inline ArrayRef<Optional<Value>> getMaybeDimAndSymbolValues() const {
+    return {values.data(), getNumDimIds() + getNumSymbolIds()};
+  }
+
   /// Sets the Value associated with the pos^th identifier.
   inline void setValue(unsigned pos, Value val) {
     assert(pos < numIds && "invalid id position");
     values[pos] = val;
   }
 
   /// Sets the Values associated with the identifiers in the range [start, end).
   void setValues(unsigned start, unsigned end, ArrayRef<Value> values) {

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

 /// 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
-/// as vectorization or loop tiling.
+/// This transformation is beneficial for a wide range of transformations such
+/// as vectorization or loop tiling: It enables additional canonicalizations
+/// inside the peeled loop body such as rewriting masked loads into unmaked
+/// loads.
 ///
 /// 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.
+///
+/// The return value indicates whether the loop was rewritten or not. Loops are
+/// not rewritten if:
+/// * Loop step size is 1 or
+/// * Loop bounds and step size are static, and step already divides the
+///   iteration space evenly.
+///
+/// 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/include/mlir/IR/AffineExpr.h

public:
   /// modified expression tree.
   AffineExpr replace(AffineExpr expr, AffineExpr replacement) const;
 
   /// Sparse replace method. If `*this` appears in `map` replaces it by
   /// `map[*this]` and return the modified expression tree. Otherwise traverse
   /// `*this` and apply replace with `map` on its subexpressions.
   AffineExpr replace(const DenseMap<AffineExpr, AffineExpr> &map) const;
 
-  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
-  AffineExpr shiftDims(unsigned numDims, unsigned shift) const;
-
-  /// Replace symbols[0 .. numSymbols - 1] by
-  ///         symbols[shift .. shift + numSymbols - 1].
-  AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift) const;
+  /// Replace dims[offset ... numDims)
+  /// by dims[offset + shift ... shift + numDims).
+  AffineExpr shiftDims(unsigned numDims, unsigned shift,
+                       unsigned offset = 0) const;
+
+  /// Replace symbols[offset ... numSymbols)
+  /// by symbols[offset + shift ... shift + numSymbols).
+  AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift,
+                          unsigned offset = 0) const;
 
   AffineExpr operator+(int64_t v) const;
   AffineExpr operator+(AffineExpr other) const;
   AffineExpr operator-() const;
   AffineExpr operator-(int64_t v) const;
   AffineExpr operator-(AffineExpr other) const;
   AffineExpr operator*(int64_t v) const;
   AffineExpr operator*(AffineExpr other) const;

mlir/include/mlir/IR/AffineMap.h

public:
   AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map) const;
 
   /// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
   /// results and return a new AffineMap with the new results and with the
   /// specified number of dims and symbols.
   AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map,
                     unsigned numResultDims, unsigned numResultSyms) const;
 
-  /// Replace dims[0 .. numDims - 1] by dims[shift .. shift + numDims - 1].
-  AffineMap shiftDims(unsigned shift) const {
-    return AffineMap::get(
-        getNumDims() + shift, getNumSymbols(),
+  /// Replace dims[offset ... numDims)
+  /// by dims[offset + shift ... shift + numDims).
+  AffineMap shiftDims(unsigned shift, unsigned offset = 0) const {
+    assert(offset <= getNumDims());
+    return AffineMap::get(getNumDims() + shift, getNumSymbols(),
                           llvm::to_vector<4>(llvm::map_range(
                               getResults(),
-            [&](AffineExpr e) { return e.shiftDims(getNumDims(), shift); })),
+                              [&](AffineExpr e) {
+                                return e.shiftDims(getNumDims(), shift, offset);
+                              })),
                           getContext());
   }
 
-  /// Replace symbols[0 .. numSymbols - 1] by
-  ///         symbols[shift .. shift + numSymbols - 1].
-  AffineMap shiftSymbols(unsigned shift) const {
+  /// Replace symbols[offset ... numSymbols)
+  /// by symbols[offset + shift ... shift + numSymbols).
+  AffineMap shiftSymbols(unsigned shift, unsigned offset = 0) const {
     return AffineMap::get(getNumDims(), getNumSymbols() + shift,
                           llvm::to_vector<4>(llvm::map_range(
                               getResults(),
                               [&](AffineExpr e) {
-                                return e.shiftSymbols(getNumSymbols(), shift);
+                                return e.shiftSymbols(getNumSymbols(), shift,
+                                                      offset);
                               })),
                           getContext());
   }
 
   /// Folds the results of the application of an affine map on the provided
   /// operands to a constant if possible.
   LogicalResult constantFold(ArrayRef<Attribute> operandConstants,
                              SmallVectorImpl<Attribute> &results) const;

mlir/lib/Analysis/AffineStructures.cpp

for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
     addEquality(other.getEquality(r));
   }
 }
 
 void FlatAffineConstraints::addLocalId(unsigned pos) {
   addId(IdKind::Local, pos);
 }
 
+unsigned FlatAffineConstraints::addLocalId() {
+  unsigned pos = getNumLocalIds();
+  addId(IdKind::Local, pos);
+  return pos;
+}
+
 void FlatAffineConstraints::addDimId(unsigned pos) {
   addId(IdKind::Dimension, pos);
 }
 
+unsigned FlatAffineConstraints::addDimId() {
+  unsigned pos = getNumDimIds();
+  addId(IdKind::Dimension, pos);
+  return pos;
+}
+
 void FlatAffineValueConstraints::addDimId(unsigned pos, Value val) {
   addId(IdKind::Dimension, pos, val);
 }
 
+unsigned FlatAffineValueConstraints::addDimId(Value val) {
+  unsigned pos = getNumDimIds();
+  addId(IdKind::Dimension, pos, val);
+  return pos;
+}
+
 void FlatAffineConstraints::addSymbolId(unsigned pos) {
   addId(IdKind::Symbol, pos);
 }
 
+unsigned FlatAffineConstraints::addSymbolId() {
+  unsigned pos = getNumSymbolIds();
+  addId(IdKind::Symbol, pos);
+  return pos;
+}
+
 void FlatAffineValueConstraints::addSymbolId(unsigned pos, Value val) {
   addId(IdKind::Symbol, pos, val);
 }
 
+unsigned FlatAffineValueConstraints::addSymbolId(Value val) {
+  unsigned pos = getNumSymbolIds();
+  addId(IdKind::Symbol, pos, val);
+  return pos;
+}
+
 unsigned FlatAffineConstraints::addId(IdKind kind, unsigned pos) {
   if (kind == IdKind::Dimension)
     assert(pos <= getNumDimIds());
   else if (kind == IdKind::Symbol)
     assert(pos <= getNumSymbolIds());
   else
     assert(pos <= getNumLocalIds());
 

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 // the following conditions:
 // *) It is a constant.
 // *) Its defining op or block arg appearance is immediately enclosed by an op
 //    with `AffineScope` trait.
 // *) It is the result of an affine.apply operation with symbol operands.
 // *) It is a result of the dim op on a memref whose corresponding size is a
 //    valid symbol.
 bool mlir::isValidSymbol(Value value) {
+  if (!value)
+    return false;
+
   // The value must be an index type.
   if (!value.getType().isIndex())
     return false;
 
   // Check that the value is a top level value.
   if (isTopLevelValue(value))
     return true;
 

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`.
+/// The return value indicates whether the loop was rewritten or not.
+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);
static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp, scf::IfOp &ifOp,
   // Build else case.
   if (!resultTypes.empty())
     ifOp.getElseBodyBuilder(b.getListener())
         .create<scf::YieldOp>(loc, forOp->getResults());
 
   return success();
 }
 
+static void unpackOptionalValues(ArrayRef<Optional<Value>> source,
+                                 SmallVector<Value> &target) {
+  target = llvm::to_vector<4>(llvm::map_range(source, [](Optional<Value> val) {
+    return val.hasValue() ? *val : Value();
+  }));
+}
+
+/// Bound an identifier `pos` in a given FlatAffineValueConstraints with
+/// constraints drawn from an affine map. Before adding the constraint, the
+/// dimensions/symbols of the affine map are aligned with `constraints`.
+/// `operands` are the SSA Value operands used with the affine map.
+/// Note: This function adds a new symbol column to the `constraints` for each
+/// dimension/symbol that exists in the affine map but not in `constraints`.
+static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints,
+                                      FlatAffineConstraints::BoundType type,
+                                      unsigned pos, AffineMap map,
+                                      ValueRange operands) {
+  SmallVector<Value> dims, syms, newSyms;
+  unpackOptionalValues(constraints.getMaybeDimValues(), dims);
+  unpackOptionalValues(constraints.getMaybeSymbolValues(), syms);
+
+  AffineMap alignedMap =
+      alignAffineMapWithValues(map, operands, dims, syms, &newSyms);
+  for (unsigned i = syms.size(); i < newSyms.size(); ++i)
+    constraints.addSymbolId(constraints.getNumSymbolIds(), newSyms[i]);
+  return constraints.addBound(type, pos, alignedMap);
+}
+
+/// This function tries to canonicalize affine.min operations by proving that
+/// its value is bounded by the same lower and upper bound. In that case, the
+/// operation can be folded away.
+///
+/// Bounds are computed by FlatAffineValueConstraints. Invariants required for
+/// finding/proving bounds should be supplied via `constraints`.
+///
+/// 1. Add dimensions for `minOp` and `minOpUb` (upper bound of `minOp`).
+/// 2. Compute an upper bound of `minOp` and bind it to `minOpUb`. SSA values
+///    that are used in `minOp` but are not part of `dims`, are added as extra
+///    symbols to the constraint set.
+/// 3. For each result of `minOp`: Add result as a dimension `r_i`. Prove that
+///    r_i >= minOpUb. If this is the case, ub(minOp) == lb(minOp) and `minOp`
+///    can be replaced with that bound.
+///
+/// In summary, the following constraints are added throughout this function.
+/// Note: `invar` are dimensions added by the caller to express the invariants.
+///
+///  invar | minOp | minOpUb | r_i | extra syms... | const |           eq/ineq
+///  ------+-------+---------+-----+---------------+-------+-------------------
+///   (various eq./ineq. constraining `invar`, added by the caller)
+///    ... |     0 |       0 |   0 |             0 |   ... |               ...
+///  ------+-------+---------+-----+---------------+-------+-------------------
+///   (various ineq. constraining `minOp` in terms of `minOp` operands (`invar`
+///    and extra `minOp` operands "extra syms" that are not in `invar`)).
+///    ... |    -1 |       0 |   0 |           ... |   ... |              >= 0
+///  ------+-------+---------+-----+---------------+-------+-------------------
+///   (set `minOpUb` to `minOp` upper bound in terms of `invar` and extra syms)
+///    ... |     0 |      -1 |   0 |           ... |   ... |               = 0
+///  ------+-------+---------+-----+---------------+-------+-------------------
+///   (for each `minOp` map result r_i: copy previous constraints, set r_i to
+///    corresponding map result, prove r_i >= minOpUb via contradiction)
+///    ... |     0 |       0 |  -1 |           ... |   ... |               = 0
+///      0 |     0 |       1 |  -1 |             0 |    -1 |              >= 0
+///
+static LogicalResult
+canonicalizeAffineMinOp(RewriterBase &rewriter, AffineMinOp minOp,
+                        FlatAffineValueConstraints constraints) {
+  RewriterBase::InsertionGuard guard(rewriter);
+  AffineMap minOpMap = minOp.getAffineMap();
+  unsigned numResults = minOpMap.getNumResults();
+
+  // Add a few extra dimensions.
+  unsigned dimMinOp = constraints.addDimId();   // `minOp`
+  unsigned dimMinOpUb = constraints.addDimId(); // `minOp` upper bound
+  unsigned resultDimStart = constraints.getNumDimIds();
+  for (unsigned i = 0; i < numResults; ++i)
+    constraints.addDimId();
+
+  // Add an inequality for each result expr_i of minOpMap: minOp <= expr_i
+  if (failed(alignAndAddBound(constraints, FlatAffineConstraints::UB, dimMinOp,
+                              minOpMap, minOp.operands())))
+    return failure();
+
+  // Try to compute an upper bound for minOp, expressed in terms of the other
+  // `dims` and extra symbols.
+  SmallVector<AffineMap> minOpValLb(1), minOpValUb(1);
+  constraints.getSliceBounds(dimMinOp, 1, minOp.getContext(), &minOpValLb,
+                             &minOpValUb);
+  // TODO: `getSliceBounds` may return multiple bounds at the moment. This is
+  // a TODO of `getSliceBounds` and not handled here.
+  if (!minOpValUb[0] || minOpValUb[0].getNumResults() != 1)
+    return failure(); // No or multiple upper bounds found.
+
+  // Add an equality: dimMinOpUb = minOpValUb[0]
+  // Add back dimension for minOp. (Was removed by `getSliceBounds`.)
+  AffineMap alignedUbMap = minOpValUb[0].shiftDims(/*shift=*/1,
+                                                   /*offset=*/dimMinOp);
+  if (failed(constraints.addBound(FlatAffineConstraints::EQ, dimMinOpUb,
+                                  alignedUbMap)))
+    return failure();
+
+  // If the constraint system is empty, there is an inconsistency. (E.g., this
+  // can happen if loop lb > ub.)
+  if (constraints.isEmpty())
+    return failure();
+
+  // Prove that each result of minOpMap has a lower bound that is equal to (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 = resultDimStart; i < resultDimStart + numResults; ++i) {
+    FlatAffineValueConstraints newConstr(constraints);
+
+    // Add an equality: r_i = expr_i
+    // Note: These equalities could have been added earlier and used to express
+    // minOp <= expr_i. However, then we run the risk that `getSliceBounds`
+    // computes minOpUb in terms of r_i dims, which is not desired.
+    if (failed(alignAndAddBound(newConstr, FlatAffineConstraints::EQ, i,
+                                minOpMap.getSubMap({i - resultDimStart}),
+                                minOp.operands())))
+      return failure();
+
+    // Add inequality: r_i < minOpUb (equiv.: minOpUb - r_i - 1 >= 0)
+    SmallVector<int64_t> ineq(newConstr.getNumCols(), 0);
+    ineq[dimMinOpUb] = 1;
+    ineq[i] = -1;
+    ineq[newConstr.getNumCols() - 1] = -1;
+    newConstr.addInequality(ineq);
+    if (!newConstr.isEmpty())
+      return failure();
+  }
+
+  // Lower and upper bound of `minOp` are equal. Replace `minOp` with its bound.
+  AffineMap newMap = alignedUbMap;
+  SmallVector<Value> newOperands;
+  unpackOptionalValues(constraints.getMaybeDimAndSymbolValues(), newOperands);
+  mlir::canonicalizeMapAndOperands(&newMap, &newOperands);
+  rewriter.setInsertionPoint(minOp);
+  rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, newMap, newOperands);
+  return success();
+}
+
+/// Try to simplify an affine.min operation `minOp` 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:
+/// * Inside the peeled loop: min(step, ub - iv) == step
+/// * Inside the scf.if operation: min(step, ub - iv) == ub - iv
+///
+/// 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.
+static LogicalResult rewritePeeledAffineOp(RewriterBase &rewriter,
+                                           AffineMinOp minOp, Value iv,
+                                           Value ub, Value step,
+                                           bool insideLoop) {
+  FlatAffineValueConstraints constraints;
+  constraints.addDimId(0, iv);
+  constraints.addDimId(1, ub);
+  constraints.addDimId(2, step);
+  if (auto constUb = getConstantIntValue(ub))
+    constraints.addBound(FlatAffineConstraints::EQ, 1, *constUb);
+  if (auto constStep = getConstantIntValue(step))
+    constraints.addBound(FlatAffineConstraints::EQ, 2, *constStep);
+
+  // 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.
+    constraints.addInequality({-1, 1, -1, 0});
+  } else {
+    // 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.)
+    constraints.addInequality({1, -1, 1, -1});
+  }
+
+  return canonicalizeAffineMinOp(rewriter, minOp, constraints);
+}
+
+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> {

mlir/lib/IR/AffineExpr.cpp

AffineExpr AffineExpr::replaceDims(ArrayRef<AffineExpr> dimReplacements) const {
   return replaceDimsAndSymbols(dimReplacements, {});
 }
 
 AffineExpr
 AffineExpr::replaceSymbols(ArrayRef<AffineExpr> symReplacements) const {
   return replaceDimsAndSymbols({}, symReplacements);
 }
 
-/// Replace symbols[0 .. numDims - 1] by symbols[shift .. shift + numDims - 1].
-AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift) const {
+/// Replace dims[offset ... numDims)
+/// by dims[offset + shift ... shift + numDims).
+AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift,
+                                 unsigned offset) const {
   SmallVector<AffineExpr, 4> dims;
-  for (unsigned idx = 0; idx < numDims; ++idx)
+  for (unsigned idx = 0; idx < offset; ++idx)
+    dims.push_back(getAffineDimExpr(idx, getContext()));
+  for (unsigned idx = offset; idx < numDims; ++idx)
     dims.push_back(getAffineDimExpr(idx + shift, getContext()));
   return replaceDimsAndSymbols(dims, {});
 }
 
-/// Replace symbols[0 .. numSymbols - 1] by
-/// symbols[shift .. shift + numSymbols - 1].
-AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift) const {
+/// Replace symbols[offset ... numSymbols)
+/// by symbols[offset + shift ... shift + numSymbols).
+AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift,
+                                    unsigned offset) const {
   SmallVector<AffineExpr, 4> symbols;
-  for (unsigned idx = 0; idx < numSymbols; ++idx)
+  for (unsigned idx = 0; idx < offset; ++idx)
+    symbols.push_back(getAffineSymbolExpr(idx, getContext()));
+  for (unsigned idx = offset; idx < numSymbols; ++idx)
     symbols.push_back(getAffineSymbolExpr(idx + shift, getContext()));
   return replaceDimsAndSymbols({}, symbols);
 }
 
 /// Sparse replace method. Return the modified expression tree.
 AffineExpr
 AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
   auto it = map.find(*this);

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] -> (-(s0 - (s0 - s1) mod s2) + s0)>
 //      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]]()[%[[UB]], %[[LB]], %[[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] -> (-(s0 - (s0 - s1) mod s2) + s0)>
+//  CHECK-DAG: #[[MAP5:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0 + 1)>
+//  CHECK-DAG: #[[MAP6:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0 - 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]] {
+//                 (affine.min folded away)
+//      CHECK:     memref.store %[[STEP]]
+//                 (affine.min folded away)
+//      CHECK:     memref.store %[[STEP]]
+//      CHECK:     %[[RES2:.*]] = affine.apply #[[MAP1]]()[%[[STEP]]]
+//      CHECK:     memref.store %[[RES2]]
+//      CHECK:     %[[RES3:.*]] = affine.min #[[MAP2]](%[[IV]])[%[[STEP]], %[[UB]]]
+//      CHECK:     memref.store %[[RES3]]
+//      CHECK:     %[[RES4:.*]] = affine.min #[[MAP3]](%[[IV]])[%[[STEP]], %[[UB]], %[[SOME_VAL]]]
+//      CHECK:     memref.store %[[RES4]]
+//      CHECK:   }
+//      CHECK:   scf.if {{.*}} {
+//      CHECK:     %[[RES_IF_0:.*]] = affine.apply #[[MAP4]]()[%[[UB]], %[[LB]], %[[STEP]]]
+//      CHECK:     memref.store %[[RES_IF_0]]
+//      CHECK:     %[[RES_IF_1:.*]] = affine.apply #[[MAP5]]()[%[[UB]], %[[LB]], %[[STEP]]]
+//      CHECK:     memref.store %[[RES_IF_1]]
+//      CHECK:     %[[RES_IF_2:.*]] = affine.apply #[[MAP5]]()[%[[UB]], %[[LB]], %[[STEP]]]
+//      CHECK:     memref.store %[[RES_IF_2]]
+//      CHECK:     %[[RES_IF_3:.*]] = affine.apply #[[MAP6]]()[%[[UB]], %[[LB]], %[[STEP]]]
+//      CHECK:     memref.store %[[RES_IF_3]]
+//      CHECK:     %[[RES_IF_4:.*]] = affine.min #[[MAP7]]()[%[[STEP]], %[[LB]], %[[UB]], %[[SOME_VAL]]]
+//      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]
+    memref.store %m0, %d[%c0] : memref<?xindex>
+
+    // Increase %ub - %iv a little bit, pattern should still apply.
+    %m1 = affine.min #map1(%ub, %iv)[%step]
+    memref.store %m1, %d[%c1] : memref<?xindex>
+
+    // Rewrite min(%ub - %iv + 1, %step + 1) to %step + 1.
+    %m2 = affine.min #map2(%ub, %iv)[%step]
+    memref.store %m2, %d[%c2] : memref<?xindex>
+
+    // 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]
+    memref.store %m3, %d[%c3] : memref<?xindex>
+
+    // 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 %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",