xref: /aosp_15_r20/external/executorch/backends/vulkan/runtime/vk_api/memory/Allocator.cpp (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 #include <executorch/backends/vulkan/runtime/vk_api/memory/Allocator.h>
10 
11 namespace vkcompute {
12 namespace vkapi {
13 
Allocator(VkInstance instance,VkPhysicalDevice physical_device,VkDevice device)14 Allocator::Allocator(
15     VkInstance instance,
16     VkPhysicalDevice physical_device,
17     VkDevice device)
18     : instance_{},
19       physical_device_(physical_device),
20       device_(device),
21       allocator_{VK_NULL_HANDLE} {
22   VmaVulkanFunctions vk_functions{};
23   vk_functions.vkGetInstanceProcAddr = vkGetInstanceProcAddr;
24   vk_functions.vkGetDeviceProcAddr = vkGetDeviceProcAddr;
25 
26   const VmaAllocatorCreateInfo allocator_create_info{
27       0u, // flags
28       physical_device_, // physicalDevice
29       device_, // device
30       0u, // preferredLargeHeapBlockSize
31       nullptr, // pAllocationCallbacks
32       nullptr, // pDeviceMemoryCallbacks
33       nullptr, // pHeapSizeLimit
34       &vk_functions, // pVulkanFunctions
35       instance, // instance
36       VK_API_VERSION_1_0, // vulkanApiVersion
37       nullptr, // pTypeExternalMemoryHandleTypes
38   };
39 
40   VK_CHECK(vmaCreateAllocator(&allocator_create_info, &allocator_));
41 }
42 
Allocator(Allocator && other)43 Allocator::Allocator(Allocator&& other) noexcept
44     : instance_(other.instance_),
45       physical_device_(other.physical_device_),
46       device_(other.device_),
47       allocator_(other.allocator_) {
48   other.allocator_ = VK_NULL_HANDLE;
49   other.device_ = VK_NULL_HANDLE;
50   other.physical_device_ = VK_NULL_HANDLE;
51   other.instance_ = VK_NULL_HANDLE;
52 }
53 
~Allocator()54 Allocator::~Allocator() {
55   if (allocator_ == VK_NULL_HANDLE) {
56     return;
57   }
58   vmaDestroyAllocator(allocator_);
59 }
60 
gpuonly_resource_create_info()61 VmaAllocationCreateInfo Allocator::gpuonly_resource_create_info() {
62   VmaAllocationCreateInfo alloc_create_info = {};
63   alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
64   alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
65   return alloc_create_info;
66 }
67 
create_allocation(const VkMemoryRequirements & memory_requirements,const VmaAllocationCreateInfo & create_info)68 Allocation Allocator::create_allocation(
69     const VkMemoryRequirements& memory_requirements,
70     const VmaAllocationCreateInfo& create_info) {
71   VmaAllocationCreateInfo alloc_create_info = create_info;
72   // Protect against using VMA_MEMORY_USAGE_AUTO_* flags when allocating memory
73   // directly, since those usage flags require that VkBufferCreateInfo and/or
74   // VkImageCreateInfo also be available.
75   switch (create_info.usage) {
76     // The logic for the below usage options are too complex, therefore prevent
77     // those from being used with direct memory allocation.
78     case VMA_MEMORY_USAGE_AUTO:
79     case VMA_MEMORY_USAGE_AUTO_PREFER_HOST:
80       VK_THROW(
81           "Only the VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE usage flag is compatible with create_allocation()");
82       break;
83     // Most of the time, VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE will simply set the
84     // DEVICE_LOCAL_BIT as a preferred memory flag. Therefore the below is a
85     // decent approximation for VMA behaviour.
86     case VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE:
87       alloc_create_info.preferredFlags = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
88       alloc_create_info.usage = VMA_MEMORY_USAGE_UNKNOWN;
89       break;
90     default:
91       break;
92   }
93 
94   return Allocation(allocator_, memory_requirements, alloc_create_info);
95 }
96 
create_image(const VkDevice device,const VkExtent3D & extents,const VkFormat image_format,const VkImageType image_type,const VkImageTiling image_tiling,const VkImageViewType image_view_type,const VulkanImage::SamplerProperties & sampler_props,VkSampler sampler,const bool allow_transfer,const bool allocate_memory)97 VulkanImage Allocator::create_image(
98     const VkDevice device,
99     const VkExtent3D& extents,
100     const VkFormat image_format,
101     const VkImageType image_type,
102     const VkImageTiling image_tiling,
103     const VkImageViewType image_view_type,
104     const VulkanImage::SamplerProperties& sampler_props,
105     VkSampler sampler,
106     const bool allow_transfer,
107     const bool allocate_memory) {
108   VkImageUsageFlags usage =
109       VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT;
110   if (allow_transfer) {
111     usage |=
112         (VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT);
113   }
114 
115   VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
116 
117   const VulkanImage::ImageProperties image_props{
118       image_type,
119       image_format,
120       extents,
121       image_tiling,
122       usage,
123   };
124 
125   const VulkanImage::ViewProperties view_props{
126       image_view_type,
127       image_format,
128   };
129 
130   const VkImageLayout initial_layout = VK_IMAGE_LAYOUT_UNDEFINED;
131 
132   return VulkanImage(
133       device,
134       allocator_,
135       alloc_create_info,
136       image_props,
137       view_props,
138       sampler_props,
139       sampler,
140       initial_layout,
141       allocate_memory);
142 }
143 
create_staging_buffer(const VkDeviceSize size)144 VulkanBuffer Allocator::create_staging_buffer(const VkDeviceSize size) {
145   const VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
146 
147   VmaAllocationCreateInfo alloc_create_info = {};
148   alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
149   alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
150 
151   // Staging buffers are accessed by both the CPU and GPU, so set the
152   // appropriate flags to indicate that the host device will be accessing
153   // the data from this buffer.
154   alloc_create_info.flags |= VMA_ALLOCATION_CREATE_HOST_ACCESS_RANDOM_BIT |
155       VMA_ALLOCATION_CREATE_MAPPED_BIT;
156   alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_HOST;
157   alloc_create_info.requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
158   alloc_create_info.preferredFlags =
159       VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_CACHED_BIT;
160 
161   return VulkanBuffer(allocator_, size, alloc_create_info, buffer_usage);
162 }
163 
create_storage_buffer(const VkDeviceSize size,const bool allocate_memory)164 VulkanBuffer Allocator::create_storage_buffer(
165     const VkDeviceSize size,
166     const bool allocate_memory) {
167   const VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
168 
169   VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
170   return VulkanBuffer(
171       allocator_, size, alloc_create_info, buffer_usage, allocate_memory);
172 }
173 
create_uniform_buffer(const VkDeviceSize size)174 VulkanBuffer Allocator::create_uniform_buffer(const VkDeviceSize size) {
175   VmaAllocationCreateInfo alloc_create_info = {};
176   alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY |
177       VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT;
178   alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO;
179 
180   VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
181 
182   return VulkanBuffer(allocator_, size, alloc_create_info, buffer_usage);
183 }
184 
185 } // namespace vkapi
186 } // namespace vkcompute
187