amd/nir: add ac_nir_opt_shared_append

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

src/amd/common/ac_nir.h

@@ -320,6 +320,9 @@ ac_nir_varying_expression_max_cost(nir_shader *producer, nir_shader *consumer);
 unsigned
 ac_nir_varying_estimate_instr_cost(nir_instr *instr);
 
+bool
+ac_nir_opt_shared_append(nir_shader *shader);
+
 #ifdef __cplusplus
 }
 #endif

src/amd/common/ac_nir_opt_shared_append.c

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2024 Valve Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "ac_nir.h"
+#include "nir_builder.h"
+
+static bool
+opt_shared_append(nir_builder *b,
+                  nir_intrinsic_instr *intrin,
+                  void *unused)
+{
+   if (intrin->intrinsic != nir_intrinsic_shared_atomic)
+      return false;
+   if (nir_intrinsic_atomic_op(intrin) != nir_atomic_op_iadd)
+      return false;
+   if (intrin->def.bit_size != 32)
+      return false;
+
+   if (!nir_src_is_const(intrin->src[0]) || !nir_src_is_const(intrin->src[1]))
+      return false;
+
+   uint32_t addr = nir_src_as_uint(intrin->src[0]) + nir_intrinsic_base(intrin);
+   int32_t data = nir_src_as_int(intrin->src[1]);
+
+   if (data != 1 && data != -1)
+      return false;
+
+   if (addr >= 65536 || addr % 4)
+      return false;
+
+   b->cursor = nir_before_instr(&intrin->instr);
+   nir_def *res;
+   if (data == 1)
+      res = nir_shared_append_amd(b, .base = addr);
+   else
+      res = nir_shared_consume_amd(b, .base = addr);
+
+   if (nir_def_is_unused(&intrin->def)) {
+      nir_instr_remove(&intrin->instr);
+      return true;
+   }
+
+   res = nir_iadd(b, res, nir_exclusive_scan(b, intrin->src[1].ssa, .reduction_op = nir_op_iadd));
+   nir_def_replace(&intrin->def, res);
+   return true;
+}
+
+
+bool
+ac_nir_opt_shared_append(nir_shader *shader)
+{
+   return nir_shader_intrinsics_pass(shader, opt_shared_append, nir_metadata_control_flow, NULL);
+}

src/amd/common/meson.build

@@ -100,6 +100,7 @@ amd_common_files = files(
   'ac_nir_meta_cs_blit.c',
   'ac_nir_meta_cs_clear_copy_buffer.c',
   'ac_nir_meta_ps_resolve.c',
+  'ac_nir_opt_shared_append.c',
   'amd_family.c',
   'ac_parse_ib.c',
   'ac_perfcounter.c',