hk: support static vertex input state

prologs inflate register pressure, so this can help a lot in the monolithic case
(together with dynamic strides). eliminates spilling from some vertex shaders in
Control that read a ton of attributes per vertex.

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

src/asahi/lib/agx_linker.h

@@ -59,6 +59,11 @@ struct agx_vs_prolog_key {
    /* Whether running as a hardware vertex shader (versus compute) */
    bool hw;
 
+   /* Whether the shader uses static vertex input and the prolog just fixes up
+    * the vertex ID.
+    */
+   bool static_vi;
+
    /* If !hw and the draw call is indexed, the index size */
    uint8_t sw_index_size_B;
 
@@ -182,10 +187,14 @@ void agx_nir_vs_prolog(struct nir_builder *b, const void *key_);
 void agx_nir_fs_epilog(struct nir_builder *b, const void *key_);
 void agx_nir_fs_prolog(struct nir_builder *b, const void *key_);
 
-bool agx_nir_lower_vs_input_to_prolog(nir_shader *s,
-                                      BITSET_WORD *attrib_components_read);
+bool agx_nir_gather_vs_inputs(nir_shader *s,
+                              BITSET_WORD *attrib_components_read);
+
+bool agx_nir_lower_vs_input_to_prolog(nir_shader *s);
 
 bool agx_nir_lower_fs_output_to_epilog(nir_shader *s,
                                        struct agx_fs_epilog_link_info *out);
 
 bool agx_nir_lower_fs_active_samples_to_register(nir_shader *s);
+
+bool agx_nir_lower_non_monolithic_uniforms(nir_shader *nir, unsigned nr);

src/asahi/lib/agx_nir_lower_vbo.c

@@ -9,11 +9,13 @@
 #include "compiler/nir/nir_format_convert.h"
 #include "util/bitset.h"
 #include "util/u_math.h"
+#include "agx_linker.h"
 #include "shader_enums.h"
 
 struct ctx {
-   struct agx_attribute *attribs;
+   const struct agx_velem_key *attribs;
    struct agx_robustness rs;
+   bool dynamic_strides;
 };
 
 static bool
@@ -25,7 +27,7 @@ is_rgb10_a2(const struct util_format_description *desc)
           desc->channel[3].shift == 30 && desc->channel[3].size == 2;
 }
 
-static enum pipe_format
+enum pipe_format
 agx_vbo_internal_format(enum pipe_format format)
 {
    const struct util_format_description *desc = util_format_description(format);
@@ -115,16 +117,15 @@ pass(struct nir_builder *b, nir_intrinsic_instr *intr, void *data)
       return false;
 
    struct ctx *ctx = data;
-   struct agx_attribute *attribs = ctx->attribs;
+   const struct agx_velem_key *attribs = ctx->attribs;
    b->cursor = nir_instr_remove(&intr->instr);
 
    nir_src *offset_src = nir_get_io_offset_src(intr);
    assert(nir_src_is_const(*offset_src) && "no attribute indirects");
    unsigned index = nir_intrinsic_base(intr) + nir_src_as_uint(*offset_src);
 
-   struct agx_attribute attrib = attribs[index];
+   struct agx_velem_key attrib = attribs[index];
    uint32_t stride = attrib.stride;
-   uint16_t offset = attrib.src_offset;
 
    const struct util_format_description *desc =
       util_format_description(attrib.format);
@@ -202,25 +203,28 @@ pass(struct nir_builder *b, nir_intrinsic_instr *intr, void *data)
    nir_def *base = nir_load_vbo_base_agx(b, buf_handle);
 
    assert((stride % interchange_align) == 0 && "must be aligned");
-   assert((offset % interchange_align) == 0 && "must be aligned");
 
    unsigned stride_el = stride / interchange_align;
-   unsigned offset_el = offset / interchange_align;
    unsigned shift = 0;
 
    /* Try to use the small shift on the load itself when possible. This can save
     * an instruction. Shifts are only available for regular interchange formats,
     * i.e. the set of formats that support masking.
     */
-   if (offset_el == 0 && (stride_el == 2 || stride_el == 4) &&
+   if ((stride_el == 2 || stride_el == 4) &&
        ail_isa_format_supports_mask((enum ail_isa_format)interchange_format)) {
 
       shift = util_logbase2(stride_el);
       stride_el = 1;
    }
 
-   nir_def *stride_offset_el =
-      nir_iadd_imm(b, nir_imul_imm(b, el, stride_el), offset_el);
+   nir_def *stride_el_def = nir_imm_int(b, stride_el);
+   if (ctx->dynamic_strides) {
+      assert(stride_el == 0);
+      stride_el_def = nir_load_vbo_stride_agx(b, buf_handle);
+   }
+
+   nir_def *stride_offset_el = nir_imul(b, el, stride_el_def);
 
    /* Fixing up the address is expected to be profitable for vec3 and above, as
     * it requires 2 instructions. It is implemented with a 64GiB carveout at the
@@ -305,8 +309,8 @@ pass(struct nir_builder *b, nir_intrinsic_instr *intr, void *data)
 }
 
 bool
-agx_nir_lower_vbo(nir_shader *shader, struct agx_attribute *attribs,
-                  struct agx_robustness robustness)
+agx_nir_lower_vbo(nir_shader *shader, const struct agx_velem_key *attribs,
+                  struct agx_robustness robustness, bool dynamic_strides)
 {
    assert(shader->info.stage == MESA_SHADER_VERTEX);
 
@@ -317,7 +321,12 @@ agx_nir_lower_vbo(nir_shader *shader, struct agx_attribute *attribs,
       robustness.level = MAX2(robustness.level, AGX_ROBUSTNESS_GL);
    }
 
-   struct ctx ctx = {.attribs = attribs, .rs = robustness};
+   struct ctx ctx = {
+      .attribs = attribs,
+      .rs = robustness,
+      .dynamic_strides = dynamic_strides,
+   };
+
    return nir_shader_intrinsics_pass(shader, pass, nir_metadata_control_flow,
                                      &ctx);
 }

src/asahi/lib/agx_nir_lower_vbo.h

@@ -17,22 +17,6 @@ extern "C" {
 #define AGX_MAX_ATTRIBS (16)
 #define AGX_MAX_VBUFS   (16)
 
-/* See pipe_vertex_element for justification on the sizes. This structure should
- * be small so it can be embedded into a shader key.
- */
-struct agx_attribute {
-   /* If instanced, Zero means all get the same value (Vulkan semantics). */
-   uint32_t divisor;
-   uint32_t stride;
-   uint16_t src_offset;
-
-   /* pipe_format, all vertex formats should be <= 255 */
-   uint8_t format;
-
-   unsigned buf   : 7;
-   bool instanced : 1;
-};
-
 enum agx_robustness_level {
    /* No robustness */
    AGX_ROBUSTNESS_DISABLED,
@@ -54,8 +38,12 @@ struct agx_robustness {
    bool soft_fault;
 };
 
-bool agx_nir_lower_vbo(nir_shader *shader, struct agx_attribute *attribs,
-                       struct agx_robustness rs);
+struct agx_velem_key;
+
+bool agx_nir_lower_vbo(nir_shader *shader, const struct agx_velem_key *attribs,
+                       struct agx_robustness rs, bool dynamic_strides);
+
+enum pipe_format agx_vbo_internal_format(enum pipe_format format);
 
 bool agx_vbo_supports_format(enum pipe_format format);
 

src/asahi/lib/agx_nir_prolog_epilog.c

@@ -56,24 +56,6 @@ agx_nir_lower_poly_stipple(nir_shader *s)
    return nir_progress(true, b->impl, nir_metadata_control_flow);
 }
 
-static bool
-lower_vbo(nir_shader *s, const struct agx_velem_key *key,
-          const struct agx_robustness rs)
-{
-   struct agx_attribute out[AGX_MAX_VBUFS];
-
-   for (unsigned i = 0; i < AGX_MAX_VBUFS; ++i) {
-      out[i] = (struct agx_attribute){
-         .divisor = key[i].divisor,
-         .stride = key[i].stride,
-         .format = key[i].format,
-         .instanced = key[i].instanced,
-      };
-   }
-
-   return agx_nir_lower_vbo(s, out, rs);
-}
-
 static int
 map_vs_part_uniform(nir_intrinsic_instr *intr, unsigned nr_attribs)
 {
@@ -148,6 +130,13 @@ lower_non_monolithic_uniforms(nir_builder *b, nir_intrinsic_instr *intr,
    }
 }
 
+bool
+agx_nir_lower_non_monolithic_uniforms(nir_shader *nir, unsigned nr)
+{
+   return nir_shader_intrinsics_pass(nir, lower_non_monolithic_uniforms,
+                                     nir_metadata_control_flow, &nr);
+}
+
 static bool
 lower_adjacency(nir_builder *b, nir_intrinsic_instr *intr, void *data)
 {
@@ -188,19 +177,22 @@ agx_nir_vs_prolog(nir_builder *b, const void *key_)
    /* First, construct a passthrough shader reading each attribute and exporting
     * the value. We also need to export vertex/instance ID in their usual regs.
     */
-   unsigned i = 0;
-   nir_def *vec = NULL;
-   unsigned vec_idx = ~0;
-   BITSET_FOREACH_SET(i, key->component_mask, AGX_MAX_ATTRIBS * 4) {
-      unsigned a = i / 4;
-      unsigned c = i % 4;
+   if (!key->static_vi) {
+      unsigned i = 0;
+      nir_def *vec = NULL;
+      unsigned vec_idx = ~0;
+      BITSET_FOREACH_SET(i, key->component_mask, AGX_MAX_ATTRIBS * 4) {
+         unsigned a = i / 4;
+         unsigned c = i % 4;
 
-      if (vec_idx != a) {
-         vec = nir_load_input(b, 4, 32, nir_imm_int(b, 0), .base = a);
-         vec_idx = a;
+         if (vec_idx != a) {
+            vec = nir_load_input(b, 4, 32, nir_imm_int(b, 0), .base = a);
+            vec_idx = a;
+         }
+
+         nir_export_agx(b, nir_channel(b, vec, c),
+                        .base = AGX_ABI_VIN_ATTRIB(i));
       }
-
-      nir_export_agx(b, nir_channel(b, vec, c), .base = AGX_ABI_VIN_ATTRIB(i));
    }
 
    if (!key->hw) {
@@ -212,12 +204,14 @@ agx_nir_vs_prolog(nir_builder *b, const void *key_)
    nir_export_agx(b, nir_load_instance_id(b), .base = AGX_ABI_VIN_INSTANCE_ID);
 
    /* Now lower the resulting program using the key */
-   lower_vbo(b->shader, key->attribs, key->robustness);
+   if (!key->static_vi) {
+      agx_nir_lower_vbo(b->shader, key->attribs, key->robustness, false);
 
-   /* Clean up redundant vertex ID loads */
-   if (!key->hw || key->adjacency) {
-      NIR_PASS(_, b->shader, nir_opt_cse);
-      NIR_PASS(_, b->shader, nir_opt_dce);
+      /* Clean up redundant vertex ID loads */
+      if (!key->hw || key->adjacency) {
+         NIR_PASS(_, b->shader, nir_opt_cse);
+         NIR_PASS(_, b->shader, nir_opt_dce);
+      }
    }
 
    if (!key->hw) {
@@ -229,13 +223,42 @@ agx_nir_vs_prolog(nir_builder *b, const void *key_)
 
    /* Finally, lower uniforms according to our ABI */
    unsigned nr = DIV_ROUND_UP(BITSET_LAST_BIT(key->component_mask), 4);
-   nir_shader_intrinsics_pass(b->shader, lower_non_monolithic_uniforms,
-                              nir_metadata_control_flow, &nr);
+   agx_nir_lower_non_monolithic_uniforms(b->shader, nr);
    b->shader->info.io_lowered = true;
 }
 
 static bool
-lower_input_to_prolog(nir_builder *b, nir_intrinsic_instr *intr, void *data)
+gather_inputs(nir_builder *b, nir_intrinsic_instr *intr, void *data)
+{
+   if (intr->intrinsic != nir_intrinsic_load_input)
+      return false;
+
+   unsigned idx = nir_src_as_uint(intr->src[0]) + nir_intrinsic_base(intr);
+   unsigned comp = nir_intrinsic_component(intr);
+
+   assert(intr->def.bit_size == 32 && "todo: push conversions up?");
+   unsigned base = 4 * idx + comp;
+
+   b->cursor = nir_before_instr(&intr->instr);
+   BITSET_WORD *comps_read = data;
+   nir_component_mask_t mask = nir_def_components_read(&intr->def);
+
+   u_foreach_bit(c, mask) {
+      BITSET_SET(comps_read, base + c);
+   }
+
+   return false;
+}
+
+bool
+agx_nir_gather_vs_inputs(nir_shader *s, BITSET_WORD *attrib_components_read)
+{
+   return nir_shader_intrinsics_pass(
+      s, gather_inputs, nir_metadata_control_flow, attrib_components_read);
+}
+
+static bool
+lower_input_to_prolog(nir_builder *b, nir_intrinsic_instr *intr, void *_data)
 {
    if (intr->intrinsic != nir_intrinsic_load_input)
       return false;
@@ -251,24 +274,15 @@ lower_input_to_prolog(nir_builder *b, nir_intrinsic_instr *intr, void *data)
       nir_load_exported_agx(b, intr->def.num_components, intr->def.bit_size,
                             .base = AGX_ABI_VIN_ATTRIB(base));
 
-   BITSET_WORD *comps_read = data;
-   nir_component_mask_t mask = nir_def_components_read(&intr->def);
-
-   u_foreach_bit(c, mask) {
-      BITSET_SET(comps_read, base + c);
-   }
-
    nir_def_replace(&intr->def, val);
    return true;
 }
 
 bool
-agx_nir_lower_vs_input_to_prolog(nir_shader *s,
-                                 BITSET_WORD *attrib_components_read)
+agx_nir_lower_vs_input_to_prolog(nir_shader *s)
 {
    return nir_shader_intrinsics_pass(s, lower_input_to_prolog,
-                                     nir_metadata_control_flow,
-                                     attrib_components_read);
+                                     nir_metadata_control_flow, NULL);
 }
 
 static bool

src/asahi/vulkan/hk_cmd_buffer.h

@@ -62,6 +62,7 @@ struct hk_root_descriptor_table {
          /* Vertex input state */
          uint64_t attrib_base[AGX_MAX_VBUFS];
          uint32_t attrib_clamps[AGX_MAX_VBUFS];
+         uint32_t attrib_strides[AGX_MAX_VBUFS];
 
          /* Pointer to the VS->TCS, VS->GS, or TES->GS buffer. */
          uint64_t vertex_output_buffer;
@@ -865,3 +866,8 @@ void hk_dispatch_precomp(struct hk_cmd_buffer *cmd, struct agx_grid grid,
 
 void hk_queue_write(struct hk_cmd_buffer *cmd, uint64_t address, uint32_t value,
                     bool after_gfx);
+
+void agx_fill_velem_keys(const struct vk_vertex_input_state *vi,
+                         uint64_t attribs_read, struct agx_velem_key *keys);
+
+struct agx_robustness hk_prolog_robustness(struct hk_device *dev);

src/asahi/vulkan/hk_cmd_draw.c

@@ -2614,6 +2614,30 @@ uses_blend_constant(const struct vk_color_blend_state *cb)
    return false;
 }
 
+void
+agx_fill_velem_keys(const struct vk_vertex_input_state *vi,
+                    uint64_t attribs_read, struct agx_velem_key *keys)
+{
+   u_foreach_bit(a, vi->attributes_valid) {
+      struct vk_vertex_attribute_state attr = vi->attributes[a];
+
+      assert(vi->bindings_valid & BITFIELD_BIT(attr.binding));
+      struct vk_vertex_binding_state binding = vi->bindings[attr.binding];
+
+      /* nir_assign_io_var_locations compacts vertex inputs, eliminating
+       * unused inputs. We need to do the same here to match the locations.
+       */
+      unsigned slot = util_bitcount64(attribs_read & BITFIELD_MASK(a));
+
+      keys[slot] = (struct agx_velem_key){
+         .format = hk_format_to_pipe_format(attr.format),
+         .stride = binding.stride,
+         .divisor = binding.divisor,
+         .instanced = binding.input_rate == VK_VERTEX_INPUT_RATE_INSTANCE,
+      };
+   }
+}
+
 static void
 hk_flush_dynamic_state(struct hk_cmd_buffer *cmd, struct hk_cs *cs,
                        uint32_t draw_id, struct agx_draw draw)
@@ -2717,6 +2741,7 @@ hk_flush_dynamic_state(struct hk_cmd_buffer *cmd, struct hk_cs *cs,
                : AGX_ROBUSTNESS_DISABLED,
 
          .prolog.robustness.soft_fault = agx_has_soft_fault(&dev->dev),
+         .prolog.static_vi = !sw_vs->info.vs.use_prolog,
       };
 
       enum mesa_prim prim = vk_conv_topology(dyn->ia.primitive_topology);
@@ -2739,31 +2764,24 @@ hk_flush_dynamic_state(struct hk_cmd_buffer *cmd, struct hk_cs *cs,
       BITSET_COPY(key.prolog.component_mask,
                   sw_vs->info.vs.attrib_components_read);
 
-      u_foreach_bit(a, dyn->vi->attributes_valid) {
-         struct vk_vertex_attribute_state attr = dyn->vi->attributes[a];
+      if (sw_vs->info.vs.use_prolog) {
+         agx_fill_velem_keys(dyn->vi, sw_vs->info.vs.attribs_read,
+                             key.prolog.attribs);
 
-         assert(dyn->vi->bindings_valid & BITFIELD_BIT(attr.binding));
-         struct vk_vertex_binding_state binding =
-            dyn->vi->bindings[attr.binding];
+         u_foreach_bit(a, dyn->vi->attributes_valid) {
+            unsigned slot =
+               util_bitcount64(sw_vs->info.vs.attribs_read & BITFIELD_MASK(a));
 
-         /* nir_assign_io_var_locations compacts vertex inputs, eliminating
-          * unused inputs. We need to do the same here to match the locations.
-          */
-         unsigned slot =
-            util_bitcount64(sw_vs->info.vs.attribs_read & BITFIELD_MASK(a));
-
-         key.prolog.attribs[slot] = (struct agx_velem_key){
-            .format = hk_format_to_pipe_format(attr.format),
-            .stride = dyn->vi_binding_strides[attr.binding],
-            .divisor = binding.divisor,
-            .instanced = binding.input_rate == VK_VERTEX_INPUT_RATE_INSTANCE,
-         };
+            key.prolog.attribs[slot].stride =
+               dyn->vi_binding_strides[dyn->vi->attributes[a].binding];
+         }
       }
 
       hk_update_fast_linked(cmd, sw_vs, &key);
    }
 
-   if (IS_DIRTY(VI) || IS_DIRTY(VI_BINDINGS_VALID) || vgt_dirty ||
+   if (IS_DIRTY(VI) || IS_DIRTY(VI_BINDINGS_VALID) ||
+       IS_DIRTY(VI_BINDING_STRIDES) || vgt_dirty ||
        (gfx->dirty & HK_DIRTY_VB)) {
 
       unsigned slot = 0;
@@ -2771,14 +2789,21 @@ hk_flush_dynamic_state(struct hk_cmd_buffer *cmd, struct hk_cs *cs,
          if (dyn->vi->attributes_valid & BITFIELD_BIT(a)) {
             struct vk_vertex_attribute_state attr = dyn->vi->attributes[a];
             struct hk_addr_range vb = gfx->vb[attr.binding];
+            enum pipe_format fmt = hk_format_to_pipe_format(attr.format);
+            enum pipe_format interchange_format = agx_vbo_internal_format(fmt);
+            unsigned interchange_align =
+               util_format_get_blocksize(interchange_format);
 
             desc->root.draw.attrib_clamps[slot] = agx_calculate_vbo_clamp(
-               vb.addr, hk_format_to_pipe_format(attr.format), vb.range,
-               dyn->vi_binding_strides[attr.binding], attr.offset,
-               &desc->root.draw.attrib_base[slot]);
+               vb.addr, fmt, vb.range, dyn->vi_binding_strides[attr.binding],
+               attr.offset, &desc->root.draw.attrib_base[slot]);
+
+            desc->root.draw.attrib_strides[slot] =
+               dyn->vi_binding_strides[attr.binding] / interchange_align;
          } else {
             desc->root.draw.attrib_base[slot] = AGX_ZERO_PAGE_ADDRESS;
             desc->root.draw.attrib_clamps[slot] = 0;
+            desc->root.draw.attrib_strides[slot] = 0;
          }
 
          ++slot;

src/asahi/vulkan/hk_nir_lower_descriptors.c

@@ -431,6 +431,10 @@ lower_uvs_index(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
    case nir_intrinsic_load_rasterization_stream:
       return lower_sysval_to_root_table(b, intrin, draw.rasterization_stream);
 
+   case nir_intrinsic_load_vbo_stride_agx:
+      return lower_sysval_to_root_table(
+         b, intrin, draw.attrib_strides[nir_src_as_uint(intrin->src[0])]);
+
    case nir_intrinsic_load_is_first_fan_agx: {
       unsigned offset = hk_root_descriptor_offset(draw.provoking);
       b->cursor = nir_instr_remove(&intrin->instr);

src/asahi/vulkan/hk_shader.c

@@ -10,6 +10,7 @@
 #include "agx_device.h"
 #include "agx_helpers.h"
 #include "agx_nir_lower_gs.h"
+#include "agx_nir_lower_vbo.h"
 #include "glsl_types.h"
 #include "libagx.h"
 #include "nir.h"
@@ -29,6 +30,7 @@
 #include "nir_intrinsics_indices.h"
 #include "nir_xfb_info.h"
 #include "shader_enums.h"
+#include "vk_graphics_state.h"
 #include "vk_nir_convert_ycbcr.h"
 #include "vk_physical_device_features.h"
 #include "vk_pipeline.h"
@@ -66,6 +68,19 @@ struct hk_fs_key {
 };
 static_assert(sizeof(struct hk_fs_key) == 4, "packed");
 
+struct hk_vs_key {
+   struct agx_velem_key attribs[32];
+   bool skip_prolog;
+   bool static_strides;
+   bool pad[2];
+};
+static_assert(sizeof(struct hk_vs_key) == 260, "packed");
+
+union hk_key {
+   struct hk_vs_key vs;
+   struct hk_fs_key fs;
+};
+
 static void
 shared_var_info(const struct glsl_type *type, unsigned *size, unsigned *align)
 {
@@ -172,6 +187,29 @@ hk_preprocess_nir(struct vk_physical_device *vk_pdev, nir_shader *nir,
    NIR_PASS(_, nir, nir_lower_compute_system_values, &csv_options);
 }
 
+static void
+hk_populate_vs_key(struct hk_vs_key *key,
+                   const struct vk_graphics_pipeline_state *state)
+{
+   memset(key, 0, sizeof(*key));
+
+   if (state == NULL || !state->vi ||
+       BITSET_TEST(state->dynamic, MESA_VK_DYNAMIC_VI) ||
+       BITSET_TEST(state->dynamic, MESA_VK_DYNAMIC_VI_BINDINGS_VALID))
+      return;
+
+   agx_fill_velem_keys(state->vi, ~0 /* compacted on use */, key->attribs);
+   key->skip_prolog = true;
+   key->static_strides =
+      !BITSET_TEST(state->dynamic, MESA_VK_DYNAMIC_VI_BINDING_STRIDES);
+
+   if (!key->static_strides) {
+      for (unsigned i = 0; i < ARRAY_SIZE(key->attribs); ++i) {
+         key->attribs[i].stride = 0;
+      }
+   }
+}
+
 static void
 hk_populate_fs_key(struct hk_fs_key *key,
                    const struct vk_graphics_pipeline_state *state)
@@ -232,7 +270,11 @@ hk_hash_graphics_state(struct vk_physical_device *device,
 {
    struct mesa_blake3 blake3_ctx;
    _mesa_blake3_init(&blake3_ctx);
-   if (state && (stages & VK_SHADER_STAGE_FRAGMENT_BIT)) {
+   if (state && (stages & VK_SHADER_STAGE_VERTEX_BIT)) {
+      struct hk_vs_key key;
+      hk_populate_vs_key(&key, state);
+      _mesa_blake3_update(&blake3_ctx, &key, sizeof(key));
+   } else if (state && (stages & VK_SHADER_STAGE_FRAGMENT_BIT)) {
       struct hk_fs_key key;
       hk_populate_fs_key(&key, state);
       _mesa_blake3_update(&blake3_ctx, &key, sizeof(key));
@@ -1013,7 +1055,7 @@ static VkResult
 hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator,
                nir_shader *nir, VkShaderCreateFlagsEXT shader_flags,
                const struct vk_pipeline_robustness_state *rs,
-               const struct hk_fs_key *fs_key, enum hk_feature_key features,
+               const union hk_key *key, enum hk_feature_key features,
                struct hk_shader *shader, gl_shader_stage sw_stage, bool hw,
                nir_xfb_info *xfb_info, unsigned set_count)
 {
@@ -1262,17 +1304,17 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info,
 
    gl_shader_stage sw_stage = nir->info.stage;
 
-   struct hk_fs_key fs_key_tmp, *fs_key = NULL;
+   union hk_key key_tmp, *key = NULL;
    if (sw_stage == MESA_SHADER_FRAGMENT) {
-      hk_populate_fs_key(&fs_key_tmp, state);
-      fs_key = &fs_key_tmp;
+      hk_populate_fs_key(&key_tmp.fs, state);
+      key = &key_tmp;
 
-      nir->info.fs.uses_sample_shading |= fs_key->force_sample_shading;
+      nir->info.fs.uses_sample_shading |= key->fs.force_sample_shading;
 
       /* Force late-Z for Z/S self-deps. TODO: There's probably a less silly way
        * to do this.
        */
-      if (fs_key->zs_self_dep) {
+      if (key->fs.zs_self_dep) {
          nir_builder b =
             nir_builder_at(nir_before_impl(nir_shader_get_entrypoint(nir)));
          nir_discard_if(&b, nir_imm_false(&b));
@@ -1280,6 +1322,9 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info,
       }
 
       NIR_PASS(_, nir, agx_nir_lower_sample_intrinsics, false);
+   } else if (sw_stage == MESA_SHADER_VERTEX) {
+      hk_populate_vs_key(&key_tmp.vs, state);
+      key = &key_tmp;
    } else if (sw_stage == MESA_SHADER_TESS_CTRL) {
       NIR_PASS(_, nir, agx_nir_lower_tcs);
    }
@@ -1371,18 +1416,49 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info,
           */
          nir_shader *clone = last ? nir : nir_shader_clone(NULL, nir);
 
-         if (sw_stage == MESA_SHADER_VERTEX) {
-            NIR_PASS(_, clone, agx_nir_lower_vs_input_to_prolog,
-                     shader->info.vs.attrib_components_read);
+         NIR_PASS(_, clone, agx_nir_gather_vs_inputs,
+                  shader->info.vs.attrib_components_read);
 
+         if (sw_stage == MESA_SHADER_VERTEX) {
+            shader->info.vs.use_prolog = !(key && key->vs.skip_prolog);
             shader->info.vs.attribs_read =
                nir->info.inputs_read >> VERT_ATTRIB_GENERIC0;
+
+            if (shader->info.vs.use_prolog) {
+               NIR_PASS(_, clone, agx_nir_lower_vs_input_to_prolog);
+            } else {
+               struct agx_velem_key attribs[AGX_MAX_ATTRIBS];
+               for (unsigned a = 0; a < AGX_MAX_ATTRIBS; ++a) {
+                  if (key->vs.attribs[a].format) {
+                     unsigned slot = util_bitcount64(
+                        shader->info.vs.attribs_read & BITFIELD_MASK(a));
+
+                     attribs[slot] = key->vs.attribs[a];
+                  }
+               }
+
+               struct agx_robustness agx_rs = {
+                  .soft_fault = agx_has_soft_fault(&dev->dev),
+
+                  /* Correctly handling GPL + pipeline-robustness requires
+                   * runtime changes, and I don't care enough to optimize this.
+                   */
+                  .level = AGX_ROBUSTNESS_D3D,
+               };
+
+               agx_nir_lower_vbo(clone, attribs, agx_rs,
+                                 !key->vs.static_strides);
+
+               unsigned nr = DIV_ROUND_UP(
+                  BITSET_LAST_BIT(shader->info.vs.attrib_components_read), 4);
+               agx_nir_lower_non_monolithic_uniforms(clone, nr);
+            }
          }
 
          /* hk_compile_nir takes ownership of the clone */
          result =
             hk_compile_nir(dev, pAllocator, clone, info->flags,
-                           info->robustness, fs_key, features, shader, sw_stage,
+                           info->robustness, key, features, shader, sw_stage,
                            hw, nir->xfb_info, info->set_layout_count);
          if (result != VK_SUCCESS) {
             hk_api_shader_destroy(&dev->vk, &obj->vk, pAllocator);
@@ -1395,8 +1471,8 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info,
 
       /* hk_compile_nir takes ownership of nir */
       result = hk_compile_nir(dev, pAllocator, nir, info->flags,
-                              info->robustness, fs_key, features, shader,
-                              sw_stage, true, NULL, info->set_layout_count);
+                              info->robustness, key, features, shader, sw_stage,
+                              true, NULL, info->set_layout_count);
       if (result != VK_SUCCESS) {
          hk_api_shader_destroy(&dev->vk, &obj->vk, pAllocator);
          return result;

src/asahi/vulkan/hk_shader.h

@@ -68,7 +68,8 @@ struct hk_shader_info {
       struct {
          uint32_t attribs_read;
          BITSET_DECLARE(attrib_components_read, AGX_MAX_ATTRIBS * 4);
-         uint8_t _pad[8];
+         bool use_prolog;
+         uint8_t _pad[7];
       } vs;
 
       struct {

src/gallium/drivers/asahi/agx_state.c

@@ -1567,8 +1567,8 @@ agx_compile_variant(struct agx_device *dev, struct pipe_context *pctx,
       if (nir->info.vs.tes_agx) {
          NIR_PASS(_, nir, agx_nir_lower_tes, key->hw);
       } else {
-         NIR_PASS(_, nir, agx_nir_lower_vs_input_to_prolog,
-                  attrib_components_read);
+         NIR_PASS(_, nir, agx_nir_gather_vs_inputs, attrib_components_read);
+         NIR_PASS(_, nir, agx_nir_lower_vs_input_to_prolog);
       }
 
       if (key->hw) {