diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c index e131bcf07211..32bef4219cf1 100644 --- a/src/asahi/compiler/agx_compile.c +++ b/src/asahi/compiler/agx_compile.c @@ -1028,7 +1028,12 @@ static void agx_emit_local_load(agx_builder *b, agx_index dst, nir_intrinsic_instr *instr) { agx_index base = agx_local_base(instr->src[0]); - agx_index index = agx_zero(); /* TODO: optimize address arithmetic */ + agx_index index = agx_zero(); + if (nir_intrinsic_base(instr)) { + unsigned unit = 4 * (instr->def.bit_size / 8); + assert(nir_intrinsic_base(instr) % unit == 0); + index = agx_immediate(nir_intrinsic_base(instr) / unit); + } assert(base.size == AGX_SIZE_16); enum agx_format format = format_for_bitsize(instr->def.bit_size); @@ -1044,7 +1049,12 @@ agx_emit_local_store(agx_builder *b, nir_intrinsic_instr *instr) { agx_index value = agx_src_index(&instr->src[0]); agx_index base = agx_local_base(instr->src[1]); - agx_index index = agx_zero(); /* TODO: optimize address arithmetic */ + agx_index index = agx_zero(); + if (nir_intrinsic_base(instr)) { + unsigned unit = 4 * (nir_src_bit_size(instr->src[0]) / 8); + assert(nir_intrinsic_base(instr) % unit == 0); + index = agx_immediate(nir_intrinsic_base(instr) / unit); + } assert(base.size == AGX_SIZE_16); enum agx_format format = format_for_bitsize(nir_src_bit_size(instr->src[0])); @@ -1697,6 +1707,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 +2048,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 +2769,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) @@ -2834,6 +2840,162 @@ agx_mem_vectorize_cb(unsigned align_mul, unsigned align_offset, return true; } +/* Peel a constant addend out of a 16-bit shared address (through the u2u16 + * narrowing and nested iadds). The lload/lstore index immediate adds 4 elements + * of the access type per unit (measured on G13G), so the multiple-of-unit part + * of the constant moves into nir_intrinsic_base (bytes) and the remainder stays + * in the address. */ +static nir_def * +agx_peel_const(nir_builder *b, nir_def *def, int64_t *c) +{ + nir_scalar s = nir_scalar_resolved(def, 0); + if (!nir_scalar_is_alu(s)) + return def; + nir_op op = nir_scalar_alu_op(s); + if (op == nir_op_u2u16 || op == nir_op_u2u32) { + nir_scalar src = nir_scalar_chase_alu_src(s, 0); + nir_def *inner = agx_peel_const(b, src.def, c); + if (inner == src.def) + return def; + return op == nir_op_u2u16 ? nir_u2u16(b, inner) : nir_u2u32(b, inner); + } + if (op == nir_op_iadd) { + nir_scalar a = nir_scalar_chase_alu_src(s, 0); + nir_scalar bb = nir_scalar_chase_alu_src(s, 1); + if (nir_scalar_is_const(bb)) { + *c += nir_scalar_as_int(bb); + return agx_peel_const(b, a.def, c); + } + if (nir_scalar_is_const(a)) { + *c += nir_scalar_as_int(a); + return agx_peel_const(b, bb.def, c); + } + } + return def; +} + +struct fold_cand { + nir_intrinsic_instr *intr; + unsigned oi; + nir_def *stripped; /* address with constant addends removed */ + nir_def *key; /* grouping key: the value under the u2u16 narrowing */ + int64_t c; /* constant addend (bytes) */ + int64_t unit; /* bytes per index unit for this access */ +}; + +/* Group the shared accesses of one block that differ only by constants. The + * lowest constant stays in the (single) base address, which is then a real, + * in-range address; every other access gets its delta as the lload/lstore + * index immediate (4 elements of the access type per unit) plus any remainder + * in the address. Folding the whole constant instead would add it to a + * possibly wrapped 16-bit base, and the hardware index add does not wrap. */ +static bool +agx_fold_shared_block(nir_builder *b, nir_block *block) +{ + struct fold_cand cands[512]; + unsigned n = 0; + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_intrinsic || n == ARRAY_SIZE(cands)) + continue; + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + unsigned oi; + if (intr->intrinsic == nir_intrinsic_load_shared) + oi = 0; + else if (intr->intrinsic == nir_intrinsic_store_shared) + oi = 1; + else + continue; + if (nir_intrinsic_base(intr) != 0) + continue; + nir_def *addr = intr->src[oi].ssa; + if (addr->bit_size != 16) + continue; + b->cursor = nir_before_instr(instr); + int64_t c = 0; + nir_def *stripped = agx_peel_const(b, addr, &c); + if (stripped == addr && c == 0) { + /* no constant to peel; still a candidate as the group anchor */ + } + unsigned bits = intr->intrinsic == nir_intrinsic_load_shared + ? intr->def.bit_size + : nir_src_bit_size(intr->src[0]); + nir_def *key = stripped; + { + nir_scalar ss = nir_scalar_resolved(stripped, 0); + if (nir_scalar_is_alu(ss) && nir_scalar_alu_op(ss) == nir_op_u2u16) + key = nir_scalar_chase_alu_src(ss, 0).def; + } + cands[n++] = (struct fold_cand){intr, oi, stripped, key, c, + 4 * (int64_t)(bits / 8)}; + } + bool progress = false; + for (unsigned i = 0; i < n; i++) { + if (!cands[i].intr) + continue; + /* collect the group sharing this stripped address */ + unsigned members[512]; + unsigned m = 0; + int64_t cmin = cands[i].c; + for (unsigned j = i; j < n; j++) { + if (cands[j].intr && cands[j].key == cands[i].key) { + members[m++] = j; + cmin = MIN2(cmin, cands[j].c); + } + } + if (m < 2) { + cands[i].intr = NULL; + continue; + } + /* the first member in block order anchors the shared base */ + struct fold_cand *first = &cands[members[0]]; + b->cursor = nir_before_instr(&first->intr->instr); + nir_def *base = cands[i].stripped; + if (base->bit_size != 16) + base = nir_u2u16(b, base); + if (cmin) + base = nir_iadd_imm(b, base, cmin); + for (unsigned k = 0; k < m; k++) { + struct fold_cand *cd = &cands[members[k]]; + int64_t d = cd->c - cmin; + int64_t fold = (d / cd->unit) * cd->unit; + int64_t rest = d - fold; + if (fold / cd->unit >= 0x10000) + continue; + nir_def *na = base; + if (rest) { + b->cursor = nir_before_instr(&cd->intr->instr); + na = nir_iadd_imm(b, base, rest); + } + nir_intrinsic_set_base(cd->intr, (int)fold); + nir_src_rewrite(&cd->intr->src[cd->oi], na); + progress = true; + cd->intr = NULL; + } + } + return progress; +} + +static bool +agx_fold_shared_offsets(nir_shader *nir) +{ + bool progress = false; + nir_foreach_function_impl(impl, nir) { + nir_builder b = nir_builder_create(impl); + nir_foreach_block(block, impl) { + progress |= agx_fold_shared_block(&b, block); + } + nir_progress(progress, impl, nir_metadata_none); + } + return progress; +} + +/* Exported for the unit tests. */ +bool +agx_nir_fold_shared_offsets(nir_shader *nir) +{ + return agx_fold_shared_offsets(nir); +} + static bool set_speculate(nir_builder *b, nir_instr *instr, UNUSED void *_) { @@ -2854,8 +3016,8 @@ set_speculate(nir_builder *b, nir_instr *instr, UNUSED void *_) } static void -agx_optimize_nir(nir_shader *nir, bool soft_fault, uint16_t *preamble_size, - uint8_t *ts_count, uint8_t *ss_count) +agx_optimize_nir(nir_shader *nir, bool soft_fault, bool fold_subformat_address, + uint16_t *preamble_size, uint8_t *ts_count, uint8_t *ss_count) { /* This runs only once up front since other optimizations don't affect it */ NIR_PASS(_, nir, nir_opt_shrink_stores, true); @@ -2907,7 +3069,7 @@ agx_optimize_nir(nir_shader *nir, bool soft_fault, uint16_t *preamble_size, * want to vectorize first since nir_opt_load_store_vectorize doesn't know * how to handle our loads. Likewise for uniform atomic optimization. */ - NIR_PASS(_, nir, agx_nir_lower_address); + NIR_PASS(_, nir, agx_nir_lower_address, fold_subformat_address); NIR_PASS(progress, nir, nir_opt_uniform_subgroup, &subgroups_options); if (progress) { @@ -2953,6 +3115,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. */ @@ -3402,10 +3569,22 @@ agx_compile_function_nir(nir_shader *nir, nir_function_impl *impl, if (agx_should_dump(nir, AGX_DBG_SHADERS)) agx_print_shader(ctx, stdout); - if (likely(!(agx_compiler_debug & AGX_DBG_NOSCHED))) { + /* The bottom-up register-pressure scheduler moves each load next to its use, + * which serializes the memory round trips of small compute kernels (the + * LLM decode chain: +1.7% tokens/s with it off, bit-identical results). + * Compute shaders on G13 therefore skip it unless AGX_SCHED_COMPUTE is set. + * Other generations keep it until they are measured. */ + bool skip_sched = (agx_compiler_debug & AGX_DBG_NOSCHED) || + (nir->info.stage == MESA_SHADER_COMPUTE && + ctx->key->dev.g13 && !getenv("AGX_SCHED_COMPUTE")); + if (likely(!skip_sched)) { AGX_PASS(ctx, agx_pressure_schedule); } + if (nir->info.stage == MESA_SHADER_COMPUTE && ctx->key->dev.hoist_loads) { + AGX_PASS(ctx, agx_hoist_loads); + } + if (agx_should_dump(nir, AGX_DBG_SHADERS)) agx_print_shader(ctx, stdout); @@ -3533,56 +3712,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 +3719,30 @@ 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, agx_nir_lower_bf16_conversions); + NIR_PASS(_, nir, nir_lower_vars_to_ssa); /* Lower large arrays to scratch and small arrays to csel */ @@ -3614,7 +3767,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 +3790,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); @@ -3747,7 +3902,8 @@ agx_compile_shader_nir(nir_shader *nir, struct agx_shader_key *key, info->push_count = key->reserved_preamble; agx_optimize_nir( - nir, key->dev.soft_fault, key->secondary ? NULL : &info->push_count, + nir, key->dev.soft_fault, key->dev.fold_subformat_address, + key->secondary ? NULL : &info->push_count, (key->secondary || !key->promote_textures) ? NULL : &info->texture_state_count, (key->secondary || !key->promote_textures) ? NULL @@ -3758,6 +3914,11 @@ agx_compile_shader_nir(nir_shader *nir, struct agx_shader_key *key, assign_coefficient_regs(nir, &info->varyings.fs); } + if (key->dev.g13 && !getenv("AGX_LOCAL_FOLD_OFF")) { + NIR_PASS(_, nir, agx_fold_shared_offsets); + NIR_PASS(_, nir, nir_opt_dce); + } + if (agx_should_dump(nir, AGX_DBG_SHADERS)) nir_print_shader(nir, stdout); diff --git a/src/asahi/compiler/agx_compile.h b/src/asahi/compiler/agx_compile.h index c8ffaa5db23b..b2ba08ca338a 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" @@ -268,6 +269,24 @@ struct agx_device_key { * kernel, but that's functionally a hardware feature. */ bool soft_fault; + + /* Is the target GPU generation 13 (M1 family)? Selects the compute + * codegen defaults that are measured on G13 only. + */ + bool g13; + + /* Hoist device loads of compute shaders above independent ALU, barriers + * and workgroup memory accesses (agx_hoist_loads). Measured on G13 and + * G14 only. + */ + bool hoist_loads; + + /* Index global accesses whose lea shift is below the format shift (the + * cooperative-matrix loads) by the scaled offset instead of a 64-bit + * address. Measured per device: on for the single-cluster G13 (M1), off + * on the multi-cluster G13C, where it slows dense GEMM. + */ + bool fold_subformat_address; }; struct agx_shader_key { @@ -317,6 +336,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 +375,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..21c81dc89cd0 100644 --- a/src/asahi/compiler/agx_compiler.h +++ b/src/asahi/compiler/agx_compiler.h @@ -1031,6 +1031,7 @@ void agx_opt_compact_constants(agx_context *ctx); void agx_opt_promote_constants(agx_context *ctx); void agx_dce(agx_context *ctx, bool partial); void agx_pressure_schedule(agx_context *ctx); +void agx_hoist_loads(agx_context *ctx); void agx_spill(agx_context *ctx, unsigned k, bool remat_only); void agx_repair_ssa(agx_context *ctx); void agx_reindex_ssa(agx_context *ctx); @@ -1096,6 +1097,13 @@ 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); +nir_def *agx_cmat_convert(struct nir_builder *b, nir_def *src, + enum glsl_base_type from, enum glsl_base_type to); +nir_def *agx_cmat_alu(struct nir_builder *b, nir_op op, enum glsl_base_type t, + nir_def *x, nir_def *y); +bool agx_nir_lower_bf16_conversions(nir_shader *shader); bool agx_nir_lower_frag_sidefx(nir_shader *s); struct agx_cycle_estimate { diff --git a/src/asahi/compiler/agx_ir.c b/src/asahi/compiler/agx_ir.c index 3eb14815d6e6..3eadfdb2cdd5 100644 --- a/src/asahi/compiler/agx_ir.c +++ b/src/asahi/compiler/agx_ir.c @@ -11,6 +11,8 @@ agx_allows_16bit_immediate(agx_instr *I) { return (I->op == AGX_OPCODE_DEVICE_LOAD) || (I->op == AGX_OPCODE_DEVICE_STORE) || + (I->op == AGX_OPCODE_LOCAL_LOAD) || + (I->op == AGX_OPCODE_LOCAL_STORE) || (I->op == AGX_OPCODE_STACK_LOAD) || (I->op == AGX_OPCODE_STACK_STORE) || (I->op == AGX_OPCODE_UNIFORM_STORE) || (I->op == AGX_OPCODE_ATOMIC) || diff --git a/src/asahi/compiler/agx_nir.h b/src/asahi/compiler/agx_nir.h index c03715027e5c..f125f653f266 100644 --- a/src/asahi/compiler/agx_nir.h +++ b/src/asahi/compiler/agx_nir.h @@ -10,7 +10,8 @@ struct nir_shader; struct nir_instr; -bool agx_nir_lower_address(struct nir_shader *shader); +bool agx_nir_lower_address(struct nir_shader *shader, + bool fold_subformat_address); bool agx_nir_lower_algebraic_late(struct nir_shader *shader); bool agx_nir_cleanup_amul(struct nir_shader *shader); bool agx_nir_fuse_lea(struct nir_shader *shader); @@ -22,6 +23,11 @@ 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_fold_shared_offsets(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_address.c b/src/asahi/compiler/agx_nir_lower_address.c index d59860f7c9de..d510d3d54eed 100644 --- a/src/asahi/compiler/agx_nir_lower_address.c +++ b/src/asahi/compiler/agx_nir_lower_address.c @@ -31,6 +31,63 @@ format_for_bitsize(unsigned bitsize) } } +/* True when the scalar is provably a multiple of 2^k. The sub-format fold + * indexes by the byte offset scaled down, which equals the byte offset only + * when its low bits are zero; an access that is aligned as a whole does not + * say that the offset alone is (base P + 2 with offset 2 for a 32-bit access + * is aligned, and scaling the offset down drops the 2). Unknown means no fold. + */ +static bool +scalar_multiple_of(nir_scalar s, unsigned k, unsigned depth) +{ + if (k == 0) + return true; + + s = nir_scalar_resolved(s.def, s.comp); + if (nir_scalar_is_const(s)) + return (nir_scalar_as_uint(s) & BITFIELD64_MASK(k)) == 0; + + if (!nir_scalar_is_alu(s) || depth == 0) + return false; + + switch (nir_scalar_alu_op(s)) { + case nir_op_ishl: { + nir_scalar amount = nir_scalar_chase_alu_src(s, 1); + nir_scalar value = nir_scalar_chase_alu_src(s, 0); + if (nir_scalar_is_const(amount)) { + unsigned a = nir_scalar_as_uint(amount); + if (a >= k) + return true; + return scalar_multiple_of(value, k - a, depth - 1); + } + return scalar_multiple_of(value, k, depth - 1); + } + case nir_op_imul: + return scalar_multiple_of(nir_scalar_chase_alu_src(s, 0), k, depth - 1) || + scalar_multiple_of(nir_scalar_chase_alu_src(s, 1), k, depth - 1); + case nir_op_iadd: + case nir_op_isub: + return scalar_multiple_of(nir_scalar_chase_alu_src(s, 0), k, depth - 1) && + scalar_multiple_of(nir_scalar_chase_alu_src(s, 1), k, depth - 1); + case nir_op_iand: { + nir_scalar mask = nir_scalar_chase_alu_src(s, 1); + if (nir_scalar_is_const(mask) && + (nir_scalar_as_uint(mask) & BITFIELD64_MASK(k)) == 0) + return true; + return scalar_multiple_of(nir_scalar_chase_alu_src(s, 0), k, depth - 1) || + scalar_multiple_of(nir_scalar_chase_alu_src(s, 1), k, depth - 1); + } + case nir_op_u2u32: + case nir_op_i2i32: + case nir_op_u2u64: + case nir_op_i2i64: + case nir_op_mov: + return scalar_multiple_of(nir_scalar_chase_alu_src(s, 0), k, depth - 1); + default: + return false; + } +} + static bool pass(struct nir_builder *b, nir_intrinsic_instr *intr, void *data) { @@ -61,12 +118,37 @@ pass(struct nir_builder *b, nir_intrinsic_instr *intr, void *data) nir_op op = nir_scalar_alu_op(base); if (op == nir_op_ulea_agx || op == nir_op_ilea_agx) { unsigned shift = nir_scalar_as_uint(nir_scalar_chase_alu_src(base, 2)); + bool sign_extend = (op == nir_op_ilea_agx); if (shift >= format_shift && shift <= max_shift) { match = (struct match){ .base = nir_scalar_chase_alu_src(base, 0), .offset = nir_scalar_chase_alu_src(base, 1), .shift = shift - format_shift, - .sign_extend = (op == nir_op_ilea_agx), + .sign_extend = sign_extend, + }; + } else if (shift < format_shift && *(const bool *)data && + scalar_multiple_of(nir_scalar_chase_alu_src(base, 1), + format_shift - shift, 4)) { + /* A byte (or half-element) offset, as the cooperative matrix + * loads produce. The offset is provably a multiple of the element + * size (scalar_multiple_of; an aligned access does not imply it): + * index by the offset scaled down instead of materializing the + * 64-bit sum (a 64-bit add plus the moves that build its + * operands). Per device + * (agx_device_key::fold_subformat_address): measured +15 % on + * dense f16 GEMM on the single-cluster M1, -13..17 % on the M1 + * Max, bit-identical on both. + */ + nir_scalar off = nir_scalar_chase_alu_src(base, 1); + nir_def *index = nir_channel(b, off.def, off.comp); + unsigned down = format_shift - shift; + index = sign_extend ? nir_ishr_imm(b, index, down) + : nir_ushr_imm(b, index, down); + match = (struct match){ + .base = nir_scalar_chase_alu_src(base, 0), + .offset = nir_get_scalar(index, 0), + .shift = 0, + .sign_extend = sign_extend, }; } } else if (op == nir_op_iadd) { @@ -151,7 +233,7 @@ pass(struct nir_builder *b, nir_intrinsic_instr *intr, void *data) } bool -agx_nir_lower_address(nir_shader *nir) +agx_nir_lower_address(nir_shader *nir, bool fold_subformat_address) { bool progress = false; @@ -169,7 +251,7 @@ agx_nir_lower_address(nir_shader *nir) /* Next, lower load/store using the lea's */ NIR_PASS(progress, nir, nir_shader_intrinsics_pass, pass, - nir_metadata_control_flow, NULL); + nir_metadata_control_flow, &fold_subformat_address); /* Finally, lower any leftover lea instructions back to ALU to let * nir_opt_algebraic simplify them from here. 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..a862e17f0c13 --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_cmat.c @@ -0,0 +1,550 @@ +/* + * 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. + * + * Note: 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]); + /* Every advertised shape has one A/B type; B converts like A below. */ + assert(da.element_type == src_desc(orig, intr->src[2]).element_type); + + /* 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; + + /* 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); + /* A/B narrower than the accumulator (fp16 or bf16 A/B, fp32 C): the HW + * form computes the products at accumulator width; mirror that. */ + for (unsigned e = 0; e < m * k; e++) { + av[e] = agx_cmat_convert(b, gather_elem(b, a, e, owner), da.element_type, + dc.element_type); + owner = owner_next(b, owner, a_cnt); + } + owner = nir_imm_int(b, 0); + for (unsigned e = 0; e < k * n; e++) { + bv[e] = agx_cmat_convert(b, gather_elem(b, mat_b, e, owner), + da.element_type, dc.element_type); + 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++) { + acc = nir_ffma(b, av[i * k + kk], bv[kk * n + j], 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); + struct glsl_cmat_description desc = src_desc(orig, intr->src[1]); + unsigned elems = desc.rows * desc.cols; + + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); + nir_deref_instr *c = kind == 2 ? nir_src_as_deref(intr->src[2]) : NULL; + for (unsigned e = 0; e < elems; e++) { + nir_def *y = kind == 0 ? NULL + : kind == 1 ? intr->src[2].ssa + : load_elem(b, c, e); + store_elem(b, dst, e, + agx_cmat_alu(b, op, desc.element_type, load_elem(b, a, e), y)); + } + 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_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++) + store_elem(b, dst, e, + agx_cmat_convert(b, sv[e], sd.element_type, dd.element_type)); + 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..fd0a7d21e7b8 --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_math.c @@ -0,0 +1,635 @@ +/* + * 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 = nir_frcp(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); + 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); +} + +/* + * 1/x, bit for bit div_rn(1, x): one Newton step on the hardware rcp, which + * flushes denormal inputs to +-inf and results below 2^-126 to +-0. Where the + * step is NaN (x zero, denormal, infinite or NaN), the hardware rcp already + * holds the answer, and NaN becomes the canonical NaN. tests/div_window.py + * checks every input. + */ +static nir_def * +rcp_rn(nir_builder *b, nir_def *x) +{ + nir_def *u = nir_frcp(b, x); + nir_def *q = nir_ffma(b, nir_ffma(b, nir_fneg(b, x), u, + nir_imm_float(b, 1.0)), u, u); + q = nir_bcsel(b, nir_fisnan(b, q), u, q); + return nir_bcsel(b, nir_fisnan(b, q), nir_imm_int(b, 0x7fc00000), q); +} + +static bool +is_one(nir_alu_instr *alu, unsigned i) +{ + nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i); + return nir_scalar_is_const(s) && nir_scalar_as_uint(s) == 0x3f800000; +} + +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 && alu->def.num_components == 1 && + is_one(alu, 0)) + res = rcp_rn(b, d); + else 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); +} + +/* + * Natural log and log2 as double-float (fp32 pair) arithmetic. + * + * x = 2^e m with m in [sqrt(1/2), sqrt(2)). With s = (m - 1) / (m + 1), + * + * log(m) = 2s + 2s^3 (1/3 + s^2/5 + s^4/7 + ...) + * + * s is kept as a pair (s, s_lo) from one correctly rounded division plus the + * fma residual, the cubic term is kept as a pair, and the terms are summed + * with Fast2Sum so the value handed to the final rounding carries about 40 + * bits. That makes the exact cases (powers of two, log(3)) come out correctly + * rounded and everything else faithful, without any table. + */ +struct log_parts { + nir_def *e; /* exponent as float */ + nir_def *s; /* 2s = leading term, as a pair (s, s_lo) */ + nir_def *s_lo; + nir_def *t_hi; /* 2 s^3 P(s^2), as a pair */ + nir_def *t_lo; +}; + +static nir_def * +fast2sum_lo(nir_builder *b, nir_def *a, nir_def *bb, nir_def *sum) +{ + /* Requires |a| >= |bb| or a == 0 */ + return nir_fsub(b, bb, nir_fsub(b, sum, a)); +} + +static struct log_parts +log_core(nir_builder *b, nir_def *x) +{ + nir_def *bits = x; + nir_def *e = nir_iadd_imm(b, nir_ushr_imm(b, bits, 23), -127); + nir_def *m = nir_ior_imm(b, nir_iand_imm(b, bits, 0x7fffff), 0x3f800000); + + /* Fold m into [sqrt(1/2), sqrt(2)) so the series converges fast and + * e ln2 never cancels more than one bit against log(m). + */ + nir_def *big = nir_fgt_imm(b, m, 1.41421353816986083984375); + m = nir_bcsel(b, big, nir_fmul_imm(b, m, 0.5), m); + e = nir_iadd(b, e, nir_b2i32(b, big)); + + nir_def *one = nir_imm_float(b, 1.0); + nir_def *two = nir_imm_float(b, 2.0); + nir_def *f = nir_fsub(b, m, one); /* exact */ + nir_def *d = nir_fadd(b, two, f); + nir_def *d_lo = fast2sum_lo(b, two, f, d); + + /* s = f / (d + d_lo) as a pair */ + nir_def *y = rcp_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_log(nir_builder *b, nir_def *x) +{ + struct log_parts l = log_core(b, x); + + /* ln2 split so e * LN2_HI is exact for |e| < 2^8 */ + nir_def *a = nir_fmul_imm(b, l.e, 0.693145751953125); + nir_def *bb = nir_fmul_imm(b, l.s, 2.0); + nir_def *h = nir_fadd(b, a, bb); + nir_def *lo = fast2sum_lo(b, a, bb, h); + nir_def *h2 = nir_fadd(b, h, l.t_hi); + nir_def *lo2 = fast2sum_lo(b, h, l.t_hi, h2); + + nir_def *small = nir_fadd(b, lo, lo2); + small = nir_fadd(b, small, l.t_lo); + small = nir_fadd(b, small, nir_fmul_imm(b, l.s_lo, 2.0)); + small = nir_ffma(b, l.e, nir_imm_float(b, 1.428606765330187e-06), small); + return nir_fadd(b, h2, small); +} + +static nir_def * +soft_log2(nir_builder *b, nir_def *x) +{ + struct log_parts l = log_core(b, x); + + /* (2s + 2s_lo + t) / ln2 with 1/ln2 as a pair */ + nir_def *il2 = nir_imm_float(b, 1.4426950216293335); + nir_def *il2_lo = nir_imm_float(b, 1.925963033500011e-08); + nir_def *bb = nir_fmul_imm(b, l.s, 2.0); + nir_def *ph = nir_fmul(b, bb, il2); + nir_def *pl = nir_ffma(b, bb, il2, nir_fneg(b, ph)); + pl = nir_ffma(b, bb, il2_lo, pl); + nir_def *th = nir_fmul(b, l.t_hi, il2); + nir_def *tl = nir_ffma(b, l.t_hi, il2, nir_fneg(b, th)); + nir_def *rest = nir_fadd(b, nir_fmul_imm(b, l.s_lo, 2.0), l.t_lo); + pl = nir_fadd(b, nir_ffma(b, rest, il2, pl), tl); + + nir_def *h = nir_fadd(b, l.e, ph); + nir_def *lo = fast2sum_lo(b, l.e, ph, h); + nir_def *h2 = nir_fadd(b, h, th); + nir_def *lo2 = fast2sum_lo(b, h, th, h2); + return nir_fadd(b, h2, nir_fadd(b, nir_fadd(b, lo, lo2), pl)); +} + +static const uint32_t LN2_F32_BITS = 0x3f317218; + +static bool +src_is_ln2(nir_alu_instr *alu, unsigned i) +{ + nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i); + return nir_scalar_is_const(s) && nir_scalar_as_uint(s) == LN2_F32_BITS; +} + +static bool +src_is_flog2(nir_alu_instr *alu, unsigned i) +{ + nir_scalar s = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), i); + return nir_scalar_is_alu(s) && nir_scalar_alu_op(s) == nir_op_flog2; +} + +/* + * nir_flog builds fmul(flog2(x), ln2); catch that before the flog2 goes and + * round once. An exact (precise) multiply keeps the rounding the source asks + * for, so a precise log2(x) * ln2 gives the same bits from a literal and from + * a buffer. + */ +static bool +lower_log(nir_builder *b, nir_alu_instr *alu, void *_) +{ + if (alu->op != nir_op_fmul || alu->def.bit_size != 32 || + alu->def.num_components != 1 || nir_alu_instr_is_exact(alu)) + return false; + + unsigned li; + if (src_is_flog2(alu, 0) && src_is_ln2(alu, 1)) + li = 0; + else if (src_is_flog2(alu, 1) && src_is_ln2(alu, 0)) + li = 1; + else + return false; + + nir_scalar log = nir_scalar_chase_alu_src(nir_get_scalar(&alu->def, 0), li); + nir_scalar xs = nir_scalar_chase_alu_src(log, 0); + + b->cursor = nir_before_instr(&alu->instr); + b->fp_math_ctrl = nir_fp_no_fast_math; + nir_def *x = nir_channel(b, xs.def, xs.comp); + nir_def *res = nir_bcsel(b, is_positive_normal(b, x), soft_log(b, x), + log_special(b, x)); + nir_def_replace(&alu->def, res); + + if (nir_def_is_unused(log.def)) + nir_instr_remove(nir_def_instr(log.def)); + + return true; +} + +static bool +lower_log2(nir_builder *b, nir_alu_instr *alu, void *_) +{ + if (alu->op != nir_op_flog2 || alu->def.bit_size != 32) + return false; + + b->cursor = nir_before_instr(&alu->instr); + b->fp_math_ctrl = nir_fp_no_fast_math; + nir_def *x = nir_ssa_for_alu_src(b, alu, 0); + nir_def *res = nir_bcsel(b, is_positive_normal(b, x), soft_log2(b, x), + log_special(b, x)); + nir_def_replace(&alu->def, res); + return true; +} + +bool +agx_nir_lower_log(nir_shader *s) +{ + bool progress = + nir_shader_alu_pass(s, lower_log, nir_metadata_control_flow, NULL); + progress |= + nir_shader_alu_pass(s, lower_log2, nir_metadata_control_flow, NULL); + return progress; +} + +/* + * Sine and cosine. Range reduction produces k = round(x 2/pi) and the fraction + * fq = x 2/pi - k in [-1/2, 1/2] quadrants; the quadrant then picks a Taylor + * polynomial for sin(pi/2 fq) or cos(pi/2 fq) and a sign. The hardware sin + * unit (sin_pt_1/sin_pt_2, nir_op_fsin_agx) is not used: it is a few ulp off, + * which is more than the reduction now loses. + * + * |x| < 2^22: Cody-Waite in radians with pi/2 in three fp32 pieces. The first + * fma is exact, the remaining rounding is tracked as a low part so fq is + * within an ulp of the true reduced argument. + * + * |x| >= 2^22: Payne-Hanek. Multiply the 24-bit mantissa by a 128-bit window + * of 2/pi aligned to the exponent; the integer part mod 4 is the quadrant and + * the top 64 fraction bits give fq. fp32 inputs never reduce below ~2^-30 of a + * quadrant, so 64 fraction bits leave more than 24 significant bits. + */ +static const double TWO_OVER_PI = 0.6366197466850281; /* fp32(2/pi) */ +static const double TWO_OVER_PI_LO = 2 / M_PI - 0.6366197466850281; +static const double PIO2_1 = 1.5707963705062866; +static const double PIO2_2 = -4.371138828673793e-08; +static const double PIO2_3 = -1.7151245100058819e-15; + +/* Bits of 2/pi, 32 per word, preceded by three zero words. */ +static const uint32_t two_over_pi_bits[10] = { + 0, 0, 0, 0xa2f9836e, 0x4e441529, 0xfc2757d1, 0xf534ddc0, + 0xdb629599, 0x3c439041, 0xfe5163ab, +}; + +struct reduced { + nir_def *k; /* int32, only the low two bits matter */ + nir_def *fq; /* about [-1/2, 1/2] in quadrants, as a pair */ + nir_def *fq_lo; +}; + +static struct reduced +reduce_cody_waite(nir_builder *b, nir_def *x) +{ + nir_def *tp = nir_imm_float(b, TWO_OVER_PI); + nir_def *c2 = nir_imm_float(b, PIO2_2); + nir_def *kf = nir_fround_even(b, nir_fmul(b, x, tp)); + nir_def *nk = nir_fneg(b, kf); + + /* r1 = x - k C1 is exact. k C2 is a TwoProduct and r1 - k C2 a TwoSum so + * nothing is lost even when r1 and k C2 nearly cancel. + */ + nir_def *r1 = nir_ffma(b, nk, nir_imm_float(b, PIO2_1), x); + nir_def *p = nir_fmul(b, kf, c2); + nir_def *p_lo = nir_ffma(b, kf, c2, nir_fneg(b, p)); + nir_def *r2 = nir_fsub(b, r1, p); + nir_def *bb = nir_fsub(b, r2, r1); + nir_def *t = nir_fadd(b, nir_fsub(b, r1, nir_fsub(b, r2, bb)), + nir_fsub(b, nir_fneg(b, p), bb)); + nir_def *r_lo = + nir_ffma(b, nk, nir_imm_float(b, PIO2_3), nir_fsub(b, t, p_lo)); + + nir_def *fq = nir_fmul(b, r2, tp); + nir_def *fq_lo = nir_ffma(b, r2, tp, nir_fneg(b, fq)); + fq_lo = nir_ffma(b, r_lo, tp, fq_lo); + fq_lo = nir_ffma(b, r2, nir_imm_float(b, TWO_OVER_PI_LO), fq_lo); + nir_def *sum = nir_fadd(b, fq, fq_lo); + + return (struct reduced){ + .k = nir_f2i32(b, kf), + .fq = sum, + .fq_lo = fast2sum_lo(b, fq, fq_lo, sum), + }; +} + +static struct reduced +reduce_payne_hanek(nir_builder *b, nir_def *ax) +{ + nir_def *e8 = nir_ushr_imm(b, ax, 23); + nir_def *m32 = + nir_ishl_imm(b, nir_ior_imm(b, nir_iand_imm(b, ax, 0x7fffff), 0x800000), + 8); + + /* Word/bit offset of the 2/pi window for this exponent */ + nir_def *off = nir_iadd_imm(b, e8, -126); + nir_def *idx = nir_ushr_imm(b, off, 5); + nir_def *sh = nir_iand_imm(b, off, 31); + nir_def *sh_r = nir_isub(b, nir_imm_int(b, 31), sh); + + nir_def *tab[ARRAY_SIZE(two_over_pi_bits)]; + for (unsigned i = 0; i < ARRAY_SIZE(tab); ++i) + tab[i] = nir_imm_int(b, two_over_pi_bits[i]); + + /* Windows n = 1..4 at weights 2^0, 2^-32, 2^-64, 2^-96 */ + nir_def *lo[5], *hi[5]; + for (unsigned n = 1; n <= 4; ++n) { + nir_def *cand[5], *cand2[5]; + for (unsigned j = 0; j < 5; ++j) { + cand[j] = tab[j + n]; + cand2[j] = tab[j + n + 1]; + } + nir_def *w0 = nir_select_from_ssa_def_array(b, cand, 5, idx); + nir_def *w1 = nir_select_from_ssa_def_array(b, cand2, 5, idx); + nir_def *win = nir_ior(b, nir_ishl(b, w0, sh), + nir_ushr_imm(b, nir_ushr(b, w1, sh_r), 1)); + lo[n] = nir_imul(b, m32, win); + hi[n] = nir_umul_high(b, m32, win); + } + + /* Carry-propagate the 96 fraction bits, keep the top 64 */ + nir_def *f2 = nir_iadd(b, lo[3], hi[4]); + nir_def *c2 = nir_b2i32(b, nir_ult(b, f2, hi[4])); + nir_def *f1a = nir_iadd(b, lo[2], hi[3]); + nir_def *c1a = nir_b2i32(b, nir_ult(b, f1a, hi[3])); + nir_def *f1 = nir_iadd(b, f1a, c2); + nir_def *c1b = nir_b2i32(b, nir_ult(b, f1, f1a)); + nir_def *k = nir_iadd(b, nir_iadd(b, lo[1], hi[2]), nir_iadd(b, c1a, c1b)); + + /* Round to nearest quadrant so fq lands in [-1/2, 1/2] */ + nir_def *half = nir_ushr_imm(b, f1, 31); + k = nir_iadd(b, k, half); + nir_def *frac = nir_pack_64_2x32_split(b, f2, f1); + nir_def *is_half = nir_ine_imm(b, half, 0); + frac = nir_bcsel(b, is_half, nir_ineg(b, frac), frac); + nir_def *mag = nir_fmul_imm(b, nir_u2f32(b, frac), 0x1p-64); + + return (struct reduced){ + .k = k, + .fq = nir_bcsel(b, is_half, nir_fneg(b, mag), mag), + .fq_lo = nir_imm_float(b, 0.0), + }; +} + +static bool +sincos_filter(const nir_instr *instr, UNUSED const void *_) +{ + if (instr->type != nir_instr_type_alu) + return false; + + nir_alu_instr *alu = nir_instr_as_alu(instr); + return (alu->op == nir_op_fsin || alu->op == nir_op_fcos) && + alu->def.bit_size <= 32; +} + +static nir_def * +lower_sincos(nir_builder *b, nir_instr *instr, UNUSED void *_) +{ + nir_alu_instr *alu = nir_instr_as_alu(instr); + b->fp_math_ctrl = nir_fp_no_fast_math; + nir_def *x = nir_mov_alu(b, alu->src[0], 1); + unsigned bit_size = x->bit_size; + if (bit_size == 16) + x = nir_f2f32(b, x); + + nir_def *ax = nir_fabs(b, x); + + nir_push_if(b, nir_fge_imm(b, ax, 0x1p22)); + struct reduced ph = reduce_payne_hanek(b, ax); + + /* Payne-Hanek reduced |x|; sin is odd, so flip k and fq for negatives */ + nir_def *neg_x = nir_flt_imm(b, x, 0.0); + ph.k = nir_bcsel(b, neg_x, nir_ineg(b, ph.k), ph.k); + ph.fq = nir_bcsel(b, neg_x, nir_fneg(b, ph.fq), ph.fq); + + nir_push_else(b, NULL); + struct reduced cw = reduce_cody_waite(b, x); + nir_pop_if(b, NULL); + + nir_def *k = nir_if_phi(b, ph.k, cw.k); + nir_def *fq = nir_if_phi(b, ph.fq, cw.fq); + nir_def *fq_lo = nir_if_phi(b, ph.fq_lo, cw.fq_lo); + + /* cos(x) = sin(x + pi/2): one more quadrant */ + if (alu->op == nir_op_fcos) + k = nir_iadd_imm(b, k, 1); + + /* Taylor series in quadrants, |fq| < 0.71 (k can be off by one), + * truncation below 2^-36. The leading sin term and the quadratic cos term + * are kept as pairs and the low part of fq feeds the derivative, so both + * land within an ulp. + */ + nir_def *t = fq; + nir_def *t_lo = fq_lo; + nir_def *w = nir_fmul(b, t, t); + nir_def *w_lo = nir_ffma(b, t, t, nir_fneg(b, w)); + + static const double sin_coef[] = { + -0.6459640860557556, 0.07969262450933456, -0.004681753925979137, + 0.00016044118092395365, -3.598843250074424e-06, 5.6921727775716136e-08, + }; + static const double cos_coef[] = { + -1.2337005138397217, 0.25366950035095215, -0.020863480865955353, + 0.0009192602592520416, -2.520204179745633e-05, 4.710874748070637e-07, + -6.386603246255618e-09, + }; + static const double cos_coef0_lo = -3.629644851343983e-08; + + nir_def *pio2 = nir_imm_float(b, PIO2_1); + nir_def *p = nir_imm_float(b, sin_coef[5]); + for (int i = 4; i >= 0; --i) + p = nir_ffma(b, p, w, nir_imm_float(b, sin_coef[i])); + + nir_def *hi = nir_fmul(b, t, pio2); + nir_def *lo = nir_ffma(b, t, pio2, nir_fneg(b, hi)); + lo = nir_ffma(b, t, nir_imm_float(b, PIO2_2), lo); + lo = nir_ffma(b, t_lo, pio2, lo); + nir_def *sn = nir_fadd(b, hi, nir_ffma(b, nir_fmul(b, t, w), p, lo)); + + nir_def *d2 = nir_imm_float(b, cos_coef[0]); + nir_def *r = nir_imm_float(b, cos_coef[6]); + for (int i = 5; i >= 1; --i) + r = nir_ffma(b, r, w, nir_imm_float(b, cos_coef[i])); + + nir_def *h = nir_fmul(b, w, d2); + nir_def *l = nir_fadd(b, nir_fadd(b, nir_ffma(b, w, d2, nir_fneg(b, h)), + nir_fmul_imm(b, w, cos_coef0_lo)), + nir_fmul(b, w_lo, d2)); + l = nir_ffma(b, nir_fmul_imm(b, t_lo, 2.0 * cos_coef[0]), t, l); + nir_def *one = nir_imm_float(b, 1.0); + nir_def *s1 = nir_fadd(b, one, h); + nir_def *e = fast2sum_lo(b, one, h, s1); + nir_def *cs = nir_fadd( + b, s1, nir_fadd(b, nir_fadd(b, e, l), nir_fmul(b, nir_fmul(b, w, w), r))); + + /* Odd quadrants take the cosine, quadrants 2 and 3 negate */ + nir_def *odd = nir_ine_imm(b, nir_iand_imm(b, k, 1), 0); + nir_def *neg = nir_ine_imm(b, nir_iand_imm(b, k, 2), 0); + nir_def *v = nir_bcsel(b, odd, cs, sn); + v = nir_bcsel(b, neg, nir_fneg(b, v), v); + + /* Inf and NaN fall into the Payne-Hanek branch with a bogus exponent */ + nir_def *nan = nir_imm_float(b, NAN); + v = nir_bcsel(b, nir_flt_imm(b, ax, INFINITY), v, nan); + + /* sin(+-0) = +-0 exactly (the reduction returns +0 for -0). 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..1fe11d9848a2 --- /dev/null +++ b/src/asahi/compiler/agx_nir_lower_simdmat.c @@ -0,0 +1,609 @@ +/* + * 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. bf16 A/B (fp32 C) have no + * hardware form: fragments stay 16-bit bf16 patterns and widen exactly to + * fp32 at the MulAdd, which then runs the fp32 A/B form. + * + * 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; +} + +/* Base-type conversion of cooperative-matrix elements. bf16 values are 16-bit + * patterns with no NIR float type; SPV_KHR_bfloat16 converts through fp32. */ +nir_def * +agx_cmat_convert(nir_builder *b, nir_def *src, enum glsl_base_type from, + enum glsl_base_type to) +{ + if (from == to) + return src; + if (from == GLSL_TYPE_BFLOAT16) + return agx_cmat_convert(b, nir_bf2f(b, src), GLSL_TYPE_FLOAT, to); + if (to == GLSL_TYPE_BFLOAT16) + return nir_f2bf(b, agx_cmat_convert(b, src, from, GLSL_TYPE_FLOAT)); + + nir_op op = nir_type_conversion_op(nir_get_nir_type_for_glsl_base_type(from), + nir_get_nir_type_for_glsl_base_type(to), + nir_rounding_mode_undef); + return nir_build_alu1(b, op, src); +} + +/* Element-wise op on elements of type t; y is NULL for a unary op. bf16 + * elements compute in fp32 and round back to bf16. */ +nir_def * +agx_cmat_alu(nir_builder *b, nir_op op, enum glsl_base_type t, nir_def *x, + nir_def *y) +{ + bool bf16 = t == GLSL_TYPE_BFLOAT16; + if (bf16) { + x = nir_bf2f(b, x); + y = y ? nir_bf2f(b, y) : NULL; + } + nir_def *r = y ? nir_build_alu2(b, op, x, y) : nir_build_alu1(b, op, x); + return bf16 ? nir_f2bf(b, r) : r; +} + +/* bf16 <-> fp32 conversions in integer ALU. bf2f is exact; f2bf rounds to + * nearest even (Vulkan allows any rounding; RNE matches what bf16 storage + * kernels already compute), and a NaN keeps its sign and top payload bits + * with the quiet bit set. */ +static bool +lower_bf16_conversion(nir_builder *b, nir_alu_instr *alu, UNUSED void *data) +{ + if (alu->op != nir_op_bf2f && alu->op != nir_op_f2bf) + return false; + + b->cursor = nir_before_instr(&alu->instr); + nir_def *x = nir_ssa_for_alu_src(b, alu, 0); + nir_def *r; + if (alu->op == nir_op_bf2f) { + r = nir_ishl_imm(b, nir_u2u32(b, x), 16); + } else { + nir_def *hi = nir_ushr_imm(b, x, 16); + nir_def *rne = nir_ushr_imm( + b, nir_iadd(b, nir_iadd_imm(b, x, 0x7fff), nir_iand_imm(b, hi, 1)), 16); + nir_def *nan = nir_ugt_imm(b, nir_iand_imm(b, x, 0x7fffffff), 0x7f800000); + r = nir_u2u16(b, nir_bcsel(b, nan, nir_ior_imm(b, hi, 0x40), rne)); + } + nir_def_replace(&alu->def, r); + return true; +} + +bool +agx_nir_lower_bf16_conversions(nir_shader *shader) +{ + return nir_shader_alu_pass(shader, lower_bf16_conversion, + nir_metadata_control_flow, NULL); +} + +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. Checked with llama.cpp: 947 of 947 MUL_MAT + * cases pass, perplexity 1.0162, prompt processing 177.1 -> 179.4 + * tokens/s. */ + nir_deref_instr *sdref = + nir_build_deref_cast(b, &deref->def, deref->modes, + glsl_scalar_type(desc.element_type), tsz); + + /* VEC2 fragment loads (on by default, AGX_HWMAT_VEC2_OFF=1 disables): load + * the contiguous (col,col+1) row-major pair as ONE i16,xy load instead of + * two i16,x scalars. The pairing gives 24 distinct fragment regs (vs ~12 + * reused) so the post-RA scheduler can hoist all loads. 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. M1 prefill +9.4%, bit-exact. + */ + const bool vec2ld = getenv("AGX_HWMAT_VEC2_OFF") == 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); + 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 */ + /* Every advertised shape has one A/B type; B is widened like A below. */ + assert(da.element_type == src_desc(intr->src[2]).element_type); + 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]); + + if (da.element_type == GLSL_TYPE_BFLOAT16) { + a = nir_bf2f(b, a); + bb = nir_bf2f(b, bb); + in_bits = 32; + } + + /* 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], + agx_cmat_alu(b, nir_intrinsic_alu_op(intr), + src_desc(intr->src[1]).element_type, 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], + agx_cmat_alu(b, nir_intrinsic_alu_op(intr), + src_desc(intr->src[1]).element_type, a, NULL)); + 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], + agx_cmat_alu(b, nir_intrinsic_alu_op(intr), + src_desc(intr->src[1]).element_type, 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]); + unsigned save = b->fp_math_ctrl; + b->fp_math_ctrl = nir_intrinsic_fp_math_ctrl(intr); + nir_def *ret = agx_cmat_convert(b, src, sd.element_type, dd.element_type); + 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_pressure_schedule.c b/src/asahi/compiler/agx_pressure_schedule.c index 57849fd6e831..57bfb1794198 100644 --- a/src/asahi/compiler/agx_pressure_schedule.c +++ b/src/asahi/compiler/agx_pressure_schedule.c @@ -269,3 +269,106 @@ agx_pressure_schedule(agx_context *ctx) ralloc_free(memctx); } + +/* + * Latency hoist (compute): move each non-coherent device_load + * up past ALU, threadgroup barriers and workgroup-memory accesses, to just + * after the last instruction that produces its sources, but never more than + * HOIST_WINDOW instructions. The load then overlaps that work instead of + * being issued right before its first use. Stops at anything that could + * order device memory: device stores, atomics, memory barriers, coherent + * loads, image ops, control flow. + */ +#define HOIST_WINDOW 64 + +/* + * A uniform_store writes at most four 16-bit uniform registers starting at its + * index and nothing else. A load may move above it unless one of the load's + * uniform sources overlaps [index, index + 16). The window is wider than the + * store, which keeps the check conservative. An index that is not an + * immediate is never crossed. + */ +static bool +hoist_uniform_store_conflict(const agx_instr *store, const agx_instr *load) +{ + if (store->src[1].type != AGX_INDEX_IMMEDIATE) + return true; + + unsigned lo = store->src[1].value, hi = lo + 16; + + agx_foreach_src(load, s) { + agx_index src = load->src[s]; + if (src.type != AGX_INDEX_UNIFORM) + continue; + + unsigned src_lo = src.value, src_hi = src_lo + agx_index_size_16(src); + if (src_lo < hi && lo < src_hi) + return true; + } + + return false; +} + +static bool +hoist_can_cross(const agx_instr *P, const agx_instr *load) +{ + if (P->op == AGX_OPCODE_THREADGROUP_BARRIER || + P->op == AGX_OPCODE_LOCAL_LOAD || P->op == AGX_OPCODE_LOCAL_STORE) + return true; + + if (P->op == AGX_OPCODE_DEVICE_LOAD) + return !P->coherent; + + if (P->op == AGX_OPCODE_UNIFORM_STORE) + return !hoist_uniform_store_conflict(P, load); + + /* Anything else must be pure: free of side effects and reorderable. The + * default schedule class is none, so the class alone also admits + * doorbell, stack_map and similar. + */ + const struct agx_opcode_info *info = &agx_opcodes_info[P->op]; + return info->schedule_class == AGX_SCHEDULE_CLASS_NONE && + info->can_reorder && info->can_eliminate; +} + +void +agx_hoist_loads(agx_context *ctx) +{ + agx_foreach_block(ctx, block) { + agx_foreach_instr_in_block_safe(block, I) { + if (instr_after_logical_end(I)) + break; + + if (I->op != AGX_OPCODE_DEVICE_LOAD || I->coherent) + continue; + + agx_instr *target = NULL; + unsigned n = 0; + + for (agx_instr *P = list_entry(I->link.prev, agx_instr, link); + &P->link != &block->instructions && n < HOIST_WINDOW; + P = list_entry(P->link.prev, agx_instr, link), ++n) { + + if (!hoist_can_cross(P, I)) + break; + + bool dep = false; + agx_foreach_ssa_dest(P, d) { + agx_foreach_ssa_src(I, s) { + if (I->src[s].value == P->dest[d].value) + dep = true; + } + } + if (dep) + break; + + target = P; + } + + if (target) { + agx_remove_instruction(I); + list_addtail(&I->link, &target->link); + } + } + } +} 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..1b163a000167 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,10 @@ if with_tests 'agx_tests', files( 'test/test-compact-constants.cpp', + 'test/test-hoist-loads.cpp', + 'test/test-lower-address.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-hoist-loads.cpp b/src/asahi/compiler/test/test-hoist-loads.cpp new file mode 100644 index 000000000000..878d1cea699e --- /dev/null +++ b/src/asahi/compiler/test/test-hoist-loads.cpp @@ -0,0 +1,183 @@ +/* + * SPDX-License-Identifier: MIT + */ + +#include "agx_builder.h" +#include "agx_compiler.h" +#include "agx_test.h" + +#include + +static void +pass(agx_context *ctx) +{ + agx_hoist_loads(ctx); + agx_reindex_ssa(ctx); +} + +#define CASE(instr, expected) INSTRUCTION_CASE(instr, expected, pass) +#define NEGCASE(instr) INSTRUCTION_CASE(instr, instr, pass) + +class HoistLoads : public testing::Test { + protected: + HoistLoads() + { + mem_ctx = ralloc_context(NULL); + } + + ~HoistLoads() + { + ralloc_free(mem_ctx); + } + + void *mem_ctx; +}; + +static agx_index +load(agx_builder *b, agx_index off, bool coherent) +{ + agx_index v = agx_temp(b->shader, AGX_SIZE_32); + agx_device_load_to(b, v, agx_uniform(0, AGX_SIZE_64), off, AGX_FORMAT_I32, 1, + 0, coherent); + return v; +} + +TEST_F(HoistLoads, AboveAluAndThreadgroupBarrier) +{ + CASE( + { + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index y = agx_fadd(b, x, x); + agx_threadgroup_barrier(b); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + agx_unit_test(b, y); + }, + { + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index v = load(b, off, false); + agx_index y = agx_fadd(b, x, x); + agx_threadgroup_barrier(b); + agx_unit_test(b, v); + agx_unit_test(b, y); + }); +} + +TEST_F(HoistLoads, StopsAtItsOwnSource) +{ + NEGCASE({ + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_index y = agx_fadd(b, x, x); + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + agx_unit_test(b, y); + }); +} + +TEST_F(HoistLoads, NotAboveDeviceStore) +{ + NEGCASE({ + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_device_store(b, x, agx_uniform(0, AGX_SIZE_64), off, AGX_FORMAT_I32, + 1, 0, false); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, NotAboveMemoryBarrier) +{ + NEGCASE({ + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_memory_barrier(b); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, NotAboveDoorbell) +{ + NEGCASE({ + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_doorbell(b, 1); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, NotAboveUniformStoreToItsBase) +{ + NEGCASE({ + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_uniform_store(b, x, agx_immediate(0), 1); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, AboveUniformStoreToOtherRegisters) +{ + CASE( + { + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_uniform_store(b, x, agx_immediate(32), 1); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + }, + { + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index v = load(b, off, false); + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_uniform_store(b, x, agx_immediate(32), 1); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, AboveUniformStoreAdjacentRegisters) +{ + CASE( + { + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_uniform_store(b, x, agx_immediate(4), 1); + agx_index v = load(b, off, false); + agx_unit_test(b, v); + }, + { + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index v = load(b, off, false); + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_uniform_store(b, x, agx_immediate(4), 1); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, NotAboveUniformStorePartialOverlap) +{ + NEGCASE({ + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_uniform_store(b, x, agx_immediate(3), 1); + agx_index v = agx_temp(b->shader, AGX_SIZE_32); + agx_device_load_to(b, v, agx_uniform(0, AGX_SIZE_64), + agx_immediate(4), AGX_FORMAT_I32, 1, 0, false); + agx_unit_test(b, v); + }); +} + +TEST_F(HoistLoads, CoherentLoadStays) +{ + NEGCASE({ + agx_index x = agx_mov_imm(b, AGX_SIZE_32, 0xcafe); + agx_index off = agx_mov_imm(b, AGX_SIZE_32, 4); + agx_index y = agx_fadd(b, x, x); + agx_index v = load(b, off, true); + agx_unit_test(b, v); + agx_unit_test(b, y); + }); +} diff --git a/src/asahi/compiler/test/test-lower-address.cpp b/src/asahi/compiler/test/test-lower-address.cpp new file mode 100644 index 000000000000..b78d4960e964 --- /dev/null +++ b/src/asahi/compiler/test/test-lower-address.cpp @@ -0,0 +1,72 @@ +/* + * Copyright 2026 Joshua Warren + * SPDX-License-Identifier: MIT + */ + +#include "nir_builder.h" + +extern "C" { +#include "agx_nir.h" +} + +#include + +/* + * A load whose address is a lea with a shift below the format's shift (the + * byte offsets the cooperative-matrix loads produce) keeps the base and, + * with fold_subformat_address, indexes by the offset scaled down to elements. + * Returns the lowered load_agx's constant index, or INT64_MIN if the load is + * not indexed by a constant. + */ +static int64_t +lowered_index(bool fold, nir_op lea, int32_t byte_offset, unsigned bit_size, + bool *same_base = nullptr) +{ + 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, "address test"); + nir_def *base = nir_load_push_constant(&b, 1, 64, nir_imm_int(&b, 0)); + nir_def *addr = nir_build_alu3(&b, lea, base, nir_imm_int(&b, byte_offset), + nir_imm_int(&b, 0)); + nir_def *v = nir_load_global(&b, 2, bit_size, addr, .align_mul = 16); + nir_store_global(&b, v, nir_imm_int64(&b, 0x200000), .align_mul = 16); + + agx_nir_lower_address(b.shader, fold); + nir_opt_constant_folding(b.shader); + nir_validate_shader(b.shader, "after agx_nir_lower_address"); + + int64_t index = INT64_MIN; + nir_foreach_block(block, b.impl) { + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_intrinsic) + continue; + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + if (intr->intrinsic != nir_intrinsic_load_agx) + continue; + if (same_base) + *same_base = intr->src[0].ssa == base; + if (nir_src_is_const(intr->src[1])) + index = nir_src_as_int(intr->src[1]); + } + } + ralloc_free(b.shader); + glsl_type_singleton_decref(); + return index; +} + +TEST(LowerAddress, FoldScalesByteOffsetToElements) +{ + bool same_base = false; + EXPECT_EQ(lowered_index(true, nir_op_ulea_agx, 24, 32, &same_base), 6); + EXPECT_TRUE(same_base); + EXPECT_EQ(lowered_index(true, nir_op_ulea_agx, 24, 16), 12); + EXPECT_EQ(lowered_index(true, nir_op_ilea_agx, -8, 32), -2); +} + +TEST(LowerAddress, NoFoldMaterializesTheSum) +{ + bool same_base = true; + EXPECT_EQ(lowered_index(false, nir_op_ulea_agx, 24, 32, &same_base), 0); + EXPECT_FALSE(same_base); +} 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..a80c57d96abc --- /dev/null +++ b/src/asahi/compiler/test/test-lower-math.cpp @@ -0,0 +1,279 @@ +/* + * Copyright 2026 Joshua Warren + * SPDX-License-Identifier: MIT + */ + +#include "nir_builder.h" + +extern "C" { +#include "agx_nir.h" +} + +#include +#include +#include +#include + +/* + * Behavioral pins for the float-control keying in agx_nir_lower_math.c. + * + * Division and reciprocal do not key on float controls. Their bits are pinned + * on hardware by tests/precise_math.py, ulp_bound.py and div_window.py; here + * the pin is that DenormFlushToZero emits exactly the default lowering. + * + * 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; + } + + /* ALU op histogram of fdiv (or frcp) lowered under one float-control mode */ + std::map lowered_division(unsigned mode, bool reciprocal) + { + nir_builder s = nir_builder_init_simple_shader(MESA_SHADER_COMPUTE, + &options, "lower fdiv"); + s.shader->info.float_controls_execution_mode = mode; + nir_def *x = nir_undef(&s, 1, 32); + if (reciprocal) + nir_frcp(&s, x); + else + nir_fdiv(&s, nir_undef(&s, 1, 32), x); + + EXPECT_TRUE(agx_nir_lower_fdiv(s.shader)); + + std::map ops; + nir_foreach_function_impl(impl, s.shader) { + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type == nir_instr_type_alu) + ops[nir_instr_as_alu(instr)->op]++; + } + } + } + ralloc_free(s.shader); + return ops; + } +}; + +const nir_shader_compiler_options LowerMathTest::options = {}; + +TEST_F(LowerMathTest, DivFtzModeUnchanged) +{ + auto ftz = lowered_division(FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP32, false); + EXPECT_EQ(ftz.count(nir_op_fdiv), 0u); + EXPECT_EQ(ftz, lowered_division(0, false)); +} + +TEST_F(LowerMathTest, FrcpFtzModeUnchanged) +{ + auto ftz = lowered_division(FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP32, true); + EXPECT_EQ(ftz.count(nir_op_frcp), 1u); + EXPECT_EQ(ftz, lowered_division(0, true)); +} + +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); +} + +/* The shared-memory offset fold must keep the sign of a constant addend. A + * 16-bit add of 0xfff0 is x - 16: the group {x, x - 16} must come out as + * constant offsets {0, -16} from x, never {0, +65520}, which the hardware + * index add (it does not wrap) would turn into a different address. */ +TEST_F(LowerMathTest, SharedFoldKeepsNegativeAddends) +{ + nir_def *x = nir_undef(&b, 1, 16); + nir_def *lo = nir_iadd_imm(&b, x, 0xfff0); + nir_load_shared(&b, 1, 32, x, .base = 0, .align_mul = 4); + nir_load_shared(&b, 1, 32, lo, .base = 0, .align_mul = 4); + + EXPECT_TRUE(agx_nir_fold_shared_offsets(b.shader)); + + std::vector totals; + nir_foreach_function_impl(impl, b.shader) { + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type != nir_instr_type_intrinsic) + continue; + nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr); + if (intr->intrinsic != nir_intrinsic_load_shared) + continue; + + int64_t total = nir_intrinsic_base(intr); + nir_scalar a = nir_scalar_resolved(intr->src[0].ssa, 0); + if (nir_scalar_is_alu(a) && nir_scalar_alu_op(a) == nir_op_iadd) { + nir_scalar k = nir_scalar_chase_alu_src(a, 1); + if (nir_scalar_is_const(k)) + total += nir_scalar_as_int(k); + } + totals.push_back(total); + } + } + } + + ASSERT_EQ(totals.size(), 2u); + std::sort(totals.begin(), totals.end()); + EXPECT_EQ(totals[0], -16); + EXPECT_EQ(totals[1], 0); +} + +/* The sub-format address fold indexes by the byte offset scaled down, which + * equals the byte offset only when the offset is a multiple of the element + * size. An offset that is not provably aligned must keep the 64-bit sum (no + * shift of the offset); a provably aligned one is folded. */ +static void +subformat_shader(nir_builder *b, nir_def *off) +{ + nir_def *base = nir_load_push_constant(b, 1, 64, nir_imm_int(b, 0)); + nir_def *addr = nir_ulea_agx(b, base, off, nir_imm_int(b, 0)); + nir_def *v = nir_load_global(b, 1, 32, addr, .align_mul = 4); + nir_store_global(b, v, nir_imm_int64(b, 0x200000), .align_mul = 4); +} + +TEST_F(LowerMathTest, SubformatFoldNeedsProvablyAlignedOffset) +{ + /* an offset loaded from memory is not provably a multiple of 4 */ + subformat_shader(&b, nir_load_push_constant(&b, 1, 32, nir_imm_int(&b, 8))); + + agx_nir_lower_address(b.shader, true); + EXPECT_EQ(count_alu(nir_op_ushr), 0u); +} + +TEST_F(LowerMathTest, SubformatFoldKeepsProvablyAlignedOffset) +{ + nir_def *off = nir_ishl_imm( + &b, nir_load_push_constant(&b, 1, 32, nir_imm_int(&b, 8)), 2); + subformat_shader(&b, off); + + agx_nir_lower_address(b.shader, true); + EXPECT_EQ(count_alu(nir_op_ushr), 1u); +} 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/lib/agx_device.c b/src/asahi/lib/agx_device.c index f158e0e7acea..248464d17e98 100644 --- a/src/asahi/lib/agx_device.c +++ b/src/asahi/lib/agx_device.c @@ -619,7 +619,7 @@ agx_open_device(void *memctx, struct agx_device *dev) .kernel_end = dev->params.vm_end, }; - uint64_t user_size = vm_create.kernel_start - user_start; + uint64_t user_size = 0; int ret = asahi_simple_ioctl(dev, DRM_IOCTL_ASAHI_VM_CREATE, &vm_create); if (ret) { @@ -628,23 +628,13 @@ agx_open_device(void *memctx, struct agx_device *dev) return false; } - /* Round the user VA window to powers-of-two... */ - user_start = util_next_power_of_two64(user_start); - user_size = util_next_power_of_two64(user_size + 1) >> 1; - - /* ...so when we cut user size in half to emulate sparse buffers... */ - user_size /= 2; - - /* ...or maybe in quarters if necessary to disambiguate */ - if (user_size == user_start) { - user_size /= 2; + if (!agx_sparse_layout(user_start, vm_create.kernel_start, + &dev->sparse_ro_offset, &user_size)) { + fprintf(stderr, "No user address window for the heap\n"); + assert(0); + return false; } - /* ...we can distinguish the top/bottom half by an address bit */ - dev->sparse_ro_offset = user_size; - assert((user_start & dev->sparse_ro_offset) == 0); - assert(((user_start + (user_size - 1)) & dev->sparse_ro_offset) == 0); - simple_mtx_init(&dev->vma_lock, mtx_plain); util_vma_heap_init(&dev->main_heap, user_start, user_size); util_vma_heap_init(&dev->usc_heap, dev->shader_base, shader_size); @@ -956,5 +946,9 @@ agx_gather_device_key(struct agx_device *dev) return (struct agx_device_key){ .needs_g13x_coherency = u_tristate_make(g13x_coh), .soft_fault = agx_has_soft_fault(dev), + .g13 = dev->params.gpu_generation == 13, + .fold_subformat_address = dev->params.gpu_generation == 13 && + dev->params.num_clusters_total == 1 && + dev->params.num_dies == 1, }; } diff --git a/src/asahi/lib/agx_device.h b/src/asahi/lib/agx_device.h index e80dd9b5a5e1..3edc7e93a8f0 100644 --- a/src/asahi/lib/agx_device.h +++ b/src/asahi/lib/agx_device.h @@ -15,6 +15,7 @@ #include "util/u_printf.h" #include "util/vma.h" #include "agx_bo.h" +#include "agx_va_layout.h" #include "agx_pack.h" #include "decode.h" #include "layout.h" diff --git a/src/asahi/lib/agx_va.c b/src/asahi/lib/agx_va.c index a2b56fc81465..d7467ea00cd6 100644 --- a/src/asahi/lib/agx_va.c +++ b/src/asahi/lib/agx_va.c @@ -5,6 +5,7 @@ #include "agx_bo.h" #include "agx_device.h" +#include "agx_va_layout.h" static struct util_vma_heap * agx_vma_heap(struct agx_device *dev, enum agx_va_flags flags) @@ -12,6 +13,34 @@ agx_vma_heap(struct agx_device *dev, enum agx_va_flags flags) return (flags & AGX_VA_USC) ? &dev->usc_heap : &dev->main_heap; } +/* + * Sparse buffers map every page twice, read-write at X and read-only at + * X + ro_offset, so the heap [user_start, user_start + size) is split by one + * address bit: ro_offset is a power of two, every read-write address lies + * below it (bit clear) and the shadow of the last one still fits below + * kernel_start. Choose the ro_offset that leaves the largest heap. On a + * 512 GiB window with the heap at 72 GiB that is 256 GiB and a 152 GiB heap; + * rounding the window and halving it twice left 64 GiB. + */ +bool +agx_sparse_layout(uint64_t user_start, uint64_t kernel_start, + uint64_t *ro_offset, uint64_t *size) +{ + *size = 0; + + for (uint64_t p = util_next_power_of_two64(user_start + 1); + p && user_start + p < kernel_start; p <<= 1) { + uint64_t s = MIN2(p - user_start, kernel_start - user_start - p); + + if (s > *size) { + *size = s; + *ro_offset = p; + } + } + + return *size != 0; +} + struct agx_va * agx_va_alloc(struct agx_device *dev, uint64_t size_B, uint64_t align_B, enum agx_va_flags flags, uint64_t fixed_va) diff --git a/src/asahi/lib/agx_va_layout.h b/src/asahi/lib/agx_va_layout.h new file mode 100644 index 000000000000..8b3c4930fdad --- /dev/null +++ b/src/asahi/lib/agx_va_layout.h @@ -0,0 +1,19 @@ +/* + * SPDX-License-Identifier: MIT + */ + +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +bool agx_sparse_layout(uint64_t user_start, uint64_t kernel_start, + uint64_t *ro_offset, uint64_t *size); + +#ifdef __cplusplus +} /* extern "C" */ +#endif diff --git a/src/asahi/lib/meson.build b/src/asahi/lib/meson.build index e50a70eff518..7a7455e8c5b7 100644 --- a/src/asahi/lib/meson.build +++ b/src/asahi/lib/meson.build @@ -76,9 +76,11 @@ if with_tests files( 'tests/test-packing.cpp', 'tests/test-tilebuffer.cpp', + 'tests/test-va-layout.cpp', ), c_args : [c_msvc_compat_args, no_override_init_args], gnu_symbol_visibility : 'hidden', + include_directories : [inc_asahi], dependencies: [idep_gtest, idep_agx_pack, idep_mesautil], link_with : [libasahi_lib, libasahi_layout], ), diff --git a/src/asahi/lib/tests/test-va-layout.cpp b/src/asahi/lib/tests/test-va-layout.cpp new file mode 100644 index 000000000000..2f3f8dd82bc9 --- /dev/null +++ b/src/asahi/lib/tests/test-va-layout.cpp @@ -0,0 +1,50 @@ +/* + * SPDX-License-Identifier: MIT + */ + +#include "agx_va_layout.h" + +#include + +static constexpr uint64_t GiB = 1ull << 30; + +/* The read-write heap, and its read-only shadow, must be disjoint, below the + * kernel window, and told apart by the single bit ro_offset. + */ +static void +check(uint64_t user_start, uint64_t kernel_start, uint64_t want_size) +{ + uint64_t ro = 0, size = 0; + + ASSERT_TRUE(agx_sparse_layout(user_start, kernel_start, &ro, &size)); + EXPECT_EQ(size, want_size); + + EXPECT_EQ(ro & (ro - 1), 0u); + EXPECT_EQ(user_start & ro, 0u); + EXPECT_EQ((user_start + size - 1) & ro, 0u); + EXPECT_LE(user_start + ro + size, kernel_start); +} + +TEST(SparseLayout, Window512GiBHeapAt72GiB) +{ + /* G13 and G14 as reported by the kernel: 408 GiB of user window; the old + * code left 64 GiB of it + */ + check(72 * GiB, 480 * GiB, 152 * GiB); +} + +TEST(SparseLayout, Window512GiBHeapAt128GiB) +{ + check(128 * GiB, 480 * GiB, 96 * GiB); +} + +TEST(SparseLayout, SmallWindow) +{ + check(8 * GiB, 120 * GiB, 48 * GiB); +} + +TEST(SparseLayout, NoRoom) +{ + uint64_t ro = 0, size = 0; + EXPECT_FALSE(agx_sparse_layout(72 * GiB, 72 * GiB, &ro, &size)); +} diff --git a/src/asahi/libagx/libagx_dgc.h b/src/asahi/libagx/libagx_dgc.h index d8266b663111..a861642cb1c8 100644 --- a/src/asahi/libagx/libagx_dgc.h +++ b/src/asahi/libagx/libagx_dgc.h @@ -379,7 +379,16 @@ 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; } @@ -390,26 +399,94 @@ 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_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; +} + +/* + * Barrier at the end of a batch in the deferred CDM flush mode. On G13G it is + * bits 4-8: the set that the deferred flush was measured with on M1 + * (token digests identical over long streams and repeated runs). Other chips + * use agx_cdm_barrier. + */ +static inline GLOBAL uint32_t * +agx_cdm_barrier_deferred(GLOBAL uint32_t *out, enum agx_chip chip) +{ + if (chip != AGX_CHIP_G13G) + return agx_cdm_barrier(out, chip); + + agx_push(out, CDM_BARRIER, cfg) { + cfg.unk_4 = true; + cfg.unk_5 = true; + cfg.unk_6 = true; + cfg.unk_7 = true; + cfg.unk_8 = true; + } + + return out; +} + +/* + * G13X candidate reduction: the designed set {unk_4, unk_5, unk_6, unk_8} + * (the bits that kept the pinned token digests of an MLX Vulkan + * backend identical in all 48 runs of a test 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_buffer.c b/src/asahi/vulkan/hk_buffer.c index 8ec78e9a7e82..58864da13108 100644 --- a/src/asahi/vulkan/hk_buffer.c +++ b/src/asahi/vulkan/hk_buffer.c @@ -114,6 +114,9 @@ hk_CreateBuffer(VkDevice device, const VkBufferCreateInfo *pCreateInfo, if (!buffer) return vk_error(dev, VK_ERROR_OUT_OF_HOST_MEMORY); + if (buffer->vk.create_flags & VK_BUFFER_CREATE_SPARSE_BINDING_BIT) + p_atomic_set(&dev->sparse_used, 1); + if (buffer->vk.size > 0 && (buffer->vk.create_flags & (VK_BUFFER_CREATE_SPARSE_BINDING_BIT | diff --git a/src/asahi/vulkan/hk_cmd_buffer.c b/src/asahi/vulkan/hk_cmd_buffer.c index b527e97a886e..8b31a0a1e7ba 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,29 @@ 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. */ + if (!HK_PERF(dev, ALWAYSCDMBARRIER) && !dev->cdm_defer_flush) { + /* 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; + + /* Keep the compute batch open across the barrier only when it is the + * LAST control stream of the command buffer and no graphics batch is + * open. Streams are submitted in list order, each ordered after every + * stream listed before it. An older compute stream left open behind a + * later graphics stream would run before that graphics work, so the + * dispatch recorded after the barrier would read unrendered data. In + * that case end it, so the next dispatch starts a new tail stream. */ + struct hk_cs *open_cs = cmd->current_cs.cs; + bool is_tail = + !open_cs || + list_last_entry(&cmd->control_streams, struct hk_cs, node) == open_cs; + if (is_tail && !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 +805,24 @@ hk_dispatch_precomp(struct hk_cmd_buffer *cmd, struct agx_grid grid, agx_usc_words_precomp(t.cpu, &prog->b, uploaded_data, data_size); + if (barrier & (AGX_PREGFX | AGX_POSTGFX)) + hk_cmd_buffer_untrack_cdm(cmd); + /* 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..db05e3bdc00c 100644 --- a/src/asahi/vulkan/hk_cmd_buffer.h +++ b/src/asahi/vulkan/hk_cmd_buffer.h @@ -108,6 +108,10 @@ struct hk_root_descriptor_table { struct { uint64_t group_count_addr; uint32_t base_group[3]; + uint32_t _pad; + + /* Address of the bound shader's large-constant data, else 0 */ + uint64_t const_data; } cs; }; @@ -121,6 +125,12 @@ struct hk_root_descriptor_table { uint8_t set_dynamic_buffer_start[HK_MAX_SETS]; }; +/* cs shares a union with draw; growing it past draw would move every later + * root offset (push constants, dynamic buffers). */ +static_assert(sizeof(((struct hk_root_descriptor_table *)0)->cs) <= + sizeof(((struct hk_root_descriptor_table *)0)->draw), + "root cs fields must fit in the draw union member"); + /* helper macro for computing root descriptor byte offsets */ #define hk_root_descriptor_offset(member) \ offsetof(struct hk_root_descriptor_table, member) @@ -295,6 +305,21 @@ 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; + + /* The command buffer recorded graphics or meta work (a render pass, an + * image copy, or a precomp launch that feeds or follows a draw). From then + * on every CDM launch is flushed as before: the consumers of those + * launches are draws and copies the app declares no barrier for, so the + * dependency tracking has nothing to key on. + */ + bool untracked; }; struct hk_cmd_push { @@ -357,6 +382,11 @@ struct hk_cs { /* Whether there is more than just the root chunk */ bool stream_linked; + /* Deferred CDM flush (hk_device.cdm_defer_flush): a launch was recorded + * without the trailing CDM barrier. It is emitted when the batch ends + * (Vulkan barrier, query, event, end of command buffer). + */ + bool cdm_flush_pending; /* Whether the sampler heap is required. Although we always must maintain the * heap for correctness, it's often not necessary since we can push lots of @@ -682,7 +712,21 @@ hk_cs_destroy(struct hk_cs *cs) free(cs); } +/* Stop dependency-tracking the CDM barrier for the rest of this command + * buffer. Launches already recorded without a trailing flush are covered by + * the pending flag, which the next launch (or the end of the stream) consumes. + */ +static inline void +hk_cmd_buffer_untrack_cdm(struct hk_cmd_buffer *cmd) +{ + if (!cmd->state.cs.untracked) { + cmd->state.cs.untracked = true; + cmd->state.cs.cdm_barrier_pending = true; + } +} + void hk_dispatch_imm_writes(struct hk_cmd_buffer *cmd, struct hk_cs *cs); +void hk_cs_finish_pending_flush(struct hk_cmd_buffer *cmd, struct hk_cs *cs); static void hk_cmd_buffer_end_compute_internal(struct hk_cmd_buffer *cmd, @@ -690,6 +734,7 @@ hk_cmd_buffer_end_compute_internal(struct hk_cmd_buffer *cmd, { if (*ptr) { struct hk_cs *cs = *ptr; + hk_cs_finish_pending_flush(cmd, cs); /* This control stream may write immediates as it ends. Queue the writes * now that we're done emitting everything else. diff --git a/src/asahi/vulkan/hk_cmd_dispatch.c b/src/asahi/vulkan/hk_cmd_dispatch.c index 7cb08bb1d3c8..2d431ce85200 100644 --- a/src/asahi/vulkan/hk_cmd_dispatch.c +++ b/src/asahi/vulkan/hk_cmd_dispatch.c @@ -39,6 +39,10 @@ hk_cmd_bind_compute_shader(struct hk_cmd_buffer *cmd, struct hk_api_shader *shader) { cmd->state.cs.shader = shader; + + /* The compute root is uploaded at every dispatch; no dirty flag needed. */ + cmd->state.cs.descriptors.root.cs.const_data = + shader ? hk_only_variant(shader)->data_addr : 0; } void @@ -48,8 +52,50 @@ 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 if (unlikely(p_atomic_read(&dev->sparse_used))) { + /* Sparse resources change page mappings behind the launches. The + * reduced G13X word leaves sparse-image atomics wrong and the GPU + * faulting, so a device that created one keeps the full word + * (bits 0-19, the barrier before the reduction). + */ + uint32_t *word = cs->current; + *word = 0xfffff | (AGX_CDM_BLOCK_TYPE_BARRIER << 29); + cs->current = word + 1; + } else if (dev->cdm_barrier_default_mask) { + uint32_t *word = cs->current; + *word = dev->cdm_barrier_default_mask | (AGX_CDM_BLOCK_TYPE_BARRIER << 29); + cs->current = word + 1; + } else { + cs->current = agx_cdm_barrier(cs->current, dev->dev.chip); + } + + if (!HK_PERF(dev, NOCDMBARRIER)) + cs->stats.flushes++; +} + +void +hk_cs_finish_pending_flush(struct hk_cmd_buffer *cmd, struct hk_cs *cs) +{ + if (cs->type == HK_CS_CDM && cs->cdm_flush_pending) { + hk_ensure_cs_has_space(cmd, cs, 0x2000); + cs->current = agx_cdm_barrier_deferred(cs->current, + hk_cmd_buffer_device(cmd)->dev.chip); + cs->stats.flushes++; + cs->cdm_flush_pending = false; + } } void @@ -61,10 +107,63 @@ 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); + return; + } + + /* Graphics or meta work in this command buffer: flush after every launch + * (and before the first one if a tracked launch is still unflushed). + */ + if (cs->cmd->state.cs.untracked) { + if (cs->cmd->state.cs.cdm_barrier_pending || cs->cdm_flush_pending) { + hk_cdm_cache_flush(dev, cs); + cs->cmd->state.cs.cdm_barrier_pending = false; + cs->cdm_flush_pending = false; + } + cs->current = + agx_cdm_launch(cs->current, dev->dev.chip, grid, wg, launch, usc); + hk_cdm_cache_flush(dev, cs); + return; + } + + /* Deferred CDM flush: no barrier between the launches of one batch. The + * batch ends at every vkCmdPipelineBarrier2 (hk_cmd_buffer.c), and + * hk_cs_finish_pending_flush emits the barrier there. + */ + if (dev->cdm_defer_flush) { + cs->current = + agx_cdm_launch(cs->current, dev->dev.chip, grid, wg, launch, usc); + cs->cdm_flush_pending = true; + 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); + /* A meta operation (an image copy with several regions, for one) writes + * through the same caches in consecutive launches that the app declared no + * dependency between; each write must reach memory before the next launch, + * as it did when every launch flushed. + */ + if (cs->cmd->in_meta) + cs->cmd->state.cs.cdm_barrier_pending = true; } void diff --git a/src/asahi/vulkan/hk_cmd_draw.c b/src/asahi/vulkan/hk_cmd_draw.c index 5823c6575fd4..6e48108d9f82 100644 --- a/src/asahi/vulkan/hk_cmd_draw.c +++ b/src/asahi/vulkan/hk_cmd_draw.c @@ -541,6 +541,7 @@ hk_CmdBeginRendering(VkCommandBuffer commandBuffer, struct hk_rendering_state *render = &cmd->state.gfx.render; memset(render, 0, sizeof(*render)); + hk_cmd_buffer_untrack_cdm(cmd); render->flags = pRenderingInfo->flags; render->area = pRenderingInfo->renderArea; diff --git a/src/asahi/vulkan/hk_cmd_meta.c b/src/asahi/vulkan/hk_cmd_meta.c index abe4eb7d9c46..480bd1179b51 100644 --- a/src/asahi/vulkan/hk_cmd_meta.c +++ b/src/asahi/vulkan/hk_cmd_meta.c @@ -100,6 +100,7 @@ hk_meta_begin(struct hk_cmd_buffer *cmd, struct hk_meta_save *save, { struct hk_descriptor_state *desc = hk_get_descriptors_state(cmd, bind_point); + hk_cmd_buffer_untrack_cdm(cmd); if (bind_point == VK_PIPELINE_BIND_POINT_GRAPHICS) { save->dynamic = cmd->vk.dynamic_graphics_state; save->_dynamic_vi = cmd->state.gfx._dynamic_vi; @@ -1359,6 +1360,12 @@ hk_meta_resolve_rendering(struct hk_cmd_buffer *cmd, hk_meta_end(cmd, &save, VK_PIPELINE_BIND_POINT_GRAPHICS); } +/* Buffer copies, updates and fills launch with no internal barrier. The + * application orders them against other work with vkCmdPipelineBarrier, which + * flushes before the next CDM launch (cdm_barrier_pending). AGX_BARRIER_ALL + * would add a flush before and after each launch, ordering work that Vulkan + * leaves unordered. + */ static void hk_cmd_copy(struct hk_cmd_buffer *cmd, uint64_t dst, uint64_t src, size_t size) { @@ -1370,7 +1377,7 @@ hk_cmd_copy(struct hk_cmd_buffer *cmd, uint64_t dst, uint64_t src, size_t size) unsigned uint4s = size / 16; unsigned bytes = uint4s * 16; - libagx_copy_uint4(cmd, agx_1d(uint4s), AGX_BARRIER_ALL, dst, src); + libagx_copy_uint4(cmd, agx_1d(uint4s), AGX_BARRIER_NONE, dst, src); dst += bytes; src += bytes; @@ -1378,7 +1385,7 @@ hk_cmd_copy(struct hk_cmd_buffer *cmd, uint64_t dst, uint64_t src, size_t size) } if (size) { - libagx_copy_uchar(cmd, agx_1d(size), AGX_BARRIER_ALL, dst, src); + libagx_copy_uchar(cmd, agx_1d(size), AGX_BARRIER_NONE, dst, src); } } @@ -1494,10 +1501,10 @@ hk_CmdFillBuffer(VkCommandBuffer commandBuffer, VkBuffer dstBuffer, vk_meta_buffer_address(&dev->vk, dstBuffer, dstOffset, dstRange); if (util_is_aligned(addr, 16) && util_is_aligned(range, 16)) { - libagx_fill_uint4(cmd, agx_2d(range / 16, 1), AGX_BARRIER_ALL, + libagx_fill_uint4(cmd, agx_2d(range / 16, 1), AGX_BARRIER_NONE, addr, 0, data, data, data, data); } else { - libagx_fill(cmd, agx_1d(range / 4), AGX_BARRIER_ALL, addr, data); + libagx_fill(cmd, agx_1d(range / 4), AGX_BARRIER_NONE, addr, data); } } diff --git a/src/asahi/vulkan/hk_device.c b/src/asahi/vulkan/hk_device.c index b3263bf0c301..9a33591f00cb 100644 --- a/src/asahi/vulkan/hk_device.c +++ b/src/asahi/vulkan/hk_device.c @@ -44,6 +44,15 @@ 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)"}, + {"trackcdmbarrier", HK_PERF_TRACKCDMBARRIER, + "Dependency-track the CDM barrier on chips other than G13"}, DEBUG_NAMED_VALUE_END }; /* clang-format on */ @@ -325,6 +334,57 @@ 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; + + /* G14X of generation 14 launches with the reduced CDM_BARRIER word 0x1FF + * unless HK_CDM_BARRIER_MASK is set in any form (=0 keeps the full word). + * M2 Max Qwen3-4B decode at 2048 tokens saves 0.54 ms per token with + * identical tokens and logprobs. hk_cdm_cache_flush applies it only after + * the perftest overrides and the sparse guard. + */ + dev->cdm_barrier_default_mask = + pdev->dev.chip == AGX_CHIP_G14X && + pdev->dev.params.gpu_generation == 14 && + !HK_PERF(dev, ALWAYSCDMBARRIER) && + !debug_get_option("HK_CDM_BARRIER_MASK", NULL) + ? 0x1FF + : 0; + + /* The dependency-tracked CDM barrier holds token-digest parity with the + * reference driver on G13 (M1 family) and on G14X of generation 14 (M2 + * Pro/Max/Ultra class; verified on M2 Max: greedy token digests of three + * LLM decode runs of 2048 tokens, CTS compute groups and the + * floating-point gates equal to the per-launch barrier). Other chips, + * including multi-cluster generation 15 and later parts that also map to + * G14X, keep the per-launch barrier until they pass the same gate; + * HK_PERFTEST=trackcdmbarrier opts them in for that run and + * HK_PERFTEST=alwayscdmbarrier forces the per-launch barrier everywhere. + */ + bool tracked_by_default = + pdev->dev.chip == AGX_CHIP_G13G || pdev->dev.chip == AGX_CHIP_G13X || + (pdev->dev.chip == AGX_CHIP_G14X && + pdev->dev.params.gpu_generation == 14); + + if (!tracked_by_default && !HK_PERF(dev, TRACKCDMBARRIER)) + dev->perftest |= HK_PERF_ALWAYSCDMBARRIER; + + /* Deferred CDM flush: one barrier at the end of each compute batch + * instead of the dependency-tracked barrier. It holds token-digest parity + * on M1 (G13G) and is faster there; M1 Max (G13X) lost digest parity with + * it. On by default on G13G only; HK_CDM_DEFER_FLUSH=0 disables it, + * HK_CDM_DEFER_FLUSH=1 enables it elsewhere for a gate run. + */ + dev->cdm_defer_flush = + !HK_PERF(dev, ALWAYSCDMBARRIER) && + debug_get_bool_option("HK_CDM_DEFER_FLUSH", + pdev->dev.chip == AGX_CHIP_G13G); + + /* 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..5afd089b13a7 100644 --- a/src/asahi/vulkan/hk_device.h +++ b/src/asahi/vulkan/hk_device.h @@ -97,7 +97,19 @@ struct hk_device { simple_mtx_t lock; } scratch; + /* Set to 1 when the app creates a sparse image or buffer; never cleared. */ + uint32_t sparse_used; + uint32_t perftest; + /* HK_CDM_BARRIER_MASK=<0x-prefixed hex>: perftest-only override of the + * per-launch CDM_BARRIER flag word (bits 0..26). 0 = unset. */ + uint32_t cdm_barrier_mask; + /* Reduced CDM_BARRIER word used by default on G14X generation 14 (see + * hk_device.c). 0 = the full barrier. */ + uint32_t cdm_barrier_default_mask; + /* HK_CDM_DEFER_FLUSH (hk_device.c): emit the CDM barrier at the end of + * each compute batch, not after vkCmdPipelineBarrier2. */ + bool cdm_defer_flush; struct { struct u_rwlock lock; @@ -114,6 +126,16 @@ 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 and the batch split at + * every vkCmdPipelineBarrier2. The default on G13 is the + * dependency-tracked barrier; other chips default to this. + */ + HK_PERF_ALWAYSCDMBARRIER = BITFIELD_BIT(8), + /* Use the dependency-tracked barrier on chips other than G13. */ + HK_PERF_TRACKCDMBARRIER = BITFIELD_BIT(9), }; #define HK_PERF(dev, flag) unlikely((dev)->perftest &HK_PERF_##flag) diff --git a/src/asahi/vulkan/hk_device_memory.c b/src/asahi/vulkan/hk_device_memory.c index dde1437a3aae..a00344f4c1f3 100644 --- a/src/asahi/vulkan/hk_device_memory.c +++ b/src/asahi/vulkan/hk_device_memory.c @@ -13,10 +13,57 @@ #include "asahi/lib/agx_bo.h" #include "util/u_atomic.h" +#include "util/os_misc.h" #include #include +/* GPU memory is system memory that the kernel cannot evict or swap, and + * heapBudget is only advisory, so the heap limit alone lets a process (or + * several) take every free page and invoke the OOM killer on the desktop. + * Refuse an allocation of 1 MiB or more that would leave less than + * max(1 GiB, 2% of RAM) available. Applications that allocate per tensor + * (tens of MiB per call) are covered. Smaller allocations always pass: + * below the reserve every vkAllocateMemory would fail, including ones the + * compositor needs. + */ +#define HK_MEMORY_RESERVE_MIN (1ull << 30) +#define HK_MEMORY_CHECK_MIN (1ull << 20) + +static bool +hk_system_memory_fits_now(uint64_t size) +{ + uint64_t available, total; + + /* If the kernel does not tell us, only the heap limit applies. */ + if (!os_get_available_system_memory(&available)) + return true; + + uint64_t reserve = HK_MEMORY_RESERVE_MIN; + if (os_get_total_physical_memory(&total)) + reserve = MAX2(reserve, total / 50); + + return available >= size + reserve; +} + +static bool +hk_system_memory_fits(struct hk_device *dev, uint64_t size) +{ + if (size < HK_MEMORY_CHECK_MIN) + return true; + + if (hk_system_memory_fits_now(size)) + return true; + + /* Freed device memory is parked in the driver's BO cache and still counts + * as used for MemAvailable. The cache is only trimmed when another BO is + * freed or an allocation fails, so refusing here would keep an application + * that just freed gigabytes from reusing them. Give the cache back first. + */ + agx_bo_cache_evict_all(&dev->dev); + return hk_system_memory_fits_now(size); +} + /* Supports opaque fd only */ const VkExternalMemoryProperties hk_opaque_fd_mem_props = { .externalMemoryFeatures = VK_EXTERNAL_MEMORY_FEATURE_EXPORTABLE_BIT | @@ -213,6 +260,13 @@ hk_AllocateMemory(VkDevice device, const VkMemoryAllocateInfo *pAllocateInfo, if (type->propertyFlags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT) flags |= AGX_BO_WRITEBACK; + if (!hk_system_memory_fits(dev, aligned_size)) { + result = vk_errorf(dev, VK_ERROR_OUT_OF_DEVICE_MEMORY, + "Not enough free system memory for %" PRIu64 + " MiB", aligned_size >> 20); + goto fail_alloc; + } + mem->bo = agx_bo_create(&dev->dev, aligned_size, 0, flags, "App memory"); if (!mem->bo) { result = vk_error(dev, VK_ERROR_OUT_OF_DEVICE_MEMORY); diff --git a/src/asahi/vulkan/hk_drirc_gen.py b/src/asahi/vulkan/hk_drirc_gen.py index f961899f123f..29ff837f1088 100644 --- a/src/asahi/vulkan/hk_drirc_gen.py +++ b/src/asahi/vulkan/hk_drirc_gen.py @@ -38,8 +38,9 @@ def declare_options(): drirc_gen.add_common_vk_options(debug_options, features_options, misc_options, valid_options=VALID_COMMON_VK_OPTIONS, - # Use 1/2 of total size to avoid swapping - defaults={"heap_memory_percent": 0.5}) + # GPU memory is unified system memory. The heap is a cap; + # hk_AllocateMemory keeps a free-memory reserve. + defaults={"heap_memory_percent": 0.90}) drirc_gen.add_common_vk_wsi_options(debug_options, performance_options) return [ diff --git a/src/asahi/vulkan/hk_image.c b/src/asahi/vulkan/hk_image.c index 9f742a37b8c4..f1bd4365e74e 100644 --- a/src/asahi/vulkan/hk_image.c +++ b/src/asahi/vulkan/hk_image.c @@ -918,6 +918,7 @@ hk_image_plane_alloc_vma(struct hk_device *dev, struct hk_image_plane *plane, assert(sparse_bound || !sparse_resident); if (sparse_bound) { + p_atomic_set(&dev->sparse_used, 1); size_t size = align(plane->layout.size_B, HK_SPARSE_ALIGN_B); plane->va = agx_va_alloc(&dev->dev, size, AIL_PAGESIZE, 0, 0); plane->addr = plane->va->addr; diff --git a/src/asahi/vulkan/hk_nir_lower_descriptors.c b/src/asahi/vulkan/hk_nir_lower_descriptors.c index 7de70007f1ca..17b9d5330165 100644 --- a/src/asahi/vulkan/hk_nir_lower_descriptors.c +++ b/src/asahi/vulkan/hk_nir_lower_descriptors.c @@ -60,27 +60,32 @@ load_root(nir_builder *b, unsigned num_components, unsigned bit_size, return load_speculatable(b, num_components, bit_size, addr, align); } +/* Large constant arrays (nir_opt_large_constants, compute only) live in the + * shader's data BO; the root table carries its address for the bound shader. + * The load may be speculated past the branch that guards its index, so the + * offset is clamped into the variable's range: an unguarded index past the + * BO would read unmapped memory. + */ static bool lower_load_constant(nir_builder *b, nir_intrinsic_instr *load, const struct lower_descriptors_ctx *ctx) { assert(load->intrinsic == nir_intrinsic_load_constant); - UNREACHABLE("todo: stick an address in the root descriptor or something"); - - uint32_t base = nir_intrinsic_base(load); - uint32_t range = nir_intrinsic_range(load); + assert(b->shader->info.stage == MESA_SHADER_COMPUTE); b->cursor = nir_before_instr(&load->instr); - nir_def *offset = nir_iadd_imm(b, load->src[0].ssa, base); - nir_def *data = nir_load_ubo( - b, load->def.num_components, load->def.bit_size, nir_imm_int(b, 0), - offset, .align_mul = nir_intrinsic_align_mul(load), - .align_offset = nir_intrinsic_align_offset(load), .range_base = base, - .range = range); - - nir_def_rewrite_uses(&load->def, data); - + uint32_t base = nir_intrinsic_base(load); + uint32_t bytes = load->def.num_components * load->def.bit_size / 8; + uint32_t last = MAX2(nir_intrinsic_range(load), bytes) - bytes; + nir_def *offset = nir_iadd_imm( + b, nir_umin(b, load->src[0].ssa, nir_imm_int(b, last)), base); + nir_def *data = load_root(b, 1, 64, + nir_imm_int(b, hk_root_descriptor_offset(cs.const_data)), 8); + nir_def *val = load_speculatable(b, load->def.num_components, load->def.bit_size, + nir_iadd(b, data, nir_u2u64(b, offset)), + nir_intrinsic_align(load)); + nir_def_replace(&load->def, val); return true; } diff --git a/src/asahi/vulkan/hk_physical_device.c b/src/asahi/vulkan/hk_physical_device.c index 80112cec02b2..bd2b29254fd6 100644 --- a/src/asahi/vulkan/hk_physical_device.c +++ b/src/asahi/vulkan/hk_physical_device.c @@ -45,8 +45,21 @@ hk_get_vk_version() return VK_MAKE_VERSION(1, 4, VK_HEADER_VERSION); } +/* Cooperative matrices hold token-digest parity with the reference driver on + * G13 (M1 family) and G14X (M2 Max). Other chips need AGX_SIMDMAT=1 until they + * pass the same gate. AGX_SIMDMAT=0 disables them everywhere. + */ +static bool +hk_cooperative_matrix_enabled(const struct agx_device *dev) +{ + return debug_get_bool_option( + "AGX_SIMDMAT", dev->chip == AGX_CHIP_G13G || dev->chip == AGX_CHIP_G13X || + dev->chip == AGX_CHIP_G14X); +} + static void -hk_get_device_extensions(const struct hk_instance *instance, +hk_get_device_extensions(const struct agx_device *dev, + const struct hk_instance *instance, struct vk_device_extension_table *ext) { *ext = (struct vk_device_extension_table){ @@ -55,6 +68,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 = hk_cooperative_matrix_enabled(dev), .KHR_copy_commands2 = true, .KHR_create_renderpass2 = true, .KHR_dedicated_allocation = true, @@ -110,6 +124,7 @@ hk_get_device_extensions(const struct hk_instance *instance, .KHR_sampler_ycbcr_conversion = true, .KHR_separate_depth_stencil_layouts = true, .KHR_shader_atomic_int64 = false, + .KHR_shader_bfloat16 = true, .KHR_shader_clock = false, .KHR_shader_draw_parameters = true, .KHR_shader_expect_assume = true, @@ -238,6 +253,17 @@ hk_get_device_features( struct vk_features *features) { *features = (struct vk_features){ + /* VK_KHR_cooperative_matrix (G13 HW 8x8x8, see + * hk_cooperative_matrix_enabled) */ + .cooperativeMatrix = hk_cooperative_matrix_enabled(dev), + .cooperativeMatrixRobustBufferAccess = false, + + /* VK_KHR_shader_bfloat16: conversions, and bf16 A/B cooperative + * matrices widened to fp32 (agx_nir_lower_simdmat.c) */ + .shaderBFloat16Type = true, + .shaderBFloat16DotProduct = false, + .shaderBFloat16CooperativeMatrix = hk_cooperative_matrix_enabled(dev), + /* Vulkan 1.0 */ .robustBufferAccess = true, .fullDrawIndexUint32 = true, @@ -820,6 +846,8 @@ hk_get_device_properties(const struct agx_device *dev, /* Vulkan 1.1 properties */ .subgroupSize = 32, + .cooperativeMatrixSupportedStages = + hk_cooperative_matrix_enabled(dev) ? VK_SHADER_STAGE_COMPUTE_BIT : 0, .subgroupSupportedStages = VK_SHADER_STAGE_COMPUTE_BIT | VK_SHADER_STAGE_ALL_GRAPHICS, .subgroupSupportedOperations = @@ -1265,7 +1293,7 @@ hk_create_drm_physical_device(struct vk_instance *_instance, &dispatch_table, &wsi_physical_device_entrypoints, false); struct vk_device_extension_table supported_extensions; - hk_get_device_extensions(instance, &supported_extensions); + hk_get_device_extensions(&pdev->dev, instance, &supported_extensions); struct vk_features supported_features; hk_get_device_features(&pdev->dev, instance, &supported_extensions, @@ -1518,3 +1546,48 @@ 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); + VK_OUTARRAY_MAKE_TYPED(VkCooperativeMatrixPropertiesKHR, out, pProperties, + pPropertyCount); + + if (hk_cooperative_matrix_enabled(&pdev->dev)) { + /* 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. + * bf16 A/B widen exactly to fp32 and run the fp32 A/B form. + */ + 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_BFLOAT16_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..96459e1d4bd1 100644 --- a/src/asahi/vulkan/hk_shader.c +++ b/src/asahi/vulkan/hk_shader.c @@ -96,11 +96,40 @@ shared_var_info(const struct glsl_type *type, unsigned *size, unsigned *align) *size = comp_size * length, *align = comp_size; } +/* Large constant arrays become compute-shader constant data; HK_LARGE_CONSTANTS=0 + * keeps the per-invocation scratch copy (A/B and bisect only). + */ +static bool +hk_large_constants(void) +{ + return debug_get_bool_option("HK_LARGE_CONSTANTS", true); +} + +/* Compute load hoist: on by default on generation 13 (M1 family), opt-in on + * generation 14 with AGX_HOIST_LOADS=1, never on other generations. + * AGX_HOIST_LOADS=0 turns it off everywhere. Part of the compiler flags, so + * toggling it never serves a shader cached under the other setting. + */ +static bool +hk_hoist_loads_enabled(const struct hk_physical_device *pdev) +{ + const unsigned gen = pdev->dev.params.gpu_generation; + + return (gen == 13 || gen == 14) && + debug_get_bool_option("AGX_HOIST_LOADS", gen == 13); +} + 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_HWMAT_VEC2_OFF") != NULL) << 34) | + ((uint64_t)(getenv("AGX_LOCAL_FOLD_OFF") != NULL) << 35) | + ((uint64_t)(getenv("AGX_SCHED_COMPUTE") != NULL) << 36) | + ((uint64_t)hk_hoist_loads_enabled(pdev) << 37) | + ((uint64_t)hk_large_constants() << 40); } const nir_shader_compiler_options * @@ -174,6 +203,18 @@ hk_preprocess_nir_internal(struct vk_physical_device *vk_pdev, nir_shader *nir) NIR_PASS(_, nir, nir_split_var_copies); NIR_PASS(_, nir, nir_split_struct_vars, nir_var_function_temp); + /* Constant arrays that nir_lower_vars_to_scratch (in agx_preprocess_nir, + * same 256 B threshold) would give every invocation a scratch copy of + * become shader constant data read through the root table instead. + * Smaller indirectly indexed arrays keep their if-else lowering. Compute + * only: graphics stages have no root field for it. + * nir_opt_large_constants cannot see through copy_deref. + */ + if (nir->info.stage == MESA_SHADER_COMPUTE && hk_large_constants()) { + NIR_PASS(_, nir, nir_lower_var_copies); + NIR_PASS(_, nir, nir_opt_large_constants, NULL, 256); + } + /* Optimize but allow copies because we haven't lowered them yet */ agx_preprocess_nir(nir); @@ -829,8 +870,6 @@ hk_lower_nir(struct hk_device *dev, nir_shader *nir, NIR_PASS(_, nir, nir_lower_explicit_io, nir_var_mem_push_const, nir_address_format_32bit_offset); - // NIR_PASS(_, nir, nir_opt_large_constants, NULL, 32); - /* Turn cache flushes into image coherency bits while we still have derefs */ NIR_PASS(_, nir, nir_lower_memory_model); @@ -937,7 +976,7 @@ hk_lower_nir(struct hk_device *dev, nir_shader *nir, nir_opt_if_optimize_phi_true_false | nir_opt_if_avoid_64bit_phis); } -static void +static VkResult hk_upload_shader(struct hk_device *dev, struct hk_shader *shader) { if (shader->b.info.has_preamble || shader->b.info.rodata.size_16) { @@ -947,11 +986,24 @@ hk_upload_shader(struct hk_device *dev, struct hk_shader *shader) shader->bo = agx_bo_create(&dev->dev, size, 0, AGX_BO_EXEC | AGX_BO_LOW_VA, "Preamble"); + if (!shader->bo) + return vk_error(dev, VK_ERROR_OUT_OF_DEVICE_MEMORY); + memcpy(agx_bo_map(shader->bo), shader->b.binary, size); shader->preamble_addr = shader->bo->va->addr + shader->b.info.preamble_offset; } + if (shader->data_size) { + shader->data_bo = + agx_bo_create(&dev->dev, shader->data_size, 0, 0, "Shader constants"); + if (!shader->data_bo) + return vk_error(dev, VK_ERROR_OUT_OF_DEVICE_MEMORY); + + memcpy(agx_bo_map(shader->data_bo), shader->data_ptr, shader->data_size); + shader->data_addr = shader->data_bo->va->addr; + } + if (!shader->linked.ht) { /* If we only have a single variant, link now. */ shader->only_linked = hk_fast_link(dev, false, shader, NULL, NULL, 0); @@ -968,6 +1020,7 @@ hk_upload_shader(struct hk_device *dev, struct hk_shader *shader) shader->b.info.sampler_state_count, false); cfg.texture_state_register_count = shader->b.info.texture_state_count; } + return VK_SUCCESS; } DERIVE_HASH_TABLE(hk_fast_link_key_vs); @@ -1178,6 +1231,9 @@ hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator, .promote_textures = true, }; + backend_key.dev.hoist_loads = + hk_hoist_loads_enabled(hk_device_physical(dev)); + /* For now, sample shading is always dynamic. Indicate that. */ if (nir->info.stage == MESA_SHADER_FRAGMENT && nir->info.fs.uses_sample_shading) @@ -1237,9 +1293,7 @@ hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator, VkResult result = hk_init_link_ht(shader, sw_stage); if (result != VK_SUCCESS) return vk_error(dev, result); - - hk_upload_shader(dev, shader); - return VK_SUCCESS; + return hk_upload_shader(dev, shader); } static const struct vk_shader_ops hk_shader_ops; @@ -1257,6 +1311,7 @@ hk_shader_destroy(struct hk_device *dev, struct hk_shader *s) free((void *)s->code_ptr); free((void *)s->data_ptr); agx_bo_unreference(&dev->dev, s->bo); + agx_bo_unreference(&dev->dev, s->data_bo); simple_mtx_destroy(&s->linked.lock); @@ -1632,8 +1687,7 @@ hk_deserialize_shader(struct hk_device *dev, struct blob_reader *blob, return vk_error(dev, VK_ERROR_INCOMPATIBLE_SHADER_BINARY_EXT); shader->b.binary = (void *)shader->code_ptr; - hk_upload_shader(dev, shader); - return VK_SUCCESS; + return hk_upload_shader(dev, shader); } static VkResult diff --git a/src/asahi/vulkan/hk_shader.h b/src/asahi/vulkan/hk_shader.h index dba428cc822e..fe3f5c9189a8 100644 --- a/src/asahi/vulkan/hk_shader.h +++ b/src/asahi/vulkan/hk_shader.h @@ -147,6 +147,9 @@ struct hk_shader { /* BO for any uploaded shader part */ struct agx_bo *bo; + /* BO holding nir->constant_data (large constant arrays), compute only */ + struct agx_bo *data_bo; + /* Cache of fast linked variants */ struct { simple_mtx_t lock; 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/compiler/nir/nir_opt_large_constants.c b/src/compiler/nir/nir_opt_large_constants.c index 54fd2f7b087d..09b9fe6f32ac 100644 --- a/src/compiler/nir/nir_opt_large_constants.c +++ b/src/compiler/nir/nir_opt_large_constants.c @@ -240,6 +240,23 @@ handle_constant_store(void *mem_ctx, struct var_info *info, } } +/* Whether a constant store leaves the collected data unchanged. Bytes no + * earlier store wrote are zero and hold undefined values, so zero is a valid + * value for them as well. + */ +static bool +constant_store_is_redundant(void *mem_ctx, struct var_info *info, + nir_deref_instr *deref, nir_def *val, + nir_component_mask_t write_mask, + glsl_type_size_align_func size_align) +{ + void *before = ralloc_memdup(NULL, info->constant_data, info->constant_data_size); + handle_constant_store(mem_ctx, info, deref, val, write_mask, size_align); + bool same = memcmp(before, info->constant_data, info->constant_data_size) == 0; + ralloc_free(before); + return same; +} + #define NIR_SMALL_CONSTANT_MAX_ABS_VALUE 255 static bool @@ -589,12 +606,20 @@ opt_large_constants_impl(nir_function_impl *impl, info->block = block; /* We only consider variables constant if they only have constant - * stores, all the stores come before any reads, and all stores - * come from the same block. We also can't handle indirect stores. + * stores, all the stores come before any reads, and every store + * outside the first store block rewrites bytes already collected, + * which changes no value (SPIR-V producers store a constant array + * into a variable that also has it as initializer). We also can't + * handle indirect stores. */ - if (!src_is_const || info->found_read || block != info->block || + if (!src_is_const || info->found_read || nir_deref_instr_has_indirect(dst_deref)) { info->is_constant = false; + } else if (block != info->block) { + if (!constant_store_is_redundant(var_infos, info, dst_deref, + intrin->src[1].ssa, write_mask, + size_align)) + info->is_constant = false; } else { handle_constant_store(var_infos, info, dst_deref, intrin->src[1].ssa, write_mask, size_align); diff --git a/src/compiler/nir/tests/opt_large_constants_tests.cpp b/src/compiler/nir/tests/opt_large_constants_tests.cpp index 0dca33f85151..123f331b9136 100644 --- a/src/compiler/nir/tests/opt_large_constants_tests.cpp +++ b/src/compiler/nir/tests/opt_large_constants_tests.cpp @@ -385,3 +385,60 @@ TEST_F(nir_large_constants_test, bcsel_vec) } )")); } + +static unsigned +count_load_deref(nir_function_impl *impl) +{ + unsigned count = 0; + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type == nir_instr_type_intrinsic && + nir_instr_as_intrinsic(instr)->intrinsic == nir_intrinsic_load_deref) + count++; + } + } + return count; +} + +/* The table is stored once in the first block and again, with the same or a + * different value at index 5, under a condition. Only the same value keeps it + * constant. + */ +static void +store_table_twice(nir_builder *b, nir_variable *array, uint32_t length, + uint32_t second_value_5) +{ + for (uint32_t i = 0; i < length; i++) + nir_store_array_var_imm(b, array, i, nir_imm_int(b, 1000003 * i), 0x1); + + nir_push_if(b, nir_ieq_imm(b, nir_load_workgroup_index(b), 0)); + for (uint32_t i = 0; i < length; i++) { + uint32_t v = i == 5 ? second_value_5 : 1000003 * i; + nir_store_array_var_imm(b, array, i, nir_imm_int(b, v), 0x1); + } + nir_pop_if(b, NULL); +} + +TEST_F(nir_large_constants_test, same_store_in_another_block) +{ + uint32_t length = 8; + array = nir_local_variable_create(b->impl, glsl_array_type(glsl_uint_type(), length, 0), "array"); + store_table_twice(b, array, length, 1000003 * 5); + + run_test(); + + EXPECT_EQ(count_load_deref(b->impl), 0); + EXPECT_EQ(b->shader->constant_data_size, length * 4); +} + +TEST_F(nir_large_constants_test, different_store_in_another_block) +{ + uint32_t length = 8; + array = nir_local_variable_create(b->impl, glsl_array_type(glsl_uint_type(), length, 0), "array"); + store_table_twice(b, array, length, 7); + + run_test(); + + EXPECT_EQ(count_load_deref(b->impl), 1); + EXPECT_EQ(b->shader->constant_data_size, 0); +} diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c index 04442865e544..1dd69a4774fe 100644 --- a/src/compiler/spirv/spirv_to_nir.c +++ b/src/compiler/spirv/spirv_to_nir.c @@ -4845,6 +4845,21 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode, SpvMemorySemanticsMask after_semantics; vtn_split_barrier_semantics(b, semantics, &before_semantics, &after_semantics); + /* MakeTexelAvailable and MakeTexelVisible apply to the texel of this very + * access, unlike MakeAvailable on a release atomic. The split above puts + * MakeAvailable before the operation and MakeVisible after it, which is + * right for atomics but leaves a MakeTexelAvailable write outside its own + * availability operation and a MakeTexelVisible read unordered after it. + * Image reads and writes carry no Acquire/Release bits, so swap the sides: + * a write is made available after it, a read is made visible before it. + */ + if (opcode == SpvOpImageRead || opcode == SpvOpImageSparseRead || + opcode == SpvOpImageWrite) { + SpvMemorySemanticsMask tmp = before_semantics; + before_semantics = after_semantics; + after_semantics = tmp; + } + if (before_semantics) vtn_emit_memory_barrier(b, scope, before_semantics); diff --git a/src/compiler/spirv/tests/avail_vis.cpp b/src/compiler/spirv/tests/avail_vis.cpp index 922694f0678f..151729abc196 100644 --- a/src/compiler/spirv/tests/avail_vis.cpp +++ b/src/compiler/spirv/tests/avail_vis.cpp @@ -134,6 +134,167 @@ TEST_F(AvailabilityVisibility, opstore_avail) EXPECT_EQ(nir_intrinsic_execution_scope(intrinsic), SCOPE_NONE); } +/* Position of the first intrinsic of the given op in the entry point, or -1. */ +static int +intrinsic_position(nir_shader *shader, nir_intrinsic_op op) +{ + int pos = 0; + nir_function_impl *impl = nir_shader_get_entrypoint(shader); + nir_foreach_block(block, impl) { + nir_foreach_instr(instr, block) { + if (instr->type == nir_instr_type_intrinsic && + nir_instr_as_intrinsic(instr)->intrinsic == op) + return pos; + pos++; + } + } + return -1; +} + +TEST_F(AvailabilityVisibility, opimageread_vis) +{ + /* + OpCapability Shader + OpCapability VulkanMemoryModel + OpCapability VulkanMemoryModelDeviceScope + OpMemoryModel Logical Vulkan + OpEntryPoint GLCompute %main "main" %img + OpExecutionMode %main LocalSize 1 1 1 + OpDecorate %img DescriptorSet 0 + OpDecorate %img Binding 0 + %void = OpTypeVoid + %fn = OpTypeFunction %void + %uint = OpTypeInt 32 0 + %int = OpTypeInt 32 1 + %v2int = OpTypeVector %int 2 + %v4uint = OpTypeVector %uint 4 + %imgtype = OpTypeImage %uint 2D 0 0 0 2 R32ui + %ptr = OpTypePointer UniformConstant %imgtype + %img = OpVariable %ptr UniformConstant + %int_0 = OpConstant %int 0 + %coord = OpConstantComposite %v2int %int_0 %int_0 + %device = OpConstant %int 1 + %zero = OpConstantNull %v4uint + %main = OpFunction %void None %fn + %l = OpLabel + %im = OpLoad %imgtype %img + %v = OpImageRead %v4uint %im %coord MakeTexelVisible|NonPrivateTexel %device + OpImageWrite %im %coord %v + OpReturn + OpFunctionEnd + */ + static const uint32_t words[] = { + 0x07230203, 0x00010500, 0x00070000, 0x00000012, 0x00000000, 0x00020011, + 0x00000001, 0x00020011, 0x000014e1, 0x00020011, 0x000014e2, 0x0003000e, + 0x00000000, 0x00000003, 0x0006000f, 0x00000005, 0x00000001, 0x6e69616d, + 0x00000000, 0x00000002, 0x00060010, 0x00000001, 0x00000011, 0x00000001, + 0x00000001, 0x00000001, 0x00040047, 0x00000002, 0x00000022, 0x00000000, + 0x00040047, 0x00000002, 0x00000021, 0x00000000, 0x00020013, 0x00000003, + 0x00030021, 0x00000004, 0x00000003, 0x00040015, 0x00000005, 0x00000020, + 0x00000000, 0x00040015, 0x00000006, 0x00000020, 0x00000001, 0x00040017, + 0x00000007, 0x00000006, 0x00000002, 0x00040017, 0x00000008, 0x00000005, + 0x00000004, 0x00090019, 0x00000009, 0x00000005, 0x00000001, 0x00000000, + 0x00000000, 0x00000000, 0x00000002, 0x00000021, 0x00040020, 0x0000000a, + 0x00000000, 0x00000009, 0x0004003b, 0x0000000a, 0x00000002, 0x00000000, + 0x0004002b, 0x00000006, 0x0000000b, 0x00000000, 0x0005002c, 0x00000007, + 0x0000000c, 0x0000000b, 0x0000000b, 0x0004002b, 0x00000006, 0x0000000d, + 0x00000001, 0x0003002e, 0x00000008, 0x0000000e, 0x00050036, 0x00000003, + 0x00000001, 0x00000000, 0x00000004, 0x000200f8, 0x0000000f, 0x0004003d, + 0x00000009, 0x00000010, 0x00000002, 0x00070062, 0x00000008, 0x00000011, + 0x00000010, 0x0000000c, 0x00000600, 0x0000000d, 0x00040063, 0x00000010, + 0x0000000c, 0x00000011, 0x000100fd, 0x00010038, + }; + + get_nir(sizeof(words) / sizeof(words[0]), words); + + nir_intrinsic_instr *intrinsic = find_intrinsic(nir_intrinsic_barrier, 0); + ASSERT_NE(intrinsic, nullptr); + + EXPECT_EQ(nir_intrinsic_memory_semantics(intrinsic), NIR_MEMORY_MAKE_VISIBLE); + EXPECT_NE(nir_intrinsic_memory_modes(intrinsic) & nir_var_image, 0); + EXPECT_EQ(nir_intrinsic_memory_scope(intrinsic), SCOPE_DEVICE); + EXPECT_EQ(nir_intrinsic_execution_scope(intrinsic), SCOPE_NONE); + + /* MakeTexelVisible makes the texel visible before it is read. */ + int barrier = intrinsic_position(shader, nir_intrinsic_barrier); + int load = intrinsic_position(shader, nir_intrinsic_image_deref_load); + ASSERT_GE(barrier, 0); + ASSERT_GE(load, 0); + EXPECT_LT(barrier, load); +} + +TEST_F(AvailabilityVisibility, opimagewrite_avail) +{ + /* + OpCapability Shader + OpCapability VulkanMemoryModel + OpCapability VulkanMemoryModelDeviceScope + OpMemoryModel Logical Vulkan + OpEntryPoint GLCompute %main "main" %img + OpExecutionMode %main LocalSize 1 1 1 + OpDecorate %img DescriptorSet 0 + OpDecorate %img Binding 0 + %void = OpTypeVoid + %fn = OpTypeFunction %void + %uint = OpTypeInt 32 0 + %int = OpTypeInt 32 1 + %v2int = OpTypeVector %int 2 + %v4uint = OpTypeVector %uint 4 + %imgtype = OpTypeImage %uint 2D 0 0 0 2 R32ui + %ptr = OpTypePointer UniformConstant %imgtype + %img = OpVariable %ptr UniformConstant + %int_0 = OpConstant %int 0 + %coord = OpConstantComposite %v2int %int_0 %int_0 + %device = OpConstant %int 1 + %zero = OpConstantNull %v4uint + %main = OpFunction %void None %fn + %l = OpLabel + %im = OpLoad %imgtype %img + %v = OpImageRead %v4uint %im %coord + OpImageWrite %im %coord %v MakeTexelAvailable|NonPrivateTexel %device + OpReturn + OpFunctionEnd + */ + static const uint32_t words[] = { + 0x07230203, 0x00010500, 0x00070000, 0x00000012, 0x00000000, 0x00020011, + 0x00000001, 0x00020011, 0x000014e1, 0x00020011, 0x000014e2, 0x0003000e, + 0x00000000, 0x00000003, 0x0006000f, 0x00000005, 0x00000001, 0x6e69616d, + 0x00000000, 0x00000002, 0x00060010, 0x00000001, 0x00000011, 0x00000001, + 0x00000001, 0x00000001, 0x00040047, 0x00000002, 0x00000022, 0x00000000, + 0x00040047, 0x00000002, 0x00000021, 0x00000000, 0x00020013, 0x00000003, + 0x00030021, 0x00000004, 0x00000003, 0x00040015, 0x00000005, 0x00000020, + 0x00000000, 0x00040015, 0x00000006, 0x00000020, 0x00000001, 0x00040017, + 0x00000007, 0x00000006, 0x00000002, 0x00040017, 0x00000008, 0x00000005, + 0x00000004, 0x00090019, 0x00000009, 0x00000005, 0x00000001, 0x00000000, + 0x00000000, 0x00000000, 0x00000002, 0x00000021, 0x00040020, 0x0000000a, + 0x00000000, 0x00000009, 0x0004003b, 0x0000000a, 0x00000002, 0x00000000, + 0x0004002b, 0x00000006, 0x0000000b, 0x00000000, 0x0005002c, 0x00000007, + 0x0000000c, 0x0000000b, 0x0000000b, 0x0004002b, 0x00000006, 0x0000000d, + 0x00000001, 0x0003002e, 0x00000008, 0x0000000e, 0x00050036, 0x00000003, + 0x00000001, 0x00000000, 0x00000004, 0x000200f8, 0x0000000f, 0x0004003d, + 0x00000009, 0x00000010, 0x00000002, 0x00050062, 0x00000008, 0x00000011, + 0x00000010, 0x0000000c, 0x00060063, 0x00000010, 0x0000000c, 0x00000011, + 0x00000500, 0x0000000d, 0x000100fd, 0x00010038, + }; + + get_nir(sizeof(words) / sizeof(words[0]), words); + + nir_intrinsic_instr *intrinsic = find_intrinsic(nir_intrinsic_barrier, 0); + ASSERT_NE(intrinsic, nullptr); + + EXPECT_EQ(nir_intrinsic_memory_semantics(intrinsic), NIR_MEMORY_MAKE_AVAILABLE); + EXPECT_NE(nir_intrinsic_memory_modes(intrinsic) & nir_var_image, 0); + EXPECT_EQ(nir_intrinsic_memory_scope(intrinsic), SCOPE_DEVICE); + EXPECT_EQ(nir_intrinsic_execution_scope(intrinsic), SCOPE_NONE); + + /* MakeTexelAvailable makes the texel available after it is written. */ + int barrier = intrinsic_position(shader, nir_intrinsic_barrier); + int store = intrinsic_position(shader, nir_intrinsic_image_deref_store); + ASSERT_GE(barrier, 0); + ASSERT_GE(store, 0); + EXPECT_GT(barrier, store); +} + TEST_F(AvailabilityVisibility, opcopymemory_visavail_both_combined) { /* 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..2970fc33a616 100644 --- a/src/vulkan/runtime/vk_drm_syncobj.c +++ b/src/vulkan/runtime/vk_drm_syncobj.c @@ -325,29 +325,80 @@ 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; + /* Never poll past the caller's deadline: a zero or already expired + * deadline (vkGetFenceStatus, short waits) skips the poll entirely. */ + const uint64_t poll_now_ns = os_time_get_nano(); + if (poll_budget_ns && poll_now_ns < abs_timeout_ns && + poll_now_ns >= p_atomic_read(&device->sync_poll_miss_expire_ns)) + poll_deadline_ns = MIN2(poll_now_ns + poll_budget_ns, abs_timeout_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/NUMERICAL-HARNESS.md b/tests/NUMERICAL-HARNESS.md new file mode 100644 index 000000000000..fdc9b6ac8a49 --- /dev/null +++ b/tests/NUMERICAL-HARNESS.md @@ -0,0 +1,26 @@ +# Numerical qualification harness + +Driver source under test: Joshua Warren’s correctness-only adaptation +`5e45716701b64a4b627e7e60df5b781c65c87f42`. CI source is the previously checked +`001e33ae855bd6b9ddc28bfc60a6bd64d71bcf91` trial tree. This branch adds test +harness fixes only; the production PR and package source pins are unchanged. + +The ULP checker uses a positive binary32 ULP at the magnitude of the correctly +rounded reference, calculated in binary64. It explicitly rejects a non-finite +result for a finite reference. Special reference values retain bitwise checks, +with NaN payloads ignored. Seven CPU controls cover both signs, a power-of-two +boundary, out-of-budget results, NaN/infinity, FLT_MAX and signed zero. + +The original checker measures binary64 spacing, gives negative distances for +negative references and can accept NaN for a finite reference. These are oracle +defects, not observations of GPU arithmetic failure. + +The following files were copied byte for byte from Joshua Warren’s public +source at `f144c7949b7d960f26677b671341d5a123057ddf` to supply the fixtures +missing from the correctness extraction. No CDM, scheduling or performance +source is imported. + +- [precise_math.comp](https://github.com/joshuaswarren/mesa-1/blob/f144c7949b7d960f26677b671341d5a123057ddf/tests/precise_math.comp), Git blob `827941d2468890243996dc8d9ab7f74005dab062`. +- [compute_runner.cpp](https://github.com/joshuaswarren/mesa-1/blob/f144c7949b7d960f26677b671341d5a123057ddf/tests/compute_runner.cpp), Git blob `e16d45d9d61252d72b2389da8c5073c781e696ca`. + +Validation: seven CPU oracle controls pass. No numerical GPU result is claimed by this commit. diff --git a/tests/cmshape2.c b/tests/cmshape2.c new file mode 100644 index 000000000000..bcddc85087b0 --- /dev/null +++ b/tests/cmshape2.c @@ -0,0 +1,398 @@ +/* + * cmshape2: extended cmshape for the cooperative matrix 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; + if (vbytes == 1) + f12.storageBuffer8BitAccess = VK_TRUE; /* the u8 case uses uint8_t storage buffers */ + 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/div_window.comp b/tests/div_window.comp new file mode 100644 index 000000000000..a622787dc435 --- /dev/null +++ b/tests/div_window.comp @@ -0,0 +1,42 @@ +#version 450 +layout(local_size_x = 64) in; +layout(std430, binding = 0) readonly buffer In { uint base; uint seed; } p; +layout(std430, binding = 1) writeonly buffer Out { uint o[]; } outp; + +uint pcg(uint v) +{ + uint s = v * 747796405u + 2891336453u; + uint w = ((s >> ((s >> 28u) + 4u)) ^ s) * 277803737u; + return (w >> 22u) ^ w; +} + +uint with_exp(uint bits, uint e) { return (bits & 0x807fffffu) | (e << 23u); } + +/* Exponents at the zero/denormal (0), normal and 2^126..inf (252..255) edges, + * plus a 46..50 group that keeps the residual of the quotient near 2^-126. */ +const uint EA[16] = uint[](0u, 1u, 2u, 46u, 47u, 48u, 49u, 50u, 126u, 127u, 128u, 252u, 253u, 254u, 255u, 255u); +const uint ED[16] = uint[](0u, 0u, 1u, 2u, 3u, 126u, 127u, 128u, 250u, 251u, 252u, 253u, 254u, 255u, 1u, 252u); + +void main() +{ + for (uint j = 0u; j < 32u; j++) { + uint i = gl_GlobalInvocationID.x * 32u + j; + uint x = p.base + i; + uint h0 = pcg(x ^ p.seed), h1 = pcg(h0 + 1u), h2 = pcg(h1 + 2u); + uint a = h0, d = h1, mode = h2 & 3u; + if (mode == 1u) { + a = with_exp(a, EA[(h2 >> 2u) & 15u]); + d = with_exp(d, ED[(h2 >> 6u) & 15u]); + } else if (mode == 2u) { + uint ed = 1u + (h2 >> 2u) % 254u; + int e = int((h2 >> 10u) % 14u) + ((h2 & 0x80000000u) != 0u ? 117 : -128); + a = with_exp(a, uint(clamp(int(ed) + e, 0, 255))); + d = with_exp(d, ed); + } else if (mode == 3u) { + a = with_exp(a, 1u + (h2 >> 2u) % 254u); + d = with_exp(d, 1u + (h2 >> 10u) % 254u); + } + outp.o[2u * i] = floatBitsToUint(uintBitsToFloat(a) / uintBitsToFloat(d)); + outp.o[2u * i + 1u] = floatBitsToUint(1.0 / uintBitsToFloat(x)); + } +} diff --git a/tests/div_window.py b/tests/div_window.py new file mode 100755 index 000000000000..1c4932d8b10a --- /dev/null +++ b/tests/div_window.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Usage: div_window.py ICD COMPUTE_RUNNER [REF_ICD] [CHUNKS]; glslangValidator and numpy are required. + +Default-path division and reciprocal against the correctly rounded (with FTZ) +quotient, bitwise, specials included: every fp32 reciprocal (64 chunks of +2^26) and the same number of hashed quotients, a quarter of them raw bits and +the rest with exponents at the zero, denormal, inf and 2^126 edges and with +quotients straddling 2^-126 and 2^128. With REF_ICD, every result must also +match that driver bitwise (for example the build before a lowering change). +CHUNKS is N (the first N of 64) or FIRST-LAST, so an interrupted sweep can +resume. +""" +import json +import os +from pathlib import Path +import subprocess +import sys +import tempfile +import numpy as np + +args = sys.argv[1:] +icd, runner = map(os.path.abspath, args[:2]) +ref = os.path.abspath(args[2]) if len(args) > 2 else None +span = args[3].split("-") if len(args) > 3 else ["64"] +chunks = range(int(span[0]), int(span[1]) + 1) if len(span) == 2 else range(int(span[0])) +here = Path(__file__).resolve().parent +N = 1 << 26 +SEED = 20261006 + +def pcg(v): + s = v * np.uint32(747796405) + np.uint32(2891336453) + w = ((s >> ((s >> np.uint32(28)) + np.uint32(4))) ^ s) * np.uint32(277803737) + return (w >> np.uint32(22)) ^ w + +EA = np.array([0, 1, 2, 46, 47, 48, 49, 50, 126, 127, 128, 252, 253, 254, 255, 255], dtype=np.uint32) +ED = np.array([0, 0, 1, 2, 3, 126, 127, 128, 250, 251, 252, 253, 254, 255, 1, 252], dtype=np.uint32) + +def operands(base): + x = np.arange(N, dtype=np.uint32) + np.uint32(base) + h0 = pcg(x ^ np.uint32(SEED)) + h1 = pcg(h0 + np.uint32(1)) + h2 = pcg(h1 + np.uint32(2)) + a, d, mode = h0.copy(), h1.copy(), h2 & 3 + exp = lambda bits, e: (bits & np.uint32(0x807fffff)) | (e.astype(np.uint32) << np.uint32(23)) + m = mode == 1 + a[m] = exp(a[m], EA[(h2[m] >> 2) & 15]) + d[m] = exp(d[m], ED[(h2[m] >> 6) & 15]) + m = mode == 2 + ed = 1 + (h2[m] >> 2) % 254 + e = ((h2[m] >> 10) % 14).astype(np.int64) + np.where(h2[m] & 0x80000000, 117, -128) + a[m] = exp(a[m], np.clip(ed.astype(np.int64) + e, 0, 255)) + d[m] = exp(d[m], ed) + m = mode == 3 + a[m] = exp(a[m], 1 + (h2[m] >> 2) % 254) + d[m] = exp(d[m], 1 + (h2[m] >> 10) % 254) + return a, d, x + +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 bad(got, want): + both_nan = ((got & 0x7fffffff) > 0x7f800000) & ((want & 0x7fffffff) > 0x7f800000) + return np.flatnonzero((got != want) & ~both_nan) + +def run(spv, driver, base, path): + path.with_suffix(".in").write_bytes(np.array([base, SEED], dtype=np.uint32).tobytes()) + env = dict(os.environ, VK_DRIVER_FILES=driver, VK_ICD_FILENAMES=driver, MESA_SHADER_CACHE_DISABLE="true") + subprocess.run(["flock", "-w", "1800", "/tmp/gpu.lock", runner, str(spv), str(path.with_suffix(".in")), + str(path), str(N * 8), str(N // 32 // 64)], env=env, check=True, timeout=2100) + return np.fromfile(path, dtype=np.uint32).reshape(N, 2) + +totals = {"div_oracle": 0, "recip_oracle": 0, "div_ref": 0, "recip_ref": 0} +examples = [] +with tempfile.TemporaryDirectory(prefix="mesa-divwin-", dir="/dev/shm" if os.path.isdir("/dev/shm") else None) as td: + td = Path(td) + spv = td / "probe.spv" + subprocess.run(["glslangValidator", "-V", str(here / "div_window.comp"), "-o", str(spv)], check=True, timeout=30) + for c in chunks: + base = c * N + out = run(spv, icd, base, td / "out.bin") + a, d, x = operands(base) + checks = {"div_oracle": bad(out[:, 0], quotient(a, d)), + "recip_oracle": bad(out[:, 1], quotient(np.full(N, 0x3f800000, dtype=np.uint32), x))} + if ref: + r = run(spv, ref, base, td / "ref.bin") + checks["div_ref"] = np.flatnonzero(out[:, 0] != r[:, 0]) + checks["recip_ref"] = np.flatnonzero(out[:, 1] != r[:, 1]) + for k, v in checks.items(): + totals[k] += len(v) + for i in v[:2]: + if len(examples) < 10: + examples.append({"check": k, "a": hex(int(a[i])), "d": hex(int(d[i])), "x": hex(int(x[i])), + "got": [hex(int(w)) for w in out[i]]}) + print(f"chunk {c}: " + " ".join(f"{k}={len(v)}" for k, v in checks.items()), flush=True) + +result = {"pairs": len(chunks) * N, "reciprocals": len(chunks) * N, "chunks": [chunks[0], chunks[-1]], + "mismatches": totals if ref else + {k: v for k, v in totals.items() if k.endswith("oracle")}, "examples": examples} +print(json.dumps(result, sort_keys=True)) +sys.exit(any(totals.values())) diff --git a/tests/large_constants.py b/tests/large_constants.py new file mode 100755 index 000000000000..43a1e810458e --- /dev/null +++ b/tests/large_constants.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Usage: large_constants.py ICD COMPUTE_RUNNER; glslangValidator and numpy are required. + +llama.cpp IQ grid pattern: a 1024-entry constant table copied into shared memory with a lane-dependent index, +a barrier, then 96 lane-dependent lookups per invocation (4096 workgroups of 32). Runs the shader with +HK_LARGE_CONSTANTS=0 and unset (shader cache off, AGX_MESA_DEBUG=shaderdb). Both runs must match the CPU +model bitwise; by default the shader must use no scratch (with HK_LARGE_CONSTANTS=0, Honeykrisp gives every +invocation a 4096 B scratch copy of the table). + +Expected output (G13C, G13G, G14C): + off: 0 mismatches of 131072, scratch 4096 + on: 0 mismatches of 131072, scratch 0 +""" +import os +from pathlib import Path +import re +import subprocess +import sys +import tempfile +import numpy as np + +icd, runner = map(os.path.abspath, sys.argv[1:3]) +N_WG, WG, LOOKUPS, SEED = 4096, 32, 96, 20261007 +table = [((i * 2654435761) >> 7) & 0xffffffff for i in range(1024)] +shader = f"""#version 450 +layout(local_size_x = {WG}) in; +layout(std430, binding = 0) readonly buffer In {{ uint seed; }} p; +layout(std430, binding = 1) writeonly buffer Out {{ uint o[]; }} outp; +const uint T[1024] = uint[]({", ".join(f"{v}u" for v in table)}); +shared uint S[1024]; +void main() +{{ + uint lid = gl_LocalInvocationIndex; + for (uint i = 0u; i < 1024u; i += {WG}u) + S[i + lid] = T[i + lid]; + barrier(); + uint g = gl_GlobalInvocationID.x, acc = 0u; + for (uint k = 0u; k < {LOOKUPS}u; k++) + acc = acc * 31u + S[(g * 7u + k * 13u + p.seed) & 1023u]; + outp.o[g] = acc; +}} +""" + +n = N_WG * WG +g = np.arange(n, dtype=np.uint64) +t = np.array(table, dtype=np.uint64) +acc = np.zeros(n, dtype=np.uint64) +for k in range(LOOKUPS): + acc = (acc * 31 + t[(g * 7 + k * 13 + SEED) & 1023]) & 0xffffffff +want = acc.astype(np.uint32) + +failed = False +with tempfile.TemporaryDirectory(prefix="large-constants-") as td: + td = Path(td) + (td / "p.comp").write_text(shader) + subprocess.run(["glslangValidator", "-V", str(td / "p.comp"), "-o", str(td / "p.spv")], check=True, + capture_output=True, timeout=60) + (td / "in.bin").write_bytes(np.array([SEED], dtype=np.uint32).tobytes()) + for arm, on in (("off", False), ("on", True)): + env = dict(os.environ, VK_DRIVER_FILES=icd, VK_ICD_FILENAMES=icd, AGX_MESA_DEBUG="shaderdb", + MESA_SHADER_CACHE_DISABLE="true") + env.pop("HK_LARGE_CONSTANTS", None) + if not on: + env["HK_LARGE_CONSTANTS"] = "0" + r = subprocess.run(["flock", "-w", "1800", "/tmp/gpu.lock", runner, str(td / "p.spv"), + str(td / "in.bin"), str(td / "out.bin"), str(n * 4), str(N_WG)], + env=env, capture_output=True, text=True, timeout=2100) + stats = [l for l in (r.stdout + r.stderr).splitlines() if l.startswith("CS shader")] + scratch = [int(m.group(1)) for l in stats for m in [re.search(r"(\d+) scratch", l)] if m] + bad = int((np.fromfile(td / "out.bin", dtype=np.uint32) != want).sum()) + print(f"{arm}: {bad} mismatches of {n}, scratch {scratch[0] if scratch else 'unknown'}") + failed |= r.returncode != 0 or bad != 0 or (on and scratch != [0]) +sys.exit(1 if failed else 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..7561bdb95d28 --- /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/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..379bc223e8dd --- /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/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 diff --git a/tests/test_ulp_bound.py b/tests/test_ulp_bound.py new file mode 100644 index 000000000000..8bce1713ad82 --- /dev/null +++ b/tests/test_ulp_bound.py @@ -0,0 +1,62 @@ +"""CPU controls for the numerical GPU oracle, without starting Vulkan.""" +import ast +from pathlib import Path +import unittest + +import numpy as np + + +# ulp_bound currently dispatches a GPU job at module scope. Extract only its +# pure checks so these controls need neither a GPU nor a Vulkan driver. +source = Path(__file__).with_name("ulp_bound.py") +tree = ast.parse(source.read_text()) +functions = ast.Module(body=[n for n in tree.body + if isinstance(n, ast.FunctionDef) + and n.name in ("quotient", "check")], type_ignores=[]) +namespace = {"np": np} +exec(compile(functions, str(source), "exec"), namespace) + + +class ULPControls(unittest.TestCase): + def measure(self, reference, result): + namespace["inp"] = np.array([[reference, 0x3f800000]], dtype=np.uint32) + namespace["out"] = np.array([[result]], dtype=np.uint32) + return namespace["check"](0, namespace["inp"][:, 0], + namespace["inp"][:, 1]) + + def test_positive_one_ulp(self): + r = self.measure(0x3f800000, 0x3f800001) + self.assertEqual(r["max_ulp"], 1.0) + self.assertEqual(r["over_2p5"], 0) + + def test_positive_over_budget(self): + self.assertEqual(self.measure(0x3f800000, 0x3f800003)["over_2p5"], 1) + + def test_negative_errors_are_counted(self): + r = self.measure(0xbf800000, 0xbf800003) + self.assertEqual(r["max_ulp"], 3.0) + self.assertEqual(r["over_2p5"], 1) + + def test_negative_power_of_two_boundary(self): + self.assertEqual(self.measure(0xbf800000, 0xbf7fffff)["max_ulp"], 0.5) + + def test_nan_and_infinity_cannot_match_finite(self): + for bits in (0x7fc00000, 0x7f800000, 0xff800000): + with self.subTest(bits=hex(bits)): + r = self.measure(0x3f800000, bits) + self.assertEqual(r["over_2p5"], 1) + self.assertEqual(r["nonfinite_mismatch"], 1) + + def test_largest_finite_has_finite_ulp(self): + r = self.measure(0x7f7fffff, 0x7f7ffffe) + self.assertEqual(r["max_ulp"], 1.0) + self.assertEqual(r["over_2p5"], 0) + + def test_specials_and_signed_zero(self): + for bits in (0, 0x80000000, 0x7f800000, 0xff800000, 0x7fc00000): + self.assertEqual(self.measure(bits, bits)["special_mismatch"], 0) + self.assertEqual(self.measure(0x80000000, 0)["special_mismatch"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ulp_bound.py b/tests/ulp_bound.py new file mode 100644 index 000000000000..bc557972e2d3 --- /dev/null +++ b/tests/ulp_bound.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +"""Usage: ulp_bound.py ICD COMPUTE_RUNNER; glslangValidator and numpy are required. + +Default-path division contract: every result within 2.5 ULP of the correctly +rounded (with FTZ) quotient, the OpFDiv budget of the Vulkan SPIR-V +environment. Two sets: the exact 262,144-pair precise_math.py regression +input (same seed and construction) and a 66,048-pair boundary set straddling +the 2^-126 normal limit (the M2 Pro failure class). +Specials (+-0, inf, NaN, flushed results) must match the oracle bitwise. +The correctly rounded contract itself lives in precise_math.py (nir_fp_exact +path); this test only pins the default. +""" +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 + +# Regression set: identical construction to precise_math.py (seed 20260913). +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] + +# Boundary set: normal denominators in [1, 2^18), quotients +# (1 + j*2^-24)*2^-126 straddling the normal limit, j in [-64, 64], +# four sign combinations. a = RN(q * d) in fp32. +rngb = np.random.default_rng(20260913) +nb = 66048 +db = ((rngb.integers(127, 145, nb, dtype=np.uint32) << 23) | + rngb.integers(0, 1 << 23, nb, dtype=np.uint32)) +j = rngb.integers(-64, 65, nb).astype(np.float64) +q64 = (1.0 + j * 2.0**-24) * 2.0**-126 +ab = (q64 * db.view(np.float32).astype(np.float64)).astype(np.float32).view(np.uint32) +ab ^= rngb.integers(0, 2, nb, dtype=np.uint32) << 31 +db ^= rngb.integers(0, 2, nb, dtype=np.uint32) << 31 +bnd = np.stack([ab, db, np.full(nb, 0x3f317218, dtype=np.uint32), + np.full(nb, 0x3f800000, dtype=np.uint32)], axis=1) +inp = np.concatenate([inp, bnd]) +ntot = len(inp) + +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 + +with tempfile.TemporaryDirectory(prefix="mesa-ulp-") 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) + env = dict(os.environ, VK_DRIVER_FILES=icd, VK_ICD_FILENAMES=icd, + MESA_SHADER_CACHE_DISABLE="true") + subprocess.run(["flock", "-w", "30", "/tmp/gpu.lock", runner, + str(td / "probe.spv"), str(td / "in.bin"), str(td / "out.bin"), + str(ntot * 6 * 4), str(ntot // 64)], env=env, check=True, + timeout=120) + out = np.fromfile(td / "out.bin", dtype=np.uint32).reshape(ntot, 6) + +def check(col, a_bits, b_bits): + got = out[:, col] + want = quotient(a_bits, b_bits) + gf = got.copy().view(np.float32).astype(np.float64) + wf = want.copy().view(np.float32).astype(np.float64) + both_nan = ((got & 0x7fffffff) > 0x7f800000) & ((want & 0x7fffffff) > 0x7f800000) + special = ~np.isfinite(wf) | (wf == 0.0) + spec_bad = np.flatnonzero(special & (got != want) & ~both_nan) + normal = np.flatnonzero(~special) + finite_result = np.isfinite(gf[normal]) + with np.errstate(all="ignore"): + # A binary32 ULP is positive for both signs. Derive its exponent + # before measuring in binary64; spacing(float64) measures the wrong + # format and spacing(negative) would hide errors. ldexp also handles + # FLT_MAX without overflowing to infinity like float32 spacing does. + exponent = np.frexp(np.abs(wf[normal]))[1] + ulp = np.ldexp(np.ones(len(normal)), np.maximum(exponent - 24, -149)) + dist = np.abs(gf[normal] - wf[normal]) / ulp + # NaN comparisons are false, so reject non-finite results explicitly when + # the oracle expected a finite value. + over = normal[(dist > 2.5) | ~finite_result] + finite_dist = dist[finite_result] + return {"max_ulp": float(finite_dist.max()) if len(finite_dist) else 0.0, + "over_2p5": len(over), + "nonfinite_mismatch": int(np.count_nonzero(~finite_result)), + "special_mismatch": len(spec_bad), + "examples": [[hex(int(v)) for v in inp[i][:2]] for i in over[:5]][:5]} + +result = {"pairs": ntot, + "div": check(0, inp[:, 0], inp[:, 1]), + "recip": check(1, np.full(ntot, 0x3f800000, dtype=np.uint32), inp[:, 1])} +print(json.dumps(result, sort_keys=True)) +sys.exit(any(r["over_2p5"] or r["special_mismatch"] for r in result.values() if isinstance(r, dict)))