nak: Rewrite nir_intrinsic_load_sample_pos and implement nir_intrinsic_load_barycentric_at_sample

nir_intrinsic_load_sample_pos was causing failures on barycentric CTS tests.

Signed-off-by: Mary Guillemard <mary.guillemard@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/26239>

src/nouveau/compiler/nak.h

@@ -39,6 +39,14 @@ struct nak_fs_key {
     * VkPipelineMultisampleStateCreateInfo::minSampleShading
     */
    bool force_sample_shading;
+
+   /**
+    * The constant buffer index and offset at which the sample locations table lives.
+    * Each sample location is two 4-bit unorm values packed into an 8-bit value
+    * with the bottom 4 bits for x and the top 4 bits for y.
+   */
+   uint8_t sample_locations_cb;
+   uint32_t sample_locations_offset;
 };
 
 void nak_postprocess_nir(nir_shader *nir, const struct nak_compiler *nak,

src/nouveau/compiler/nak_nir.c

@@ -610,6 +610,27 @@ load_interpolated_input(nir_builder *b, unsigned num_components, uint32_t addr,
    }
 }
 
+static nir_def *
+load_sample_pos_at(nir_builder *b, nir_def *sample_id,
+                   const struct nak_fs_key *fs_key)
+{
+   nir_def *loc = nir_load_ubo(b, 1, 64,
+                               nir_imm_int(b, fs_key->sample_locations_cb),
+                               nir_imm_int(b, fs_key->sample_locations_offset),
+                               .align_mul = 8,
+                               .align_offset = 0,
+                               .range = fs_key->sample_locations_offset + 8);
+
+   /* Yay little endian */
+   loc = nir_ushr(b, loc, nir_imul_imm(b, sample_id, 8));
+   nir_def *loc_x_u4 = nir_iand_imm(b, loc, 0xf);
+   nir_def *loc_y_u4 = nir_iand_imm(b, nir_ushr_imm(b, loc, 4), 0xf);
+   nir_def *loc_u4 = nir_vec2(b, loc_x_u4, loc_y_u4);
+   nir_def *result = nir_fmul_imm(b, nir_i2f32(b, loc_u4), 1.0 / 16.0);
+
+   return result;
+}
+
 struct lower_fs_input_ctx {
    const struct nak_compiler *nak;
    const struct nak_fs_key *fs_key;
@@ -630,8 +651,7 @@ lower_fs_input_intrin(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
    }
 
    case nir_intrinsic_load_frag_coord:
-   case nir_intrinsic_load_point_coord:
+   case nir_intrinsic_load_point_coord: {
-   case nir_intrinsic_load_sample_pos: {
       b->cursor = nir_before_instr(&intrin->instr);
 
       const enum nak_interp_loc interp_loc =
@@ -655,9 +675,6 @@ lower_fs_input_intrin(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
          break;
       case nir_intrinsic_load_point_coord:
          break;
-      case nir_intrinsic_load_sample_pos:
-         coord = nir_ffract(b, coord);
-         break;
       default:
          unreachable("Unknown intrinsic");
       }
@@ -715,10 +732,20 @@ lower_fs_input_intrin(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
       nir_def *offset = NULL;
       enum nak_interp_loc interp_loc;
       switch (bary->intrinsic) {
-      case nir_intrinsic_load_barycentric_at_offset: {
+      case nir_intrinsic_load_barycentric_at_offset:
+      case nir_intrinsic_load_barycentric_at_sample: {
          interp_loc = NAK_INTERP_LOC_OFFSET;
 
-         nir_def *offset_f = bary->src[0].ssa;
+         nir_def *offset_f;
+
+         if (bary->intrinsic == nir_intrinsic_load_barycentric_at_sample) {
+            nir_def *sample_id = bary->src[0].ssa;
+            nir_def *sample_pos = load_sample_pos_at(b, sample_id, ctx->fs_key);
+            offset_f = nir_fadd_imm(b, sample_pos, -0.5);
+         } else {
+            offset_f = bary->src[0].ssa;
+         }
+
          offset_f = nir_fclamp(b, offset_f, nir_imm_float(b, -0.5),
                                nir_imm_float(b, 0.437500));
          nir_def *offset_fixed =
@@ -773,6 +800,18 @@ lower_fs_input_intrin(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
       return true;
    }
 
+   case nir_intrinsic_load_sample_pos: {
+      b->cursor = nir_before_instr(&intrin->instr);
+
+      nir_def *sample_id = nir_load_sample_id(b);
+      nir_def *sample_pos = load_sample_pos_at(b, sample_id, ctx->fs_key);
+
+      nir_def_rewrite_uses(&intrin->def, sample_pos);
+      nir_instr_remove(&intrin->instr);
+
+      return true;
+   }
+
    default:
       return false;
    }

src/nouveau/vulkan/nvk_graphics_pipeline.c

@@ -4,6 +4,7 @@
  */
 #include "nvk_pipeline.h"
 
+#include "nvk_cmd_buffer.h"
 #include "nvk_device.h"
 #include "nvk_physical_device.h"
 #include "nvk_shader.h"
@@ -48,6 +49,9 @@ nvk_populate_fs_key(struct nak_fs_key *key,
 {
    memset(key, 0, sizeof(*key));
 
+   key->sample_locations_cb = 0;
+   key->sample_locations_offset = nvk_root_descriptor_offset(draw.sample_locations);
+
    if (state->pipeline_flags &
        VK_PIPELINE_CREATE_DEPTH_STENCIL_ATTACHMENT_FEEDBACK_LOOP_BIT_EXT)
       key->zs_self_dep = true;