From f96d162ddb32a204c31d573b327e35eceeb05d4f Mon Sep 17 00:00:00 2001
From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com>
Date: Thu, 8 May 2025 10:40:57 -0700
Subject: [PATCH] [ET-VK] Removing descriptor pool intialization from
 DescriptorPool ctor.

Descriptor pool is initialized with proper capacity while ComputeGraph is build, in ComputeGraph::prepare function.
Thus, initializing descriptor pool in constructor just add throwaway work.

Differential Revision: [D74347030](https://our.internmc.facebook.com/intern/diff/D74347030/)

[ghstack-poisoned]
---
 backends/vulkan/runtime/api/Context.cpp       | 22 ++++++++++---------
 backends/vulkan/runtime/api/Context.h         |  2 ++
 backends/vulkan/runtime/vk_api/Descriptor.cpp |  6 +----
 .../vulkan/test/vulkan_compute_api_test.cpp   |  3 +++
 4 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp
index df3d43ee901..b8281ca6d4a 100644
--- a/backends/vulkan/runtime/api/Context.cpp
+++ b/backends/vulkan/runtime/api/Context.cpp
@@ -235,15 +235,6 @@ Context* context() {
           8u, // cmdPoolBatchSize
       };
 
-      const vkapi::DescriptorPoolConfig descriptor_pool_config{
-          VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorPoolMaxSets
-          VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorUniformBufferCount
-          VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorStorageBufferCount
-          VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorCombinedSamplerCount
-          VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorStorageImageCount
-          32u, // descriptorPileSizes
-      };
-
       const vkapi::QueryPoolConfig query_pool_config{
           VULKAN_QUERY_POOL_SIZE, // maxQueryCount
           256u, // initialReserveSize
@@ -252,7 +243,7 @@ Context* context() {
       const ContextConfig config{
           cmd_submit_frequency,
           cmd_config,
-          descriptor_pool_config,
+          {},
           query_pool_config,
       };
 
@@ -266,6 +257,17 @@ Context* context() {
   return context.get();
 }
 
+vkapi::DescriptorPoolConfig default_descriptor_pool_config() {
+  return {
+      VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorPoolMaxSets
+      VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorUniformBufferCount
+      VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorStorageBufferCount
+      VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorCombinedSamplerCount
+      VULKAN_DESCRIPTOR_POOL_SIZE, // descriptorStorageImageCount
+      32u, // descriptorPileSizes
+  };
+}
+
 #ifdef VULKAN_DEBUG
 
 #ifdef VK_KHR_pipeline_executable_properties
diff --git a/backends/vulkan/runtime/api/Context.h b/backends/vulkan/runtime/api/Context.h
index 6cfbc64f141..8a96a7d2e73 100644
--- a/backends/vulkan/runtime/api/Context.h
+++ b/backends/vulkan/runtime/api/Context.h
@@ -267,6 +267,8 @@ bool available();
 // a static local variable.
 Context* context();
 
+vkapi::DescriptorPoolConfig default_descriptor_pool_config();
+
 namespace detail {
 
 inline void arg_is_empty(
diff --git a/backends/vulkan/runtime/vk_api/Descriptor.cpp b/backends/vulkan/runtime/vk_api/Descriptor.cpp
index b42ade6ea02..3e5c357c871 100644
--- a/backends/vulkan/runtime/vk_api/Descriptor.cpp
+++ b/backends/vulkan/runtime/vk_api/Descriptor.cpp
@@ -270,11 +270,7 @@ DescriptorPool::DescriptorPool(
       pool_(VK_NULL_HANDLE),
       config_(config),
       mutex_{},
-      piles_{} {
-  if (config.descriptor_pool_max_sets > 0) {
-    init(config);
-  }
-}
+      piles_{} {}
 
 DescriptorPool::~DescriptorPool() {
   if (pool_ == VK_NULL_HANDLE) {
diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp
index 143e6704889..1ff3937e07e 100644
--- a/backends/vulkan/test/vulkan_compute_api_test.cpp
+++ b/backends/vulkan/test/vulkan_compute_api_test.cpp
@@ -87,6 +87,9 @@ class VulkanComputeAPITest : public ::testing::Test {
   void SetUp() override {
     // Make sure we are starting with a clean slate
     EXPECT_TRUE(get_vma_allocation_count() == 0);
+    if (!context()->descriptor_pool()) {
+      context()->descriptor_pool().init(default_descriptor_pool_config());
+    }
   }
 
   void TearDown() override {