[Mlir-commits] [mlir] [mlir][scf] Identify single iteration due to equal ub and step (PR #177194)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: None (zbenzion)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/177194.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Utils/StaticValueUtils.cpp (+4-1) 
- (modified) mlir/test/Dialect/SCF/canonicalize.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index bc9d8a2496b4b..065d1d3545b1c 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -346,10 +346,13 @@ std::optional<APInt> constantTripCount(
   std::optional<std::pair<APInt, bool>> maybeLbCst = getConstantAPIntValue(lb);
   std::optional<std::pair<APInt, bool>> maybeUbCst = getConstantAPIntValue(ub);
   if (maybeLbCst) {
+    APSInt lbCst(maybeLbCst->first, /*isUnsigned=*/!isSigned);
+    if (lbCst.isZero() && step == ub)
+      return APInt(bitwidth, 1);
+
     // If one of the bounds is not a constant, we can't compute the trip count.
     if (!maybeUbCst)
       return std::nullopt;
-    APSInt lbCst(maybeLbCst->first, /*isUnsigned=*/!isSigned);
     APSInt ubCst(maybeUbCst->first, /*isUnsigned=*/!isSigned);
     if (ubCst <= lbCst) {
       LDBG() << "constantTripCount is 0 because ub <= lb (" << lbCst << "("
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index e770f595bd262..230ea843e7057 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -789,6 +789,25 @@ func.func @replace_single_iteration_loop_unsigned_cmp() {
 
 // -----
 
+// CHECK-LABEL: @replace_single_iteration_loop_ub_equal_step
+func.func @replace_single_iteration_loop_ub_equal_step(%ub_step : index) {
+  // CHECK: %[[LB:.*]] = arith.constant 0
+  %c0 = arith.constant 0 : index
+  // CHECK: %[[INIT:.*]] = "test.init"
+  %init = "test.init"() : () -> i32
+  // CHECK-NOT: scf.for
+  // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
+  %0 = scf.for %i = %c0 to %ub_step step %ub_step iter_args(%arg = %init) -> (i32) {
+    %1 = "test.op"(%i, %arg) : (index, i32) -> i32
+    scf.yield %1 : i32
+  }
+  // CHECK: "test.consume"(%[[VAL]])
+  "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"

``````````

</details>


https://github.com/llvm/llvm-project/pull/177194