asahi: Add passes to lower sample intrinsics

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_sample_intrinsics.c

@@ -0,0 +1,122 @@
+/*
+ * Copyright 2023 Valve Corporation
+ * Copyright 2023 Alyssa Rosenzweig
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "agx_tilebuffer.h"
+#include "nir_builder.h"
+
+static nir_ssa_def *
+mask_by_sample_id(nir_builder *b, nir_ssa_def *mask)
+{
+   nir_ssa_def *id_mask =
+      nir_ishl(b, nir_imm_intN_t(b, 1, mask->bit_size), nir_load_sample_id(b));
+   return nir_iand(b, mask, id_mask);
+}
+
+static bool
+lower_to_sample(nir_builder *b, nir_instr *instr, void *_)
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+   b->cursor = nir_before_instr(instr);
+
+   switch (intr->intrinsic) {
+   case nir_intrinsic_load_sample_pos: {
+      /* Lower sample positions to decode the packed fixed-point register:
+       *
+       *    uint32_t packed = load_sample_positions();
+       *    uint32_t shifted = packed >> (sample_id * 8);
+       *
+       *    for (i = 0; i < 2; ++i) {
+       *       uint8_t nibble = (shifted >> (i * 4)) & 0xF;
+       *       xy[component] = ((float)nibble) / 16.0;
+       *    }
+       */
+      nir_ssa_def *packed = nir_load_sample_positions_agx(b);
+
+      /* The n'th sample is the in the n'th byte of the register */
+      nir_ssa_def *shifted = nir_ushr(
+         b, packed, nir_u2u32(b, nir_imul_imm(b, nir_load_sample_id(b), 8)));
+
+      nir_ssa_def *xy[2];
+      for (unsigned i = 0; i < 2; ++i) {
+         /* Get the appropriate nibble */
+         nir_ssa_def *nibble =
+            nir_iand_imm(b, nir_ushr_imm(b, shifted, i * 4), 0xF);
+
+         /* Convert it from fixed point to float */
+         xy[i] = nir_fmul_imm(b, nir_u2f16(b, nibble), 1.0 / 16.0);
+
+         /* Upconvert if necessary */
+         xy[i] = nir_f2fN(b, xy[i], nir_dest_bit_size(intr->dest));
+      }
+
+      /* Collect and rewrite */
+      nir_ssa_def_rewrite_uses(&intr->dest.ssa, nir_vec2(b, xy[0], xy[1]));
+      nir_instr_remove(instr);
+      return true;
+   }
+
+   case nir_intrinsic_load_sample_mask_in: {
+      /* In OpenGL, gl_SampleMaskIn is only supposed to have the single bit set
+       * of the sample currently being shaded when sample shading is used. Mask
+       * by the sample ID to make that happen.
+       */
+      b->cursor = nir_after_instr(instr);
+      nir_ssa_def *old = &intr->dest.ssa;
+      nir_ssa_def *lowered = mask_by_sample_id(b, old);
+      nir_ssa_def_rewrite_uses_after(old, lowered, lowered->parent_instr);
+      return true;
+   }
+
+   case nir_intrinsic_load_barycentric_sample: {
+      /* Lower fragment varyings with "sample" interpolation to
+       * interpolateAtSample() with the sample ID
+       */
+      b->cursor = nir_after_instr(instr);
+      nir_ssa_def *old = &intr->dest.ssa;
+
+      nir_ssa_def *lowered = nir_load_barycentric_at_sample(
+         b, nir_dest_bit_size(intr->dest), nir_load_sample_id(b),
+         .interp_mode = nir_intrinsic_interp_mode(intr));
+
+      nir_ssa_def_rewrite_uses_after(old, lowered, lowered->parent_instr);
+      return true;
+   }
+
+   default:
+      return false;
+   }
+}
+
+/*
+ * In a fragment shader using sample shading, lower intrinsics like
+ * load_sample_position to variants in terms of load_sample_id. Except for a
+ * possible API bit to force sample shading in shaders that don't otherwise need
+ * it, this pass does not depend on the shader key. In particular, it does not
+ * depend on the sample count. So it runs on fragment shaders at compile-time.
+ * The load_sample_id intrinsics themselves are lowered later, with different
+ * lowerings for monolithic vs epilogs.
+ *
+ * Note that fragment I/O (like store_local_pixel_agx and sample_mask_agx) does
+ * not get lowered here, because that lowering is different for monolithic vs
+ * FS epilogs even though there's no dependency on sample count.
+ */
+bool
+agx_nir_lower_sample_intrinsics(nir_shader *shader)
+{
+   /* If sample shading is disabled, the unlowered shader will broadcast pixel
+    * values across the sample (the default). By definition, there are no sample
+    * position or sample barycentrics, as these trigger sample shading.
+    */
+   if (!shader->info.fs.uses_sample_shading)
+      return false;
+
+   return nir_shader_instructions_pass(
+      shader, lower_to_sample,
+      nir_metadata_block_index | nir_metadata_dominance, NULL);
+}

src/asahi/lib/agx_tilebuffer.h

@@ -60,6 +60,8 @@ struct agx_msaa_state {
 bool agx_nir_lower_monolithic_msaa(struct nir_shader *shader,
                                    struct agx_msaa_state *state);
 
+bool agx_nir_lower_sample_intrinsics(struct nir_shader *shader);
+
 void agx_usc_tilebuffer(struct agx_usc_builder *b,
                         struct agx_tilebuffer_layout *tib);
 

src/asahi/lib/meson.build

@@ -12,6 +12,7 @@ libasahi_lib_files = files(
   'agx_meta.c',
   'agx_tilebuffer.c',
   'agx_nir_lower_msaa.c',
+  'agx_nir_lower_sample_intrinsics.c',
   'agx_nir_lower_tilebuffer.c',
   'agx_nir_lower_vbo.c',
   'agx_ppp.h',