[mlir] Introduce AffineMinSCF folding as a pattern

[nicolasvasilache]

Summary

This revision adds a folding pattern to replace affine.min ops by the actual min value, when it can be determined statically from the strides and bounds of enclosing scf loop .

This matches the type of expressions that Linalg produces during tiling and simplifies boundary checks. For now Linalg depends both on Affine and SCF but they do not depend on each other, so the pattern is added there.
In the future this will move to a more appropriate place when it is determined.

The canonicalization of AffineMinOp operations in the context of enclosing scf.for and scf.parallel proceeds by:

1. building an affine map where uses of the induction variable of a loop are replaced by %lb + %step * floordiv(%iv - %lb, %step) expressions.
2. checking if any of the results of this affine map divides all the other results (in which case it is also guaranteed to be the min).
3. replacing the AffineMinOp by the result of (2).

The algorithm is functional in simple parametric tiling cases by using semi-affine maps. However simplifications of such semi-affine maps are not yet available and the canonicalization does not succeed yet.

[rriddle]

mlir/include/mlir/IR/AffineExpr.h

18

I don't think this is necessary.

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

357

nit: Please move the lambda out of the if, it is making it difficult to see the control flow structure.

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

309

Please run clang-format on this.

[bondhugula]

Can you rewrite the summary a bit? 'canonicalize' is a general term - you are rather folding/simplifying; replacing canonicalize/canonicalization will greatly clarify comments / commit summary. Condition (1) isn't fully spelt out - just replacing IVs would be something totally different. The choice for (2) (divides the results as opposed to min of the results isn't clear either - it is for catching a simple subset?).

Based on the commit summary, this is all about folding it to one of the results. Then, implementing it in the fold hook on AffineMinOp would be the right thing. You'll be able to do faster/more simplification via the rewriter.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

506

Nit: by the result of (2).

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

240–242

Could you rephrase this? As written, it doesn't parse right. It's not scf.for %iv = %lb to %ub step %step that's being substituted!

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

33–34

Can you add a TODO or better file a feature request?

34

It's a 2-d map. How could it evaluate to 0? Did you just mean for the first result dimension?

106

Likewise.

[bondhugula]

Looks like this pattern could just work on LoopLikeOpInterface?

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

268–271

Looks like this whole pattern can just be registered on LoopLikeOpInterface? - by just making the interface return lb + step * iv.floorDiv(step) expressions with the necessary operands (dims + symbols). If it can't, because the lower bound is a max, etc. for a AffineForOp, just fail to match.

270

I'm missing something - did you mean lb + step * (iv - lb).floorDiv(step)?

[nicolasvasilache]

Address review comments.

[nicolasvasilache]

Changed the commit message.

Re implementing it in the fold hook on AffineMinOp would be the right thing. You'll be able to do faster/more simplification via the rewriter.,
I actually need to control it programmatically for now and add it as part of: getLinalgTilingCanonicalizationPatterns.
I do not see how to get folding hooks exposed programmatically at this point: they only activate as part of a call to the canonicalization pass.

Lastly, for the day this moves into Affine, it would make Affine depend on SCF. Are you ok with that ?

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

268–271

I wouldn't think so based on today's needs: the part about deciding between constant, symbol and dimension and their position is quite specific to this particular rewrite.
At the very best I could see a refactoring if we ever see a second user for this functionality with similar specifics.
Tentatively marking as solved.

270

Very true, thank you for catching!

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

33–34

https://bugs.llvm.org/show_bug.cgi?id=46367 for the missing canonicalization.

[ftynse]

mlir/include/mlir/IR/AffineExpr.h

123

I don't understand what "sparse" means in this context. Does this function replace any occurrence of the expr subtree within this with replacement?

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

240

Nit: corresponds

268–271

Looks like this whole pattern can just be registered on LoopLikeOpInterface?

We don't (yet) have patterns that can work on interfaces. That being said, I had a similar thought -- this transformation is useful for any kind of loops. There is another pending patch that could benefit from if lb/ub/step are constants. That being said, this refactoring may be out of scope for this patch.

275

Nit: "instead in place", chose one

334

llvm::to_vector<4>(minOp.getAffineMap().getResults()) ?

368

Nit: please use symmetric braces

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

14

I'm puzzled. This seems to check for a store of an _index_ constant into a memref<i64>, since it explicitly checks for the absence of index_cast. This looks incorrect.

[nicolasvasilache]

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

14

nice catch !

The issue is that the test pattern application is done via greedily applying the pattern which triggers foldings.
There is also a constant 2: i64 that is the one being used and it happens that its SSA name is a prefix of the one for constant 2: index.
The CHECK only matches that prefix.

Rewriting to ensure the folding does not happen which will keep the index_cast.

[nicolasvasilache]

mlir/include/mlir/IR/AffineExpr.h

123

Yes, added more doc to clarify.

[nicolasvasilache]

Address review.

[nicolasvasilache]

@bondhugula I believe I addressed your comments?

[bondhugula]

In D82009#2099175, @nicolasvasilache wrote:

Changed the commit message.

Re implementing it in the fold hook on AffineMinOp would be the right thing. You'll be able to do faster/more simplification via the rewriter.,
I actually need to control it programmatically for now and add it as part of: getLinalgTilingCanonicalizationPatterns.
I do not see how to get folding hooks exposed programmatically at this point: they only activate as part of a call to the canonicalization pass.

Oh that's straightforward :) The folding pattern will automatically get called when you run applyPatternsAndFoldGreedily, which you are anyway calling. There wouldn't be anything to add to patterns for it. The only thing is that: with it being in the folder, there wouldn't be a way to prevent it from being run currently.

[bondhugula]

In D82009#2099175, @nicolasvasilache wrote:

Changed the commit message.

Lastly, for the day this moves into Affine, it would make Affine depend on SCF. Are you ok with that ?

I didn't understand this part. What moves into affine? Affine already depends on LoopLikeOpInterface, and so when this pattern is moved to LoopLikeOpInterface, it should be usable with any more deps.

[bondhugula]

In D82009#2100866, @nicolasvasilache wrote:

@bondhugula I believe I addressed your comments?

Just added a few more comments on the folding part. But do finally rely on @ftynse's approval for going ahead with the commit.

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

270

Is this covered via test cases now?

[nicolasvasilache]

Upon deeper testing, some incorrect assumptions were made when using modulos.
Switch to an implementation in which we explicitly build out the min expression.

[ftynse]

Please add a test where the constant is not the first dimension in the map.

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

249

Nit: return positivePath ? min : max and drop the outer braces?

278

Typo: matchine

279

I don't see lb, ub and step in the argument list. Did you mean lbVal and co? In this case, how do you store AffineDimExpr into a Value?

339

Consider bailing out early if exprs.empty() or documenting the assumption

368

Can substitute just return the resulting map? It deconstructs the map into the vector-of-expressions, and here the map is reconstructed back...

394

Shouldn't this be continue? Otherwise you stop after the first expression.

402

If the map is known to be symbolic (the check at the entry of the loop), it shouldn't have any dims, so we can just spare them here. I don't know if further simplification is posisble.

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

60

I think you actually want to check for affine.min, index_cast will always be there

76

Idem

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

394

Thanks for catching, added a new test.

[nicolasvasilache]

Address review comments.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 struct LinalgCopyVTWForwardingPattern
     : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
                                 PatternRewriter &rewriter) const override;
 };
 
+/// Canonicalize AffineMinOp operations in the context of enclosing scf.for and
+/// scf.parallel by:
+///   1. building an affine map where uses of the induction variable of a loop
+///   are replaced by a `%lb + %step * floorDiv(%iv - %lb, %step)` expression.
+///   2. checking whether any of the results of this affine map divides all the
+///   other results (in which case it is also guaranteed to be the min).
+///   3. replacing the AffineMinOp by the result of (2).
+// TODO: move to a more appropriate place when it is determined. For now Linalg
+// depends both on Affine and SCF but they do not depend on each other.
+struct AffineMinSCFCanonicalizationPattern
+    : public OpRewritePattern<AffineMinOp> {
+  using OpRewritePattern<AffineMinOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineMinOp minOp,
+                                PatternRewriter &rewriter) const override;
+};
+
 //===----------------------------------------------------------------------===//
 // Support for staged pattern application.
 //===----------------------------------------------------------------------===//
 /// Helper function to allow applying rewrite patterns, interleaved with more
 /// global transformations, in a staged fashion:
 ///   1. the first stage consists of a list of OwningRewritePatternList. Each
 ///   OwningRewritePatternList in this list is applied once, in order.
 ///   2. the second stage consists of a single OwningRewritePattern that is
 ///   applied greedily until convergence.
 ///   3. the third stage consists of applying a lambda, generally used for
 ///   non-local transformation effects. This allows creating custom fused
 ///   transformations where patterns can be ordered and applied at a finer
 ///   granularity than a sequence of traditional compiler passes.
 LogicalResult applyStagedPatterns(
     Operation *op, ArrayRef<OwningRewritePatternList> stage1Patterns,
     const OwningRewritePatternList &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda = nullptr);
+
 } // namespace linalg
 } // namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_

mlir/include/mlir/IR/AffineExpr.h

 // symbols, including ceildiv/floordiv/mod by a constant integer.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_AFFINE_EXPR_H
 #define MLIR_IR_AFFINE_EXPR_H
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/Support/Casting.h"
 #include <type_traits>
 
 namespace mlir {
 
 class MLIRContext;
 class AffineMap;
 class IntegerSet;
public:
   /// Return true if the affine expression involves AffineDimExpr `position`.
   bool isFunctionOfDim(unsigned position) const;
 
   /// Walk all of the AffineExpr's in this expression in postorder.
   void walk(std::function<void(AffineExpr)> callback) const;
 
   /// This method substitutes any uses of dimensions and symbols (e.g.
   /// dim#0 with dimReplacements[0]) and returns the modified expression tree.
+  /// This is a dense replacement method: a replacement must be specified for
+  /// every single dim and symbol.
   AffineExpr replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                    ArrayRef<AffineExpr> symReplacements) const;
 
+  /// Sparse replace method. Replace `expr` by `replacement` and return the
+  /// 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;
+
   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;
   AffineExpr floorDiv(uint64_t v) const;

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

 #define DEBUG_TYPE "linalg-transforms"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
+
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 // Marker used as attribute name in generated Linalg rewriting transformations.
 const StringLiteral mlir::linalg::LinalgTransforms::kLinalgTransformMarker =
     "__internal_linalg_transform__";
 
 mlir::linalg::LinalgMarker::LinalgMarker(ArrayRef<Identifier> matchDisjunction,
if (stage3Lambda) {
       if (failed(stage3Lambda(op)))
         return failure();
       LLVM_DEBUG(DBGS() << "After 3rd stage, iter : " << iteration << "\n"
                         << *op);
     }
   }
   return success();
 }
+
+/// Substitute the AffineExprDim at position `dimIdx`, which corresponds to a
+/// loop induction variable (e.g. scf.for %iv = %lb to %ub step %step) by the
+/// AffineExpr representing `%lb + %step * floorDiv(%iv - %lb, %step)` such
+/// that:
+/// 1. the AffineExpr for %lb is either an AffineConstantExpr or an
+///    AffineDimExpr depending on whether the value is constant or not.
+/// 2. the AffineExpr for %step is either an AffineConstantExpr or an
+///    AffineSymbolExpr depending on whether the value is constant or not.
+static void substituteLoop(unsigned dimIdx, Value lbVal, Value ubVal,
+                           Value stepVal, SmallVectorImpl<AffineExpr> &exprs,
+                           SmallVectorImpl<Value> &dims,
+                           SmallVectorImpl<Value> &symbols) {
+  MLIRContext *ctx = lbVal.getContext();
+
+  // 1. maybe add a new dim for the `lb`.
+  auto lbConstant = lbVal.getDefiningOp<ConstantIndexOp>();
+  AffineExpr lb = lbConstant ? getAffineConstantExpr(lbConstant.getValue(), ctx)
+                             : getAffineDimExpr(dims.size(), ctx);
+  if (!lbConstant)
+    dims.push_back(lbVal);
+
+  // 2. maybe add a new symbol for the `step`.
+  auto stepConstant = stepVal.getDefiningOp<ConstantIndexOp>();
+  AffineExpr step = stepConstant
+                        ? getAffineConstantExpr(stepConstant.getValue(), ctx)
+                        : getAffineSymbolExpr(symbols.size(), ctx);
+  if (!stepConstant)
+    symbols.push_back(stepVal);
+
+  // 3. Rewrite `exprs` in place by replacing `dim[dimIdx]` by `lb + step * iv`.
+  AffineExpr iv = getAffineDimExpr(dimIdx, ctx);
+  for (auto &e : exprs)
+    e = e.replace(iv, lb + step * (iv - lb).floorDiv(step));
+}
+
+/// Traverse the `dims` and substitute linear expressions in place of induction
+/// variables in `exprs`.
+static void substitute(SmallVectorImpl<AffineExpr> &exprs,
+                       SmallVectorImpl<Value> &dims,
+                       SmallVectorImpl<Value> &symbols) {
+  assert(!exprs.empty() && "Unexpected empty exprs");
+  LLVM_DEBUG(llvm::interleaveComma(dims, DBGS() << "Start subst with dims: "));
+  LLVM_DEBUG(llvm::dbgs() << "\n");
+
+  // Note: `dims` and `symbols` grow as we iterate, upper bound is dynamic.
+  for (unsigned dimIdx = 0; dimIdx < dims.size(); ++dimIdx) {
+    Value dim = dims[dimIdx];
+    LLVM_DEBUG(DBGS() << "Subst: " << dim << "\n");
+
+    // Replace dim @ pos[dimIdx] by `%lb + %step * new_dim`
+    // Where new dim / symbols are added depending on whether the values are
+    // static or not.
+    if (auto forOp = scf::getForInductionVarOwner(dim)) {
+      substituteLoop(dimIdx, forOp.lowerBound(), forOp.upperBound(),
+                     forOp.step(), exprs, dims, symbols);
+      continue;
+    }
+    if (auto parallelForOp = scf::getParallelForInductionVarOwner(dim)) {
+      for (unsigned idx = 0, e = parallelForOp.getNumLoops(); idx < e; ++idx)
+        substituteLoop(dimIdx, parallelForOp.lowerBound()[idx],
+                       parallelForOp.upperBound()[idx],
+                       parallelForOp.step()[idx], exprs, dims, symbols);
+      continue;
+    }
+  }
+
+  // Cleanup and simplify the results.
+  SmallVector<Value, 4> operands(dims.begin(), dims.end());
+  operands.append(symbols.begin(), symbols.end());
+  auto map = AffineMap::get(dims.size(), symbols.size(), exprs,
+                            exprs.front().getContext());
+  // Pull in affine.apply operations and compose them fully into the result.
+  fullyComposeAffineMapAndOperands(&map, &operands);
+  canonicalizeMapAndOperands(&map, &operands);
+  map = simplifyAffineMap(map);
+  // Assign the results.
+  exprs.assign(map.getResults().begin(), map.getResults().end());
+  dims.assign(operands.begin(), operands.begin() + map.getNumDims());
+  symbols.assign(operands.begin() + map.getNumDims(), operands.end());
+}
+
+LogicalResult AffineMinSCFCanonicalizationPattern::matchAndRewrite(
+    AffineMinOp minOp, PatternRewriter &rewriter) const {
+  // At least one loop is needed to canonicalize affine.min + SCF.
+  auto isLoopLike = [](Value v) {
+    return scf::getParallelForInductionVarOwner(v) ||
+           scf::getForInductionVarOwner(v);
+  };
+  if (llvm::none_of(minOp.getDimOperands(), isLoopLike))
+    return failure();
+
+  LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
+                    << "\n");
+
+  auto exprs = llvm::to_vector<4>(minOp.getAffineMap().getResults());
+  SmallVector<Value, 4> dims(minOp.getDimOperands()),
+      symbols(minOp.getSymbolOperands());
+  substitute(exprs, dims, symbols);
+
+  MLIRContext *ctx = minOp.getContext();
+  auto map = AffineMap::get(dims.size(), symbols.size(), exprs, ctx);
+  LLVM_DEBUG(DBGS() << "Resulting map: " << map << "\n");
+
+  // Check whether any of the expressions divides all expressions. In which case
+  // it is guaranteed to be the min.
+  for (auto e : map.getResults()) {
+    LLVM_DEBUG(DBGS() << "Candidate mod: " << e << "\n");
+    if (!e.isSymbolicOrConstant())
+      continue;
+
+    LLVM_DEBUG(DBGS() << "Check whether mod: " << e << " is zero\n");
+    SmallVector<AffineExpr, 4> modExprs;
+    for (auto ee : map.getResults())
+      modExprs.push_back(ee % e);
+
+    AffineMap modMap = simplifyAffineMap(
+        AffineMap::get(map.getNumDims(), map.getNumSymbols(), modExprs, ctx));
+    LLVM_DEBUG(DBGS() << "simplified modMap: " << modMap << "\n");
+
+    auto isZero = [](AffineExpr e) {
+      if (auto cst = e.dyn_cast<AffineConstantExpr>())
+        return cst.getValue() == 0;
+      return false;
+    };
+    if (llvm::all_of(modMap.getResults(), isZero)) {
+      if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
+        rewriter.replaceOpWithNewOp<ConstantIndexOp>(minOp, cst.getValue());
+      } else {
+        auto resultMap =
+            AffineMap::get(map.getNumDims(), map.getNumSymbols(), {e}, ctx);
+        SmallVector<Value, 4> resultOperands = dims;
+        resultOperands.append(symbols.begin(), symbols.end());
+        canonicalizeMapAndOperands(&resultMap, &resultOperands);
+        resultMap = simplifyAffineMap(resultMap);
+        rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, resultMap,
+                                                   resultOperands);
+      }
+      return success();
+    }
+  }
+
+  return failure();
+}

mlir/lib/IR/AffineExpr.cpp

case AffineExprKind::Mod:
     auto newRHS = rhs.replaceDimsAndSymbols(dimReplacements, symReplacements);
     if (newLHS == lhs && newRHS == rhs)
       return *this;
     return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
   }
   llvm_unreachable("Unknown AffineExpr");
 }
 
+/// Sparse replace method. Return the modified expression tree.
+AffineExpr
+AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const {
+  auto it = map.find(*this);
+  if (it != map.end())
+    return it->second;
+  switch (getKind()) {
+  default:
+    return *this;
+  case AffineExprKind::Add:
+  case AffineExprKind::Mul:
+  case AffineExprKind::FloorDiv:
+  case AffineExprKind::CeilDiv:
+  case AffineExprKind::Mod:
+    auto binOp = cast<AffineBinaryOpExpr>();
+    auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
+    auto newLHS = lhs.replace(map);
+    auto newRHS = rhs.replace(map);
+    if (newLHS == lhs && newRHS == rhs)
+      return *this;
+    return getAffineBinaryOpExpr(getKind(), newLHS, newRHS);
+  }
+  llvm_unreachable("Unknown AffineExpr");
+}
+
+/// Sparse replace method. Return the modified expression tree.
+AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const {
+  DenseMap<AffineExpr, AffineExpr> map;
+  map.insert(std::make_pair(expr, replacement));
+  return replace(map);
+}
+
 /// Returns true if this expression is made out of only symbols and
 /// constants (no dimensional identifiers).
 bool AffineExpr::isSymbolicOrConstant() const {
   switch (getKind()) {
   case AffineExprKind::Constant:
     return true;
   case AffineExprKind::DimId:
     return false;

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

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-affine-min-scf-canonicalization-patterns | FileCheck %s
+
+// CHECK-LABEL: scf_for
+func @scf_for(%A : memref<i64>, %step : index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c7 = constant 7 : index
+  %c4 = constant 4 : index
+
+  //      CHECK: scf.for
+  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
+  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
+  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
+  scf.for %i = %c0 to %c4 step %c2 {
+    %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  //      CHECK: scf.for
+  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
+  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
+  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
+  scf.for %i = %c1 to %c7 step %c2 {
+    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  // This example should simplify but affine_map is currently missing
+  // semi-affine canonicalizations.
+  // This affine map does not currently evaluate to 0:
+  //   (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * 42) mod s0)
+  // Revisit when support is added.
+  // CHECK: scf.for
+  // CHECK: index_cast
+  %ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
+  scf.for %i = %c0 to %ub step %step {
+    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub)
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  // This example should simplify but affine_map is currently missing
+  // semi-affine canonicalizations.
+  // This affine map does not currently evaluate to 0:
+  //   (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * s0) mod s0)
+  // Revisit when support is added.
+  // CHECK: scf.for
+  // CHECK: index_cast
+  %ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
+  scf.for %i = %c0 to %ub2 step %step {
+    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  return
+}
+
+// CHECK-LABEL: scf_parallel
+func @scf_parallel(%A : memref<i64>, %step : index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c7 = constant 7 : index
+  %c4 = constant 4 : index
+
+  // CHECK: scf.parallel
+  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
+  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
+  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
+  scf.parallel (%i) = (%c0) to (%c4) step (%c2) {
+    %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  // CHECK: scf.parallel
+  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
+  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
+  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
+  scf.parallel (%i) = (%c1) to (%c7) step (%c2) {
+    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  // This example should simplify but affine_map is currently missing
+  // semi-affine canonicalizations.
+  // This affine map does not currently evaluate to (0, 0):
+  //   (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * 42) mod s0)
+  // TODO: Revisit when support is added.
+  // CHECK: scf.parallel
+  // CHECK: index_cast
+  %ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
+  scf.parallel (%i) = (%c0) to (%ub) step (%step) {
+    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub)
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  // This example should simplify but affine_map is currently missing
+  // semi-affine canonicalizations.
+  // This affine map does not currently evaluate to (0, 0):
+  //   (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * s0) mod s0)
+  // TODO: Revisit when support is added.
+  // CHECK: scf.parallel
+  // CHECK: index_cast
+  %ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
+  scf.parallel (%i) = (%c0) to (%ub2) step (%step) {
+    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  return
+}

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

struct TestLinalgTransforms
   Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
                                     llvm::cl::desc("Test promotion options"),
                                     llvm::cl::init(false)};
   Option<bool> testVectorTransferForwardingPatterns{
       *this, "test-vector-transfer-forwarding-patterns",
       llvm::cl::desc(
           "Test a fused pass that forwards linalg.copy to vector.transfer"),
       llvm::cl::init(false)};
+  Option<bool> testAffineMinSCFCanonicalizationPatterns{
+      *this, "test-affine-min-scf-canonicalization-patterns",
+      llvm::cl::desc("Test affine-min + scf canonicalization patterns."),
+      llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//

 
 static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
   OwningRewritePatternList forwardPattern;
   forwardPattern.insert<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
   forwardPattern.insert<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
   applyPatternsAndFoldGreedily(funcOp, forwardPattern);
 }
 
+static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
+  OwningRewritePatternList foldPattern;
+  foldPattern.insert<AffineMinSCFCanonicalizationPattern>(funcOp.getContext());
+  // Explicitly walk and apply the pattern locally to avoid more general folding
+  // on the rest of the IR.
+  funcOp.walk([&foldPattern](AffineMinOp minOp) {
+    applyOpPatternsAndFold(minOp, foldPattern);
+  });
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {
   auto lambda = [&](void *) {
     getFunction().walk([](LinalgOp op) {
       op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
   };
   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
 
   if (testPromotionOptions) {
     OwningRewritePatternList patterns;
     fillPromotionCallBackPatterns(&getContext(), patterns);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
     return;
   }
   if (testPatterns)
     return applyPatterns(getFunction());
   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
     return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
+  if (testAffineMinSCFCanonicalizationPatterns)
+    return applyAffineMinSCFCanonicalizationPatterns(getFunction());
 }
 
 namespace mlir {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace mlir