[Mlir-commits] [mlir] [mlir][scf] Fully unroll SCF/Affine loops (PR #215220)
Adam Siemieniuk
llvmlistbot at llvm.org
Mon Aug 10 01:52:26 PDT 2026
https://github.com/adam-smnk created https://github.com/llvm/llvm-project/pull/215220
Adds a new transform op that fully unrolls given loops. Also, updates 'loop.unroll' documentation to better reflect its functionality.
A new op is added to avoid overloading and changing the default behavior of the other existing unroll ops.
On its own, the new op complements the existing two transform ops and mirrors available SCF/Affine utils.
Assisted-by: Copilot
>From 380755773b07598bea6330bc11986260aa54412e Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Fri, 7 Aug 2026 17:45:02 +0200
Subject: [PATCH] [mlir][scf] Fully unroll SCF/Affine loops
Adds a new transform op that fully unrolls given loops.
Also, updates 'loop.unroll' documentation to better reflect its
functionality.
A new op is added to avoid overloading and changing the default
behavior of the other existing unroll ops.
On its own, the new op complements the existing two transform ops and
mirrors available SCF/Affine utils.
Assisted-by: Copilot
---
.../SCF/TransformOps/SCFTransformOps.td | 35 ++++++++-
.../SCF/TransformOps/SCFTransformOps.cpp | 23 ++++++
mlir/python/mlir/dialects/transform/loop.py | 18 +++++
.../Dialect/SCF/transform-ops-invalid.mlir | 21 ++++++
mlir/test/Dialect/SCF/transform-ops.mlir | 49 +++++++++++++
.../python/integration/dialects/transform.py | 73 ++++++++++++++++++-
6 files changed, 215 insertions(+), 4 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
index f97225477ef8b..a9ac0c3dad76b 100644
--- a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
+++ b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
@@ -302,8 +302,9 @@ def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
let summary = "Unrolls the given loop with the given unroll factor";
let description = [{
Unrolls each loop associated with the given handle to have up to the given
- 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.
+ number of loop body copies per iteration. If the unroll factor exceeds the
+ loop trip count, the main unrolled loop may have zero iterations and a
+ cleanup loop may execute all remaining iterations.
#### Return modes
@@ -330,6 +331,36 @@ def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
}];
}
+def LoopUnrollFullOp : Op<Transform_Dialect, "loop.unroll_full",
+ [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+ TransformOpInterface, TransformEachOpTrait]> {
+ let summary = "Fully unrolls the given loop";
+ let description = [{
+ Fully unrolls each loop associated with the given handle.
+
+ #### Return modes
+
+ This operation ignores non-`scf.for`, non-`affine.for` ops and drops them
+ in the return. If all the operations referred to by the `target` operand
+ unroll properly, the transform succeeds. Otherwise the transform produces a
+ silenceable failure.
+
+ Does not return handles as the loop is removed after a full unrolling.
+ }];
+
+ let arguments = (ins TransformHandleTypeInterface:$target);
+
+ let assemblyFormat = "$target attr-dict `:` type($target)";
+
+ let extraClassDeclaration = [{
+ ::mlir::DiagnosedSilenceableFailure applyToOne(
+ ::mlir::transform::TransformRewriter &rewriter,
+ ::mlir::Operation *target,
+ ::mlir::transform::ApplyToEachResultList &results,
+ ::mlir::transform::TransformState &state);
+ }];
+}
+
def LoopUnrollAndJamOp : Op<Transform_Dialect, "loop.unroll_and_jam",
[FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
TransformOpInterface, TransformEachOpTrait]> {
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index d4e46ceb946d5..1d2ec1ad0e0dc 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -403,6 +403,29 @@ transform::LoopUnrollOp::applyToOne(transform::TransformRewriter &rewriter,
return DiagnosedSilenceableFailure::success();
}
+//===----------------------------------------------------------------------===//
+// LoopUnrollFullOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::LoopUnrollFullOp::applyToOne(
+ transform::TransformRewriter &rewriter, Operation *op,
+ transform::ApplyToEachResultList &results,
+ transform::TransformState &state) {
+ LogicalResult result(failure());
+ if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op))
+ result = loopUnrollFull(scfFor);
+ else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
+ result = loopUnrollFull(affineFor);
+ else
+ return emitSilenceableError()
+ << "failed to fully unroll, incorrect type of payload";
+
+ if (failed(result))
+ return emitSilenceableError() << "failed to fully unroll";
+
+ return DiagnosedSilenceableFailure::success();
+}
+
//===----------------------------------------------------------------------===//
// LoopUnrollAndJamOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/python/mlir/dialects/transform/loop.py b/mlir/python/mlir/dialects/transform/loop.py
index c4770b1c4067e..88ac3ecab2a63 100644
--- a/mlir/python/mlir/dialects/transform/loop.py
+++ b/mlir/python/mlir/dialects/transform/loop.py
@@ -125,3 +125,21 @@ def __init__(
ip=ip,
loc=loc,
)
+
+
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class LoopUnrollFullOp(LoopUnrollFullOp):
+ """Extension for LoopUnrollFullOp."""
+
+ def __init__(
+ self,
+ target: Union[Operation, Value],
+ *,
+ ip=None,
+ loc=None,
+ ):
+ super().__init__(
+ _get_op_result_or_value(target),
+ ip=ip,
+ loc=loc,
+ )
diff --git a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
index 742b8a2861839..01b57a8ca56fa 100644
--- a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
@@ -41,6 +41,27 @@ module attributes {transform.with_named_sequence} {
// -----
+func.func @loop_unroll_full_unsupported_dynamic_trip_count(%upper_bound: index) {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ scf.for %i = %c0 to %upper_bound step %c2 {
+ arith.addi %i, %i : index
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["arith.addi"]} 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">
+ // expected-error @below {{failed to fully unroll}}
+ transform.loop.unroll_full %1 : !transform.op<"scf.for">
+ transform.yield
+ }
+}
+
+// -----
+
func.func @loop_unroll_and_jam_unsupported_trip_count_not_multiple_of_factor() {
%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 c8be2d2b3a883..ef7d5543e2d5c 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -168,6 +168,55 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-LABEL: @loop_unroll_full_op
+func.func @loop_unroll_full_op(%arg0: tensor<4xf32>) -> f32 {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c2 = arith.constant 2 : index
+ %c0_f32 = arith.constant 0.0 : f32
+ // CHECK-NOT: scf.for
+ // CHECK-COUNT-2: arith.addf
+ %res = scf.for %i = %c0 to %c4 step %c2 iter_args(%acc = %c0_f32) -> f32 {
+ %val = tensor.extract %arg0[%i] : tensor<4xf32>
+ %add = arith.addf %acc, %val : f32
+ scf.yield %add : f32
+ }
+ return %res : f32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.loop.unroll_full %0 : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @loop_unroll_full_op
+func.func @loop_unroll_full_op(%arg0: tensor<4xf32>) -> f32 {
+ %c0_f32 = arith.constant 0.0 : f32
+ // CHECK-NOT: affine.for
+ // CHECK-COUNT-4: arith.addf
+ %res = affine.for %i = 0 to 4 iter_args(%acc = %c0_f32) -> f32 {
+ %val = tensor.extract %arg0[%i] : tensor<4xf32>
+ %add = arith.addf %acc, %val : f32
+ affine.yield %add : f32
+ }
+ return %res : f32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["affine.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.loop.unroll_full %0 : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: @loop_unroll_and_jam_op
func.func @loop_unroll_and_jam_op() {
// CHECK: %[[VAL_0:.*]] = arith.constant 0 : index
diff --git a/mlir/test/python/integration/dialects/transform.py b/mlir/test/python/integration/dialects/transform.py
index 303274a8f8828..d89e88050f632 100644
--- a/mlir/test/python/integration/dialects/transform.py
+++ b/mlir/test/python/integration/dialects/transform.py
@@ -2,7 +2,7 @@
from mlir.passmanager import PassManager
from mlir.ir import Context, Location, Module, InsertionPoint, UnitAttr
-from mlir.dialects import scf, pdl, func, arith, linalg
+from mlir.dialects import scf, pdl, func, arith, linalg, tensor
from mlir.dialects.transform import (
get_parent_op,
apply_patterns_canonicalization,
@@ -10,7 +10,7 @@
any_op_t,
)
from mlir.dialects.transform.structured import structured_match
-from mlir.dialects.transform.loop import loop_unroll
+from mlir.dialects.transform.loop import loop_unroll, loop_unroll_full
from mlir.dialects.transform.extras import named_sequence, apply_patterns
from mlir.extras import types as T
from mlir.dialects.builtin import module, ModuleOp
@@ -157,3 +157,72 @@ def pats():
# CHECK: return %[[VAL_3]] : tensor<3x3xf32>
# CHECK: }
print(module_)
+
+
+# CHECK-LABEL: TEST: test_loop_unroll_full
+ at construct_and_print_in_module
+def test_loop_unroll_full(module_):
+ # CHECK-LABEL: func.func @loop_unroll_full_op(
+ # CHECK-SAME: %[[VAL_0:.*]]: tensor<4xf32>) -> f32 {
+ # CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
+ # CHECK: %[[VAL_2:.*]] = arith.constant 4 : index
+ # CHECK: %[[VAL_3:.*]] = arith.constant 2 : index
+ # CHECK: %[[VAL_4:.*]] = arith.constant 0.000000e+00 : f32
+ # CHECK: %[[VAL_5:.*]] = scf.for %[[VAL_6:.*]] = %[[VAL_1]] to %[[VAL_2]] step
+ # CHECK-SAME: %[[VAL_3]] iter_args(%[[VAL_7:.*]] = %[[VAL_4]]) -> (f32) {
+ # CHECK: %[[VAL_8:.*]] = tensor.extract %[[VAL_0]][%[[VAL_6]]] : tensor<4xf32>
+ # CHECK: %[[VAL_9:.*]] = arith.addf %[[VAL_7]], %[[VAL_8]] : f32
+ # CHECK: scf.yield %[[VAL_9]] : f32
+ # CHECK: }
+ # CHECK: return %[[VAL_5]] : f32
+ # CHECK: }
+ @func.func(T.tensor(4, T.f32()))
+ def loop_unroll_full_op(arg0):
+ c0 = arith.constant(T.index(), 0)
+ c4 = arith.constant(T.index(), 4)
+ c2 = arith.constant(T.index(), 2)
+ c0_f32 = arith.constant(T.f32(), 0.0)
+
+ for i, acc, res in scf.for_(c0, c4, c2, [c0_f32]):
+ val = tensor.extract(arg0, [i])
+ add = arith.addf(acc, val)
+ scf.yield_([add])
+
+ return res
+
+ # CHECK-LABEL: module attributes {transform.with_named_sequence} {
+ # CHECK: transform.named_sequence @__transform_main(%[[VAL_0:.*]]: !transform.any_op) {
+ # CHECK: %[[VAL_1:.*]] = transform.structured.match ops{["scf.for"]} in %[[VAL_0]] : (!transform.any_op) -> !transform.any_op
+ # CHECK: transform.loop.unroll_full %[[VAL_1]] : !transform.any_op
+ # CHECK: transform.yield
+ # CHECK: }
+ # CHECK: }
+ @module(attrs={"transform.with_named_sequence": UnitAttr.get()})
+ def mod():
+ @named_sequence("__transform_main", [any_op_t()], [])
+ def basic(target: any_op_t()):
+ loop = structured_match(any_op_t(), target, ops=["scf.for"])
+ loop_unroll_full(loop)
+
+ print(module_)
+
+ pm = PassManager.parse("builtin.module(transform-interpreter)")
+ pm.run(module_.operation)
+
+ # CHECK-LABEL: func.func @loop_unroll_full_op(
+ # CHECK-SAME: %[[VAL_0:.*]]: tensor<4xf32>) -> f32 {
+ # CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
+ # CHECK: %[[VAL_2:.*]] = arith.constant 4 : index
+ # CHECK: %[[VAL_3:.*]] = arith.constant 2 : index
+ # CHECK: %[[VAL_4:.*]] = arith.constant 0.000000e+00 : f32
+ # CHECK: %[[VAL_5:.*]] = arith.constant 4 : index
+ # CHECK: %[[VAL_6:.*]] = tensor.extract %[[VAL_0]][%[[VAL_1]]] : tensor<4xf32>
+ # CHECK: %[[VAL_7:.*]] = arith.addf %[[VAL_4]], %[[VAL_6]] : f32
+ # CHECK: %[[VAL_8:.*]] = arith.constant 1 : index
+ # CHECK: %[[VAL_9:.*]] = arith.muli %[[VAL_3]], %[[VAL_8]] : index
+ # CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_1]], %[[VAL_9]] : index
+ # CHECK: %[[VAL_11:.*]] = tensor.extract %[[VAL_0]][%[[VAL_10]]] : tensor<4xf32>
+ # CHECK: %[[VAL_12:.*]] = arith.addf %[[VAL_7]], %[[VAL_11]] : f32
+ # CHECK: return %[[VAL_12]] : f32
+ # CHECK: }
+ print(module_)
More information about the Mlir-commits
mailing list