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..fa45541ecacb --- /dev/null +++ b/src/asahi/compiler/test/test-cmat-accumulator.cpp @@ -0,0 +1,89 @@ +/* + * 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); + nir_divergence_analysis(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, + * 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->def.divergent; + } + } + } + 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); +} 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: