nir: add pass to lower atomic arithmetic to a loop with cmpxchg.

Signed-off-by: Amber Harmonia <amber@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27776>

src/compiler/nir/meson.build

@@ -294,6 +294,7 @@ files_libnir = files(
   'nir_worklist.c',
   'nir_worklist.h',
   'nir_xfb_info.h',
+  'nir_lower_atomics.c'
 )
 
 # When a file is too large, -Wmisleading-indentation will give a note about

src/compiler/nir/nir.h

@@ -6043,6 +6043,8 @@ bool nir_lower_is_helper_invocation(nir_shader *shader);
 
 bool nir_lower_single_sampled(nir_shader *shader);
 
+bool nir_lower_atomics(nir_shader *shader, nir_instr_filter_cb filter);
+
 typedef struct nir_lower_subgroups_options {
    /* In addition to the boolean lowering options below, this optional callback
     * will filter instructions for lowering if non-NULL. The data passed will be

src/compiler/nir/nir_lower_atomics.c

@@ -0,0 +1,127 @@
+/*
+ * Copyright © 2024 Igalia
+ *
+ * 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 "nir.h"
+#include "nir_builder.h"
+#include "nir_builder_opcodes.h"
+#include "nir_intrinsics.h"
+
+/*
+ * Convert atomic arithmetic to regular arithmetic along with cmpxchg
+ * by repeating the operation until the result is expected.
+ *
+ * eg:
+ * atomicAdd(a[0], 1) ->
+ *
+ * uint expected = a[0];
+ * while (true) {
+ *    uint before = expected;
+ *    expected += 1;
+ *    uint original = atomicCompareExchange(a[0], before, expected);
+ *    if (original == before) {break;}
+ *    expected = original;
+ * }
+ */
+
+static nir_def *
+build_atomic(nir_builder *b, nir_intrinsic_instr *intr, bool ssbo)
+{
+   nir_def *load = ssbo ? nir_load_ssbo(b, 1, intr->def.bit_size, intr->src[0].ssa,
+                                        intr->src[1].ssa,
+                                        .align_mul = intr->def.bit_size / 8,
+                                        .align_offset = 0)
+                        : nir_load_global(b, intr->src[0].ssa, 8, 1, intr->def.bit_size);
+   nir_def *data = ssbo ? intr->src[2].ssa : intr->src[1].ssa;
+   nir_loop *loop = nir_push_loop(b);
+   nir_def *xchg;
+   {
+      nir_phi_instr *phi = nir_phi_instr_create(b->shader);
+      nir_def_init(&phi->instr, &phi->def, 1, intr->def.bit_size);
+      nir_phi_instr_add_src(phi, load->parent_instr->block, load);
+      nir_def *before = &phi->def;
+      nir_def *expected = nir_build_alu2(
+         b, nir_atomic_op_to_alu(nir_intrinsic_atomic_op(intr)), before, data);
+      nir_alu_instr* op = nir_instr_as_alu(expected->parent_instr);
+      op->exact = true;
+      op->fp_fast_math = 0;
+      if (ssbo) {
+         xchg = nir_ssbo_atomic_swap(b, intr->def.bit_size, intr->src[0].ssa, intr->src[1].ssa,
+                                     before, expected,
+                                     .atomic_op = nir_atomic_op_cmpxchg);
+      } else {
+         xchg =
+            nir_global_atomic_swap(b, intr->def.bit_size, intr->src[0].ssa, before, expected,
+                                   .atomic_op = nir_atomic_op_cmpxchg);
+      }
+      nir_break_if(b, nir_ieq(b, xchg, before));
+      nir_phi_instr_add_src(phi, nir_loop_last_block(loop), xchg);
+      b->cursor = nir_before_block(nir_loop_first_block(loop));
+      nir_builder_instr_insert(b, &phi->instr);
+   }
+   nir_pop_loop(b, loop);
+   return xchg;
+}
+
+static bool
+lower_atomics(struct nir_builder *b, nir_intrinsic_instr *intr,
+              void *supported)
+{
+   nir_instr_filter_cb supported_cb = supported;
+
+   if (intr->intrinsic != nir_intrinsic_ssbo_atomic &&
+       intr->intrinsic != nir_intrinsic_global_atomic)
+      return false;
+   if (supported_cb(&intr->instr, NULL))
+      return false;
+   b->cursor = nir_before_instr(&intr->instr);
+   switch (nir_intrinsic_atomic_op(intr)) {
+   case nir_atomic_op_imin:
+   case nir_atomic_op_umin:
+   case nir_atomic_op_imax:
+   case nir_atomic_op_umax:
+   case nir_atomic_op_iand:
+   case nir_atomic_op_ior:
+   case nir_atomic_op_ixor:
+   case nir_atomic_op_fadd:
+   case nir_atomic_op_fmin:
+   case nir_atomic_op_fmax:
+   case nir_atomic_op_iadd: {
+      nir_def_replace(&intr->def, build_atomic(
+                                     b, intr, intr->intrinsic == nir_intrinsic_ssbo_atomic));
+      return true;
+   }
+   case nir_atomic_op_cmpxchg:
+   case nir_atomic_op_xchg:
+      return false;
+   case nir_atomic_op_fcmpxchg: /* unimplemented */
+   default:
+      unreachable("Invalid nir_atomic_op");
+   }
+}
+
+bool
+nir_lower_atomics(nir_shader *shader, nir_instr_filter_cb supported)
+{
+   return nir_shader_intrinsics_pass(shader, lower_atomics,
+                                     nir_metadata_none, supported);
+}