nir: Take a shader and variable mode in nir_assign_io_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/amd/vulkan/radv_shader.c b/src/amd/vulkan/radv_shader.c
index 9d5eb0598be119b4bf0b79a5d903541f4b538f42..749167e17236c18a125ac9b3d570206ba2809deb 100644 (file)
--- a/src/amd/vulkan/radv_shader.c
+++ b/src/amd/vulkan/radv_shader.c
@@ -648,7 +648,7 @@ void
 radv_lower_fs_io(nir_shader *nir)
 {
        NIR_PASS_V(nir, lower_view_index);
-       nir_assign_io_var_locations(&nir->inputs, &nir->num_inputs,
+       nir_assign_io_var_locations(nir, nir_var_shader_in, &nir->num_inputs,
                                    MESA_SHADER_FRAGMENT);
 
        NIR_PASS_V(nir, nir_lower_io, nir_var_shader_in, type_size_vec4, 0);

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 280f441665377f42ca5a1ec6f160660e204b002d..dca19596632f4b65256e28aa8b2ea6e2b230f094 100644 (file)
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -3999,7 +3999,8 @@ bool nir_link_opt_varyings(nir_shader *producer, nir_shader *consumer);
 bool nir_lower_amul(nir_shader *shader,
                     int (*type_size)(const struct glsl_type *, bool));
 
-void nir_assign_io_var_locations(struct exec_list *var_list,
+void nir_assign_io_var_locations(nir_shader *shader,
+                                 nir_variable_mode mode,
                                  unsigned *size,
                                  gl_shader_stage stage);
 

diff --git a/src/compiler/nir/nir_linking_helpers.c b/src/compiler/nir/nir_linking_helpers.c
index 8047d122ffe57b7194f4e078c12eb04bc06ca7b1..c73633cba415bc83606f960d0ecea4aca0d8cf40 100644 (file)
--- a/src/compiler/nir/nir_linking_helpers.c
+++ b/src/compiler/nir/nir_linking_helpers.c
@@ -1075,30 +1075,30 @@ insert_sorted(struct exec_list *var_list, nir_variable *new_var)
 }
 
 static void
-sort_varyings(struct exec_list *var_list)
+sort_varyings(nir_shader *shader, nir_variable_mode mode,
+              struct exec_list *sorted_list)
 {
-   struct exec_list new_list;
-   exec_list_make_empty(&new_list);
-   nir_foreach_variable_safe(var, var_list) {
+   exec_list_make_empty(sorted_list);
+   nir_foreach_variable_with_modes_safe(var, shader, mode) {
       exec_node_remove(&var->node);
-      insert_sorted(&new_list, var);
+      insert_sorted(sorted_list, var);
    }
-   exec_list_move_nodes_to(&new_list, var_list);
 }
 
 void
-nir_assign_io_var_locations(struct exec_list *var_list, unsigned *size,
-                            gl_shader_stage stage)
+nir_assign_io_var_locations(nir_shader *shader, nir_variable_mode mode,
+                            unsigned *size, gl_shader_stage stage)
 {
    unsigned location = 0;
    unsigned assigned_locations[VARYING_SLOT_TESS_MAX];
    uint64_t processed_locs[2] = {0};
 
-   sort_varyings(var_list);
+   struct exec_list io_vars;
+   sort_varyings(shader, mode, &io_vars);
 
    int UNUSED last_loc = 0;
    bool last_partial = false;
-   nir_foreach_variable(var, var_list) {
+   nir_foreach_variable(var, &io_vars) {
       const struct glsl_type *type = var->type;
       if (nir_is_per_vertex_io(var, stage) || var->data.per_view) {
          assert(glsl_type_is_array(type));
@@ -1201,6 +1201,8 @@ nir_assign_io_var_locations(struct exec_list *var_list, unsigned *size,
    if (last_partial)
       location++;
 
+   struct exec_list *var_list = nir_variable_list_for_mode(shader, mode);
+   exec_list_append(var_list, &io_vars);
    *size = location;
 }
 

diff --git a/src/freedreno/vulkan/tu_shader.c b/src/freedreno/vulkan/tu_shader.c
index a98a967359b9a3210b266a426cbeab10f3dea939..7f17706cdc90c9844efe822abc3e82d4923315a5 100644 (file)
--- a/src/freedreno/vulkan/tu_shader.c
+++ b/src/freedreno/vulkan/tu_shader.c
@@ -764,8 +764,8 @@ tu_shader_create(struct tu_device *dev,
 
    NIR_PASS_V(nir, nir_lower_io_arrays_to_elements_no_indirects, false);
 
-   nir_assign_io_var_locations(&nir->inputs, &nir->num_inputs, stage);
-   nir_assign_io_var_locations(&nir->outputs, &nir->num_outputs, stage);
+   nir_assign_io_var_locations(nir, nir_var_shader_in, &nir->num_inputs, stage);
+   nir_assign_io_var_locations(nir, nir_var_shader_out, &nir->num_outputs, stage);
 
    NIR_PASS_V(nir, nir_lower_system_values);
    NIR_PASS_V(nir, nir_lower_frexp);

diff --git a/src/mesa/state_tracker/st_glsl_to_nir.cpp b/src/mesa/state_tracker/st_glsl_to_nir.cpp
index bc8592d22e292bccf1253e91a8f415170d78ed18..bdd9cddf98c3c6f03c1e5c11e594de60d7712589 100644 (file)
--- a/src/mesa/state_tracker/st_glsl_to_nir.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_nir.cpp
@@ -853,28 +853,28 @@ void
 st_nir_assign_varying_locations(struct st_context *st, nir_shader *nir)
 {
    if (nir->info.stage == MESA_SHADER_VERTEX) {
-      nir_assign_io_var_locations(&nir->outputs,
+      nir_assign_io_var_locations(nir, nir_var_shader_out,
                                   &nir->num_outputs,
                                   nir->info.stage);
       st_nir_fixup_varying_slots(st, &nir->outputs);
    } else if (nir->info.stage == MESA_SHADER_GEOMETRY ||
               nir->info.stage == MESA_SHADER_TESS_CTRL ||
               nir->info.stage == MESA_SHADER_TESS_EVAL) {
-      nir_assign_io_var_locations(&nir->inputs,
+      nir_assign_io_var_locations(nir, nir_var_shader_in,
                                   &nir->num_inputs,
                                   nir->info.stage);
       st_nir_fixup_varying_slots(st, &nir->inputs);
 
-      nir_assign_io_var_locations(&nir->outputs,
+      nir_assign_io_var_locations(nir, nir_var_shader_out,
                                   &nir->num_outputs,
                                   nir->info.stage);
       st_nir_fixup_varying_slots(st, &nir->outputs);
    } else if (nir->info.stage == MESA_SHADER_FRAGMENT) {
-      nir_assign_io_var_locations(&nir->inputs,
+      nir_assign_io_var_locations(nir, nir_var_shader_in,
                                   &nir->num_inputs,
                                   nir->info.stage);
       st_nir_fixup_varying_slots(st, &nir->inputs);
-      nir_assign_io_var_locations(&nir->outputs,
+      nir_assign_io_var_locations(nir, nir_var_shader_out,
                                   &nir->num_outputs,
                                   nir->info.stage);
    } else if (nir->info.stage == MESA_SHADER_COMPUTE) {