nir: Add nir_lower_helper_writes pass

This NIR pass lowers stores in fragment shaders to:

   if (!gl_HelperInvocaton) {
      store();
   }

This implements the API requirement that helper invocations do not have visible
side effects, and the lowering is required on any hardware that cannot directly
mask helper invocation's side effects. The pass was originally written for
Midgard (which has this issue) but is also needed for Asahi. Let's share the
code, and fix it while we're at it.

Changes from the Midgard pass:

1. Add an option to only lower atomics.

   AGX hardware can mask helper invocations for "plain" stores but not for
   atomics.  Accordingly, the AGX compiler wants this lowering for atomics but
   not store_global. By contrast, Midgard cannot mask any stores and needs the
   lowering for all store intrinsics. Add an option to the common pass to
   accommodate both cases.

   This is an optimization for AGX. It is not required for correctness, this
   lowering is always legal.

2. Fix dominance issues.

   It's invalid to have NIR like

      if ... {
         ssa_1 = ...
      }

      foo ssa_1

   Instead we need to rewrite as

      if ... {
         ssa_1 = ...
      } else {
         ssa_2 = undef
      }
      ssa_3 = phi ssa_1, ssa_2

      foo ssa_3

   By default, neither nir_validate nor the backends check this, so this doesn't
   currently fix a (known) real bug. But it's still invalid and fails validation
   with NIR_DEBUG=validate_ssa_dominance.

   Fix this in lower_helper_writes for intrinsics that return data (atomics).

3. Assert that the pass is run only for fragment shaders. This encourages
   backends to be judicious about which passes they call instead of just
   throwing everything in a giant lower everything spaghetti.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Italo Nicola <italonicola@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21413>

src/compiler/nir/meson.build

@@ -162,6 +162,7 @@ files_libnir = files(
   'nir_lower_global_vars_to_local.c',
   'nir_lower_goto_ifs.c',
   'nir_lower_gs_intrinsics.c',
+  'nir_lower_helper_writes.c',
   'nir_lower_load_const_to_scalar.c',
   'nir_lower_locals_to_regs.c',
   'nir_lower_idiv.c',

src/compiler/nir/nir.h

@@ -5684,6 +5684,7 @@ bool nir_lower_samplers(nir_shader *shader);
 bool nir_lower_cl_images(nir_shader *shader, bool lower_image_derefs, bool lower_sampler_derefs);
 bool nir_dedup_inline_samplers(nir_shader *shader);
 bool nir_lower_ssbo(nir_shader *shader);
+bool nir_lower_helper_writes(nir_shader *shader, bool lower_plain_stores);
 
 typedef struct nir_lower_printf_options {
    bool treat_doubles_as_floats : 1;

src/compiler/nir/nir_lower_helper_writes.c

@@ -0,0 +1,134 @@
+/*
+ * Copyright (C) 2020-2021 Collabora, Ltd.
+ * Copyright © 2020 Valve Corporation
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice (including the next
+ * paragraph) shall be included in all copies or substantial portions of the
+ * Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ */
+
+#include "compiler/nir/nir.h"
+#include "compiler/nir/nir_builder.h"
+
+static bool
+lower(nir_builder *b, nir_instr *instr, void *data)
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+   bool *lower_plain_stores = data;
+
+   switch (intr->intrinsic) {
+   case nir_intrinsic_global_atomic_add:
+   case nir_intrinsic_global_atomic_and:
+   case nir_intrinsic_global_atomic_comp_swap:
+   case nir_intrinsic_global_atomic_exchange:
+   case nir_intrinsic_global_atomic_fadd:
+   case nir_intrinsic_global_atomic_fcomp_swap:
+   case nir_intrinsic_global_atomic_fmax:
+   case nir_intrinsic_global_atomic_fmin:
+   case nir_intrinsic_global_atomic_imax:
+   case nir_intrinsic_global_atomic_imin:
+   case nir_intrinsic_global_atomic_or:
+   case nir_intrinsic_global_atomic_umax:
+   case nir_intrinsic_global_atomic_umin:
+   case nir_intrinsic_global_atomic_xor:
+   case nir_intrinsic_image_atomic_add:
+   case nir_intrinsic_image_atomic_and:
+   case nir_intrinsic_image_atomic_comp_swap:
+   case nir_intrinsic_image_atomic_dec_wrap:
+   case nir_intrinsic_image_atomic_exchange:
+   case nir_intrinsic_image_atomic_fadd:
+   case nir_intrinsic_image_atomic_imax:
+   case nir_intrinsic_image_atomic_imin:
+   case nir_intrinsic_image_atomic_inc_wrap:
+   case nir_intrinsic_image_atomic_or:
+   case nir_intrinsic_image_atomic_umax:
+   case nir_intrinsic_image_atomic_umin:
+   case nir_intrinsic_image_atomic_xor:
+      break;
+   case nir_intrinsic_store_global:
+   case nir_intrinsic_image_store:
+      if (!(*lower_plain_stores))
+         return false;
+      else
+         break;
+   default:
+      return false;
+   }
+
+   b->cursor = nir_before_instr(instr);
+   bool has_dest = nir_intrinsic_infos[intr->intrinsic].has_dest;
+   nir_ssa_def *undef = NULL;
+
+   nir_ssa_def *helper = nir_load_helper_invocation(b, 1);
+   nir_push_if(b, nir_inot(b, helper));
+   nir_instr_remove(instr);
+   nir_builder_instr_insert(b, instr);
+
+   /* Per the spec, it does not matter what we return for helper threads.
+    * Represent this by an ssa_undef in the hopes the backend will be clever
+    * enough to optimize out the phi.
+    *
+    *    Fragment shader helper invocations execute the same shader code as
+    *    non-helper invocations, but will not have side effects that modify the
+    *    framebuffer or other shader-accessible memory. In particular:
+    *
+    *       ...
+    *
+    *       Atomic operations to image, buffer, or atomic counter variables
+    *       performed by helper invocations have no effect on the underlying
+    *       image or buffer memory. The values returned by such atomic
+    *       operations are undefined.
+    */
+   if (has_dest) {
+      nir_push_else(b, NULL);
+      undef = nir_ssa_undef(b, nir_dest_num_components(intr->dest),
+                            nir_dest_bit_size(intr->dest));
+   }
+
+   nir_pop_if(b, NULL);
+
+   if (has_dest) {
+      assert(intr->dest.is_ssa);
+      nir_ssa_def *phi = nir_if_phi(b, &intr->dest.ssa, undef);
+
+      /* We can't use nir_ssa_def_rewrite_uses_after on phis, so use the global
+       * version and fixup the phi manually
+       */
+      nir_ssa_def_rewrite_uses(&intr->dest.ssa, phi);
+
+      nir_instr *phi_instr = phi->parent_instr;
+      nir_phi_instr *phi_as_phi = nir_instr_as_phi(phi_instr);
+      nir_phi_src *phi_src = nir_phi_get_src_from_block(phi_as_phi,
+                                                        instr->block);
+      nir_instr_rewrite_src_ssa(phi->parent_instr, &phi_src->src,
+                                &intr->dest.ssa);
+   }
+
+   return true;
+}
+
+bool
+nir_lower_helper_writes(nir_shader *shader, bool lower_plain_stores)
+{
+   assert(shader->info.stage == MESA_SHADER_FRAGMENT);
+   return nir_shader_instructions_pass(shader, lower, nir_metadata_none,
+                                       &lower_plain_stores);
+}

src/panfrost/midgard/midgard_nir_lower_helper_writes.c

@@ -27,7 +27,7 @@
 #include "midgard_nir.h"
 
 static bool
-nir_lower_helper_writes(nir_builder *b, nir_instr *instr, UNUSED void *data)
+pass(nir_builder *b, nir_instr *instr, UNUSED void *data)
 {
    if (instr->type != nir_instr_type_intrinsic)
       return false;
@@ -86,6 +86,5 @@ midgard_nir_lower_helper_writes(nir_shader *shader)
    if (shader->info.stage != MESA_SHADER_FRAGMENT)
       return false;
 
-   return nir_shader_instructions_pass(shader, nir_lower_helper_writes,
-                                       nir_metadata_none, NULL);
+   return nir_shader_instructions_pass(shader, pass, nir_metadata_none, NULL);
 }