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

[llvmlistbot at llvm.org]

Author: Veera
Date: 2026-01-10T02:19:50-08:00
New Revision: b2d90d4da03417b441fff0cc5cedc1c29799e657

URL: https://github.com/llvm/llvm-project/commit/b2d90d4da03417b441fff0cc5cedc1c29799e657
DIFF: https://github.com/llvm/llvm-project/commit/b2d90d4da03417b441fff0cc5cedc1c29799e657.diff

LOG: [MLIR][SCF] Fix Loop Trip Count Calculation for Unsigned Values (#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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Utils/StaticValueUtils.cpp
    mlir/test/Dialect/SCF/canonicalize.mlir

Removed: 
    


################################################################################
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 d5d0aee3bbe25..365c0e1d5c86f 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -762,6 +762,33 @@ 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:           %[[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
+  %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"