nir: add nir_variable::index and nir_index_vars

This will be useful as a deterministic identifier/index for the variable.

v2: fix comment style

Signed-off-by: Rhys Perry <[email protected]>
Reviewed-by: Connor Abbott <[email protected]> (v1)

2 files changed, 41 insertions, 0 deletions

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index d7bdc7ed4ca..8a100e18dd8 100644
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -1787,6 +1787,39 @@ nir_index_instrs(nir_function_impl *impl)
    return index;
 }
 
+static void
+index_var_list(struct exec_list *list)
+{
+   unsigned next_index = 0;
+   nir_foreach_variable(var, list)
+      var->index = next_index++;
+}
+
+void
+nir_index_vars(nir_shader *shader, nir_function_impl *impl, nir_variable_mode modes)
+{
+   if ((modes & nir_var_function_temp) && impl)
+      index_var_list(&impl->locals);
+
+   if (modes & nir_var_shader_temp)
+      index_var_list(&shader->globals);
+
+   if (modes & nir_var_shader_in)
+      index_var_list(&shader->inputs);
+
+   if (modes & nir_var_shader_out)
+      index_var_list(&shader->outputs);
+
+   if (modes & (nir_var_uniform | nir_var_mem_ubo | nir_var_mem_ssbo))
+      index_var_list(&shader->uniforms);
+
+   if (modes & nir_var_mem_shared)
+      index_var_list(&shader->shared);
+
+   if (modes & nir_var_system_value)
+      index_var_list(&shader->system_values);
+}
+
 static nir_instr *
 cursor_next_instr(nir_cursor cursor)
 {
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 4db85b3fafd..e88d3f8df96 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -504,6 +504,12 @@ typedef struct nir_variable {
       };
    } data;
 
+   /**
+    * Identifier for this variable generated by nir_index_vars() that is unique
+    * among other variables in the same exec_list.
+    */
+   unsigned index;
+
    /* Number of nir_variable_data members */
    uint16_t num_members;
 
@@ -3340,6 +3346,8 @@ unsigned nir_index_instrs(nir_function_impl *impl);
 
 void nir_index_blocks(nir_function_impl *impl);
 
+void nir_index_vars(nir_shader *shader, nir_function_impl *impl, nir_variable_mode modes);
+
 void nir_print_shader(nir_shader *shader, FILE *fp);
 void nir_print_shader_annotated(nir_shader *shader, FILE *fp, struct hash_table *errors);
 void nir_print_instr(const nir_instr *instr, FILE *fp);