honeykrisp: M1 compute speedups and the deferred CDM barrier on G13G (stacked on #3, #4, #5)

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,

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 {

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) ||

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);

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)

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.