asahi/lib: Move alpha_to_one and alpha_to_coverage lowering to common code.

Reviewed-by: Iago Toral Quiroga <itoral@igalia.com>
Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33942>

src/compiler/nir/meson.build

@@ -115,6 +115,7 @@ files_libnir = files(
   'nir_loop_analyze.h',
   'nir_lower_alu.c',
   'nir_lower_alu_width.c',
+  'nir_lower_alpha.c',
   'nir_lower_alpha_test.c',
   'nir_lower_amul.c',
   'nir_lower_array_deref_of_vec.c',

src/compiler/nir/nir.h

@@ -5222,6 +5222,12 @@ bool nir_lower_vec_to_regs(nir_shader *shader, nir_instr_writemask_filter_cb cb,
 bool nir_lower_alpha_test(nir_shader *shader, enum compare_func func,
                           bool alpha_to_one,
                           const gl_state_index16 *alpha_ref_state_tokens);
+
+bool nir_lower_alpha_to_coverage(nir_shader *shader,
+                                 uint8_t nr_samples);
+
+bool nir_lower_alpha_to_one(nir_shader *shader);
+
 bool nir_lower_alu(nir_shader *shader);
 
 bool nir_lower_flrp(nir_shader *shader, unsigned lowering_mask,

src/compiler/nir/nir_gather_info.c

@@ -806,7 +806,7 @@ gather_intrinsic_info(nir_intrinsic_instr *instr, nir_shader *shader,
       shader->info.outputs_written |= BITFIELD64_BIT(FRAG_RESULT_SAMPLE_MASK);
       break;
 
-   case nir_intrinsic_discard_agx:
+   case nir_intrinsic_demote_samples:
       shader->info.fs.uses_discard = true;
       break;
 

src/compiler/nir/nir_intrinsics.py

@@ -1120,6 +1120,10 @@ barycentric("coord_at_offset", 3, [2])
 intrinsic("load_sample_pos_from_id", src_comp=[1], dest_comp=2,
           flags=[CAN_ELIMINATE, CAN_REORDER])
 
+# Demote a subset of samples given by a specified sample mask. This acts like a
+# per-sample demote, or an inverted accumulating gl_SampleMask write.
+intrinsic("demote_samples", src_comp=[1])
+
 intrinsic("load_persp_center_rhw_ir3", dest_comp=1,
           flags=[CAN_ELIMINATE, CAN_REORDER])
 
@@ -2173,12 +2177,6 @@ load("sysval_agx", [], [DESC_SET, BINDING, FLAGS], [CAN_REORDER, CAN_ELIMINATE])
 # documented elsewhere as they are too complicated for this comment.
 intrinsic("sample_mask_agx", src_comp=[1, 1])
 
-# Discard a subset of samples given by a specified sample mask. This acts like a
-# per-sample discard, or an inverted accumulating gl_SampleMask write. The
-# compiler will lower to sample_mask_agx, but that lowering is nontrivial as
-# sample_mask_agx also triggers depth/stencil testing.
-intrinsic("discard_agx", src_comp=[1])
-
 # For a given row of the polygon stipple given as an integer source in [0, 31],
 # load the 32-bit stipple pattern for that row.
 intrinsic("load_polygon_stipple_agx", src_comp=[1], dest_comp=1, bit_sizes=[32],

src/compiler/nir/nir_lower_alpha.c

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2023 Alyssa Rosenzweig
+ * Copyright 2021 Intel Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include "nir_builder.h"
+
+/*
+ * Lower alpha-to-coverage to sample_mask and some math. May run on either a
+ * monolithic pixel shader or a fragment epilogue.
+ */
+bool
+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);
+
+   /* 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.
+    *
+    * Similarly, if there are less than 4 components, alpha is undefined.
+    */
+   nir_def *rgba = store ? store->src[0].ssa : NULL;
+   if (!rgba || rgba->num_components < 4) {
+      return nir_no_progress(impl);
+   }
+
+   nir_builder _b = nir_builder_at(nir_before_instr(&store->instr));
+   nir_builder *b = &_b;
+
+   /* 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_def *alpha = nir_channel(b, rgba, 3);
+   nir_def *bits = nir_f2u32(b, nir_fmul_imm(b, alpha, nr_samples));
+   nir_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_demote_samples(b, nir_inot(b, mask));
+   shader->info.fs.uses_discard = true;
+   return nir_progress(true, impl, nir_metadata_control_flow);
+}
+
+/*
+ * 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.
+ */
+bool
+nir_lower_alpha_to_one(nir_shader *shader)
+{
+   bool progress = false;
+
+   /* 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_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_def *rgba = intr->src[0].ssa;
+      if (rgba->num_components < 4)
+         continue;
+
+      nir_builder b = nir_builder_at(nir_before_instr(instr));
+      nir_def *rgb1 = nir_vector_insert_imm(
+         &b, rgba, nir_imm_floatN_t(&b, 1.0, rgba->bit_size), 3);
+
+      nir_src_rewrite(&intr->src[0], rgb1);
+      progress = true;
+   }
+
+   return nir_progress(progress, impl, nir_metadata_control_flow);
+}