[Impeller] Reland: one descriptor pool per frame. (#55960)

reverted #55939 due to postsubmit failures.

It turns out the lifecycle is not exactly the same as the command pool. Instead, use a thread id map on the context to keep the right thread running and simplify shutdown. Leaving CI yaml edit for now to prove it works :)

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/context_vk.cc

@@ -3,9 +3,12 @@
 // found in the LICENSE file.
 
 #include "impeller/renderer/backend/vulkan/context_vk.h"
+#include <thread>
+#include <unordered_map>
 
 #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"
 
@@ -30,6 +33,7 @@
 #include "impeller/renderer/backend/vulkan/command_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/command_queue_vk.h"
 #include "impeller/renderer/backend/vulkan/debug_report_vk.h"
+#include "impeller/renderer/backend/vulkan/descriptor_pool_vk.h"
 #include "impeller/renderer/backend/vulkan/fence_waiter_vk.h"
 #include "impeller/renderer/backend/vulkan/gpu_tracer_vk.h"
 #include "impeller/renderer/backend/vulkan/resource_manager_vk.h"
@@ -497,8 +501,22 @@ std::shared_ptr<CommandBuffer> ContextVK::CreateCommandBuffer() const {
     return nullptr;
   }
 
+  // look up a cached descriptor pool for the current frame and reuse it
+  // if it exists, otherwise create a new pool.
+  DescriptorPoolMap::iterator current_pool =
+      cached_descriptor_pool_.find(std::this_thread::get_id());
+  std::shared_ptr<DescriptorPoolVK> descriptor_pool;
+  if (current_pool == cached_descriptor_pool_.end()) {
+    descriptor_pool =
+        (cached_descriptor_pool_[std::this_thread::get_id()] =
+             std::make_shared<DescriptorPoolVK>(weak_from_this()));
+  } else {
+    descriptor_pool = current_pool->second;
+  }
+
   auto tracked_objects = std::make_shared<TrackedObjectsVK>(
-      weak_from_this(), std::move(tls_pool), GetGPUTracer()->CreateGPUProbe());
+      weak_from_this(), std::move(tls_pool), std::move(descriptor_pool),
+      GetGPUTracer()->CreateGPUProbe());
   auto queue = GetGraphicsQueue();
 
   if (!tracked_objects || !tracked_objects->IsValid() || !queue) {
@@ -517,10 +535,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)  //
       ));
 }
 
@@ -651,6 +668,7 @@ void ContextVK::InitializeCommonlyUsedShadersIfNeeded() const {
 }
 
 void ContextVK::DisposeThreadLocalCachedResources() {
+  cached_descriptor_pool_.erase(std::this_thread::get_id());
   command_pool_recycler_->Dispose();
 }
 

impeller/renderer/backend/vulkan/context_vk.h

@@ -10,7 +10,6 @@
 #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"
@@ -40,6 +39,7 @@ class SurfaceContextVK;
 class GPUTracerVK;
 class DescriptorPoolRecyclerVK;
 class CommandQueueVK;
+class DescriptorPoolVK;
 
 class ContextVK final : public Context,
                         public BackendCast<ContextVK, Context>,
@@ -224,12 +224,17 @@ class ContextVK final : public Context,
   std::shared_ptr<const Capabilities> device_capabilities_;
   std::shared_ptr<FenceWaiterVK> fence_waiter_;
   std::shared_ptr<ResourceManagerVK> resource_manager_;
+  std::shared_ptr<DescriptorPoolRecyclerVK> descriptor_pool_recycler_;
   std::shared_ptr<CommandPoolRecyclerVK> command_pool_recycler_;
   std::string device_name_;
   std::shared_ptr<fml::ConcurrentMessageLoop> raster_message_loop_;
   std::shared_ptr<GPUTracerVK> gpu_tracer_;
-  std::shared_ptr<DescriptorPoolRecyclerVK> descriptor_pool_recycler_;
   std::shared_ptr<CommandQueue> command_queue_vk_;
+
+  using DescriptorPoolMap =
+      std::unordered_map<std::thread::id, std::shared_ptr<DescriptorPoolVK>>;
+
+  mutable DescriptorPoolMap cached_descriptor_pool_;
   bool should_disable_surface_control_ = false;
   bool should_batch_cmd_buffers_ = false;
   std::vector<std::shared_ptr<CommandBuffer>> pending_command_buffers_;

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/surface_context_vk.cc

@@ -86,7 +86,7 @@ std::unique_ptr<Surface> SurfaceContextVK::AcquireNextSurface() {
     impeller::PipelineLibraryVK::Cast(*pipeline_library)
         .DidAcquireSurfaceFrame();
   }
-  parent_->GetCommandPoolRecycler()->Dispose();
+  parent_->DisposeThreadLocalCachedResources();
   parent_->GetResourceAllocator()->DebugTraceMemoryStatistics();
   return surface;
 }

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

@@ -20,6 +20,7 @@ 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 +44,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_;