radv/aco: Set I/O variable locations outside ACO.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Acked-by: Daniel Schürmann <daniel@schuermann.dev>
Reviewed-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/6865>
diff --git a/src/amd/compiler/aco_instruction_selection_setup.cpp b/src/amd/compiler/aco_instruction_selection_setup.cpp
index d01bb2e..9baee2a 100644
--- a/src/amd/compiler/aco_instruction_selection_setup.cpp
+++ b/src/amd/compiler/aco_instruction_selection_setup.cpp
@@ -470,16 +470,6 @@
 void
 setup_vs_variables(isel_context *ctx, nir_shader *nir)
 {
-   nir_foreach_shader_in_variable(variable, nir)
-   {
-      variable->data.driver_location = variable->data.location;
-   }
-   nir_foreach_shader_out_variable(variable, nir)
-   {
-      if (ctx->stage == vertex_vs || ctx->stage == ngg_vertex_gs)
-         variable->data.driver_location = variable->data.location;
-   }
-
    if (ctx->stage == vertex_vs || ctx->stage == ngg_vertex_gs) {
       radv_vs_output_info *outinfo = &ctx->program->info->vs.outinfo;
       setup_vs_output_info(ctx, nir, outinfo->export_prim_id,
@@ -501,10 +491,6 @@
    if (ctx->stage == vertex_geometry_gs || ctx->stage == tess_eval_geometry_gs)
       ctx->program->config->lds_size = ctx->program->info->gs_ring_info.lds_size; /* Already in units of the alloc granularity */
 
-   nir_foreach_shader_out_variable(variable, nir) {
-      variable->data.driver_location = variable->data.location;
-   }
-
    if (ctx->stage == vertex_geometry_gs)
       ctx->program->info->gs.es_type = MESA_SHADER_VERTEX;
    else if (ctx->stage == tess_eval_geometry_gs)
@@ -568,11 +554,6 @@
    ctx->tcs_num_patches = ctx->args->options->key.tes.num_patches;
    ctx->tcs_num_outputs = ctx->program->info->tes.num_linked_inputs;
 
-   nir_foreach_shader_out_variable(variable, nir) {
-      if (ctx->stage == tess_eval_vs || ctx->stage == ngg_tess_eval_gs)
-         variable->data.driver_location = variable->data.location;
-   }
-
    if (ctx->stage == tess_eval_vs || ctx->stage == ngg_tess_eval_gs) {
       radv_vs_output_info *outinfo = &ctx->program->info->tes.outinfo;
       setup_vs_output_info(ctx, nir, outinfo->export_prim_id,
@@ -585,11 +566,6 @@
 {
    switch (nir->info.stage) {
    case MESA_SHADER_FRAGMENT: {
-      nir_foreach_shader_out_variable(variable, nir)
-      {
-         int idx = variable->data.location + variable->data.index;
-         variable->data.driver_location = idx;
-      }
       break;
    }
    case MESA_SHADER_COMPUTE: {
diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index 8f89d0b..600ce8f 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -2263,17 +2263,24 @@
 }
 
 static void
-radv_set_linked_driver_locations(struct radv_pipeline *pipeline, nir_shader **shaders,
-                                 struct radv_shader_info infos[MESA_SHADER_STAGES])
+radv_set_driver_locations(struct radv_pipeline *pipeline, nir_shader **shaders,
+                              struct radv_shader_info infos[MESA_SHADER_STAGES])
 {
-	bool has_tess = shaders[MESA_SHADER_TESS_CTRL];
-	bool has_gs = shaders[MESA_SHADER_GEOMETRY];
+	if (shaders[MESA_SHADER_FRAGMENT]) {
+		nir_foreach_shader_out_variable(var, shaders[MESA_SHADER_FRAGMENT])
+		{
+			var->data.driver_location = var->data.location + var->data.index;
+		}
+	}
 
-	if (!has_tess && !has_gs)
+	if (!shaders[MESA_SHADER_VERTEX])
 		return;
 
+	bool has_tess = shaders[MESA_SHADER_TESS_CTRL];
+	bool has_gs = shaders[MESA_SHADER_GEOMETRY];
 	unsigned vs_info_idx = MESA_SHADER_VERTEX;
 	unsigned tes_info_idx = MESA_SHADER_TESS_EVAL;
+	unsigned last_vtg_stage = MESA_SHADER_VERTEX;
 
 	if (pipeline->device->physical_device->rad_info.chip_class >= GFX9) {
 		/* These are merged into the next stage */
@@ -2281,6 +2288,10 @@
 		tes_info_idx = has_gs ? MESA_SHADER_GEOMETRY : MESA_SHADER_TESS_EVAL;
 	}
 
+	nir_foreach_shader_in_variable(var, shaders[MESA_SHADER_VERTEX]) {
+		var->data.driver_location = var->data.location;
+	}
+
 	if (has_tess) {
 		nir_linked_io_var_info vs2tcs =
 			nir_assign_linked_io_var_locations(shaders[MESA_SHADER_VERTEX], shaders[MESA_SHADER_TESS_CTRL]);
@@ -2300,6 +2311,9 @@
 
 			infos[tes_info_idx].tes.num_linked_outputs = tes2gs.num_linked_io_vars;
 			infos[MESA_SHADER_GEOMETRY].gs.num_linked_inputs = tes2gs.num_linked_io_vars;
+			last_vtg_stage = MESA_SHADER_GEOMETRY;
+		} else {
+			last_vtg_stage = MESA_SHADER_TESS_EVAL;
 		}
 	} else if (has_gs) {
 		nir_linked_io_var_info vs2gs =
@@ -2307,6 +2321,11 @@
 
 		infos[vs_info_idx].vs.num_linked_outputs = vs2gs.num_linked_io_vars;
 		infos[MESA_SHADER_GEOMETRY].gs.num_linked_inputs = vs2gs.num_linked_io_vars;
+		last_vtg_stage = MESA_SHADER_GEOMETRY;
+	}
+
+	nir_foreach_shader_out_variable(var, shaders[last_vtg_stage]) {
+		var->data.driver_location = var->data.location;
 	}
 }
 
@@ -2937,7 +2956,7 @@
 	if (!(flags & VK_PIPELINE_CREATE_DISABLE_OPTIMIZATION_BIT))
 		radv_link_shaders(pipeline, nir);
 
-	radv_set_linked_driver_locations(pipeline, nir, infos);
+	radv_set_driver_locations(pipeline, nir, infos);
 
 	for (int i = 0; i < MESA_SHADER_STAGES; ++i) {
 		if (nir[i]) {
@@ -4653,10 +4672,10 @@
 			vs_size = pipeline->shaders[MESA_SHADER_TESS_EVAL]->info.wave_size;
 		else if (pipeline->shaders[MESA_SHADER_VERTEX])
 			vs_size = pipeline->shaders[MESA_SHADER_VERTEX]->info.wave_size;
-		
+
 		if (radv_pipeline_has_ngg(pipeline))
 			gs_size = vs_size;
-			
+
 		/* legacy GS only supports Wave64 */
 		stages |= S_028B54_HS_W32_EN(hs_size == 32 ? 1 : 0) |
 			  S_028B54_GS_W32_EN(gs_size == 32 ? 1 : 0) |
@@ -5237,7 +5256,7 @@
 		} else {
 			ret += 1u;
 		}
-		
+
 	}
 	return ret;
 }