[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))
+        return failure();
+    }
+    b.setInsertionPoint(forOp);
+    auto loc = forOp.getLoc();
+    Value lbPlusStep = b.createOrFold<arith::AddIOp>(loc, forOp.getLowerBound(),
+                                                     forOp.getStep());
+    splitBound =
----------------
matthias-springer wrote:

Coding agent analysis suggests that the split bound should be at least the original LB. Otherwise, a loop with zero iterations may execute an iteration after peeling. This seems to be a problem of both the unsigned code path and the pre-existing signed code path.

E.g.:
```
  func.func @repro() {
    %lb = arith.constant 9223372036854775807 : i64
    %ub = arith.constant -2 : i64
    %step = arith.constant 9223372036854775806 : i64
    scf.for %iv = %lb to %ub step %step : i64 {
      func.call @effect(%iv) : (i64) -> ()
    }
    return
  }
```

After peeling:
```
  %split = arith.constant -3 : i64

  scf.for %iv = %lb to %split step %step : i64 {
    func.call @effect(%iv) : (i64) -> ()
  }
  scf.for %iv = %split to %ub step %step : i64 {
    func.call @effect(%iv) : (i64) -> ()
  }
```

I don't have time to investigate this in detail right now. But maybe this should be fixed before adding the unsigned case.  (And then write the unsigned code path in the same way.)


We should also double-check if the C++ API makes (and documents) any assumptions such as "the to-be-peeled loop must have at least one iteration". Maybe it's a non-issue in that case.


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