Every optimisation you'll do on a production MLIR backend is one of: tile, fuse, pack, vectorize, unroll, distribute to threads. In MLIR you can drive these two ways: (a) inside a C++ pass with hardcoded heuristics (a vendor backend’s TilingPass.cpp, FusionPass.cpp; IREE LLVMCPUTile.cpp); (b) from a TRANSFORM SCRIPT: IR that manipulates IR. The payload (your kernel) and the schedule (how to optimise it) are separate. (b) is how you EXPERIMENT: change a tile size in a text file, rerun, time it, no rebuild. Once a schedule wins, it gets baked into a C++ pass or heuristic. IREE also lets you ship transform scripts per dispatch. Track 04 uses transform scripts for every "with X / without X" measurement. Anatomy of the script at the bottom of this file: transform.structured.match ops{["linalg.matmul"]} -> a HANDLE to payload ops transform.structured.tile_using_forall ... [8, 32] -> loops over 8x32 tiles transform.structured.fuse_into_containing_op -> pull producers into the loop Handles are like pointers into the payload. A transform CONSUMES a handle and returns new ones (the tiled op, the loop); using a consumed handle is an error, the transform-dialect version of use-after-free.
$ mlir-opt 03_transform_dialect.mlir | grep -E '= linalg\.' | grep -oE 'linalg\.(fill|matmul|generic)|-> tensor<[0-9x]+xf32>' | paste - - | sort | uniq -c
1 linalg.fill -> tensor<64x256xf32>
1 linalg.generic
1 linalg.matmul -> tensor<64x256xf32>$ mlir-opt 03_transform_dialect.mlir --transform-interpreter --test-transform-dialect-erase-schedule --canonicalize --cse > out/tiled.mlir && sed -n '/scf.forall/,/parallel_insert_slice/p' out/tiled.mlir
%1 = scf.forall (%arg3, %arg4) in (8, 8) shared_outs(%arg5 = %0) -> (tensor<64x256xf32>) {
%2 = affine.apply #map(%arg3)
%3 = affine.apply #map1(%arg4)
%extracted_slice = tensor.extract_slice %arg0[%2, 0] [8, 128] [1, 1] : tensor<64x128xf32> to tensor<8x128xf32>
%extracted_slice_0 = tensor.extract_slice %arg1[0, %3] [128, 32] [1, 1] : tensor<128x256xf32> to tensor<128x32xf32>
%extracted_slice_1 = tensor.extract_slice %arg5[%2, %3] [8, 32] [1, 1] : tensor<64x256xf32> to tensor<8x32xf32>
%4 = linalg.fill ins(%cst : f32) outs(%extracted_slice_1 : tensor<8x32xf32>) -> tensor<8x32xf32>
%5 = linalg.matmul ins(%extracted_slice, %extracted_slice_0 : tensor<8x128xf32>, tensor<128x32xf32>) outs(%4 : tensor<8x32xf32>) -> tensor<8x32xf32>
%extracted_slice_2 = tensor.extract_slice %arg2[%3] [32] [1] : tensor<256xf32> to tensor<32xf32>
%6 = linalg.generic {indexing_maps = [#map2, #map3, #map2], iterator_types = ["parallel", "parallel"]} ins(%5, %extracted_slice_2 : tensor<8x32xf32>, tensor<32xf32>) outs(%extracted_slice_1 : tensor<8x32xf32>) {
^bb0(%in: f32, %in_3: f32, %out: f32):
%7 = arith.addf %in, %in_3 : f32
%8 = arith.maximumf %7, %cst : f32
linalg.yield %8 : f32
} -> tensor<8x32xf32>
scf.forall.in_parallel {
tensor.parallel_insert_slice %6 into %arg5[%2, %3] [8, 32] [1, 1] : tensor<8x32xf32> into tensor<64x256xf32>
READ IT LINE BY LINE — this is the most important IR in the whole lab:
scf.forall (%arg3, %arg4) in (8, 8) 64 independent tiles (64/8 x 256/32).
"forall" = parallel; later lowered to
threads (IREE: workgroups; vendor
backends: their own pass) or a plain loop.
extract_slice %arg0[%2, 0] [8, 128] the matmul tile needs 8 ROWS of A (all of K)...
extract_slice %arg1[0, %3] [128, 32] ...and 32 COLUMNS of B. Tiling math done for you
from the indexing maps (lab 01).
linalg.fill / linalg.matmul 8x32 the producer now computes ONLY this tile.
linalg.generic 8x32 bias+relu consumes it immediately, while hot in cache.
parallel_insert_slice each tile writes its piece of the result.
The 64x256 matmul result is never materialised in memory. Before, it was
written out (64 KiB) and read back by relu. After, each 8x32 tile (1 KiB)
stays in L1/registers. That's FUSION, and on memory-bound kernels (most of
what Triton users write) it's the biggest single win.$ grep -E '= linalg\.' out/tiled.mlir | grep -oE 'linalg\.(fill|matmul|generic)|ins\([^)]*\)|-> tensor<[0-9x]+xf32>' | sed -E 's/ins\(%[^:]*: /ins(/' | paste -sd' ' | sed 's/ linalg/\nlinalg/g'
linalg.fill ins(f32) -> tensor<8x32xf32>
linalg.matmul ins(tensor<8x128xf32>, tensor<128x32xf32>) -> tensor<8x32xf32>
linalg.generic ins(tensor<8x32xf32>, tensor<32xf32>)
Every op now works on tiles: 8x128 of A times 128x32 of B -> 8x32.$ sed 's/fuse_into_containing_op %fill into %l1/fuse_into_containing_op %fill into %relu/' 03_transform_dialect.mlir | mlir-opt --transform-interpreter 2>&1 >/dev/null | grep -E 'error|note' | cut -c1-150 | head -3
<stdin>:125:20: error: op uses a handle invalidated by a previously executed transform op
<stdin>:116:13: note: handle to invalidated ops
<stdin>:120:21: note: invalidated by this transform op that consumes its operand #0 and invalidates all handles to payload IR entities associated with
tile_using_forall CONSUMED %relu (the old untiled op no longer exists) and
gave back %tiled/%loop. Using %relu afterwards is a use-after-free, and
the interpreter catches it with three pointers: the bad use, where the
handle was made, and who invalidated it. (Note fuse_into_containing_op
only READS its loop operand; reusing %loop there would be legal. Each
op's docs say which operands it consumes.)bias + relu applied to a matmul result (the classic "linear layer")
1. Change tile_sizes to [16, 64] and [4, 256]. How many tiles? What slice of B does each read? 2. Delete the two fuse ops. Where do the matmul and fill end up? (Outside the loop, full size: tiling without fusion.) 3. Add `transform.structured.tile_using_for %mm_f tile_sizes [0, 0, 32]` after fusion to also tile K. What does the loop carry now? ============================================================
#id = affine_map<(i, j) -> (i, j)>
#col = affine_map<(i, j) -> (j)>
func.func @linear_relu(%a: tensor<64x128xf32>, %b: tensor<128x256xf32>, %bias: tensor<256xf32>) -> tensor<64x256xf32> {
%z = arith.constant 0.0 : f32
%e = tensor.empty() : tensor<64x256xf32>
%init = linalg.fill ins(%z : f32) outs(%e : tensor<64x256xf32>) -> tensor<64x256xf32>
%mm = linalg.matmul ins(%a, %b : tensor<64x128xf32>, tensor<128x256xf32>) outs(%init : tensor<64x256xf32>) -> tensor<64x256xf32>
%r = linalg.generic {indexing_maps = [#id, #col, #id], iterator_types = ["parallel", "parallel"]}
ins(%mm, %bias : tensor<64x256xf32>, tensor<256xf32>) outs(%e : tensor<64x256xf32>) {
^bb0(%x: f32, %bb: f32, %o: f32):
%s = arith.addf %x, %bb : f32
%m = arith.maximumf %s, %z : f32
linalg.yield %m : f32
} -> tensor<64x256xf32>
return %r : tensor<64x256xf32>
}
// SCHEDULE: a separate module of transform ops that says HOW to optimise the payload.
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
// 1. find the ops
%relu = transform.structured.match ops{["linalg.generic"]} in %root : (!transform.any_op) -> !transform.any_op
%mm = transform.structured.match ops{["linalg.matmul"]} in %root : (!transform.any_op) -> !transform.any_op
%fill = transform.structured.match ops{["linalg.fill"]} in %root : (!transform.any_op) -> !transform.any_op
// 2. tile the CONSUMER (relu) into 8x32 tiles with a parallel loop (scf.forall)
%tiled, %loop = transform.structured.tile_using_forall %relu tile_sizes [8, 32]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
// 3. pull the producers INTO the loop: each tile computes only the matmul slice it needs
%mm_f, %l1 = transform.structured.fuse_into_containing_op %mm into %loop
: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
%fill_f, %l2 = transform.structured.fuse_into_containing_op %fill into %l1
: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
}A lab is done when you can say what it showed on your machine — not when you have read it.
It decides which problems come first, which drills open, which labs you see and how the Learn surface is organised. Change it any time from the switch in the top bar, or press Escape to look around first.