nir,radv: pass the number of samples to load_sample_positions_amd

This will be used to lower it when it's dynamic.

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18677>

src/amd/vulkan/radv_nir_lower_abi.c

@@ -255,21 +255,16 @@ lower_abi_instr(nir_builder *b, nir_instr *instr, void *state)
       break;
    }
    case nir_intrinsic_load_sample_positions_amd: {
-      uint32_t sample_pos_offset = RING_PS_SAMPLE_POSITIONS * 16;
+      uint32_t sample_pos_offset = (RING_PS_SAMPLE_POSITIONS * 16) - 8;
 
       nir_ssa_def *ring_offsets = ac_nir_load_arg(b, &s->args->ac, s->args->ring_offsets);
       nir_ssa_def *addr = nir_pack_64_2x32(b, ring_offsets);
       nir_ssa_def *sample_id = nir_umin(b, intrin->src[0].ssa, nir_imm_int(b, 7));
       nir_ssa_def *offset = nir_ishl_imm(b, sample_id, 3); /* 2 floats containing samplepos.xy */
 
-      if (s->pl_key->ps.num_samples == 2) {
-         sample_pos_offset += 1 << 3;
-      } else if (s->pl_key->ps.num_samples == 4) {
-         sample_pos_offset += 3 << 3;
-      } else {
-         assert(s->pl_key->ps.num_samples == 8);
-         sample_pos_offset += 7 << 3;
-      }
+      nir_const_value *const_num_samples = nir_src_as_const_value(intrin->src[1]);
+      assert(const_num_samples);
+      sample_pos_offset += (const_num_samples->u32 << 3);
 
       replacement = nir_load_global_amd(b, 2, 32, addr, offset,
                                         .base = sample_pos_offset, .access = ACCESS_NON_WRITEABLE);

src/amd/vulkan/radv_shader.c

@@ -596,6 +596,7 @@ radv_lower_fs_intrinsics(nir_shader *nir, const struct radv_pipeline_stage *fs_s
             break;
          }
          case nir_intrinsic_load_barycentric_at_sample: {
+            nir_ssa_def *num_samples = nir_load_rasterization_samples_amd(&b);
             nir_ssa_def *new_dest;
 
             if (!key->ps.num_samples) {
@@ -603,7 +604,8 @@ radv_lower_fs_intrinsics(nir_shader *nir, const struct radv_pipeline_stage *fs_s
                   nir_load_barycentric_pixel(&b, 32,
                                              .interp_mode = nir_intrinsic_interp_mode(intrin));
             } else {
-               nir_ssa_def *sample_pos = nir_load_sample_positions_amd(&b, 32, intrin->src[0].ssa);
+               nir_ssa_def *sample_pos =
+                  nir_load_sample_positions_amd(&b, 32, intrin->src[0].ssa, num_samples);
 
                /* sample_pos -= 0.5 */
                sample_pos = nir_fsub(&b, sample_pos, nir_imm_float(&b, 0.5f));

src/compiler/nir/nir_intrinsics.py

@@ -1303,8 +1303,8 @@ store("global_amd", [1, 1], indices=[BASE, ACCESS, ALIGN_MUL, ALIGN_OFFSET, WRIT
 # Same as shared_atomic_add, but with GDS. src[] = {store_val, gds_addr, m0}
 intrinsic("gds_atomic_add_amd",  src_comp=[1, 1, 1], dest_comp=1, indices=[BASE])
 
-# src[] = { sample_id }
-intrinsic("load_sample_positions_amd", src_comp=[1], dest_comp=2, flags=[CAN_ELIMINATE, CAN_REORDER])
+# src[] = { sample_id, num_samples }
+intrinsic("load_sample_positions_amd", src_comp=[1, 1], dest_comp=2, flags=[CAN_ELIMINATE, CAN_REORDER])
 
 # Descriptor where TCS outputs are stored for TES
 system_value("ring_tess_offchip_amd", 4)