[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

315

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

511

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

59

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

75

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 either the min (i.e. `%lb`) of the max
+///   (i.e. `%lb + %step * floordiv(%ub -1 - %lb, %step)`) expression, depending
+///   on whether the induction variable is used with a positive or negative
+///   coefficient.
+///   2. checking whether any of the results of this affine map is known to be
+///   greater than all other results.
+///   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;
+
   /// Replace symbols[0 .. numDims - 1] by
   /// symbols[shift .. shift + numDims - 1].
   AffineExpr shiftSymbols(unsigned numSymbols, unsigned shift) const;
 
   AffineExpr operator+(int64_t v) const;
   AffineExpr operator+(AffineExpr other) const;
   AffineExpr operator-() const;
   AffineExpr operator-(int64_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();
 }
+
+/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
+/// been replaced by either:
+///  - `min` if `positivePath` is true when we reach an occurrence of `dim`
+///  - `max` if `positivePath` is true when we reach an occurrence of `dim`
+/// `positivePath` is negated each time we hit a multiplicative or divisive
+/// binary op with a constant negative coefficient.
+static AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
+                               AffineExpr max, bool positivePath = true) {
+  if (e == dim) {
+    if (positivePath)
+      return min;
+    if (!positivePath)
+      return max;
+  }
+  if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) {
+    AffineExpr lhs = bin.getLHS();
+    AffineExpr rhs = bin.getRHS();
+    if (bin.getKind() == mlir::AffineExprKind::Add)
+      return substWithMin(lhs, dim, min, max, positivePath) +
+             substWithMin(rhs, dim, min, max, positivePath);
+
+    auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>();
+    auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>();
+    if (c1 && c1.getValue() < 0)
+      return getAffineBinaryOpExpr(
+          bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
+    if (c2 && c2.getValue() < 0)
+      return getAffineBinaryOpExpr(
+          bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
+    return getAffineBinaryOpExpr(
+        bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
+        substWithMin(rhs, dim, min, max, positivePath));
+  }
+  return e;
+}
+
+/// Gieven a `forOp`, append its lower and upper bounds to `dims` and its step
+/// to `symbols`. Also create new AffineDimExpr and AffineSymbolExpr with
+/// positions matchine the newly appended values and store them in `lb`, `ub`
+/// and `step` respectively.
+static AffineExpr substituteLoopInExpr(AffineExpr expr, unsigned dimIdx,
+                                       Value lbVal, Value ubVal, Value stepVal,
+                                       SmallVectorImpl<Value> &dims,
+                                       SmallVectorImpl<Value> &symbols) {
+  MLIRContext *ctx = lbVal.getContext();
+  AffineExpr lb = getAffineDimExpr(dims.size(), ctx);
+  dims.push_back(lbVal);
+  AffineExpr ub = getAffineDimExpr(dims.size(), ctx);
+  dims.push_back(ubVal);
+  AffineExpr step = getAffineSymbolExpr(symbols.size(), ctx);
+  symbols.push_back(stepVal);
+  LLVM_DEBUG(DBGS() << "Before: " << expr << "\n");
+  AffineExpr d = getAffineDimExpr(dimIdx, expr.getContext());
+  AffineExpr ee =
+      substWithMin(expr, d, lb, lb + step * ((ub - 1) - lb).floorDiv(step));
+  LLVM_DEBUG(DBGS() << "After: " << expr << "\n");
+  return ee;
+}
+
+/// Traverse the `dims` and substitute known min or max expressions in place of
+/// induction variables in `exprs`.
+static void substitute(SmallVectorImpl<AffineExpr> &exprs,
+                       SmallVectorImpl<Value> &dims,
+                       SmallVectorImpl<Value> &symbols) {
+  for (AffineExpr &expr : exprs) {
+    bool substituted = true;
+    while (substituted) {
+      substituted = false;
+      for (unsigned dimIdx = 0; dimIdx < dims.size(); ++dimIdx) {
+        Value dim = dims[dimIdx];
+        LLVM_DEBUG(DBGS() << "Subst: " << dim << "\n");
+        AffineExpr substitutedExpr;
+        if (auto forOp = scf::getForInductionVarOwner(dim))
+          substitutedExpr = substituteLoopInExpr(
+              expr, dimIdx, forOp.lowerBound(), forOp.upperBound(),
+              forOp.step(), dims, symbols);
+
+        if (auto parallelForOp = scf::getParallelForInductionVarOwner(dim))
+          for (unsigned idx = 0, e = parallelForOp.getNumLoops(); idx < e;
+               ++idx)
+            substitutedExpr = substituteLoopInExpr(
+                expr, dimIdx, parallelForOp.lowerBound()[idx],
+                parallelForOp.upperBound()[idx], parallelForOp.step()[idx],
+                dims, symbols);
+
+        if (!substitutedExpr)
+          continue;
+
+        substituted = (substitutedExpr != expr);
+        expr = substitutedExpr;
+      }
+    }
+
+    // Cleanup and simplify the results.
+    // This needs to happen outside of the loop iterating on dims.size() since
+    // it modifies dims.
+    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());
+
+    LLVM_DEBUG(DBGS() << "Map to simplify: " << map << "\n");
+
+    // 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());
+
+    LLVM_DEBUG(DBGS() << "Map simplified: " << map << "\n");
+  }
+}
+
+LogicalResult AffineMinSCFCanonicalizationPattern::matchAndRewrite(
+    AffineMinOp minOp, PatternRewriter &rewriter) const {
+  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, when subtracted from all other
+  // expressions, produces only >= 0 constants. If so, it is the min.
+  for (auto e : map.getResults()) {
+    LLVM_DEBUG(DBGS() << "Candidate min: " << e << "\n");
+    if (!e.isSymbolicOrConstant())
+      continue;
+
+    auto isNonPositive = [](AffineExpr e) {
+      if (auto cst = e.dyn_cast<AffineConstantExpr>())
+        return cst.getValue() < 0;
+      return true;
+    };
+
+    // Build the subMap and check everything is statically known to be
+    // positive.
+    SmallVector<AffineExpr, 4> subExprs;
+    subExprs.reserve(map.getNumResults());
+    for (auto ee : map.getResults())
+      subExprs.push_back(ee - e);
+    AffineMap subMap = simplifyAffineMap(
+        AffineMap::get(map.getNumDims(), map.getNumSymbols(), subExprs, ctx));
+    LLVM_DEBUG(DBGS() << "simplified subMap: " << subMap << "\n");
+    if (llvm::any_of(subMap.getResults(), isNonPositive))
+      return failure();
+
+    // Static min found.
+    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

 /// Replace symbols[0 .. numDims - 1] by symbols[shift .. shift + numDims - 1].
 AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift) const {
   SmallVector<AffineExpr, 4> symbols;
   for (unsigned idx = 0; 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);
+  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
+  %c16 = constant 16 : index
+  %c1024 = constant 1024 : 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 should not canonicalize because: 4 - %i may take the value 1 < 2.
+  //     CHECK:   scf.for
+  //     CHECK:     affine.min
+  //     CHECK:     index_cast
+  scf.for %i = %c1 to %c4 step %c2 {
+    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c4]
+    %2 = index_cast %1: index to i64
+    store %2, %A[]: memref<i64>
+  }
+
+  // This should not canonicalize because: 16 - %i may take the value 15 < 1024.
+  //     CHECK:   scf.for
+  //     CHECK:     affine.min
+  //     CHECK:     index_cast
+  scf.for %i = %c1 to %c16 step %c1024 {
+    %1 = affine.min affine_map<(d0) -> (1024, 16 - d0)> (%i)
+    %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: `((s0 * 42 - 1) floordiv s0) * s0`
+  // should evaluate to 41 * s0.
+  // Note that this may require positivity assumptions on `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, d1 - d2)> (%step, %ub, %i)
+    %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 example should simplify but affine_map is currently missing
+  // semi-affine canonicalizations: ` -(((s0 * s0 - 1) floordiv s0) * s0)`
+  // should evaluate to (s0 - 1) * s0.
+  // Note that this may require positivity assumptions on `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

Option<bool> testVectorTransferForwardingPatterns{
       llvm::cl::desc(
           "Test a fused pass that forwards linalg.copy to vector.transfer"),
       llvm::cl::init(false)};
   Option<bool> testGenericToVectorPattern{
       *this, "test-contraction-to-vector-patterns",
       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
                      "in vector.contract form"),
       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 applyContractionToVectorPatterns(FuncOp funcOp) {
   OwningRewritePatternList patterns;
   patterns.insert<LinalgVectorizationPattern<BatchMatmulOp>,
                   LinalgVectorizationPattern<MatmulOp>,
                   LinalgVectorizationPattern<MatvecOp>,
                   LinalgVectorizationPattern<DotOp>,
                   LinalgVectorizationPattern<GenericOp>>(funcOp.getContext());
   applyPatternsAndFoldGreedily(funcOp, patterns);
 }
 
+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};
void TestLinalgTransforms::runOnFunction() {
   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
     return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
   if (testGenericToVectorPattern)
     return applyContractionToVectorPatterns(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