Re-use descriptors in descriptor pool

* If an application allocates from and resets descriptor pools at very high
  frequency the overhead of freeing and reallocating those descriptor sets can
  be high. Instead use the descriptor pool as a pool for children and look up
  the freelist first for an existing descriptor set before trying to allocate a
  new one.

renderdoc/driver/vulkan/vk_common.h

@@ -450,6 +450,11 @@ struct DescriptorSetSlot
 
 struct BindingStorage
 {
+  BindingStorage() = default;
+  // disallow copy
+  BindingStorage(const BindingStorage &) = delete;
+  BindingStorage &operator=(const BindingStorage &) = delete;
+
   ~BindingStorage() { clear(); }
   bytebuf inlineBytes;
   rdcarray<DescriptorSetSlot *> binds;
@@ -467,6 +472,18 @@ struct BindingStorage
     memset(elems.data(), 0, elems.byteSize());
   }
 
+  void copy(DescriptorSetSlot *&slots, uint32_t &slotCount, byte *&inlineData, size_t &inlineSize)
+  {
+    slotCount = elems.count();
+
+    slots = new DescriptorSetSlot[slotCount];
+    memcpy(slots, elems.data(), sizeof(DescriptorSetSlot) * slotCount);
+
+    inlineSize = inlineBytes.size();
+    inlineData = AllocAlignedBuffer(inlineSize);
+    memcpy(inlineData, inlineBytes.data(), inlineSize);
+  }
+
 private:
   rdcarray<DescriptorSetSlot> elems;
   friend struct DescSetLayout;

renderdoc/driver/vulkan/vk_core.cpp

@@ -1557,6 +1557,8 @@ void WrappedVulkan::StartFrameCapture(void *dev, void *wnd)
 
   RDCLOG("Starting capture");
 
+  Atomic::Dec32(&m_ReuseEnabled);
+
   m_CaptureTimer.Restart();
 
   GetResourceManager()->ResetCaptureStartTime();
@@ -2060,6 +2062,8 @@ bool WrappedVulkan::EndFrameCapture(void *dev, void *wnd)
 
   m_CmdBufferRecords.clear();
 
+  Atomic::Inc32(&m_ReuseEnabled);
+
   GetResourceManager()->ResetLastWriteTimes();
 
   GetResourceManager()->ResetLastPartialUseTimes();
@@ -2109,6 +2113,8 @@ bool WrappedVulkan::DiscardFrameCapture(void *dev, void *wnd)
     }
   }
 
+  Atomic::Inc32(&m_ReuseEnabled);
+
   SAFE_DELETE(m_HeaderChunk);
 
   // delete cmd buffers now - had to keep them alive until after serialiser flush.

renderdoc/driver/vulkan/vk_core.h

@@ -295,6 +295,8 @@ private:
   CaptureState m_State;
   bool m_AppControlledCapture = false;
 
+  volatile int32_t m_ReuseEnabled = 1;
+
   PerformanceTimer m_CaptureTimer;
 
   bool m_MarkedActive = false;

renderdoc/driver/vulkan/vk_initstate.cpp

@@ -57,37 +57,8 @@ bool WrappedVulkan::Prepare_InitialState(WrappedVkRes *res)
 
     if((layout.flags & VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT_KHR) == 0)
     {
-      for(size_t i = 0; i < layout.bindings.size(); i++)
-      {
-        if(layout.bindings[i].descriptorType == VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT)
-          initialContents.numDescriptors++;
-        else
-          initialContents.numDescriptors += layout.bindings[i].descriptorCount;
-      }
-
-      initialContents.descriptorSlots = new DescriptorSetSlot[initialContents.numDescriptors];
-      RDCEraseMem(initialContents.descriptorSlots,
-                  sizeof(DescriptorSetSlot) * initialContents.numDescriptors);
-
-      uint32_t e = 0;
-      for(size_t i = 0; i < layout.bindings.size(); i++)
-      {
-        if(layout.bindings[i].descriptorType == VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT)
-        {
-          initialContents.descriptorSlots[e++] = record->descInfo->data.binds[i][0];
-        }
-        else
-        {
-          for(uint32_t b = 0; b < layout.bindings[i].descriptorCount; b++)
-          {
-            initialContents.descriptorSlots[e++] = record->descInfo->data.binds[i][b];
-          }
-        }
-      }
-
-      initialContents.inlineData = AllocAlignedBuffer(record->descInfo->data.inlineBytes.size());
-      memcpy(initialContents.inlineData, record->descInfo->data.inlineBytes.data(),
-             record->descInfo->data.inlineBytes.size());
+      record->descInfo->data.copy(initialContents.descriptorSlots, initialContents.numDescriptors,
+                                  initialContents.inlineData, initialContents.inlineByteSize);
     }
     else
     {
@@ -555,22 +526,7 @@ uint64_t WrappedVulkan::GetSize_InitialState(ResourceId id, const VkInitialConte
 {
   if(initial.type == eResDescriptorSet)
   {
-    VkResourceRecord *record = GetResourceManager()->GetResourceRecord(id);
-
-    RDCASSERT(record->descInfo && record->descInfo->layout);
-    const DescSetLayout &layout = *record->descInfo->layout;
-
-    uint32_t NumBindings = 0;
-
-    for(size_t i = 0; i < layout.bindings.size(); i++)
-    {
-      if(layout.bindings[i].descriptorType == VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT)
-        NumBindings++;
-      else
-        NumBindings += layout.bindings[i].descriptorCount;
-    }
-
-    return 32 + NumBindings * sizeof(DescriptorSetSlot) + layout.inlineByteSize;
+    return 32 + initial.numDescriptors * sizeof(DescriptorSetSlot) + initial.inlineByteSize;
   }
   else if(initial.type == eResBuffer)
   {
@@ -604,8 +560,8 @@ static rdcliteral NameOfType(VkResourceType type)
 }
 
 template <typename SerialiserType>
-bool WrappedVulkan::Serialise_InitialState(SerialiserType &ser, ResourceId id,
-                                           VkResourceRecord *record, const VkInitialContents *initial)
+bool WrappedVulkan::Serialise_InitialState(SerialiserType &ser, ResourceId id, VkResourceRecord *,
+                                           const VkInitialContents *initial)
 {
   bool ret = true;
 
@@ -631,20 +587,10 @@ bool WrappedVulkan::Serialise_InitialState(SerialiserType &ser, ResourceId id,
     // while writing, fetching binding information from prepared initial contents
     if(ser.IsWriting())
     {
-      RDCASSERT(record->descInfo && record->descInfo->layout);
-      const DescSetLayout &layout = *record->descInfo->layout;
-
       Bindings = initial->descriptorSlots;
+      NumBindings = initial->numDescriptors;
 
-      for(size_t i = 0; i < layout.bindings.size(); i++)
-      {
-        if(layout.bindings[i].descriptorType == VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT)
-          NumBindings++;
-        else
-          NumBindings += layout.bindings[i].descriptorCount;
-      }
-
-      InlineData.assign(initial->inlineData, layout.inlineByteSize);
+      InlineData.assign(initial->inlineData, initial->inlineByteSize);
     }
 
     SERIALISE_ELEMENT_ARRAY(Bindings, NumBindings);

renderdoc/driver/vulkan/vk_manager.h

@@ -151,6 +151,7 @@ struct VkInitialContents
   VkDescriptorBufferInfo *descriptorInfo;
   VkWriteDescriptorSetInlineUniformBlockEXT *inlineInfo;
   byte *inlineData;
+  size_t inlineByteSize;
   uint32_t numDescriptors;
 
   // for plain resources, we store the resource type and memory allocation details of the contents
@@ -325,6 +326,20 @@ public:
     return id;
   }
 
+  template <typename realtype>
+  ResourceId WrapReusedResource(VkResourceRecord *record, realtype &obj)
+  {
+    RDCASSERT(obj != VK_NULL_HANDLE);
+
+    typename UnwrapHelper<realtype>::Outer *wrapped =
+        (typename UnwrapHelper<realtype>::Outer *)record->Resource;
+    wrapped->real = obj;
+
+    obj = realtype((uint64_t)wrapped);
+
+    return wrapped->id;
+  }
+
   template <typename realtype>
   void ReleaseWrappedResource(realtype obj, bool clearID = false)
   {

renderdoc/driver/vulkan/vk_resources.cpp

@@ -3760,6 +3760,9 @@ VkResourceRecord::~VkResourceRecord()
   if(resType == eResPipelineLayout)
     SAFE_DELETE(pipeLayoutInfo);
 
+  if(resType == eResDescriptorPool)
+    SAFE_DELETE(descPoolInfo);
+
   if(resType == eResDescUpdateTemplate)
     SAFE_DELETE(descTemplateInfo);
 }
@@ -3777,7 +3780,7 @@ void VkResourceRecord::MarkImageFrameReferenced(VkResourceRecord *img, const Ima
   ImageSubresourceRange range2(range);
 
   FrameRefType maxRef = MarkImageReferenced(cmdInfo->imageStates, id, img->resInfo->imageInfo,
-                                            range2, pool->queueFamilyIndex, refType);
+                                            range2, VK_QUEUE_FAMILY_IGNORED, refType);
 
   // maintain the reference type of the image itself as the maximum reference type of any
   // subresource
@@ -3822,7 +3825,7 @@ void VkResourceRecord::MarkImageViewFrameReferenced(VkResourceRecord *view, cons
   imgRange.Sanitise(view->resInfo->imageInfo);
 
   FrameRefType maxRef = MarkImageReferenced(cmdInfo->imageStates, img, view->resInfo->imageInfo,
-                                            imgRange, pool->queueFamilyIndex, refType);
+                                            imgRange, VK_QUEUE_FAMILY_IGNORED, refType);
 
   // maintain the reference type of the image itself as the maximum reference type of any
   // subresource

renderdoc/driver/vulkan/vk_resources.h

@@ -1074,6 +1074,11 @@ struct PipelineLayoutData
   rdcarray<DescSetLayout> layouts;
 };
 
+struct DescPoolInfo
+{
+  rdcarray<VkResourceRecord *> freelist;
+};
+
 struct MemMapState
 {
   VkBuffer wholeMemBuf = VK_NULL_HANDLE;
@@ -2164,7 +2169,7 @@ public:
 
     FrameRefType maxRef =
         MarkImageReferenced(descInfo->bindImageStates, view->baseResource, view->resInfo->imageInfo,
-                            ImageSubresourceRange(imgRange), pool->queueFamilyIndex, refType);
+                            ImageSubresourceRange(imgRange), VK_QUEUE_FAMILY_IGNORED, refType);
 
     p.second = ComposeFrameRefsDisjoint(p.second, maxRef);
   }
@@ -2266,6 +2271,7 @@ public:
     PipelineLayoutData *pipeLayoutInfo;      // only for pipeline layouts
     DescriptorSetData *descInfo;             // only for descriptor sets and descriptor set layouts
     DescUpdateTemplate *descTemplateInfo;    // only for descriptor update templates
+    DescPoolInfo *descPoolInfo;              // only for descriptor pools
     uint32_t queueFamilyIndex;               // only for queues and command pools
   };
 

renderdoc/driver/vulkan/wrappers/vk_descriptor_funcs.cpp

@@ -23,6 +23,10 @@
  ******************************************************************************/
 
 #include "../vk_core.h"
+#include "core/settings.h"
+
+RDOC_DEBUG_CONFIG(bool, Vulkan_Debug_AllowDescriptorSetReuse, true,
+                  "Allow the re-use of descriptor sets via vkResetDescriptorPool.");
 
 template <>
 VkDescriptorSetLayoutCreateInfo WrappedVulkan::UnwrapInfo(const VkDescriptorSetLayoutCreateInfo *info)
@@ -243,6 +247,8 @@ VkResult WrappedVulkan::vkCreateDescriptorPool(VkDevice device,
 
       VkResourceRecord *record = GetResourceManager()->AddResourceRecord(*pDescriptorPool);
       record->AddChunk(chunk);
+
+      record->descPoolInfo = new DescPoolInfo;
     }
     else
     {
@@ -477,54 +483,106 @@ VkResult WrappedVulkan::vkAllocateDescriptorSets(VkDevice device,
 
   for(uint32_t i = 0; i < pAllocateInfo->descriptorSetCount; i++)
   {
-    ResourceId id = GetResourceManager()->WrapResource(Unwrap(device), pDescriptorSets[i]);
+    VkResourceRecord *poolrecord = NULL;
+    VkResourceRecord *layoutRecord = NULL;
+
+    ResourceId id;
+    VkResourceRecord *record = NULL;
+    bool exactReuse = false;
 
     if(IsCaptureMode(m_State))
     {
-      Chunk *chunk = NULL;
+      layoutRecord = GetRecord(pAllocateInfo->pSetLayouts[i]);
+      poolrecord = GetRecord(pAllocateInfo->descriptorPool);
 
+      if(Atomic::CmpExch32(&m_ReuseEnabled, 1, 1) == 1)
       {
-        CACHE_THREAD_SERIALISER();
+        rdcarray<VkResourceRecord *> &freelist = poolrecord->descPoolInfo->freelist;
 
-        VkDescriptorSetAllocateInfo info = *pAllocateInfo;
-        info.descriptorSetCount = 1;
-        info.pSetLayouts += i;
+        if(!freelist.empty())
+        {
+          DescSetLayout *search = layoutRecord->descInfo->layout;
 
-        SCOPED_SERIALISE_CHUNK(VulkanChunk::vkAllocateDescriptorSets);
-        Serialise_vkAllocateDescriptorSets(ser, device, &info, &pDescriptorSets[i]);
+          // try to find an exact layout match, then we don't need to re-initialise the descriptor
+          // set.
+          auto it = std::lower_bound(freelist.begin(), freelist.end(), search,
+                                     [](VkResourceRecord *a, DescSetLayout *search) {
+                                       return a->descInfo->layout < search;
+                                     });
 
-        chunk = scope.Get();
+          if(it != freelist.end() && (*it)->descInfo->layout == layoutRecord->descInfo->layout)
+          {
+            record = freelist.takeAt(it - freelist.begin());
+            exactReuse = true;
+          }
+          else
+          {
+            record = freelist.back();
+            freelist.pop_back();
+          }
+
+          if(!exactReuse)
+            record->DeleteChunks();
+        }
       }
+    }
 
-      VkResourceRecord *record = GetResourceManager()->AddResourceRecord(pDescriptorSets[i]);
-      record->AddChunk(chunk);
-
-      ResourceId layoutID = GetResID(pAllocateInfo->pSetLayouts[i]);
-      VkResourceRecord *layoutRecord = GetRecord(pAllocateInfo->pSetLayouts[i]);
-
-      VkResourceRecord *poolrecord = GetRecord(pAllocateInfo->descriptorPool);
+    if(record)
+      id = GetResourceManager()->WrapReusedResource(record, pDescriptorSets[i]);
+    else
+      id = GetResourceManager()->WrapResource(Unwrap(device), pDescriptorSets[i]);
 
+    if(IsCaptureMode(m_State))
+    {
+      if(record == NULL)
       {
+        record = GetResourceManager()->AddResourceRecord(pDescriptorSets[i]);
+
         poolrecord->LockChunks();
         poolrecord->pooledChildren.push_back(record);
         poolrecord->UnlockChunks();
+
+        record->pool = poolrecord;
+
+        // only mark descriptor set as dirty if it's not a push descriptor layout
+        if((layoutRecord->descInfo->layout->flags &
+            VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT_KHR) == 0)
+        {
+          GetResourceManager()->MarkDirtyResource(id);
+        }
+
+        record->descInfo = new DescriptorSetData();
       }
 
-      record->pool = poolrecord;
-
-      record->AddParent(poolrecord);
-      record->AddParent(GetResourceManager()->GetResourceRecord(layoutID));
-
-      // only mark descriptor set as dirty if it's not a push descriptor layout
-      if((layoutRecord->descInfo->layout->flags &
-          VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT_KHR) == 0)
+      if(!exactReuse)
       {
-        GetResourceManager()->MarkDirtyResource(id);
-      }
+        Chunk *chunk = NULL;
 
-      record->descInfo = new DescriptorSetData();
-      record->descInfo->layout = layoutRecord->descInfo->layout;
-      record->descInfo->layout->CreateBindingsArray(record->descInfo->data);
+        {
+          CACHE_THREAD_SERIALISER();
+
+          VkDescriptorSetAllocateInfo info = *pAllocateInfo;
+          info.descriptorSetCount = 1;
+          info.pSetLayouts += i;
+
+          SCOPED_SERIALISE_CHUNK(VulkanChunk::vkAllocateDescriptorSets);
+          Serialise_vkAllocateDescriptorSets(ser, device, &info, &pDescriptorSets[i]);
+
+          chunk = scope.Get();
+        }
+        record->AddChunk(chunk);
+
+        record->FreeParents(GetResourceManager());
+        record->AddParent(poolrecord);
+        record->AddParent(layoutRecord);
+
+        record->descInfo->layout = layoutRecord->descInfo->layout;
+        record->descInfo->layout->CreateBindingsArray(record->descInfo->data);
+      }
+      else
+      {
+        record->descInfo->data.reset();
+      }
     }
     else
     {
@@ -562,18 +620,51 @@ VkResult WrappedVulkan::vkResetDescriptorPool(VkDevice device, VkDescriptorPool
   // need to free all child descriptor pools. Application is responsible for
   // ensuring no concurrent use with alloc/free from this pool, the same as
   // for DestroyDescriptorPool.
-  if(IsCaptureMode(m_State))
   {
-    VkResourceRecord *record = GetRecord(descriptorPool);
+    // don't reset while capture transition lock is held, so that we can't reset and potentially
+    // reuse a record we might be preparing. We do this here rather than in vkAllocateDescriptorSets
+    // where we actually modify the record, since that's much higher frequency
+    SCOPED_READLOCK(m_CapTransitionLock);
 
-    // delete all of the children
-    for(auto it = record->pooledChildren.begin(); it != record->pooledChildren.end(); ++it)
+    if(IsCaptureMode(m_State))
     {
-      // unset record->pool so we don't recurse
-      (*it)->pool = NULL;
-      GetResourceManager()->ReleaseWrappedResource((VkDescriptorSet)(uint64_t)(*it)->Resource, true);
+      VkResourceRecord *record = GetRecord(descriptorPool);
+
+      if(Vulkan_Debug_AllowDescriptorSetReuse())
+      {
+        for(auto it = record->pooledChildren.begin(); it != record->pooledChildren.end(); ++it)
+        {
+          ((WrappedVkNonDispRes *)(*it)->Resource)->real = RealVkRes(0x123456);
+          (*it)->descInfo->data.reset();
+          (*it)->descInfo->bindFrameRefs.clear();
+          (*it)->descInfo->bindMemRefs.clear();
+          (*it)->descInfo->bindImageStates.clear();
+          (*it)->descInfo->backgroundFrameRefs.clear();
+        }
+
+        record->descPoolInfo->freelist.assign(record->pooledChildren);
+
+        // sort by layout
+        std::sort(record->descPoolInfo->freelist.begin(), record->descPoolInfo->freelist.end(),
+                  [](VkResourceRecord *a, VkResourceRecord *b) {
+                    return a->descInfo->layout < b->descInfo->layout;
+                  });
+      }
+      else
+      {
+        // if descriptor set re-use is banned, we can simply free all the sets immediately without
+        // adding them to the free list and that will effectively disallow re-use.
+        for(auto it = record->pooledChildren.begin(); it != record->pooledChildren.end(); ++it)
+        {
+          // unset record->pool so we don't recurse
+          (*it)->pool = NULL;
+          GetResourceManager()->ReleaseWrappedResource((VkDescriptorSet)(uint64_t)(*it)->Resource,
+                                                       true);
+        }
+
+        record->pooledChildren.clear();
+      }
     }
-    record->pooledChildren.clear();
   }
 
   return ObjDisp(device)->ResetDescriptorPool(Unwrap(device), Unwrap(descriptorPool), flags);