agx: vectorize SSBOs

this was missed due to the lowering, and mitigates a lot of stats weirdness with
the address mode rework.

total instructions in shared programs: 2755170 -> 2751399 (-0.14%)
instructions in affected programs: 16323 -> 12552 (-23.10%)
helped: 71
HURT: 0
helped stats (abs) min: 10 max: 178 x̄: 53.11 x̃: 42
helped stats (rel) min: 2.04% max: 50.00% x̄: 34.73% x̃: 40.79%
95% mean confidence interval for instructions value: -60.94 -45.28
95% mean confidence interval for instructions %-change: -37.81% -31.65%
Instructions are helped.

total alu in shared programs: 2169888 -> 2168281 (-0.07%)
alu in affected programs: 9547 -> 7940 (-16.83%)
helped: 71
HURT: 0
helped stats (abs) min: 5 max: 90 x̄: 22.63 x̃: 16
helped stats (rel) min: 1.02% max: 43.33% x̄: 25.39% x̃: 29.41%
95% mean confidence interval for alu value: -26.33 -18.93
95% mean confidence interval for alu %-change: -27.91% -22.87%
Alu are helped.

total fscib in shared programs: 2165597 -> 2163990 (-0.07%)
fscib in affected programs: 9547 -> 7940 (-16.83%)
helped: 71
HURT: 0
helped stats (abs) min: 5 max: 90 x̄: 22.63 x̃: 16
helped stats (rel) min: 1.02% max: 43.33% x̄: 25.39% x̃: 29.41%
95% mean confidence interval for fscib value: -26.33 -18.93
95% mean confidence interval for fscib %-change: -27.91% -22.87%
Fscib are helped.

total bytes in shared programs: 21517750 -> 21489352 (-0.13%)
bytes in affected programs: 126270 -> 97872 (-22.49%)
helped: 71
HURT: 0
helped stats (abs) min: 80 max: 1084 x̄: 399.97 x̃: 324
helped stats (rel) min: 1.77% max: 50.57% x̄: 35.07% x̃: 42.31%
95% mean confidence interval for bytes value: -455.66 -344.28
95% mean confidence interval for bytes %-change: -38.34% -31.79%
Bytes are helped.

total regs in shared programs: 864490 -> 865162 (0.08%)
regs in affected programs: 4567 -> 5239 (14.71%)
helped: 4
HURT: 61
helped stats (abs) min: 6 max: 6 x̄: 6.00 x̃: 6
helped stats (rel) min: 4.51% max: 5.13% x̄: 4.82% x̃: 4.82%
HURT stats (abs)   min: 2 max: 24 x̄: 11.41 x̃: 12
HURT stats (rel)   min: 1.98% max: 82.35% x̄: 21.05% x̃: 16.00%
95% mean confidence interval for regs value: 8.52 12.16
95% mean confidence interval for regs %-change: 14.91% 24.00%
Regs are HURT.

total threads in shared programs: 27613056 -> 27613312 (<.01%)
threads in affected programs: 3200 -> 3456 (8.00%)
helped: 4
HURT: 0
helped stats (abs) min: 64 max: 64 x̄: 64.00 x̃: 64
helped stats (rel) min: 7.69% max: 8.33% x̄: 8.01% x̃: 8.01%
95% mean confidence interval for threads value: 64.00 64.00
95% mean confidence interval for threads %-change: 7.42% 8.60%
Threads are helped.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31964>

src/asahi/compiler/agx_compile.c

@@ -2830,11 +2830,11 @@ agx_optimize_loop_nir(nir_shader *nir)
    } while (progress);
 }
 
-static bool
-mem_vectorize_cb(unsigned align_mul, unsigned align_offset, unsigned bit_size,
-                 unsigned num_components, unsigned hole_size,
-                 nir_intrinsic_instr *low, nir_intrinsic_instr *high,
-                 void *data)
+bool
+agx_mem_vectorize_cb(unsigned align_mul, unsigned align_offset,
+                     unsigned bit_size, unsigned num_components,
+                     unsigned hole_size, nir_intrinsic_instr *low,
+                     nir_intrinsic_instr *high, void *data)
 {
    if (hole_size)
       return false;
@@ -3003,7 +3003,7 @@ agx_optimize_nir(nir_shader *nir, bool soft_fault, unsigned *preamble_size)
    NIR_PASS(_, nir, nir_opt_load_store_vectorize,
             &(const nir_load_store_vectorize_options){
                .modes = nir_var_mem_global | nir_var_mem_constant,
-               .callback = mem_vectorize_cb,
+               .callback = agx_mem_vectorize_cb,
             });
    NIR_PASS(_, nir, nir_lower_pack);
 

src/asahi/compiler/agx_compile.h

@@ -263,6 +263,10 @@ bool agx_nir_lower_interpolation(nir_shader *s);
 bool agx_nir_lower_cull_distance_vs(struct nir_shader *s);
 bool agx_nir_lower_cull_distance_fs(struct nir_shader *s,
                                     unsigned nr_distances);
+bool agx_mem_vectorize_cb(unsigned align_mul, unsigned align_offset,
+                          unsigned bit_size, unsigned num_components,
+                          unsigned hole_size, nir_intrinsic_instr *low,
+                          nir_intrinsic_instr *high, void *data);
 
 void agx_compile_shader_nir(nir_shader *nir, struct agx_shader_key *key,
                             struct util_debug_callback *debug,

src/gallium/drivers/asahi/agx_state.c

@@ -1868,6 +1868,19 @@ agx_shader_initialize(struct agx_device *dev, struct agx_uncompiled_shader *so,
       so->info.cull_distance_size = nir->info.cull_distance_array_size;
    }
 
+   /* Vectorize SSBOs before lowering them, since it is significantly harder to
+    * vectorize the lowered code.
+    */
+   NIR_PASS(_, nir, nir_lower_load_const_to_scalar);
+   NIR_PASS(_, nir, nir_copy_prop);
+
+   NIR_PASS(
+      _, nir, nir_opt_load_store_vectorize,
+      &(const nir_load_store_vectorize_options){
+         .modes = nir_var_mem_global | nir_var_mem_constant | nir_var_mem_ssbo,
+         .callback = agx_mem_vectorize_cb,
+      });
+
    NIR_PASS(_, nir, agx_nir_lower_texture, true);
    NIR_PASS(_, nir, nir_lower_ssbo, NULL);