[Mlir-commits] [mlir] [mlir][scf] support front peeling for unsigned loop (PR #233238)

[Matthias Springer]

================
@@ -233,51 +233,66 @@ LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
 /// original lower bound plus the step (i.e. skips the first iteration).
 LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
                                                    ForOp &firstIteration) {
-  // TODO: Add support for unsigned loops. The current logic
-  // uses `affine.min` and `arith.index_cast`, which is incorrect for
-  // unsigned comparisons.
-  if (forOp.getUnsignedCmp())
-    return failure();
-
   RewriterBase::InsertionGuard guard(b);
-  auto lbInt = getConstantIntValue(forOp.getLowerBound());
-  auto ubInt = getConstantIntValue(forOp.getUpperBound());
-  auto stepInt = getConstantIntValue(forOp.getStep());
-
-  // Peeling is not needed if there is at most one iteration.
-  // Per the scf.for semantics, the trip count is
-  // max(0, ceil((UB - LB) / Step)), so for a positive step there is more than
-  // one iteration exactly when UB - LB > Step.
-  if (lbInt && ubInt && stepInt && *ubInt - *lbInt <= *stepInt)
-    return failure();
-
-  // The peeling bound (%lb + %step) is computed with affine.apply, which
-  // accepts only index operands. %ub does not feed into this bound, so only
-  // %lb and %step need to be constant to guarantee the affine.apply (see below)
-  // folds away before its (non-index) operand types matter.
-  if ((!lbInt || !stepInt) && !forOp.getInductionVar().getType().isIndex())
-    return failure();
-
-  AffineExpr lbSymbol, stepSymbol;
-  bindSymbols(b.getContext(), lbSymbol, stepSymbol);
+  Value splitBound;
+  // Check whether we have a unsigned loop
+  if (forOp.getUnsignedCmp()) {
+    IntegerAttr lbAttr, ubAttr, stepAttr;
+    if (matchPattern(forOp.getLowerBound(), m_Constant(&lbAttr)) &&
+        matchPattern(forOp.getUpperBound(), m_Constant(&ubAttr)) &&
+        matchPattern(forOp.getStep(), m_Constant(&stepAttr))) {
+
+      // Get the constants as APInt at the loop’s own bit width
+      APInt lb = lbAttr.getValue();
+      APInt ub = ubAttr.getValue();
+      APInt step = stepAttr.getValue();
+
+      // bail out if ub <= lb or ub - lb <= step (unsigned)
+      if (ub.ule(lb) || (ub - lb).ule(step))
----------------
matthias-springer wrote:

Can the subtraction overflow? In such a case, we should probably also abort. (Same for the signed case.)

https://github.com/llvm/llvm-project/pull/233238