hk: Expose the vertex param buffer to other stages

For vertex shaders, it comes from the preamble.  For geometry and
tessellation shaders it comes straight from the root table.

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/vulkan/hk_cmd_draw.c

@@ -1849,7 +1849,7 @@ build_fs_prolog(nir_builder *b, const void *key)
    agx_nir_fs_prolog(b, key);
 
    /* Lower load_stat_query_address_agx, needed for FS statistics */
-   NIR_PASS(_, b->shader, hk_lower_uvs_index, 0);
+   NIR_PASS(_, b->shader, hk_lower_uvs_index, MESA_SHADER_FRAGMENT, 0);
    NIR_PASS(_, b->shader, nir_shader_intrinsics_pass, lower_fs_root,
             nir_metadata_control_flow, NULL);
 }

src/asahi/vulkan/hk_nir_lower_descriptors.c

@@ -388,10 +388,15 @@ translate_pipeline_stat_bit(enum pipe_statistics_query_index pipe)
    }
 }
 
+struct lower_uvs_args {
+   mesa_shader_stage sw_stage;
+   unsigned nr_vbos;
+};
+
 static bool
 lower_uvs_index(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
 {
-   unsigned *nr_vbos = data;
+   const struct lower_uvs_args *args = data;
 
    switch (intrin->intrinsic) {
    case nir_intrinsic_load_uvs_index_agx: {
@@ -428,6 +433,19 @@ lower_uvs_index(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
    case nir_intrinsic_load_tess_param_buffer_poly:
       return lower_sysval_to_root_table(b, intrin, draw.tess_params);
 
+   case nir_intrinsic_load_vertex_param_buffer_poly:
+      if (args->sw_stage == MESA_SHADER_VERTEX) {
+         b->cursor = nir_instr_remove(&intrin->instr);
+
+         unsigned base = AGX_ABI_VUNI_VERTEX_PARAMS(args->nr_vbos);
+         nir_def *val = nir_load_preamble(b, 1, 64, .base = base);
+         nir_def_rewrite_uses(&intrin->def, val);
+
+         return true;
+      } else {
+         return lower_sysval_to_root_table(b, intrin, draw.vertex_params);
+      }
+
    case nir_intrinsic_load_rasterization_stream:
       return lower_sysval_to_root_table(b, intrin, draw.rasterization_stream);
 
@@ -454,22 +472,17 @@ lower_uvs_index(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
    case nir_intrinsic_load_base_vertex:
    case nir_intrinsic_load_first_vertex:
    case nir_intrinsic_load_base_instance:
-   case nir_intrinsic_load_draw_id:
-   case nir_intrinsic_load_vertex_param_buffer_poly: {
+   case nir_intrinsic_load_draw_id: {
       b->cursor = nir_instr_remove(&intrin->instr);
 
-      unsigned base = AGX_ABI_VUNI_FIRST_VERTEX(*nr_vbos);
+      unsigned base = AGX_ABI_VUNI_FIRST_VERTEX(args->nr_vbos);
       unsigned size = 32;
 
       if (intrin->intrinsic == nir_intrinsic_load_base_instance) {
-         base = AGX_ABI_VUNI_BASE_INSTANCE(*nr_vbos);
+         base = AGX_ABI_VUNI_BASE_INSTANCE(args->nr_vbos);
       } else if (intrin->intrinsic == nir_intrinsic_load_draw_id) {
-         base = AGX_ABI_VUNI_DRAW_ID(*nr_vbos);
+         base = AGX_ABI_VUNI_DRAW_ID(args->nr_vbos);
          size = 16;
-      } else if (intrin->intrinsic ==
-                 nir_intrinsic_load_vertex_param_buffer_poly) {
-         base = AGX_ABI_VUNI_VERTEX_PARAMS(*nr_vbos);
-         size = 64;
       }
 
       nir_def *val = nir_load_preamble(b, 1, size, .base = base);
@@ -529,10 +542,14 @@ lower_uvs_index(nir_builder *b, nir_intrinsic_instr *intrin, void *data)
 }
 
 bool
-hk_lower_uvs_index(nir_shader *s, unsigned nr_vbos)
+hk_lower_uvs_index(nir_shader *s, mesa_shader_stage sw_stage, unsigned nr_vbos)
 {
+   struct lower_uvs_args args = {
+      .sw_stage = sw_stage,
+      .nr_vbos = nr_vbos,
+   };
    return nir_shader_intrinsics_pass(s, lower_uvs_index,
-                                     nir_metadata_control_flow, &nr_vbos);
+                                     nir_metadata_control_flow, &args);
 }
 
 static bool

src/asahi/vulkan/hk_shader.c

@@ -1153,7 +1153,7 @@ hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator,
    shader->info.set_count = set_count;
 
    /* XXX: rename */
-   NIR_PASS(_, nir, hk_lower_uvs_index, nr_vbos);
+   NIR_PASS(_, nir, hk_lower_uvs_index, sw_stage, nr_vbos);
    NIR_PASS(_, nir, nir_shader_intrinsics_pass, lower_uniforms,
             nir_metadata_control_flow, &root);
 

src/asahi/vulkan/hk_shader.h

@@ -342,7 +342,8 @@ hk_buffer_addr_format(VkPipelineRobustnessBufferBehaviorEXT robustness)
    }
 }
 
-bool hk_lower_uvs_index(nir_shader *s, unsigned vs_uniform_base);
+bool
+hk_lower_uvs_index(nir_shader *s, mesa_shader_stage sw_stage, unsigned nr_vbos);
 
 bool
 hk_nir_lower_descriptors(nir_shader *nir,