asahi/hk: opt-in 8x8 cooperative matrices on AGX3B

src/asahi/compiler/agx3_matrix_core.c

@@ -1,14 +1,81 @@
 #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. 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 ||

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);

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];

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]),

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',

src/asahi/compiler/test/test-agx3b-cooperative-matrix.cpp

@@ -0,0 +1,143 @@
+/* Copyright 2026 Sean Dolbec; SPDX-License-Identifier: MIT */
+#include <gtest/gtest.h>
+#include <vector>
+#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<uint8_t> &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<uint8_t> *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<nir_def *> 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<const uint8_t *>(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<uint8_t> 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<uint8_t> 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<uint8_t> 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<uint8_t> code;
+   EXPECT_FALSE(compile_mma8(AGX_CHIP_G17P, nir_type_uint16, 0, &code));
+}