nir: add nir_opt_frag_coord_to_pixel_coord

Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31864>

src/compiler/nir/meson.build

@@ -230,6 +230,7 @@ files_libnir = files(
   'nir_opt_dead_cf.c',
   'nir_opt_dead_write_vars.c',
   'nir_opt_find_array_copies.c',
+  'nir_opt_frag_coord_to_pixel_coord.c',
   'nir_opt_fragdepth.c',
   'nir_opt_gcm.c',
   'nir_opt_generate_bfi.c',

src/compiler/nir/nir.h

@@ -6089,6 +6089,7 @@ nir_shader *nir_create_passthrough_gs(const nir_shader_compiler_options *options
 
 bool nir_lower_fragcolor(nir_shader *shader, unsigned max_cbufs);
 bool nir_lower_fragcoord_wtrans(nir_shader *shader);
+bool nir_opt_frag_coord_to_pixel_coord(nir_shader *shader);
 bool nir_lower_frag_coord_to_pixel_coord(nir_shader *shader);
 bool nir_lower_viewport_transform(nir_shader *shader);
 bool nir_lower_uniforms_to_ubo(nir_shader *shader, bool dword_packed, bool load_vec4);

src/compiler/nir/nir_opt_frag_coord_to_pixel_coord.c

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2024 Valve Corpoation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include "nir_builder.h"
+#include "nir_builder_opcodes.h"
+
+/**
+ * If load_frag_coord.xy is only used by conversions to integer,
+ * replace it with load_pixel_coord.
+ */
+
+static bool
+opt_frag_pos(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *data)
+{
+   if (intr->intrinsic != nir_intrinsic_load_frag_coord)
+      return false;
+
+   /* Don't increase precision. */
+   if (intr->def.bit_size != 32)
+      return false;
+
+   /* Check if xy are only used by casts to integers. */
+   nir_foreach_use(use, &intr->def) {
+      if (nir_src_is_if(use))
+         return false;
+
+      unsigned mask = nir_src_components_read(use);
+
+      if (!(mask & 0x3))
+         continue;
+
+      /* Don't handle instructions that read x/y and z/w for simplicity. */
+      if (mask & ~0x3)
+         return false;
+
+      nir_instr *use_instr = nir_src_parent_instr(use);
+
+      if (use_instr->type != nir_instr_type_alu)
+         return false;
+
+      switch (nir_instr_as_alu(use_instr)->op) {
+      case nir_op_f2i8:
+      case nir_op_f2i16:
+      case nir_op_f2i32:
+      case nir_op_f2i64:
+      case nir_op_f2u8:
+      case nir_op_f2u16:
+      case nir_op_f2u32:
+      case nir_op_f2u64:
+         continue;
+      default:
+         return false;
+      }
+   }
+
+   b->cursor = nir_before_instr(&intr->instr);
+   nir_def *pixel_coord = nir_load_pixel_coord(b);
+
+   nir_foreach_use_safe(use, &intr->def) {
+      unsigned mask = nir_src_components_read(use);
+
+      if (!(mask & 0x3))
+         continue;
+
+      nir_src_rewrite(use, pixel_coord);
+
+      nir_alu_instr *use_instr = nir_instr_as_alu(nir_src_parent_instr(use));
+
+      /* load_frag_coord is always positive, so we should never sign extend here. */
+      nir_alu_type dst_type = nir_type_uint | use_instr->def.bit_size;
+      use_instr->op = nir_type_conversion_op(nir_type_uint16, dst_type, nir_rounding_mode_undef);
+   }
+
+   return true;
+}
+
+bool
+nir_opt_frag_coord_to_pixel_coord(nir_shader *shader)
+{
+   return nir_shader_intrinsics_pass(shader, opt_frag_pos,
+                                     nir_metadata_control_flow,
+                                     NULL);
+}