Fix various descriptor handling bugs

src/Pipeline/SpirvShader.cpp

@@ -1334,20 +1334,23 @@ namespace sw
 				auto setLayout = routine->pipelineLayout->getDescriptorSetLayout(d.DescriptorSet);
 				int bindingOffset = static_cast<int>(setLayout->getBindingOffset(d.Binding, arrayIndex));
 
-				Pointer<Byte> bufferInfo = set.base + bindingOffset; // VkDescriptorBufferInfo*
-				Pointer<Byte> buffer = *Pointer<Pointer<Byte>>(bufferInfo + OFFSET(VkDescriptorBufferInfo, buffer)); // vk::Buffer*
-				Pointer<Byte> data = *Pointer<Pointer<Byte>>(buffer + vk::Buffer::DataOffset); // void*
-				Int offset = *Pointer<Int>(bufferInfo + OFFSET(VkDescriptorBufferInfo, offset));
-				Int size = *Pointer<Int>(buffer + vk::Buffer::DataSize); // void*
+				Pointer<Byte> descriptor = set.base + bindingOffset; // BufferDescriptor*
+				Pointer<Byte> data = *Pointer<Pointer<Byte>>(descriptor + OFFSET(vk::BufferDescriptor, ptr)); // void*
+				Int size = *Pointer<Int>(descriptor + OFFSET(vk::BufferDescriptor, sizeInBytes));
 				if (setLayout->isBindingDynamic(d.Binding))
 				{
 					uint32_t dynamicBindingIndex =
 						routine->pipelineLayout->getDynamicOffsetBase(d.DescriptorSet) +
 						setLayout->getDynamicDescriptorOffset(d.Binding) +
 						arrayIndex;
-					offset += routine->descriptorDynamicOffsets[dynamicBindingIndex];
+					Int offset = routine->descriptorDynamicOffsets[dynamicBindingIndex];
+					Int robustnessSize = *Pointer<Int>(descriptor + OFFSET(vk::BufferDescriptor, robustnessSize));
+					return SIMD::Pointer(data + offset, Min(size, robustnessSize - offset));
+				}
+				else
+				{
+					return SIMD::Pointer(data, size);
 				}
-				return SIMD::Pointer(data + offset, size - offset);
 			}
 
 			default:

src/Vulkan/VkBuffer.cpp

@@ -21,9 +21,6 @@
 namespace vk
 {
 
-const int Buffer::DataOffset = static_cast<int>(offsetof(Buffer, memory));
-const int Buffer::DataSize = static_cast<int>(offsetof(Buffer, size));
-
 Buffer::Buffer(const VkBufferCreateInfo* pCreateInfo, void* mem) :
 	flags(pCreateInfo->flags), size(pCreateInfo->size), usage(pCreateInfo->usage),
 	sharingMode(pCreateInfo->sharingMode)

src/Vulkan/VkBuffer.hpp

@@ -40,11 +40,6 @@ public:
 	inline VkDeviceSize getSize() const { return size; }
 	uint8_t* end() const;
 
-	// DataOffset is the offset in bytes from the Buffer to the pointer to the
-	// buffer's data memory.
-	static const int DataOffset;
-	static const int DataSize;
-
 private:
 	void*                 memory = nullptr;
 	VkBufferCreateFlags   flags = 0;

src/Vulkan/VkDescriptorPool.cpp

@@ -36,13 +36,12 @@ void DescriptorPool::destroy(const VkAllocationCallbacks* pAllocator)
 
 size_t DescriptorPool::ComputeRequiredAllocationSize(const VkDescriptorPoolCreateInfo* pCreateInfo)
 {
-	size_t size = 0;
+	size_t size = pCreateInfo->maxSets * sizeof(DescriptorSetHeader);
 
 	for(uint32_t i = 0; i < pCreateInfo->poolSizeCount; i++)
 	{
 		size += pCreateInfo->pPoolSizes[i].descriptorCount *
-		        (sizeof(DescriptorSetHeader) +
-		         DescriptorSetLayout::GetDescriptorSize(pCreateInfo->pPoolSizes[i].type));
+		        DescriptorSetLayout::GetDescriptorSize(pCreateInfo->pPoolSizes[i].type);
 	}
 
 	return size;

src/Vulkan/VkDescriptorSetLayout.cpp

@@ -17,6 +17,7 @@
 #include "VkDescriptorSet.hpp"
 #include "VkSampler.hpp"
 #include "VkImageView.hpp"
+#include "VkBuffer.hpp"
 #include "VkBufferView.hpp"
 #include "System/Types.hpp"
 
@@ -65,6 +66,7 @@ DescriptorSetLayout::DescriptorSetLayout(const VkDescriptorSetLayoutCreateInfo*
 		bindingOffsets[i] = offset;
 		offset += bindings[i].descriptorCount * GetDescriptorSize(bindings[i].descriptorType);
 	}
+	ASSERT_MSG(offset == getDescriptorSetDataSize(), "offset: %d, size: %d", int(offset), int(getDescriptorSetDataSize()));
 }
 
 void DescriptorSetLayout::destroy(const VkAllocationCallbacks* pAllocator)
@@ -89,44 +91,32 @@ size_t DescriptorSetLayout::ComputeRequiredAllocationSize(const VkDescriptorSetL
 
 size_t DescriptorSetLayout::GetDescriptorSize(VkDescriptorType type)
 {
-	size_t size = 0;
-
 	switch(type)
 	{
 	case VK_DESCRIPTOR_TYPE_SAMPLER:
 	case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
 	case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE:
-		size = sizeof(SampledImageDescriptor);
-		break;
+	case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
+		return sizeof(SampledImageDescriptor);
 	case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
 	case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
-		size = sizeof(StorageImageDescriptor);
-		break;
 	case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
-		size = sizeof(VkDescriptorImageInfo);
-		break;
-	case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
-		size = sizeof(VkBufferView);
-		break;
+		return sizeof(StorageImageDescriptor);
 	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:
-		size = sizeof(VkDescriptorBufferInfo);
-		break;
+		return sizeof(BufferDescriptor);
 	default:
 		UNIMPLEMENTED("Unsupported Descriptor Type");
 		return 0;
 	}
-
-	// Aligning each descriptor to 16 bytes allows for more efficient vector accesses in the shaders.
-	return sw::align<16>(size);  // TODO(b/123244275): Eliminate by using a custom alignas(16) struct for each desctriptor.
 }
 
 size_t DescriptorSetLayout::getDescriptorSetAllocationSize() const
 {
 	// vk::DescriptorSet has a layout member field.
-	return sizeof(vk::DescriptorSetHeader) + getDescriptorSetDataSize();
+	return sw::align<alignof(DescriptorSet)>(OFFSET(DescriptorSet, data) + getDescriptorSetDataSize());
 }
 
 size_t DescriptorSetLayout::getDescriptorSetDataSize() const
@@ -258,8 +248,11 @@ void SampledImageDescriptor::updateSampler(const vk::Sampler *sampler)
 {
 	this->sampler = sampler;
 
-	texture.minLod = sw::clamp(sampler->minLod, 0.0f, (float)(sw::MAX_TEXTURE_LOD));
-	texture.maxLod = sw::clamp(sampler->maxLod, 0.0f, (float)(sw::MAX_TEXTURE_LOD));
+	if (sampler)
+	{
+		texture.minLod = sw::clamp(sampler->minLod, 0.0f, (float) (sw::MAX_TEXTURE_LOD));
+		texture.maxLod = sw::clamp(sampler->maxLod, 0.0f, (float) (sw::MAX_TEXTURE_LOD));
+	}
 }
 
 void DescriptorSetLayout::WriteDescriptorSet(DescriptorSet *dstSet, VkDescriptorUpdateTemplateEntry const &entry, char const *src)
@@ -274,7 +267,27 @@ void DescriptorSetLayout::WriteDescriptorSet(DescriptorSet *dstSet, VkDescriptor
 
 	ASSERT(reinterpret_cast<intptr_t>(memToWrite) % 16 == 0);  // Each descriptor must be 16-byte aligned.
 
-	if(entry.descriptorType == VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER)
+	if (entry.descriptorType == VK_DESCRIPTOR_TYPE_SAMPLER)
+	{
+		SampledImageDescriptor *imageSampler = reinterpret_cast<SampledImageDescriptor*>(memToWrite);
+
+		for(uint32_t i = 0; i < entry.descriptorCount; i++)
+		{
+			auto update = reinterpret_cast<VkDescriptorImageInfo const *>(src + entry.offset + entry.stride * i);
+			// "All consecutive bindings updated via a single VkWriteDescriptorSet structure, except those with a
+			//  descriptorCount of zero, must all either use immutable samplers or must all not use immutable samplers."
+			if (!binding.pImmutableSamplers)
+			{
+				imageSampler[i].updateSampler(vk::Cast(update->sampler));
+			}
+		}
+	}
+	else if (entry.descriptorType == VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER)
+	{
+		UNIMPLEMENTED("VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER");
+	}
+	else if (entry.descriptorType == VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER ||
+	   		 entry.descriptorType == VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE)
 	{
 		SampledImageDescriptor *imageSampler = reinterpret_cast<SampledImageDescriptor*>(memToWrite);
 
@@ -486,16 +499,20 @@ void DescriptorSetLayout::WriteDescriptorSet(DescriptorSet *dstSet, VkDescriptor
 			descriptor[i].sizeInBytes = bufferView->getRangeInBytes();
 		}
 	}
-	else
+	else if (entry.descriptorType == VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER ||
+			 entry.descriptorType == VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC ||
+			 entry.descriptorType == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER ||
+			 entry.descriptorType == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC)
 	{
-		// 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.
-		for (auto i = 0u; i < entry.descriptorCount; i++)
-			memcpy(memToWrite + typeSize * i, src + entry.offset + entry.stride * i, typeSize);
+		auto descriptor = reinterpret_cast<BufferDescriptor *>(memToWrite);
+		for (uint32_t i = 0; i < entry.descriptorCount; i++)
+		{
+			auto update = reinterpret_cast<VkDescriptorBufferInfo const *>(src + entry.offset + entry.stride * i);
+			auto buffer = Cast(update->buffer);
+			descriptor[i].ptr = buffer->getOffsetPointer(update->offset);
+			descriptor[i].sizeInBytes = (update->range == VK_WHOLE_SIZE) ? buffer->getSize() - update->offset : update->range;
+			descriptor[i].robustnessSize = buffer->getSize() - update->offset;
+		}
 	}
 }
 
@@ -514,7 +531,7 @@ void DescriptorSetLayout::WriteDescriptorSet(const VkWriteDescriptorSet& writeDe
 	case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
 	case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
 		ptr = writeDescriptorSet.pTexelBufferView;
-		e.stride = sizeof(*VkWriteDescriptorSet::pTexelBufferView);
+		e.stride = sizeof(VkBufferView);
 		break;
 
 	case VK_DESCRIPTOR_TYPE_SAMPLER:

src/Vulkan/VkDescriptorSetLayout.hpp

@@ -48,6 +48,13 @@ struct alignas(16) StorageImageDescriptor
 	int sizeInBytes;
 };
 
+struct alignas(16) BufferDescriptor
+{
+	void *ptr;
+	int sizeInBytes;		// intended size of the bound region -- slides along with dynamic offsets
+	int robustnessSize;		// total accessible size from static offset -- does not move with dynamic offset
+};
+
 class DescriptorSetLayout : public Object<DescriptorSetLayout, VkDescriptorSetLayout>
 {
 public: