[mlir][scf] Add general affine.min canonicalization pattern

[springerm]

Summary

This canonicalization simplifies affine.min operations inside "for loop"-like operations (e.g., scf.for and scf.parallel) based on two invariants:

• iv >= lb
• iv < lb + step * ((ub - lb - 1) floorDiv step) + 1

This commit adds a new pass canonicalize-scf-affine-min (instead of being a canonicalization pattern) to avoid dependencies between the Affine dialect and the SCF dialect.

Depends On D107222

[springerm]

rebase

[ftynse]

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

20–27

Why do we need a separate pass for this? Can't this be done as part of regular -canonicalize? If there is a good reason, please state it in the commit description. If this is only useful for isolated testing (without other canonicalizations), please move this under test/lib.

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

404

Nit: this syntax is uncommon and hard to parse visually. Could you please move declarators without initializer to a separate line?

421

Nit: please move the lambda out of the function call.

422

438–439

Don't we want to use a constant instead of dimLb here if we know LB is constant?

485

Nit: let's move this lambda out of the call too.

501

Nit: return failure() here. If this function is ever extended, it will needlessly try matching other op kinds instead of bailing out, leading to surprising behaviors.

[springerm]

rebase

[springerm]

rebase

[springerm]

address comments

[springerm]

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

20–27

As far as I know, canonicalization patterns must be defined along with the ops of that dialect.

So this could be a canonicalization pattern of the SCF dialect. But then there is a dependency of the SCF dialect on the affine dialect. I think this is not good. Similarly, we probably don't want a dependency of the affine dialect on the SCF dialect.

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

438–439

Doesn't really matter. The current implementation is a bit shorter (fewer lines of code).

[springerm]

update test case

[bondhugula]

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

31

Missing doc comment.

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

40

for in backticks please.

[springerm]

rebase

[springerm]

rebase

[ftynse]

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

20–27

Ah yes, the infamous cross-dialect canonicalization problem. Could you add what you said as a comment above the pass definition?

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

449

Could we assert (if the call should never fail) or propagate errors (if there is a hypothetical future where it may) instead if ignoring the return value?

545

Let's propagate errors here and signalPassFailure if patterns couldn't apply. It's not impossible that further (failable) patterns are added above without changing the error propagation.

[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();
 
+/// Creates a pass that canonicalizes affine.min ops in scf.for loops with
+/// known lower and upper bounds.
+std::unique_ptr<Pass> createAffineMinSCFCanonicalizationPass();
+
 /// 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();
 
 /// Creates a pass which tiles innermost parallel loops.

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()";
+  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"];
 }
 
 def SCFForLoopSpecialization

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

 //===- Transforms.h - SCF dialect transformation utilities ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines transformations on SCF operations.
 //
 //===----------------------------------------------------------------------===//
 
 #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 ConversionTarget;
 struct LogicalResult;
 class MLIRContext;
 class Region;
 class RewriterBase;
 class TypeConverter;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 class Operation;
+class Value;
 
 namespace scf {
 
 class IfOp;
 class ForOp;
 class ParallelOp;
 class ForOp;
 
+/// Try to canonicalize an affine.min operation 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`.
+///
+/// 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);
+
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
 /// into one scf.parallel operations. Uses a naive aliasing and dependency
 /// analysis.
 void naivelyFuseParallelOps(Region &region);
 
 /// Rewrite a for loop with bounds/step that potentially do not divide evenly
 /// into a for loop where the step divides the iteration space evenly, followed
 /// by an scf.if for the last (partial) iteration (if any). This transformation

 /// scf.for %I = %C0 to %N - 2 {
 ///  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,
+/// 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

LogicalResult mlir::scf::peelAndCanonicalizeForLoop(RewriterBase &rewriter,
   ifOp.walk([&](AffineMinOp minOp) {
     (void)rewritePeeledAffineOp(rewriter, minOp, splitBound, ub, forOp.step(),
                                 /*insideLoop=*/false);
   });
 
   return success();
 }
 
+/// Canonicalize AffineMinOp operations in the context of for loops with a known
+/// range. Call `canonicalizeAffineMinOp` 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) {
+  FlatAffineValueConstraints constraints;
+  DenseSet<Value> allIvs;
+
+  // Find all iteration variables among `minOp`'s operands add constrain them.
+  for (Value operand : minOp.operands()) {
+    // Skip duplicate ivs.
+    if (llvm::find(allIvs, operand) != allIvs.end())
+      continue;
+
+    // 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;
+    allIvs.insert(iv);
+
+    // FlatAffineConstraints does not support semi-affine expressions.
+    // Therefore, only constant step values are supported.
+    auto stepInt = getConstantIntValue(step);
+    if (!stepInt)
+      continue;
+
+    unsigned dimIv = constraints.addDimId(iv);
+    unsigned dimLb = constraints.addDimId(lb);
+    unsigned dimUb = constraints.addDimId(ub);
+
+    // If loop lower/upper bounds are constant: Add EQ constraint.
+    Optional<int64_t> lbInt = getConstantIntValue(lb);
+    Optional<int64_t> ubInt = getConstantIntValue(ub);
+    if (lbInt)
+      constraints.addBound(FlatAffineConstraints::EQ, dimLb, *lbInt);
+    if (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 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
+    auto map = AffineMap::get(
+        /*dimCount=*/constraints.getNumDimIds(),
+        /*symbolCount=*/constraints.getNumSymbolIds(), /*result=*/ivUb);
+
+    (void)constraints.addBound(FlatAffineConstraints::UB, dimIv, map);
+  }
+
+  return canonicalizeAffineMinOp(rewriter, minOp, constraints);
+}
+
 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;
+
+  LogicalResult matchAndRewrite(AffineMinOp minOp,
+                                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);
+  }
+};
 } // namespace
 
 namespace {
 struct ParallelLoopSpecialization
     : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
   void runOnFunction() override {
     getFunction().walk(
         [](ParallelOp op) { specializeParallelLoopForUnrolling(op); });
void runOnFunction() override {
     RewritePatternSet patterns(ctx);
     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> {
+  void runOnFunction() override {
+    FuncOp funcOp = getFunction();
+    MLIRContext *ctx = funcOp.getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<AffineMinSCFCanonicalizationPattern>(ctx);
+    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  }
+};
 } // namespace
 
+std::unique_ptr<Pass> mlir::createAffineMinSCFCanonicalizationPass() {
+  return std::make_unique<AffineMinSCFCanonicalization>();
+}
+
 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());
+}

mlir/test/Dialect/SCF/canonicalize-scf-affine-min.mlir

+// RUN: mlir-opt %s -canonicalize-scf-affine-min -split-input-file | FileCheck %s
+
+// Note: This is mostly a copy of test/Dialect/Linalg/fold-affine-min-scf.mlir
+
+// 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_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
+  %c3 = constant 3 : index
+  %c4 = constant 4 : index
+  %c6 = constant 6 : index
+
+  scf.for %i = %c0 to %c4 step %c2 {
+    scf.for %j = %c0 to %c6 step %c3 {
+      %1 = affine.min affine_map<(d0, d1, d2, d3)[] -> (5, d1 + d3 - d0 - d2)> (%i, %c4, %j, %c6)
+      %2 = index_cast %1: index to i64
+      memref.store %2, %A[]: memref<i64>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_not_canonicalizable_1
+//       CHECK:   scf.for
+//       CHECK:     affine.min
+//       CHECK:     index_cast
+func @scf_for_not_canonicalizable_1(%A : memref<i64>) {
+  // This should not canonicalize because: 4 - %i may take the value 1 < 2.
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c4 = constant 4 : index
+
+  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
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_canonicalize_partly
+//       CHECK:   scf.for
+//       CHECK:     affine.apply
+//       CHECK:     index_cast
+func @scf_for_canonicalize_partly(%A : memref<i64>) {
+  // This should canonicalize only partly: 256 - %i <= 256.
+  %c1 = constant 1 : index
+  %c16 = constant 16 : index
+  %c256 = constant 256 : index
+
+  scf.for %i = %c1 to %c256 step %c16 {
+    %1 = affine.min affine_map<(d0) -> (256, 256 - d0)> (%i)
+    %2 = index_cast %1: index to i64
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_not_canonicalizable_2
+//       CHECK: scf.for
+//       CHECK:   affine.min
+//       CHECK:   index_cast
+func @scf_for_not_canonicalizable_2(%A : memref<i64>, %step : index) {
+  // 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.
+  %c0 = constant 0 : index
+
+  %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
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_not_canonicalizable_3
+//       CHECK: scf.for
+//       CHECK:   affine.min
+//       CHECK:   index_cast
+func @scf_for_not_canonicalizable_3(%A : memref<i64>, %step : index) {
+  // 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.
+  %c0 = constant 0 : index
+
+  %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
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_for_invalid_loop
+//       CHECK: scf.for
+//       CHECK:   affine.min
+//       CHECK:   index_cast
+func @scf_for_invalid_loop(%A : memref<i64>, %step : index) {
+  // This is an invalid loop. It should not be touched by the canonicalization
+  // pattern.
+  %c1 = constant 1 : index
+  %c7 = constant 7 : index
+  %c256 = constant 256 : index
+
+  scf.for %i = %c256 to %c1 step %c1 {
+    %1 = affine.min affine_map<(d0)[s0] -> (s0 + d0, 0)> (%i)[%c7]
+    %2 = index_cast %1: index to i64
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_parallel_canonicalize_min_1
+//       CHECK:   %[[C2:.*]] = constant 2 : i64
+//       CHECK:   scf.parallel
+//  CHECK-NEXT:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
+func @scf_parallel_canonicalize_min_1(%A : memref<i64>) {
+  %c0 = constant 0 : index
+  %c2 = constant 2 : index
+  %c4 = constant 4 : index
+
+  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
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_parallel_canonicalize_min_2
+//       CHECK:   %[[C2:.*]] = constant 2 : i64
+//       CHECK:   scf.parallel
+//  CHECK-NEXT:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
+func @scf_parallel_canonicalize_min_2(%A : memref<i64>) {
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c7 = constant 7 : index
+
+  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
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}