nir: Take a nir_shader and variable mode in assign_var_locations

Reviewed-by: Gert Wollny <gert.wollny@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5966>

5 files changed, 17 insertions, 22 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 7c04a773fbc..280f4416653 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -3982,7 +3982,8 @@ void nir_gather_ssa_types(nir_function_impl *impl,
                           BITSET_WORD *float_types,
                           BITSET_WORD *int_types);
 
-void nir_assign_var_locations(struct exec_list *var_list, unsigned *size,
+void nir_assign_var_locations(nir_shader *shader, nir_variable_mode mode,
+                              unsigned *size,
                               int (*type_size)(const struct glsl_type *, bool));
 
 /* Some helpers to do very simple linking */
diff --git a/src/compiler/nir/nir_lower_io.c b/src/compiler/nir/nir_lower_io.c
index b004c62b81e..f8c3fe61535 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -121,19 +121,13 @@ shared_atomic_for_deref(nir_intrinsic_op deref_op)
 }
 
 void
-nir_assign_var_locations(struct exec_list *var_list, unsigned *size,
+nir_assign_var_locations(nir_shader *shader, nir_variable_mode mode,
+                         unsigned *size,
                          int (*type_size)(const struct glsl_type *, bool))
 {
    unsigned location = 0;
 
-   nir_foreach_variable(var, var_list) {
-      /*
-       * UBOs have their own address spaces, so don't count them towards the
-       * number of global uniforms
-       */
-      if (var->data.mode == nir_var_mem_ubo || var->data.mode == nir_var_mem_ssbo)
-         continue;
-
+   nir_foreach_variable_with_modes(var, shader, mode) {
       var->data.driver_location = location;
       bool bindless_type_size = var->data.mode == nir_var_shader_in ||
                                 var->data.mode == nir_var_shader_out ||
diff --git a/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c b/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
index c9dc7beff36..294fb61e752 100644
--- a/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
+++ b/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
@@ -147,7 +147,7 @@ load_glsl(unsigned num_files, char* const* files, gl_shader_stage stage)
 
 	switch (stage) {
 	case MESA_SHADER_VERTEX:
-		nir_assign_var_locations(&nir->inputs,
+		nir_assign_var_locations(nir, nir_var_shader_in,
 				&nir->num_inputs,
 				ir3_glsl_type_size);
 
@@ -155,18 +155,18 @@ load_glsl(unsigned num_files, char* const* files, gl_shader_stage stage)
 		NIR_PASS_V(nir, nir_lower_global_vars_to_local);
 
 		sort_varyings(&nir->outputs);
-		nir_assign_var_locations(&nir->outputs,
+		nir_assign_var_locations(nir, nir_var_shader_out,
 				&nir->num_outputs,
 				ir3_glsl_type_size);
 		fixup_varying_slots(&nir->outputs);
 		break;
 	case MESA_SHADER_FRAGMENT:
 		sort_varyings(&nir->inputs);
-		nir_assign_var_locations(&nir->inputs,
+		nir_assign_var_locations(nir, nir_var_shader_in,
 				&nir->num_inputs,
 				ir3_glsl_type_size);
 		fixup_varying_slots(&nir->inputs);
-		nir_assign_var_locations(&nir->outputs,
+		nir_assign_var_locations(nir, nir_var_shader_out,
 				&nir->num_outputs,
 				ir3_glsl_type_size);
 		break;
@@ -177,7 +177,7 @@ load_glsl(unsigned num_files, char* const* files, gl_shader_stage stage)
 		errx(1, "unhandled shader stage: %d", stage);
 	}
 
-	nir_assign_var_locations(&nir->uniforms,
+	nir_assign_var_locations(nir, nir_var_uniform,
 			&nir->num_uniforms,
 			ir3_glsl_type_size);
 
diff --git a/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c b/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c
index a05caf6e306..044bbcccf7d 100644
--- a/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c
+++ b/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c
@@ -140,30 +140,30 @@ load_glsl(unsigned num_files, char* const* files, gl_shader_stage stage)
 
    switch (stage) {
    case MESA_SHADER_VERTEX:
-      nir_assign_var_locations(&nir->inputs, &nir->num_inputs,
+      nir_assign_var_locations(nir, nir_var_shader_in, &nir->num_inputs,
                                st_glsl_type_size);
 
       /* Re-lower global vars, to deal with any dead VS inputs. */
       NIR_PASS_V(nir, nir_lower_global_vars_to_local);
 
       sort_varyings(&nir->outputs);
-      nir_assign_var_locations(&nir->outputs, &nir->num_outputs,
+      nir_assign_var_locations(nir, nir_var_shader_out, &nir->num_outputs,
                                st_glsl_type_size);
       fixup_varying_slots(&nir->outputs);
       break;
    case MESA_SHADER_FRAGMENT:
       sort_varyings(&nir->inputs);
-      nir_assign_var_locations(&nir->inputs, &nir->num_inputs,
+      nir_assign_var_locations(nir, nir_var_shader_in, &nir->num_inputs,
                                st_glsl_type_size);
       fixup_varying_slots(&nir->inputs);
-      nir_assign_var_locations(&nir->outputs, &nir->num_outputs,
+      nir_assign_var_locations(nir, nir_var_shader_out, &nir->num_outputs,
                                st_glsl_type_size);
       break;
    default:
       errx(1, "unhandled shader stage: %d", stage);
    }
 
-   nir_assign_var_locations(&nir->uniforms,
+   nir_assign_var_locations(nir, nir_var_uniform,
                             &nir->num_uniforms,
                             st_glsl_type_size);
 
diff --git a/src/mesa/drivers/dri/i965/brw_program.c b/src/mesa/drivers/dri/i965/brw_program.c
index 6bb6f84182d..b3db2cc61bb 100644
--- a/src/mesa/drivers/dri/i965/brw_program.c
+++ b/src/mesa/drivers/dri/i965/brw_program.c
@@ -62,11 +62,11 @@ static bool
 brw_nir_lower_uniforms(nir_shader *nir, bool is_scalar)
 {
    if (is_scalar) {
-      nir_assign_var_locations(&nir->uniforms, &nir->num_uniforms,
+      nir_assign_var_locations(nir, nir_var_uniform, &nir->num_uniforms,
                                type_size_scalar_bytes);
       return nir_lower_io(nir, nir_var_uniform, type_size_scalar_bytes, 0);
    } else {
-      nir_assign_var_locations(&nir->uniforms, &nir->num_uniforms,
+      nir_assign_var_locations(nir, nir_var_uniform, &nir->num_uniforms,
                                type_size_vec4_bytes);
       return nir_lower_io(nir, nir_var_uniform, type_size_vec4_bytes, 0);
    }