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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+2) 
- (modified) mlir/test/Dialect/SCF/transform-ops.mlir (+26) 


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

``````````

</details>


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