radv: switch to ac_nir_lower_sample_mask_in

load_use_sample_mask_in_amd will be fully implemented in a later commit.

This is already more efficient because of using helper_invocation
in the pass.

Reviewed-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41768>

src/amd/common/nir/ac_nir_lower_intrinsics_to_args.c

@@ -205,6 +205,9 @@ lower_intrinsic_to_arg(nir_builder *b, nir_intrinsic_instr *intrin, void *state)
    case nir_intrinsic_load_frag_coord_w_rcp:
       replacement = ac_nir_load_arg(b, s->args, s->args->frag_pos[3]);
       break;
+   case nir_intrinsic_load_sample_mask_in:
+      replacement = ac_nir_load_arg(b, s->args, s->args->sample_coverage);
+      break;
    case nir_intrinsic_load_local_invocation_id: {
       unsigned num_bits[3];
       nir_def *vec[3];

src/amd/vulkan/nir/radv_nir_lower_abi.c

@@ -444,6 +444,12 @@ lower_abi_instr(nir_builder *b, nir_intrinsic_instr *intrin, void *state)
    case nir_intrinsic_load_use_float_frag_coord_xy_amd:
       replacement = nir_ine_imm(b, GET_SGPR_FIELD_NIR(s->args->ps_state, PS_STATE_USE_FLOAT_FRAG_COORD_XY), 0);
       break;
+   case nir_intrinsic_load_ps_iter_mask_amd:
+      replacement = GET_SGPR_FIELD_NIR(s->args->ps_state, PS_STATE_PS_ITER_MASK);
+      break;
+   case nir_intrinsic_load_use_sample_mask_in_amd:
+      replacement = nir_imm_true(b);
+      break;
    default:
       progress = false;
       break;

src/amd/vulkan/nir/radv_nir_lower_fs_intrinsics.c

@@ -25,28 +25,9 @@ pass(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
 
    const struct ctx *ctx = data;
    const struct radv_graphics_state_key *gfx_state = ctx->gfx_state;
-   const struct radv_shader_info *info = &ctx->fs_stage->info;
    const struct radv_shader_args *args = &ctx->fs_stage->args;
 
    switch (intrin->intrinsic) {
-   case nir_intrinsic_load_sample_mask_in: {
-      nir_def *sample_coverage = nir_load_vector_arg_amd(b, 1, .base = args->ac.sample_coverage.arg_index);
-
-      nir_def *def = NULL;
-      if (info->ps.uses_sample_shading || gfx_state->ms.sample_shading_enable) {
-         /* gl_SampleMaskIn[0] = (SampleCoverage & (PsIterMask << gl_SampleID)). */
-         nir_def *ps_state = nir_load_scalar_arg_amd(b, 1, .base = args->ps_state.arg_index);
-         nir_def *ps_iter_mask =
-            nir_ubfe_imm(b, ps_state, PS_STATE_PS_ITER_MASK__SHIFT, util_bitcount(PS_STATE_PS_ITER_MASK__MASK));
-         nir_def *sample_id = nir_load_sample_id(b);
-         def = nir_iand(b, sample_coverage, nir_ishl(b, ps_iter_mask, sample_id));
-      } else {
-         def = sample_coverage;
-      }
-
-      nir_def_replace(&intrin->def, def);
-      return true;
-   }
    case nir_intrinsic_load_frag_coord_z: {
       if (!gfx_state->adjust_frag_coord_z)
          return false;

src/amd/vulkan/radv_pipeline_graphics.c

@@ -1690,9 +1690,16 @@ radv_generate_graphics_state_key(const struct radv_compiler_info *compiler_info,
 
    if (state->ms) {
       key.ms.sample_shading_enable = state->ms->sample_shading_enable;
+      key.ms.max_sample_shading_enable = state->ms->sample_shading_enable && state->ms->min_sample_shading == 1;
+
       if (!BITSET_TEST(state->dynamic, MESA_VK_DYNAMIC_MS_RASTERIZATION_SAMPLES) &&
           state->ms->rasterization_samples > 1) {
          key.ms.rasterization_samples = state->ms->rasterization_samples;
+
+         if (state->ms->sample_shading_enable) {
+            key.ms.ps_iter_samples =
+               util_next_power_of_two(ceilf(state->ms->rasterization_samples * state->ms->min_sample_shading));
+         }
       }
    }
 
@@ -2569,6 +2576,22 @@ radv_graphics_shaders_compile(const struct radv_compiler_info *compiler_info, st
                !gfx_state->vrs_may_be_enabled && !gfx_state->ms.sample_shading_enable &&
                   !stages[MESA_SHADER_FRAGMENT].nir->info.fs.uses_sample_shading);
 
+      ac_nir_lower_sample_mask_in_options lower_sample_mask_in_options = {0};
+
+      if (stages[MESA_SHADER_FRAGMENT].nir->info.fs.uses_sample_shading || gfx_state->ms.max_sample_shading_enable) {
+         lower_sample_mask_in_options.behavior = ac_nir_lower_samplemask_sample_shading_max;
+      } else if (gfx_state->ms.sample_shading_enable) {
+         lower_sample_mask_in_options.behavior = ac_nir_lower_samplemask_sample_shading_partial;
+         lower_sample_mask_in_options.ps_iter_samples = gfx_state->ms.ps_iter_samples;
+      } else if (!gfx_state->vrs_may_be_enabled && !gfx_state->dynamic_rasterization_samples &&
+                 gfx_state->ms.rasterization_samples == 0) {
+         lower_sample_mask_in_options.behavior = ac_nir_lower_samplemask_1sample_no_vrs;
+      } else {
+         lower_sample_mask_in_options.behavior = ac_nir_lower_samplemask_unknown_states_no_sample_shading;
+      }
+
+      NIR_PASS(_, stages[MESA_SHADER_FRAGMENT].nir, ac_nir_lower_sample_mask_in, &lower_sample_mask_in_options);
+
       /* Lower the view index to map on the layer. */
       NIR_PASS(_, stages[MESA_SHADER_FRAGMENT].nir, radv_nir_lower_view_index);
    }

src/amd/vulkan/radv_shader.h

@@ -170,9 +170,11 @@ struct radv_graphics_state_key {
    } rs;
 
    struct {
-      bool sample_shading_enable;
-      bool alpha_to_coverage_via_mrtz; /* GFX11+ */
+      bool sample_shading_enable : 1;
+      bool max_sample_shading_enable : 1;
+      bool alpha_to_coverage_via_mrtz : 1; /* GFX11+ */
       uint8_t rasterization_samples;
+      uint8_t ps_iter_samples; /* 0 if dynamic */
    } ms;
 
    struct vs {