freedreno/computerator: add support for UBOs

Exposed via a new @ubo header that works just like @buf. UBOs can be
accessed through an index starting at 0 for the first one.

Signed-off-by: Job Noorman <jnoorman@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36073>

src/freedreno/computerator/a6xx.cc

@@ -356,15 +356,39 @@ cs_const_emit(struct fd_ringbuffer *ring, struct kernel *kernel,
    }
 }
 
+static unsigned
+kernel_num_bufs(struct kernel *kernel, enum kernel_buf_type buf_type)
+{
+   unsigned num_bufs = 0;
+
+   for (unsigned i = 0; i < kernel->num_bufs; i++) {
+      if (kernel->buf_types[i] == buf_type) {
+         num_bufs++;
+      }
+   }
+
+   return num_bufs;
+}
+
 template<chip CHIP>
 static void
 cs_ibo_emit(struct fd_ringbuffer *ring, struct fd_submit *submit,
             struct kernel *kernel)
 {
+   unsigned num_bufs = kernel_num_bufs(kernel, KERNEL_BUF_UAV);
+
+   if (num_bufs == 0) {
+      return;
+   }
+
    struct fd_ringbuffer *state = fd_submit_new_ringbuffer(
       submit, kernel->num_bufs * 16 * 4, FD_RINGBUFFER_STREAMING);
 
    for (unsigned i = 0; i < kernel->num_bufs; i++) {
+      if (kernel->buf_types[i] != KERNEL_BUF_UAV) {
+         continue;
+      }
+
       /* size is encoded with low 15b in WIDTH and high bits in HEIGHT,
        * in units of elements:
        */
@@ -396,7 +420,7 @@ cs_ibo_emit(struct fd_ringbuffer *ring, struct fd_submit *submit,
                      CP_LOAD_STATE6_0_STATE_TYPE(ST6_UAV) |
                      CP_LOAD_STATE6_0_STATE_SRC(SS6_INDIRECT) |
                      CP_LOAD_STATE6_0_STATE_BLOCK(SB6_CS_SHADER) |
-                     CP_LOAD_STATE6_0_NUM_UNIT(kernel->num_bufs));
+                     CP_LOAD_STATE6_0_NUM_UNIT(num_bufs));
    OUT_RB(ring, state);
 
    if (CHIP == A6XX) {
@@ -407,11 +431,41 @@ cs_ibo_emit(struct fd_ringbuffer *ring, struct fd_submit *submit,
    OUT_RB(ring, state);
 
    OUT_PKT4(ring, REG_A6XX_SP_CS_USIZE, 1);
-   OUT_RING(ring, kernel->num_bufs);
+   OUT_RING(ring, num_bufs);
 
    fd_ringbuffer_del(state);
 }
 
+static void
+cs_ubo_emit(struct fd_ringbuffer *ring, struct kernel *kernel)
+{
+   unsigned num_bufs = kernel_num_bufs(kernel, KERNEL_BUF_UBO);
+
+   if (num_bufs == 0) {
+      return;
+   }
+
+   for (unsigned i = 0, offset = 0; i < kernel->num_bufs; i++) {
+      if (kernel->buf_types[i] != KERNEL_BUF_UBO) {
+         continue;
+      }
+
+      OUT_PKT7(ring, CP_LOAD_STATE6_FRAG, 5);
+      OUT_RING(ring, CP_LOAD_STATE6_0_DST_OFF(offset) |
+                        CP_LOAD_STATE6_0_STATE_TYPE(ST6_UBO) |
+                        CP_LOAD_STATE6_0_STATE_SRC(SS6_DIRECT) |
+                        CP_LOAD_STATE6_0_STATE_BLOCK(SB6_CS_SHADER) |
+                        CP_LOAD_STATE6_0_NUM_UNIT(1));
+      OUT_RING(ring, CP_LOAD_STATE6_1_EXT_SRC_ADDR(0));
+      OUT_RING(ring, CP_LOAD_STATE6_2_EXT_SRC_ADDR_HI(0));
+      unsigned size_vec4s = DIV_ROUND_UP(kernel->buf_sizes[i], 4);
+      OUT_RELOC(ring, kernel->bufs[i], 0,
+                (uint64_t)A6XX_UBO_1_SIZE(size_vec4s) << 32, 0);
+
+      offset++;
+   }
+}
+
 template<chip CHIP>
 static inline unsigned
 event_write(struct fd_ringbuffer *ring, struct kernel *kernel,
@@ -485,6 +539,7 @@ a6xx_emit_grid(struct kernel *kernel, uint32_t grid[3],
    cs_program_emit<CHIP>(ring, kernel);
    cs_const_emit<CHIP>(ring, kernel, grid);
    cs_ibo_emit<CHIP>(ring, submit, kernel);
+   cs_ubo_emit(ring, kernel);
 
    OUT_PKT7(ring, CP_SET_MARKER, 1);
    OUT_RING(ring, A6XX_CP_SET_MARKER_0_MODE(RM6_COMPUTE));

src/freedreno/computerator/examples/ubo.asm

@@ -0,0 +1,29 @@
+@localsize 1, 1, 1
+@ubo 4 1, 2, 3, 4 ; UBO 0
+@ubo 4 5, 6, 7, 8 ; UBO 1
+@buf 4
+@invocationid(r0.x)
+@branchstack 1
+@earlypreamble
+
+shps #main
+getone #main
+
+mova1 a1.x, (r)0
+(ss)nop
+
+; copy UBO 0 to c0
+ldc.4.k.imm c[a1.x], 0, 0
+
+(sy)(ss)shpe
+
+main:
+
+; load UBO 1 to r1
+(jp)ldc.offset0.4.imm r1.x, 0, 1
+
+(sy)(rpt3)add.u r1.x, (r)r1.x, (r)c0.x
+(rpt2)nop
+
+stib.b.untyped.1d.u32.4.imm r1.x, r0.x, 0
+end

src/freedreno/computerator/ir3_asm.cc

@@ -42,6 +42,8 @@ ir3_asm_assemble(struct ir3_compiler *c, FILE *in)
    kernel->base.local_size[1] = v->local_size[1];
    kernel->base.local_size[2] = v->local_size[2];
    kernel->base.num_bufs = kernel->info.num_bufs;
+   memcpy(kernel->base.buf_types, kernel->info.buf_types,
+          sizeof(kernel->base.buf_sizes));
    memcpy(kernel->base.buf_sizes, kernel->info.buf_sizes,
           sizeof(kernel->base.buf_sizes));
    memcpy(kernel->base.buf_addr_regs, kernel->info.buf_addr_regs,

src/freedreno/computerator/main.cc

@@ -267,7 +267,8 @@ main(int argc, char **argv)
    printf("localsize: %dx%dx%d\n", kernel->local_size[0], kernel->local_size[1],
           kernel->local_size[2]);
    for (int i = 0; i < kernel->num_bufs; i++) {
-      printf("buf[%d]: size=%u\n", i, kernel->buf_sizes[i]);
+      printf("buf[%d]: size=%u, type=%s\n", i, kernel->buf_sizes[i],
+             kernel->buf_types[i] == KERNEL_BUF_UBO ? "UBO" : "SSBO");
       kernel->bufs[i] = fd_bo_new(dev, kernel->buf_sizes[i] * 4, 0, "buf[%d]", i);
 
       if (kernel->buf_init_data[i]) {

src/freedreno/computerator/main.h

@@ -16,12 +16,15 @@
 #include "adreno_common.xml.h"
 #include "adreno_pm4.xml.h"
 
+#include "ir3/ir3_assembler.h"
+
 #define MAX_BUFS 4
 
 struct kernel {
    /* filled in by backend when shader is assembled: */
    uint32_t local_size[3];
    uint32_t num_bufs;
+   enum kernel_buf_type buf_types[MAX_BUFS];
    uint32_t buf_sizes[MAX_BUFS]; /* size in dwords */
    uint32_t buf_addr_regs[MAX_BUFS];
    uint32_t *buf_init_data[MAX_BUFS];

src/freedreno/ir3/ir3_assembler.h

@@ -15,8 +15,14 @@
 extern "C" {
 #endif
 
+enum kernel_buf_type {
+   KERNEL_BUF_UAV,
+   KERNEL_BUF_UBO,
+};
+
 struct ir3_kernel_info {
    uint32_t num_bufs;
+   enum kernel_buf_type buf_types[MAX_BUFS];
    uint32_t buf_sizes[MAX_BUFS]; /* size in dwords */
    uint32_t buf_addr_regs[MAX_BUFS];
    uint32_t *buf_init_data[MAX_BUFS];

src/freedreno/ir3/ir3_lexer.l

@@ -138,6 +138,7 @@ static int parse_reg(const char *str)
 "@localsize"                      return TOKEN(T_A_LOCALSIZE);
 "@const"                          return TOKEN(T_A_CONST);
 "@buf"                            return TOKEN(T_A_BUF);
+"@ubo"                            return TOKEN(T_A_UBO);
 "@invocationid"                   return TOKEN(T_A_INVOCATIONID);
 "@wgid"                           return TOKEN(T_A_WGID);
 "@numwg"                          return TOKEN(T_A_NUMWG);

src/freedreno/ir3/ir3_parser.y

@@ -403,6 +403,7 @@ static void print_token(FILE *file, int type, YYSTYPE value)
 %token <tok> T_A_LOCALSIZE
 %token <tok> T_A_CONST
 %token <tok> T_A_BUF
+%token <tok> T_A_UBO
 %token <tok> T_A_INVOCATIONID
 %token <tok> T_A_WGID
 %token <tok> T_A_NUMWG
@@ -796,6 +797,7 @@ static void print_token(FILE *file, int type, YYSTYPE value)
 %type <tok> cat5_opc cat5_samp cat5_tex cat5_type
 %type <type> type
 %type <unum> const_val cat6_src_shift
+%type <num> buf_type
 
 %error-verbose
 
@@ -853,9 +855,13 @@ buf_header_addr_reg:
 }
 |
 
-buf_header:        T_A_BUF const_val {
+buf_type: T_A_BUF { $$ = KERNEL_BUF_UAV; }
+|         T_A_UBO { $$ = KERNEL_BUF_UBO; }
+
+buf_header:        buf_type const_val {
                        int idx = info->num_bufs++;
                        assert(idx < MAX_BUFS);
+                       info->buf_types[idx] = $1;
                        info->buf_sizes[idx] = $2;
 } buf_header_addr_reg buf_header_init_vals