ac: use Z_EXPORT_FORMAT=32_AR for Z + Alpha mrtz exports

This should be faster than 32_ABGR.

Also, stencil exports are changed from UINT16_ABGR to 32_GR,
which should have no effect on performance.

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

src/amd/common/ac_shader_util.c

@@ -17,20 +17,29 @@
 unsigned ac_get_spi_shader_z_format(bool writes_z, bool writes_stencil, bool writes_samplemask,
                                     bool writes_mrt0_alpha)
 {
-   if (writes_z || writes_mrt0_alpha) {
-      /* Z needs 32 bits. */
-      if (writes_samplemask || writes_mrt0_alpha)
+   /* RGBA = (Z, stencil, samplemask, mrt0_alpha).
+    * Both stencil and sample mask need only 16 bits.
+    */
+   if (writes_mrt0_alpha) {
+      if (writes_stencil || writes_samplemask)
          return V_028710_SPI_SHADER_32_ABGR;
-      else if (writes_stencil)
-         return V_028710_SPI_SHADER_32_GR;
       else
-         return V_028710_SPI_SHADER_32_R;
-   } else if (writes_stencil || writes_samplemask) {
-      /* Both stencil and sample mask need only 16 bits. */
-      return V_028710_SPI_SHADER_UINT16_ABGR;
-   } else {
-      return V_028710_SPI_SHADER_ZERO;
+         return V_028710_SPI_SHADER_32_AR;
    }
+
+   if (writes_samplemask) {
+      if (writes_z)
+         return V_028710_SPI_SHADER_32_ABGR;
+      else
+         return V_028710_SPI_SHADER_UINT16_ABGR;
+   }
+
+   if (writes_stencil)
+      return V_028710_SPI_SHADER_32_GR;
+   else if (writes_z)
+      return V_028710_SPI_SHADER_32_R;
+   else
+      return V_028710_SPI_SHADER_ZERO;
 }
 
 unsigned ac_get_cb_shader_mask(unsigned spi_shader_col_format)

src/amd/common/nir/ac_nir_lower_ps_late.c

@@ -275,18 +275,28 @@ emit_ps_mrtz_export(nir_builder *b, lower_ps_state *s, nir_def *mrtz_alpha)
       }
 
       if (s->stencil) {
+         assert(format == V_028710_SPI_SHADER_32_GR ||
+                format == V_028710_SPI_SHADER_32_ABGR);
          outputs[1] = s->stencil;
          write_mask |= 0x2;
       }
 
       if (s->sample_mask) {
+         assert(format == V_028710_SPI_SHADER_32_ABGR);
          outputs[2] = s->sample_mask;
          write_mask |= 0x4;
       }
 
       if (mrtz_alpha) {
-         outputs[3] = mrtz_alpha;
-         write_mask |= 0x8;
+         assert(format == V_028710_SPI_SHADER_32_AR ||
+                format == V_028710_SPI_SHADER_32_ABGR);
+         if (format == V_028710_SPI_SHADER_32_AR && s->options->gfx_level >= GFX10) {
+            outputs[1] = mrtz_alpha;
+            write_mask |= 0x2;
+         } else {
+            outputs[3] = mrtz_alpha;
+            write_mask |= 0x8;
+         }
       }
    }
 

src/amd/compiler/aco_instruction_selection.cpp

@@ -10626,8 +10626,12 @@ export_fs_mrtz(isel_context* ctx, const struct aco_ps_epilog_info* info, Temp de
       values[i] = Operand(v1);
    }
 
+   const unsigned format =
+      ac_get_spi_shader_z_format(depth.id(), stencil.id(), samplemask.id(), alpha.id());
+   assert(format != V_028710_SPI_SHADER_ZERO);
+
    /* Both stencil and sample mask only need 16-bits. */
-   if (!depth.id() && !alpha.id() && (stencil.id() || samplemask.id())) {
+   if (format == V_028710_SPI_SHADER_UINT16_ABGR) {
       compr = ctx->program->gfx_level < GFX11; /* COMPR flag */
 
       if (stencil.id()) {
@@ -10648,16 +10652,19 @@ export_fs_mrtz(isel_context* ctx, const struct aco_ps_epilog_info* info, Temp de
       }
 
       if (stencil.id()) {
+         assert(format == V_028710_SPI_SHADER_32_GR || format == V_028710_SPI_SHADER_32_ABGR);
          values[1] = Operand(stencil);
          enabled_channels |= 0x2;
       }
 
       if (samplemask.id()) {
+         assert(format == V_028710_SPI_SHADER_32_ABGR);
          values[2] = Operand(samplemask);
          enabled_channels |= 0x4;
       }
 
       if (alpha.id()) {
+         assert(format == V_028710_SPI_SHADER_32_AR || format == V_028710_SPI_SHADER_32_ABGR);
          assert(ctx->program->gfx_level >= GFX11 || info->alpha_to_one);
          values[3] = Operand(alpha);
          enabled_channels |= 0x8;

src/amd/llvm/ac_llvm_build.c

@@ -3546,16 +3546,26 @@ void ac_export_mrt_z(struct ac_llvm_context *ctx, LLVMValueRef depth, LLVMValueR
          mask |= 0x1;
       }
       if (stencil) {
+         assert(format == V_028710_SPI_SHADER_32_GR ||
+                format == V_028710_SPI_SHADER_32_ABGR);
          args->out[1] = stencil;
          mask |= 0x2;
       }
       if (samplemask) {
+         assert(format == V_028710_SPI_SHADER_32_ABGR);
          args->out[2] = samplemask;
          mask |= 0x4;
       }
       if (mrt0_alpha) {
-         args->out[3] = mrt0_alpha;
-         mask |= 0x8;
+         assert(format == V_028710_SPI_SHADER_32_AR ||
+                format == V_028710_SPI_SHADER_32_ABGR);
+         if (format == V_028710_SPI_SHADER_32_AR && ctx->gfx_level >= GFX10) {
+            args->out[1] = mrt0_alpha;
+            mask |= 0x2;
+         } else {
+            args->out[3] = mrt0_alpha;
+            mask |= 0x8;
+         }
       }
    }