Epilogue concept

Where the accumulator is finally correct

The K loop ends, the accumulator tile holds complete sums in registers, and everything the output needs happens before the single store:

tile loop:  acc = 0
            for k: acc += A[i, k] * B[k, j]
            acc = epilogue(acc, bias[j])   // registers only
            C[i, j] = acc                  // one store

Bias, activation, scaling and dtype conversion all belong in that gap. The store happens once because after the epilogue there is nothing left to compute.

Why it cannot move earlier

Inside the K loop the accumulator holds a partial sum. relu of a partial sum is not the relu of the sum, and a bias added on every K iteration is added K times. Both mistakes type-check, both produce finite numbers, and both happen to be correct when K is 1.

What it removes

The unfused chain touches the output once per stage: the matmul writes 4 bytes per f32 element, the bias pass reads and writes 8, and relu reads and writes 8, for 20 bytes per element against 4 in the fused path. Those are traffic counts, not timings, and they are an upper bound on what the memory system sees: if the intermediate stays in cache, most of those bytes never reach DRAM.

In MLIR

A linalg.generic body has no after-the-reduction phase; it runs once per iteration of the whole nest, reduction dimension included. So the epilogue is a second op on the tile, fused at tile level so the intermediate is a tile rather than a tensor, or it is written into the microkernel. This is where the intrinsics path stops being a detour and starts being the place where the schedule is expressible.

Things that break

The epilogue is folded into the matmul’s own body. The reference check passes at K=1 and fails at K=3, with the bias multiplied by K. Run the same shape twice, once with K=1 and once with K=3, and compare against the unfused chain: agreement at both means the epilogue is outside the reduction, and agreement only at K=1 identifies this bug in one step.