nir: Add a late texcoord replacement pass

Add a second NIR pass for lowering point/texture coordinate replacement (i.e.
point sprites). Why a second one? The current pass works on derefs/variables,
which is good for drivers that don't lower I/O at all (like Zink, where the pass
originates). However, it is problematic for hardware drivers: the inputs to this
pass depend on the shader key, so we want to run the pass as late as possible to
minimize the cost of building/compiling the associated shader variants. In
particular, we need to be able to lower point sprites after lowering I/O if we
would like to lower I/O when preprocessing NIR.

The logic for early lowering and late lowering is considerably different (the
late lowering is a lot simpler), so I've split this out into a second pass
rather than trying to weld them together into one.

This pass will be used on Asahi, which currently uses the early pass. It may be
useful for other drivers as well. (Actually, it's been shipping on Asahi for a
little while now, just hasn't been sent upstream yet.)

Tested with Neverball.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Emma Anholt <emma@anholt.net>
Acked-by: Asahi Lina <lina@asahilina.net>
Acked-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21065>

src/compiler/nir/meson.build

@@ -205,6 +205,7 @@ files_libnir = files(
   'nir_lower_tex_shadow.c',
   'nir_lower_tex.c',
   'nir_lower_texcoord_replace.c',
+  'nir_lower_texcoord_replace_late.c',
   'nir_lower_to_source_mods.c',
   'nir_lower_two_sided_color.c',
   'nir_lower_undef_to_zero.c',

src/compiler/nir/nir.h

@@ -5534,6 +5534,9 @@ bool nir_lower_point_size(nir_shader *shader, float min, float max);
 void nir_lower_texcoord_replace(nir_shader *s, unsigned coord_replace,
                                 bool point_coord_is_sysval, bool yinvert);
 
+void nir_lower_texcoord_replace_late(nir_shader *s, unsigned coord_replace,
+                                     bool point_coord_is_sysval);
+
 typedef enum {
    nir_lower_interpolation_at_sample = (1 << 1),
    nir_lower_interpolation_at_offset = (1 << 2),

src/compiler/nir/nir_lower_texcoord_replace_late.c

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2022 Alyssa Rosenzweig
+ * Copyright 2020 Collabora, Ltd.
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include "nir_builder.h"
+#include "nir_deref.h"
+
+struct opts {
+   unsigned coord_replace;
+   bool point_coord_is_sysval;
+};
+
+static nir_ssa_def *
+nir_channel_or_undef(nir_builder *b, nir_ssa_def *def, signed int channel)
+{
+   if (channel >= 0 && channel < def->num_components)
+      return nir_channel(b, def, channel);
+   else
+      return nir_ssa_undef(b, def->bit_size, 1);
+}
+
+static bool
+pass(nir_builder *b, nir_instr *instr, void *data)
+{
+   struct opts *opts = data;
+
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+   if (intr->intrinsic != nir_intrinsic_load_interpolated_input &&
+       intr->intrinsic != nir_intrinsic_load_input)
+      return false;
+
+   nir_src *offset = nir_get_io_offset_src(intr);
+   assert(nir_src_is_const(*offset) && "no indirects supported");
+
+   nir_io_semantics sem = nir_intrinsic_io_semantics(intr);
+   unsigned location = sem.location + nir_src_as_uint(*offset);
+   signed component = nir_intrinsic_component(intr);
+
+   if (location < VARYING_SLOT_TEX0 || location > VARYING_SLOT_TEX7)
+      return false;
+
+   if (!(opts->coord_replace & BITFIELD_BIT(location - VARYING_SLOT_TEX0)))
+      return false;
+
+   b->cursor = nir_before_instr(instr);
+   nir_ssa_def *channels[4] = {
+      NULL, NULL,
+      nir_imm_float(b, 0.0),
+      nir_imm_float(b, 1.0)
+   };
+
+   if (opts->point_coord_is_sysval) {
+      nir_ssa_def *pntc = nir_load_point_coord(b);
+
+      b->cursor = nir_after_instr(instr);
+      channels[0] = nir_channel(b, pntc, 0);
+      channels[1] = nir_channel(b, pntc, 1);
+   } else {
+      sem.location = VARYING_SLOT_PNTC;
+      nir_instr_rewrite_src_ssa(instr, offset, nir_imm_int(b, 0));
+      nir_intrinsic_set_io_semantics(intr, sem);
+      nir_ssa_def *raw = &intr->dest.ssa;
+
+      b->cursor = nir_after_instr(instr);
+      channels[0] = nir_channel_or_undef(b, raw, 0 - component);
+      channels[1] = nir_channel_or_undef(b, raw, 1 - component);
+   }
+
+   nir_ssa_def *res = nir_vec(b, &channels[component], intr->num_components);
+   nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, res,
+                                  res->parent_instr);
+   return true;
+}
+
+void
+nir_lower_texcoord_replace_late(nir_shader *s, unsigned coord_replace,
+                                bool point_coord_is_sysval)
+{
+   assert(s->info.stage == MESA_SHADER_FRAGMENT);
+   assert(coord_replace != 0);
+
+   uint64_t replace_mask = (((uint64_t) coord_replace) << VARYING_SLOT_TEX0);
+
+   /* If no relevant texcoords are read, there's nothing to do */
+   if (!(s->info.inputs_read & replace_mask))
+      return;
+
+   /* Otherwise, we're going to replace these texcoord reads with a PNTC read */
+   s->info.inputs_read &= ~(((uint64_t) coord_replace) << VARYING_SLOT_TEX0);
+
+   if (!point_coord_is_sysval)
+      s->info.inputs_read |= BITFIELD64_BIT(VARYING_SLOT_PNTC);
+
+   nir_shader_instructions_pass(s, pass,
+         nir_metadata_block_index | nir_metadata_dominance,
+         &(struct opts) {
+            .coord_replace = coord_replace,
+            .point_coord_is_sysval = point_coord_is_sysval,
+         });
+}