honeykrisp: gate cooperative matrices and the tracked CDM barrier to G13 (stacked on #3 and #4)

src/asahi/compiler/agx_compile.c

@@ -1697,6 +1697,24 @@ agx_emit_intrinsic(agx_builder *b, nir_intrinsic_instr *instr)
       return agx_quad_ballot_to(b, dst, agx_src_index(&instr->src[0]));
    }
 
+   case nir_intrinsic_simd_matrix_fmadd_agx: {
+      /* Apple G13 hardware 8x8x8 matrix MAC: D = A*B + C, SIMD-group-wide.
+       * Each operand is this lane's 2 fragment elements: a 32-bit register
+       * pair for fp32, one 32-bit register (two halves) for fp16. The
+       * accumulator precision picks the opcode; A/B carry their own size
+       * bits, so fmadd32 with fp16 A/B is the mixed-precision form.
+       */
+      agx_index a = agx_src_index(&instr->src[0]);
+      agx_index bm = agx_src_index(&instr->src[1]);
+      agx_index c = agx_src_index(&instr->src[2]);
+      if (instr->def.bit_size == 16)
+         agx_simd_matrix_fmadd16_to(b, dst, a, bm, c);
+      else
+         agx_simd_matrix_fmadd32_to(b, dst, a, bm, c);
+      agx_emit_cached_split(b, dst, 2);
+      return NULL;
+   }
+
    case nir_intrinsic_doorbell_agx: {
       return agx_doorbell(b, nir_src_as_uint(instr->src[0]));
    }
@@ -2020,6 +2038,24 @@ agx_emit_alu(agx_builder *b, nir_alu_instr *instr)
       return agx_bfeil_to(b, dst, i0, s0, s1, m);
    }
 
+   case nir_op_ubfe: {
+      /* SM5 semantics: offset and width are both taken modulo 32. The width is
+       * a constant here because agx_nir_fuse_algebraic_late is the only thing
+       * that forms ubfe (has_bfe is not advertised, so NIR never generates
+       * it), and that pass masks the offset for the same reason
+       * ubitfield_extract needed it: bfeil does not wrap on the hardware.
+       */
+      unsigned m = nir_alu_src_as_uint(instr->src[2]) & 0x1F;
+
+      /* Unlike ubitfield_extract, a width of zero reads as zero here rather
+       * than meaning "all 32 bits", so it cannot share the masking above.
+       */
+      if (m == 0)
+         return agx_mov_imm_to(b, dst, 0);
+
+      return agx_bfeil_to(b, dst, i0, s0, s1, m);
+   }
+
    case nir_op_bcsel:
       return agx_icmpsel_to(b, dst, s0, i0, s2, s1, AGX_ICOND_UEQ);
 
@@ -2723,46 +2759,6 @@ agx_calc_stats(agx_context *ctx, unsigned size, struct agx2_stats *stats)
    }
 }
 
-static bool
-agx_lower_sincos_filter(const nir_instr *instr, UNUSED const void *_)
-{
-   if (instr->type != nir_instr_type_alu)
-      return false;
-
-   nir_alu_instr *alu = nir_instr_as_alu(instr);
-   return alu->op == nir_op_fsin || alu->op == nir_op_fcos;
-}
-
-/* Sine and cosine are implemented via the sin_pt_1 and sin_pt_2 opcodes for
- * heavy lifting. sin_pt_2 implements sinc in the first quadrant, expressed in
- * turns (sin (tau x) / x), while sin_pt_1 implements a piecewise sign/offset
- * fixup to transform a quadrant angle [0, 4] to [-1, 1]. The NIR opcode
- * fsin_agx models the fixup, sinc, and multiply to obtain sine, so we just
- * need to change units from radians to quadrants modulo turns. Cosine is
- * implemented by shifting by one quadrant: cos(x) = sin(x + tau/4).
- */
-
-static nir_def *
-agx_lower_sincos_impl(struct nir_builder *b, nir_instr *instr, UNUSED void *_)
-{
-   nir_alu_instr *alu = nir_instr_as_alu(instr);
-   nir_def *x = nir_mov_alu(b, alu->src[0], 1);
-   nir_def *turns = nir_fmul_imm(b, x, M_1_PI * 0.5f);
-
-   if (alu->op == nir_op_fcos)
-      turns = nir_fadd_imm(b, turns, 0.25f);
-
-   nir_def *quadrants = nir_fmul_imm(b, nir_ffract(b, turns), 4.0);
-   return nir_fsin_agx(b, quadrants);
-}
-
-static bool
-agx_lower_sincos(nir_shader *shader)
-{
-   return nir_shader_lower_instructions(shader, agx_lower_sincos_filter,
-                                        agx_lower_sincos_impl, NULL);
-}
-
 static bool
 agx_lower_front_face(struct nir_builder *b, nir_intrinsic_instr *intr,
                      UNUSED void *data)
@@ -2953,6 +2949,11 @@ agx_optimize_nir(nir_shader *nir, bool soft_fault, uint16_t *preamble_size,
       } while (progress);
    }
 
+   /* int64 lowering and the perspective divide emit fdiv after preprocessing;
+    * the backend has no fdiv, so lower what is left before it is scheduled.
+    */
+   NIR_PASS(_, nir, agx_nir_lower_fdiv_late);
+
    /* Reassociate before forming preambles because it makes preambles more
     * effective. Clean up after.
     */
@@ -3533,63 +3534,36 @@ 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)
 {
    if (!nir)
       return;
 
+   /* Lower VK_KHR_cooperative_matrix before vars_to_scratch can spill the cmat
+    * temps. The G13 matrix tile spans all 32 lanes of a subgroup, so the HW
+    * path requires fully-populated subgroups: compute stage, static workgroup
+    * size, size a multiple of 32. Everything else lowers through the software
+    * path (agx_nir_lower_cmat.c) -- a valid shader is never rejected. On by
+    * default; AGX_SIMDMAT=0 disables it.
+    */
+   if (agx_simdmat_enabled()) {
+      bool full_subgroups =
+         nir->info.stage == MESA_SHADER_COMPUTE &&
+         !nir->info.workgroup_size_variable &&
+         nir->info.workgroup_size[0] && nir->info.workgroup_size[1] &&
+         nir->info.workgroup_size[2] &&
+         (((uint64_t)nir->info.workgroup_size[0] *
+              nir->info.workgroup_size[1] * nir->info.workgroup_size[2]) %
+          32) == 0;
+
+      if (full_subgroups)
+         NIR_PASS(_, nir, agx_nir_lower_simdmat, 32);
+      else
+         NIR_PASS(_, nir, agx_nir_lower_cmat);
+   }
+
    NIR_PASS(_, nir, nir_lower_vars_to_ssa);
 
    /* Lower large arrays to scratch and small arrays to csel */
@@ -3614,7 +3588,7 @@ agx_preprocess_nir(nir_shader *nir)
    NIR_PASS(_, nir, nir_lower_alu);
    NIR_PASS(_, nir, nir_lower_load_const_to_scalar);
    NIR_PASS(_, nir, nir_lower_flrp, 16 | 32 | 64, false);
-   NIR_PASS(_, nir, agx_lower_sincos);
+   NIR_PASS(_, nir, agx_nir_lower_sincos);
    NIR_PASS(_, nir, nir_shader_intrinsics_pass, agx_lower_front_face,
             nir_metadata_control_flow, NULL);
    NIR_PASS(_, nir, agx_nir_lower_subgroups);
@@ -3637,9 +3611,11 @@ agx_preprocess_nir(nir_shader *nir)
 
    NIR_PASS(_, nir, nir_lower_idiv, &idiv_options);
 
-   /* Has to run after nir_lower_idiv */
-   NIR_PASS(_, nir, nir_shader_alu_pass, agx_nir_lower_fdiv,
-            nir_metadata_control_flow, NULL);
+   /* Has to run after nir_lower_idiv, and before the log lowering so the
+    * reciprocals it builds are not refined twice.
+    */
+   NIR_PASS(_, nir, agx_nir_lower_fdiv);
+   NIR_PASS(_, nir, agx_nir_lower_log);
 
    NIR_PASS(_, nir, nir_opt_deref);
    NIR_PASS(_, nir, nir_lower_vars_to_ssa);

src/asahi/compiler/agx_compile.h

@@ -7,6 +7,7 @@
 
 #include "compiler/nir/nir.h"
 #include "util/shader_stats.h"
+#include "util/u_debug.h"
 #include "util/u_dynarray.h"
 #include "util/u_tristate.h"
 #include "shader_enums.h"
@@ -317,6 +318,15 @@ struct agx_shader_key {
 uint64_t agx_gather_texcoords(nir_shader *nir);
 
 void agx_preprocess_nir(nir_shader *nir);
+
+/* VK_KHR_cooperative_matrix on the G13 simd_matrix unit, with a software
+ * fallback for partial subgroups. On by default; AGX_SIMDMAT=0 disables it.
+ */
+static inline bool
+agx_simdmat_enabled(void)
+{
+   return debug_get_bool_option("AGX_SIMDMAT", true);
+}
 bool agx_nir_lower_discard_zs_emit(nir_shader *s);
 bool agx_nir_lower_sample_mask(nir_shader *s);
 bool agx_nir_lower_interpolation(nir_shader *s);
@@ -347,7 +357,6 @@ agx_round_registers(unsigned halfregs)
 }
 
 static const nir_shader_compiler_options agx_nir_options = {
-   .lower_fdiv = true,
    .float_mul_add16 = nir_float_muladd_support_has_ffma | nir_float_muladd_support_fuse,
    .float_mul_add32 = nir_float_muladd_support_has_ffma | nir_float_muladd_support_fuse,
    .lower_flrp16 = true,

src/asahi/compiler/agx_compiler.h

@@ -1096,6 +1096,8 @@ bool agx_nir_opt_preamble(nir_shader *s, unsigned *sizes);
 bool agx_nir_lower_load_mask(nir_shader *shader);
 bool agx_nir_lower_ubo(nir_shader *shader);
 bool agx_nir_lower_shared_bitsize(nir_shader *shader);
+bool agx_nir_lower_simdmat(nir_shader *shader, unsigned subgroup_size);
+bool agx_nir_lower_cmat(nir_shader *shader);
 bool agx_nir_lower_frag_sidefx(nir_shader *s);
 
 struct agx_cycle_estimate {

src/asahi/compiler/agx_nir.h

@@ -22,6 +22,10 @@ bool agx_nir_lower_layer(struct nir_shader *s);
 bool agx_nir_lower_clip_distance(struct nir_shader *s);
 bool agx_nir_lower_subgroups(struct nir_shader *s);
 bool agx_nir_lower_fminmax(struct nir_shader *s);
+bool agx_nir_lower_fdiv(struct nir_shader *s);
+bool agx_nir_lower_fdiv_late(struct nir_shader *s);
+bool agx_nir_lower_log(struct nir_shader *s);
+bool agx_nir_lower_sincos(struct nir_shader *s);
 
 bool agx_nir_lower_texture_early(struct nir_shader *s, bool support_lod_bias);
 bool agx_nir_lower_texture(struct nir_shader *s);

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)