Shape specialization

Guard, then call a clone

Specialization is a dispatch decision at the boundary, not a transformation inside the kernel. Clone the function once per hot shape, keep the generic version, and choose between them on values you already have:

if (m == 512 && n == 512 && k == 512)
  matmul_512(a, b, c);
else
  matmul_dyn(a, b, c, m, n, k);

The static clone has tensor<512x512xf32> in its signature, so every tile bound, stride and vector length is a constant again, and the pipeline that produced the fast path applies unchanged. The generic clone keeps the dynamic types and the tail machinery from the previous nodes.

Which shapes are hot

Specialize the shapes you measured, not the shapes you imagine. A histogram of (M, N, K) over the workloads you care about gives a coverage curve, and the histogram tells you how many entries are worth a clone. Every clone costs code size, compile time and a second place to fix a bug, so the set stays small.

In MLIR the clone is a second func.func with refined types, and the dispatch is an scf.if or a C-level branch. If the guard proves the extents, the tensor.cast from tensor<?x?xf32> to tensor<512x512xf32> carries the same fact in the type system; it is still a promise, so the guard and the cast have to agree.

Check whether your pinned release ships a specialization pass before writing one; the usual route is a cloned function selected by your driver.

Things that break

The fallback is the path everything else takes, and it is the one that rots because tests only exercise the hot shape. A guard that refines M and N but not K leaves a dynamic reduction. And specializing every observed shape turns a small kernel into a large binary with an instruction-cache problem, which shows up as a regression at small sizes.