From a8b0011539c2d97adf5bde1c3eb20b283f869750 Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Wed, 11 Nov 2026 20:28:24 -0600 Subject: [PATCH 1/2] agx: lower cooperative matrix multiply-add with a half accumulator The flash attention shader of llama.cpp accumulates the P times V product in a half precision matrix, so its cmat_muladd has half A and B and a half accumulator. lower_matrix_instr asserts that the accumulator is float or bfloat, so compiling that shader aborts on a build with assertions. The lowering already widens the accumulator to 32 bits for the hardware multiply-add and converts the result back to the result type, so a half accumulator needs no new code: allow it in the assertion. Add a unit test that lowers a 16x16x16 half multiply-add with a float and with a half accumulator and checks that no cmat_muladd is left and that the hardware multiply-add takes a 32-bit accumulator. The half case aborts in the assertion without the change. Signed-off-by: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> --- src/asahi/compiler/agx_nir_lower_cmat.c | 3 +- src/asahi/compiler/meson.build | 1 + .../compiler/test/test-cmat-accumulator.cpp | 86 +++++++++++++++++++ 3 files changed, 89 insertions(+), 1 deletion(-) create mode 100644 src/asahi/compiler/test/test-cmat-accumulator.cpp diff --git a/src/asahi/compiler/agx_nir_lower_cmat.c b/src/asahi/compiler/agx_nir_lower_cmat.c index eb7e0aabebe5..1645737fee5e 100644 --- a/src/asahi/compiler/agx_nir_lower_cmat.c +++ b/src/asahi/compiler/agx_nir_lower_cmat.c @@ -242,7 +242,8 @@ lower_matrix_instr(nir_builder *b, nir_instr *instr, struct hash_table *types) (input_b == input || (matrix_desc(intr->src[1]).rows == 16 && ((input == GLSL_TYPE_FLOAT16 && input_b == GLSL_TYPE_BFLOAT16) || (input == GLSL_TYPE_BFLOAT16 && input_b == GLSL_TYPE_FLOAT16)))) && - (accum == GLSL_TYPE_FLOAT || accum == GLSL_TYPE_BFLOAT16) && result == accum); + (accum == GLSL_TYPE_FLOAT || accum == GLSL_TYPE_FLOAT16 || + accum == GLSL_TYPE_BFLOAT16) && result == accum); nir_def *a = load_matrix(b, intr->src[1]); nir_def *bb = load_matrix(b, intr->src[2]); nir_def *c = convert_matrix_value(b, load_matrix(b, intr->src[3]), diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index 2b766a4f14d3..f0a4a0d0f25d 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -185,6 +185,7 @@ if with_tests 'test/test-compact-constants.cpp', 'test/test-minifloat.cpp', 'test/test-optimizer.cpp', + 'test/test-cmat-accumulator.cpp', 'test/test-pressure-cse.cpp', 'test/test-rte-fma.cpp', 'test/test-border-index.cpp', diff --git a/src/asahi/compiler/test/test-cmat-accumulator.cpp b/src/asahi/compiler/test/test-cmat-accumulator.cpp new file mode 100644 index 000000000000..54a1b11634a0 --- /dev/null +++ b/src/asahi/compiler/test/test-cmat-accumulator.cpp @@ -0,0 +1,86 @@ +/* + * SPDX-License-Identifier: MIT + */ + +#include "agx_compile.h" +extern "C" { +#include "agx_nir.h" +} +#include "compiler/glsl_types.h" +#include "nir.h" +#include "nir_builder.h" +#include + +static const glsl_type * +cmat_type(enum glsl_base_type element, enum glsl_cmat_use use) +{ + struct glsl_cmat_description desc = { + .rows = 16, + .cols = 16, + .element_type = (uint16_t)element, + .scope = SCOPE_SUBGROUP, + .use = (uint8_t)use, + }; + return glsl_cmat_type(&desc); +} + +static nir_deref_instr * +cmat_var(nir_builder *b, enum glsl_base_type element, enum glsl_cmat_use use, + const char *name) +{ + return nir_build_deref_var( + b, nir_local_variable_create(b->impl, cmat_type(element, use), name)); +} + +/* A 16x16x16 multiply-add of half inputs with the given accumulator type, as + * the flash attention shaders of llama.cpp write it for the PV product. + */ +static unsigned +lowered_muladd_count(enum glsl_base_type accum) +{ + glsl_type_singleton_init_or_ref(); + nir_builder b = nir_builder_init_simple_shader( + MESA_SHADER_COMPUTE, &agx_nir_options, "cmat"); + b.shader->info.cs.has_cooperative_matrix = true; + + nir_deref_instr *a = cmat_var(&b, GLSL_TYPE_FLOAT16, GLSL_CMAT_USE_A, "a"); + nir_deref_instr *bb = cmat_var(&b, GLSL_TYPE_FLOAT16, GLSL_CMAT_USE_B, "b"); + nir_deref_instr *c = cmat_var(&b, accum, GLSL_CMAT_USE_ACCUMULATOR, "c"); + nir_deref_instr *d = cmat_var(&b, accum, GLSL_CMAT_USE_ACCUMULATOR, "d"); + + nir_cmat_construct(&b, &a->def, nir_imm_float16(&b, 1.0f)); + nir_cmat_construct(&b, &bb->def, nir_imm_float16(&b, 1.0f)); + nir_cmat_construct(&b, &c->def, + accum == GLSL_TYPE_FLOAT16 ? nir_imm_float16(&b, 0.0f) + : nir_imm_float(&b, 0.0f)); + nir_cmat_muladd(&b, &d->def, &a->def, &bb->def, &c->def); + + agx_nir_lower_cmat(b.shader); + + unsigned left = 0, agx = 0; + nir_foreach_function_impl(impl, b.shader) { + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_intrinsic) + continue; + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + left += intr->intrinsic == nir_intrinsic_cmat_muladd; + /* The hardware accumulates in 32 bits whatever the matrix type. */ + agx += intr->intrinsic == nir_intrinsic_cmat_muladd_agx && + intr->src[2].ssa->bit_size == 32; + } + } + } + ralloc_free(b.shader); + return left == 0 ? agx : 0; +} + +TEST(CmatLower, HalfInputsFloatAccumulator) +{ + EXPECT_EQ(lowered_muladd_count(GLSL_TYPE_FLOAT), 1u); +} + +TEST(CmatLower, HalfInputsHalfAccumulator) +{ + EXPECT_EQ(lowered_muladd_count(GLSL_TYPE_FLOAT16), 1u); +} From be59f0fc4acb7be988a3765262771ba31e327085 Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:07:08 -0600 Subject: [PATCH 2/2] nir/divergence: treat cmat_muladd_agx as divergent agx_nir_lower_cmat turns every cooperative matrix multiply-add into cmat_muladd_agx, and the divergence analysis run on the lowered shader has no case for it. With assertions enabled the first shader that uses cooperative matrices dies in visit_intrinsic with "Unhandled intrinsic", for example every matrix multiply pipeline of llama.cpp. Each lane holds its own part of the result, which is how the other cooperative matrix multiply-add intrinsics are classified, so list it with them. The CmatLower tests now run the divergence analysis after the lowering and expect the multiply-add result to be divergent. They abort on the unhandled intrinsic without this change. Signed-off-by: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> --- src/asahi/compiler/test/test-cmat-accumulator.cpp | 7 +++++-- src/compiler/nir/nir_divergence_analysis.c | 1 + 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/asahi/compiler/test/test-cmat-accumulator.cpp b/src/asahi/compiler/test/test-cmat-accumulator.cpp index 54a1b11634a0..fa45541ecacb 100644 --- a/src/asahi/compiler/test/test-cmat-accumulator.cpp +++ b/src/asahi/compiler/test/test-cmat-accumulator.cpp @@ -56,6 +56,7 @@ lowered_muladd_count(enum glsl_base_type accum) nir_cmat_muladd(&b, &d->def, &a->def, &bb->def, &c->def); agx_nir_lower_cmat(b.shader); + nir_divergence_analysis(b.shader); unsigned left = 0, agx = 0; nir_foreach_function_impl(impl, b.shader) { @@ -65,9 +66,11 @@ lowered_muladd_count(enum glsl_base_type accum) continue; nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); left += intr->intrinsic == nir_intrinsic_cmat_muladd; - /* The hardware accumulates in 32 bits whatever the matrix type. */ + /* The hardware accumulates in 32 bits whatever the matrix type, + * and every lane holds its own part of the result. + */ agx += intr->intrinsic == nir_intrinsic_cmat_muladd_agx && - intr->src[2].ssa->bit_size == 32; + intr->src[2].ssa->bit_size == 32 && intr->def.divergent; } } } 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: