[Mlir-commits] [mlir] [mlir][scf] support front peeling for unsigned loop (PR #233238)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 20 02:53:34 PDT 2026
https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/233238
>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/3] [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/3] [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
>From 59cef9baf76a6ca8138eacfb979734e5f7346d6b Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Tue, 20 Oct 2026 15:23:06 +0530
Subject: [PATCH 3/3] [vector][scf] improve bounds handling for unsigned,
aceept front peeling and signed peeling and add tests for edge cases
---
.../SCF/Transforms/LoopSpecialization.cpp | 79 ++++++------
.../Dialect/SCF/for-loop-peeling-front.mlir | 117 ++++++++++++++++--
.../Dialect/SCF/transform-ops-invalid.mlir | 68 ++++++++++
mlir/test/Dialect/SCF/transform-ops.mlir | 98 +--------------
4 files changed, 216 insertions(+), 146 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index a3032217d33963..b54fa6f75c051f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -234,7 +234,14 @@ LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
ForOp &firstIteration) {
RewriterBase::InsertionGuard guard(b);
+ Value lbVal = forOp.getLowerBound();
+ Value ubVal = forOp.getUpperBound();
+ Value stepVal = forOp.getStep();
+ auto loc = forOp.getLoc();
+
Value splitBound;
+ b.setInsertionPoint(forOp);
+
// Check whether we have a unsigned loop
if (forOp.getUnsignedCmp()) {
IntegerAttr lbAttr, ubAttr, stepAttr;
@@ -247,51 +254,41 @@ LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
APInt ub = ubAttr.getValue();
APInt step = stepAttr.getValue();
- // bail out if ub <= lb or ub - lb <= step (unsigned)
+ // bail out if ub <= lb or ub - lb <= step (unsigned constants)
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();
- // 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();
+ // split = max(lb, min(lb + step, ub)): the first loop runs one iteration,
+ // or none if the loop is empty. If lb + step wraps, the loop has at most
+ // one iteration and split falls back to lb, so the main loop does the work.
+ // use unsigned ops (MinUIOp, MaxUIOp)
+ Value sum = b.createOrFold<arith::AddIOp>(loc, lbVal, stepVal);
+ Value capped = b.createOrFold<arith::MinUIOp>(loc, sum, ubVal);
+ splitBound = b.createOrFold<arith::MaxUIOp>(loc, lbVal, capped);
+ } else {
+ auto lbInt = getConstantIntValue(lbVal);
+ auto ubInt = getConstantIntValue(ubVal);
+ auto stepInt = getConstantIntValue(stepVal);
+
+ // Peeling is not needed if there is at most one iteration. scf.for requires
+ // a positive step, so the trip count is ceil((UB - LB) / Step) when UB > LB
+ // and 0 otherwise. With a positive step, "UB - LB <= Step" therefore covers
+ // both the empty loop and the single-iteration loop. Bail out if the
+ // subtraction overflows, since we can't reason about the trip count then.
+ if (lbInt && ubInt && stepInt) {
+ int64_t diff;
+ if (llvm::SubOverflow(*ubInt, *lbInt, diff) || diff <= *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, 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);
+ // split = max(lb, min(lb + step, ub)): the first loop runs one iteration,
+ // or none if the loop is empty. If lb + step wraps, the loop has at most
+ // one iteration and split falls back to lb, so the main loop does the work.
+ // use signed ops (MinSIOp, MaxSIOp)
+ Value sum = b.createOrFold<arith::AddIOp>(loc, lbVal, stepVal);
+ Value capped = b.createOrFold<arith::MinSIOp>(loc, sum, ubVal);
+ splitBound = b.createOrFold<arith::MaxSIOp>(loc, lbVal, capped);
}
// Peel the first iteration.
firstIteration = cast<ForOp>(b.clone(*forOp.getOperation()));
@@ -318,7 +315,7 @@ struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
LogicalResult matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const override {
- if (forOp.getUnsignedCmp())
+ if (!peelFront && forOp.getUnsignedCmp())
return rewriter.notifyMatchFailure(forOp,
"unsigned loops are not supported");
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
index 13a98309993adc..f05292c914f3e9 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-front.mlir
@@ -87,7 +87,8 @@ func.func @static_two_iterations_ub_used_in_loop(%arg0: memref<1xi32>) -> i32 {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// The peeled first iteration is bounded by %ub as well as by %lb + %step, so
// it stays empty when the source loop has no iterations.
-// CHECK: %[[SPLIT:.*]] = affine.min {{.*}}[%[[C0]], %[[C4]], %[[UB]]]
+// CHECK: %[[MIN:.*]] = arith.minsi %[[UB]], %[[C4]] : index
+// CHECK: %[[SPLIT:.*]] = arith.maxsi %[[MIN]], %{{.*}} : index
// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[SPLIT]] step %[[C4]] {
// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[UB]], %[[IV]])[%[[C4]]]
// CHECK: %[[LOAD:.*]] = memref.load %[[MEMREF]][]
@@ -120,20 +121,21 @@ func.func @no_loop_results(%ub : index, %d : memref<i32>) {
// -----
-// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s0 + s1, s2)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
// CHECK: func @fully_dynamic_bounds(
// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32
-// CHECK: %[[NEW_UB:.*]] = affine.min #[[MAP0]]()[%[[LB]], %[[STEP]], %[[UB]]]
-// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]]
+// CHECK: %[[SUM:.*]] = arith.addi %[[LB]], %[[STEP]] : index
+// CHECK: %[[CAPPED:.*]] = arith.minsi %[[SUM]], %[[UB]] : index
+// CHECK: %[[SPLIT:.*]] = arith.maxsi %[[LB]], %[[CAPPED]] : index
+// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[SPLIT]]
// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {
// CHECK: %[[MIN:.*]] = affine.min #[[MAP1]](%[[UB]], %[[IV]])[%[[STEP]]]
// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32
// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32
// CHECK: scf.yield %[[ADD]]
// CHECK: }
-// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[NEW_UB]] to %[[UB]]
+// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[SPLIT]] to %[[UB]]
// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) {
// CHECK: %[[REM:.*]] = affine.min #[[MAP1]](%[[UB]], %[[IV2]])[%[[STEP]]]
// CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]]
@@ -155,6 +157,26 @@ func.func @fully_dynamic_bounds(%lb : index, %ub: index, %step: index) -> i32 {
// -----
+// CHECK: func @fully_dynamic_bounds_unsigned(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[UB:.*]]: i32, %[[STEP:.*]]: i32
+// CHECK: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: %[[SUM:.*]] = arith.addi %[[LB]], %[[STEP]] : i32
+// CHECK: %[[CAPPED:.*]] = arith.minui %[[SUM]], %[[UB]] : i32
+// CHECK: %[[SPLIT:.*]] = arith.maxui %[[LB]], %[[CAPPED]] : i32
+// CHECK: %[[FIRST:.*]] = scf.for unsigned %[[IV:.*]] = %[[LB]] to %[[SPLIT]]
+// CHECK: %[[RESULT:.*]] = scf.for unsigned %[[IV2:.*]] = %[[SPLIT]] to %[[UB]]
+// CHECK: return %[[RESULT]]
+func.func @fully_dynamic_bounds_unsigned(%lb : i32, %ub: i32, %step: i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %r = scf.for unsigned %iv = %lb to %ub step %step iter_args(%arg = %c0) -> (i32) : i32 {
+ %0 = arith.addi %arg, %iv : i32
+ scf.yield %0 : i32
+ }
+ return %r : i32
+}
+
+// -----
+
// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)>
// CHECK: func @two_iteration_example(
// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32
@@ -229,10 +251,11 @@ func.func @no_peeling_front() -> i32 {
// CHECK-SAME: %[[LB:[^:]*]]: i32,
// CHECK-SAME: %[[N:[^:]*]]: i32) -> i32 {
// CHECK: %[[C2:.*]] = arith.constant 2 : i32
-// CHECK: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[N]] step %[[C2]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (i32) : i32 {
-// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[IV]] : i32
-// CHECK: scf.yield %[[ADD]] : i32
-// CHECK: }
+// CHECK: %[[SUM:.*]] = arith.addi %[[LB]], %[[C2]] : i32
+// CHECK: %[[CAPPED:.*]] = arith.minsi %[[SUM]], %[[N]] : i32
+// CHECK: %[[SPLIT:.*]] = arith.maxsi %[[LB]], %[[CAPPED]] : i32
+// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[SPLIT]] step %[[C2]] iter_args(%[[ACC:.*]] = %[[INIT]]) -> (i32) : i32 {
+// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[SPLIT]] to %[[N]] step %[[C2]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) : i32 {
// CHECK: return %[[RESULT]] : i32
func.func @non_index_loop_bounds(%init: i32, %lb: i32, %n: i32) -> i32 {
%c2 = arith.constant 2 : i32
@@ -267,3 +290,79 @@ func.func @peel_front_step_near_ub(%out: memref<2xindex>) {
}
return
}
+
+// -----
+
+// dynamic ub >= 2^63 (negative as signed). Split must use minui, clamped by maxui.
+// CHECK-LABEL: func @dynamic_ub_high_bit
+// CHECK-SAME: %[[UB:.*]]: index
+// CHECK-DAG: %[[C:.*]] = arith.constant 4611686018427387905 : index
+// CHECK: %[[MIN:.*]] = arith.minui %[[UB]], %[[C]] : index
+// CHECK: %[[SPLIT:.*]] = arith.maxui %[[MIN]], %{{.*}} : 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
+}
+
+// -----
+
+// 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
+}
+
+// -----
+
+// unsigned, empty loop (lb > ub) whose lb + step wraps to 90.
+// Without the clamp the main loop would start at 90. The split must stay at lb.
+// CHECK-LABEL: func @unsigned_lb_plus_step_wraps
+// CHECK-NOT: to %c90
+// CHECK-NOT: scf.for unsigned %{{.*}} = %c90
+// CHECK: scf.for unsigned %{{.*}} = %{{[-a-z0-9_]+}} to %[[SPLIT:[-a-z0-9_]+]] step
+// CHECK: scf.for unsigned %{{.*}} = %[[SPLIT]] to %{{.*}} step
+func.func @unsigned_lb_plus_step_wraps(%ub: i64, %m: memref<?xi64>) {
+ %c0 = arith.constant 0 : index
+ %lb = arith.constant -10 : i64
+ %step = arith.constant 100 : i64
+ scf.for unsigned %i = %lb to %ub step %step : i64 {
+ memref.store %i, %m[%c0] : memref<?xi64>
+ }
+ return
+}
+
+// -----
+
+// signed i64, empty loop (lb > ub) whose lb + step wraps to -3
+// The split must stay at lb.
+// CHECK-LABEL: func @signed_lb_plus_step_wraps
+// CHECK-NOT: to %c-3
+// CHECK-NOT: scf.for %{{.*}} = %c-3
+// CHECK: scf.for %{{.*}} = %{{[-a-z0-9_]+}} to %[[SPLIT:[-a-z0-9_]+]] step
+// CHECK: scf.for %{{.*}} = %[[SPLIT]] to %{{.*}} step
+func.func @signed_lb_plus_step_wraps(%ub: i64, %m: memref<?xi64>) {
+ %c0 = arith.constant 0 : index
+ %lb = arith.constant 9223372036854775807 : i64
+ %step = arith.constant 9223372036854775806 : i64
+ scf.for %i = %lb to %ub step %step : i64 {
+ memref.store %i, %m[%c0] : memref<?xi64>
+ }
+ return
+}
diff --git a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
index cd0f0eb8cb5554..996a7c462cea34 100644
--- a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
@@ -282,3 +282,71 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// signed constant bounds where ub - lb overflows int64. Must refuse.
+func.func @signed_sub_overflow(%m: memref<?xi64>) {
+ %c0 = arith.constant 0 : index
+ %lb = arith.constant 9223372036854775807 : i64
+ %ub = arith.constant -2 : i64
+ %step = arith.constant 9223372036854775806 : i64
+ scf.for %i = %lb to %ub step %step : i64 {
+ memref.store %i, %m[%c0] : memref<?xi64>
+ }
+ 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 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
+ }
+}
\ No newline at end of file
diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index 954b8321bcee40..6c7b4b6c50165e 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -75,98 +75,6 @@ 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
- }
-}
// -----
@@ -176,10 +84,8 @@ func.func @loop_peel_first_iter_op() {
// CHECK: %[[C41:.+]] = arith.constant 41
// CHECK: %[[C5:.+]] = arith.constant 5
// CHECK: %[[C5_0:.+]] = arith.constant 5
- // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C5_0]] step %[[C5]]
- // CHECK: arith.addi
- // CHECK: scf.for %{{.+}} = %[[C5_0]] to %[[C41]] step %[[C5]]
- // CHECK: arith.addi
+ // CHECK: scf.for %{{.+}} = %{{.+}} to %[[SPLIT:.+]] step %[[STEP:.+]]
+ // CHECK: scf.for %{{.+}} = %[[SPLIT]] to %{{.+}} step %[[STEP]]
%0 = arith.constant 0 : index
%1 = arith.constant 41 : index
%2 = arith.constant 5 : index
More information about the Mlir-commits
mailing list