diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index e131bcf07211..9bbfb6e97620 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -1697,6 +1697,24 @@ 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 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_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; + } + case nir_intrinsic_doorbell_agx: { return agx_doorbell(b, nir_src_as_uint(instr->src[0])); } @@ -2020,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); @@ -2723,46 +2759,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) @@ -2953,6 +2949,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. */ @@ -3533,56 +3534,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) @@ -3590,6 +3541,29 @@ agx_preprocess_nir(nir_shader *nir) if (!nir) return; + /* 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()) { + 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); /* Lower large arrays to scratch and small arrays to csel */ @@ -3614,7 +3588,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); @@ -3637,9 +3611,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..246952739821 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, with a software + * fallback for partial subgroups. 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); @@ -347,7 +357,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_compiler.h b/src/asahi/compiler/agx_compiler.h index 45460934cf0f..1d9ee7d976f1 100644 --- a/src/asahi/compiler/agx_compiler.h +++ b/src/asahi/compiler/agx_compiler.h @@ -1096,6 +1096,8 @@ 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_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.h b/src/asahi/compiler/agx_nir.h index c03715027e5c..719ece490c92 100644 --- a/src/asahi/compiler/agx_nir.h +++ b/src/asahi/compiler/agx_nir.h @@ -22,6 +22,10 @@ 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_fdiv_late(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_algebraic.py b/src/asahi/compiler/agx_nir_algebraic.py index b1e036459d21..6f657ff052e6 100644 --- a/src/asahi/compiler/agx_nir_algebraic.py +++ b/src/asahi/compiler/agx_nir_algebraic.py @@ -141,11 +141,24 @@ ('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. 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), (1 << bits) - 1), - ('ubitfield_extract', a, b, bits)) + (('iand', ('ushr', 'a@32', '#b(is_ult_32)'), (1 << bits) - 1), + ('ubfe', a, b, bits)), + (('iand', ('ushr', 'a@32', 'b(is_not_const)'), (1 << bits) - 1), + ('ubfe', a, ('iand', b, 31), bits)), ]) # (x * y) + s = (x * y) + (s << 0) 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..c7ea9821035d --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_cmat.c @@ -0,0 +1,578 @@ +/* + * 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) +{ + const unsigned bits = ptr->def.bit_size; + nir_def *row_comps = nir_udiv_imm( + 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, nir_u2uN(b, minor, bits))); +} + +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/agx_nir_lower_math.c b/src/asahi/compiler/agx_nir_lower_math.c new file mode 100644 index 000000000000..fc9d1a1e4569 --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_math.c @@ -0,0 +1,546 @@ +/* + * 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. + */ + +/* The caller keeps x near one so the Newton residual cannot flush. */ +static nir_def * +rcp_normal(nir_builder *b, nir_def *x) +{ + nir_def *u = nir_frcp(b, x); + return nir_ffma(b, nir_ffma(b, nir_fneg(b, x), u, + nir_imm_float(b, 1.0)), u, u); +} + +/* 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 *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 +lower_fdiv(nir_builder *b, nir_alu_instr *alu, void *data) +{ + 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) + 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) +{ + 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); +} + +/* + * 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 without a lookup 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_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); + 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_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 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) +{ + return nir_shader_alu_pass(s, lower_log2, nir_metadata_control_flow, NULL); +} + +/* + * 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). 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); + + 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/agx_nir_lower_simdmat.c b/src/asahi/compiler/agx_nir_lower_simdmat.c new file mode 100644 index 000000000000..01f27f4372fd --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_simdmat.c @@ -0,0 +1,570 @@ +/* + * 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). + * + * 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 + * 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. + * + * 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" +#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 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); + + 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). Verified MUL_MAT 947/947 + PPL 1.0162, + * pp512 177.1 -> 179.4. */ + nir_deref_instr *sdref = + nir_build_deref_cast(b, &deref->def, deref->modes, + glsl_scalar_type(desc.element_type), tsz); + + /* 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 + * 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 && width_matched; + nir_deref_instr *v2dref = + 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). 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 = + 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; + /* 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 && 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. */ + 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]); + + /* 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 + * (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. 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; +} + +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; +} + +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; +} + +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); + 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); + 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_register_allocate.c b/src/asahi/compiler/agx_register_allocate.c index 87a0335a4a79..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) { @@ -1014,11 +1020,34 @@ 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, 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; @@ -1034,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 (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; + } + + 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. */ @@ -1132,7 +1175,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/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..e54d5bff7bda 100644 --- a/src/asahi/compiler/meson.build +++ b/src/asahi/compiler/meson.build @@ -14,8 +14,11 @@ 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', + 'agx_nir_lower_cmat.c', 'agx_nir_lower_texture.c', 'agx_nir_opt_preamble.c', 'agx_lower_64bit.c', @@ -106,6 +109,8 @@ 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', 'test/test-lower-pseudo.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(); +} 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); +} diff --git a/src/asahi/isa/AGX2.xml b/src/asahi/isa/AGX2.xml index f3e6733725fb..04ac235e449d 100644 --- a/src/asahi/isa/AGX2.xml +++ b/src/asahi/isa/AGX2.xml @@ -979,6 +979,28 @@ + + 1101111 + 1 + 1 + 1 + + + + + + + + 1101111 + 0 + 1 + 1 + + + + + + 111010 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 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..ea79700896bc 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); } @@ -773,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 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..8591d5859dbf 100644 --- a/src/asahi/vulkan/hk_device.c +++ b/src/asahi/vulkan/hk_device.c @@ -44,6 +44,13 @@ static const struct debug_named_value hk_perf_options[] = { {"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 +332,14 @@ 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; + + /* 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/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) diff --git a/src/asahi/vulkan/hk_physical_device.c b/src/asahi/vulkan/hk_physical_device.c index 80112cec02b2..bb04ae12eaa1 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 = agx_simdmat_enabled(), .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, AGX_SIMDMAT=0 disables) */ + .cooperativeMatrix = agx_simdmat_enabled(), + .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 = + agx_simdmat_enabled() ? VK_SHADER_STAGE_COMPUTE_BIT : 0, .subgroupSupportedStages = VK_SHADER_STAGE_COMPUTE_BIT | VK_SHADER_STAGE_ALL_GRAPHICS, .subgroupSupportedOperations = @@ -1518,3 +1525,47 @@ 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 (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 + * 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; + } + } + } + } + + return vk_outarray_status(&out); +} 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 * 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)) 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..9e220c0fee3f 100644 --- a/src/compiler/nir/nir_intrinsics.py +++ b/src/compiler/nir/nir_intrinsics.py @@ -2472,6 +2472,14 @@ 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_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. 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=[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]) 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); 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