nir: commonize barycentric intrinsic opt pass

Introduces an opt pass that attempts to optimize
load_barycentric_at_{sample,offset} with simpler load_barycentric_*
equivalents where possible, and optionally lowers
load_barycentric_at_sample to load_barycentric_at_offset with a position
derived from the sample ID instead.

Signed-off-by: Simon Perretta <simon.perretta@imgtec.com>
Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37658>

src/compiler/nir/meson.build

@@ -251,6 +251,7 @@ else
   'nir_normalize_cubemap_coords.c',
   'nir_opt_access.c',
   'nir_opt_barriers.c',
+  'nir_opt_barycentric.c',
   'nir_opt_call.c',
   'nir_opt_clip_cull_const.c',
   'nir_opt_combine_stores.c',

src/compiler/nir/nir.h

@@ -6770,6 +6770,8 @@ bool nir_lower_cooperative_matrix_flexible_dimensions(nir_shader *shader, unsign
 
 bool nir_unlower_io_to_vars(nir_shader *nir, bool keep_intrinsics);
 
+bool nir_opt_barycentric(nir_shader *shader, bool lower_sample_to_pos);
+
 #include "nir_inline_helpers.h"
 
 #ifdef __cplusplus

src/compiler/nir/nir_opt_barycentric.c

@@ -0,0 +1,173 @@
+/*
+ * Copyright © 2025 Imagination Technologies Ltd.
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include "nir_builder.h"
+
+/* This pass attempts to optimize load_barycentric_at_{sample,offset} with
+ * simpler load_barycentric_* equivalents where possible, and optionally
+ * lowers load_barycentric_at_sample to load_barycentric_at_offset with a
+ * position derived from the sample ID instead.
+ */
+
+static bool
+opt_bary_at_sample(nir_builder *b, nir_intrinsic_instr *intr, bool lower_sample_to_pos)
+{
+   /* Check for and handle simple replacement cases:
+    * - Sample num is current sample.
+    */
+   enum glsl_interp_mode interp_mode = nir_intrinsic_interp_mode(intr);
+   nir_intrinsic_instr *sample = nir_src_as_intrinsic(intr->src[0]);
+
+   assert(interp_mode != INTERP_MODE_FLAT);
+   if (sample && sample->intrinsic == nir_intrinsic_load_sample_id) {
+      nir_def *repl = nir_load_barycentric_sample(
+         b,
+         intr->def.bit_size,
+         .interp_mode = interp_mode);
+      nir_def_replace(&intr->def, repl);
+      nir_instr_free(&intr->instr);
+      return true;
+   }
+
+   if (!lower_sample_to_pos)
+      return false;
+
+   /* Turn the sample id into a position. */
+   nir_def *offset =
+      nir_load_sample_pos_from_id(b, intr->def.bit_size, intr->src[0].ssa);
+   offset = nir_fadd_imm(b, offset, -0.5f);
+
+   nir_def *repl = nir_load_barycentric_at_offset(
+      b,
+      intr->def.bit_size,
+      offset,
+      .interp_mode = interp_mode);
+
+   nir_def_replace(&intr->def, repl);
+   nir_instr_free(&intr->instr);
+   return true;
+}
+
+static bool
+src_is_vec2_sample_pos_minus_half(nir_src src)
+{
+   nir_alu_instr *alu = nir_src_as_alu_instr(src);
+   if (!alu || alu->op != nir_op_vec2)
+      return false;
+
+   /* Check both vec2 components. */
+   for (unsigned u = 0; u < 2; ++u) {
+      nir_scalar comp = nir_get_scalar(&alu->def, u);
+      comp = nir_scalar_chase_movs(comp);
+
+      if (!nir_scalar_is_alu(comp))
+         return false;
+
+      /* Look for fadd(sample_pos.x/y, -0.5f) or fsub(sample_pos.x/y, +0.5f) */
+      nir_op op = nir_scalar_alu_op(comp);
+      if (op != nir_op_fadd && op != nir_op_fsub)
+         return false;
+
+      float half_val = op == nir_op_fadd ? -0.5f : +0.5f;
+      unsigned sample_pos_srcn = ~0U;
+      unsigned half_srcn = ~0U;
+
+      /* Check both fadd/fsub sources. */
+      for (unsigned n = 0; n < 2; ++n) {
+         nir_scalar src = nir_scalar_chase_alu_src(comp, n);
+
+         if (nir_scalar_is_intrinsic(src) &&
+             nir_scalar_intrinsic_op(src) == nir_intrinsic_load_sample_pos) {
+            sample_pos_srcn = n;
+         } else if (nir_scalar_is_const(src) &&
+                    nir_scalar_as_const_value(src).f32 == half_val) {
+            half_srcn = n;
+         }
+      }
+
+      /* One or more operands not found. */
+      if (sample_pos_srcn == ~0U || half_srcn == ~0U)
+         return false;
+
+      /* fsub is not commutative. */
+      if (op == nir_op_fsub && (sample_pos_srcn != 0 || half_srcn != 1))
+         return false;
+
+      /* vec2.{x,y} needs to be referencing load_sample_pos.{x,y}. */
+      nir_scalar sample_pos_src =
+         nir_scalar_chase_alu_src(comp, sample_pos_srcn);
+      if (sample_pos_src.comp != u)
+         return false;
+   }
+
+   return true;
+}
+
+static bool
+opt_bary_at_offset(nir_builder *b, nir_intrinsic_instr *intr)
+{
+   /* Check for and handle simple replacement cases:
+    * - Flat interpolation - don't care about offset, will get consumed.
+    * - Offset is zero.
+    */
+   enum glsl_interp_mode interp_mode = nir_intrinsic_interp_mode(intr);
+   nir_src src = intr->src[0];
+
+   assert(interp_mode != INTERP_MODE_FLAT);
+   if (nir_src_is_const(src) && !nir_src_comp_as_int(src, 0) &&
+       !nir_src_comp_as_int(src, 1)) {
+      nir_def *repl = nir_load_barycentric_pixel(
+         b,
+         intr->def.bit_size,
+         .interp_mode = interp_mode);
+      nir_def_replace(&intr->def, repl);
+      nir_instr_free(&intr->instr);
+      return true;
+   }
+
+   /* Offset is vec2(sample_pos - 0.5f). */
+   if (src_is_vec2_sample_pos_minus_half(src)) {
+      nir_def *repl = nir_load_barycentric_sample(
+         b,
+         intr->def.bit_size,
+         .interp_mode = interp_mode);
+      nir_def_replace(&intr->def, repl);
+      nir_instr_free(&intr->instr);
+      return true;
+   }
+
+   return false;
+}
+
+static bool
+opt_bary(nir_builder *b, nir_intrinsic_instr *intr, void *cb_data)
+{
+   bool *lower_sample_to_pos = cb_data;
+   b->cursor = nir_before_instr(&intr->instr);
+
+   switch (intr->intrinsic) {
+   case nir_intrinsic_load_barycentric_at_sample:
+      return opt_bary_at_sample(b, intr, *lower_sample_to_pos);
+
+   case nir_intrinsic_load_barycentric_at_offset:
+      return opt_bary_at_offset(b, intr);
+
+   default:
+      break;
+   }
+
+   return false;
+}
+
+bool
+nir_opt_barycentric(nir_shader *shader, bool lower_sample_to_pos)
+{
+   return nir_shader_intrinsics_pass(shader,
+                                     opt_bary,
+                                     nir_metadata_control_flow,
+                                     &lower_sample_to_pos);
+}

src/imagination/pco/pco_nir_pvfio.c

@@ -1312,151 +1312,6 @@ bool pco_nir_link_clip_cull_vars(nir_shader *producer, nir_shader *consumer)
    return true;
 }
 
-static bool lower_bary_at_sample(nir_builder *b, nir_intrinsic_instr *intr)
-{
-   /* Check for and handle simple replacement cases:
-    * - Flat interpolation - don't care about sample num, will get consumed.
-    * - Sample num is current sample.
-    */
-   enum glsl_interp_mode interp_mode = nir_intrinsic_interp_mode(intr);
-   nir_intrinsic_instr *sample = nir_src_as_intrinsic(intr->src[0]);
-
-   if (interp_mode == INTERP_MODE_FLAT ||
-       (sample && sample->intrinsic == nir_intrinsic_load_sample_id)) {
-      nir_def *repl = nir_load_barycentric_sample(
-         b,
-         intr->def.bit_size,
-         .interp_mode = nir_intrinsic_interp_mode(intr));
-      nir_def_replace(&intr->def, repl);
-      nir_instr_free(&intr->instr);
-      return true;
-   }
-
-   /* Turn the sample id into a position. */
-   nir_def *offset =
-      nir_load_sample_pos_from_id(b, intr->def.bit_size, intr->src[0].ssa);
-   offset = nir_fadd_imm(b, offset, -0.5f);
-
-   nir_def *repl = nir_load_barycentric_at_offset(
-      b,
-      intr->def.bit_size,
-      offset,
-      .interp_mode = nir_intrinsic_interp_mode(intr));
-
-   nir_def_replace(&intr->def, repl);
-   nir_instr_free(&intr->instr);
-   return true;
-}
-
-static bool src_is_vec2_sample_pos_minus_half(nir_src src)
-{
-   nir_alu_instr *alu = nir_src_as_alu_instr(src);
-   if (!alu || alu->op != nir_op_vec2)
-      return false;
-
-   /* Check both vec2 components. */
-   for (unsigned u = 0; u < 2; ++u) {
-      nir_scalar comp = nir_get_scalar(&alu->def, u);
-      comp = nir_scalar_chase_movs(comp);
-
-      if (!nir_scalar_is_alu(comp))
-         return false;
-
-      /* Look for fadd(sample_pos.x/y, -0.5f) or fsub(sample_pos.x/y, +0.5f) */
-      nir_op op = nir_scalar_alu_op(comp);
-      if (op != nir_op_fadd && op != nir_op_fsub)
-         return false;
-
-      float half_val = op == nir_op_fadd ? -0.5f : +0.5f;
-      unsigned sample_pos_srcn = ~0U;
-      unsigned half_srcn = ~0U;
-
-      /* Check both fadd/fsub sources. */
-      for (unsigned n = 0; n < 2; ++n) {
-         nir_scalar src = nir_scalar_chase_alu_src(comp, n);
-
-         if (nir_scalar_is_intrinsic(src) &&
-             nir_scalar_intrinsic_op(src) == nir_intrinsic_load_sample_pos) {
-            sample_pos_srcn = n;
-         } else if (nir_scalar_is_const(src) &&
-                    nir_scalar_as_const_value(src).f32 == half_val) {
-            half_srcn = n;
-         }
-      }
-
-      /* One or more operands not found. */
-      if (sample_pos_srcn == ~0U || half_srcn == ~0U)
-         return false;
-
-      /* fsub is not commutative. */
-      if (op == nir_op_fsub && (sample_pos_srcn != 0 || half_srcn != 1))
-         return false;
-
-      /* vec2.{x,y} needs to be referencing load_sample_pos.{x,y}. */
-      nir_scalar sample_pos_src =
-         nir_scalar_chase_alu_src(comp, sample_pos_srcn);
-      if (sample_pos_src.comp != u)
-         return false;
-   }
-
-   return true;
-}
-
-static bool lower_bary_at_offset(nir_builder *b, nir_intrinsic_instr *intr)
-{
-   /* Check for and handle simple replacement cases:
-    * - Flat interpolation - don't care about offset, will get consumed.
-    * - Offset is zero.
-    * - sample_pos - 0.5f.
-    */
-   enum glsl_interp_mode interp_mode = nir_intrinsic_interp_mode(intr);
-   nir_src src = intr->src[0];
-
-   if (interp_mode == INTERP_MODE_FLAT ||
-       (nir_src_is_const(src) && !nir_src_comp_as_int(src, 0) &&
-        !nir_src_comp_as_int(src, 1))) {
-      nir_def *repl = nir_load_barycentric_pixel(
-         b,
-         intr->def.bit_size,
-         .interp_mode = nir_intrinsic_interp_mode(intr));
-      nir_def_replace(&intr->def, repl);
-      nir_instr_free(&intr->instr);
-      return true;
-   }
-
-   if (src_is_vec2_sample_pos_minus_half(src)) {
-      nir_def *repl = nir_load_barycentric_sample(
-         b,
-         intr->def.bit_size,
-         .interp_mode = nir_intrinsic_interp_mode(intr));
-      nir_def_replace(&intr->def, repl);
-      nir_instr_free(&intr->instr);
-      return true;
-   }
-
-   /* Non-zero offsets handled in lower_interp. */
-   return false;
-}
-
-static bool
-lower_bary(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *cb_data)
-{
-   b->cursor = nir_before_instr(&intr->instr);
-
-   switch (intr->intrinsic) {
-   case nir_intrinsic_load_barycentric_at_sample:
-      return lower_bary_at_sample(b, intr);
-
-   case nir_intrinsic_load_barycentric_at_offset:
-      return lower_bary_at_offset(b, intr);
-
-   default:
-      break;
-   }
-
-   return false;
-}
-
 static nir_def *alu_iter(nir_builder *b,
                          nir_def *coords,
                          unsigned component,
@@ -1576,11 +1431,7 @@ bool pco_nir_lower_interpolation(nir_shader *shader, pco_fs_data *fs)
 {
    bool progress = false;
 
-   progress |= nir_shader_intrinsics_pass(shader,
-                                          lower_bary,
-                                          nir_metadata_control_flow,
-                                          NULL);
-
+   progress |= nir_opt_barycentric(shader, true);
    progress |= nir_shader_intrinsics_pass(shader,
                                           lower_interp,
                                           nir_metadata_control_flow,