radv: add shader arguments for dynamic patch control points

This introduces two new user SGPRS:

- tcs_offchip_layout: input patch size and number of patches in TCS
- tes_num_patches: number of patches in TES

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18344>

src/amd/vulkan/radv_shader.h

@@ -158,9 +158,11 @@ enum radv_ud_index {
    AC_UD_CS_TASK_IB,
    AC_UD_CS_MAX_UD,
    AC_UD_GS_MAX_UD,
+   AC_UD_TCS_OFFCHIP_LAYOUT = AC_UD_VS_MAX_UD,
    AC_UD_TCS_MAX_UD,
+   AC_UD_TES_NUM_PATCHES = AC_UD_SHADER_START,
    AC_UD_TES_MAX_UD,
-   AC_UD_MAX_UD = AC_UD_TCS_MAX_UD,
+   AC_UD_MAX_UD = AC_UD_CS_MAX_UD,
 };
 
 struct radv_stream_output {

src/amd/vulkan/radv_shader_args.c

@@ -91,6 +91,17 @@ count_vs_user_sgprs(const struct radv_shader_info *info)
    return count;
 }
 
+static uint8_t
+count_tes_user_sgprs(const struct radv_pipeline_key *key)
+{
+   unsigned count = 0;
+
+   if (key->dynamic_patch_control_points)
+      count++; /* tes_num_patches */
+
+   return count;
+}
+
 static uint8_t
 count_ms_user_sgprs(const struct radv_shader_info *info)
 {
@@ -152,6 +163,7 @@ static void
 allocate_user_sgprs(enum amd_gfx_level gfx_level, const struct radv_shader_info *info,
                     struct radv_shader_args *args, gl_shader_stage stage, bool has_previous_stage,
                     gl_shader_stage previous_stage, bool needs_view_index, bool has_ngg_query,
+                    const struct radv_pipeline_key *key,
                     struct user_sgpr_info *user_sgpr_info)
 {
    uint8_t user_sgpr_count = 0;
@@ -196,8 +208,11 @@ allocate_user_sgprs(enum amd_gfx_level gfx_level, const struct radv_shader_info
          if (previous_stage == MESA_SHADER_VERTEX)
             user_sgpr_count += count_vs_user_sgprs(info);
       }
+      if (key->dynamic_patch_control_points)
+         user_sgpr_count += 1; /* tcs_offchip_layout */
       break;
    case MESA_SHADER_TESS_EVAL:
+      count_tes_user_sgprs(key);
       break;
    case MESA_SHADER_GEOMETRY:
       if (has_previous_stage) {
@@ -206,6 +221,8 @@ allocate_user_sgprs(enum amd_gfx_level gfx_level, const struct radv_shader_info
 
          if (previous_stage == MESA_SHADER_VERTEX) {
             user_sgpr_count += count_vs_user_sgprs(info);
+         } else if (previous_stage == MESA_SHADER_TESS_EVAL) {
+            user_sgpr_count += count_tes_user_sgprs(key);
          } else if (previous_stage == MESA_SHADER_MESH) {
             user_sgpr_count += count_ms_user_sgprs(info);
          }
@@ -555,7 +572,7 @@ radv_declare_shader_args(enum amd_gfx_level gfx_level, const struct radv_pipelin
       args->user_sgprs_locs.shader_data[i].sgpr_idx = -1;
 
    allocate_user_sgprs(gfx_level, info, args, stage, has_previous_stage, previous_stage,
-                       needs_view_index, has_ngg_query, &user_sgpr_info);
+                       needs_view_index, has_ngg_query, key, &user_sgpr_info);
 
    if (args->explicit_scratch_args) {
       ac_add_arg(&args->ac, AC_ARG_SGPR, 2, AC_ARG_CONST_DESC_PTR, &args->ring_offsets);
@@ -671,6 +688,10 @@ radv_declare_shader_args(enum amd_gfx_level gfx_level, const struct radv_pipelin
             ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.view_index);
          }
 
+         if (key->dynamic_patch_control_points) {
+            ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->tcs_offchip_layout);
+         }
+
          ac_add_arg(&args->ac, AC_ARG_VGPR, 1, AC_ARG_INT, &args->ac.tcs_patch_id);
          ac_add_arg(&args->ac, AC_ARG_VGPR, 1, AC_ARG_INT, &args->ac.tcs_rel_ids);
 
@@ -682,6 +703,10 @@ radv_declare_shader_args(enum amd_gfx_level gfx_level, const struct radv_pipelin
             ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.view_index);
          }
 
+         if (key->dynamic_patch_control_points) {
+            ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->tcs_offchip_layout);
+         }
+
          ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.tess_offchip_offset);
          ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.tcs_factor_offset);
          if (args->explicit_scratch_args) {
@@ -700,6 +725,9 @@ radv_declare_shader_args(enum amd_gfx_level gfx_level, const struct radv_pipelin
       if (needs_view_index)
          ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.view_index);
 
+      if (key->dynamic_patch_control_points)
+         ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->tes_num_patches);
+
       if (info->tes.as_es) {
          ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.tess_offchip_offset);
          ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, NULL);
@@ -744,6 +772,9 @@ radv_declare_shader_args(enum amd_gfx_level gfx_level, const struct radv_pipelin
             ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.view_index);
          }
 
+         if (previous_stage == MESA_SHADER_TESS_EVAL && key->dynamic_patch_control_points)
+            ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->tes_num_patches);
+
          if (info->force_vrs_per_vertex) {
             ac_add_arg(&args->ac, AC_ARG_SGPR, 1, AC_ARG_INT, &args->ac.force_vrs_rates);
          }
@@ -861,15 +892,24 @@ radv_declare_shader_args(enum amd_gfx_level gfx_level, const struct radv_pipelin
    case MESA_SHADER_TESS_CTRL:
       if (args->ac.view_index.used)
          set_loc_shader(args, AC_UD_VIEW_INDEX, &user_sgpr_idx, 1);
+
+      if (args->tcs_offchip_layout.used)
+         set_loc_shader(args, AC_UD_TCS_OFFCHIP_LAYOUT, &user_sgpr_idx, 1);
       break;
    case MESA_SHADER_TESS_EVAL:
       if (args->ac.view_index.used)
          set_loc_shader(args, AC_UD_VIEW_INDEX, &user_sgpr_idx, 1);
+
+      if (args->tes_num_patches.used)
+         set_loc_shader(args, AC_UD_TES_NUM_PATCHES, &user_sgpr_idx, 1);
       break;
    case MESA_SHADER_GEOMETRY:
       if (args->ac.view_index.used)
          set_loc_shader(args, AC_UD_VIEW_INDEX, &user_sgpr_idx, 1);
 
+      if (args->tes_num_patches.used)
+         set_loc_shader(args, AC_UD_TES_NUM_PATCHES, &user_sgpr_idx, 1);
+
       if (args->ac.force_vrs_rates.used)
          set_loc_shader(args, AC_UD_FORCE_VRS_RATES, &user_sgpr_idx, 1);
 

src/amd/vulkan/radv_shader_args.h

@@ -65,6 +65,15 @@ struct radv_shader_args {
    /* PS epilogs */
    struct ac_arg ps_epilog_inputs[MAX_RTS];
 
+   /* TCS */
+   /* # [0:5] = the number of patch control points
+    * # [6:13] = the number of tessellation patches
+    */
+   struct ac_arg tcs_offchip_layout;
+
+   /* TES */
+   struct ac_arg tes_num_patches;
+
    struct radv_userdata_locations user_sgprs_locs;
    unsigned num_user_sgprs;