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/vk_api/Shader.h> 16 17 #include <executorch/backends/vulkan/runtime/vk_api/memory/Buffer.h> 18 #include <executorch/backends/vulkan/runtime/vk_api/memory/Image.h> 19 20 #include <unordered_map> 21 22 namespace vkcompute { 23 namespace vkapi { 24 25 /* 26 * Stores the binding information of a Vulkan Buffer so that the buffer can be 27 * bound at a later time. This struct should only be used if the buffer to be 28 * bound is guaranteed to be active at the time of binding. 29 */ 30 struct BufferBindInfo final { 31 VkBuffer handle; 32 VkDeviceSize offset; 33 VkDeviceSize range; 34 35 BufferBindInfo(); 36 BufferBindInfo(const VulkanBuffer& buffer_p); 37 }; 38 39 struct ParamsBindList final { 40 std::vector<BufferBindInfo> bind_infos; 41 42 ParamsBindList() = default; 43 ParamsBindList(std::initializer_list<const BufferBindInfo> init_list); 44 45 void append(const BufferBindInfo& bind_info); 46 void append(const ParamsBindList& other); 47 }; 48 49 class DescriptorSet final { 50 public: 51 explicit DescriptorSet(VkDevice, VkDescriptorSet, ShaderLayout::Signature); 52 53 DescriptorSet(const DescriptorSet&) = delete; 54 DescriptorSet& operator=(const DescriptorSet&) = delete; 55 56 DescriptorSet(DescriptorSet&&) noexcept; 57 DescriptorSet& operator=(DescriptorSet&&) noexcept; 58 59 ~DescriptorSet() = default; 60 61 struct ResourceBinding final { 62 uint32_t binding_idx; 63 VkDescriptorType descriptor_type; 64 bool is_image; 65 66 union { 67 VkDescriptorBufferInfo buffer_info; 68 VkDescriptorImageInfo image_info; 69 } resource_info; 70 }; 71 72 private: 73 VkDevice device_; 74 VkDescriptorSet handle_; 75 ShaderLayout::Signature shader_layout_signature_; 76 std::vector<ResourceBinding> bindings_; 77 78 public: 79 DescriptorSet& bind(const uint32_t, const BufferBindInfo&); 80 DescriptorSet& bind(const uint32_t, const VulkanBuffer&); 81 DescriptorSet& bind(const uint32_t, const VulkanImage&); 82 83 VkDescriptorSet get_bind_handle() const; 84 85 private: 86 void add_binding(const ResourceBinding& resource); 87 }; 88 89 class DescriptorSetPile final { 90 public: 91 DescriptorSetPile( 92 const uint32_t, 93 VkDescriptorSetLayout, 94 VkDevice, 95 VkDescriptorPool); 96 97 DescriptorSetPile(const DescriptorSetPile&) = delete; 98 DescriptorSetPile& operator=(const DescriptorSetPile&) = delete; 99 100 DescriptorSetPile(DescriptorSetPile&&) = default; 101 DescriptorSetPile& operator=(DescriptorSetPile&&) = default; 102 103 ~DescriptorSetPile() = default; 104 105 private: 106 uint32_t pile_size_; 107 VkDescriptorSetLayout set_layout_; 108 VkDevice device_; 109 VkDescriptorPool pool_; 110 std::vector<VkDescriptorSet> descriptors_; 111 size_t in_use_; 112 113 public: 114 VkDescriptorSet get_descriptor_set(); 115 116 private: 117 void allocate_new_batch(); 118 }; 119 120 struct DescriptorPoolConfig final { 121 // Overall Pool capacity 122 uint32_t descriptor_pool_max_sets; 123 // DescriptorCounts by type 124 uint32_t descriptor_uniform_buffer_count; 125 uint32_t descriptor_storage_buffer_count; 126 uint32_t descriptor_combined_sampler_count; 127 uint32_t descriptor_storage_image_count; 128 // Pile size for pre-allocating descriptor sets 129 uint32_t descriptor_pile_sizes; 130 }; 131 132 class DescriptorPool final { 133 public: 134 explicit DescriptorPool(VkDevice, const DescriptorPoolConfig&); 135 136 DescriptorPool(const DescriptorPool&) = delete; 137 DescriptorPool& operator=(const DescriptorPool&) = delete; 138 139 DescriptorPool(DescriptorPool&&) = delete; 140 DescriptorPool& operator=(DescriptorPool&&) = delete; 141 142 ~DescriptorPool(); 143 144 private: 145 VkDevice device_; 146 VkDescriptorPool pool_; 147 DescriptorPoolConfig config_; 148 // New Descriptors 149 std::mutex mutex_; 150 std::unordered_map<VkDescriptorSetLayout, DescriptorSetPile> piles_; 151 152 public: 153 operator bool() const { 154 return (pool_ != VK_NULL_HANDLE); 155 } 156 157 void init(const DescriptorPoolConfig& config); 158 159 DescriptorSet get_descriptor_set( 160 VkDescriptorSetLayout handle, 161 const ShaderLayout::Signature& signature); 162 163 void flush(); 164 }; 165 166 } // namespace vkapi 167 } // namespace vkcompute 168