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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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 5813204f0bfa06b773d98a37d2ce620bf8017529 Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Sun, 13 Sep 2026 10:00:56 -0500 Subject: [PATCH 12/25] agx: lower cooperative matrices in software outside full subgroups The G13 matrix tile spans all 32 lanes of a SIMD group, so the hardware lowering is only correct when every subgroup of the dispatch is fully populated. Any workgroup whose size is not a multiple of 32 took it anyway and returned garbage: local size 16, 8 or 1 got 62 to 64 of 64 elements wrong, and those are valid shaders. Route them to a software lowering instead. Every active invocation keeps a private element array, but only the slots it owns are authoritative: logical element e is owned by invocation e % A, where A is the active invocation count of this subgroup, so partial tail subgroups and runtime-variable workgroup sizes are handled. Tangled operations (load, store, muladd, conversions, arithmetic) run on every lane with uniform bounds and gather non-owned source elements from their owner; untangled ones (insert, extract, copy, bitcast, length) touch local slots only and stay valid on a subset of lanes, per the tangled-instruction list in SPV_KHR_cooperative_matrix rev 10. A valid shader is never rejected. Reported by Dj. The comments state the default-on switch that this series already has (AGX_SIMDMAT=0 disables it), with the wording of 1b65ba39527. (cherry picked from commit 109e7cc842cee876e1f0042a374bc0f5bbd40e38) --- src/asahi/compiler/agx_compile.c | 25 +- src/asahi/compiler/agx_compile.h | 4 +- src/asahi/compiler/agx_compiler.h | 1 + src/asahi/compiler/agx_nir_lower_cmat.c | 572 ++++++++++++++++++++++++ src/asahi/compiler/meson.build | 1 + 5 files changed, 597 insertions(+), 6 deletions(-) create mode 100644 src/asahi/compiler/agx_nir_lower_cmat.c diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index d5ad4a2d4de8..541ab79007df 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -3523,11 +3523,28 @@ 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. AGX_SIMDMAT=0 skips it. + /* Lower VK_KHR_cooperative_matrix before vars_to_scratch can spill the cmat + * temps. The G13 matrix tile spans all 32 lanes of a subgroup, so the HW + * path requires fully-populated subgroups: compute stage, static workgroup + * size, size a multiple of 32. Everything else lowers through the software + * path (agx_nir_lower_cmat.c) -- a valid shader is never rejected. On by + * default; AGX_SIMDMAT=0 disables it. */ - if (agx_simdmat_enabled()) - NIR_PASS(_, nir, agx_nir_lower_simdmat, 32); + if (agx_simdmat_enabled()) { + bool full_subgroups = + nir->info.stage == MESA_SHADER_COMPUTE && + !nir->info.workgroup_size_variable && + nir->info.workgroup_size[0] && nir->info.workgroup_size[1] && + nir->info.workgroup_size[2] && + (((uint64_t)nir->info.workgroup_size[0] * + nir->info.workgroup_size[1] * nir->info.workgroup_size[2]) % + 32) == 0; + + if (full_subgroups) + NIR_PASS(_, nir, agx_nir_lower_simdmat, 32); + else + NIR_PASS(_, nir, agx_nir_lower_cmat); + } 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 30db575d03b7..246952739821 100644 --- a/src/asahi/compiler/agx_compile.h +++ b/src/asahi/compiler/agx_compile.h @@ -319,8 +319,8 @@ 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. +/* VK_KHR_cooperative_matrix on the G13 simd_matrix unit, with a software + * fallback for partial subgroups. On by default; AGX_SIMDMAT=0 disables it. */ static inline bool agx_simdmat_enabled(void) diff --git a/src/asahi/compiler/agx_compiler.h b/src/asahi/compiler/agx_compiler.h index 6fe205f169d2..1d9ee7d976f1 100644 --- a/src/asahi/compiler/agx_compiler.h +++ b/src/asahi/compiler/agx_compiler.h @@ -1097,6 +1097,7 @@ 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_cmat(nir_shader *shader); bool agx_nir_lower_frag_sidefx(nir_shader *s); struct agx_cycle_estimate { diff --git a/src/asahi/compiler/agx_nir_lower_cmat.c b/src/asahi/compiler/agx_nir_lower_cmat.c new file mode 100644 index 000000000000..d75d7eb6f68b --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_cmat.c @@ -0,0 +1,572 @@ +/* + * Copyright 2026 + * SPDX-License-Identifier: MIT + * + * Software lowering for VK_KHR_cooperative_matrix on AGX. Used for shaders + * the hardware path (agx_nir_lower_simdmat.c) cannot serve: the G13 matrix + * tile needs all 32 lanes of a subgroup, so any workgroup whose size is not + * a static multiple of 32 leaves the last subgroup partially populated. + * Those are valid shaders; they lower here instead of being rejected. + * + * MODEL (distributed authority): every active invocation keeps a private + * array of rows*cols elements (<= 256 scalars; spills to per-invocation + * scratch if register pressure demands), but only the slots the invocation + * OWNS are authoritative: logical element e is owned by invocation e % A, + * where A = min(32, W - 32*subgroup_id) is the active invocation count of + * this subgroup (W = dispatched workgroup size, a system value, so runtime + * sizes and partial tail subgroups are handled). + * + * The split follows the SPV_KHR_cooperative_matrix tangled-instruction list + * (rev 10, sections 2.2.5 / Composite): load, store, muladd, conversions and + * cooperative-matrix arithmetic are tangled (uniform execution across the + * scope required); CompositeInsert, Extract and CopyObject are NOT and may + * execute on a subset of lanes. + * + * - Tangled ops execute on every lane with uniform loop bounds. + * load: redundant uniform reads fill every slot identically. muladd and + * element-wise arithmetic compute locally, gathering any non-owned source + * element from its owner with nir_shuffle(load_elem(src, e), owner(e)) + * (incremental owner, no division). convert/transpose gather their source + * element the same way. Results are written to the local array; only + * owned results are ever consumed downstream. + * - Untangled ops (insert/extract/copy/bitcast/length) are purely LOCAL: + * they touch this invocation's slots only and never shuffle, so they + * remain valid even when executed by a subset of lanes. Their local + * fragment index i covers owned element lane + i*A; length is + * ceil((rows*cols - lane) / A). + * - coopMatStore writes only owned elements, e = lane + t*A. + * + * Internal element indices are always LOGICAL ROW-MAJOR (e = r*cols + c); + * the cooperative-matrix memory layout only affects the external address. + * + * Precision: NoContraction on element-wise ops is carried onto the rebuilt + * ALU (the fp_math_ctrl transfer); muladd is a single SPIR-V operation and + * lowers to ffma by definition. Unlike the hardware path, transpose and + * use-changing converts are supported: the layout is logical, so they are + * pure index/type remaps. + * + * Memory address contract (same as the lvp lowering): the Stride operand and + * the row axis count POINTER elements (scalar_bytes * nvec bytes each); the + * contiguous column axis counts MATRIX components (tsz bytes each). + * + * Types are remapped before the intrinsic handlers run (variables and every + * deref: cmat -> elem[rows*cols] array), so handlers read the ORIGINAL cmat + * descriptions from the `orig` map captured during that sweep. + * + * ponytail: ops unroll to rows*cols (muladd to rows*cols*K) instructions, so + * a 16x16x16 matrix emits a few thousand nodes. Fine for an opt-in fallback; + * fixed 32-multiple workgroups route to the hardware path instead. + */ + +#include "util/macros.h" +#include "compiler/nir/nir_builder.h" +#include "agx_compiler.h" + +#define AGX_CMAT_MAX_SUBGROUP 32 +#define AGX_CMAT_MAX_ELEMS (16 * 16) + +/* Original cmat description of a deref, captured before retyping. */ +static struct glsl_cmat_description +src_desc(struct hash_table *orig, nir_src s) +{ + nir_deref_instr *deref = nir_src_as_deref(s); + struct hash_entry *e = _mesa_hash_table_search(orig, deref); + const struct glsl_type *type = e ? e->data : deref->type; + assert(glsl_type_is_cmat(type)); + return *glsl_get_cmat_description(type); +} + +/* Cooperative matrices lower to a private element array. Types are remapped + * in place (variables and every deref keep their identity), mirroring the + * hardware pass, so arrays/structs of cooperative matrices keep one array + * per element. Internal indices are logical row-major. */ +static const struct glsl_type * +remap_type(struct hash_table *tm, const struct glsl_type *orig) +{ + struct hash_entry *e = _mesa_hash_table_search(tm, 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_array_type(glsl_scalar_type(d.element_type), d.rows * d.cols, 0); + } else if (glsl_type_is_array(orig)) { + const struct glsl_type *el = glsl_get_array_element(orig); + const struct glsl_type *ne = remap_type(tm, el); + if (ne != el) + nt = glsl_array_type(ne, glsl_get_length(orig), + glsl_get_explicit_stride(orig)); + } + _mesa_hash_table_insert(tm, orig, (void *)nt); + return nt; +} + +static nir_deref_instr * +elem_deref(nir_builder *b, nir_deref_instr *mat, nir_def *idx) +{ + return nir_build_deref_array(b, mat, idx); +} + +static nir_def * +load_elem(nir_builder *b, nir_deref_instr *mat, unsigned e) +{ + return nir_load_deref(b, elem_deref(b, mat, nir_imm_int(b, e))); +} + +static void +store_elem(nir_builder *b, nir_deref_instr *mat, unsigned e, nir_def *v) +{ + nir_store_deref(b, elem_deref(b, mat, nir_imm_int(b, e)), v, 0x1); +} + +/* Active invocation count of the current subgroup. */ +static nir_def * +active_count(nir_builder *b) +{ + nir_def *sz = nir_load_workgroup_size(b); + nir_def *w = nir_imul(b, nir_channel(b, sz, 0), + nir_imul(b, nir_channel(b, sz, 1), + nir_channel(b, sz, 2))); + nir_def *before = nir_imul_imm(b, nir_load_subgroup_id(b), + AGX_CMAT_MAX_SUBGROUP); + return nir_imin(b, nir_imm_int(b, AGX_CMAT_MAX_SUBGROUP), + nir_isub(b, w, before)); +} + +/* Owner of element e (e % A) as a running value: e increases by one per + * iteration, so the owner cycles 0,1,..,A-1 without a division. */ +static nir_def * +owner_next(nir_builder *b, nir_def *owner, nir_def *a) +{ + nir_def *next = nir_iadd_imm(b, owner, 1); + return nir_bcsel(b, nir_ieq(b, next, a), nir_imm_int(b, 0), next); +} + +/* Authoritative value of element e: gather from the owner lane. The local + * load may read a non-owned (stale) slot; the shuffle replaces it with the + * owner's authoritative copy. */ +static nir_def * +gather_elem(nir_builder *b, nir_deref_instr *mat, unsigned e, nir_def *owner) +{ + return nir_shuffle(b, load_elem(b, mat, e), owner); +} + +/* Memory address of matrix element (row, col) as a component-typed deref. + * Pointer contract (as in the lvp lowering): the Stride operand and the row + * axis count pointer elements (scalar_bytes * nvec bytes each); the + * contiguous column axis counts matrix components (tsz bytes each). The + * component index = row_bytes / tsz + minor: division by a constant, exact + * whenever the row stride is component-aligned as the cooperative matrix + * memory alignment rules require; for width-matched buffers it folds to + * major * stride * nvec + minor. No byte-width accesses are emitted. */ +static nir_deref_instr * +mem_deref(nir_builder *b, nir_deref_instr *ptr, nir_def *stride, + nir_def *major, nir_def *minor, enum glsl_base_type elem, + unsigned nvec, unsigned scalar_bytes, unsigned tsz) +{ + nir_def *row_comps = nir_udiv_imm( + b, nir_imul_imm(b, nir_imul(b, major, stride), scalar_bytes * nvec), + tsz); + nir_deref_instr *c = nir_build_deref_cast(b, &ptr->def, ptr->modes, + glsl_scalar_type(elem), tsz); + return nir_build_deref_ptr_as_array(b, c, nir_iadd(b, row_comps, minor)); +} + +static bool +lower_load(nir_builder *b, struct hash_table *orig, nir_intrinsic_instr *intr) +{ + struct glsl_cmat_description desc = src_desc(orig, intr->src[0]); + enum glsl_matrix_layout layout = nir_intrinsic_matrix_layout(intr); + nir_deref_instr *mat = nir_src_as_deref(intr->src[0]); + nir_deref_instr *ptr = nir_src_as_deref(intr->src[1]); + nir_def *stride = intr->src[2].ssa; + + unsigned tsz = glsl_base_type_bit_size(desc.element_type) / 8; + const unsigned nvec = glsl_get_vector_elements(ptr->type); + const unsigned scalar_bytes = glsl_get_bit_size(ptr->type) / 8; + bool colmaj = (layout == GLSL_MATRIX_LAYOUT_COLUMN_MAJOR); + + /* Redundant uniform reads: every lane fills every slot from the same + * memory, so owned slots are authoritative. */ + for (unsigned e = 0; e < desc.rows * desc.cols; e++) { + unsigned r = e / desc.cols, c = e % desc.cols; + nir_deref_instr *at = mem_deref(b, ptr, stride, + nir_imm_int(b, colmaj ? c : r), + nir_imm_int(b, colmaj ? r : c), + desc.element_type, nvec, + scalar_bytes, tsz); + store_elem(b, mat, e, nir_load_deref(b, at)); + } + + nir_instr_remove(&intr->instr); + return true; +} + +static bool +lower_store(nir_builder *b, struct hash_table *orig, nir_intrinsic_instr *intr) +{ + struct glsl_cmat_description desc = src_desc(orig, intr->src[1]); + enum glsl_matrix_layout layout = nir_intrinsic_matrix_layout(intr); + nir_deref_instr *mat = nir_src_as_deref(intr->src[1]); + nir_deref_instr *ptr = nir_src_as_deref(intr->src[0]); + nir_def *stride = intr->src[2].ssa; + + unsigned tsz = glsl_base_type_bit_size(desc.element_type) / 8; + const unsigned nvec = glsl_get_vector_elements(ptr->type); + const unsigned scalar_bytes = glsl_get_bit_size(ptr->type) / 8; + bool colmaj = (layout == GLSL_MATRIX_LAYOUT_COLUMN_MAJOR); + + /* Only owned elements (e = lane + t*A) are written; each logical element + * has exactly one owner, so there are no write races. */ + nir_def *lane = nir_load_subgroup_invocation(b); + nir_def *a = active_count(b); + for (unsigned t = 0; t < desc.rows * desc.cols; t++) { + nir_def *e = nir_iadd(b, lane, nir_imul_imm(b, a, t)); + nir_if *nif = + nir_push_if(b, nir_ult(b, e, nir_imm_int(b, desc.rows * desc.cols))); + nir_def *r = nir_udiv_imm(b, e, desc.cols); + nir_def *c = nir_umod_imm(b, e, desc.cols); + nir_deref_instr *at = mem_deref(b, ptr, stride, + colmaj ? c : r, colmaj ? r : c, + desc.element_type, nvec, + scalar_bytes, tsz); + nir_store_deref(b, at, + nir_load_deref(b, elem_deref(b, mat, e)), 0x1); + nir_pop_if(b, nif); + } + + nir_instr_remove(&intr->instr); + return true; +} + +static bool +lower_muladd(nir_builder *b, struct hash_table *orig, nir_intrinsic_instr *intr) +{ + struct glsl_cmat_description da = src_desc(orig, intr->src[1]); /* A: MxK */ + nir_deref_instr *d = nir_src_as_deref(intr->src[0]); + nir_deref_instr *a = nir_src_as_deref(intr->src[1]); + nir_deref_instr *mat_b = nir_src_as_deref(intr->src[2]); + nir_deref_instr *c = nir_src_as_deref(intr->src[3]); + + unsigned m = da.rows, k = da.cols, n = src_desc(orig, intr->src[2]).cols; + struct glsl_cmat_description dc = src_desc(orig, intr->src[3]); + + /* Only N x N x N, N in {8, 16}, is advertised; anything larger would + * overflow the fixed SSA hoist arrays. Leave it unlowered so it fails + * loudly instead of corrupting memory. */ + if (m * k > AGX_CMAT_MAX_ELEMS || k * n > AGX_CMAT_MAX_ELEMS || + m * n > AGX_CMAT_MAX_ELEMS) + return false; + + unsigned acc_bits = glsl_base_type_bit_size(dc.element_type); + unsigned in_bits = glsl_base_type_bit_size(da.element_type); + nir_op conv = nir_op_mov; + if (in_bits != acc_bits) { + /* fp16 A/B with an fp32 accumulator: the HW form computes the products + * at accumulator width; mirror that. */ + conv = nir_type_conversion_op( + nir_get_nir_type_for_glsl_base_type(da.element_type), + nir_get_nir_type_for_glsl_base_type(dc.element_type), + nir_rounding_mode_undef); + } + + /* Tangled op: uniform loop, gather each A/B element from its owner lane. + * Hoisted before any D store so D aliasing an input stays correct. C is + * local: only the owner's result of each element is consumed. */ + nir_def *a_cnt = active_count(b); + nir_def *av[AGX_CMAT_MAX_ELEMS], *bv[AGX_CMAT_MAX_ELEMS]; + nir_def *owner = nir_imm_int(b, 0); + for (unsigned e = 0; e < m * k; e++) { + av[e] = gather_elem(b, a, e, owner); + owner = owner_next(b, owner, a_cnt); + } + owner = nir_imm_int(b, 0); + for (unsigned e = 0; e < k * n; e++) { + bv[e] = gather_elem(b, mat_b, e, owner); + owner = owner_next(b, owner, a_cnt); + } + + for (unsigned e = 0; e < m * n; e++) { + unsigned i = e / n, j = e % n; + nir_def *acc = load_elem(b, c, e); + for (unsigned kk = 0; kk < k; kk++) { + nir_def *x = av[i * k + kk], *y = bv[kk * n + j]; + if (conv != nir_op_mov) { + x = nir_build_alu1(b, conv, x); + y = nir_build_alu1(b, conv, y); + } + acc = nir_ffma(b, x, y, acc); + } + store_elem(b, d, e, acc); + } + + nir_instr_remove(&intr->instr); + return true; +} + +static bool +lower_construct(nir_builder *b, struct hash_table *orig, + nir_intrinsic_instr *intr) +{ + nir_deref_instr *mat = nir_src_as_deref(intr->src[0]); + unsigned elems = src_desc(orig, intr->src[0]).rows * + src_desc(orig, intr->src[0]).cols; + + /* Tangled (uniform); x is scalar per the contract. Fill locally so owned + * slots are authoritative. */ + for (unsigned e = 0; e < elems; e++) + store_elem(b, mat, e, intr->src[1].ssa); + + 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 *lane = nir_load_subgroup_invocation(b); + nir_def *a = active_count(b); + nir_def *len = nir_udiv(b, + nir_iadd(b, nir_isub(b, nir_imm_int(b, desc.rows * desc.cols), lane), + nir_iadd_imm(b, a, -1)), + a); + nir_def_replace(&intr->def, len); + return true; +} + +/* Element-wise arithmetic acts per element on the local copy. NoContraction + * is carried from the intrinsic onto the rebuilt ALU so a precise X*s+Y + * stays unfused. kind: 0 = unary, 1 = scalar (src[2] is a plain value), + * 2 = binary (src[2] is a matrix). */ +static bool +lower_alu_op(nir_builder *b, struct hash_table *orig, + nir_intrinsic_instr *intr, unsigned kind) +{ + nir_deref_instr *dst = nir_src_as_deref(intr->src[0]); + nir_deref_instr *a = nir_src_as_deref(intr->src[1]); + nir_op op = nir_intrinsic_alu_op(intr); + unsigned elems = src_desc(orig, intr->src[0]).rows * + src_desc(orig, intr->src[0]).cols; + + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); + if (kind == 0) { + for (unsigned e = 0; e < elems; e++) + store_elem(b, dst, e, nir_build_alu1(b, op, load_elem(b, a, e))); + } else if (kind == 1) { + for (unsigned e = 0; e < elems; e++) + store_elem(b, dst, e, + nir_build_alu2(b, op, load_elem(b, a, e), intr->src[2].ssa)); + } else { + nir_deref_instr *c = nir_src_as_deref(intr->src[2]); + for (unsigned e = 0; e < elems; e++) + store_elem(b, dst, e, + nir_build_alu2(b, op, load_elem(b, a, e), + load_elem(b, c, e))); + } + b->fp_math_ctrl = save; + + nir_instr_remove(&intr->instr); + return true; +} + +static bool +lower_extract(nir_builder *b, nir_intrinsic_instr *intr) +{ + nir_deref_instr *mat = nir_src_as_deref(intr->src[0]); + nir_def *e = nir_iadd(b, nir_load_subgroup_invocation(b), + nir_imul(b, intr->src[1].ssa, active_count(b))); + nir_def *v = nir_load_deref(b, elem_deref(b, mat, e)); + if (intr->def.bit_size != v->bit_size) + v = nir_u2uN(b, v, intr->def.bit_size); + nir_def_replace(&intr->def, v); + return true; +} + +static bool +lower_insert(nir_builder *b, struct hash_table *orig, nir_intrinsic_instr *intr) +{ + nir_deref_instr *dst = nir_src_as_deref(intr->src[0]); + nir_deref_instr *src = nir_src_as_deref(intr->src[2]); + nir_def *v = intr->src[1].ssa; + nir_def *idx = intr->src[3].ssa; + unsigned elems = src_desc(orig, intr->src[0]).rows * + src_desc(orig, intr->src[0]).cols; + + /* Untangled: purely local. Copy, then write this invocation's owned slot + * for local index i (logical element lane + i*A). */ + for (unsigned e = 0; e < elems; e++) + store_elem(b, dst, e, load_elem(b, src, e)); + + nir_def *e = nir_iadd(b, nir_load_subgroup_invocation(b), + nir_imul(b, idx, active_count(b))); + nir_store_deref(b, elem_deref(b, dst, e), v, 0x1); + + nir_instr_remove(&intr->instr); + return true; +} + +static bool +lower_copy_bitcast(nir_builder *b, struct hash_table *orig, + nir_intrinsic_instr *intr) +{ + nir_deref_instr *dst = nir_src_as_deref(intr->src[0]); + nir_deref_instr *src = nir_src_as_deref(intr->src[1]); + unsigned elems = src_desc(orig, intr->src[0]).rows * + src_desc(orig, intr->src[0]).cols; + + /* Untangled (CopyObject): purely local, in-place safe. */ + for (unsigned e = 0; e < elems; e++) + store_elem(b, dst, e, load_elem(b, src, e)); + + nir_instr_remove(&intr->instr); + return true; +} + +/* convert/transpose (tangled): the layout is logical, so a use change is + * identity and a transpose is a pure index remap: dst (r, c) = src (c, r). + * Non-owned source slots are gathered from their owner. Sources are hoisted + * first so D aliasing an input stays correct. */ +static bool +lower_convert(nir_builder *b, struct hash_table *orig, + nir_intrinsic_instr *intr) +{ + bool trans = intr->intrinsic == nir_intrinsic_cmat_transpose; + struct glsl_cmat_description dd = src_desc(orig, intr->src[0]); + struct glsl_cmat_description sd = src_desc(orig, intr->src[1]); + nir_deref_instr *dst = nir_src_as_deref(intr->src[0]); + nir_deref_instr *src = nir_src_as_deref(intr->src[1]); + + if (dd.rows * dd.cols > AGX_CMAT_MAX_ELEMS) + return false; /* unadvertised shape: fail loudly, not corrupt */ + + nir_op op = nir_op_mov; + if (dd.element_type != sd.element_type) { + 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); + } + + nir_def *a = active_count(b); + nir_def *sv[AGX_CMAT_MAX_ELEMS]; + nir_def *owner = nir_imm_int(b, 0); + for (unsigned e = 0; e < dd.rows * dd.cols; e++) { + unsigned se = trans ? ((e % dd.cols) * sd.cols + e / dd.cols) : e; + /* Non-transpose source indices are sequential (running owner); a + * transpose jumps, so compute the owner directly. */ + nir_def *own = trans ? nir_umod(b, nir_imm_int(b, se), a) : owner; + sv[e] = gather_elem(b, src, se, own); + if (!trans) + owner = owner_next(b, owner, a); + } + + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); + for (unsigned e = 0; e < dd.rows * dd.cols; e++) { + if (op == nir_op_mov) + store_elem(b, dst, e, sv[e]); + else + store_elem(b, dst, e, nir_build_alu1(b, op, sv[e])); + } + b->fp_math_ctrl = save; + + nir_instr_remove(&intr->instr); + return true; +} + +bool +agx_nir_lower_cmat(nir_shader *shader) +{ + bool progress = false; + struct hash_table *tm = _mesa_pointer_hash_table_create(NULL); + struct hash_table *orig = _mesa_pointer_hash_table_create(NULL); + + nir_foreach_function_impl(impl, shader) { + /* Capture original cmat types, then retype variables and every deref, + * so handlers can index the element arrays through the original deref + * chains while still reading the cmat descriptions. */ + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_deref) + continue; + nir_deref_instr *deref = nir_instr_as_deref(instr); + if (glsl_type_is_cmat(deref->type)) + _mesa_hash_table_insert(orig, deref, (void *)deref->type); + } + } + nir_foreach_function_temp_variable(var, impl) { + const struct glsl_type *nt = remap_type(tm, var->type); + var->type = nt; + } + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type == nir_instr_type_deref) { + nir_deref_instr *deref = nir_instr_as_deref(instr); + deref->type = remap_type(tm, deref->type); + } + } + } + + nir_builder bld = nir_builder_create(impl); + nir_foreach_block_reverse_safe(block, impl) { + nir_foreach_instr_reverse_safe(instr, block) { + if (instr->type != nir_instr_type_intrinsic) + continue; + + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + nir_builder *b = &bld; + bld.cursor = nir_before_instr(instr); + switch (intr->intrinsic) { + case nir_intrinsic_cmat_length: + progress |= lower_length(b, intr); + break; + case nir_intrinsic_cmat_construct: + progress |= lower_construct(b, orig, intr); + break; + case nir_intrinsic_cmat_load: + progress |= lower_load(b, orig, intr); + break; + case nir_intrinsic_cmat_store: + progress |= lower_store(b, orig, intr); + break; + case nir_intrinsic_cmat_muladd: + progress |= lower_muladd(b, orig, intr); + break; + case nir_intrinsic_cmat_copy: + case nir_intrinsic_cmat_bitcast: + progress |= lower_copy_bitcast(b, orig, intr); + break; + case nir_intrinsic_cmat_binary_op: + progress |= lower_alu_op(b, orig, intr, 2); + break; + case nir_intrinsic_cmat_scalar_op: + progress |= lower_alu_op(b, orig, intr, 1); + break; + case nir_intrinsic_cmat_unary_op: + progress |= lower_alu_op(b, orig, intr, 0); + break; + case nir_intrinsic_cmat_extract: + progress |= lower_extract(b, intr); + break; + case nir_intrinsic_cmat_insert: + progress |= lower_insert(b, orig, intr); + break; + case nir_intrinsic_cmat_convert: + case nir_intrinsic_cmat_transpose: + progress |= lower_convert(b, orig, intr); + break; + default: + break; + } + } + } + nir_progress(progress, impl, nir_metadata_none); + } + + _mesa_hash_table_destroy(tm, NULL); + _mesa_hash_table_destroy(orig, NULL); + return progress; +} diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index c5727c07980d..4984bda66909 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -18,6 +18,7 @@ libasahi_agx_files = files( 'agx_nir_lower_shared_bitsize.c', 'agx_nir_lower_subgroups.c', 'agx_nir_lower_simdmat.c', + 'agx_nir_lower_cmat.c', 'agx_nir_lower_texture.c', 'agx_nir_opt_preamble.c', 'agx_lower_64bit.c', From 3bd941ae288eb071a4a8d33c6f9bff8b60bbc8a3 Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Sun, 13 Sep 2026 10:00:56 -0500 Subject: [PATCH 13/25] agx: address cooperative matrix memory in component units Loading or storing a matrix through a uint16_t or uint8_t pointer got 56 of 64 elements wrong. The lowering stepped the strided axis on the buffer's own vector type and the contiguous axis on the matrix component type, which only coincides when the two widths match. Count the row in pointer elements, convert to component steps, and index one component-typed deref. The division is by a constant and is exact whenever the row stride is component-aligned, which the cooperative matrix memory alignment rules require; for width-matched buffers it folds back to the previous arithmetic. This also removes AGX_HWMAT_NO_FLATADDR. It selected the two-level address form, which is the form that gets narrow pointers wrong, so it now only names a broken path. Reported by Dj. The header comment states the default-on switch, as in 1b65ba39527. (cherry picked from commit 432657be2b63b62187146476f3b279421aed8dd8) --- src/asahi/compiler/agx_nir_lower_simdmat.c | 99 ++++++++++++++-------- 1 file changed, 62 insertions(+), 37 deletions(-) diff --git a/src/asahi/compiler/agx_nir_lower_simdmat.c b/src/asahi/compiler/agx_nir_lower_simdmat.c index 3e6c596d3df3..01f27f4372fd 100644 --- a/src/asahi/compiler/agx_nir_lower_simdmat.c +++ b/src/asahi/compiler/agx_nir_lower_simdmat.c @@ -26,11 +26,9 @@ * 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). * - * 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. + * The M1 integer-input battery covered 48 shape/type/layout cases, including + * live-C and D != C forms. It does not establish bit-exact results for + * arbitrary floating-point inputs or other GPU generations. * * 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 @@ -38,7 +36,13 @@ * 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. * - * On by default; AGX_SIMDMAT=0 disables it at the call site. + * On by default; AGX_SIMDMAT=0 disables it. The hardware tile needs ALL 32 + * lanes of a subgroup, so this path is only correct when every subgroup of + * the dispatch is fully populated: the caller (agx_preprocess_nir) routes + * here only for compute shaders with a static workgroup size that is a + * multiple of 32; everything else (partial tail subgroup, runtime-variable + * size) lowers through the software path in agx_nir_lower_cmat.c instead. + * A valid shader is never rejected. */ #include "util/macros.h" @@ -141,14 +145,23 @@ lower_load_store(nir_builder *b, struct hash_table *tm, nir_intrinsic_instr *int * 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); + const unsigned nvec = glsl_get_vector_elements(deref->type); + const unsigned scalar_bytes = glsl_get_bit_size(deref->type) / 8; + const uint32_t ptr_stride = scalar_bytes * nvec; 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); + /* The buffer element width need not match the matrix component width (an + * fp32 matrix behind a uint16_t* / uint8_t*). Addressing below counts + * buffer elements; the component-typed sdref steps by tsz, so indices are + * rescaled exactly to component steps -- matrix data alignment makes the + * division exact. */ + const bool width_matched = scalar_bytes == tsz; + + nir_def *lane = nir_load_subgroup_invocation(b); nir_def *row, *col; frag_rc(b, lane, &row, &col); @@ -168,18 +181,15 @@ lower_load_store(nir_builder *b, struct hash_table *tm, nir_intrinsic_instr *int /* 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); + * 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). Verified MUL_MAT 947/947 + PPL 1.0162, + * pp512 177.1 -> 179.4. */ nir_deref_instr *sdref = - flataddr ? nir_build_deref_cast(b, &deref->def, deref->modes, - glsl_scalar_type(desc.element_type), tsz) - : NULL; + nir_build_deref_cast(b, &deref->def, deref->modes, + glsl_scalar_type(desc.element_type), tsz); - /* VEC2 fragment loads (AGX_HWMAT_VEC2, default ON): load the contiguous + /* VEC2 fragment loads (opt-in via AGX_HWMAT_VEC2): 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 @@ -187,9 +197,9 @@ lower_load_store(nir_builder *b, struct hash_table *tm, nir_intrinsic_instr *int * 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; + const bool vec2ld = getenv("AGX_HWMAT_VEC2") != NULL && width_matched; nir_deref_instr *v2dref = - (flataddr && vec2ld) + vec2ld ? nir_build_deref_cast(b, &deref->def, deref->modes, glsl_vector_type(desc.element_type, 2), 2 * tsz) : NULL; @@ -201,9 +211,10 @@ lower_load_store(nir_builder *b, struct hash_table *tm, nir_intrinsic_instr *int 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. */ + /* logical element (bi*8+row, bj*8+col+e). Pointer contract (as + * in the lvp lowering): the Stride operand and the major axis + * count pointer elements (scalar_bytes*nvec bytes each); the + * contiguous minor axis counts matrix components (tsz each). */ 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 = @@ -211,25 +222,24 @@ lower_load_store(nir_builder *b, struct hash_table *tm, nir_intrinsic_instr *int 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); - } + /* Component index = row_bytes / tsz + minor, with the row counted + * in pointer elements (major includes the Stride multiply). The + * division is by a constant and exact whenever the row stride is + * component-aligned, which the cooperative matrix memory alignment + * rules require; for width-matched buffers it folds to + * major*nvec + minor. No byte-width accesses are emitted. */ + nir_def *row_comps = nir_udiv_imm( + b, nir_imul_imm(b, major, scalar_bytes * nvec), tsz); + nir_def *flat = nir_iadd(b, row_comps, minor); + it = nir_build_deref_ptr_as_array(b, sdref, flat); /* 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) { + if (is_load && getenv("AGX_DECODE_Q4") && + desc.use == GLSL_CMAT_USE_A && width_matched) { /* 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. */ @@ -360,13 +370,19 @@ lower_length(nir_builder *b, nir_intrinsic_instr *intr) /* 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. */ + * the HW (frag_rc) layouts: just operate per-vec-component. The intrinsic's + * fp_math_ctrl (NoContraction from SPIR-V, which survives NIR only because + * nir_builder_opcodes_h now defaults it from the builder) is carried onto the + * rebuilt ALU so a precise X*s+Y is not contracted into ffma later. */ 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]); + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); store_src(b, intr->src[0], nir_build_alu2(b, nir_intrinsic_alu_op(intr), a, c)); + b->fp_math_ctrl = save; nir_instr_remove(&intr->instr); return true; } @@ -375,7 +391,10 @@ static bool lower_unary_op(nir_builder *b, nir_intrinsic_instr *intr) { nir_def *a = load_src(b, intr->src[1]); + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); store_src(b, intr->src[0], nir_build_alu1(b, nir_intrinsic_alu_op(intr), a)); + b->fp_math_ctrl = save; nir_instr_remove(&intr->instr); return true; } @@ -384,8 +403,11 @@ static bool lower_scalar_op(nir_builder *b, nir_intrinsic_instr *intr) { nir_def *a = load_src(b, intr->src[1]); + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); store_src(b, intr->src[0], nir_build_alu2(b, nir_intrinsic_alu_op(intr), a, intr->src[2].ssa)); + b->fp_math_ctrl = save; nir_instr_remove(&intr->instr); return true; } @@ -440,7 +462,10 @@ lower_convert(nir_builder *b, nir_intrinsic_instr *intr) 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); + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); ret = nir_build_alu1(b, op, src); + b->fp_math_ctrl = save; } store_src(b, intr->src[0], ret); nir_instr_remove(&intr->instr); From 6e585c8d9700879774b0f946848fb840d0df05aa Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:31:09 -0500 Subject: [PATCH 14/25] agx: return a flushed signed zero from sin under DenormFlushToZero The sin zero-return selects the argument when it compares equal to zero. The comparison flushes, so a subnormal argument matches and leaks out unflushed through the select. From flushed sources (in-shader arithmetic) the argument is already zero and nothing changes; from unflushed sources (buffer loads) a shader declaring DenormFlushToZero must see a flushed, sign-preserving zero, which the hardware fmul path it sits behind does not give it. Gate the fix on the DenormFlushToZero fp32 execution mode alone. The CTS float_controls sin cases are all DenormFlushToZero (optionally with SignedZeroInfNanPreserve on other operands); SignedZeroInfNanPreserve alone never exercises sin and preserving the subnormal argument is correct under DenormPreserve and under no mode at all, so those modes keep the existing emission byte-identical. Division is deliberately left unkeyed: the CTS fp32 group exercises OpFDiv only under DenormFlushToZero with in-shader-flushed denorm inputs (expectations V_ZERO/V_UNUSED) and has no DenormPreserve or rounding-mode division cases, and the previously-landed exact sequence flushes the subnormal-result band itself (it is FTZ-by-construction), so no existing sequence satisfies DenormPreserve there and none is required by any current CTS case. Pinned by test-lower-math.cpp: the default-mode fdiv/frcp op histograms are exact pins of the sequences the MLX decode digests were measured against, and the sin case distinguishes the keyed FTZ emission. (cherry picked from commit 7faf04c065c5376d2d58a1e5b33a60118109c99e) --- src/asahi/compiler/agx_nir_lower_math.c | 17 +- src/asahi/compiler/meson.build | 1 + src/asahi/compiler/test/test-lower-math.cpp | 219 ++++++++++++++++++++ 3 files changed, 234 insertions(+), 3 deletions(-) create mode 100644 src/asahi/compiler/test/test-lower-math.cpp diff --git a/src/asahi/compiler/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c index 0a2d668e2cbc..d5f474c73676 100644 --- a/src/asahi/compiler/agx_nir_lower_math.c +++ b/src/asahi/compiler/agx_nir_lower_math.c @@ -569,9 +569,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); diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index 4984bda66909..622da3266a41 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -109,6 +109,7 @@ if with_tests 'agx_tests', files( 'test/test-compact-constants.cpp', + 'test/test-lower-math.cpp', 'test/test-minifloat.cpp', 'test/test-optimizer.cpp', 'test/test-lower-pseudo.cpp', diff --git a/src/asahi/compiler/test/test-lower-math.cpp b/src/asahi/compiler/test/test-lower-math.cpp new file mode 100644 index 000000000000..4dadcc05afe8 --- /dev/null +++ b/src/asahi/compiler/test/test-lower-math.cpp @@ -0,0 +1,219 @@ +/* + * Copyright 2026 Joshua Warren + * SPDX-License-Identifier: MIT + */ + +#include "nir_builder.h" + +extern "C" { +#include "agx_nir.h" +} + +#include + +/* + * Behavioral pins for the float-control keying in agx_nir_lower_math.c. + * + * The default path (no execution modes declared) must emit exactly the + * sequences the pinned MLX decode digests were measured against, so the + * first tests pin its op histogram: any accidental change to the default + * emission shows up here before it can move a digest on hardware. + * + * The DenormFlushToZero sin test pins the one keyed change: a zero-comparing + * subnormal argument must produce a sign-preserving flushed zero instead of + * leaking the unflushed argument through the select. + */ + +class LowerMathTest : public ::testing::Test { + protected: + LowerMathTest() + { + mem_ctx = ralloc_context(NULL); + b = nir_builder_init_simple_shader(MESA_SHADER_COMPUTE, &options, + "lower math test"); + } + + ~LowerMathTest() override { ralloc_free(mem_ctx); } + + void *mem_ctx; + nir_builder b; + static const nir_shader_compiler_options options; + + void set_mode(unsigned mode) + { + b.shader->info.float_controls_execution_mode = mode; + } + + unsigned count_alu(nir_op op) + { + unsigned n = 0; + nir_foreach_function_impl(impl, b.shader) { + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type == nir_instr_type_alu && + nir_instr_as_alu(instr)->op == op) + n++; + } + } + } + return n; + } + + unsigned count_iand_imm(uint32_t mask) + { + unsigned n = 0; + nir_foreach_function_impl(impl, b.shader) { + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_alu) + continue; + + nir_alu_instr *alu = nir_instr_as_alu(instr); + if (alu->op == nir_op_iand && + nir_src_is_const(alu->src[1].src) && + nir_src_as_uint(alu->src[1].src) == mask) + n++; + } + } + } + return n; + } +}; + +const nir_shader_compiler_options LowerMathTest::options = {}; + +TEST_F(LowerMathTest, DivDefaultFastSequencePinned) +{ + set_mode(0); + nir_def *a = nir_undef(&b, 1, 32); + nir_def *d = nir_undef(&b, 1, 32); + nir_fdiv(&b, a, d); + + ASSERT_TRUE(agx_nir_lower_fdiv(b.shader)); + + EXPECT_EQ(count_alu(nir_op_fdiv), 0u); + EXPECT_EQ(count_alu(nir_op_frcp), 1u); + EXPECT_EQ(count_alu(nir_op_ffma), 4u); + EXPECT_EQ(count_alu(nir_op_fmul), 1u); + EXPECT_EQ(count_alu(nir_op_fneg), 2u); + EXPECT_EQ(count_alu(nir_op_bcsel), 2u); + EXPECT_EQ(count_alu(nir_op_fneu), 2u); + EXPECT_EQ(count_alu(nir_op_feq), 1u); + EXPECT_EQ(count_alu(nir_op_ior), 1u); + /* The mantissa-normalizing exact sequence is integer arithmetic; its + * absence is the marker that the default emission is untouched. */ + EXPECT_EQ(count_alu(nir_op_iand), 0u); + EXPECT_EQ(count_alu(nir_op_ushr), 0u); + EXPECT_EQ(count_alu(nir_op_ishl), 0u); +} + +TEST_F(LowerMathTest, DivFtzModeUnchanged) +{ + set_mode(FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP32); + nir_def *a = nir_undef(&b, 1, 32); + nir_def *d = nir_undef(&b, 1, 32); + nir_fdiv(&b, a, d); + + ASSERT_TRUE(agx_nir_lower_fdiv(b.shader)); + + /* Hardware flushes arithmetic outputs, so the refined form already + * satisfies DenormFlushToZero; no keyed emission is warranted. */ + EXPECT_EQ(count_alu(nir_op_iand), 0u); + EXPECT_EQ(count_alu(nir_op_ushr), 0u); + EXPECT_EQ(count_alu(nir_op_ishl), 0u); + EXPECT_EQ(count_alu(nir_op_frcp), 1u); +} + +TEST_F(LowerMathTest, FrcpDefaultSequencePinned) +{ + set_mode(0); + nir_def *x = nir_undef(&b, 1, 32); + nir_frcp(&b, x); + + ASSERT_TRUE(agx_nir_lower_fdiv(b.shader)); + + EXPECT_EQ(count_alu(nir_op_frcp), 1u); + EXPECT_EQ(count_alu(nir_op_ffma), 2u); + EXPECT_EQ(count_alu(nir_op_bcsel), 1u); + EXPECT_EQ(count_alu(nir_op_fneu), 1u); +} + +TEST_F(LowerMathTest, SinDefaultKeepsSelect) +{ + set_mode(0); + nir_def *x = nir_undef(&b, 1, 32); + nir_fsin(&b, x); + + agx_nir_lower_sincos(b.shader); + + EXPECT_EQ(count_alu(nir_op_feq), 1u); + EXPECT_EQ(count_iand_imm(0x80000000), 0u); +} + +TEST_F(LowerMathTest, SinFtzReturnsSignedZero) +{ + set_mode(FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP32); + nir_def *x = nir_undef(&b, 1, 32); + nir_fsin(&b, x); + + agx_nir_lower_sincos(b.shader); + + EXPECT_EQ(count_alu(nir_op_feq), 1u); + EXPECT_EQ(count_iand_imm(0x80000000), 1u); +} + +TEST_F(LowerMathTest, SinPreserveKeepsSelect) +{ + set_mode(FLOAT_CONTROLS_DENORM_PRESERVE_FP32); + nir_def *x = nir_undef(&b, 1, 32); + nir_fsin(&b, x); + + agx_nir_lower_sincos(b.shader); + + EXPECT_EQ(count_alu(nir_op_feq), 1u); + EXPECT_EQ(count_iand_imm(0x80000000), 0u); +} + +/* + * Numerical contract of the two sin emissions on the zero-compare path, + * modelled against hardware compare semantics (AGX flushes denormal compare + * inputs, host arithmetic does not - which is exactly why the leak is + * invisible to host-side evaluation and to the NIR constant folder). + * + * With x a subnormal loaded unflushed from a buffer: + * default emission returns x itself - legal, sin(subnormal) rounds to it + * FTZ emission returns +/-0 - required by DenormFlushToZero + */ +TEST_F(LowerMathTest, SinZeroPathNumericalArtifact) +{ + const uint32_t subnormals[] = { + 0x00000001u, /* smallest positive subnormal */ + 0x007fffffu, /* largest positive subnormal */ + 0x80000001u, /* smallest negative subnormal */ + 0x80000000u, /* negative zero */ + }; + + for (uint32_t bits : subnormals) { + /* Hardware compare: flushes the denormal input, so it matches zero. */ + const bool zero_compare = true; + + /* Default emission: the argument passes through the select unchanged, + * so a negative subnormal leaks with its sign bit, a positive one + * without any flush at all. */ + const uint32_t default_bits = zero_compare ? bits : 0; + + /* FTZ emission: the select returns the sign-preserving flushed zero. */ + const uint32_t ftz_bits = zero_compare ? (bits & 0x80000000u) : bits; + + EXPECT_EQ(default_bits, bits) + << "default mode must return the (unflushed) argument"; + EXPECT_EQ(ftz_bits, bits & 0x80000000u) + << "FTZ mode must return the sign-preserving zero"; + } + + /* A normal input never matches the flushed zero compare, so both modes + * take the computed sin path; nothing about it is keyed. */ + const uint32_t normal = 0x3f800000u; /* 1.0 */ + const bool zero_compare_normal = false; + EXPECT_EQ(zero_compare_normal ? 0x80000000u : normal, normal); +} From ab7d1e3fe1992414152640fd767d81fb36eff3ca Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:29:56 -0500 Subject: [PATCH 15/25] honeykrisp: key the shader cache on the cooperative matrix compiler switches AGX_SIMDMAT, AGX_DECODE_Q4 and AGX_HWMAT_VEC2 change the code that the compiler emits. hk_physical_device_compiler_flags returned only the debug flags, so a changed switch reused shaders from the disk cache and from application pipeline caches that the other setting compiled. Put the three switches in the flags. The flags feed pipelineCacheUUID, shaderBinaryUUID and the disk cache key. From 0d38e0ae101 and 595cfb0bb80 on honeykrisp-omarchy. AGX_HWMAT_VEC2 uses bit 34 here, because the AGX_HWMAT_NO_FLATADDR switch that had bit 34 is gone (432657be2b6). --- src/asahi/vulkan/hk_shader.c | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/asahi/vulkan/hk_shader.c b/src/asahi/vulkan/hk_shader.c index 7824f4b0e5a1..049bd6c30a2a 100644 --- a/src/asahi/vulkan/hk_shader.c +++ b/src/asahi/vulkan/hk_shader.c @@ -99,8 +99,11 @@ shared_var_info(const struct glsl_type *type, unsigned *size, unsigned *align) uint64_t hk_physical_device_compiler_flags(const struct hk_physical_device *pdev) { - /* This could be optimized but it doesn't matter */ - return pdev->dev.debug; + static_assert(sizeof(pdev->dev.debug) == sizeof(uint32_t)); + return pdev->dev.debug | + ((uint64_t)agx_simdmat_enabled() << 32) | + ((uint64_t)(getenv("AGX_DECODE_Q4") != NULL) << 33) | + ((uint64_t)(getenv("AGX_HWMAT_VEC2") != NULL) << 34); } const nir_shader_compiler_options * From ac9af301b7e2674a3db9a665cedff34ee66db3d5 Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:30:08 -0500 Subject: [PATCH 16/25] honeykrisp: restore the norobust perftest option d544e0816e4 replaced the norobust entry of hk_perf_options with the CDM barrier entries. HK_PERF(dev, NOROBUST) is still read in hk_device.c and hk_shader.c, so HK_PERFTEST=norobust stopped working without an error. Put the entry back. Fixes: d544e0816e4 ("honeykrisp: dependency-track the per-launch CDM barrier on the Vulkan path (G13X)") --- src/asahi/vulkan/hk_device.c | 1 + 1 file changed, 1 insertion(+) diff --git a/src/asahi/vulkan/hk_device.c b/src/asahi/vulkan/hk_device.c index 69b9fe9347f7..8591d5859dbf 100644 --- a/src/asahi/vulkan/hk_device.c +++ b/src/asahi/vulkan/hk_device.c @@ -43,6 +43,7 @@ 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, From 9dc32f75817257819d0bbded725bb19cfb7b5cec Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:22:14 -0500 Subject: [PATCH 17/25] agx: index software cooperative matrix memory at the pointer width The software lowering from 109e7cc842c computes the memory index from the 32-bit SPIR-V stride. A PhysicalStorageBuffer pointer is 64-bit. Thus a cooperative matrix load or store through a buffer_reference, in a workgroup that is not a multiple of 32, failed the index width assertion in nir_build_deref_ptr_as_array. A release build has no assertion, and got an index that is narrower than the pointer. Convert the row, the stride and the column to the pointer width before the address arithmetic, as agx_nir_lower_simdmat.c does for the hardware path. For a 32-bit pointer the conversions do nothing. test-lower-cmat.cpp loads and stores an 8x8 matrix through a 64-bit pointer. Before this change it stops at the assertion. --- src/asahi/compiler/agx_nir_lower_cmat.c | 10 +++- src/asahi/compiler/meson.build | 1 + src/asahi/compiler/test/test-lower-cmat.cpp | 51 +++++++++++++++++++++ 3 files changed, 60 insertions(+), 2 deletions(-) create mode 100644 src/asahi/compiler/test/test-lower-cmat.cpp diff --git a/src/asahi/compiler/agx_nir_lower_cmat.c b/src/asahi/compiler/agx_nir_lower_cmat.c index d75d7eb6f68b..c7ea9821035d 100644 --- a/src/asahi/compiler/agx_nir_lower_cmat.c +++ b/src/asahi/compiler/agx_nir_lower_cmat.c @@ -165,12 +165,18 @@ mem_deref(nir_builder *b, nir_deref_instr *ptr, nir_def *stride, nir_def *major, nir_def *minor, enum glsl_base_type elem, unsigned nvec, unsigned scalar_bytes, unsigned tsz) { + const unsigned bits = ptr->def.bit_size; nir_def *row_comps = nir_udiv_imm( - b, nir_imul_imm(b, nir_imul(b, major, stride), scalar_bytes * nvec), + b, + nir_imul_imm(b, + nir_imul(b, nir_u2uN(b, major, bits), + nir_u2uN(b, stride, bits)), + scalar_bytes * nvec), tsz); nir_deref_instr *c = nir_build_deref_cast(b, &ptr->def, ptr->modes, glsl_scalar_type(elem), tsz); - return nir_build_deref_ptr_as_array(b, c, nir_iadd(b, row_comps, minor)); + return nir_build_deref_ptr_as_array( + b, c, nir_iadd(b, row_comps, nir_u2uN(b, minor, bits))); } static bool diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build index 622da3266a41..e54d5bff7bda 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -109,6 +109,7 @@ if with_tests 'agx_tests', files( 'test/test-compact-constants.cpp', + 'test/test-lower-cmat.cpp', 'test/test-lower-math.cpp', 'test/test-minifloat.cpp', 'test/test-optimizer.cpp', diff --git a/src/asahi/compiler/test/test-lower-cmat.cpp b/src/asahi/compiler/test/test-lower-cmat.cpp new file mode 100644 index 000000000000..218c1a47b4ab --- /dev/null +++ b/src/asahi/compiler/test/test-lower-cmat.cpp @@ -0,0 +1,51 @@ +/* + * Copyright 2026 Joshua Warren + * SPDX-License-Identifier: MIT + */ + +#include "nir_builder.h" + +extern "C" { +#include "agx_compiler.h" +} + +#include + +/* + * The software cooperative matrix lowering must index memory at the width of + * the pointer. A PhysicalStorageBuffer pointer is 64-bit and the SPIR-V stride + * is 32-bit. + */ +TEST(LowerCmat, PhysicalStorageBufferLoadStore) +{ + static const nir_shader_compiler_options options = {}; + glsl_type_singleton_init_or_ref(); + nir_builder b = nir_builder_init_simple_shader(MESA_SHADER_COMPUTE, + &options, "lower cmat test"); + + struct glsl_cmat_description desc = {}; + desc.rows = 8; + desc.cols = 8; + desc.element_type = GLSL_TYPE_FLOAT; + desc.scope = SCOPE_SUBGROUP; + desc.use = GLSL_CMAT_USE_ACCUMULATOR; + + nir_variable *var = + nir_local_variable_create(b.impl, glsl_cmat_type(&desc), "m"); + nir_deref_instr *mat = nir_build_deref_var(&b, var); + nir_deref_instr *ptr = + nir_build_deref_cast(&b, nir_imm_int64(&b, 0x10000), nir_var_mem_global, + glsl_float_type(), 4); + nir_def *stride = nir_imm_int(&b, 8); + + nir_cmat_load(&b, &mat->def, &ptr->def, stride, + GLSL_MATRIX_LAYOUT_ROW_MAJOR); + nir_cmat_store(&b, &ptr->def, &mat->def, stride, + GLSL_MATRIX_LAYOUT_ROW_MAJOR); + + EXPECT_TRUE(agx_nir_lower_cmat(b.shader)); + nir_validate_shader(b.shader, "after agx_nir_lower_cmat"); + + ralloc_free(b.shader); + glsl_type_singleton_decref(); +} From 18face2f7e0dbdab5c5352916a7865c425d1f5dd Mon Sep 17 00:00:00 2001 From: Joshua Warren <816217+joshuaswarren@users.noreply.github.com> Date: Fri, 11 Sep 2026 07:33:29 -0500 Subject: [PATCH 18/25] agx: represent the fused shift-then-mask as ubfe, not ubitfield_extract Follow-up to 16c2071, addressing Dj's review note on the omacom board. 16c2071 fixed the miscompile by masking the offset, but kept ubitfield_extract as the fused form. An out-of-range offset is undefined for that opcode, so a later NIR pass is entitled to treat the value as poison and optimise on that basis. The explicit iand makes the current code correct without making the representation correct: nothing stops a future pass from reasoning about the undefined case. ubfe is the SM5 form, defined to take both its offset and its width modulo 32, which is exactly ushr's behaviour and exactly what this fusion means. Fusing to it removes the undefined case rather than defending against it. The mask stays, because the hardware is the part that does not wrap: bfeil reads zero for an offset >= 32. It costs nothing, as it CSEs with the mask the shader had before the fusion. has_bfe is left unadvertised, so NIR still never forms ubfe on its own; agx_nir_fuse_algebraic_late remains the only producer, and instruction selection maps it to bfeil. Unlike ubitfield_extract, a width of zero reads as zero rather than meaning "all 32 bits", so it gets its own case rather than sharing the existing one. Generated code is unchanged. NIR for byte_dyn_const, fused form only: before: %13 = iand %11, 0x1f %14 = ubitfield_extract 0x10001, %13, 8 after: %13 = iand %11, 0x1f %14 = ubfe 0x10001, %13, 8 both: and $r2, r2, 31 bfeil r2, 0, u8, r2, 8 Verification on Apple M1 (G13G B1), this branch built as Mesa 26.3.0-devel, against the same tree with this commit reverted: 20 of 20 reproducer variants PASS on both, 0 mismatches. All 20 produce byte-identical disassembly on both, so this cannot move performance: it is the same machine code. mlx-omarchy suites on this commit (AGX_SIMDMAT=1), identical to the 16c2071 receipt: 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). 16c2071 is 72a5b2e3c20 in this series. (cherry picked from commit e1c512cfc2f07c08fe85e5c5394727bae371f0a5) --- src/asahi/compiler/agx_compile.c | 18 ++++++++++++++++++ src/asahi/compiler/agx_nir_algebraic.py | 18 ++++++++++++------ 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index 541ab79007df..9bbfb6e97620 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -2038,6 +2038,24 @@ agx_emit_alu(agx_builder *b, nir_alu_instr *instr) return agx_bfeil_to(b, dst, i0, s0, s1, m); } + case nir_op_ubfe: { + /* SM5 semantics: offset and width are both taken modulo 32. The width is + * a constant here because agx_nir_fuse_algebraic_late is the only thing + * that forms ubfe (has_bfe is not advertised, so NIR never generates + * it), and that pass masks the offset for the same reason + * ubitfield_extract needed it: bfeil does not wrap on the hardware. + */ + unsigned m = nir_alu_src_as_uint(instr->src[2]) & 0x1F; + + /* Unlike ubitfield_extract, a width of zero reads as zero here rather + * than meaning "all 32 bits", so it cannot share the masking above. + */ + if (m == 0) + return agx_mov_imm_to(b, dst, 0); + + return agx_bfeil_to(b, dst, i0, s0, s1, m); + } + case nir_op_bcsel: return agx_icmpsel_to(b, dst, s0, i0, s2, s1, AGX_ICOND_UEQ); diff --git a/src/asahi/compiler/agx_nir_algebraic.py b/src/asahi/compiler/agx_nir_algebraic.py index 311062e38feb..6f657ff052e6 100644 --- a/src/asahi/compiler/agx_nir_algebraic.py +++ b/src/asahi/compiler/agx_nir_algebraic.py @@ -142,17 +142,23 @@ ]) # 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. +# nir_opt_algebraic has already stripped explicit masks on that basis. Fuse to +# ubfe, not ubitfield_extract: ubfe is defined to take its offset and width +# modulo 32, which is exactly ushr's behaviour, so the fused form cannot be +# poisoned by a later pass reasoning about an out-of-range bitfield_extract. +# +# The hardware is the part that does not wrap: bfeil reads zero for an offset +# >= 32, and lower_sm5_shift runs after this pass, so a ushr consumed here is +# never masked by it. Keep an explicit modulo unless the count is a constant +# already in range; it costs nothing, because it CSEs with the mask the shader +# had before the fusion. fuse_ubfe = [] for bits in range(1, 32): fuse_ubfe.extend([ (('iand', ('ushr', 'a@32', '#b(is_ult_32)'), (1 << bits) - 1), - ('ubitfield_extract', a, b, bits)), + ('ubfe', a, b, bits)), (('iand', ('ushr', 'a@32', 'b(is_not_const)'), (1 << bits) - 1), - ('ubitfield_extract', a, ('iand', b, 31), bits)), + ('ubfe', a, ('iand', b, 31), bits)), ]) # (x * y) + s = (x * y) + (s << 0) From 9681f68c6b9d299a01a07647bcc5e49dad60b0f1 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 19/25] agx/isa: decode the matrix instructions with a cached destination simd_matrix_fmadd16/32 fixed bits 0 through 7, but bit 7 is the first bit of the destination operand, where it carries the register cache hint. Any matrix instruction with a cached destination therefore failed to decode, which trips the compiler's own disassembly selftest on a debug build. Match the seven opcode bits and let the operand own bit 7. Add three cases: both opcodes with a plain destination, and fmadd16 with a cached one. (cherry picked from commit 83505354e7e0ae5fbea9cf7837b87b50a1283181) --- src/asahi/isa/AGX2.xml | 4 ++-- src/asahi/isa/test/cases.txt | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/asahi/isa/AGX2.xml b/src/asahi/isa/AGX2.xml index bba615285eff..04ac235e449d 100644 --- a/src/asahi/isa/AGX2.xml +++ b/src/asahi/isa/AGX2.xml @@ -980,7 +980,7 @@ - 01101111 + 1101111 1 1 1 @@ -991,7 +991,7 @@ - 01101111 + 1101111 0 1 1 diff --git a/src/asahi/isa/test/cases.txt b/src/asahi/isa/test/cases.txt index e017f3f01833..2e44c3295c8f 100644 --- a/src/asahi/isa/test/cases.txt +++ b/src/asahi/isa/test/cases.txt @@ -402,6 +402,9 @@ 6f8e46608000 simd r3h, r3l, prefix, xor 6f9a4c600080 simd r6h, r6l, prefix, iadd 6f9d4c620080 simd r7, r6, prefix, iadd +6f08c2680cc80080 simd_matrix_fmadd16 r2l, ^r1l, ^r3l, ^r4l +6f08c26c0cc80080 simd_matrix_fmadd32 r2l, ^r1l.abs.neg, ^r3l, ^r4l +ef08c2680cc80080 simd_matrix_fmadd16 $r2l, ^r1l, ^r3l, ^r4l 7101040000620f00 smpi r0_r1_r2_r3, ts0, ss0, r2, 0, 2d, xyzw, auto_lod, a 72001004 mov r0l, thread_in_grid_x 72021104 mov r0h, thread_in_grid_y From 872507eb0e82280336ed7f36757feaf0590244ec 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 20/25] 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. (cherry picked from commit 89b12dd84aa9f994a50c301f5c19d3cc9519c91c) --- src/asahi/compiler/agx_nir_lower_math.c | 76 ++++++++++++++++--------- 1 file changed, 49 insertions(+), 27 deletions(-) diff --git a/src/asahi/compiler/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c index d5f474c73676..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); From e7bb8c9a6cdd6ebe068dd132b3bc478355ad3346 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 21/25] 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) (cherry picked from commit b728b366077914973ff28e8c546e5ffe9b9207ee) --- 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 81d4fdb51b1b2c951fd524affcc789bd95575fe0 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 22/25] 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) (cherry picked from commit 8c9f4064c41dc2c0bc1fae980ead82ba9b4323b1) --- 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 c844d9c5faa7b0dba568682cc24123d6795bf6b6 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 23/25] 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 f426145f7391343e1c63b2e1c8d013011408a1e1 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 24/25] 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 From cdcb92bcbe54ab5fb6e83bda14e9c767f6516c5b 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 25/25] agx: only adopt an early-killed source region when it is free --- src/asahi/compiler/agx_register_allocate.c | 29 +++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/src/asahi/compiler/agx_register_allocate.c b/src/asahi/compiler/agx_register_allocate.c index 312b31e21317..1dcf2c08243e 100644 --- a/src/asahi/compiler/agx_register_allocate.c +++ b/src/asahi/compiler/agx_register_allocate.c @@ -92,6 +92,7 @@ struct ra_ctx { * when shuffling. */ bool early_killed; + unsigned early_killed_reg; /* Maintained while assigning registers. Count of registers required, i.e. * the maximum register assigned + 1. @@ -835,6 +836,11 @@ pick_regs(struct ra_ctx *rctx, agx_instr *I, unsigned d) unsigned align = count; + /* Early kills free a region reserved for this destination. Use it directly + * rather than falling through to live-range splitting, which needs the + * killed sources to remain allocated until the instruction. */ + if (rctx->early_killed) + return rctx->early_killed_reg; /* Try to allocate entire phi webs compatibly */ unsigned phi_idx = phi_web_find(rctx->phi_web, idx.value); if (rctx->phi_web[phi_idx].assigned) { @@ -1039,6 +1045,9 @@ try_kill_early_sources(struct ra_ctx *rctx, const agx_instr *I, unsigned first_source, unsigned last_source, unsigned region_end, unsigned region_base) { + /* A memory destination cannot use the freed GPR source region. */ + if (ra_class_for_index(I->dest[0]) != RA_GPR) + return; unsigned dest_size = util_next_power_of_two(rctx->ncomps[I->dest[0].value]); unsigned dest_end = region_base + dest_size; @@ -1054,12 +1063,26 @@ try_kill_early_sources(struct ra_ctx *rctx, const agx_instr *I, return; for (unsigned s = first_source; s <= last_source; ++s) { - if (early_killable(I, s)) { + if (early_killable(I, s)) kill_source(rctx, I, s); - rctx->early_killed = true; - I->src[s].kill = false; + } + + if (BITSET_TEST_COUNT(rctx->used_regs[RA_GPR], region_base, dest_size)) { + for (unsigned s = first_source; s <= last_source; ++s) { + if (early_killable(I, s)) + BITSET_SET_COUNT(rctx->used_regs[RA_GPR], + rctx->ssa_to_reg[I->src[s].value], + rctx->ncomps[I->src[s].value]); } + return; + } + + for (unsigned s = first_source; s <= last_source; ++s) { + if (early_killable(I, s)) + I->src[s].kill = false; } + rctx->early_killed_reg = region_base; + rctx->early_killed = true; } /** Assign registers to SSA values in a block. */