[Mlir-commits] [mlir] [mlir][scf] fix unsigned trip count with nuw in `constantTripCount` (PR #235272)

[llvmlistbot at llvm.org]

https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/235272

>From 783b1520f1327f31f464f66faac5d84be9a756f1 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 16:18:11 +0530
Subject: [PATCH 1/3] [mlir][scf] fix unsigned trip count handling with
 non-negative upper bounds and update tests

---
 mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 5 ++---
 mlir/test/Dialect/SCF/trip_count.mlir       | 4 ++--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index ee1086e70597c2..d99aa47f4efe14 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -392,7 +392,7 @@ std::optional<APInt> constantTripCount(
 
     /// Non-constant bound, let's try to compute the difference between the
     /// upper and lower bound
-    std::optional<llvm::APSInt> maybeDiff =
+    std::optional<llvm::APInt> maybeDiff =
         computeUbMinusLb(cast<Value>(lb), cast<Value>(ub), isSigned);
     if (!maybeDiff)
       return std::nullopt;
@@ -400,7 +400,7 @@ std::optional<APInt> constantTripCount(
   }
   LDBG() << "constantTripCount: " << (isSigned ? "isSigned" : "isUnsigned")
          << ", ub-lb: " << diff << "(" << diff.getBitWidth() << "b)";
-  if (diff.isNegative()) {
+  if (isSigned && diff.isNegative()) {
     LDBG() << "constantTripCount is 0 because ub-lb diff is negative";
     return APInt(bitwidth, 0);
   }
@@ -424,7 +424,6 @@ std::optional<APInt> constantTripCount(
   llvm::APInt remainder = diff.urem(stepCst);
   if (!remainder.isZero())
     tripCount = tripCount + 1;
-
   LDBG() << "constantTripCount found: " << tripCount;
   return tripCount;
 }
diff --git a/mlir/test/Dialect/SCF/trip_count.mlir b/mlir/test/Dialect/SCF/trip_count.mlir
index 7e74988b350197..732db0ead10f42 100644
--- a/mlir/test/Dialect/SCF/trip_count.mlir
+++ b/mlir/test/Dialect/SCF/trip_count.mlir
@@ -663,7 +663,7 @@ func.func @trip_count_arith_add_negative_nuw_loop_unsigned(%lb : i32) -> i32 {
   %c0_i32 = arith.constant 0 : i32
   %c4_i32 = arith.constant 4 : i32
   %c-16_i32 = arith.constant -16 : i32
-  // CHECK: "test.trip-count" = 0
+  // CHECK: "test.trip-count" = 1073741820
   %ub = arith.addi %lb, %c-16_i32 overflow<nuw> : i32
   %1 = scf.for unsigned %arg0 = %lb to %ub step %c4_i32 iter_args(%arg1 = %c0_i32) -> (i32)  : i32 {
     scf.yield %arg0 : i32
@@ -677,7 +677,7 @@ func.func @trip_count_arith_add_negative_nuw_loop_unsigned(%lb : i32) -> i32 {
 func.func @trip_count_arith_add_negative_nuw_loop_unsigned_step_dyn(%lb : i32, %step : i32) -> i32 {
   %c0_i32 = arith.constant 0 : i32
   %c-16_i32 = arith.constant -16 : i32
-  // CHECK: "test.trip-count" = 0
+  // CHECK: "test.trip-count" = "none"
   %ub = arith.addi %lb, %c-16_i32 overflow<nuw> : i32
   %1 = scf.for unsigned %arg0 = %lb to %ub step %step iter_args(%arg1 = %c0_i32) -> (i32)  : i32 {
     scf.yield %arg0 : i32

>From 7d197099baa9cc2be6dc922cb3f779f9c78fd428 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 16:33:24 +0530
Subject: [PATCH 2/3] fix: change APInt to APSInt for trip count difference
 computation

---
 mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index d99aa47f4efe14..9d5161c5abb463 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -392,7 +392,7 @@ std::optional<APInt> constantTripCount(
 
     /// Non-constant bound, let's try to compute the difference between the
     /// upper and lower bound
-    std::optional<llvm::APInt> maybeDiff =
+    std::optional<llvm::APSInt> maybeDiff =
         computeUbMinusLb(cast<Value>(lb), cast<Value>(ub), isSigned);
     if (!maybeDiff)
       return std::nullopt;

>From 75dff0f21b63861c07a75f577618c66dba299238 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Sat, 24 Oct 2026 16:40:55 +0530
Subject: [PATCH 3/3] fix: add missing line

---
 mlir/lib/Dialect/Utils/StaticValueUtils.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 9d5161c5abb463..3be4d0166abe41 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -424,6 +424,7 @@ std::optional<APInt> constantTripCount(
   llvm::APInt remainder = diff.urem(stepCst);
   if (!remainder.isZero())
     tripCount = tripCount + 1;
+
   LDBG() << "constantTripCount found: " << tripCount;
   return tripCount;
 }