xref: /aosp_15_r20/external/executorch/backends/vulkan/runtime/vk_api/memory/Allocator.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #pragma once
10 
11 // @lint-ignore-every CLANGTIDY facebook-hte-BadMemberName
12 
13 #include <executorch/backends/vulkan/runtime/vk_api/vk_api.h>
14 
15 #include <executorch/backends/vulkan/runtime/utils/VecUtils.h>
16 
17 #include <executorch/backends/vulkan/runtime/vk_api/memory/vma_api.h>
18 
19 #include <executorch/backends/vulkan/runtime/vk_api/memory/Allocation.h>
20 #include <executorch/backends/vulkan/runtime/vk_api/memory/Buffer.h>
21 #include <executorch/backends/vulkan/runtime/vk_api/memory/Image.h>
22 
23 namespace vkcompute {
24 namespace vkapi {
25 
26 constexpr VmaAllocationCreateFlags DEFAULT_ALLOCATION_STRATEGY =
27     VMA_ALLOCATION_CREATE_STRATEGY_MIN_MEMORY_BIT;
28 
29 class Allocator final {
30  public:
31   explicit Allocator(
32       VkInstance instance,
33       VkPhysicalDevice physical_device,
34       VkDevice device);
35 
36   Allocator(const Allocator&) = delete;
37   Allocator& operator=(const Allocator&) = delete;
38 
39   Allocator(Allocator&&) noexcept;
40   Allocator& operator=(Allocator&&) = delete;
41 
42   ~Allocator();
43 
44  private:
45   VkInstance instance_;
46   VkPhysicalDevice physical_device_;
47   VkDevice device_;
48   VmaAllocator allocator_;
49 
50  public:
51   VmaAllocationCreateInfo gpuonly_resource_create_info();
52 
53   Allocation create_allocation(
54       const VkMemoryRequirements& memory_requirements,
55       const VmaAllocationCreateInfo& create_info);
56 
57   VulkanImage create_image(
58       const VkDevice,
59       const VkExtent3D&,
60       const VkFormat,
61       const VkImageType,
62       const VkImageTiling,
63       const VkImageViewType,
64       const VulkanImage::SamplerProperties&,
65       VkSampler,
66       const bool allow_transfer = false,
67       const bool allocate_memory = true);
68 
69   VulkanBuffer create_staging_buffer(const VkDeviceSize);
70 
71   VulkanBuffer create_storage_buffer(
72       const VkDeviceSize,
73       const bool allocate_memory = true);
74 
75   /*
76    * Create a uniform buffer with a specified size
77    */
78   VulkanBuffer create_uniform_buffer(const VkDeviceSize);
79 
80   /*
81    * Create a uniform buffer containing the data in an arbitrary struct
82    */
83   template <typename Block>
84   VulkanBuffer create_params_buffer(const Block& block);
85 
get_memory_statistics()86   VmaTotalStatistics get_memory_statistics() const {
87     VmaTotalStatistics stats = {};
88     vmaCalculateStatistics(allocator_, &stats);
89     return stats;
90   }
91 };
92 
93 //
94 // Impl
95 //
96 
97 template <typename Block>
create_params_buffer(const Block & block)98 inline VulkanBuffer Allocator::create_params_buffer(const Block& block) {
99   VulkanBuffer uniform_buffer = create_uniform_buffer(sizeof(Block));
100 
101   // Fill the uniform buffer with data in block
102   {
103     MemoryMap mapping(uniform_buffer, MemoryAccessType::WRITE);
104     Block* data_ptr = mapping.template data<Block>();
105 
106     *data_ptr = block;
107   }
108 
109   return uniform_buffer;
110 }
111 
112 } // namespace vkapi
113 } // namespace vkcompute
114