nir: Add a ptr_as_array deref type

These correspond directly to SPIR-V's OpPtrAccessChain.  As such, they
treat whatever their parent gives them as if it's the first element in
some array and dereferences that array.  If the parent is, itself, an
array deref, then the two indices can just be added together to get the
final array deref.  However, it can also be used in cases where what you
have is a dereference to some random vec2 value somewhere.  In this
case, we require a cast before the ptr_as_array and use the ptr_stride
field in the cast to provide a stride for the ptr_as_array derefs.

Reviewed-by: Alejandro Piñeiro <apinheiro@igalia.com>
Reviewed-by: Caio Marcelo de Oliveira Filho <caio.oliveira@intel.com>

src/compiler/nir/nir.c

@@ -460,7 +460,8 @@ nir_deref_instr_create(nir_shader *shader, nir_deref_type deref_type)
    if (deref_type != nir_deref_type_var)
       src_init(&instr->parent);
 
-   if (deref_type == nir_deref_type_array)
+   if (deref_type == nir_deref_type_array ||
+       deref_type == nir_deref_type_ptr_as_array)
       src_init(&instr->arr.index);
 
    dest_init(&instr->dest);
@@ -1067,7 +1068,8 @@ visit_deref_instr_src(nir_deref_instr *instr,
          return false;
    }
 
-   if (instr->deref_type == nir_deref_type_array) {
+   if (instr->deref_type == nir_deref_type_array ||
+       instr->deref_type == nir_deref_type_ptr_as_array) {
       if (!visit_src(&instr->arr.index, cb, state))
          return false;
    }

src/compiler/nir/nir.h

@@ -998,6 +998,7 @@ typedef enum {
    nir_deref_type_var,
    nir_deref_type_array,
    nir_deref_type_array_wildcard,
+   nir_deref_type_ptr_as_array,
    nir_deref_type_struct,
    nir_deref_type_cast,
 } nir_deref_type;
@@ -1031,6 +1032,10 @@ typedef struct {
       struct {
          unsigned index;
       } strct;
+
+      struct {
+         unsigned ptr_stride;
+      } cast;
    };
 
    /** Destination to store the resulting "pointer" */
@@ -1078,6 +1083,8 @@ bool nir_deref_instr_has_indirect(nir_deref_instr *instr);
 
 bool nir_deref_instr_remove_if_unused(nir_deref_instr *instr);
 
+unsigned nir_deref_instr_ptr_as_array_stride(nir_deref_instr *instr);
+
 typedef struct {
    nir_instr instr;
 

src/compiler/nir/nir_builder.h

@@ -811,6 +811,31 @@ nir_build_deref_array(nir_builder *build, nir_deref_instr *parent,
    return deref;
 }
 
+static inline nir_deref_instr *
+nir_build_deref_ptr_as_array(nir_builder *build, nir_deref_instr *parent,
+                             nir_ssa_def *index)
+{
+   assert(parent->deref_type == nir_deref_type_array ||
+          parent->deref_type == nir_deref_type_ptr_as_array ||
+          parent->deref_type == nir_deref_type_cast);
+
+   nir_deref_instr *deref =
+      nir_deref_instr_create(build->shader, nir_deref_type_ptr_as_array);
+
+   deref->mode = parent->mode;
+   deref->type = parent->type;
+   deref->parent = nir_src_for_ssa(&parent->dest.ssa);
+   deref->arr.index = nir_src_for_ssa(index);
+
+   nir_ssa_dest_init(&deref->instr, &deref->dest,
+                     parent->dest.ssa.num_components,
+                     parent->dest.ssa.bit_size, NULL);
+
+   nir_builder_instr_insert(build, &deref->instr);
+
+   return deref;
+}
+
 static inline nir_deref_instr *
 nir_build_deref_array_wildcard(nir_builder *build, nir_deref_instr *parent)
 {
@@ -858,7 +883,8 @@ nir_build_deref_struct(nir_builder *build, nir_deref_instr *parent,
 
 static inline nir_deref_instr *
 nir_build_deref_cast(nir_builder *build, nir_ssa_def *parent,
-                     nir_variable_mode mode, const struct glsl_type *type)
+                     nir_variable_mode mode, const struct glsl_type *type,
+                     unsigned ptr_stride)
 {
    nir_deref_instr *deref =
       nir_deref_instr_create(build->shader, nir_deref_type_cast);
@@ -866,6 +892,7 @@ nir_build_deref_cast(nir_builder *build, nir_ssa_def *parent,
    deref->mode = mode;
    deref->type = type;
    deref->parent = nir_src_for_ssa(parent);
+   deref->cast.ptr_stride = ptr_stride;
 
    nir_ssa_dest_init(&deref->instr, &deref->dest,
                      parent->num_components, parent->bit_size, NULL);

src/compiler/nir/nir_clone.c

@@ -311,15 +311,19 @@ clone_deref_instr(clone_state *state, const nir_deref_instr *deref)
       break;
 
    case nir_deref_type_array:
+   case nir_deref_type_ptr_as_array:
       __clone_src(state, &nderef->instr,
                   &nderef->arr.index, &deref->arr.index);
       break;
 
    case nir_deref_type_array_wildcard:
-   case nir_deref_type_cast:
       /* Nothing to do */
       break;
 
+   case nir_deref_type_cast:
+      nderef->cast.ptr_stride = deref->cast.ptr_stride;
+      break;
+
    default:
       unreachable("Invalid instruction deref type");
    }

src/compiler/nir/nir_deref.c

@@ -111,7 +111,8 @@ nir_deref_instr_has_indirect(nir_deref_instr *instr)
       if (instr->deref_type == nir_deref_type_cast)
          return true;
 
-      if (instr->deref_type == nir_deref_type_array &&
+      if ((instr->deref_type == nir_deref_type_array ||
+           instr->deref_type == nir_deref_type_ptr_as_array) &&
           !nir_src_is_const(instr->arr.index))
          return true;
 
@@ -121,6 +122,23 @@ nir_deref_instr_has_indirect(nir_deref_instr *instr)
    return false;
 }
 
+unsigned
+nir_deref_instr_ptr_as_array_stride(nir_deref_instr *deref)
+{
+   assert(deref->deref_type == nir_deref_type_ptr_as_array);
+   nir_deref_instr *parent = nir_deref_instr_parent(deref);
+   switch (parent->deref_type) {
+   case nir_deref_type_array:
+      return glsl_get_explicit_stride(nir_deref_instr_parent(parent)->type);
+   case nir_deref_type_ptr_as_array:
+      return nir_deref_instr_ptr_as_array_stride(parent);
+   case nir_deref_type_cast:
+      return parent->cast.ptr_stride;
+   default:
+      unreachable("Invalid parent for ptr_as_array deref");
+   }
+}
+
 static unsigned
 type_get_array_stride(const struct glsl_type *elem_type,
                       glsl_type_size_align_func size_align)

src/compiler/nir/nir_instr_set.c

@@ -96,12 +96,16 @@ hash_deref(uint32_t hash, const nir_deref_instr *instr)
       break;
 
    case nir_deref_type_array:
+   case nir_deref_type_ptr_as_array:
       hash = hash_src(hash, &instr->arr.index);
       break;
 
+   case nir_deref_type_cast:
+      hash = HASH(hash, instr->cast.ptr_stride);
+      break;
+
    case nir_deref_type_var:
    case nir_deref_type_array_wildcard:
-   case nir_deref_type_cast:
       /* Nothing to do */
       break;
 
@@ -351,13 +355,18 @@ nir_instrs_equal(const nir_instr *instr1, const nir_instr *instr2)
          break;
 
       case nir_deref_type_array:
+      case nir_deref_type_ptr_as_array:
          if (!nir_srcs_equal(deref1->arr.index, deref2->arr.index))
             return false;
          break;
 
+      case nir_deref_type_cast:
+         if (deref1->cast.ptr_stride != deref2->cast.ptr_stride)
+            return false;
+         break;
+
       case nir_deref_type_var:
       case nir_deref_type_array_wildcard:
-      case nir_deref_type_cast:
          /* Nothing to do */
          break;
 

src/compiler/nir/nir_opt_copy_propagate.c

@@ -224,7 +224,8 @@ copy_prop_instr(nir_instr *instr)
             progress = true;
       }
 
-      if (deref->deref_type == nir_deref_type_array) {
+      if (deref->deref_type == nir_deref_type_array ||
+          deref->deref_type == nir_deref_type_ptr_as_array) {
          while (copy_prop_src(&deref->arr.index, instr, NULL, 1))
             progress = true;
       }

src/compiler/nir/nir_print.c

@@ -636,7 +636,8 @@ print_deref_link(const nir_deref_instr *instr, bool whole_chain, print_state *st
               glsl_get_struct_elem_name(parent->type, instr->strct.index));
       break;
 
-   case nir_deref_type_array: {
+   case nir_deref_type_array:
+   case nir_deref_type_ptr_as_array: {
       nir_const_value *const_index = nir_src_as_const_value(instr->arr.index);
       if (const_index) {
          fprintf(fp, "[%u]", const_index->u32[0]);
@@ -678,6 +679,9 @@ print_deref_instr(nir_deref_instr *instr, print_state *state)
    case nir_deref_type_cast:
       fprintf(fp, " = deref_cast ");
       break;
+   case nir_deref_type_ptr_as_array:
+      fprintf(fp, " = deref_ptr_as_array ");
+      break;
    default:
       unreachable("Invalid deref instruction type");
    }

src/compiler/nir/nir_serialize.c

@@ -438,11 +438,15 @@ write_deref(write_ctx *ctx, const nir_deref_instr *deref)
       break;
 
    case nir_deref_type_array:
+   case nir_deref_type_ptr_as_array:
       write_src(ctx, &deref->arr.index);
       break;
 
-   case nir_deref_type_array_wildcard:
    case nir_deref_type_cast:
+      blob_write_uint32(ctx->blob, deref->cast.ptr_stride);
+      break;
+
+   case nir_deref_type_array_wildcard:
       /* Nothing to do */
       break;
 
@@ -475,11 +479,15 @@ read_deref(read_ctx *ctx)
       break;
 
    case nir_deref_type_array:
+   case nir_deref_type_ptr_as_array:
       read_src(ctx, &deref->arr.index, &deref->instr);
       break;
 
-   case nir_deref_type_array_wildcard:
    case nir_deref_type_cast:
+      deref->cast.ptr_stride = blob_read_uint32(ctx->blob);
+      break;
+
+   case nir_deref_type_array_wildcard:
       /* Nothing to do */
       break;
 

src/compiler/nir/nir_validate.c

@@ -470,6 +470,19 @@ validate_deref_instr(nir_deref_instr *instr, validate_state *state)
          }
          break;
 
+      case nir_deref_type_ptr_as_array:
+         /* ptr_as_array derefs must have a parent that is either an array,
+          * ptr_as_array, or cast.  If the parent is a cast, we get the stride
+          * information (if any) from the cast deref.
+          */
+         validate_assert(state,
+                         parent->deref_type == nir_deref_type_array ||
+                         parent->deref_type == nir_deref_type_ptr_as_array ||
+                         parent->deref_type == nir_deref_type_cast);
+         validate_src(&instr->arr.index, state,
+                      nir_dest_bit_size(instr->dest), 1);
+         break;
+
       default:
          unreachable("Invalid deref instruction type");
       }

src/compiler/spirv/vtn_cfg.c

@@ -884,7 +884,7 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
                glsl_get_bare_type(b->func->type->return_type->type);
             nir_deref_instr *ret_deref =
                nir_build_deref_cast(&b->nb, nir_load_param(&b->nb, 0),
-                                    nir_var_local, ret_type);
+                                    nir_var_local, ret_type, 0);
             vtn_local_store(b, src, ret_deref);
          }
 

src/compiler/spirv/vtn_variables.c

@@ -1627,7 +1627,7 @@ vtn_pointer_from_ssa(struct vtn_builder *b, nir_ssa_def *ssa,
       assert(!vtn_pointer_is_external_block(b, ptr));
       const struct glsl_type *deref_type = ptr_type->deref->type;
       ptr->deref = nir_build_deref_cast(&b->nb, ssa, nir_mode,
-                                        glsl_get_bare_type(deref_type));
+                                        glsl_get_bare_type(deref_type), 0);
    }
 
    return ptr;