descriptorBindingVariableDescriptorCount support

This CL adds support for descriptorBindingVariableDescriptorCount,
which allows the implementation to support descriptor sets with a
variable-sized last binding.

Tests: dEQP-VK.binding_model.descriptorset_random.*
Bug: b/236957543
Change-Id: Ifc83ed711c7dbef066b8a0d708f2b8c66da0c589
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/67748
Tested-by: Alexis Hétu <sugoi@google.com>
Commit-Queue: Alexis Hétu <sugoi@google.com>
Presubmit-Ready: Alexis Hétu <sugoi@google.com>
Kokoro-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Nicolas Capens <nicolascapens@google.com>
diff --git a/src/Vulkan/VkDescriptorPool.cpp b/src/Vulkan/VkDescriptorPool.cpp
index 8b98d4b..3eb1662 100644
--- a/src/Vulkan/VkDescriptorPool.cpp
+++ b/src/Vulkan/VkDescriptorPool.cpp
@@ -64,14 +64,17 @@
 	return size;
 }
 
-VkResult DescriptorPool::allocateSets(uint32_t descriptorSetCount, const VkDescriptorSetLayout *pSetLayouts, VkDescriptorSet *pDescriptorSets)
+VkResult DescriptorPool::allocateSets(uint32_t descriptorSetCount, const VkDescriptorSetLayout *pSetLayouts, VkDescriptorSet *pDescriptorSets, const VkDescriptorSetVariableDescriptorCountAllocateInfo *variableDescriptorCountAllocateInfo)
 {
+	const uint32_t *variableDescriptorCounts =
+	    (variableDescriptorCountAllocateInfo && (variableDescriptorCountAllocateInfo->descriptorSetCount == descriptorSetCount)) ? variableDescriptorCountAllocateInfo->pDescriptorCounts : nullptr;
+
 	// FIXME (b/119409619): use an allocator here so we can control all memory allocations
 	std::unique_ptr<size_t[]> layoutSizes(new size_t[descriptorSetCount]);
 	for(uint32_t i = 0; i < descriptorSetCount; i++)
 	{
 		pDescriptorSets[i] = VK_NULL_HANDLE;
-		layoutSizes[i] = vk::Cast(pSetLayouts[i])->getDescriptorSetAllocationSize();
+		layoutSizes[i] = vk::Cast(pSetLayouts[i])->getDescriptorSetAllocationSize(variableDescriptorCounts ? variableDescriptorCounts[i] : 0);
 	}
 
 	VkResult result = allocateSets(&(layoutSizes[0]), descriptorSetCount, pDescriptorSets);
@@ -79,7 +82,7 @@
 	{
 		for(uint32_t i = 0; i < descriptorSetCount; i++)
 		{
-			vk::Cast(pSetLayouts[i])->initialize(vk::Cast(pDescriptorSets[i]));
+			vk::Cast(pSetLayouts[i])->initialize(vk::Cast(pDescriptorSets[i]), variableDescriptorCounts ? variableDescriptorCounts[i] : 0);
 		}
 	}
 	return result;
diff --git a/src/Vulkan/VkDescriptorPool.hpp b/src/Vulkan/VkDescriptorPool.hpp
index 7428835..cc372fb 100644
--- a/src/Vulkan/VkDescriptorPool.hpp
+++ b/src/Vulkan/VkDescriptorPool.hpp
@@ -28,7 +28,7 @@
 
 	static size_t ComputeRequiredAllocationSize(const VkDescriptorPoolCreateInfo *pCreateInfo);
 
-	VkResult allocateSets(uint32_t descriptorSetCount, const VkDescriptorSetLayout *pSetLayouts, VkDescriptorSet *pDescriptorSets);
+	VkResult allocateSets(uint32_t descriptorSetCount, const VkDescriptorSetLayout *pSetLayouts, VkDescriptorSet *pDescriptorSets, const VkDescriptorSetVariableDescriptorCountAllocateInfo *variableDescriptorCountAllocateInfo);
 	void freeSets(uint32_t descriptorSetCount, const VkDescriptorSet *pDescriptorSets);
 	VkResult reset();
 
diff --git a/src/Vulkan/VkDescriptorSetLayout.cpp b/src/Vulkan/VkDescriptorSetLayout.cpp
index 0c5c8bf..5fe5a40 100644
--- a/src/Vulkan/VkDescriptorSetLayout.cpp
+++ b/src/Vulkan/VkDescriptorSetLayout.cpp
@@ -85,7 +85,7 @@
 		offset += bindings[i].descriptorCount * GetDescriptorSize(bindings[i].descriptorType);
 	}
 
-	ASSERT_MSG(offset == getDescriptorSetDataSize(), "offset: %d, size: %d", int(offset), int(getDescriptorSetDataSize()));
+	ASSERT_MSG(offset == getDescriptorSetDataSize(0), "offset: %d, size: %d", int(offset), int(getDescriptorSetDataSize(0)));
 }
 
 void DescriptorSetLayout::destroy(const VkAllocationCallbacks *pAllocator)
@@ -143,24 +143,31 @@
 	       type == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
 }
 
-size_t DescriptorSetLayout::getDescriptorSetAllocationSize() const
+size_t DescriptorSetLayout::getDescriptorSetAllocationSize(uint32_t variableDescriptorCount) const
 {
 	// vk::DescriptorSet has a header with a pointer to the layout.
-	return sw::align<alignof(DescriptorSet)>(sizeof(DescriptorSetHeader) + getDescriptorSetDataSize());
+	return sw::align<alignof(DescriptorSet)>(sizeof(DescriptorSetHeader) + getDescriptorSetDataSize(variableDescriptorCount));
 }
 
-size_t DescriptorSetLayout::getDescriptorSetDataSize() const
+size_t DescriptorSetLayout::getDescriptorSetDataSize(uint32_t variableDescriptorCount) const
 {
 	size_t size = 0;
 	for(uint32_t i = 0; i < bindingsArraySize; i++)
 	{
-		size += bindings[i].descriptorCount * GetDescriptorSize(bindings[i].descriptorType);
+		uint32_t descriptorCount = bindings[i].descriptorCount;
+
+		if((i == (bindingsArraySize - 1)) && (variableDescriptorCount > 0))
+		{
+			descriptorCount = variableDescriptorCount;
+		}
+
+		size += descriptorCount * GetDescriptorSize(bindings[i].descriptorType);
 	}
 
 	return size;
 }
 
-void DescriptorSetLayout::initialize(DescriptorSet *descriptorSet)
+void DescriptorSetLayout::initialize(DescriptorSet *descriptorSet, uint32_t variableDescriptorCount)
 {
 	ASSERT(descriptorSet->header.layout == nullptr);
 
@@ -172,9 +179,16 @@
 	{
 		size_t descriptorSize = GetDescriptorSize(bindings[i].descriptorType);
 
+		uint32_t descriptorCount = bindings[i].descriptorCount;
+
+		if((i == (bindingsArraySize - 1)) && (variableDescriptorCount > 0))
+		{
+			descriptorCount = variableDescriptorCount;
+		}
+
 		if(bindings[i].immutableSamplers)
 		{
-			for(uint32_t j = 0; j < bindings[i].descriptorCount; j++)
+			for(uint32_t j = 0; j < descriptorCount; j++)
 			{
 				SampledImageDescriptor *imageSamplerDescriptor = reinterpret_cast<SampledImageDescriptor *>(data);
 				imageSamplerDescriptor->samplerId = bindings[i].immutableSamplers[j]->id;
@@ -190,7 +204,7 @@
 			case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
 			case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE:
 			case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
-				for(uint32_t j = 0; j < bindings[i].descriptorCount; j++)
+				for(uint32_t j = 0; j < descriptorCount; j++)
 				{
 					SampledImageDescriptor *imageSamplerDescriptor = reinterpret_cast<SampledImageDescriptor *>(data);
 					imageSamplerDescriptor->memoryOwner = nullptr;
@@ -200,7 +214,7 @@
 			case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
 			case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
 			case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
-				for(uint32_t j = 0; j < bindings[i].descriptorCount; j++)
+				for(uint32_t j = 0; j < descriptorCount; j++)
 				{
 					StorageImageDescriptor *storageImage = reinterpret_cast<StorageImageDescriptor *>(data);
 					storageImage->memoryOwner = nullptr;
@@ -211,10 +225,10 @@
 			case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER:
 			case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC:
 			case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC:
-				data += bindings[i].descriptorCount * descriptorSize;
+				data += descriptorCount * descriptorSize;
 				break;
 			case VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK:
-				data += bindings[i].descriptorCount;
+				data += descriptorCount;
 				break;
 			default:
 				UNSUPPORTED("Unsupported Descriptor Type: %d", int(bindings[i].descriptorType));
@@ -277,7 +291,7 @@
 	ASSERT(bindingNumber < bindingsArraySize);
 	*typeSize = GetDescriptorSize(bindings[bindingNumber].descriptorType);
 	size_t byteOffset = bindings[bindingNumber].offset + (*typeSize * arrayElement);
-	ASSERT(((*typeSize * count) + byteOffset) <= getDescriptorSetDataSize());  // Make sure the operation will not go out of bounds
+	ASSERT(((*typeSize * count) + byteOffset) <= getDescriptorSetDataSize(0));  // Make sure the operation will not go out of bounds
 
 	return descriptorSet->getDataAddress() + byteOffset;
 }
diff --git a/src/Vulkan/VkDescriptorSetLayout.hpp b/src/Vulkan/VkDescriptorSetLayout.hpp
index 80786a9..10ac3b3 100644
--- a/src/Vulkan/VkDescriptorSetLayout.hpp
+++ b/src/Vulkan/VkDescriptorSetLayout.hpp
@@ -106,10 +106,10 @@
 
 	static void WriteDescriptorSet(Device *device, DescriptorSet *dstSet, VkDescriptorUpdateTemplateEntry const &entry, char const *src);
 
-	void initialize(DescriptorSet *descriptorSet);
+	void initialize(DescriptorSet *descriptorSet, uint32_t variableDescriptorCount);
 
 	// Returns the total size of the descriptor set in bytes.
-	size_t getDescriptorSetAllocationSize() const;
+	size_t getDescriptorSetAllocationSize(uint32_t variableDescriptorCount) const;
 
 	// Returns the byte offset from the base address of the descriptor set for
 	// the given binding number.
@@ -136,7 +136,7 @@
 
 private:
 	uint8_t *getDescriptorPointer(DescriptorSet *descriptorSet, uint32_t bindingNumber, uint32_t arrayElement, uint32_t count, size_t *typeSize) const;
-	size_t getDescriptorSetDataSize() const;
+	size_t getDescriptorSetDataSize(uint32_t variableDescriptorCount) const;
 	static bool isDynamic(VkDescriptorType type);
 
 	const VkDescriptorSetLayoutCreateFlags flags;
diff --git a/src/Vulkan/VkDevice.cpp b/src/Vulkan/VkDevice.cpp
index 6d074dd..caed186 100644
--- a/src/Vulkan/VkDevice.cpp
+++ b/src/Vulkan/VkDevice.cpp
@@ -19,6 +19,7 @@
 #include "VkFence.hpp"
 #include "VkQueue.hpp"
 #include "VkSemaphore.hpp"
+#include "VkStringify.hpp"
 #include "VkTimelineSemaphore.hpp"
 #include "Debug/Context.hpp"
 #include "Debug/Server.hpp"
@@ -362,6 +363,59 @@
 
 	// We have no "strange" limitations to enforce beyond the device limits, so we can safely always claim support.
 	pSupport->supported = VK_TRUE;
+
+	if(pCreateInfo->bindingCount > 0)
+	{
+		bool hasVariableSizedDescriptor = false;
+
+		const VkBaseInStructure *layoutInfo = reinterpret_cast<const VkBaseInStructure *>(pCreateInfo->pNext);
+		while(layoutInfo && !hasVariableSizedDescriptor)
+		{
+			if(layoutInfo->sType == VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO)
+			{
+				const VkDescriptorSetLayoutBindingFlagsCreateInfo *bindingFlagsCreateInfo =
+				    reinterpret_cast<const VkDescriptorSetLayoutBindingFlagsCreateInfo *>(layoutInfo);
+
+				for(uint32_t i = 0; i < bindingFlagsCreateInfo->bindingCount; i++)
+				{
+					if(bindingFlagsCreateInfo->pBindingFlags[i] & VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT)
+					{
+						hasVariableSizedDescriptor = true;
+						break;
+					}
+				}
+			}
+			else
+			{
+				UNSUPPORTED("layoutInfo->sType = %s", vk::Stringify(layoutInfo->sType).c_str());
+			}
+
+			layoutInfo = layoutInfo->pNext;
+		}
+
+		const auto &highestNumberedBinding = pCreateInfo->pBindings[pCreateInfo->bindingCount - 1];
+
+		VkBaseOutStructure *layoutSupport = reinterpret_cast<VkBaseOutStructure *>(pSupport->pNext);
+		while(layoutSupport)
+		{
+			if(layoutSupport->sType == VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_LAYOUT_SUPPORT)
+			{
+				VkDescriptorSetVariableDescriptorCountLayoutSupport *variableDescriptorCountLayoutSupport =
+				    reinterpret_cast<VkDescriptorSetVariableDescriptorCountLayoutSupport *>(layoutSupport);
+
+				// If the VkDescriptorSetLayoutCreateInfo structure does not include a variable-sized descriptor,
+				// [...] then maxVariableDescriptorCount is set to zero.
+				variableDescriptorCountLayoutSupport->maxVariableDescriptorCount =
+				    hasVariableSizedDescriptor ? ((highestNumberedBinding.descriptorType == VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK) ? vk::MAX_INLINE_UNIFORM_BLOCK_SIZE : vk::MAX_UPDATE_AFTER_BIND_DESCRIPTORS) : 0;
+			}
+			else
+			{
+				UNSUPPORTED("layoutSupport->sType = %s", vk::Stringify(layoutSupport->sType).c_str());
+			}
+
+			layoutSupport = layoutSupport->pNext;
+		}
+	}
 }
 
 void Device::updateDescriptorSets(uint32_t descriptorWriteCount, const VkWriteDescriptorSet *pDescriptorWrites,
diff --git a/src/Vulkan/VkPhysicalDevice.cpp b/src/Vulkan/VkPhysicalDevice.cpp
index 27aec3c..499f1ee 100644
--- a/src/Vulkan/VkPhysicalDevice.cpp
+++ b/src/Vulkan/VkPhysicalDevice.cpp
@@ -261,7 +261,7 @@
 	features->descriptorBindingStorageTexelBufferUpdateAfterBind = VK_TRUE;
 	features->descriptorBindingUpdateUnusedWhilePending = VK_TRUE;
 	features->descriptorBindingPartiallyBound = VK_TRUE;
-	features->descriptorBindingVariableDescriptorCount = VK_FALSE;
+	features->descriptorBindingVariableDescriptorCount = VK_TRUE;
 	features->runtimeDescriptorArray = VK_TRUE;
 }
 
diff --git a/src/Vulkan/libVulkan.cpp b/src/Vulkan/libVulkan.cpp
index 14af446..43fb688 100644
--- a/src/Vulkan/libVulkan.cpp
+++ b/src/Vulkan/libVulkan.cpp
@@ -2465,14 +2465,24 @@
 	TRACE("(VkDevice device = %p, const VkDescriptorSetAllocateInfo* pAllocateInfo = %p, VkDescriptorSet* pDescriptorSets = %p)",
 	      device, pAllocateInfo, pDescriptorSets);
 
+	const VkDescriptorSetVariableDescriptorCountAllocateInfo *variableDescriptorCountAllocateInfo = nullptr;
+
 	auto extInfo = reinterpret_cast<VkBaseInStructure const *>(pAllocateInfo->pNext);
 	while(extInfo)
 	{
-		UNSUPPORTED("pAllocateInfo->pNext sType = %s", vk::Stringify(extInfo->sType).c_str());
+		switch(extInfo->sType)
+		{
+		case VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO:
+			variableDescriptorCountAllocateInfo = reinterpret_cast<VkDescriptorSetVariableDescriptorCountAllocateInfo const *>(extInfo);
+			break;
+		default:
+			UNSUPPORTED("pAllocateInfo->pNext sType = %s", vk::Stringify(extInfo->sType).c_str());
+			break;
+		}
 		extInfo = extInfo->pNext;
 	}
 
-	return vk::Cast(pAllocateInfo->descriptorPool)->allocateSets(pAllocateInfo->descriptorSetCount, pAllocateInfo->pSetLayouts, pDescriptorSets);
+	return vk::Cast(pAllocateInfo->descriptorPool)->allocateSets(pAllocateInfo->descriptorSetCount, pAllocateInfo->pSetLayouts, pDescriptorSets, variableDescriptorCountAllocateInfo);
 }
 
 VKAPI_ATTR VkResult VKAPI_CALL vkFreeDescriptorSets(VkDevice device, VkDescriptorPool descriptorPool, uint32_t descriptorSetCount, const VkDescriptorSet *pDescriptorSets)
@@ -4192,6 +4202,21 @@
 	TRACE("(VkDevice device = %p, const VkDescriptorSetLayoutCreateInfo* pCreateInfo = %p, VkDescriptorSetLayoutSupport* pSupport = %p)",
 	      device, pCreateInfo, pSupport);
 
+	VkBaseOutStructure *layoutSupport = reinterpret_cast<VkBaseOutStructure *>(pSupport->pNext);
+	while(layoutSupport)
+	{
+		switch(layoutSupport->sType)
+		{
+		case VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_LAYOUT_SUPPORT:
+			break;
+		default:
+			UNSUPPORTED("pSupport->pNext sType = %s", vk::Stringify(layoutSupport->sType).c_str());
+			break;
+		}
+
+		layoutSupport = layoutSupport->pNext;
+	}
+
 	vk::Cast(device)->getDescriptorSetLayoutSupport(pCreateInfo, pSupport);
 }