[Impeller] cache descriptor set layouts. (#164952)

Currently we reset the entire descriptor pool(s) each frame.

We can actually reuse the allocated descriptor sets, as that is safe to
do as soon as the cmd buffer that references them has been submitted
(AFAIK). This provides a minor speed up for a large amount of small
drawing ops, but not much beyond that.

engine/src/flutter/impeller/renderer/backend/vulkan/command_buffer_vk.cc

@@ -160,13 +160,14 @@ bool CommandBufferVK::Track(const std::shared_ptr<const Texture>& texture) {
 
 fml::StatusOr<vk::DescriptorSet> CommandBufferVK::AllocateDescriptorSets(
     const vk::DescriptorSetLayout& layout,
+    PipelineKey pipeline_key,
     const ContextVK& context) {
   if (!IsValid()) {
     return fml::Status(fml::StatusCode::kUnknown, "command encoder invalid");
   }
 
-  return tracked_objects_->GetDescriptorPool().AllocateDescriptorSets(layout,
-                                                                      context);
+  return tracked_objects_->GetDescriptorPool().AllocateDescriptorSets(
+      layout, pipeline_key, context);
 }
 
 void CommandBufferVK::PushDebugGroup(std::string_view label) const {

engine/src/flutter/impeller/renderer/backend/vulkan/command_buffer_vk.h

@@ -74,6 +74,7 @@ class CommandBufferVK final
   /// @brief Allocate a new descriptor set for the given [layout].
   fml::StatusOr<vk::DescriptorSet> AllocateDescriptorSets(
       const vk::DescriptorSetLayout& layout,
+      PipelineKey pipeline_key,
       const ContextVK& context);
 
   // Visible for testing.

engine/src/flutter/impeller/renderer/backend/vulkan/compute_pass_vk.cc

@@ -58,7 +58,8 @@ void ComputePassVK::SetPipeline(
   pipeline_layout_ = pipeline_vk.GetPipelineLayout();
 
   auto descriptor_result = command_buffer_->AllocateDescriptorSets(
-      pipeline_vk.GetDescriptorSetLayout(), ContextVK::Cast(*context_));
+      pipeline_vk.GetDescriptorSetLayout(), pipeline_vk.GetPipelineKey(),
+      ContextVK::Cast(*context_));
   if (!descriptor_result.ok()) {
     return;
   }

engine/src/flutter/impeller/renderer/backend/vulkan/compute_pipeline_vk.cc

@@ -12,12 +12,14 @@ ComputePipelineVK::ComputePipelineVK(
     const ComputePipelineDescriptor& desc,
     vk::UniquePipeline pipeline,
     vk::UniquePipelineLayout layout,
-    vk::UniqueDescriptorSetLayout descriptor_set_layout)
+    vk::UniqueDescriptorSetLayout descriptor_set_layout,
+    PipelineKey pipeline_key)
     : Pipeline(std::move(library), desc),
       device_holder_(std::move(device_holder)),
       pipeline_(std::move(pipeline)),
       layout_(std::move(layout)),
-      descriptor_set_layout_(std::move(descriptor_set_layout)) {
+      descriptor_set_layout_(std::move(descriptor_set_layout)),
+      pipeline_key_(pipeline_key) {
   is_valid_ = pipeline_ && layout_ && descriptor_set_layout_;
 }
 
@@ -51,4 +53,8 @@ const vk::DescriptorSetLayout& ComputePipelineVK::GetDescriptorSetLayout()
   return *descriptor_set_layout_;
 }
 
+PipelineKey ComputePipelineVK::GetPipelineKey() const {
+  return pipeline_key_;
+}
+
 }  // namespace impeller

engine/src/flutter/impeller/renderer/backend/vulkan/compute_pipeline_vk.h

@@ -24,7 +24,8 @@ class ComputePipelineVK final
                     const ComputePipelineDescriptor& desc,
                     vk::UniquePipeline pipeline,
                     vk::UniquePipelineLayout layout,
-                    vk::UniqueDescriptorSetLayout descriptor_set_layout);
+                    vk::UniqueDescriptorSetLayout descriptor_set_layout,
+                    PipelineKey pipeline_key);
 
   // |Pipeline|
   ~ComputePipelineVK() override;
@@ -35,6 +36,10 @@ class ComputePipelineVK final
 
   const vk::DescriptorSetLayout& GetDescriptorSetLayout() const;
 
+  /// @brief Retrieve the unique identifier for this pipeline's descriptor set
+  ///        layout.
+  PipelineKey GetPipelineKey() const;
+
  private:
   friend class PipelineLibraryVK;
 
@@ -42,6 +47,7 @@ class ComputePipelineVK final
   vk::UniquePipeline pipeline_;
   vk::UniquePipelineLayout layout_;
   vk::UniqueDescriptorSetLayout descriptor_set_layout_;
+  const PipelineKey pipeline_key_;
   bool is_valid_ = false;
 
   // |Pipeline|

engine/src/flutter/impeller/renderer/backend/vulkan/context_vk.cc

@@ -546,9 +546,8 @@ std::shared_ptr<CommandBuffer> ContextVK::CreateCommandBuffer() const {
     DescriptorPoolMap::iterator current_pool =
         cached_descriptor_pool_.find(std::this_thread::get_id());
     if (current_pool == cached_descriptor_pool_.end()) {
-      descriptor_pool =
-          (cached_descriptor_pool_[std::this_thread::get_id()] =
-               std::make_shared<DescriptorPoolVK>(weak_from_this()));
+      descriptor_pool = (cached_descriptor_pool_[std::this_thread::get_id()] =
+                             descriptor_pool_recycler_->GetDescriptorPool());
     } else {
       descriptor_pool = current_pool->second;
     }

engine/src/flutter/impeller/renderer/backend/vulkan/descriptor_pool_vk.cc

@@ -7,9 +7,6 @@
 #include <optional>
 
 #include "impeller/base/validation.h"
-#include "impeller/renderer/backend/vulkan/resource_manager_vk.h"
-#include "vulkan/vulkan_enums.hpp"
-#include "vulkan/vulkan_handles.hpp"
 
 namespace impeller {
 
@@ -29,44 +26,20 @@ static const constexpr DescriptorPoolSize kDefaultBindingSize =
         .subpass_bindings = 4u  // Subpass Bindings
     };
 
-// Holds the command pool in a background thread, recyling it when not in use.
-class BackgroundDescriptorPoolVK final {
- public:
-  BackgroundDescriptorPoolVK(BackgroundDescriptorPoolVK&&) = default;
-
-  explicit BackgroundDescriptorPoolVK(
-      vk::UniqueDescriptorPool&& pool,
-      std::weak_ptr<DescriptorPoolRecyclerVK> recycler)
-      : pool_(std::move(pool)), recycler_(std::move(recycler)) {}
-
-  ~BackgroundDescriptorPoolVK() {
-    auto const recycler = recycler_.lock();
-
-    // Not only does this prevent recycling when the context is being destroyed,
-    // but it also prevents the destructor from effectively being called twice;
-    // once for the original BackgroundCommandPoolVK() and once for the moved
-    // BackgroundCommandPoolVK().
-    if (!recycler) {
-      return;
-    }
-
-    recycler->Reclaim(std::move(pool_));
-  }
-
- private:
-  BackgroundDescriptorPoolVK(const BackgroundDescriptorPoolVK&) = delete;
-
-  BackgroundDescriptorPoolVK& operator=(const BackgroundDescriptorPoolVK&) =
-      delete;
-
-  vk::UniqueDescriptorPool pool_;
-  uint32_t allocated_capacity_;
-  std::weak_ptr<DescriptorPoolRecyclerVK> recycler_;
-};
-
 DescriptorPoolVK::DescriptorPoolVK(std::weak_ptr<const ContextVK> context)
     : context_(std::move(context)) {}
 
+void DescriptorPoolVK::Destroy() {
+  pools_.clear();
+}
+
+DescriptorPoolVK::DescriptorPoolVK(std::weak_ptr<const ContextVK> context,
+                                   DescriptorCacheMap descriptor_sets,
+                                   std::vector<vk::UniqueDescriptorPool> pools)
+    : context_(std::move(context)),
+      descriptor_sets_(std::move(descriptor_sets)),
+      pools_(std::move(pools)) {}
+
 DescriptorPoolVK::~DescriptorPoolVK() {
   if (pools_.empty()) {
     return;
@@ -81,19 +54,21 @@ DescriptorPoolVK::~DescriptorPoolVK() {
     return;
   }
 
-  for (auto i = 0u; i < pools_.size(); i++) {
-    auto reset_pool_when_dropped =
-        BackgroundDescriptorPoolVK(std::move(pools_[i]), recycler);
-
-    UniqueResourceVKT<BackgroundDescriptorPoolVK> pool(
-        context->GetResourceManager(), std::move(reset_pool_when_dropped));
-  }
-  pools_.clear();
+  recycler->Reclaim(std::move(descriptor_sets_), std::move(pools_));
 }
 
 fml::StatusOr<vk::DescriptorSet> DescriptorPoolVK::AllocateDescriptorSets(
     const vk::DescriptorSetLayout& layout,
+    PipelineKey pipeline_key,
     const ContextVK& context_vk) {
+  DescriptorCacheMap::iterator existing = descriptor_sets_.find(pipeline_key);
+  if (existing != descriptor_sets_.end() && !existing->second.unused.empty()) {
+    auto descriptor_set = existing->second.unused.back();
+    existing->second.unused.pop_back();
+    existing->second.used.push_back(descriptor_set);
+    return descriptor_set;
+  }
+
   if (pools_.empty()) {
     CreateNewPool(context_vk);
   }
@@ -111,6 +86,9 @@ fml::StatusOr<vk::DescriptorSet> DescriptorPoolVK::AllocateDescriptorSets(
     set_info.setDescriptorPool(pools_.back().get());
     result = context_vk.GetDevice().allocateDescriptorSets(&set_info, &set);
   }
+  auto lookup_result =
+      descriptor_sets_.try_emplace(pipeline_key, DescriptorCache{});
+  lookup_result.first->second.used.push_back(set);
 
   if (result != vk::Result::eSuccess) {
     VALIDATION_LOG << "Could not allocate descriptor sets: "
@@ -130,30 +108,35 @@ fml::Status DescriptorPoolVK::CreateNewPool(const ContextVK& context_vk) {
   return fml::Status();
 }
 
-void DescriptorPoolRecyclerVK::Reclaim(vk::UniqueDescriptorPool&& pool) {
+void DescriptorPoolRecyclerVK::Reclaim(
+    DescriptorCacheMap descriptor_sets,
+    std::vector<vk::UniqueDescriptorPool> pools) {
   // Reset the pool on a background thread.
   auto strong_context = context_.lock();
   if (!strong_context) {
     return;
   }
-  auto device = strong_context->GetDevice();
-  device.resetDescriptorPool(pool.get());
 
-  // Move the pool to the recycled list.
-  Lock recycled_lock(recycled_mutex_);
-
-  if (recycled_.size() < kMaxRecycledPools) {
-    recycled_.push_back(std::move(pool));
-    return;
+  for (auto& [_, cache] : descriptor_sets) {
+    cache.unused.insert(cache.unused.end(), cache.used.begin(),
+                        cache.used.end());
+    cache.used.clear();
   }
+
+  // Move the pool to the recycled list. If more than 32 pool are
+  // cached then delete the newest entry.
+  Lock recycled_lock(recycled_mutex_);
+  while (recycled_.size() >= kMaxRecycledPools) {
+    auto& back_entry = recycled_.back();
+    back_entry->Destroy();
+    recycled_.pop_back();
+  }
+  recycled_.push_back(std::make_shared<DescriptorPoolVK>(
+      context_, std::move(descriptor_sets), std::move(pools)));
 }
 
 vk::UniqueDescriptorPool DescriptorPoolRecyclerVK::Get() {
   // Recycle a pool with a matching minumum capcity if it is available.
-  auto recycled_pool = Reuse();
-  if (recycled_pool.has_value()) {
-    return std::move(recycled_pool.value());
-  }
   return Create();
 }
 
@@ -187,15 +170,17 @@ vk::UniqueDescriptorPool DescriptorPoolRecyclerVK::Create() {
   return std::move(pool);
 }
 
-std::optional<vk::UniqueDescriptorPool> DescriptorPoolRecyclerVK::Reuse() {
-  Lock lock(recycled_mutex_);
-  if (recycled_.empty()) {
-    return std::nullopt;
+std::shared_ptr<DescriptorPoolVK>
+DescriptorPoolRecyclerVK::GetDescriptorPool() {
+  {
+    Lock recycled_lock(recycled_mutex_);
+    if (!recycled_.empty()) {
+      auto result = recycled_.back();
+      recycled_.pop_back();
+      return result;
+    }
   }
-
-  auto recycled = std::move(recycled_[recycled_.size() - 1]);
-  recycled_.pop_back();
-  return recycled;
+  return std::make_shared<DescriptorPoolVK>(context_);
 }
 
 }  // namespace impeller

engine/src/flutter/impeller/renderer/backend/vulkan/descriptor_pool_vk.h

@@ -6,13 +6,22 @@
 #define FLUTTER_IMPELLER_RENDERER_BACKEND_VULKAN_DESCRIPTOR_POOL_VK_H_
 
 #include <cstdint>
+#include <unordered_map>
 
 #include "fml/status_or.h"
 #include "impeller/renderer/backend/vulkan/context_vk.h"
-#include "vulkan/vulkan_handles.hpp"
+#include "impeller/renderer/pipeline.h"
 
 namespace impeller {
 
+/// Used and un-used descriptor sets.
+struct DescriptorCache {
+  std::vector<vk::DescriptorSet> unused;
+  std::vector<vk::DescriptorSet> used;
+};
+
+using DescriptorCacheMap = std::unordered_map<PipelineKey, DescriptorCache>;
+
 //------------------------------------------------------------------------------
 /// @brief      A per-frame descriptor pool. Descriptors
 ///             from this pool don't need to be freed individually. Instead, the
@@ -28,16 +37,26 @@ class DescriptorPoolVK {
  public:
   explicit DescriptorPoolVK(std::weak_ptr<const ContextVK> context);
 
+  DescriptorPoolVK(std::weak_ptr<const ContextVK> context,
+                   DescriptorCacheMap descriptor_sets,
+                   std::vector<vk::UniqueDescriptorPool> pools);
+
   ~DescriptorPoolVK();
 
   fml::StatusOr<vk::DescriptorSet> AllocateDescriptorSets(
       const vk::DescriptorSetLayout& layout,
+      PipelineKey pipeline_key,
       const ContextVK& context_vk);
 
  private:
+  friend class DescriptorPoolRecyclerVK;
+
   std::weak_ptr<const ContextVK> context_;
+  DescriptorCacheMap descriptor_sets_;
   std::vector<vk::UniqueDescriptorPool> pools_;
 
+  void Destroy();
+
   fml::Status CreateNewPool(const ContextVK& context_vk);
 
   DescriptorPoolVK(const DescriptorPoolVK&) = delete;
@@ -68,17 +87,16 @@ class DescriptorPoolRecyclerVK final
   ///             the necessary capacity.
   vk::UniqueDescriptorPool Get();
 
-  /// @brief      Returns the descriptor pool to be reset on a background
-  ///             thread.
-  ///
-  /// @param[in]  pool The pool to recycler.
-  void Reclaim(vk::UniqueDescriptorPool&& pool);
+  std::shared_ptr<DescriptorPoolVK> GetDescriptorPool();
+
+  void Reclaim(DescriptorCacheMap descriptor_sets,
+               std::vector<vk::UniqueDescriptorPool> pools);
 
  private:
   std::weak_ptr<ContextVK> context_;
 
   Mutex recycled_mutex_;
-  std::vector<vk::UniqueDescriptorPool> recycled_ IPLR_GUARDED_BY(
+  std::vector<std::shared_ptr<DescriptorPoolVK>> recycled_ IPLR_GUARDED_BY(
       recycled_mutex_);
 
   /// @brief      Creates a new |vk::CommandPool|.
@@ -86,11 +104,6 @@ class DescriptorPoolRecyclerVK final
   /// @returns    Returns a |std::nullopt| if a pool could not be created.
   vk::UniqueDescriptorPool Create();
 
-  /// @brief      Reuses a recycled |vk::CommandPool|, if available.
-  ///
-  /// @returns    Returns a |std::nullopt| if a pool was not available.
-  std::optional<vk::UniqueDescriptorPool> Reuse();
-
   DescriptorPoolRecyclerVK(const DescriptorPoolRecyclerVK&) = delete;
 
   DescriptorPoolRecyclerVK& operator=(const DescriptorPoolRecyclerVK&) = delete;

engine/src/flutter/impeller/renderer/backend/vulkan/descriptor_pool_vk_unittests.cc

@@ -3,21 +3,18 @@
 // found in the LICENSE file.
 
 #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"
 
 namespace impeller {
 namespace testing {
 
 TEST(DescriptorPoolRecyclerVKTest, GetDescriptorPoolRecyclerCreatesNewPools) {
-  auto const context = MockVulkanContextBuilder().Build();
+  std::shared_ptr<ContextVK> context = MockVulkanContextBuilder().Build();
 
-  auto const pool1 = context->GetDescriptorPoolRecycler()->Get();
-  auto const pool2 = context->GetDescriptorPoolRecycler()->Get();
+  vk::UniqueDescriptorPool pool1 = context->GetDescriptorPoolRecycler()->Get();
+  vk::UniqueDescriptorPool pool2 = context->GetDescriptorPoolRecycler()->Get();
 
   // The two descriptor pools should be different.
   EXPECT_NE(pool1.get(), pool2.get());
@@ -26,37 +23,20 @@ TEST(DescriptorPoolRecyclerVKTest, GetDescriptorPoolRecyclerCreatesNewPools) {
 }
 
 TEST(DescriptorPoolRecyclerVKTest, ReclaimMakesDescriptorPoolAvailable) {
-  auto const context = MockVulkanContextBuilder().Build();
+  std::shared_ptr<ContextVK> context = MockVulkanContextBuilder().Build();
 
   {
     // Fetch a pool (which will be created).
-    auto pool = DescriptorPoolVK(context);
-    pool.AllocateDescriptorSets({}, *context);
+    DescriptorPoolVK pool = DescriptorPoolVK(context);
+    pool.AllocateDescriptorSets({}, /*pipeline_key=*/0, *context);
   }
 
-  // There is a chance that the first death rattle item below is destroyed in
-  // the same reclaim cycle as the pool allocation above. These items are placed
-  // into a std::vector and free'd, which may free in reverse order. That would
-  // imply that the death rattle and subsequent waitable event fires before the
-  // pool is reset. To work around this, we can either manually remove items
-  // from the vector or use two death rattles.
-  for (auto i = 0u; i < 2; i++) {
-    // Add something to the resource manager and have it notify us when it's
-    // destroyed. That should give us a non-flaky signal that the pool has been
-    // reclaimed as well.
-    auto waiter = fml::AutoResetWaitableEvent();
-    auto rattle = fml::ScopedCleanupClosure([&waiter]() { waiter.Signal(); });
-    {
-      UniqueResourceVKT<fml::ScopedCleanupClosure> resource(
-          context->GetResourceManager(), std::move(rattle));
-    }
-    waiter.Wait();
-  }
-
-  auto const pool = context->GetDescriptorPoolRecycler()->Get();
+  std::shared_ptr<DescriptorPoolVK> pool =
+      context->GetDescriptorPoolRecycler()->GetDescriptorPool();
 
   // Now check that we only ever created one pool.
-  auto const called = GetMockVulkanFunctions(context->GetDevice());
+  std::shared_ptr<std::vector<std::string>> called =
+      GetMockVulkanFunctions(context->GetDevice());
   EXPECT_EQ(
       std::count(called->begin(), called->end(), "vkCreateDescriptorPool"), 1u);
 
@@ -64,58 +44,39 @@ TEST(DescriptorPoolRecyclerVKTest, ReclaimMakesDescriptorPoolAvailable) {
 }
 
 TEST(DescriptorPoolRecyclerVKTest, ReclaimDropsDescriptorPoolIfSizeIsExceeded) {
-  auto const context = MockVulkanContextBuilder().Build();
+  std::shared_ptr<ContextVK> context = MockVulkanContextBuilder().Build();
 
   // Create 33 pools
   {
     std::vector<std::unique_ptr<DescriptorPoolVK>> pools;
-    for (auto i = 0u; i < 33; i++) {
-      auto pool = std::make_unique<DescriptorPoolVK>(context);
-      pool->AllocateDescriptorSets({}, *context);
+    for (size_t i = 0u; i < 33; i++) {
+      std::unique_ptr<DescriptorPoolVK> pool =
+          std::make_unique<DescriptorPoolVK>(context);
+      pool->AllocateDescriptorSets({}, /*pipeline_key=*/0, *context);
       pools.push_back(std::move(pool));
     }
   }
 
-  // See note above.
-  for (auto i = 0u; i < 2; i++) {
-    auto waiter = fml::AutoResetWaitableEvent();
-    auto rattle = fml::ScopedCleanupClosure([&waiter]() { waiter.Signal(); });
-    {
-      UniqueResourceVKT<fml::ScopedCleanupClosure> resource(
-          context->GetResourceManager(), std::move(rattle));
-    }
-    waiter.Wait();
-  }
-
-  auto const called = GetMockVulkanFunctions(context->GetDevice());
+  std::shared_ptr<std::vector<std::string>> called =
+      GetMockVulkanFunctions(context->GetDevice());
   EXPECT_EQ(
       std::count(called->begin(), called->end(), "vkCreateDescriptorPool"),
       33u);
-  EXPECT_EQ(std::count(called->begin(), called->end(), "vkResetDescriptorPool"),
-            33u);
 
   // Now create 33 more descriptor pools and observe that only one more is
   // allocated.
   {
-    std::vector<std::unique_ptr<DescriptorPoolVK>> pools;
-    for (auto i = 0u; i < 33; i++) {
-      auto pool = std::make_unique<DescriptorPoolVK>(context);
-      pool->AllocateDescriptorSets({}, *context);
+    std::vector<std::shared_ptr<DescriptorPoolVK>> pools;
+    for (size_t i = 0u; i < 33; i++) {
+      std::shared_ptr<DescriptorPoolVK> pool =
+          context->GetDescriptorPoolRecycler()->GetDescriptorPool();
+      pool->AllocateDescriptorSets({}, /*pipeline_key=*/0, *context);
       pools.push_back(std::move(pool));
     }
   }
 
-  for (auto i = 0u; i < 2; i++) {
-    auto waiter = fml::AutoResetWaitableEvent();
-    auto rattle = fml::ScopedCleanupClosure([&waiter]() { waiter.Signal(); });
-    {
-      UniqueResourceVKT<fml::ScopedCleanupClosure> resource(
-          context->GetResourceManager(), std::move(rattle));
-    }
-    waiter.Wait();
-  }
-
-  auto const called_twice = GetMockVulkanFunctions(context->GetDevice());
+  std::shared_ptr<std::vector<std::string>> called_twice =
+      GetMockVulkanFunctions(context->GetDevice());
   // 32 of the descriptor pools were recycled, so only one more is created.
   EXPECT_EQ(
       std::count(called->begin(), called->end(), "vkCreateDescriptorPool"),
@@ -125,10 +86,10 @@ TEST(DescriptorPoolRecyclerVKTest, ReclaimDropsDescriptorPoolIfSizeIsExceeded) {
 }
 
 TEST(DescriptorPoolRecyclerVKTest, MultipleCommandBuffersShareDescriptorPool) {
-  auto const context = MockVulkanContextBuilder().Build();
+  std::shared_ptr<ContextVK> context = MockVulkanContextBuilder().Build();
 
-  auto cmd_buffer_1 = context->CreateCommandBuffer();
-  auto cmd_buffer_2 = context->CreateCommandBuffer();
+  std::shared_ptr<CommandBuffer> cmd_buffer_1 = context->CreateCommandBuffer();
+  std::shared_ptr<CommandBuffer> cmd_buffer_2 = context->CreateCommandBuffer();
 
   CommandBufferVK& vk_1 = CommandBufferVK::Cast(*cmd_buffer_1);
   CommandBufferVK& vk_2 = CommandBufferVK::Cast(*cmd_buffer_2);
@@ -138,7 +99,7 @@ TEST(DescriptorPoolRecyclerVKTest, MultipleCommandBuffersShareDescriptorPool) {
   // Resetting resources creates a new pool.
   context->DisposeThreadLocalCachedResources();
 
-  auto cmd_buffer_3 = context->CreateCommandBuffer();
+  std::shared_ptr<CommandBuffer> cmd_buffer_3 = context->CreateCommandBuffer();
   CommandBufferVK& vk_3 = CommandBufferVK::Cast(*cmd_buffer_3);
 
   EXPECT_NE(&vk_1.GetDescriptorPool(), &vk_3.GetDescriptorPool());
@@ -146,5 +107,31 @@ TEST(DescriptorPoolRecyclerVKTest, MultipleCommandBuffersShareDescriptorPool) {
   context->Shutdown();
 }
 
+TEST(DescriptorPoolRecyclerVKTest, DescriptorsAreRecycled) {
+  std::shared_ptr<ContextVK> context = MockVulkanContextBuilder().Build();
+
+  {
+    DescriptorPoolVK pool = DescriptorPoolVK(context);
+    pool.AllocateDescriptorSets({}, /*pipeline_key=*/0, *context);
+  }
+
+  // Should reuse the same descriptor set allocated above.
+  std::shared_ptr<DescriptorPoolVK> pool =
+      context->GetDescriptorPoolRecycler()->GetDescriptorPool();
+  pool->AllocateDescriptorSets({}, /*pipeline_key=*/0, *context);
+
+  std::shared_ptr<std::vector<std::string>> called =
+      GetMockVulkanFunctions(context->GetDevice());
+  EXPECT_EQ(
+      std::count(called->begin(), called->end(), "vkAllocateDescriptorSets"),
+      1);
+
+  // Should allocate a new descriptor set.
+  pool->AllocateDescriptorSets({}, /*pipeline_key=*/0, *context);
+  EXPECT_EQ(
+      std::count(called->begin(), called->end(), "vkAllocateDescriptorSets"),
+      2);
+}
+
 }  // namespace testing
 }  // namespace impeller

engine/src/flutter/impeller/renderer/backend/vulkan/pipeline_library_vk.cc

@@ -4,23 +4,16 @@
 
 #include "impeller/renderer/backend/vulkan/pipeline_library_vk.h"
 
-#include <chrono>
 #include <cstdint>
-#include <optional>
-#include <sstream>
 
 #include "flutter/fml/container.h"
 #include "flutter/fml/trace_event.h"
 #include "impeller/base/promise.h"
-#include "impeller/base/timing.h"
 #include "impeller/base/validation.h"
 #include "impeller/renderer/backend/vulkan/context_vk.h"
 #include "impeller/renderer/backend/vulkan/formats_vk.h"
 #include "impeller/renderer/backend/vulkan/pipeline_vk.h"
 #include "impeller/renderer/backend/vulkan/shader_function_vk.h"
-#include "impeller/renderer/backend/vulkan/vertex_descriptor_vk.h"
-#include "vulkan/vulkan_core.h"
-#include "vulkan/vulkan_enums.hpp"
 
 namespace impeller {
 
@@ -50,7 +43,8 @@ bool PipelineLibraryVK::IsValid() const {
 }
 
 std::unique_ptr<ComputePipelineVK> PipelineLibraryVK::CreateComputePipeline(
-    const ComputePipelineDescriptor& desc) {
+    const ComputePipelineDescriptor& desc,
+    PipelineKey pipeline_key) {
   TRACE_EVENT0("flutter", __FUNCTION__);
   vk::ComputePipelineCreateInfo pipeline_info;
 
@@ -151,8 +145,8 @@ std::unique_ptr<ComputePipelineVK> PipelineLibraryVK::CreateComputePipeline(
       desc,                              //
       std::move(pipeline),               //
       std::move(pipeline_layout.value),  //
-      std::move(descs_layout)            //
-  );
+      std::move(descs_layout),           //
+      pipeline_key);
 }
 
 // |PipelineLibrary|
@@ -179,7 +173,8 @@ PipelineFuture<PipelineDescriptor> PipelineLibraryVK::GetPipeline(
 
   auto weak_this = weak_from_this();
 
-  auto generation_task = [descriptor, weak_this, promise]() {
+  PipelineKey next_key = pipeline_key_++;
+  auto generation_task = [descriptor, weak_this, promise, next_key]() {
     auto thiz = weak_this.lock();
     if (!thiz) {
       promise->set_value(nullptr);
@@ -191,7 +186,8 @@ PipelineFuture<PipelineDescriptor> PipelineLibraryVK::GetPipeline(
     promise->set_value(PipelineVK::Create(
         descriptor,                                            //
         PipelineLibraryVK::Cast(*thiz).device_holder_.lock(),  //
-        weak_this                                              //
+        weak_this,                                             //
+        next_key                                               //
         ));
   };
 
@@ -208,7 +204,7 @@ PipelineFuture<PipelineDescriptor> PipelineLibraryVK::GetPipeline(
 PipelineFuture<ComputePipelineDescriptor> PipelineLibraryVK::GetPipeline(
     ComputePipelineDescriptor descriptor,
     bool async) {
-  Lock lock(compute_pipelines_mutex_);
+  Lock lock(pipelines_mutex_);
   if (auto found = compute_pipelines_.find(descriptor);
       found != compute_pipelines_.end()) {
     return found->second;
@@ -230,7 +226,8 @@ PipelineFuture<ComputePipelineDescriptor> PipelineLibraryVK::GetPipeline(
 
   auto weak_this = weak_from_this();
 
-  auto generation_task = [descriptor, weak_this, promise]() {
+  PipelineKey next_key = pipeline_key_++;
+  auto generation_task = [descriptor, weak_this, promise, next_key]() {
     auto self = weak_this.lock();
     if (!self) {
       promise->set_value(nullptr);
@@ -239,8 +236,8 @@ PipelineFuture<ComputePipelineDescriptor> PipelineLibraryVK::GetPipeline(
       return;
     }
 
-    auto pipeline =
-        PipelineLibraryVK::Cast(*self).CreateComputePipeline(descriptor);
+    auto pipeline = PipelineLibraryVK::Cast(*self).CreateComputePipeline(
+        descriptor, next_key);
     if (!pipeline) {
       promise->set_value(nullptr);
       VALIDATION_LOG << "Could not create pipeline: " << descriptor.GetLabel();

engine/src/flutter/impeller/renderer/backend/vulkan/pipeline_library_vk.h

@@ -15,6 +15,7 @@
 #include "impeller/renderer/backend/vulkan/pipeline_cache_vk.h"
 #include "impeller/renderer/backend/vulkan/pipeline_vk.h"
 #include "impeller/renderer/backend/vulkan/vk.h"
+#include "impeller/renderer/pipeline.h"
 #include "impeller/renderer/pipeline_library.h"
 
 namespace impeller {
@@ -42,10 +43,9 @@ class PipelineLibraryVK final
   std::shared_ptr<fml::ConcurrentTaskRunner> worker_task_runner_;
   Mutex pipelines_mutex_;
   PipelineMap pipelines_ IPLR_GUARDED_BY(pipelines_mutex_);
-  Mutex compute_pipelines_mutex_;
-  ComputePipelineMap compute_pipelines_ IPLR_GUARDED_BY(
-      compute_pipelines_mutex_);
+  ComputePipelineMap compute_pipelines_ IPLR_GUARDED_BY(pipelines_mutex_);
   std::atomic_size_t frames_acquired_ = 0u;
+  PipelineKey pipeline_key_ IPLR_GUARDED_BY(pipelines_mutex_) = 1;
   bool is_valid_ = false;
   bool cache_dirty_ = false;
 
@@ -75,7 +75,8 @@ class PipelineLibraryVK final
       std::shared_ptr<const ShaderFunction> function) override;
 
   std::unique_ptr<ComputePipelineVK> CreateComputePipeline(
-      const ComputePipelineDescriptor& desc);
+      const ComputePipelineDescriptor& desc,
+      PipelineKey pipeline_key);
 
   void PersistPipelineCacheToDisk();
 

engine/src/flutter/impeller/renderer/backend/vulkan/pipeline_vk.cc

@@ -465,6 +465,7 @@ std::unique_ptr<PipelineVK> PipelineVK::Create(
     const PipelineDescriptor& desc,
     const std::shared_ptr<DeviceHolderVK>& device_holder,
     const std::weak_ptr<PipelineLibrary>& weak_library,
+    PipelineKey pipeline_key,
     std::shared_ptr<SamplerVK> immutable_sampler) {
   TRACE_EVENT1("flutter", "PipelineVK::Create", "Name", desc.GetLabel().data());
 
@@ -510,6 +511,7 @@ std::unique_ptr<PipelineVK> PipelineVK::Create(
       std::move(render_pass),              //
       std::move(pipeline_layout.value()),  //
       std::move(descs_layout.value()),     //
+      pipeline_key,                        //
       std::move(immutable_sampler)         //
       ));
   if (!pipeline_vk->IsValid()) {
@@ -526,6 +528,7 @@ PipelineVK::PipelineVK(std::weak_ptr<DeviceHolderVK> device_holder,
                        vk::UniqueRenderPass render_pass,
                        vk::UniquePipelineLayout layout,
                        vk::UniqueDescriptorSetLayout descriptor_set_layout,
+                       PipelineKey pipeline_key,
                        std::shared_ptr<SamplerVK> immutable_sampler)
     : Pipeline(std::move(library), desc),
       device_holder_(std::move(device_holder)),
@@ -533,7 +536,8 @@ PipelineVK::PipelineVK(std::weak_ptr<DeviceHolderVK> device_holder,
       render_pass_(std::move(render_pass)),
       layout_(std::move(layout)),
       descriptor_set_layout_(std::move(descriptor_set_layout)),
-      immutable_sampler_(std::move(immutable_sampler)) {
+      immutable_sampler_(std::move(immutable_sampler)),
+      pipeline_key_(pipeline_key) {
   is_valid_ = pipeline_ && render_pass_ && layout_ && descriptor_set_layout_;
 }
 
@@ -578,7 +582,8 @@ std::shared_ptr<PipelineVK> PipelineVK::CreateVariantForImmutableSamplers(
     return nullptr;
   }
   return (immutable_sampler_variants_[cache_key] =
-              Create(desc_, device_holder, library_, immutable_sampler));
+              Create(desc_, device_holder, library_, pipeline_key_,
+                     immutable_sampler));
 }
 
 }  // namespace impeller

engine/src/flutter/impeller/renderer/backend/vulkan/pipeline_vk.h

@@ -10,10 +10,7 @@
 
 #include "impeller/base/backend_cast.h"
 #include "impeller/base/thread.h"
-#include "impeller/core/texture.h"
 #include "impeller/renderer/backend/vulkan/device_holder_vk.h"
-#include "impeller/renderer/backend/vulkan/formats_vk.h"
-#include "impeller/renderer/backend/vulkan/pipeline_cache_vk.h"
 #include "impeller/renderer/backend/vulkan/sampler_vk.h"
 #include "impeller/renderer/backend/vulkan/vk.h"
 #include "impeller/renderer/backend/vulkan/yuv_conversion_vk.h"
@@ -33,6 +30,7 @@ class PipelineVK final
       const PipelineDescriptor& desc,
       const std::shared_ptr<DeviceHolderVK>& device_holder,
       const std::weak_ptr<PipelineLibrary>& weak_library,
+      PipelineKey pipeline_key,
       std::shared_ptr<SamplerVK> immutable_sampler = {});
 
   // |Pipeline|
@@ -47,6 +45,8 @@ class PipelineVK final
   std::shared_ptr<PipelineVK> CreateVariantForImmutableSamplers(
       const std::shared_ptr<SamplerVK>& immutable_sampler) const;
 
+  PipelineKey GetPipelineKey() const { return pipeline_key_; }
+
  private:
   friend class PipelineLibraryVK;
 
@@ -62,6 +62,7 @@ class PipelineVK final
   vk::UniquePipelineLayout layout_;
   vk::UniqueDescriptorSetLayout descriptor_set_layout_;
   std::shared_ptr<SamplerVK> immutable_sampler_;
+  const PipelineKey pipeline_key_;
   mutable Mutex immutable_sampler_variants_mutex_;
   mutable ImmutableSamplerVariants immutable_sampler_variants_ IPLR_GUARDED_BY(
       immutable_sampler_variants_mutex_);
@@ -74,6 +75,7 @@ class PipelineVK final
              vk::UniqueRenderPass render_pass,
              vk::UniquePipelineLayout layout,
              vk::UniqueDescriptorSetLayout descriptor_set_layout,
+             PipelineKey pipeline_key,
              std::shared_ptr<SamplerVK> immutable_sampler);
 
   // |Pipeline|

engine/src/flutter/impeller/renderer/backend/vulkan/render_pass_vk.cc

@@ -489,7 +489,8 @@ fml::Status RenderPassVK::Draw() {
   const auto& pipeline_vk = PipelineVK::Cast(*pipeline_);
 
   auto descriptor_result = command_buffer_->AllocateDescriptorSets(
-      pipeline_vk.GetDescriptorSetLayout(), context_vk);
+      pipeline_vk.GetDescriptorSetLayout(), pipeline_vk.GetPipelineKey(),
+      context_vk);
   if (!descriptor_result.ok()) {
     return fml::Status(fml::StatusCode::kAborted,
                        "Could not allocate descriptor sets.");

engine/src/flutter/impeller/renderer/pipeline.h

@@ -18,6 +18,8 @@
 
 namespace impeller {
 
+using PipelineKey = uint64_t;
+
 class PipelineLibrary;
 template <typename PipelineDescriptor_>
 class Pipeline;