@@ -133,6 +133,8 @@ This will generate two files, `MyExtension.h.inc` and `MyExtension.cpp.inc`, tha
133133``` c++
134134// In MyExtension.cpp.
135135
136+ #include " MyExtension.h."
137+
136138#define GET_OP_CLASSES
137139#include " MyExtension.cpp.inc"
138140
@@ -245,7 +247,8 @@ must be modified with the provided rewriter.
245247 return diag;
246248 }
247249
248- updateCallee(call, getNewTarget());
250+ // Use rewriter to modify the callee in place.
251+ rewriter.modifyOpInPlace(call, [&]() { call.setCallee(getNewTarget()); });
249252 }
250253
251254 // If everything went well, return success.
@@ -263,7 +266,7 @@ void ChangeCallTargetOp::getEffects(
263266 // Indicate that the `call` handle is only read by this operation because the
264267 // associated operation is not erased but rather modified in-place, so the
265268 // reference to it remains valid.
266- onlyReadsHandle(getCall (), effects);
269+ onlyReadsHandle(this->getOperation()->getOpOperands().front (), effects);
267270
268271 // Indicate that the payload is modified by this operation.
269272 modifiesPayload(effects);
@@ -288,67 +291,80 @@ After registering the extension, it becomes possible to use our new operation in
288291```mlir
289292module attributes {transform.with_named_sequence} {
290293 transform.named_sequence @__transform_main(
291- %arg0: !transform.any_op,
292- %arg1: !transform.op<"linalg.matmul">,
293- %arg2: !transform.op<"linalg.elementwise">) {
294+ %arg0: !transform.any_op,
295+ %arg1: !transform.op<"linalg.matmul">,
296+ %arg2: !transform.op<"linalg.elementwise">) {
294297 // Since the %arg2 handle is associated with both elementwise operations,
295298 // we need to split it into two handles so we can target only the second
296299 // elementwise operation.
297300 %add, %max = transform.split_handle %arg2
298301 : (!transform.op<"linalg.elementwise">)
299- -> (!transform.any_op, !transform.any_op)
302+ -> (!transform.any_op, !transform.any_op)
300303
301304 // The actual tiling transformation takes tile sizes as attributes. It
302305 // produces a handle to the loop generated during tiling.
303- %loop , %tiled = transform.structured.tile_using_forall %max
306+ %tiled , %loop = transform.structured.tile_using_forall %max
304307 tile_sizes [8, 32]
305308 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
306309
307310 // We can now fuse the other operations into the loop. Here, we fuse
308- // operations one-by-one. This requires the operation that is being fused
309- // to define the value used within the loop, so the order of such fusions
310- // is important. We could also use "transform.merge_handles" to obtain
311- // a single handle to all operations and give it to
312- // `fuse_into_containing_op` that would take care of the ordering in this
313- // case.
314- %add_fused = transform.structured.fuse_into_containing_op %add into %loop
315- : (!transform.any_op, !transform.any_op) -> !transform.any_op
316- %matmul_fused = transform.structured.fuse_into_containing_op %arg1
317- into %loop
318- : (!transform.op<"linalg.matmul">, !transform.any_op)
319- -> !transform.any_op
311+ // operations one by one. This requires the operation that is being fused to
312+ // define the value used within the loop, so the order of such fusions is
313+ // important. We could also use "transform.merge_handles" to obtain a single
314+ // handle to all operations and give it to `fuse_into_containing_op` that
315+ // would take care of the ordering in this case.
316+ %add_fused, %loop_0 =
317+ transform.structured.fuse_into_containing_op %add into %loop
318+ : (!transform.any_op, !transform.any_op)
319+ -> (!transform.any_op, !transform.any_op)
320+ %matmul_fused, %loop_1 =
321+ transform.structured.fuse_into_containing_op %arg1 into %loop_0
322+ : (!transform.op<"linalg.matmul">, !transform.any_op)
323+ -> (!transform.any_op, !transform.any_op)
320324
321325 // Tile again to get the desired size. Note that this time this tiles the
322326 // "add" operation and fuses matmul into the loop, but doesn't affect the
323327 // "max" operation. This illustrates the precise targeting with the
324328 // transform dialect. Otherwise, it is difficult to differentiate "add" and
325329 // "max", both of which having the same kind.
326- %loop_2, %tiled_2 = transform.structured.tile_using_forall %add_fused
327- tile_sizes [4, 4]
328- : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
329- %matmul_fused_2 = transform.structured.fuse_into_containing_op %matmul_fused
330- into %loop_2
331- : (!transform.any_op, !transform.any_op) -> !transform.any_op
330+ %tiled_2, %loop_2 =
331+ transform.structured.tile_using_forall %add_fused tile_sizes [4, 4]
332+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
333+ %matmul_fused_2, %loop_3 =
334+ transform.structured.fuse_into_containing_op %matmul_fused into %loop_2
335+ : (!transform.any_op, !transform.any_op)
336+ -> (!transform.any_op, !transform.any_op)
332337
333338 // Since outlining is currently only implemented for region-holding
334339 // operations such as loops, use tiling to size 1 to materialize the outer
335340 // loop that is going to be outlined.
336- %outline_target, %_ = transform.structured.tile_using_forall %tiled_2 tile_sizes [1]
337- : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
338- transform.structured.fuse_into_containing_op %matmul_fused_2 into %outline_target
339- : (!transform.any_op, !transform.any_op) -> !transform.any_op
341+ %_, %outline_target =
342+ transform.structured.tile_using_forall %tiled_2 tile_sizes [1]
343+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
344+ transform.structured.fuse_into_containing_op %matmul_fused_2
345+ into %outline_target
346+ : (!transform.any_op, !transform.any_op)
347+ -> (!transform.any_op, !transform.any_op)
340348 %func, %call = transform.loop.outline %outline_target
341349 {func_name = "outlined"}
342- : (!transform.any_op) -> (!transform.any_op, !transform.any_op )
350+ : (!transform.any_op) -> (!transform.any_op, !transform.op<"func.call"> )
343351
344352 // Rewrite the call target.
345- transform.my.change_call_target %call, "microkernel" : !transform.any_op
346-
353+ transform.my.change_call_target %call, "microkernel" : !transform.op<"func.call">
347354 transform.yield
348355 }
349356}
350357```
351358
359+ When you run it with the interpreter, it produces the following error.
360+
361+ ```
362+ sequence.mlir:7:8: error: 'func.call' op 'microkernel' does not reference a valid function
363+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<add> ins(%0, %arg2 : tensor<512x512xf32>, tensor<512x512xf32>) outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32>
364+ ^
365+ sequence.mlir:7:8: note: see current operation: %39 = "func.call"(%32, %33, %34, %36, %37) <{callee = @microkernel}> : (tensor<4x512xf32>, tensor<512x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4xf32>
366+ ```
367+
352368## Appendix: Autogenerated Documentation
353369
354370[ include "Tutorials/transform/MyExtensionCh2.md"]
0 commit comments