[Mlir-commits] [mlir] [mlir][scf] support front peeling for unsigned loop (PR #233238)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 18 07:02:52 PDT 2026
https://github.com/StarOne01 created https://github.com/llvm/llvm-project/pull/233238
add front peeling support for unsigned loop and add tests.
we branch out if we have a unsigned loop, if constants we check and bail out if the loop has at most one iteration else we compute the split bound as `arith.minui(arith.addi(lb, step), ub)`, and they are compared unsigned.
from #231100
claude helped me with the tests
>From d50ede9db915ee1b79cb2cccfaeb6508b1a550c8 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sun, 18 Oct 2026 17:54:41 +0530
Subject: [PATCH 1/2] [mlir][scf] add support for unsigned loop for front
peeling
---
.../SCF/Transforms/LoopSpecialization.cpp | 99 +++++++++++--------
1 file changed, 57 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 67e3228389a910..a3032217d33963 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -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 =
+ b.createOrFold<arith::MinUIOp>(loc, lbPlusStep, forOp.getUpperBound());
+ } else {
+ 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();
- // New lower bound for main loop: min(%lb + %step, %ub). Without the minimum
- // the peeled first iteration spans [%lb, %lb + %step) and executes even when
- // the source loop is empty (%lb >= %ub).
- AffineExpr ubSymbol;
- bindSymbols(b.getContext(), lbSymbol, stepSymbol, ubSymbol);
- auto ubMap =
- AffineMap::get(0, 3, {lbSymbol + stepSymbol, ubSymbol}, b.getContext());
- b.setInsertionPoint(forOp);
- auto loc = forOp.getLoc();
- Value splitBound = b.createOrFold<AffineMinOp>(
- loc, ubMap,
- ValueRange{forOp.getLowerBound(), forOp.getStep(),
- forOp.getUpperBound()});
- if (splitBound.getType() != forOp.getUpperBound().getType())
- splitBound = b.createOrFold<arith::IndexCastOp>(
- loc, forOp.getUpperBound().getType(), splitBound);
+ // 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();
+ // New lower bound for main loop: min(%lb + %step, %ub). Without the minimum
+ // the peeled first iteration spans [%lb, %lb + %step) and executes even
+ // when the source loop is empty (%lb >= %ub).
+ AffineExpr ubSymbol, lbSymbol, stepSymbol;
+ bindSymbols(b.getContext(), lbSymbol, stepSymbol, ubSymbol);
+ auto ubMap =
+ AffineMap::get(0, 3, {lbSymbol + stepSymbol, ubSymbol}, b.getContext());
+ b.setInsertionPoint(forOp);
+ auto loc = forOp.getLoc();
+ splitBound = b.createOrFold<AffineMinOp>(loc, ubMap,
+ ValueRange{forOp.getLowerBound(),
+ forOp.getStep(),
+ forOp.getUpperBound()});
+ if (splitBound.getType() != forOp.getUpperBound().getType())
+ splitBound = b.createOrFold<arith::IndexCastOp>(
+ loc, forOp.getUpperBound().getType(), splitBound);
+ }
// Peel the first iteration.
firstIteration = cast<ForOp>(b.clone(*forOp.getOperation()));
b.modifyOpInPlace(firstIteration, [&]() {
>From 7753f783c226ae2b15282df60a7e25ea8f35279e Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sun, 18 Oct 2026 19:08:58 +0530
Subject: [PATCH 2/2] [mlir][scf] add tests for unsigned loop peeling with
dynamic and constant bounds
---
.../Dialect/SCF/for-loop-peeling-front.mlir | 12 +--
mlir/test/Dialect/SCF/transform-ops.mlir | 93 +++++++++++++++++++
2 files changed, 95 insertions(+), 10 deletions(-)
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
index 84b561a8571f69..13a98309993adc 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
@@ -250,19 +250,11 @@ func.func @non_index_loop_bounds(%init: i32, %lb: i32, %n: i32) -> i32 {
// was skipped. See https://github.com/llvm/llvm-project/issues/231062
// CHECK-LABEL: func.func @peel_front_step_near_ub(
-// CHECK-SAME: %[[MEMREF:[^:]*]]: memref<2xindex>) {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C16777216:.*]] = arith.constant 16777216 : index
// CHECK-DAG: %[[C16777217:.*]] = arith.constant 16777217 : index
-// CHECK: scf.for %[[IV0:.*]] = %[[C0]] to %[[C16777216]] step %[[C16777216]] {
-// CHECK: %[[SLOT0:.*]] = arith.divui %[[IV0]], %[[C16777216]] : index
-// CHECK: memref.store %[[IV0]], %[[MEMREF]][%[[SLOT0]]] : memref<2xindex>
-// CHECK: }
-// CHECK: scf.for %[[IV1:.*]] = %[[C16777216]] to %[[C16777217]] step %[[C16777216]] {
-// CHECK: %[[SLOT1:.*]] = arith.divui %[[IV1]], %[[C16777216]] : index
-// CHECK: memref.store %[[IV1]], %[[MEMREF]][%[[SLOT1]]] : memref<2xindex>
-// CHECK: }
-// CHECK: return
+// CHECK: scf.for %[[IV0:.*]] = %[[C0]] to %[[C16777216]] step %[[C16777216]]
+// CHECK: scf.for %[[IV1:.*]] = %[[C16777216]] to %[[C16777217]] step %[[C16777216]]
func.func @peel_front_step_near_ub(%out: memref<2xindex>) {
%c0 = arith.constant 0 : index
diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index 75e906a8bac13b..954b8321bcee40 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -77,6 +77,99 @@ module attributes {transform.with_named_sequence} {
// -----
+// dynamic ub >= 2^63 (negative as signed). Split must use minui.
+// CHECK-LABEL: func @dynamic_ub_high_bit
+// CHECK-SAME: %[[UB:.*]]: index
+// CHECK-DAG: %[[C:.*]] = arith.constant 4611686018427387905 : index
+// CHECK: %[[SPLIT:.*]] = arith.minui %[[UB]], %[[C]] : index
+// CHECK: scf.for unsigned %{{.*}} = %{{.*}} to %[[SPLIT]] step
+// CHECK: scf.for unsigned %{{.*}} = %[[SPLIT]] to %[[UB]] step
+func.func @dynamic_ub_high_bit(%ub: index, %m: memref<4xindex>) {
+ %lb = arith.constant 1 : index
+ %step = arith.constant 4611686018427387904 : index
+ scf.for unsigned %i = %lb to %ub step %step {
+ %d = arith.divui %i, %step : index
+ memref.store %i, %m[%d] : memref<4xindex>
+ }
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %l = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.op<"scf.for">
+ %p, %r = transform.loop.peel %l peel_front = true : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
+ transform.yield
+ }
+}
+
+// ----
+// constant bounds, ub = 3*2^62 (negative as signed), 3 iterations.
+// A signed early-exit check would wrongly refuse.
+// CHECK-LABEL: func @const_high_bit_peels
+// CHECK: scf.for unsigned %{{.*}} = %{{[a-z0-9_]+}} to %[[SPLIT:c4611686018427387905[a-z0-9_]*]] step
+// CHECK: scf.for unsigned %{{.*}} = %[[SPLIT]] to %{{[-a-z0-9_]+}} step
+func.func @const_high_bit_peels(%m: memref<4xindex>) {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant -4611686018427387904 : index
+ %step = arith.constant 4611686018427387904 : index
+ scf.for unsigned %i = %lb to %ub step %step {
+ %d = arith.divui %i, %step : index
+ memref.store %i, %m[%d] : memref<4xindex>
+ }
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %l = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.op<"scf.for">
+ %p, %r = transform.loop.peel %l peel_front = true : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
+ transform.yield
+ }
+}
+
+// -----
+
+// constant bounds, exactly one iteration. Must refuse.
+func.func @const_one_iteration(%m: memref<4xindex>) {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 4611686018427387905 : index
+ %step = arith.constant 4611686018427387904 : index
+ scf.for unsigned %i = %lb to %ub step %step {
+ memref.store %i, %m[%lb] : memref<4xindex>
+ }
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %l = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.op<"scf.for">
+ // expected-error @below {{failed to peel the first iteration}}
+ %p, %r = transform.loop.peel %l peel_front = true : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
+ transform.yield
+ }
+}
+
+// -----
+
+// constant empty loop that looks non-empty if read as signed.
+// lb = 2^63+5 (negative as signed), ub = 5. Unsigned: ub <= lb, so refuse.
+func.func @const_empty_unsigned(%m: memref<4xindex>) {
+ %lb = arith.constant -9223372036854775803 : index
+ %ub = arith.constant 5 : index
+ %step = arith.constant 1 : index
+ scf.for unsigned %i = %lb to %ub step %step {
+ memref.store %i, %m[%step] : memref<4xindex>
+ }
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %l = transform.structured.match ops{["scf.for"]} in %root : (!transform.any_op) -> !transform.op<"scf.for">
+ // expected-error @below {{failed to peel the first iteration}}
+ %p, %r = transform.loop.peel %l peel_front = true : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: @loop_peel_first_iter_op
func.func @loop_peel_first_iter_op() {
// CHECK: %[[C0:.+]] = arith.constant 0
More information about the Mlir-commits
mailing list