vectorize descriptor loads

src/amd/vulkan/radv_shader.c

@@ -199,6 +199,15 @@ radv_optimize_nir(struct nir_shader *shader, bool optimize_conservatively, bool
    NIR_PASS(progress, shader, nir_opt_move, nir_move_load_ubo);
 }
 
+static bool
+smem_vectorize_callback(unsigned align_mul, unsigned align_offset, unsigned bit_size,
+                        unsigned num_components, nir_intrinsic_instr *low,
+                        nir_intrinsic_instr *high, void *data)
+{
+   return low->intrinsic == nir_intrinsic_load_smem_amd && num_components <= 16 && bit_size == 32 &&
+          align_mul % 4 == 0 && align_offset % 4 == 0;
+}
+
 void
 radv_optimize_nir_algebraic(nir_shader *nir, bool opt_offsets)
 {
@@ -212,6 +221,16 @@ radv_optimize_nir_algebraic(nir_shader *nir, bool opt_offsets)
       NIR_PASS(more_algebraic, nir, nir_opt_algebraic);
    }
 
+   /* vectorize descriptor loads */ // TODO: skip for radv_lower_ngg()
+   NIR_PASS_V(nir, nir_opt_cse);
+   nir_load_store_vectorize_options smem_vectorize_opts = {
+      .modes = nir_var_mem_global,
+      .callback = smem_vectorize_callback,
+      .robust_modes = 0,
+   };
+   if (nir_opt_load_store_vectorize(nir, &smem_vectorize_opts))
+      NIR_PASS_V(nir, nir_lower_alu_to_scalar, NULL, NULL);
+
    if (opt_offsets) {
       static const nir_opt_offsets_options offset_options = {
          .uniform_max = 0,

src/compiler/nir/nir_opt_load_store_vectorize.c

@@ -82,6 +82,7 @@ case nir_intrinsic_##op: {\
    STORE(nir_var_mem_shared, shared, -1, 1, -1, 0)
    LOAD(nir_var_mem_global, global, -1, 0, -1)
    STORE(nir_var_mem_global, global, -1, 1, -1, 0)
+   LOAD(nir_var_mem_global, smem_amd, 0, 1, -1)
    ATOMIC(nir_var_mem_ssbo, ssbo, add, 0, 1, -1, 2)
    ATOMIC(nir_var_mem_ssbo, ssbo, imin, 0, 1, -1, 2)
    ATOMIC(nir_var_mem_ssbo, ssbo, umin, 0, 1, -1, 2)