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

[llvmlistbot at llvm.org]

Author: Alessandro Potenza
Date: 2026-08-23T11:29:00+02:00
New Revision: 8132d4da712113b30255b4adff8fe12682835990

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

LOG: [mlir][scf] Do not read non-constant loop bounds when unrolling (#217392)

`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.

Assisted-by: Claude (Anthropic)

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Utils/Utils.cpp
    mlir/test/Dialect/SCF/loop-unroll.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 36cddd5cf8bd7..2350f705a7ed4 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