radv: Handle GFX9 merged shaders in radv_flush_constants()

This was not previously handled correctly. For example,
push_constant_stages might only contain MESA_SHADER_VERTEX because
only that stage was changed by CmdPushConstants or
CmdBindDescriptorSets.

In that case, if vertex has been merged with tess control, then the
push constant address wouldn't be updated since
pipeline->shaders[MESA_SHADER_VERTEX] would be NULL.

Use radv_get_shader() instead of getting the shader directly so that
we get the right shader if merged. Also, skip emitting the address
redundantly - if two merged stages are set in push_constant_stages
this change would have made the address get emitted twice.

Signed-off-by: Alex Smith <[email protected]>
Cc: "18.1" <[email protected]>
Reviewed-by: Bas Nieuwenhuizen <[email protected]>

1 files changed, 8 insertions, 1 deletions

diff --git a/src/amd/vulkan/radv_cmd_buffer.c b/src/amd/vulkan/radv_cmd_buffer.c
index e9b1dffc1cd..043b4a2f44a 100644
--- a/src/amd/vulkan/radv_cmd_buffer.c
+++ b/src/amd/vulkan/radv_cmd_buffer.c
@@ -1585,6 +1585,7 @@ radv_flush_constants(struct radv_cmd_buffer *cmd_buffer,
 					 ? cmd_buffer->state.compute_pipeline
 					 : cmd_buffer->state.pipeline;
 	struct radv_pipeline_layout *layout = pipeline->layout;
+	struct radv_shader_variant *shader, *prev_shader;
 	unsigned offset;
 	void *ptr;
 	uint64_t va;
@@ -1609,10 +1610,16 @@ radv_flush_constants(struct radv_cmd_buffer *cmd_buffer,
 	MAYBE_UNUSED unsigned cdw_max = radeon_check_space(cmd_buffer->device->ws,
 	                                                   cmd_buffer->cs, MESA_SHADER_STAGES * 4);
 
+	prev_shader = NULL;
 	radv_foreach_stage(stage, stages) {
-		if (pipeline->shaders[stage]) {
+		shader = radv_get_shader(pipeline, stage);
+
+		/* Avoid redundantly emitting the address for merged stages. */
+		if (shader && shader != prev_shader) {
 			radv_emit_userdata_address(cmd_buffer, pipeline, stage,
 						   AC_UD_PUSH_CONSTANTS, va);
+
+			prev_shader = shader;
 		}
 	}