radv: Use smaller esgs_itemsize for ACO.

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

4 files changed, 30 insertions, 23 deletions

diff --git a/src/amd/compiler/aco_instruction_selection_setup.cpp b/src/amd/compiler/aco_instruction_selection_setup.cpp
index 04dbe8f7e03..28645017b77 100644
--- a/src/amd/compiler/aco_instruction_selection_setup.cpp
+++ b/src/amd/compiler/aco_instruction_selection_setup.cpp
@@ -796,14 +796,6 @@ setup_vs_variables(isel_context *ctx, nir_shader *nir)
       radv_vs_output_info *outinfo = &ctx->program->info->vs.outinfo;
       setup_vs_output_info(ctx, nir, outinfo->export_prim_id,
                            ctx->options->key.vs_common_out.export_clip_dists, outinfo);
-   } else if (ctx->stage == vertex_geometry_gs || ctx->stage == vertex_es) {
-      /* TODO: radv_nir_shader_info_pass() already sets this but it's larger
-       * than it needs to be in order to set it better, we have to improve
-       * radv_nir_shader_info_pass() because gfx9_get_gs_info() uses
-       * esgs_itemsize and has to be done before compilation
-       */
-      /* radv_es_output_info *outinfo = &ctx->program->info->vs.es_info;
-      outinfo->esgs_itemsize = util_bitcount64(ctx->output_masks[nir->info.stage]) * 16u; */
    } else if (ctx->stage == vertex_ls) {
       ctx->tcs_num_inputs = ctx->program->info->vs.num_linked_outputs;
    }
diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index ae584328ab4..33b93e7d0c9 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -2617,7 +2617,8 @@ radv_fill_shader_info(struct radv_pipeline *pipeline,
 		radv_nir_shader_info_pass(nir[MESA_SHADER_FRAGMENT],
 					  pipeline->layout,
 					  &keys[MESA_SHADER_FRAGMENT],
-					  &infos[MESA_SHADER_FRAGMENT]);
+					  &infos[MESA_SHADER_FRAGMENT],
+					  pipeline->device->physical_device->use_aco);
 
 		/* TODO: These are no longer used as keys we should refactor this */
 		keys[MESA_SHADER_VERTEX].vs_common_out.export_prim_id =
@@ -2668,7 +2669,8 @@ radv_fill_shader_info(struct radv_pipeline *pipeline,
 		for (int i = 0; i < 2; i++) {
 			radv_nir_shader_info_pass(combined_nir[i],
 						  pipeline->layout, &key,
-						  &infos[MESA_SHADER_TESS_CTRL]);
+						  &infos[MESA_SHADER_TESS_CTRL],
+						  pipeline->device->physical_device->use_aco);
 		}
 
 		keys[MESA_SHADER_TESS_EVAL].tes.num_patches =
@@ -2691,7 +2693,8 @@ radv_fill_shader_info(struct radv_pipeline *pipeline,
 			radv_nir_shader_info_pass(combined_nir[i],
 						  pipeline->layout,
 						  &keys[pre_stage],
-						  &infos[MESA_SHADER_GEOMETRY]);
+						  &infos[MESA_SHADER_GEOMETRY],
+						  pipeline->device->physical_device->use_aco);
 		}
 
 		filled_stages |= (1 << pre_stage);
@@ -2716,7 +2719,7 @@ radv_fill_shader_info(struct radv_pipeline *pipeline,
 
 		radv_nir_shader_info_init(&infos[i]);
 		radv_nir_shader_info_pass(nir[i], pipeline->layout,
-					  &keys[i], &infos[i]);
+					  &keys[i], &infos[i], pipeline->device->physical_device->use_aco);
 	}
 
 	for (int i = 0; i < MESA_SHADER_STAGES; i++) {
@@ -2975,7 +2978,7 @@ void radv_create_shaders(struct radv_pipeline *pipeline,
 
 			radv_nir_shader_info_pass(nir[MESA_SHADER_GEOMETRY],
 						  pipeline->layout, &key,
-						  &info);
+						  &info, pipeline->device->physical_device->use_aco);
 			info.wave_size = 64; /* Wave32 not supported. */
 			info.ballot_bit_size = 64;
 
diff --git a/src/amd/vulkan/radv_private.h b/src/amd/vulkan/radv_private.h
index ec4b45235c6..a54f0147fbc 100644
--- a/src/amd/vulkan/radv_private.h
+++ b/src/amd/vulkan/radv_private.h
@@ -2366,7 +2366,8 @@ struct radv_shader_variant_key;
 void radv_nir_shader_info_pass(const struct nir_shader *nir,
 			       const struct radv_pipeline_layout *layout,
 			       const struct radv_shader_variant_key *key,
-			       struct radv_shader_info *info);
+			       struct radv_shader_info *info,
+			       bool use_aco);
 
 void radv_nir_shader_info_init(struct radv_shader_info *info);
 
diff --git a/src/amd/vulkan/radv_shader_info.c b/src/amd/vulkan/radv_shader_info.c
index dfccba5600c..c651f216280 100644
--- a/src/amd/vulkan/radv_shader_info.c
+++ b/src/amd/vulkan/radv_shader_info.c
@@ -656,7 +656,8 @@ void
 radv_nir_shader_info_pass(const struct nir_shader *nir,
 			  const struct radv_pipeline_layout *layout,
 			  const struct radv_shader_variant_key *key,
-			  struct radv_shader_info *info)
+			  struct radv_shader_info *info,
+			  bool use_aco)
 {
 	struct nir_function *func =
 		(struct nir_function *)exec_list_get_head_const(&nir->functions);
@@ -808,17 +809,27 @@ radv_nir_shader_info_pass(const struct nir_shader *nir,
 	    key->vs_common_out.as_es) {
 		struct radv_es_output_info *es_info =
 			nir->info.stage == MESA_SHADER_VERTEX ? &info->vs.es_info : &info->tes.es_info;
-		uint32_t max_output_written = 0;
 
-		uint64_t output_mask = nir->info.outputs_written;
-		while (output_mask) {
-			const int i = u_bit_scan64(&output_mask);
-			unsigned param_index = shader_io_get_unique_index(i);
+		if (use_aco) {
+			/* The outputs don't contain gaps, se we can use the number of outputs */
+			uint32_t num_outputs_written = nir->info.stage == MESA_SHADER_VERTEX
+				? info->vs.num_linked_outputs
+				: info->tes.num_linked_outputs;
+			es_info->esgs_itemsize = num_outputs_written * 16;
+		} else {
+			/* The outputs may contain gaps, use the highest output index + 1 */
+			uint32_t max_output_written = 0;
+			uint64_t output_mask = nir->info.outputs_written;
+
+			while (output_mask) {
+				const int i = u_bit_scan64(&output_mask);
+				unsigned param_index = shader_io_get_unique_index(i);
+
+				max_output_written = MAX2(param_index, max_output_written);
+			}
 
-			max_output_written = MAX2(param_index, max_output_written);
+			es_info->esgs_itemsize = (max_output_written + 1) * 16;
 		}
-
-		es_info->esgs_itemsize = (max_output_written + 1) * 16;
 	}
 
 	info->float_controls_mode = nir->info.float_controls_execution_mode;