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

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.

Differential Revision: https://reviews.llvm.org/D107731

Author: matthias-springer

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

@@ -28,6 +28,10 @@ std::unique_ptr<Pass> createForLoopSpecializationPass();
 /// 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();
 

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

@@ -17,6 +17,16 @@ def SCFBufferize : FunctionPass<"scf-bufferize"> {
   let dependentDialects = ["memref::MemRefDialect"];
 }
 
+// Note: Making this a canonicalization pattern would require a dependency
+// of the SCF dialect on the Affine dialect or vice versa.
+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.";

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

@@ -13,10 +13,12 @@
 #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;
@@ -26,6 +28,7 @@ class TypeConverter;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 class Operation;
+class Value;
 
 namespace scf {
 
@@ -34,6 +37,21 @@ 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.
@@ -149,6 +167,11 @@ struct PipeliningOption {
 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
 

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

@@ -382,6 +382,77 @@ LogicalResult mlir::scf::peelAndCanonicalizeForLoop(RewriterBase &rewriter,
   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);
+
+    if (failed(constraints.addBound(FlatAffineConstraints::UB, dimIv, map)))
+      return failure();
+  }
+
+  return canonicalizeAffineMinOp(rewriter, minOp, constraints);
+}
+
 static constexpr char kPeeledLoopLabel[] = "__peeled_loop__";
 static constexpr char kPartialIterationLabel[] = "__partial_iteration__";
 
@@ -423,6 +494,39 @@ struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
   /// the direct parent.
   bool skipPartial;
 };
+
+/// 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 {
@@ -456,8 +560,24 @@ struct ForLoopPeeling : public SCFForLoopPeelingBase<ForLoopPeeling> {
     });
   }
 };
+
+struct AffineMinSCFCanonicalization
+    : public AffineMinSCFCanonicalizationBase<AffineMinSCFCanonicalization> {
+  void runOnFunction() override {
+    FuncOp funcOp = getFunction();
+    MLIRContext *ctx = funcOp.getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<AffineMinSCFCanonicalizationPattern>(ctx);
+    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
+      signalPassFailure();
+  }
+};
 } // namespace
 
+std::unique_ptr<Pass> mlir::createAffineMinSCFCanonicalizationPass() {
+  return std::make_unique<AffineMinSCFCanonicalization>();
+}
+
 std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {
   return std::make_unique<ParallelLoopSpecialization>();
 }
@@ -469,3 +589,8 @@ std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() {
 std::unique_ptr<Pass> mlir::createForLoopPeelingPass() {
   return std::make_unique<ForLoopPeeling>();
 }
+
+void mlir::scf::populateSCFLoopBodyCanonicalizationPatterns(
+    RewritePatternSet &patterns) {
+  patterns.insert<AffineMinSCFCanonicalizationPattern>(patterns.getContext());
+}