asahi: eliminate num_workgroups for VS->GS + VS->TCS

this is a step towards simpler geom/tess uniform management.

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

src/asahi/lib/agx_nir_lower_gs.c

@@ -1378,11 +1378,19 @@ lower_vs_before_gs(nir_builder *b, nir_intrinsic_instr *intr, void *data)
     */
    nir_def *mask = nir_imm_int64(b, b->shader->info.outputs_written);
 
-   nir_def *linear_id =
-      nir_iadd(b,
-               nir_imul(b, load_instance_id(b),
-                        nir_channel(b, nir_load_num_workgroups(b), 0)),
-               load_primitive_id(b));
+   nir_def *nr_verts;
+   if (b->shader->info.stage == MESA_SHADER_VERTEX) {
+      nr_verts =
+         libagx_input_vertices(b, nir_load_input_assembly_buffer_agx(b));
+   } else {
+      /* TODO: Do something similar for tessellation, load_num_workgroups is
+       * annoying in a software graphics shader.
+       */
+      nr_verts = nir_channel(b, nir_load_num_workgroups(b), 0);
+   }
+
+   nir_def *linear_id = nir_iadd(b, nir_imul(b, load_instance_id(b), nr_verts),
+                                 load_primitive_id(b));
 
    nir_def *addr = libagx_vertex_output_address(
       b, nir_load_vs_output_buffer_agx(b), mask, linear_id, location);

src/asahi/lib/shaders/geometry.cl

@@ -466,6 +466,8 @@ libagx_gs_setup_indirect(global struct agx_geometry_params *p,
    uint vertex_count = in_draw[0];
    uint instance_count = in_draw[1];
 
+   ia->verts_per_instance = vertex_count;
+
    /* Calculate number of primitives input into the GS */
    uint prim_per_instance = u_decomposed_prims_for_vertices(mode, vertex_count);
    p->input_primitives = prim_per_instance * instance_count;
@@ -586,3 +588,9 @@ libagx_vertex_output_address(uintptr_t buffer, uint64_t mask, uint vtx,
 {
    return buffer + libagx_tcs_in_offs(vtx, location, mask);
 }
+
+unsigned
+libagx_input_vertices(constant struct agx_ia_state *ia)
+{
+   return ia->verts_per_instance;
+}

src/asahi/lib/shaders/geometry.h

@@ -55,6 +55,11 @@ struct agx_ia_state {
    /* When unrolling primitive restart, output draw descriptors */
    GLOBAL(uint) out_draws;
 
+   /* Number of vertices per instance. Written by CPU for direct draw, indirect
+    * setup kernel for indirect. This is used for VS->GS and VS->TCS indexing.
+    */
+   uint32_t verts_per_instance;
+
    /* Input: maximum draw count, count is clamped to this */
    uint32_t max_draws;
 
@@ -76,7 +81,7 @@ struct agx_ia_state {
    /* The index size (1, 2, 4) or 0 if drawing without an index buffer. */
    uint32_t index_size_B;
 } PACKED;
-AGX_STATIC_ASSERT(sizeof(struct agx_ia_state) == 18 * 4);
+AGX_STATIC_ASSERT(sizeof(struct agx_ia_state) == 19 * 4);
 
 struct agx_geometry_params {
    /* Persistent (cross-draw) geometry state */

src/gallium/drivers/asahi/agx_state.c

@@ -3835,8 +3835,8 @@ agx_batch_geometry_state(struct agx_batch *batch)
 static void
 agx_upload_ia_params(struct agx_batch *batch, const struct pipe_draw_info *info,
                      const struct pipe_draw_indirect_info *indirect,
-                     uint64_t input_index_buffer, size_t index_buffer_size_B,
-                     uint64_t unroll_output)
+                     uint32_t count, uint64_t input_index_buffer,
+                     size_t index_buffer_size_B, uint64_t unroll_output)
 {
    struct agx_ia_state ia = {
       .heap = agx_batch_geometry_state(batch),
@@ -3853,6 +3853,8 @@ agx_upload_ia_params(struct agx_batch *batch, const struct pipe_draw_info *info,
       agx_batch_reads(batch, rsrc);
 
       ia.draws = rsrc->bo->ptr.gpu + indirect->offset;
+   } else {
+      ia.verts_per_instance = count;
    }
 
    batch->uniforms.input_assembly =
@@ -3866,8 +3868,8 @@ agx_batch_geometry_params(struct agx_batch *batch, uint64_t input_index_buffer,
                           const struct pipe_draw_start_count_bias *draw,
                           const struct pipe_draw_indirect_info *indirect)
 {
-   agx_upload_ia_params(batch, info, indirect, input_index_buffer,
-                        index_buffer_size_B, 0);
+   agx_upload_ia_params(batch, info, indirect, draw ? draw->count : 0,
+                        input_index_buffer, index_buffer_size_B, 0);
 
    struct agx_geometry_params params = {
       .state = agx_batch_geometry_state(batch),
@@ -4125,7 +4127,7 @@ agx_draw_without_restart(struct agx_batch *batch,
       &batch->pool, 5 * sizeof(uint32_t) * indirect->draw_count, 4,
       &out_draws_rsrc.bo);
 
-   agx_upload_ia_params(batch, info, indirect, ib, ib_extent, out_draws.gpu);
+   agx_upload_ia_params(batch, info, indirect, 0, ib, ib_extent, out_draws.gpu);
 
    /* Unroll the index buffer for each draw */
    const struct pipe_grid_info grid_setup = {
@@ -4451,7 +4453,8 @@ agx_draw_patches(struct agx_context *ctx, const struct pipe_draw_info *info,
    if (info->index_size)
       ib = agx_index_buffer_ptr(batch, info, draws, &ib_extent);
 
-   agx_upload_ia_params(batch, info, indirect, ib, ib_extent, 0);
+   agx_upload_ia_params(batch, info, indirect, draws ? draws->count : 0, ib,
+                        ib_extent, 0);
    agx_upload_draw_params(batch, indirect, draws, info);
 
    /* Setup parameters */