[mlir][TD] Fix the order of return handles (#76929)

Replace (in tests and docs):

    %forall, %tiled = transform.structured.tile_using_forall

with (updated order of return handles):

    %tiled, %forall = transform.structured.tile_using_forall

Similar change is applied to (in the TD tutorial):

    transform.structured.fuse_into_containing_op

This update makes sure that the tests/documentation are consistent with
the Op specifications. Follow-up for #67320 which updated the order of
the return handles for `tile_using_forall`.
This commit is contained in:
Andrzej Warzyński 2024-01-04 12:54:16 +00:00 committed by GitHub
parent e147dcbcbc
commit ca5d34ec71
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 15 additions and 15 deletions

View File

@ -261,7 +261,7 @@ transform.sequence failures(propagate) {
// The actual tiling transformation takes tile sizes as attributes. It
// produces a handle to the loop generated during tiling.
%loop, %tiled_max =
%tiled_max, %loop =
transform.structured.tile_using_forall %max tile_sizes [8, 32]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
@ -271,12 +271,12 @@ transform.sequence failures(propagate) {
// important. We could also use "transform.merge_handles" to obtain a single
// handle to all operations and give it to `fuse_into_containing_op` that
// would take care of the ordering in this case.
%loop_0, %tiled_add =
%add_fused, %loop_0 =
transform.structured.fuse_into_containing_op %add into %loop
: (!transform.any_op, !transform.any_op)
-> (!transform.any_op, !transform.any_op)
%loop_1, %tiled_matmul =
transform.structured.fuse_into_containing_op %arg1 into %loop
%matmul_fused, %loop_1 =
transform.structured.fuse_into_containing_op %arg1 into %loop_0
: (!transform.op<"linalg.matmul">, !transform.any_op)
-> (!transform.any_op, !transform.any_op)
@ -304,7 +304,7 @@ transform.sequence failures(propagate) {
// The actual tiling transformation takes tile sizes as attributes. It
// produces a handle to the loop generated during tiling.
%loop, %tiled = transform.structured.tile_using_forall %max tile_sizes [8, 32]
%tiled, %loop = transform.structured.tile_using_forall %max tile_sizes [8, 32]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
// We can now fuse the other operations into the loop. Here, we fuse
@ -318,7 +318,7 @@ transform.sequence failures(propagate) {
: (!transform.any_op, !transform.any_op)
-> (!transform.any_op, !transform.any_op)
%matmul_fused, %loop_1 =
transform.structured.fuse_into_containing_op %arg1 into %loop
transform.structured.fuse_into_containing_op %arg1 into %loop_0
: (!transform.op<"linalg.matmul">, !transform.any_op)
-> (!transform.any_op, !transform.any_op)
@ -327,7 +327,7 @@ transform.sequence failures(propagate) {
// "max" operation. This illustrates the precise targeting with the transform
// dialect. Otherwise, it is difficult to differentiate "add" and "max", both
// of which having the same kind.
%loop_2, %tiled_2 =
%tiled_2, %loop_2 =
transform.structured.tile_using_forall %add_fused tile_sizes [4, 4]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%matmul_fused_2, %loop_3 =
@ -338,7 +338,7 @@ transform.sequence failures(propagate) {
// Since outlining is currently only implemented for region-holding operations
// such as loops, use tiling to size 1 to materialize the outer loop that is
// going to be outlined.
%outline_target, %_ =
%_, %outline_target =
transform.structured.tile_using_forall %tiled_2 tile_sizes [1]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.structured.fuse_into_containing_op %matmul_fused_2

View File

@ -1911,8 +1911,8 @@ def TileUsingForallOp :
tiled operations, which can all be empty.
These two returned handles point to:
- the new scf.forall op,
- the tiled op that implements TilingInterface.
- the tiled op that implements TilingInterface,
- the new scf.forall op.
#### Example using `num_threads`

View File

@ -144,7 +144,7 @@ func.func @map_nested_forall_to_threads_not_buffer(%x: tensor<32x32xf32>, %y: te
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
%matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op
%forall, %tiled = transform.structured.tile_using_forall %matmul num_threads [2, 3, 1] (mapping = [ #gpu.thread<y>, #gpu.thread<x>, #gpu.thread<z> ] )
%tiled, %forall = transform.structured.tile_using_forall %matmul num_threads [2, 3, 1] (mapping = [ #gpu.thread<y>, #gpu.thread<x>, #gpu.thread<z> ] )
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!transform.any_op) -> !transform.any_op
// expected-error @below {{only bufferized scf.forall can be mapped}}

View File

@ -389,7 +389,7 @@ module attributes {transform.with_named_sequence} {
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%forall, %tiled_generic = transform.structured.tile_using_forall %0 num_threads [7]
%tiled_generic, %forall = transform.structured.tile_using_forall %0 num_threads [7]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
@ -445,7 +445,7 @@ module attributes {transform.with_named_sequence} {
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%IN_MAT2: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %IN_MAT2 : (!transform.any_op) -> !transform.any_op
%forall, %tiled_generic = transform.structured.tile_using_forall %0 num_threads [4]
%tiled_generic, %forall = transform.structured.tile_using_forall %0 num_threads [4]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}

View File

@ -19,7 +19,7 @@ transform.sequence failures(propagate) {
%arg2: !transform.op<"linalg.elemwise_binary">):
// The actual tiling transformation takes tile sizes as attributes.
// expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
%loop, %tiled = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
%tiled, %loop = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
: (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)
// This is trying to use an invalidated handle leading to undefined behavior.
@ -64,7 +64,7 @@ transform.sequence failures(propagate) {
// The actual tiling transformation takes tile sizes as attributes.
// expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with this operand and entities nested in them}}
%loop, %tiled = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
%tiled, %loop = transform.structured.tile_using_forall %arg1 tile_sizes [4, 32]
: (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)
// Consuming an operand invalidates the consumed handle and any other handle that is