pco: fix split-type vertex attrib allocations/nir vars

Signed-off-by: Simon Perretta <simon.perretta@imgtec.com>
Acked-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36412>

src/imagination/.clang-format

@@ -240,6 +240,7 @@ ForEachMacros: [
   'nir_foreach_use',
   'nir_foreach_use_safe',
   'nir_foreach_variable_with_modes',
+  'nir_foreach_variable_with_modes_safe',
   'pco_foreach_block_in_func',
   'pco_foreach_block_in_func_from',
   'pco_foreach_block_in_func_from_rev',

src/imagination/pco/pco_nir.c

@@ -58,6 +58,8 @@ static const nir_shader_compiler_options nir_options = {
    .scalarize_ddx = true,
 
    .max_unroll_iterations = 16,
+
+   .io_options = nir_io_vectorizer_ignores_types,
 };
 
 /**
@@ -760,13 +762,14 @@ void pco_lower_nir(pco_ctx *ctx, nir_shader *nir, pco_data *data)
             glsl_type_size,
             nir_lower_io_lower_64bit_to_32);
 
-   if (nir->info.stage == MESA_SHADER_FRAGMENT) {
-      NIR_PASS(_, nir, nir_lower_io_to_scalar, nir_var_shader_out, NULL, NULL);
-      NIR_PASS(_, nir, nir_copy_prop);
-      NIR_PASS(_, nir, nir_opt_dce);
-      NIR_PASS(_, nir, nir_opt_cse);
-      NIR_PASS(_, nir, nir_opt_vectorize_io, nir_var_shader_out, false);
-   }
+   nir_variable_mode vec_modes = nir->info.stage == MESA_SHADER_FRAGMENT
+                                    ? nir_var_shader_out
+                                    : nir_var_shader_in;
+   NIR_PASS(_, nir, nir_lower_io_to_scalar, vec_modes, NULL, NULL);
+   NIR_PASS(_, nir, nir_copy_prop);
+   NIR_PASS(_, nir, nir_opt_dce);
+   NIR_PASS(_, nir, nir_opt_cse);
+   NIR_PASS(_, nir, nir_opt_vectorize_io, vec_modes, false);
 
    NIR_PASS(_, nir, nir_opt_dce);
    NIR_PASS(_, nir, nir_opt_constant_folding);
@@ -865,13 +868,12 @@ void pco_lower_nir(pco_ctx *ctx, nir_shader *nir, pco_data *data)
       NIR_PASS(_, nir, nir_opt_cse);
    } while (progress);
 
-   nir_variable_mode vec_modes = nir_var_shader_in;
+   vec_modes = nir_var_shader_in;
    /* Fragment shader needs scalar writes after pfo. */
    if (nir->info.stage != MESA_SHADER_FRAGMENT)
       vec_modes |= nir_var_shader_out;
 
-   if (nir->info.stage == MESA_SHADER_FRAGMENT)
-      NIR_PASS(_, nir, nir_opt_vectorize_io, vec_modes, false);
+   NIR_PASS(_, nir, nir_opt_vectorize_io, vec_modes, false);
 
    /* Special case for frag coords:
     * - x,y come from (non-consecutive) special regs - always scalar.

src/imagination/pco/pco_nir_pvfio.c

@@ -1253,9 +1253,21 @@ static nir_def *lower_pvi(nir_builder *b, nir_instr *instr, void *cb_data)
    if (output->num_components > intr->def.num_components)
       output = nir_trim_vector(b, output, intr->def.num_components);
 
-   /* Update the type of the stored variable. */
-   nir_variable *var =
-      nir_find_variable_with_location(b->shader, nir_var_shader_in, location);
+   /* Update the type of the stored variable, remove any fractional vars. */
+   nir_variable *var = NULL;
+   nir_foreach_variable_with_modes_safe (iter_var,
+                                         b->shader,
+                                         nir_var_shader_in) {
+      if (iter_var->data.location != location)
+         continue;
+
+      if (!iter_var->data.location_frac) {
+         var = iter_var;
+         continue;
+      }
+
+      exec_node_remove(&iter_var->node);
+   }
    assert(var);
 
    unsigned format_dwords =