[Mesa-dev] [PATCH 16/28] nir/linker: handle non-ubo uses of vulkan_resource_index

[Alejandro Piñeiro]

From: Neil Roberts <nroberts at igalia.com>

In order to replicate the behaviour of lower_ubo_reference_visitor,
the lowering code should search the list of blocks in
ShaderStorageBlocks for the matching binding whenever a non-ubo usage
of the resource index is encountered.

The intended usage of the vulkan_resource_index is determined by
searching for an intrinsic which uses the result. Unfortunately some
other lower passes can add instructions to perform arithmetic on the
result so the search needs to be performed recursively on the result
of those.

Signed-off-by: Neil Roberts <nroberts at igalia.com>
Signed-off-by: Alejandro Piñeiro <apinheiro at igalia.com>
---
 .../glsl/gl_nir_lower_vulkan_resource_index.c      | 55 +++++++++++++++++++---
 1 file changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/compiler/glsl/gl_nir_lower_vulkan_resource_index.c b/src/compiler/glsl/gl_nir_lower_vulkan_resource_index.c
index 92ee3dd707a..561d2a03de2 100644
--- a/src/compiler/glsl/gl_nir_lower_vulkan_resource_index.c
+++ b/src/compiler/glsl/gl_nir_lower_vulkan_resource_index.c
@@ -37,12 +37,10 @@
  */
 
 static unsigned
-find_block_by_binding(struct gl_linked_shader *linked_shader,
+find_block_by_binding(unsigned num_blocks,
+                      struct gl_uniform_block **blocks,
                       unsigned binding)
 {
-   unsigned num_blocks = linked_shader->Program->info.num_ubos;
-   struct gl_uniform_block **blocks = linked_shader->Program->sh.UniformBlocks;
-
    for (unsigned i = 0; i < num_blocks; i++) {
       if (blocks[i]->Binding == binding)
          return i;
@@ -51,6 +49,35 @@ find_block_by_binding(struct gl_linked_shader *linked_shader,
    unreachable("No block found with the given binding");
 }
 
+static bool
+find_intrinsic_usage(nir_ssa_def *def,
+                     bool *is_ubo_usage)
+{
+   nir_foreach_use_safe(use_src, def) {
+      if (use_src->parent_instr->type == nir_instr_type_alu) {
+         nir_alu_instr *alu = nir_instr_as_alu(use_src->parent_instr);
+
+         if (find_intrinsic_usage(&alu->dest.dest.ssa, is_ubo_usage))
+            return true;
+
+         continue;
+      }
+
+      if (use_src->parent_instr->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *intr = nir_instr_as_intrinsic(use_src->parent_instr);
+
+      if (intr == NULL)
+         continue;
+
+      *is_ubo_usage = intr->intrinsic == nir_intrinsic_load_ubo;
+      return true;
+   }
+
+   return false;
+}
+
 static bool
 convert_block(nir_block *block,
               struct gl_linked_shader *linked_shader,
@@ -67,13 +94,29 @@ convert_block(nir_block *block,
       if (res_index->intrinsic != nir_intrinsic_vulkan_resource_index)
          continue;
 
+      bool is_ubo_usage;
+      if (!find_intrinsic_usage(&res_index->dest.ssa, &is_ubo_usage))
+         continue;
+
       b->cursor = nir_after_instr(instr);
 
       /* The descriptor set should always be zero for GL */
       assert(nir_intrinsic_desc_set(res_index) == 0);
-
       unsigned binding = nir_intrinsic_binding(res_index);
-      unsigned block = find_block_by_binding(linked_shader, binding);
+
+      unsigned num_blocks;
+      struct gl_uniform_block **blocks;
+
+      if (is_ubo_usage) {
+         num_blocks = linked_shader->Program->info.num_ubos;
+         blocks = linked_shader->Program->sh.UniformBlocks;
+      } else {
+         num_blocks = linked_shader->Program->info.num_ssbos;
+         blocks = linked_shader->Program->sh.ShaderStorageBlocks;
+      }
+
+      unsigned block = find_block_by_binding(num_blocks, blocks, binding);
+
       nir_ssa_def *surface =
          nir_iadd(b,
                   nir_imm_int(b, block),
-- 
2.14.1