From cbbe1186cb5f7aa964ee03b8aed14a2ef92a6baf Mon Sep 17 00:00:00 2001 From: Sean Dolbec Date: Sun, 1 Nov 2026 19:32:58 -0500 Subject: [PATCH 1/4] agx: encode the AGX3B 8x8 matrix multiply-accumulate from G16 witnesses agx3_matrix_encode_mma8 only packed the 14-byte AGX3C (M5) form. The AGX3B table has the 12-byte f32matmad8x8 that Apple's compiler emits for simdgroup_float8x8 / simdgroup_half8x8 multiply-accumulate on applegpu_g16g (M4), so encode it as well: the same anonymous fields as on AGX3C (2 D, 18 A, 35 B, 52 C, 26/43/60 operand types, 27/29, 44/46, 61/63 cache and discard, 11 D cache, 59 = 1 for a multiply without C), all operands through the register-pair form (r0..r95), and six wait slots instead of eight. A and B may differ in type, as Metal's half/float mixes do, through a new optional b_input_type; AGX3C still takes one input type and refuses a different B. The witnesses are metal-tt -arch applegpu_g16g builds of MSL kernels that load every A, B and C fragment from, and store every D to, a distinct offset, so each MAC's operand registers are known from the surrounding loads and stores rather than from the MAC itself. All 177 MACs found that way (123 fp32, 51 fp16 A and B, three simdgroup_multiply and mixed-type examples) reproduce byte for byte; the test carries the 54 that cover every cache, discard, wait and type combination seen, plus refusals. This adds the encoding only. Nothing selects it yet: the compiler still asserts AGX3C for cmat_muladd_agx and the packer refuses MATMAD8 off AGX3C, so no shader changes. Outside this tree, the fp32-operand form drives a cooperative-matrix lowering on an M4 Mac mini (T8132) with another kernel and Mesa tree, where Vulkan GEMMs match the CPU bit for bit from 8x8x8 to 1024x1024x1024. The fp16-operand forms are from Metal output only: that lowering widens fp16 to fp32 before the instruction. Signed-off-by: Sean Dolbec Co-Authored-By: Claude Opus 5.5 --- src/asahi/compiler/agx3_matrix_core.c | 68 +++++++- src/asahi/compiler/agx3_matrix_core.h | 4 + .../test/test-native-form-regressions.cpp | 157 ++++++++++++++++++ 3 files changed, 228 insertions(+), 1 deletion(-) diff --git a/src/asahi/compiler/agx3_matrix_core.c b/src/asahi/compiler/agx3_matrix_core.c index 7df8f7cfefe0..012c36411e97 100644 --- a/src/asahi/compiler/agx3_matrix_core.c +++ b/src/asahi/compiler/agx3_matrix_core.c @@ -1,12 +1,78 @@ #include "agx3_matrix_core.h" #include "agx3_binder.h" +/* + * The AGX3B (G16/G17P) 8x8x8 multiply-accumulate with an fp32 accumulator, + * the 12-byte f32matmad8x8 form. Field assignments come from Apple's + * applegpu_g16g output for MSL simdgroup_float8x8 and simdgroup_half8x8 + * kernels (metal-tt), whose operand registers are known from the + * surrounding loads and stores; see the AGX3BMma8 witnesses. + * + * Every operand is a register named in 16-bit units, so an even value: D, C + * and an fp32 A or B are register pairs, an fp16 A or B (two halves) is one + * register. Apple names them all through the register-pair form (25, 42, 59 + * = 2), which reaches r0..r95, and gives each operand its own type (26, 43, + * 60). With c_zero (simdgroup_multiply) the C operand is absent (59 = 1, its + * register field zero). The wait mask has six slots here, not eight. + */ +static bool +agx3_matrix_encode_mma8_agx3b(const struct agx3_matrix_mma8_desc *d, + uint8_t out[14], uint8_t *length) +{ + unsigned a_type = d->input_type ? d->input_type : AGX3_MATRIX_INPUT_F16; + unsigned b_type = d->b_input_type ? d->b_input_type : a_type; + if ((a_type != AGX3_MATRIX_INPUT_F16 && a_type != AGX3_MATRIX_INPUT_F32) || + (b_type != AGX3_MATRIX_INPUT_F16 && b_type != AGX3_MATRIX_INPUT_F32) || + (d->dest & 1) || d->dest > 190 || (d->a & 1) || d->a > 190 || + (d->b & 1) || d->b > 190 || d->wait_mask > 0x3f || + (!d->c_zero && ((d->c & 1) || d->c > 190)) || + (d->c_zero && (d->c || d->cache[2] || d->discard[2]))) + return false; + + for (unsigned i = 0; i < 3; ++i) { + if (d->cache[i] && d->discard[i]) + return false; + } + + const struct agx3_field_kv base[] = { + {9, 2}, {10, 2}, {25, 2}, {42, 2}, {60, 2}, + }; + struct agx3_field_kv fields[21]; + unsigned count = 0; + fields[count++] = (struct agx3_field_kv){26, a_type}; + fields[count++] = (struct agx3_field_kv){43, b_type}; + fields[count++] = (struct agx3_field_kv){2, d->dest}; + fields[count++] = (struct agx3_field_kv){11, d->dest_cache}; + fields[count++] = (struct agx3_field_kv){18, d->a}; + fields[count++] = (struct agx3_field_kv){35, d->b}; + fields[count++] = (struct agx3_field_kv){52, d->c}; + fields[count++] = (struct agx3_field_kv){59, d->c_zero ? 1 : 2}; + fields[count++] = (struct agx3_field_kv){27, d->cache[0]}; + fields[count++] = (struct agx3_field_kv){29, d->discard[0]}; + fields[count++] = (struct agx3_field_kv){44, d->cache[1]}; + fields[count++] = (struct agx3_field_kv){46, d->discard[1]}; + fields[count++] = (struct agx3_field_kv){61, d->cache[2]}; + fields[count++] = (struct agx3_field_kv){63, d->discard[2]}; + for (unsigned i = 0; i < 6; ++i) + fields[count++] = (struct agx3_field_kv){69 + i, + (d->wait_mask >> i) & 1}; + + return agx3_bind_compose(AGX_ISA_ENCODING_AGX3B, "f32matmad8x8", 0, base, + sizeof(base) / sizeof(base[0]), fields, count, + out, length); +} + bool agx3_matrix_encode_mma8(enum agx_isa_encoding encoding, const struct agx3_matrix_mma8_desc *d, uint8_t out[14], uint8_t *length) { - if (!d) + if (!d || !out || !length) + return false; + if (encoding == AGX_ISA_ENCODING_AGX3B) + return agx3_matrix_encode_mma8_agx3b(d, out, length); + /* AGX3C takes one input type for A and B. */ + if (d->b_input_type && d->b_input_type != d->input_type) return false; unsigned type = d->input_type ? d->input_type : AGX3_MATRIX_INPUT_F16; unsigned align_mask = type == AGX3_MATRIX_INPUT_F16 ? 3 : 1; diff --git a/src/asahi/compiler/agx3_matrix_core.h b/src/asahi/compiler/agx3_matrix_core.h index e521a200852c..1c44526b6338 100644 --- a/src/asahi/compiler/agx3_matrix_core.h +++ b/src/asahi/compiler/agx3_matrix_core.h @@ -46,8 +46,12 @@ struct agx3_matrix_mma8_desc { bool discard[3]; uint8_t wait_mask; enum agx3_matrix_input_type input_type; + /* B's type when it differs from A's (AGX3B only); 0 means input_type. */ + enum agx3_matrix_input_type b_input_type; }; +/* AGX3C: the 14-byte form. AGX3B (G16/G17P): the 12-byte fp32-accumulator + * form, fp16 or fp32 A and B. */ bool agx3_matrix_encode_mma8(enum agx_isa_encoding encoding, const struct agx3_matrix_mma8_desc *desc, uint8_t out[14], uint8_t *length); diff --git a/src/asahi/compiler/test/test-native-form-regressions.cpp b/src/asahi/compiler/test/test-native-form-regressions.cpp index 5e69bd6b6792..72e1353368a7 100644 --- a/src/asahi/compiler/test/test-native-form-regressions.cpp +++ b/src/asahi/compiler/test/test-native-form-regressions.cpp @@ -5172,6 +5172,163 @@ TEST(NativeForms, M5MmaNativeOperandTypeWitnesses) } } +/* G16 (applegpu_g16g) f32matmad8x8 from metal-tt builds of MSL + * simdgroup_float8x8 / simdgroup_half8x8 kernels. The operand registers come + * from the loads and stores around each MAC, not from the MAC itself; the + * selection covers every cache, discard, wait and operand-type combination + * seen in 177 such MACs, plus simdgroup_multiply (no C) and mixed fp16/fp32 + * operands. Registers are in 16-bit units. */ +TEST(NativeForms, AGX3BMma8MetalWitnesses) +{ + struct Case { agx3_matrix_mma8_desc desc; uint8_t bytes[12]; }; + const Case cases[] = { + {{4,4,8,12,false,false,{false,false,false},{true,true,true},0x20,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x56, 0x02, 0x00, 0x04, 0x08, 0x09, 0xd4, 0x43, 0x24, 0x01}}, + {{22,18,8,22,false,false,{true,true,false},{false,false,true},0x20,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x56, 0x0b, 0x00, 0x92, 0x09, 0x07, 0x54, 0x42, 0x24, 0x01}}, + {{16,18,26,34,false,false,{false,true,false},{true,false,true},0x01,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x12, 0x54, 0x08, 0x00, 0x12, 0x1b, 0x22, 0xd4, 0x42, 0x24, 0x01}}, + {{10,4,8,12,false,false,{true,false,false},{false,true,true},0x02,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x22, 0x54, 0x05, 0x00, 0x84, 0x08, 0x09, 0x54, 0x43, 0x24, 0x01}}, + {{4,4,26,30,false,false,{false,false,false},{true,true,true},0x04,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x42, 0x54, 0x02, 0x00, 0x04, 0x1a, 0x0f, 0xd4, 0x43, 0x24, 0x01}}, + {{28,24,8,28,false,false,{true,true,false},{false,false,true},0x01,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x12, 0x54, 0x0e, 0x00, 0x98, 0x09, 0x0d, 0x54, 0x42, 0x24, 0x01}}, + {{22,24,38,46,false,false,{false,false,false},{true,false,true},0x02,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x22, 0x54, 0x0b, 0x00, 0x18, 0x26, 0x2b, 0xd4, 0x42, 0x24, 0x01}}, + {{16,4,50,18,false,false,{true,false,false},{false,true,true},0x04,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x42, 0x54, 0x08, 0x00, 0x84, 0x32, 0x06, 0x54, 0x43, 0x24, 0x01}}, + {{10,4,8,12,false,false,{true,false,false},{false,true,true},0x08,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x82, 0x54, 0x05, 0x00, 0x84, 0x08, 0x09, 0x54, 0x43, 0x24, 0x01}}, + {{4,4,38,42,false,false,{false,false,false},{true,true,true},0x10,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x55, 0x02, 0x00, 0x04, 0x26, 0x2a, 0xd4, 0x43, 0x24, 0x01}}, + {{34,24,68,34,false,false,{true,true,false},{false,false,true},0x04,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x42, 0x54, 0x11, 0x00, 0x98, 0x45, 0x22, 0x54, 0x42, 0x24, 0x01}}, + {{28,24,8,28,false,false,{true,true,false},{false,false,true},0x00,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x54, 0x0e, 0x00, 0x98, 0x09, 0x0d, 0x54, 0x42, 0x24, 0x01}}, + {{22,24,56,64,false,false,{false,false,false},{true,false,true},0x00,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x54, 0x0b, 0x00, 0x18, 0x38, 0x10, 0xd4, 0x42, 0x24, 0x01}}, + {{16,4,68,18,false,false,{true,false,false},{false,true,true},0x10,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x55, 0x08, 0x00, 0x84, 0x44, 0x06, 0x54, 0x43, 0x24, 0x01}}, + {{10,4,8,12,false,false,{true,false,false},{false,true,true},0x00,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x54, 0x05, 0x00, 0x84, 0x08, 0x09, 0x54, 0x43, 0x24, 0x01}}, + {{4,4,56,60,false,false,{false,false,false},{true,true,true},0x08,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x82, 0x54, 0x02, 0x00, 0x04, 0x38, 0x2d, 0xd4, 0x43, 0x24, 0x01}}, + {{82,78,8,82,false,false,{true,false,false},{false,false,true},0x00,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x54, 0x29, 0x00, 0xce, 0x08, 0x16, 0x54, 0x42, 0x24, 0x01}}, + {{58,54,8,58,false,false,{true,false,false},{false,false,true},0x10,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x55, 0x1d, 0x00, 0xb6, 0x08, 0x2e, 0x54, 0x42, 0x24, 0x01}}, + {{46,30,116,46,false,false,{true,true,false},{false,false,true},0x08,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x82, 0x54, 0x17, 0x00, 0x9e, 0x75, 0x2b, 0x54, 0x42, 0x24, 0x01}}, + {{34,30,8,34,false,false,{true,false,false},{false,false,true},0x04,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x42, 0x54, 0x11, 0x00, 0x9e, 0x08, 0x22, 0x54, 0x42, 0x24, 0x01}}, + {{4,4,98,102,false,false,{false,false,false},{true,true,true},0x01,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x12, 0x54, 0x02, 0x00, 0x04, 0x62, 0x33, 0xd4, 0x43, 0x24, 0x01}}, + {{52,36,110,52,false,false,{true,true,false},{false,false,true},0x10,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x55, 0x1a, 0x00, 0xa4, 0x6f, 0x25, 0x54, 0x42, 0x24, 0x01}}, + {{40,36,8,40,false,false,{true,false,false},{false,false,true},0x08,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x82, 0x54, 0x14, 0x00, 0xa4, 0x08, 0x28, 0x54, 0x42, 0x24, 0x01}}, + {{28,30,122,130,false,false,{false,false,false},{true,false,true},0x01,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x12, 0x54, 0x0e, 0x00, 0x1e, 0x7a, 0x82, 0xd4, 0x42, 0x24, 0x01}}, + {{52,42,158,52,false,false,{true,false,false},{false,false,true},0x02,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x22, 0x54, 0x1a, 0x00, 0xaa, 0x9e, 0x25, 0x54, 0x42, 0x24, 0x01}}, + {{172,144,24,160,false,false,{true,true,false},{false,false,true},0x02,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x22, 0x54, 0x56, 0x00, 0x91, 0x19, 0xa0, 0x54, 0x42, 0x24, 0x01}}, + {{154,144,182,140,false,false,{false,true,false},{true,false,true},0x00,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x54, 0x4d, 0x00, 0x11, 0xb7, 0x89, 0xd4, 0x42, 0x24, 0x01}}, + {{66,66,182,70,false,false,{false,true,false},{true,false,true},0x00,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x54, 0x21, 0x00, 0x42, 0xb7, 0x13, 0xd4, 0x42, 0x24, 0x01}}, + {{40,0,182,40,false,false,{false,false,false},{true,false,true},0x08,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x82, 0x54, 0x14, 0x00, 0x00, 0xb6, 0x28, 0xd4, 0x42, 0x24, 0x01}}, + {{10,4,8,12,false,false,{false,false,false},{false,true,true},0x20,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x56, 0x05, 0x00, 0x04, 0x08, 0x09, 0x54, 0x43, 0x24, 0x01}}, + {{4,4,182,186,false,false,{false,false,false},{true,true,true},0x01,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x12, 0x54, 0x02, 0x00, 0x04, 0xb6, 0xae, 0xd4, 0x43, 0x24, 0x01}}, + {{4,16,4,0,false,false,{false,false,false},{true,true,true},0x20,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x56, 0x02, 0x00, 0x10, 0x04, 0x00, 0xca, 0x43, 0x24, 0x01}}, + {{22,32,30,16,false,false,{true,true,false},{false,false,true},0x20,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x56, 0x0b, 0x00, 0xa0, 0x1f, 0x04, 0x4a, 0x42, 0x24, 0x01}}, + {{16,32,28,12,false,false,{false,true,false},{true,false,true},0x01,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x12, 0x54, 0x08, 0x00, 0x20, 0x1d, 0x09, 0xca, 0x42, 0x24, 0x01}}, + {{10,26,30,8,false,false,{true,false,false},{false,true,true},0x02,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x22, 0x54, 0x05, 0x00, 0x9a, 0x1e, 0x08, 0x4a, 0x43, 0x24, 0x01}}, + {{4,26,28,4,false,false,{false,false,false},{true,true,true},0x04,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x42, 0x54, 0x02, 0x00, 0x1a, 0x1c, 0x01, 0xca, 0x43, 0x24, 0x01}}, + {{28,46,42,20,false,false,{true,true,false},{false,false,true},0x01,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x12, 0x54, 0x0e, 0x00, 0xae, 0x2b, 0x05, 0x4a, 0x42, 0x24, 0x01}}, + {{22,46,44,16,false,false,{false,true,false},{true,false,true},0x02,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x22, 0x54, 0x0b, 0x00, 0x2e, 0x2d, 0x04, 0xca, 0x42, 0x24, 0x01}}, + {{16,38,40,12,false,false,{true,false,false},{false,true,true},0x04,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x42, 0x54, 0x08, 0x00, 0xa6, 0x28, 0x09, 0x4a, 0x43, 0x24, 0x01}}, + {{10,38,42,8,false,false,{true,false,false},{false,true,true},0x08,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x82, 0x54, 0x05, 0x00, 0xa6, 0x2a, 0x08, 0x4a, 0x43, 0x24, 0x01}}, + {{4,38,44,4,false,false,{false,false,false},{true,true,true},0x10,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x55, 0x02, 0x00, 0x26, 0x2c, 0x01, 0xca, 0x43, 0x24, 0x01}}, + {{34,62,58,24,false,false,{true,true,false},{false,false,true},0x04,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x42, 0x54, 0x11, 0x00, 0xbe, 0x3b, 0x0c, 0x4a, 0x42, 0x24, 0x01}}, + {{28,62,60,20,false,false,{true,true,false},{false,false,true},0x00,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x54, 0x0e, 0x00, 0xbe, 0x3d, 0x05, 0x4a, 0x42, 0x24, 0x01}}, + {{22,62,56,16,false,false,{false,true,false},{true,false,true},0x00,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x54, 0x0b, 0x00, 0x3e, 0x39, 0x04, 0xca, 0x42, 0x24, 0x01}}, + {{16,50,58,12,false,false,{true,false,false},{false,true,true},0x10,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x55, 0x08, 0x00, 0xb2, 0x3a, 0x09, 0x4a, 0x43, 0x24, 0x01}}, + {{10,50,60,8,false,false,{true,false,false},{false,true,true},0x00,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x54, 0x05, 0x00, 0xb2, 0x3c, 0x08, 0x4a, 0x43, 0x24, 0x01}}, + {{4,50,56,4,false,false,{false,false,false},{true,true,true},0x08,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x82, 0x54, 0x02, 0x00, 0x32, 0x38, 0x01, 0xca, 0x43, 0x24, 0x01}}, + {{58,106,102,40,false,false,{true,true,false},{false,false,true},0x10,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x55, 0x1d, 0x00, 0xea, 0x67, 0x28, 0x4a, 0x42, 0x24, 0x01}}, + {{46,108,86,32,false,false,{true,true,false},{false,false,true},0x08,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x82, 0x54, 0x17, 0x00, 0xec, 0x57, 0x20, 0x4a, 0x42, 0x24, 0x01}}, + {{4,74,80,4,false,false,{false,false,false},{true,true,true},0x01,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x12, 0x54, 0x02, 0x00, 0x4a, 0x50, 0x01, 0xca, 0x43, 0x24, 0x01}}, + {{64,64,80,44,false,false,{false,false,false},{true,false,true},0x00,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x54, 0x20, 0x00, 0x40, 0x50, 0x29, 0xca, 0x42, 0x24, 0x01}}, + {{4,4,8,0,true,false,{false,false,false},{true,true,false},0x20,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x56, 0x02, 0x00, 0x04, 0x08, 0x00, 0xd4, 0x41, 0x24, 0x00}}, + {{4,14,0,4,false,false,{false,false,false},{true,true,true},0x20,AGX3_MATRIX_INPUT_F16,AGX3_MATRIX_INPUT_F32}, + {0xcf, 0x02, 0x56, 0x02, 0x00, 0x0e, 0x00, 0x01, 0xd2, 0x43, 0x24, 0x01}}, + {{4,4,12,8,false,false,{false,false,false},{true,true,true},0x20,AGX3_MATRIX_INPUT_F32,AGX3_MATRIX_INPUT_F16}, + {0xcf, 0x02, 0x56, 0x02, 0x00, 0x04, 0x0c, 0x08, 0xcc, 0x43, 0x24, 0x01}}, + }; + for (const auto &c : cases) { + uint8_t bytes[14], length; + ASSERT_TRUE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3B, &c.desc, + bytes, &length)); + ASSERT_EQ(length, 12u); + EXPECT_EQ(std::vector(bytes, bytes + length), + std::vector(c.bytes, c.bytes + 12)); + } + + /* Fail closed outside what the witnesses cover. */ + agx3_matrix_mma8_desc d = cases[0].desc; + uint8_t bytes[14], length; + d.wait_mask = 0x40; /* six wait slots on AGX3B */ + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3B, &d, bytes, &length)); + d = cases[0].desc; + d.input_type = d.b_input_type = AGX3_MATRIX_INPUT_BF16; + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3B, &d, bytes, &length)); + d = cases[0].desc; + d.dest = 192; + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3B, &d, bytes, &length)); + d = cases[0].desc; + d.a = 5; + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3B, &d, bytes, &length)); + d = cases[0].desc; + d.c_zero = true; + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3B, &d, bytes, &length)); + + /* AGX3C keeps a single input type. */ + d = {}; + d.dest = 24; + d.a = 2; + d.b = 20; + d.c_zero = true; + d.input_type = AGX3_MATRIX_INPUT_F16; + d.b_input_type = AGX3_MATRIX_INPUT_F32; + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3C, &d, bytes, &length)); +} + TEST(NativeForms, M5GemmMixedOperandExecutionWitnesses) { agx3_matrix_gemm16_desc d = {}; From aebf435c6b09a74f1b87b16783a76beabfac7fd0 Mon Sep 17 00:00:00 2001 From: Sean Dolbec Date: Sun, 1 Nov 2026 22:27:28 -0500 Subject: [PATCH 2/4] agx: compare the AGX3C B type against A's resolved type agx3_matrix_encode_mma8 compared b_input_type with the raw input_type, so on AGX3C a descriptor that left input_type at its fp16 default of 0 and named B as fp16 was refused, although both inputs are fp16. Default A's type first and compare B against that. Shader callers leave B unset and are unaffected. AGX3BMma8MetalWitnesses now checks that naming B's type, equal to A's (including A left at 0), packs the same AGX3C bytes as leaving it unset, and that the default fp16 A with an fp32 B is still refused. Signed-off-by: Sean Dolbec Co-Authored-By: Claude Opus 5.5 --- src/asahi/compiler/agx3_matrix_core.c | 7 ++++--- .../test/test-native-form-regressions.cpp | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/asahi/compiler/agx3_matrix_core.c b/src/asahi/compiler/agx3_matrix_core.c index 012c36411e97..002dbf220e6d 100644 --- a/src/asahi/compiler/agx3_matrix_core.c +++ b/src/asahi/compiler/agx3_matrix_core.c @@ -71,10 +71,11 @@ agx3_matrix_encode_mma8(enum agx_isa_encoding encoding, return false; if (encoding == AGX_ISA_ENCODING_AGX3B) return agx3_matrix_encode_mma8_agx3b(d, out, length); - /* AGX3C takes one input type for A and B. */ - if (d->b_input_type && d->b_input_type != d->input_type) - return false; + /* AGX3C takes one input type for A and B. Compare B against A's type + * after A's 0 has defaulted to fp16. */ unsigned type = d->input_type ? d->input_type : AGX3_MATRIX_INPUT_F16; + if (d->b_input_type && d->b_input_type != type) + return false; unsigned align_mask = type == AGX3_MATRIX_INPUT_F16 ? 3 : 1; unsigned source_max = type == AGX3_MATRIX_INPUT_F32 ? 252 : 254; if (encoding != AGX_ISA_ENCODING_AGX3C || !out || !length || diff --git a/src/asahi/compiler/test/test-native-form-regressions.cpp b/src/asahi/compiler/test/test-native-form-regressions.cpp index 72e1353368a7..857afde9212f 100644 --- a/src/asahi/compiler/test/test-native-form-regressions.cpp +++ b/src/asahi/compiler/test/test-native-form-regressions.cpp @@ -5327,6 +5327,27 @@ TEST(NativeForms, AGX3BMma8MetalWitnesses) d.input_type = AGX3_MATRIX_INPUT_F16; d.b_input_type = AGX3_MATRIX_INPUT_F32; EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3C, &d, bytes, &length)); + + /* ...but naming B's type explicitly, equal to A's, is the same + * instruction, including when A's type is the default 0 (fp16). */ + uint8_t plain[14], plain_length; + for (agx3_matrix_input_type a : {(agx3_matrix_input_type)0, + AGX3_MATRIX_INPUT_F16, + AGX3_MATRIX_INPUT_F32}) { + SCOPED_TRACE(a); + d.input_type = a; + d.b_input_type = (agx3_matrix_input_type)0; + ASSERT_TRUE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3C, &d, plain, + &plain_length)); + d.b_input_type = a ? a : AGX3_MATRIX_INPUT_F16; + ASSERT_TRUE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3C, &d, bytes, + &length)); + EXPECT_EQ(std::vector(bytes, bytes + length), + std::vector(plain, plain + plain_length)); + } + d.input_type = (agx3_matrix_input_type)0; + d.b_input_type = AGX3_MATRIX_INPUT_F32; + EXPECT_FALSE(agx3_matrix_encode_mma8(AGX_ISA_ENCODING_AGX3C, &d, bytes, &length)); } TEST(NativeForms, M5GemmMixedOperandExecutionWitnesses) From 84a8fcb08bb343e17c7c400fc619419238dc7537 Mon Sep 17 00:00:00 2001 From: Sean Dolbec Date: Sun, 1 Nov 2026 20:54:10 -0500 Subject: [PATCH 3/4] nir/divergence: Handle cmat_muladd_agx Like the other drivers' cooperative-matrix multiply-adds, its result is per invocation. Without a case, a shader that reaches divergence analysis with cmat_muladd_agx still in it stops on the "Unhandled intrinsic" assertion in builds with assertions enabled (release builds already treat unknown intrinsics as divergent). An AGX3B compute shader compiled through agx_preprocess_nir and agx_compile_shader_nir hits it. Signed-off-by: Sean Dolbec Co-Authored-By: Claude Opus 5.5 --- src/compiler/nir/nir_divergence_analysis.c | 1 + 1 file changed, 1 insertion(+) diff --git a/src/compiler/nir/nir_divergence_analysis.c b/src/compiler/nir/nir_divergence_analysis.c index e05bc59ff977..52e70df4d8cd 100644 --- a/src/compiler/nir/nir_divergence_analysis.c +++ b/src/compiler/nir/nir_divergence_analysis.c @@ -1094,6 +1094,7 @@ visit_intrinsic(nir_intrinsic_instr *instr, struct divergence_state *state) case nir_intrinsic_load_ray_triangle_vertex_positions: case nir_intrinsic_cmat_extract: case nir_intrinsic_cmat_muladd_amd: + case nir_intrinsic_cmat_muladd_agx: case nir_intrinsic_cmat_muladd_pan: case nir_intrinsic_dpas_intel: case nir_intrinsic_convert_cmat_intel: From c96a7c0de91e613dca1e2ebd72277ea3cc7c9b20 Mon Sep 17 00:00:00 2001 From: Sean Dolbec Date: Sun, 1 Nov 2026 20:54:35 -0500 Subject: [PATCH 4/4] asahi/hk: opt-in 8x8 cooperative matrices on AGX3B agx_nir_lower_cmat already produces the per-lane fragments of Apple's simdgroup_float8x8 layout, and the previous commit encodes the AGX3B f32matmad8x8. Connect them on AGX3B: - cmat_muladd_agx is emitted off AGX3C as well. The packer takes MATMAD8 on AGX3B with fp16 or fp32 A and B and an fp32 accumulator; bf16 and the 16x16 form have no AGX3B encoding and are refused, so such a shader fails to compile instead of being miscompiled. D's cache hint stays clear on AGX3B, as in every G16 witness. The register-pair form reaches halves 0..191, the whole AGX3B register file, so the allocator needs no new bound. - hk advertises VK_KHR_cooperative_matrix on G17P (and G16G, which is not admitted) only with ASAHI_AGX3B_COOPERATIVE_MATRIX=1: 8x8x8 with fp32 or fp16 A and B and fp32 C. No bf16 and no 16x16 there; M5 is unchanged. The fragment layout is the one this lowering already uses on M5. Another Mesa tree's lowering with the same 8x8 layout, emitting the fp32-operand form of this instruction on an M4 (T8132, GravityLinux's Apple9 backend and kernel), matches the CPU bit for bit for Vulkan GEMMs from 8x8x8 to 1024x1024x1024, row- and column-major, and took MLX 4-bit matmul from 0.8-0.9 to 1.3-1.6 TFLOP/s and 512-token prefill of 4-bit models by 1.1-2.1x, with identical greedy text. That lowering widened fp16 to fp32 first, so the fp16-operand form here has only Metal output behind it. Nothing has run on G17P, hence the opt-in. AGX3BCooperativeMatrix in agx_tests compiles a compute shader for G17P and G16G and finds exactly one f32matmad8x8 with fp32 and with fp16 operands, packs it with the register file full, and refuses bf16. Signed-off-by: Sean Dolbec Co-Authored-By: Claude Opus 5.5 --- src/asahi/compiler/agx3_pack_binary.c | 11 +- src/asahi/compiler/agx_compile.c | 4 +- src/asahi/compiler/meson.build | 1 + .../test/test-agx3b-cooperative-matrix.cpp | 143 ++++++++++++++++++ src/asahi/vulkan/hk_physical_device.c | 11 +- src/asahi/vulkan/hk_physical_device.h | 13 ++ src/asahi/vulkan/hk_shader.c | 3 +- 7 files changed, 179 insertions(+), 7 deletions(-) create mode 100644 src/asahi/compiler/test/test-agx3b-cooperative-matrix.cpp diff --git a/src/asahi/compiler/agx3_pack_binary.c b/src/asahi/compiler/agx3_pack_binary.c index d53ab77bd109..71969085047a 100644 --- a/src/asahi/compiler/agx3_pack_binary.c +++ b/src/asahi/compiler/agx3_pack_binary.c @@ -10219,7 +10219,11 @@ agx3_pack_binary(agx_context *ctx, struct util_dynarray *emission) if (I->op == AGX_OPCODE_MATMAD8_AGX3) { if ((I->matrix_type & 15) != ((I->matrix_type >> 4) & 15)) AGX3_REFUSE(I); - if (enc != AGX_ISA_ENCODING_AGX3C || ctx->is_preamble || + /* AGX3B: the 12-byte f32matmad8x8 (fp16 or fp32 A and B, six + * wait slots); agx3_matrix_encode_mma8 refuses other types. */ + if ((enc != AGX_ISA_ENCODING_AGX3C && + enc != AGX_ISA_ENCODING_AGX3B) || + ctx->is_preamble || ctx->stage != MESA_SHADER_COMPUTE || I->nr_dests != 1 || I->nr_srcs != 3 || I->saturate || wait > 0xff || !agx3_src_is_register(I->dest[0]) || I->dest[0].memory || @@ -10232,9 +10236,12 @@ agx3_pack_binary(agx_context *ctx, struct util_dynarray *emission) .a = I->src[0].value, .b = I->src[1].value, .c = I->src[2].value, - .dest_cache = I->dest[0].cache, + /* None of the G16 witnesses sets D's cache hint; leave it + * clear there. */ + .dest_cache = enc == AGX_ISA_ENCODING_AGX3C && I->dest[0].cache, .wait_mask = wait, .input_type = I->matrix_type & 15, + .b_input_type = (I->matrix_type >> 4) & 15, }; for (unsigned s = 0; s < 3; ++s) { const agx_index src = I->src[s]; diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index 5443f4e509f0..e6f2055e2601 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -2796,7 +2796,9 @@ agx_emit_intrinsic(agx_builder *b, nir_intrinsic_instr *instr) ? AGX3_MATRIX_INPUT_BF16 : nir_intrinsic_src_type2(instr) == nir_type_float32 ? AGX3_MATRIX_INPUT_F32 : AGX3_MATRIX_INPUT_F16; type |= type_b << 4; - assert(b->shader->key->dev.encoding == AGX_ISA_ENCODING_AGX3C && + /* AGX3C has the 8x8 and 16x16 forms. AGX3B packs only the 8x8 form, + * with fp16 or fp32 A and B; the packer refuses anything else. */ + assert(b->shader->key->dev.encoding != AGX_ISA_ENCODING_AGX2 && stage == MESA_SHADER_COMPUTE && !b->shader->is_preamble); if (instr->def.num_components == 8) { agx_matmad16_agx3_to(b, dst, agx_src_index(&instr->src[0]), diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index 2b766a4f14d3..c6813d2b602d 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -215,6 +215,7 @@ if with_tests 'test/test-negated-minifloat.cpp', 'test/test-dual-absolute-select.cpp', 'test/test-native-form-regressions.cpp', + 'test/test-agx3b-cooperative-matrix.cpp', 'test/test-inline-mrt.cpp', 'test/test-mrt-terminal-tail.cpp', 'test/test-agx3-cprog-spill-certificate.cpp', diff --git a/src/asahi/compiler/test/test-agx3b-cooperative-matrix.cpp b/src/asahi/compiler/test/test-agx3b-cooperative-matrix.cpp new file mode 100644 index 000000000000..ae42cbdf68fe --- /dev/null +++ b/src/asahi/compiler/test/test-agx3b-cooperative-matrix.cpp @@ -0,0 +1,143 @@ +/* Copyright 2026 Sean Dolbec; SPDX-License-Identifier: MIT */ +#include +#include +#include "agx_compile.h" +#include "agx_test.h" +#include "compiler/nir/nir_builder.h" +extern "C" { +#include "agx_nir.h" +} + +/* + * cmat_muladd_agx on AGX3B compiles to the 12-byte f32matmad8x8 (see + * NativeForms.AGX3BMma8MetalWitnesses for the encoding): fp32 or fp16 A and + * B, an fp32 accumulator. bf16 has no AGX3B form and is refused rather than + * miscompiled. The register-pair form reaches halves 0..191, which is the + * whole AGX3B register file (AGX3B_NUM_REGS), so full pressure needs no extra + * register-allocation constraint. + */ + +/* Constant bits of every f32matmad8x8 witness, and its A/B type byte (bit 7 + * of byte 8 is A's discard hint). */ +static unsigned +count_mma8(const std::vector &code, uint8_t types) +{ + unsigned n = 0; + for (size_t i = 0; i + 12 <= code.size(); ++i) { + const uint8_t *c = &code[i]; + if (c[0] == 0xcf && (c[2] & 0xfc) == 0x54 && c[4] == 0x00 && + (c[9] & 0xfc) == 0x40 && c[10] == 0x24 && (c[11] & 0xfe) == 0 && + (c[8] & 0x7f) == types) + ++n; + } + return n; +} + +static nir_def * +load_global(nir_builder *b, nir_def *addr) +{ + nir_def *def = nir_load_global(b, 1, 32, addr); + nir_intrinsic_set_access( + nir_def_as_intrinsic(def), + (gl_access_qualifier)(ACCESS_NON_WRITEABLE | ACCESS_CAN_REORDER)); + nir_intrinsic_set_align(nir_def_as_intrinsic(def), 4, 0); + return def; +} + +/* live > 0 keeps that many loaded values live across the multiply so the + * register file fills past half 192. */ +static bool +compile_mma8(agx_chip chip, nir_alu_type type, unsigned live, + std::vector *code) +{ + nir_builder b = nir_builder_init_simple_shader( + MESA_SHADER_COMPUTE, &agx_nir_options, "AGX3B mma8"); + const unsigned bits = nir_alu_type_get_type_size(type); + nir_def *a = agx_test_load_preamble(&b, 2, bits, 4); + nir_def *bm = agx_test_load_preamble(&b, 2, bits, 8); + nir_def *c = agx_test_load_preamble(&b, 2, 32, 12); + + std::vector values; + if (live) { + nir_def *base = agx_test_load_preamble(&b, 1, 64, 16); + nir_def *sum = nir_imm_float(&b, 0.0f); + for (unsigned i = 0; i < live; ++i) { + values.push_back(load_global(&b, nir_iadd_imm(&b, base, 4 * i))); + sum = nir_fadd(&b, sum, values.back()); + } + c = nir_vec2(&b, nir_fadd(&b, nir_channel(&b, c, 0), sum), + nir_channel(&b, c, 1)); + } + + nir_def *d = nir_cmat_muladd_agx(&b, a, bm, c); + nir_intrinsic_set_src_type(nir_def_as_intrinsic(d), type); + nir_intrinsic_set_src_type2(nir_def_as_intrinsic(d), type); + + nir_def *x = nir_channel(&b, d, 0), *y = nir_channel(&b, d, 1); + for (unsigned i = 0; i < values.size(); ++i) + y = nir_ffma(&b, values[i], x, y); + + nir_intrinsic_instr *store = nir_image_store( + &b, nir_imm_intN_t(&b, 1, 16), nir_imm_ivec4(&b, 0, 0, 0, 0), + nir_imm_intN_t(&b, 0, 16), nir_vec4(&b, x, y, x, y), + nir_imm_intN_t(&b, 0, 16)); + nir_intrinsic_set_image_dim(store, GLSL_SAMPLER_DIM_2D); + nir_intrinsic_set_format(store, PIPE_FORMAT_R32G32B32A32_FLOAT); + nir_intrinsic_set_src_type(store, nir_type_float32); + + agx_shader_key key = {}; + key.dev.encoding = AGX_ISA_ENCODING_AGX3B; + key.dev.chip = chip; + key.no_preamble = true; + key.reserved_preamble = 24; + key.has_scratch = live != 0; + agx_preprocess_nir(b.shader); + agx_shader_part part = {}; + const bool ok = agx_compile_shader_nir(b.shader, &key, &part); + if (ok) { + const uint8_t *bytes = static_cast(part.binary); + code->assign(bytes, bytes + part.info.binary_size); + } + free(part.binary); + ralloc_free(b.shader); + return ok; +} + +TEST(AGX3BCooperativeMatrix, Fp32OperandsUseMma8) +{ + for (agx_chip chip : {AGX_CHIP_G17P, AGX_CHIP_G16G}) { + SCOPED_TRACE(chip); + std::vector code; + ASSERT_TRUE(compile_mma8(chip, nir_type_float32, 0, &code)); + EXPECT_EQ(count_mma8(code, 0x54), 1u); + } +} + +TEST(AGX3BCooperativeMatrix, Fp16OperandsUseMma8) +{ + for (agx_chip chip : {AGX_CHIP_G17P, AGX_CHIP_G16G}) { + SCOPED_TRACE(chip); + std::vector code; + ASSERT_TRUE(compile_mma8(chip, nir_type_float16, 0, &code)); + EXPECT_EQ(count_mma8(code, 0x4a), 1u); + } +} + +/* With the register file full across the multiply (spilling), the shader + * still packs. */ +TEST(AGX3BCooperativeMatrix, FullRegisterFilePacks) +{ + for (nir_alu_type type : {nir_type_float32, nir_type_float16}) { + SCOPED_TRACE(type); + std::vector code; + ASSERT_TRUE(compile_mma8(AGX_CHIP_G17P, type, 110, &code)); + EXPECT_EQ(count_mma8(code, type == nir_type_float32 ? 0x54 : 0x4a), 1u); + } +} + +/* bf16 operands have no AGX3B encoding: refuse the shader. */ +TEST(AGX3BCooperativeMatrix, Bf16IsRefused) +{ + std::vector code; + EXPECT_FALSE(compile_mma8(AGX_CHIP_G17P, nir_type_uint16, 0, &code)); +} diff --git a/src/asahi/vulkan/hk_physical_device.c b/src/asahi/vulkan/hk_physical_device.c index ac48a0805b15..6bb5cdfa5318 100644 --- a/src/asahi/vulkan/hk_physical_device.c +++ b/src/asahi/vulkan/hk_physical_device.c @@ -1951,10 +1951,13 @@ hk_create_drm_physical_device(struct vk_instance *_instance, supported_extensions.KHR_pipeline_library = native_rt; } supported_extensions.KHR_cooperative_matrix = - agx_chip_g17_m5(pdev->dev.chip) && - debug_get_bool_option("ASAHI_M5_COOPERATIVE_MATRIX", true); + (agx_chip_g17_m5(pdev->dev.chip) && + debug_get_bool_option("ASAHI_M5_COOPERATIVE_MATRIX", true)) || + hk_agx3b_cooperative_matrix(&pdev->dev); + /* AGX3B has no bfloat16 matrix form. */ supported_extensions.KHR_shader_bfloat16 = + agx_chip_g17_m5(pdev->dev.chip) && supported_extensions.KHR_cooperative_matrix && debug_get_bool_option("ASAHI_M5_BFLOAT16", true); @@ -2393,7 +2396,9 @@ hk_GetPhysicalDeviceCooperativeMatrixPropertiesKHR( .scope = VK_SCOPE_SUBGROUP_KHR, }; } - unsigned max_size = (pdev->dev.params.features & DRM_ASAHI_FEATURE_COMPUTE_MATRIX) ? 16 : 8; + /* 16x16 is the M5 kernel-gated form; AGX3B has only 8x8x8. */ + unsigned max_size = agx_chip_g17_m5(pdev->dev.chip) && + (pdev->dev.params.features & DRM_ASAHI_FEATURE_COMPUTE_MATRIX) ? 16 : 8; for (unsigned size = max_size; size >= 8; size /= 2) { const VkComponentTypeKHR types[] = {VK_COMPONENT_TYPE_FLOAT16_KHR, VK_COMPONENT_TYPE_BFLOAT16_KHR}; diff --git a/src/asahi/vulkan/hk_physical_device.h b/src/asahi/vulkan/hk_physical_device.h index daad7d945ba0..6bf134af5ac5 100644 --- a/src/asahi/vulkan/hk_physical_device.h +++ b/src/asahi/vulkan/hk_physical_device.h @@ -9,6 +9,7 @@ #include "asahi/lib/agx_device.h" #include +#include "util/u_debug.h" #include "hk_private.h" #include "vk_physical_device.h" #include "vk_sync.h" @@ -103,6 +104,18 @@ hk_physical_device_instance(struct hk_physical_device *pdev) return (struct hk_instance *)pdev->vk.instance; } +/* VK_KHR_cooperative_matrix through the AGX3B f32matmad8x8: 8x8x8 with fp16 + * or fp32 A and B and an fp32 accumulator. The encoding comes from G16G + * (M4) Metal output, and G17P shares the AGX3B table, but it has not run on + * an AGX3B GPU with this driver, so it stays opt-in + * (ASAHI_AGX3B_COOPERATIVE_MATRIX=1) until it has. */ +static inline bool +hk_agx3b_cooperative_matrix(const struct agx_device *dev) +{ + return (dev->chip == AGX_CHIP_G17P || dev->chip == AGX_CHIP_G16G) && + debug_get_bool_option("ASAHI_AGX3B_COOPERATIVE_MATRIX", false); +} + VkResult hk_create_drm_physical_device(struct vk_instance *vk_instance, struct _drmDevice *drm_device, struct vk_physical_device **pdev_out); diff --git a/src/asahi/vulkan/hk_shader.c b/src/asahi/vulkan/hk_shader.c index 77d3dbdf13bc..800bf2307e88 100644 --- a/src/asahi/vulkan/hk_shader.c +++ b/src/asahi/vulkan/hk_shader.c @@ -2994,7 +2994,8 @@ hk_lower_nir(struct hk_device *dev, nir_shader *nir, NIR_PASS(_, nir, nir_lower_explicit_io, nir_var_mem_push_const, nir_address_format_32bit_offset); - if (agx_chip_g17_m5(dev->dev.chip)) + if (agx_chip_g17_m5(dev->dev.chip) || + hk_agx3b_cooperative_matrix(&dev->dev)) NIR_PASS(_, nir, agx_nir_lower_cmat); /* Turn cache flushes into image coherency bits while we still have derefs */