nir: Add a new optimization for acquire/release atomics & barriers

Some shaders contain back-to-back atomic accesses in SPIR-V with
AcquireRelease semantics.  In NIR, we translate these to a release
memory barrier, the atomic, then an acquire memory barrier.

This results in a lot of unnecessary memory barriers in the middle
of the sequence of atomics:

   0. Release barrier
   1. Atomic
   2. Acquire barrier
   3. Release barrier
   4. Atomic
   5. Acquire barrier
   6. Release barrier
   7. Atomic
   8. Acquire barrier

In the absence of loads/stores, and when the atomic destinations are
unused, these barriers in-between atomics shouldn't be required.

This optimization pass would drop them (lines 2-3 and 5-6 above) while
leaving the first and last barriers (0 and 8), so the sequence remains
synchronized against other access elsewhere in the program.

One common example where this occurs is a sequence of min and max
atomics to clamp a certain memory location's value within a range.

Reviewed-by: Caio Oliveira <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33504>

src/compiler/nir/nir.h

@@ -6033,6 +6033,7 @@ typedef bool (*nir_combine_barrier_cb)(
 bool nir_opt_combine_barriers(nir_shader *shader,
                               nir_combine_barrier_cb combine_cb,
                               void *data);
+bool nir_opt_acquire_release_barriers(nir_shader *shader, mesa_scope max_scope);
 bool nir_opt_barrier_modes(nir_shader *shader);
 
 bool nir_minimize_call_live_states(nir_shader *shader);

src/compiler/nir/nir_opt_barriers.c

@@ -95,6 +95,199 @@ nir_opt_combine_barriers(nir_shader *shader,
    return progress;
 }
 
+/** If \p instr is a nir_intrinsic_barrier, returns it, else NULL. */
+static nir_intrinsic_instr *
+instr_as_barrier(nir_instr *instr)
+{
+   if (instr && instr->type == nir_instr_type_intrinsic) {
+      nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+      return intrin->intrinsic == nir_intrinsic_barrier ? intrin : NULL;
+   }
+   return NULL;
+}
+
+/**
+ * Return true if \p atomic is surrounded by a pattern:
+ *
+ *    1. Release barrier
+ *    2. Atomic operation
+ *    3. Acquire barrier
+ *
+ * where all three have the same mode, both barriers have the same scope,
+ * and that scope is \p max_scope or narrower.
+ *
+ * For simplicity, we require the barriers to have exactly the one mode
+ * used by the atomic, so that we don't have to compare many barriers for
+ * other side effects they may have.  nir_opt_barrier_modes() can be used
+ * to help reduce unnecessary barrier modes.
+ */
+static bool
+is_acquire_release_atomic(nir_intrinsic_instr *atomic, mesa_scope max_scope)
+{
+   assert(atomic->intrinsic == nir_intrinsic_deref_atomic ||
+          atomic->intrinsic == nir_intrinsic_deref_atomic_swap);
+
+   nir_deref_instr *atomic_deref = nir_src_as_deref(atomic->src[0]);
+
+   nir_intrinsic_instr *prev =
+      instr_as_barrier(nir_instr_prev(&atomic->instr));
+   nir_intrinsic_instr *next =
+      instr_as_barrier(nir_instr_next(&atomic->instr));
+
+   if (!prev || !next)
+      return false;
+
+   return nir_intrinsic_memory_semantics(prev) == NIR_MEMORY_RELEASE &&
+          nir_intrinsic_memory_semantics(next) == NIR_MEMORY_ACQUIRE &&
+          nir_intrinsic_memory_modes(prev) == atomic_deref->modes &&
+          nir_intrinsic_memory_modes(next) == atomic_deref->modes &&
+          nir_intrinsic_memory_scope(prev) <= max_scope &&
+          nir_intrinsic_memory_scope(prev) == nir_intrinsic_memory_scope(next);
+}
+
+/**
+ * Remove redundant barriers between sequences of atomics.
+ *
+ * Some shaders contain back-to-back atomic accesses in SPIR-V with
+ * AcquireRelease semantics.  In NIR, we translate these to a release
+ * memory barrier, the atomic, then an acquire memory barrier.
+ *
+ * This results in a lot of unnecessary memory barriers in the
+ * middle of the sequence of atomics:
+ *
+ *    1a. Release memory barrier
+ *    1b. Atomic
+ *    1c. Acquire memory barrier
+ *    ...
+ *    2a. Release memory barrier
+ *    2b. Atomic
+ *    2c. Acquire memory barrier
+ *    ...
+ *
+ * We pattern match for <release, atomic, acquire> instruction triplets,
+ * and when we find back-to-back occurrences of that pattern, we eliminate
+ * the barriers in-between the atomics (1c and 2a above):
+ *
+ *    1. Release memory barrier
+ *    2. Atomic
+ *    ...
+ *    m. Atomic
+ *    n. Acquire memory barrier
+ *
+ * Some requirements:
+ * - The atomics' destinations must be unused (so their only effect is
+ *   to update the associated memory store)
+ * - Matched barriers must impact the atomic's memory mode.
+ * - All barriers must have have identical scope no wider than \p max_scope
+ *   (beyond that, removing synchronization could be observable).
+ *
+ * And for simplicity:
+ * - Barrier modes must be exactly the mode of the atomics (otherwise we'd
+ *   have to take care to preserve side-effects for other modes).
+ * - Barriers must appear directly before/after the instruction (easier
+ *   pattern matching, and it's what we generate for the SPIR-V construct)
+ *
+ * Other instructions are allowed to be present between the atomics, so
+ * long as they don't affect the relevant memory mode.  Loads/stores or
+ * atomics not matching this pattern are not allowed (we stop matching).
+ * For example, this allows calculating the value to be used as the next
+ * atomic's operand to appear in-between the two.
+ */
+static bool
+nir_opt_acquire_release_barriers_impl(nir_function_impl *impl,
+                                      mesa_scope max_scope)
+{
+   bool progress = false;
+   nir_intrinsic_instr *last_atomic = NULL;
+
+   nir_foreach_block(block, impl) {
+      last_atomic = NULL;
+
+      nir_foreach_instr_safe(instr, block) {
+         if (instr->type != nir_instr_type_intrinsic)
+            continue;
+
+         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+
+         switch (intrin->intrinsic) {
+         case nir_intrinsic_load_deref:
+         case nir_intrinsic_load_deref_block_intel:
+         case nir_intrinsic_store_deref:
+         case nir_intrinsic_store_deref_block_intel:
+            if (last_atomic) {
+               /* If there is a load/store of the same mode as our matched
+                * atomic, then abandon our pattern match.
+                */
+               nir_deref_instr *ref = nir_src_as_deref(intrin->src[0]);
+               nir_deref_instr *lastdr = nir_src_as_deref(last_atomic->src[0]);
+               if (nir_deref_mode_may_be(ref, lastdr->modes))
+                  last_atomic = NULL;
+            }
+            break;
+
+         case nir_intrinsic_deref_atomic:
+         case nir_intrinsic_deref_atomic_swap:
+            if (nir_def_is_unused(&intrin->def) &&
+                is_acquire_release_atomic(intrin, max_scope)) {
+
+               if (!last_atomic) {
+                  last_atomic = intrin;
+               } else {
+                  nir_intrinsic_instr *last_acquire =
+                     nir_instr_as_intrinsic(nir_instr_next(&last_atomic->instr));
+                  nir_intrinsic_instr *this_release =
+                     nir_instr_as_intrinsic(nir_instr_prev(&intrin->instr));
+                  assert(last_acquire->intrinsic == nir_intrinsic_barrier);
+                  assert(this_release->intrinsic == nir_intrinsic_barrier);
+
+                  /* Verify that this atomic's barrier modes/scopes match
+                   * the last atomic's modes/scope.  (Note that we already
+                   * verified that each atomic's pair of barriers match
+                   * each other, so we can compare against either here.)
+                   */
+                  if (nir_intrinsic_memory_modes(last_acquire) ==
+                      nir_intrinsic_memory_modes(this_release) &&
+                      nir_intrinsic_memory_scope(last_acquire) ==
+                      nir_intrinsic_memory_scope(this_release)) {
+                     progress = true;
+                     nir_instr_remove(&last_acquire->instr);
+                     nir_instr_remove(&this_release->instr);
+                  }
+
+                  /* Regardless of progress, continue matching from here */
+                  last_atomic = intrin;
+               }
+            } else {
+               /* Abandon our pattern match, this is another kind of access */
+               last_atomic = NULL;
+            }
+            break;
+
+         default:
+            /* Ignore instructions that don't affect this kind of memory */
+            break;
+         }
+      }
+   }
+
+   nir_progress(progress, impl, nir_metadata_control_flow |
+                                nir_metadata_live_defs);
+
+   return progress;
+}
+
+bool
+nir_opt_acquire_release_barriers(nir_shader *shader, mesa_scope max_scope)
+{
+   bool progress = false;
+
+   nir_foreach_function_impl(impl, shader) {
+      progress |= nir_opt_acquire_release_barriers_impl(impl, max_scope);
+   }
+
+   return progress;
+}
+
 static bool
 barrier_happens_before(const nir_instr *a, const nir_instr *b)
 {