Tile-level fusion

Tile in, tile out

Tile-level fusion keeps the producer’s tile loop and puts the consumer’s work inside it, so the intermediate is a tile-sized buffer rather than a tensor. In the transform dialect:

%tile, %loops:2 = transform.structured.tile_using_for %matmul tile_sizes [32, 32, 0]
transform.structured.fuse_into_containing_op %epilogue into %loops#1

The result counts and handle types are release-bound; take the current form from the Transform dialect page and from transform.structured.fuse_into_containing_op in mlir-opt --help, not from a snippet in a blog post.

What changes

Memory: the intermediate becomes 32 by 32 by 4 bytes instead of M * N * 4, so it stays in L1 or L2 and the consumer’s read hits the line the producer just wrote. It does not disappear. Bufferization still allocates the tile buffer unless the consumer reads the accumulator memref directly, which is the p2-12 check.

Tails: the consumer inherits the producer’s clamped extents. An elementwise consumer is fine because it has the same iteration space. A consumer with its own bounds, such as a bias load indexed by the tile’s column offset, must clamp with the same value. A consumer that reduces needs the full extent and is out of scope here.

What it buys in the schedule

The producer’s K loop stays in registers; only the completed tile is written; the consumer reads that tile from cache and writes the final output once. The win is the traversal, not the arithmetic, so it shows up as bytes rather than flops.

Things that break

Producer and consumer disagree about the tail. The consumer processes the full tile while the producer filled only the valid part, so the epilogue reads lanes the producer never wrote. The values are finite garbage, so the reference check fails on a handful of elements at ragged shapes only. Fill the tile buffer with NaN in a debug build: garbage lanes surface immediately. The other failure is fusing at the wrong level so the tile buffer is hoisted out of the loop, which turns it back into a whole-tensor intermediate with an extra copy.