[mlir][SCF] Generalize AffineMinSCFCanonicalization to min/max ops

[springerm]

Summary

• Add support for affine.max ops to SCF loop peeling pattern.
• Add support for affine.max ops to AffineMinSCFCanonicalizationPattern.
• Rename AffineMinSCFCanonicalizationPattern to AffineOpSCFCanonicalizationPattern.
• Rename AffineMinSCFCanonicalization pass to SCFAffineOpCanonicalization.

Depends On D107732

[springerm]

rebase

[springerm]

rebase

[ftynse]

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

52–56

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

207

319

Nit: update minOpUb to opBound here? Maybe also spell out the max case?

mlir/test/Dialect/SCF/canonicalize-affine-op.mlir

32

Could we also have a test where the canonicalization does not kick in, preferably an adjacent one, e.g., using %c3.

[springerm]

address comments

mlir/include/mlir/Dialect/SCF/Passes.h

 /// Creates a pass that specializes for loop for unrolling and
 /// vectorization.
 std::unique_ptr<Pass> createForLoopSpecializationPass();
 
 /// Creates a pass that peels for loops at their upper bounds for
 /// better vectorization.
 std::unique_ptr<Pass> createForLoopPeelingPass();
 
-std::unique_ptr<Pass> createAffineMinSCFCanonicalizationPass();
+/// Creates a pass that canonicalizes affine.min and affine.max operations
+/// inside of scf.for loops.
+std::unique_ptr<Pass> createSCFAffineOpCanonicalizationPass();
 
 /// Creates a loop fusion pass which fuses parallel loops.
 std::unique_ptr<Pass> createParallelLoopFusionPass();
 
 /// Creates a pass that specializes parallel loop for unrolling and
 /// vectorization.
 std::unique_ptr<Pass> createParallelLoopSpecializationPass();
 

mlir/include/mlir/Dialect/SCF/Passes.td

 include "mlir/Pass/PassBase.td"
 
 def SCFBufferize : FunctionPass<"scf-bufferize"> {
   let summary = "Bufferize the scf dialect.";
   let constructor = "mlir::createSCFBufferizePass()";
   let dependentDialects = ["memref::MemRefDialect"];
 }
 
-def AffineMinSCFCanonicalization
-    : FunctionPass<"canonicalize-scf-affine-min"> {
-  let summary = "Canonicalize affine.min ops in the context of SCF loops with "
-                "known bounds";
-  let constructor = "mlir::createAffineMinSCFCanonicalizationPass()";
+def SCFAffineOpCanonicalization
+    : FunctionPass<"canonicalize-scf-affine-op"> {
+  let summary = "Canonicalize affine.min and affine.max ops in the context of "
+                "SCF loops with known bounds";
+  let constructor = "mlir::createSCFAffineOpCanonicalizationPass()";
   let dependentDialects = ["AffineDialect"];
 }
 
 def SCFForLoopPeeling
     : FunctionPass<"for-loop-peeling"> {
   let summary = "Peel `for` loops at their upper bounds.";
   let constructor = "mlir::createForLoopPeelingPass()";
   let dependentDialects = ["AffineDialect"];

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

 #ifndef MLIR_DIALECT_SCF_TRANSFORMS_H_
 #define MLIR_DIALECT_SCF_TRANSFORMS_H_
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 
-class AffineMinOp;
+class AffineMap;
 class ConversionTarget;
 struct LogicalResult;
 class MLIRContext;
 class Region;
 class RewriterBase;
 class TypeConverter;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 class Operation;
 class Value;
+class ValueRange;
 
 namespace scf {
 
 class IfOp;
 class ForOp;
 class ParallelOp;
 class ForOp;
 
-/// Try to canonicalize an affine.min operation in the context of for loops with
+/// Match "for loop"-like operations: If the first parameter is an iteration
+/// variable, return lower/upper bounds via the second/third parameter and the
+/// step size via the last parameter. The function should return `success` in
+/// that case. If the first parameter is not an iteration variable, return
+/// `failure`.
+using LoopMatcherFn =
+    function_ref<LogicalResult(Value, Value &, Value &, Value &)>;
+
+/// Try to canonicalize an min/max operations in the context of for loops with
 /// a known range.
 ///
-/// `loopMatcher` is used to retrieve loop bounds and step size for a given
-/// iteration variable: If the first parameter is an iteration variable, return
-/// lower/upper bounds via the second/third parameter and the step size via the
-/// last parameter. The function should return `success` in that case. If the
-/// first parameter is not an iteration variable, return `failure`.
+/// `map is the body of the min/max operation and `operands` are the SSA values
+/// that the dimensions and symbols are bound to; dimensions are listed first.
+/// If `isMin`, the operation is a min operation; otherwise, a max operation.
+/// `loopMatcher` is used to retrieve loop bounds and the step size for a given
+/// iteration variable.
 ///
 /// Note: `loopMatcher` allows this function to be used with any "for loop"-like
 /// operation (scf.for, scf.parallel and even ops defined in other dialects).
-LogicalResult canonicalizeAffineMinOpInLoop(
-    AffineMinOp minOp, RewriterBase &rewriter,
-    function_ref<LogicalResult(Value, Value &, Value &, Value &)> loopMatcher);
+LogicalResult canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,
+                                         AffineMap map, ValueRange operands,
+                                         bool isMin, LoopMatcherFn loopMatcher);
 
 /// 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

 ///   (loop body)
 /// }
 /// scf.if %newUb < %ub {
 ///   (loop body)
 /// }
 /// ```
 ///
 /// 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.
+/// and affine.max ops in the body of the loop and the scf.if, taking advantage
+/// of the fact that the peeled loop has only "full" iterations and the scf.if
+/// is always a partial iteration (if any). 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

 ///  S0(I+2) S1(I+1) S2(I)       // Pipelined kernel
 /// }
 /// S1(N) S2(N-1)                // Epilogue
 /// S2(N)                        // Epilogue
 void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns,
                                        const PipeliningOption &options);
 
 /// Populate patterns for canonicalizing operations inside SCF loop bodies.
-/// At the moment, only affine.min computations with iteration variables,
+/// At the moment, only affine.min/max computations with iteration variables,
 /// loop bounds and loop steps are canonicalized.
 void populateSCFLoopBodyCanonicalizationPatterns(RewritePatternSet &patterns);
 
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_TRANSFORMS_H_

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

static LogicalResult alignAndAddBound(FlatAffineConstraints &constraints,
   AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, &newSyms);
   for (unsigned i = 0; i < newSyms.size() - syms.size(); ++i)
     constraints.addSymbolId(constraints.getNumSymbolIds());
   std::swap(syms, newSyms);
   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
+/// This function tries to canonicalize min/max operations by proving that their
+/// value is bounded by the same lower and upper bound. In that case, the
 /// operation can be replaced with the bound.
 ///
 /// Bounds are computed by FlatAffineConstraints. Invariants required for
 /// finding/proving bounds should be supplied via `constraintBuilder` by adding
 /// constraints to the provided FlatAffineConstraints. Only the dimensions in
 /// `dims` and constants can be used when adding constraints. Adding new
 /// dimensions/symbols is not allowed. (However, local columns may be added.)
 ///
 /// 1. Set up a constraint system with the dimensions passed as `dims`.
 /// 2. Call the builder, which may add new constraints.
-/// 3. Add dimensions for `minOp` and `minOpUb` (upper bound of `minOp`).
-/// 4. 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.
-/// 5. 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.
+/// 3. Add dimensions for `op` and `opBound` (lower or upper bound of `op`).
+/// 4. Compute an upper bound of `op` (in case of `isMin`) or a lower bound (in
+///    case of `!isMin`) and bind it to `opBound`. SSA values that are used in
+///    `op` but are not part of `dims`, are added as extra symbols.
+/// 5. For each result of `op`: Add result as a dimension `r_i`. Prove that:
+///    * If `isMin`: r_i >= opBound
+///    * If `isMax`: r_i <= opBound
+///    If this is the case, ub(op) == lb(op).
+/// 6. Replace `op` can be replaced with `opBound`.
 ///
 /// In summary, the following constraints are added throughout this function:
+/// (Showing only the case where `isMin`.)
 ///
-///   dims | minOp | minOpUb | r_i | extra syms... | const |           eq/ineq
+///   dims |    op | opBound | r_i | extra syms... | const |           eq/ineq
 ///  ------+-------+---------+-----+---------------+-------+-------------------
 ///   (various eq./ineq. constraining `dims`, added by `constraintBuilder`)
 ///    ... |     0 |       0 |   0 |             0 |   ... |               ...
 ///  ------+-------+---------+-----+---------------+-------+-------------------
-///   (various ineq. constraining `minOp` in terms of `minOp` operands (`dims`
-///    and extra `minOp` operands "extra syms" that are not in `dims`)).
+///   (various ineq. constraining `op` in terms of `op` operands (`dims` and
+///    extra `op` operands "extra syms" that are not in `dims`)).
 ///    ... |    -1 |       0 |   0 |           ... |   ... |              >= 0
 ///  ------+-------+---------+-----+---------------+-------+-------------------
-///   (set `minOpUb` to `minOp` upper bound in terms of `dims` and "extra syms")
+///   (set `opBound` to `op` upper bound in terms of `dims` 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)
+///   (for each `op` map result r_i: set r_i to corresponding map result,
+///    prove that r_i >= minOpUb via contradiction)
 ///    ... |     0 |       0 |  -1 |           ... |   ... |               = 0
 ///      0 |     0 |       1 |  -1 |             0 |    -1 |              >= 0
 ///
-static LogicalResult canonicalizeAffineMinOp(
-    RewriterBase &rewriter, AffineMinOp minOp, ValueRange dims,
+static LogicalResult canonicalizeMinMaxOp(
+    RewriterBase &rewriter, Operation *op, AffineMap map, ValueRange operands,
+    bool isMin, ValueRange dims,
     function_ref<LogicalResult(FlatAffineConstraints &)> constraintsBuilder) {
   RewriterBase::InsertionGuard guard(rewriter);
-  AffineMap minOpMap = minOp.getAffineMap();
-  unsigned numResults = minOpMap.getNumResults();
+  unsigned numResults = map.getNumResults();
   FlatAffineConstraints constraints;
   // Keep track of SSA values of dims (if any), so that affine maps can be
   // aligned with the dims in `constraints`.
   SmallVector<Value> dimValues, symValues;
 
   /// Add an SSA value as a dimension to the constraint system. If the SSA value
   /// is a constant, set the dimension to the constant value.
   auto addDim = [&](Value value = {}) {
     unsigned dimId = constraints.getNumDimIds();
     constraints.addDimId(dimId);
     if (auto constInt = getConstantIntValue(value))
       constraints.addBound(FlatAffineConstraints::EQ, dimId, *constInt);
     dimValues.push_back(value);
     return dimId;
   };
 
   // Set up constraint system and call builder.
   for (Value value : dims)
     addDim(value);
   if (failed(constraintsBuilder(constraints)))
     return failure();
 
   // Add extra dimensions.
-  unsigned dimMinOp = addDim();   // `minOp`
-  unsigned dimMinOpUb = addDim(); // `minOp` upper bound
+  unsigned dimOp = addDim();      // `op`
+  unsigned dimOpBound = addDim(); // `op` lower/upper bound
   unsigned resultDimStart = constraints.getNumDimIds();
   for (unsigned i = 0; i < numResults; ++i)
     addDim();
 
-  // Add an inequality for each result expr_i of minOpMap: minOp <= expr_i
-  if (failed(alignAndAddBound(constraints, FlatAffineConstraints::UB, dimMinOp,
-                              minOpMap, minOp.operands(), dimValues,
-                              symValues)))
+  // Add an inequality for each result expr_i of map:
+  // isMin: op <= expr_i, !isMin: op >= expr_i
+  auto boundType =
+      isMin ? FlatAffineConstraints::UB : FlatAffineConstraints::LB;
+  if (failed(alignAndAddBound(constraints, boundType, dimOp, map, operands,
+                              dimValues, symValues)))
     return failure();
 
-  // Try to compute an upper bound for minOp, expressed in terms of the other
+  // Try to compute a lower/upper bound for op, expressed in terms of the other
   // `dims` and extra symbols.
-  SmallVector<AffineMap> minOpValLb(1), minOpValUb(1);
-  constraints.getSliceBounds(dimMinOp, 1, minOp.getContext(), &minOpValLb,
-                             &minOpValUb);
+  SmallVector<AffineMap> opLb(1), opUb(1);
+  constraints.getSliceBounds(dimOp, 1, rewriter.getContext(), &opLb, &opUb);
+  AffineMap boundMap = isMin ? opUb[0] : opLb[0];
   // 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.
+  if (!boundMap || boundMap.getNumResults() != 1)
+    return failure(); // No or multiple 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)))
+  // Add an equality: Set dimOpBound to computed bound.
+  // Add back dimension for op. (Was removed by `getSliceBounds`.)
+  AffineMap alignedBoundMap = boundMap.shiftDims(/*shift=*/1, /*offset=*/dimOp);
+  if (failed(constraints.addBound(FlatAffineConstraints::EQ, dimOpBound,
+                                  alignedBoundMap)))
     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
+  // In the case of `isMin` (`!isMin` is inversed):
+  // Prove that each result of `map` has a lower bound that is equal to (or
+  // greater than) the upper bound of `op` (`dimOpBound`). In that case, `op`
+  // can be replaced with the bound. I.e., prove that for each result
   // expr_i (represented by dimension r_i):
   //
-  // r_i >= minOpUb
+  // r_i >= opBound
   //
   // To prove this inequality, add its negation to the constraint set and prove
   // that the constraint set is empty.
   unsigned symValuesSize = symValues.size();
   (void)symValuesSize;
   for (unsigned i = resultDimStart; i < resultDimStart + numResults; ++i) {
     assert(!constraints.isEmpty() && "constraint set should not be empty yet");
 
     // 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(constraints, FlatAffineConstraints::EQ, i,
-                                minOpMap.getSubMap({i - resultDimStart}),
-                                minOp.operands(), dimValues, symValues)))
+                                map.getSubMap({i - resultDimStart}), operands,
+                                dimValues, symValues)))
       return failure();
     assert(symValues.size() == symValuesSize && "no new symbols expected");
 
     // Add inequality: r_i < minOpUb (equiv.: minOpUb - r_i - 1 >= 0)
     SmallVector<int64_t> ineq(constraints.getNumCols(), 0);
-    ineq[dimMinOpUb] = 1;
-    ineq[i] = -1;
+    ineq[dimOpBound] = isMin ? 1 : -1;
+    ineq[i] = isMin ? -1 : 1;
     ineq[constraints.getNumCols() - 1] = -1;
     constraints.addInequality(ineq);
     if (!constraints.isEmpty())
       return failure();
 
-    // Check each result one-by-one, so remove r_i < minOpUb ineq. again.
+    // Check each result one-by-one, so remove inequality again.
     constraints.removeInequality(constraints.getNumInequalities() - 1);
   }
 
-  // Lower and upper bound of `minOp` are equal. Replace `minOp` with its upper
-  // bound. `dimValues` and `symValues` may have "empty" Values. These must be
-  // filtered out from the list of Values and from the affine map.
+  // Lower and upper bound of `op` are equal. Replace `op` with its upper bound.
+  // `dimValues` and `symValues` may have "empty" Values. These must be filtered
+  // out from the list of Values and from the affine map.
   SmallVector<Value> newOperands;
   SmallVector<AffineExpr> dimReplacements, symReplacements;
   unsigned numDims = 0, numSyms = 0;
   for (Value val : dimValues) {
     if (val) {
       newOperands.push_back(val);
       dimReplacements.push_back(rewriter.getAffineDimExpr(numDims++));
     } else {
       dimReplacements.push_back(rewriter.getAffineDimExpr(numDims));
     }
   }
   for (Value val : symValues) {
     if (val) {
       newOperands.push_back(val);
       symReplacements.push_back(rewriter.getAffineSymbolExpr(numSyms++));
     } else {
       symReplacements.push_back(rewriter.getAffineSymbolExpr(numSyms));
     }
   }
-  AffineMap newMap = alignedUbMap.replaceDimsAndSymbols(
+  AffineMap newMap = alignedBoundMap.replaceDimsAndSymbols(
       dimReplacements, symReplacements, numDims, numSyms);
-  rewriter.setInsertionPoint(minOp);
-  rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, newMap, newOperands);
+  rewriter.setInsertionPoint(op);
+  rewriter.replaceOpWithNewOp<AffineApplyOp>(op, 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):
+/// Try to simplify a min/max operation `op` after loop peeling. This function
+/// can simplify min/max 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
+/// min/max 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
 ///
+/// Returns `success` if the given operation was replaced by a new operation;
+/// `failure` otherwise.
+///
 /// 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,
+/// `insideLoop` must be true for min/max ops inside the loop and false for
+/// affine.min ops inside the scf.for op. For an explanation of the other
+/// parameters, see comment of `canonicalizeMinMaxOpInLoop`.
+static LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter,
+                                           Operation *op, AffineMap map,
+                                           ValueRange operands, bool isMin,
+                                           Value iv, Value ub, Value step,
                                            bool insideLoop) {
   auto constraintsBuilder = [&](FlatAffineConstraints &constraints) {
     // 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 success();
   };
 
-  return canonicalizeAffineMinOp(rewriter, minOp,
-                                 /*dims=*/ValueRange{iv, ub, step},
-                                 constraintsBuilder);
+  return canonicalizeMinMaxOp(rewriter, op, map, operands, isMin,
 }
 
+template <typename OpTy, bool IsMin>
+static void
+rewriteAffineOpAfterPeeling(RewriterBase &rewriter, ForOp forOp, scf::IfOp ifOp,
+                            Value iv, Value splitBound, Value ub, Value step) {
+  forOp.walk([&](OpTy affineOp) {
+    (void)rewritePeeledMinMaxOp(rewriter, affineOp, affineOp.getAffineMap(),
+                                affineOp.operands(), IsMin, iv, ub, step,
+                                /*insideLoop=*/true);
+  });
+  ifOp.walk([&](OpTy affineOp) {
+    (void)rewritePeeledMinMaxOp(rewriter, affineOp, affineOp.getAffineMap(),
+                                affineOp.operands(), IsMin, splitBound, ub,
+                                step, /*insideLoop=*/false);
+  });
+}
+
 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);
-  });
+  // Rewrite affine.min and affine.max ops.
+  Value iv = forOp.getInductionVar(), step = forOp.step();
+  rewriteAffineOpAfterPeeling<AffineMinOp, /*IsMin=*/true>(
+      rewriter, forOp, ifOp, iv, splitBound, ub, step);
+  rewriteAffineOpAfterPeeling<AffineMaxOp, /*IsMin=*/false>(
+      rewriter, forOp, ifOp, iv, splitBound, ub, step);
 
   return success();
 }
 
-/// Canonicalize AffineMinOp operations in the context of for loops with a known
-/// range. Call `canonicalizeAffineMinOp` and add the following constraints to
+/// Canonicalize min/max operations in the context of for loops with a known
+/// range. Call `canonicalizeMinMaxOp` and add the following constraints to
 /// the constraint system (along with the missing dimensions):
 ///
 /// * iv >= lb
 /// * iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
 ///
 /// Note: Due to limitations of FlatAffineConstraints, only constant step sizes
 /// are currently supported.
-LogicalResult mlir::scf::canonicalizeAffineMinOpInLoop(
-    AffineMinOp minOp, RewriterBase &rewriter,
-    function_ref<LogicalResult(Value, Value &, Value &, Value &)> loopMatcher) {
+LogicalResult
+mlir::scf::canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,
+                                      AffineMap map, ValueRange operands,
+                                      bool isMin, LoopMatcherFn loopMatcher) {
   SmallVector<Value> dims;
   SmallVector<int64_t> steps;
 
   // Find all iteration variables among `minOp`'s operands add constrain them.
-  for (Value operand : minOp.operands()) {
+  for (Value operand : operands) {
     // If `operand` is an iteration variable: Find corresponding loop
     // bounds and step.
     Value iv = operand;
     Value lb, ub, step;
     if (failed(loopMatcher(operand, lb, ub, step)))
       continue;
 
     // FlatAffineConstraints does not support semi-affine expressions.
for (int i = 0, e = dims.size(); i < e; i += 3) {
       unsigned dimIv = i, dimLb = i + 1, dimUb = i + 2;
       Value lb = dims[dimLb], ub = dims[dimUb];
       int64_t step = steps[i / 3];
 
       // If loop lower/upper bounds are constant: Add constraint.
       Optional<int64_t> lbInt = getConstantIntValue(lb);
       Optional<int64_t> ubInt = getConstantIntValue(ub);
       if (lbInt)
-        constraints.setIdToConstant(i + 1, *lbInt);
+        constraints.addBound(FlatAffineConstraints::EQ, dimLb, *lbInt);
       if (ubInt)
-        constraints.setIdToConstant(i + 2, *ubInt);
+        constraints.addBound(FlatAffineConstraints::EQ, dimUb, *ubInt);
 
       // iv >= lb (equiv.: iv - lb >= 0)
       SmallVector<int64_t> ineqLb(constraints.getNumCols(), 0);
       ineqLb[dimIv] = 1;
       ineqLb[dimLb] = -1;
       constraints.addInequality(ineqLb);
 
       // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
       AffineExpr exprLb = lbInt ? rewriter.getAffineConstantExpr(*lbInt)
                                 : rewriter.getAffineDimExpr(dimLb);
       AffineExpr exprUb = ubInt ? rewriter.getAffineConstantExpr(*ubInt)
                                 : rewriter.getAffineDimExpr(dimUb);
       AffineExpr ivUb =
           exprLb + 1 + (step * ((exprUb - exprLb - 1).floorDiv(step)));
       auto map = AffineMap::get(
           /*dimCount=*/constraints.getNumDimIds(),
           /*symbolCount=*/constraints.getNumSymbolIds(), /*result=*/ivUb);
 
-      (void)constraints.addLowerOrUpperBound(dimIv, map, /*eq=*/false,
-                                             /*lower=*/false);
+      (void)constraints.addBound(FlatAffineConstraints::UB, dimIv, map);
     }
 
     return success();
   };
 
-  return canonicalizeAffineMinOp(rewriter, minOp, dims, constraintsBuilder);
+  return canonicalizeMinMaxOp(rewriter, op, map, operands, isMin, dims,
+                              constraintsBuilder);
 }
 
 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();
     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();
   }
 };
 
-/// Canonicalize AffineMinOp operations in the context of scf.for and
-/// scf.parallel loops with a known range.
-struct AffineMinSCFCanonicalizationPattern
-    : public OpRewritePattern<AffineMinOp> {
-  using OpRewritePattern<AffineMinOp>::OpRewritePattern;
+/// Canonicalize AffineMinOp/AffineMaxOp operations in the context of scf.for
+/// and scf.parallel loops with a known range.
+template <typename OpTy, bool IsMin>
+struct AffineOpSCFCanonicalizationPattern : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(AffineMinOp minOp,
+  LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
     auto loopMatcher = [](Value iv, Value &lb, Value &ub, Value &step) {
       if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
         lb = forOp.lowerBound();
         ub = forOp.upperBound();
         step = forOp.step();
         return success();
       }
       if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
         for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
           if (parOp.getInductionVars()[idx] == iv) {
             lb = parOp.lowerBound()[idx];
             ub = parOp.upperBound()[idx];
             step = parOp.step()[idx];
             return success();
           }
         }
         return failure();
       }
       return failure();
     };
 
-    return scf::canonicalizeAffineMinOpInLoop(minOp, rewriter, loopMatcher);
+    return scf::canonicalizeMinMaxOpInLoop(rewriter, op, op.getAffineMap(),
+                                           op.operands(), IsMin, loopMatcher);
   }
 };
 } // namespace
 
 namespace {
 struct ParallelLoopSpecialization
     : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
   void runOnFunction() override {
void runOnFunction() override {
     patterns.add<ForLoopPeelingPattern>(ctx);
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 
     // Drop the marker.
     funcOp.walk([](ForOp op) { op->removeAttr(kPeeledLoopLabel); });
   }
 };
 
-struct AffineMinSCFCanonicalization
-    : public AffineMinSCFCanonicalizationBase<AffineMinSCFCanonicalization> {
+struct SCFAffineOpCanonicalization
+    : public SCFAffineOpCanonicalizationBase<SCFAffineOpCanonicalization> {
   void runOnFunction() override {
     FuncOp funcOp = getFunction();
     MLIRContext *ctx = funcOp.getContext();
     RewritePatternSet patterns(ctx);
-    patterns.add<AffineMinSCFCanonicalizationPattern>(ctx);
+    scf::populateSCFLoopBodyCanonicalizationPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   }
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::createAffineMinSCFCanonicalizationPass() {
-  return std::make_unique<AffineMinSCFCanonicalization>();
+std::unique_ptr<Pass> mlir::createSCFAffineOpCanonicalizationPass() {
+  return std::make_unique<SCFAffineOpCanonicalization>();
 }
 
 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>();
 }
 
 void mlir::scf::populateSCFLoopBodyCanonicalizationPatterns(
     RewritePatternSet &patterns) {
-  patterns.insert<AffineMinSCFCanonicalizationPattern>(patterns.getContext());
+  MLIRContext *ctx = patterns.getContext();
+  patterns
+      .insert<AffineOpSCFCanonicalizationPattern<AffineMinOp, /*IsMin=*/true>,
+              AffineOpSCFCanonicalizationPattern<AffineMaxOp, /*IsMin=*/false>>(
+          ctx);
 }

mlir/test/Dialect/SCF/canonicalize-affine-op.mlir

-// RUN: mlir-opt %s -canonicalize-scf-affine-min -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -canonicalize-scf-affine-op -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @scf_for_canonicalize_min
 //       CHECK:   %[[C2:.*]] = constant 2 : i64
 //       CHECK:   scf.for
 //       CHECK:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
 func @scf_for_canonicalize_min(%A : memref<i64>) {
   %c0 = constant 0 : index
   %c2 = constant 2 : index
   %c4 = constant 4 : index
 
   scf.for %i = %c0 to %c4 step %c2 {
     %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
     %2 = index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @scf_for_canonicalize_max
+//       CHECK:   %[[Cneg2:.*]] = constant -2 : i64
+//       CHECK:   scf.for
+//       CHECK:     memref.store %[[Cneg2]], %{{.*}}[] : memref<i64>
+func @scf_for_canonicalize_max(%A : memref<i64>) {
+  %c0 = constant 0 : index
+  %c2 = constant 2 : index
+  %c4 = constant 4 : index
+
+  scf.for %i = %c0 to %c4 step %c2 {
+    %1 = affine.max affine_map<(d0, d1)[] -> (-2, -(d1 - d0))> (%i, %c4)
+    %2 = index_cast %1: index to i64
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func @scf_for_loop_nest_canonicalize_min
 //       CHECK:   %[[C5:.*]] = constant 5 : i64
 //       CHECK:   scf.for
 //       CHECK:     scf.for
 //       CHECK:       memref.store %[[C5]], %{{.*}}[] : memref<i64>
 func @scf_for_loop_nest_canonicalize_min(%A : memref<i64>) {
   %c0 = constant 0 : index
   %c2 = constant 2 : index

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

scf.for %iv = %lb to %ub step %step {
     %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
+// Test rewriting of affine.min/max 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.
+// does not rewrite 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-DAG: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (d0 - s0)>
+//  CHECK-DAG: #[[MAP5:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2))>
+//  CHECK-DAG: #[[MAP6:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2) + 1)>
+//  CHECK-DAG: #[[MAP7:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - (s1 - s0) mod s2) - 1)>
+//  CHECK-DAG: #[[MAP8:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0, s2 - (s2 - (s2 - s1) mod s0), s3)>
+//  CHECK-DAG: #[[MAP9:.*]] = affine_map<()[s0] -> (-s0)>
+//      CHECK: func @test_affine_op_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:     %[[RES4:.*]] = affine.min #[[MAP3]](%[[IV]])[%[[STEP]], %[[UB]], %[[SOME_VAL]]]
+//      CHECK:     %[[RES5:.*]] = affine.apply #[[MAP4]](%[[IV]])[%[[UB]]]
 //      CHECK:     memref.store %[[STEP]]
 //      CHECK:     memref.store %[[STEP]]
 //      CHECK:     memref.store %[[RES2]]
 //      CHECK:     memref.store %[[RES3]]
 //      CHECK:     memref.store %[[RES4]]
+//      CHECK:     memref.store %[[RES5]]
 //      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:     %[[RES_IF_0:.*]] = affine.apply #[[MAP5]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_1:.*]] = affine.apply #[[MAP6]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_2:.*]] = affine.apply #[[MAP6]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_3:.*]] = affine.apply #[[MAP7]]()[%[[LB]], %[[UB]], %[[STEP]]]
+//      CHECK:     %[[RES_IF_4:.*]] = affine.min #[[MAP8]]()[%[[STEP]], %[[LB]], %[[UB]], %[[SOME_VAL]]]
+//      CHECK:     %[[RES_IF_5:.*]] = affine.apply #[[MAP9]]()[%[[STEP]]]
 //      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]]
+//      CHECK:     memref.store %[[RES_IF_5]]
 #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) {
+#map5 = affine_map<(d0, d1)[s0] -> (-s0, -d0 + d1)>
+func @test_affine_op_rewrite(%lb : index, %ub: index,
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
+  %c5 = constant 5 : 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]
+    // Rewrite max(-%ub + %iv, -%step) to -%ub + %iv (and -%step in the scf.if).
+    %m5 = affine.min #map5(%ub, %iv)[%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>
+    memref.store %m5, %d[%c5] : memref<?xindex>
   }
   return
 }