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>
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 7c04a77..280f441 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -3982,7 +3982,8 @@
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 b004c62..f8c3fe6 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -121,19 +121,13 @@
}
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 c9dc7be..294fb61 100644
--- a/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
+++ b/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
@@ -147,7 +147,7 @@
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 @@
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 @@
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 a05caf6..044bbcc 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 @@
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 6bb6f84..b3db2cc 100644
--- a/src/mesa/drivers/dri/i965/brw_program.c
+++ b/src/mesa/drivers/dri/i965/brw_program.c
@@ -62,11 +62,11 @@
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);
}
