poly: Fetch the index size from a sysval

On asahi, we can still specialize based on the shader key and get
everything folded.  But this gives drivers the option to make it
dynamic if they wish.

Co-authored-by: Mary Guillemard <mary.guillemard@collabora.com>
Acked-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Reviewed-by: Mary Guillemard <mary@mary.zone>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/38404>

src/asahi/lib/agx_nir_prolog_epilog.c

@@ -162,7 +162,7 @@ lower_adjacency(nir_builder *b, nir_intrinsic_instr *intr, void *data)
       UNREACHABLE("unknown");
    }
 
-   id = poly_nir_load_vertex_id(b, id, key->sw_index_size_B);
+   id = poly_nir_load_vertex_id(b, id);
 
    nir_def_replace(&intr->def, id);
    return true;
@@ -216,11 +216,14 @@ agx_nir_vs_prolog(nir_builder *b, const void *key_)
    }
 
    if (!key->hw) {
-      poly_nir_lower_sw_vs(b->shader, key->sw_index_size_B);
+      b->cursor = nir_before_impl(nir_shader_get_entrypoint(b->shader));
+      poly_nir_lower_sw_vs(b->shader);
    } else if (key->adjacency) {
       nir_shader_intrinsics_pass(b->shader, lower_adjacency,
                                  nir_metadata_control_flow, (void *)key);
    }
+   nir_inline_sysval(b->shader, nir_intrinsic_load_index_size_poly,
+                     key->sw_index_size_B);
 
    /* Finally, lower uniforms according to our ABI */
    unsigned nr = DIV_ROUND_UP(BITSET_LAST_BIT(key->component_mask), 4);

src/compiler/nir/nir_intrinsics.py

@@ -1431,6 +1431,11 @@ system_value("ro_sink_address_poly", 1, bit_sizes=[64])
 # mesa_prim for the input topology (in a geometry shader)
 system_value("input_topology_poly", 1)
 
+# Size of an index in the index buffer in bytes, or zero for no indexing.
+# This is modeled as a sysval so it can be constant folded by drivers based
+# on a shader key if desired.
+system_value("index_size_poly", 1, bit_sizes=[32])
+
 # Mask of VS->TCS, VS->GS, or TES->GS outputs. This is modelled as a sysval
 # so it can be dynamic with shader objects or constant folded with monolithic.
 system_value("vs_outputs_poly", 1, bit_sizes=[64])

src/poly/cl/geometry.cl

@@ -378,6 +378,12 @@ poly_vertex_output_address(constant struct poly_vertex_params *p,
       ((uintptr_t)poly_tcs_in_offs_el(vtx, location, mask)) * 16;
 }
 
+unsigned
+poly_index_size(constant struct poly_vertex_params *p)
+{
+   return p->index_size_B;
+}
+
 unsigned
 poly_input_vertices(constant struct poly_vertex_params *p)
 {

src/poly/geometry.h

@@ -182,6 +182,9 @@ struct poly_vertex_params {
    /* Index buffer if present. */
    uint64_t index_buffer;
 
+   /* Size of an index in the index buffer, in bytes */
+   uint32_t index_size_B;
+
    /* Size of the bound index buffer for bounds checking */
    uint32_t index_buffer_range_el;
 
@@ -190,13 +193,15 @@ struct poly_vertex_params {
     */
    uint32_t verts_per_instance;
 
+   uint32_t _pad;
+
    /* Output buffer for vertex data */
    uint64_t output_buffer;
 
    /* Mask of outputs present in the output buffer */
    uint64_t outputs;
 } PACKED;
-static_assert(sizeof(struct poly_vertex_params) == 8 * 4);
+static_assert(sizeof(struct poly_vertex_params) == 10 * 4);
 
 static inline uint
 poly_index_buffer_range_el(uint size_el, uint offset_el)

src/poly/nir/poly_nir.h

@@ -15,10 +15,9 @@ struct nir_def *poly_load_per_vertex_input(struct nir_builder *b,
                                            nir_intrinsic_instr *intr,
                                            struct nir_def *vertex);
 
-nir_def *poly_nir_load_vertex_id(struct nir_builder *b, nir_def *id,
-                                 unsigned index_size_B);
+nir_def *poly_nir_load_vertex_id(struct nir_builder *b, nir_def *id);
 
-bool poly_nir_lower_sw_vs(struct nir_shader *s, unsigned index_size_B);
+bool poly_nir_lower_sw_vs(struct nir_shader *s);
 
 bool poly_nir_lower_vs_before_gs(struct nir_shader *vs);
 

src/poly/nir/poly_nir_lower_sysvals.c

@@ -11,6 +11,13 @@ static bool
 lower_sysvals_intr(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 {
    switch (intr->intrinsic) {
+   case nir_intrinsic_load_index_size_poly: {
+      b->cursor = nir_before_instr(&intr->instr);
+      nir_def *vp = nir_load_vertex_param_buffer_poly(b);
+      nir_def_replace(&intr->def, poly_index_size(b, vp));
+      return true;
+   }
+
    case nir_intrinsic_load_vs_outputs_poly: {
       b->cursor = nir_before_instr(&intr->instr);
       nir_def *vp = nir_load_vertex_param_buffer_poly(b);

src/poly/nir/poly_nir_lower_vs.c

@@ -15,32 +15,35 @@
  * the topology, which happens in the geometry shader.
  */
 nir_def *
-poly_nir_load_vertex_id(nir_builder *b, nir_def *id, unsigned index_size_B)
+poly_nir_load_vertex_id(nir_builder *b, nir_def *id)
 {
    /* If drawing with an index buffer, pull the vertex ID. Otherwise, the
     * vertex ID is just the index as-is.
     */
-   if (index_size_B) {
+   nir_def *index_size = nir_load_index_size_poly(b);
+   nir_def *index_buffer_id;
+   nir_if *index_size_present = nir_push_if(b, nir_ine_imm(b, index_size, 0));
+   {
       nir_def *p = nir_load_vertex_param_buffer_poly(b);
-      id = poly_load_index_buffer(b, p, id, nir_imm_int(b, index_size_B));
+      index_buffer_id = poly_load_index_buffer(b, p, id, index_size);
    }
+   nir_pop_if(b, index_size_present);
+   nir_def *effective_id = nir_if_phi(b, index_buffer_id, id);
 
    /* Add the "start", either an index bias or a base vertex. This must happen
     * after indexing for proper index bias behaviour.
     */
-   return nir_iadd(b, id, nir_load_first_vertex(b));
+   return nir_iadd(b, effective_id, nir_load_first_vertex(b));
 }
 
 static bool
 lower(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 {
-   unsigned *index_size_B = data;
    b->cursor = nir_before_instr(&intr->instr);
 
    if (intr->intrinsic == nir_intrinsic_load_vertex_id) {
       nir_def *id = nir_channel(b, nir_load_global_invocation_id(b, 32), 0);
-      nir_def_replace(&intr->def,
-                      poly_nir_load_vertex_id(b, id, *index_size_B));
+      nir_def_replace(&intr->def, poly_nir_load_vertex_id(b, id));
       return true;
    } else if (intr->intrinsic == nir_intrinsic_load_instance_id) {
       nir_def_replace(&intr->def,
@@ -52,8 +55,7 @@ lower(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 }
 
 bool
-poly_nir_lower_sw_vs(nir_shader *s, unsigned index_size_B)
+poly_nir_lower_sw_vs(nir_shader *s)
 {
-   return nir_shader_intrinsics_pass(s, lower, nir_metadata_control_flow,
-                                     &index_size_B);
+   return nir_shader_intrinsics_pass(s, lower, nir_metadata_control_flow, NULL);
 }