VideoBackends:Vulkan: Allocate descriptor pools as needed

4 files changed, 115 insertions, 95 deletions

diff --git a/Source/Core/VideoBackends/Vulkan/CommandBufferManager.cpp b/Source/Core/VideoBackends/Vulkan/CommandBufferManager.cpp
index 8f6bfc2929..246e0ad213 100644
--- a/Source/Core/VideoBackends/Vulkan/CommandBufferManager.cpp
+++ b/Source/Core/VideoBackends/Vulkan/CommandBufferManager.cpp
@@ -95,34 +95,6 @@ bool CommandBufferManager::CreateCommandBuffers()
     }
   }
 
-  for (VkDescriptorPool& descriptor_pool : m_descriptor_pools)
-  {
-    // TODO: A better way to choose the number of descriptors.
-    const std::array<VkDescriptorPoolSize, 5> pool_sizes{{
-        {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 500000},
-        {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 500000},
-        {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 16},
-        {VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 16384},
-        {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 16384},
-    }};
-
-    const VkDescriptorPoolCreateInfo pool_create_info = {
-        VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
-        nullptr,
-        0,
-        100000,  // tweak this
-        static_cast<u32>(pool_sizes.size()),
-        pool_sizes.data(),
-    };
-
-    res = vkCreateDescriptorPool(device, &pool_create_info, nullptr, &descriptor_pool);
-    if (res != VK_SUCCESS)
-    {
-      LOG_VULKAN_ERROR(res, "vkCreateDescriptorPool failed: ");
-      return false;
-    }
-  }
-
   res = vkCreateSemaphore(device, &semaphore_create_info, nullptr, &m_present_semaphore);
   if (res != VK_SUCCESS)
   {
@@ -160,29 +132,85 @@ void CommandBufferManager::DestroyCommandBuffers()
       vkDestroyFence(device, resources.fence, nullptr);
   }
 
-  for (VkDescriptorPool descriptor_pool : m_descriptor_pools)
+  for (FrameResources& resources : m_frame_resources)
   {
-    if (descriptor_pool != VK_NULL_HANDLE)
+    for (VkDescriptorPool descriptor_pool : resources.descriptor_pools)
+    {
       vkDestroyDescriptorPool(device, descriptor_pool, nullptr);
+    }
   }
 
   vkDestroySemaphore(device, m_present_semaphore, nullptr);
 }
 
-VkDescriptorSet CommandBufferManager::AllocateDescriptorSet(VkDescriptorSetLayout set_layout)
+VkDescriptorPool CommandBufferManager::CreateDescriptorPool(u32 max_descriptor_sets)
 {
-  VkDescriptorSetAllocateInfo allocate_info = {VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
-                                               nullptr, GetCurrentDescriptorPool(), 1, &set_layout};
+  /*
+   * Worst case descriptor counts according to the descriptor layout created in ObjectCache.cpp:
+   * UNIFORM_BUFFER_DYNAMIC: 3
+   * COMBINED_IMAGE_SAMPLER: 18
+   * STORAGE_BUFFER: 2
+   * UNIFORM_TEXEL_BUFFER: 3
+   * STORAGE_IMAGE: 1
+   */
+  const std::array<VkDescriptorPoolSize, 5> pool_sizes{{
+      {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, max_descriptor_sets * 3},
+      {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, max_descriptor_sets * 18},
+      {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, max_descriptor_sets * 2},
+      {VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, max_descriptor_sets * 3},
+      {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, max_descriptor_sets * 1},
+  }};
+
+  const VkDescriptorPoolCreateInfo pool_create_info = {
+      VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO, nullptr,           0, max_descriptor_sets,
+      static_cast<u32>(pool_sizes.size()),           pool_sizes.data(),
+  };
 
-  VkDescriptorSet descriptor_set;
-  VkResult res =
-      vkAllocateDescriptorSets(g_vulkan_context->GetDevice(), &allocate_info, &descriptor_set);
+  VkDevice device = g_vulkan_context->GetDevice();
+  VkDescriptorPool descriptor_pool = VK_NULL_HANDLE;
+  VkResult res = vkCreateDescriptorPool(device, &pool_create_info, nullptr, &descriptor_pool);
   if (res != VK_SUCCESS)
   {
-    // Failing to allocate a descriptor set is not a fatal error, we can
-    // recover by moving to the next command buffer.
+    LOG_VULKAN_ERROR(res, "vkCreateDescriptorPool failed: ");
     return VK_NULL_HANDLE;
   }
+  return descriptor_pool;
+}
+
+VkDescriptorSet CommandBufferManager::AllocateDescriptorSet(VkDescriptorSetLayout set_layout)
+{
+  VkDescriptorSet descriptor_set = VK_NULL_HANDLE;
+  FrameResources& resources = GetCurrentFrameResources();
+
+  if (!resources.descriptor_pools.empty()) [[likely]]
+  {
+    VkDescriptorSetAllocateInfo allocate_info = {
+        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO, nullptr,
+        resources.descriptor_pools[resources.current_descriptor_pool_index], 1, &set_layout};
+
+    VkResult res =
+        vkAllocateDescriptorSets(g_vulkan_context->GetDevice(), &allocate_info, &descriptor_set);
+    if (res != VK_SUCCESS &&
+        resources.descriptor_pools.size() > resources.current_descriptor_pool_index + 1)
+    {
+      // Mark the next descriptor set as active and try again.
+      resources.current_descriptor_pool_index++;
+      descriptor_set = AllocateDescriptorSet(set_layout);
+    }
+  }
+
+  if (descriptor_set == VK_NULL_HANDLE) [[unlikely]]
+  {
+    VkDescriptorPool descriptor_pool = CreateDescriptorPool(DESCRIPTOR_SETS_PER_POOL);
+    m_descriptor_set_count += DESCRIPTOR_SETS_PER_POOL;
+    if (descriptor_pool == VK_NULL_HANDLE) [[unlikely]]
+      return VK_NULL_HANDLE;
+
+    resources.descriptor_pools.push_back(descriptor_pool);
+    resources.current_descriptor_pool_index =
+        static_cast<u32>(resources.descriptor_pools.size()) - 1;
+    descriptor_set = AllocateDescriptorSet(set_layout);
+  }
 
   return descriptor_set;
 }
@@ -348,11 +376,29 @@ void CommandBufferManager::SubmitCommandBuffer(bool submit_on_worker_thread,
       cmd_buffer_index = (cmd_buffer_index + 1) % NUM_COMMAND_BUFFERS;
     }
 
-    // Reset the descriptor pool
-    VkResult res =
-        vkResetDescriptorPool(g_vulkan_context->GetDevice(), GetCurrentDescriptorPool(), 0);
-    if (res != VK_SUCCESS)
-      LOG_VULKAN_ERROR(res, "vkResetDescriptorPool failed: ");
+    // Reset the descriptor pools
+    FrameResources& frame_resources = GetCurrentFrameResources();
+
+    if (frame_resources.descriptor_pools.size() == 1) [[likely]]
+    {
+      VkResult res = vkResetDescriptorPool(g_vulkan_context->GetDevice(),
+                                           frame_resources.descriptor_pools[0], 0);
+      if (res != VK_SUCCESS)
+        LOG_VULKAN_ERROR(res, "vkResetDescriptorPool failed: ");
+    }
+    else [[unlikely]]
+    {
+      for (VkDescriptorPool descriptor_pool : frame_resources.descriptor_pools)
+      {
+        vkDestroyDescriptorPool(g_vulkan_context->GetDevice(), descriptor_pool, nullptr);
+      }
+      frame_resources.descriptor_pools.clear();
+      VkDescriptorPool descriptor_pool = CreateDescriptorPool(m_descriptor_set_count);
+      if (descriptor_pool != VK_NULL_HANDLE) [[likely]]
+        frame_resources.descriptor_pools.push_back(descriptor_pool);
+    }
+
+    frame_resources.current_descriptor_pool_index = 0;
   }
 
   // Switch to next cmdbuffer.
diff --git a/Source/Core/VideoBackends/Vulkan/CommandBufferManager.h b/Source/Core/VideoBackends/Vulkan/CommandBufferManager.h
index 45297a63f0..85e0ec8741 100644
--- a/Source/Core/VideoBackends/Vulkan/CommandBufferManager.h
+++ b/Source/Core/VideoBackends/Vulkan/CommandBufferManager.h
@@ -43,7 +43,6 @@ public:
     const CmdBufferResources& cmd_buffer_resources = m_command_buffers[m_current_cmd_buffer];
     return cmd_buffer_resources.command_buffers[1];
   }
-  VkDescriptorPool GetCurrentDescriptorPool() const { return m_descriptor_pools[m_current_frame]; }
   // Allocates a descriptors set from the pool reserved for the current frame.
   VkDescriptorSet AllocateDescriptorSet(VkDescriptorSetLayout set_layout);
 
@@ -105,6 +104,10 @@ private:
                            u32 present_image_index);
   void BeginCommandBuffer();
 
+  VkDescriptorPool CreateDescriptorPool(u32 descriptor_sizes);
+
+  const u32 DESCRIPTOR_SETS_PER_POOL = 1024;
+
   struct CmdBufferResources
   {
     // [0] - Init (upload) command buffer, [1] - draw command buffer
@@ -120,6 +123,14 @@ private:
     std::vector<std::function<void()>> cleanup_resources;
   };
 
+  struct FrameResources
+  {
+    std::vector<VkDescriptorPool> descriptor_pools;
+    u32 current_descriptor_pool_index = 0;
+  };
+
+  FrameResources& GetCurrentFrameResources() { return m_frame_resources[m_current_frame]; }
+
   CmdBufferResources& GetCurrentCmdBufferResources()
   {
     return m_command_buffers[m_current_cmd_buffer];
@@ -128,7 +139,7 @@ private:
   u64 m_next_fence_counter = 1;
   u64 m_completed_fence_counter = 0;
 
-  std::array<VkDescriptorPool, NUM_FRAMES_IN_FLIGHT> m_descriptor_pools;
+  std::array<FrameResources, NUM_FRAMES_IN_FLIGHT> m_frame_resources;
   std::array<CmdBufferResources, NUM_COMMAND_BUFFERS> m_command_buffers;
   u32 m_current_frame = 0;
   u32 m_current_cmd_buffer = 0;
@@ -150,6 +161,7 @@ private:
   Common::Flag m_last_present_failed;
   VkResult m_last_present_result = VK_SUCCESS;
   bool m_use_threaded_submission = false;
+  u32 m_descriptor_set_count = 0;
 };
 
 extern std::unique_ptr<CommandBufferManager> g_command_buffer_mgr;
diff --git a/Source/Core/VideoBackends/Vulkan/StateTracker.cpp b/Source/Core/VideoBackends/Vulkan/StateTracker.cpp
index 686bbce973..d9b03940ce 100644
--- a/Source/Core/VideoBackends/Vulkan/StateTracker.cpp
+++ b/Source/Core/VideoBackends/Vulkan/StateTracker.cpp
@@ -356,18 +356,7 @@ bool StateTracker::Bind()
     EndRenderPass();
 
   // Get a new descriptor set if any parts have changed
-  if (!UpdateDescriptorSet())
-  {
-    // We can fail to allocate descriptors if we exhaust the pool for this command buffer.
-    WARN_LOG_FMT(VIDEO, "Failed to get a descriptor set, executing buffer");
-    Renderer::GetInstance()->ExecuteCommandBuffer(false, false);
-    if (!UpdateDescriptorSet())
-    {
-      // Something strange going on.
-      ERROR_LOG_FMT(VIDEO, "Failed to get descriptor set, skipping draw");
-      return false;
-    }
-  }
+  UpdateDescriptorSet();
 
   // Start render pass if not already started
   if (!InRenderPass())
@@ -416,18 +405,7 @@ bool StateTracker::BindCompute()
                       m_compute_shader->GetComputePipeline());
   }
 
-  if (!UpdateComputeDescriptorSet())
-  {
-    WARN_LOG_FMT(VIDEO, "Failed to get a compute descriptor set, executing buffer");
-    Renderer::GetInstance()->ExecuteCommandBuffer(false, false);
-    if (!UpdateComputeDescriptorSet())
-    {
-      // Something strange going on.
-      ERROR_LOG_FMT(VIDEO, "Failed to get descriptor set, skipping dispatch");
-      return false;
-    }
-  }
-
+  UpdateComputeDescriptorSet();
   m_dirty_flags &= ~DIRTY_FLAG_COMPUTE_SHADER;
   return true;
 }
@@ -464,15 +442,15 @@ void StateTracker::EndClearRenderPass()
   EndRenderPass();
 }
 
-bool StateTracker::UpdateDescriptorSet()
+void StateTracker::UpdateDescriptorSet()
 {
   if (m_pipeline->GetUsage() != AbstractPipelineUsage::Utility)
-    return UpdateGXDescriptorSet();
+    UpdateGXDescriptorSet();
   else
-    return UpdateUtilityDescriptorSet();
+    UpdateUtilityDescriptorSet();
 }
 
-bool StateTracker::UpdateGXDescriptorSet()
+void StateTracker::UpdateGXDescriptorSet()
 {
   const size_t MAX_DESCRIPTOR_WRITES = NUM_UBO_DESCRIPTOR_SET_BINDINGS +  // UBO
                                        1 +                                // Samplers
@@ -484,8 +462,6 @@ bool StateTracker::UpdateGXDescriptorSet()
   {
     m_gx_descriptor_sets[0] = g_command_buffer_mgr->AllocateDescriptorSet(
         g_object_cache->GetDescriptorSetLayout(DESCRIPTOR_SET_LAYOUT_STANDARD_UNIFORM_BUFFERS));
-    if (m_gx_descriptor_sets[0] == VK_NULL_HANDLE)
-      return false;
 
     for (size_t i = 0; i < NUM_UBO_DESCRIPTOR_SET_BINDINGS; i++)
     {
@@ -514,8 +490,6 @@ bool StateTracker::UpdateGXDescriptorSet()
   {
     m_gx_descriptor_sets[1] = g_command_buffer_mgr->AllocateDescriptorSet(
         g_object_cache->GetDescriptorSetLayout(DESCRIPTOR_SET_LAYOUT_STANDARD_SAMPLERS));
-    if (m_gx_descriptor_sets[1] == VK_NULL_HANDLE)
-      return false;
 
     writes[num_writes++] = {VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                             nullptr,
@@ -541,8 +515,6 @@ bool StateTracker::UpdateGXDescriptorSet()
     m_gx_descriptor_sets[2] =
         g_command_buffer_mgr->AllocateDescriptorSet(g_object_cache->GetDescriptorSetLayout(
             DESCRIPTOR_SET_LAYOUT_STANDARD_SHADER_STORAGE_BUFFERS));
-    if (m_gx_descriptor_sets[2] == VK_NULL_HANDLE)
-      return false;
 
     writes[num_writes++] = {
         VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, m_gx_descriptor_sets[2], 0,      0, 1,
@@ -591,11 +563,9 @@ bool StateTracker::UpdateGXDescriptorSet()
                             m_bindings.gx_ubo_offsets.data());
     m_dirty_flags &= ~DIRTY_FLAG_GX_UBO_OFFSETS;
   }
-
-  return true;
 }
 
-bool StateTracker::UpdateUtilityDescriptorSet()
+void StateTracker::UpdateUtilityDescriptorSet()
 {
   // Max number of updates - UBO, Samplers, TexelBuffer
   std::array<VkWriteDescriptorSet, 3> dswrites;
@@ -606,8 +576,6 @@ bool StateTracker::UpdateUtilityDescriptorSet()
   {
     m_utility_descriptor_sets[0] = g_command_buffer_mgr->AllocateDescriptorSet(
         g_object_cache->GetDescriptorSetLayout(DESCRIPTOR_SET_LAYOUT_UTILITY_UNIFORM_BUFFER));
-    if (!m_utility_descriptor_sets[0])
-      return false;
 
     dswrites[writes++] = {VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                           nullptr,
@@ -627,8 +595,6 @@ bool StateTracker::UpdateUtilityDescriptorSet()
   {
     m_utility_descriptor_sets[1] = g_command_buffer_mgr->AllocateDescriptorSet(
         g_object_cache->GetDescriptorSetLayout(DESCRIPTOR_SET_LAYOUT_UTILITY_SAMPLERS));
-    if (!m_utility_descriptor_sets[1])
-      return false;
 
     dswrites[writes++] = {VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                           nullptr,
@@ -672,11 +638,9 @@ bool StateTracker::UpdateUtilityDescriptorSet()
                             1, m_utility_descriptor_sets.data(), 1, &m_bindings.utility_ubo_offset);
     m_dirty_flags &= ~(DIRTY_FLAG_DESCRIPTOR_SETS | DIRTY_FLAG_UTILITY_UBO_OFFSET);
   }
-
-  return true;
 }
 
-bool StateTracker::UpdateComputeDescriptorSet()
+void StateTracker::UpdateComputeDescriptorSet()
 {
   // Max number of updates - UBO, Samplers, TexelBuffer, Image
   std::array<VkWriteDescriptorSet, 4> dswrites;
@@ -741,8 +705,6 @@ bool StateTracker::UpdateComputeDescriptorSet()
                             &m_compute_descriptor_set, 1, &m_bindings.utility_ubo_offset);
     m_dirty_flags &= ~DIRTY_FLAG_COMPUTE_DESCRIPTOR_SET;
   }
-
-  return true;
 }
 
 }  // namespace Vulkan
diff --git a/Source/Core/VideoBackends/Vulkan/StateTracker.h b/Source/Core/VideoBackends/Vulkan/StateTracker.h
index bb7311d932..eb5e4676de 100644
--- a/Source/Core/VideoBackends/Vulkan/StateTracker.h
+++ b/Source/Core/VideoBackends/Vulkan/StateTracker.h
@@ -116,10 +116,10 @@ private:
   // If not, ends the render pass if it is a clear render pass.
   bool IsViewportWithinRenderArea() const;
 
-  bool UpdateDescriptorSet();
-  bool UpdateGXDescriptorSet();
-  bool UpdateUtilityDescriptorSet();
-  bool UpdateComputeDescriptorSet();
+  void UpdateDescriptorSet();
+  void UpdateGXDescriptorSet();
+  void UpdateUtilityDescriptorSet();
+  void UpdateComputeDescriptorSet();
 
   // Which bindings/state has to be updated before the next draw.
   u32 m_dirty_flags = 0;