[Mlir-commits] [mlir] [mlir][scf] Do not read non-constant loop bounds when unrolling (PR #217392)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Alessandro Potenza (alepot55)

<details>
<summary>Changes</summary>

`loopUnrollByFactor` asserts `expected constant loop bound` (`SCF/Utils/Utils.cpp`) when the loop bounds are not constant. `getStaticTripCount` takes the constant path whenever `constantTripCount` answers, but `constantTripCount` answers on three shapes where the bounds themselves are **not** constant:

- `lb == ub` (same `Value`), giving 0 iterations,
- `lb == 0` and `ub == step`, giving 1 iteration,
- `ub` a constant offset from a non-constant `lb`, via `computeUbMinusLb` under `nsw`.

In all three the code then reads the bounds as constants and asserts.

The fix takes the constant path only when `lb`, `ub` and `step` are all constant, and otherwise falls through to the dynamic path that already handles them. The other two callers of `getStaticTripCount` in that file use only the count, never the bounds, so the defect is confined to this one.

Fixes #<!-- -->203858.

Tests: three cases in `mlir/test/Dialect/SCF/loop-unroll.mlir`, one per shape. Each aborts on current `main` (`0ed130af5`) without the fix and passes with it; all ten `RUN` lines of that file are green.

Assisted-by: Claude (Anthropic)

AI-assisted, disclosed per the LLVM AI Tool Use Policy.


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


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+10) 
- (modified) mlir/test/Dialect/SCF/loop-unroll.mlir (+60) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c158e624002bd..6209f8579506c 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -386,6 +386,16 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
   bool generateEpilogueLoop = true;
 
   std::optional<APInt> constTripCount = forOp.getStaticTripCount();
+  // A static trip count does not imply constant bounds: it is also known when
+  // the lower and the upper bound are the same value (zero iterations), when
+  // the lower bound is zero and the upper bound is the step (one iteration),
+  // and when the upper bound is a constant offset from a non-constant lower
+  // bound. The computation below reads all three bounds as constants, so fall
+  // back to the dynamic case unless they are.
+  if (constTripCount && !(getConstantAPIntValue(forOp.getLowerBound()) &&
+                          getConstantAPIntValue(forOp.getUpperBound()) &&
+                          getConstantAPIntValue(step)))
+    constTripCount = std::nullopt;
   if (constTripCount) {
     // Constant loop bounds computation.
     bool isUnsignedLoop = forOp.getUnsignedCmp();
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index a047d4661eb95..89d86b09cddfb 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -700,3 +700,63 @@ func.func @static_loop_unroll_by_3_no_promote_epilogue(%arg0 : memref<?xf32>) {
 //  PROMOTE-BY-3: memref.store
 
 
+
+// -----
+
+// `scf.for` has a statically known trip count whenever the lower and the upper
+// bound are the same value, even when that value is not a constant. Unrolling
+// must not read the bounds as constants in that case.
+func.func @unroll_static_trip_count_dynamic_equal_bounds(%arg0: index, %arg1: memref<?xf32>) {
+  %0 = arith.constant 7.0 : f32
+  scf.for %i0 = %arg0 to %arg0 step %arg0 {
+    memref.store %0, %arg1[%i0] : memref<?xf32>
+  }
+  return
+}
+// UNROLL-BY-2-LABEL: func @unroll_static_trip_count_dynamic_equal_bounds
+//       UNROLL-BY-2:   scf.for
+//       UNROLL-BY-2:     memref.store
+//       UNROLL-BY-2:     memref.store
+//       UNROLL-BY-2:   scf.for
+//       UNROLL-BY-2:     memref.store
+
+// -----
+
+// Same, for the single-iteration case recognised from a zero lower bound and an
+// upper bound equal to the step. Here the lower bound is a constant but the
+// other two are not.
+func.func @unroll_static_trip_count_dynamic_ub_eq_step(%arg0: index, %arg1: memref<?xf32>) {
+  %0 = arith.constant 7.0 : f32
+  %c0 = arith.constant 0 : index
+  scf.for %i0 = %c0 to %arg0 step %arg0 {
+    memref.store %0, %arg1[%i0] : memref<?xf32>
+  }
+  return
+}
+// UNROLL-BY-2-LABEL: func @unroll_static_trip_count_dynamic_ub_eq_step
+//       UNROLL-BY-2:   scf.for
+//       UNROLL-BY-2:     memref.store
+//       UNROLL-BY-2:     memref.store
+//       UNROLL-BY-2:   scf.for
+//       UNROLL-BY-2:     memref.store
+
+// -----
+
+// Same, for an upper bound that is a constant offset from a non-constant lower
+// bound. The trip count is known, the lower and upper bound are not constants.
+func.func @unroll_static_trip_count_ub_offset_from_lb(%arg0: index, %arg1: memref<?xf32>) {
+  %0 = arith.constant 7.0 : f32
+  %c4 = arith.constant 4 : index
+  %c16 = arith.constant 16 : index
+  %ub = arith.addi %arg0, %c16 overflow<nsw> : index
+  scf.for %i0 = %arg0 to %ub step %c4 {
+    memref.store %0, %arg1[%i0] : memref<?xf32>
+  }
+  return
+}
+// UNROLL-BY-2-LABEL: func @unroll_static_trip_count_ub_offset_from_lb
+//       UNROLL-BY-2:   scf.for
+//       UNROLL-BY-2:     memref.store
+//       UNROLL-BY-2:     memref.store
+//       UNROLL-BY-2:   scf.for
+//       UNROLL-BY-2:     memref.store

``````````

</details>


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