diff --git a/src/amd/common/ac_nir.h b/src/amd/common/ac_nir.h
index ed7906afaaf..80ce91d39a0 100644
--- a/src/amd/common/ac_nir.h
+++ b/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;
diff --git a/src/amd/common/ac_nir_lower_ps_early.c b/src/amd/common/ac_nir_lower_ps_early.c
new file mode 100644
index 00000000000..34227ad50fd
--- /dev/null
+++ b/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);
+}
diff --git a/src/amd/common/ac_nir_lower_ps.c b/src/amd/common/ac_nir_lower_ps_late.c
similarity index 82%
rename from src/amd/common/ac_nir_lower_ps.c
rename to src/amd/common/ac_nir_lower_ps_late.c
index 447ff6bbbc4..e37715ae83b 100644
--- a/src/amd/common/ac_nir_lower_ps.c
+++ b/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);
diff --git a/src/amd/common/ac_shader_util.c b/src/amd/common/ac_shader_util.c
index 50e873b2f7e..3bb2b4fb5b9 100644
--- a/src/amd/common/ac_shader_util.c
+++ b/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:
diff --git a/src/amd/common/meson.build b/src/amd/common/meson.build
index 8717245a4d0..d16179d229d 100644
--- a/src/amd/common/meson.build
+++ b/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',
diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index a0921695ef8..6addff1d4af 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/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)) {
diff --git a/src/gallium/drivers/radeonsi/si_shader.c b/src/gallium/drivers/radeonsi/si_shader.c
index 7be1e9d3c6f..5bc4a38c3a3 100644
--- a/src/gallium/drivers/radeonsi/si_shader.c
+++ b/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);