From 27c06dbcd3af468933170e9f8c738b53b53e0277 Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Mon, 22 Nov 2021 13:33:14 +0200
Subject: [PATCH] Set push constants only for enabled shader stages

---
 source/backends/vulkan/pipelinestate_backend.cpp | 2 +-
 source/backends/vulkan/program_backend.cpp       | 8 ++++++--
 source/backends/vulkan/program_backend.h         | 1 +
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/source/backends/vulkan/pipelinestate_backend.cpp b/source/backends/vulkan/pipelinestate_backend.cpp
index 444d5f66..4f4f5600 100644
--- a/source/backends/vulkan/pipelinestate_backend.cpp
+++ b/source/backends/vulkan/pipelinestate_backend.cpp
@@ -369,7 +369,7 @@ void VulkanPipelineState::apply(VkCommandBuffer command_buffer) const
 		if(first_block.used && first_block.binding==ReflectData::PUSH_CONSTANT)
 		{
 			const UniformBlock &pc_block = *first_block.block;
-			vk.CmdPushConstants(command_buffer, self.shprog->layout_handle, VK_SHADER_STAGE_ALL,
+			vk.CmdPushConstants(command_buffer, self.shprog->layout_handle, self.shprog->stage_flags,
 				pc_block.get_offset(), pc_block.get_data_size(), pc_block.get_data_pointer());
 		}
 	}
diff --git a/source/backends/vulkan/program_backend.cpp b/source/backends/vulkan/program_backend.cpp
index 61b5efae..60e96698 100644
--- a/source/backends/vulkan/program_backend.cpp
+++ b/source/backends/vulkan/program_backend.cpp
@@ -19,6 +19,7 @@ VulkanProgram::VulkanProgram():
 VulkanProgram::VulkanProgram(VulkanProgram &&other):
 	device(other.device),
 	n_stages(other.n_stages),
+	stage_flags(other.stage_flags),
 	creation_info(move(other.creation_info)),
 	desc_set_layout_handles(move(other.desc_set_layout_handles)),
 	layout_handle(other.layout_handle)
@@ -86,8 +87,11 @@ void VulkanProgram::add_spirv_stages(const SpirVModule &mod, const map<string, i
 	i = 0;
 	for(const SpirVModule::EntryPoint &e: entry_points)
 	{
+		unsigned stage_bit = get_vulkan_stage(e.stage);
+		stage_flags |= stage_bit;
+
 		stage_infos[i].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
-		stage_infos[i].stage = static_cast<VkShaderStageFlagBits>(get_vulkan_stage(e.stage));
+		stage_infos[i].stage = static_cast<VkShaderStageFlagBits>(stage_bit);
 		stage_infos[i].module = handle_cast<::VkShaderModule>(mod.handle);
 		strcpy(name_ptr, e.name.c_str());
 		stage_infos[i].pName = name_ptr;
@@ -142,7 +146,7 @@ void VulkanProgram::finalize_uniforms()
 	}
 
 	VkPushConstantRange push_const_range = { };
-	push_const_range.stageFlags = VK_SHADER_STAGE_ALL;
+	push_const_range.stageFlags = stage_flags;
 	push_const_range.offset = 0;
 	push_const_range.size = (push_const_block ? push_const_block->data_size : 0);
 
diff --git a/source/backends/vulkan/program_backend.h b/source/backends/vulkan/program_backend.h
index 3bdf3b73..c01a5361 100644
--- a/source/backends/vulkan/program_backend.h
+++ b/source/backends/vulkan/program_backend.h
@@ -22,6 +22,7 @@ protected:
 
 	Device &device;
 	unsigned n_stages = 0;
+	unsigned stage_flags = 0;
 	std::vector<char> creation_info;
 	std::vector<VkDescriptorSetLayout> desc_set_layout_handles;
 	VkPipelineLayout layout_handle = 0;
-- 
2.43.0