nir/vulkan: Add a descriptor type to vulkan resource intrinsics

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

4 files changed, 56 insertions, 5 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 553410b92d1..8db7cc44aef 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1278,6 +1278,11 @@ typedef enum {
    NIR_INTRINSIC_ALIGN_MUL = 17,
    NIR_INTRINSIC_ALIGN_OFFSET = 18,
 
+   /**
+    * The Vulkan descriptor type for a vulkan_resource_[re]index intrinsic.
+    */
+   NIR_INTRINSIC_DESC_TYPE = 19,
+
    NIR_INTRINSIC_NUM_INDEX_FLAGS,
 
 } nir_intrinsic_index_flag;
@@ -1378,6 +1383,7 @@ INTRINSIC_IDX_ACCESSORS(access, ACCESS, enum gl_access_qualifier)
 INTRINSIC_IDX_ACCESSORS(format, FORMAT, unsigned)
 INTRINSIC_IDX_ACCESSORS(align_mul, ALIGN_MUL, unsigned)
 INTRINSIC_IDX_ACCESSORS(align_offset, ALIGN_OFFSET, unsigned)
+INTRINSIC_IDX_ACCESSORS(desc_type, DESC_TYPE, unsigned)
 
 static inline void
 nir_intrinsic_set_align(nir_intrinsic_instr *intrin,
diff --git a/src/compiler/nir/nir_intrinsics.py b/src/compiler/nir/nir_intrinsics.py
index e449764492d..3735fde2915 100644
--- a/src/compiler/nir/nir_intrinsics.py
+++ b/src/compiler/nir/nir_intrinsics.py
@@ -112,6 +112,8 @@ FORMAT = "NIR_INTRINSIC_FORMAT"
 # Offset or address alignment
 ALIGN_MUL = "NIR_INTRINSIC_ALIGN_MUL"
 ALIGN_OFFSET = "NIR_INTRINSIC_ALIGN_OFFSET"
+# The vulkan descriptor type for vulkan_resource_index
+DESC_TYPE = "NIR_INTRINSIC_DESC_TYPE"
 
 #
 # Possible flags:
@@ -358,9 +360,10 @@ image("store_raw_intel", src_comp=[1, 0])
 # corresponds to the tuple (set, binding, index) and computes an index
 # corresponding to tuple (set, binding, idx + src1).
 intrinsic("vulkan_resource_index", src_comp=[1], dest_comp=1,
-          indices=[DESC_SET, BINDING], flags=[CAN_ELIMINATE, CAN_REORDER])
-intrinsic("vulkan_resource_reindex", src_comp=[1, 1], dest_comp=1,
+          indices=[DESC_SET, BINDING, DESC_TYPE],
           flags=[CAN_ELIMINATE, CAN_REORDER])
+intrinsic("vulkan_resource_reindex", src_comp=[1, 1], dest_comp=1,
+          indices=[DESC_TYPE], flags=[CAN_ELIMINATE, CAN_REORDER])
 
 # variable atomic intrinsics
 #
diff --git a/src/compiler/nir/nir_print.c b/src/compiler/nir/nir_print.c
index a73b9d37ae5..696839a69b1 100644
--- a/src/compiler/nir/nir_print.c
+++ b/src/compiler/nir/nir_print.c
@@ -28,6 +28,7 @@
 #include "nir.h"
 #include "compiler/shader_enums.h"
 #include "util/half_float.h"
+#include "vulkan/vulkan_core.h"
 #include <stdio.h>
 #include <stdlib.h>
 #include <inttypes.h> /* for PRIx64 macro */
@@ -705,6 +706,26 @@ print_deref_instr(nir_deref_instr *instr, print_state *state)
    }
 }
 
+static const char *
+vulkan_descriptor_type_name(VkDescriptorType type)
+{
+   switch (type) {
+   case VK_DESCRIPTOR_TYPE_SAMPLER: return "sampler";
+   case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER: return "texture+sampler";
+   case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE: return "texture";
+   case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE: return "image";
+   case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER: return "texture-buffer";
+   case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER: return "image-buffer";
+   case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER: return "UBO";
+   case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER: return "SSBO";
+   case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC: return "UBO";
+   case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC: return "SSBO";
+   case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT: return "input-att";
+   case VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT: return "inline-UBO";
+   default: return "unknown";
+   }
+}
+
 static void
 print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
 {
@@ -756,6 +777,7 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
       [NIR_INTRINSIC_FORMAT] = "format",
       [NIR_INTRINSIC_ALIGN_MUL] = "align_mul",
       [NIR_INTRINSIC_ALIGN_OFFSET] = "align_offset",
+      [NIR_INTRINSIC_DESC_TYPE] = "desc_type",
    };
    for (unsigned idx = 1; idx < NIR_INTRINSIC_NUM_INDEX_FLAGS; idx++) {
       if (!info->index_map[idx])
@@ -789,6 +811,9 @@ print_intrinsic_instr(nir_intrinsic_instr *instr, print_state *state)
       } else if (idx == NIR_INTRINSIC_IMAGE_ARRAY) {
          bool array = nir_intrinsic_image_dim(instr);
          fprintf(fp, " image_dim=%s", array ? "true" : "false");
+      } else if (idx == NIR_INTRINSIC_DESC_TYPE) {
+         VkDescriptorType desc_type = nir_intrinsic_desc_type(instr);
+         fprintf(fp, " desc_type=%s", vulkan_descriptor_type_name(desc_type));
       } else {
          unsigned off = info->index_map[idx] - 1;
          assert(index_name[idx]);  /* forgot to update index_name table? */
diff --git a/src/compiler/spirv/vtn_variables.c b/src/compiler/spirv/vtn_variables.c
index 3041fa47ec1..c177952d7e1 100644
--- a/src/compiler/spirv/vtn_variables.c
+++ b/src/compiler/spirv/vtn_variables.c
@@ -28,6 +28,7 @@
 #include "vtn_private.h"
 #include "spirv_info.h"
 #include "nir_deref.h"
+#include <vulkan/vulkan_core.h>
 
 static struct vtn_access_chain *
 vtn_access_chain_create(struct vtn_builder *b, unsigned length)
@@ -131,6 +132,19 @@ vtn_nir_deref_pointer_dereference(struct vtn_builder *b,
    return ptr;
 }
 
+static VkDescriptorType
+vk_desc_type_for_mode(struct vtn_builder *b, enum vtn_variable_mode mode)
+{
+   switch (mode) {
+   case vtn_variable_mode_ubo:
+      return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+   case vtn_variable_mode_ssbo:
+      return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+   default:
+      vtn_fail("Invalid mode for vulkan_resource_index");
+   }
+}
+
 static nir_ssa_def *
 vtn_variable_resource_index(struct vtn_builder *b, struct vtn_variable *var,
                             nir_ssa_def *desc_array_index)
@@ -146,6 +160,7 @@ vtn_variable_resource_index(struct vtn_builder *b, struct vtn_variable *var,
    instr->src[0] = nir_src_for_ssa(desc_array_index);
    nir_intrinsic_set_desc_set(instr, var->descriptor_set);
    nir_intrinsic_set_binding(instr, var->binding);
+   nir_intrinsic_set_desc_type(instr, vk_desc_type_for_mode(b, var->mode));
 
    nir_ssa_dest_init(&instr->instr, &instr->dest, 1, 32, NULL);
    nir_builder_instr_insert(&b->nb, &instr->instr);
@@ -154,14 +169,15 @@ vtn_variable_resource_index(struct vtn_builder *b, struct vtn_variable *var,
 }
 
 static nir_ssa_def *
-vtn_resource_reindex(struct vtn_builder *b, nir_ssa_def *base_index,
-                     nir_ssa_def *offset_index)
+vtn_resource_reindex(struct vtn_builder *b, enum vtn_variable_mode mode,
+                     nir_ssa_def *base_index, nir_ssa_def *offset_index)
 {
    nir_intrinsic_instr *instr =
       nir_intrinsic_instr_create(b->nb.shader,
                                  nir_intrinsic_vulkan_resource_reindex);
    instr->src[0] = nir_src_for_ssa(base_index);
    instr->src[1] = nir_src_for_ssa(offset_index);
+   nir_intrinsic_set_desc_type(instr, vk_desc_type_for_mode(b, mode));
 
    nir_ssa_dest_init(&instr->instr, &instr->dest, 1, 32, NULL);
    nir_builder_instr_insert(&b->nb, &instr->instr);
@@ -239,7 +255,8 @@ vtn_ssa_offset_pointer_dereference(struct vtn_builder *b,
             vtn_access_link_as_ssa(b, deref_chain->link[0], 1, 32);
          idx++;
 
-         block_index = vtn_resource_reindex(b, block_index, offset_index);
+         block_index = vtn_resource_reindex(b, base->mode,
+                                            block_index, offset_index);
       }
    }