radv: prepare the PS epilog key for exporting MRTZ on RDNA3

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/26413>

src/amd/compiler/aco_shader_info.h

@@ -73,11 +73,12 @@ struct aco_ps_epilog_info {
 
    bool mrt0_is_dual_src;
 
+   bool alpha_to_coverage_via_mrtz;
+
    /* OpenGL only */
    uint16_t color_types;
    bool clamp_color;
    bool alpha_to_one;
-   bool alpha_to_coverage_via_mrtz;
    bool skip_null_export;
    unsigned broadcast_last_cbuf;
    enum compare_func alpha_func;

src/amd/vulkan/radv_aco_shader_info.h

@@ -115,6 +115,7 @@ radv_aco_convert_ps_epilog_key(struct aco_ps_epilog_info *aco_info, const struct
    ASSIGN_FIELD(color_is_int8);
    ASSIGN_FIELD(color_is_int10);
    ASSIGN_FIELD(mrt0_is_dual_src);
+   ASSIGN_FIELD(alpha_to_coverage_via_mrtz);
 
    memcpy(aco_info->colors, radv_args->colors, sizeof(aco_info->colors));
 

src/amd/vulkan/radv_cmd_buffer.c

@@ -4174,6 +4174,7 @@ radv_emit_color_blend(struct radv_cmd_buffer *cmd_buffer)
 static struct radv_shader_part *
 lookup_ps_epilog(struct radv_cmd_buffer *cmd_buffer)
 {
+   const struct radv_shader *ps = cmd_buffer->state.shaders[MESA_SHADER_FRAGMENT];
    const struct radv_rendering_state *render = &cmd_buffer->state.render;
    const struct radv_dynamic_state *d = &cmd_buffer->state.dynamic;
    struct radv_device *device = cmd_buffer->device;
@@ -4207,6 +4208,14 @@ lookup_ps_epilog(struct radv_cmd_buffer *cmd_buffer)
       state.need_src_alpha |= 0x1;
    }
 
+   if (ps && ps->info.ps.exports_mrtz_via_epilog) {
+      assert(device->physical_device->rad_info.gfx_level >= GFX11);
+      state.export_depth = ps->info.ps.writes_z;
+      state.export_stencil = ps->info.ps.writes_stencil;
+      state.export_sample_mask = ps->info.ps.writes_sample_mask;
+      state.alpha_to_coverage_via_mrtz = d->vk.ms.alpha_to_coverage_enable;
+   }
+
    struct radv_ps_epilog_key key = radv_generate_ps_epilog_key(device, &state, true);
    return radv_shader_part_cache_get(device, &device->ps_epilogs, &cmd_buffer->ps_epilogs, &key);
 }

src/amd/vulkan/radv_pipeline_graphics.c

@@ -1738,6 +1738,10 @@ radv_generate_ps_epilog_key(const struct radv_device *device, const struct radv_
    key.color_is_int10 = device->physical_device->rad_info.gfx_level < GFX8 ? is_int10 : 0;
    key.enable_mrt_output_nan_fixup = device->instance->enable_mrt_output_nan_fixup ? is_float32 : 0;
    key.mrt0_is_dual_src = state->mrt0_is_dual_src;
+   key.export_depth = state->export_depth;
+   key.export_stencil = state->export_stencil;
+   key.export_sample_mask = state->export_sample_mask;
+   key.alpha_to_coverage_via_mrtz = state->alpha_to_coverage_via_mrtz;
 
    return key;
 }

src/amd/vulkan/radv_private.h

@@ -2019,6 +2019,10 @@ struct radv_ps_epilog_state {
    uint32_t color_blend_enable;
 
    bool mrt0_is_dual_src;
+   bool export_depth;
+   bool export_stencil;
+   bool export_sample_mask;
+   bool alpha_to_coverage_via_mrtz;
    uint8_t need_src_alpha;
 };
 

src/amd/vulkan/radv_shader.h

@@ -74,6 +74,10 @@ struct radv_ps_epilog_key {
    uint8_t enable_mrt_output_nan_fixup;
 
    bool mrt0_is_dual_src;
+   bool export_depth;
+   bool export_stencil;
+   bool export_sample_mask;
+   bool alpha_to_coverage_via_mrtz;
 };
 
 struct radv_pipeline_key {