nir/opt_uniform_atomics: add fs atomics predicated? flag

on agx (and mali), we predicate atomics on "if (!helper)", so doing so again in
this pass is redundant. and would cause a problem since we'd then have to lower
the "is helper inv?" flag late. so just skip the extra lowering code.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Acked-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30488>

src/amd/compiler/aco_instruction_selection_setup.cpp

@@ -279,7 +279,7 @@ init_context(isel_context* ctx, nir_shader* shader)
    ctx->ub_config.max_workgroup_size[2] = 2048;
 
    nir_divergence_analysis(shader);
-   if (nir_opt_uniform_atomics(shader) && nir_lower_int64(shader))
+   if (nir_opt_uniform_atomics(shader, false) && nir_lower_int64(shader))
       nir_divergence_analysis(shader);
 
    apply_nuw_to_offsets(ctx, impl);

src/compiler/nir/nir.h

@@ -6718,7 +6718,7 @@ bool nir_opt_undef(nir_shader *shader);
 
 bool nir_lower_undef_to_zero(nir_shader *shader);
 
-bool nir_opt_uniform_atomics(nir_shader *shader);
+bool nir_opt_uniform_atomics(nir_shader *shader, bool fs_atomics_predicated);
 
 bool nir_opt_uniform_subgroup(nir_shader *shader,
                               const nir_lower_subgroups_options *);

src/compiler/nir/nir_opt_uniform_atomics.c

@@ -287,10 +287,11 @@ optimize_atomic(nir_builder *b, nir_intrinsic_instr *intrin, bool return_prev)
 }
 
 static void
-optimize_and_rewrite_atomic(nir_builder *b, nir_intrinsic_instr *intrin)
+optimize_and_rewrite_atomic(nir_builder *b, nir_intrinsic_instr *intrin,
+                            bool fs_atomics_predicated)
 {
    nir_if *helper_nif = NULL;
-   if (b->shader->info.stage == MESA_SHADER_FRAGMENT) {
+   if (b->shader->info.stage == MESA_SHADER_FRAGMENT && !fs_atomics_predicated) {
       nir_def *helper = nir_is_helper_invocation(b, 1);
       helper_nif = nir_push_if(b, nir_inot(b, helper));
    }
@@ -320,7 +321,7 @@ optimize_and_rewrite_atomic(nir_builder *b, nir_intrinsic_instr *intrin)
 }
 
 static bool
-opt_uniform_atomics(nir_function_impl *impl)
+opt_uniform_atomics(nir_function_impl *impl, bool fs_atomics_predicated)
 {
    bool progress = false;
    nir_builder b = nir_builder_create(impl);
@@ -346,7 +347,7 @@ opt_uniform_atomics(nir_function_impl *impl)
             continue;
 
          b.cursor = nir_before_instr(instr);
-         optimize_and_rewrite_atomic(&b, intrin);
+         optimize_and_rewrite_atomic(&b, intrin, fs_atomics_predicated);
          progress = true;
       }
    }
@@ -355,7 +356,7 @@ opt_uniform_atomics(nir_function_impl *impl)
 }
 
 bool
-nir_opt_uniform_atomics(nir_shader *shader)
+nir_opt_uniform_atomics(nir_shader *shader, bool fs_atomics_predicated)
 {
    bool progress = false;
 
@@ -371,7 +372,7 @@ nir_opt_uniform_atomics(nir_shader *shader)
    nir_foreach_function_impl(impl, shader) {
       nir_metadata_require(impl, nir_metadata_block_index);
 
-      if (opt_uniform_atomics(impl)) {
+      if (opt_uniform_atomics(impl, fs_atomics_predicated)) {
          progress = true;
          nir_metadata_preserve(impl, nir_metadata_none);
       } else {

src/intel/compiler/brw_nir.c

@@ -1790,7 +1790,7 @@ brw_postprocess_nir(nir_shader *nir, const struct brw_compiler *compiler,
       .lower_subgroup_masks = true,
    };
 
-   if (OPT(nir_opt_uniform_atomics)) {
+   if (OPT(nir_opt_uniform_atomics, false)) {
       OPT(nir_lower_subgroups, &subgroups_options);
 
       OPT(nir_opt_algebraic_before_lower_int64);

src/intel/compiler/elk/elk_nir.c

@@ -1472,7 +1472,7 @@ elk_postprocess_nir(nir_shader *nir, const struct elk_compiler *compiler,
     */
    bool opt_uniform_atomic_stage_allowed = devinfo->ver >= 8;
 
-   if (opt_uniform_atomic_stage_allowed && OPT(nir_opt_uniform_atomics)) {
+   if (opt_uniform_atomic_stage_allowed && OPT(nir_opt_uniform_atomics, false)) {
       const nir_lower_subgroups_options subgroups_options = {
          .ballot_bit_size = 32,
          .ballot_components = 1,