[Mlir-commits] [mlir] [mlir][scf] Fix the scf.while result reconstructed by upliftWhileToForLoop (PR #219617)
Samarth Narang
llvmlistbot at llvm.org
Fri Aug 28 19:37:11 PDT 2026
https://github.com/snarang181 created https://github.com/llvm/llvm-project/pull/219617
`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.
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.
Fixes https://github.com/llvm/llvm-project/issues/219616
>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] [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
More information about the Mlir-commits
mailing list