ac/nir: split ac_nir_lower_ps into 2 passes

It's split into ac_nir_lower_ps_early ac_nir_lower_ps_late.

ac_nir_lower_ps_early doesn't generate any AMD specific intrinsics except
some system values and is mainly an optimization pass with some lowering.
The new change here is that it also eliminates output components not needed
by spi_shader_col_format.

ac_nir_lower_ps_late lowers output stores to exports and does the bc_optimize
thing.

Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32966>

src/amd/common/ac_nir.h

@@ -271,41 +271,59 @@ ac_nir_lower_legacy_gs(nir_shader *nir,
                        ac_nir_gs_output_info *output_info);
 
 typedef struct {
-   enum radeon_family family;
-   enum amd_gfx_level gfx_level;
-
-   bool use_aco;
-   bool uses_discard;
-   bool alpha_to_coverage_via_mrtz;
-   bool dual_src_blend_swizzle;
-   unsigned spi_shader_col_format;
-   unsigned color_is_int8;
-   unsigned color_is_int10;
-
-   bool bc_optimize_for_persp;
-   bool bc_optimize_for_linear;
+   /* This is a pre-link pass. It should only eliminate code and do lowering that mostly doesn't
+    * generate AMD-specific intrinsics.
+    */
+   /* System values. */
    bool force_persp_sample_interp;
    bool force_linear_sample_interp;
    bool force_persp_center_interp;
    bool force_linear_center_interp;
    unsigned ps_iter_samples;
 
-   /* OpenGL only */
-   bool clamp_color;
-   bool alpha_to_one;
-   enum compare_func alpha_func;
+   /* Outputs. */
+   bool clamp_color;                /* GL only */
+   bool alpha_test_alpha_to_one;    /* GL only, this only affects alpha test */
+   enum compare_func alpha_func;    /* GL only */
+   bool keep_alpha_for_mrtz;        /* this prevents killing alpha based on spi_shader_col_format_hint */
+   unsigned spi_shader_col_format_hint; /* this only shrinks and eliminates output stores */
    bool kill_z;
    bool kill_stencil;
    bool kill_samplemask;
+} ac_nir_lower_ps_early_options;
+
+void
+ac_nir_lower_ps_early(nir_shader *nir, const ac_nir_lower_ps_early_options *options);
+
+typedef struct {
+   /* This is a post-link pass. It shouldn't eliminate any code and it shouldn't affect shader_info
+    * (those should be done in the early pass).
+    */
+   enum amd_gfx_level gfx_level;
+   enum radeon_family family;
+   bool use_aco;
+
+   /* System values. */
+   bool bc_optimize_for_persp;
+   bool bc_optimize_for_linear;
+
+   /* Exports. */
+   bool uses_discard;
+   bool alpha_to_coverage_via_mrtz;
+   bool dual_src_blend_swizzle;
+   unsigned spi_shader_col_format;
+   unsigned color_is_int8;
+   unsigned color_is_int10;
+   bool alpha_to_one;
 
    /* Vulkan only */
    unsigned enable_mrt_output_nan_fixup;
    bool no_color_export;
    bool no_depth_export;
-} ac_nir_lower_ps_options;
+} ac_nir_lower_ps_late_options;
 
 void
-ac_nir_lower_ps(nir_shader *nir, const ac_nir_lower_ps_options *options);
+ac_nir_lower_ps_late(nir_shader *nir, const ac_nir_lower_ps_late_options *options);
 
 typedef struct {
    enum amd_gfx_level gfx_level;

src/amd/common/ac_nir_lower_ps_early.c

@@ -0,0 +1,360 @@
+/*
+ * Copyright 2024 Advanced Micro Devices, Inc.
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+/* This is a pre-link lowering and optimization pass that modifies the shader for the purpose
+ * of gathering accurate shader_info and determining hw registers. It should be run before
+ * linking passes and it doesn't produce AMD intrinsics that would break linking passes.
+ * Some of the options come from dynamic state.
+ */
+
+#include "ac_nir.h"
+#include "sid.h"
+#include "nir_builder.h"
+#include "nir_builtin_builder.h"
+
+typedef struct {
+   const ac_nir_lower_ps_early_options *options;
+
+   nir_variable *persp_center;
+   nir_variable *persp_centroid;
+   nir_variable *persp_sample;
+   nir_variable *linear_center;
+   nir_variable *linear_centroid;
+   nir_variable *linear_sample;
+   bool lower_load_barycentric;
+
+   bool seen_color0_alpha;
+} lower_ps_early_state;
+
+static void
+create_interp_param(nir_builder *b, lower_ps_early_state *s)
+{
+   if (s->options->force_persp_sample_interp) {
+      s->persp_center =
+         nir_local_variable_create(b->impl, glsl_vec_type(2), "persp_center");
+   }
+
+   if (s->options->force_persp_sample_interp ||
+       s->options->force_persp_center_interp) {
+      s->persp_centroid =
+         nir_local_variable_create(b->impl, glsl_vec_type(2), "persp_centroid");
+   }
+
+   if (s->options->force_persp_center_interp) {
+      s->persp_sample =
+         nir_local_variable_create(b->impl, glsl_vec_type(2), "persp_sample");
+   }
+
+   if (s->options->force_linear_sample_interp) {
+      s->linear_center =
+         nir_local_variable_create(b->impl, glsl_vec_type(2), "linear_center");
+   }
+
+   if (s->options->force_linear_sample_interp ||
+       s->options->force_linear_center_interp) {
+      s->linear_centroid =
+         nir_local_variable_create(b->impl, glsl_vec_type(2), "linear_centroid");
+   }
+
+   if (s->options->force_linear_center_interp) {
+      s->linear_sample =
+         nir_local_variable_create(b->impl, glsl_vec_type(2), "linear_sample");
+   }
+
+   s->lower_load_barycentric =
+      s->persp_center || s->persp_centroid || s->persp_sample ||
+      s->linear_center || s->linear_centroid || s->linear_sample;
+}
+
+static void
+init_interp_param(nir_builder *b, lower_ps_early_state *s)
+{
+   b->cursor = nir_before_cf_list(&b->impl->body);
+
+   if (s->options->force_persp_sample_interp) {
+      nir_def *sample =
+         nir_load_barycentric_sample(b, 32, .interp_mode = INTERP_MODE_SMOOTH);
+      nir_store_var(b, s->persp_center, sample, 0x3);
+      nir_store_var(b, s->persp_centroid, sample, 0x3);
+   }
+
+   if (s->options->force_linear_sample_interp) {
+      nir_def *sample =
+         nir_load_barycentric_sample(b, 32, .interp_mode = INTERP_MODE_NOPERSPECTIVE);
+      nir_store_var(b, s->linear_center, sample, 0x3);
+      nir_store_var(b, s->linear_centroid, sample, 0x3);
+   }
+
+   if (s->options->force_persp_center_interp) {
+      nir_def *center =
+         nir_load_barycentric_pixel(b, 32, .interp_mode = INTERP_MODE_SMOOTH);
+      nir_store_var(b, s->persp_sample, center, 0x3);
+      nir_store_var(b, s->persp_centroid, center, 0x3);
+   }
+
+   if (s->options->force_linear_center_interp) {
+      nir_def *center =
+         nir_load_barycentric_pixel(b, 32, .interp_mode = INTERP_MODE_NOPERSPECTIVE);
+      nir_store_var(b, s->linear_sample, center, 0x3);
+      nir_store_var(b, s->linear_centroid, center, 0x3);
+   }
+}
+
+static bool
+rewrite_ps_load_barycentric(nir_builder *b, nir_intrinsic_instr *intrin, lower_ps_early_state *s)
+{
+   enum glsl_interp_mode mode = nir_intrinsic_interp_mode(intrin);
+   nir_variable *var = NULL;
+
+   switch (mode) {
+   case INTERP_MODE_NONE:
+   case INTERP_MODE_SMOOTH:
+      switch (intrin->intrinsic) {
+      case nir_intrinsic_load_barycentric_pixel:
+         var = s->persp_center;
+         break;
+      case nir_intrinsic_load_barycentric_centroid:
+         var = s->persp_centroid;
+         break;
+      case nir_intrinsic_load_barycentric_sample:
+         var = s->persp_sample;
+         break;
+      default:
+         break;
+      }
+      break;
+
+   case INTERP_MODE_NOPERSPECTIVE:
+      switch (intrin->intrinsic) {
+      case nir_intrinsic_load_barycentric_pixel:
+         var = s->linear_center;
+         break;
+      case nir_intrinsic_load_barycentric_centroid:
+         var = s->linear_centroid;
+         break;
+      case nir_intrinsic_load_barycentric_sample:
+         var = s->linear_sample;
+         break;
+      default:
+         break;
+      }
+      break;
+
+   default:
+      break;
+   }
+
+   if (!var)
+      return false;
+
+   b->cursor = nir_before_instr(&intrin->instr);
+
+   nir_def *replacement = nir_load_var(b, var);
+   nir_def_replace(&intrin->def, replacement);
+   return true;
+}
+
+static bool
+optimize_lower_ps_outputs(nir_builder *b, nir_intrinsic_instr *intrin, lower_ps_early_state *s)
+{
+   unsigned slot = nir_intrinsic_io_semantics(intrin).location;
+
+   switch (slot) {
+   case FRAG_RESULT_DEPTH:
+      if (!s->options->kill_z)
+         return false;
+      nir_instr_remove(&intrin->instr);
+      return true;
+
+   case FRAG_RESULT_STENCIL:
+      if (!s->options->kill_stencil)
+         return false;
+      nir_instr_remove(&intrin->instr);
+      return true;
+
+   case FRAG_RESULT_SAMPLE_MASK:
+      if (!s->options->kill_samplemask)
+         return false;
+      nir_instr_remove(&intrin->instr);
+      return true;
+   }
+
+   unsigned writemask = nir_intrinsic_write_mask(intrin);
+   unsigned component = nir_intrinsic_component(intrin);
+   unsigned color_index = (slot >= FRAG_RESULT_DATA0 ? slot - FRAG_RESULT_DATA0 : 0) +
+                          nir_intrinsic_io_semantics(intrin).dual_source_blend_index;
+   nir_def *value = intrin->src[0].ssa;
+   bool progress = false;
+
+   b->cursor = nir_before_instr(&intrin->instr);
+
+   /* Clamp color. */
+   if (s->options->clamp_color) {
+      value = nir_fsat(b, value);
+      progress = true;
+   }
+
+   /* Alpha test. */
+   if (color_index == 0 && s->options->alpha_func != COMPARE_FUNC_ALWAYS &&
+       (writemask << component) & BITFIELD_BIT(3)) {
+      assert(!s->seen_color0_alpha);
+      s->seen_color0_alpha = true;
+
+      if (s->options->alpha_func == COMPARE_FUNC_NEVER) {
+         nir_discard(b);
+      } else {
+         nir_def *ref = nir_load_alpha_reference_amd(b);
+         ref = nir_convert_to_bit_size(b, ref, nir_type_float, value->bit_size);
+         nir_def *alpha = s->options->alpha_test_alpha_to_one ?
+                             nir_imm_floatN_t(b, 1, value->bit_size) :
+                             nir_channel(b, value, 3 - component);
+         nir_def *cond = nir_compare_func(b, s->options->alpha_func, alpha, ref);
+         nir_discard_if(b, nir_inot(b, cond));
+      }
+      progress = true;
+   }
+
+   /* Trim the src according to the format and writemask. */
+   unsigned cb_shader_mask = ac_get_cb_shader_mask(s->options->spi_shader_col_format_hint);
+   unsigned format_mask;
+
+   if (slot == FRAG_RESULT_COLOR) {
+      /* cb_shader_mask is 0 for disabled color buffers, so combine all of them. */
+      format_mask = 0;
+      for (unsigned i = 0; i < 8; i++)
+         format_mask |= (cb_shader_mask >> (i * 4)) & 0xf;
+   } else {
+      format_mask = (cb_shader_mask >> (color_index * 4)) & 0xf;
+   }
+
+   if (s->options->keep_alpha_for_mrtz && color_index == 0)
+      format_mask |= BITFIELD_BIT(3);
+
+   writemask = (format_mask >> component) & writemask;
+   nir_intrinsic_set_write_mask(intrin, writemask);
+
+   /* Empty writemask. */
+   if (!writemask) {
+      nir_instr_remove(&intrin->instr);
+      return true;
+   }
+
+   /* Trim the src to the last bit of writemask. */
+   unsigned num_components = util_last_bit(writemask);
+
+   if (num_components != value->num_components) {
+      assert(num_components < value->num_components);
+      value = nir_trim_vector(b, value, num_components);
+      progress = true;
+   }
+
+   /* Replace disabled channels in a non-contiguous writemask with undef. */
+   if (!util_is_power_of_two_nonzero(writemask + 1)) {
+      u_foreach_bit(i, BITFIELD_MASK(num_components) & ~writemask) {
+         value = nir_vector_insert_imm(b, value, nir_undef(b, 1, value->bit_size), i);
+         progress = true;
+      }
+   }
+
+   if (progress && intrin->src[0].ssa != value) {
+      nir_src_rewrite(&intrin->src[0], value);
+      intrin->num_components = value->num_components;
+   } else {
+      assert(intrin->src[0].ssa == value);
+   }
+
+   return progress;
+}
+
+static bool
+lower_ps_load_sample_mask_in(nir_builder *b, nir_intrinsic_instr *intrin, lower_ps_early_state *s)
+{
+   /* Section 15.2.2 (Shader Inputs) of the OpenGL 4.5 (Core Profile) spec
+    * says:
+    *
+    *    "When per-sample shading is active due to the use of a fragment
+    *     input qualified by sample or due to the use of the gl_SampleID
+    *     or gl_SamplePosition variables, only the bit for the current
+    *     sample is set in gl_SampleMaskIn. When state specifies multiple
+    *     fragment shader invocations for a given fragment, the sample
+    *     mask for any single fragment shader invocation may specify a
+    *     subset of the covered samples for the fragment. In this case,
+    *     the bit corresponding to each covered sample will be set in
+    *     exactly one fragment shader invocation."
+    *
+    * The samplemask loaded by hardware is always the coverage of the
+    * entire pixel/fragment, so mask bits out based on the sample ID.
+    */
+
+   b->cursor = nir_before_instr(&intrin->instr);
+
+   uint32_t ps_iter_mask = ac_get_ps_iter_mask(s->options->ps_iter_samples);
+   nir_def *sampleid = nir_load_sample_id(b);
+   nir_def *submask = nir_ishl(b, nir_imm_int(b, ps_iter_mask), sampleid);
+
+   nir_def *sample_mask = nir_load_sample_mask_in(b);
+   nir_def *replacement = nir_iand(b, sample_mask, submask);
+
+   nir_def_replace(&intrin->def, replacement);
+   return true;
+}
+
+static bool
+lower_ps_intrinsic(nir_builder *b, nir_instr *instr, void *state)
+{
+   lower_ps_early_state *s = (lower_ps_early_state *)state;
+
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+
+   switch (intrin->intrinsic) {
+   case nir_intrinsic_store_output:
+      return optimize_lower_ps_outputs(b, intrin, s);
+   case nir_intrinsic_load_barycentric_pixel:
+   case nir_intrinsic_load_barycentric_centroid:
+   case nir_intrinsic_load_barycentric_sample:
+      if (s->lower_load_barycentric)
+         return rewrite_ps_load_barycentric(b, intrin, s);
+      break;
+   case nir_intrinsic_load_sample_mask_in:
+      if (s->options->ps_iter_samples > 1)
+         return lower_ps_load_sample_mask_in(b, intrin, s);
+      break;
+   default:
+      break;
+   }
+
+   return false;
+}
+
+void
+ac_nir_lower_ps_early(nir_shader *nir, const ac_nir_lower_ps_early_options *options)
+{
+   assert(nir->info.stage == MESA_SHADER_FRAGMENT);
+   nir_function_impl *impl = nir_shader_get_entrypoint(nir);
+
+   nir_builder builder = nir_builder_create(impl);
+   nir_builder *b = &builder;
+
+   lower_ps_early_state state = {
+      .options = options,
+   };
+
+   create_interp_param(b, &state);
+
+   nir_shader_instructions_pass(nir, lower_ps_intrinsic,
+                                nir_metadata_control_flow,
+                                &state);
+
+   /* This must be after lower_ps_intrinsic. */
+   init_interp_param(b, &state);
+
+   /* Cleanup local variables, as RADV won't do this. */
+   if (state.lower_load_barycentric)
+      nir_lower_vars_to_ssa(nir);
+}

src/amd/common/ac_nir_lower_ps_late.c

@@ -4,13 +4,19 @@
  * SPDX-License-Identifier: MIT
  */
 
+/* This is a post-link lowering pass that lowers intrinsics to AMD-specific ones and thus breaks
+ * shader_info gathering.
+ *
+ * It lowers output stores to exports and inserts the bc_optimize conditional.
+ */
+
 #include "ac_nir.h"
 #include "sid.h"
 #include "nir_builder.h"
 #include "nir_builtin_builder.h"
 
 typedef struct {
-   const ac_nir_lower_ps_options *options;
+   const ac_nir_lower_ps_late_options *options;
 
    nir_variable *persp_center;
    nir_variable *persp_centroid;
@@ -41,43 +47,17 @@ typedef struct {
 static void
 create_interp_param(nir_builder *b, lower_ps_state *s)
 {
-   if (s->options->force_persp_sample_interp) {
-      s->persp_center =
-         nir_local_variable_create(b->impl, glsl_vec_type(2), "persp_center");
-   }
-
-   if (s->options->bc_optimize_for_persp ||
-       s->options->force_persp_sample_interp ||
-       s->options->force_persp_center_interp) {
+   if (s->options->bc_optimize_for_persp) {
       s->persp_centroid =
          nir_local_variable_create(b->impl, glsl_vec_type(2), "persp_centroid");
    }
 
-   if (s->options->force_persp_center_interp) {
-      s->persp_sample =
-         nir_local_variable_create(b->impl, glsl_vec_type(2), "persp_sample");
-   }
-
-   if (s->options->force_linear_sample_interp) {
-      s->linear_center =
-         nir_local_variable_create(b->impl, glsl_vec_type(2), "linear_center");
-   }
-
-   if (s->options->bc_optimize_for_linear ||
-       s->options->force_linear_sample_interp ||
-       s->options->force_linear_center_interp) {
+   if (s->options->bc_optimize_for_linear) {
       s->linear_centroid =
          nir_local_variable_create(b->impl, glsl_vec_type(2), "linear_centroid");
    }
 
-   if (s->options->force_linear_center_interp) {
-      s->linear_sample =
-         nir_local_variable_create(b->impl, glsl_vec_type(2), "linear_sample");
-   }
-
-   s->lower_load_barycentric =
-      s->persp_center || s->persp_centroid || s->persp_sample ||
-      s->linear_center || s->linear_centroid || s->linear_sample;
+   s->lower_load_barycentric = s->persp_centroid || s->linear_centroid;
 }
 
 static void
@@ -112,34 +92,6 @@ init_interp_param(nir_builder *b, lower_ps_state *s)
          nir_store_var(b, s->linear_centroid, value, 0x3);
       }
    }
-
-   if (s->options->force_persp_sample_interp) {
-      nir_def *sample =
-         nir_load_barycentric_sample(b, 32, .interp_mode = INTERP_MODE_SMOOTH);
-      nir_store_var(b, s->persp_center, sample, 0x3);
-      nir_store_var(b, s->persp_centroid, sample, 0x3);
-   }
-
-   if (s->options->force_linear_sample_interp) {
-      nir_def *sample =
-         nir_load_barycentric_sample(b, 32, .interp_mode = INTERP_MODE_NOPERSPECTIVE);
-      nir_store_var(b, s->linear_center, sample, 0x3);
-      nir_store_var(b, s->linear_centroid, sample, 0x3);
-   }
-
-   if (s->options->force_persp_center_interp) {
-      nir_def *center =
-         nir_load_barycentric_pixel(b, 32, .interp_mode = INTERP_MODE_SMOOTH);
-      nir_store_var(b, s->persp_sample, center, 0x3);
-      nir_store_var(b, s->persp_centroid, center, 0x3);
-   }
-
-   if (s->options->force_linear_center_interp) {
-      nir_def *center =
-         nir_load_barycentric_pixel(b, 32, .interp_mode = INTERP_MODE_NOPERSPECTIVE);
-      nir_store_var(b, s->linear_sample, center, 0x3);
-      nir_store_var(b, s->linear_centroid, center, 0x3);
-   }
 }
 
 static bool
@@ -216,18 +168,15 @@ gather_ps_store_output(nir_builder *b, nir_intrinsic_instr *intrin, lower_ps_sta
       switch (slot) {
       case FRAG_RESULT_DEPTH:
          assert(comp == 0);
-         if (!s->options->kill_z)
-            s->depth = chan;
+         s->depth = chan;
          break;
       case FRAG_RESULT_STENCIL:
          assert(comp == 0);
-         if (!s->options->kill_stencil)
-            s->stencil = chan;
+         s->stencil = chan;
          break;
       case FRAG_RESULT_SAMPLE_MASK:
          assert(comp == 0);
-         if (!s->options->kill_samplemask)
-            s->sample_mask = chan;
+         s->sample_mask = chan;
          break;
       case FRAG_RESULT_COLOR:
          s->color[color_index][comp] = chan;
@@ -262,39 +211,6 @@ gather_ps_store_output(nir_builder *b, nir_intrinsic_instr *intrin, lower_ps_sta
    return true;
 }
 
-static bool
-lower_ps_load_sample_mask_in(nir_builder *b, nir_intrinsic_instr *intrin, lower_ps_state *s)
-{
-   /* Section 15.2.2 (Shader Inputs) of the OpenGL 4.5 (Core Profile) spec
-    * says:
-    *
-    *    "When per-sample shading is active due to the use of a fragment
-    *     input qualified by sample or due to the use of the gl_SampleID
-    *     or gl_SamplePosition variables, only the bit for the current
-    *     sample is set in gl_SampleMaskIn. When state specifies multiple
-    *     fragment shader invocations for a given fragment, the sample
-    *     mask for any single fragment shader invocation may specify a
-    *     subset of the covered samples for the fragment. In this case,
-    *     the bit corresponding to each covered sample will be set in
-    *     exactly one fragment shader invocation."
-    *
-    * The samplemask loaded by hardware is always the coverage of the
-    * entire pixel/fragment, so mask bits out based on the sample ID.
-    */
-
-   b->cursor = nir_before_instr(&intrin->instr);
-
-   uint32_t ps_iter_mask = ac_get_ps_iter_mask(s->options->ps_iter_samples);
-   nir_def *sampleid = nir_load_sample_id(b);
-   nir_def *submask = nir_ishl(b, nir_imm_int(b, ps_iter_mask), sampleid);
-
-   nir_def *sample_mask = nir_load_sample_mask_in(b);
-   nir_def *replacement = nir_iand(b, sample_mask, submask);
-
-   nir_def_replace(&intrin->def, replacement);
-   return true;
-}
-
 static bool
 lower_ps_intrinsic(nir_builder *b, nir_instr *instr, void *state)
 {
@@ -314,10 +230,6 @@ lower_ps_intrinsic(nir_builder *b, nir_instr *instr, void *state)
       if (s->lower_load_barycentric)
          return lower_ps_load_barycentric(b, intrin, s);
       break;
-   case nir_intrinsic_load_sample_mask_in:
-      if (s->options->ps_iter_samples > 1)
-         return lower_ps_load_sample_mask_in(b, intrin, s);
-      break;
    default:
       break;
    }
@@ -325,35 +237,6 @@ lower_ps_intrinsic(nir_builder *b, nir_instr *instr, void *state)
    return false;
 }
 
-static void
-emit_ps_color_clamp_and_alpha_test(nir_builder *b, lower_ps_state *s)
-{
-   u_foreach_bit (slot, s->colors_written) {
-      if (s->options->clamp_color) {
-         for (int i = 0; i < 4; i++) {
-            if (s->color[slot][i])
-               s->color[slot][i] = nir_fsat(b, s->color[slot][i]);
-         }
-      }
-
-      if (s->options->alpha_to_one)
-         s->color[slot][3] = nir_imm_floatN_t(b, 1, nir_alu_type_get_type_size(s->color_type[slot]));
-
-      if (slot == 0) {
-         if (s->options->alpha_func == COMPARE_FUNC_ALWAYS) {
-            /* always pass, do nothing */
-         } else if (s->options->alpha_func == COMPARE_FUNC_NEVER) {
-            nir_discard(b);
-         } else if (s->color[slot][3]) {
-            nir_def *ref = nir_load_alpha_reference_amd(b);
-            nir_def *cond =
-               nir_compare_func(b, s->options->alpha_func, s->color[slot][3], ref);
-            nir_discard_if(b, nir_inot(b, cond));
-         }
-      }
-   }
-}
-
 static void
 emit_ps_mrtz_export(nir_builder *b, lower_ps_state *s, nir_def *mrtz_alpha)
 {
@@ -750,7 +633,10 @@ export_ps_outputs(nir_builder *b, lower_ps_state *s)
    if (!s->options->no_depth_export && s->options->alpha_to_coverage_via_mrtz)
       mrtz_alpha = s->color[0][3];
 
-   emit_ps_color_clamp_and_alpha_test(b, s);
+   u_foreach_bit (slot, s->colors_written) {
+      if (s->options->alpha_to_one)
+         s->color[slot][3] = nir_imm_floatN_t(b, 1, nir_alu_type_get_type_size(s->color_type[slot]));
+   }
 
    if (!s->options->no_depth_export)
       emit_ps_mrtz_export(b, s, mrtz_alpha);
@@ -833,8 +719,9 @@ export_ps_outputs(nir_builder *b, lower_ps_state *s)
 }
 
 void
-ac_nir_lower_ps(nir_shader *nir, const ac_nir_lower_ps_options *options)
+ac_nir_lower_ps_late(nir_shader *nir, const ac_nir_lower_ps_late_options *options)
 {
+   assert(nir->info.stage == MESA_SHADER_FRAGMENT);
    nir_function_impl *impl = nir_shader_get_entrypoint(nir);
 
    nir_builder builder = nir_builder_create(impl);

src/amd/common/ac_shader_util.c

@@ -359,6 +359,10 @@ unsigned ac_get_cb_shader_mask(unsigned spi_shader_col_format)
 {
    unsigned i, cb_shader_mask = 0;
 
+   /* If the format is ~0, it means we want a full mask. */
+   if (spi_shader_col_format == ~0)
+      return ~0;
+
    for (i = 0; i < 8; i++) {
       switch ((spi_shader_col_format >> (i * 4)) & 0xf) {
       case V_028714_SPI_SHADER_ZERO:

src/amd/common/meson.build

@@ -95,7 +95,8 @@ amd_common_files = files(
   'ac_nir_lower_tess_io_to_mem.c',
   'ac_nir_lower_tex.c',
   'ac_nir_lower_ngg.c',
-  'ac_nir_lower_ps.c',
+  'ac_nir_lower_ps_early.c',
+  'ac_nir_lower_ps_late.c',
   'ac_nir_meta.h',
   'ac_nir_meta_cs_blit.c',
   'ac_nir_meta_cs_clear_copy_buffer.c',

src/amd/vulkan/radv_pipeline.c

@@ -462,42 +462,48 @@ radv_postprocess_nir(struct radv_device *device, const struct radv_graphics_stat
          NIR_PASS_V(stage->nir, ac_nir_lower_legacy_gs, false, false, &gs_out_info);
       }
    } else if (stage->stage == MESA_SHADER_FRAGMENT) {
-      ac_nir_lower_ps_options options = {
+      ac_nir_lower_ps_early_options early_options = {
+         .alpha_func = COMPARE_FUNC_ALWAYS,
+         .spi_shader_col_format_hint = ~0,
+      };
+
+      ac_nir_lower_ps_late_options late_options = {
          .gfx_level = gfx_level,
          .family = pdev->info.family,
          .use_aco = !radv_use_llvm_for_stage(pdev, stage->stage),
-         .uses_discard = true,
-         .alpha_func = COMPARE_FUNC_ALWAYS,
-         .no_color_export = stage->info.ps.has_epilog,
-         .no_depth_export = stage->info.ps.exports_mrtz_via_epilog,
-
          .bc_optimize_for_persp = G_0286CC_PERSP_CENTER_ENA(stage->info.ps.spi_ps_input_ena) &&
                                   G_0286CC_PERSP_CENTROID_ENA(stage->info.ps.spi_ps_input_ena),
          .bc_optimize_for_linear = G_0286CC_LINEAR_CENTER_ENA(stage->info.ps.spi_ps_input_ena) &&
                                    G_0286CC_LINEAR_CENTROID_ENA(stage->info.ps.spi_ps_input_ena),
+         .uses_discard = true,
+         .no_color_export = stage->info.ps.has_epilog,
+         .no_depth_export = stage->info.ps.exports_mrtz_via_epilog,
+
       };
 
-      if (!options.no_color_export) {
-         options.dual_src_blend_swizzle = gfx_state->ps.epilog.mrt0_is_dual_src && gfx_level >= GFX11;
-         options.color_is_int8 = gfx_state->ps.epilog.color_is_int8;
-         options.color_is_int10 = gfx_state->ps.epilog.color_is_int10;
-         options.enable_mrt_output_nan_fixup =
+      if (!late_options.no_color_export) {
+         late_options.dual_src_blend_swizzle = gfx_state->ps.epilog.mrt0_is_dual_src && gfx_level >= GFX11;
+         late_options.color_is_int8 = gfx_state->ps.epilog.color_is_int8;
+         late_options.color_is_int10 = gfx_state->ps.epilog.color_is_int10;
+         late_options.enable_mrt_output_nan_fixup =
             gfx_state->ps.epilog.enable_mrt_output_nan_fixup && !stage->nir->info.internal;
          /* Need to filter out unwritten color slots. */
-         options.spi_shader_col_format = gfx_state->ps.epilog.spi_shader_col_format & stage->info.ps.colors_written;
-         options.alpha_to_one = gfx_state->ps.epilog.alpha_to_one;
+         early_options.spi_shader_col_format_hint = late_options.spi_shader_col_format =
+            gfx_state->ps.epilog.spi_shader_col_format & stage->info.ps.colors_written;
+         late_options.alpha_to_one = gfx_state->ps.epilog.alpha_to_one;
       }
 
-      if (!options.no_depth_export) {
+      if (!late_options.no_depth_export) {
          /* Compared to gfx_state.ps.alpha_to_coverage_via_mrtz,
           * radv_shader_info.ps.writes_mrt0_alpha need any depth/stencil/sample_mask exist.
           * ac_nir_lower_ps() require this field to reflect whether alpha via mrtz is really
           * present.
           */
-         options.alpha_to_coverage_via_mrtz = stage->info.ps.writes_mrt0_alpha;
+         early_options.keep_alpha_for_mrtz = late_options.alpha_to_coverage_via_mrtz = stage->info.ps.writes_mrt0_alpha;
       }
 
-      NIR_PASS_V(stage->nir, ac_nir_lower_ps, &options);
+      NIR_PASS_V(stage->nir, ac_nir_lower_ps_early, &early_options);
+      NIR_PASS_V(stage->nir, ac_nir_lower_ps_late, &late_options);
    }
 
    if (radv_shader_should_clear_lds(device, stage->nir)) {

src/gallium/drivers/radeonsi/si_shader.c

@@ -2523,33 +2523,41 @@ static struct nir_shader *si_get_nir_shader(struct si_shader *shader, struct si_
       shader->info.num_ps_inputs = info.num_inputs;
       shader->info.ps_colors_read = info.colors_read;
 
-      ac_nir_lower_ps_options options = {
+      ac_nir_lower_ps_early_options early_options = {
+         .force_persp_sample_interp = key->ps.part.prolog.force_persp_sample_interp,
+         .force_linear_sample_interp = key->ps.part.prolog.force_linear_sample_interp,
+         .force_persp_center_interp = key->ps.part.prolog.force_persp_center_interp,
+         .force_linear_center_interp = key->ps.part.prolog.force_linear_center_interp,
+         .ps_iter_samples = 1 << key->ps.part.prolog.samplemask_log_ps_iter,
+
+         .clamp_color = key->ps.part.epilog.clamp_color,
+         .alpha_test_alpha_to_one = key->ps.part.epilog.alpha_to_one,
+         .alpha_func = key->ps.part.epilog.alpha_func,
+         .keep_alpha_for_mrtz = key->ps.part.epilog.alpha_to_coverage_via_mrtz,
+         .spi_shader_col_format_hint = key->ps.part.epilog.spi_shader_col_format,
+         .kill_z = key->ps.part.epilog.kill_z,
+         .kill_stencil = key->ps.part.epilog.kill_stencil,
+         .kill_samplemask = key->ps.part.epilog.kill_samplemask,
+      };
+
+      NIR_PASS_V(nir, ac_nir_lower_ps_early, &early_options);
+
+      ac_nir_lower_ps_late_options late_options = {
          .gfx_level = sel->screen->info.gfx_level,
          .family = sel->screen->info.family,
          .use_aco = nir->info.use_aco_amd,
+         .bc_optimize_for_persp = key->ps.part.prolog.bc_optimize_for_persp,
+         .bc_optimize_for_linear = key->ps.part.prolog.bc_optimize_for_linear,
          .uses_discard = si_shader_uses_discard(shader),
          .alpha_to_coverage_via_mrtz = key->ps.part.epilog.alpha_to_coverage_via_mrtz,
          .dual_src_blend_swizzle = key->ps.part.epilog.dual_src_blend_swizzle,
          .spi_shader_col_format = key->ps.part.epilog.spi_shader_col_format,
          .color_is_int8 = key->ps.part.epilog.color_is_int8,
          .color_is_int10 = key->ps.part.epilog.color_is_int10,
-         .clamp_color = key->ps.part.epilog.clamp_color,
          .alpha_to_one = key->ps.part.epilog.alpha_to_one,
-         .alpha_func = key->ps.part.epilog.alpha_func,
-         .kill_z = key->ps.part.epilog.kill_z,
-         .kill_stencil = key->ps.part.epilog.kill_stencil,
-         .kill_samplemask = key->ps.part.epilog.kill_samplemask,
-
-         .bc_optimize_for_persp = key->ps.part.prolog.bc_optimize_for_persp,
-         .bc_optimize_for_linear = key->ps.part.prolog.bc_optimize_for_linear,
-         .force_persp_sample_interp = key->ps.part.prolog.force_persp_sample_interp,
-         .force_linear_sample_interp = key->ps.part.prolog.force_linear_sample_interp,
-         .force_persp_center_interp = key->ps.part.prolog.force_persp_center_interp,
-         .force_linear_center_interp = key->ps.part.prolog.force_linear_center_interp,
-         .ps_iter_samples = 1 << key->ps.part.prolog.samplemask_log_ps_iter,
       };
 
-      NIR_PASS_V(nir, ac_nir_lower_ps, &options);
+      NIR_PASS_V(nir, ac_nir_lower_ps_late, &late_options);
 
       if (key->ps.part.prolog.poly_stipple)
          NIR_PASS_V(nir, si_nir_emit_polygon_stipple, args);