[Mlir-commits] [mlir] [MLIR][SCF] Fix Loop Trip Count Calculation for Unsigned Values (PR #175301)

[llvmlistbot at llvm.org]

https://github.com/veera-sivarajan created https://github.com/llvm/llvm-project/pull/175301

Previously, loop trip count was calculated using signed division
and remainder for both signed and unsigned values.

For loops like:
```
scf.for unsigned 0 to -100 step 2147483647 : i32 {}
```
This resulted in a trip count of 1 when the actual trip count is 2.

This PR fixes it by using unsigned division and remainder for
unsigned values.

>From f55e001c19d2eab921c3facc289e0c456ca26324 Mon Sep 17 00:00:00 2001
From: Veera <sveera.2001 at gmail.com>
Date: Sat, 10 Jan 2026 08:29:53 +0000
Subject: [PATCH 1/2] Add Test

---
 mlir/test/Dialect/SCF/canonicalize.mlir | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index d5d0aee3bbe25..b475a66771d06 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -762,6 +762,28 @@ func.func @replace_single_iteration_const_diff(%arg0 : index) {
 
 // -----
 
+func.func @replace_single_iteration_loop_unsigned_cmp() {
+// CHECK-LABEL:   func.func @replace_single_iteration_loop_unsigned_cmp() {
+// CHECK:           %[[CONSTANT_0:.*]] = arith.constant 0 : i32
+// CHECK:           %[[VAL_0:.*]] = "test.init"() : () -> i32
+// CHECK:           %[[VAL_1:.*]] = "test.op"(%[[CONSTANT_0]], %[[VAL_0]]) : (i32, i32) -> i32
+// CHECK:           "test.consume"(%[[VAL_1]]) : (i32) -> ()
+// CHECK:           return
+// CHECK:         }
+  %lowerBound = arith.constant 0 : i32
+  %upperBound = arith.constant -100 : i32
+  %step = arith.constant 2147483647 : i32
+  %init = "test.init"() : () -> i32
+  %0 = scf.for unsigned %i = %lowerBound to %upperBound step %step iter_args(%arg = %init) -> (i32) : i32 {
+    %1 = "test.op"(%i, %arg) : (i32, i32) -> i32
+    scf.yield %1 : i32
+  }
+  "test.consume"(%0) : (i32) -> ()
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @remove_empty_parallel_loop
 func.func @remove_empty_parallel_loop(%lb: index, %ub: index, %s: index) {
   // CHECK: %[[INIT:.*]] = "test.init"

>From 50b4233d7419f5154ae788ed8d50aca9db0bfd38 Mon Sep 17 00:00:00 2001
From: Veera <sveera.2001 at gmail.com>
Date: Sat, 10 Jan 2026 09:02:34 +0000
Subject: [PATCH 2/2] [MLIR][SCF] Fix Loop Trip Count Calculation for Unsigned
 Values

Previously, loop trip count was calculated using signed division
and remainder for both signed and unsigned values.

This resulted in incorrect trip count for loops like:
```
scf.for unsigned 0 to -100 step 2147483647 {}
```

This PR fixes it by using unsigned division and remainder for
unsigned values.
---
 mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 8 +++-----
 mlir/test/Dialect/SCF/canonicalize.mlir     | 9 +++++++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 59f068c205cf3..bc9d8a2496b4b 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -351,8 +351,6 @@ std::optional<APInt> constantTripCount(
       return std::nullopt;
     APSInt lbCst(maybeLbCst->first, /*isUnsigned=*/!isSigned);
     APSInt ubCst(maybeUbCst->first, /*isUnsigned=*/!isSigned);
-    if (!maybeUbCst)
-      return std::nullopt;
     if (ubCst <= lbCst) {
       LDBG() << "constantTripCount is 0 because ub <= lb (" << lbCst << "("
              << lbCst.getBitWidth() << ") <= " << ubCst << "("
@@ -385,9 +383,9 @@ std::optional<APInt> constantTripCount(
     return std::nullopt;
   }
   auto &stepCst = maybeStepCst->first;
-  llvm::APInt tripCount = diff.sdiv(stepCst);
-  llvm::APInt r = diff.srem(stepCst);
-  if (!r.isZero())
+  llvm::APInt tripCount = isSigned ? diff.sdiv(stepCst) : diff.udiv(stepCst);
+  llvm::APInt remainder = isSigned ? diff.srem(stepCst) : diff.urem(stepCst);
+  if (!remainder.isZero())
     tripCount = tripCount + 1;
   LDBG() << "constantTripCount found: " << tripCount;
   return tripCount;
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index b475a66771d06..365c0e1d5c86f 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -765,9 +765,14 @@ func.func @replace_single_iteration_const_diff(%arg0 : index) {
 func.func @replace_single_iteration_loop_unsigned_cmp() {
 // CHECK-LABEL:   func.func @replace_single_iteration_loop_unsigned_cmp() {
 // CHECK:           %[[CONSTANT_0:.*]] = arith.constant 0 : i32
+// CHECK:           %[[CONSTANT_1:.*]] = arith.constant -100 : i32
+// CHECK:           %[[CONSTANT_2:.*]] = arith.constant 2147483647 : i32
 // CHECK:           %[[VAL_0:.*]] = "test.init"() : () -> i32
-// CHECK:           %[[VAL_1:.*]] = "test.op"(%[[CONSTANT_0]], %[[VAL_0]]) : (i32, i32) -> i32
-// CHECK:           "test.consume"(%[[VAL_1]]) : (i32) -> ()
+// CHECK:           %[[FOR_0:.*]] = scf.for unsigned %[[VAL_1:.*]] = %[[CONSTANT_0]] to %[[CONSTANT_1]] step %[[CONSTANT_2]] iter_args(%[[VAL_2:.*]] = %[[VAL_0]]) -> (i32)  : i32 {
+// CHECK:             %[[VAL_3:.*]] = "test.op"(%[[VAL_1]], %[[VAL_2]]) : (i32, i32) -> i32
+// CHECK:             scf.yield %[[VAL_3]] : i32
+// CHECK:           }
+// CHECK:           "test.consume"(%[[FOR_0]]) : (i32) -> ()
 // CHECK:           return
 // CHECK:         }
   %lowerBound = arith.constant 0 : i32