IREE LAB #01 — Using the reference compiler: same matmul, IREE vs our schedule

TRACK: iree | DIFFICULTY: 3 | NEEDS: ../.venv (iree-base-compiler / -runtime 3.11), tools/mca_loop.sh FILES: matmul.mlir (tensor version of 04's payload), run_iree.sh

What this teaches

IREE (with its ARM backend) is the reference this lab measures against. Concretely, that means: for a kernel you're optimising in a production MLIR backend, compile the same computation with IREE for the same CPU, and use it as a target number (how fast can MLIR-generated code be here?) and a source of techniques (which strategy, which tile sizes, which lowering did it pick?). This lab does that for the 512³ matmul from track 04, on x86 (measured) and AArch64 (compiled, then modelled with llvm-mca).

STEP 1 — Three IREE builds, measured under the same conditions as lab 04/01

./run_iree.sh

Real output (pinned P-core, single thread, --iree-llvmcpu-target-cpu=host):

iree_default                 median     233 ms     1.15 GFLOP/s
iree_datatiling              median    6.03 ms    44.52 GFLOP/s
iree_dt_ukernels             median    6.01 ms    44.66 GFLOP/s

And ours, from lab 04/01: naive 1.16, best transform schedule ~30–46 (noisy; 4×16 tile around 42–46).

Three findings, each worth knowing before you ever cite "IREE's number":

1. IREE's default is as slow as a naive loop on this input. The fast path is data tiling (--iree-opt-data-tiling), which is off by default in this release. A reference number without its flags is meaningless.
2. With data tiling, IREE and our hand schedule land in the same range. The techniques in track 04 are the real ones.
3. Ukernels (--iree-llvmcpu-enable-ukernels=all) add nothing here: for f32 on x86 the codegen path is already as good. (For int8 on Arm with i8mm, the ukernels are the hand-written mmt4d_arm_64_i8mm.c, lab 05/04.)

STEP 2 — See WHAT IREE decided (not just how fast it was)

../.venv/bin/iree-compile matmul.mlir --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu \
  --iree-llvmcpu-target-cpu=host --iree-hal-dump-executable-configurations-to=out/cfg -o /dev/null
grep -oE 'translation_info<pipeline = [A-Za-z]+|lowering_config = #iree_cpu[^}]*' out/cfg/*.mlir

lowering_config = #iree_cpu.lowering_config<cache_parallel = [64, 64, 0], distribution = [64, 64, 0], vector_common_parallel = [1, 1, 0], vector_reduction = [0, 0, 8]>
translation_info<pipeline = CPUDoubleTilingExpert

Read the default decision: 64×64 cache tiles, then vector tiles of [1, 1] over (m, n) with a vector reduction of 8 along k: every output element is a dot product over K, reading B down a column. That's lab 04/01's naive access pattern, vectorized. Hence 1.15 GFLOP/s.

With data tiling the dispatch becomes Mmt4dTilingExpert (plus two CPUDataTiling "encoding" dispatches, the packs), exactly the pack → mmt4d → unpack structure from lab 05/04.

These configs come from KernelDispatch.cpp: $IREE_LLVMCPU/KernelDispatch.cpp (see tools/env.sh). The other dump flags worth knowing: --iree-hal-dump-executable-sources-to=DIR (the dispatch as linalg, before codegen) and --iree-hal-dump-executable-intermediates-to=DIR (.ll before and after LLVM opt, .o, .s).

STEP 3 — AArch64: IREE's microkernel vs ours, in the same model

../.venv/bin/iree-compile matmul.mlir --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu \
  --iree-llvmcpu-target-triple=aarch64-unknown-linux-gnu --iree-llvmcpu-target-cpu=neoverse-v1 \
  --iree-opt-data-tiling=true --iree-hal-dump-executable-intermediates-to=out/a64_true -o out/mm_a64_true.vmfb
MCA_FUNC=matmul_dispatch_0_matmul_512x512x512_f32 ../tools/mca_loop.sh \
  out/a64_true/module_matmul_linked_embedded_elf_arm_64.s -mtriple=aarch64 -mcpu=neoverse-v1

Results (Neoverse-V1 model, 4 NEON FMA pipes → 4.0 fmla/cycle peak):

What to take from it:

• IREE picked an 8×8 f32 tile for Neoverse-V1 (16 fmla v.4s by-element on 16 accumulators), the same tile llvm-mca ranked best in lab 05/03. Independent confirmation that the method works.
• It does it in 23 instructions vs our 29: after packing, A and B are read contiguously (a few ldp q), with no per-row address arithmetic. That's what packing buys even when the FMA count is equal.
• The model's blind spot: IREE's default loop scores 3.9 fmla/cycle in llvm-mca, yet the same strategy measured 39× slower on x86 (STEP 1). llvm-mca assumes every load hits L1; the default walks B down its columns and misses constantly. The model ranks loop bodies; only a measurement ranks memory access patterns. Use both.

The workflow to reuse on a production MLIR backend

1. Express the kernel as linalg on tensors (triton-shared's output is almost that).
2. iree-compile it for the same target, with the flags that matter (data tiling, ukernels, target cpu/features), and dump configurations.
3. Compare: strategy, tile sizes, inner-loop instruction mix (llvm-mca), and measured time on the device.
4. The gap between IREE and a production MLIR backend is your to-do list, and IREE's source for that strategy (KernelDispatch.cpp, Passes.cpp) is the recipe.

Try this yourself

1. Add --iree-llvmcpu-target-cpu-features=+sve (with neoverse-v1). Does IREE use SVE for this matmul?
2. Change the matmul to i8×i8→i32 and compile with +i8mm and data tiling. Find smmla in the .s.
3. Run ./run_iree.sh with --task_topology_max_group_count=4 (4 threads). How does IREE scale?