asahi/gs: factor out output info

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

src/asahi/lib/agx_nir_lower_gs.c

@@ -1193,7 +1193,7 @@ calculate_max_indices(enum mesa_prim prim, unsigned verts, signed static_verts,
 bool
 agx_nir_lower_gs(nir_shader *gs, bool rasterizer_discard, nir_shader **gs_count,
                  nir_shader **gs_copy, nir_shader **pre_gs,
-                 enum mesa_prim *out_mode, unsigned *out_count_words)
+                 struct agx_gs_info *info)
 {
    /* Lower I/O as assumed by the rest of GS lowering */
    if (gs->xfb_info != NULL) {
@@ -1386,8 +1386,11 @@ agx_nir_lower_gs(nir_shader *gs, bool rasterizer_discard, nir_shader **gs_count,
       gs->info.gs.invocations, gs_state.max_indices);
 
    /* Signal what primitive we want to draw the GS Copy VS with */
-   *out_mode = gs->info.gs.output_primitive;
-   *out_count_words = gs_state.count_stride_el;
+   *info = (struct agx_gs_info){
+      .mode = gs->info.gs.output_primitive,
+      .count_words = gs_state.count_stride_el,
+   };
+
    return true;
 }
 

src/asahi/lib/agx_nir_lower_gs.h

@@ -30,10 +30,17 @@ bool agx_nir_lower_sw_vs(struct nir_shader *s, unsigned index_size_B);
 
 bool agx_nir_lower_vs_before_gs(struct nir_shader *vs);
 
+struct agx_gs_info {
+   /* Output primitive mode for geometry shaders */
+   enum mesa_prim mode;
+
+   /* Number of words per primitive in the count buffer */
+   unsigned count_words;
+};
+
 bool agx_nir_lower_gs(struct nir_shader *gs, 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);
+                      struct nir_shader **pre_gs, struct agx_gs_info *info);
 
 bool agx_nir_lower_tcs(struct nir_shader *tcs);
 

src/asahi/vulkan/hk_cmd_draw.c

@@ -1094,7 +1094,7 @@ hk_rast_prim(struct hk_cmd_buffer *cmd)
    struct vk_dynamic_graphics_state *dyn = &cmd->vk.dynamic_graphics_state;
 
    if (gs != NULL) {
-      return gs->variants[HK_GS_VARIANT_RAST].info.gs.out_prim;
+      return gs->variants[HK_GS_VARIANT_RAST].info.gs.mode;
    } else {
       switch (dyn->ia.primitive_topology) {
       case VK_PRIMITIVE_TOPOLOGY_LINE_LIST_WITH_ADJACENCY:

src/asahi/vulkan/hk_shader.c

@@ -1131,6 +1131,7 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info,
    /* Compile all variants up front */
    if (sw_stage == MESA_SHADER_GEOMETRY) {
       for (unsigned rast_disc = 0; rast_disc < 2; ++rast_disc) {
+         struct hk_shader *main_variant = hk_main_gs_variant(obj, rast_disc);
          struct hk_shader *count_variant = hk_count_gs_variant(obj, rast_disc);
          bool last = (rast_disc + 1) == 2;
 
@@ -1138,20 +1139,20 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info,
           * original NIR for the last stage.
           */
          nir_shader *clone = last ? nir : nir_shader_clone(NULL, nir);
-
-         enum mesa_prim out_prim = MESA_PRIM_MAX;
          nir_shader *count = NULL, *rast = NULL, *pre_gs = NULL;
 
          NIR_PASS(_, clone, agx_nir_lower_gs, rast_disc, &count, &rast, &pre_gs,
-                  &out_prim, &count_variant->info.gs.count_words);
+                  &count_variant->info.gs);
 
          if (!rast_disc) {
             struct hk_shader *shader = &obj->variants[HK_GS_VARIANT_RAST];
 
             hk_lower_hw_vs(rast, shader);
-            shader->info.gs.out_prim = out_prim;
+            shader->info.gs = count_variant->info.gs;
          }
 
+         main_variant->info.gs = count_variant->info.gs;
+
          struct {
             nir_shader *in;
             struct hk_shader *out;

src/asahi/vulkan/hk_shader.h

@@ -10,6 +10,7 @@
 #include "asahi/compiler/agx_compile.h"
 #include "util/macros.h"
 #include "agx_linker.h"
+#include "agx_nir_lower_gs.h"
 #include "agx_nir_lower_vbo.h"
 #include "agx_pack.h"
 #include "agx_usc.h"
@@ -96,11 +97,7 @@ struct hk_shader_info {
          struct hk_tess_info info;
       } tess;
 
-      struct {
-         unsigned count_words;
-         enum mesa_prim out_prim;
-         uint8_t _pad[27];
-      } gs;
+      struct agx_gs_info gs;
 
       /* Used to initialize the union for other stages */
       uint8_t _pad[32];

src/gallium/drivers/asahi/agx_disk_cache.c

@@ -74,8 +74,7 @@ write_shader(struct blob *blob, const struct agx_compiled_shader *binary,
                     sizeof(binary->push[0]) * binary->push_range_count);
 
    if (is_root_gs) {
-      blob_write_uint32(blob, binary->gs_count_words);
-      blob_write_uint32(blob, binary->gs_output_mode);
+      blob_write_bytes(blob, &binary->gs, sizeof(binary->gs));
       write_shader(blob, binary->pre_gs, false);
 
       blob_write_uint8(blob, binary->gs_copy != NULL);
@@ -126,8 +125,7 @@ read_shader(struct agx_screen *screen, struct blob_reader *blob,
                    sizeof(binary->push[0]) * binary->push_range_count);
 
    if (is_root && uncompiled->type == PIPE_SHADER_GEOMETRY) {
-      binary->gs_count_words = blob_read_uint32(blob);
-      binary->gs_output_mode = blob_read_uint32(blob);
+      blob_copy_bytes(blob, &binary->gs, sizeof(binary->gs));
       binary->pre_gs = read_shader(screen, blob, uncompiled, false);
 
       if (blob_read_uint8(blob))

src/gallium/drivers/asahi/agx_state.c

@@ -1587,10 +1587,9 @@ agx_compile_variant(struct agx_device *dev, struct pipe_context *pctx,
    nir_shader *nir = nir_deserialize(NULL, &agx_nir_options, &reader);
 
    /* Auxiliary programs */
-   enum mesa_prim gs_out_prim = MESA_PRIM_MAX;
+   struct agx_gs_info gs_info = {0};
    uint64_t outputs = 0;
    struct agx_fs_epilog_link_info epilog_key = {false};
-   unsigned gs_out_count_words = 0;
    nir_shader *gs_count = NULL;
    nir_shader *gs_copy = NULL;
    nir_shader *pre_gs = NULL;
@@ -1638,7 +1637,7 @@ agx_compile_variant(struct agx_device *dev, struct pipe_context *pctx,
       struct asahi_gs_shader_key *key = &key_->gs;
 
       NIR_PASS(_, nir, agx_nir_lower_gs, key->rasterizer_discard, &gs_count,
-               &gs_copy, &pre_gs, &gs_out_prim, &gs_out_count_words);
+               &gs_copy, &pre_gs, &gs_info);
    } else if (nir->info.stage == MESA_SHADER_FRAGMENT) {
       struct asahi_fs_shader_key *key = &key_->fs;
 
@@ -1733,8 +1732,7 @@ agx_compile_variant(struct agx_device *dev, struct pipe_context *pctx,
       compiled->gs_copy->uvs = uvs;
    }
 
-   compiled->gs_output_mode = gs_out_prim;
-   compiled->gs_count_words = gs_out_count_words;
+   compiled->gs = gs_info;
    compiled->b.info.outputs = outputs;
 
    ralloc_free(nir);
@@ -4039,7 +4037,7 @@ agx_batch_geometry_params(struct agx_batch *batch, uint64_t input_index_buffer,
     */
    batch->uniforms.vertex_outputs = batch->ctx->vs->b.info.outputs;
    params.input_mask = batch->uniforms.vertex_outputs;
-   params.count_buffer_stride = batch->ctx->gs->gs_count_words * 4;
+   params.count_buffer_stride = batch->ctx->gs->gs.count_words * 4;
 
    if (indirect) {
       batch->uniforms.vertex_output_buffer_ptr =
@@ -4171,7 +4169,7 @@ agx_launch_gs_prerast(struct agx_batch *batch,
       agx_launch(batch, grid_gs, wg, gs->gs_count, NULL, PIPE_SHADER_GEOMETRY,
                  0);
 
-      libagx_prefix_sum_geom(batch, agx_1d(1024 * gs->gs_count_words),
+      libagx_prefix_sum_geom(batch, agx_1d(1024 * gs->gs.count_words),
                              AGX_BARRIER_ALL, gp);
    }
 
@@ -5074,7 +5072,7 @@ agx_draw_vbo(struct pipe_context *pctx, const struct pipe_draw_info *info,
 
       /* Setup to rasterize the GS results */
       info_gs = (struct pipe_draw_info){
-         .mode = ctx->gs->gs_output_mode,
+         .mode = ctx->gs->gs.mode,
          .index_size = 4,
          .primitive_restart = true,
          .restart_index = ~0,

src/gallium/drivers/asahi/agx_state.h

@@ -33,6 +33,7 @@
 #include "util/u_range.h"
 #include "agx_bg_eot.h"
 #include "agx_helpers.h"
+#include "agx_nir_lower_gs.h"
 #include "agx_nir_texture.h"
 
 #ifdef __GLIBC__
@@ -248,11 +249,7 @@ struct agx_compiled_shader {
    struct agx_compiled_shader *gs_count, *pre_gs;
    struct agx_compiled_shader *gs_copy;
 
-   /* Output primitive mode for geometry shaders */
-   enum mesa_prim gs_output_mode;
-
-   /* Number of words per primitive in the count buffer */
-   unsigned gs_count_words;
+   struct agx_gs_info gs;
 
    /* Logical shader stage used for descriptor access. This may differ from the
     * physical shader stage of the compiled shader, for example when executing a