nir: add helpers for chasing resource bindings

Signed-off-by: Rhys Perry <pendingchaos02@gmail.com>
Reviewed-by: Daniel Schürmann <daniel@schuermann.dev>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7291>

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index 693d119..6e75b0b 100644 (file)
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -2517,3 +2517,107 @@ nir_get_shader_call_payload_src(nir_intrinsic_instr *call)
       return NULL;
    }
 }
+
+nir_binding nir_chase_binding(nir_src rsrc)
+{
+   nir_binding res = {0};
+   if (rsrc.ssa->parent_instr->type == nir_instr_type_deref) {
+      const struct glsl_type *type = glsl_without_array(nir_src_as_deref(rsrc)->type);
+      bool is_image = glsl_type_is_image(type) || glsl_type_is_sampler(type);
+      while (rsrc.ssa->parent_instr->type == nir_instr_type_deref) {
+         nir_deref_instr *deref = nir_src_as_deref(rsrc);
+
+         if (deref->deref_type == nir_deref_type_var) {
+            res.success = true;
+            res.var = deref->var;
+            res.desc_set = deref->var->data.descriptor_set;
+            res.binding = deref->var->data.binding;
+            return res;
+         } else if (deref->deref_type == nir_deref_type_array && is_image) {
+            if (res.num_indices == ARRAY_SIZE(res.indices))
+               return (nir_binding){0};
+            res.indices[res.num_indices++] = deref->arr.index;
+         }
+
+         rsrc = deref->parent;
+      }
+   }
+
+   /* Skip copies and trimming. Trimming can appear as nir_op_mov instructions
+    * when removing the offset from addresses. We also consider nir_op_is_vec()
+    * instructions to skip trimming of vec2_index_32bit_offset addresses after
+    * lowering ALU to scalar.
+    */
+   while (true) {
+      nir_alu_instr *alu = nir_src_as_alu_instr(rsrc);
+      nir_intrinsic_instr *intrin = nir_src_as_intrinsic(rsrc);
+      if (alu && alu->op == nir_op_mov) {
+         for (unsigned i = 0; i < alu->dest.dest.ssa.num_components; i++) {
+            if (alu->src[0].swizzle[i] != i)
+               return (nir_binding){0};
+         }
+         rsrc = alu->src[0].src;
+      } else if (alu && nir_op_is_vec(alu->op)) {
+         for (unsigned i = 0; i < nir_op_infos[alu->op].num_inputs; i++) {
+            if (alu->src[i].swizzle[0] != i || alu->src[i].src.ssa != alu->src[0].src.ssa)
+               return (nir_binding){0};
+         }
+         rsrc = alu->src[0].src;
+      } else if (intrin && intrin->intrinsic == nir_intrinsic_read_first_invocation) {
+         /* The caller might want to be aware if only the first invocation of
+          * the indices are used.
+          */
+         res.read_first_invocation = true;
+         rsrc = intrin->src[0];
+      } else {
+         break;
+      }
+   }
+
+   if (nir_src_is_const(rsrc)) {
+      /* GL binding model after deref lowering */
+      res.success = true;
+      res.binding = nir_src_as_uint(rsrc);
+      return res;
+   }
+
+   /* otherwise, must be Vulkan binding model after deref lowering or GL bindless */
+
+   nir_intrinsic_instr *intrin = nir_src_as_intrinsic(rsrc);
+   if (!intrin)
+      return (nir_binding){0};
+
+   /* skip load_vulkan_descriptor */
+   if (intrin->intrinsic == nir_intrinsic_load_vulkan_descriptor) {
+      intrin = nir_src_as_intrinsic(intrin->src[0]);
+      if (!intrin)
+         return (nir_binding){0};
+   }
+
+   if (intrin->intrinsic != nir_intrinsic_vulkan_resource_index)
+      return (nir_binding){0};
+
+   assert(res.num_indices == 0);
+   res.success = true;
+   res.desc_set = nir_intrinsic_desc_set(intrin);
+   res.binding = nir_intrinsic_binding(intrin);
+   res.num_indices = 1;
+   res.indices[0] = intrin->src[0];
+   return res;
+}
+
+nir_variable *nir_get_binding_variable(nir_shader *shader, nir_binding binding)
+{
+   if (!binding.success)
+      return NULL;
+
+   if (binding.var)
+      return binding.var;
+
+   nir_foreach_variable_with_modes(var, shader, nir_var_mem_ubo | nir_var_mem_ssbo) {
+      if (var->data.descriptor_set == binding.desc_set && var->data.binding == binding.binding)
+         return var;
+   }
+
+   return NULL;
+}

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index fa9c3cd..9b2debc 100644 (file)
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -2543,6 +2543,21 @@ nir_ssa_scalar_chase_alu_src(nir_ssa_scalar s, unsigned alu_src_idx)
 }
 
 
+typedef struct {
+   bool success;
+
+   nir_variable *var;
+   unsigned desc_set;
+   unsigned binding;
+   unsigned num_indices;
+   nir_src indices[4];
+   bool read_first_invocation;
+} nir_binding;
+
+nir_binding nir_chase_binding(nir_src rsrc);
+nir_variable *nir_get_binding_variable(struct nir_shader *shader, nir_binding binding);
+
+
 /*
  * Control flow
  *