Remove unnecessary member structs from Allocation struct to reduce vTensor size (#5392)

Summary:
Pull Request resolved: #5392

## Context

Recently, the size of `vTensor` has grown dramatically, due to adding multiple metadata UBO member variables to store Tensor metadata. The size of the `vTensor` object was over 2KB on 64-bit machines.

One large contributor to the size of `vTensor` was the `Allocation` class, which was storing various structs as class members. As a result of keeping these structs in memory, the size of each instance of `Allocation` was over 200 B.

These structs are not necessary to keep in memory, thus they can be removed. Once removed, the size of `vTensor` is now halved to be 1KB.

The size of `vTensor` can be further reduced by sharing metadata UBO member variables and being more conservative with keeping tensor metadata in memory.
ghstack-source-id: 242810657
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D62660980

fbshipit-source-id: ac4ae82ba7a560c162a4f3aea488647af434a332

Author: SS-JIA

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

@@ -27,6 +27,8 @@ class StagingBuffer final {
   size_t nbytes_;
   vkapi::VulkanBuffer vulkan_buffer_;
 
+  void* mapped_data_;
+
  public:
   StagingBuffer(
       Context* context_p,
@@ -37,7 +39,8 @@ class StagingBuffer final {
         numel_(numel),
         nbytes_(element_size(dtype_) * numel_),
         vulkan_buffer_(
-            context_p_->adapter_ptr()->vma().create_staging_buffer(nbytes_)) {}
+            context_p_->adapter_ptr()->vma().create_staging_buffer(nbytes_)),
+        mapped_data_(nullptr) {}
 
   StagingBuffer(const StagingBuffer&) = delete;
   StagingBuffer& operator=(const StagingBuffer&) = delete;
@@ -58,7 +61,10 @@ class StagingBuffer final {
   }
 
   inline void* data() {
-    return vulkan_buffer_.allocation_info().pMappedData;
+    if (!mapped_data_) {
+      mapped_data_ = vulkan_buffer_.allocation_info().pMappedData;
+    }
+    return mapped_data_;
   }
 
   inline size_t numel() {

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

@@ -436,17 +436,6 @@ size_t vTensor::staging_buffer_numel() const {
   return padded_numel_;
 }
 
-VmaAllocationCreateInfo vTensor::get_allocation_create_info() const {
-  switch (storage_type()) {
-    case utils::kBuffer:
-      return storage_.buffer_.allocation_create_info();
-    case utils::kTexture2D:
-    case utils::kTexture3D:
-      return storage_.image_.allocation_create_info();
-  }
-  return {};
-}
-
 VkMemoryRequirements vTensor::get_memory_requirements() const {
   switch (storage_type()) {
     case utils::kBuffer:

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

@@ -15,47 +15,27 @@ namespace vkcompute {
 void SharedObject::add_user(ComputeGraph* const graph, const ValueRef idx) {
   vTensorPtr t = graph->get_tensor(idx);
 
-  //
   // Aggregate Memory Requirements
-  //
-
   const VkMemoryRequirements mem_reqs = t->get_memory_requirements();
   aggregate_memory_requirements.size =
       std::max(mem_reqs.size, aggregate_memory_requirements.size);
   aggregate_memory_requirements.alignment =
       std::max(mem_reqs.alignment, aggregate_memory_requirements.alignment);
   aggregate_memory_requirements.memoryTypeBits |= mem_reqs.memoryTypeBits;
 
-  //
-  // Aggregate Allocation Create Info
-  //
-
-  const VmaAllocationCreateInfo create_info = t->get_allocation_create_info();
-  // Clear out CREATE_STRATEGY bit flags in case of conflict
-  VmaAllocationCreateFlags clear_mask = ~VMA_ALLOCATION_CREATE_STRATEGY_MASK;
-  VmaAllocationCreateFlags create_flags = create_info.flags & clear_mask;
-  // Use the default allocation strategy
-  aggregate_create_info.flags =
-      create_flags | vkapi::DEFAULT_ALLOCATION_STRATEGY;
-
-  // Set the usage flag if it is currently not set
-  if (aggregate_create_info.usage == VMA_MEMORY_USAGE_UNKNOWN) {
-    aggregate_create_info.usage = create_info.usage;
-  }
-  // Otherwise check that there is no conflict regarding usage
-  VK_CHECK_COND(aggregate_create_info.usage == create_info.usage);
-  aggregate_create_info.requiredFlags |= create_info.requiredFlags;
-  aggregate_create_info.preferredFlags |= create_info.preferredFlags;
-
   users.emplace_back(idx);
 }
 
 void SharedObject::allocate(ComputeGraph* const graph) {
   if (aggregate_memory_requirements.size == 0) {
     return;
   }
+
+  VmaAllocationCreateInfo alloc_create_info =
+      graph->context()->adapter_ptr()->vma().gpuonly_resource_create_info();
+
   allocation = graph->context()->adapter_ptr()->vma().create_allocation(
-      aggregate_memory_requirements, aggregate_create_info);
+      aggregate_memory_requirements, alloc_create_info);
 }
 
 void SharedObject::bind_users(ComputeGraph* const graph) {

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

@@ -28,7 +28,6 @@ struct SharedObject {
   explicit SharedObject() = default;
 
   VkMemoryRequirements aggregate_memory_requirements;
-  VmaAllocationCreateInfo aggregate_create_info;
   std::vector<ValueRef> users;
   vkapi::Allocation allocation;
 

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

@@ -26,58 +26,37 @@ namespace vkcompute {
 namespace vkapi {
 
 Allocation::Allocation()
-    : memory_requirements{},
-      create_info{},
-      allocator(VK_NULL_HANDLE),
-      allocation(VK_NULL_HANDLE),
-      allocation_info({}),
-      is_copy_(false) {}
+    : allocator(VK_NULL_HANDLE), allocation(VK_NULL_HANDLE), is_copy_(false) {}
 
 Allocation::Allocation(
     VmaAllocator vma_allocator,
     const VkMemoryRequirements& mem_props,
     const VmaAllocationCreateInfo& create_info)
-    : memory_requirements(mem_props),
-      create_info(create_info),
-      allocator(vma_allocator),
-      allocation(VK_NULL_HANDLE),
-      allocation_info({}),
-      is_copy_(false) {
+    : allocator(vma_allocator), allocation(VK_NULL_HANDLE), is_copy_(false) {
   VK_CHECK(vmaAllocateMemory(
-      allocator, &memory_requirements, &create_info, &allocation, nullptr));
+      allocator, &mem_props, &create_info, &allocation, nullptr));
 }
 
 Allocation::Allocation(const Allocation& other) noexcept
-    : memory_requirements(other.memory_requirements),
-      create_info(other.create_info),
-      allocator(other.allocator),
+    : allocator(other.allocator),
       allocation(other.allocation),
-      allocation_info(other.allocation_info),
       is_copy_(true) {}
 
 Allocation::Allocation(Allocation&& other) noexcept
-    : memory_requirements(other.memory_requirements),
-      create_info(other.create_info),
-      allocator(other.allocator),
+    : allocator(other.allocator),
       allocation(other.allocation),
-      allocation_info(other.allocation_info),
       is_copy_(other.is_copy_) {
   other.allocation = VK_NULL_HANDLE;
-  other.allocation_info = {};
 }
 
 Allocation& Allocation::operator=(Allocation&& other) noexcept {
   VmaAllocation tmp_allocation = allocation;
 
-  memory_requirements = other.memory_requirements;
-  create_info = other.create_info;
   allocator = other.allocator;
   allocation = other.allocation;
-  allocation_info = other.allocation_info;
   is_copy_ = other.is_copy_;
 
   other.allocation = tmp_allocation;
-  other.allocation_info = {};
 
   return *this;
 }

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

@@ -55,15 +55,10 @@ struct Allocation final {
 
   ~Allocation();
 
-  VkMemoryRequirements memory_requirements;
-  // The properties this allocation was created with
-  VmaAllocationCreateInfo create_info;
   // The allocator object this was allocated from
   VmaAllocator allocator;
   // Handles to the allocated memory
   VmaAllocation allocation;
-  // Information about the allocated memory
-  VmaAllocationInfo allocation_info;
 
  private:
   // Indicates whether this class instance is a copy of another class instance,

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

@@ -58,6 +58,13 @@ Allocator::~Allocator() {
   vmaDestroyAllocator(allocator_);
 }
 
+VmaAllocationCreateInfo Allocator::gpuonly_resource_create_info() {
+  VmaAllocationCreateInfo alloc_create_info = {};
+  alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
+  alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
+  return alloc_create_info;
+}
+
 Allocation Allocator::create_allocation(
     const VkMemoryRequirements& memory_requirements,
     const VmaAllocationCreateInfo& create_info) {
@@ -103,9 +110,7 @@ VulkanImage Allocator::create_image(
         (VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT);
   }
 
-  VmaAllocationCreateInfo alloc_create_info = {};
-  alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
-  alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
+  VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
 
   const VulkanImage::ImageProperties image_props{
       image_type,
@@ -157,10 +162,7 @@ VulkanBuffer Allocator::create_storage_buffer(
     const bool allocate_memory) {
   const VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
 
-  VmaAllocationCreateInfo alloc_create_info = {};
-  alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
-  alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
-
+  VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
   return VulkanBuffer(
       allocator_, size, alloc_create_info, buffer_usage, allocate_memory);
 }

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

@@ -48,6 +48,8 @@ class Allocator final {
   VmaAllocator allocator_;
 
  public:
+  VmaAllocationCreateInfo gpuonly_resource_create_info();
+
   Allocation create_allocation(
       const VkMemoryRequirements& memory_requirements,
       const VmaAllocationCreateInfo& create_info);

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

@@ -58,16 +58,14 @@ VulkanBuffer::VulkanBuffer(
       nullptr, // pQueueFamilyIndices
   };
 
-  memory_.create_info = allocation_create_info;
-
   if (allocate_memory) {
     VK_CHECK(vmaCreateBuffer(
         allocator_,
         &buffer_create_info,
         &allocation_create_info,
         &handle_,
         &(memory_.allocation),
-        &(memory_.allocation_info)));
+        nullptr));
   } else {
     VmaAllocatorInfo allocator_info{};
     vmaGetAllocatorInfo(allocator_, &allocator_info);
@@ -137,6 +135,12 @@ VulkanBuffer::~VulkanBuffer() {
   }
 }
 
+VmaAllocationInfo VulkanBuffer::allocation_info() const {
+  VmaAllocationInfo info;
+  vmaGetAllocationInfo(allocator_, memory_.allocation, &info);
+  return info;
+}
+
 VkMemoryRequirements VulkanBuffer::get_memory_requirements() const {
   VkMemoryRequirements memory_requirements;
   vkGetBufferMemoryRequirements(this->device(), handle_, &memory_requirements);

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

@@ -114,13 +114,7 @@ class VulkanBuffer final {
     return memory_.allocation;
   }
 
-  inline VmaAllocationInfo allocation_info() const {
-    return memory_.allocation_info;
-  }
-
-  inline VmaAllocationCreateInfo allocation_create_info() const {
-    return VmaAllocationCreateInfo(memory_.create_info);
-  }
+  VmaAllocationInfo allocation_info() const;
 
   inline VkBuffer handle() const {
     return handle_;

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

@@ -159,8 +159,6 @@ VulkanImage::VulkanImage(
       layout_, // initialLayout
   };
 
-  memory_.create_info = allocation_create_info;
-
   if (allocate_memory) {
     VK_CHECK(vmaCreateImage(
         allocator_,

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

@@ -169,10 +169,6 @@ class VulkanImage final {
     return memory_.allocation;
   }
 
-  inline VmaAllocationCreateInfo allocation_create_info() const {
-    return VmaAllocationCreateInfo(memory_.create_info);
-  }
-
   inline VkFormat format() const {
     return image_properties_.image_format;
   }

backends/vulkan/test/utils/test_utils.cpp

@@ -495,8 +495,10 @@ void submit_to_gpu() {
 }
 
 vkapi::Allocation allocate_memory_for(const api::vTensor& vten) {
+  VmaAllocationCreateInfo alloc_create_info =
+      api::context()->adapter_ptr()->vma().gpuonly_resource_create_info();
   return api::context()->adapter_ptr()->vma().create_allocation(
-      vten.get_memory_requirements(), vten.get_allocation_create_info());
+      vten.get_memory_requirements(), alloc_create_info);
 }
 
 VmaTotalStatistics get_vma_stats() {

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -1066,11 +1066,11 @@ TEST_F(VulkanComputeAPITest, print_object_sizes) {
   // can alert ourselves to any significant changes in the sizes of these
   // objects by checking the `sizeof()` the class against some loose thresholds.
 
-  // Current known size on 64 bit system: 2064 B
-  EXPECT_TRUE(sizeof(vTensor) < 2200);
-  // Current known size on 64 bit system: 2080 B
-  EXPECT_TRUE(sizeof(Value) < 2200);
-  // Current known size on 64 bit system: 240 B
+  // Current known size on 64 bit system: 1040 B
+  EXPECT_TRUE(sizeof(vTensor) < 1200);
+  // Current known size on 64 bit system: 1056 B
+  EXPECT_TRUE(sizeof(Value) < 1200);
+  // Current known size on 64 bit system: 120 B
   EXPECT_TRUE(sizeof(StagingBuffer) < 500);
   // Current known size on 64 bit system: 384 B
   EXPECT_TRUE(sizeof(ComputeGraph) < 500);