Uniquely identify sampler state

src/Vulkan/VkDevice.cpp

@@ -59,6 +59,44 @@ void Device::SamplingRoutineCache::updateConstCache()
 	cache.updateConstCache();
 }
 
+Device::SamplerIndexer::~SamplerIndexer()
+{
+	ASSERT(map.empty());
+}
+
+uint32_t Device::SamplerIndexer::index(const SamplerState &samplerState)
+{
+	std::lock_guard<std::mutex> lock(mutex);
+
+	auto it = map.find(samplerState);
+
+	if(it != map.end())
+	{
+		it->second.count++;
+		return it->second.id;
+	}
+
+	nextID++;
+
+	map.emplace(samplerState, Identifier{ nextID, 1 });
+
+	return nextID;
+}
+
+void Device::SamplerIndexer::remove(const SamplerState &samplerState)
+{
+	std::lock_guard<std::mutex> lock(mutex);
+
+	auto it = map.find(samplerState);
+	ASSERT(it != map.end());
+
+	auto count = --it->second.count;
+	if(count == 0)
+	{
+		map.erase(it);
+	}
+}
+
 Device::Device(const VkDeviceCreateInfo *pCreateInfo, void *mem, PhysicalDevice *physicalDevice, const VkPhysicalDeviceFeatures *enabledFeatures, const std::shared_ptr<marl::Scheduler> &scheduler)
     : physicalDevice(physicalDevice)
     , queues(reinterpret_cast<Queue *>(mem))
@@ -99,6 +137,7 @@ Device::Device(const VkDeviceCreateInfo *pCreateInfo, void *mem, PhysicalDevice
 	// FIXME (b/119409619): use an allocator here so we can control all memory allocations
 	blitter.reset(new sw::Blitter());
 	samplingRoutineCache.reset(new SamplingRoutineCache());
+	samplerIndexer.reset(new SamplerIndexer());
 
 #ifdef ENABLE_VK_DEBUGGER
 	static auto port = getenv("VK_DEBUGGER_PORT");
@@ -279,4 +318,14 @@ std::mutex &Device::getSamplingRoutineCacheMutex()
 	return samplingRoutineCacheMutex;
 }
 
+uint32_t Device::indexSampler(const SamplerState &samplerState)
+{
+	return samplerIndexer->index(samplerState);
+}
+
+void Device::removeSampler(const SamplerState &samplerState)
+{
+	samplerIndexer->remove(samplerState);
+}
+
 }  // namespace vk

src/Vulkan/VkDevice.hpp

@@ -16,8 +16,11 @@
 #define VK_DEVICE_HPP_
 
 #include "VkObject.hpp"
+#include "VkSampler.hpp"
 #include "Device/LRUCache.hpp"
 #include "Reactor/Routine.hpp"
+
+#include <map>
 #include <memory>
 #include <mutex>
 
@@ -98,6 +101,30 @@ public:
 	rr::Routine *findInConstCache(const SamplingRoutineCache::Key &key) const;
 	void updateSamplingRoutineConstCache();
 
+	class SamplerIndexer
+	{
+	public:
+		~SamplerIndexer();
+
+		uint32_t index(const SamplerState &samplerState);
+		void remove(const SamplerState &samplerState);
+
+	private:
+		struct Identifier
+		{
+			uint32_t id;
+			uint32_t count;  // Number of samplers sharing this state identifier.
+		};
+
+		std::map<SamplerState, Identifier> map;  // guarded by mutex
+		std::mutex mutex;
+
+		uint32_t nextID = 0;
+	};
+
+	uint32_t indexSampler(const SamplerState &samplerState);
+	void removeSampler(const SamplerState &samplerState);
+
 	std::shared_ptr<vk::dbg::Context> getDebuggerContext() const
 	{
 #ifdef ENABLE_VK_DEBUGGER
@@ -118,7 +145,9 @@ private:
 	typedef char ExtensionName[VK_MAX_EXTENSION_NAME_SIZE];
 	ExtensionName *extensions = nullptr;
 	const VkPhysicalDeviceFeatures enabledFeatures = {};
+
 	std::shared_ptr<marl::Scheduler> scheduler;
+	std::unique_ptr<SamplerIndexer> samplerIndexer;
 
 #ifdef ENABLE_VK_DEBUGGER
 	struct

src/Vulkan/VkSampler.cpp

@@ -16,8 +16,6 @@
 
 namespace vk {
 
-std::atomic<uint32_t> Sampler::nextID(1);
-
 SamplerState::SamplerState(const VkSamplerCreateInfo *pCreateInfo, const vk::SamplerYcbcrConversion *ycbcrConversion)
     : Memset(this, 0)
     , magFilter(pCreateInfo->magFilter)
@@ -44,8 +42,9 @@ SamplerState::SamplerState(const VkSamplerCreateInfo *pCreateInfo, const vk::Sam
 	}
 }
 
-Sampler::Sampler(const VkSamplerCreateInfo *pCreateInfo, void *mem, const vk::SamplerYcbcrConversion *ycbcrConversion)
-    : SamplerState(pCreateInfo, ycbcrConversion)
+Sampler::Sampler(const VkSamplerCreateInfo *pCreateInfo, void *mem, const SamplerState &samplerState, uint32_t samplerID)
+    : SamplerState(samplerState)
+    , id(samplerID)
 {
 }
 

src/Vulkan/VkSampler.hpp

@@ -58,17 +58,14 @@ struct SamplerState : sw::Memset<SamplerState>
 class Sampler : public Object<Sampler, VkSampler>, public SamplerState
 {
 public:
-	Sampler(const VkSamplerCreateInfo *pCreateInfo, void *mem, const vk::SamplerYcbcrConversion *ycbcrConversion);
+	Sampler(const VkSamplerCreateInfo *pCreateInfo, void *mem, const SamplerState &samplerState, uint32_t samplerID);
 
 	static size_t ComputeRequiredAllocationSize(const VkSamplerCreateInfo *pCreateInfo)
 	{
 		return 0;
 	}
 
-	const uint32_t id = nextID++;
-
-private:
-	static std::atomic<uint32_t> nextID;
+	const uint32_t id = 0;
 };
 
 class SamplerYcbcrConversion : public Object<SamplerYcbcrConversion, VkSamplerYcbcrConversion>

src/Vulkan/libVulkan.cpp

@@ -1857,7 +1857,18 @@ VKAPI_ATTR VkResult VKAPI_CALL vkCreateSampler(VkDevice device, const VkSamplerC
 		extensionCreateInfo = extensionCreateInfo->pNext;
 	}
 
-	return vk::Sampler::Create(pAllocator, pCreateInfo, pSampler, ycbcrConversion);
+	vk::SamplerState samplerState(pCreateInfo, ycbcrConversion);
+	uint32_t samplerID = vk::Cast(device)->indexSampler(samplerState);
+
+	VkResult result = vk::Sampler::Create(pAllocator, pCreateInfo, pSampler, samplerState, samplerID);
+
+	if(*pSampler == VK_NULL_HANDLE)
+	{
+		ASSERT(result != VK_SUCCESS);
+		vk::Cast(device)->removeSampler(samplerState);
+	}
+
+	return result;
 }
 
 VKAPI_ATTR void VKAPI_CALL vkDestroySampler(VkDevice device, VkSampler sampler, const VkAllocationCallbacks *pAllocator)
@@ -1865,7 +1876,12 @@ VKAPI_ATTR void VKAPI_CALL vkDestroySampler(VkDevice device, VkSampler sampler,
 	TRACE("(VkDevice device = %p, VkSampler sampler = %p, const VkAllocationCallbacks* pAllocator = %p)",
 	      device, static_cast<void *>(sampler), pAllocator);
 
-	vk::destroy(sampler, pAllocator);
+	if(sampler != VK_NULL_HANDLE)
+	{
+		vk::Cast(device)->removeSampler(*vk::Cast(sampler));
+
+		vk::destroy(sampler, pAllocator);
+	}
 }
 
 VKAPI_ATTR VkResult VKAPI_CALL vkCreateDescriptorSetLayout(VkDevice device, const VkDescriptorSetLayoutCreateInfo *pCreateInfo, const VkAllocationCallbacks *pAllocator, VkDescriptorSetLayout *pSetLayout)