asahi: Add alpha-to-coverage (and alpha-to-one) lowering

This should probably be shared code but meh.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23480>

src/asahi/lib/agx_nir_lower_alpha.c

@@ -0,0 +1,122 @@
+/*
+ * Copyright 2023 Alyssa Rosenzweig
+ * Copyright 2021 Intel Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "agx_tilebuffer.h"
+#include "nir_builder.h"
+
+#define ALL_SAMPLES (0xFF)
+
+/*
+ * Lower alpha-to-coverage to sample_mask and some math. May run on either a
+ * monolithic pixel shader or a fragment epilogue.
+ */
+void
+agx_nir_lower_alpha_to_coverage(nir_shader *shader, uint8_t nr_samples)
+{
+   /* nir_lower_io_to_temporaries ensures that stores are in the last block */
+   nir_function_impl *impl = nir_shader_get_entrypoint(shader);
+   nir_block *block = nir_impl_last_block(impl);
+
+   nir_builder _b;
+   nir_builder_init(&_b, impl);
+   nir_builder *b = &_b;
+
+   /* The store is probably at the end of the block, so search in reverse. */
+   nir_intrinsic_instr *store = NULL;
+   nir_foreach_instr_reverse(instr, block) {
+      if (instr->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+      if (intr->intrinsic != nir_intrinsic_store_output)
+         continue;
+
+      nir_io_semantics sem = nir_intrinsic_io_semantics(intr);
+      if (sem.location != FRAG_RESULT_DATA0)
+         continue;
+      if (sem.dual_source_blend_index != 0)
+         continue;
+
+      store = intr;
+      break;
+   }
+
+   /* If render target 0 isn't written, the alpha value input to
+    * alpha-to-coverage is undefined. We assume that the alpha would be 1.0,
+    * which would effectively disable alpha-to-coverage, skipping the lowering.
+    */
+   if (!store)
+      return;
+
+   /* Similarly, if there are less than 4 components, alpha is undefined */
+   nir_ssa_def *rgba = store->src[0].ssa;
+   if (rgba->num_components < 4)
+      return;
+
+   b->cursor = nir_before_instr(&store->instr);
+
+   /* Calculate a coverage mask (alpha * nr_samples) bits set. The way we do
+    * this isn't particularly clever:
+    *
+    *    # of bits = (unsigned int) (alpha * nr_samples)
+    *    mask = (1 << (# of bits)) - 1
+    */
+   nir_ssa_def *alpha = nir_channel(b, rgba, 3);
+   nir_ssa_def *bits = nir_f2u32(b, nir_fmul_imm(b, alpha, nr_samples));
+   nir_ssa_def *mask =
+      nir_iadd_imm(b, nir_ishl(b, nir_imm_intN_t(b, 1, 16), bits), -1);
+
+   /* Discard samples that aren't covered */
+   nir_sample_mask_agx(b, nir_imm_intN_t(b, ALL_SAMPLES, 16), mask);
+   shader->info.outputs_written |= BITFIELD64_BIT(FRAG_RESULT_SAMPLE_MASK);
+}
+
+/*
+ * Modify the inputs to store_output instructions in a pixel shader when
+ * alpha-to-one is used. May run on either a monolithic pixel shader or a
+ * fragment epilogue.
+ */
+void
+agx_nir_lower_alpha_to_one(nir_shader *shader)
+{
+   /* nir_lower_io_to_temporaries ensures that stores are in the last block */
+   nir_function_impl *impl = nir_shader_get_entrypoint(shader);
+   nir_block *block = nir_impl_last_block(impl);
+
+   nir_builder _b;
+   nir_builder_init(&_b, impl);
+   nir_builder *b = &_b;
+
+   nir_foreach_instr(instr, block) {
+      if (instr->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+      if (intr->intrinsic != nir_intrinsic_store_output)
+         continue;
+
+      /* The OpenGL spec is a bit confusing here, but seemingly alpha-to-one
+       * applies to all render targets. Piglit
+       * ext_framebuffer_multisample-draw-buffers-alpha-to-one checks this.
+       *
+       * Even more confusingly, it seems to apply to dual-source blending too.
+       * ext_framebuffer_multisample-alpha-to-one-dual-src-blend checks this.
+       */
+      nir_io_semantics sem = nir_intrinsic_io_semantics(intr);
+      if (sem.location < FRAG_RESULT_DATA0)
+         continue;
+
+      nir_ssa_def *rgba = intr->src[0].ssa;
+      if (rgba->num_components < 4)
+         continue;
+
+      b->cursor = nir_before_instr(instr);
+      nir_ssa_def *rgb1 = nir_vector_insert_imm(
+         b, rgba, nir_imm_floatN_t(b, 1.0, rgba->bit_size), 3);
+
+      nir_instr_rewrite_src_ssa(instr, &intr->src[0], rgb1);
+   }
+}

src/asahi/lib/agx_tilebuffer.h

@@ -62,6 +62,11 @@ bool agx_nir_lower_monolithic_msaa(struct nir_shader *shader,
 
 bool agx_nir_lower_sample_intrinsics(struct nir_shader *shader);
 
+void agx_nir_lower_alpha_to_coverage(struct nir_shader *shader,
+                                     uint8_t nr_samples);
+
+void agx_nir_lower_alpha_to_one(struct nir_shader *shader);
+
 void agx_usc_tilebuffer(struct agx_usc_builder *b,
                         struct agx_tilebuffer_layout *tib);
 

src/asahi/lib/meson.build

@@ -11,6 +11,7 @@ libasahi_lib_files = files(
   'agx_formats.c',
   'agx_meta.c',
   'agx_tilebuffer.c',
+  'agx_nir_lower_alpha.c',
   'agx_nir_lower_msaa.c',
   'agx_nir_lower_sample_intrinsics.c',
   'agx_nir_lower_tilebuffer.c',