[Mlir-commits] [mlir] 38517d9 - [mlir][scf] Fully unroll SCF/Affine loops (#215220)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 17 12:55:13 PDT 2026
Author: Adam Siemieniuk
Date: 2026-08-17T21:55:09+02:00
New Revision: 38517d91d15c6e0d969d1201367dfddaab9ae7ac
URL: https://github.com/llvm/llvm-project/commit/38517d91d15c6e0d969d1201367dfddaab9ae7ac
DIFF: https://github.com/llvm/llvm-project/commit/38517d91d15c6e0d969d1201367dfddaab9ae7ac.diff
LOG: [mlir][scf] Fully unroll SCF/Affine loops (#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
Added:
Modified:
mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
mlir/python/mlir/dialects/transform/loop.py
mlir/test/Dialect/SCF/transform-ops-invalid.mlir
mlir/test/Dialect/SCF/transform-ops.mlir
mlir/test/python/integration/dialects/transform.py
Removed:
################################################################################
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