[Mlir-commits] [mlir] [mlir][DialectUtils] Fix 0 step handling in `constantTripCount` (PR #177329)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-scf

Author: Matthias Springer (matthias-springer)

<details>
<summary>Changes</summary>

A step size of "zero" does not indicate "zero iterations". It may indicate an infinite number of iterations.

This commit makes some transformations more conservative. We used to fold away some loops with step size 0 and that's now no longer the case.

Relation discussion: https://discourse.llvm.org/t/infinite-loops-and-dead-code/89530


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


3 Files Affected:

- (modified) mlir/lib/Dialect/Utils/StaticValueUtils.cpp (+11-3) 
- (modified) mlir/test/Dialect/SCF/canonicalize.mlir (+4-3) 
- (modified) mlir/test/Dialect/SCF/for-loop-peeling.mlir (+8-2) 


``````````diff
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index bc9d8a2496b4b..7fb0d4e9710f8 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -316,8 +316,12 @@ std::optional<APInt> constantTripCount(
            << lb;
     return std::nullopt;
   }
-  if (lb == ub)
+  if (lb == ub) {
+    // Fast path: LB == UB. The loop has zero iterations.
+    // Note: LB and UB could match at runtime, even though they are different
+    // SSA values. That case cannot be detected here.
     return APInt(bitwidth, 0);
+  }
 
   std::optional<std::pair<APInt, bool>> maybeStepCst =
       getConstantAPIntValue(step);
@@ -326,8 +330,12 @@ std::optional<APInt> constantTripCount(
     auto &stepCst = maybeStepCst->first;
     assert(static_cast<int>(stepCst.getBitWidth()) == bitwidth &&
            "step must have the same bitwidth as lb and ub");
-    if (stepCst.isZero())
-      return stepCst;
+    if (stepCst.isZero()) {
+      // Step is zero. If LB and UB match, we have zero iterations. Otherwise,
+      // we have an infinite number of iterations. We cannot tell for sure which
+      // case applies, so the static trip count is unknown.
+      return std::nullopt;
+    }
     if (stepCst.isNegative())
       return APInt(bitwidth, 0);
   }
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index e770f595bd262..3674680d80683 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -2187,13 +2187,14 @@ func.func @index_switch_fold_no_res() {
 
 // -----
 
-// Step 0 is invalid, the loop is eliminated.
+// Step size 0: The loop has an infinite number of iterations.
 // CHECK-LABEL: func @scf_for_all_step_size_0()
-//       CHECK-NOT:   scf.forall
+//       CHECK:   scf.forall (%[[arg0:.*]]) = (0) to (1) step (0)
+//       CHECK:     vector.print %[[arg0]]
 func.func @scf_for_all_step_size_0()  {
   %x = arith.constant 0 : index
   scf.forall (%i, %j) = (0, 4) to (1, 5) step (%x, 8) {
-    vector.print %x : index
+    vector.print %i : index
     scf.forall.in_parallel {}
   }
   return
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index 084576625f32c..6caa4bd50bad8 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -357,9 +357,15 @@ func.func @regression(%arg0: memref<i64>, %arg1: index) {
 // -----
 
 // Regression test: Make sure that we do not crash.
-// The step is 0, the loop will be eliminated.
+
 // CHECK-LABEL: func @zero_step(
-//       CHECK-NOT:   scf.for
+//       CHECK:   %[[c0:.*]] = arith.constant 0
+//       CHECK:   %[[c1:.*]] = arith.constant 1
+//       CHECK:   %[[poison:.*]] = ub.poison
+//       CHECK:   scf.for %{{.*}} = %[[c0]] to %[[poison]] step %[[c0]]
+//       CHECK:     arith.index_cast
+//       CHECK:   scf.for %{{.*}} = %[[poison]] to %[[c1]] step %[[c0]]
+//       CHECK:     arith.index_cast
 func.func @zero_step(%arg0: memref<i64>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index

``````````

</details>


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