[Mlir-commits] [mlir] 0d2c7d7 - [mlir][scf] Fix crash in loopUnrollJamByFactor for zero trip count loops (#184620)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-06T12:09:15+01:00
New Revision: 0d2c7d71744108fcfcbfb6258994da2c008314ba

URL: https://github.com/llvm/llvm-project/commit/0d2c7d71744108fcfcbfb6258994da2c008314ba
DIFF: https://github.com/llvm/llvm-project/commit/0d2c7d71744108fcfcbfb6258994da2c008314ba.diff

LOG: [mlir][scf] Fix crash in loopUnrollJamByFactor for zero trip count loops (#184620)

When loopUnrollJamByFactor is called on a loop with a zero trip count
(e.g., lb == ub), the function correctly detects that the unroll factor
exceeds the trip count and adjusts it to match. However, capping the
factor at 0 leads to unsigned underflow when computing the SmallVector
size for `operandMaps(unrollJamFactor - 1)`, resulting in an attempt to
allocate ~2^64 elements and an immediate crash.

Add an early return when the trip count is zero, matching the pattern
already used by loopUnrollFull. A loop that executes zero times needs no
unrolling.

Fixes #103703

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Utils/Utils.cpp
    mlir/test/Dialect/SCF/transform-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 2dd7df09e0f22..5c175a753854d 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -549,6 +549,8 @@ LogicalResult mlir::loopUnrollJamByFactor(scf::ForOp forOp,
     return failure();
   }
   uint64_t tripCountValue = tripCount->getZExtValue();
+  if (tripCountValue == 0)
+    return success();
   if (unrollJamFactor > tripCountValue) {
     LDBG() << "unroll and jam factor is greater than trip count, set factor to "
               "trip "

diff  --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index d9445182769e7..0e5a4a98a98ea 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -581,3 +581,29 @@ module attributes {transform.with_named_sequence} {
      transform.yield
  }
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/103703
+// unroll_and_jam on a loop with zero trip count used to crash because the
+// unroll factor was set to 0 (= tripCount) and then used as a SmallVector size
+// with unsigned underflow.
+
+// CHECK-LABEL: @loop_unroll_and_jam_zero_trip_count
+// CHECK:       %[[C0:.*]] = arith.constant 0 : index
+// CHECK:       scf.for %{{.*}} = %[[C0]] to %[[C0]] step %[[C0]]
+func.func @loop_unroll_and_jam_zero_trip_count() {
+  %c0 = arith.constant 0 : index
+  scf.for %arg0 = %c0 to %c0 step %c0 {
+    %0 = arith.addi %arg0, %arg0 : index
+  }
+  return
+}
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["arith.addi"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
+    transform.loop.unroll_and_jam %1 {factor = 3 : i64} : !transform.op<"scf.for">
+    transform.yield
+  }
+}