nir/spirv: Add support for SPV_KHR_variable_pointers

Reviewed-by: Iago Toral Quiroga <[email protected]>

5 files changed, 108 insertions, 10 deletions

diff --git a/src/compiler/spirv/nir_spirv.h b/src/compiler/spirv/nir_spirv.h
index 7f16866b493..83577fb5d23 100644
--- a/src/compiler/spirv/nir_spirv.h
+++ b/src/compiler/spirv/nir_spirv.h
@@ -51,6 +51,7 @@ struct nir_spirv_supported_extensions {
    bool image_write_without_format;
    bool int64;
    bool multiview;
+   bool variable_pointers;
 };
 
 nir_function *spirv_to_nir(const uint32_t *words, size_t word_count,
diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index 6e35f83a421..4b9c1218cf2 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -185,6 +185,13 @@ vtn_ssa_value(struct vtn_builder *b, uint32_t value_id)
    case vtn_value_type_ssa:
       return val->ssa;
 
+   case vtn_value_type_pointer:
+      assert(val->pointer->ptr_type && val->pointer->ptr_type->type);
+      struct vtn_ssa_value *ssa =
+         vtn_create_ssa_value(b, val->pointer->ptr_type->type);
+      ssa->def = vtn_pointer_to_ssa(b, val->pointer);
+      return ssa;
+
    default:
       unreachable("Invalid type for an SSA value");
    }
@@ -861,9 +868,16 @@ vtn_handle_type(struct vtn_builder *b, SpvOp opcode,
          vtn_value(b, w[3], vtn_value_type_type)->type;
 
       val->type->base_type = vtn_base_type_pointer;
-      val->type->type = NULL;
       val->type->storage_class = storage_class;
       val->type->deref = deref_type;
+
+      if (storage_class == SpvStorageClassUniform ||
+          storage_class == SpvStorageClassStorageBuffer) {
+         /* These can actually be stored to nir_variables and used as SSA
+          * values so they need a real glsl_type.
+          */
+         val->type->type = glsl_vector_type(GLSL_TYPE_UINT, 2);
+      }
       break;
    }
 
@@ -1387,7 +1401,8 @@ vtn_handle_function_call(struct vtn_builder *b, SpvOp opcode,
    for (unsigned i = 0; i < call->num_params; i++) {
       unsigned arg_id = w[4 + i];
       struct vtn_value *arg = vtn_untyped_value(b, arg_id);
-      if (arg->value_type == vtn_value_type_pointer) {
+      if (arg->value_type == vtn_value_type_pointer &&
+          arg->pointer->ptr_type->type == NULL) {
          nir_deref_var *d = vtn_pointer_to_deref(b, arg->pointer);
          call->params[i] = nir_deref_var_clone(d, call);
       } else {
@@ -2769,6 +2784,11 @@ vtn_handle_preamble_instruction(struct vtn_builder *b, SpvOp opcode,
          spv_check_supported(multiview, cap);
          break;
 
+      case SpvCapabilityVariablePointersStorageBuffer:
+      case SpvCapabilityVariablePointers:
+         spv_check_supported(variable_pointers, cap);
+         break;
+
       default:
          unreachable("Unhandled capability");
       }
@@ -3153,6 +3173,19 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
       break;
    }
 
+   case SpvOpSelect: {
+      /* Handle OpSelect up-front here because it needs to be able to handle
+       * pointers and not just regular vectors and scalars.
+       */
+      struct vtn_type *res_type = vtn_value(b, w[1], vtn_value_type_type)->type;
+      struct vtn_ssa_value *ssa = vtn_create_ssa_value(b, res_type->type);
+      ssa->def = nir_bcsel(&b->nb, vtn_ssa_value(b, w[3])->def,
+                                   vtn_ssa_value(b, w[4])->def,
+                                   vtn_ssa_value(b, w[5])->def);
+      vtn_push_ssa(b, w[2], res_type, ssa);
+      break;
+   }
+
    case SpvOpSNegate:
    case SpvOpFNegate:
    case SpvOpNot:
@@ -3210,7 +3243,6 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
    case SpvOpBitwiseOr:
    case SpvOpBitwiseXor:
    case SpvOpBitwiseAnd:
-   case SpvOpSelect:
    case SpvOpIEqual:
    case SpvOpFOrdEqual:
    case SpvOpFUnordEqual:
diff --git a/src/compiler/spirv/vtn_cfg.c b/src/compiler/spirv/vtn_cfg.c
index c81a62d7a8e..03c452cb31a 100644
--- a/src/compiler/spirv/vtn_cfg.c
+++ b/src/compiler/spirv/vtn_cfg.c
@@ -52,7 +52,8 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
       func->num_params = func_type->length;
       func->params = ralloc_array(b->shader, nir_parameter, func->num_params);
       for (unsigned i = 0; i < func->num_params; i++) {
-         if (func_type->params[i]->base_type == vtn_base_type_pointer) {
+         if (func_type->params[i]->base_type == vtn_base_type_pointer &&
+             func_type->params[i]->type == NULL) {
             func->params[i].type = func_type->params[i]->deref->type;
          } else {
             func->params[i].type = func_type->params[i]->type;
@@ -82,7 +83,7 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
       assert(b->func_param_idx < b->func->impl->num_params);
       nir_variable *param = b->func->impl->params[b->func_param_idx++];
 
-      if (type->base_type == vtn_base_type_pointer) {
+      if (type->base_type == vtn_base_type_pointer && type->type == NULL) {
          struct vtn_variable *vtn_var = rzalloc(b, struct vtn_variable);
          vtn_var->type = type->deref;
          vtn_var->var = param;
diff --git a/src/compiler/spirv/vtn_private.h b/src/compiler/spirv/vtn_private.h
index 8d745bb4a03..84584620fc1 100644
--- a/src/compiler/spirv/vtn_private.h
+++ b/src/compiler/spirv/vtn_private.h
@@ -369,6 +369,13 @@ struct vtn_pointer {
    struct nir_ssa_def *offset;
 };
 
+static inline bool
+vtn_pointer_uses_ssa_offset(struct vtn_pointer *ptr)
+{
+   return ptr->mode == vtn_variable_mode_ubo ||
+          ptr->mode == vtn_variable_mode_ssbo;
+}
+
 struct vtn_variable {
    enum vtn_variable_mode mode;
 
@@ -501,6 +508,12 @@ struct vtn_builder {
    bool has_loop_continue;
 };
 
+nir_ssa_def *
+vtn_pointer_to_ssa(struct vtn_builder *b, struct vtn_pointer *ptr);
+struct vtn_pointer *
+vtn_pointer_from_ssa(struct vtn_builder *b, nir_ssa_def *ssa,
+                     struct vtn_type *ptr_type);
+
 static inline struct vtn_value *
 vtn_push_value(struct vtn_builder *b, uint32_t value_id,
                enum vtn_value_type value_type)
@@ -517,8 +530,14 @@ static inline struct vtn_value *
 vtn_push_ssa(struct vtn_builder *b, uint32_t value_id,
              struct vtn_type *type, struct vtn_ssa_value *ssa)
 {
-   struct vtn_value *val = vtn_push_value(b, value_id, vtn_value_type_ssa);
-   val->ssa = ssa;
+   struct vtn_value *val;
+   if (type->base_type == vtn_base_type_pointer) {
+      val = vtn_push_value(b, value_id, vtn_value_type_pointer);
+      val->pointer = vtn_pointer_from_ssa(b, ssa->def, type);
+   } else {
+      val = vtn_push_value(b, value_id, vtn_value_type_ssa);
+      val->ssa = ssa;
+   }
    return val;
 }
 
diff --git a/src/compiler/spirv/vtn_variables.c b/src/compiler/spirv/vtn_variables.c
index a9e2dbfaa04..4432e72e54a 100644
--- a/src/compiler/spirv/vtn_variables.c
+++ b/src/compiler/spirv/vtn_variables.c
@@ -223,8 +223,7 @@ vtn_pointer_dereference(struct vtn_builder *b,
                         struct vtn_pointer *base,
                         struct vtn_access_chain *deref_chain)
 {
-   if (base->mode == vtn_variable_mode_ubo ||
-       base->mode == vtn_variable_mode_ssbo) {
+   if (vtn_pointer_uses_ssa_offset(base)) {
       return vtn_ssa_offset_pointer_dereference(b, base, deref_chain);
    } else {
       return vtn_access_chain_pointer_dereference(b, base, deref_chain);
@@ -1478,6 +1477,53 @@ vtn_storage_class_to_mode(SpvStorageClass class,
    return mode;
 }
 
+nir_ssa_def *
+vtn_pointer_to_ssa(struct vtn_builder *b, struct vtn_pointer *ptr)
+{
+   /* This pointer needs to have a pointer type with actual storage */
+   assert(ptr->ptr_type);
+   assert(ptr->ptr_type->type);
+
+   if (ptr->offset && ptr->block_index) {
+      return nir_vec2(&b->nb, ptr->block_index, ptr->offset);
+   } else {
+      /* If we don't have an offset or block index, then we must be a pointer
+       * to the variable itself.
+       */
+      assert(!ptr->offset && !ptr->block_index);
+
+      /* We can't handle a pointer to an array of descriptors because we have
+       * no way of knowing later on that we need to add to update the block
+       * index when dereferencing.
+       */
+      assert(ptr->var && ptr->var->type->base_type == vtn_base_type_struct);
+
+      return nir_vec2(&b->nb, vtn_variable_resource_index(b, ptr->var, NULL),
+                              nir_imm_int(&b->nb, 0));
+   }
+}
+
+struct vtn_pointer *
+vtn_pointer_from_ssa(struct vtn_builder *b, nir_ssa_def *ssa,
+                     struct vtn_type *ptr_type)
+{
+   assert(ssa->num_components == 2 && ssa->bit_size == 32);
+   assert(ptr_type->base_type == vtn_base_type_pointer);
+   assert(ptr_type->deref->base_type != vtn_base_type_pointer);
+   /* This pointer type needs to have actual storage */
+   assert(ptr_type->type);
+
+   struct vtn_pointer *ptr = rzalloc(b, struct vtn_pointer);
+   ptr->mode = vtn_storage_class_to_mode(ptr_type->storage_class,
+                                         ptr_type, NULL);
+   ptr->type = ptr_type->deref;
+   ptr->ptr_type = ptr_type;
+   ptr->block_index = nir_channel(&b->nb, ssa, 0);
+   ptr->offset = nir_channel(&b->nb, ssa, 1);
+
+   return ptr;
+}
+
 static bool
 is_per_vertex_inout(const struct vtn_variable *var, gl_shader_stage stage)
 {
@@ -1503,7 +1549,6 @@ vtn_create_variable(struct vtn_builder *b, struct vtn_value *val,
 {
    assert(ptr_type->base_type == vtn_base_type_pointer);
    struct vtn_type *type = ptr_type->deref;
-   assert(type->base_type != vtn_base_type_pointer);
 
    struct vtn_type *without_array = type;
    while(glsl_type_is_array(without_array->type))