Vulkan: Recycle descriptor pools

This PR adds recycling of descriptor pools. Currently descriptor pools
are discarded when full or context is flushed. This PR allows
descriptor pools to be discarded for reuse.
It is also more conservative and only discard
Descriptor pools when they are full or fragmented.

When using the Vulkan backend a small amount of descriptor memory can leak. Even
when we clean up all resources, drivers can still keep data around on the
GPU. Eventually this can lead to out of memory issues depending on how
the GPU driver actually manages descriptor sets.

When the descriptor sets of the descriptor pool aren't used anymore
the VKDiscardPool will recycle the pools back to its original VKDescriptorPools.
It needs to be the same instance as descriptor pools/sets are owned by
a single thread.

Pull Request: https://projects.blender.org/blender/blender/pulls/144992

source/blender/gpu/vulkan/vk_context.cc

@@ -172,9 +172,6 @@ TimelineValue VKContext::flush_render_graph(RenderGraphFlushFlags flags,
   }
   VKDevice &device = VKBackend::get().device;
   descriptor_set_get().upload_descriptor_sets();
-  if (!device.extensions_get().descriptor_buffer) {
-    descriptor_pools_get().discard(*this);
-  }
   TimelineValue timeline = device.render_graph_submit(
       &render_graph_.value().get(),
       discard_pool,

source/blender/gpu/vulkan/vk_descriptor_pools.cc

@@ -12,38 +12,37 @@
 #include "vk_device.hh"
 
 namespace blender::gpu {
-VKDescriptorPools::VKDescriptorPools() {}
 
 VKDescriptorPools::~VKDescriptorPools()
 {
   const VKDevice &device = VKBackend::get().device;
-  for (const VkDescriptorPool vk_descriptor_pool : pools_) {
+  for (const VkDescriptorPool vk_descriptor_pool : recycled_pools_) {
     vkDestroyDescriptorPool(device.vk_handle(), vk_descriptor_pool, nullptr);
   }
+  recycled_pools_.clear();
+  if (vk_descriptor_pool_ != VK_NULL_HANDLE) {
+    vkDestroyDescriptorPool(device.vk_handle(), vk_descriptor_pool_, nullptr);
+    vk_descriptor_pool_ = VK_NULL_HANDLE;
+  }
 }
 
 void VKDescriptorPools::init(const VKDevice &device)
 {
-  BLI_assert(pools_.is_empty());
-  add_new_pool(device);
+  ensure_pool(device);
 }
 
-void VKDescriptorPools::discard(VKContext &context)
+void VKDescriptorPools::ensure_pool(const VKDevice &device)
 {
-  const VKDevice &device = VKBackend::get().device;
-  VKDiscardPool &discard_pool = context.discard_pool;
-
-  for (const VkDescriptorPool vk_descriptor_pool : pools_) {
-    discard_pool.discard_descriptor_pool(vk_descriptor_pool);
+  if (vk_descriptor_pool_ != VK_NULL_HANDLE) {
+    return;
   }
-  pools_.clear();
 
-  add_new_pool(device);
-  active_pool_index_ = 0;
-}
+  std::scoped_lock lock(mutex_);
+  if (!recycled_pools_.is_empty()) {
+    vk_descriptor_pool_ = recycled_pools_.pop_last();
+    return;
+  }
 
-void VKDescriptorPools::add_new_pool(const VKDevice &device)
-{
   Vector<VkDescriptorPoolSize> pool_sizes = {
       {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, POOL_SIZE_STORAGE_BUFFER},
       {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, POOL_SIZE_STORAGE_IMAGE},
@@ -56,44 +55,32 @@ void VKDescriptorPools::add_new_pool(const VKDevice &device)
   pool_info.maxSets = POOL_SIZE_DESCRIPTOR_SETS;
   pool_info.poolSizeCount = pool_sizes.size();
   pool_info.pPoolSizes = pool_sizes.data();
-  VkDescriptorPool descriptor_pool = VK_NULL_HANDLE;
-  VkResult result = vkCreateDescriptorPool(
-      device.vk_handle(), &pool_info, nullptr, &descriptor_pool);
-  UNUSED_VARS(result);
-  pools_.append(descriptor_pool);
+  vkCreateDescriptorPool(device.vk_handle(), &pool_info, nullptr, &vk_descriptor_pool_);
 }
 
-VkDescriptorPool VKDescriptorPools::active_pool_get()
+void VKDescriptorPools::discard_active_pool(VKContext &context)
 {
-  BLI_assert(!pools_.is_empty());
-  return pools_[active_pool_index_];
+  context.discard_pool.discard_descriptor_pool_for_reuse(vk_descriptor_pool_, this);
+  vk_descriptor_pool_ = VK_NULL_HANDLE;
 }
 
-void VKDescriptorPools::activate_next_pool()
+void VKDescriptorPools::recycle(VkDescriptorPool vk_descriptor_pool)
 {
-  BLI_assert(!is_last_pool_active());
-  active_pool_index_ += 1;
-}
-
-void VKDescriptorPools::activate_last_pool()
-{
-  active_pool_index_ = pools_.size() - 1;
-}
-
-bool VKDescriptorPools::is_last_pool_active()
-{
-  return active_pool_index_ == pools_.size() - 1;
+  const VKDevice &device = VKBackend::get().device;
+  vkResetDescriptorPool(device.vk_handle(), vk_descriptor_pool, 0);
+  std::scoped_lock lock(mutex_);
+  recycled_pools_.append(vk_descriptor_pool);
 }
 
 VkDescriptorSet VKDescriptorPools::allocate(const VkDescriptorSetLayout descriptor_set_layout)
 {
   BLI_assert(descriptor_set_layout != VK_NULL_HANDLE);
+  BLI_assert(vk_descriptor_pool_ != VK_NULL_HANDLE);
   const VKDevice &device = VKBackend::get().device;
 
   VkDescriptorSetAllocateInfo allocate_info = {};
-  VkDescriptorPool pool = active_pool_get();
   allocate_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
-  allocate_info.descriptorPool = pool;
+  allocate_info.descriptorPool = vk_descriptor_pool_;
   allocate_info.descriptorSetCount = 1;
   allocate_info.pSetLayouts = &descriptor_set_layout;
   VkDescriptorSet vk_descriptor_set = VK_NULL_HANDLE;
@@ -101,12 +88,10 @@ VkDescriptorSet VKDescriptorPools::allocate(const VkDescriptorSetLayout descript
       device.vk_handle(), &allocate_info, &vk_descriptor_set);
 
   if (ELEM(result, VK_ERROR_OUT_OF_POOL_MEMORY, VK_ERROR_FRAGMENTED_POOL)) {
-    if (is_last_pool_active()) {
-      add_new_pool(device);
-      activate_last_pool();
-    }
-    else {
-      activate_next_pool();
+    {
+      VKContext &context = *VKContext::get();
+      discard_active_pool(context);
+      ensure_pool(device);
     }
     return allocate(descriptor_set_layout);
   }

source/blender/gpu/vulkan/vk_descriptor_pools.hh

@@ -21,9 +21,6 @@ class VKDevice;
  * In Vulkan a pool is constructed with a fixed size per resource type. When more resources are
  * needed it a next pool should be created. VKDescriptorPools will keep track of those pools and
  * construct new pools when the previous one is exhausted.
- *
- * At the beginning of a new frame the descriptor pools are reset. This will start allocating
- * again from the first descriptor pool in order to use freed space from previous pools.
  */
 class VKDescriptorPools {
   /**
@@ -40,30 +37,40 @@ class VKDescriptorPools {
   static constexpr uint32_t POOL_SIZE_UNIFORM_TEXEL_BUFFER = 100;
   static constexpr uint32_t POOL_SIZE_INPUT_ATTACHMENT = 100;
 
-  Vector<VkDescriptorPool> pools_;
-  int64_t active_pool_index_ = 0;
+  /**
+   * Unused recycled pools.
+   *
+   * When a pool is full it is being discarded (for reuse). After all descriptor sets of the pool
+   * are unused the descriptor pool can be reused.
+   * Note: descriptor pools/sets are pinned to a single thread so the pools should always return to
+   * the instance it was created on.
+   */
+  Vector<VkDescriptorPool> recycled_pools_;
+  /** Active descriptor pool. Should always be a valid handle. */
+  VkDescriptorPool vk_descriptor_pool_ = VK_NULL_HANDLE;
+  Mutex mutex_;
 
  public:
-  VKDescriptorPools();
   ~VKDescriptorPools();
 
   void init(const VKDevice &vk_device);
 
+  /**
+   * Allocate a new descriptor set.
+   *
+   * When the active descriptor pool is full it is discarded and another descriptor pool is
+   * ensured.
+   */
   VkDescriptorSet allocate(const VkDescriptorSetLayout descriptor_set_layout);
 
   /**
-   * Discard all existing pools and re-initializes this instance.
-   *
-   * This is a fix to ensure that resources will not be rewritten. Eventually we should discard the
-   * resource pools for reuse.
+   * Recycle a previous discarded descriptor pool.
    */
-  void discard(VKContext &vk_context);
+  void recycle(VkDescriptorPool vk_descriptor_pool);
 
  private:
-  VkDescriptorPool active_pool_get();
-  void activate_next_pool();
-  void activate_last_pool();
-  bool is_last_pool_active();
   void add_new_pool(const VKDevice &device);
+  void discard_active_pool(VKContext &vk_context);
+  void ensure_pool(const VKDevice &device);
 };
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_device.cc

@@ -77,6 +77,8 @@ void VKDevice::deinit()
   samplers_.free();
   GPU_SHADER_FREE_SAFE(vk_backbuffer_blit_sh_);
 
+  orphaned_data_render.deinit(*this);
+  orphaned_data.deinit(*this);
   {
     while (!thread_data_.is_empty()) {
       VKThreadData *thread_data = thread_data_.pop_last();
@@ -87,8 +89,6 @@ void VKDevice::deinit()
   pipelines.write_to_disk();
   pipelines.free_data();
   descriptor_set_layouts_.deinit();
-  orphaned_data_render.deinit(*this);
-  orphaned_data.deinit(*this);
   vmaDestroyPool(mem_allocator_, vma_pools.external_memory);
   vmaDestroyAllocator(mem_allocator_);
   mem_allocator_ = VK_NULL_HANDLE;

source/blender/gpu/vulkan/vk_resource_pool.cc

@@ -98,10 +98,11 @@ void VKDiscardPool::discard_render_pass(VkRenderPass vk_render_pass)
   render_passes_.append_timeline(timeline_, vk_render_pass);
 }
 
-void VKDiscardPool::discard_descriptor_pool(VkDescriptorPool vk_descriptor_pool)
+void VKDiscardPool::discard_descriptor_pool_for_reuse(VkDescriptorPool vk_descriptor_pool,
+                                                      VKDescriptorPools *descriptor_pools)
 {
   std::scoped_lock mutex(mutex_);
-  descriptor_pools_.append_timeline(timeline_, vk_descriptor_pool);
+  descriptor_pools_.append_timeline(timeline_, std::pair(vk_descriptor_pool, descriptor_pools));
 }
 
 void VKDiscardPool::destroy_discarded_resources(VKDevice &device, bool force)
@@ -147,11 +148,10 @@ void VKDiscardPool::destroy_discarded_resources(VKDevice &device, bool force)
     vkDestroyRenderPass(device.vk_handle(), vk_render_pass, nullptr);
   });
 
-  /* TODO: Introduce reuse_old as the allocations can all be reused by resetting the pool. */
-  descriptor_pools_.remove_old(current_timeline, [&](VkDescriptorPool vk_descriptor_pool) {
-    vkResetDescriptorPool(device.vk_handle(), vk_descriptor_pool, 0);
-    vkDestroyDescriptorPool(device.vk_handle(), vk_descriptor_pool, nullptr);
-  });
+  descriptor_pools_.remove_old(
+      current_timeline, [&](std::pair<VkDescriptorPool, VKDescriptorPools *> descriptor_pool) {
+        descriptor_pool.second->recycle(descriptor_pool.first);
+      });
 }
 
 VKDiscardPool &VKDiscardPool::discard_pool_get()

source/blender/gpu/vulkan/vk_resource_pool.hh

@@ -88,7 +88,7 @@ class VKDiscardPool {
   TimelineResources<VkPipelineLayout> pipeline_layouts_;
   TimelineResources<VkRenderPass> render_passes_;
   TimelineResources<VkFramebuffer> framebuffers_;
-  TimelineResources<VkDescriptorPool> descriptor_pools_;
+  TimelineResources<std::pair<VkDescriptorPool, VKDescriptorPools *>> descriptor_pools_;
 
   Mutex mutex_;
 
@@ -106,7 +106,8 @@ class VKDiscardPool {
   void discard_pipeline_layout(VkPipelineLayout vk_pipeline_layout);
   void discard_framebuffer(VkFramebuffer vk_framebuffer);
   void discard_render_pass(VkRenderPass vk_render_pass);
-  void discard_descriptor_pool(VkDescriptorPool vk_descriptor_pool);
+  void discard_descriptor_pool_for_reuse(VkDescriptorPool vk_descriptor_pool,
+                                         VKDescriptorPools *descriptor_pools);
 
   /**
    * Move discarded resources from src_pool into this.