From 2789cee0c0a96058a1bf7447dfe56bfcd2fc19c0 Mon Sep 17 00:00:00 2001
From: Georg Lehmann <dadschoorse@gmail.com>
Date: Sat, 7 Sep 2024 14:49:47 +0200
Subject: [PATCH] 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                   |  3 ++
 src/amd/common/ac_nir_opt_shared_append.c | 56 +++++++++++++++++++++++
 src/amd/common/meson.build                |  1 +
 3 files changed, 60 insertions(+)
 create mode 100644 src/amd/common/ac_nir_opt_shared_append.c

diff --git a/src/amd/common/ac_nir.h b/src/amd/common/ac_nir.h
index ab6ff6d6c5f..15d5555bef1 100644
--- a/src/amd/common/ac_nir.h
+++ b/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
diff --git a/src/amd/common/ac_nir_opt_shared_append.c b/src/amd/common/ac_nir_opt_shared_append.c
new file mode 100644
index 00000000000..c97af7ebfd9
--- /dev/null
+++ b/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);
+}
diff --git a/src/amd/common/meson.build b/src/amd/common/meson.build
index 74329ffa5dd..9fb3fdbecc3 100644
--- a/src/amd/common/meson.build
+++ b/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',