[Mlir-commits] [mlir] 1e70ab5 - [mlir][transform] Update transform.loop.peel (reland #67482)
Andrzej Warzynski
llvmlistbot at llvm.org
Thu Sep 28 07:35:55 PDT 2023
Author: Andrzej Warzynski
Date: 2023-09-28T14:35:46Z
New Revision: 1e70ab5f0d54d49c2c14c0c80ec19c82e2701db4
URL: https://github.com/llvm/llvm-project/commit/1e70ab5f0d54d49c2c14c0c80ec19c82e2701db4
DIFF: https://github.com/llvm/llvm-project/commit/1e70ab5f0d54d49c2c14c0c80ec19c82e2701db4.diff
LOG: [mlir][transform] Update transform.loop.peel (reland #67482)
This patch updates `transform.loop.peel` so that this Op returns two
rather than one handle:
* one for the peeled loop, and
* one for the remainder loop.
Also, following this change this Op will fail if peeling fails. This is
consistent with other similar Ops that also fail if no transformation
takes place.
Relands #67482 with an extra fix for transform_loop_ext.py
Added:
Modified:
mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
mlir/python/mlir/dialects/_loop_transform_ops_ext.py
mlir/test/Dialect/Linalg/transform-op-fuse.mlir
mlir/test/Dialect/SCF/transform-ops-invalid.mlir
mlir/test/Dialect/SCF/transform-ops.mlir
mlir/test/python/dialects/transform_loop_ext.py
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
index 207a004c54ef5af..700a29139a35b10 100644
--- a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
+++ b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
@@ -142,23 +142,22 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
This operation ignores non-scf::ForOp ops and drops them in the return.
- This operation always succeeds and returns the scf::ForOp with the
- postcondition: "the loop trip count is divisible by the step".
- This operation may return the same unmodified loop handle when peeling did
- not modify the IR (i.e. the loop trip count was already divisible).
+ This operation returns two scf::ForOp Ops, with the first Op satisfying
+ the postcondition: "the loop trip count is divisible by the step". The
+ second loop Op contains the remaining iteration. Note that even though the
+ Payload IR modification may be performed in-place, this operation consumes
+ the operand handle and produces a new one.
- Note that even though the Payload IR modification may be performed
- in-place, this operation consumes the operand handle and produces a new
- one.
+ #### Return Modes
- TODO: Return both the peeled loop and the remainder loop.
+ Produces a definite failure if peeling fails.
}];
let arguments =
(ins Transform_ScfForOp:$target,
DefaultValuedAttr<BoolAttr, "false">:$fail_if_already_divisible);
- // TODO: Return both the peeled loop and the remainder loop.
- let results = (outs TransformHandleTypeInterface:$transformed);
+ let results = (outs TransformHandleTypeInterface:$peeled_loop,
+ TransformHandleTypeInterface:$remainder_loop);
let assemblyFormat =
"$target attr-dict `:` functional-type(operands, results)";
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index d7e8c38478ced1a..65d503d7c4ad8b8 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -226,14 +226,16 @@ transform::LoopPeelOp::applyToOne(transform::TransformRewriter &rewriter,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
scf::ForOp result;
- // This helper returns failure when peeling does not occur (i.e. when the IR
- // is not modified). This is not a failure for the op as the postcondition:
- // "the loop trip count is divisible by the step"
- // is valid.
LogicalResult status =
scf::peelForLoopAndSimplifyBounds(rewriter, target, result);
- // TODO: Return both the peeled loop and the remainder loop.
- results.push_back(failed(status) ? target : result);
+ if (failed(status)) {
+ DiagnosedSilenceableFailure diag = emitSilenceableError()
+ << "failed to peel";
+ return diag;
+ }
+ results.push_back(target);
+ results.push_back(result);
+
return DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/python/mlir/dialects/_loop_transform_ops_ext.py b/mlir/python/mlir/dialects/_loop_transform_ops_ext.py
index 3536d45ab7369c4..1cdb2b9e77b5afe 100644
--- a/mlir/python/mlir/dialects/_loop_transform_ops_ext.py
+++ b/mlir/python/mlir/dialects/_loop_transform_ops_ext.py
@@ -66,7 +66,8 @@ class LoopPeelOp:
def __init__(
self,
- result_type: Type,
+ main_loop_type: Type,
+ remainder_loop_type: Type,
target: Union[Operation, Value],
*,
fail_if_already_divisible: Union[bool, BoolAttr] = False,
@@ -74,7 +75,8 @@ def __init__(
loc=None,
):
super().__init__(
- result_type,
+ main_loop_type,
+ remainder_loop_type,
_get_op_result_or_value(target),
fail_if_already_divisible=(
fail_if_already_divisible
diff --git a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
index d3772e9c04eef52..72c61661b55efd3 100644
--- a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
@@ -48,7 +48,7 @@ transform.sequence failures(propagate) {
%0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]}
: (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)
- transform.loop.peel %loops#0 : (!transform.op<"scf.for">) -> !transform.any_op
+ transform.loop.peel %loops#0 : (!transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
}
// -----
diff --git a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
index b69ea0f5975cd80..e8212f500e693ad 100644
--- a/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops-invalid.mlir
@@ -59,3 +59,23 @@ transform.sequence failures(propagate) {
// expected-error @below {{failed to outline}}
transform.loop.outline %0 {func_name = "foo"} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
}
+
+// -----
+
+func.func @test_loops_do_not_get_peeled() {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 40 : index
+ %step = arith.constant 5 : index
+ scf.for %i = %lb to %ub step %step {
+ arith.addi %i, %i : index
+ }
+ return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+ %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.loop.get_parent_for %0 : (!transform.any_op) -> !transform.op<"scf.for">
+ // expected-error @below {{failed to peel}}
+ transform.loop.peel %1 : (!transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
+}
diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index 7718ebcd1fc6ae0..043998711f64e5f 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -87,16 +87,18 @@ transform.sequence failures(propagate) {
// CHECK-LABEL: @loop_peel_op
func.func @loop_peel_op() {
// CHECK: %[[C0:.+]] = arith.constant 0
- // CHECK: %[[C42:.+]] = arith.constant 42
+ // CHECK: %[[C41:.+]] = arith.constant 41
// CHECK: %[[C5:.+]] = arith.constant 5
// CHECK: %[[C40:.+]] = arith.constant 40
// CHECK: scf.for %{{.+}} = %[[C0]] to %[[C40]] step %[[C5]]
// CHECK: arith.addi
- // CHECK: scf.for %{{.+}} = %[[C40]] to %[[C42]] step %[[C5]]
+ // CHECK: scf.for %{{.+}} = %[[C40]] to %[[C41]] step %[[C5]]
// CHECK: arith.addi
%0 = arith.constant 0 : index
- %1 = arith.constant 42 : index
+ %1 = arith.constant 41 : index
%2 = arith.constant 5 : index
+ // expected-remark @below {{main loop}}
+ // expected-remark @below {{remainder loop}}
scf.for %i = %0 to %1 step %2 {
arith.addi %i, %i : index
}
@@ -107,7 +109,10 @@ transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1 = transform.loop.get_parent_for %0 : (!transform.any_op) -> !transform.op<"scf.for">
- transform.loop.peel %1 : (!transform.op<"scf.for">) -> !transform.any_op
+ %main_loop, %remainder = transform.loop.peel %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
+ // Make sure
+ transform.test_print_remark_at_operand %main_loop, "main loop" : !transform.op<"scf.for">
+ transform.test_print_remark_at_operand %remainder, "remainder loop" : !transform.op<"scf.for">
}
// -----
diff --git a/mlir/test/python/dialects/transform_loop_ext.py b/mlir/test/python/dialects/transform_loop_ext.py
index 95339b2363e67c0..daec6707d6743b9 100644
--- a/mlir/test/python/dialects/transform_loop_ext.py
+++ b/mlir/test/python/dialects/transform_loop_ext.py
@@ -59,7 +59,7 @@ def loopPeel():
transform.OperationType.get("scf.for"),
)
with InsertionPoint(sequence.body):
- loop.LoopPeelOp(pdl.OperationType.get(), sequence.bodyTarget)
+ loop.LoopPeelOp(transform.AnyOpType.get(), transform.AnyOpType.get(), sequence.bodyTarget)
transform.YieldOp()
# CHECK-LABEL: TEST: loopPeel
# CHECK: = transform.loop.peel %
More information about the Mlir-commits
mailing list