spirv: Handle any bit size in vector_insert/extract

This crops up both in the actual SPIR-V VectorInsert/Extract opcodes as
well as various places where we deal with vector derefs.

Cc: [email protected]
Reviewed-by: Alejandro Piñeiro <[email protected]>
Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>

3 files changed, 15 insertions, 11 deletions

diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index 2433d8638cc..341a1cb17ce 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -2913,13 +2913,19 @@ vtn_vector_insert(struct vtn_builder *b, nir_ssa_def *src, nir_ssa_def *insert,
    return &vec->dest.dest.ssa;
 }
 
+static nir_ssa_def *
+nir_ieq_imm(nir_builder *b, nir_ssa_def *x, uint64_t i)
+{
+   return nir_ieq(b, x, nir_imm_intN_t(b, i, x->bit_size));
+}
+
 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(&b->nb, index, nir_imm_int(&b->nb, i)),
+      dest = nir_bcsel(&b->nb, nir_ieq_imm(&b->nb, index, i),
                        vtn_vector_extract(b, src, i), dest);
 
    return dest;
@@ -2931,7 +2937,7 @@ vtn_vector_insert_dynamic(struct vtn_builder *b, nir_ssa_def *src,
 {
    nir_ssa_def *dest = vtn_vector_insert(b, src, insert, 0);
    for (unsigned i = 1; i < src->num_components; i++)
-      dest = nir_bcsel(&b->nb, nir_ieq(&b->nb, index, nir_imm_int(&b->nb, i)),
+      dest = nir_bcsel(&b->nb, nir_ieq_imm(&b->nb, index, i),
                        vtn_vector_insert(b, src, insert, i), dest);
 
    return dest;
diff --git a/src/compiler/spirv/vtn_glsl450.c b/src/compiler/spirv/vtn_glsl450.c
index bca4889d772..04807700b72 100644
--- a/src/compiler/spirv/vtn_glsl450.c
+++ b/src/compiler/spirv/vtn_glsl450.c
@@ -882,10 +882,9 @@ handle_glsl450_interpolation(struct vtn_builder *b, enum GLSLstd450 opcode,
 
    if (vec_array_deref) {
       assert(vec_deref);
-      nir_const_value *const_index = nir_src_as_const_value(vec_deref->arr.index);
-      if (const_index) {
+      if (nir_src_is_const(vec_deref->arr.index)) {
          val->ssa->def = vtn_vector_extract(b, &intrin->dest.ssa,
-                                            const_index->u32[0]);
+                                            nir_src_as_uint(vec_deref->arr.index));
       } else {
          val->ssa->def = vtn_vector_extract_dynamic(b, &intrin->dest.ssa,
                                                     vec_deref->arr.index.ssa);
diff --git a/src/compiler/spirv/vtn_variables.c b/src/compiler/spirv/vtn_variables.c
index 55721fc36e3..025c6a23734 100644
--- a/src/compiler/spirv/vtn_variables.c
+++ b/src/compiler/spirv/vtn_variables.c
@@ -511,9 +511,9 @@ vtn_local_load(struct vtn_builder *b, nir_deref_instr *src)
 
    if (src_tail != src) {
       val->type = src->type;
-      nir_const_value *const_index = nir_src_as_const_value(src->arr.index);
-      if (const_index)
-         val->def = vtn_vector_extract(b, val->def, const_index->u32[0]);
+      if (nir_src_is_const(src->arr.index))
+         val->def = vtn_vector_extract(b, val->def,
+                                       nir_src_as_uint(src->arr.index));
       else
          val->def = vtn_vector_extract_dynamic(b, val->def, src->arr.index.ssa);
    }
@@ -531,10 +531,9 @@ vtn_local_store(struct vtn_builder *b, struct vtn_ssa_value *src,
       struct vtn_ssa_value *val = vtn_create_ssa_value(b, dest_tail->type);
       _vtn_local_load_store(b, true, dest_tail, val);
 
-      nir_const_value *const_index = nir_src_as_const_value(dest->arr.index);
-      if (const_index)
+      if (nir_src_is_const(dest->arr.index))
          val->def = vtn_vector_insert(b, val->def, src->def,
-                                      const_index->u32[0]);
+                                      nir_src_as_uint(dest->arr.index));
       else
          val->def = vtn_vector_insert_dynamic(b, val->def, src->def,
                                               dest->arr.index.ssa);