radeonsi/gfx11: export alpha through mrtz for alpha-to-coverage if mrtz is there

If both mrtz and alpha-to-coverage are enabled, the alpha channel must
be exported through mrtz.

Acked-by: Pierre-Eric Pelloux-Prayer <pierre-eric.pelloux-prayer@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16328>

src/amd/llvm/ac_llvm_build.c

@@ -4232,7 +4232,8 @@ LLVMValueRef ac_build_call(struct ac_llvm_context *ctx, LLVMValueRef func, LLVMV
 }
 
 void ac_export_mrt_z(struct ac_llvm_context *ctx, LLVMValueRef depth, LLVMValueRef stencil,
-                     LLVMValueRef samplemask, bool is_last, struct ac_export_args *args)
+                     LLVMValueRef samplemask, LLVMValueRef mrtz_alpha, bool is_last,
+                     struct ac_export_args *args)
 {
    unsigned mask = 0;
    unsigned format = ac_get_spi_shader_z_format(depth != NULL, stencil != NULL, samplemask != NULL);
@@ -4271,6 +4272,17 @@ void ac_export_mrt_z(struct ac_llvm_context *ctx, LLVMValueRef depth, LLVMValueR
          args->out[1] = samplemask;
          mask |= ctx->chip_class >= GFX11 ? 0x2 : 0xc;
       }
+      if (mrtz_alpha) {
+         /* MRT0 alpha should be in Y[31:16] if alpha-to-coverage is enabled and MRTZ is present. */
+         assert(ctx->chip_class >= GFX11);
+         mrtz_alpha = LLVMBuildFPTrunc(ctx->builder, mrtz_alpha, ctx->f16, "");
+         mrtz_alpha = ac_to_integer(ctx, mrtz_alpha);
+         mrtz_alpha = LLVMBuildZExt(ctx->builder, mrtz_alpha, ctx->i32, "");
+         mrtz_alpha = LLVMBuildShl(ctx->builder, mrtz_alpha, LLVMConstInt(ctx->i32, 16, 0), "");
+         args->out[1] = LLVMBuildOr(ctx->builder, ac_to_integer(ctx, args->out[1]), mrtz_alpha, "");
+         args->out[1] = ac_to_float(ctx, args->out[1]);
+         mask |= 0x2;
+      }
    } else {
       if (depth) {
          args->out[0] = depth;
@@ -4284,6 +4296,10 @@ void ac_export_mrt_z(struct ac_llvm_context *ctx, LLVMValueRef depth, LLVMValueR
          args->out[2] = samplemask;
          mask |= 0x4;
       }
+      if (mrtz_alpha) {
+         args->out[3] = mrtz_alpha;
+         mask |= 0x8;
+      }
    }
 
    /* GFX6 (except OLAND and HAINAN) has a bug that it only looks

src/amd/llvm/ac_llvm_build.h

@@ -555,7 +555,8 @@ LLVMValueRef ac_build_atomic_cmp_xchg(struct ac_llvm_context *ctx, LLVMValueRef
                                       LLVMValueRef cmp, LLVMValueRef val, const char *sync_scope);
 
 void ac_export_mrt_z(struct ac_llvm_context *ctx, LLVMValueRef depth, LLVMValueRef stencil,
-                     LLVMValueRef samplemask, bool is_last, struct ac_export_args *args);
+                     LLVMValueRef samplemask, LLVMValueRef mrtz_alpha, bool is_last,
+                     struct ac_export_args *args);
 
 void ac_build_sendmsg_gs_alloc_req(struct ac_llvm_context *ctx, LLVMValueRef wave_id,
                                    LLVMValueRef vtx_cnt, LLVMValueRef prim_cnt);

src/amd/vulkan/radv_nir_to_llvm.c

@@ -1775,7 +1775,7 @@ radv_export_mrt_z(struct radv_shader_context *ctx, LLVMValueRef depth, LLVMValue
 {
    struct ac_export_args args;
 
-   ac_export_mrt_z(&ctx->ac, depth, stencil, samplemask, true, &args);
+   ac_export_mrt_z(&ctx->ac, depth, stencil, samplemask, NULL, true, &args);
 
    ac_build_export(&ctx->ac, &args);
 }

src/gallium/drivers/radeonsi/si_shader.c

@@ -1282,6 +1282,7 @@ static void si_dump_shader_key(const struct si_shader *shader, FILE *f)
       fprintf(f, "  epilog.last_cbuf = %u\n", key->ps.part.epilog.last_cbuf);
       fprintf(f, "  epilog.alpha_func = %u\n", key->ps.part.epilog.alpha_func);
       fprintf(f, "  epilog.alpha_to_one = %u\n", key->ps.part.epilog.alpha_to_one);
+      fprintf(f, "  epilog.alpha_to_coverage_via_mrtz = %u\n", key->ps.part.epilog.alpha_to_coverage_via_mrtz);
       fprintf(f, "  epilog.clamp_color = %u\n", key->ps.part.epilog.clamp_color);
       fprintf(f, "  mono.poly_line_smoothing = %u\n", key->ps.mono.poly_line_smoothing);
       fprintf(f, "  mono.interpolate_at_sample_force_center = %u\n",

src/gallium/drivers/radeonsi/si_shader.h

@@ -576,6 +576,7 @@ struct si_ps_epilog_bits {
    unsigned last_cbuf : 3;
    unsigned alpha_func : 3;
    unsigned alpha_to_one : 1;
+   unsigned alpha_to_coverage_via_mrtz : 1;  /* gfx11+ */
    unsigned clamp_color : 1;
 };
 

src/gallium/drivers/radeonsi/si_shader_llvm_ps.c

@@ -851,7 +851,7 @@ void si_llvm_build_ps_epilog(struct si_shader_context *ctx, union si_shader_part
 {
    int i;
    struct si_ps_exports exp = {};
-   LLVMValueRef color[8][4];
+   LLVMValueRef color[8][4] = {};
 
    memset(&ctx->args, 0, sizeof(ctx->args));
 
@@ -899,10 +899,14 @@ void si_llvm_build_ps_epilog(struct si_shader_context *ctx, union si_shader_part
       si_llvm_build_clamp_alpha_test(ctx, color[write_i], write_i);
    }
 
+   LLVMValueRef mrtz_alpha =
+      key->ps_epilog.states.alpha_to_coverage_via_mrtz ? color[0][3] : NULL;
+
    /* Prepare the mrtz export. */
    if (key->ps_epilog.writes_z ||
        key->ps_epilog.writes_stencil ||
-       key->ps_epilog.writes_samplemask) {
+       key->ps_epilog.writes_samplemask ||
+       mrtz_alpha) {
       LLVMValueRef depth = NULL, stencil = NULL, samplemask = NULL;
       unsigned vgpr_index = ctx->args.num_sgprs_used +
                             util_bitcount(key->ps_epilog.colors_written) * 4;
@@ -914,7 +918,8 @@ void si_llvm_build_ps_epilog(struct si_shader_context *ctx, union si_shader_part
       if (key->ps_epilog.writes_samplemask)
          samplemask = LLVMGetParam(ctx->main_fn, vgpr_index++);
 
-      ac_export_mrt_z(&ctx->ac, depth, stencil, samplemask, false, &exp.args[exp.num++]);
+      ac_export_mrt_z(&ctx->ac, depth, stencil, samplemask, mrtz_alpha, false,
+                      &exp.args[exp.num++]);
    }
 
    /* Prepare color exports. */

src/gallium/drivers/radeonsi/si_state_shaders.cpp

@@ -2266,8 +2266,15 @@ void si_ps_key_update_blend_rasterizer(struct si_context *sctx)
    union si_shader_key *key = &sctx->shader.ps.key;
    struct si_state_blend *blend = sctx->queued.named.blend;
    struct si_state_rasterizer *rs = sctx->queued.named.rasterizer;
+   struct si_shader_selector *ps = sctx->shader.ps.cso;
+
+   if (!ps)
+      return;
 
    key->ps.part.epilog.alpha_to_one = blend->alpha_to_one && rs->multisample_enable;
+   key->ps.part.epilog.alpha_to_coverage_via_mrtz =
+      sctx->chip_class >= GFX11 && blend->alpha_to_coverage && rs->multisample_enable &&
+      (ps->info.writes_z || ps->info.writes_stencil || ps->info.writes_samplemask);
 }
 
 void si_ps_key_update_rasterizer(struct si_context *sctx)