[Mlir-commits] [mlir] [mlir][scf] Reject unsafe unroll-and-jam effects (PR #228138)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 10:03:38 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-scf
@llvm/pr-subscribers-mlir
Author: Ziliang Zhang (ziliangzl)
<details>
<summary>Changes</summary>
Unroll-and-jam may reorder operations across outer loop iterations and violate loop-carried memory dependences.
This change conservatively rejects `scf.for` loops with non-read or unknown effects when an inner loop would be jammed.
For example, this loop is rejected because jamming may reorder a write with a dependent read from another outer iteration:
```text
for i:
for j:
A[i, j] = A[i - 1, j + 1]
```
Read-only nested loops remain supported:
```text
for i:
for j:
sum += A[j, i]
```
A loop without an inner loop is equivalent to ordinary unrolling, so writes remain supported:
```text
for i:
A[i] = i
```
This is intentionally conservative and does not perform alias or dependence analysis. Consequently, nested loops containing writes are rejected even when their accesses are actually independent. A future change may admit such cases when independence can be proven.
---
Full diff: https://github.com/llvm/llvm-project/pull/228138.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td (+5)
- (modified) mlir/include/mlir/Dialect/SCF/Utils/Utils.h (+3-1)
- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+32)
- (modified) mlir/test/Dialect/SCF/transform-ops-invalid.mlir (+32)
- (modified) mlir/test/Dialect/SCF/transform-ops.mlir (+31-14)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
index 4f463bb02fb77..61a7b0defbb57 100644
--- a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
+++ b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
@@ -370,6 +370,11 @@ def LoopUnrollAndJamOp : Op<Transform_Dialect, "loop.unroll_and_jam",
number of loop body copies per iteration. If the unroll factor is larger
than the loop trip count, the latter is used as the unroll factor instead.
+ For `scf.for`, loops with effects other than reads are currently rejected
+ conservatively when an inner `scf.for` would be jammed. This prevents the
+ transform from reordering potentially conflicting effects until a more
+ precise dependence analysis can prove that the reordering is legal.
+
#### Return modes
This operation ignores non-`scf.for`, non-`affine.for` ops and drops them
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index 4cc7e62e6fd58..b2d70b348eb4d 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -146,7 +146,9 @@ LogicalResult loopUnrollFull(scf::ForOp forOp);
/// due to invalid unroll factors. In case of unroll factor of 1, the function
/// bails out without doing anything (returns success). Currently, only constant
/// trip count that are divided by the unroll factor is supported. Currently,
-/// for operations with results are not supported.
+/// for operations with results are not supported. If an inner `scf.for` is
+/// jammed, the loop must have no recursively nested memory effects other than
+/// reads. Loops with other or unknown effects are rejected conservatively.
LogicalResult loopUnrollJamByFactor(scf::ForOp forOp, uint64_t unrollFactor);
/// Materialize bounds and step of a zero-based and unit-step loop derived by
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 7ba0a796c8631..a5828f280e5eb 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -600,6 +600,27 @@ static bool areInnerBoundsInvariant(scf::ForOp forOp) {
return !walkResult.wasInterrupted();
}
+/// Return true if unroll-and-jam does not reorder any effects or if all
+/// reordered effects are reads. Operations without a memory effect interface
+/// are treated conservatively.
+static bool hasSafeEffectsForUnrollAndJam(scf::ForOp forOp) {
+ // Without a directly nested loop, unroll-and-jam degenerates to ordinary
+ // unrolling and preserves the original order of operation instances.
+ if (llvm::none_of(forOp.getBody()->without_terminator(),
+ [](Operation &op) { return isa<scf::ForOp>(op); }))
+ return true;
+
+ std::optional<SmallVector<MemoryEffects::EffectInstance>> effects =
+ getEffectsRecursively(forOp);
+ if (!effects)
+ return false;
+
+ return llvm::all_of(*effects,
+ [](const MemoryEffects::EffectInstance &effect) {
+ return isa<MemoryEffects::Read>(effect.getEffect());
+ });
+}
+
/// Unrolls and jams this loop by the specified factor.
LogicalResult mlir::loopUnrollJamByFactor(scf::ForOp forOp,
uint64_t unrollJamFactor) {
@@ -647,6 +668,17 @@ LogicalResult mlir::loopUnrollJamByFactor(scf::ForOp forOp,
if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
return success();
+ // Jamming changes the relative order of operations from different outer
+ // loop iterations. Reject effects that may not commute unless no inner loop
+ // is actually jammed. This check is intentionally conservative: a more
+ // precise analysis can admit writes after proving that accesses do not
+ // conflict across the reordered iterations.
+ if (unrollJamFactor > 1 && !hasSafeEffectsForUnrollAndJam(forOp)) {
+ LDBG() << "failed to unroll and jam: loop has effects that may be "
+ "reordered";
+ return failure();
+ }
+
// Gather all sub-blocks to jam upon the loop being unrolled.
JamBlockGatherer<scf::ForOp> jbg;
jbg.walk(forOp);
diff --git a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
index 20549581029eb..725d1d95fc046 100644
--- a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
@@ -84,6 +84,38 @@ module attributes {transform.with_named_sequence} {
// -----
+func.func @loop_unroll_and_jam_reordered_memory_dependence(
+ %arr: memref<5x5xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c5 = arith.constant 5 : index
+ scf.for %i = %c1 to %c5 step %c1 {
+ scf.for %j = %c0 to %c4 step %c1 {
+ %im1 = arith.subi %i, %c1 : index
+ %jp1 = arith.addi %j, %c1 : index
+ %value = memref.load %arr[%im1, %jp1] : memref<5x5xi32>
+ memref.store %value, %arr[%i, %j] : memref<5x5xi32>
+ }
+ } {unroll_jam}
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root: !transform.any_op {transform.readonly}) {
+ %loop = transform.structured.match ops{["scf.for"]}
+ attributes{unroll_jam} in %root
+ : (!transform.any_op) -> !transform.op<"scf.for">
+ // expected-error @below {{failed to unroll and jam}}
+ transform.loop.unroll_and_jam %loop factor = 2
+ : !transform.op<"scf.for">
+ transform.yield
+ }
+}
+
+// -----
+
func.func @loop_unroll_and_jam_unsupported_loop_with_results() -> index {
%c0 = arith.constant 0 : index
%c40 = arith.constant 40 : index
diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index 75e906a8bac13..c8ed510c99c49 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -255,6 +255,35 @@ module attributes {transform.with_named_sequence} {
// -----
+// Unroll-and-jam without an inner loop preserves operation order and may
+// therefore unroll loops with write effects.
+// CHECK-LABEL: @loop_unroll_and_jam_no_inner_loop_with_write
+func.func @loop_unroll_and_jam_no_inner_loop_with_write(
+ %arg0: memref<4xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ scf.for %i = %c0 to %c4 step %c1 {
+ // CHECK-COUNT-2: memref.store
+ memref.store %i, %arg0[%i] : memref<4xindex>
+ } {unroll_jam}
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root: !transform.any_op {transform.readonly}) {
+ %loop = transform.structured.match ops{["scf.for"]}
+ attributes{unroll_jam} in %root
+ : (!transform.any_op) -> !transform.op<"scf.for">
+ transform.loop.unroll_and_jam %loop factor = 2
+ : !transform.op<"scf.for">
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: @loop_unroll_and_jam_op
// CHECK: %[[VAL_0:.*]]: memref<96x128xi8, 3>, %[[VAL_1:.*]]: memref<128xi8, 3>) {
func.func private @loop_unroll_and_jam_op(%arg0: memref<96x128xi8, 3>, %arg1: memref<128xi8, 3>) {
@@ -302,25 +331,13 @@ func.func private @loop_unroll_and_jam_op(%arg0: memref<96x128xi8, 3>, %arg1: me
// CHECK: scf.yield %[[SUM_0]], %[[SUM_1]], %[[SUM_2]], %[[SUM_3]]
scf.yield %4 : i8
}
- memref.store %sum, %arg1[%arg2] : memref<128xi8, 3>
- // CHECK: memref.store %[[VAL_39:.*]]#0, %[[VAL_1]]{{\[}}%[[OUTER_I]]]
- // CHECK: %[[ONE_2:.*]] = arith.constant 1
- // CHECK: %[[INC_STORE1:.*]] = arith.addi %[[OUTER_I]], %[[ONE_2]]
- // CHECK: memref.store %[[VAL_39]]#1, %[[VAL_1]]{{\[}}%[[INC_STORE1]]]
- // CHECK: %[[TWO_2:.*]] = arith.constant 2
- // CHECK: %[[INC_STORE2:.*]] = arith.addi %[[OUTER_I]], %[[TWO_2]]
- // CHECK: memref.store %[[VAL_39]]#2, %[[VAL_1]]{{\[}}%[[INC_STORE2]]]
- // CHECK: %[[THREE_2:.*]] = arith.constant 3
- // CHECK: %[[INC_STORE3:.*]] = arith.addi %[[OUTER_I]], %[[THREE_2]]
- // CHECK: memref.store %[[VAL_39]]#3, %[[VAL_1]]{{\[}}%[[INC_STORE3]]]
- }
+ } {unroll_jam}
return
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
- %0 = transform.structured.match ops{["memref.store"]} in %arg1 : (!transform.any_op) -> !transform.any_op
- %1 = transform.get_parent_op %0 <op_name = "scf.for"> : (!transform.any_op) -> !transform.op<"scf.for">
+ %1 = transform.structured.match ops{["scf.for"]} attributes{unroll_jam} in %arg1 : (!transform.any_op) -> !transform.op<"scf.for">
transform.loop.unroll_and_jam %1 factor = 4 : !transform.op<"scf.for">
transform.yield
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/228138
More information about the Mlir-commits
mailing list