AARCH64 LAB #04 — int8 matmul on smmla: pack → mmt4d → unpack (IREE's data tiling, in miniature)

TRACK: aarch64 | DIFFICULTY: 4 | NEEDS: mlir-opt, clang, qemu, tools/a64_runtime.c FILE: 04_i8mm_mmt4d.mlir (a copy of upstream mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir)

What this teaches

Quantized models (int8 weights and activations, int32 accumulation) are the bread and butter of on-device inference, and Armv8.6's FEAT_I8MM has an instruction for them: smmla v.4s, a.16b, b.16b computes a 2×8 · 8×2 int8 block product into a 2×2 int32 tile. That's 32 multiply-adds in one instruction. To use it, the data must already be laid out in those tiny tiles. That's what packing does:

C[M][N] = A[M][K] · B[K][N]                        (row-major, arbitrary sizes)
   | linalg.pack A -> A'[M/4][K/8][4][8]           (pad to multiples, tile)
   | linalg.pack B -> B'[N/4][K/8][4][8]           (tiled AND transposed: the "t" in mmt4d)
   v linalg.mmt4d  A', B' -> C'[M/4][N/4][4][4]    (inner 4x4x8 block = contiguous, static, no masks)
   | linalg.unpack C' -> C[M][N]

This is exactly IREE's "data tiling" (mmt4d + ukernels in runtime/src/iree/builtins/ukernel/arch/arm_64/mmt4d_arm_64_i8mm.c) and the reason IREE is the team's reference. Packing costs a pass over A and B, but the inner kernel becomes perfectly regular: static shapes, contiguous loads, no masks, and one smmla per block.

STEP 1 — Compile the upstream test with our toolchain

mlir-opt 04_i8mm_mmt4d.mlir -transform-interpreter -test-transform-dialect-erase-schedule \
  -cse -canonicalize -convert-vector-to-scf -o out/i8.t.mlir
grep -oE 'arm_neon\.[a-z.]+' out/i8.t.mlir | sort | uniq -c                 # -> 4 arm_neon.intr.smmla
mlir-opt out/i8.t.mlir -convert-vector-to-llvm="enable-arm-neon enable-arm-i8mm" -test-lower-to-llvm -o out/i8.llvm.mlir
mlir-translate --mlir-to-llvmir out/i8.llvm.mlir -o out/i8.ll
clang --target=aarch64-linux-gnu -O2 -march=armv8.6-a+i8mm -c out/i8.ll -o out/i8.o -Wno-override-module
clang --target=aarch64-linux-gnu -O2 -march=armv8.6-a+i8mm -S out/i8.ll -o - -Wno-override-module | grep -cE '^\s+smmla'   # -> 8

The key transform is LowerContractionToNeonI8MMPattern (enabled by enable-arm-i8mm): a vector.contract on vector<4x8xi8> x vector<4x8xi8> -> vector<4x4xi32> becomes arm_neon.intr.smmla ops. It only matches when the shapes are exactly right, which is why the packing exists.

STEP 2 — Link against a tiny AArch64 runtime and run

The test calls printMemrefI32, printString and memrefCopy from MLIR's runner-utils, which only exist here as x86 C++ libraries:

llvm-nm -u out/i8.o          # malloc memcpy memrefCopy printMemrefI32 printString

tools/a64_runtime.c implements those in plain C (the signatures mirror mlir/lib/ExecutionEngine/CRunnerUtils.cpp):

a64cc -O2 out/i8.o ../tools/a64_runtime.c -o out/i8
a64run ./out/i8

RESULT FROM linalg.mmt4d:
Unranked Memref rank = 2 data =
[[193,   200,   207,   214,   221,   228,   235,   242,   249,   256,   263,   270,   277],
 [194,   201,   208,   ...
RESULT FROM linalg.matmul:
[[193,   200,   207,   ...

Both match the // CHECK: lines in the upstream test (3·4·16 + C[i][j]). Always compute a reference variant alongside the optimized one; the upstream test does exactly that.

STEP 3 — Feature presence is not feature use

a64run -cpu cortex-a76 ./out/i8

qemu: uncaught target signal 4 (Illegal instruction) - core dumped

Cortex-A76 has no I8MM. The binary was compiled for +i8mm unconditionally, so it just dies. A shipped compiler must either:

• compile for the exact device (-mcpu= of the phone's cores; JIT compilers like Triton can query the CPU at runtime), or
• emit multiple variants and dispatch on HWCAP (see ~/debug_lab/aarch64/04_dispatch.c); IREE ukernels do this (mmt4d_arm_64_entry_point.c picks i8mm / dotprod / base).

Try this yourself

1. Change the RHS inner tile from [4, 8] to [4, 4]. Does the i8mm pattern still fire? (Count smmla.)
2. Build without enable-arm-i8mm. What does the contraction lower to? Count instructions.
3. Read mmt4d_arm_64_i8mm.c in $IREE_SRC/runtime/src/iree/builtins/ukernel/arch/arm_64/ and find the vmmlaq_s32 intrinsic calls: the C-level version of what this MLIR generates.