intel/compiler: call brw_nir_adjust_payload from brw_postprocess_nir

Calling anything after nir_trivialize_registers() risks undoing some of
its work.
In this case, brw_nir_adjust_payload() will do a constant folding pass
if any payload adjusting happened, and that can turn a bunch of
@store_regs into basically noops.

Fixes dEQP-VK.subgroups.*task

Reviewed-by: Kenneth Graunke <kenneth@whitecape.org>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24325>

src/intel/compiler/brw_mesh.cpp

@@ -228,7 +228,7 @@ brw_nir_adjust_task_payload_offsets(nir_shader *nir)
                                        NULL);
 }
 
-static void
+void
 brw_nir_adjust_payload(nir_shader *shader, const struct brw_compiler *compiler)
 {
    /* Adjustment of task payload offsets must be performed *after* last pass
@@ -331,8 +331,6 @@ brw_compile_task(const struct brw_compiler *compiler,
       brw_postprocess_nir(shader, compiler, debug_enabled,
                           key->base.robust_buffer_access);
 
-      brw_nir_adjust_payload(shader, compiler);
-
       v[simd] = std::make_unique<fs_visitor>(compiler, &params->base,
                                              &key->base,
                                              &prog_data->base.base,
@@ -1494,8 +1492,6 @@ brw_compile_mesh(const struct brw_compiler *compiler,
       brw_postprocess_nir(shader, compiler, debug_enabled,
                           key->base.robust_buffer_access);
 
-      brw_nir_adjust_payload(shader, compiler);
-
       v[simd] = std::make_unique<fs_visitor>(compiler, &params->base,
                                              &key->base,
                                              &prog_data->base.base,

src/intel/compiler/brw_nir.c

@@ -1742,6 +1742,16 @@ brw_postprocess_nir(nir_shader *nir, const struct brw_compiler *compiler,
       brw_nir_analyze_boolean_resolves(nir);
 
    OPT(nir_opt_dce);
+
+   /* The mesh stages require this pass to be called at the last minute,
+    * but if anything is done by it, it will also constant fold, and that
+    * undoes the work done by nir_trivialize_registers, so call it right
+    * before that one instead.
+    */
+   if (nir->info.stage == MESA_SHADER_MESH ||
+       nir->info.stage == MESA_SHADER_TASK)
+      brw_nir_adjust_payload(nir, compiler);
+
    nir_trivialize_registers(nir);
    nir_sweep(nir);
 

src/intel/compiler/brw_nir.h

@@ -286,6 +286,8 @@ nir_ssa_def *brw_nir_load_global_const(nir_builder *b,
 const struct glsl_type *brw_nir_get_var_type(const struct nir_shader *nir,
                                              nir_variable *var);
 
+void brw_nir_adjust_payload(nir_shader *shader, const struct brw_compiler *compiler);
+
 #ifdef __cplusplus
 }
 #endif