From 57e0dbe55b424477fcfa9c4b4416f272fc4948c8 Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Fri, 2 Dec 2022 21:34:33 -0500
Subject: [PATCH] asahi: Implement load_ssbo_address/get_ssbo_size

More uniforms that get pushed.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21062>
---
 .../drivers/asahi/agx_nir_lower_sysvals.c     |  6 ++++++
 src/gallium/drivers/asahi/agx_state.h         |  4 ++++
 src/gallium/drivers/asahi/agx_uniforms.c      | 20 +++++++++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/src/gallium/drivers/asahi/agx_nir_lower_sysvals.c b/src/gallium/drivers/asahi/agx_nir_lower_sysvals.c
index 77568aca2a5..050bab9af5f 100644
--- a/src/gallium/drivers/asahi/agx_nir_lower_sysvals.c
+++ b/src/gallium/drivers/asahi/agx_nir_lower_sysvals.c
@@ -66,6 +66,12 @@ pass(struct nir_builder *b, nir_instr *instr, void *data)
    case nir_intrinsic_load_blend_const_color_a_float:
       ptr = &u->fs.blend_constant[3];
       break;
+   case nir_intrinsic_load_ssbo_address:
+      ptr = &u->ssbo_base[nir_src_as_uint(intr->src[0])];
+      break;
+   case nir_intrinsic_get_ssbo_size:
+      ptr = &u->ssbo_size[nir_src_as_uint(intr->src[0])];
+      break;
    default:
       return false;
    }
diff --git a/src/gallium/drivers/asahi/agx_state.h b/src/gallium/drivers/asahi/agx_state.h
index 3d728ea8a40..6133f4de40f 100644
--- a/src/gallium/drivers/asahi/agx_state.h
+++ b/src/gallium/drivers/asahi/agx_state.h
@@ -75,6 +75,10 @@ struct PACKED agx_draw_uniforms {
    /* Uniform buffer objects */
    uint64_t ubo_base[PIPE_MAX_CONSTANT_BUFFERS];
 
+   /* Shader storage buffer objects */
+   uint64_t ssbo_base[PIPE_MAX_SHADER_BUFFERS];
+   uint32_t ssbo_size[PIPE_MAX_SHADER_BUFFERS];
+
    union {
       struct {
          /* Vertex buffer object bases, if present */
diff --git a/src/gallium/drivers/asahi/agx_uniforms.c b/src/gallium/drivers/asahi/agx_uniforms.c
index 4d2c5b509a6..1149cc38ca9 100644
--- a/src/gallium/drivers/asahi/agx_uniforms.c
+++ b/src/gallium/drivers/asahi/agx_uniforms.c
@@ -39,6 +39,21 @@ agx_const_buffer_ptr(struct agx_batch *batch, struct pipe_constant_buffer *cb)
    }
 }
 
+static uint64_t
+agx_shader_buffer_ptr(struct agx_batch *batch, struct pipe_shader_buffer *sb)
+{
+   if (sb->buffer) {
+      struct agx_resource *rsrc = agx_resource(sb->buffer);
+
+      /* Assume SSBOs are written. TODO: Optimize read-only SSBOs */
+      agx_batch_writes(batch, rsrc);
+
+      return rsrc->bo->ptr.gpu + sb->buffer_offset;
+   } else {
+      return 0;
+   }
+}
+
 static uint64_t
 agx_vertex_buffer_ptr(struct agx_batch *batch, unsigned vbo)
 {
@@ -68,6 +83,11 @@ agx_upload_uniforms(struct agx_batch *batch, uint64_t textures,
       uniforms.ubo_base[cb] = agx_const_buffer_ptr(batch, &st->cb[cb]);
    }
 
+   u_foreach_bit(cb, st->ssbo_mask) {
+      uniforms.ssbo_base[cb] = agx_shader_buffer_ptr(batch, &st->ssbo[cb]);
+      uniforms.ssbo_size[cb] = st->ssbo[cb].buffer_size;
+   }
+
    if (stage == PIPE_SHADER_VERTEX) {
       u_foreach_bit(vbo, ctx->vb_mask) {
          uniforms.vs.vbo_base[vbo] = agx_vertex_buffer_ptr(batch, vbo);