OPTIMIZATION LAB #02 — elementwise fusion: the 8× hiding in "just 3 simple ops"

TRACK: optimization | DIFFICULTY: 3 | NEEDS: mlir-opt, clang, perf or /usr/bin/time FILES: fusion.mlir (payload), bench_fusion.c (harness), run_fusion.sh (driver)

What this teaches

Most Triton kernels people write, and most of a transformer outside the matmuls (bias, activation, norm, residual add), are memory-bound: a few FLOPs per byte loaded. For those, speed = bytes moved, and the compiler's job is to move each byte once. The payload is y = relu(x*s + b) over 16 M floats, written the way a frontend emits it: one linalg.generic per op.

Two independent Xs:

STEP 1 — Run the 2×2

./run_fusion.sh

Real run (one P-core, 64 MiB in + 64 MiB out):

A_unfused                  generics=3 allocs=1 copies=1
A_unfused                  min   66.31 ms   p50  187.62 ms   useful   2.0 GB/s   OK
B_fused                    generics=1 allocs=1 copies=1
B_fused                    min   40.12 ms   p50   43.06 ms   useful   3.3 GB/s   OK
C_unfused_elim             generics=3 allocs=1 copies=0
C_unfused_elim             min   45.02 ms   p50   48.57 ms   useful   3.0 GB/s   OK
D_fused_elim               generics=1 allocs=0 copies=0
D_fused_elim               min    8.35 ms   p50    8.82 ms   useful  16.1 GB/s   OK

Neither X alone is worth much (1.5–1.7×). Both together give 8×. When optimisations interact like this, testing them one at a time, each against the baseline, would have told you "fusion is a small win" and you'd have moved on. Always ablate from the fully-optimised end too: turn each X off in D.

STEP 2 — Read the buffers, not the timings

$ grep -E 'memref.alloc|memref.copy|linalg.generic' out/A_unfused.buf.mlir | sed -E 's/\{indexing_maps.*\]\} //; s/ \{$//; s/^ *//'
%alloc = memref.alloc() alignment = 64 : memref<16777216xf32>
linalg.generic ins(%arg0 : memref<16777216xf32>) outs(%alloc : memref<16777216xf32>)
linalg.generic ins(%alloc : memref<16777216xf32>) outs(%alloc : memref<16777216xf32>)
linalg.generic ins(%alloc : memref<16777216xf32>) outs(%alloc : memref<16777216xf32>)
memref.copy %alloc, %arg1 : memref<16777216xf32> to memref<16777216xf32>
$ grep -E 'memref.alloc|memref.copy|linalg.generic' out/D_fused_elim.buf.mlir | sed -E 's/\{indexing_maps.*\]\} //; s/ \{$//; s/^ *//'
linalg.generic ins(%arg0 : memref<16777216xf32>) outs(%arg1 : memref<16777216xf32>)

Count full passes over 64 MiB arrays:

• A: read x → write tmp; read tmp → write tmp; read tmp → write tmp; read tmp → write y (copy). 8 passes, plus a 64 MiB allocation.
• D: read x → write y. 2 passes, no allocation. The minimum possible.

STEP 3 — Why B is only 1.6× better, when it removes 4 of the 8 passes

Only D has no memref.alloc. Count page faults:

$ for v in A_unfused B_fused C_unfused_elim D_fused_elim; do printf "%-16s " $v; /usr/bin/time -f "minor_faults=%R" ./out/$v 10 $v 2>&1 | grep -o 'minor_faults=.*'; done
A_unfused        minor_faults=213085
B_fused          minor_faults=213083
C_unfused_elim   minor_faults=213082
D_fused_elim     minor_faults=32849

That's 11 calls (1 warm-up + 10 timed). (213 k − 33 k) / 11 ≈ 16 384 faults per call = 64 MiB / 4 KiB pages. The temp buffer is malloc'd and free'd on every call. Blocks that big come from mmap, so each call gets fresh pages and page-faults every one of them. A page fault costs roughly 1 µs, so that's ~16 ms per call of pure OS overhead. That's more than D's entire runtime.

So the real ranking of costs in this kernel:

1. Allocating inside the kernel (page faults): the biggest cost.
2. Extra passes over memory: the second biggest.
3. The arithmetic: irrelevant.

This is why a vendor MLIR backend has dedicated copy-elimination passes, why IREE's runtime pre-plans every buffer, and why "grep the final IR for memref.alloc inside the kernel" is on every perf checklist.

STEP 4 — What fusion did to the IR

$ mlir-opt fusion.mlir --linalg-fuse-elementwise-ops | sed -n '/linalg.generic/,/} ->/p'
    %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%0 : tensor<16777216xf32>) outs(%1 : tensor<16777216xf32>) {
    ^bb0(%in: f32, %out: f32):
      %3 = arith.mulf %in, %arg2 : f32
      %4 = arith.addf %3, %arg3 : f32
      %5 = arith.maximumf %4, %cst : f32
      linalg.yield %5 : f32
    } -> tensor<16777216xf32>

Three bodies became one: mulf → addf → maximumf on the same scalar, in registers. The producer's body was inlined into the consumer because the maps line up (identity → identity). Change one map to a transpose and the pass refuses: fusion legality comes straight from the indexing maps (lab 03_linalg/01).

Try this yourself

1. Try to make the allocator reuse pages: MALLOC_MMAP_THRESHOLD_=1000000000 /usr/bin/time -f %R ./out/B_fused. On this machine the faults stay at ~213 k. glibc caps that threshold at 32 MiB on 64-bit, so a 64 MiB block is always mmap'd. Now shrink N to 4 M floats (16 MiB) in fusion.mlir and bench_fusion.c, rebuild, and repeat with and without the variable. The lesson: don't count on the allocator to save a kernel that allocates. Remove the allocation.
2. Make one intermediate a broadcast (b becomes a tensor<16777216xf32> input). Does it still fuse into one generic?
3. Add a reduction at the end (sum(y)). Does linalg-fuse-elementwise-ops fuse elementwise into a reduction? Look at the result.