Descriptor sets update mechanism

src/Vulkan/VkDescriptorPool.cpp

@@ -37,7 +37,9 @@ size_t DescriptorPool::ComputeRequiredAllocationSize(const VkDescriptorPoolCreat
 
 	for(uint32_t i = 0; i < pCreateInfo->poolSizeCount; i++)
 	{
-		size += pCreateInfo->pPoolSizes[i].descriptorCount * DescriptorSetLayout::GetDescriptorSize(pCreateInfo->pPoolSizes[i].type);
+		size += pCreateInfo->pPoolSizes[i].descriptorCount *
+		        (sizeof(DescriptorSetLayout*) +
+		         DescriptorSetLayout::GetDescriptorSize(pCreateInfo->pPoolSizes[i].type));
 	}
 
 	return size;
@@ -52,7 +54,15 @@ VkResult DescriptorPool::allocateSets(uint32_t descriptorSetCount, const VkDescr
 		layoutSizes[i] = Cast(pSetLayouts[i])->getSize();
 	}
 
-	return allocateSets(&(layoutSizes[0]), descriptorSetCount, pDescriptorSets);
+	VkResult result = allocateSets(&(layoutSizes[0]), descriptorSetCount, pDescriptorSets);
+	if(result == VK_SUCCESS)
+	{
+		for(uint32_t i = 0; i < descriptorSetCount; i++)
+		{
+			Cast(pSetLayouts[i])->initialize(pDescriptorSets[i]);
+		}
+	}
+	return result;
 }
 
 VkDescriptorSet DescriptorPool::findAvailableMemory(size_t size)

src/Vulkan/VkDescriptorSetLayout.cpp

@@ -13,11 +13,23 @@
 // limitations under the License.
 
 #include "VkDescriptorSetLayout.hpp"
+#include <algorithm>
 #include <cstring>
 
 namespace
 {
 
+struct DescriptorSet
+{
+	vk::DescriptorSetLayout* layout;
+	uint8_t data[];
+};
+
+static inline DescriptorSet* Cast(VkDescriptorSet object)
+{
+	return reinterpret_cast<DescriptorSet*>(object);
+}
+
 static bool UsesImmutableSamplers(const VkDescriptorSetLayoutBinding& binding)
 {
 	return (((binding.descriptorType == VK_DESCRIPTOR_TYPE_SAMPLER) ||
@@ -33,16 +45,19 @@ namespace vk
 DescriptorSetLayout::DescriptorSetLayout(const VkDescriptorSetLayoutCreateInfo* pCreateInfo, void* mem) :
 	flags(pCreateInfo->flags), bindingCount(pCreateInfo->bindingCount), bindings(reinterpret_cast<VkDescriptorSetLayoutBinding*>(mem))
 {
-	char* host_memory = static_cast<char*>(mem) + bindingCount * sizeof(VkDescriptorSetLayoutBinding);
+	uint8_t* hostMemory = static_cast<uint8_t*>(mem) + bindingCount * sizeof(VkDescriptorSetLayoutBinding);
+	bindingOffsets = reinterpret_cast<size_t*>(hostMemory);
+	hostMemory += bindingCount * sizeof(size_t);
 
+	size_t offset = 0;
 	for(uint32_t i = 0; i < bindingCount; i++)
 	{
 		bindings[i] = pCreateInfo->pBindings[i];
 		if(UsesImmutableSamplers(bindings[i]))
 		{
 			size_t immutableSamplersSize = bindings[i].descriptorCount * sizeof(VkSampler);
-			bindings[i].pImmutableSamplers = reinterpret_cast<const VkSampler*>(host_memory);
-			host_memory += immutableSamplersSize;
+			bindings[i].pImmutableSamplers = reinterpret_cast<const VkSampler*>(hostMemory);
+			hostMemory += immutableSamplersSize;
 			memcpy(const_cast<VkSampler*>(bindings[i].pImmutableSamplers),
 			       pCreateInfo->pBindings[i].pImmutableSamplers,
 			       immutableSamplersSize);
@@ -51,6 +66,8 @@ DescriptorSetLayout::DescriptorSetLayout(const VkDescriptorSetLayoutCreateInfo*
 		{
 			bindings[i].pImmutableSamplers = nullptr;
 		}
+		bindingOffsets[i] = offset;
+		offset += bindings[i].descriptorCount * GetDescriptorSize(bindings[i].descriptorType);
 	}
 }
 
@@ -71,7 +88,7 @@ void DescriptorSetLayout::destroy(const VkAllocationCallbacks* pAllocator)
 
 size_t DescriptorSetLayout::ComputeRequiredAllocationSize(const VkDescriptorSetLayoutCreateInfo* pCreateInfo)
 {
-	size_t allocationSize = pCreateInfo->bindingCount * sizeof(VkDescriptorSetLayoutBinding);
+	size_t allocationSize = pCreateInfo->bindingCount * (sizeof(VkDescriptorSetLayoutBinding) + sizeof(size_t));
 
 	for(uint32_t i = 0; i < pCreateInfo->bindingCount; i++)
 	{
@@ -92,14 +109,16 @@ size_t DescriptorSetLayout::GetDescriptorSize(VkDescriptorType type)
 	case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
 	case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE:
 	case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
+	case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
+		return sizeof(VkDescriptorImageInfo);
 	case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
 	case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
+		return sizeof(VkBufferView);
 	case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
 	case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER:
 	case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC:
 	case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC:
-	case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
-		return sizeof(void*); // FIXME(b/123244275) : Compute actual required size for each desciptor type
+		return sizeof(VkDescriptorBufferInfo);
 	default:
 		UNIMPLEMENTED("Unsupported Descriptor Type");
 	}
@@ -117,4 +136,119 @@ size_t DescriptorSetLayout::getSize() const
 	return size;
 }
 
+uint32_t DescriptorSetLayout::getBindingIndex(uint32_t binding) const
+{
+	for(uint32_t i = 0; i < bindingCount; i++)
+	{
+		if(binding == bindings[i].binding)
+		{
+			return i;
+		}
+	}
+
+	ASSERT(false); // Bindings should always be found
+	return 0;
+}
+
+void DescriptorSetLayout::initialize(VkDescriptorSet vkDescriptorSet)
+{
+	// Use a pointer to this descriptor set layout as the descriptor set's header
+	DescriptorSet* descriptorSet = ::Cast(vkDescriptorSet);
+	descriptorSet->layout = this;
+	uint8_t* mem = descriptorSet->data;
+
+	for(uint32_t i = 0; i < bindingCount; i++)
+	{
+		size_t typeSize = GetDescriptorSize(bindings[i].descriptorType);
+		if(UsesImmutableSamplers(bindings[i]))
+		{
+			for(uint32_t j = 0; j < bindings[i].descriptorCount; j++)
+			{
+				VkDescriptorImageInfo* imageInfo = reinterpret_cast<VkDescriptorImageInfo*>(mem);
+				imageInfo->sampler = bindings[i].pImmutableSamplers[j];
+				mem += typeSize;
+			}
+		}
+		else
+		{
+			mem += bindings[i].descriptorCount * typeSize;
+		}
+	}
+}
+
+uint8_t* DescriptorSetLayout::getOffsetPointer(VkDescriptorSet descriptorSet, uint32_t binding, uint32_t arrayElement, uint32_t count, size_t* typeSize) const
+{
+	uint32_t index = getBindingIndex(binding);
+	*typeSize = GetDescriptorSize(bindings[index].descriptorType);
+	size_t byteOffset = bindingOffsets[index] + (*typeSize * arrayElement);
+	ASSERT(((*typeSize * count) + byteOffset) <= getSize()); // Make sure the operation will not go out of bounds
+	return &(::Cast(descriptorSet)->data[byteOffset]);
+}
+
+const uint8_t* DescriptorSetLayout::GetInputData(const VkWriteDescriptorSet& descriptorWrites)
+{
+	switch(descriptorWrites.descriptorType)
+	{
+	case VK_DESCRIPTOR_TYPE_SAMPLER:
+	case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
+	case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE:
+	case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
+	case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
+		return reinterpret_cast<const uint8_t*>(descriptorWrites.pImageInfo);
+	case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
+	case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
+		return reinterpret_cast<const uint8_t*>(descriptorWrites.pTexelBufferView);
+		break;
+	case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
+	case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER:
+	case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC:
+	case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC:
+		return reinterpret_cast<const uint8_t*>(descriptorWrites.pBufferInfo);
+		break;
+	default:
+		UNIMPLEMENTED();
+		return nullptr;
+	}
+}
+
+void DescriptorSetLayout::WriteDescriptorSet(const VkWriteDescriptorSet& descriptorWrites)
+{
+	DescriptorSet* dstSet = ::Cast(descriptorWrites.dstSet);
+	DescriptorSetLayout* dstLayout = dstSet->layout;
+	ASSERT(dstLayout);
+
+	size_t typeSize = 0;
+	uint8_t* memToWrite = dstLayout->getOffsetPointer(descriptorWrites.dstSet, descriptorWrites.dstBinding, descriptorWrites.dstArrayElement, descriptorWrites.descriptorCount, &typeSize);
+
+	// If the dstBinding has fewer than descriptorCount array elements remaining
+	// starting from dstArrayElement, then the remainder will be used to update
+	// the subsequent binding - dstBinding+1 starting at array element zero. If
+	// a binding has a descriptorCount of zero, it is skipped. This behavior
+	// applies recursively, with the update affecting consecutive bindings as
+	// needed to update all descriptorCount descriptors.
+	size_t writeSize = typeSize * descriptorWrites.descriptorCount;
+	memcpy(memToWrite, DescriptorSetLayout::GetInputData(descriptorWrites), writeSize);
+}
+
+void DescriptorSetLayout::CopyDescriptorSet(const VkCopyDescriptorSet& descriptorCopies)
+{
+	DescriptorSet* srcSet = ::Cast(descriptorCopies.srcSet);
+	DescriptorSetLayout* srcLayout = srcSet->layout;
+	ASSERT(srcLayout);
+
+	DescriptorSet* dstSet = ::Cast(descriptorCopies.dstSet);
+	DescriptorSetLayout* dstLayout = dstSet->layout;
+	ASSERT(dstLayout);
+
+	size_t srcTypeSize = 0;
+	uint8_t* memToRead = srcLayout->getOffsetPointer(descriptorCopies.srcSet, descriptorCopies.srcBinding, descriptorCopies.srcArrayElement, descriptorCopies.descriptorCount, &srcTypeSize);
+
+	size_t dstTypeSize = 0;
+	uint8_t* memToWrite = dstLayout->getOffsetPointer(descriptorCopies.dstSet, descriptorCopies.dstBinding, descriptorCopies.dstArrayElement, descriptorCopies.descriptorCount, &dstTypeSize);
+
+	ASSERT(srcTypeSize == dstTypeSize);
+	size_t writeSize = dstTypeSize * descriptorCopies.descriptorCount;
+	memcpy(memToWrite, memToRead, writeSize);
+}
+
 } // namespace vk
\ No newline at end of file

src/Vulkan/VkDescriptorSetLayout.hpp

@@ -30,13 +30,21 @@ public:
 	static size_t ComputeRequiredAllocationSize(const VkDescriptorSetLayoutCreateInfo* pCreateInfo);
 
 	static size_t GetDescriptorSize(VkDescriptorType type);
+	static void WriteDescriptorSet(const VkWriteDescriptorSet& descriptorWrites);
+	static void CopyDescriptorSet(const VkCopyDescriptorSet& descriptorCopies);
 
+	void initialize(VkDescriptorSet descriptorSet);
 	size_t getSize() const;
 
 private:
+	uint32_t getBindingIndex(uint32_t binding) const;
+	uint8_t* getOffsetPointer(VkDescriptorSet descriptorSet, uint32_t binding, uint32_t arrayElement, uint32_t count, size_t* typeSize) const;
+	static const uint8_t* GetInputData(const VkWriteDescriptorSet& descriptorWrites);
+
 	VkDescriptorSetLayoutCreateFlags flags;
 	uint32_t                         bindingCount;
 	VkDescriptorSetLayoutBinding*    bindings;
+	size_t*                          bindingOffsets;
 };
 
 static inline DescriptorSetLayout* Cast(VkDescriptorSetLayout object)

src/Vulkan/VkDevice.cpp

@@ -16,6 +16,7 @@
 
 #include "VkConfig.h"
 #include "VkDebug.hpp"
+#include "VkDescriptorSetLayout.hpp"
 #include "VkQueue.hpp"
 
 #include <new> // Must #include this to use "placement new"
@@ -99,4 +100,18 @@ void Device::getDescriptorSetLayoutSupport(const VkDescriptorSetLayoutCreateInfo
 	pSupport->supported = VK_FALSE;
 }
 
+void Device::updateDescriptorSets(uint32_t descriptorWriteCount, const VkWriteDescriptorSet* pDescriptorWrites,
+                                  uint32_t descriptorCopyCount, const VkCopyDescriptorSet* pDescriptorCopies)
+{
+	for(uint32_t i = 0; i < descriptorWriteCount; i++)
+	{
+		DescriptorSetLayout::WriteDescriptorSet(pDescriptorWrites[i]);
+	}
+
+	for(uint32_t i = 0; i < descriptorCopyCount; i++)
+	{
+		DescriptorSetLayout::CopyDescriptorSet(pDescriptorCopies[i]);
+	}
+}
+
 } // namespace vk

src/Vulkan/VkDevice.hpp

@@ -44,6 +44,8 @@ public:
 	void getDescriptorSetLayoutSupport(const VkDescriptorSetLayoutCreateInfo* pCreateInfo,
 	                                   VkDescriptorSetLayoutSupport* pSupport) const;
 	VkPhysicalDevice getPhysicalDevice() const { return physicalDevice; }
+	void updateDescriptorSets(uint32_t descriptorWriteCount, const VkWriteDescriptorSet* pDescriptorWrites,
+	                          uint32_t descriptorCopyCount, const VkCopyDescriptorSet* pDescriptorCopies);
 
 private:
 	VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;

src/Vulkan/libVulkan.cpp

@@ -1189,8 +1189,10 @@ VKAPI_ATTR VkResult VKAPI_CALL vkFreeDescriptorSets(VkDevice device, VkDescripto
 
 VKAPI_ATTR void VKAPI_CALL vkUpdateDescriptorSets(VkDevice device, uint32_t descriptorWriteCount, const VkWriteDescriptorSet* pDescriptorWrites, uint32_t descriptorCopyCount, const VkCopyDescriptorSet* pDescriptorCopies)
 {
-	TRACE("()");
-	UNIMPLEMENTED();
+	TRACE("(VkDevice device = 0x%X, uint32_t descriptorWriteCount = %d, const VkWriteDescriptorSet* pDescriptorWrites = 0x%X, uint32_t descriptorCopyCount = %d, const VkCopyDescriptorSet* pDescriptorCopies = 0x%X)",
+	      device, descriptorWriteCount, pDescriptorWrites, descriptorCopyCount, pDescriptorCopies);
+
+	vk::Cast(device)->updateDescriptorSets(descriptorWriteCount, pDescriptorWrites, descriptorCopyCount, pDescriptorCopies);
 }
 
 VKAPI_ATTR VkResult VKAPI_CALL vkCreateFramebuffer(VkDevice device, const VkFramebufferCreateInfo* pCreateInfo, const VkAllocationCallbacks* pAllocator, VkFramebuffer* pFramebuffer)