Record descriptor byte patterns in trie with DescriptorSetSlot value

renderdoc/driver/vulkan/vk_common.cpp

@@ -1308,13 +1308,14 @@ void DescriptorSetSlot::SetImage(VkDescriptorType writeType, const VkDescriptorI
     sampler = GetResID(imInfo.sampler);
   if(type != DescriptorSlotType::Sampler)
     resource = GetResID(imInfo.imageView);
-  imageLayout = convert(imInfo.imageLayout);
+  imageLayoutOrFormat = convert(imInfo.imageLayout);
 }
 
 void DescriptorSetSlot::SetTexelBuffer(VkDescriptorType writeType, ResourceId id)
 {
   type = convert(writeType);
   resource = id;
+  imageLayoutOrFormat = DescriptorSlotImageLayout::Undefined;
 }
 
 void DescriptorSetSlot::SetAccelerationStructure(VkDescriptorType writeType,
@@ -1324,6 +1325,71 @@ void DescriptorSetSlot::SetAccelerationStructure(VkDescriptorType writeType,
   resource = GetResID(accelerationStructure);
 }
 
+void DescriptorSetSlot::SetDescriptor(WrappedVulkan *driver, const VkDescriptorGetInfoEXT &desc)
+{
+  type = convert(desc.type);
+  switch(desc.type)
+  {
+    case VK_DESCRIPTOR_TYPE_SAMPLER:
+    {
+      sampler = desc.data.pSampler ? GetResID(*desc.data.pSampler) : ResourceId();
+      break;
+    }
+    case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
+      // ignore the sampler part
+    case VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT:
+    case VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE:
+    case VK_DESCRIPTOR_TYPE_STORAGE_IMAGE:
+    {
+      // sampled/storage/input attachment are identical in the union. Since the type forms part of
+      // our this logic can be done in common
+      if(desc.data.pCombinedImageSampler)
+      {
+        resource = GetResID(desc.data.pCombinedImageSampler->imageView);
+        sampler = GetResID(desc.data.pCombinedImageSampler->sampler);
+        imageLayoutOrFormat = convert(desc.data.pCombinedImageSampler->imageLayout);
+      }
+      break;
+    }
+    case VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER:
+    case VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER:
+    case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
+    case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER:
+    {
+      // uniform/storage are identical in the union. Since the type forms part of our this
+      // logic can be done in common
+      if(desc.data.pUniformBuffer)
+      {
+        ResourceId id;
+        driver->GetResIDFromAddr(desc.data.pUniformBuffer->address, resource, offset);
+        range = desc.data.pUniformBuffer->range;
+        // we only expect texel buffers with the simple formats that come from vulkan base which are less than 256
+        imageLayoutOrFormat = DescriptorSlotImageLayout(desc.data.pUniformTexelBuffer->format & 0xff);
+        RDCASSERT(uint32_t(desc.data.pUniformTexelBuffer->format) < 0xff,
+                  desc.data.pUniformTexelBuffer->format);
+      }
+      break;
+    }
+    case VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR:
+    {
+      if(desc.data.accelerationStructure)
+        resource = driver->GetASFromAddr(desc.data.accelerationStructure);
+      break;
+    }
+    case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC:
+    case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC:
+    case VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK:
+    case VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_NV:
+    case VK_DESCRIPTOR_TYPE_SAMPLE_WEIGHT_IMAGE_QCOM:
+    case VK_DESCRIPTOR_TYPE_BLOCK_MATCH_IMAGE_QCOM:
+    case VK_DESCRIPTOR_TYPE_MUTABLE_EXT:
+    case VK_DESCRIPTOR_TYPE_PARTITIONED_ACCELERATION_STRUCTURE_NV:
+    case VK_DESCRIPTOR_TYPE_MAX_ENUM:
+      RDCERR("Invalid descriptor type passed to vkGetDescriptorEXT");
+      break;
+  }
+}
+
 void AddBindFrameRef(DescriptorBindRefs &refs, ResourceId id, FrameRefType ref)
 {
   if(id == ResourceId())

renderdoc/driver/vulkan/vk_common.h

@@ -811,6 +811,7 @@ struct DescriptorSetSlot
   void SetTexelBuffer(VkDescriptorType writeType, ResourceId id);
   void SetAccelerationStructure(VkDescriptorType writeType,
                                 VkAccelerationStructureKHR accelerationStructure);
+  void SetDescriptor(WrappedVulkan *driver, const VkDescriptorGetInfoEXT &desc);
 
   // 48-bit truncated VK_WHOLE_SIZE
   static const VkDeviceSize WholeSizeRange = 0xFFFFFFFFFFFF;
@@ -828,7 +829,9 @@ struct DescriptorSetSlot
   // normal descriptors.
   DescriptorSlotType type : 8;
   // used for images, the image layout
-  DescriptorSlotImageLayout imageLayout : 8;
+  // aliased with format for descriptor-based texel buffers - to preserve bitfield packing we use
+  // the more useful enum of the image layout and will do a direct cast for formats
+  DescriptorSlotImageLayout imageLayoutOrFormat : 8;
 
 #if GCC_WORKAROUND
 #pragma pack(pop)

renderdoc/driver/vulkan/vk_core.cpp

@@ -103,6 +103,40 @@ void VkInitParams::Set(const VkInstanceCreateInfo *pCreateInfo, ResourceId inst)
   InstanceID = inst;
 }
 
+bool DescriptorTrieNode::operator==(const DescriptorTrieNode &o) const
+{
+  // allow samplers to alias as drivers may deduplicate these. We will still verify the match by
+  // checking that the descriptor bytes for the resulting sampler comes out the same.
+  if(type == DescriptorSlotType::Sampler && o.type == DescriptorSlotType::Sampler)
+    return true;
+
+  // allow a NULL descriptor of different types to alias
+  if(resource == ResourceId() && sampler == ResourceId() && o.resource == ResourceId() &&
+     sampler == ResourceId())
+    return true;
+
+  // allow a NULL combined image/sampler to alias with a sampler
+  if((type == DescriptorSlotType::Sampler && o.type == DescriptorSlotType::SampledImage) ||
+     (type == DescriptorSlotType::SampledImage && o.type == DescriptorSlotType::Sampler))
+  {
+    if(resource == ResourceId() && o.resource == ResourceId())
+      return true;
+  }
+
+  if(type != o.type || resource != o.resource || sampler != o.sampler || offset != o.offset)
+    return false;
+
+  // deliberately allow imageLayout differences to be considered equal still - some drivers are
+  // likely to ignore imageLayout for the descriptor bytes even if the feature is not enabled
+
+  if((range & rangeToleranceMask) != (o.range & rangeToleranceMask))
+    return false;
+
+  return true;
+}
+
+uint64_t DescriptorTrieNode::rangeToleranceMask = ~0ULL;
+
 WrappedVulkan::WrappedVulkan()
 {
   RenderDoc::Inst().RegisterMemoryRegion(this, sizeof(WrappedVulkan));
@@ -5383,6 +5417,12 @@ WrappedVulkan::CommandBufferNode *WrappedVulkan::GetCommandBufferPartialSubmissi
   return NULL;
 }
 
+ResourceId WrappedVulkan::GetASFromAddr(VkDeviceAddress addr)
+{
+  SCOPED_LOCK(m_ASLookupByAddrLock);
+  return m_ASLookupByAddr[addr];
+}
+
 size_t WrappedVulkan::DescriptorDataSize(VkDescriptorType type)
 {
   size_t ret = 0;
@@ -5414,6 +5454,11 @@ size_t WrappedVulkan::DescriptorDataSize(VkDescriptorType type)
   return ret;
 }
 
+void WrappedVulkan::RegisterDescriptor(const bytebuf &key, const DescriptorSetSlot &data)
+{
+  m_DescriptorLookup.fallback.insert(key, data);
+}
+
 bool WrappedVulkan::IsPartialRenderPassActive()
 {
   for(const CommandBufferNode &cmdNode : m_Partial.partialStack)

renderdoc/driver/vulkan/vk_core.h

@@ -26,6 +26,7 @@
 
 #include "common/timing.h"
 #include "core/gpu_address_range_tracker.h"
+#include "core/rdcbytetrie.h"
 #include "serialise/serialiser.h"
 #include "vk_acceleration_structure.h"
 #include "vk_common.h"
@@ -281,6 +282,19 @@ struct UserDebugUtilsCallbackData
   VkDebugUtilsMessengerEXT realObject;
 };
 
+struct DescriptorTrieNode : DescriptorSetSlot
+{
+  DescriptorTrieNode() = default;
+  DescriptorTrieNode(const DescriptorSetSlot &slot) : DescriptorSetSlot(slot), _trie(0) {}
+  uint16_t _trie;
+
+  // this equality operator allows a tolerance of range to account for implementations that drop
+  // lower bits off sizes (the alignment requirements on offsets takes care of that)
+  bool operator==(const DescriptorTrieNode &o) const;
+
+  static uint64_t rangeToleranceMask;
+};
+
 class WrappedVulkan : public IFrameCapturer
 {
 private:
@@ -453,6 +467,16 @@ private:
   Threading::CriticalSection m_ASLookupByAddrLock;
   rdcflatmap<VkDeviceAddress, ResourceId> m_ASLookupByAddr;
 
+  struct DescriptorLookups
+  {
+    // overall lookup of all descriptors by bytes, fallback in case any others don't work - we
+    // expect this to always hit
+    rdcbytetrie<DescriptorTrieNode> fallback;
+  };
+  DescriptorLookups m_DescriptorLookup;
+
+  void RegisterDescriptor(const bytebuf &key, const DescriptorSetSlot &data);
+
   bool m_NULLDescriptorPatternSaved = false;
   bool m_IgnoreLayoutForDescriptors = false;
 
@@ -1293,6 +1317,8 @@ public:
   void UntrackBufferAddress(VkDevice device, VkBuffer buffer);
   void GetResIDFromAddr(GPUAddressRange::Address addr, ResourceId &id, uint64_t &offs);
 
+  ResourceId GetASFromAddr(VkDeviceAddress addr);
+
   EventFlags GetEventFlags(uint32_t eid) { return m_EventFlags[eid]; }
   rdcarray<EventUsage> GetUsage(ResourceId id) { return m_ResourceUses[id]; }
   // return the pre-selected device and queue

renderdoc/driver/vulkan/vk_info.cpp

@@ -739,7 +739,7 @@ bool CreateDescriptorWritesForSlotData(WrappedVulkan *vk, rdcarray<VkWriteDescri
         else
           writeImage[arrayIdx].sampler = VK_NULL_HANDLE;
 
-        writeImage[arrayIdx].imageLayout = convert(slots[slot].imageLayout);
+        writeImage[arrayIdx].imageLayout = convert(slots[slot].imageLayoutOrFormat);
 
         // if we're not updating a SAMPLER descriptor fill in immutable samplers so that
         // our

renderdoc/driver/vulkan/vk_replay.cpp

@@ -2377,7 +2377,7 @@ void VulkanReplay::FillDescriptor(Descriptor &dstel, const DescriptorSetSlot &sr
         dstel.firstSlice = dstel.numSlices = 0;
 
       // cheeky hack, store image layout enum in byteOffset as it's not used for images
-      dstel.byteOffset = convert(srcel.imageLayout);
+      dstel.byteOffset = convert(srcel.imageLayoutOrFormat);
 
       dstel.minLODClamp = c.m_ImageView[viewid].minLOD;
     }

renderdoc/driver/vulkan/vk_serialise.cpp

@@ -5602,14 +5602,14 @@ void DoSerialise(SerialiserType &ser, DescriptorSetSlot &el)
     {
       VkImageLayout imageLayout;
       if(ser.IsWriting())
-        imageLayout = convert(el.imageLayout);
+        imageLayout = convert(el.imageLayoutOrFormat);
       SERIALISE_ELEMENT(imageLayout);
       if(ser.IsReading())
-        el.imageLayout = convert(imageLayout);
+        el.imageLayoutOrFormat = convert(imageLayout);
     }
     else if(ser.IsReading())
     {
-      el.imageLayout = DescriptorSlotImageLayout::Undefined;
+      el.imageLayoutOrFormat = DescriptorSlotImageLayout::Undefined;
     }
 
     // serialise buffer range, for buffer types and inline block
@@ -5670,7 +5670,7 @@ void DoSerialise(SerialiserType &ser, DescriptorSetSlot &el)
     {
       el.resource = imageInfo.imageView;
       el.sampler = imageInfo.sampler;
-      el.imageLayout = convert(imageInfo.imageLayout);
+      el.imageLayoutOrFormat = convert(imageInfo.imageLayout);
     }
     else
     {

renderdoc/driver/vulkan/vk_state.cpp

@@ -1137,7 +1137,7 @@ void VulkanRenderState::BindDescriptorSet(WrappedVulkan *vk, const DescSetLayout
 
         for(uint32_t a = 0; a < push.descriptorCount; a++)
         {
-          dst[a].imageLayout = convert(slots[a].imageLayout);
+          dst[a].imageLayout = convert(slots[a].imageLayoutOrFormat);
           dst[a].sampler = Unwrap(rm->GetCurrentHandle<VkSampler>(slots[a].sampler));
           dst[a].imageView = Unwrap(rm->GetCurrentHandle<VkImageView>(slots[a].resource));
 

renderdoc/driver/vulkan/wrappers/vk_descriptor_funcs.cpp

@@ -2008,6 +2008,11 @@ bool WrappedVulkan::Serialise_vkGetDescriptorEXT(SerialiserType &ser, VkDevice d
                        GetPhysDeviceCompatString(false, false).c_str());
       return false;
     }
+
+    DescriptorSetSlot descriptorData = {};
+    descriptorData.SetDescriptor(this, DescriptorInfo);
+
+    RegisterDescriptor(replayDescriptor, descriptorData);
   }
 
   return true;