[Mesa-dev] [PATCH] radv: keep a stage mask per pipeline. (v2)

[Dave Airlie]

From: Dave Airlie <airlied at redhat.com>

This should reduce some pointless loops.

v2: fix missing check which causes crashes with compute shaders

Signed-off-by: Dave Airlie <airlied at redhat.com>
---
 src/amd/vulkan/radv_cmd_buffer.c | 53 +++++++++++-----------------------------
 src/amd/vulkan/radv_pipeline.c   |  2 ++
 src/amd/vulkan/radv_private.h    |  1 +
 3 files changed, 17 insertions(+), 39 deletions(-)

diff --git a/src/amd/vulkan/radv_cmd_buffer.c b/src/amd/vulkan/radv_cmd_buffer.c
index 833c3eb3f0d..7357eadae39 100644
--- a/src/amd/vulkan/radv_cmd_buffer.c
+++ b/src/amd/vulkan/radv_cmd_buffer.c
@@ -682,17 +682,10 @@ static void
 radv_emit_shaders_prefetch(struct radv_cmd_buffer *cmd_buffer,
 			   struct radv_pipeline *pipeline)
 {
-	radv_emit_shader_prefetch(cmd_buffer,
-				  pipeline->shaders[MESA_SHADER_VERTEX]);
-	radv_emit_shader_prefetch(cmd_buffer,
-				  pipeline->shaders[MESA_SHADER_TESS_CTRL]);
-	radv_emit_shader_prefetch(cmd_buffer,
-				  pipeline->shaders[MESA_SHADER_TESS_EVAL]);
-	radv_emit_shader_prefetch(cmd_buffer,
-				  pipeline->shaders[MESA_SHADER_GEOMETRY]);
+	radv_foreach_stage(stage, pipeline->stage_mask)
+		radv_emit_shader_prefetch(cmd_buffer,
+					  pipeline->shaders[stage]);
 	radv_emit_shader_prefetch(cmd_buffer, pipeline->gs_copy_shader);
-	radv_emit_shader_prefetch(cmd_buffer,
-				  pipeline->shaders[MESA_SHADER_FRAGMENT]);
 }
 
 static void
@@ -1603,12 +1596,11 @@ radv_emit_descriptor_set_userdata(struct radv_cmd_buffer *cmd_buffer,
 				  struct radv_descriptor_set *set,
 				  unsigned idx)
 {
-	if (cmd_buffer->state.pipeline) {
+	if (cmd_buffer->state.pipeline && (stages & cmd_buffer->state.pipeline->stage_mask)) {
 		radv_foreach_stage(stage, stages) {
-			if (cmd_buffer->state.pipeline->shaders[stage])
-				emit_stage_descriptor_set_userdata(cmd_buffer, cmd_buffer->state.pipeline,
-								   idx, set->va,
-								   stage);
+			emit_stage_descriptor_set_userdata(cmd_buffer, cmd_buffer->state.pipeline,
+							   idx, set->va,
+							   stage);
 		}
 	}
 
@@ -1658,25 +1650,10 @@ radv_flush_indirect_descriptor_sets(struct radv_cmd_buffer *cmd_buffer)
 	va += offset;
 
 	if (cmd_buffer->state.pipeline) {
-		if (cmd_buffer->state.pipeline->shaders[MESA_SHADER_VERTEX])
-			radv_emit_userdata_address(cmd_buffer, cmd_buffer->state.pipeline, MESA_SHADER_VERTEX,
-						   AC_UD_INDIRECT_DESCRIPTOR_SETS, va);
-
-		if (cmd_buffer->state.pipeline->shaders[MESA_SHADER_FRAGMENT])
-			radv_emit_userdata_address(cmd_buffer, cmd_buffer->state.pipeline, MESA_SHADER_FRAGMENT,
-						   AC_UD_INDIRECT_DESCRIPTOR_SETS, va);
-
-		if (radv_pipeline_has_gs(cmd_buffer->state.pipeline))
-			radv_emit_userdata_address(cmd_buffer, cmd_buffer->state.pipeline, MESA_SHADER_GEOMETRY,
-						   AC_UD_INDIRECT_DESCRIPTOR_SETS, va);
-
-		if (radv_pipeline_has_tess(cmd_buffer->state.pipeline))
-			radv_emit_userdata_address(cmd_buffer, cmd_buffer->state.pipeline, MESA_SHADER_TESS_CTRL,
-						   AC_UD_INDIRECT_DESCRIPTOR_SETS, va);
-
-		if (radv_pipeline_has_tess(cmd_buffer->state.pipeline))
-			radv_emit_userdata_address(cmd_buffer, cmd_buffer->state.pipeline, MESA_SHADER_TESS_EVAL,
+		radv_foreach_stage(stage, cmd_buffer->state.pipeline->stage_mask) {
+			radv_emit_userdata_address(cmd_buffer, cmd_buffer->state.pipeline, stage,
 						   AC_UD_INDIRECT_DESCRIPTOR_SETS, va);
+		}
 	}
 
 	if (cmd_buffer->state.compute_pipeline)
@@ -1751,10 +1728,8 @@ radv_flush_constants(struct radv_cmd_buffer *cmd_buffer,
 	                                                   cmd_buffer->cs, MESA_SHADER_STAGES * 4);
 
 	radv_foreach_stage(stage, stages) {
-		if (pipeline->shaders[stage]) {
-			radv_emit_userdata_address(cmd_buffer, pipeline, stage,
-						   AC_UD_PUSH_CONSTANTS, va);
-		}
+		radv_emit_userdata_address(cmd_buffer, pipeline, stage,
+					   AC_UD_PUSH_CONSTANTS, va);
 	}
 
 	cmd_buffer->push_constant_stages &= ~stages;
@@ -1819,9 +1794,9 @@ radv_upload_graphics_shader_descriptors(struct radv_cmd_buffer *cmd_buffer, bool
 	if (!radv_cmd_buffer_update_vertex_descriptors(cmd_buffer, pipeline_is_dirty))
 		return false;
 
-	radv_flush_descriptors(cmd_buffer, VK_SHADER_STAGE_ALL_GRAPHICS);
+	radv_flush_descriptors(cmd_buffer, cmd_buffer->state.pipeline->stage_mask);
 	radv_flush_constants(cmd_buffer, cmd_buffer->state.pipeline,
-			     VK_SHADER_STAGE_ALL_GRAPHICS);
+			     cmd_buffer->state.pipeline->stage_mask);
 
 	return true;
 }
diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index 5895a76ff4b..f059fa43f99 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -1994,6 +1994,8 @@ radv_pipeline_init(struct radv_pipeline *pipeline,
 	for (uint32_t i = 0; i < pCreateInfo->stageCount; i++) {
 		gl_shader_stage stage = ffs(pCreateInfo->pStages[i].stage) - 1;
 		pStages[stage] = &pCreateInfo->pStages[i];
+
+		pipeline->stage_mask |= (1 << stage);
 	}
 
 	radv_create_shaders(pipeline, device, cache, 
diff --git a/src/amd/vulkan/radv_private.h b/src/amd/vulkan/radv_private.h
index c59c8b8a732..27746783902 100644
--- a/src/amd/vulkan/radv_private.h
+++ b/src/amd/vulkan/radv_private.h
@@ -1127,6 +1127,7 @@ struct radv_pipeline {
 
 	bool                                         needs_data_cache;
 	bool					     need_indirect_descriptor_sets;
+	uint8_t                                      stage_mask;
 	struct radv_shader_variant *                 shaders[MESA_SHADER_STAGES];
 	struct radv_shader_variant *gs_copy_shader;
 	VkShaderStageFlags                           active_stages;
-- 
2.14.2