[Mlir-commits] [mlir] [mlir][scf] add support for unsigned loop for back peeling (PR #234435)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 16:56:19 PDT 2026
https://github.com/StarOne01 created https://github.com/llvm/llvm-project/pull/234435
- add support for unsigned loop for back peeling
- fix sub overflow bug in the signed path
- remove unsigned operation check in dispatcher
- return early if unsigned in `peelForLoopAndSimplifyBounds` since affine.min/max does not hold for unsigned loops
- add tests for edge cases in unsigned, sub overflow and no bound simplification for unsigned
from https://github.com/llvm/llvm-project/pull/231100
claude helped me with the tests and correctness.
>From 6a001aee185927ac296e246d63959815c409d37d Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Thu, 22 Oct 2026 05:21:30 +0530
Subject: [PATCH] [mlir][scf] add support for unsigned loop for back peeling
- add support for unsigned loop for back peeling
- fix sub overflow bug in the signed path
- remove unsigned operation check in dispatcher
- return early if unsigned in `peelForLoopAndSimplifyBounds` since affine.min/max does not hold for unsigned loops
- add tests for edge cases in unsigned, sub overflow and no bound simplification for unsigned
---
.../SCF/Transforms/LoopSpecialization.cpp | 120 ++++++++-----
mlir/test/Dialect/SCF/for-loop-peeling.mlir | 166 ++++++++++++++++++
2 files changed, 244 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 0a1e8ffc8a046b..46bf6883fcd89a 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -126,52 +126,86 @@ static void specializeForLoopForUnrolling(ForOp op) {
static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
ForOp &partialIteration, Value &splitBound) {
RewriterBase::InsertionGuard guard(b);
- auto lbInt = getConstantIntValue(forOp.getLowerBound());
- auto ubInt = getConstantIntValue(forOp.getUpperBound());
- auto stepInt = getConstantIntValue(forOp.getStep());
+ Value lbVal = forOp.getLowerBound();
+ Value ubVal = forOp.getUpperBound();
+ Value stepVal = forOp.getStep();
+ b.setInsertionPoint(forOp);
+ auto loc = forOp.getLoc();
- // No specialization necessary if step size is 1. Also bail out in case of an
- // invalid zero or negative step which might have happened during folding.
- if (stepInt && *stepInt <= 1)
- return failure();
+ // Check whether we have a unsigned loop
+ if (forOp.getUnsignedCmp()) {
+ IntegerAttr lbAttr, ubAttr, stepAttr;
+ if (matchPattern(lbVal, m_Constant(&lbAttr)) &&
+ matchPattern(ubVal, m_Constant(&ubAttr)) &&
+ matchPattern(stepVal, m_Constant(&stepAttr))) {
- // No specialization necessary if step already divides upper bound evenly.
- // Fast path: lb, ub and step are constants.
- if (lbInt && ubInt && stepInt && (*ubInt - *lbInt) % *stepInt == 0)
- return failure();
+ // Get the constants as APInt at the loop’s own bit width
+ APInt lb = lbAttr.getValue();
+ APInt ub = ubAttr.getValue();
+ APInt step = stepAttr.getValue();
- // Only the dynamic path computes the peeling bound with affine.apply, which
- // accepts only index operands.
- if ((!lbInt || !ubInt || !stepInt) &&
- !forOp.getInductionVar().getType().isIndex())
- return failure();
+ if (step.ule(1) || ub.ule(lb) || (ub - lb).urem(step).isZero())
+ return failure();
+ }
+
+ // split = ub - ((ub - min(lb, ub)) remui step). The min makes an empty loop
+ // (ub <= lb) give a zero remainder, so split = ub and both loops stay
+ // empty. The wrapping subtraction is exact for ub > lb, even above the
+ // signed range.
+ Value lbEff = b.createOrFold<arith::MinUIOp>(loc, lbVal, ubVal);
+ Value diff = b.createOrFold<arith::SubIOp>(loc, ubVal, lbEff);
+ Value rem = b.createOrFold<arith::RemUIOp>(loc, diff, stepVal);
+ splitBound = b.createOrFold<arith::SubIOp>(loc, ubVal, rem);
- // Slow path: Examine the ops that define lb, ub and step.
- AffineExpr sym0, sym1, sym2;
- bindSymbols(b.getContext(), sym0, sym1, sym2);
- SmallVector<Value> operands{forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep()};
- AffineMap map = AffineMap::get(0, 3, {(sym1 - sym0) % sym2});
- affine::fullyComposeAffineMapAndOperands(&map, &operands);
- if (auto constExpr = dyn_cast<AffineConstantExpr>(map.getResult(0)))
- if (constExpr.getValue() == 0)
+ } else {
+ auto lbInt = getConstantIntValue(lbVal);
+ auto ubInt = getConstantIntValue(ubVal);
+ auto stepInt = getConstantIntValue(stepVal);
+ // No specialization necessary if step size is 1. Also bail out in case of
+ // an invalid zero or negative step which might have happened during
+ // folding.
+ if (stepInt && *stepInt <= 1)
return failure();
- // New upper bound: max(%lb, %ub - (%ub - %lb) mod %step). Affine `mod` is
- // non-negative, so without the clamp an empty loop (%lb >= %ub) gets a split
- // bound below %lb and its partial iteration executes.
- auto modMap = AffineMap::get(0, 3, {sym1 - ((sym1 - sym0) % sym2), sym0},
- b.getContext());
- b.setInsertionPoint(forOp);
- auto loc = forOp.getLoc();
- splitBound = b.createOrFold<AffineMaxOp>(loc, modMap,
- ValueRange{forOp.getLowerBound(),
- forOp.getUpperBound(),
- forOp.getStep()});
- if (splitBound.getType() != forOp.getLowerBound().getType())
- splitBound = b.createOrFold<arith::IndexCastOp>(
- loc, forOp.getLowerBound().getType(), splitBound);
+ // No specialization necessary if step already divides upper bound evenly.
+ // Fast path: lb, ub and step are constants.
+ if (lbInt && ubInt && stepInt) {
+ int64_t diff;
+ // bail out if the subtraction overflows
+ if (llvm::SubOverflow(*ubInt, *lbInt, diff) || diff % *stepInt == 0)
+ return failure();
+ }
+
+ // Only the dynamic path computes the peeling bound with affine.apply, which
+ // accepts only index operands.
+ if ((!lbInt || !ubInt || !stepInt) &&
+ !forOp.getInductionVar().getType().isIndex())
+ return failure();
+
+ // Slow path: Examine the ops that define lb, ub and step.
+ AffineExpr sym0, sym1, sym2;
+ bindSymbols(b.getContext(), sym0, sym1, sym2);
+ SmallVector<Value> operands{forOp.getLowerBound(), forOp.getUpperBound(),
+ forOp.getStep()};
+ AffineMap map = AffineMap::get(0, 3, {(sym1 - sym0) % sym2});
+ affine::fullyComposeAffineMapAndOperands(&map, &operands);
+ if (auto constExpr = dyn_cast<AffineConstantExpr>(map.getResult(0)))
+ if (constExpr.getValue() == 0)
+ return failure();
+ // New upper bound: max(%lb, %ub - (%ub - %lb) mod %step). Affine `mod` is
+ // non-negative, so without the clamp an empty loop (%lb >= %ub) gets a
+ // split bound below %lb and its partial iteration executes.
+ auto modMap = AffineMap::get(0, 3, {sym1 - ((sym1 - sym0) % sym2), sym0},
+ b.getContext());
+ splitBound = b.createOrFold<AffineMaxOp>(loc, modMap,
+ ValueRange{forOp.getLowerBound(),
+ forOp.getUpperBound(),
+ forOp.getStep()});
+ if (splitBound.getType() != forOp.getLowerBound().getType())
+ splitBound = b.createOrFold<arith::IndexCastOp>(
+ loc, forOp.getLowerBound().getType(), splitBound);
+ }
// Create ForOp for partial iteration.
b.setInsertionPointAfter(forOp);
partialIteration = cast<ForOp>(b.clone(*forOp.getOperation()));
@@ -220,6 +254,11 @@ LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
if (failed(peelForLoop(rewriter, forOp, partialIteration, splitBound)))
return failure();
+ // The affine.min/max simplification relies on signed reasoning, which does
+ // not hold for unsigned loops, so leave those ops untouched.
+ if (forOp.getUnsignedCmp())
+ return success();
+
// Rewrite affine.min and affine.max ops.
rewriteAffineOpAfterPeeling(rewriter, forOp, partialIteration, previousUb);
@@ -314,9 +353,6 @@ struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
LogicalResult matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const override {
- if (!peelFront && forOp.getUnsignedCmp())
- return rewriter.notifyMatchFailure(forOp,
- "unsigned loops are not supported");
// Do not peel already peeled loops.
if (forOp->hasDiscardableAttr(kPeeledLoopLabel))
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index 86d40f176e63e0..1906b7039ec160 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -399,3 +399,169 @@ func.func @non_index_loop_bounds(%init: i32, %n: i32) -> i32 {
}
return %r : i32
}
+
+// -----
+
+// Fully dynamic i32 bounds: nothing can fold. Unsigned min/sub/remui chain,
+// iter_args threaded from the main loop into the partial loop.
+
+// CHECK-LABEL: func @unsigned_back_peel_dynamic(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32, %[[STEP:.*]]: i32
+// CHECK: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: %[[LBEFF:.*]] = arith.minui %[[LB]], %[[UB]] : i32
+// CHECK: %[[DIFF:.*]] = arith.subi %[[UB]], %[[LBEFF]] : i32
+// CHECK: %[[REM:.*]] = arith.remui %[[DIFF]], %[[STEP]] : i32
+// CHECK: %[[SPLIT:.*]] = arith.subi %[[UB]], %[[REM]] : i32
+// CHECK: %[[MAIN:.*]] = scf.for unsigned %[[IV:.*]] = %[[LB]] to %[[SPLIT]] step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0]]) -> (i32) : i32 {
+// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[IV]] : i32
+// CHECK: scf.yield %[[ADD]] : i32
+// CHECK: }
+// CHECK: %[[RES:.*]] = scf.for unsigned %[[IV2:.*]] = %[[SPLIT]] to %[[UB]] step %[[STEP]] iter_args(%[[ACC2:.*]] = %[[MAIN]]) -> (i32) : i32 {
+// CHECK: %[[ADD2:.*]] = arith.addi %[[ACC2]], %[[IV2]] : i32
+// CHECK: scf.yield %[[ADD2]] : i32
+// CHECK: }
+// CHECK: return %[[RES]] : i32
+func.func @unsigned_back_peel_dynamic(%lb: i32, %ub: i32, %step: i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %r = scf.for unsigned %iv = %lb to %ub step %step iter_args(%acc = %c0) -> (i32) : i32 {
+ %0 = arith.addi %acc, %iv : i32
+ scf.yield %0 : i32
+ }
+ return %r : i32
+}
+
+// -----
+
+// Constants that only make sense unsigned: lb = 0x7FFFFFF8, ub = 0xFFFFFFF9
+// (printed -7), step 4. (ub - lb) = 2147483649, remainder 1, so split =
+// 0xFFFFFFF8 (-8). The main loop runs 536870912 iterations. The partial loop
+// is exactly one iteration, so the canonicalizer inlines it with iv = split.
+
+func.func private @effect(i32)
+
+// CHECK-LABEL: func @unsigned_back_peel_above_signed_max
+// CHECK-DAG: %[[LB:.*]] = arith.constant 2147483640 : i32
+// CHECK-DAG: %[[STEP:.*]] = arith.constant 4 : i32
+// CHECK-DAG: %[[SPLIT:.*]] = arith.constant -8 : i32
+// CHECK: scf.for unsigned %{{.*}} = %[[LB]] to %[[SPLIT]] step %[[STEP]] : i32 {
+// CHECK: call @effect(%{{.*}}) : (i32) -> ()
+// CHECK: }
+// CHECK: call @effect(%[[SPLIT]]) : (i32) -> ()
+func.func @unsigned_back_peel_above_signed_max() {
+ %lb = arith.constant 2147483640 : i32
+ %ub = arith.constant -7 : i32
+ %step = arith.constant 4 : i32
+ scf.for unsigned %iv = %lb to %ub step %step : i32 {
+ func.call @effect(%iv) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+// Empty unsigned loop: lb = 0xFFFFFFFF, ub = 0 (ub <= lb). Signed code would
+// read this as -1 to 0, which has one iteration. After peeling and
+// canonicalization no iteration may remain.
+
+func.func private @effect(i32)
+
+// CHECK-LABEL: func @unsigned_back_peel_empty_loop
+// CHECK-NOT: call @effect
+// CHECK: return
+func.func @unsigned_back_peel_empty_loop() {
+ %lb = arith.constant -1 : i32
+ %ub = arith.constant 0 : i32
+ %step = arith.constant 2 : i32
+ scf.for unsigned %iv = %lb to %ub step %step : i32 {
+ func.call @effect(%iv) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+// Step divides the range (0x80000000 - 0 = 2^29 * 4): nothing to peel, so one
+// loop stays. 0x80000000 prints as -2147483648, which signed code would read
+// as an empty loop.
+
+func.func private @effect(i32)
+
+// CHECK-LABEL: func @unsigned_back_peel_step_divides_range
+// CHECK: scf.for unsigned
+// CHECK: call @effect
+// CHECK: }
+// CHECK-NOT: scf.for
+// CHECK-NOT: call @effect
+func.func @unsigned_back_peel_step_divides_range() {
+ %lb = arith.constant 0 : i32
+ %ub = arith.constant -2147483648 : i32
+ %step = arith.constant 4 : i32
+ scf.for unsigned %iv = %lb to %ub step %step : i32 {
+ func.call @effect(%iv) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+// Constant step of 1 with dynamic bounds: step 1 always divides the range,
+// so there is nothing to peel. One loop stays and no split bound is computed.
+
+func.func private @effect(i32)
+
+// CHECK-LABEL: func @unsigned_back_peel_step_one_dynamic_bounds(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32
+// CHECK: %[[STEP:.*]] = arith.constant 1 : i32
+// CHECK: scf.for unsigned %{{.*}} = %[[LB]] to %[[UB]] step %[[STEP]] : i32 {
+// CHECK: call @effect
+// CHECK: }
+// CHECK-NOT: scf.for
+// CHECK-NOT: arith.remui
+func.func @unsigned_back_peel_step_one_dynamic_bounds(%lb: i32, %ub: i32) {
+ %step = arith.constant 1 : i32
+ scf.for unsigned %iv = %lb to %ub step %step : i32 {
+ func.call @effect(%iv) : (i32) -> ()
+ }
+ return
+}
+
+// -----
+
+// Unsigned loop with an affine.min in the body: the loop is peeled, but the
+// affine.min is not simplified in either loop (the simplification is signed).
+
+#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
+func.func private @use(index)
+
+// CHECK-LABEL: func @unsigned_no_bound_simplification
+// CHECK: scf.for unsigned
+// CHECK: affine.min
+// CHECK: scf.for unsigned
+// CHECK: affine.min
+func.func @unsigned_no_bound_simplification(%lb: index, %ub: index, %step: index) {
+ scf.for unsigned %iv = %lb to %ub step %step {
+ %s = affine.min #map(%ub, %iv)[%step]
+ func.call @use(%s) : (index) -> ()
+ }
+ return
+}
+
+// -----
+
+// Signed empty loop whose range is not a multiple of the step: lb=10, ub=5,
+// step=4. Without the clamp to lb, split was 2 and the partial loop ran iv=2.
+
+func.func private @effect(index)
+
+// CHECK-LABEL: func @signed_empty_loop_not_executed
+// CHECK-NOT: call @effect
+// CHECK: return
+func.func @signed_empty_loop_not_executed() {
+ %lb = arith.constant 10 : index
+ %ub = arith.constant 5 : index
+ %step = arith.constant 4 : index
+ scf.for %iv = %lb to %ub step %step {
+ func.call @effect(%iv) : (index) -> ()
+ }
+ return
+}
\ No newline at end of file
More information about the Mlir-commits
mailing list