From fa321e0cc4a4e92e5e9296af317d475c2305482c Mon Sep 17 00:00:00 2001
From: Mike Blumenkrantz <michael.blumenkrantz@gmail.com>
Date: Tue, 4 Aug 2020 14:28:50 -0400
Subject: [PATCH] zink: start supporting atomic shader ops

this handles atomic_add, which is what atomic counters use

Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/8330>
---
 .../drivers/zink/nir_to_spirv/nir_to_spirv.c  | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
index b092ade9c3e..4d03e11f017 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
@@ -2038,6 +2038,50 @@ emit_interpolate(struct ntv_context *ctx, nir_intrinsic_instr *intr)
    store_dest(ctx, &intr->dest, result, nir_type_uint);
 }
 
+static void
+emit_atomic_intrinsic(struct ntv_context *ctx, nir_intrinsic_instr *intr)
+{
+   SpvId result = 0;
+   SpvId ssbo;
+   SpvId param;
+   SpvId dest_type = get_dest_type(ctx, &intr->dest, nir_type_uint32);
+
+   nir_const_value *const_block_index = nir_src_as_const_value(intr->src[0]);
+   assert(const_block_index); // no dynamic indexing for now
+   ssbo = ctx->ssbos[const_block_index->u32];
+   param = get_src(ctx, &intr->src[2]);
+
+   SpvId pointer_type = spirv_builder_type_pointer(&ctx->builder,
+                                                   SpvStorageClassStorageBuffer,
+                                                   dest_type);
+   SpvId uint_type = get_uvec_type(ctx, 32, 1);
+   /* an id of the array stride in bytes */
+   SpvId vec4_size = emit_uint_const(ctx, 32, sizeof(uint32_t) * 4);
+   /* an id of an array member in bytes */
+   SpvId uint_size = emit_uint_const(ctx, 32, sizeof(uint32_t));
+   SpvId member = emit_uint_const(ctx, 32, 0);
+   SpvId offset = get_src(ctx, &intr->src[1]);
+   SpvId vec_offset = emit_binop(ctx, SpvOpUDiv, uint_type, offset, vec4_size);
+   SpvId vec_member_offset = emit_binop(ctx, SpvOpUDiv, uint_type,
+                                        emit_binop(ctx, SpvOpUMod, uint_type, offset, vec4_size),
+                                        uint_size);
+   SpvId indices[3] = { member, vec_offset, vec_member_offset };
+   SpvId ptr = spirv_builder_emit_access_chain(&ctx->builder, pointer_type,
+                                               ssbo, indices,
+                                               ARRAY_SIZE(indices));
+   switch (intr->intrinsic) {
+   case nir_intrinsic_ssbo_atomic_add:
+      result = emit_atomic(ctx, SpvOpAtomicIAdd, dest_type, ptr, param);
+      break;
+   default:
+      fprintf(stderr, "emit_atomic_intrinsic: not implemented (%s)\n",
+              nir_intrinsic_infos[intr->intrinsic].name);
+      unreachable("unsupported intrinsic");
+   }
+   assert(result);
+   store_dest(ctx, &intr->dest, result, nir_type_uint);
+}
+
 static void
 emit_intrinsic(struct ntv_context *ctx, nir_intrinsic_instr *intr)
 {
@@ -2145,6 +2189,10 @@ emit_intrinsic(struct ntv_context *ctx, nir_intrinsic_instr *intr)
       emit_interpolate(ctx, intr);
       break;
 
+   case nir_intrinsic_ssbo_atomic_add:
+      emit_atomic_intrinsic(ctx, intr);
+      break;
+
    default:
       fprintf(stderr, "emit_intrinsic: not implemented (%s)\n",
               nir_intrinsic_infos[intr->intrinsic].name);