vulkan: move cached PSO ownership to shaders.

Fixes issues with state corruption when a shader is created with the same memory address as a recently deleted shader.

src/modules/graphics/vulkan/Graphics.cpp

@@ -2315,13 +2315,14 @@ void Graphics::prepareDraw(const VertexAttributes &attributes, const BufferBindi
 	if (!renderPassState.active)
 		startRenderPass();
 
-	usedShadersInFrame.insert((dynamic_cast<Shader*>(Shader::current)));
+	auto s = dynamic_cast<Shader*>(Shader::current);
+
+	usedShadersInFrame.insert(s);
 
 	GraphicsPipelineConfiguration configuration{};
 
 	configuration.renderPass = renderPassState.beginInfo.renderPass;
 	configuration.vertexAttributes = attributes;
-	configuration.shader = (Shader*)Shader::current;
 	configuration.wireFrame = states.back().wireframe;
 	configuration.blendState = states.back().blend;
 	configuration.colorChannelMask = states.back().colorMask;
@@ -2341,11 +2342,15 @@ void Graphics::prepareDraw(const VertexAttributes &attributes, const BufferBindi
 		configuration.dynamicState.cullmode = cullmode;
 	}
 
-	configuration.shader->setMainTex(texture);
-
-	ensureGraphicsPipelineConfiguration(configuration);
+	VkPipeline pipeline = s->getCachedGraphicsPipeline(this, configuration);
+	if (pipeline != renderPassState.pipeline)
+	{
+		vkCmdBindPipeline(commandBuffers.at(currentFrame), VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
+		renderPassState.pipeline = pipeline;
+	}
 
-	configuration.shader->cmdPushDescriptorSets(commandBuffers.at(currentFrame), VK_PIPELINE_BIND_POINT_GRAPHICS);
+	s->setMainTex(texture);
+	s->cmdPushDescriptorSets(commandBuffers.at(currentFrame), VK_PIPELINE_BIND_POINT_GRAPHICS);
 
 	VkBuffer vkbuffers[BufferBindings::MAX];
 	VkDeviceSize vkoffsets[BufferBindings::MAX];
@@ -2640,7 +2645,6 @@ void Graphics::cleanupUnusedObjects()
 {
 	eraseUnusedObjects(renderPasses, renderPassUsages, vkDestroyRenderPass, device);
 	eraseUnusedObjects(framebuffers, framebufferUsages, vkDestroyFramebuffer, device);
-	eraseUnusedObjects(graphicsPipelines, pipelineUsages, vkDestroyPipeline, device);
 }
 
 void Graphics::requestSwapchainRecreation()
@@ -2670,17 +2674,17 @@ VkSampler Graphics::getCachedSampler(const SamplerState &samplerState)
 	}
 }
 
-VkPipeline Graphics::createGraphicsPipeline(GraphicsPipelineConfiguration &configuration)
+VkPipeline Graphics::createGraphicsPipeline(Shader *shader, const GraphicsPipelineConfiguration &configuration)
 {
 	VkGraphicsPipelineCreateInfo pipelineInfo{};
 	pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
 
-	auto &shaderStages = configuration.shader->getShaderStages();
+	auto &shaderStages = shader->getShaderStages();
 
 	std::vector<VkVertexInputBindingDescription> bindingDescriptions;
 	std::vector<VkVertexInputAttributeDescription> attributeDescriptions;
 
-	createVulkanVertexFormat(configuration.shader, configuration.vertexAttributes, bindingDescriptions, attributeDescriptions);
+	createVulkanVertexFormat(shader, configuration.vertexAttributes, bindingDescriptions, attributeDescriptions);
 
 	VkPipelineVertexInputStateCreateInfo vertexInputInfo{};
 	vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
@@ -2812,7 +2816,7 @@ VkPipeline Graphics::createGraphicsPipeline(GraphicsPipelineConfiguration &confi
 	pipelineInfo.pMultisampleState = &multisampling;
 	pipelineInfo.pColorBlendState = &colorBlending;
 	pipelineInfo.pDynamicState = &dynamicState;
-	pipelineInfo.layout = configuration.shader->getGraphicsPipelineLayout();
+	pipelineInfo.layout = shader->getGraphicsPipelineLayout();
 	pipelineInfo.subpass = 0;
 	pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
 	pipelineInfo.basePipelineIndex = -1;
@@ -2824,28 +2828,6 @@ VkPipeline Graphics::createGraphicsPipeline(GraphicsPipelineConfiguration &confi
 	return graphicsPipeline;
 }
 
-void Graphics::ensureGraphicsPipelineConfiguration(GraphicsPipelineConfiguration &configuration)
-{
-	auto it = graphicsPipelines.find(configuration);
-	if (it != graphicsPipelines.end())
-	{
-		if (it->second != renderPassState.pipeline)
-		{
-			vkCmdBindPipeline(commandBuffers.at(currentFrame), VK_PIPELINE_BIND_POINT_GRAPHICS, it->second);
-			renderPassState.pipeline = it->second;
-			pipelineUsages[it->second] = true;
-		}
-	}
-	else
-	{
-		VkPipeline pipeline = createGraphicsPipeline(configuration);
-		graphicsPipelines.insert({configuration, pipeline});
-		vkCmdBindPipeline(commandBuffers.at(currentFrame), VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
-		renderPassState.pipeline = pipeline;
-		pipelineUsages[pipeline] = true;
-	}
-}
-
 VkSampleCountFlagBits Graphics::getMsaaCount(int requestedMsaa) const
 {
 	VkPhysicalDeviceProperties physicalDeviceProperties;
@@ -3112,10 +3094,6 @@ void Graphics::cleanup()
 	for (const auto &entry : framebuffers)
 		vkDestroyFramebuffer(device, entry.second, nullptr);
 	framebuffers.clear();
-		
-	for (auto const &p : graphicsPipelines)
-		vkDestroyPipeline(device, p.second, nullptr);
-	graphicsPipelines.clear();
 
 	vkDestroyCommandPool(device, commandPool, nullptr);
 	vkDestroyPipelineCache(device, pipelineCache, nullptr);

src/modules/graphics/vulkan/Graphics.h

@@ -170,46 +170,6 @@ struct OptionalDeviceExtensions
 	bool spirv14 = false;
 };
 
-struct GraphicsPipelineConfiguration
-{
-	VkRenderPass renderPass;
-	VertexAttributes vertexAttributes;
-	Shader *shader = nullptr;
-	bool wireFrame;
-	BlendState blendState;
-	ColorChannelMask colorChannelMask;
-	VkSampleCountFlagBits msaaSamples;
-	uint32_t numColorAttachments;
-	PrimitiveType primitiveType;
-
-	struct DynamicState
-	{
-		CullMode cullmode = CULL_NONE;
-		Winding winding = WINDING_MAX_ENUM;
-		StencilAction stencilAction = STENCIL_MAX_ENUM;
-		CompareMode stencilCompare = COMPARE_MAX_ENUM;
-		DepthState depthState{};
-	} dynamicState;
-
-	GraphicsPipelineConfiguration()
-	{
-		memset(this, 0, sizeof(GraphicsPipelineConfiguration));
-	}
-
-	bool operator==(const GraphicsPipelineConfiguration &other) const
-	{
-		return memcmp(this, &other, sizeof(GraphicsPipelineConfiguration)) == 0;
-	}
-};
-
-struct GraphicsPipelineConfigurationHasher
-{
-	size_t operator() (const GraphicsPipelineConfiguration &configuration) const
-	{
-		return XXH32(&configuration, sizeof(GraphicsPipelineConfiguration), 0);
-	}
-};
-
 struct QueueFamilyIndices
 {
 	Optional<uint32_t> graphicsFamily;
@@ -315,6 +275,8 @@ public:
 	int getVsync() const;
 	void mapLocalUniformData(void *data, size_t size, VkDescriptorBufferInfo &bufferInfo);
 
+	VkPipeline createGraphicsPipeline(Shader *shader, const GraphicsPipelineConfiguration &configuration);
+
 	uint32 getDeviceApiVersion() const { return deviceApiVersion; }
 
 protected:
@@ -349,7 +311,6 @@ private:
 	void createDefaultShaders();
 	VkRenderPass createRenderPass(RenderPassConfiguration &configuration);
 	VkRenderPass getRenderPass(RenderPassConfiguration &configuration);
-	VkPipeline createGraphicsPipeline(GraphicsPipelineConfiguration &configuration);
 	void createColorResources();
 	VkFormat findSupportedFormat(const std::vector<VkFormat> &candidates, VkImageTiling tiling, VkFormatFeatureFlags features);
 	VkFormat findDepthFormat();
@@ -364,7 +325,6 @@ private:
 	void beginFrame();
 	void startRecordingGraphicsCommands();
 	void endRecordingGraphicsCommands();
-	void ensureGraphicsPipelineConfiguration(GraphicsPipelineConfiguration &configuration);
 	void createVulkanVertexFormat(
 		Shader *shader,
 		const VertexAttributes &attributes, 
@@ -412,10 +372,8 @@ private:
 	VkPipelineCache pipelineCache = VK_NULL_HANDLE;
 	std::unordered_map<RenderPassConfiguration, VkRenderPass, RenderPassConfigurationHasher> renderPasses;
 	std::unordered_map<FramebufferConfiguration, VkFramebuffer, FramebufferConfigurationHasher> framebuffers;
-	std::unordered_map<GraphicsPipelineConfiguration, VkPipeline, GraphicsPipelineConfigurationHasher> graphicsPipelines;
 	std::unordered_map<VkRenderPass, bool> renderPassUsages;
 	std::unordered_map<VkFramebuffer, bool> framebufferUsages;
-	std::unordered_map<VkPipeline, bool> pipelineUsages;
 	std::unordered_map<uint64, VkSampler> samplers;
 	VkCommandPool commandPool = VK_NULL_HANDLE;
 	std::vector<VkCommandBuffer> commandBuffers;

src/modules/graphics/vulkan/Shader.cpp

@@ -167,6 +167,7 @@ static bool usesLocalUniformData(const graphics::Shader::UniformInfo *info)
 
 Shader::Shader(StrongRef<love::graphics::ShaderStage> stages[], const CompileOptions &options)
 	: graphics::Shader(stages, options)
+	, builtinUniformInfo()
 {
 	auto gfx = Module::getInstance<Graphics>(Module::ModuleType::M_GRAPHICS);
 	vgfx = dynamic_cast<Graphics*>(gfx);
@@ -199,7 +200,8 @@ void Shader::unloadVolatile()
 	if (shaderModules.empty())
 		return;
 
-	vgfx->queueCleanUp([shaderModules = std::move(shaderModules), device = device, descriptorSetLayout = descriptorSetLayout, pipelineLayout = pipelineLayout, descriptorPools = descriptorPools, computePipeline = computePipeline](){
+	vgfx->queueCleanUp([shaderModules = std::move(shaderModules), device = device, descriptorSetLayout = descriptorSetLayout, pipelineLayout = pipelineLayout,
+		descriptorPools = descriptorPools, computePipeline = computePipeline, graphicsPipelines = std::move(graphicsPipelines)](){
 		for (const auto &pools : descriptorPools)
 		{
 			for (const auto pool : pools)
@@ -211,6 +213,8 @@ void Shader::unloadVolatile()
 		vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
 		if (computePipeline != VK_NULL_HANDLE)
 			vkDestroyPipeline(device, computePipeline, nullptr);
+		for (const auto& kvp : graphicsPipelines)
+			vkDestroyPipeline(device, kvp.second, nullptr);
 	});
 
 	shaderModules.clear();
@@ -1199,6 +1203,18 @@ VkDescriptorSet Shader::allocateDescriptorSet()
 	}
 }
 
+VkPipeline Shader::getCachedGraphicsPipeline(Graphics *vgfx, const GraphicsPipelineConfiguration &configuration)
+{
+	auto it = graphicsPipelines.find(configuration);
+	if (it != graphicsPipelines.end())
+		return it->second;
+
+	VkPipeline pipeline = vgfx->createGraphicsPipeline(this, configuration);
+	graphicsPipelines.insert({ configuration, pipeline });
+	
+	return pipeline;
+}
+
 } // vulkan
 } // graphics
 } // love

src/modules/graphics/vulkan/Shader.h

@@ -29,6 +29,7 @@
 // Libraries
 #include "VulkanWrapper.h"
 #include "libraries/spirv_cross/spirv_reflect.hpp"
+#include "libraries/xxHash/xxhash.h"
 
 // C++
 #include <map>
@@ -45,6 +46,45 @@ namespace graphics
 namespace vulkan
 {
 
+struct GraphicsPipelineConfiguration
+{
+	VkRenderPass renderPass;
+	VertexAttributes vertexAttributes;
+	bool wireFrame;
+	BlendState blendState;
+	ColorChannelMask colorChannelMask;
+	VkSampleCountFlagBits msaaSamples;
+	uint32_t numColorAttachments;
+	PrimitiveType primitiveType;
+
+	struct DynamicState
+	{
+		CullMode cullmode = CULL_NONE;
+		Winding winding = WINDING_MAX_ENUM;
+		StencilAction stencilAction = STENCIL_MAX_ENUM;
+		CompareMode stencilCompare = COMPARE_MAX_ENUM;
+		DepthState depthState{};
+	} dynamicState;
+
+	GraphicsPipelineConfiguration()
+	{
+		memset(this, 0, sizeof(GraphicsPipelineConfiguration));
+	}
+
+	bool operator==(const GraphicsPipelineConfiguration &other) const
+	{
+		return memcmp(this, &other, sizeof(GraphicsPipelineConfiguration)) == 0;
+	}
+};
+
+struct GraphicsPipelineConfigurationHasher
+{
+	size_t operator() (const GraphicsPipelineConfiguration &configuration) const
+	{
+		return XXH32(&configuration, sizeof(GraphicsPipelineConfiguration), 0);
+	}
+};
+
 class Graphics;
 
 class Shader final
@@ -95,6 +135,8 @@ public:
 
 	void setMainTex(graphics::Texture *texture);
 
+	VkPipeline getCachedGraphicsPipeline(Graphics *vgfx, const GraphicsPipelineConfiguration &configuration);
+
 private:
 	void compileShaders();
 	void createDescriptorSetLayout();
@@ -107,10 +149,10 @@ private:
 	void setTextureDescriptor(const UniformInfo *info, love::graphics::Texture *texture, int index);
 	void setBufferDescriptor(const UniformInfo *info, love::graphics::Buffer *buffer, int index);
 
-	VkPipeline computePipeline;
+	VkPipeline computePipeline = VK_NULL_HANDLE;
 
-	VkDescriptorSetLayout descriptorSetLayout;
-	VkPipelineLayout pipelineLayout;
+	VkDescriptorSetLayout descriptorSetLayout = VK_NULL_HANDLE;
+	VkPipelineLayout pipelineLayout = VK_NULL_HANDLE;
 	std::vector<VkDescriptorPoolSize> descriptorPoolSizes;
 
 	std::vector<std::vector<VkDescriptorPool>> descriptorPools;
@@ -124,7 +166,7 @@ private:
 	std::vector<VkShaderModule> shaderModules;
 
 	Graphics *vgfx = nullptr;
-	VkDevice device;
+	VkDevice device = VK_NULL_HANDLE;
 
 	bool isCompute = false;
 	bool resourceDescriptorsDirty = false;
@@ -135,13 +177,15 @@ private:
 	std::unique_ptr<StreamBuffer> uniformBufferObjectBuffer;
 	std::vector<uint8> localUniformData;
 	std::vector<uint8> localUniformStagingData;
-	uint32_t localUniformLocation;
+	uint32_t localUniformLocation = 0;
 	OptionalInt builtinUniformDataOffset;
 
 	std::unordered_map<std::string, AttributeInfo> attributes;
 
-	uint32_t currentFrame;
-	uint32_t currentDescriptorPool;
+	std::unordered_map<GraphicsPipelineConfiguration, VkPipeline, GraphicsPipelineConfigurationHasher> graphicsPipelines;
+
+	uint32_t currentFrame = 0;
+	uint32_t currentDescriptorPool = 0;
 };
 
 }