[Impeller] one descriptor pool per frame. (#55939)

Creating descriptor pools is expensive, no need to have more than one per frame.

The lifecycle of the descriptor pool is exactly the same as the cmd pool, which indicates that we can probably do some consolidation/refactoring in the future.

Fixes flutter/flutter#157115

Author: Jonah Williams

impeller/renderer/backend/vulkan/command_buffer_vk.cc

@@ -22,12 +22,10 @@ namespace impeller {
 CommandBufferVK::CommandBufferVK(
     std::weak_ptr<const Context> context,
     std::weak_ptr<const DeviceHolderVK> device_holder,
-    std::shared_ptr<TrackedObjectsVK> tracked_objects,
-    std::shared_ptr<FenceWaiterVK> fence_waiter)
+    std::shared_ptr<TrackedObjectsVK> tracked_objects)
     : CommandBuffer(std::move(context)),
       device_holder_(std::move(device_holder)),
-      tracked_objects_(std::move(tracked_objects)),
-      fence_waiter_(std::move(fence_waiter)) {}
+      tracked_objects_(std::move(tracked_objects)) {}
 
 CommandBufferVK::~CommandBufferVK() = default;
 
@@ -212,4 +210,8 @@ void CommandBufferVK::InsertDebugMarker(std::string_view label) const {
   }
 }
 
+DescriptorPoolVK& CommandBufferVK::GetDescriptorPool() const {
+  return tracked_objects_->GetDescriptorPool();
+}
+
 }  // namespace impeller

impeller/renderer/backend/vulkan/command_buffer_vk.h

@@ -8,6 +8,7 @@
 #include "fml/status_or.h"
 #include "impeller/base/backend_cast.h"
 #include "impeller/renderer/backend/vulkan/command_queue_vk.h"
+#include "impeller/renderer/backend/vulkan/descriptor_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/device_holder_vk.h"
 #include "impeller/renderer/backend/vulkan/texture_source_vk.h"
 #include "impeller/renderer/backend/vulkan/tracked_objects_vk.h"
@@ -81,18 +82,19 @@ class CommandBufferVK final
   // Visible for testing.
   bool IsTracking(const std::shared_ptr<const Texture>& texture) const;
 
+  // Visible for testing.
+  DescriptorPoolVK& GetDescriptorPool() const;
+
  private:
   friend class ContextVK;
   friend class CommandQueueVK;
 
   std::weak_ptr<const DeviceHolderVK> device_holder_;
   std::shared_ptr<TrackedObjectsVK> tracked_objects_;
-  std::shared_ptr<FenceWaiterVK> fence_waiter_;
 
   CommandBufferVK(std::weak_ptr<const Context> context,
                   std::weak_ptr<const DeviceHolderVK> device_holder,
-                  std::shared_ptr<TrackedObjectsVK> tracked_objects,
-                  std::shared_ptr<FenceWaiterVK> fence_waiter);
+                  std::shared_ptr<TrackedObjectsVK> tracked_objects);
 
   // |CommandBuffer|
   void SetLabel(std::string_view label) const override;

impeller/renderer/backend/vulkan/command_pool_vk.cc

@@ -9,6 +9,7 @@
 #include <utility>
 
 #include "impeller/renderer/backend/vulkan/context_vk.h"
+#include "impeller/renderer/backend/vulkan/descriptor_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/resource_manager_vk.h"
 
 #include "impeller/renderer/backend/vulkan/vk.h"  // IWYU pragma: keep.
@@ -155,10 +156,6 @@ void CommandPoolVK::Destroy() {
   collected_buffers_.clear();
 }
 
-// Associates a resource with a thread and context.
-using CommandPoolMap =
-    std::unordered_map<uint64_t, std::shared_ptr<CommandPoolVK>>;
-
 // CommandPoolVK Lifecycle:
 // 1. End of frame will reset the command pool (clearing this on a thread).
 //    There will still be references to the command pool from the uncompleted
@@ -169,17 +166,47 @@ using CommandPoolMap =
 //    available for reuse ("recycle").
 static thread_local std::unique_ptr<CommandPoolMap> tls_command_pool_map;
 
+struct WeakThreadLocalData {
+  std::weak_ptr<CommandPoolVK> command_pool;
+  std::weak_ptr<DescriptorPoolVK> descriptor_pool;
+};
+
 // Map each context to a list of all thread-local command pools associated
 // with that context.
 static Mutex g_all_pools_map_mutex;
-static std::unordered_map<
-    const ContextVK*,
-    std::vector<std::weak_ptr<CommandPoolVK>>> g_all_pools_map
+static std::unordered_map<const ContextVK*,
+                          std::vector<WeakThreadLocalData>> g_all_pools_map
     IPLR_GUARDED_BY(g_all_pools_map_mutex);
 
+std::shared_ptr<DescriptorPoolVK> CommandPoolRecyclerVK::GetDescriptorPool() {
+  const auto& strong_context = context_.lock();
+  if (!strong_context) {
+    return nullptr;
+  }
+
+  // If there is a resource in used for this thread and context, return it.
+  if (!tls_command_pool_map.get()) {
+    tls_command_pool_map.reset(new CommandPoolMap());
+  }
+  CommandPoolMap& pool_map = *tls_command_pool_map.get();
+  auto const hash = strong_context->GetHash();
+  auto const it = pool_map.find(hash);
+  if (it != pool_map.end()) {
+    return it->second.descriptor_pool;
+  }
+
+  const auto& result =
+      InitializeThreadLocalResources(strong_context, pool_map, hash);
+  if (result.has_value()) {
+    return result->descriptor_pool;
+  }
+
+  return nullptr;
+}
+
 // TODO(matanlurey): Return a status_or<> instead of nullptr when we have one.
 std::shared_ptr<CommandPoolVK> CommandPoolRecyclerVK::Get() {
-  auto const strong_context = context_.lock();
+  const auto& strong_context = context_.lock();
   if (!strong_context) {
     return nullptr;
   }
@@ -192,25 +219,40 @@ std::shared_ptr<CommandPoolVK> CommandPoolRecyclerVK::Get() {
   auto const hash = strong_context->GetHash();
   auto const it = pool_map.find(hash);
   if (it != pool_map.end()) {
-    return it->second;
+    return it->second.command_pool;
   }
 
-  // Otherwise, create a new resource and return it.
+  const auto& result =
+      InitializeThreadLocalResources(strong_context, pool_map, hash);
+  if (result.has_value()) {
+    return result->command_pool;
+  }
+
+  return nullptr;
+}
+
+std::optional<ThreadLocalData>
+CommandPoolRecyclerVK::InitializeThreadLocalResources(
+    const std::shared_ptr<ContextVK>& context,
+    CommandPoolMap& pool_map,
+    uint64_t pool_key) {
   auto data = Create();
   if (!data || !data->pool) {
-    return nullptr;
+    return std::nullopt;
   }
 
-  auto const resource = std::make_shared<CommandPoolVK>(
+  auto command_pool = std::make_shared<CommandPoolVK>(
       std::move(data->pool), std::move(data->buffers), context_);
-  pool_map.emplace(hash, resource);
-
+  auto descriptor_pool = std::make_shared<DescriptorPoolVK>(context_);
   {
     Lock all_pools_lock(g_all_pools_map_mutex);
-    g_all_pools_map[strong_context.get()].push_back(resource);
+    g_all_pools_map[context.get()].push_back(WeakThreadLocalData{
+        .command_pool = command_pool, .descriptor_pool = descriptor_pool});
   }
 
-  return resource;
+  return pool_map[pool_key] =
+             ThreadLocalData{.command_pool = std::move(command_pool),
+                             .descriptor_pool = std::move(descriptor_pool)};
 }
 
 // TODO(matanlurey): Return a status_or<> instead of nullopt when we have one.
@@ -293,14 +335,13 @@ void CommandPoolRecyclerVK::DestroyThreadLocalPools(const ContextVK* context) {
   Lock all_pools_lock(g_all_pools_map_mutex);
   auto found = g_all_pools_map.find(context);
   if (found != g_all_pools_map.end()) {
-    for (auto& weak_pool : found->second) {
-      auto pool = weak_pool.lock();
-      if (!pool) {
-        continue;
+    for (auto& [command_pool, desc_pool] : found->second) {
+      const auto& strong_pool = command_pool.lock();
+      if (strong_pool) {
+        // Delete all objects held by this pool.  The destroyed pool will still
+        // remain in its thread's TLS map until that thread exits.
+        strong_pool->Destroy();
       }
-      // Delete all objects held by this pool.  The destroyed pool will still
-      // remain in its thread's TLS map until that thread exits.
-      pool->Destroy();
     }
     g_all_pools_map.erase(found);
   }

impeller/renderer/backend/vulkan/command_pool_vk.h

@@ -10,6 +10,7 @@
 #include <utility>
 
 #include "impeller/base/thread.h"
+#include "impeller/renderer/backend/vulkan/descriptor_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/vk.h"  // IWYU pragma: keep.
 #include "vulkan/vulkan_handles.hpp"
 
@@ -76,6 +77,14 @@ class CommandPoolVK final {
       pool_mutex_);
 };
 
+struct ThreadLocalData {
+  std::shared_ptr<CommandPoolVK> command_pool;
+  std::shared_ptr<DescriptorPoolVK> descriptor_pool;
+};
+
+// Associates a resource with a thread and context.
+using CommandPoolMap = std::unordered_map<uint64_t, ThreadLocalData>;
+
 //------------------------------------------------------------------------------
 /// @brief      Creates and manages the lifecycle of |vk::CommandPool| objects.
 ///
@@ -128,6 +137,11 @@ class CommandPoolRecyclerVK final
   /// @warning    Returns a |nullptr| if a pool could not be created.
   std::shared_ptr<CommandPoolVK> Get();
 
+  /// @brief      Gets a descriptor pool for the current thread.
+  ///
+  /// @warning    Returns a |nullptr| if a pool could not be created.
+  std::shared_ptr<DescriptorPoolVK> GetDescriptorPool();
+
   /// @brief      Returns a command pool to be reset on a background thread.
   ///
   /// @param[in]  pool The pool to recycler.
@@ -153,6 +167,13 @@ class CommandPoolRecyclerVK final
   /// @returns    Returns a |std::nullopt| if a pool was not available.
   std::optional<RecycledData> Reuse();
 
+  /// Create a DescriptorPoolVK and a CommandPoolVK and stash them in the TLS
+  /// map.
+  std::optional<ThreadLocalData> InitializeThreadLocalResources(
+      const std::shared_ptr<ContextVK>& context,
+      CommandPoolMap& pool_map,
+      uint64_t pool_key);
+
   CommandPoolRecyclerVK(const CommandPoolRecyclerVK&) = delete;
 
   CommandPoolRecyclerVK& operator=(const CommandPoolRecyclerVK&) = delete;

impeller/renderer/backend/vulkan/context_vk.cc

@@ -6,6 +6,7 @@
 
 #include "fml/concurrent_message_loop.h"
 #include "impeller/renderer/backend/vulkan/command_queue_vk.h"
+#include "impeller/renderer/backend/vulkan/descriptor_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/render_pass_builder_vk.h"
 #include "impeller/renderer/render_target.h"
 
@@ -495,9 +496,14 @@ std::shared_ptr<CommandBuffer> ContextVK::CreateCommandBuffer() const {
   if (!tls_pool) {
     return nullptr;
   }
+  auto tls_desc_pool = recycler->GetDescriptorPool();
+  if (!tls_desc_pool) {
+    return nullptr;
+  }
 
   auto tracked_objects = std::make_shared<TrackedObjectsVK>(
-      weak_from_this(), std::move(tls_pool), GetGPUTracer()->CreateGPUProbe());
+      weak_from_this(), std::move(tls_pool), tls_desc_pool,
+      GetGPUTracer()->CreateGPUProbe());
   auto queue = GetGraphicsQueue();
 
   if (!tracked_objects || !tracked_objects->IsValid() || !queue) {
@@ -516,10 +522,9 @@ std::shared_ptr<CommandBuffer> ContextVK::CreateCommandBuffer() const {
       tracked_objects->GetCommandBuffer());
 
   return std::shared_ptr<CommandBufferVK>(new CommandBufferVK(
-      shared_from_this(),          //
-      GetDeviceHolder(),           //
-      std::move(tracked_objects),  //
-      GetFenceWaiter()             //
+      shared_from_this(),         //
+      GetDeviceHolder(),          //
+      std::move(tracked_objects)  //
       ));
 }
 

impeller/renderer/backend/vulkan/context_vk.h

@@ -10,11 +10,9 @@
 #include "flutter/fml/concurrent_message_loop.h"
 #include "flutter/fml/mapping.h"
 #include "flutter/fml/unique_fd.h"
-#include "fml/thread.h"
 #include "impeller/base/backend_cast.h"
 #include "impeller/base/strings.h"
 #include "impeller/core/formats.h"
-#include "impeller/renderer/backend/vulkan/command_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/device_holder_vk.h"
 #include "impeller/renderer/backend/vulkan/driver_info_vk.h"
 #include "impeller/renderer/backend/vulkan/pipeline_library_vk.h"

impeller/renderer/backend/vulkan/descriptor_pool_vk_unittests.cc

@@ -5,6 +5,7 @@
 #include "flutter/testing/testing.h"  // IWYU pragma: keep.
 #include "fml/closure.h"
 #include "fml/synchronization/waitable_event.h"
+#include "impeller/renderer/backend/vulkan/command_buffer_vk.h"
 #include "impeller/renderer/backend/vulkan/descriptor_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/resource_manager_vk.h"
 #include "impeller/renderer/backend/vulkan/test/mock_vulkan.h"
@@ -123,5 +124,27 @@ TEST(DescriptorPoolRecyclerVKTest, ReclaimDropsDescriptorPoolIfSizeIsExceeded) {
   context->Shutdown();
 }
 
+TEST(DescriptorPoolRecyclerVKTest, MultipleCommandBuffersShareDescriptorPool) {
+  auto const context = MockVulkanContextBuilder().Build();
+
+  auto cmd_buffer_1 = context->CreateCommandBuffer();
+  auto cmd_buffer_2 = context->CreateCommandBuffer();
+
+  CommandBufferVK& vk_1 = CommandBufferVK::Cast(*cmd_buffer_1);
+  CommandBufferVK& vk_2 = CommandBufferVK::Cast(*cmd_buffer_2);
+
+  EXPECT_EQ(&vk_1.GetDescriptorPool(), &vk_2.GetDescriptorPool());
+
+  // Resetting resources creates a new pool.
+  context->DisposeThreadLocalCachedResources();
+
+  auto cmd_buffer_3 = context->CreateCommandBuffer();
+  CommandBufferVK& vk_3 = CommandBufferVK::Cast(*cmd_buffer_3);
+
+  EXPECT_NE(&vk_1.GetDescriptorPool(), &vk_3.GetDescriptorPool());
+
+  context->Shutdown();
+}
+
 }  // namespace testing
 }  // namespace impeller

impeller/renderer/backend/vulkan/test/mock_vulkan_unittests.cc

@@ -4,6 +4,7 @@
 
 #include "flutter/testing/testing.h"  // IWYU pragma: keep
 #include "gtest/gtest.h"
+#include "impeller/renderer/backend/vulkan/command_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/test/mock_vulkan.h"
 #include "vulkan/vulkan_enums.hpp"
 

impeller/renderer/backend/vulkan/tracked_objects_vk.cc

@@ -4,15 +4,17 @@
 
 #include "impeller/renderer/backend/vulkan/tracked_objects_vk.h"
 
+#include "impeller/renderer/backend/vulkan/command_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/gpu_tracer_vk.h"
 
 namespace impeller {
 
 TrackedObjectsVK::TrackedObjectsVK(
     const std::weak_ptr<const ContextVK>& context,
     const std::shared_ptr<CommandPoolVK>& pool,
+    std::shared_ptr<DescriptorPoolVK> descriptor_pool,
     std::unique_ptr<GPUProbe> probe)
-    : desc_pool_(context), probe_(std::move(probe)) {
+    : desc_pool_(std::move(descriptor_pool)), probe_(std::move(probe)) {
   if (!pool) {
     return;
   }
@@ -78,7 +80,7 @@ vk::CommandBuffer TrackedObjectsVK::GetCommandBuffer() const {
 }
 
 DescriptorPoolVK& TrackedObjectsVK::GetDescriptorPool() {
-  return desc_pool_;
+  return *desc_pool_;
 }
 
 GPUProbe& TrackedObjectsVK::GetGPUProbe() const {

impeller/renderer/backend/vulkan/tracked_objects_vk.h

@@ -14,12 +14,15 @@
 
 namespace impeller {
 
+class CommandPoolVK;
+
 /// @brief A per-frame object used to track resource lifetimes and allocate
 ///        command buffers and descriptor sets.
 class TrackedObjectsVK {
  public:
   explicit TrackedObjectsVK(const std::weak_ptr<const ContextVK>& context,
                             const std::shared_ptr<CommandPoolVK>& pool,
+                            std::shared_ptr<DescriptorPoolVK> descriptor_pool,
                             std::unique_ptr<GPUProbe> probe);
 
   ~TrackedObjectsVK();
@@ -43,7 +46,7 @@ class TrackedObjectsVK {
   GPUProbe& GetGPUProbe() const;
 
  private:
-  DescriptorPoolVK desc_pool_;
+  std::shared_ptr<DescriptorPoolVK> desc_pool_;
   // `shared_ptr` since command buffers have a link to the command pool.
   std::shared_ptr<CommandPoolVK> pool_;
   vk::UniqueCommandBuffer buffer_;