libagx: add kernel for incrementing CS counter

for indirect dispatch

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30051>

src/asahi/lib/agx_nir_lower_gs.c

@@ -1605,3 +1605,9 @@ agx_nir_tess_setup_indirect(nir_builder *b, const void *data)
 
    libagx_tess_setup_indirect(b, params, with_counts, point_mode);
 }
+
+void
+agx_nir_increment_cs_invocations(nir_builder *b, const void *data)
+{
+   libagx_increment_cs_invocations(b, nir_load_preamble(b, 1, 64, .base = 0));
+}

src/asahi/lib/agx_nir_lower_gs.h

@@ -81,3 +81,5 @@ uint64_t agx_tcs_per_vertex_outputs(const struct nir_shader *nir);
 unsigned agx_tcs_output_stride(const struct nir_shader *nir);
 
 void agx_nir_tess_setup_indirect(struct nir_builder *b, const void *data);
+
+void agx_nir_increment_cs_invocations(struct nir_builder *b, const void *data);

src/asahi/lib/shaders/query.cl

@@ -50,3 +50,10 @@ libagx_copy_xfb_counters(constant struct libagx_xfb_counter_copy *push)
 
    *(push->dest[i]) = push->src[i] ? *(push->src[i]) : 0;
 }
+
+void
+libagx_increment_cs_invocations(constant struct libagx_cs_invocation_params *p)
+{
+   *(p->statistic) += libagx_cs_invocations(p->local_size_threads, p->grid[0],
+                                            p->grid[1], p->grid[2]);
+}

src/asahi/lib/shaders/query.h

@@ -27,3 +27,21 @@ struct libagx_xfb_counter_copy {
    GLOBAL(uint32_t) dest[4];
    GLOBAL(uint32_t) src[4];
 };
+
+struct libagx_cs_invocation_params {
+   /* Pointer to the indirect dispatch grid */
+   GLOBAL(uint32_t) grid;
+
+   /* Pointer to the compute shader invocation statistic */
+   GLOBAL(uint32_t) statistic;
+
+   /* Local workgroup size in threads */
+   uint32_t local_size_threads;
+};
+
+static inline uint32_t
+libagx_cs_invocations(uint32_t local_size_threads, uint32_t x, uint32_t y,
+                      uint32_t z)
+{
+   return local_size_threads * x * y * z;
+}