nir/lower_ssbo: Add an option to conditionally lower loads

On Mali(Valhall), we have a way to load SSBO data without going through
an SSBO index -> global address translation, so let's provide a way
to tell nir_lower_ssbo() when it shouldn't lower loads.

Signed-off-by: Boris Brezillon <boris.brezillon@collabora.com>
Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Acked-by: Eric R. Smith <eric.smith@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31164>

src/compiler/nir/nir.h

@@ -6637,7 +6637,13 @@ bool nir_rematerialize_derefs_in_use_blocks_impl(nir_function_impl *impl);
 bool nir_lower_samplers(nir_shader *shader);
 bool nir_lower_cl_images(nir_shader *shader, bool lower_image_derefs, bool lower_sampler_derefs);
 bool nir_dedup_inline_samplers(nir_shader *shader);
-bool nir_lower_ssbo(nir_shader *shader);
+
+typedef struct nir_lower_ssbo_options {
+   bool native_loads;
+} nir_lower_ssbo_options;
+
+bool nir_lower_ssbo(nir_shader *shader, const nir_lower_ssbo_options *opts);
+
 bool nir_lower_helper_writes(nir_shader *shader, bool lower_plain_stores);
 
 typedef struct nir_lower_printf_options {

src/compiler/nir/nir_lower_ssbo.c

@@ -22,13 +22,18 @@ calc_address(nir_builder *b, nir_intrinsic_instr *intr)
 }
 
 static bool
-pass(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *data)
+pass(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 {
+   const nir_lower_ssbo_options *opts = data;
+
    b->cursor = nir_before_instr(&intr->instr);
 
    nir_def *def = NULL;
    switch (intr->intrinsic) {
    case nir_intrinsic_load_ssbo:
+      if (opts && opts->native_loads)
+         return false;
+
       def = nir_build_load_global(b, intr->def.num_components,
                                   intr->def.bit_size, calc_address(b, intr),
                                   .align_mul = nir_intrinsic_align_mul(intr),
@@ -66,9 +71,9 @@ pass(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *data)
 }
 
 bool
-nir_lower_ssbo(nir_shader *shader)
+nir_lower_ssbo(nir_shader *shader, const nir_lower_ssbo_options *opts)
 {
    return nir_shader_intrinsics_pass(shader, pass,
                                      nir_metadata_control_flow,
-                                     NULL);
+                                     (void *)opts);
 }

src/gallium/drivers/asahi/agx_state.c

@@ -1862,7 +1862,7 @@ agx_shader_initialize(struct agx_device *dev, struct agx_uncompiled_shader *so,
    }
 
    NIR_PASS(_, nir, agx_nir_lower_texture);
-   NIR_PASS(_, nir, nir_lower_ssbo);
+   NIR_PASS(_, nir, nir_lower_ssbo, NULL);
 
    agx_preprocess_nir(nir, dev->libagx);
 

src/panfrost/compiler/bifrost_compile.c

@@ -5085,7 +5085,7 @@ bifrost_preprocess_nir(nir_shader *nir, unsigned gpu_id)
               bi_lower_load_push_const_with_dyn_offset,
               nir_metadata_control_flow, NULL);
 
-   NIR_PASS_V(nir, nir_lower_ssbo);
+   NIR_PASS_V(nir, nir_lower_ssbo, NULL);
    NIR_PASS_V(nir, pan_lower_sample_pos);
    NIR_PASS_V(nir, nir_lower_bit_size, bi_lower_bit_size, NULL);
    NIR_PASS_V(nir, nir_lower_64bit_phis);

src/panfrost/midgard/midgard_compile.c

@@ -359,7 +359,7 @@ midgard_preprocess_nir(nir_shader *nir, unsigned gpu_id)
       NIR_PASS_V(nir, pan_nir_lower_store_component);
    }
 
-   NIR_PASS_V(nir, nir_lower_ssbo);
+   NIR_PASS_V(nir, nir_lower_ssbo, NULL);
    NIR_PASS_V(nir, pan_nir_lower_zs_store);
 
    NIR_PASS_V(nir, nir_lower_frexp);