Avoid recompiling identical SPIR-V code

src/Device/PixelProcessor.hpp

@@ -34,7 +34,7 @@ namespace sw
 		{
 			unsigned int computeHash();
 
-			int shaderID;
+			uint64_t shaderID;
 
 			VkCompareOp depthCompareMode;
 			bool depthWriteEnable;

src/Pipeline/SpirvShader.cpp

@@ -366,16 +366,16 @@ namespace sw
 
 	} // namespace SIMD
 
-	std::atomic<int> SpirvShader::serialCounter(1);    // Start at 1, 0 is invalid shader.
-
 	SpirvShader::SpirvShader(
-			VkPipelineShaderStageCreateInfo const *createInfo,
+			uint32_t codeSerialID,
+			VkShaderStageFlagBits pipelineStage,
+			const char *entryPointName,
 			InsnStore const &insns,
 			vk::RenderPass *renderPass,
 			uint32_t subpassIndex)
 				: insns{insns}, inputs{MAX_INTERFACE_COMPONENTS},
 				outputs{MAX_INTERFACE_COMPONENTS},
-				serialID{serialCounter++}, modes{}
+				codeSerialID(codeSerialID), modes{}
 	{
 		ASSERT(insns.size() > 0);
 
@@ -411,7 +411,7 @@ namespace sw
 				auto id = Object::ID(insn.word(2));
 				auto name = insn.string(3);
 				auto stage = executionModelToStage(executionModel);
-				if (stage == createInfo->stage && strcmp(name, createInfo->pName) == 0)
+				if (stage == pipelineStage && strcmp(name, entryPointName) == 0)
 				{
 					ASSERT_MSG(entryPointFunctionId == 0, "Duplicate entry point with name '%s' and stage %d", name, int(stage));
 					entryPointFunctionId = id;
@@ -982,7 +982,7 @@ namespace sw
 			}
 		}
 
-		ASSERT_MSG(entryPointFunctionId != 0, "Entry point '%s' not found", createInfo->pName);
+		ASSERT_MSG(entryPointFunctionId != 0, "Entry point '%s' not found", entryPointName);
 		AssignBlockFields();
 	}
 

src/Pipeline/SpirvShader.hpp

@@ -540,15 +540,19 @@ namespace sw
 
 		static_assert(sizeof(ImageInstruction) == sizeof(uint32_t), "ImageInstruction must be 32-bit");
 
-		int getSerialID() const
+		// This method is for retrieving an ID that uniquely identifies the
+		// shader entry point represented by this object.
+		uint64_t getSerialID() const
 		{
-			return serialID;
+			return  ((uint64_t)entryPointBlockId.value() << 32) | codeSerialID;
 		}
 
-		SpirvShader(VkPipelineShaderStageCreateInfo const *createInfo,
-					InsnStore const &insns,
-					vk::RenderPass *renderPass,
-					uint32_t subpassIndex);
+		SpirvShader(uint32_t codeSerialID,
+		            VkShaderStageFlagBits stage,
+		            const char *entryPointName,
+		            InsnStore const &insns,
+		            vk::RenderPass *renderPass,
+		            uint32_t subpassIndex);
 
 		struct Modes
 		{
@@ -740,8 +744,7 @@ namespace sw
 		}
 
 	private:
-		const int serialID;
-		static std::atomic<int> serialCounter;
+		const uint32_t codeSerialID;
 		Modes modes;
 		HandleMap<Type> types;
 		HandleMap<Object> defs;

src/Vulkan/VkPipeline.cpp

@@ -152,8 +152,7 @@ unsigned char getNumberOfChannels(VkFormat format)
 	return 0;
 }
 
-// preprocessSpirv applies and freezes specializations into constants, inlines
-// all functions and performs constant folding.
+// preprocessSpirv applies and freezes specializations into constants, and inlines all functions.
 std::vector<uint32_t> preprocessSpirv(
 		std::vector<uint32_t> const &code,
 		VkSpecializationInfo const *specializationInfo)
@@ -451,12 +450,16 @@ void GraphicsPipeline::compileShaders(const VkAllocationCallbacks* pAllocator, c
 			UNIMPLEMENTED("pStage->flags");
 		}
 
-		auto module = vk::Cast(pStage->module);
+		const ShaderModule *module = vk::Cast(pStage->module);
 		auto code = preprocessSpirv(module->getCode(), pStage->pSpecializationInfo);
 
+		// If the pipeline has specialization constants, assume they're unique and
+		// use a new serial ID so the shader gets recompiled.
+		uint32_t codeSerialID = (pStage->pSpecializationInfo ? ShaderModule::nextSerialID() : module->getSerialID());
+
 		// FIXME (b/119409619): use an allocator here so we can control all memory allocations
 		// TODO: also pass in any pipeline state which will affect shader compilation
-		auto spirvShader = new sw::SpirvShader{pStage, code, vk::Cast(pCreateInfo->renderPass), pCreateInfo->subpass};
+		auto spirvShader = new sw::SpirvShader(codeSerialID, pStage->stage, pStage->pName, code, vk::Cast(pCreateInfo->renderPass), pCreateInfo->subpass);
 
 		switch (pStage->stage)
 		{
@@ -542,16 +545,21 @@ size_t ComputePipeline::ComputeRequiredAllocationSize(const VkComputePipelineCre
 
 void ComputePipeline::compileShaders(const VkAllocationCallbacks* pAllocator, const VkComputePipelineCreateInfo* pCreateInfo)
 {
-	auto module = vk::Cast(pCreateInfo->stage.module);
+	auto &stage = pCreateInfo->stage;
+	const ShaderModule *module = vk::Cast(stage.module);
 
-	auto code = preprocessSpirv(module->getCode(), pCreateInfo->stage.pSpecializationInfo);
+	auto code = preprocessSpirv(module->getCode(), stage.pSpecializationInfo);
 
 	ASSERT_OR_RETURN(code.size() > 0);
 
 	ASSERT(shader == nullptr);
 
-	// FIXME(b/119409619): use allocator.
-	shader = new sw::SpirvShader(&pCreateInfo->stage, code, nullptr, 0);
+	// If the pipeline has specialization constants, assume they're unique and
+	// use a new serial ID so the shader gets recompiled.
+	uint32_t codeSerialID = (stage.pSpecializationInfo ? ShaderModule::nextSerialID() : module->getSerialID());
+
+	// TODO(b/119409619): use allocator.
+	shader = new sw::SpirvShader(codeSerialID, stage.stage, stage.pName, code, nullptr, 0);
 	vk::DescriptorSet::Bindings descriptorSets;  // FIXME(b/129523279): Delay code generation until invoke time.
 	program = new sw::ComputeProgram(shader, layout, descriptorSets);
 	program->generate();

src/Vulkan/VkShaderModule.cpp

@@ -19,7 +19,10 @@
 namespace vk
 {
 
-ShaderModule::ShaderModule(const VkShaderModuleCreateInfo* pCreateInfo, void* mem) : code(reinterpret_cast<uint32_t*>(mem))
+std::atomic<uint32_t> ShaderModule::serialCounter(1);    // Start at 1, 0 is invalid shader.
+
+ShaderModule::ShaderModule(const VkShaderModuleCreateInfo* pCreateInfo, void* mem)
+	: serialID(nextSerialID()), code(reinterpret_cast<uint32_t*>(mem))
 {
 	memcpy(code, pCreateInfo->pCode, pCreateInfo->codeSize);
 	wordCount = static_cast<uint32_t>(pCreateInfo->codeSize / sizeof(uint32_t));

src/Vulkan/VkShaderModule.hpp

@@ -16,6 +16,8 @@
 #define VK_SHADER_MODULE_HPP_
 
 #include "VkObject.hpp"
+
+#include <atomic>
 #include <vector>
 
 namespace rr
@@ -37,7 +39,13 @@ public:
 	// guts' operations, and this copy.
 	std::vector<uint32_t> getCode() const { return std::vector<uint32_t>{ code, code + wordCount };}
 
+	uint32_t getSerialID() const { return serialID; }
+	static uint32_t nextSerialID() { return serialCounter++; }
+
 private:
+	const uint32_t serialID;
+	static std::atomic<uint32_t> serialCounter;
+
 	uint32_t* code = nullptr;
 	uint32_t wordCount = 0;
 };