radv: lower load_sample_mask_in in NIR

No fossils-db changes.

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/15450>

src/amd/compiler/aco_instruction_selection.cpp

@@ -7015,34 +7015,6 @@ visit_store_scratch(isel_context* ctx, nir_intrinsic_instr* instr)
    }
 }
 
-void
-visit_load_sample_mask_in(isel_context* ctx, nir_intrinsic_instr* instr)
-{
-   uint8_t log2_ps_iter_samples;
-   if (ctx->program->info->ps.uses_sample_shading) {
-      log2_ps_iter_samples = util_logbase2(ctx->options->key.ps.num_samples);
-   } else {
-      log2_ps_iter_samples = ctx->options->key.ps.log2_ps_iter_samples;
-   }
-
-   Builder bld(ctx->program, ctx->block);
-
-   Temp dst = get_ssa_temp(ctx, &instr->dest.ssa);
-
-   if (log2_ps_iter_samples) {
-      /* gl_SampleMaskIn[0] = (SampleCoverage & (1 << gl_SampleID)). */
-      Temp sample_id =
-         bld.vop3(aco_opcode::v_bfe_u32, bld.def(v1), get_arg(ctx, ctx->args->ac.ancillary),
-                  Operand::c32(8u), Operand::c32(4u));
-      Temp mask = bld.vop2(aco_opcode::v_lshlrev_b32, bld.def(v1), sample_id,
-                           bld.copy(bld.def(v1), Operand::c32(1u)));
-      bld.vop2(aco_opcode::v_and_b32, Definition(dst), mask,
-               get_arg(ctx, ctx->args->ac.sample_coverage));
-   } else {
-      bld.copy(Definition(dst), get_arg(ctx, ctx->args->ac.sample_coverage));
-   }
-}
-
 void
 visit_emit_vertex_with_counter(isel_context* ctx, nir_intrinsic_instr* instr)
 {
@@ -8029,10 +8001,6 @@ visit_intrinsic(isel_context* ctx, nir_intrinsic_instr* instr)
                get_arg(ctx, ctx->args->ac.ancillary), Operand::c32(8u), Operand::c32(4u));
       break;
    }
-   case nir_intrinsic_load_sample_mask_in: {
-      visit_load_sample_mask_in(ctx, instr);
-      break;
-   }
    case nir_intrinsic_read_first_invocation: {
       Temp src = get_ssa_temp(ctx, instr->src[0].ssa);
       Temp dst = get_ssa_temp(ctx, &instr->dest.ssa);

src/amd/compiler/aco_instruction_selection_setup.cpp

@@ -635,7 +635,6 @@ init_context(isel_context* ctx, nir_shader* shader)
                case nir_intrinsic_load_scalar_arg_amd:
                case nir_intrinsic_load_smem_amd: type = RegType::sgpr; break;
                case nir_intrinsic_load_sample_id:
-               case nir_intrinsic_load_sample_mask_in:
                case nir_intrinsic_load_input:
                case nir_intrinsic_load_output:
                case nir_intrinsic_load_input_vertex:

src/amd/vulkan/radv_nir_to_llvm.c

@@ -239,32 +239,6 @@ load_sample_position(struct ac_shader_abi *abi, LLVMValueRef sample_id)
    return result;
 }
 
-static LLVMValueRef
-load_sample_mask_in(struct ac_shader_abi *abi)
-{
-   struct radv_shader_context *ctx = radv_shader_context_from_abi(abi);
-   uint8_t log2_ps_iter_samples;
-
-   if (ctx->shader_info->ps.uses_sample_shading) {
-      log2_ps_iter_samples = util_logbase2(ctx->options->key.ps.num_samples);
-   } else {
-      log2_ps_iter_samples = ctx->options->key.ps.log2_ps_iter_samples;
-   }
-
-   LLVMValueRef result, sample_id;
-   if (log2_ps_iter_samples) {
-      /* gl_SampleMaskIn[0] = (SampleCoverage & (1 << gl_SampleID)). */
-      sample_id = ac_unpack_param(&ctx->ac, ac_get_arg(&ctx->ac, ctx->args->ac.ancillary), 8, 4);
-      sample_id = LLVMBuildShl(ctx->ac.builder, LLVMConstInt(ctx->ac.i32, 1, false), sample_id, "");
-      result = LLVMBuildAnd(ctx->ac.builder, sample_id,
-                            ac_get_arg(&ctx->ac, ctx->args->ac.sample_coverage), "");
-   } else {
-      result = ac_get_arg(&ctx->ac, ctx->args->ac.sample_coverage);
-   }
-
-   return result;
-}
-
 static void gfx10_ngg_gs_emit_vertex(struct radv_shader_context *ctx, unsigned stream,
                                      LLVMValueRef vertexidx, LLVMValueRef *addrs);
 
@@ -2202,7 +2176,6 @@ ac_translate_nir_to_llvm(struct ac_llvm_compiler *ac_llvm,
          ctx.abi.load_inputs = radv_load_vs_inputs;
       } else if (shaders[shader_idx]->info.stage == MESA_SHADER_FRAGMENT) {
          ctx.abi.load_sample_position = load_sample_position;
-         ctx.abi.load_sample_mask_in = load_sample_mask_in;
       }
 
       if (shaders[shader_idx]->info.stage == MESA_SHADER_VERTEX && info->is_ngg &&

src/amd/vulkan/radv_pipeline.c

@@ -4217,6 +4217,11 @@ radv_create_shaders(struct radv_pipeline *pipeline, struct radv_pipeline_layout
    struct radv_shader_args args[MESA_VULKAN_SHADER_STAGES] = {{{{{0}}}}};
    radv_declare_pipeline_args(device, args, nir, infos, pipeline_key);
 
+   if (nir[MESA_SHADER_FRAGMENT]) {
+      NIR_PASS_V(nir[MESA_SHADER_FRAGMENT], radv_lower_fs_intrinsics, &infos[MESA_SHADER_FRAGMENT],
+                 &args[MESA_SHADER_FRAGMENT], pipeline_key);
+   }
+
    for (int i = 0; i < MESA_VULKAN_SHADER_STAGES; ++i) {
       if (nir[i]) {
          radv_start_feedback(stage_feedbacks[i]);

src/amd/vulkan/radv_shader.c

@@ -468,6 +468,61 @@ radv_force_primitive_shading_rate(nir_shader *nir, struct radv_device *device)
    return progress;
 }
 
+bool
+radv_lower_fs_intrinsics(nir_shader *nir, const struct radv_shader_info *info,
+                         const struct radv_shader_args *args, const struct radv_pipeline_key *key)
+{
+   nir_function_impl *impl = nir_shader_get_entrypoint(nir);
+   bool progress = false;
+
+   nir_builder b;
+   nir_builder_init(&b, impl);
+
+   nir_foreach_block(block, impl) {
+      nir_foreach_instr_safe(instr, block) {
+         if (instr->type != nir_instr_type_intrinsic)
+            continue;
+
+         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+         b.cursor = nir_after_instr(&intrin->instr);
+
+         switch (intrin->intrinsic) {
+         case nir_intrinsic_load_sample_mask_in: {
+            uint8_t log2_ps_iter_samples;
+
+            if (info->ps.uses_sample_shading) {
+               log2_ps_iter_samples = util_logbase2(key->ps.num_samples);
+            } else {
+               log2_ps_iter_samples = key->ps.log2_ps_iter_samples;
+            }
+
+            nir_ssa_def *sample_coverage =
+               nir_load_vector_arg_amd(&b, 1, .base = args->ac.sample_coverage.arg_index);
+
+            nir_ssa_def *def = NULL;
+            if (log2_ps_iter_samples) {
+               /* gl_SampleMaskIn[0] = (SampleCoverage & (1 << gl_SampleID)). */
+               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));
+            } else {
+               def = sample_coverage;
+            }
+
+            nir_ssa_def_rewrite_uses(&intrin->dest.ssa, def);
+
+            nir_instr_remove(instr);
+            progress = true;
+            break;
+         }
+         default:
+            break;
+         }
+      }
+   }
+
+   return progress;
+}
+
 nir_shader *
 radv_shader_compile_to_nir(struct radv_device *device, struct vk_shader_module *module,
                            const char *entrypoint_name, gl_shader_stage stage,

src/amd/vulkan/radv_shader.h

@@ -682,4 +682,8 @@ void radv_get_nir_options(struct radv_physical_device *device);
 
 bool radv_force_primitive_shading_rate(nir_shader *nir, struct radv_device *device);
 
+bool radv_lower_fs_intrinsics(nir_shader *nir, const struct radv_shader_info *info,
+                              const struct radv_shader_args *args,
+                              const struct radv_pipeline_key *key);
+
 #endif