Prevent accessing deleted ImageView objects

src/Device/Renderer.cpp

@@ -245,6 +245,7 @@ void Renderer::draw(const sw::Context *context, VkIndexType indexType, unsigned
 	}
 
 	DrawData *data = draw->data;
+	draw->device = device;
 	draw->occlusionQuery = occlusionQuery;
 	draw->batchDataPool = &batchDataPool;
 	draw->numPrimitives = count;
@@ -390,7 +391,7 @@ void Renderer::draw(const sw::Context *context, VkIndexType indexType, unsigned
 
 	draw->events = events;
 
-	vk::DescriptorSet::PrepareForSampling(draw->descriptorSetObjects, draw->pipelineLayout);
+	vk::DescriptorSet::PrepareForSampling(draw->descriptorSetObjects, draw->pipelineLayout, device);
 
 	DrawCall::run(draw, &drawTickets, clusterQueues);
 }
@@ -439,7 +440,7 @@ void DrawCall::teardown()
 
 	if(containsImageWrite)
 	{
-		vk::DescriptorSet::ContentsChanged(descriptorSetObjects, pipelineLayout);
+		vk::DescriptorSet::ContentsChanged(descriptorSetObjects, pipelineLayout, device);
 	}
 }
 

src/Device/Renderer.hpp

@@ -164,6 +164,7 @@ struct DrawCall
 	SetupFunction setupPrimitives;
 	SetupProcessor::State setupState;
 
+	vk::Device *device;
 	vk::ImageView *renderTarget[RENDERTARGETS];
 	vk::ImageView *depthBuffer;
 	vk::ImageView *stencilBuffer;

src/Pipeline/ComputeProgram.cpp

@@ -37,8 +37,9 @@ enum
 
 namespace sw {
 
-ComputeProgram::ComputeProgram(SpirvShader const *shader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets)
-    : shader(shader)
+ComputeProgram::ComputeProgram(vk::Device *device, SpirvShader const *shader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets)
+    : device(device)
+    , shader(shader)
     , pipelineLayout(pipelineLayout)
     , descriptorSets(descriptorSets)
 {
@@ -301,7 +302,7 @@ void ComputeProgram::run(
 
 	if(shader->containsImageWrite())
 	{
-		vk::DescriptorSet::ContentsChanged(descriptorSetObjects, pipelineLayout);
+		vk::DescriptorSet::ContentsChanged(descriptorSetObjects, pipelineLayout, device);
 	}
 }
 

src/Pipeline/ComputeProgram.hpp

@@ -24,8 +24,9 @@
 #include <functional>
 
 namespace vk {
+class Device;
 class PipelineLayout;
-}
+}  // namespace vk
 
 namespace sw {
 
@@ -45,7 +46,7 @@ class ComputeProgram : public Coroutine<SpirvShader::YieldResult(
                            int32_t subgroupCount)>
 {
 public:
-	ComputeProgram(SpirvShader const *spirvShader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets);
+	ComputeProgram(vk::Device *device, SpirvShader const *spirvShader, vk::PipelineLayout const *pipelineLayout, const vk::DescriptorSet::Bindings &descriptorSets);
 
 	virtual ~ComputeProgram();
 
@@ -79,6 +80,7 @@ protected:
 		const Constants *constants;
 	};
 
+	vk::Device *const device;
 	SpirvShader const *const shader;
 	vk::PipelineLayout const *const pipelineLayout;
 	const vk::DescriptorSet::Bindings &descriptorSets;

src/Vulkan/VkDescriptorSet.cpp

@@ -13,12 +13,13 @@
 // limitations under the License.
 
 #include "VkDescriptorSet.hpp"
+#include "VkDevice.hpp"
 #include "VkImageView.hpp"
 #include "VkPipelineLayout.hpp"
 
 namespace vk {
 
-void DescriptorSet::ParseDescriptors(const Array &descriptorSets, const PipelineLayout *layout, NotificationType notificationType)
+void DescriptorSet::ParseDescriptors(const Array &descriptorSets, const PipelineLayout *layout, Device *device, NotificationType notificationType)
 {
 	if(layout)
 	{
@@ -62,11 +63,11 @@ void DescriptorSet::ParseDescriptors(const Array &descriptorSets, const Pipeline
 					{
 						if(notificationType == PREPARE_FOR_SAMPLING)
 						{
-							memoryOwner->prepareForSampling();
+							device->prepareForSampling(memoryOwner);
 						}
 						else if((notificationType == CONTENTS_CHANGED) && (type == VK_DESCRIPTOR_TYPE_STORAGE_IMAGE))
 						{
-							memoryOwner->contentsChanged();
+							device->contentsChanged(memoryOwner);
 						}
 					}
 					descriptorMemory += descriptorSize;
@@ -76,14 +77,14 @@ void DescriptorSet::ParseDescriptors(const Array &descriptorSets, const Pipeline
 	}
 }
 
-void DescriptorSet::ContentsChanged(const Array &descriptorSets, const PipelineLayout *layout)
+void DescriptorSet::ContentsChanged(const Array &descriptorSets, const PipelineLayout *layout, Device *device)
 {
-	ParseDescriptors(descriptorSets, layout, CONTENTS_CHANGED);
+	ParseDescriptors(descriptorSets, layout, device, CONTENTS_CHANGED);
 }
 
-void DescriptorSet::PrepareForSampling(const Array &descriptorSets, const PipelineLayout *layout)
+void DescriptorSet::PrepareForSampling(const Array &descriptorSets, const PipelineLayout *layout, Device *device)
 {
-	ParseDescriptors(descriptorSets, layout, PREPARE_FOR_SAMPLING);
+	ParseDescriptors(descriptorSets, layout, device, PREPARE_FOR_SAMPLING);
 }
 
 }  // namespace vk
\ No newline at end of file

src/Vulkan/VkDescriptorSet.hpp

@@ -25,6 +25,7 @@
 namespace vk {
 
 class DescriptorSetLayout;
+class Device;
 class PipelineLayout;
 
 struct alignas(16) DescriptorSetHeader
@@ -40,8 +41,8 @@ public:
 	using Bindings = std::array<uint8_t *, vk::MAX_BOUND_DESCRIPTOR_SETS>;
 	using DynamicOffsets = std::array<uint32_t, vk::MAX_DESCRIPTOR_SET_COMBINED_BUFFERS_DYNAMIC>;
 
-	static void ContentsChanged(const Array &descriptorSets, const PipelineLayout *layout);
-	static void PrepareForSampling(const Array &descriptorSets, const PipelineLayout *layout);
+	static void ContentsChanged(const Array &descriptorSets, const PipelineLayout *layout, Device *device);
+	static void PrepareForSampling(const Array &descriptorSets, const PipelineLayout *layout, Device *device);
 
 	DescriptorSetHeader header;
 	alignas(16) uint8_t data[1];
@@ -52,7 +53,7 @@ private:
 		CONTENTS_CHANGED,
 		PREPARE_FOR_SAMPLING
 	};
-	static void ParseDescriptors(const Array &descriptorSets, const PipelineLayout *layout, NotificationType notificationType);
+	static void ParseDescriptors(const Array &descriptorSets, const PipelineLayout *layout, Device *device, NotificationType notificationType);
 };
 
 inline DescriptorSet *Cast(VkDescriptorSet object)

src/Vulkan/VkDevice.cpp

@@ -325,4 +325,54 @@ VkResult Device::setDebugUtilsObjectTag(const VkDebugUtilsObjectTagInfoEXT *pTag
 	return VK_SUCCESS;
 }
 
+void Device::registerImageView(ImageView *imageView)
+{
+	if(imageView != nullptr)
+	{
+		marl::lock lock(imageViewSetMutex);
+		imageViewSet.insert(imageView);
+	}
+}
+
+void Device::unregisterImageView(ImageView *imageView)
+{
+	if(imageView != nullptr)
+	{
+		marl::lock lock(imageViewSetMutex);
+		auto it = imageViewSet.find(imageView);
+		if(it != imageViewSet.end())
+		{
+			imageViewSet.erase(it);
+		}
+	}
+}
+
+void Device::prepareForSampling(ImageView *imageView)
+{
+	if(imageView != nullptr)
+	{
+		marl::lock lock(imageViewSetMutex);
+
+		auto it = imageViewSet.find(imageView);
+		if(it != imageViewSet.end())
+		{
+			imageView->prepareForSampling();
+		}
+	}
+}
+
+void Device::contentsChanged(ImageView *imageView)
+{
+	if(imageView != nullptr)
+	{
+		marl::lock lock(imageViewSetMutex);
+
+		auto it = imageViewSet.find(imageView);
+		if(it != imageViewSet.end())
+		{
+			imageView->contentsChanged();
+		}
+	}
+}
+
 }  // namespace vk

src/Vulkan/VkDevice.hpp

@@ -15,7 +15,7 @@
 #ifndef VK_DEVICE_HPP_
 #define VK_DEVICE_HPP_
 
-#include "VkObject.hpp"
+#include "VkImageView.hpp"
 #include "VkSampler.hpp"
 #include "Reactor/Routine.hpp"
 #include "System/LRUCache.hpp"
@@ -26,6 +26,7 @@
 #include <map>
 #include <memory>
 #include <unordered_map>
+#include <unordered_set>
 
 namespace marl {
 class Scheduler;
@@ -67,6 +68,11 @@ public:
 	const VkPhysicalDeviceFeatures &getEnabledFeatures() const { return enabledFeatures; }
 	sw::Blitter *getBlitter() const { return blitter.get(); }
 
+	void registerImageView(ImageView *imageView);
+	void unregisterImageView(ImageView *imageView);
+	void prepareForSampling(ImageView *imageView);
+	void contentsChanged(ImageView *imageView);
+
 	class SamplingRoutineCache
 	{
 	public:
@@ -177,6 +183,9 @@ private:
 	std::unique_ptr<SamplingRoutineCache> samplingRoutineCache;
 	std::unique_ptr<SamplerIndexer> samplerIndexer;
 
+	marl::mutex imageViewSetMutex;
+	std::unordered_set<ImageView *> imageViewSet GUARDED_BY(imageViewSetMutex);
+
 #ifdef ENABLE_VK_DEBUGGER
 	struct
 	{

src/Vulkan/VkPipeline.cpp

@@ -122,13 +122,13 @@ std::shared_ptr<sw::SpirvShader> createShader(
 	                                         code, key.getRenderPass(), key.getSubpassIndex(), robustBufferAccess, dbgctx);
 }
 
-std::shared_ptr<sw::ComputeProgram> createProgram(const vk::PipelineCache::ComputeProgramKey &key)
+std::shared_ptr<sw::ComputeProgram> createProgram(vk::Device *device, const vk::PipelineCache::ComputeProgramKey &key)
 {
 	MARL_SCOPED_EVENT("createProgram");
 
 	vk::DescriptorSet::Bindings descriptorSets;  // FIXME(b/129523279): Delay code generation until invoke time.
 	// TODO(b/119409619): use allocator.
-	auto program = std::make_shared<sw::ComputeProgram>(key.getShader(), key.getLayout(), descriptorSets);
+	auto program = std::make_shared<sw::ComputeProgram>(device, key.getShader(), key.getLayout(), descriptorSets);
 	program->generate();
 	program->finalize();
 	return program;
@@ -138,7 +138,7 @@ std::shared_ptr<sw::ComputeProgram> createProgram(const vk::PipelineCache::Compu
 
 namespace vk {
 
-Pipeline::Pipeline(PipelineLayout *layout, const Device *device)
+Pipeline::Pipeline(PipelineLayout *layout, Device *device)
     : layout(layout)
     , device(device)
     , robustBufferAccess(device->getEnabledFeatures().robustBufferAccess)
@@ -153,7 +153,7 @@ void Pipeline::destroy(const VkAllocationCallbacks *pAllocator)
 	vk::release(static_cast<VkPipelineLayout>(*layout), pAllocator);
 }
 
-GraphicsPipeline::GraphicsPipeline(const VkGraphicsPipelineCreateInfo *pCreateInfo, void *mem, const Device *device)
+GraphicsPipeline::GraphicsPipeline(const VkGraphicsPipelineCreateInfo *pCreateInfo, void *mem, Device *device)
     : Pipeline(vk::Cast(pCreateInfo->layout), device)
 {
 	context.robustBufferAccess = robustBufferAccess;
@@ -581,7 +581,7 @@ bool GraphicsPipeline::hasDynamicState(VkDynamicState dynamicState) const
 	return (dynamicStateFlags & (1 << dynamicState)) != 0;
 }
 
-ComputePipeline::ComputePipeline(const VkComputePipelineCreateInfo *pCreateInfo, void *mem, const Device *device)
+ComputePipeline::ComputePipeline(const VkComputePipelineCreateInfo *pCreateInfo, void *mem, Device *device)
     : Pipeline(vk::Cast(pCreateInfo->layout), device)
 {
 }
@@ -615,14 +615,14 @@ void ComputePipeline::compileShaders(const VkAllocationCallbacks *pAllocator, co
 
 		const PipelineCache::ComputeProgramKey programKey(shader.get(), layout);
 		program = pPipelineCache->getOrCreateComputeProgram(programKey, [&] {
-			return createProgram(programKey);
+			return createProgram(device, programKey);
 		});
 	}
 	else
 	{
 		shader = createShader(shaderKey, module, robustBufferAccess, device->getDebuggerContext());
 		const PipelineCache::ComputeProgramKey programKey(shader.get(), layout);
-		program = createProgram(programKey);
+		program = createProgram(device, programKey);
 	}
 }
 

src/Vulkan/VkPipeline.hpp

@@ -42,7 +42,7 @@ class Device;
 class Pipeline
 {
 public:
-	Pipeline(PipelineLayout *layout, const Device *device);
+	Pipeline(PipelineLayout *layout, Device *device);
 	virtual ~Pipeline() = default;
 
 	operator VkPipeline()
@@ -69,7 +69,7 @@ public:
 
 protected:
 	PipelineLayout *layout = nullptr;
-	Device const *const device;
+	Device *const device;
 
 	const bool robustBufferAccess = true;
 };
@@ -79,7 +79,7 @@ class GraphicsPipeline : public Pipeline, public ObjectBase<GraphicsPipeline, Vk
 public:
 	GraphicsPipeline(const VkGraphicsPipelineCreateInfo *pCreateInfo,
 	                 void *mem,
-	                 const Device *device);
+	                 Device *device);
 	virtual ~GraphicsPipeline() = default;
 
 	void destroyPipeline(const VkAllocationCallbacks *pAllocator) override;
@@ -120,7 +120,7 @@ private:
 class ComputePipeline : public Pipeline, public ObjectBase<ComputePipeline, VkPipeline>
 {
 public:
-	ComputePipeline(const VkComputePipelineCreateInfo *pCreateInfo, void *mem, const Device *device);
+	ComputePipeline(const VkComputePipelineCreateInfo *pCreateInfo, void *mem, Device *device);
 	virtual ~ComputePipeline() = default;
 
 	void destroyPipeline(const VkAllocationCallbacks *pAllocator) override;

src/Vulkan/libVulkan.cpp

@@ -1771,7 +1771,13 @@ VKAPI_ATTR VkResult VKAPI_CALL vkCreateImageView(VkDevice device, const VkImageV
 		extensionCreateInfo = extensionCreateInfo->pNext;
 	}
 
-	return vk::ImageView::Create(pAllocator, pCreateInfo, pView, ycbcrConversion);
+	VkResult result = vk::ImageView::Create(pAllocator, pCreateInfo, pView, ycbcrConversion);
+	if(result == VK_SUCCESS)
+	{
+		vk::Cast(device)->registerImageView(vk::Cast(*pView));
+	}
+
+	return result;
 }
 
 VKAPI_ATTR void VKAPI_CALL vkDestroyImageView(VkDevice device, VkImageView imageView, const VkAllocationCallbacks *pAllocator)
@@ -1779,6 +1785,7 @@ VKAPI_ATTR void VKAPI_CALL vkDestroyImageView(VkDevice device, VkImageView image
 	TRACE("(VkDevice device = %p, VkImageView imageView = %p, const VkAllocationCallbacks* pAllocator = %p)",
 	      device, static_cast<void *>(imageView), pAllocator);
 
+	vk::Cast(device)->unregisterImageView(vk::Cast(imageView));
 	vk::destroy(imageView, pAllocator);
 }