radv: fix gl_SampleMaskIn for sample shading

When sample shading, we need

  gl_SampleMaskIn = SampleCoverage & (PsIterMask << gl_SampleID);

Add a new shader arg, ps_iter_mask, to pass PsIterMask to ps.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23265>

src/amd/vulkan/nir/radv_nir_lower_fs_intrinsics.c

@@ -55,9 +55,11 @@ radv_nir_lower_fs_intrinsics(nir_shader *nir, const struct radv_pipeline_stage *
 
             nir_ssa_def *def = NULL;
             if (info->ps.uses_sample_shading || key->ps.sample_shading_enable) {
-               /* gl_SampleMaskIn[0] = (SampleCoverage & (1 << gl_SampleID)). */
+               /* gl_SampleMaskIn[0] = (SampleCoverage & (PsIterMask << gl_SampleID)). */
+               nir_ssa_def *ps_iter_mask =
+                  nir_load_scalar_arg_amd(&b, 1, .base = args->ps_iter_mask.arg_index);
                nir_ssa_def *sample_id = nir_load_sample_id(&b);
-               def = nir_iand(&b, sample_coverage, nir_ishl(&b, nir_imm_int(&b, 1u), sample_id));
+               def = nir_iand(&b, sample_coverage, nir_ishl(&b, ps_iter_mask, sample_id));
             } else {
                def = sample_coverage;
             }

src/amd/vulkan/radv_cmd_buffer.c

@@ -2658,6 +2658,7 @@ radv_emit_rasterization_samples(struct radv_cmd_buffer *cmd_buffer)
    const struct radv_physical_device *pdevice = cmd_buffer->device->physical_device;
    const struct radv_shader *ps = cmd_buffer->state.shaders[MESA_SHADER_FRAGMENT];
    unsigned rasterization_samples = radv_get_rasterization_samples(cmd_buffer);
+   unsigned ps_iter_samples = radv_get_ps_iter_samples(cmd_buffer);
    const struct radv_rendering_state *render = &cmd_buffer->state.render;
    const struct radv_dynamic_state *d = &cmd_buffer->state.dynamic;
    unsigned spi_baryc_cntl = S_0286E0_FRONT_FACE_ALL_BITS(1);
@@ -2682,13 +2683,9 @@ radv_emit_rasterization_samples(struct radv_cmd_buffer *cmd_buffer)
    if (!d->sample_location.count)
       radv_emit_default_sample_locations(cmd_buffer->cs, rasterization_samples);
 
-   if (rasterization_samples > 1) {
-      unsigned ps_iter_samples = radv_get_ps_iter_samples(cmd_buffer);
-
-      if (ps_iter_samples > 1) {
-         spi_baryc_cntl |= S_0286E0_POS_FLOAT_LOCATION(2);
-         pa_sc_mode_cntl_1 |= S_028A4C_PS_ITER_SAMPLE(1);
-      }
+   if (ps_iter_samples > 1) {
+      spi_baryc_cntl |= S_0286E0_POS_FLOAT_LOCATION(2);
+      pa_sc_mode_cntl_1 |= S_028A4C_PS_ITER_SAMPLE(1);
    }
 
    if (pdevice->rad_info.gfx_level >= GFX10_3 &&
@@ -2739,6 +2736,14 @@ radv_emit_rasterization_samples(struct radv_cmd_buffer *cmd_buffer)
          uint32_t base_reg = cmd_buffer->state.shaders[MESA_SHADER_FRAGMENT]->info.user_data_0;
          radeon_set_sh_reg(cmd_buffer->cs, base_reg + loc->sgpr_idx * 4, rasterization_samples);
       }
+
+      loc =
+         radv_get_user_sgpr(cmd_buffer->state.shaders[MESA_SHADER_FRAGMENT], AC_UD_PS_ITER_MASK);
+      if (loc->sgpr_idx != -1) {
+         const uint16_t ps_iter_mask = ac_get_ps_iter_mask(ps_iter_samples);
+         uint32_t base_reg = cmd_buffer->state.shaders[MESA_SHADER_FRAGMENT]->info.user_data_0;
+         radeon_set_sh_reg(cmd_buffer->cs, base_reg + loc->sgpr_idx * 4, ps_iter_mask);
+      }
    }
 }
 

src/amd/vulkan/radv_shader.h

@@ -170,6 +170,7 @@ enum radv_ud_index {
    AC_UD_PS_EPILOG_PC,
    AC_UD_PS_NUM_SAMPLES,
    AC_UD_PS_LINE_RAST_MODE,
+   AC_UD_PS_ITER_MASK,
    AC_UD_PS_MAX_UD,
    AC_UD_CS_GRID_SIZE = AC_UD_SHADER_START,
    AC_UD_CS_SBT_DESCRIPTORS,

src/amd/vulkan/radv_shader_args.c

@@ -647,6 +647,10 @@ declare_shader_args(const struct radv_device *device, const struct radv_pipeline
       if (key->dynamic_line_rast_mode)
          add_ud_arg(args, 1, AC_ARG_INT, &args->ps_line_rast_mode, AC_UD_PS_LINE_RAST_MODE);
 
+      if (info->ps.reads_sample_mask_in && (info->ps.uses_sample_shading ||
+                                            key->ps.sample_shading_enable))
+         add_ud_arg(args, 1, AC_ARG_INT, &args->ps_iter_mask, AC_UD_PS_ITER_MASK);
+
       ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.prim_mask);
       if (args->explicit_scratch_args && gfx_level < GFX11) {
          ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.scratch_offset);

src/amd/vulkan/radv_shader_args.h

@@ -61,6 +61,7 @@ struct radv_shader_args {
    struct ac_arg ps_epilog_pc;
    struct ac_arg ps_num_samples;
    struct ac_arg ps_line_rast_mode;
+   struct ac_arg ps_iter_mask;
 
    struct ac_arg prolog_inputs;
    struct ac_arg vs_inputs[MAX_VERTEX_ATTRIBS];