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

[Alessandro Potenza]

https://github.com/alepot55 created https://github.com/llvm/llvm-project/pull/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.

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.


>From 86bbca68972c9012a07518dc9864af58b357b278 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Tue, 18 Aug 2026 11:28:09 +0200
Subject: [PATCH] [mlir][scf] Do not read non-constant loop bounds when
 unrolling

`loopUnrollByFactor` takes the constant-bound path whenever
`getStaticTripCount` returns a value, and then reads the lower bound, the
upper bound and the step as constants. A static trip count does not imply
constant bounds: `constantTripCount` also succeeds when the lower and the
upper bound are the same value, when the lower bound is zero and the upper
bound is the step, and when the upper bound is a constant offset from a
non-constant lower bound. In those cases the read asserts with
`expected constant loop bound`.

Take that path only when all three bounds are constants, and fall back to
the existing dynamic computation otherwise.

Fixes #203858.
---
 mlir/lib/Dialect/SCF/Utils/Utils.cpp   | 10 +++++
 mlir/test/Dialect/SCF/loop-unroll.mlir | 60 ++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

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