nir/builder: Add a vector extract helper

This one's a tiny bit better than what we had in spirv_to_nir because it
emits a binary tree rather than a linear walk.  It also doesn't leave
around unneeded bcsel instructions for a constant index and returns an
undef for constant OOB access.

Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>

2 files changed, 30 insertions, 6 deletions

diff --git a/src/compiler/nir/nir_builder.h b/src/compiler/nir/nir_builder.h
index 4b4d1eef462..9662cd2a217 100644
--- a/src/compiler/nir/nir_builder.h
+++ b/src/compiler/nir/nir_builder.h
@@ -570,6 +570,35 @@ nir_channels(nir_builder *b, nir_ssa_def *def, nir_component_mask_t mask)
 }
 
 static inline nir_ssa_def *
+_nir_vector_extract_helper(nir_builder *b, nir_ssa_def *vec, nir_ssa_def *c,
+                           unsigned start, unsigned end)
+{
+   if (start == end - 1) {
+      return nir_channel(b, vec, start);
+   } else {
+      unsigned mid = start + (end - start) / 2;
+      return nir_bcsel(b, nir_ilt(b, c, nir_imm_int(b, mid)),
+                       _nir_vector_extract_helper(b, vec, c, start, mid),
+                       _nir_vector_extract_helper(b, vec, c, mid, end));
+   }
+}
+
+static inline nir_ssa_def *
+nir_vector_extract(nir_builder *b, nir_ssa_def *vec, nir_ssa_def *c)
+{
+   nir_src c_src = nir_src_for_ssa(c);
+   if (nir_src_is_const(c_src)) {
+      unsigned c_const = nir_src_as_uint(c_src);
+      if (c_const < vec->num_components)
+         return nir_channel(b, vec, c_const);
+      else
+         return nir_ssa_undef(b, 1, vec->bit_size);
+   } else {
+      return _nir_vector_extract_helper(b, vec, c, 0, vec->num_components);
+   }
+}
+
+static inline nir_ssa_def *
 nir_i2i(nir_builder *build, nir_ssa_def *x, unsigned dest_bit_size)
 {
    if (x->bit_size == dest_bit_size)
diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index da0f2dcd805..3b7f8d8a9ac 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -3047,12 +3047,7 @@ nir_ssa_def *
 vtn_vector_extract_dynamic(struct vtn_builder *b, nir_ssa_def *src,
                            nir_ssa_def *index)
 {
-   nir_ssa_def *dest = vtn_vector_extract(b, src, 0);
-   for (unsigned i = 1; i < src->num_components; i++)
-      dest = nir_bcsel(&b->nb, nir_ieq_imm(&b->nb, index, i),
-                       vtn_vector_extract(b, src, i), dest);
-
-   return dest;
+   return nir_vector_extract(&b->nb, src, nir_i2i(&b->nb, index, 32));
 }
 
 nir_ssa_def *