Record information needed during replay for fast descriptor lookups

* For palettised image/sampler descriptors this includes initialising the
  palette with the opaque data, and for regular image descriptors this means
  recording the address range of images. Buffers we already have their addresses
  which is all we need.

renderdoc/driver/vulkan/vk_common.cpp

@@ -1325,6 +1325,32 @@ void DescriptorSetSlot::SetAccelerationStructure(VkDescriptorType writeType,
   resource = GetResID(accelerationStructure);
 }
 
+void DescriptorSetSlot::SetSampler(ResourceId samplerId)
+{
+  type = DescriptorSlotType::Sampler;
+  sampler = samplerId;
+}
+
+void DescriptorSetSlot::SetImageSampler(VkDescriptorType descType, ResourceId imageView,
+                                        ResourceId samplerId, VkImageLayout layout)
+{
+  type = convert(descType);
+  resource = imageView;
+  sampler = samplerId;
+  imageLayoutOrFormat = convert(layout);
+}
+
+void DescriptorSetSlot::SetBuffer(VkDescriptorType descType, ResourceId buffer,
+                                  uint64_t startOffset, uint64_t size, VkFormat format)
+{
+  type = convert(descType);
+  resource = buffer;
+  offset = startOffset;
+  range = size;
+  imageLayoutOrFormat = DescriptorSlotImageLayout(format & 0xff);
+  RDCASSERT(uint32_t(format) < 0xff, format);
+}
+
 void DescriptorSetSlot::SetDescriptor(WrappedVulkan *driver, const VkDescriptorGetInfoEXT &desc)
 {
   type = convert(desc.type);

renderdoc/driver/vulkan/vk_common.h

@@ -811,7 +811,13 @@ struct DescriptorSetSlot
   void SetTexelBuffer(VkDescriptorType writeType, ResourceId id);
   void SetAccelerationStructure(VkDescriptorType writeType,
                                 VkAccelerationStructureKHR accelerationStructure);
+
   void SetDescriptor(WrappedVulkan *driver, const VkDescriptorGetInfoEXT &desc);
+  void SetSampler(ResourceId samplerId);
+  void SetImageSampler(VkDescriptorType descType, ResourceId imageView, ResourceId samplerId,
+                       VkImageLayout layout);
+  void SetBuffer(VkDescriptorType descType, ResourceId buffer, uint64_t startOffset, uint64_t size,
+                 VkFormat format);
 
   // 48-bit truncated VK_WHOLE_SIZE
   static const VkDeviceSize WholeSizeRange = 0xFFFFFFFFFFFF;
@@ -906,6 +912,118 @@ DECLARE_REFLECTION_STRUCT(DescriptorSetSlot);
 constexpr uint64_t FixedOpaqueDescriptorCaptureSize = 16;
 constexpr uint64_t MaxDescriptorSize = 256;
 
+// pointers are considered to be 48-bit only, as some descriptors only store those and no-one
+// uses the upper bits relevantly
+//
+// byteSize is obvious, elemSize is the size in elements (N byte texels for texel buffers, bytes
+// for everything else)
+//
+// sizes with a numerical suffix indicates alignment e.g byteSize64 = AlignUp(byteSize, 64)
+//
+// we only care about enough identifiable information is to 'read' a descriptor. We do _not_
+// expect to be able to fully predict the bits in a descriptor. Some of these formats leave no
+// bits unspecified but many have implementation-defined bits, we are only using this for lookup.
+enum class BufferDescriptorFormat
+{
+  UnknownBufferDescriptor = 0,
+
+  // 8 bytes:
+  //    uint64[0] = pointer
+  Pointer_8,
+
+  // 8 bytes:
+  //    bottom 45 = pointer>>4
+  //    top_19 = byteSize16>>4
+  Packed_4519_Aligned16_8,
+
+  // 8 bytes:
+  //    bottom 45 = pointer>>4
+  //    top_19 = byteSize256>>4
+  Packed_4519_Aligned256_8,
+
+  // 16 bytes:
+  //    uint64[0] = pointer
+  //    uint64[1] = elemSize
+  Pointer_ElemSize_16,
+
+  // 16 bytes:
+  //    uint64[0] = pointer/texelSize - special handling for NPOT texelSize (12 bytes)
+  //    uint64[1] = elemSize
+  PointerDivided_ElemSize_16,
+
+  // 16 bytes:
+  //    uint64[0] = pointer
+  Pointer0_16,
+
+  // 32 bytes:
+  //    uint64[0] = byteSize<<32
+  //    uint64[1] = pointer
+  ByteSize0_Pointer1_32,
+
+  // 32 bytes:
+  //    uint64[1] = pointer
+  Pointer1_32,
+
+  // 64 bytes:
+  //    uint64[4] = pointer
+  //    uint64[5] = byteSize << 32
+  Pointer4_ByteSize5_Unaligned_64,
+
+  // 64 bytes:
+  //    uint64[4] = pointer
+  //    uint64[5] = byteSize64 << 32
+  Pointer4_ByteSize5_Aligned_64,
+
+  // 64 bytes:
+  //    uint64[1] = bitScattered(elemSize-1) - complex bit scattering and masking
+  //    uint64[4] = pointer
+  ElemSizeScattered1_Pointer4_64,
+
+  // 64*n bytes: for n=1 or n=2 repeated with different strides (4, 2, 4+2, 2+4, 2+1, 1+2)
+  //    uint64[0] = (byteSize >> stride) << 32
+  //    uint64[1] = ((pointer&0x3f) >> stride) << 16
+  //    uint64[2] = pointer64
+  Strided4_MultiDescriptor_64,
+  Strided2_MultiDescriptor_64,
+  Strided1_MultiDescriptor_64,
+
+  // 64 bytes:
+  //    uint64[0] = elemSize << 32
+  //    uint64[2] = pointer
+  ElemSize0_Pointer2_64,
+
+  // 64 bytes:
+  //    uint64[2] = pointer
+  Pointer2_64,
+};
+
+//
+enum class ImageDescriptorFormat
+{
+  UnknownImageDescriptor = 0,
+
+  // 4 bytes:
+  //   bottom 20 bits = imageViewIndex (stored in opaque capture data, sampled/storage/input), 0
+  //   if absent top 12 bits = samplerIndex (stored in opaque capture data), 0 if absent
+  Indexed2012,
+
+  // 32 bytes:
+  //    uint64[0] = pointer>>8
+  PointerShifted_32,
+
+  // 64 bytes:
+  //    uint64[0] = pointer>>8
+  PointerShifted_64,
+
+  // 64 bytes:
+  //    uint64[2] = pointer
+  Pointer2_64,
+
+  // 64 bytes:
+  //    uint64[4] = pointer
+  Pointer4_64,
+};
+
 #define NUM_VK_IMAGE_ASPECTS 4
 #define VK_ACCESS_ALL_READ_BITS                                                        \
   (VK_ACCESS_INDIRECT_COMMAND_READ_BIT | VK_ACCESS_INDEX_READ_BIT |                    \

renderdoc/driver/vulkan/vk_core.cpp

@@ -5466,6 +5466,65 @@ void WrappedVulkan::RegisterDescriptor(const bytebuf &key, const DescriptorSetSl
     if(data.resource != ResourceId() && !m_DescriptorLookup.texelFormats.contains(fmt))
       m_DescriptorLookup.texelFormats.push_back(fmt);
   }
+  else if(data.type == DescriptorSlotType::CombinedImageSampler ||
+          data.type == DescriptorSlotType::SampledImage ||
+          data.type == DescriptorSlotType::StorageImage ||
+          data.type == DescriptorSlotType::InputAttachment)
+  {
+    VkImageLayout layout = convert(data.imageLayoutOrFormat);
+
+    if(layout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL ||
+       layout == VK_IMAGE_LAYOUT_DEPTH_READ_ONLY_STENCIL_ATTACHMENT_OPTIMAL ||
+       layout == VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_STENCIL_READ_ONLY_OPTIMAL ||
+       layout == VK_IMAGE_LAYOUT_DEPTH_READ_ONLY_OPTIMAL ||
+       layout == VK_IMAGE_LAYOUT_STENCIL_READ_ONLY_OPTIMAL)
+    {
+      if(!m_DescriptorLookup.depthImageLayouts.contains(layout))
+        m_DescriptorLookup.depthImageLayouts.push_back(layout);
+    }
+    else
+    {
+      if(!m_DescriptorLookup.generalImageLayouts.contains(layout))
+      {
+        m_DescriptorLookup.generalImageLayouts.push_back(layout);
+        // keep the list sorted so that rare/niche layouts like feedback loop or local read are tried last
+        std::sort(m_DescriptorLookup.generalImageLayouts.begin(),
+                  m_DescriptorLookup.generalImageLayouts.end());
+      }
+    }
+  }
+
+  size_t combinedSize = m_DescriptorBufferProperties.combinedImageSamplerDescriptorSize;
+  size_t sampledSize = m_DescriptorBufferProperties.sampledImageDescriptorSize;
+  size_t samplerSize = m_DescriptorBufferProperties.samplerDescriptorSize;
+
+  // if this is just a sampler descriptor, store it directly (unless we're not using indexed)
+  if(data.type == DescriptorSlotType::Sampler &&
+     m_DescriptorLookup.sampled != ImageDescriptorFormat::Indexed2012)
+  {
+    m_DescriptorLookup.samplers.insert(key, data);
+  }
+  // if this is a combined descriptor but it looks like it's an image+sampler (which is common) and
+  // we're not indexed, store the second part as sampler bytes. This may be wrong, but that's fine
+  // and worst case we pollute the samplers lookup and fail to do a fast lookup of this descriptor
+  else if(data.type == DescriptorSlotType::CombinedImageSampler &&
+          combinedSize == m_DescriptorLookup.combinedSamplerOffset + samplerSize)
+  {
+    DescriptorSetSlot samplerData;
+    samplerData.SetSampler(data.sampler);
+
+    m_DescriptorLookup.samplers.insert(
+        {key.data() + m_DescriptorLookup.combinedSamplerOffset, samplerSize}, samplerData);
+  }
+
+  if((data.type == DescriptorSlotType::InputAttachment ||
+      data.type == DescriptorSlotType::StorageImage || data.type == DescriptorSlotType::SampledImage ||
+      data.type == DescriptorSlotType::CombinedImageSampler) &&
+     m_DescriptorLookup.sampled != ImageDescriptorFormat::Indexed2012)
+  {
+    m_CreationInfo.m_Image[m_CreationInfo.m_ImageView[data.resource].image].viewDescriptors.push_back(
+        {bytebuf(key.data(), sampledSize), data.resource});
+  }
 }
 
 bool WrappedVulkan::IsPartialRenderPassActive()

renderdoc/driver/vulkan/vk_core.h

@@ -24,6 +24,7 @@
 
 #pragma once
 
+#include "common/formatting.h"
 #include "common/timing.h"
 #include "core/gpu_address_range_tracker.h"
 #include "core/rdcbytetrie.h"
@@ -469,14 +470,43 @@ private:
 
   struct DescriptorLookups
   {
+    BufferDescriptorFormat uniformBuffer = BufferDescriptorFormat::UnknownBufferDescriptor;
+    BufferDescriptorFormat storageBuffer = BufferDescriptorFormat::UnknownBufferDescriptor;
+    BufferDescriptorFormat uniformTexelBuffer = BufferDescriptorFormat::UnknownBufferDescriptor;
+    BufferDescriptorFormat storageTexelBuffer = BufferDescriptorFormat::UnknownBufferDescriptor;
+
+    BufferDescriptorFormat accelStructure = BufferDescriptorFormat::UnknownBufferDescriptor;
+
+    ImageDescriptorFormat sampled = ImageDescriptorFormat::UnknownImageDescriptor;
+    ImageDescriptorFormat storage = ImageDescriptorFormat::UnknownImageDescriptor;
+
+    uint32_t combinedSamplerOffset = 0;
+
     // overall lookup of all descriptors by bytes, fallback in case any others don't work - we
     // expect this to always hit
     rdcbytetrie<DescriptorTrieNode> fallback;
 
+    // lookup with only samplers, as we expect for non-indexed descriptors this will be hit often
+    rdcbytetrie<DescriptorTrieNode> samplers;
+
+    // for implementations where image descriptors are expected to contain a pointer to the image.
+    // We use a _range_ tracker here because some descriptors like depth/stencil or planar formats
+    // can contain base addresses different to the simple base of the image
+    GPUAddressRangeTracker imageAddresses;
+
+    // for NV-style palettised sampler/image view descriptors. These will be resized to the max size
+    // (0xfff / 0xfffff respectively) and can be used for direct indexed lookup
+    rdcarray<ResourceId> samplerPalette;
+    rdcarray<ResourceId> imageViewPalette;
+
     // unique texel formats. So that if we fast identify a buffer via address+size we can iterate
     // over all of these if we know it's a texel buffer. The expectation is this is short so we
     // don't have to store this per-buffer but globally and can just try different possibilities.
     rdcarray<VkFormat> texelFormats;
+
+    // unique image layouts. In case image layout affects the descriptor bits
+    rdcarray<VkImageLayout> generalImageLayouts;
+    rdcarray<VkImageLayout> depthImageLayouts;
   };
   DescriptorLookups m_DescriptorLookup;
 

renderdoc/driver/vulkan/vk_info.cpp

@@ -2404,6 +2404,8 @@ void VulkanCreationInfo::Image::Init(VulkanResourceManager *resourceMan, VulkanC
     creationFlags |= TextureCategory::ShaderReadWrite;
 
   cube = (pCreateInfo->flags & VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT) ? true : false;
+
+  address = 0;
 }
 
 void VulkanCreationInfo::Sampler::Init(VulkanResourceManager *resourceMan, VulkanCreationInfo &info,
@@ -2540,6 +2542,8 @@ void VulkanCreationInfo::ImageView::Init(VulkanResourceManager *resourceMan, Vul
   {
     minLOD = minLODInfo->minLod;
   }
+
+  isDepthImage = !!(info.m_Image[image].creationFlags & TextureCategory::DepthTarget);
 }
 
 void VulkanCreationInfo::ShaderModule::Init(VulkanResourceManager *resourceMan,

renderdoc/driver/vulkan/vk_info.h

@@ -643,7 +643,22 @@ struct VulkanCreationInfo
     bool cube;
     TextureCategory creationFlags;
 
+    VkDeviceAddress address;
     VkMemoryRequirements mrq;
+
+    rdcarray<rdcpair<bytebuf, ResourceId>> viewDescriptors;
+
+    ResourceId getViewFromDescriptor(const byte *descriptorBytes, size_t descriptorSize)
+    {
+      for(auto it = viewDescriptors.begin(); it != viewDescriptors.end(); ++it)
+      {
+        if(it->first.size() == descriptorSize &&
+           memcmp(it->first.data(), descriptorBytes, descriptorSize) == 0)
+          return it->second;
+      }
+
+      return ResourceId();
+    }
   };
   std::unordered_map<ResourceId, Image> m_Image;
 
@@ -711,6 +726,8 @@ struct VulkanCreationInfo
     VkImageSubresourceRange range;
     VkComponentMapping componentMapping;
 
+    bool isDepthImage;
+
     // VkImageViewMinLodCreateInfoEXT
     float minLOD;
   };

renderdoc/driver/vulkan/wrappers/vk_misc_funcs.cpp

@@ -669,6 +669,29 @@ bool WrappedVulkan::Serialise_vkCreateSampler(SerialiserType &ser, VkDevice devi
 
         m_CreationInfo.m_Sampler[live].Init(GetResourceManager(), m_CreationInfo, &CreateInfo);
       }
+
+      // if we're using indexed descriptors then look for the opaque info
+      if(DescriptorBuffers() && m_DescriptorLookup.sampled == ImageDescriptorFormat::Indexed2012)
+      {
+        // if we have opaque capture data that's 8 bytes and isn't 0, assume it's the image's
+        // address. If we guess wrong here this won't be bad necessarily it would just break the
+        // fast descriptor lookup for images
+        VkOpaqueCaptureDescriptorDataCreateInfoEXT *opaque =
+            (VkOpaqueCaptureDescriptorDataCreateInfoEXT *)FindNextStruct(
+                &CreateInfo, VK_STRUCTURE_TYPE_OPAQUE_CAPTURE_DESCRIPTOR_DATA_CREATE_INFO_EXT);
+        if(opaque && m_DescriptorBufferProperties.samplerCaptureReplayDescriptorDataSize == 4)
+        {
+          uint32_t idx = *(uint32_t *)opaque->opaqueCaptureDescriptorData;
+
+          if(idx)
+          {
+            if(idx < m_DescriptorLookup.samplerPalette.size())
+              m_DescriptorLookup.samplerPalette[idx] = GetResID(samp);
+            else
+              RDCERR("Invalid saved index %u", idx);
+          }
+        }
+      }
     }
 
     AddResource(Sampler, ResourceType::Sampler, "Sampler");

renderdoc/driver/vulkan/wrappers/vk_resource_funcs.cpp

@@ -449,6 +449,11 @@ bool WrappedVulkan::Serialise_vkAllocateMemory(SerialiserType &ser, VkDevice dev
             VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
         };
 
+        if(DescriptorBuffers())
+        {
+          bufInfo.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
+        }
+
         ret = ObjDisp(device)->CreateBuffer(Unwrap(device), &bufInfo, NULL, &buf);
         RDCASSERTEQUAL(ret, VK_SUCCESS);
 
@@ -1698,6 +1703,40 @@ bool WrappedVulkan::Serialise_vkBindImageMemory(SerialiserType &ser, VkDevice de
     m_CreationInfo.m_Memory[GetResID(memory)].BindMemory(
         memoryOffset, mrq.size,
         imgInfo.linear ? VulkanCreationInfo::Memory::Linear : VulkanCreationInfo::Memory::Tiled);
+
+    // try to determine image's address via best effort for descriptor buffer lookups if needed
+    if(DescriptorBuffers() && m_DescriptorLookup.sampled != ImageDescriptorFormat::Indexed2012)
+    {
+      VkDeviceAddress addr = imgInfo.address;
+
+      if(addr == 0)
+      {
+        const VulkanCreationInfo::Memory &memInfo = m_CreationInfo.m_Memory[GetResID(memory)];
+
+        if(memInfo.wholeMemBuf != VK_NULL_HANDLE)
+        {
+          VkBufferDeviceAddressInfo getInfo = {
+              VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
+              NULL,
+              Unwrap(memInfo.wholeMemBuf),
+          };
+
+          addr = ObjDisp(device)->GetBufferDeviceAddress(Unwrap(device), &getInfo) + memoryOffset;
+        }
+        else if(memInfo.opaqueAddr)
+        {
+          addr = memInfo.opaqueAddr + memoryOffset;
+          RDCWARN("Using opaque address %llx to estimate as base address for image %s", addr,
+                  ToStr(GetResID(image)).c_str());
+        }
+      }
+
+      if(addr != 0)
+        m_DescriptorLookup.imageAddresses.AddTo(
+            {addr, addr + mrq.size, addr + mrq.size, GetResID(image)});
+      else
+        RDCLOG("Couldn't get base address for image %s", ToStr(GetResID(image)).c_str());
+    }
   }
 
   return true;
@@ -2427,6 +2466,20 @@ bool WrappedVulkan::Serialise_vkCreateImage(SerialiserType &ser, VkDevice device
       m_CreationInfo.m_Image[live].Init(GetResourceManager(), m_CreationInfo, &CreateInfo,
                                         memoryRequirements);
 
+      // if we have opaque capture data that's 8 bytes and isn't 0, assume it's the image's address.
+      // If we guess wrong here this won't be bad necessarily it would just break the fast
+      // descriptor lookup for images
+      VkOpaqueCaptureDescriptorDataCreateInfoEXT *opaque =
+          (VkOpaqueCaptureDescriptorDataCreateInfoEXT *)FindNextStruct(
+              &CreateInfo, VK_STRUCTURE_TYPE_OPAQUE_CAPTURE_DESCRIPTOR_DATA_CREATE_INFO_EXT);
+      if(opaque && m_DescriptorBufferProperties.imageCaptureReplayDescriptorDataSize == 8)
+      {
+        uint64_t *ptr = (uint64_t *)opaque->opaqueCaptureDescriptorData;
+
+        if(*ptr != 0)
+          m_CreationInfo.m_Image[live].address = *ptr;
+      }
+
       bool inserted = false;
       auto state = InsertImageState(img, live, CreateInfo, eFrameRef_Unknown, &inserted);
       if(!inserted)
@@ -2978,6 +3031,37 @@ bool WrappedVulkan::Serialise_vkCreateImageView(SerialiserType &ser, VkDevice de
       }
     }
 
+    // if we're using indexed descriptors then look for the opaque info
+    if(DescriptorBuffers() && m_DescriptorLookup.sampled == ImageDescriptorFormat::Indexed2012)
+    {
+      // if we have opaque capture data that's 8 bytes and isn't 0, assume it's the image's address.
+      // If we guess wrong here this won't be bad necessarily it would just break the fast
+      // descriptor lookup for images
+      VkOpaqueCaptureDescriptorDataCreateInfoEXT *opaque =
+          (VkOpaqueCaptureDescriptorDataCreateInfoEXT *)FindNextStruct(
+              &CreateInfo, VK_STRUCTURE_TYPE_OPAQUE_CAPTURE_DESCRIPTOR_DATA_CREATE_INFO_EXT);
+      if(opaque && m_DescriptorBufferProperties.imageViewCaptureReplayDescriptorDataSize >= 4)
+      {
+        const uint32_t numIndices =
+            uint32_t(m_DescriptorBufferProperties.imageViewCaptureReplayDescriptorDataSize / 4);
+        uint32_t *ptr = (uint32_t *)opaque->opaqueCaptureDescriptorData;
+
+        // these would be expected to be sampled, storage (if possible), input (if possible)
+        // but we only need to know the image view itself as we always know the descriptor type when
+        // decoding due to varying size etc
+        for(uint32_t i = 0; i < numIndices; i++)
+        {
+          if(ptr[i])
+          {
+            if(ptr[0] < m_DescriptorLookup.imageViewPalette.size())
+              m_DescriptorLookup.imageViewPalette[ptr[i]] = GetResID(view);
+            else
+              RDCERR("Invalid saved index %u", ptr[i]);
+          }
+        }
+      }
+    }
+
     AddResource(View, ResourceType::View, "Image View");
     DerivedResource(device, View);
     DerivedResource(CreateInfo.image, View);
@@ -3292,6 +3376,42 @@ bool WrappedVulkan::Serialise_vkBindImageMemory2(SerialiserType &ser, VkDevice d
         m_CreationInfo.m_Memory[GetResID(bindInfo.memory)].BindMemory(
             bindInfo.memoryOffset, mrq.size,
             imgInfo.linear ? VulkanCreationInfo::Memory::Linear : VulkanCreationInfo::Memory::Tiled);
+
+        // try to determine image's address via best effort for descriptor buffer lookups if needed
+        if(DescriptorBuffers() && m_DescriptorLookup.sampled != ImageDescriptorFormat::Indexed2012)
+        {
+          VkDeviceAddress addr = imgInfo.address;
+
+          if(addr == 0)
+          {
+            const VulkanCreationInfo::Memory &memInfo =
+                m_CreationInfo.m_Memory[GetResID(bindInfo.memory)];
+
+            if(memInfo.wholeMemBuf != VK_NULL_HANDLE)
+            {
+              VkBufferDeviceAddressInfo getInfo = {
+                  VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
+                  NULL,
+                  Unwrap(memInfo.wholeMemBuf),
+              };
+
+              addr = ObjDisp(device)->GetBufferDeviceAddress(Unwrap(device), &getInfo) +
+                     bindInfo.memoryOffset;
+            }
+            else if(memInfo.opaqueAddr)
+            {
+              addr = memInfo.opaqueAddr + bindInfo.memoryOffset;
+              RDCWARN("Using opaque address to estimate as base address for image %s",
+                      ToStr(GetResID(bindInfo.image)).c_str());
+            }
+          }
+
+          if(addr != 0)
+            m_DescriptorLookup.imageAddresses.AddTo(
+                {addr, addr + mrq.size, addr + mrq.size, GetResID(bindInfo.image)});
+          else
+            RDCLOG("Couldn't get base address for image %s", ToStr(GetResID(bindInfo.image)).c_str());
+        }
       }
     }