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

[Mehdi Amini]

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

>From f23b0aaa6bb1baae13875c45d3c18ab487a69017 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 06:08:18 -0800
Subject: [PATCH] [mlir][scf] Fix crash in loopUnrollJamByFactor for zero trip
 count loops

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
---
 mlir/lib/Dialect/SCF/Utils/Utils.cpp     |  2 ++
 mlir/test/Dialect/SCF/transform-ops.mlir | 26 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+)

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
+  }
+}