nir/lower_io: Add a new buffer_array_length intrinsic and lowering

Reviewed-by: Kristian H. Kristensen <[email protected]>
Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>

2 files changed, 45 insertions, 0 deletions

diff --git a/src/compiler/nir/nir_intrinsics.py b/src/compiler/nir/nir_intrinsics.py
index a6c74dc2543..d88e4ef7d45 100644
--- a/src/compiler/nir/nir_intrinsics.py
+++ b/src/compiler/nir/nir_intrinsics.py
@@ -156,6 +156,10 @@ intrinsic("interp_deref_at_sample", src_comp=[1, 1], dest_comp=0,
 intrinsic("interp_deref_at_offset", src_comp=[1, 2], dest_comp=0,
           flags=[CAN_ELIMINATE, CAN_REORDER])
 
+# Gets the length of an unsized array at the end of a buffer
+intrinsic("deref_buffer_array_length", src_comp=[-1], dest_comp=1,
+          flags=[CAN_ELIMINATE, CAN_REORDER])
+
 # Ask the driver for the size of a given buffer. It takes the buffer index
 # as source.
 intrinsic("get_buffer_size", src_comp=[-1], dest_comp=1,
diff --git a/src/compiler/nir/nir_lower_io.c b/src/compiler/nir/nir_lower_io.c
index c9db2de0ba0..786f2951286 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -940,6 +940,38 @@ lower_explicit_io_access(nir_builder *b, nir_intrinsic_instr *intrin,
    nir_instr_remove(&intrin->instr);
 }
 
+static void
+lower_explicit_io_array_length(nir_builder *b, nir_intrinsic_instr *intrin,
+                               nir_address_format addr_format)
+{
+   b->cursor = nir_after_instr(&intrin->instr);
+
+   nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+
+   assert(glsl_type_is_array(deref->type));
+   assert(glsl_get_length(deref->type) == 0);
+   unsigned stride = glsl_get_explicit_stride(deref->type);
+   assert(stride > 0);
+
+   assert(addr_format == nir_address_format_32bit_index_offset);
+   nir_ssa_def *addr = &deref->dest.ssa;
+   nir_ssa_def *index = addr_to_index(b, addr, addr_format);
+   nir_ssa_def *offset = addr_to_offset(b, addr, addr_format);
+
+   nir_intrinsic_instr *bsize =
+      nir_intrinsic_instr_create(b->shader, nir_intrinsic_get_buffer_size);
+   bsize->src[0] = nir_src_for_ssa(index);
+   nir_ssa_dest_init(&bsize->instr, &bsize->dest, 1, 32, NULL);
+   nir_builder_instr_insert(b, &bsize->instr);
+
+   nir_ssa_def *arr_size =
+      nir_idiv(b, nir_isub(b, &bsize->dest.ssa, offset),
+                  nir_imm_int(b, stride));
+
+   nir_ssa_def_rewrite_uses(&intrin->dest.ssa, nir_src_for_ssa(arr_size));
+   nir_instr_remove(&intrin->instr);
+}
+
 static bool
 nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes,
                            nir_address_format addr_format)
@@ -992,6 +1024,15 @@ nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes,
                break;
             }
 
+            case nir_intrinsic_deref_buffer_array_length: {
+               nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+               if (deref->mode & modes) {
+                  lower_explicit_io_array_length(&b, intrin, addr_format);
+                  progress = true;
+               }
+               break;
+            }
+
             default:
                break;
             }