asahi: make gs topology dynamic

even with shobjs, we know the class of topology statically, so we just need to
select between the (up to) 3 compatible topologies, and luckily there are common
subexpressions we can factor out when calculating all 3 at once.

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

src/asahi/lib/agx_nir_lower_gs.c

@@ -168,7 +168,7 @@ load_instance_id(nir_builder *b)
 }
 
 static bool
-lower_gs_inputs(nir_builder *b, nir_intrinsic_instr *intr, void *key)
+lower_gs_inputs(nir_builder *b, nir_intrinsic_instr *intr, void *_)
 {
    if (intr->intrinsic != nir_intrinsic_load_per_vertex_input)
       return false;
@@ -178,9 +178,10 @@ lower_gs_inputs(nir_builder *b, nir_intrinsic_instr *intr, void *key)
 
    nir_def *location = nir_iadd_imm(b, intr->src[1].ssa, sem.location);
 
-   /* Calculate the vertex ID we're pulling, based on the topology */
+   /* Calculate the vertex ID we're pulling, based on the topology class */
    nir_def *vert_in_prim = intr->src[0].ssa;
-   nir_def *vertex = agx_vertex_id_for_topology(b, vert_in_prim, key);
+   nir_def *vertex = agx_vertex_id_for_topology_class(
+      b, vert_in_prim, b->shader->info.gs.input_primitive);
 
    /* The unrolled vertex ID uses the input_vertices, which differs from what
     * our load_num_vertices will return (vertices vs primitives).
@@ -318,6 +319,8 @@ lower_id(nir_builder *b, nir_intrinsic_instr *intr, void *data)
       id = nir_channel(b, nir_load_num_workgroups(b), 0);
    else if (intr->intrinsic == nir_intrinsic_load_flat_mask)
       id = load_geometry_param(b, flat_outputs);
+   else if (intr->intrinsic == nir_intrinsic_load_input_topology_agx)
+      id = load_geometry_param(b, input_topology);
    else if (intr->intrinsic == nir_intrinsic_load_provoking_last) {
       id = nir_b2b32(
          b, libagx_is_provoking_last(b, nir_load_input_assembly_buffer_agx(b)));
@@ -423,6 +426,7 @@ lower_to_gs_rast(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 
    case nir_intrinsic_load_flat_mask:
    case nir_intrinsic_load_provoking_last:
+   case nir_intrinsic_load_input_topology_agx:
       /* Lowering the same in both GS variants */
       return lower_id(b, intr, data);
 
@@ -1118,10 +1122,9 @@ link_libagx(nir_shader *nir, const nir_shader *libagx)
 
 bool
 agx_nir_lower_gs(nir_shader *gs, const nir_shader *libagx,
-                 struct agx_ia_key *ia, bool rasterizer_discard,
-                 nir_shader **gs_count, nir_shader **gs_copy,
-                 nir_shader **pre_gs, enum mesa_prim *out_mode,
-                 unsigned *out_count_words)
+                 bool rasterizer_discard, nir_shader **gs_count,
+                 nir_shader **gs_copy, nir_shader **pre_gs,
+                 enum mesa_prim *out_mode, unsigned *out_count_words)
 {
    /* Collect output component counts so we can size the geometry output buffer
     * appropriately, instead of assuming everything is vec4.
@@ -1145,7 +1148,7 @@ agx_nir_lower_gs(nir_shader *gs, const nir_shader *libagx,
    }
 
    NIR_PASS(_, gs, nir_shader_intrinsics_pass, lower_gs_inputs,
-            nir_metadata_block_index | nir_metadata_dominance, ia);
+            nir_metadata_block_index | nir_metadata_dominance, NULL);
 
    /* Lower geometry shader writes to contain all of the required counts, so we
     * know where in the various buffers we should write vertices.

src/asahi/lib/agx_nir_lower_gs.h

@@ -25,9 +25,9 @@ struct agx_lower_output_to_var_state {
 bool agx_lower_output_to_var(struct nir_builder *b, struct nir_instr *instr,
                              void *data);
 
-struct nir_def *agx_vertex_id_for_topology(struct nir_builder *b,
-                                           struct nir_def *vert,
-                                           struct agx_ia_key *key);
+struct nir_def *agx_vertex_id_for_topology_class(struct nir_builder *b,
+                                                 struct nir_def *vert,
+                                                 enum mesa_prim clas);
 
 bool agx_nir_lower_ia(struct nir_shader *s, struct agx_ia_key *ia);
 
@@ -36,10 +36,9 @@ bool agx_nir_lower_vs_before_gs(struct nir_shader *vs,
                                 unsigned index_size_B, uint64_t *outputs);
 
 bool agx_nir_lower_gs(struct nir_shader *gs, const struct nir_shader *libagx,
-                      struct agx_ia_key *ia, bool rasterizer_discard,
-                      struct nir_shader **gs_count, struct nir_shader **gs_copy,
-                      struct nir_shader **pre_gs, enum mesa_prim *out_mode,
-                      unsigned *out_count_words);
+                      bool rasterizer_discard, struct nir_shader **gs_count,
+                      struct nir_shader **gs_copy, struct nir_shader **pre_gs,
+                      enum mesa_prim *out_mode, unsigned *out_count_words);
 
 void agx_nir_prefix_sum_gs(struct nir_builder *b, const void *data);
 

src/asahi/lib/agx_nir_lower_ia.c

@@ -20,62 +20,48 @@
  * lives in CL library routines.
  */
 
-/*
- * Sync with geometry.cl, this is preferred to avoid NIR needing to chew through
- * the massive switch statement (bad for compile time).
- */
 nir_def *
-agx_vertex_id_for_topology(nir_builder *b, nir_def *vert,
-                           struct agx_ia_key *key)
+agx_vertex_id_for_topology_class(nir_builder *b, nir_def *vert,
+                                 enum mesa_prim cls)
 {
    nir_def *prim = nir_load_primitive_id(b);
    nir_def *flatshade_first = nir_ieq_imm(b, nir_load_provoking_last(b), 0);
+   nir_def *nr = nir_load_num_vertices(b);
+   nir_def *topology = nir_load_input_topology_agx(b);
 
-   switch (key->mode) {
+   switch (cls) {
    case MESA_PRIM_POINTS:
       return prim;
 
    case MESA_PRIM_LINES:
+      return libagx_vertex_id_for_line_class(b, topology, prim, vert, nr);
+
    case MESA_PRIM_TRIANGLES:
+      return libagx_vertex_id_for_tri_class(b, topology, prim, vert,
+                                            flatshade_first);
+
    case MESA_PRIM_LINES_ADJACENCY:
+      return libagx_vertex_id_for_line_adj_class(b, topology, prim, vert);
+
    case MESA_PRIM_TRIANGLES_ADJACENCY:
-      return nir_iadd(
-         b, nir_imul_imm(b, prim, mesa_vertices_per_prim(key->mode)), vert);
-
-   case MESA_PRIM_LINE_LOOP:
-      return libagx_vertex_id_for_line_loop(b, prim, vert,
-                                            nir_load_num_vertices(b));
-
-   case MESA_PRIM_LINE_STRIP:
-   case MESA_PRIM_LINE_STRIP_ADJACENCY:
-      return nir_iadd(b, prim, vert);
-
-   case MESA_PRIM_TRIANGLE_STRIP: {
-      return nir_iadd(
-         b, prim,
-         libagx_map_vertex_in_tri_strip(b, prim, vert, flatshade_first));
-   }
-
-   case MESA_PRIM_TRIANGLE_FAN:
-      return libagx_vertex_id_for_tri_fan(b, prim, vert, flatshade_first);
-
-   case MESA_PRIM_TRIANGLE_STRIP_ADJACENCY:
-      return libagx_vertex_id_for_tri_strip_adj(
-         b, prim, vert, nir_load_num_vertices(b), flatshade_first);
-
-   case MESA_PRIM_PATCHES:
-      return nir_iadd(b, nir_imul(b, prim, nir_load_patch_vertices_in(b)),
-                      nir_load_invocation_id(b));
+      return libagx_vertex_id_for_tri_adj_class(b, topology, prim, vert, nr,
+                                                flatshade_first);
 
    default:
-      unreachable("invalid mode");
+      unreachable("invalid topology class");
    }
 }
 
 static nir_def *
 load_vertex_id(nir_builder *b, struct agx_ia_key *key)
 {
-   nir_def *id = agx_vertex_id_for_topology(b, NULL, key);
+   assert(key->mode == MESA_PRIM_POINTS || key->mode == MESA_PRIM_PATCHES);
+   nir_def *id = nir_load_primitive_id(b);
+
+   if (key->mode == MESA_PRIM_PATCHES) {
+      id = nir_iadd(b, nir_imul(b, id, nir_load_patch_vertices_in(b)),
+                    nir_load_invocation_id(b));
+   }
 
    /* If drawing with an index buffer, pull the vertex ID. Otherwise, the
     * vertex ID is just the index as-is.

src/asahi/lib/shaders/geometry.cl

@@ -60,6 +60,20 @@ libagx_vertex_id_for_line_loop(uint prim, uint vert, uint num_prims)
       return prim + vert;
 }
 
+uint
+libagx_vertex_id_for_line_class(enum mesa_prim mode, uint prim, uint vert,
+                                uint num_prims)
+{
+   /* Line list, line strip, or line loop */
+   if (mode == MESA_PRIM_LINE_LOOP && prim == (num_prims - 1) && vert == 1)
+      return 0;
+
+   if (mode == MESA_PRIM_LINES)
+      prim *= 2;
+
+   return prim + vert;
+}
+
 uint
 libagx_vertex_id_for_tri_fan(uint prim, uint vert, bool flatshade_first)
 {
@@ -79,6 +93,44 @@ libagx_vertex_id_for_tri_fan(uint prim, uint vert, bool flatshade_first)
    return (vert == 0) ? 0 : prim + vert;
 }
 
+uint
+libagx_vertex_id_for_tri_class(enum mesa_prim mode, uint prim, uint vert,
+                               bool flatshade_first)
+{
+   if (flatshade_first && mode == MESA_PRIM_TRIANGLE_FAN) {
+      vert = vert + 1;
+      vert = (vert == 3) ? 0 : vert;
+   }
+
+   if (mode == MESA_PRIM_TRIANGLE_FAN && vert == 0)
+      return 0;
+
+   if (mode == MESA_PRIM_TRIANGLES)
+      prim *= 3;
+
+   /* Triangle list, triangle strip, or triangle fan */
+   if (mode == MESA_PRIM_TRIANGLE_STRIP) {
+      unsigned pv = flatshade_first ? 0 : 2;
+
+      bool even = (prim & 1) == 0;
+      bool provoking = vert == pv;
+
+      vert = ((provoking || even) ? vert : ((3 - pv) - vert));
+   }
+
+   return prim + vert;
+}
+
+uint
+libagx_vertex_id_for_line_adj_class(enum mesa_prim mode, uint prim, uint vert)
+{
+   /* Line list adj or line strip adj */
+   if (mode == MESA_PRIM_LINES_ADJACENCY)
+      prim *= 4;
+
+   return prim + vert;
+}
+
 uint
 libagx_vertex_id_for_tri_strip_adj(uint prim, uint vert, uint num_prims,
                                    bool flatshade_first)
@@ -125,7 +177,19 @@ libagx_vertex_id_for_tri_strip_adj(uint prim, uint vert, uint num_prims,
    return (prim * 2) + offset;
 }
 
-/* Sync with agx_nir_lower_ia.c, this is for the restart unrolling */
+uint
+libagx_vertex_id_for_tri_adj_class(enum mesa_prim mode, uint prim, uint vert,
+                                   uint nr, bool flatshade_first)
+{
+   /* Tri adj list or tri adj strip */
+   if (mode == MESA_PRIM_TRIANGLE_STRIP_ADJACENCY) {
+      return libagx_vertex_id_for_tri_strip_adj(prim, vert, nr,
+                                                flatshade_first);
+   } else {
+      return (6 * prim) + vert;
+   }
+}
+
 uint
 libagx_vertex_id_for_topology(enum mesa_prim mode, bool flatshade_first,
                               uint prim, uint vert, uint num_prims)

src/asahi/lib/shaders/geometry.h

@@ -162,8 +162,13 @@ struct agx_geometry_params {
     * allocating counts.
     */
    uint32_t count_buffer_stride;
+
+   /* Dynamic input topology. Must be compatible with the geometry shader's
+    * layout() declared input class.
+    */
+   uint32_t input_topology;
 } PACKED;
-AGX_STATIC_ASSERT(sizeof(struct agx_geometry_params) == 82 * 4);
+AGX_STATIC_ASSERT(sizeof(struct agx_geometry_params) == 83 * 4);
 
 struct agx_tess_params {
    /* Persistent (cross-draw) geometry state */

src/compiler/nir/nir_intrinsics.py

@@ -1741,6 +1741,9 @@ intrinsic("load_fb_layers_v3d", dest_comp=1, flags=[CAN_ELIMINATE, CAN_REORDER])
 # Equivalent to popcount(ballot(true) & ((1 << subgroup_invocation) - 1))
 system_value("active_subgroup_invocation_agx", 1)
 
+# mesa_prim for the input topology (in a geometry shader)
+system_value("input_topology_agx", 1)
+
 # Load a bindless sampler handle mapping a binding table sampler.
 intrinsic("load_sampler_handle_agx", [1], 1, [],
           flags=[CAN_ELIMINATE, CAN_REORDER],

src/gallium/drivers/asahi/agx_state.c

@@ -1942,13 +1942,8 @@ agx_compile_variant(struct agx_device *dev, struct pipe_context *pctx,
 
       NIR_PASS(_, nir, nir_lower_io_to_scalar, nir_var_shader_out, NULL, NULL);
 
-      struct agx_ia_key ia = {
-         .mode = key->mode,
-      };
-
-      NIR_PASS(_, nir, agx_nir_lower_gs, dev->libagx, &ia,
-               key->rasterizer_discard, &gs_count, &gs_copy, &pre_gs,
-               &gs_out_prim, &gs_out_count_words);
+      NIR_PASS(_, nir, agx_nir_lower_gs, dev->libagx, key->rasterizer_discard,
+               &gs_count, &gs_copy, &pre_gs, &gs_out_prim, &gs_out_count_words);
    } else if (nir->info.stage == MESA_SHADER_FRAGMENT) {
       struct asahi_fs_shader_key *key = &key_->fs;
 
@@ -2543,7 +2538,6 @@ agx_update_gs(struct agx_context *ctx, const struct pipe_draw_info *info,
       rast_prim(gs->gs_mode, ctx->rast->base.fill_front);
 
    struct asahi_gs_shader_key key = {
-      .mode = info->mode,
       .rasterizer_discard = ctx->rast->base.rasterizer_discard,
 
       /* TODO: Deduplicate */
@@ -4108,6 +4102,7 @@ agx_batch_geometry_params(struct agx_batch *batch, uint64_t input_index_buffer,
       .indirect_desc = batch->geom_indirect,
       .flat_outputs =
          batch->ctx->stage[PIPE_SHADER_FRAGMENT].shader->info.inputs_flat_shaded,
+      .input_topology = info->mode,
    };
 
    for (unsigned i = 0; i < ARRAY_SIZE(batch->ctx->streamout.targets); ++i) {

src/gallium/drivers/asahi/agx_state.h

@@ -507,17 +507,14 @@ struct asahi_gs_shader_key {
    /* Rasterizer shader key */
    uint64_t outputs_flat_shaded;
    uint64_t outputs_linear_shaded;
-
-   /* Input assembly key */
-   enum mesa_prim mode;
-
-   /* Rasterizer shader key */
    bool clip_halfz;
    bool fixed_point_size;
 
    /* If true, this GS is run only for its side effects (including XFB) */
    bool rasterizer_discard;
+   bool padding[5];
 };
+static_assert(sizeof(struct asahi_gs_shader_key) == 24, "no holes");
 
 union asahi_shader_key {
    struct asahi_vs_shader_key vs;