[Mlir-commits] [mlir] [mlir][scf] Fix the scf.while result reconstructed by upliftWhileToForLoop (PR #219617)
Samarth Narang
llvmlistbot at llvm.org
Mon Aug 31 13:46:33 PDT 2026
https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/219617
>From ad6bd17be4734a04d9ea93382fd497a4622e0d62 Mon Sep 17 00:00:00 2001
From: sanarang <sanarang at nvidia.com>
Date: Sat, 29 Aug 2026 02:35:26 +0000
Subject: [PATCH 1/2] [mlir][scf] Fix the scf.while result reconstructed by
upliftWhileToForLoop
`scf.condition` forwards its operands to the after region when the condition
is true and to the `scf.while` results when it is false. The result is
therefore the first induction value that fails the condition,
`lb + tripCount * step`.
upliftWhileToForLoop computed `lb + (tripCount - 1) * step`, i.e. the last
executed induction value. For `lb = 0, ub = 5, step = 2` the loop returns 6 but
the uplifted code returned 4. Loops whose condition is false on entry were
worse: a zero-trip loop returned `lb - step`, a value the loop never produces.
Compute `lb + tripCount * step` and clamp the trip count at zero, matching the
loop trip-count definition, so a loop that never executes reports `lb`. A plain
removal of the decrement would still be wrong for spans where the signed
ceil-division truncates to a negative count.
The existing CHECK lines encoded the decremented form; they are updated to the
clamped one, and a new test covers an `scf.while` whose induction-variable
result is used.
Fixes https://github.com/llvm/llvm-project/issues/219616
---
.../SCF/Transforms/UpliftWhileToFor.cpp | 14 ++++-
mlir/test/Dialect/SCF/uplift-while.mlir | 51 ++++++++++++++++---
2 files changed, 56 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp b/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp
index ec1044aaa42ac..34d46d4dafe03 100644
--- a/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp
@@ -241,11 +241,23 @@ FailureOr<scf::ForOp> mlir::scf::upliftWhileToForLoop(RewriterBase &rewriter,
one = arith::ConstantIntOp::create(rewriter, loc, step.getType(), 1);
}
+ // `scf.while` returns the value `scf.condition` forwards when the condition
+ // is false, i.e. the first induction value that fails the condition, which is
+ // `lb + tripCount * step`. Following the loop trip-count definition, the
+ // count is clamped at zero so that a loop whose condition is false on entry
+ // reports `lb` rather than a value it never produced.
+ Value zero;
+ if (isa<IndexType>(step.getType())) {
+ zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ } else {
+ zero = arith::ConstantIntOp::create(rewriter, loc, step.getType(), 0);
+ }
+
Value stepDec = arith::SubIOp::create(rewriter, loc, step, one);
Value len = arith::SubIOp::create(rewriter, loc, ub, lb);
len = arith::AddIOp::create(rewriter, loc, len, stepDec);
len = arith::DivSIOp::create(rewriter, loc, len, step);
- len = arith::SubIOp::create(rewriter, loc, len, one);
+ len = arith::MaxSIOp::create(rewriter, loc, len, zero);
Value res = arith::MulIOp::create(rewriter, loc, len, step);
res = arith::AddIOp::create(rewriter, loc, lb, res);
diff --git a/mlir/test/Dialect/SCF/uplift-while.mlir b/mlir/test/Dialect/SCF/uplift-while.mlir
index cbe2ce5076ad2..aa10b64cafe84 100644
--- a/mlir/test/Dialect/SCF/uplift-while.mlir
+++ b/mlir/test/Dialect/SCF/uplift-while.mlir
@@ -16,7 +16,8 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
// CHECK-LABEL: func @uplift_while
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
-// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {
// CHECK: "test.test1"(%[[I]]) : (index) -> ()
// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index
@@ -25,7 +26,7 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index
-// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index
+// CHECK: %[[R5:.*]] = arith.maxsi %[[R4]], %[[C0]] : index
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index
// CHECK: return %[[R7]] : index
@@ -48,7 +49,8 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
// CHECK-LABEL: func @uplift_while
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
-// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {
// CHECK: "test.test1"(%[[I]]) : (index) -> ()
// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index
@@ -57,7 +59,7 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index
-// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index
+// CHECK: %[[R5:.*]] = arith.maxsi %[[R4]], %[[C0]] : index
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index
// CHECK: return %[[R7]] : index
@@ -80,7 +82,8 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
// CHECK-LABEL: func @uplift_while
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
-// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {
// CHECK: "test.test1"(%[[I]]) : (index) -> ()
// CHECK: %[[INC:.*]] = arith.addi %[[STEP]], %[[I]] : index
@@ -89,7 +92,7 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index
-// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index
+// CHECK: %[[R5:.*]] = arith.maxsi %[[R4]], %[[C0]] : index
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index
// CHECK: return %[[R7]] : index
@@ -142,7 +145,8 @@ func.func @uplift_while(%arg0: i64, %arg1: i64, %arg2: i64) -> i64 {
// CHECK-LABEL: func @uplift_while
// CHECK-SAME: (%[[BEGIN:.*]]: i64, %[[END:.*]]: i64, %[[STEP:.*]]: i64) -> i64
-// CHECK: %[[C1:.*]] = arith.constant 1 : i64
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i64
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i64
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] : i64 {
// CHECK: "test.test1"(%[[I]]) : (i64) -> ()
// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : i64
@@ -151,7 +155,7 @@ func.func @uplift_while(%arg0: i64, %arg1: i64, %arg2: i64) -> i64 {
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : i64
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : i64
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : i64
-// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : i64
+// CHECK: %[[R5:.*]] = arith.maxsi %[[R4]], %[[C0]] : i64
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : i64
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : i64
// CHECK: return %[[R7]] : i64
@@ -185,3 +189,34 @@ func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> (i32, f32)
// CHECK: %[[T2:.*]] = "test.test2"(%[[ARG2]]) : (f32) -> f32
// CHECK: scf.yield %[[T1]], %[[T2]] : i32, f32
// CHECK: return %[[RES]]#0, %[[RES]]#1 : i32, f32
+
+// -----
+
+// The `scf.while` result is the value `scf.condition` forwards when the
+// condition is false, i.e. `lb + tripCount * step`, and the trip count is
+// clamped at zero so a loop whose condition is false on entry yields `lb`.
+func.func @uplift_while_iv_result(%begin: index, %end: index, %step: index) -> index {
+ %0 = scf.while (%iv = %begin) : (index) -> index {
+ %cond = arith.cmpi slt, %iv, %end : index
+ scf.condition(%cond) %iv : index
+ } do {
+ ^bb0(%iv: index):
+ %next = arith.addi %iv, %step : index
+ scf.yield %next : index
+ }
+ return %0 : index
+}
+
+// CHECK-LABEL: func @uplift_while_iv_result
+// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[SDEC:.*]] = arith.subi %[[STEP]], %[[C1]] : index
+// CHECK: %[[SPAN:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
+// CHECK: %[[NUM:.*]] = arith.addi %[[SPAN]], %[[SDEC]] : index
+// CHECK: %[[DIV:.*]] = arith.divsi %[[NUM]], %[[STEP]] : index
+// Trip count is clamped at zero, and NOT decremented.
+// CHECK: %[[TC:.*]] = arith.maxsi %[[DIV]], %[[C0]] : index
+// CHECK: %[[OFF:.*]] = arith.muli %[[TC]], %[[STEP]] : index
+// CHECK: %[[RES:.*]] = arith.addi %[[BEGIN]], %[[OFF]] : index
+// CHECK: return %[[RES]] : index
>From 04ed3292f543ed066827523227b7086c311fdcd2 Mon Sep 17 00:00:00 2001
From: sanarang <sanarang at nvidia.com>
Date: Mon, 31 Aug 2026 20:43:16 +0000
Subject: [PATCH 2/2] Address review: explain why the exit value is not the
last executed index
Spell out that scf.condition forwards to the results only once the condition
fails, so the escaping value comes from the increment on the final,
non-executing iteration and sits one step past anything the body saw.
Also record the strictly-positive step precondition explicitly. The matcher
takes step from whatever operand the matched arith.addi supplies and never
constrains its sign, so a non-positive step breaks the ceil-div and produces an
scf.for that violates its own contract. Documenting rather than asserting: a
non-positive step is reachable today, so an assert would fire on input the
pattern currently accepts.
---
.../lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp | 13 ++++++++-----
1 file changed, 8 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp b/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp
index 34d46d4dafe03..80f2ccf9a309f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/UpliftWhileToFor.cpp
@@ -241,11 +241,14 @@ FailureOr<scf::ForOp> mlir::scf::upliftWhileToForLoop(RewriterBase &rewriter,
one = arith::ConstantIntOp::create(rewriter, loc, step.getType(), 1);
}
- // `scf.while` returns the value `scf.condition` forwards when the condition
- // is false, i.e. the first induction value that fails the condition, which is
- // `lb + tripCount * step`. Following the loop trip-count definition, the
- // count is clamped at zero so that a loop whose condition is false on entry
- // reports `lb` rather than a value it never produced.
+ // Not the last executed induction value: `scf.condition` forwards to the
+ // results only once the condition fails, so the escaping value is one step
+ // past anything the body saw, `lb + tripCount * step`, clamped at zero trips
+ // so an entry-false loop reports `lb`.
+ //
+ // Assumes a strictly positive `step`. The matcher does not check this, so a
+ // non-positive step breaks the ceil-div and also builds an `scf.for` that
+ // violates its own contract; pre-existing gap, not repairable here.
Value zero;
if (isa<IndexType>(step.getType())) {
zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
More information about the Mlir-commits
mailing list