mesa/st, nir: commonize unlower_io_to_vars pass

Signed-off-by: Simon Perretta <simon.perretta@imgtec.com>
Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37540>

src/compiler/nir/meson.build

@@ -332,6 +332,7 @@ else
   'nir_to_lcssa.c',
   'nir_tcs_info.h',
   'nir_trivialize_registers.c',
+  'nir_unlower_io_to_vars.c',
   'nir_use_dominance.c',
   'nir_validate.c',
   'nir_vla.h',

src/compiler/nir/nir.h

@@ -6771,6 +6771,8 @@ void nir_gather_output_clipper_var_groups(nir_shader *nir,
 
 bool nir_lower_cooperative_matrix_flexible_dimensions(nir_shader *shader, unsigned m_gran, unsigned n_gran, unsigned k_gran);
 
+bool nir_unlower_io_to_vars(nir_shader *nir, bool keep_intrinsics);
+
 #include "nir_inline_helpers.h"
 
 #ifdef __cplusplus

src/compiler/nir/nir_unlower_io_to_vars.c

@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: MIT
  */
 
-#include "st_nir.h"
 #include "nir_builder.h"
 
 struct io_desc {
@@ -676,8 +675,11 @@ unlower_io_to_vars(nir_builder *b, nir_intrinsic_instr *intr, void *opaque)
    return true;
 }
 
+/* If keep_intrinsics is set, the pass will skip unlowering lowered I/O
+ * intrinsics to derefs, and only (re)create I/O variables from them.
+ */
 bool
-st_nir_unlower_io_to_vars(nir_shader *nir)
+nir_unlower_io_to_vars(nir_shader *nir, bool keep_intrinsics)
 {
    if (nir->info.stage == MESA_SHADER_COMPUTE)
       return false;
@@ -691,7 +693,12 @@ st_nir_unlower_io_to_vars(nir_shader *nir)
    assert(!(nir->options->io_options &
             nir_io_mix_convergent_flat_with_interpolated));
 
-   nir_foreach_variable_with_modes(var, nir, nir_var_shader_in | nir_var_shader_out) {
+   nir_foreach_variable_with_modes_safe(var, nir, nir_var_shader_in | nir_var_shader_out) {
+      if (keep_intrinsics) {
+         exec_node_remove(&var->node);
+         continue;
+      }
+
       UNREACHABLE("the shader should have no IO variables");
    }
 
@@ -719,11 +726,13 @@ st_nir_unlower_io_to_vars(nir_shader *nir)
    }
 
    /* Unlower IO using the created variables. */
-   ASSERTED bool lower_progress =
-      nir_shader_intrinsics_pass(nir, unlower_io_to_vars,
-                                 nir_metadata_control_flow, NULL);
-   assert(lower_progress);
-   nir->info.io_lowered = false;
+   if (!keep_intrinsics) {
+      ASSERTED bool lower_progress =
+         nir_shader_intrinsics_pass(nir, unlower_io_to_vars,
+                                    nir_metadata_control_flow, NULL);
+      assert(lower_progress);
+      nir->info.io_lowered = false;
+   }
 
    /* Count IO variables. */
    nir->num_inputs = 0;

src/mesa/meson.build

@@ -350,7 +350,6 @@ files_libmesa = files(
   'state_tracker/st_nir_lower_point_size_mov.c',
   'state_tracker/st_nir_lower_position_invariant.c',
   'state_tracker/st_nir_lower_tex_src_plane.c',
-  'state_tracker/st_nir_unlower_io_to_vars.c',
   'state_tracker/st_pbo.c',
   'state_tracker/st_pbo_compute.c',
   'state_tracker/st_pbo.h',

src/mesa/state_tracker/st_nir.h

@@ -111,8 +111,6 @@ bool st_nir_lower_position_invariant(struct nir_shader *s,
                                      struct gl_program_parameter_list *paramList,
                                      bool packed_driver_uniform_storage);
 
-bool st_nir_unlower_io_to_vars(struct nir_shader *nir);
-
 #ifdef __cplusplus
 }
 #endif

src/mesa/state_tracker/st_nir_builtins.c

@@ -64,7 +64,7 @@ st_nir_finish_builtin_nir(struct st_context *st, nir_shader *nir)
 
    if (nir->info.io_lowered &&
        !(nir->options->io_options & nir_io_has_intrinsics)) {
-      NIR_PASS(_, nir, st_nir_unlower_io_to_vars);
+      NIR_PASS(_, nir, nir_unlower_io_to_vars, false);
       gl_nir_opts(nir);
    }
 

src/mesa/state_tracker/st_program.c

@@ -873,7 +873,7 @@ st_create_common_variant(struct st_context *st,
        * are still counted as enabled IO, which breaks things.
        */
       NIR_PASS(_, state.ir.nir, nir_opt_dce);
-      NIR_PASS(_, state.ir.nir, st_nir_unlower_io_to_vars);
+      NIR_PASS(_, state.ir.nir, nir_unlower_io_to_vars, false);
 
       if (state.ir.nir->info.stage == MESA_SHADER_TESS_CTRL &&
           state.ir.nir->options->compact_arrays &&
@@ -1256,7 +1256,7 @@ st_create_fp_variant(struct st_context *st,
        * are still counted as enabled IO, which breaks things.
        */
       NIR_PASS(_, state.ir.nir, nir_opt_dce);
-      NIR_PASS(_, state.ir.nir, st_nir_unlower_io_to_vars);
+      NIR_PASS(_, state.ir.nir, nir_unlower_io_to_vars, false);
       gl_nir_opts(state.ir.nir);
       finalize = true;
    }