diff --git a/src/asahi/compiler/agx3_matrix_core.c b/src/asahi/compiler/agx3_matrix_core.c index 7df8f7cfefe0..002dbf220e6d 100644 --- a/src/asahi/compiler/agx3_matrix_core.c +++ b/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 || 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/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/compiler/test/test-native-form-regressions.cpp b/src/asahi/compiler/test/test-native-form-regressions.cpp index 5e69bd6b6792..857afde9212f 100644 --- a/src/asahi/compiler/test/test-native-form-regressions.cpp +++ b/src/asahi/compiler/test/test-native-form-regressions.cpp @@ -5172,6 +5172,184 @@ 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)); + + /* ...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) { agx3_matrix_gemm16_desc d = {}; 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 */ 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: