nir/lower_ssbo: Extend the load_ssbo_address intrinsic to pass an offset

On Mali(Valhall), the bounds checking can be done when in hardware, but
for this to work properly, we need to pass the offset to the
nir_load_ssbo_address() intrinsic.

Add an offset source to the intrinsic, and adjust the lowering pass
to conditionally lower the offset addition.

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

@@ -6640,6 +6640,7 @@ bool nir_dedup_inline_samplers(nir_shader *shader);
 
 typedef struct nir_lower_ssbo_options {
    bool native_loads;
+   bool native_offset;
 } nir_lower_ssbo_options;
 
 bool nir_lower_ssbo(nir_shader *shader, const nir_lower_ssbo_options *opts);

src/compiler/nir/nir_intrinsics.py

@@ -1162,8 +1162,8 @@ load("per_primitive_input", [1], [BASE, COMPONENT, DEST_TYPE, IO_SEMANTICS], [CA
 
 # src[] = { buffer_index, offset }.
 load("ssbo", [-1, 1], [ACCESS, ALIGN_MUL, ALIGN_OFFSET], [CAN_ELIMINATE])
-# src[] = { buffer_index }
-load("ssbo_address", [1], [], [CAN_ELIMINATE, CAN_REORDER])
+# src[] = { buffer_index, offset }
+load("ssbo_address", [1, 1], [], [CAN_ELIMINATE, CAN_REORDER])
 # src[] = { offset }.
 load("output", [1], [BASE, RANGE, COMPONENT, DEST_TYPE, IO_SEMANTICS], flags=[CAN_ELIMINATE])
 # src[] = { vertex, offset }.

src/compiler/nir/nir_lower_ssbo.c

@@ -13,12 +13,20 @@
  */
 
 static nir_def *
-calc_address(nir_builder *b, nir_intrinsic_instr *intr)
+calc_address(nir_builder *b, nir_intrinsic_instr *intr,
+             const nir_lower_ssbo_options *opts)
 {
    unsigned index_src = intr->intrinsic == nir_intrinsic_store_ssbo ? 1 : 0;
-   nir_def *base = nir_load_ssbo_address(b, 1, 64, intr->src[index_src].ssa);
+   bool lower_offset = !opts || !opts->native_offset;
+   nir_def *offset = nir_get_io_offset_src(intr)->ssa;
+   nir_def *addr =
+      nir_load_ssbo_address(b, 1, 64, intr->src[index_src].ssa,
+                            lower_offset ? nir_imm_int(b, 0) : offset);
 
-   return nir_iadd(b, base, nir_u2u64(b, nir_get_io_offset_src(intr)->ssa));
+   if (lower_offset)
+      addr = nir_iadd(b, addr, nir_u2u64(b, offset));
+
+   return addr;
 }
 
 static bool
@@ -35,26 +43,30 @@ pass(nir_builder *b, nir_intrinsic_instr *intr, void *data)
          return false;
 
       def = nir_build_load_global(b, intr->def.num_components,
-                                  intr->def.bit_size, calc_address(b, intr),
+                                  intr->def.bit_size,
+                                  calc_address(b, intr, opts),
                                   .align_mul = nir_intrinsic_align_mul(intr),
                                   .align_offset = nir_intrinsic_align_offset(intr));
       break;
 
    case nir_intrinsic_store_ssbo:
-      nir_build_store_global(b, intr->src[0].ssa, calc_address(b, intr),
+      nir_build_store_global(b, intr->src[0].ssa,
+                             calc_address(b, intr, opts),
                              .align_mul = nir_intrinsic_align_mul(intr),
                              .align_offset = nir_intrinsic_align_offset(intr),
                              .write_mask = nir_intrinsic_write_mask(intr));
       break;
 
    case nir_intrinsic_ssbo_atomic:
-      def = nir_global_atomic(b, intr->def.bit_size, calc_address(b, intr),
+      def = nir_global_atomic(b, intr->def.bit_size,
+                              calc_address(b, intr, opts),
                               intr->src[2].ssa,
                               .atomic_op = nir_intrinsic_atomic_op(intr));
       break;
 
    case nir_intrinsic_ssbo_atomic_swap:
-      def = nir_global_atomic_swap(b, intr->def.bit_size, calc_address(b, intr),
+      def = nir_global_atomic_swap(b, intr->def.bit_size,
+                                   calc_address(b, intr, opts),
                                    intr->src[2].ssa, intr->src[3].ssa,
                                    .atomic_op = nir_intrinsic_atomic_op(intr));
       break;

src/gallium/drivers/asahi/agx_nir_lower_sysvals.c

@@ -170,6 +170,7 @@ lower_intrinsic(nir_builder *b, nir_intrinsic_instr *intr,
       return load_sysval_root(
          b, 1, 64, &u->pipeline_statistics[nir_intrinsic_base(intr)]);
    case nir_intrinsic_load_ssbo_address:
+      assert(nir_src_as_uint(intr->src[1]) == 0);
       return load_sysval_indirect(b, 1, 64, stage_table(b), &s->ssbo_base,
                                   intr->src[0].ssa);
    case nir_intrinsic_get_ubo_size:

src/gallium/drivers/panfrost/pan_nir_lower_sysvals.c

@@ -58,6 +58,7 @@ sysval_for_intrinsic(nir_intrinsic_instr *intr, unsigned *offset)
 {
    switch (intr->intrinsic) {
    case nir_intrinsic_load_ssbo_address:
+      assert(nir_src_as_uint(intr->src[1]) == 0);
       return PAN_SYSVAL(SSBO, nir_src_as_uint(intr->src[0]));
    case nir_intrinsic_get_ssbo_size:
       *offset = 8;

src/gallium/drivers/radeonsi/si_nir_lower_resource.c

@@ -299,6 +299,7 @@ static bool lower_resource_intrinsic(nir_builder *b, nir_intrinsic_instr *intrin
       break;
    }
    case nir_intrinsic_load_ssbo_address: {
+      assert(nir_src_as_uint(intrin->src[1]) == 0);
       nir_def *desc = load_ssbo_desc(b, &intrin->src[0], s);
       nir_def *lo = nir_channel(b, desc, 0);
       nir_def *hi = nir_i2i32(b, nir_u2u16(b, nir_channel(b, desc, 1)));