[techmore]

Builds on #26 (the AGX3B 8x8 encoding); its first commit is #26's.

What

agx_nir_lower_cmat already turns cooperative matrices into per-lane fragments in Apple's simdgroup_float8x8 layout and emits cmat_muladd_agx. Today only M5 uses it: agx_compile.c asserts AGX3C, the packer refuses MATMAD8_AGX3 everywhere else, and hk advertises the extension only on M5. This PR wires up the AGX3B encoding from PR 1:

1. nir/divergence: Handle cmat_muladd_agx. The result is per invocation, like the AMD, Panfrost and NVIDIA variants. Without this case, an assertion-enabled build stops on "Unhandled intrinsic" when a shader containing it reaches divergence analysis. A G17P compute shader compiled through agx_preprocess_nir + agx_compile_shader_nir reaches it, so the new test hits this.
2. asahi/hk: opt-in 8x8 cooperative matrices on AGX3B.
   • cmat_muladd_agx is emitted off AGX3C too. On AGX3B the packer takes MATMAD8 with fp16 or fp32 A and B and an fp32 accumulator, and sets b_input_type.
   • bf16 and the 16x16 form have no AGX3B encoding. They are refused, so such a shader fails to compile rather than being miscompiled.
   • D's cache hint stays clear on AGX3B, as it is in every G16 witness. Apple never sets it in 177 MACs.
   • The register-pair form reaches halves 0..191. That is the whole AGX3B file (AGX3B_NUM_REGS), so the allocator needs no new bound.
   • hk advertises VK_KHR_cooperative_matrix on G17P only with ASAHI_AGX3B_COOPERATIVE_MATRIX=1. G16G is also listed, but it is not admitted. The advertised configurations are 8x8x8 with fp32 A/B and with fp16 A/B, both with fp32 C. KHR_shader_bfloat16 is now explicitly M5-only, and 16x16 is listed only on M5 with DRM_ASAHI_FEATURE_COMPUTE_MATRIX. M5 behaviour does not change.

Why

Cooperative matrices are the biggest single gain we found on M4. All numbers below are from our receipts. They were measured on the M4 Mac mini (J773g, T8132/G16G) with GravityLinux's Apple9 backend and a Gravity kernel, not with this tree. That stack emits this same instruction, fp32 operands, from a lowering with the same 8x8 lane layout as agx_nir_lower_cmat.

Workload	FMA route	Cooperative matrix
MLX 4-bit matmul, bf16 activations, 512 rows (8192 x 3072)	889 GFLOP/s	1631 GFLOP/s
MLX 4-bit matmul, f16 activations, 512 rows	855 GFLOP/s	1396 GFLOP/s
Prefill 512 tok, Llama 3.2 1B 4-bit	307.5 tok/s	475.0 tok/s
Prefill 512 tok, granite 4.0 1B 4-bit	168.8 tok/s	355.1 tok/s
Prefill 512 tok, Qwen2.5 1.5B 4-bit	246.5 tok/s	380.1 tok/s
Prefill 512 tok, Mistral 7B 4-bit	56.9 tok/s	91.2 tok/s

• Vulkan GEMMs with one 8x8 tile per subgroup use integer-valued data, so every sum is exact. They match the CPU bit for bit from 8x8x8 to 1024x1024x1024, including non-square shapes and column-major loads and stores.
• Greedy text is identical with the matrix route on and off for 13 models. There were no GPU faults.

The fp16-operand form was never run on hardware: that stack widened fp16 to fp32 first. Here it rests on Metal output only (PR 1).

Tests run (2026-11-01, macOS, Apple Silicon host, Apple clang 21)

• New AGX3BCooperativeMatrix suite in agx_tests (test/test-agx3b-cooperative-matrix.cpp). Each test builds a compute shader around cmat_muladd_agx and compiles it to binary:
  • Fp32OperandsUseMma8, Fp16OperandsUseMma8: for G17P and G16G, exactly one f32matmad8x8 with the right A/B type byte and D's cache hint clear.
  • FullRegisterFilePacks: 110 loaded values live across the multiply, with spilling. It still packs.
  • Bf16IsRefused: agx_compile_shader_nir returns false.
• agx_tests on this branch: 507 tests, 506 passed, 1 skipped. On the base: 502 tests, 501 passed, 1 skipped.
• Together with PR 3, merged locally: 514 tests, 513 passed, 1 skipped. All 17 compiler, packing and layout suites that build without libdrm also pass under meson test (agx_tests, asahi-agx3-packing, the agx3_*/gen1 suites, asahi_layout_tests, agx2_disasm).
• Not built here: hk, because it needs Linux and libdrm. The hk hunks (hk_physical_device.[ch], hk_shader.c) are compile-checked only by the Linux run below.
• shim-probe/cmat_f32.comp and cmat_f16.comp compile with glslc (--target-env=vulkan1.3). cmat-probe.c passes -Wall -Wextra -Werror -fsyntax-only against the tree's Vulkan headers.

Remaining test plan

• Linux, drm-shim, no GPU (prepared in mini-run/run-drm-shim.sh; probe sources are in shim-probe/). Build the tree with drm-shim and hk, assertions on, then:
  • meson test (all suites), drm-smoke.py, pipelines.py.
  • The coopmat probe: g17p with the variable unset or =0 must not advertise the extension.
  • g17p with =1 must list only 8x8x8 f32/f32 and f16/f16 with f32 C, and create both compute pipelines: fp32 row-major, and fp16 with a column-major B and C.
  • g15s and g13g with =1 must not advertise the extension.
• Hardware (G17P / A18 Pro), opt-in:
  • dEQP-VK.compute.pipeline.cooperative_matrix.* with ASAHI_AGX3B_COOPERATIVE_MATRIX=1.
  • An exact GEMM at 8x8x8 to 1024x1024x1024, fp32 and fp16, row- and column-major.
  • Then the MLX matmul/prefill A/B above.
  • Things to watch: the G17P kernel may need a matrix-state flag the way M5's 16x16 does (DRM_ASAHI_COMPUTE_USES_MATRIX). On M4 with the Gravity kernel the 8x8 form ran without one.
• M4 (G16G) stays refused by admission, so this only becomes testable on omacom's stack once T8132 is admitted.

Risks and limitations

• Off by default. With the variable unset, the only change on AGX3B is that the compiler no longer asserts on a cmat_muladd_agx that cannot be reached there.
• Mixed fp16/fp32 A/B is encodable (PR 1) but not advertised, since agx_nir_lower_cmat only mixes types in 16x16.
• The divergence-analysis fix is in common NIR code. It only adds a case to the existing divergent list.

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The changes since the previous review appear safe to merge.

Summary

Adds opt-in 8x8 cooperative matrices on AGX3B.

• AGX3B shaders can compile 8x8 cooperative matrix operations.
• hk advertises AGX3B cooperative matrices only when opted in.

No new actionable issues were found in the changes since the previous review. No previous review threads were supplied. This review checked source code; it did not run builds or tests.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[AGX3B opt-in] --> B[Advertise 8x8 configurations]
  B --> C[Lower cooperative matrices]
  C --> D[cmat_muladd_agx]
  D --> E[Encode AGX3B instruction]

Loading

_{Reviews (2) · Last reviewed commit: "asahi/hk: opt-in 8x8 cooperative matrice..." · Reviewed by Greptile}

[techmore]

Rebased onto the updated #26 (head c96a7c0), which fixes the AGX3C B-type check Greptile flagged there. This PR's own changes are unchanged. All three branches merged together: 516 of 517 agx_tests pass; the one skip also skips on the base.

🤖 Generated with Claude Code

[techmore]

Aurora downstream draft: aurora-silicon#192