[ET-VK] custom memory pools

backends/vulkan/runtime/api/Context.cpp

@@ -40,6 +40,10 @@ Context::Context(size_t adapter_i, const ContextConfig& config)
       cmd_mutex_{},
       cmd_(VK_NULL_HANDLE, 0u),
       submit_count_{0u},
+      // Custom memory pools
+      custom_vma_pool_(
+          adapter_p_->vma().handle(),
+          adapter_p_->num_memory_types()),
       // Memory Management
       buffer_clearlist_mutex_{},
       buffers_to_clear_{},
@@ -60,6 +64,13 @@ Context::~Context() {
   }
 }
 
+vkapi::MemoryPoolManager* Context::get_custom_memory_pool_ptr() {
+  if (config_.use_custom_vma_pools) {
+    return &custom_vma_pool_;
+  }
+  return nullptr;
+}
+
 void Context::initialize_querypool() {
   querypool_.initialize(adapter_p_);
 }

backends/vulkan/runtime/api/Context.h

@@ -13,6 +13,8 @@
 #include <executorch/backends/vulkan/runtime/utils/MacroUtils.h>
 #include <executorch/backends/vulkan/runtime/utils/VecUtils.h>
 
+#include <executorch/backends/vulkan/runtime/vk_api/memory/Pool.h>
+
 #include <executorch/backends/vulkan/runtime/vk_api/Adapter.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Command.h>
 #include <executorch/backends/vulkan/runtime/vk_api/Descriptor.h>
@@ -29,6 +31,7 @@ struct ContextConfig final {
   vkapi::CommandPoolConfig cmd_pool_config;
   vkapi::DescriptorPoolConfig descriptor_pool_config;
   vkapi::QueryPoolConfig query_pool_config;
+  bool use_custom_vma_pools;
 };
 
 //
@@ -69,6 +72,9 @@ class Context final {
   std::mutex cmd_mutex_;
   vkapi::CommandBuffer cmd_;
   uint32_t submit_count_;
+  // Custom memory pool that can be used to allocate resources that are used in
+  // this command stream
+  vkapi::MemoryPoolManager custom_vma_pool_;
   // Memory Management
   std::mutex buffer_clearlist_mutex_;
   std::vector<vkapi::VulkanBuffer> buffers_to_clear_;
@@ -78,6 +84,10 @@ class Context final {
   VkImageTiling preferred_image_tiling_;
 
  public:
+  inline const ContextConfig& config() const {
+    return config_;
+  }
+
   // Adapter access
 
   inline vkapi::Adapter* adapter_ptr() {
@@ -130,6 +140,8 @@ class Context final {
     return preferred_image_tiling_;
   }
 
+  vkapi::MemoryPoolManager* get_custom_memory_pool_ptr();
+
   /*
    * By default, the querypool attached to a Context instance is uninitialized.
    * This function triggers the querypool to be created via vkCreateQueryPool.

backends/vulkan/runtime/api/api.h

@@ -30,3 +30,4 @@
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Allocator.h>
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Buffer.h>
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Image.h>
+#include <executorch/backends/vulkan/runtime/vk_api/memory/Pool.h>

backends/vulkan/runtime/api/containers/ParamsBuffer.cpp

@@ -36,7 +36,8 @@ ParamsBuffer::ParamsBuffer(const ParamsBuffer& other)
     : context_p_(other.context_p_), vulkan_buffer_{} {
   if (other.vulkan_buffer_) {
     vulkan_buffer_ = context_p_->adapter_ptr()->vma().create_uniform_buffer(
-        other.vulkan_buffer_.mem_size());
+        other.vulkan_buffer_.mem_size(),
+        context_p_->get_custom_memory_pool_ptr());
 
     memcpy_to_buffer(other.vulkan_buffer_, vulkan_buffer_);
   }
@@ -55,7 +56,8 @@ ParamsBuffer& ParamsBuffer::operator=(const ParamsBuffer& other) {
 
     if (other.vulkan_buffer_) {
       vulkan_buffer_ = context_p_->adapter_ptr()->vma().create_uniform_buffer(
-          other.vulkan_buffer_.mem_size());
+          other.vulkan_buffer_.mem_size(),
+          context_p_->get_custom_memory_pool_ptr());
 
       memcpy_to_buffer(other.vulkan_buffer_, vulkan_buffer_);
     }

backends/vulkan/runtime/api/containers/ParamsBuffer.h

@@ -35,8 +35,9 @@ class ParamsBuffer final {
   // constructor from the one above.
   ParamsBuffer(Context* context_p, const VkDeviceSize nbytes, const bool unused)
       : context_p_(context_p),
-        vulkan_buffer_(
-            context_p_->adapter_ptr()->vma().create_uniform_buffer(nbytes)) {}
+        vulkan_buffer_(context_p_->adapter_ptr()->vma().create_uniform_buffer(
+            nbytes,
+            context_p->get_custom_memory_pool_ptr())) {}
 
   ParamsBuffer(const ParamsBuffer&);
   ParamsBuffer& operator=(const ParamsBuffer&);

backends/vulkan/runtime/api/containers/StagingBuffer.h

@@ -35,7 +35,8 @@ class StagingBuffer final {
       : context_p_(context_p),
         dtype_(dtype),
         vulkan_buffer_(context_p_->adapter_ptr()->vma().create_staging_buffer(
-            element_size(dtype_) * numel)),
+            element_size(dtype_) * numel,
+            context_p_->get_custom_memory_pool_ptr())),
         mapped_data_(nullptr) {}
 
   StagingBuffer(const StagingBuffer&) = delete;

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -292,7 +292,8 @@ vkapi::VulkanImage allocate_image(
       sampler_props,
       sampler,
       /*allow_transfer = */ true,
-      /*allocate_memory = */ allocate_memory);
+      /*allocate_memory = */ allocate_memory,
+      /*pool_manager = */ context_ptr->get_custom_memory_pool_ptr());
 }
 
 vkapi::VulkanBuffer allocate_buffer(
@@ -301,8 +302,6 @@ vkapi::VulkanBuffer allocate_buffer(
     const utils::StorageType storage_type,
     const vkapi::ScalarType dtype,
     const bool allocate_memory) {
-  vkapi::Adapter* adapter_ptr = context_ptr->adapter_ptr();
-
   switch (storage_type) {
     case utils::kBuffer:
       break;
@@ -313,8 +312,10 @@ vkapi::VulkanBuffer allocate_buffer(
 
   VK_CHECK_COND(numel <= context_ptr->adapter_ptr()->max_buffer_numel());
 
-  return adapter_ptr->vma().create_storage_buffer(
-      element_size(dtype) * numel, allocate_memory);
+  return context_ptr->adapter_ptr()->vma().create_storage_buffer(
+      element_size(dtype) * numel,
+      allocate_memory,
+      context_ptr->get_custom_memory_pool_ptr());
 }
 
 vTensorStorage::vTensorStorage(

backends/vulkan/runtime/graph/GraphConfig.cpp

@@ -40,6 +40,7 @@ GraphConfig::GraphConfig() {
       cmd_config,
       descriptor_pool_config,
       query_pool_config,
+      false,
   };
 
   // Empirically selected safety factor. If descriptor pools start running out

backends/vulkan/runtime/graph/containers/SharedObject.cpp

@@ -39,7 +39,9 @@ void SharedObject::allocate(ComputeGraph* const graph) {
       graph->context()->adapter_ptr()->vma().gpuonly_resource_create_info();
 
   allocation = graph->context()->adapter_ptr()->vma().create_allocation(
-      aggregate_memory_requirements, alloc_create_info);
+      aggregate_memory_requirements,
+      alloc_create_info,
+      graph->context()->get_custom_memory_pool_ptr());
 }
 
 void SharedObject::bind_users(ComputeGraph* const graph) {

backends/vulkan/runtime/vk_api/Adapter.h

@@ -110,6 +110,10 @@ class Adapter final {
     return physical_device_.has_unified_memory;
   }
 
+  inline uint32_t num_memory_types() const {
+    return physical_device_.memory_properties.memoryTypeCount;
+  }
+
   inline uint32_t num_compute_queues() const {
     return physical_device_.num_compute_queues;
   }

backends/vulkan/runtime/vk_api/memory/Allocator.cpp

@@ -67,7 +67,8 @@ VmaAllocationCreateInfo Allocator::gpuonly_resource_create_info() {
 
 Allocation Allocator::create_allocation(
     const VkMemoryRequirements& memory_requirements,
-    const VmaAllocationCreateInfo& create_info) {
+    const VmaAllocationCreateInfo& create_info,
+    MemoryPoolManager* pool_manager) {
   VmaAllocationCreateInfo alloc_create_info = create_info;
   // Protect against using VMA_MEMORY_USAGE_AUTO_* flags when allocating memory
   // directly, since those usage flags require that VkBufferCreateInfo and/or
@@ -90,6 +91,11 @@ Allocation Allocator::create_allocation(
     default:
       break;
   }
+  if (pool_manager) {
+    uint32_t memory_type_idx =
+        pool_manager->get_memory_type_idx(alloc_create_info);
+    alloc_create_info.pool = pool_manager->get_memory_pool(memory_type_idx);
+  }
 
   return Allocation(allocator_, memory_requirements, alloc_create_info);
 }
@@ -104,7 +110,8 @@ VulkanImage Allocator::create_image(
     const VulkanImage::SamplerProperties& sampler_props,
     VkSampler sampler,
     const bool allow_transfer,
-    const bool allocate_memory) {
+    const bool allocate_memory,
+    MemoryPoolManager* pool_manager) {
   VkImageUsageFlags usage =
       VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT;
   if (allow_transfer) {
@@ -129,6 +136,15 @@ VulkanImage Allocator::create_image(
 
   const VkImageLayout initial_layout = VK_IMAGE_LAYOUT_UNDEFINED;
 
+  if (pool_manager) {
+    VkImageCreateInfo image_create_info = vkapi::generate_image_create_info(
+        image_type, image_format, extents, image_tiling, usage, initial_layout);
+
+    uint32_t memory_type_idx =
+        pool_manager->get_memory_type_idx(alloc_create_info, image_create_info);
+    alloc_create_info.pool = pool_manager->get_memory_pool(memory_type_idx);
+  }
+
   return VulkanImage(
       device,
       allocator_,
@@ -141,7 +157,26 @@ VulkanImage Allocator::create_image(
       allocate_memory);
 }
 
-VulkanBuffer Allocator::create_staging_buffer(const VkDeviceSize size) {
+VmaPool get_memory_pool_for_buffer(
+    MemoryPoolManager* pool_manager,
+    const VkDeviceSize size,
+    const VkBufferUsageFlags buffer_usage,
+    const VmaAllocationCreateInfo& alloc_create_info) {
+  if (pool_manager) {
+    VkBufferCreateInfo buffer_create_info =
+        vkapi::generate_buffer_create_info(size, buffer_usage);
+
+    uint32_t memory_type_idx = pool_manager->get_memory_type_idx(
+        alloc_create_info, buffer_create_info);
+
+    return pool_manager->get_memory_pool(memory_type_idx);
+  }
+  return VK_NULL_HANDLE; // Return a default value if pool_manager is null
+}
+
+VulkanBuffer Allocator::create_staging_buffer(
+    const VkDeviceSize size,
+    MemoryPoolManager* pool_manager) {
   const VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
 
   VmaAllocationCreateInfo alloc_create_info = {};
@@ -159,26 +194,37 @@ VulkanBuffer Allocator::create_staging_buffer(const VkDeviceSize size) {
   alloc_create_info.preferredFlags =
       VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_CACHED_BIT;
 
+  alloc_create_info.pool = get_memory_pool_for_buffer(
+      pool_manager, size, buffer_usage, alloc_create_info);
+
   return VulkanBuffer(allocator_, size, alloc_create_info, buffer_usage);
 }
 
 VulkanBuffer Allocator::create_storage_buffer(
     const VkDeviceSize size,
-    const bool allocate_memory) {
+    const bool allocate_memory,
+    MemoryPoolManager* pool_manager) {
   const VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
 
   VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
+  alloc_create_info.pool = get_memory_pool_for_buffer(
+      pool_manager, size, buffer_usage, alloc_create_info);
+
   return VulkanBuffer(
       allocator_, size, alloc_create_info, buffer_usage, allocate_memory);
 }
 
-VulkanBuffer Allocator::create_uniform_buffer(const VkDeviceSize size) {
+VulkanBuffer Allocator::create_uniform_buffer(
+    const VkDeviceSize size,
+    MemoryPoolManager* pool_manager) {
   VmaAllocationCreateInfo alloc_create_info = {};
   alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY |
       VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT;
   alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO;
 
   VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
+  alloc_create_info.pool = get_memory_pool_for_buffer(
+      pool_manager, size, buffer_usage, alloc_create_info);
 
   return VulkanBuffer(allocator_, size, alloc_create_info, buffer_usage);
 }

backends/vulkan/runtime/vk_api/memory/Allocator.h

@@ -19,6 +19,7 @@
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Allocation.h>
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Buffer.h>
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Image.h>
+#include <executorch/backends/vulkan/runtime/vk_api/memory/Pool.h>
 
 namespace vkcompute {
 namespace vkapi {
@@ -48,11 +49,16 @@ class Allocator final {
   VmaAllocator allocator_;
 
  public:
+  inline VmaAllocator handle() {
+    return allocator_;
+  }
+
   VmaAllocationCreateInfo gpuonly_resource_create_info();
 
   Allocation create_allocation(
       const VkMemoryRequirements& memory_requirements,
-      const VmaAllocationCreateInfo& create_info);
+      const VmaAllocationCreateInfo& create_info,
+      MemoryPoolManager* pool_manager = nullptr);
 
   VulkanImage create_image(
       const VkDevice,
@@ -64,18 +70,24 @@ class Allocator final {
       const VulkanImage::SamplerProperties&,
       VkSampler,
       const bool allow_transfer = false,
-      const bool allocate_memory = true);
+      const bool allocate_memory = true,
+      MemoryPoolManager* pool_manager = nullptr);
 
-  VulkanBuffer create_staging_buffer(const VkDeviceSize);
+  VulkanBuffer create_staging_buffer(
+      const VkDeviceSize,
+      MemoryPoolManager* pool_manager = nullptr);
 
   VulkanBuffer create_storage_buffer(
       const VkDeviceSize,
-      const bool allocate_memory = true);
+      const bool allocate_memory = true,
+      MemoryPoolManager* pool_manager = nullptr);
 
   /*
    * Create a uniform buffer with a specified size
    */
-  VulkanBuffer create_uniform_buffer(const VkDeviceSize);
+  VulkanBuffer create_uniform_buffer(
+      const VkDeviceSize,
+      MemoryPoolManager* pool_manager = nullptr);
 
   /*
    * Create a uniform buffer containing the data in an arbitrary struct

backends/vulkan/runtime/vk_api/memory/Buffer.cpp

@@ -11,6 +11,21 @@
 namespace vkcompute {
 namespace vkapi {
 
+VkBufferCreateInfo generate_buffer_create_info(
+    VkDeviceSize size,
+    VkBufferUsageFlags usage) {
+  return VkBufferCreateInfo{
+      VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, // sType
+      nullptr, // pNext
+      0u, // flags
+      size, // size
+      usage, // usage
+      VK_SHARING_MODE_EXCLUSIVE, // sharingMode
+      0u, // queueFamilyIndexCount
+      nullptr, // pQueueFamilyIndices
+  };
+}
+
 //
 // VulkanBuffer
 //
@@ -42,16 +57,8 @@ VulkanBuffer::VulkanBuffer(
     buffer_properties_.mem_range = 1u;
   }
 
-  const VkBufferCreateInfo buffer_create_info{
-      VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, // sType
-      nullptr, // pNext
-      0u, // flags
-      buffer_properties_.size, // size
-      usage, // usage
-      VK_SHARING_MODE_EXCLUSIVE, // sharingMode
-      0u, // queueFamilyIndexCount
-      nullptr, // pQueueFamilyIndices
-  };
+  const VkBufferCreateInfo buffer_create_info =
+      generate_buffer_create_info(buffer_properties_.size, usage);
 
   if (allocate_memory) {
     VK_CHECK(vmaCreateBuffer(