[mlir][SCF][Utils][NFC] Make some utils public for better reuse

[springerm]

Summary

These functions will be used in a subsequent change. Also some minor refactoring.

mlir/include/mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h

 /// Match "for loop"-like operations: If the first parameter is an iteration
 /// variable, return lower/upper bounds via the second/third parameter and the
 /// step size via the last parameter. The function should return `success` in
 /// that case. If the first parameter is not an iteration variable, return
 /// `failure`.
 using LoopMatcherFn = function_ref<LogicalResult(
     Value, OpFoldResult &, OpFoldResult &, OpFoldResult &)>;
 
+/// Match "for loop"-like operations from the SCF dialect.
+LogicalResult matchForLikeLoop(Value iv, OpFoldResult &lb, OpFoldResult &ub,
+                               OpFoldResult &step);
+
+/// Populate the given constraint set with induction variable constraints of a
+/// "for" loop with the given range and step.
+LogicalResult addLoopRangeConstraints(FlatAffineValueConstraints &cstr,
+                                      Value iv, OpFoldResult lb,
+                                      OpFoldResult ub, OpFoldResult step);
+
 /// Try to canonicalize the given affine.min/max operation in the context of
 /// for `loops` with a known range.
 ///
 /// `loopMatcher` is used to retrieve loop bounds and the step size for a given
 /// iteration variable.
 ///
 /// Note: `loopMatcher` allows this function to be used with any "for loop"-like
 /// operation (scf.for, scf.parallel and even ops defined in other dialects).

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

 /// Canonicalize AffineMinOp/AffineMaxOp operations in the context of scf.for
 /// and scf.parallel loops with a known range.
 template <typename OpTy>
 struct AffineOpSCFCanonicalizationPattern : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy op,
                                 PatternRewriter &rewriter) const override {
-    auto loopMatcher = [](Value iv, OpFoldResult &lb, OpFoldResult &ub,
-                          OpFoldResult &step) {
-      if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
-        lb = forOp.getLowerBound();
-        ub = forOp.getUpperBound();
-        step = forOp.getStep();
-        return success();
-      }
-      if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
-        for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
-          if (parOp.getInductionVars()[idx] == iv) {
-            lb = parOp.getLowerBound()[idx];
-            ub = parOp.getUpperBound()[idx];
-            step = parOp.getStep()[idx];
-            return success();
-          }
-        }
-        return failure();
-      }
-      if (scf::ForallOp forallOp = scf::getForallOpThreadIndexOwner(iv)) {
-        for (int64_t idx = 0; idx < forallOp.getRank(); ++idx) {
-          if (forallOp.getInductionVar(idx) == iv) {
-            lb = forallOp.getMixedLowerBound()[idx];
-            ub = forallOp.getMixedUpperBound()[idx];
-            step = forallOp.getMixedStep()[idx];
-            return success();
-          }
-        }
-        return failure();
-      }
-      return failure();
-    };
-
-    return scf::canonicalizeMinMaxOpInLoop(rewriter, op, loopMatcher);
+    return scf::canonicalizeMinMaxOpInLoop(rewriter, op, scf::matchForLikeLoop);
   }
 };
 
 struct SCFForLoopCanonicalization
     : public impl::SCFForLoopCanonicalizationBase<SCFForLoopCanonicalization> {
   void runOnOperation() override {
     auto *parentOp = getOperation();
     MLIRContext *ctx = parentOp->getContext();

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

 //===- AffineCanonicalizationUtils.cpp - Affine Canonicalization in SCF ---===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Utility functions to canonicalize affine ops within SCF op regions.
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
-#include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "mlir-scf-affine-utils"
 
 using namespace mlir;
 using namespace presburger;
 
+LogicalResult scf::matchForLikeLoop(Value iv, OpFoldResult &lb,
+                                    OpFoldResult &ub, OpFoldResult &step) {
+  if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
+    lb = forOp.getLowerBound();
+    ub = forOp.getUpperBound();
+    step = forOp.getStep();
+    return success();
+  }
+  if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {
+    for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {
+      if (parOp.getInductionVars()[idx] == iv) {
+        lb = parOp.getLowerBound()[idx];
+        ub = parOp.getUpperBound()[idx];
+        step = parOp.getStep()[idx];
+        return success();
+      }
+    }
+    return failure();
+  }
+  if (scf::ForallOp forallOp = scf::getForallOpThreadIndexOwner(iv)) {
+    for (int64_t idx = 0; idx < forallOp.getRank(); ++idx) {
+      if (forallOp.getInductionVar(idx) == iv) {
+        lb = forallOp.getMixedLowerBound()[idx];
+        ub = forallOp.getMixedUpperBound()[idx];
+        step = forallOp.getMixedStep()[idx];
+        return success();
+      }
+    }
+    return failure();
+  }
+  return failure();
+}
+
 static FailureOr<AffineApplyOp>
 canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op,
                      FlatAffineValueConstraints constraints) {
   RewriterBase::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(op);
   FailureOr<AffineValueMap> simplified =
       mlir::simplifyConstrainedMinMaxOp(op, std::move(constraints));
   if (failed(simplified))
     return failure();
   return rewriter.replaceOpWithNewOp<AffineApplyOp>(
       op, simplified->getAffineMap(), simplified->getOperands());
 }
 
-static LogicalResult
-addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
-                        OpFoldResult lb, OpFoldResult ub, OpFoldResult step,
-                        RewriterBase &rewriter) {
+LogicalResult scf::addLoopRangeConstraints(FlatAffineValueConstraints &cstr,
+                                           Value iv, OpFoldResult lb,
+                                           OpFoldResult ub, OpFoldResult step) {
+  Builder b(iv.getContext());
+
   // IntegerPolyhedron does not support semi-affine expressions.
   // Therefore, only constant step values are supported.
   auto stepInt = getConstantIntValue(step);
   if (!stepInt)
     return failure();
 
-  unsigned dimIv = constraints.appendDimVar(iv);
+  unsigned dimIv = cstr.appendDimVar(iv);
   auto lbv = lb.dyn_cast<Value>();
-  unsigned symLb = lbv ? constraints.appendSymbolVar(lbv)
-                       : constraints.appendSymbolVar(/*num=*/1);
+  unsigned symLb =
+      lbv ? cstr.appendSymbolVar(lbv) : cstr.appendSymbolVar(/*num=*/1);
   auto ubv = ub.dyn_cast<Value>();
-  unsigned symUb = ubv ? constraints.appendSymbolVar(ubv)
-                       : constraints.appendSymbolVar(/*num=*/1);
+  unsigned symUb =
+      ubv ? cstr.appendSymbolVar(ubv) : cstr.appendSymbolVar(/*num=*/1);
 
   // If loop lower/upper bounds are constant: Add EQ constraint.
   std::optional<int64_t> lbInt = getConstantIntValue(lb);
   std::optional<int64_t> ubInt = getConstantIntValue(ub);
   if (lbInt)
-    constraints.addBound(IntegerPolyhedron::EQ, symLb, *lbInt);
+    cstr.addBound(IntegerPolyhedron::EQ, symLb, *lbInt);
   if (ubInt)
-    constraints.addBound(IntegerPolyhedron::EQ, symUb, *ubInt);
+    cstr.addBound(IntegerPolyhedron::EQ, symUb, *ubInt);
 
   // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
-  SmallVector<int64_t> ineqLb(constraints.getNumCols(), 0);
+  SmallVector<int64_t> ineqLb(cstr.getNumCols(), 0);
   ineqLb[dimIv] = 1;
   ineqLb[symLb] = -1;
-  constraints.addInequality(ineqLb);
+  cstr.addInequality(ineqLb);
 
   // Upper bound
   AffineExpr ivUb;
   if (lbInt && ubInt && (*lbInt + *stepInt >= *ubInt)) {
     // The loop has at most one iteration.
     // iv < lb + 1
     // TODO: Try to derive this constraint by simplifying the expression in
     // the else-branch.
-    ivUb =
-        rewriter.getAffineSymbolExpr(symLb - constraints.getNumDimVars()) + 1;
+    ivUb = b.getAffineSymbolExpr(symLb - cstr.getNumDimVars()) + 1;
   } else {
     // The loop may have more than one iteration.
     // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
     AffineExpr exprLb =
-        lbInt
-            ? rewriter.getAffineConstantExpr(*lbInt)
-            : rewriter.getAffineSymbolExpr(symLb - constraints.getNumDimVars());
+        lbInt ? b.getAffineConstantExpr(*lbInt)
+              : b.getAffineSymbolExpr(symLb - cstr.getNumDimVars());
     AffineExpr exprUb =
-        ubInt
-            ? rewriter.getAffineConstantExpr(*ubInt)
-            : rewriter.getAffineSymbolExpr(symUb - constraints.getNumDimVars());
+        ubInt ? b.getAffineConstantExpr(*ubInt)
+              : b.getAffineSymbolExpr(symUb - cstr.getNumDimVars());
     ivUb = exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
   }
   auto map = AffineMap::get(
-      /*dimCount=*/constraints.getNumDimVars(),
-      /*symbolCount=*/constraints.getNumSymbolVars(), /*result=*/ivUb);
+      /*dimCount=*/cstr.getNumDimVars(),
+      /*symbolCount=*/cstr.getNumSymbolVars(), /*result=*/ivUb);
 
-  return constraints.addBound(IntegerPolyhedron::UB, dimIv, map);
+  return cstr.addBound(IntegerPolyhedron::UB, dimIv, map);
 }
 
 /// Canonicalize min/max operations in the context of for loops with a known
 /// range. Call `canonicalizeMinMaxOp` and add the following constraints to
 /// the constraint system (along with the missing dimensions):
 ///
 /// * iv >= lb
 /// * iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
for (Value operand : op->getOperands()) {
     // If `operand` is an iteration variable: Find corresponding loop
     // bounds and step.
     Value iv = operand;
     OpFoldResult lb, ub, step;
     if (failed(loopMatcher(operand, lb, ub, step)))
       continue;
     allIvs.insert(iv);
 
-    if (failed(
-            addLoopRangeConstraints(constraints, iv, lb, ub, step, rewriter)))
+    if (failed(addLoopRangeConstraints(constraints, iv, lb, ub, step)))
       return failure();
   }
 
   return canonicalizeMinMaxOp(rewriter, op, constraints);
 }
 
 /// Try to simplify the given affine.min/max operation `op` after loop peeling.
 /// This function can simplify min/max operations such as (ub is the previous