From 5aba080c3eb6a2f470e90e98f4a4c449b2404fe6 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 7 Sep 2026 17:18:41 -0500
Subject: [PATCH 01/16] honeykrisp: hardware 8x8x8 fp32 cooperative matrix via
simd_matrix_fmadd32
Port the GPU-verified HW matrix path from the agx-llm-perf showcase:
opcodes+packer (byte-exact vs applegpu), NIR intrinsic bridge, agx_nir_lower_simdmat
lowering, validator widths, and gated advertisement of an 8x8x8 fp32 coopmat shape.
Everything is behind AGX_SIMDMAT; default Honeykrisp behavior is unchanged.
---
src/asahi/compiler/agx_compile.c | 18 +
src/asahi/compiler/agx_compiler.h | 1 +
src/asahi/compiler/agx_nir_lower_simdmat.c | 544 +++++++++++++++++++++
src/asahi/compiler/agx_opcodes.py | 7 +
src/asahi/compiler/agx_pack.c | 7 +
src/asahi/compiler/agx_validate.c | 9 +
src/asahi/compiler/meson.build | 1 +
src/asahi/vulkan/hk_physical_device.c | 38 ++
src/compiler/nir/nir_divergence_analysis.c | 7 +
src/compiler/nir/nir_intrinsics.py | 6 +
10 files changed, 638 insertions(+)
create mode 100644 src/asahi/compiler/agx_nir_lower_simdmat.c
diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index e131bcf07211..8904beab6840 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -1697,6 +1697,17 @@ agx_emit_intrinsic(agx_builder *b, nir_intrinsic_instr *instr)
return agx_quad_ballot_to(b, dst, agx_src_index(&instr->src[0]));
}
+ case nir_intrinsic_simd_matrix_fmadd_agx: {
+ /* Apple G13 hardware 8x8x8 matrix MAC: D = A*B + C, SIMD-group-wide.
+ * Each operand is this lane's 2 fragment elements (a 32-bit register pair).
+ */
+ agx_simd_matrix_fmadd32_to(b, dst, agx_src_index(&instr->src[0]),
+ agx_src_index(&instr->src[1]),
+ agx_src_index(&instr->src[2]));
+ agx_emit_cached_split(b, dst, 2);
+ return NULL;
+ }
+
case nir_intrinsic_doorbell_agx: {
return agx_doorbell(b, nir_src_as_uint(instr->src[0]));
}
@@ -3590,6 +3601,13 @@ agx_preprocess_nir(nir_shader *nir)
if (!nir)
return;
+ /* Lower VK_KHR_cooperative_matrix to the G13 HW matrix instruction before
+ * vars_to_scratch can spill the 8x8 cmat temps. Gated so default Honeykrisp
+ * does not compile experimental coopmat shaders.
+ */
+ if (getenv("AGX_SIMDMAT"))
+ NIR_PASS(_, nir, agx_nir_lower_simdmat, 32);
+
NIR_PASS(_, nir, nir_lower_vars_to_ssa);
/* Lower large arrays to scratch and small arrays to csel */
diff --git a/src/asahi/compiler/agx_compiler.h b/src/asahi/compiler/agx_compiler.h
index 45460934cf0f..6fe205f169d2 100644
--- a/src/asahi/compiler/agx_compiler.h
+++ b/src/asahi/compiler/agx_compiler.h
@@ -1096,6 +1096,7 @@ bool agx_nir_opt_preamble(nir_shader *s, unsigned *sizes);
bool agx_nir_lower_load_mask(nir_shader *shader);
bool agx_nir_lower_ubo(nir_shader *shader);
bool agx_nir_lower_shared_bitsize(nir_shader *shader);
+bool agx_nir_lower_simdmat(nir_shader *shader, unsigned subgroup_size);
bool agx_nir_lower_frag_sidefx(nir_shader *s);
struct agx_cycle_estimate {
diff --git a/src/asahi/compiler/agx_nir_lower_simdmat.c b/src/asahi/compiler/agx_nir_lower_simdmat.c
new file mode 100644
index 000000000000..bff4d74dd1b7
--- /dev/null
+++ b/src/asahi/compiler/agx_nir_lower_simdmat.c
@@ -0,0 +1,544 @@
+/*
+ * Copyright 2026
+ * SPDX-License-Identifier: MIT
+ *
+ * Lower cooperative-matrix to the Apple G13 HARDWARE matrix instruction
+ * (simd_matrix_fmadd16/32), instead of the software shuffle+FMA path
+ * (agx_nir_lower_cmat.c, ~12x slower). One SIMD-group-wide 8x8x8 tile = ONE
+ * instruction: D = A*B + C, row-major, no transpose.
+ *
+ * The instruction was RE'd from Metal's simdgroup_multiply_accumulate codegen
+ * (dougallj/applegpu SimdMatrixFMadd32InstructionDesc; macOS phase). Encoding +
+ * packer are in agx_opcodes.py / agx_pack.c (byte-exact verified). This pass
+ * provides the NIR->instruction lowering.
+ *
+ * FRAGMENT LAYOUT (one 8x8 sub-block across 32 lanes, 2 consecutive elems/lane):
+ * row = ((lane>>4)&1)*4 + ((lane>>1)&3)
+ * col = ((lane>>3)&1)*4 + (lane&1)*2 (lane holds (row,col) and (row,col+1))
+ * fp32: the 2 elements are a register pair (vec2). VERIFIED 2/2 vs golden_vectors
+ * on GPU (8x8 fp32). See FRAGMENT_LAYOUT.md.
+ *
+ * SIZES: N x N x N for N in {8, 16}. 8x8 = one HW tile. 16x16 is tiled into a
+ * 2x2 grid of 8x8 sub-blocks (nblk = N/8): each lane holds nblk*nblk register
+ * pairs (vec(N*N/32): vec2 for 8x8, vec8 for 16x16). Sub-block (bi,bj) lives in
+ * vec channels [sb*2, sb*2+1] (sb = bi*nblk+bj) and maps to matrix rows
+ * [bi*8,bi*8+8) x cols [bj*8,bj*8+8) using the verified 8x8 layout above.
+ * muladd does blocked GEMM: D[di][dj] = C[di][dj] + sum_k A[di][k]*B[k][dj],
+ * i.e. nblk^3 HW fmadds (8 for 16x16).
+ *
+ * !!! UNVERIFIED for N=16 / fp16 (no GPU available to the impl fork): the 8x8
+ * fp32 sub-block is GPU-verified, but (a) that the 4 sub-blocks tile the 16x16
+ * in this row-major-block order and (b) that fmadd16 uses the SAME fragment
+ * layout as fmadd32 are ASSUMPTIONS. Discriminating oracles: an 8x8 fp16 golden
+ * (gpu_check16) isolates (b); a 16x16 golden isolates (a). If 16x16 output is
+ * block-permuted vs golden, fix the sb<->(bi,bj) mapping in subblk()/load_store.
+ *
+ * Precision: operands carry their bit size (fp16 A/B -> fmadd16; fp32 -> fmadd32,
+ * selected in agx_compile.c). For an fp16xfp16->fp32 (f32acc) shape, A/B are
+ * up-converted to fp32 here so the accumulate uses the verified fmadd32 layout.
+ *
+ * Gated by AGX_SIMDMAT (8x8) or AGX_COOPMAT+AGX_HWMAT (16x16) at the call site.
+ */
+
+#include "util/macros.h"
+#include "compiler/nir/nir_builder.h"
+#include "agx_compiler.h"
+
+#define SIMDMAT_SUBGROUP 32
+
+/* rows*cols / subgroup_size. For 8x8 over 32 lanes = 2 elements/lane. */
+static unsigned
+simdmat_len(struct glsl_cmat_description desc)
+{
+ return (desc.rows * desc.cols) / SIMDMAT_SUBGROUP;
+}
+
+static const struct glsl_type *
+remap_type(struct hash_table *mapping, const struct glsl_type *orig)
+{
+ struct hash_entry *e = _mesa_hash_table_search(mapping, orig);
+ if (e)
+ return e->data;
+
+ const struct glsl_type *nt = orig;
+ if (glsl_type_is_cmat(orig)) {
+ struct glsl_cmat_description d = *glsl_get_cmat_description(orig);
+ nt = glsl_vector_type(d.element_type, simdmat_len(d)); /* vec2 fp32 */
+ } else if (glsl_type_is_array(orig)) {
+ const struct glsl_type *el = glsl_get_array_element(orig);
+ const struct glsl_type *ne = remap_type(mapping, el);
+ if (ne != el)
+ nt = glsl_array_type(ne, glsl_get_length(orig),
+ glsl_get_explicit_stride(orig));
+ }
+ _mesa_hash_table_insert(mapping, orig, (void *)nt);
+ return nt;
+}
+
+static struct glsl_cmat_description
+src_desc(nir_src s)
+{
+ return *glsl_get_cmat_description(nir_src_as_deref(s)->type);
+}
+
+static nir_def *
+load_src(nir_builder *b, nir_src s)
+{
+ nir_deref_instr *d = nir_src_as_deref(s);
+ struct glsl_cmat_description desc = *glsl_get_cmat_description(d->type);
+ return nir_build_load_deref(b, simdmat_len(desc),
+ glsl_base_type_bit_size(desc.element_type),
+ &d->def, 0);
+}
+
+static void
+store_src(nir_builder *b, nir_src s, nir_def *v)
+{
+ nir_store_deref(b, nir_src_as_deref(s), v,
+ nir_component_mask(v->num_components));
+}
+
+/* lane -> logical (row,col) of the fragment's first element (2nd is col+1). */
+static void
+frag_rc(nir_builder *b, nir_def *lane, nir_def **row, nir_def **col)
+{
+ *row = nir_iadd(b,
+ nir_imul_imm(b, nir_iand_imm(b, nir_ushr_imm(b, lane, 4), 1), 4),
+ nir_iand_imm(b, nir_ushr_imm(b, lane, 1), 3));
+ *col = nir_iadd(b,
+ nir_imul_imm(b, nir_iand_imm(b, nir_ushr_imm(b, lane, 3), 1), 4),
+ nir_imul_imm(b, nir_iand_imm(b, lane, 1), 2));
+}
+
+static bool
+lower_load_store(nir_builder *b, struct hash_table *tm, nir_intrinsic_instr *intr)
+{
+ const bool is_load = intr->intrinsic == nir_intrinsic_cmat_load;
+ struct glsl_cmat_description desc = src_desc(intr->src[!is_load]);
+ enum glsl_matrix_layout layout = nir_intrinsic_matrix_layout(intr);
+ nir_deref_instr *cmat_deref = nir_src_as_deref(intr->src[!is_load]);
+ nir_deref_instr *deref = nir_src_as_deref(intr->src[is_load]);
+ nir_def *stride = intr->src[2].ssa;
+
+ /* For inline Q4 decode: the coopMatLoad pointer already points at the tile's
+ * base element a[base_off]; capture base_off (buffer-element units) so the
+ * decode can form the ABSOLUTE linear weight index = base_off + flat. */
+ nir_def *base_off = NULL;
+ {
+ nir_deref_instr *od = nir_src_as_deref(intr->src[is_load]);
+ while (od && od->deref_type == nir_deref_type_cast)
+ od = nir_deref_instr_parent(od);
+ if (od && (od->deref_type == nir_deref_type_ptr_as_array ||
+ od->deref_type == nir_deref_type_array))
+ base_off = od->arr.index.ssa;
+ }
+
+ unsigned tsz = glsl_base_type_bit_size(desc.element_type) / 8;
+ /* Keep the buffer's OWN element type (may be a vecN -- e.g. mul_mm's vec2
+ * shared tiles) so the caller-supplied stride/index, which are in
+ * buffer-element units, address correctly. We descend to a scalar only for
+ * the contiguous within-row index below. Casting to scalar here (as a
+ * previous version did) made the major stride off by the buffer vec width,
+ * which silently corrupted every vec2-backed load/store. */
+ const uint32_t ptr_stride =
+ glsl_get_bit_size(deref->type) / 8 * glsl_get_vector_elements(deref->type);
+ deref = nir_build_deref_cast(b, &deref->def, deref->modes, deref->type,
+ ptr_stride);
+ const struct glsl_type *ct = remap_type(tm, cmat_deref->type);
+ cmat_deref =
+ nir_build_deref_cast(b, &cmat_deref->def, cmat_deref->modes, ct, 0);
+
+ nir_def *lane = nir_load_subgroup_invocation(b);
+ nir_def *row, *col;
+ frag_rc(b, lane, &row, &col);
+
+ unsigned idx_bits = deref->def.bit_size;
+ nir_def *str = nir_u2uN(b, stride, idx_bits);
+ unsigned nblk = desc.rows / 8; /* 1 for 8x8, 2 for 16x16 */
+ unsigned len = simdmat_len(desc); /* nblk*nblk*2 */
+ nir_def *vals[8];
+
+ if (!is_load) {
+ nir_def *src = load_src(b, intr->src[!is_load]);
+ for (unsigned e = 0; e < len; e++)
+ vals[e] = nir_channel(b, src, e);
+ }
+
+ /* Flatten the two-level (major ptr_as_array -> cast -> minor ptr_as_array)
+ * address into a single scalar index = major*nvec + minor on a scalar-cast
+ * deref built ONCE. Same byte address (major steps whole vecN elements =
+ * nvec scalars; minor steps scalars), but one address computation per element
+ * instead of two -- the lowering's address arithmetic is the matmul's #1 cost
+ * (pt7). Default ON (opt out via AGX_HWMAT_NO_FLATADDR); verified MUL_MAT
+ * 947/947 + PPL 1.0162, pp512 177.1 -> 179.4. */
+ const bool flataddr = getenv("AGX_HWMAT_NO_FLATADDR") == NULL;
+ const unsigned nvec = glsl_get_vector_elements(deref->type);
+ nir_deref_instr *sdref =
+ flataddr ? nir_build_deref_cast(b, &deref->def, deref->modes,
+ glsl_scalar_type(desc.element_type), tsz)
+ : NULL;
+
+ /* VEC2 fragment loads (AGX_HWMAT_VEC2, default ON): load the contiguous
+ * (col,col+1) row-major pair as ONE i16,xy load instead of two i16,x scalars.
+ * pt11 found this perf-neutral ALONE, but macOS RE of Apple's pp512=270 kernel
+ * (RECIPE 2026-06-22) shows Apple emits paired i16,xy loads, ALL hoisted into one
+ * block before a tight 8x4 simd_matrix burst. The pairing is the ENABLER: it gives
+ * 24 distinct fragment regs (vs ~12 reused) so the post-RA scheduler can hoist all
+ * loads (kills the 42 waits -> ~1). flat is even (frag_rc col is always x2) so the
+ * pair is vec2-aligned. Row-major (A) only; col-major (B) pair is strided. */
+ const bool vec2ld = getenv("AGX_HWMAT_VEC2") != NULL;
+ nir_deref_instr *v2dref =
+ (flataddr && vec2ld)
+ ? nir_build_deref_cast(b, &deref->def, deref->modes,
+ glsl_vector_type(desc.element_type, 2), 2 * tsz)
+ : NULL;
+
+ /* Walk the nblk x nblk grid of 8x8 sub-blocks; each lane handles its 2
+ * fragment elements per sub-block. Sub-block (bi,bj) -> matrix rows/cols
+ * [bi*8,..) x [bj*8,..), stored at vec channels [sb*2, sb*2+1]. */
+ for (unsigned bi = 0; bi < nblk; bi++) {
+ for (unsigned bj = 0; bj < nblk; bj++) {
+ unsigned sb = bi * nblk + bj;
+ for (unsigned e = 0; e < 2; e++) {
+ /* logical element (bi*8+row, bj*8+col+e). The strided (major) axis
+ * is addressed in buffer-element units on the still-vecN deref; the
+ * contiguous (minor) axis in scalar units after casting down. */
+ bool colmaj = (layout == GLSL_MATRIX_LAYOUT_COLUMN_MAJOR);
+ nir_def *lr = nir_u2uN(b, nir_iadd_imm(b, row, bi * 8), idx_bits);
+ nir_def *lc =
+ nir_u2uN(b, nir_iadd_imm(b, col, bj * 8 + e), idx_bits);
+ nir_def *major = nir_imul(b, colmaj ? lc : lr, str);
+ nir_def *minor = colmaj ? lr : lc;
+ nir_deref_instr *it;
+ nir_def *flat = NULL;
+ if (flataddr) {
+ /* flat (scalar units) = major*nvec + minor */
+ flat = nir_iadd(b, nir_imul_imm(b, major, nvec), minor);
+ it = nir_build_deref_ptr_as_array(b, sdref, flat);
+ } else {
+ it = nir_build_deref_ptr_as_array(b, deref, major);
+ it = nir_build_deref_cast(b, &it->def, deref->modes,
+ glsl_scalar_type(desc.element_type), tsz);
+ it = nir_build_deref_ptr_as_array(b, it, minor);
+ }
+ /* INLINE Q4_0 DECODE (env-gated, A-operand): instead of loading f16,
+ * decode a quantized weight directly from a 4-byte-aligned padded-Q4_0
+ * buffer (block = 20 bytes = 10 u16: [d:f16][16B qs][2B pad]) using the
+ * linear weight index `flat`. Eliminates the f16 dequant-staging phase
+ * that bounds the Q4 matmul (LAB_NOTEBOOK pt43-45). PROTOTYPE: assumes
+ * deref base = buffer start (flat = absolute linear index). */
+ if (is_load && flataddr && getenv("AGX_DECODE_Q4") &&
+ desc.use == GLSL_CMAT_USE_A) {
+ /* AMORTIZED: decode BOTH elements (e=0,1) of this sub-block from ONE
+ * scale + ONE qbyte load (they are the lo/hi nibble of the same byte
+ * in the padded-Q4 format). Only run at e==0; e==1 is filled here. */
+ if (e == 0) {
+ nir_deref_instr *u16d = nir_build_deref_cast(
+ b, &deref->def, deref->modes, glsl_uint16_t_type(), 2);
+ nir_def *aidx = base_off
+ ? nir_iadd(b, nir_u2uN(b, base_off, idx_bits), flat) : flat;
+ nir_def *blk = nir_ushr_imm(b, aidx, 5); /* /32 */
+ nir_def *within = nir_iand_imm(b, aidx, 31); /* %32 (even) */
+ nir_def *base = nir_imul_imm(b, blk, 10); /* u16 units */
+ nir_def *d = nir_load_deref(
+ b, nir_build_deref_ptr_as_array(b, sdref, base));
+ nir_def *bb = nir_ushr_imm(b, within, 1); /* byte index 0..15 */
+ nir_def *qoff = nir_iadd(b, nir_iadd_imm(b, base, 1),
+ nir_u2uN(b, nir_ushr_imm(b, bb, 1), idx_bits));
+ nir_def *qv32 = nir_u2u32(b, nir_load_deref(
+ b, nir_build_deref_ptr_as_array(b, u16d, qoff)));
+ /* byte = (bb&1)? hi8 : lo8 */
+ nir_def *byte = nir_bcsel(b, nir_ine_imm(b, nir_iand_imm(b, bb, 1), 0),
+ nir_ushr_imm(b, qv32, 8),
+ nir_iand_imm(b, qv32, 0xFF));
+ nir_def *nlo = nir_iand_imm(b, byte, 0xF); /* e=0 */
+ nir_def *nhi = nir_iand_imm(b, nir_ushr_imm(b, byte, 4), 0xF); /* e=1 */
+ nir_def *flo = nir_i2f16(b, nir_iadd_imm(b, nir_u2u32(b, nlo), -8));
+ nir_def *fhi = nir_i2f16(b, nir_iadd_imm(b, nir_u2u32(b, nhi), -8));
+ vals[sb * 2 + 0] = nir_fmul(b, flo, d);
+ vals[sb * 2 + 1] = nir_fmul(b, fhi, d);
+ }
+ } else if (is_load && v2dref && !colmaj) {
+ /* VEC2: load the contiguous (col,col+1) row-major pair as one
+ * i16,xy. flat is even -> vec2 index = flat>>1. Fill both elems at
+ * e==0; e==1 is a no-op. Gives paired loads + distinct regs so the
+ * post-RA scheduler can hoist them (RECIPE 2026-06-22). */
+ if (e == 0) {
+ nir_def *pv = nir_load_deref(
+ b, nir_build_deref_ptr_as_array(b, v2dref,
+ nir_ushr_imm(b, flat, 1)));
+ vals[sb * 2 + 0] = nir_channel(b, pv, 0);
+ vals[sb * 2 + 1] = nir_channel(b, pv, 1);
+ }
+ } else if (is_load)
+ vals[sb * 2 + e] = nir_load_deref(b, it);
+ else
+ nir_store_deref(b, it, vals[sb * 2 + e], 0x1);
+ }
+ }
+ }
+
+ if (is_load)
+ nir_store_deref(b, cmat_deref, nir_vec(b, vals, len),
+ nir_component_mask(len));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+/* Extract 8x8 sub-block sb (its vec2 register pair) from a packed fragment. */
+static nir_def *
+subblk(nir_builder *b, nir_def *v, unsigned sb)
+{
+ return nir_vec2(b, nir_channel(b, v, sb * 2), nir_channel(b, v, sb * 2 + 1));
+}
+
+static bool
+lower_muladd(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ /* src order: dst, A, B, C. A/B element type may differ from the accumulator
+ * (f32acc shape: fp16 A/B, fp32 C/D). */
+ struct glsl_cmat_description da = src_desc(intr->src[1]); /* A */
+ struct glsl_cmat_description dc = src_desc(intr->src[3]); /* C / accumulator */
+ unsigned nblk = da.rows / 8;
+ unsigned acc_bits = glsl_base_type_bit_size(dc.element_type);
+ unsigned in_bits = glsl_base_type_bit_size(da.element_type);
+
+ nir_def *a = load_src(b, intr->src[1]);
+ nir_def *bb = load_src(b, intr->src[2]);
+ nir_def *c = load_src(b, intr->src[3]);
+
+ /* Accumulate in the C type. If A/B are narrower (fp16 -> fp32 acc),
+ * up-convert so the HW op uses the GPU-verified fmadd32 fragment layout. */
+ if (in_bits != acc_bits) {
+ a = nir_f2fN(b, a, acc_bits);
+ bb = nir_f2fN(b, bb, acc_bits);
+ }
+
+ /* Blocked GEMM over the 8x8 sub-blocks: D[di][dj] = C[di][dj] +
+ * sum_k A[di][k] * B[k][dj]. Each term is ONE HW matrix fmadd
+ * (nblk^3 total: 1 for 8x8, 8 for 16x16). */
+ unsigned len = simdmat_len(dc);
+ nir_def *dch[8];
+ for (unsigned di = 0; di < nblk; di++) {
+ for (unsigned dj = 0; dj < nblk; dj++) {
+ unsigned sb = di * nblk + dj;
+ nir_def *acc = subblk(b, c, sb);
+ for (unsigned k = 0; k < nblk; k++) {
+ nir_def *as = subblk(b, a, di * nblk + k);
+ nir_def *bs = subblk(b, bb, k * nblk + dj);
+ acc = nir_simd_matrix_fmadd_agx(b, acc_bits, as, bs, acc);
+ }
+ dch[sb * 2] = nir_channel(b, acc, 0);
+ dch[sb * 2 + 1] = nir_channel(b, acc, 1);
+ }
+ }
+
+ store_src(b, intr->src[0], nir_vec(b, dch, len));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_construct(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ nir_deref_instr *dst = nir_src_as_deref(intr->src[0]);
+ struct glsl_cmat_description desc = *glsl_get_cmat_description(dst->type);
+ nir_def *r = nir_replicate(b, intr->src[1].ssa, simdmat_len(desc));
+ nir_store_deref(b, dst, r, nir_component_mask(r->num_components));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_length(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ struct glsl_cmat_description desc = nir_intrinsic_cmat_desc(intr);
+ nir_def_replace(&intr->def, nir_imm_int(b, simdmat_len(desc)));
+ return true;
+}
+
+/* Element-wise ops on the packed fragment. binary/scalar/unary/extract/insert/
+ * bitcast never move data between lanes, so they act identically on the SW and
+ * the HW (frag_rc) layouts: just operate per-vec-component. */
+static bool
+lower_binary_op(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ nir_def *a = load_src(b, intr->src[1]);
+ nir_def *c = load_src(b, intr->src[2]);
+ store_src(b, intr->src[0], nir_build_alu2(b, nir_intrinsic_alu_op(intr), a, c));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_unary_op(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ nir_def *a = load_src(b, intr->src[1]);
+ store_src(b, intr->src[0], nir_build_alu1(b, nir_intrinsic_alu_op(intr), a));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_scalar_op(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ nir_def *a = load_src(b, intr->src[1]);
+ store_src(b, intr->src[0],
+ nir_build_alu2(b, nir_intrinsic_alu_op(intr), a, intr->src[2].ssa));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_extract(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ nir_def *mat = load_src(b, intr->src[0]);
+ nir_def_replace(&intr->def, nir_vector_extract(b, mat, intr->src[1].ssa));
+ return true;
+}
+
+static bool
+lower_insert(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ nir_def *mat = load_src(b, intr->src[2]);
+ store_src(b, intr->src[0],
+ nir_vector_insert(b, mat, intr->src[1].ssa, intr->src[3].ssa));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_bitcast(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ store_src(b, intr->src[0], load_src(b, intr->src[1]));
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+/* convert: base-type conversion is element-wise (layout-safe). A use-change
+ * (A<->B<->Accumulator) or transpose would need a HW-layout-aware cross-lane
+ * relayout, which we don't implement yet -- assert so it's loud, not silent
+ * corruption. mul_mm's accumulator convert is base-type only (same use). */
+static bool
+lower_convert(nir_builder *b, nir_intrinsic_instr *intr)
+{
+ struct glsl_cmat_description dd = src_desc(intr->src[0]);
+ struct glsl_cmat_description sd = src_desc(intr->src[1]);
+ enum glsl_cmat_use du = dd.use, su = sd.use;
+ if (du == GLSL_CMAT_USE_ACCUMULATOR)
+ du = GLSL_CMAT_USE_A;
+ if (su == GLSL_CMAT_USE_ACCUMULATOR)
+ su = GLSL_CMAT_USE_A;
+ assert(du == su && intr->intrinsic != nir_intrinsic_cmat_transpose &&
+ "HWMAT: cmat use-change/transpose not implemented");
+
+ nir_def *src = load_src(b, intr->src[1]);
+ nir_def *ret = src;
+ if (dd.element_type != sd.element_type) {
+ nir_op op = nir_type_conversion_op(
+ nir_get_nir_type_for_glsl_base_type(sd.element_type),
+ nir_get_nir_type_for_glsl_base_type(dd.element_type),
+ nir_rounding_mode_undef);
+ ret = nir_build_alu1(b, op, src);
+ }
+ store_src(b, intr->src[0], ret);
+ nir_instr_remove(&intr->instr);
+ return true;
+}
+
+static bool
+lower_impl(nir_function_impl *impl, struct hash_table *tm)
+{
+ bool progress = false;
+
+ nir_foreach_function_temp_variable(var, impl) {
+ const struct glsl_type *nt = remap_type(tm, var->type);
+ if (nt != var->type) {
+ var->type = nt;
+ progress = true;
+ }
+ }
+
+ nir_builder b = nir_builder_create(impl);
+ nir_foreach_block_reverse_safe(block, impl) {
+ nir_foreach_instr_reverse_safe(instr, block) {
+ b.cursor = nir_before_instr(instr);
+
+ if (instr->type == nir_instr_type_deref) {
+ nir_deref_instr *deref = nir_instr_as_deref(instr);
+ const struct glsl_type *nt = remap_type(tm, deref->type);
+ if (nt != deref->type) {
+ deref->type = nt;
+ progress = true;
+ }
+ continue;
+ }
+ if (instr->type != nir_instr_type_intrinsic)
+ continue;
+
+ nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+ switch (intr->intrinsic) {
+ case nir_intrinsic_cmat_length:
+ progress |= lower_length(&b, intr);
+ break;
+ case nir_intrinsic_cmat_construct:
+ progress |= lower_construct(&b, intr);
+ break;
+ case nir_intrinsic_cmat_load:
+ case nir_intrinsic_cmat_store:
+ progress |= lower_load_store(&b, tm, intr);
+ break;
+ case nir_intrinsic_cmat_muladd:
+ progress |= lower_muladd(&b, intr);
+ break;
+ case nir_intrinsic_cmat_copy:
+ nir_build_copy_deref(&b, intr->src[0].ssa, intr->src[1].ssa);
+ nir_instr_remove(instr);
+ progress = true;
+ break;
+ case nir_intrinsic_cmat_binary_op:
+ progress |= lower_binary_op(&b, intr);
+ break;
+ case nir_intrinsic_cmat_unary_op:
+ progress |= lower_unary_op(&b, intr);
+ break;
+ case nir_intrinsic_cmat_scalar_op:
+ progress |= lower_scalar_op(&b, intr);
+ break;
+ case nir_intrinsic_cmat_extract:
+ progress |= lower_extract(&b, intr);
+ break;
+ case nir_intrinsic_cmat_insert:
+ progress |= lower_insert(&b, intr);
+ break;
+ case nir_intrinsic_cmat_bitcast:
+ progress |= lower_bitcast(&b, intr);
+ break;
+ case nir_intrinsic_cmat_convert:
+ case nir_intrinsic_cmat_transpose:
+ progress |= lower_convert(&b, intr);
+ break;
+ default:
+ break;
+ }
+ }
+ }
+
+ return nir_progress(progress, impl, nir_metadata_none);
+}
+
+bool
+agx_nir_lower_simdmat(nir_shader *shader, unsigned subgroup_size)
+{
+ assert(subgroup_size == SIMDMAT_SUBGROUP);
+ bool progress = false;
+ struct hash_table *tm = _mesa_pointer_hash_table_create(NULL);
+
+ nir_foreach_function_impl(impl, shader) {
+ if (lower_impl(impl, tm))
+ progress = true;
+ }
+
+ _mesa_hash_table_destroy(tm, NULL);
+ return progress;
+}
diff --git a/src/asahi/compiler/agx_opcodes.py b/src/asahi/compiler/agx_opcodes.py
index f5afa4623e91..1d7f33dff8f9 100644
--- a/src/asahi/compiler/agx_opcodes.py
+++ b/src/asahi/compiler/agx_opcodes.py
@@ -217,6 +217,13 @@ def iunop(name, opcode):
op("fadd", encoding = (0x2A, 4, 6), srcs = 2, is_float = True)
op("ffma", encoding = (0x3A, 6, 8), srcs = 3, is_float = True)
op("fmul", encoding = (0x1A, 4, 6), srcs = 2, is_float = True)
+
+# Apple G13 hardware matrix multiply-accumulate: D = A*B + C over an 8x8x8 tile,
+# SIMD-group-wide (RE'd from Metal simdgroup_multiply_accumulate codegen).
+# opcode 0x6F, bit27 fixed, bit26 = precision (1 = fp32, 0 = fp16).
+# A/B/C/D are paired-register operands; bit63 is applied in agx_pack.c.
+op("simd_matrix_fmadd32", encoding = (0x6F | (0x1 << 26) | (0x1 << 27), 8, _), srcs = 3, is_float = True)
+op("simd_matrix_fmadd16", encoding = (0x6F | (0x1 << 27), 8, _), srcs = 3, is_float = True)
op("hadd", encoding = (0x26, 4, 6), srcs = 2, is_float = True)
op("hfma", encoding = (0x36, 6, 8), srcs = 3, is_float = True)
op("hmul", encoding = (0x16, 4, 6), srcs = 2, is_float = True)
diff --git a/src/asahi/compiler/agx_pack.c b/src/asahi/compiler/agx_pack.c
index 6236c3cf4f3d..18c183d22feb 100644
--- a/src/asahi/compiler/agx_pack.c
+++ b/src/asahi/compiler/agx_pack.c
@@ -510,6 +510,13 @@ agx_pack_alu(struct util_dynarray *emission, agx_instr *I)
else if (info.immediates & (AGX_IMMEDIATE_FCOND | AGX_IMMEDIATE_ICOND))
raw |= (uint64_t)(I->fcond) << 61;
+ /* Apple's hardware matrix multiply-accumulate sets bit 63; the generic
+ * ALU packer doesn't, so set it here.
+ */
+ if (I->op == AGX_OPCODE_SIMD_MATRIX_FMADD32 ||
+ I->op == AGX_OPCODE_SIMD_MATRIX_FMADD16)
+ raw |= (uint64_t)1 << 63;
+
/* Determine length bit */
unsigned length = encoding.length_short;
if (I->op == AGX_OPCODE_MOV_IMM && I->dest[0].size == AGX_SIZE_16)
diff --git a/src/asahi/compiler/agx_validate.c b/src/asahi/compiler/agx_validate.c
index 788b5e5ecb5e..87c4e1b36ee8 100644
--- a/src/asahi/compiler/agx_validate.c
+++ b/src/asahi/compiler/agx_validate.c
@@ -185,6 +185,10 @@ agx_write_registers(const agx_instr *I, unsigned d)
return 6;
case AGX_OPCODE_COLLECT:
return I->nr_srcs * agx_size_align_16(I->src[0].size);
+ case AGX_OPCODE_SIMD_MATRIX_FMADD32:
+ case AGX_OPCODE_SIMD_MATRIX_FMADD16:
+ /* paired-register fragment (2 elements/lane) */
+ return agx_index_size_16(I->dest[d]);
default:
return size;
}
@@ -352,6 +356,11 @@ agx_read_registers(const agx_instr *I, unsigned s)
else
return size;
+ case AGX_OPCODE_SIMD_MATRIX_FMADD32:
+ case AGX_OPCODE_SIMD_MATRIX_FMADD16:
+ /* paired-register fragment operands (2 elements/lane) */
+ return agx_index_size_16(I->src[s]);
+
default:
return size;
}
diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build
index cce0dbf8f1d2..62430fe3b641 100644
--- a/src/asahi/compiler/meson.build
+++ b/src/asahi/compiler/meson.build
@@ -16,6 +16,7 @@ libasahi_agx_files = files(
'agx_nir_lower_interpolation.c',
'agx_nir_lower_shared_bitsize.c',
'agx_nir_lower_subgroups.c',
+ 'agx_nir_lower_simdmat.c',
'agx_nir_lower_texture.c',
'agx_nir_opt_preamble.c',
'agx_lower_64bit.c',
diff --git a/src/asahi/vulkan/hk_physical_device.c b/src/asahi/vulkan/hk_physical_device.c
index 80112cec02b2..6f6d99599c77 100644
--- a/src/asahi/vulkan/hk_physical_device.c
+++ b/src/asahi/vulkan/hk_physical_device.c
@@ -55,6 +55,7 @@ hk_get_device_extensions(const struct hk_instance *instance,
.KHR_bind_memory2 = true,
.KHR_buffer_device_address = true,
.KHR_calibrated_timestamps = true,
+ .KHR_cooperative_matrix = getenv("AGX_SIMDMAT") != NULL,
.KHR_copy_commands2 = true,
.KHR_create_renderpass2 = true,
.KHR_dedicated_allocation = true,
@@ -238,6 +239,10 @@ hk_get_device_features(
struct vk_features *features)
{
*features = (struct vk_features){
+ /* VK_KHR_cooperative_matrix (G13 HW 8x8x8, gated by AGX_SIMDMAT) */
+ .cooperativeMatrix = getenv("AGX_SIMDMAT") != NULL,
+ .cooperativeMatrixRobustBufferAccess = false,
+
/* Vulkan 1.0 */
.robustBufferAccess = true,
.fullDrawIndexUint32 = true,
@@ -820,6 +825,8 @@ hk_get_device_properties(const struct agx_device *dev,
/* Vulkan 1.1 properties */
.subgroupSize = 32,
+ .cooperativeMatrixSupportedStages =
+ getenv("AGX_SIMDMAT") ? VK_SHADER_STAGE_COMPUTE_BIT : 0,
.subgroupSupportedStages =
VK_SHADER_STAGE_COMPUTE_BIT | VK_SHADER_STAGE_ALL_GRAPHICS,
.subgroupSupportedOperations =
@@ -1518,3 +1525,34 @@ hk_GetPhysicalDeviceMultisamplePropertiesEXT(
pMultisampleProperties->maxSampleLocationGridSize = (VkExtent2D){0, 0};
}
}
+
+VKAPI_ATTR VkResult VKAPI_CALL
+hk_GetPhysicalDeviceCooperativeMatrixPropertiesKHR(
+ VkPhysicalDevice physicalDevice, uint32_t *pPropertyCount,
+ VkCooperativeMatrixPropertiesKHR *pProperties)
+{
+ VK_FROM_HANDLE(hk_physical_device, pdev, physicalDevice);
+ (void)pdev;
+ VK_OUTARRAY_MAKE_TYPED(VkCooperativeMatrixPropertiesKHR, out, pProperties,
+ pPropertyCount);
+
+ if (getenv("AGX_SIMDMAT")) {
+ /* Hardware 8x8x8 fp32 matrix MAC (simd_matrix_fmadd32). */
+ vk_outarray_append_typed(VkCooperativeMatrixPropertiesKHR, &out, p)
+ {
+ p->sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR;
+ p->pNext = NULL;
+ p->MSize = 8;
+ p->NSize = 8;
+ p->KSize = 8;
+ p->AType = VK_COMPONENT_TYPE_FLOAT32_KHR;
+ p->BType = VK_COMPONENT_TYPE_FLOAT32_KHR;
+ p->CType = VK_COMPONENT_TYPE_FLOAT32_KHR;
+ p->ResultType = VK_COMPONENT_TYPE_FLOAT32_KHR;
+ p->saturatingAccumulation = VK_FALSE;
+ p->scope = VK_SCOPE_SUBGROUP_KHR;
+ }
+ }
+
+ return vk_outarray_status(&out);
+}
diff --git a/src/compiler/nir/nir_divergence_analysis.c b/src/compiler/nir/nir_divergence_analysis.c
index cac39fe1fd32..877491f44ab0 100644
--- a/src/compiler/nir/nir_divergence_analysis.c
+++ b/src/compiler/nir/nir_divergence_analysis.c
@@ -217,6 +217,13 @@ visit_intrinsic(nir_intrinsic_instr *instr, struct divergence_state *state)
(state->options & nir_divergence_across_subgroups);
break;
+ case nir_intrinsic_simd_matrix_fmadd_agx:
+ /* Apple G13 hardware matrix MAC: each lane gets its own fragment of the
+ * 8x8 result -> divergent.
+ */
+ is_divergent = true;
+ break;
+
case nir_intrinsic_as_uniform:
case nir_intrinsic_read_invocation:
case nir_intrinsic_read_first_invocation:
diff --git a/src/compiler/nir/nir_intrinsics.py b/src/compiler/nir/nir_intrinsics.py
index c80a232f3c3a..b6c434d8f1b9 100644
--- a/src/compiler/nir/nir_intrinsics.py
+++ b/src/compiler/nir/nir_intrinsics.py
@@ -2472,6 +2472,12 @@ def store(name, srcs, indices=[], flags=[]):
# Equivalent to popcount(ballot(true))
intrinsic("load_active_subgroup_count_agx", dest_comp=1, flags=[CAN_ELIMINATE])
+# Apple G13 hardware matrix multiply-accumulate (simd_matrix_fmadd32). One
+# SIMD-group-wide 8x8x8 tile: D = A*B + C. Per-invocation each operand is the
+# lane's 2 fragment elements (a register pair). src order: A, B, C.
+intrinsic("simd_matrix_fmadd_agx", src_comp=[2, 2, 2], dest_comp=2,
+ bit_sizes=[32], flags=[CAN_ELIMINATE])
+
# Like ballot() but only within a quad.
intrinsic("quad_ballot_agx", src_comp=[1], dest_comp=1, flags=[CAN_ELIMINATE])
From 81a9c6113ad6b623eadf7e9d78304a56e45ec6d6 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 7 Sep 2026 17:27:35 -0500
Subject: [PATCH 02/16] honeykrisp: match generated
nir_simd_matrix_fmadd_agx(fp32) signature
---
src/asahi/compiler/agx_nir_lower_simdmat.c | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/src/asahi/compiler/agx_nir_lower_simdmat.c b/src/asahi/compiler/agx_nir_lower_simdmat.c
index bff4d74dd1b7..5b4e030ba8cd 100644
--- a/src/asahi/compiler/agx_nir_lower_simdmat.c
+++ b/src/asahi/compiler/agx_nir_lower_simdmat.c
@@ -326,7 +326,7 @@ lower_muladd(nir_builder *b, nir_intrinsic_instr *intr)
for (unsigned k = 0; k < nblk; k++) {
nir_def *as = subblk(b, a, di * nblk + k);
nir_def *bs = subblk(b, bb, k * nblk + dj);
- acc = nir_simd_matrix_fmadd_agx(b, acc_bits, as, bs, acc);
+ acc = nir_simd_matrix_fmadd_agx(b, as, bs, acc);
}
dch[sb * 2] = nir_channel(b, acc, 0);
dch[sb * 2 + 1] = nir_channel(b, acc, 1);
From 023b9b8c5ff567f17821bf32fb8b8cb1d75ef40e Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 7 Sep 2026 17:40:49 -0500
Subject: [PATCH 03/16] agx/isa: disassemble simd_matrix_fmadd32/16 so debug
builds can compile coopmat shaders
---
src/asahi/isa/AGX2.xml | 22 ++++++++++++++++++++++
1 file changed, 22 insertions(+)
diff --git a/src/asahi/isa/AGX2.xml b/src/asahi/isa/AGX2.xml
index f3e6733725fb..6407e8e78801 100644
--- a/src/asahi/isa/AGX2.xml
+++ b/src/asahi/isa/AGX2.xml
@@ -979,6 +979,28 @@
+
+ 101111
+ 1
+ 1
+ 1
+
+
+
+
+
+
+
+ 101111
+ 0
+ 1
+ 1
+
+
+
+
+
+
111010
From 9f495d41cb04795a086e024c2910f2d109e1940c Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 7 Sep 2026 17:42:25 -0500
Subject: [PATCH 04/16] agx/isa: include bit6 of opcode 0x6f in simd_matrix
exact mask
---
src/asahi/isa/AGX2.xml | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/src/asahi/isa/AGX2.xml b/src/asahi/isa/AGX2.xml
index 6407e8e78801..bba615285eff 100644
--- a/src/asahi/isa/AGX2.xml
+++ b/src/asahi/isa/AGX2.xml
@@ -980,7 +980,7 @@
- 101111
+ 01101111
1
1
1
@@ -991,7 +991,7 @@
- 101111
+ 01101111
0
1
1
From 5bc10a6006d999c30c13cadf02fdfa49141b73b5 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Tue, 8 Sep 2026 07:55:39 -0500
Subject: [PATCH 05/16] honeykrisp: advertise 16x16x16 and fp16 cooperative
matrix shapes
VK_KHR_cooperative_matrix on Honeykrisp only advertised 8x8x8 fp32 even
though agx_nir_lower_simdmat.c already tiles a 16x16x16 coopmat into a
2x2x2 grid of simd_matrix_fmadd ops, and the G13 matrix unit takes fp16
operands. Now advertised (subgroup scope, N in {8, 16}, NxNxN):
fp32 A/B, fp32 C/D simd_matrix_fmadd32 with 32-bit pairs
fp16 A/B, fp32 C/D simd_matrix_fmadd32 with 16-bit A/B pairs
fp16 A/B, fp16 C/D simd_matrix_fmadd16
The mixed form is the one Apple's own Metal mul_mm kernel emits
(simd_matrix_fmadd32 r15_r16, r49l_r49h, r46l_r46h, r15_r16); the generic
ALU packer already encodes the per-source size bits, so the lowering
passes A/B at their own width instead of up-converting to fp32, and the
nir intrinsic gains a 16-bit dest form that selects fmadd16.
Bug found while verifying: with fp16 A and fp32 D the register allocator's
early-kill let D (a 32-bit pair) land on top of A (a 16-bit pair inside
that pair). The hardware then produced zeros for that fmadd: 16x16x16
fp16/fp32 failed in every input/layout variant, always sub-block (1,1),
whose last fmadd was the only one with D overlapping A. Apple's compiler
only ever emits D == C, so the RA now keeps A and B live until D is
placed (D may still reuse C).
Verification on Apple M1 (G13G B1), receipts/cmshape.c, one coopMatMulAdd
per run, exact host reference, integer-valued inputs bit-exact and random
floats within tolerance, identity/ones/randint/randf x row/col-major:
stock Mesa 26.1.7: FAIL, extension not exposed
5bb2b28c (before): 8x8x8 fp32 8/8 PASS; 16x16x16 not advertised
this commit, no RA guard: 40/48 (16x16x16 fp16/fp32: 8/8 FAIL)
this commit: 48/48 PASS, plus 48/48 randint over 4 seeds
mlx-omarchy Release tests (primitive 98, reduce_ops 34, runtime 40,
matmul_family 14 cases) pass identically on the stock ICD and on this
driver with AGX_SIMDMAT=1.
---
src/asahi/compiler/agx_compile.c | 15 +++++---
src/asahi/compiler/agx_nir_lower_simdmat.c | 35 +++++++++---------
src/asahi/compiler/agx_register_allocate.c | 24 +++++++++++--
src/asahi/vulkan/hk_physical_device.c | 41 ++++++++++++++--------
src/compiler/nir/nir_intrinsics.py | 8 +++--
5 files changed, 83 insertions(+), 40 deletions(-)
diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index 8904beab6840..bc2468809444 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -1699,11 +1699,18 @@ agx_emit_intrinsic(agx_builder *b, nir_intrinsic_instr *instr)
case nir_intrinsic_simd_matrix_fmadd_agx: {
/* Apple G13 hardware 8x8x8 matrix MAC: D = A*B + C, SIMD-group-wide.
- * Each operand is this lane's 2 fragment elements (a 32-bit register pair).
+ * Each operand is this lane's 2 fragment elements: a 32-bit register
+ * pair for fp32, one 32-bit register (two halves) for fp16. The
+ * accumulator precision picks the opcode; A/B carry their own size
+ * bits, so fmadd32 with fp16 A/B is the mixed-precision form.
*/
- agx_simd_matrix_fmadd32_to(b, dst, agx_src_index(&instr->src[0]),
- agx_src_index(&instr->src[1]),
- agx_src_index(&instr->src[2]));
+ agx_index a = agx_src_index(&instr->src[0]);
+ agx_index bm = agx_src_index(&instr->src[1]);
+ agx_index c = agx_src_index(&instr->src[2]);
+ if (instr->def.bit_size == 16)
+ agx_simd_matrix_fmadd16_to(b, dst, a, bm, c);
+ else
+ agx_simd_matrix_fmadd32_to(b, dst, a, bm, c);
agx_emit_cached_split(b, dst, 2);
return NULL;
}
diff --git a/src/asahi/compiler/agx_nir_lower_simdmat.c b/src/asahi/compiler/agx_nir_lower_simdmat.c
index 5b4e030ba8cd..44c0202c30b8 100644
--- a/src/asahi/compiler/agx_nir_lower_simdmat.c
+++ b/src/asahi/compiler/agx_nir_lower_simdmat.c
@@ -26,18 +26,19 @@
* muladd does blocked GEMM: D[di][dj] = C[di][dj] + sum_k A[di][k]*B[k][dj],
* i.e. nblk^3 HW fmadds (8 for 16x16).
*
- * !!! UNVERIFIED for N=16 / fp16 (no GPU available to the impl fork): the 8x8
- * fp32 sub-block is GPU-verified, but (a) that the 4 sub-blocks tile the 16x16
- * in this row-major-block order and (b) that fmadd16 uses the SAME fragment
- * layout as fmadd32 are ASSUMPTIONS. Discriminating oracles: an 8x8 fp16 golden
- * (gpu_check16) isolates (b); a 16x16 golden isolates (a). If 16x16 output is
- * block-permuted vs golden, fix the sb<->(bi,bj) mapping in subblk()/load_store.
+ * GPU-verified exact (identity, ones, random-integer and random-float inputs,
+ * row- and column-major, 2026-09-08 on M1 G13G) for every advertised shape:
+ * N in {8,16} x {fp32 A/B/C, fp16 A/B + fp32 C, fp16 A/B/C}. The 16x16
+ * sub-block order above and the fp16 fragment layout (same lane mapping, two
+ * halves in one 32-bit register) are what the hardware uses.
*
- * Precision: operands carry their bit size (fp16 A/B -> fmadd16; fp32 -> fmadd32,
- * selected in agx_compile.c). For an fp16xfp16->fp32 (f32acc) shape, A/B are
- * up-converted to fp32 here so the accumulate uses the verified fmadd32 layout.
+ * Precision: the accumulator type selects the opcode (fp16 C -> fmadd16, fp32
+ * C -> fmadd32, in agx_compile.c); A/B are passed at their own width. Verified
+ * operand forms: fmadd32 with fp32 or fp16 A/B (the fp16 A/B, fp32 C form is
+ * what Apple's own mul_mm emits), fmadd16 with fp16 A/B. fp32 A/B with an fp16
+ * accumulator is not advertised and not lowered.
*
- * Gated by AGX_SIMDMAT (8x8) or AGX_COOPMAT+AGX_HWMAT (16x16) at the call site.
+ * Gated by AGX_SIMDMAT at the call site.
*/
#include "util/macros.h"
@@ -307,12 +308,12 @@ lower_muladd(nir_builder *b, nir_intrinsic_instr *intr)
nir_def *bb = load_src(b, intr->src[2]);
nir_def *c = load_src(b, intr->src[3]);
- /* Accumulate in the C type. If A/B are narrower (fp16 -> fp32 acc),
- * up-convert so the HW op uses the GPU-verified fmadd32 fragment layout. */
- if (in_bits != acc_bits) {
- a = nir_f2fN(b, a, acc_bits);
- bb = nir_f2fN(b, bb, acc_bits);
- }
+ /* The HW op accumulates in the C type and takes A/B at their own width
+ * (fp16 A/B with fp32 C is the mixed-precision form Apple's own mul_mm
+ * emits: simd_matrix_fmadd32 rN_rN+1, rMl_rMh, rKl_rKh, rN_rN+1). Only a
+ * wider A/B than C (fp32 A/B, fp16 acc) needs a convert, and no shape
+ * advertises that. */
+ assert(in_bits <= acc_bits);
/* Blocked GEMM over the 8x8 sub-blocks: D[di][dj] = C[di][dj] +
* sum_k A[di][k] * B[k][dj]. Each term is ONE HW matrix fmadd
@@ -326,7 +327,7 @@ lower_muladd(nir_builder *b, nir_intrinsic_instr *intr)
for (unsigned k = 0; k < nblk; k++) {
nir_def *as = subblk(b, a, di * nblk + k);
nir_def *bs = subblk(b, bb, k * nblk + dj);
- acc = nir_simd_matrix_fmadd_agx(b, as, bs, acc);
+ acc = nir_simd_matrix_fmadd_agx(b, acc_bits, as, bs, acc);
}
dch[sb * 2] = nir_channel(b, acc, 0);
dch[sb * 2 + 1] = nir_channel(b, acc, 1);
diff --git a/src/asahi/compiler/agx_register_allocate.c b/src/asahi/compiler/agx_register_allocate.c
index 87a0335a4a79..312b31e21317 100644
--- a/src/asahi/compiler/agx_register_allocate.c
+++ b/src/asahi/compiler/agx_register_allocate.c
@@ -1014,6 +1014,26 @@ kill_source(struct ra_ctx *rctx, const agx_instr *I, unsigned s)
BITSET_CLEAR_COUNT(rctx->used_regs[cls], reg, count);
}
+/* A killed source may be freed before the destination is allocated so the
+ * destination can reuse its registers. The G13 matrix unit reads its A/B
+ * fragments while it writes D: with fp16 A and fp32 D the pair D = A*B+C where
+ * D overlaps A returned zeros (16x16 f16/f32 coopmat, sub-block 3). Apple's
+ * compiler only ever emits D == C, so keep A and B live until D is placed.
+ */
+static bool
+early_killable(const agx_instr *I, unsigned s)
+{
+ if (!I->src[s].kill || I->src[s].memory)
+ return false;
+
+ if ((I->op == AGX_OPCODE_SIMD_MATRIX_FMADD32 ||
+ I->op == AGX_OPCODE_SIMD_MATRIX_FMADD16) &&
+ s < 2)
+ return false;
+
+ return true;
+}
+
static void
try_kill_early_sources(struct ra_ctx *rctx, const agx_instr *I,
unsigned first_source, unsigned last_source,
@@ -1034,7 +1054,7 @@ try_kill_early_sources(struct ra_ctx *rctx, const agx_instr *I,
return;
for (unsigned s = first_source; s <= last_source; ++s) {
- if (I->src[s].kill && !I->src[s].memory) {
+ if (early_killable(I, s)) {
kill_source(rctx, I, s);
rctx->early_killed = true;
I->src[s].kill = false;
@@ -1132,7 +1152,7 @@ agx_ra_assign_local(struct ra_ctx *rctx)
unsigned start = 0;
agx_foreach_ssa_src(I, s) {
- if (I->src[s].kill && !I->src[s].memory) {
+ if (early_killable(I, s)) {
unsigned reg = rctx->ssa_to_reg[I->src[s].value];
if (start == end || end != reg) {
diff --git a/src/asahi/vulkan/hk_physical_device.c b/src/asahi/vulkan/hk_physical_device.c
index 6f6d99599c77..0ecd271033ce 100644
--- a/src/asahi/vulkan/hk_physical_device.c
+++ b/src/asahi/vulkan/hk_physical_device.c
@@ -1537,20 +1537,33 @@ hk_GetPhysicalDeviceCooperativeMatrixPropertiesKHR(
pPropertyCount);
if (getenv("AGX_SIMDMAT")) {
- /* Hardware 8x8x8 fp32 matrix MAC (simd_matrix_fmadd32). */
- vk_outarray_append_typed(VkCooperativeMatrixPropertiesKHR, &out, p)
- {
- p->sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR;
- p->pNext = NULL;
- p->MSize = 8;
- p->NSize = 8;
- p->KSize = 8;
- p->AType = VK_COMPONENT_TYPE_FLOAT32_KHR;
- p->BType = VK_COMPONENT_TYPE_FLOAT32_KHR;
- p->CType = VK_COMPONENT_TYPE_FLOAT32_KHR;
- p->ResultType = VK_COMPONENT_TYPE_FLOAT32_KHR;
- p->saturatingAccumulation = VK_FALSE;
- p->scope = VK_SCOPE_SUBGROUP_KHR;
+ /* G13 hardware 8x8x8 matrix MAC (simd_matrix_fmadd16/32). The unit
+ * takes fp16 or fp32 A/B and accumulates in fp32 (fmadd32) or fp16
+ * (fmadd16); 16x16x16 is a 2x2x2 tiling of 8x8x8 ops in
+ * agx_nir_lower_simdmat.c. Every shape here is GPU-verified exact.
+ */
+ static const struct {
+ VkComponentTypeKHR ab, c;
+ } types[] = {
+ {VK_COMPONENT_TYPE_FLOAT32_KHR, VK_COMPONENT_TYPE_FLOAT32_KHR},
+ {VK_COMPONENT_TYPE_FLOAT16_KHR, VK_COMPONENT_TYPE_FLOAT32_KHR},
+ {VK_COMPONENT_TYPE_FLOAT16_KHR, VK_COMPONENT_TYPE_FLOAT16_KHR},
+ };
+ for (unsigned n = 8; n <= 16; n *= 2) {
+ for (unsigned t = 0; t < ARRAY_SIZE(types); t++) {
+ vk_outarray_append_typed(VkCooperativeMatrixPropertiesKHR, &out, p)
+ {
+ p->MSize = n;
+ p->NSize = n;
+ p->KSize = n;
+ p->AType = types[t].ab;
+ p->BType = types[t].ab;
+ p->CType = types[t].c;
+ p->ResultType = types[t].c;
+ p->saturatingAccumulation = VK_FALSE;
+ p->scope = VK_SCOPE_SUBGROUP_KHR;
+ }
+ }
}
}
diff --git a/src/compiler/nir/nir_intrinsics.py b/src/compiler/nir/nir_intrinsics.py
index b6c434d8f1b9..9e220c0fee3f 100644
--- a/src/compiler/nir/nir_intrinsics.py
+++ b/src/compiler/nir/nir_intrinsics.py
@@ -2472,11 +2472,13 @@ def store(name, srcs, indices=[], flags=[]):
# Equivalent to popcount(ballot(true))
intrinsic("load_active_subgroup_count_agx", dest_comp=1, flags=[CAN_ELIMINATE])
-# Apple G13 hardware matrix multiply-accumulate (simd_matrix_fmadd32). One
+# Apple G13 hardware matrix multiply-accumulate (simd_matrix_fmadd16/32). One
# SIMD-group-wide 8x8x8 tile: D = A*B + C. Per-invocation each operand is the
-# lane's 2 fragment elements (a register pair). src order: A, B, C.
+# lane's 2 fragment elements (a register pair). src order: A, B, C. The dest
+# bit size is the accumulator precision (16 -> fmadd16, 32 -> fmadd32) and
+# must match C; A and B may independently be 16- or 32-bit.
intrinsic("simd_matrix_fmadd_agx", src_comp=[2, 2, 2], dest_comp=2,
- bit_sizes=[32], flags=[CAN_ELIMINATE])
+ bit_sizes=[16, 32], flags=[CAN_ELIMINATE])
# Like ballot() but only within a quad.
intrinsic("quad_ballot_agx", src_comp=[1], dest_comp=1, flags=[CAN_ELIMINATE])
From 600f1670f833f4664fd8e691c62a96746c6cdda6 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Tue, 8 Sep 2026 07:56:37 -0500
Subject: [PATCH 06/16] agx: keep the mod-32 shift count when fusing ushr+iand
into ubfe
Bug: (word >> s) & mask with a non-constant s returns zero on Honeykrisp
whenever s has any bit set above bit 4 (s >= 32). mlx-omarchy hit this as
the "dynamic shift-then-mask byte extraction" miscompile family: bool
scatter dropping writes across word lanes, select picking the wrong
operand under broadcast conditions, reduce_general's load_truthy reading
zero, masked_scatter's scan undercounting (s_scan[255] = 32, not 128).
Root cause: agx_nir_fuse_algebraic_late rewrites
iand(ushr(a, b), (1 << bits) - 1) -> ubitfield_extract(a, b, bits)
NIR's ushr uses only the low 5 bits of its count, and by this point the
main nir_opt_algebraic has already relied on that to strip explicit masks
(ushr(a, iand(b, 31)) -> ushr(a, b); ushr(a, ishl(iand(b, 3), 3)) ->
ushr(a, ishl(b, 3))). ubitfield_extract has no such wrap: an offset >= 32
is undefined in NIR and bfeil returns 0 for it on the hardware.
lower_sm5_shift, which masks the count of every surviving shift, runs
after the fusion and never sees the ushr it consumed.
Fix: fuse directly only when the count is a constant below 32; otherwise
fuse with an explicit iand(b, 31) on the offset. Same instruction count as
the unfused shift+and, and the mask CSEs with one the shader already had.
Verification on Apple M1 (G13G B1), standalone Vulkan compute reproducer
(20 variants, 1024 lanes each, host-compared):
stock Mesa 26.1.7: 9 of 16 base variants FAIL (byte_dyn,
byte_dyn_const, byte_dyn_gid, scan_dyn, scan_dyn_const, scan_dyn_rand,
select_bcast, select_dyn, shift_dyn_mask), plus all 4 boundary
variants FAIL (count exactly 32: 1016/1024 wrong; 63: 516/1024;
uniform 40: 1020/1024; divergent 8/40: 509/1024, only the 40 lanes).
this commit: 20 of 20 PASS, 0 mismatches.
mlx-omarchy suites, stock vs this commit (AGX_SIMDMAT=1), identical:
primitive 98/98 (2699373 assertions), reduce_ops 34/34 (7152),
runtime 40/40 (6307), indexing_ops 53/53 (4093), select_layout 11/11
(1908), scatter_determinism 21/21 (127).
---
src/asahi/compiler/agx_nir_algebraic.py | 11 +++++++++--
1 file changed, 9 insertions(+), 2 deletions(-)
diff --git a/src/asahi/compiler/agx_nir_algebraic.py b/src/asahi/compiler/agx_nir_algebraic.py
index b1e036459d21..311062e38feb 100644
--- a/src/asahi/compiler/agx_nir_algebraic.py
+++ b/src/asahi/compiler/agx_nir_algebraic.py
@@ -141,11 +141,18 @@
('extr_agx', a, b, start, 0)),
])
+# ushr uses only the low 5 bits of its count (SM5 semantics), and the main
+# nir_opt_algebraic has already stripped explicit masks on that basis.
+# ubitfield_extract (bfeil) does not wrap: an offset >= 32 reads as zero on
+# the hardware. lower_sm5_shift runs after this pass, so a ushr fused here is
+# never masked by it; keep the modulo unless the count is a constant in range.
fuse_ubfe = []
for bits in range(1, 32):
fuse_ubfe.extend([
- (('iand', ('ushr', 'a@32', b), (1 << bits) - 1),
- ('ubitfield_extract', a, b, bits))
+ (('iand', ('ushr', 'a@32', '#b(is_ult_32)'), (1 << bits) - 1),
+ ('ubitfield_extract', a, b, bits)),
+ (('iand', ('ushr', 'a@32', 'b(is_not_const)'), (1 << bits) - 1),
+ ('ubitfield_extract', a, ('iand', b, 31), bits)),
])
# (x * y) + s = (x * y) + (s << 0)
From caf7649e9ce75d7e5849feca7c05dc19384e45d2 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Tue, 8 Sep 2026 08:17:37 -0500
Subject: [PATCH 07/16] agx: correctly rounded fdiv/frcp, faithful log/log2,
exact sin/cos reduction
Three fp32 rounding defects on Honeykrisp (Apple M1 G13G, Mesa 26.1.7)
that mlx-omarchy patches around in its shaders:
a) GLSL "/" was lowered by nir_opt_algebraic (lower_fdiv) to
a * frcp(b), and frcp was only Newton-refined when the instruction
was `precise`. Even with an exact reciprocal a * RN(1/b) is not
RN(a/b): 80/25 and 10/25 come out one ulp low, 6.55859375/0.030016447
one ulp high. Measured on 1008 pairs: 299 mismatches (298 at 1 ulp,
1 at 2 ulp), 1/b: 204 mismatches; 1M random pairs: 305792 mismatches.
fma() itself is IEEE (0/1000 mismatches), the mlx-omarchy note
blaming fma contraction was a misdiagnosis.
b) log(x) was flog2(x) * ln2 on the hardware log2, which is only
absolutely accurate: log(3) = 1.0986121892929077 instead of
1.0986123085021973, 2792/4107 mismatches on [0.5, 64] with errors up
to 357 ulp near 1 (log(1.0000001) = 1.65e-7 vs 1.0e-7); log2 had
2905/4107 mismatches.
c) sin/cos reduced the argument as fract(x / 2pi) in fp32, so the
reduced angle lost ulp(x) of precision: sin(3.0) 14 ulp off,
sin(1e5) 1.06e-3 absolute error, sin(1e8) = 0.0, sin(-1e-3) 571 ulp
off, sin(pi_f32) = 0, and the hardware sin_pt_2 core itself is 2 to
3 ulp off (cos(1e-3) 3 ulp).
The lowering moves to a new the AGX compiler pass file agx_nir_lower_math.c,
built only from the exact fp32 fma and marked no-fast-math so
nir_opt_algebraic cannot reassociate it:
- lower_fdiv is dropped from agx_nir_options; fdiv32 becomes Markstein
division (y = refined rcp, q = a y, r = fma(-b, q, a), q + r y) with
the NaN/zero guard keeping q for b = 0, overflow and signed zero.
frcp32 is always refined. fdiv16 keeps a * frcp(b).
- fmul(flog2(x), ln2) as emitted by nir_flog, and any remaining flog2,
become a double-float evaluation of log(m) = 2s + 2s^3 P(s^2) with
s = (m-1)/(m+1) kept as a pair, m folded into [sqrt(1/2), sqrt(2)),
e ln2 split hi/lo and Fast2Sum accumulation (about 40 bits before the
final rounding). Zero/subnormal/negative/inf/NaN are handled by bit
tests, not the hardware log2.
- fsin/fcos: Cody-Waite in radians with pi/2 in three pieces, the k C2
product and its subtraction as TwoProduct/TwoSum so the reduced
argument is a pair accurate to an ulp, for |x| < 2^22; an integer
Payne-Hanek (24-bit mantissa times a 128-bit window of 2/pi, 64
fraction bits kept) for |x| >= 2^22 up to FLT_MAX. The quadrant then
selects a Taylor polynomial in quadrant units (sin degree 13, cos
degree 14, leading terms as pairs, low part of the argument fed to
the derivative). The hardware sin_pt_1/sin_pt_2 core (nir_op_fsin_agx)
is no longer used by the lowering. fp16 sin/cos go through the fp32
path.
Verification (receipts/2026-09-08-precise-math.json, standalone Vulkan
compute reproducer in receipts/precise-math/):
- division: 1008-pair table incl. 80/25, 10/25, 1000 random: 0
mismatches (stock 299); 1/b 0 (stock 204); 1M random pairs 0
mismatches for both a/b and 1/b; b = 0, inf, NaN, overflow, flushed
subnormal and -0/5 = -0 match host IEEE.
- log: 4107 points on [0.5, 64] plus pinned exact cases: 0 mismatches
(stock 2792), log(3) = 1.0986123085021973; log2 0 (stock 2905);
1M random on [1e-37, 1e38]: log 3 at 1 ulp, log2 14 at 1 ulp,
rest exact.
- sin/cos at 1e5, 5e5, 1e6, 5e6, 1e8: all 0 ulp from the double
reference except cos(5e6) at 1 ulp (stock: 2.8e5 to 1.06e9 ulp);
1310 mixed points max 1 ulp; 1M random over [1e-3, 1e9] and [-10, 10]:
sin 42346 and cos 41184 at 1 ulp, rest exact, max 1 ulp (stock
>= 8 ulp for 627689 / 512816 of them).
- mlx-omarchy suites on the same Release tree, stock vs fixed:
primitive 98/98, eq_math 7/7, fast_ops 30/30, runtime 40/40,
reduce_ops 34/34 both; complex_ops 30/33 both with the same three
pre-existing cases (failed assertions 31 -> 29). The upstream
ops_tests log(3) pin through libmlx: FAIL on stock, PASS on fixed,
as do 80/25, 10/25 and sin/cos up to the 1e5 limit mlx enforces.
- timing: 16M-element elementwise dispatch is memory bound at ~2.45 ms
on both drivers (sin 2.474 -> 2.463 ms median, within spread). A
32-deep in-register chain shows the ALU cost: sin 6.64 -> 39.6 ms,
log 2.57 -> 34.8 ms, div 2.84 -> 7.24 ms, mul unchanged.
---
src/asahi/compiler/agx_compile.c | 100 +----
src/asahi/compiler/agx_compile.h | 1 -
src/asahi/compiler/agx_nir.h | 3 +
src/asahi/compiler/agx_nir_lower_math.c | 567 ++++++++++++++++++++++++
src/asahi/compiler/meson.build | 1 +
5 files changed, 577 insertions(+), 95 deletions(-)
create mode 100644 src/asahi/compiler/agx_nir_lower_math.c
diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index bc2468809444..add80493133d 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -2741,46 +2741,6 @@ agx_calc_stats(agx_context *ctx, unsigned size, struct agx2_stats *stats)
}
}
-static bool
-agx_lower_sincos_filter(const nir_instr *instr, UNUSED const void *_)
-{
- if (instr->type != nir_instr_type_alu)
- return false;
-
- nir_alu_instr *alu = nir_instr_as_alu(instr);
- return alu->op == nir_op_fsin || alu->op == nir_op_fcos;
-}
-
-/* Sine and cosine are implemented via the sin_pt_1 and sin_pt_2 opcodes for
- * heavy lifting. sin_pt_2 implements sinc in the first quadrant, expressed in
- * turns (sin (tau x) / x), while sin_pt_1 implements a piecewise sign/offset
- * fixup to transform a quadrant angle [0, 4] to [-1, 1]. The NIR opcode
- * fsin_agx models the fixup, sinc, and multiply to obtain sine, so we just
- * need to change units from radians to quadrants modulo turns. Cosine is
- * implemented by shifting by one quadrant: cos(x) = sin(x + tau/4).
- */
-
-static nir_def *
-agx_lower_sincos_impl(struct nir_builder *b, nir_instr *instr, UNUSED void *_)
-{
- nir_alu_instr *alu = nir_instr_as_alu(instr);
- nir_def *x = nir_mov_alu(b, alu->src[0], 1);
- nir_def *turns = nir_fmul_imm(b, x, M_1_PI * 0.5f);
-
- if (alu->op == nir_op_fcos)
- turns = nir_fadd_imm(b, turns, 0.25f);
-
- nir_def *quadrants = nir_fmul_imm(b, nir_ffract(b, turns), 4.0);
- return nir_fsin_agx(b, quadrants);
-}
-
-static bool
-agx_lower_sincos(nir_shader *shader)
-{
- return nir_shader_lower_instructions(shader, agx_lower_sincos_filter,
- agx_lower_sincos_impl, NULL);
-}
-
static bool
agx_lower_front_face(struct nir_builder *b, nir_intrinsic_instr *intr,
UNUSED void *data)
@@ -3551,56 +3511,6 @@ agx_compile_function_nir(nir_shader *nir, nir_function_impl *impl,
return offset;
}
-/*
- * The hardware frcp instruction is sometimes off by 1 ULP. For correctly
- * rounded frcp, a refinement step is required. This routine has been
- * exhaustively tested with a modified math_bruteforce.
- *
- * While Khronos APIs allow 2.5 ULP error for divides, nir_lower_idiv relies on
- * correctly rounded frcp. This is therefore load bearing for integer division
- * on all APIs.
- */
-static nir_def *
-libagx_frcp(nir_builder *b, nir_def *x)
-{
- nir_def *u = nir_frcp(b, x);
-
- /* Do 1 Newton-Raphson refinement step.
- *
- * Define f(u) = xu - 1. Then f(u) = 0 iff u = 1/x. Newton's method gives:
- *
- * u_2 = u - f(u) / f'(u) = u - (xu - 1) / x
- *
- * Our original guess is close, so we approximate (1 / x) by u:
- *
- * u_2 = u - u(xu - 1) = u + u(1 - xu)
- * = fma(fma(-x, u, 1), u, u)
- */
- nir_def *one = nir_imm_float(b, 1.0);
- nir_def *u_2 = nir_ffma(b, nir_ffma(b, nir_fneg(b, x), u, one), u, u);
-
- /* If the original value was infinite, frcp will generate the correct zero.
- * However, the Newton-Raphson step would multiply 0 * Inf and get a NaN. So
- * skip the refinement step for infinite inputs. We do this backwards,
- * checking whether the refined result is NaN, since we can implement this
- * check in a single fcmpsel instruction. The other case where the refinement
- * is NaN is a NaN input, in which skipping refinement is acceptable.
- */
- return nir_bcsel(b, nir_fisnan(b, u_2), u, u_2);
-}
-
-static bool
-agx_nir_lower_fdiv(nir_builder *b, nir_alu_instr *alu, void *_)
-{
- if (alu->op != nir_op_frcp || !nir_alu_instr_is_exact(alu) ||
- alu->def.bit_size != 32)
- return false;
-
- b->cursor = nir_before_instr(&alu->instr);
- nir_def_replace(&alu->def, libagx_frcp(b, nir_ssa_for_alu_src(b, alu, 0)));
- return true;
-}
-
/* Preprocess NIR independent of shader state */
void
agx_preprocess_nir(nir_shader *nir)
@@ -3639,7 +3549,7 @@ agx_preprocess_nir(nir_shader *nir)
NIR_PASS(_, nir, nir_lower_alu);
NIR_PASS(_, nir, nir_lower_load_const_to_scalar);
NIR_PASS(_, nir, nir_lower_flrp, 16 | 32 | 64, false);
- NIR_PASS(_, nir, agx_lower_sincos);
+ NIR_PASS(_, nir, agx_nir_lower_sincos);
NIR_PASS(_, nir, nir_shader_intrinsics_pass, agx_lower_front_face,
nir_metadata_control_flow, NULL);
NIR_PASS(_, nir, agx_nir_lower_subgroups);
@@ -3662,9 +3572,11 @@ agx_preprocess_nir(nir_shader *nir)
NIR_PASS(_, nir, nir_lower_idiv, &idiv_options);
- /* Has to run after nir_lower_idiv */
- NIR_PASS(_, nir, nir_shader_alu_pass, agx_nir_lower_fdiv,
- nir_metadata_control_flow, NULL);
+ /* Has to run after nir_lower_idiv, and before the log lowering so the
+ * reciprocals it builds are not refined twice.
+ */
+ NIR_PASS(_, nir, agx_nir_lower_fdiv);
+ NIR_PASS(_, nir, agx_nir_lower_log);
NIR_PASS(_, nir, nir_opt_deref);
NIR_PASS(_, nir, nir_lower_vars_to_ssa);
diff --git a/src/asahi/compiler/agx_compile.h b/src/asahi/compiler/agx_compile.h
index c8ffaa5db23b..58706acd2f77 100644
--- a/src/asahi/compiler/agx_compile.h
+++ b/src/asahi/compiler/agx_compile.h
@@ -347,7 +347,6 @@ agx_round_registers(unsigned halfregs)
}
static const nir_shader_compiler_options agx_nir_options = {
- .lower_fdiv = true,
.float_mul_add16 = nir_float_muladd_support_has_ffma | nir_float_muladd_support_fuse,
.float_mul_add32 = nir_float_muladd_support_has_ffma | nir_float_muladd_support_fuse,
.lower_flrp16 = true,
diff --git a/src/asahi/compiler/agx_nir.h b/src/asahi/compiler/agx_nir.h
index c03715027e5c..1fead739f933 100644
--- a/src/asahi/compiler/agx_nir.h
+++ b/src/asahi/compiler/agx_nir.h
@@ -22,6 +22,9 @@ bool agx_nir_lower_layer(struct nir_shader *s);
bool agx_nir_lower_clip_distance(struct nir_shader *s);
bool agx_nir_lower_subgroups(struct nir_shader *s);
bool agx_nir_lower_fminmax(struct nir_shader *s);
+bool agx_nir_lower_fdiv(struct nir_shader *s);
+bool agx_nir_lower_log(struct nir_shader *s);
+bool agx_nir_lower_sincos(struct nir_shader *s);
bool agx_nir_lower_texture_early(struct nir_shader *s, bool support_lod_bias);
bool agx_nir_lower_texture(struct nir_shader *s);
diff --git a/src/asahi/compiler/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c
new file mode 100644
index 000000000000..3e9b923264f7
--- /dev/null
+++ b/src/asahi/compiler/agx_nir_lower_math.c
@@ -0,0 +1,567 @@
+/*
+ * Copyright 2021 Alyssa Rosenzweig
+ * Copyright 2026 Joshua Warren
+ * SPDX-License-Identifier: MIT
+ */
+
+#include
+#include "compiler/nir/nir.h"
+#include "compiler/nir/nir_builder.h"
+#include "compiler/nir/nir_builtin_builder.h"
+#include "agx_nir.h"
+
+/*
+ * AGX transcendental lowering with fp32 rounding contracts.
+ *
+ * The hardware gives an approximate rcp (off by 1 ulp), an approximate log2
+ * (absolute error only, so relative error explodes near 1) and a sin unit that
+ * takes an argument in quadrants. Everything here is built from the exact fp32
+ * fma, so the sequences are marked no-fast-math to keep nir_opt_algebraic from
+ * reassociating or contracting them.
+ */
+
+/*
+ * Correctly rounded 1/x. One Newton-Raphson step on the hardware rcp:
+ *
+ * u_2 = u + u(1 - xu) = fma(fma(-x, u, 1), u, u)
+ *
+ * exhaustively verified upstream with a modified math_bruteforce. If u is
+ * infinite (x zero or flushed), the refinement is NaN, so keep u.
+ */
+static nir_def *
+rcp_rn(nir_builder *b, nir_def *x)
+{
+ nir_def *u = nir_frcp(b, x);
+ nir_def *one = nir_imm_float(b, 1.0);
+ nir_def *u_2 = nir_ffma(b, nir_ffma(b, nir_fneg(b, x), u, one), u, u);
+ return nir_bcsel(b, nir_fisnan(b, u_2), u, u_2);
+}
+
+/*
+ * Correctly rounded a/b (Markstein). With y = RN(1/b) and q = RN(a y), the
+ * residual r = a - b q is exact in an fma and q + r y rounds to RN(a/b) for
+ * normal operands. If the refinement produces NaN (b zero, q infinite), q
+ * already holds the IEEE result for that case. A zero q is also final (a zero
+ * or the quotient flushed) and keeps its sign, which the refinement would
+ * lose.
+ */
+static nir_def *
+div_rn(nir_builder *b, nir_def *a, nir_def *d)
+{
+ nir_def *y = rcp_rn(b, d);
+ nir_def *q = nir_fmul(b, a, y);
+ nir_def *r = nir_ffma(b, nir_fneg(b, d), q, a);
+ nir_def *q_2 = nir_ffma(b, r, y, q);
+ nir_def *keep = nir_ior(b, nir_fisnan(b, q_2), nir_feq_imm(b, q, 0.0));
+ return nir_bcsel(b, keep, q, q_2);
+}
+
+static bool
+lower_fdiv(nir_builder *b, nir_alu_instr *alu, void *_)
+{
+ if (alu->op != nir_op_frcp && alu->op != nir_op_fdiv)
+ return false;
+
+ if (alu->def.bit_size == 64)
+ return false;
+
+ b->cursor = nir_before_instr(&alu->instr);
+ b->fp_math_ctrl = nir_fp_no_fast_math;
+
+ nir_def *res;
+ if (alu->op == nir_op_frcp) {
+ if (alu->def.bit_size != 32)
+ return false;
+
+ res = rcp_rn(b, nir_ssa_for_alu_src(b, alu, 0));
+ } else {
+ nir_def *a = nir_ssa_for_alu_src(b, alu, 0);
+ nir_def *d = nir_ssa_for_alu_src(b, alu, 1);
+
+ if (alu->def.bit_size == 32)
+ res = div_rn(b, a, d);
+ else
+ res = nir_fmul(b, a, nir_frcp(b, d));
+ }
+
+ nir_def_replace(&alu->def, res);
+ return true;
+}
+
+bool
+agx_nir_lower_fdiv(nir_shader *s)
+{
+ return nir_shader_alu_pass(s, lower_fdiv, nir_metadata_control_flow, NULL);
+}
+
+/*
+ * Natural log and log2 as double-float (fp32 pair) arithmetic.
+ *
+ * x = 2^e m with m in [sqrt(1/2), sqrt(2)). With s = (m - 1) / (m + 1),
+ *
+ * log(m) = 2s + 2s^3 (1/3 + s^2/5 + s^4/7 + ...)
+ *
+ * s is kept as a pair (s, s_lo) from one correctly rounded division plus the
+ * fma residual, the cubic term is kept as a pair, and the terms are summed
+ * with Fast2Sum so the value handed to the final rounding carries about 40
+ * bits. That makes the exact cases (powers of two, log(3)) come out correctly
+ * rounded and everything else faithful, without any table.
+ */
+struct log_parts {
+ nir_def *e; /* exponent as float */
+ nir_def *s; /* 2s = leading term, as a pair (s, s_lo) */
+ nir_def *s_lo;
+ nir_def *t_hi; /* 2 s^3 P(s^2), as a pair */
+ nir_def *t_lo;
+};
+
+static nir_def *
+fast2sum_lo(nir_builder *b, nir_def *a, nir_def *bb, nir_def *sum)
+{
+ /* Requires |a| >= |bb| or a == 0 */
+ return nir_fsub(b, bb, nir_fsub(b, sum, a));
+}
+
+static struct log_parts
+log_core(nir_builder *b, nir_def *x)
+{
+ nir_def *bits = x;
+ nir_def *e = nir_iadd_imm(b, nir_ushr_imm(b, bits, 23), -127);
+ nir_def *m = nir_ior_imm(b, nir_iand_imm(b, bits, 0x7fffff), 0x3f800000);
+
+ /* Fold m into [sqrt(1/2), sqrt(2)) so the series converges fast and
+ * e ln2 never cancels more than one bit against log(m).
+ */
+ nir_def *big = nir_fgt_imm(b, m, 1.41421353816986083984375);
+ m = nir_bcsel(b, big, nir_fmul_imm(b, m, 0.5), m);
+ e = nir_iadd(b, e, nir_b2i32(b, big));
+
+ nir_def *one = nir_imm_float(b, 1.0);
+ nir_def *two = nir_imm_float(b, 2.0);
+ nir_def *f = nir_fsub(b, m, one); /* exact */
+ nir_def *d = nir_fadd(b, two, f);
+ nir_def *d_lo = fast2sum_lo(b, two, f, d);
+
+ /* s = f / (d + d_lo) as a pair */
+ nir_def *y = rcp_rn(b, d);
+ nir_def *s = nir_fmul(b, f, y);
+ nir_def *r = nir_ffma(b, nir_fneg(b, s), d, f);
+ s = nir_ffma(b, r, y, s);
+ r = nir_ffma(b, nir_fneg(b, s), d, f);
+ nir_def *s_lo =
+ nir_fmul(b, nir_fsub(b, r, nir_fmul(b, s, d_lo)), y);
+
+ /* u = s^3 as a pair, including the 3 s^2 s_lo cross term */
+ nir_def *w = nir_fmul(b, s, s);
+ nir_def *w_lo = nir_ffma(b, s, s, nir_fneg(b, w));
+ nir_def *u = nir_fmul(b, s, w);
+ nir_def *u_lo = nir_fadd(b, nir_ffma(b, s, w, nir_fneg(b, u)),
+ nir_fmul(b, s, w_lo));
+ u_lo = nir_ffma(b, nir_fmul_imm(b, s_lo, 3.0), w, u_lo);
+
+ /* P(w) = 1/3 + w Q(w) as a pair, 1/3 split hi/lo */
+ nir_def *q = nir_imm_float(b, 1.0 / 13.0);
+ q = nir_ffma(b, q, w, nir_imm_float(b, 1.0 / 11.0));
+ q = nir_ffma(b, q, w, nir_imm_float(b, 1.0 / 9.0));
+ q = nir_ffma(b, q, w, nir_imm_float(b, 1.0 / 7.0));
+ q = nir_ffma(b, q, w, nir_imm_float(b, 1.0 / 5.0));
+ nir_def *third = nir_imm_float(b, 0.3333333432674408);
+ nir_def *v = nir_ffma(b, w, q, nir_imm_float(b, -9.934107758624577e-09));
+ nir_def *p = nir_fadd(b, third, v);
+ nir_def *p_lo = fast2sum_lo(b, third, v, p);
+
+ nir_def *th = nir_fmul(b, u, p);
+ nir_def *tl = nir_fadd(b, nir_fadd(b, nir_ffma(b, u, p, nir_fneg(b, th)),
+ nir_fmul(b, u, p_lo)),
+ nir_fmul(b, u_lo, p));
+
+ return (struct log_parts){
+ .e = nir_i2f32(b, e),
+ .s = s,
+ .s_lo = s_lo,
+ .t_hi = nir_fmul(b, two, th),
+ .t_lo = nir_fmul(b, two, tl),
+ };
+}
+
+/* Positive, finite, not flushed: the only inputs the series handles. */
+static nir_def *
+is_positive_normal(nir_builder *b, nir_def *x)
+{
+ return nir_ult_imm(b, nir_iadd_imm(b, x, -0x00800000),
+ 0x7f7fffff - 0x00800000 + 1);
+}
+
+/*
+ * Everything else: +-0 and flushed subnormals give -inf, negatives and NaN
+ * give NaN, +inf gives +inf. Same answer for log and log2.
+ */
+static nir_def *
+log_special(nir_builder *b, nir_def *x)
+{
+ nir_def *zero = nir_ult_imm(b, nir_iand_imm(b, x, 0x7fffffff), 0x00800000);
+ nir_def *nan = nir_ugt_imm(b, x, 0x7f800000);
+ return nir_bcsel(b, zero, nir_imm_float(b, -INFINITY),
+ nir_bcsel(b, nan, nir_imm_float(b, NAN),
+ nir_imm_float(b, INFINITY)));
+}
+
+static nir_def *
+soft_log(nir_builder *b, nir_def *x)
+{
+ struct log_parts l = log_core(b, x);
+
+ /* ln2 split so e * LN2_HI is exact for |e| < 2^8 */
+ nir_def *a = nir_fmul_imm(b, l.e, 0.693145751953125);
+ nir_def *bb = nir_fmul_imm(b, l.s, 2.0);
+ nir_def *h = nir_fadd(b, a, bb);
+ nir_def *lo = fast2sum_lo(b, a, bb, h);
+ nir_def *h2 = nir_fadd(b, h, l.t_hi);
+ nir_def *lo2 = fast2sum_lo(b, h, l.t_hi, h2);
+
+ nir_def *small = nir_fadd(b, lo, lo2);
+ small = nir_fadd(b, small, l.t_lo);
+ small = nir_fadd(b, small, nir_fmul_imm(b, l.s_lo, 2.0));
+ small = nir_ffma(b, l.e, nir_imm_float(b, 1.428606765330187e-06), small);
+ return nir_fadd(b, h2, small);
+}
+
+static nir_def *
+soft_log2(nir_builder *b, nir_def *x)
+{
+ struct log_parts l = log_core(b, x);
+
+ /* (2s + 2s_lo + t) / ln2 with 1/ln2 as a pair */
+ nir_def *il2 = nir_imm_float(b, 1.4426950216293335);
+ nir_def *il2_lo = nir_imm_float(b, 1.925963033500011e-08);
+ nir_def *bb = nir_fmul_imm(b, l.s, 2.0);
+ nir_def *ph = nir_fmul(b, bb, il2);
+ nir_def *pl = nir_ffma(b, bb, il2, nir_fneg(b, ph));
+ pl = nir_ffma(b, bb, il2_lo, pl);
+ nir_def *th = nir_fmul(b, l.t_hi, il2);
+ nir_def *tl = nir_ffma(b, l.t_hi, il2, nir_fneg(b, th));
+ nir_def *rest = nir_fadd(b, nir_fmul_imm(b, l.s_lo, 2.0), l.t_lo);
+ pl = nir_fadd(b, nir_ffma(b, rest, il2, pl), tl);
+
+ nir_def *h = nir_fadd(b, l.e, ph);
+ nir_def *lo = fast2sum_lo(b, l.e, ph, h);
+ nir_def *h2 = nir_fadd(b, h, th);
+ nir_def *lo2 = fast2sum_lo(b, h, th, h2);
+ return nir_fadd(b, h2, nir_fadd(b, nir_fadd(b, lo, lo2), pl));
+}
+
+static const uint32_t LN2_F32_BITS = 0x3f317218;
+
+static bool
+src_is_ln2(nir_alu_instr *alu, unsigned i)
+{
+ nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i);
+ return nir_scalar_is_const(s) && nir_scalar_as_uint(s) == LN2_F32_BITS;
+}
+
+static bool
+src_is_flog2(nir_alu_instr *alu, unsigned i)
+{
+ nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i);
+ return nir_scalar_is_alu(s) && nir_scalar_alu_op(s) == nir_op_flog2;
+}
+
+/* nir_flog builds fmul(flog2(x), ln2); catch that before the flog2 goes. */
+static bool
+lower_log(nir_builder *b, nir_alu_instr *alu, void *_)
+{
+ if (alu->op != nir_op_fmul || alu->def.bit_size != 32 ||
+ alu->def.num_components != 1)
+ return false;
+
+ unsigned li;
+ if (src_is_flog2(alu, 0) && src_is_ln2(alu, 1))
+ li = 0;
+ else if (src_is_flog2(alu, 1) && src_is_ln2(alu, 0))
+ li = 1;
+ else
+ return false;
+
+ nir_scalar log = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), li);
+ nir_scalar xs = nir_scalar_chase_alu_src(log, 0);
+
+ b->cursor = nir_before_instr(&alu->instr);
+ b->fp_math_ctrl = nir_fp_no_fast_math;
+ nir_def *x = nir_channel(b, xs.def, xs.comp);
+ nir_def *res = nir_bcsel(b, is_positive_normal(b, x), soft_log(b, x),
+ log_special(b, x));
+ nir_def_replace(&alu->def, res);
+
+ if (nir_def_is_unused(log.def))
+ nir_instr_remove(nir_def_instr(log.def));
+
+ return true;
+}
+
+static bool
+lower_log2(nir_builder *b, nir_alu_instr *alu, void *_)
+{
+ if (alu->op != nir_op_flog2 || alu->def.bit_size != 32)
+ return false;
+
+ b->cursor = nir_before_instr(&alu->instr);
+ b->fp_math_ctrl = nir_fp_no_fast_math;
+ nir_def *x = nir_ssa_for_alu_src(b, alu, 0);
+ nir_def *res = nir_bcsel(b, is_positive_normal(b, x), soft_log2(b, x),
+ log_special(b, x));
+ nir_def_replace(&alu->def, res);
+ return true;
+}
+
+bool
+agx_nir_lower_log(nir_shader *s)
+{
+ bool progress =
+ nir_shader_alu_pass(s, lower_log, nir_metadata_control_flow, NULL);
+ progress |=
+ nir_shader_alu_pass(s, lower_log2, nir_metadata_control_flow, NULL);
+ return progress;
+}
+
+/*
+ * Sine and cosine. Range reduction produces k = round(x 2/pi) and the fraction
+ * fq = x 2/pi - k in [-1/2, 1/2] quadrants; the quadrant then picks a Taylor
+ * polynomial for sin(pi/2 fq) or cos(pi/2 fq) and a sign. The hardware sin
+ * unit (sin_pt_1/sin_pt_2, nir_op_fsin_agx) is not used: it is a few ulp off,
+ * which is more than the reduction now loses.
+ *
+ * |x| < 2^22: Cody-Waite in radians with pi/2 in three fp32 pieces. The first
+ * fma is exact, the remaining rounding is tracked as a low part so fq is
+ * within an ulp of the true reduced argument.
+ *
+ * |x| >= 2^22: Payne-Hanek. Multiply the 24-bit mantissa by a 128-bit window
+ * of 2/pi aligned to the exponent; the integer part mod 4 is the quadrant and
+ * the top 64 fraction bits give fq. fp32 inputs never reduce below ~2^-30 of a
+ * quadrant, so 64 fraction bits leave more than 24 significant bits.
+ */
+static const double TWO_OVER_PI = 0.6366197466850281; /* fp32(2/pi) */
+static const double TWO_OVER_PI_LO = 2 / M_PI - 0.6366197466850281;
+static const double PIO2_1 = 1.5707963705062866;
+static const double PIO2_2 = -4.371138828673793e-08;
+static const double PIO2_3 = -1.7151245100058819e-15;
+
+/* Bits of 2/pi, 32 per word, preceded by three zero words. */
+static const uint32_t two_over_pi_bits[10] = {
+ 0, 0, 0, 0xa2f9836e, 0x4e441529, 0xfc2757d1, 0xf534ddc0,
+ 0xdb629599, 0x3c439041, 0xfe5163ab,
+};
+
+struct reduced {
+ nir_def *k; /* int32, only the low two bits matter */
+ nir_def *fq; /* about [-1/2, 1/2] in quadrants, as a pair */
+ nir_def *fq_lo;
+};
+
+static struct reduced
+reduce_cody_waite(nir_builder *b, nir_def *x)
+{
+ nir_def *tp = nir_imm_float(b, TWO_OVER_PI);
+ nir_def *c2 = nir_imm_float(b, PIO2_2);
+ nir_def *kf = nir_fround_even(b, nir_fmul(b, x, tp));
+ nir_def *nk = nir_fneg(b, kf);
+
+ /* r1 = x - k C1 is exact. k C2 is a TwoProduct and r1 - k C2 a TwoSum so
+ * nothing is lost even when r1 and k C2 nearly cancel.
+ */
+ nir_def *r1 = nir_ffma(b, nk, nir_imm_float(b, PIO2_1), x);
+ nir_def *p = nir_fmul(b, kf, c2);
+ nir_def *p_lo = nir_ffma(b, kf, c2, nir_fneg(b, p));
+ nir_def *r2 = nir_fsub(b, r1, p);
+ nir_def *bb = nir_fsub(b, r2, r1);
+ nir_def *t = nir_fadd(b, nir_fsub(b, r1, nir_fsub(b, r2, bb)),
+ nir_fsub(b, nir_fneg(b, p), bb));
+ nir_def *r_lo =
+ nir_ffma(b, nk, nir_imm_float(b, PIO2_3), nir_fsub(b, t, p_lo));
+
+ nir_def *fq = nir_fmul(b, r2, tp);
+ nir_def *fq_lo = nir_ffma(b, r2, tp, nir_fneg(b, fq));
+ fq_lo = nir_ffma(b, r_lo, tp, fq_lo);
+ fq_lo = nir_ffma(b, r2, nir_imm_float(b, TWO_OVER_PI_LO), fq_lo);
+ nir_def *sum = nir_fadd(b, fq, fq_lo);
+
+ return (struct reduced){
+ .k = nir_f2i32(b, kf),
+ .fq = sum,
+ .fq_lo = fast2sum_lo(b, fq, fq_lo, sum),
+ };
+}
+
+static struct reduced
+reduce_payne_hanek(nir_builder *b, nir_def *ax)
+{
+ nir_def *e8 = nir_ushr_imm(b, ax, 23);
+ nir_def *m32 =
+ nir_ishl_imm(b, nir_ior_imm(b, nir_iand_imm(b, ax, 0x7fffff), 0x800000),
+ 8);
+
+ /* Word/bit offset of the 2/pi window for this exponent */
+ nir_def *off = nir_iadd_imm(b, e8, -126);
+ nir_def *idx = nir_ushr_imm(b, off, 5);
+ nir_def *sh = nir_iand_imm(b, off, 31);
+ nir_def *sh_r = nir_isub(b, nir_imm_int(b, 31), sh);
+
+ nir_def *tab[ARRAY_SIZE(two_over_pi_bits)];
+ for (unsigned i = 0; i < ARRAY_SIZE(tab); ++i)
+ tab[i] = nir_imm_int(b, two_over_pi_bits[i]);
+
+ /* Windows n = 1..4 at weights 2^0, 2^-32, 2^-64, 2^-96 */
+ nir_def *lo[5], *hi[5];
+ for (unsigned n = 1; n <= 4; ++n) {
+ nir_def *cand[5], *cand2[5];
+ for (unsigned j = 0; j < 5; ++j) {
+ cand[j] = tab[j + n];
+ cand2[j] = tab[j + n + 1];
+ }
+ nir_def *w0 = nir_select_from_ssa_def_array(b, cand, 5, idx);
+ nir_def *w1 = nir_select_from_ssa_def_array(b, cand2, 5, idx);
+ nir_def *win = nir_ior(b, nir_ishl(b, w0, sh),
+ nir_ushr_imm(b, nir_ushr(b, w1, sh_r), 1));
+ lo[n] = nir_imul(b, m32, win);
+ hi[n] = nir_umul_high(b, m32, win);
+ }
+
+ /* Carry-propagate the 96 fraction bits, keep the top 64 */
+ nir_def *f2 = nir_iadd(b, lo[3], hi[4]);
+ nir_def *c2 = nir_b2i32(b, nir_ult(b, f2, hi[4]));
+ nir_def *f1a = nir_iadd(b, lo[2], hi[3]);
+ nir_def *c1a = nir_b2i32(b, nir_ult(b, f1a, hi[3]));
+ nir_def *f1 = nir_iadd(b, f1a, c2);
+ nir_def *c1b = nir_b2i32(b, nir_ult(b, f1, f1a));
+ nir_def *k = nir_iadd(b, nir_iadd(b, lo[1], hi[2]), nir_iadd(b, c1a, c1b));
+
+ /* Round to nearest quadrant so fq lands in [-1/2, 1/2] */
+ nir_def *half = nir_ushr_imm(b, f1, 31);
+ k = nir_iadd(b, k, half);
+ nir_def *frac = nir_pack_64_2x32_split(b, f2, f1);
+ nir_def *is_half = nir_ine_imm(b, half, 0);
+ frac = nir_bcsel(b, is_half, nir_ineg(b, frac), frac);
+ nir_def *mag = nir_fmul_imm(b, nir_u2f32(b, frac), 0x1p-64);
+
+ return (struct reduced){
+ .k = k,
+ .fq = nir_bcsel(b, is_half, nir_fneg(b, mag), mag),
+ .fq_lo = nir_imm_float(b, 0.0),
+ };
+}
+
+static bool
+sincos_filter(const nir_instr *instr, UNUSED const void *_)
+{
+ if (instr->type != nir_instr_type_alu)
+ return false;
+
+ nir_alu_instr *alu = nir_instr_as_alu(instr);
+ return (alu->op == nir_op_fsin || alu->op == nir_op_fcos) &&
+ alu->def.bit_size <= 32;
+}
+
+static nir_def *
+lower_sincos(nir_builder *b, nir_instr *instr, UNUSED void *_)
+{
+ nir_alu_instr *alu = nir_instr_as_alu(instr);
+ b->fp_math_ctrl = nir_fp_no_fast_math;
+ nir_def *x = nir_mov_alu(b, alu->src[0], 1);
+ unsigned bit_size = x->bit_size;
+ if (bit_size == 16)
+ x = nir_f2f32(b, x);
+
+ nir_def *ax = nir_fabs(b, x);
+
+ nir_push_if(b, nir_fge_imm(b, ax, 0x1p22));
+ struct reduced ph = reduce_payne_hanek(b, ax);
+
+ /* Payne-Hanek reduced |x|; sin is odd, so flip k and fq for negatives */
+ nir_def *neg_x = nir_flt_imm(b, x, 0.0);
+ ph.k = nir_bcsel(b, neg_x, nir_ineg(b, ph.k), ph.k);
+ ph.fq = nir_bcsel(b, neg_x, nir_fneg(b, ph.fq), ph.fq);
+
+ nir_push_else(b, NULL);
+ struct reduced cw = reduce_cody_waite(b, x);
+ nir_pop_if(b, NULL);
+
+ nir_def *k = nir_if_phi(b, ph.k, cw.k);
+ nir_def *fq = nir_if_phi(b, ph.fq, cw.fq);
+ nir_def *fq_lo = nir_if_phi(b, ph.fq_lo, cw.fq_lo);
+
+ /* cos(x) = sin(x + pi/2): one more quadrant */
+ if (alu->op == nir_op_fcos)
+ k = nir_iadd_imm(b, k, 1);
+
+ /* Taylor series in quadrants, |fq| < 0.71 (k can be off by one),
+ * truncation below 2^-36. The leading sin term and the quadratic cos term
+ * are kept as pairs and the low part of fq feeds the derivative, so both
+ * land within an ulp.
+ */
+ nir_def *t = fq;
+ nir_def *t_lo = fq_lo;
+ nir_def *w = nir_fmul(b, t, t);
+ nir_def *w_lo = nir_ffma(b, t, t, nir_fneg(b, w));
+
+ static const double sin_coef[] = {
+ -0.6459640860557556, 0.07969262450933456, -0.004681753925979137,
+ 0.00016044118092395365, -3.598843250074424e-06, 5.6921727775716136e-08,
+ };
+ static const double cos_coef[] = {
+ -1.2337005138397217, 0.25366950035095215, -0.020863480865955353,
+ 0.0009192602592520416, -2.520204179745633e-05, 4.710874748070637e-07,
+ -6.386603246255618e-09,
+ };
+ static const double cos_coef0_lo = -3.629644851343983e-08;
+
+ nir_def *pio2 = nir_imm_float(b, PIO2_1);
+ nir_def *p = nir_imm_float(b, sin_coef[5]);
+ for (int i = 4; i >= 0; --i)
+ p = nir_ffma(b, p, w, nir_imm_float(b, sin_coef[i]));
+
+ nir_def *hi = nir_fmul(b, t, pio2);
+ nir_def *lo = nir_ffma(b, t, pio2, nir_fneg(b, hi));
+ lo = nir_ffma(b, t, nir_imm_float(b, PIO2_2), lo);
+ lo = nir_ffma(b, t_lo, pio2, lo);
+ nir_def *sn = nir_fadd(b, hi, nir_ffma(b, nir_fmul(b, t, w), p, lo));
+
+ nir_def *d2 = nir_imm_float(b, cos_coef[0]);
+ nir_def *r = nir_imm_float(b, cos_coef[6]);
+ for (int i = 5; i >= 1; --i)
+ r = nir_ffma(b, r, w, nir_imm_float(b, cos_coef[i]));
+
+ nir_def *h = nir_fmul(b, w, d2);
+ nir_def *l = nir_fadd(b, nir_fadd(b, nir_ffma(b, w, d2, nir_fneg(b, h)),
+ nir_fmul_imm(b, w, cos_coef0_lo)),
+ nir_fmul(b, w_lo, d2));
+ l = nir_ffma(b, nir_fmul_imm(b, t_lo, 2.0 * cos_coef[0]), t, l);
+ nir_def *one = nir_imm_float(b, 1.0);
+ nir_def *s1 = nir_fadd(b, one, h);
+ nir_def *e = fast2sum_lo(b, one, h, s1);
+ nir_def *cs = nir_fadd(
+ b, s1, nir_fadd(b, nir_fadd(b, e, l), nir_fmul(b, nir_fmul(b, w, w), r)));
+
+ /* Odd quadrants take the cosine, quadrants 2 and 3 negate */
+ nir_def *odd = nir_ine_imm(b, nir_iand_imm(b, k, 1), 0);
+ nir_def *neg = nir_ine_imm(b, nir_iand_imm(b, k, 2), 0);
+ nir_def *v = nir_bcsel(b, odd, cs, sn);
+ v = nir_bcsel(b, neg, nir_fneg(b, v), v);
+
+ /* Inf and NaN fall into the Payne-Hanek branch with a bogus exponent */
+ nir_def *nan = nir_imm_float(b, NAN);
+ v = nir_bcsel(b, nir_flt_imm(b, ax, INFINITY), v, nan);
+
+ /* sin(+-0) = +-0 exactly (the reduction returns +0 for -0) */
+ if (alu->op == nir_op_fsin)
+ v = nir_bcsel(b, nir_feq_imm(b, x, 0.0), x, v);
+
+ if (bit_size == 16)
+ v = nir_f2f16(b, v);
+
+ return v;
+}
+
+bool
+agx_nir_lower_sincos(nir_shader *s)
+{
+ return nir_shader_lower_instructions(s, sincos_filter, lower_sincos, NULL);
+}
diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build
index 62430fe3b641..c5727c07980d 100644
--- a/src/asahi/compiler/meson.build
+++ b/src/asahi/compiler/meson.build
@@ -14,6 +14,7 @@ libasahi_agx_files = files(
'agx_nir_lower_sample_mask.c',
'agx_nir_lower_discard_zs_emit.c',
'agx_nir_lower_interpolation.c',
+ 'agx_nir_lower_math.c',
'agx_nir_lower_shared_bitsize.c',
'agx_nir_lower_subgroups.c',
'agx_nir_lower_simdmat.c',
From 9994e65d9528994597a4544a058caee81e3e27cc Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Tue, 8 Sep 2026 13:54:03 +0000
Subject: [PATCH 08/16] agx: lower the fdiv that int64 lowering and the
perspective divide emit late
Symptom: any shader converting fp32 to a 64-bit integer aborts in the
backend with "Unhandled ALU op fdiv" (agx_compile.c agx_emit_alu); the
mlx-omarchy suite omarchy_copy_offset_tests crashes on "dtype converting
copies cover the integer family and bool". Fragment shaders with
perspective-correct varyings hit the same path.
Cause: 2022ddeae30 dropped lower_fdiv from the compiler options so
nir_opt_algebraic no longer splits fdiv into fmul/frcp, and only ran
agx_nir_lower_fdiv in agx_preprocess_nir. nir_lower_int64 (f2u64/f2i64:
fdiv by 2^32, plus the frem that lower_fmod turns into a second fdiv) and
agx_nir_lower_interpolation run after preprocessing, so their fdiv reached
the backend, which has no fdiv.
Fix: agx_nir_lower_fdiv_late runs in agx_optimize_nir after the int64
cleanup loop and lowers only fdiv, leaving frcp alone so the reciprocals
built by the first pass and by the log lowering are not refined twice.
Verification: hk/precise-math driver alone reproduced the abort on jwm1;
with this commit omarchy_copy_offset_tests passes 25/25 (see the
honeykrisp-omarchy integration receipt).
---
src/asahi/compiler/agx_compile.c | 5 +++++
src/asahi/compiler/agx_nir.h | 1 +
src/asahi/compiler/agx_nir_lower_math.c | 25 ++++++++++++++++++++++---
3 files changed, 28 insertions(+), 3 deletions(-)
diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index add80493133d..8cc96b31edc3 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -2931,6 +2931,11 @@ agx_optimize_nir(nir_shader *nir, bool soft_fault, uint16_t *preamble_size,
} while (progress);
}
+ /* int64 lowering and the perspective divide emit fdiv after preprocessing;
+ * the backend has no fdiv, so lower what is left before it is scheduled.
+ */
+ NIR_PASS(_, nir, agx_nir_lower_fdiv_late);
+
/* Reassociate before forming preambles because it makes preambles more
* effective. Clean up after.
*/
diff --git a/src/asahi/compiler/agx_nir.h b/src/asahi/compiler/agx_nir.h
index 1fead739f933..719ece490c92 100644
--- a/src/asahi/compiler/agx_nir.h
+++ b/src/asahi/compiler/agx_nir.h
@@ -23,6 +23,7 @@ bool agx_nir_lower_clip_distance(struct nir_shader *s);
bool agx_nir_lower_subgroups(struct nir_shader *s);
bool agx_nir_lower_fminmax(struct nir_shader *s);
bool agx_nir_lower_fdiv(struct nir_shader *s);
+bool agx_nir_lower_fdiv_late(struct nir_shader *s);
bool agx_nir_lower_log(struct nir_shader *s);
bool agx_nir_lower_sincos(struct nir_shader *s);
diff --git a/src/asahi/compiler/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c
index 3e9b923264f7..0a2d668e2cbc 100644
--- a/src/asahi/compiler/agx_nir_lower_math.c
+++ b/src/asahi/compiler/agx_nir_lower_math.c
@@ -57,9 +57,11 @@ div_rn(nir_builder *b, nir_def *a, nir_def *d)
}
static bool
-lower_fdiv(nir_builder *b, nir_alu_instr *alu, void *_)
+lower_fdiv(nir_builder *b, nir_alu_instr *alu, void *data)
{
- if (alu->op != nir_op_frcp && alu->op != nir_op_fdiv)
+ bool fdiv_only = *(bool *)data;
+
+ if (alu->op != nir_op_fdiv && (alu->op != nir_op_frcp || fdiv_only))
return false;
if (alu->def.bit_size == 64)
@@ -91,7 +93,24 @@ lower_fdiv(nir_builder *b, nir_alu_instr *alu, void *_)
bool
agx_nir_lower_fdiv(nir_shader *s)
{
- return nir_shader_alu_pass(s, lower_fdiv, nir_metadata_control_flow, NULL);
+ bool fdiv_only = false;
+ return nir_shader_alu_pass(s, lower_fdiv, nir_metadata_control_flow,
+ &fdiv_only);
+}
+
+/*
+ * Passes that run after agx_preprocess_nir still emit fdiv: nir_lower_int64
+ * (f2u64/f2i64 divide by 2^32, and lower_fmod turns the frem into a second
+ * fdiv) and agx_nir_lower_interpolation (perspective divide). Lower only fdiv
+ * here so the reciprocals built by the first pass and by the log lowering are
+ * not refined a second time.
+ */
+bool
+agx_nir_lower_fdiv_late(nir_shader *s)
+{
+ bool fdiv_only = true;
+ return nir_shader_alu_pass(s, lower_fdiv, nir_metadata_control_flow,
+ &fdiv_only);
}
/*
From 488793ee40cc4a73d7d5e5ce5e1558d6f112792a Mon Sep 17 00:00:00 2001
From: Joshua Warren
Date: Tue, 8 Sep 2026 09:05:43 -0500
Subject: [PATCH 09/16] honeykrisp: enable VK_KHR_cooperative_matrix by default
(AGX_SIMDMAT=0 disables)
Replace the five getenv("AGX_SIMDMAT") checks (agx_preprocess_nir lowering,
the extension table, the feature bit, the supported stages, and the
properties query) with agx_simdmat_enabled(), which reads the variable through
debug_get_bool_option with a default of true. Setting AGX_SIMDMAT=0 restores
the previous behaviour.
The 25/25 mlx-omarchy battery and every reproducer in
receipts/2026-09-08-honeykrisp-omarchy-integration.json ran with the gate on
(AGX_SIMDMAT=1), so the default now matches the verified configuration and a
system-installed driver exposes coopmat to unmodified wheels.
---
src/asahi/compiler/agx_compile.c | 5 ++---
src/asahi/compiler/agx_compile.h | 10 ++++++++++
src/asahi/compiler/agx_nir_lower_simdmat.c | 2 +-
src/asahi/vulkan/hk_physical_device.c | 10 +++++-----
4 files changed, 18 insertions(+), 9 deletions(-)
diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index 8cc96b31edc3..d5ad4a2d4de8 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -3524,10 +3524,9 @@ agx_preprocess_nir(nir_shader *nir)
return;
/* Lower VK_KHR_cooperative_matrix to the G13 HW matrix instruction before
- * vars_to_scratch can spill the 8x8 cmat temps. Gated so default Honeykrisp
- * does not compile experimental coopmat shaders.
+ * vars_to_scratch can spill the 8x8 cmat temps. AGX_SIMDMAT=0 skips it.
*/
- if (getenv("AGX_SIMDMAT"))
+ if (agx_simdmat_enabled())
NIR_PASS(_, nir, agx_nir_lower_simdmat, 32);
NIR_PASS(_, nir, nir_lower_vars_to_ssa);
diff --git a/src/asahi/compiler/agx_compile.h b/src/asahi/compiler/agx_compile.h
index 58706acd2f77..30db575d03b7 100644
--- a/src/asahi/compiler/agx_compile.h
+++ b/src/asahi/compiler/agx_compile.h
@@ -7,6 +7,7 @@
#include "compiler/nir/nir.h"
#include "util/shader_stats.h"
+#include "util/u_debug.h"
#include "util/u_dynarray.h"
#include "util/u_tristate.h"
#include "shader_enums.h"
@@ -317,6 +318,15 @@ struct agx_shader_key {
uint64_t agx_gather_texcoords(nir_shader *nir);
void agx_preprocess_nir(nir_shader *nir);
+
+/* VK_KHR_cooperative_matrix on the G13 simd_matrix unit. On by default;
+ * AGX_SIMDMAT=0 disables it.
+ */
+static inline bool
+agx_simdmat_enabled(void)
+{
+ return debug_get_bool_option("AGX_SIMDMAT", true);
+}
bool agx_nir_lower_discard_zs_emit(nir_shader *s);
bool agx_nir_lower_sample_mask(nir_shader *s);
bool agx_nir_lower_interpolation(nir_shader *s);
diff --git a/src/asahi/compiler/agx_nir_lower_simdmat.c b/src/asahi/compiler/agx_nir_lower_simdmat.c
index 44c0202c30b8..3e6c596d3df3 100644
--- a/src/asahi/compiler/agx_nir_lower_simdmat.c
+++ b/src/asahi/compiler/agx_nir_lower_simdmat.c
@@ -38,7 +38,7 @@
* what Apple's own mul_mm emits), fmadd16 with fp16 A/B. fp32 A/B with an fp16
* accumulator is not advertised and not lowered.
*
- * Gated by AGX_SIMDMAT at the call site.
+ * On by default; AGX_SIMDMAT=0 disables it at the call site.
*/
#include "util/macros.h"
diff --git a/src/asahi/vulkan/hk_physical_device.c b/src/asahi/vulkan/hk_physical_device.c
index 0ecd271033ce..bb04ae12eaa1 100644
--- a/src/asahi/vulkan/hk_physical_device.c
+++ b/src/asahi/vulkan/hk_physical_device.c
@@ -55,7 +55,7 @@ hk_get_device_extensions(const struct hk_instance *instance,
.KHR_bind_memory2 = true,
.KHR_buffer_device_address = true,
.KHR_calibrated_timestamps = true,
- .KHR_cooperative_matrix = getenv("AGX_SIMDMAT") != NULL,
+ .KHR_cooperative_matrix = agx_simdmat_enabled(),
.KHR_copy_commands2 = true,
.KHR_create_renderpass2 = true,
.KHR_dedicated_allocation = true,
@@ -239,8 +239,8 @@ hk_get_device_features(
struct vk_features *features)
{
*features = (struct vk_features){
- /* VK_KHR_cooperative_matrix (G13 HW 8x8x8, gated by AGX_SIMDMAT) */
- .cooperativeMatrix = getenv("AGX_SIMDMAT") != NULL,
+ /* VK_KHR_cooperative_matrix (G13 HW 8x8x8, AGX_SIMDMAT=0 disables) */
+ .cooperativeMatrix = agx_simdmat_enabled(),
.cooperativeMatrixRobustBufferAccess = false,
/* Vulkan 1.0 */
@@ -826,7 +826,7 @@ hk_get_device_properties(const struct agx_device *dev,
/* Vulkan 1.1 properties */
.subgroupSize = 32,
.cooperativeMatrixSupportedStages =
- getenv("AGX_SIMDMAT") ? VK_SHADER_STAGE_COMPUTE_BIT : 0,
+ agx_simdmat_enabled() ? VK_SHADER_STAGE_COMPUTE_BIT : 0,
.subgroupSupportedStages =
VK_SHADER_STAGE_COMPUTE_BIT | VK_SHADER_STAGE_ALL_GRAPHICS,
.subgroupSupportedOperations =
@@ -1536,7 +1536,7 @@ hk_GetPhysicalDeviceCooperativeMatrixPropertiesKHR(
VK_OUTARRAY_MAKE_TYPED(VkCooperativeMatrixPropertiesKHR, out, pProperties,
pPropertyCount);
- if (getenv("AGX_SIMDMAT")) {
+ if (agx_simdmat_enabled()) {
/* G13 hardware 8x8x8 matrix MAC (simd_matrix_fmadd16/32). The unit
* takes fp16 or fp32 A/B and accumulates in fp32 (fmadd32) or fp16
* (fmadd16); 16x16x16 is a 2x2x2 tiling of 8x8x8 ops in
From 0e1105bb55bd9bc9b641943b73ac712496c2cd39 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 28 Sep 2026 23:16:42 -0500
Subject: [PATCH 10/16] honeykrisp: dependency-track the per-launch CDM barrier
on the Vulkan path (G13X)
hk_dispatch_with_usc_launch flushed the full CDM barrier after every launch
unconditionally, ~4.6-5.1 us per dispatch on compute streams of hundreds of
small dependent dispatches. Key the barrier on what the app declared instead:
vkCmdPipelineBarrier2 marks a pending dependency and the next CDM launch
carries the barrier and clears it. Launches the app left unordered have no
Vulkan-defined ordering to enforce, so their drain is skipped. The flag lives
on the command buffer so it survives CS splits and merge_control_streams;
hk_EndCommandBuffer lands an unconsumed barrier on the trailing edge, and
driver-internal ordering (query availability writes, hk_queue_write) keeps its
explicit flush. The compute batch also stays open across a compute-to-compute
vkCmdPipelineBarrier2 when no graphics batch is open.
On G13X the barrier carries the measured designed set {unk_4,5,6,8} + USC
cache invalidate plus bits 0-2 (0x17f). Bit bisection showed every one of the
designed bits is required for bit-exact results: removing any corrupts, and a
USC-invalidate-only barrier is fast but wrong (320/320 greedy token flips).
G13G/G14 keep the previous kitchen-sink barrier.
New perftests: HK_PERFTEST=alwayscdmbarrier (restore the unconditional flush),
nocdmbarrier, usccdmbarrier, designedusccdmbarrier, and HK_CDM_BARRIER_MASK
(exact flag word) for bisecting.
Measured on an M1 Max (G13C), MLX-on-Vulkan Qwen3.5-class 4-bit greedy decode:
the designed-set trim 63.50 -> 68.34 tok/s (10 interleaved paired reps, +7.7%,
identical token digests); with the app declaring only real dependencies
(MLX dependency gating on) the tracked barrier adds +2.9% (81.3 -> 83.7 tok/s),
and 16 interleaved runs of the 10-prompt x 10-pass greedy contract produced
identical token digests in both arms.
---
src/asahi/libagx/libagx_dgc.h | 91 +++++++++++++++++++++++++-----
src/asahi/vulkan/hk_cmd_buffer.c | 34 +++++++++--
src/asahi/vulkan/hk_cmd_buffer.h | 7 +++
src/asahi/vulkan/hk_cmd_dispatch.c | 46 +++++++++++++--
src/asahi/vulkan/hk_device.c | 10 +++-
src/asahi/vulkan/hk_device.h | 9 +++
6 files changed, 175 insertions(+), 22 deletions(-)
diff --git a/src/asahi/libagx/libagx_dgc.h b/src/asahi/libagx/libagx_dgc.h
index d8266b663111..6b5bba235b42 100644
--- a/src/asahi/libagx/libagx_dgc.h
+++ b/src/asahi/libagx/libagx_dgc.h
@@ -379,9 +379,29 @@ agx_cdm_barrier(GLOBAL uint32_t *out, enum agx_chip chip)
// cfg.unk_20 = true;
// cfg.unk_24 = true; if clustered?
if (chip == AGX_CHIP_G13X) {
+ /* G13X (t600x, M1 Pro/Max, G13C cores): restore bits 0-2
+ * (PBE/texture flush) alongside unk_4 + usc_cache_inval.
+ * Bits 0-2 are required for cross-launch coherency when
+ * control streams are not split and across dependency-skipped
+ * launches ({0,1,2,4,5,6,8}+usc_cache_inval = 0x17f). */
+ cfg.unk_0 = true;
+ cfg.unk_1 = true;
+ cfg.unk_2 = true;
cfg.unk_4 = true;
+ cfg.usc_cache_inval = true;
// cfg.unk_26 = true;
}
+ if (chip == AGX_CHIP_G13G) {
+ /* G13G (t8103, M1): the kitchen-sink bits below cost ~10 us per
+ * dispatch in real dependent compute chains (Qwen decode: ~2.1 ms
+ * of a ~9.8 ms token). Bits 4-8 hold the mlx-omarchy pinned
+ * generated-ID digests (48/48 interleaved runs, both legs) and
+ * the omarchy runtime suite (22 cases / 6189 assertions), and
+ * measure +3.05% ctx1053 decode on jwm1. Other chips keep the
+ * full sink until measured there. */
+ cfg.unk_4 = true;
+ cfg.unk_7 = true;
+ }
/* With multiple launches in the same CDM stream, we can get cache
* coherency (? or sync?) issues. We hit this with blits, which need - in
@@ -390,26 +410,71 @@ agx_cdm_barrier(GLOBAL uint32_t *out, enum agx_chip chip)
* let's just set these after every launch to be safe. We can revisit in
* the future when we figure out what the bits mean.
*/
+ if (chip != AGX_CHIP_G13G && chip != AGX_CHIP_G13X) {
+ cfg.unk_0 = true;
+ cfg.unk_1 = true;
+ cfg.unk_2 = true;
+ cfg.usc_cache_inval = true;
+ cfg.unk_4 = true;
+ cfg.unk_5 = true;
+ cfg.unk_6 = true;
+ cfg.unk_7 = true;
+ cfg.unk_8 = true;
+ cfg.unk_9 = true;
+ cfg.unk_10 = true;
+ cfg.unk_11 = true;
+ cfg.unk_12 = true;
+ cfg.unk_13 = true;
+ cfg.unk_14 = true;
+ cfg.unk_15 = true;
+ cfg.unk_16 = true;
+ cfg.unk_17 = true;
+ cfg.unk_18 = true;
+ cfg.unk_19 = true;
+ }
+ }
+
+ return out;
+}
+
+/*
+ * Minimal variant of agx_cdm_barrier: invalidate the USC cache only. Used by
+ * the Vulkan driver for launches whose results are consumed by another compute
+ * launch through the same (L2-backed) storage path, where the PBE/texture
+ * cache maintenance of the full barrier is not required.
+ */
+static inline GLOBAL uint32_t *
+agx_cdm_barrier_usc(GLOBAL uint32_t *out)
+{
+ agx_push(out, CDM_BARRIER, cfg) {
+ cfg.usc_cache_inval = true;
+ }
+
+ return out;
+}
+
+/*
+ * G13X candidate reduction: the designed set {unk_4, unk_5, unk_6, unk_8}
+ * (the bits that held mlx-omarchy pinned digests 48/48 in the termA
+ * battery) plus the USC cache invalidate (covers uniform/texture-state
+ * changes between launches; measured 4.6 us/launch of the 20.4 us
+ * sink+cache cost on real dependent chains). The full-sink unk_0..unk_19
+ * block (PBE/texture cache maintenance, measured 15.8 us/launch) is
+ * dropped. Perftest arm designedusccdmbarrier; NOT a default-emission
+ * proposal until it holds both decode legs and the digest protocol.
+ */
+static inline GLOBAL uint32_t *
+agx_cdm_barrier_designed_usc(GLOBAL uint32_t *out)
+{
+ agx_push(out, CDM_BARRIER, cfg) {
cfg.unk_0 = true;
cfg.unk_1 = true;
cfg.unk_2 = true;
- cfg.usc_cache_inval = true;
cfg.unk_4 = true;
cfg.unk_5 = true;
cfg.unk_6 = true;
- cfg.unk_7 = true;
cfg.unk_8 = true;
- cfg.unk_9 = true;
- cfg.unk_10 = true;
- cfg.unk_11 = true;
- cfg.unk_12 = true;
- cfg.unk_13 = true;
- cfg.unk_14 = true;
- cfg.unk_15 = true;
- cfg.unk_16 = true;
- cfg.unk_17 = true;
- cfg.unk_18 = true;
- cfg.unk_19 = true;
+ cfg.usc_cache_inval = true;
}
return out;
diff --git a/src/asahi/vulkan/hk_cmd_buffer.c b/src/asahi/vulkan/hk_cmd_buffer.c
index b527e97a886e..cb3c8b73182b 100644
--- a/src/asahi/vulkan/hk_cmd_buffer.c
+++ b/src/asahi/vulkan/hk_cmd_buffer.c
@@ -303,6 +303,19 @@ hk_EndCommandBuffer(VkCommandBuffer commandBuffer)
"must end rendering before ending the command buffer");
perf_debug(cmd, "End command buffer");
+
+ /* A barrier declared with no following launch in this command buffer still
+ * orders this buffer's writes against whatever runs next (the queue's next
+ * submission). Land it on the trailing edge of the open compute stream;
+ * with no trailing edge there is no local work to order.
+ */
+ if (cmd->state.cs.cdm_barrier_pending) {
+ cmd->state.cs.cdm_barrier_pending = false;
+ if (cmd->current_cs.cs) {
+ hk_cdm_cache_flush(hk_cmd_buffer_device(cmd), cmd->current_cs.cs);
+ }
+ }
+
hk_cmd_buffer_end_compute(cmd);
hk_cmd_buffer_end_compute_internal(cmd, &cmd->current_cs.post_gfx);
@@ -336,10 +349,23 @@ hk_CmdPipelineBarrier2(VkCommandBuffer commandBuffer,
perf_debug(cmd, "Pipeline barrier");
- /* The big hammer. We end both compute and graphics batches. Ending compute
- * here is necessary to properly handle graphics->compute dependencies.
- *
- * XXX: perf. */
+ /* The app declared a memory dependency: the next CDM launch carries the
+ * full barrier set (hk_cmd_dispatch.c). Compute launches before the
+ * barrier were already flushed by the previous launch's gating.
+ */
+ cmd->state.cs.cdm_barrier_pending = true;
+
+ /* Outside a render pass only a CDM control stream can be open; the
+ * pending flag above orders it against the next launch. Ending the
+ * compute batch here would only produce the stream merge_control_streams
+ * stitches back together at hk_EndCommandBuffer, plus a stream link and
+ * a fresh chunk per barrier; keep it open. The batches must end when a
+ * graphics batch is open: the dependency then spans batch types and the
+ * kernel orders the resulting commands (agx_cmd_header barriers). */
+ if (!cmd->current_cs.gfx && !cmd->current_cs.pre_gfx &&
+ !cmd->current_cs.post_gfx)
+ return;
+
hk_cmd_buffer_end_compute(cmd);
hk_cmd_buffer_end_graphics(cmd);
}
diff --git a/src/asahi/vulkan/hk_cmd_buffer.h b/src/asahi/vulkan/hk_cmd_buffer.h
index b8f16d3f3aca..7ade2e7595ee 100644
--- a/src/asahi/vulkan/hk_cmd_buffer.h
+++ b/src/asahi/vulkan/hk_cmd_buffer.h
@@ -295,6 +295,13 @@ struct hk_graphics_state {
struct hk_compute_state {
struct hk_descriptor_state descriptors;
struct hk_api_shader *shader;
+
+ /* Set by hk_CmdPipelineBarrier2, consumed by the next CDM launch, which
+ * then carries the full designed barrier set (dependency-tracked
+ * per-launch barrier). Lives on the command buffer, not the hk_cs, so it
+ * survives CS splits and merge_control_streams.
+ */
+ bool cdm_barrier_pending;
};
struct hk_cmd_push {
diff --git a/src/asahi/vulkan/hk_cmd_dispatch.c b/src/asahi/vulkan/hk_cmd_dispatch.c
index 7cb08bb1d3c8..e3f61cc716ac 100644
--- a/src/asahi/vulkan/hk_cmd_dispatch.c
+++ b/src/asahi/vulkan/hk_cmd_dispatch.c
@@ -48,10 +48,28 @@ hk_cdm_cache_flush(struct hk_device *dev, struct hk_cs *cs)
assert(cs->current + AGX_CDM_BARRIER_LENGTH < cs->end &&
"caller must ensure space");
- cs->current = agx_cdm_barrier(cs->current, dev->dev.chip);
- cs->stats.flushes++;
+ if (unlikely(dev->cdm_barrier_mask)) {
+ /* Perftest only: exact CDM_BARRIER flag word from the environment. */
+ uint32_t *word = cs->current;
+ *word = dev->cdm_barrier_mask | (AGX_CDM_BLOCK_TYPE_BARRIER << 29);
+ cs->current = word + 1;
+ } else if (HK_PERF(dev, NOCDMBARRIER)) {
+ /* Perftest only: launch with no cache maintenance at all. */
+ } else if (HK_PERF(dev, DESIGNEDUSCCDMBARRIER)) {
+ /* Perftest only: designed set {4,5,6,8} + USC cache invalidate. */
+ cs->current = agx_cdm_barrier_designed_usc(cs->current);
+ } else if (HK_PERF(dev, USCCDMBARRIER)) {
+ /* Perftest only: keep the USC cache invalidate, drop the rest. */
+ cs->current = agx_cdm_barrier_usc(cs->current);
+ } else {
+ cs->current = agx_cdm_barrier(cs->current, dev->dev.chip);
+ }
+
+ if (!HK_PERF(dev, NOCDMBARRIER))
+ cs->stats.flushes++;
}
+
void
hk_dispatch_with_usc_launch(struct hk_device *dev, struct hk_cs *cs,
struct agx_cdm_launch_word_0_packed launch,
@@ -61,10 +79,30 @@ hk_dispatch_with_usc_launch(struct hk_device *dev, struct hk_cs *cs,
hk_ensure_cs_has_space(cs->cmd, cs, 0x2000 /* TODO */);
cs->stats.cmds++;
+ if (HK_PERF(dev, ALWAYSCDMBARRIER)) {
+ cs->current =
+ agx_cdm_launch(cs->current, dev->dev.chip, grid, wg, launch, usc);
+ hk_cdm_cache_flush(dev, cs);
+ cs->cmd->state.cs.cdm_barrier_pending = false;
+ return;
+ }
+
+ /* Dependency-tracked per-launch CDM barrier. A vkCmdPipelineBarrier
+ * recorded since the previous launch declares a memory dependency, so the
+ * next launch carries the full designed barrier set. Between launches the
+ * app left unordered, Vulkan defines no ordering, so the barrier enforces
+ * nothing and the drain it costs is skipped. Driver-internal ordering that
+ * does not ride on vkCmdPipelineBarrier (query availability writes) calls
+ * hk_cdm_cache_flush explicitly and is unaffected.
+ */
+ if (cs->cmd->state.cs.cdm_barrier_pending) {
+ assert(cs->type == HK_CS_CDM);
+ hk_cdm_cache_flush(dev, cs);
+ cs->cmd->state.cs.cdm_barrier_pending = false;
+ }
+
cs->current =
agx_cdm_launch(cs->current, dev->dev.chip, grid, wg, launch, usc);
-
- hk_cdm_cache_flush(dev, cs);
}
void
diff --git a/src/asahi/vulkan/hk_device.c b/src/asahi/vulkan/hk_device.c
index b3263bf0c301..ea6fe0b6a04f 100644
--- a/src/asahi/vulkan/hk_device.c
+++ b/src/asahi/vulkan/hk_device.c
@@ -43,7 +43,13 @@ static const struct debug_named_value hk_perf_options[] = {
{"noborder", HK_PERF_NOBORDER, "Disable custom border colour emulation"},
{"nobarrier", HK_PERF_NOBARRIER,"Ignore pipeline barriers"},
{"batch", HK_PERF_BATCH, "Batch submissions"},
- {"norobust", HK_PERF_NOROBUST, "Disable robustness"},
+ {"nocdmbarrier", HK_PERF_NOCDMBARRIER,
+ "Omit the per-launch CDM barrier entirely"},
+ {"usccdmbarrier", HK_PERF_USCCDMBARRIER,
+ "Reduce the per-launch CDM barrier to a USC cache invalidate"},
+ {"designedusccdmbarrier", HK_PERF_DESIGNEDUSCCDMBARRIER,
+ "G13X designed set {0,1,2,4,5,6,8} plus USC cache invalidate (0x17f)"},
+ {"alwayscdmbarrier", HK_PERF_ALWAYSCDMBARRIER, "Unconditional per-launch CDM barrier (disable dependency tracking)"},
DEBUG_NAMED_VALUE_END
};
/* clang-format on */
@@ -325,6 +331,8 @@ hk_CreateDevice(VkPhysicalDevice physicalDevice,
}
dev->perftest = debug_get_flags_option("HK_PERFTEST", hk_perf_options, 0);
+ dev->cdm_barrier_mask =
+ debug_get_num_option("HK_CDM_BARRIER_MASK", 0) & 0x07ffffff;
if (instance->drirc.misc.disable_border_emulation) {
dev->perftest |= HK_PERF_NOBORDER;
diff --git a/src/asahi/vulkan/hk_device.h b/src/asahi/vulkan/hk_device.h
index e6c71aa1f484..4238cc1e941f 100644
--- a/src/asahi/vulkan/hk_device.h
+++ b/src/asahi/vulkan/hk_device.h
@@ -98,6 +98,9 @@ struct hk_device {
} scratch;
uint32_t perftest;
+ /* HK_CDM_BARRIER_MASK=: perftest-only override of the per-launch
+ * CDM_BARRIER flag word (bits 0..26). 0 = unset. */
+ uint32_t cdm_barrier_mask;
struct {
struct u_rwlock lock;
@@ -114,6 +117,12 @@ enum hk_perftest {
HK_PERF_NOBARRIER = BITFIELD_BIT(2),
HK_PERF_BATCH = BITFIELD_BIT(3),
HK_PERF_NOROBUST = BITFIELD_BIT(4),
+ HK_PERF_NOCDMBARRIER = BITFIELD_BIT(5),
+ HK_PERF_USCCDMBARRIER = BITFIELD_BIT(6),
+ HK_PERF_DESIGNEDUSCCDMBARRIER = BITFIELD_BIT(7),
+ /* Restore the unconditional per-launch CDM barrier (the default is the
+ * dependency-tracked barrier keyed on vkCmdPipelineBarrier2). */
+ HK_PERF_ALWAYSCDMBARRIER = BITFIELD_BIT(8),
};
#define HK_PERF(dev, flag) unlikely((dev)->perftest &HK_PERF_##flag)
From be3f3a4b15c077b70c8eeb47d862d643915a742a Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 28 Sep 2026 23:16:42 -0500
Subject: [PATCH 11/16] hk: opt-in bounded poll before blocking syncobj waits
(HK_SUBMIT_POLL_US)
HK_SUBMIT_POLL_US= (default 0 = off) makes vk_drm_syncobj_wait_many poll
with a zero timeout for up to that long before falling back to the blocking
timeline wait with the original deadline. Short submit->signal round trips
(one submit + one wait per step of a stepwise decoder loop) pay a host wake-up
premium under the sleeping wait; the poll observes the signal the moment the
kernel publishes it. After a budget expires without seeing the signal the poll
is disabled for budget << misses (capped at 1024x), so long waits (shader cache
compile storms, whole-frame jobs) pay at most one budget per backoff period;
any poll hit resets the state.
Semantics preserved: same ioctl with timeout 0 during the poll budget, original
absolute deadline on the fallback, ETIME still maps to VK_TIMEOUT, non-ETIME
errors surface immediately, and drivers that do not set sync_wait_poll_us take
the identical single-call path.
Measured on an M1 Max (G13C): a Parakeet TDT decoder loop (278 submit+wait
round trips) 341 -> 254 ms (-25%) at HK_SUBMIT_POLL_US=2000 (poll budget sweep
0/100/200/500/1000/2000/5000 us: 347/314/292/266/258/248/250 ms), transcript
unchanged; LLM decode/prefill throughput and token digests unchanged.
---
src/asahi/vulkan/hk_device.c | 6 ++
src/vulkan/runtime/vk_device.h | 19 ++++++
src/vulkan/runtime/vk_drm_syncobj.c | 93 ++++++++++++++++++++++-------
3 files changed, 96 insertions(+), 22 deletions(-)
diff --git a/src/asahi/vulkan/hk_device.c b/src/asahi/vulkan/hk_device.c
index ea6fe0b6a04f..69b9fe9347f7 100644
--- a/src/asahi/vulkan/hk_device.c
+++ b/src/asahi/vulkan/hk_device.c
@@ -334,6 +334,12 @@ hk_CreateDevice(VkPhysicalDevice physicalDevice,
dev->cdm_barrier_mask =
debug_get_num_option("HK_CDM_BARRIER_MASK", 0) & 0x07ffffff;
+ /* Bounded poll before blocking syncobj waits (see vk_device.sync_wait_poll_us).
+ * Off by default; pays CPU to skip wake-up latency on short submit->signal
+ * round trips. */
+ dev->vk.sync_wait_poll_us =
+ CLAMP(debug_get_num_option("HK_SUBMIT_POLL_US", 0), 0, 1000000);
+
if (instance->drirc.misc.disable_border_emulation) {
dev->perftest |= HK_PERF_NOBORDER;
}
diff --git a/src/vulkan/runtime/vk_device.h b/src/vulkan/runtime/vk_device.h
index fb6221ad9de3..c84fa682c7ff 100644
--- a/src/vulkan/runtime/vk_device.h
+++ b/src/vulkan/runtime/vk_device.h
@@ -143,6 +143,25 @@ struct vk_device {
/** Shader vtable for VK_EXT_shader_object and common pipelines */
const struct vk_device_shader_ops *shader_ops;
+ /** Bounded poll before blocking syncobj waits, in microseconds
+ *
+ * Zero (the default) disables the poll. When nonzero, common syncobj
+ * waits first poll with a zero timeout for up to this long before
+ * falling back to the blocking wait with the original deadline,
+ * trading CPU for wake-up latency on short waits.
+ */
+ uint32_t sync_wait_poll_us;
+
+ /** Miss backoff state for the bounded syncobj poll (see sync_wait_poll_us)
+ *
+ * sync_poll_misses counts consecutive polls whose budget expired
+ * without seeing the signal; sync_poll_miss_expire_ns is the absolute
+ * time until which polling is disabled after such a miss. Accessed
+ * with the generic p_atomic helpers (64-bit, 8-byte aligned).
+ */
+ uint64_t sync_poll_misses;
+ uint64_t sync_poll_miss_expire_ns;
+
/** Acceleration structure build vtable for common BVH building. */
const struct vk_acceleration_structure_build_ops *as_build_ops;
diff --git a/src/vulkan/runtime/vk_drm_syncobj.c b/src/vulkan/runtime/vk_drm_syncobj.c
index 57de98eaffb0..4af92852df62 100644
--- a/src/vulkan/runtime/vk_drm_syncobj.c
+++ b/src/vulkan/runtime/vk_drm_syncobj.c
@@ -325,29 +325,78 @@ vk_drm_syncobj_wait_many(struct vk_device *device,
if (!(wait_flags & VK_SYNC_WAIT_ANY))
syncobj_wait_flags |= DRM_SYNCOBJ_WAIT_FLAGS_WAIT_ALL;
- int err;
- if (wait_count == 0) {
- err = 0;
- } else if (wait_flags & VK_SYNC_WAIT_PENDING) {
- /* We always use a timeline wait for WAIT_PENDING, even for binary
- * syncobjs because the non-timeline wait doesn't support
- * DRM_SYNCOBJ_WAIT_FLAGS_WAIT_AVAILABLE.
+ int err = 0;
+
+ /* When the driver opted in, poll with a zero timeout first: short
+ * waits that are submitted-and-signaled within a few hundred
+ * microseconds pay a wake-up latency premium under the blocking wait,
+ * and a bounded poll observes the signal the moment the kernel
+ * publishes it. Once the poll budget is exhausted, fall through to
+ * the blocking wait with the original deadline.
+ *
+ * A poll whose budget expires unsignaled means this wait is a long
+ * one (compile-time sync storms, frame fences); polling again would
+ * only burn the budget on every subsequent long wait. Back off
+ * multiplicatively after each miss and re-arm on the next hit.
+ */
+ const uint64_t poll_budget_ns =
+ (uint64_t)device->sync_wait_poll_us * 1000;
+ uint64_t poll_deadline_ns = 0;
+ if (poll_budget_ns &&
+ os_time_get_nano() >=
+ p_atomic_read(&device->sync_poll_miss_expire_ns))
+ poll_deadline_ns = os_time_get_nano() + poll_budget_ns;
+
+ while (wait_count > 0) {
+ const bool poll =
+ poll_deadline_ns && os_time_get_nano() < poll_deadline_ns;
+ const uint64_t timeout_ns = poll ? 0 : abs_timeout_ns;
+
+ if (wait_flags & VK_SYNC_WAIT_PENDING) {
+ /* We always use a timeline wait for WAIT_PENDING, even for binary
+ * syncobjs because the non-timeline wait doesn't support
+ * DRM_SYNCOBJ_WAIT_FLAGS_WAIT_AVAILABLE.
+ */
+ err = device->sync->timeline_wait(device->sync, handles, wait_values,
+ wait_count, timeout_ns,
+ syncobj_wait_flags |
+ DRM_SYNCOBJ_WAIT_FLAGS_WAIT_AVAILABLE,
+ NULL /* first_signaled */);
+ } else if (has_timeline) {
+ err = device->sync->timeline_wait(device->sync, handles, wait_values,
+ wait_count, timeout_ns,
+ syncobj_wait_flags,
+ NULL /* first_signaled */);
+ } else {
+ err = device->sync->wait(device->sync, handles,
+ wait_count, timeout_ns,
+ syncobj_wait_flags,
+ NULL /* first_signaled */);
+ }
+
+ /* A zero-timeout wait is a poll iteration: success is done, a
+ * timeout means poll again until the budget runs out, and any
+ * other error surfaces immediately.
*/
- err = device->sync->timeline_wait(device->sync, handles, wait_values,
- wait_count, abs_timeout_ns,
- syncobj_wait_flags |
- DRM_SYNCOBJ_WAIT_FLAGS_WAIT_AVAILABLE,
- NULL /* first_signaled */);
- } else if (has_timeline) {
- err = device->sync->timeline_wait(device->sync, handles, wait_values,
- wait_count, abs_timeout_ns,
- syncobj_wait_flags,
- NULL /* first_signaled */);
- } else {
- err = device->sync->wait(device->sync, handles,
- wait_count, abs_timeout_ns,
- syncobj_wait_flags,
- NULL /* first_signaled */);
+ if (err == 0) {
+ if (poll_deadline_ns) {
+ p_atomic_set(&device->sync_poll_misses, 0);
+ p_atomic_set(&device->sync_poll_miss_expire_ns, 0);
+ }
+ break;
+ }
+ if (!poll || errno != ETIME)
+ break;
+
+ if (os_time_get_nano() >= poll_deadline_ns) {
+ /* Budget ran out unsignaled: this is a long wait. */
+ const uint32_t misses =
+ (uint32_t)p_atomic_inc_return(&device->sync_poll_misses);
+ const uint64_t backoff_ns =
+ poll_budget_ns << MIN2(misses, 10);
+ p_atomic_set(&device->sync_poll_miss_expire_ns,
+ os_time_get_nano() + backoff_ns);
+ }
}
STACK_ARRAY_FINISH(handles);
From 89b12dd84aa9f994a50c301f5c19d3cc9519c91c Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 5 Oct 2026 11:52:50 -0500
Subject: [PATCH 12/16] agx: flush the exact quotient in precise fp32 division,
not the correction step
Dj stress-testing an M2 Pro found a third division failure class the
original review missed (Basecamp 10277024959, 2026-09-13): when the exact
quotient sits just below the normal boundary, the GPU zeroes the
correction step's intermediate value and the Markstein step is silently
skipped, giving 1 ulp errors -- 18,382 wrong results over 1M random
pairs. The first commit in this series refined a * rcp(b) in place, so
the flush landed on an intermediate, not on the result.
Divide the normalized mantissas instead and restore the exponent
afterwards, so the refinement always runs on operands near one and the
flush happens on the exact quotient, before rounding. The rounding
decision is integer: the half-ulp point is exact by construction, and the
residual of the second Markstein step is exact by the fma contract, so
the comparison against it decides the tie without a third refinement.
The reciprocal takes the same path, and the log2 series divisor (near
two) keeps the unguarded Newton form.
This re-lands the division part of the fix first made on 2026-09-13 and
pulled with it during the 2026-09-20 A/B isolation windows; it was never
re-applied to the omacom/mesa stack. Verified against an exact model over
262,144 pairs covering edge bit patterns and the binade boundaries, and
by an independent exact-rational checker; see the receipt in the PR.
---
src/asahi/compiler/agx_nir_lower_math.c | 93 +++++++++++++++++--------
1 file changed, 63 insertions(+), 30 deletions(-)
diff --git a/src/asahi/compiler/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c
index 0a2d668e2cbc..09b764c1aa1e 100644
--- a/src/asahi/compiler/agx_nir_lower_math.c
+++ b/src/asahi/compiler/agx_nir_lower_math.c
@@ -20,40 +20,62 @@
* reassociating or contracting them.
*/
-/*
- * Correctly rounded 1/x. One Newton-Raphson step on the hardware rcp:
- *
- * u_2 = u + u(1 - xu) = fma(fma(-x, u, 1), u, u)
- *
- * exhaustively verified upstream with a modified math_bruteforce. If u is
- * infinite (x zero or flushed), the refinement is NaN, so keep u.
- */
+/* The caller keeps x near one so the Newton residual cannot flush. */
static nir_def *
-rcp_rn(nir_builder *b, nir_def *x)
+rcp_normal(nir_builder *b, nir_def *x)
{
nir_def *u = nir_frcp(b, x);
- nir_def *one = nir_imm_float(b, 1.0);
- nir_def *u_2 = nir_ffma(b, nir_ffma(b, nir_fneg(b, x), u, one), u, u);
- return nir_bcsel(b, nir_fisnan(b, u_2), u, u_2);
+ return nir_ffma(b, nir_ffma(b, nir_fneg(b, x), u,
+ nir_imm_float(b, 1.0)), u, u);
}
-/*
- * Correctly rounded a/b (Markstein). With y = RN(1/b) and q = RN(a y), the
- * residual r = a - b q is exact in an fma and q + r y rounds to RN(a/b) for
- * normal operands. If the refinement produces NaN (b zero, q infinite), q
- * already holds the IEEE result for that case. A zero q is also final (a zero
- * or the quotient flushed) and keeps its sign, which the refinement would
- * lose.
- */
+/* Normalize before refinement; exponent restoration must flush the exact
+ * quotient before rounding, including values just below the normal boundary. */
static nir_def *
div_rn(nir_builder *b, nir_def *a, nir_def *d)
{
- nir_def *y = rcp_rn(b, d);
- nir_def *q = nir_fmul(b, a, y);
- nir_def *r = nir_ffma(b, nir_fneg(b, d), q, a);
- nir_def *q_2 = nir_ffma(b, r, y, q);
- nir_def *keep = nir_ior(b, nir_fisnan(b, q_2), nir_feq_imm(b, q, 0.0));
- return nir_bcsel(b, keep, q, q_2);
+ nir_def *aa = nir_iand_imm(b, a, 0x7fffffff);
+ nir_def *dd = nir_iand_imm(b, d, 0x7fffffff);
+ nir_def *ea = nir_ushr_imm(b, aa, 23);
+ nir_def *ed = nir_ushr_imm(b, dd, 23);
+ nir_def *ma = nir_ior_imm(b, nir_iand_imm(b, aa, 0x7fffff), 0x3f800000);
+ nir_def *md = nir_ior_imm(b, nir_iand_imm(b, dd, 0x7fffff), 0x3f800000);
+ nir_def *y = rcp_normal(b, md);
+ nir_def *q = nir_fmul(b, ma, y);
+ nir_def *r = nir_ffma(b, nir_fneg(b, md), q, ma);
+ q = nir_ffma(b, r, y, q);
+ r = nir_ffma(b, nir_fneg(b, md), q, ma);
+ nir_def *half = nir_ishl_imm(b, nir_iadd_imm(b, nir_ushr_imm(b, q, 23), -24), 23);
+ nir_def *up = nir_fmul(b, md, half);
+ nir_def *down = nir_bcsel(b, nir_ieq_imm(b, nir_iand_imm(b, q, 0x7fffff), 0),
+ nir_fmul_imm(b, up, 0.5), up);
+ nir_def *odd = nir_ine_imm(b, nir_iand_imm(b, q, 1), 0);
+ nir_def *inc = nir_ior(b, nir_flt(b, up, r), nir_iand(b, nir_feq(b, up, r), odd));
+ nir_def *nr = nir_fneg(b, r);
+ nir_def *dec = nir_ior(b, nir_flt(b, down, nr), nir_iand(b, nir_feq(b, down, nr), odd));
+ q = nir_isub(b, nir_iadd(b, q, nir_b2i32(b, inc)), nir_b2i32(b, dec));
+ nir_def *e = nir_isub(b, ea, ed);
+ nir_def *out = nir_iadd(b, q, nir_ishl_imm(b, e, 23));
+ nir_def *under = nir_ior(b, nir_ilt_imm(b, e, -126),
+ nir_iand(b, nir_ieq_imm(b, e, -126), nir_ult(b, ma, md)));
+ out = nir_bcsel(b, under, nir_imm_int(b, 0), out);
+ nir_def *over = nir_ige_imm(b, nir_iadd(b, e, nir_ushr_imm(b, q, 23)), 255);
+ out = nir_bcsel(b, over, nir_imm_int(b, 0x7f800000), out);
+ nir_def *az = nir_ieq_imm(b, ea, 0), *dz = nir_ieq_imm(b, ed, 0);
+ nir_def *ai = nir_ieq_imm(b, ea, 255), *di = nir_ieq_imm(b, ed, 255);
+ out = nir_bcsel(b, nir_ior(b, az, di), nir_imm_int(b, 0), out);
+ out = nir_bcsel(b, nir_ior(b, dz, ai), nir_imm_int(b, 0x7f800000), out);
+ out = nir_ior(b, out, nir_iand_imm(b, nir_ixor(b, a, d), 0x80000000));
+ nir_def *invalid = nir_ior(b, nir_ugt_imm(b, aa, 0x7f800000),
+ nir_ugt_imm(b, dd, 0x7f800000));
+ invalid = nir_ior(b, invalid, nir_ior(b, nir_iand(b, az, dz), nir_iand(b, ai, di)));
+ return nir_bcsel(b, invalid, nir_imm_int(b, 0x7fc00000), out);
+}
+
+static nir_def *
+rcp_rn(nir_builder *b, nir_def *x)
+{
+ return div_rn(b, nir_imm_float(b, 1.0), x);
}
static bool
@@ -162,7 +184,7 @@ log_core(nir_builder *b, nir_def *x)
nir_def *d_lo = fast2sum_lo(b, two, f, d);
/* s = f / (d + d_lo) as a pair */
- nir_def *y = rcp_rn(b, d);
+ nir_def *y = rcp_normal(b, d);
nir_def *s = nir_fmul(b, f, y);
nir_def *r = nir_ffma(b, nir_fneg(b, s), d, f);
s = nir_ffma(b, r, y, s);
@@ -569,9 +591,20 @@ lower_sincos(nir_builder *b, nir_instr *instr, UNUSED void *_)
nir_def *nan = nir_imm_float(b, NAN);
v = nir_bcsel(b, nir_flt_imm(b, ax, INFINITY), v, nan);
- /* sin(+-0) = +-0 exactly (the reduction returns +0 for -0) */
- if (alu->op == nir_op_fsin)
- v = nir_bcsel(b, nir_feq_imm(b, x, 0.0), x, v);
+ /* sin(+-0) = +-0 exactly (the reduction returns +0 for -0). The zero
+ * comparison flushes its input, so a subnormal argument matches; under
+ * DenormFlushToZero the argument from an unflushed source (a buffer load)
+ * must not leak out through the select, so return the sign-preserving
+ * zero. Without the mode, preserving the subnormal argument is legal and
+ * sin(subnormal) rounds to it anyway. */
+ if (alu->op == nir_op_fsin) {
+ if (nir_is_denorm_flush_to_zero(
+ b->shader->info.float_controls_execution_mode, 32))
+ v = nir_bcsel(b, nir_feq_imm(b, x, 0.0),
+ nir_iand_imm(b, x, 0x80000000), v);
+ else
+ v = nir_bcsel(b, nir_feq_imm(b, x, 0.0), x, v);
+ }
if (bit_size == 16)
v = nir_f2f16(b, v);
From b728b366077914973ff28e8c546e5ffe9b9207ee Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Sun, 13 Sep 2026 10:00:55 -0500
Subject: [PATCH 13/16] nir: pass FP_MATH_CTRL to the builder generator
nir_builder_opcodes_h.py renders the intrinsic builders from a mako
template that references FP_MATH_CTRL, but the index was never imported
or passed to render(), so mako resolved the name to UNDEFINED and
silently skipped the builder-defaulting block for every FP_MATH_CTRL
intrinsic. NoContraction on cooperative-matrix operations was therefore
dropped in every driver, not just this one, and the generated header is
missing the twelve fp_math_ctrl defaulting blocks.
Import the index and hand it to the template. The generated header grows
exactly 24 lines, twelve two-line defaulting blocks.
(cherry picked from commit 864b4e88941110fb2f68d15168c9280094a5275c)
---
src/compiler/nir/nir_builder_opcodes_h.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/src/compiler/nir/nir_builder_opcodes_h.py b/src/compiler/nir/nir_builder_opcodes_h.py
index ffbc9245c4e8..e37032aee285 100644
--- a/src/compiler/nir/nir_builder_opcodes_h.py
+++ b/src/compiler/nir/nir_builder_opcodes_h.py
@@ -226,7 +226,7 @@ def get_intrinsic_bitsize(opcode):
#endif /* _NIR_BUILDER_OPCODES_ */"""
from nir_opcodes import opcodes, type_size, type_base_type
-from nir_intrinsics import INTR_OPCODES, WRITE_MASK, ALIGN_MUL, IO_SEMANTICS, DEST_TYPE, SRC_TYPE
+from nir_intrinsics import INTR_OPCODES, WRITE_MASK, ALIGN_MUL, IO_SEMANTICS, DEST_TYPE, SRC_TYPE, FP_MATH_CTRL
from mako.template import Template
# List of intrinsics that also need a nir_build_ prefixed factory macro.
@@ -257,4 +257,5 @@ def get_intrinsic_bitsize(opcode):
IO_SEMANTICS=IO_SEMANTICS,
DEST_TYPE=DEST_TYPE,
SRC_TYPE=SRC_TYPE,
+ FP_MATH_CTRL=FP_MATH_CTRL,
build_prefixed_intrinsics=build_prefixed_intrinsics))
From 8c9f4064c41dc2c0bc1fae980ead82ba9b4323b1 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Sun, 13 Sep 2026 10:00:57 -0500
Subject: [PATCH 14/16] tests: cooperative matrix and precise math regressions
for AGX
cmshape2 walks the advertised matrix shapes, layouts, workgroup sizes
and pointer widths against a host oracle, and checks that NoContraction
survives and that a divergent component write is not lost. run_matrix2
drives the 30 cases that cover the fixes in this series.
precise_math compares division, reciprocal, log and log2 against an
exact model over 262,144 inputs, including the binade boundaries and the
literal-versus-buffer pair for ln2. compute_runner is the small Vulkan
dispatch harness it uses.
(cherry picked from commit fb6237546e3f95dc644199840532f9bf5d48491b)
---
tests/cmshape2.c | 396 +++++++++++++++++++++++++++++++++++++++
tests/compute_runner.cpp | 196 +++++++++++++++++++
tests/precise_math.comp | 17 ++
tests/precise_math.py | 72 +++++++
tests/run_matrix2.sh | 43 +++++
5 files changed, 724 insertions(+)
create mode 100644 tests/cmshape2.c
create mode 100644 tests/compute_runner.cpp
create mode 100644 tests/precise_math.comp
create mode 100644 tests/precise_math.py
create mode 100755 tests/run_matrix2.sh
diff --git a/tests/cmshape2.c b/tests/cmshape2.c
new file mode 100644
index 000000000000..2eaab9d1f4af
--- /dev/null
+++ b/tests/cmshape2.c
@@ -0,0 +1,396 @@
+/*
+ * cmshape2: extended cmshape for the honeykrisp-omarchy coopmat fixes.
+ * Adds, without touching the original cmshape:
+ * arg 11: lsize local_size_x (default 32; 16/8/1/48 exercise partial or
+ * multi-subgroup workgroups -> software cmat path)
+ * arg 12: elem memory VIEW element type for all four buffers:
+ * f16 | f32 (default: same as matrix type) | u16 | u8
+ * (narrow/wide pointer forms: fp32 matrix behind uint16_t*)
+ * mode "precise": D = precise(C * s + C) element-wise with NoContraction;
+ * PASS requires the unfused fp32 reference bit-exactly.
+ * Buffer BYTES stay component-contiguous in the matrix's own format; only the
+ * shader's view type and the coopMatLoad stride units change (stride is in
+ * view-element units, so strides scale by matrix_bytes / view_bytes).
+ *
+ * usage: cmshape2 M N K TA TB TC TD row|col identity|ones|randint|randf|precise [seed] [lsize] [elem]
+ * exit: 0 pass, 1 mismatch, 2 extension missing, 3 shape not advertised,
+ * 4 setup error.
+ */
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+typedef _Float16 f16;
+static void die(const char *msg, int code) {
+ printf("{\"result\":\"FAIL\",\"reason\":\"%s\"}\n", msg);
+ exit(code);
+}
+static int is16(const char *t) { return strcmp(t, "f16") == 0; }
+static VkComponentTypeKHR ctype(const char *t) {
+ return is16(t) ? VK_COMPONENT_TYPE_FLOAT16_KHR : VK_COMPONENT_TYPE_FLOAT32_KHR;
+}
+static const char *gltype(const char *t) { return is16(t) ? "float16_t" : "float32_t"; }
+static double getv(const void *p, int f16_, size_t i) {
+ return f16_ ? (double)((const f16 *)p)[i] : (double)((const float *)p)[i];
+}
+static void setv(void *p, int f16_, size_t i, double v) {
+ if (f16_) ((f16 *)p)[i] = (f16)v; else ((float *)p)[i] = (float)v;
+}
+static uint32_t rng_state;
+static uint32_t rnd(void) {
+ rng_state ^= rng_state << 13; rng_state ^= rng_state >> 17; rng_state ^= rng_state << 5;
+ return rng_state;
+}
+int main(int argc, char **argv) {
+ if (argc < 10) { fprintf(stderr, "usage: see header\n"); return 4; }
+ int M = atoi(argv[1]), N = atoi(argv[2]), K = atoi(argv[3]);
+ const char *TA = argv[4], *TB = argv[5], *TC = argv[6], *TD = argv[7];
+ int colmaj = strcmp(argv[8], "col") == 0;
+ const char *mode = argv[9];
+ rng_state = argc > 10 ? (uint32_t)atoi(argv[10]) : 0x9E3779B9u;
+ if (!rng_state) rng_state = 1;
+ int lsize = argc > 11 ? atoi(argv[11]) : 32;
+ const char *elem = argc > 12 ? argv[12] : NULL;
+ if (!elem || !strcmp(elem, "")) elem = TA;
+ if (lsize <= 0 || lsize > 1024) die("bad lsize", 4);
+ int precise = !strcmp(mode, "precise");
+ int divergent = !strcmp(mode, "divergent");
+ int subgroups = (lsize + 31) / 32;
+ if (precise && !(strcmp(TA, "f32") == 0 && M == N && M == K))
+ die("precise needs f32 and M==K==N", 4);
+ if (divergent && strcmp(TA, "f16")) die("divergent needs f16 A", 4);
+
+ /* view element: bytes per view element, GLSL type, required extension */
+ unsigned vbytes; const char *vgltype; const char *vexts;
+ if (!strcmp(elem, "f16")) { vbytes = 2; vgltype = "float16_t"; vexts = ""; }
+ else if (!strcmp(elem, "f32")) { vbytes = 4; vgltype = "float32_t"; vexts = ""; }
+ else if (!strcmp(elem, "u16")) { vbytes = 2; vgltype = "uint16_t";
+ vexts = "#extension GL_EXT_shader_explicit_arithmetic_types_int16 : enable\n"; }
+ else if (!strcmp(elem, "u8")) { vbytes = 1; vgltype = "uint8_t";
+ vexts = "#extension GL_EXT_shader_explicit_arithmetic_types_int8 : enable\n"; }
+ else die("bad elem", 4);
+ int mbytes = is16(TA) ? 2 : 4; /* matrix component bytes */
+ if (mbytes % vbytes) die("elem must divide matrix component size", 4);
+ int ascale = mbytes / vbytes;
+ int bscale = (is16(TB) ? 2 : 4) / vbytes;
+ int cscale = (is16(TC) ? 2 : 4) / vbytes;
+ int dscale = (is16(TD) ? 2 : 4) / vbytes;
+
+ /* ---- instance / device ---- */
+ VkApplicationInfo ai = {VK_STRUCTURE_TYPE_APPLICATION_INFO};
+ ai.pApplicationName = "cmshape2"; ai.apiVersion = VK_API_VERSION_1_3;
+ VkInstanceCreateInfo ici = {VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO};
+ ici.pApplicationInfo = &ai;
+ VkInstance inst;
+ if (vkCreateInstance(&ici, NULL, &inst) != VK_SUCCESS) die("vkCreateInstance", 4);
+ uint32_t npd = 0;
+ vkEnumeratePhysicalDevices(inst, &npd, NULL);
+ if (!npd) die("no physical device", 4);
+ VkPhysicalDevice pds[8];
+ if (npd > 8) npd = 8;
+ vkEnumeratePhysicalDevices(inst, &npd, pds);
+ VkPhysicalDevice pd = pds[0];
+ VkPhysicalDeviceProperties props;
+ vkGetPhysicalDeviceProperties(pd, &props);
+ uint32_t nxt = 0;
+ vkEnumerateDeviceExtensionProperties(pd, NULL, &nxt, NULL);
+ VkExtensionProperties *exts = calloc(nxt, sizeof *exts);
+ vkEnumerateDeviceExtensionProperties(pd, NULL, &nxt, exts);
+ int have_cm = 0;
+ for (uint32_t i = 0; i < nxt; i++)
+ if (!strcmp(exts[i].extensionName, VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME)) have_cm = 1;
+ free(exts);
+ printf("{\"tool\":\"cmshape2\",\"device\":\"%s\",\"driver\":\"", props.deviceName);
+ {
+ VkPhysicalDeviceDriverProperties drv = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DRIVER_PROPERTIES};
+ VkPhysicalDeviceProperties2 p2 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, &drv};
+ vkGetPhysicalDeviceProperties2(pd, &p2);
+ printf("%s\",\"shape\":\"%dx%dx%d %s %s %s %s %s %s l%d %s\"", drv.driverInfo, M, N, K, TA, TB, TC, TD,
+ colmaj ? "col" : "row", mode, lsize, elem);
+ }
+ if (!have_cm) { printf(","); die("VK_KHR_cooperative_matrix not exposed", 2); }
+ PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR getcm =
+ (PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR)vkGetInstanceProcAddr(
+ inst, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR");
+ if (!getcm) { printf(","); die("no vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR", 2); }
+ uint32_t ncm = 0;
+ getcm(pd, &ncm, NULL);
+ VkCooperativeMatrixPropertiesKHR *cm = calloc(ncm ? ncm : 1, sizeof *cm);
+ for (uint32_t i = 0; i < ncm; i++) cm[i].sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR;
+ getcm(pd, &ncm, cm);
+ int advertised = 0;
+ for (uint32_t i = 0; i < ncm; i++) {
+ if (cm[i].MSize == (uint32_t)M && cm[i].NSize == (uint32_t)N && cm[i].KSize == (uint32_t)K &&
+ cm[i].AType == ctype(TA) && cm[i].BType == ctype(TB) && cm[i].CType == ctype(TC) &&
+ cm[i].ResultType == ctype(TD) && cm[i].scope == VK_SCOPE_SUBGROUP_KHR)
+ advertised = 1;
+ }
+ free(cm);
+ if (!advertised) { printf(","); die("shape not advertised", 3); }
+ uint32_t nqf = 0;
+ vkGetPhysicalDeviceQueueFamilyProperties(pd, &nqf, NULL);
+ VkQueueFamilyProperties qf[8];
+ if (nqf > 8) nqf = 8;
+ vkGetPhysicalDeviceQueueFamilyProperties(pd, &nqf, qf);
+ uint32_t qfi = UINT32_MAX;
+ for (uint32_t i = 0; i < nqf; i++)
+ if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { qfi = i; break; }
+ if (qfi == UINT32_MAX) die("no compute queue", 4);
+ float prio = 1.0f;
+ VkDeviceQueueCreateInfo qci = {VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO};
+ qci.queueFamilyIndex = qfi; qci.queueCount = 1; qci.pQueuePriorities = &prio;
+ VkPhysicalDeviceCooperativeMatrixFeaturesKHR cmf = {
+ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR};
+ cmf.cooperativeMatrix = VK_TRUE;
+ VkPhysicalDeviceVulkan12Features f12 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &cmf};
+ f12.shaderFloat16 = VK_TRUE;
+ f12.vulkanMemoryModel = VK_TRUE;
+ f12.vulkanMemoryModelDeviceScope = VK_TRUE;
+ VkPhysicalDeviceVulkan11Features f11 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &f12};
+ f11.storageBuffer16BitAccess = VK_TRUE;
+ const char *devext[] = {VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME};
+ VkDeviceCreateInfo dci = {VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f11};
+ dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &qci;
+ dci.enabledExtensionCount = 1; dci.ppEnabledExtensionNames = devext;
+ VkDevice dev;
+ if (vkCreateDevice(pd, &dci, NULL, &dev) != VK_SUCCESS) die("vkCreateDevice", 4);
+ VkQueue queue;
+ vkGetDeviceQueue(dev, qfi, 0, &queue);
+
+ /* ---- shader ---- */
+ int sa = colmaj ? M : K, sb = colmaj ? K : N, sc = colmaj ? M : N;
+ const char *lay = colmaj ? "gl_CooperativeMatrixLayoutColumnMajor" : "gl_CooperativeMatrixLayoutRowMajor";
+ char body[1024];
+ if (precise) {
+ /* C is loaded, then D = C * s + C with NoContraction (from `precise`). */
+ snprintf(body, sizeof body,
+ " coopMatLoad(C, c, 0, %d, %s);\n"
+ " precise coopmat<%s, gl_ScopeSubgroup, %d, %d, gl_MatrixUseAccumulator> D =\n"
+ " C * 1.0078125 + C;\n"
+ " coopMatStore(D, d, doff, %d, %s);\n",
+ sc * cscale, lay, gltype(TC), M, N, sc * dscale, lay);
+ } else {
+ snprintf(body, sizeof body,
+ " coopMatLoad(A, a, 0, %d, %s);\n"
+ " coopMatLoad(B, b, 0, %d, %s);\n"
+ " coopMatLoad(C, c, 0, %d, %s);\n"
+ "%s"
+ " coopmat<%s, gl_ScopeSubgroup, %d, %d, gl_MatrixUseAccumulator> D = coopMatMulAdd(A, B, C);\n"
+ " coopMatStore(D, d, doff, %d, %s);\n",
+ sa * ascale, lay, sb * bscale, lay, sc * cscale, lay,
+ divergent ? " if (gl_SubgroupInvocationID == 0u) A[0] = float16_t(0.0);\n" : "",
+ gltype(TD), M, N, sc * dscale, lay);
+ }
+ char src[4096];
+ snprintf(src, sizeof src,
+ "#version 450\n"
+ "#extension GL_KHR_cooperative_matrix : enable\n"
+ "#extension GL_KHR_memory_scope_semantics : enable\n"
+ "#extension GL_KHR_shader_subgroup_basic : enable\n"
+ "%s"
+ "#extension GL_EXT_shader_explicit_arithmetic_types_float16 : enable\n"
+ "#extension GL_EXT_shader_explicit_arithmetic_types_float32 : enable\n"
+ "layout(local_size_x=%d) in;\n"
+ "layout(binding=0) readonly buffer InA { %s a[]; };\n"
+ "layout(binding=1) readonly buffer InB { %s b[]; };\n"
+ "layout(binding=2) readonly buffer InC { %s c[]; };\n"
+ "layout(binding=3) writeonly buffer OutD { %s d[]; };\n"
+ "void main() {\n"
+ " uint doff = gl_SubgroupID * %du;\n"
+ " coopmat<%s, gl_ScopeSubgroup, %d, %d, gl_MatrixUseA> A;\n"
+ " coopmat<%s, gl_ScopeSubgroup, %d, %d, gl_MatrixUseB> B;\n"
+ " coopmat<%s, gl_ScopeSubgroup, %d, %d, gl_MatrixUseAccumulator> C;\n"
+ "%s"
+ "}\n",
+ vexts, lsize,
+ vgltype, vgltype, vgltype, vgltype, M * N * dscale,
+ gltype(TA), M, K, gltype(TB), K, N, gltype(TC), M, N,
+ body);
+ char path[256], spvpath[256], cmd[1024];
+ snprintf(path, sizeof path, "/tmp/cmshape2_%d.comp", (int)getpid());
+ snprintf(spvpath, sizeof spvpath, "/tmp/cmshape2_%d.spv", (int)getpid());
+ FILE *f = fopen(path, "w");
+ if (!f) die("write comp", 4);
+ fputs(src, f);
+ fclose(f);
+ snprintf(cmd, sizeof cmd, "glslangValidator -V --target-env vulkan1.3 %s -o %s > %s.log 2>&1", path,
+ spvpath, path);
+ if (system(cmd) != 0) die("glslangValidator failed", 4);
+ f = fopen(spvpath, "rb");
+ if (!f) die("open spv", 4);
+ fseek(f, 0, SEEK_END);
+ long spvlen = ftell(f);
+ fseek(f, 0, SEEK_SET);
+ uint32_t *spv = malloc(spvlen);
+ if (fread(spv, 1, spvlen, f) != (size_t)spvlen) die("read spv", 4);
+ fclose(f);
+ VkShaderModuleCreateInfo smci = {VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO};
+ smci.codeSize = spvlen; smci.pCode = spv;
+ VkShaderModule mod;
+ if (vkCreateShaderModule(dev, &smci, NULL, &mod) != VK_SUCCESS)
+ die("vkCreateShaderModule", 4);
+
+ /* ---- descriptors ---- */
+ VkDescriptorSetLayoutBinding bnd[4];
+ for (int i = 0; i < 4; i++) {
+ bnd[i] = (VkDescriptorSetLayoutBinding){i, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, NULL};
+ }
+ VkDescriptorSetLayoutCreateInfo dslci = {VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO};
+ dslci.bindingCount = 4; dslci.pBindings = bnd;
+ VkDescriptorSetLayout dsl;
+ if (vkCreateDescriptorSetLayout(dev, &dslci, NULL, &dsl) != VK_SUCCESS) die("dsl", 4);
+ VkPipelineLayoutCreateInfo plci = {VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO};
+ plci.setLayoutCount = 1; plci.pSetLayouts = &dsl;
+ VkPipelineLayout pl;
+ if (vkCreatePipelineLayout(dev, &plci, NULL, &pl) != VK_SUCCESS) die("pl", 4);
+ VkComputePipelineCreateInfo cpci = {VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO};
+ cpci.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
+ cpci.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
+ cpci.stage.module = mod; cpci.stage.pName = "main";
+ cpci.layout = pl;
+ VkPipeline pipe;
+ if (vkCreateComputePipelines(dev, VK_NULL_HANDLE, 1, &cpci, NULL, &pipe) != VK_SUCCESS)
+ die("vkCreateComputePipelines", 4);
+ VkDescriptorPoolSize ps = {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 4};
+ VkDescriptorPoolCreateInfo dpci = {VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO};
+ dpci.maxSets = 1; dpci.poolSizeCount = 1; dpci.pPoolSizes = &ps;
+ VkDescriptorPool pool;
+ if (vkCreateDescriptorPool(dev, &dpci, NULL, &pool) != VK_SUCCESS) die("pool", 4);
+ VkDescriptorSetAllocateInfo dsai = {VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO};
+ dsai.descriptorPool = pool; dsai.descriptorSetCount = 1; dsai.pSetLayouts = &dsl;
+ VkDescriptorSet set;
+ if (vkAllocateDescriptorSets(dev, &dsai, &set) != VK_SUCCESS) die("set", 4);
+
+ /* ---- buffers ---- */
+ VkPhysicalDeviceMemoryProperties mp;
+ vkGetPhysicalDeviceMemoryProperties(pd, &mp);
+ size_t nelem[4] = {(size_t)M * K, (size_t)K * N, (size_t)M * N, (size_t)M * N * subgroups};
+ int e16[4] = {is16(TA), is16(TB), is16(TC), is16(TD)};
+ VkBuffer buf[4];
+ VkDeviceMemory mem[4];
+ void *map[4];
+ for (int i = 0; i < 4; i++) {
+ VkDeviceSize sz = nelem[i] * (e16[i] ? 2 : 4) + 64;
+ VkBufferCreateInfo bci = {VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO};
+ bci.size = sz; bci.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
+ if (vkCreateBuffer(dev, &bci, NULL, &buf[i]) != VK_SUCCESS) die("vkCreateBuffer", 4);
+ VkMemoryRequirements req;
+ vkGetBufferMemoryRequirements(dev, buf[i], &req);
+ uint32_t mt = UINT32_MAX;
+ for (uint32_t t = 0; t < mp.memoryTypeCount; t++) {
+ VkMemoryPropertyFlags want = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
+ if ((req.memoryTypeBits & (1u << t)) && (mp.memoryTypes[t].propertyFlags & want) == want) { mt = t; break; }
+ }
+ if (mt == UINT32_MAX) die("no host-visible memory", 4);
+ VkMemoryAllocateInfo mai = {VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO};
+ mai.allocationSize = req.size; mai.memoryTypeIndex = mt;
+ if (vkAllocateMemory(dev, &mai, NULL, &mem[i]) != VK_SUCCESS) die("vkAllocateMemory", 4);
+ vkBindBufferMemory(dev, buf[i], mem[i], 0);
+ vkMapMemory(dev, mem[i], 0, VK_WHOLE_SIZE, 0, &map[i]);
+ memset(map[i], 0, sz);
+ }
+ double *A = calloc(M * K, sizeof(double)), *B = calloc(K * N, sizeof(double)), *C = calloc(M * N, sizeof(double));
+ int identity = !strcmp(mode, "identity"), ones = !strcmp(mode, "ones") || divergent, randf = !strcmp(mode, "randf") || precise;
+ if (identity && M != K) die("identity needs M==K", 4);
+ for (int r = 0; r < M; r++)
+ for (int k = 0; k < K; k++)
+ A[r * K + k] = identity ? (r == k) : ones ? 1.0 : randf ? ((int)(rnd() % 2001) - 1000) / 1000.0 : (double)((int)(rnd() % 7) - 3);
+ for (int k = 0; k < K; k++)
+ for (int c = 0; c < N; c++)
+ B[k * N + c] = ones ? 1.0 : randf ? ((int)(rnd() % 2001) - 1000) / 1000.0 : (double)((int)(rnd() % 7) - 3);
+ for (int r = 0; r < M; r++)
+ for (int c = 0; c < N; c++)
+ C[r * N + c] = (identity || ones) ? 0.0 : randf ? ((int)(rnd() % 2001) - 1000) / 1000.0 : (double)((int)(rnd() % 7) - 3);
+ for (int r = 0; r < M; r++)
+ for (int k = 0; k < K; k++) {
+ size_t idx = colmaj ? (size_t)k * M + r : (size_t)r * K + k;
+ setv(map[0], e16[0], idx, A[r * K + k]);
+ A[r * K + k] = getv(map[0], e16[0], idx);
+ }
+ for (int k = 0; k < K; k++)
+ for (int c = 0; c < N; c++) {
+ size_t idx = colmaj ? (size_t)c * K + k : (size_t)k * N + c;
+ setv(map[1], e16[1], idx, B[k * N + c]);
+ B[k * N + c] = getv(map[1], e16[1], idx);
+ }
+ for (int r = 0; r < M; r++)
+ for (int c = 0; c < N; c++) {
+ size_t idx = colmaj ? (size_t)c * M + r : (size_t)r * N + c;
+ setv(map[2], e16[2], idx, C[r * N + c]);
+ C[r * N + c] = getv(map[2], e16[2], idx);
+ }
+ for (size_t i = 0; i < nelem[3]; i++) setv(map[3], e16[3], i, -12345.0);
+ (void)A; (void)B;
+
+ /* ---- dispatch ---- */
+ VkDescriptorBufferInfo dbi[4];
+ VkWriteDescriptorSet w[4];
+ for (int i = 0; i < 4; i++) {
+ dbi[i] = (VkDescriptorBufferInfo){buf[i], 0, VK_WHOLE_SIZE};
+ w[i] = (VkWriteDescriptorSet){VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, NULL, set, (uint32_t)i, 0, 1,
+ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, NULL, &dbi[i], NULL};
+ }
+ vkUpdateDescriptorSets(dev, 4, w, 0, NULL);
+ VkCommandPoolCreateInfo cpc = {VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO};
+ cpc.queueFamilyIndex = qfi;
+ VkCommandPool cpool;
+ if (vkCreateCommandPool(dev, &cpc, NULL, &cpool) != VK_SUCCESS) die("cpool", 4);
+ VkCommandBufferAllocateInfo cbai = {VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO};
+ cbai.commandPool = cpool; cbai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cbai.commandBufferCount = 1;
+ VkCommandBuffer cb;
+ if (vkAllocateCommandBuffers(dev, &cbai, &cb) != VK_SUCCESS) die("cb", 4);
+ VkCommandBufferBeginInfo cbbi = {VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO};
+ vkBeginCommandBuffer(cb, &cbbi);
+ vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe);
+ vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pl, 0, 1, &set, 0, NULL);
+ vkCmdDispatch(cb, 1, 1, 1);
+ vkEndCommandBuffer(cb);
+ VkFenceCreateInfo fci = {VK_STRUCTURE_TYPE_FENCE_CREATE_INFO};
+ VkFence fence;
+ vkCreateFence(dev, &fci, NULL, &fence);
+ VkSubmitInfo si = {VK_STRUCTURE_TYPE_SUBMIT_INFO};
+ si.commandBufferCount = 1; si.pCommandBuffers = &cb;
+ if (vkQueueSubmit(queue, 1, &si, fence) != VK_SUCCESS) die("submit", 4);
+ if (vkWaitForFences(dev, 1, &fence, VK_TRUE, 10000000000ull) != VK_SUCCESS) die("fence timeout", 4);
+
+ /* ---- check ---- */
+ int bad = 0, shown = 0;
+ for (int sg = 0; sg < subgroups; sg++) {
+ int reduced = 0;
+ for (int r = 0; r < M; r++)
+ for (int c = 0; c < N; c++) {
+ double ref;
+ if (precise) {
+ /* unfused fp32: round(c * s), then round(sum) */
+ float prod = (float)C[r * N + c] * 1.0078125f;
+ ref = (double)((float)(prod + (float)C[r * N + c]));
+ } else {
+ ref = C[r * N + c];
+ for (int k = 0; k < K; k++) ref += A[r * K + k] * B[k * N + c];
+ ref = (float)ref;
+ }
+ size_t idx = colmaj ? (size_t)c * M + r : (size_t)r * N + c;
+ idx += (size_t)sg * M * N;
+ double got = getv(map[3], e16[3], idx);
+ if (divergent) {
+ reduced += got == K - 1;
+ ref = got == K - 1 ? K - 1 : K;
+ }
+ double err = fabs(got - ref);
+ double tol = randf && !precise ? (e16[3] ? 2e-2 : 1e-4) : 0.0;
+ if (!(err <= tol)) {
+ bad++;
+ if (shown < 300) { fprintf(stderr, "mismatch D[%d][%d]: got %g want %g\n", r, c, got, ref); shown++; }
+ }
+ }
+ if (divergent && reduced != N) bad++;
+ }
+ printf(",\"elements\":%d,\"mismatches\":%d,\"result\":\"%s\"}\n", M * N * subgroups, bad,
+ bad ? "FAIL" : "PASS");
+ remove(path);
+ remove(spvpath);
+ return bad ? 1 : 0;
+}
diff --git a/tests/compute_runner.cpp b/tests/compute_runner.cpp
new file mode 100644
index 000000000000..e16d45d9d612
--- /dev/null
+++ b/tests/compute_runner.cpp
@@ -0,0 +1,196 @@
+// Build: c++ -O2 compute_runner.cpp -lvulkan -o compute_runner
+// Usage: compute_runner shader.spv in.bin out.bin out_bytes groups
+#include
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#define CK(x) \
+ do { \
+ VkResult r_ = (x); \
+ if (r_ != VK_SUCCESS) { \
+ std::fprintf(stderr, "%s failed: %d\n", #x, (int)r_); \
+ std::exit(2); \
+ } \
+ } while (0)
+
+static std::string slurp(const char* p) {
+ std::ifstream f(p, std::ios::binary);
+ if (!f) { std::fprintf(stderr, "open %s\n", p); std::exit(2); }
+ std::ostringstream ss;
+ ss << f.rdbuf();
+ return ss.str();
+}
+
+struct Buf { VkBuffer b; VkDeviceMemory m; void* map; };
+
+static VkPhysicalDeviceMemoryProperties g_mp;
+static Buf mkbuf(VkDevice dev, size_t size) {
+ Buf r{};
+ VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO};
+ bi.size = size;
+ bi.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
+ CK(vkCreateBuffer(dev, &bi, nullptr, &r.b));
+ VkMemoryRequirements mr;
+ vkGetBufferMemoryRequirements(dev, r.b, &mr);
+ uint32_t ti = UINT32_MAX;
+ for (uint32_t i = 0; i < g_mp.memoryTypeCount; ++i) {
+ VkMemoryPropertyFlags want = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
+ VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
+ if ((mr.memoryTypeBits & (1u << i)) &&
+ (g_mp.memoryTypes[i].propertyFlags & want) == want) { ti = i; break; }
+ }
+ if (ti == UINT32_MAX) { std::fprintf(stderr, "no host memtype\n"); std::exit(2); }
+ VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO};
+ ai.allocationSize = mr.size;
+ ai.memoryTypeIndex = ti;
+ CK(vkAllocateMemory(dev, &ai, nullptr, &r.m));
+ CK(vkBindBufferMemory(dev, r.b, r.m, 0));
+ CK(vkMapMemory(dev, r.m, 0, size, 0, &r.map));
+ return r;
+}
+
+int main(int argc, char** argv) {
+ if (argc < 6) { std::fprintf(stderr, "usage\n"); return 2; }
+ std::string spv = slurp(argv[1]);
+ std::string in = slurp(argv[2]);
+ size_t out_bytes = std::strtoull(argv[4], nullptr, 10);
+ uint32_t groups = std::strtoul(argv[5], nullptr, 10);
+
+ VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO};
+ app.apiVersion = VK_API_VERSION_1_1;
+ VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO};
+ ici.pApplicationInfo = &app;
+ VkInstance inst;
+ CK(vkCreateInstance(&ici, nullptr, &inst));
+ uint32_t n = 1;
+ VkPhysicalDevice pd;
+ CK(vkEnumeratePhysicalDevices(inst, &n, &pd));
+ VkPhysicalDeviceProperties props;
+ vkGetPhysicalDeviceProperties(pd, &props);
+ vkGetPhysicalDeviceMemoryProperties(pd, &g_mp);
+ std::fprintf(stderr, "device: %s driver %u.%u.%u\n", props.deviceName,
+ VK_VERSION_MAJOR(props.driverVersion),
+ VK_VERSION_MINOR(props.driverVersion),
+ VK_VERSION_PATCH(props.driverVersion));
+
+ uint32_t qc = 0;
+ vkGetPhysicalDeviceQueueFamilyProperties(pd, &qc, nullptr);
+ std::vector qf(qc);
+ vkGetPhysicalDeviceQueueFamilyProperties(pd, &qc, qf.data());
+ uint32_t qi = 0;
+ for (; qi < qc; ++qi) if (qf[qi].queueFlags & VK_QUEUE_COMPUTE_BIT) break;
+ float prio = 1.0f;
+ VkDeviceQueueCreateInfo dqi{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO};
+ dqi.queueFamilyIndex = qi;
+ dqi.queueCount = 1;
+ dqi.pQueuePriorities = &prio;
+ VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO};
+ dci.queueCreateInfoCount = 1;
+ dci.pQueueCreateInfos = &dqi;
+ VkDevice dev;
+ CK(vkCreateDevice(pd, &dci, nullptr, &dev));
+ VkQueue q;
+ vkGetDeviceQueue(dev, qi, 0, &q);
+
+ Buf bin = mkbuf(dev, in.size());
+ Buf bout = mkbuf(dev, out_bytes);
+ std::memcpy(bin.map, in.data(), in.size());
+ std::memset(bout.map, 0, out_bytes);
+
+ VkDescriptorSetLayoutBinding binds[2]{};
+ for (int i = 0; i < 2; ++i) {
+ binds[i].binding = i;
+ binds[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+ binds[i].descriptorCount = 1;
+ binds[i].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
+ }
+ VkDescriptorSetLayoutCreateInfo dsl{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO};
+ dsl.bindingCount = 2;
+ dsl.pBindings = binds;
+ VkDescriptorSetLayout layout;
+ CK(vkCreateDescriptorSetLayout(dev, &dsl, nullptr, &layout));
+ VkPipelineLayoutCreateInfo pli{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO};
+ pli.setLayoutCount = 1;
+ pli.pSetLayouts = &layout;
+ VkPipelineLayout pl;
+ CK(vkCreatePipelineLayout(dev, &pli, nullptr, &pl));
+ VkShaderModuleCreateInfo smi{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO};
+ smi.codeSize = spv.size();
+ smi.pCode = (const uint32_t*)spv.data();
+ VkShaderModule sm;
+ CK(vkCreateShaderModule(dev, &smi, nullptr, &sm));
+ VkComputePipelineCreateInfo cpi{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO};
+ cpi.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
+ cpi.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
+ cpi.stage.module = sm;
+ cpi.stage.pName = "main";
+ cpi.layout = pl;
+ VkPipeline pipe;
+ CK(vkCreateComputePipelines(dev, VK_NULL_HANDLE, 1, &cpi, nullptr, &pipe));
+
+ VkDescriptorPoolSize ps{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 2};
+ VkDescriptorPoolCreateInfo dpi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO};
+ dpi.maxSets = 1;
+ dpi.poolSizeCount = 1;
+ dpi.pPoolSizes = &ps;
+ VkDescriptorPool pool;
+ CK(vkCreateDescriptorPool(dev, &dpi, nullptr, &pool));
+ VkDescriptorSetAllocateInfo dsa{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO};
+ dsa.descriptorPool = pool;
+ dsa.descriptorSetCount = 1;
+ dsa.pSetLayouts = &layout;
+ VkDescriptorSet ds;
+ CK(vkAllocateDescriptorSets(dev, &dsa, &ds));
+ VkDescriptorBufferInfo dbi[2] = {{bin.b, 0, VK_WHOLE_SIZE}, {bout.b, 0, VK_WHOLE_SIZE}};
+ VkWriteDescriptorSet w[2]{};
+ for (int i = 0; i < 2; ++i) {
+ w[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
+ w[i].dstSet = ds;
+ w[i].dstBinding = i;
+ w[i].descriptorCount = 1;
+ w[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+ w[i].pBufferInfo = &dbi[i];
+ }
+ vkUpdateDescriptorSets(dev, 2, w, 0, nullptr);
+
+ VkCommandPoolCreateInfo cpci{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO};
+ cpci.queueFamilyIndex = qi;
+ VkCommandPool cp;
+ CK(vkCreateCommandPool(dev, &cpci, nullptr, &cp));
+ VkCommandBufferAllocateInfo cba{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO};
+ cba.commandPool = cp;
+ cba.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
+ cba.commandBufferCount = 1;
+ VkCommandBuffer cb;
+ CK(vkAllocateCommandBuffers(dev, &cba, &cb));
+ VkCommandBufferBeginInfo cbb{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO};
+ CK(vkBeginCommandBuffer(cb, &cbb));
+ vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe);
+ vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pl, 0, 1, &ds, 0, nullptr);
+ vkCmdDispatch(cb, groups, 1, 1);
+ VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER};
+ barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
+ barrier.dstAccessMask = VK_ACCESS_HOST_READ_BIT;
+ vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
+ VK_PIPELINE_STAGE_HOST_BIT, 0, 1, &barrier, 0, nullptr, 0, nullptr);
+ CK(vkEndCommandBuffer(cb));
+
+ VkFenceCreateInfo fci{VK_STRUCTURE_TYPE_FENCE_CREATE_INFO};
+ VkFence fence;
+ CK(vkCreateFence(dev, &fci, nullptr, &fence));
+ VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO};
+ si.commandBufferCount = 1;
+ si.pCommandBuffers = &cb;
+ CK(vkQueueSubmit(q, 1, &si, fence));
+ CK(vkWaitForFences(dev, 1, &fence, VK_TRUE, 60000000000ull));
+ std::ofstream of(argv[3], std::ios::binary);
+ of.write((const char*)bout.map, out_bytes);
+ vkDeviceWaitIdle(dev);
+ return 0;
+}
diff --git a/tests/precise_math.comp b/tests/precise_math.comp
new file mode 100644
index 000000000000..827941d24688
--- /dev/null
+++ b/tests/precise_math.comp
@@ -0,0 +1,17 @@
+#version 450
+layout(local_size_x = 64) in;
+layout(std430, binding = 0) readonly buffer In { vec4 v[]; } inp;
+layout(std430, binding = 1) writeonly buffer Out { float o[]; } outp;
+void main() {
+ uint i = gl_GlobalInvocationID.x;
+ vec4 x = inp.v[i];
+ precise float t = log2(x.w);
+ precise float literal = t * 0.6931471805599453;
+ precise float from_buffer = t * x.z;
+ outp.o[i * 6 + 0] = x.x / x.y;
+ outp.o[i * 6 + 1] = 1.0 / x.y;
+ outp.o[i * 6 + 2] = literal;
+ outp.o[i * 6 + 3] = from_buffer;
+ outp.o[i * 6 + 4] = t;
+ outp.o[i * 6 + 5] = log(x.w);
+}
diff --git a/tests/precise_math.py b/tests/precise_math.py
new file mode 100644
index 000000000000..7e322251c73e
--- /dev/null
+++ b/tests/precise_math.py
@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""Usage: precise_math.py ICD COMPUTE_RUNNER; glslangValidator and numpy are required."""
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import numpy as np
+
+icd, runner = map(os.path.abspath, sys.argv[1:])
+here = Path(__file__).resolve().parent
+rng = np.random.default_rng(20260913)
+n = 262144
+inp = rng.integers(0, 2**32, (n, 4), dtype=np.uint32)
+edges = np.array([0, 0x80000000, 1, 0x7fffff, 0x800000, 0x800001,
+ 0x80800000, 0x3f800000, 0x3f7fffff, 0x3f800001,
+ 0x40000000, 0x7f7fffff, 0xff7fffff, 0x7f800000,
+ 0xff800000, 0x7fc00000], dtype=np.uint32)
+inp[:256, 0] = np.repeat(edges, len(edges))
+inp[:256, 1] = np.tile(edges, len(edges))
+inp[256:65536, 0] = (inp[256:65536, 0] & 0x807fffff) | (rng.integers(1, 40, 65280, dtype=np.uint32) << 23)
+inp[256:65536, 1] = (inp[256:65536, 1] & 0x807fffff) | (rng.integers(80, 160, 65280, dtype=np.uint32) << 23)
+inp[:, 2] = 0x3f317218
+inp[:, 3] = rng.integers(0x00800000, 0x7f800000, n, dtype=np.uint32)
+inp[:7, 3] = [0x3f800000, 0x40000000, 0x40400000, 0x3f7fffff, 0x3f800001, 0x00800000, 0x7f7fffff]
+
+def quotient(a, b):
+ a = a.copy()
+ b = b.copy()
+ a[(a & 0x7fffffff) < 0x800000] &= 0x80000000
+ b[(b & 0x7fffffff) < 0x800000] &= 0x80000000
+ with np.errstate(all="ignore"):
+ exact = a.view(np.float32).astype(np.float64) / b.view(np.float32).astype(np.float64)
+ rounded = exact.astype(np.float32).view(np.uint32)
+ flushed = np.abs(exact) < 2.0**-126
+ rounded[flushed] = (a[flushed] ^ b[flushed]) & 0x80000000
+ return rounded
+
+def mismatch(got, want):
+ both_nan = ((got & 0x7fffffff) > 0x7f800000) & ((want & 0x7fffffff) > 0x7f800000)
+ return np.flatnonzero((got != want) & ~both_nan)
+
+with tempfile.TemporaryDirectory(prefix="mesa-math-") as td:
+ td = Path(td)
+ inp.tofile(td / "in.bin")
+ subprocess.run(["glslangValidator", "-V", str(here / "precise_math.comp"), "-o", str(td / "probe.spv")], check=True, timeout=30)
+ env = dict(os.environ, VK_DRIVER_FILES=icd, VK_ICD_FILENAMES=icd, MESA_SHADER_CACHE_DISABLE="true")
+ subprocess.run(["flock", "-w", "30", "/tmp/m1-gpu.lock", runner,
+ str(td / "probe.spv"), str(td / "in.bin"), str(td / "out.bin"),
+ str(n * 6 * 4), str(n // 64)], env=env, check=True, timeout=120)
+ out = np.fromfile(td / "out.bin", dtype=np.uint32).reshape(n, 6)
+
+checks = {
+ "division": mismatch(out[:, 0], quotient(inp[:, 0], inp[:, 1])),
+ "reciprocal": mismatch(out[:, 1], quotient(np.full(n, 0x3f800000, dtype=np.uint32), inp[:, 1])),
+ "literal_vs_buffer": mismatch(out[:, 2], out[:, 3]),
+ "precise_multiply": mismatch(out[:, 2], (out[:, 4].copy().view(np.float32) * np.float32(np.log(2.0))).view(np.uint32)),
+}
+x = inp[:, 3].copy().view(np.float32).astype(np.float64)
+for name, column, ref in [("log", 5, np.log(x)), ("log2", 4, np.log2(x))]:
+ want = ref.astype(np.float32).view(np.int32).astype(np.int64)
+ got = out[:, column].copy().view(np.int32).astype(np.int64)
+ ulp = np.abs(want - got)
+ checks[name + "_over_1ulp"] = np.flatnonzero(ulp > 1)
+ print(name, "max_ulp", int(ulp.max()))
+result = {"pairs": n, "failures": {k: len(v) for k, v in checks.items()}}
+for name, indices in checks.items():
+ if len(indices):
+ print(name, [{"input": [hex(int(v)) for v in inp[i]], "output": [hex(int(v)) for v in out[i]]} for i in indices[:5]])
+print(json.dumps(result, sort_keys=True))
+sys.exit(any(len(v) for v in checks.values()))
diff --git a/tests/run_matrix2.sh b/tests/run_matrix2.sh
new file mode 100755
index 000000000000..b8e8740ffad4
--- /dev/null
+++ b/tests/run_matrix2.sh
@@ -0,0 +1,43 @@
+#!/bin/sh
+# Build: cc -O2 -ffp-contract=off cmshape2.c -lvulkan -lm -o cmshape2
+# Usage: run_matrix2.sh /absolute/path/to/icd.json
+ICD="${1:?usage: run_matrix2.sh }"
+cd "$(dirname "$0")" || exit 1
+pass=0; fail=0
+run() {
+ desc="$1"; shift
+ out=$(VK_DRIVER_FILES="$ICD" VK_ICD_FILENAMES="$ICD" AGX_SIMDMAT=1 MESA_SHADER_CACHE_DISABLE=true flock -w 30 /tmp/m1-gpu.lock timeout --kill-after=5s 120 ./cmshape2 "$@" 2>&1)
+ rc=$?
+ echo "$out"
+ if [ $rc -eq 0 ]; then pass=$((pass+1)); else fail=$((fail+1)); echo "^^ FAIL: $desc" >&2; fi
+}
+
+echo "=== HW baseline (lsize 32, must always pass) ==="
+for shape in "8 8 8" "16 16 16"; do
+ for types in "f32 f32 f32 f32" "f16 f16 f32 f32" "f16 f16 f16 f16"; do
+ for lay in row col; do
+ run "hw $shape $types $lay" $shape $types $lay randint 7 32
+ done
+ done
+done
+
+echo "=== workgroup sizes (SW path after fix) ==="
+for lsize in 1 8 16 48 33; do
+ run "w $lsize f32" 8 8 8 f32 f32 f32 f32 row randint 7 $lsize
+ run "w $lsize f16" 16 16 16 f16 f16 f16 f16 row randint 7 $lsize
+done
+
+echo "=== pointer forms ==="
+run "ptr u16 f16" 8 8 8 f16 f16 f16 f16 row randint 7 32 u16
+run "ptr u16 f32" 8 8 8 f32 f32 f32 f32 row randint 7 32 u16
+run "ptr u8 f16" 8 8 8 f16 f16 f16 f16 row randint 7 32 u8
+run "ptr u16 f32 col" 16 16 16 f32 f32 f32 f32 col randint 7 32 u16
+
+echo "=== precise contraction (NoContraction must survive) ==="
+run "precise l32" 8 8 8 f32 f32 f32 f32 row precise 7 32 f32
+run "precise l16" 8 8 8 f32 f32 f32 f32 row precise 7 16 f32
+echo "=== divergent component updates ==="
+run "divergent l16" 8 8 8 f16 f16 f32 f32 row divergent 7 16 u16
+run "divergent l33" 8 8 8 f16 f16 f32 f32 col divergent 7 33 u16
+echo "SUMMARY pass=$pass fail=$fail"
+test "$fail" -eq 0
From edf760cb4b14f797d1bd06f552bd6ed29b1eb723 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 5 Oct 2026 12:05:12 -0500
Subject: [PATCH 15/16] agx: lower log() the same way regardless of where ln2
comes from
Dj's stress runs found (Basecamp 10277024959, 2026-09-13) that a precise
log2(x) * ln2 gave different bits depending on whether ln2 was a literal
or read from a buffer, on 26.8% of inputs. The stack carries an amended
form of the log fix in which a multiply of a hardware flog2 by the fp32
ln2 constant is rewritten to the correctly rounded soft log. That reopens
the finding from the other side: the literal form now gets the
single-rounding soft log, while the buffer form still computes
soft log2 and multiplies, so the two forms of the same source expression
differ by 1 ulp on 25.3% of inputs (66,332 of 262,144 in the
literal_vs_buffer reproducer), and the precise multiply check fails the
same way.
Drop the literal-only rewrite. log2 is lowered to the double-float soft
log2 everywhere, and log() keeps the one rounded multiply the source
asks for, from a literal or from a buffer: the two forms now agree bit
for bit on every input, powers of two still come out exact because the
series returns the exponent unchanged, and log stays within 1 ulp on the
262,144-input reproducer.
(cherry picked from commit 155e743bd26b9a00e7c23797f7ccc88a6ed50c0f)
---
src/asahi/compiler/agx_nir_lower_math.c | 79 +------------------------
1 file changed, 3 insertions(+), 76 deletions(-)
diff --git a/src/asahi/compiler/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c
index 09b764c1aa1e..fc9d1a1e4569 100644
--- a/src/asahi/compiler/agx_nir_lower_math.c
+++ b/src/asahi/compiler/agx_nir_lower_math.c
@@ -136,7 +136,7 @@ agx_nir_lower_fdiv_late(nir_shader *s)
}
/*
- * Natural log and log2 as double-float (fp32 pair) arithmetic.
+ * Log2 as double-float (fp32 pair) arithmetic.
*
* x = 2^e m with m in [sqrt(1/2), sqrt(2)). With s = (m - 1) / (m + 1),
*
@@ -145,8 +145,7 @@ agx_nir_lower_fdiv_late(nir_shader *s)
* s is kept as a pair (s, s_lo) from one correctly rounded division plus the
* fma residual, the cubic term is kept as a pair, and the terms are summed
* with Fast2Sum so the value handed to the final rounding carries about 40
- * bits. That makes the exact cases (powers of two, log(3)) come out correctly
- * rounded and everything else faithful, without any table.
+ * bits without a lookup table.
*/
struct log_parts {
nir_def *e; /* exponent as float */
@@ -247,26 +246,6 @@ log_special(nir_builder *b, nir_def *x)
nir_imm_float(b, INFINITY)));
}
-static nir_def *
-soft_log(nir_builder *b, nir_def *x)
-{
- struct log_parts l = log_core(b, x);
-
- /* ln2 split so e * LN2_HI is exact for |e| < 2^8 */
- nir_def *a = nir_fmul_imm(b, l.e, 0.693145751953125);
- nir_def *bb = nir_fmul_imm(b, l.s, 2.0);
- nir_def *h = nir_fadd(b, a, bb);
- nir_def *lo = fast2sum_lo(b, a, bb, h);
- nir_def *h2 = nir_fadd(b, h, l.t_hi);
- nir_def *lo2 = fast2sum_lo(b, h, l.t_hi, h2);
-
- nir_def *small = nir_fadd(b, lo, lo2);
- small = nir_fadd(b, small, l.t_lo);
- small = nir_fadd(b, small, nir_fmul_imm(b, l.s_lo, 2.0));
- small = nir_ffma(b, l.e, nir_imm_float(b, 1.428606765330187e-06), small);
- return nir_fadd(b, h2, small);
-}
-
static nir_def *
soft_log2(nir_builder *b, nir_def *x)
{
@@ -291,54 +270,6 @@ soft_log2(nir_builder *b, nir_def *x)
return nir_fadd(b, h2, nir_fadd(b, nir_fadd(b, lo, lo2), pl));
}
-static const uint32_t LN2_F32_BITS = 0x3f317218;
-
-static bool
-src_is_ln2(nir_alu_instr *alu, unsigned i)
-{
- nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i);
- return nir_scalar_is_const(s) && nir_scalar_as_uint(s) == LN2_F32_BITS;
-}
-
-static bool
-src_is_flog2(nir_alu_instr *alu, unsigned i)
-{
- nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i);
- return nir_scalar_is_alu(s) && nir_scalar_alu_op(s) == nir_op_flog2;
-}
-
-/* nir_flog builds fmul(flog2(x), ln2); catch that before the flog2 goes. */
-static bool
-lower_log(nir_builder *b, nir_alu_instr *alu, void *_)
-{
- if (alu->op != nir_op_fmul || alu->def.bit_size != 32 ||
- alu->def.num_components != 1)
- return false;
-
- unsigned li;
- if (src_is_flog2(alu, 0) && src_is_ln2(alu, 1))
- li = 0;
- else if (src_is_flog2(alu, 1) && src_is_ln2(alu, 0))
- li = 1;
- else
- return false;
-
- nir_scalar log = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), li);
- nir_scalar xs = nir_scalar_chase_alu_src(log, 0);
-
- b->cursor = nir_before_instr(&alu->instr);
- b->fp_math_ctrl = nir_fp_no_fast_math;
- nir_def *x = nir_channel(b, xs.def, xs.comp);
- nir_def *res = nir_bcsel(b, is_positive_normal(b, x), soft_log(b, x),
- log_special(b, x));
- nir_def_replace(&alu->def, res);
-
- if (nir_def_is_unused(log.def))
- nir_instr_remove(nir_def_instr(log.def));
-
- return true;
-}
-
static bool
lower_log2(nir_builder *b, nir_alu_instr *alu, void *_)
{
@@ -357,11 +288,7 @@ lower_log2(nir_builder *b, nir_alu_instr *alu, void *_)
bool
agx_nir_lower_log(nir_shader *s)
{
- bool progress =
- nir_shader_alu_pass(s, lower_log, nir_metadata_control_flow, NULL);
- progress |=
- nir_shader_alu_pass(s, lower_log2, nir_metadata_control_flow, NULL);
- return progress;
+ return nir_shader_alu_pass(s, lower_log2, nir_metadata_control_flow, NULL);
}
/*
From c56e08634bba83c00b7345494a4cef9bdf681686 Mon Sep 17 00:00:00 2001
From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com>
Date: Mon, 5 Oct 2026 15:11:27 -0500
Subject: [PATCH 16/16] hk: order internal full-barrier precomp launches with
surrounding CDM work
---
src/asahi/vulkan/hk_cmd_buffer.c | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/src/asahi/vulkan/hk_cmd_buffer.c b/src/asahi/vulkan/hk_cmd_buffer.c
index cb3c8b73182b..ea79700896bc 100644
--- a/src/asahi/vulkan/hk_cmd_buffer.c
+++ b/src/asahi/vulkan/hk_cmd_buffer.c
@@ -799,9 +799,22 @@ hk_dispatch_precomp(struct hk_cmd_buffer *cmd, struct agx_grid grid,
agx_usc_words_precomp(t.cpu, &prog->b, uploaded_data, data_size);
+ /* An internal full-barrier launch orders in both directions: drain the
+ * producers before it, and it drains before the next CDM launch. The
+ * producer side is not tracked by cdm_barrier_pending, since an internal
+ * launch may consume the output of a regular launch (e.g. the TCS of the
+ * same draw) that sets no flag. */
+ if (barrier & AGX_BARRIER_ALL) {
+ hk_cdm_cache_flush(dev, cs);
+ cmd->state.cs.cdm_barrier_pending = false;
+ }
+
hk_dispatch_with_usc_launch(dev, cs, prog->b.launch,
agx_usc_addr(&dev->dev, t.gpu), grid,
prog->b.workgroup);
+
+ if (barrier & AGX_BARRIER_ALL)
+ cmd->state.cs.cdm_barrier_pending = true;
}
void