[ET-VK] Allocate memory for weight and activation tensors lazily

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

@@ -897,6 +897,16 @@ VkMemoryRequirements vTensor::get_memory_requirements() const {
   return {};
 }
 
+bool vTensor::memory_is_bound() const {
+  switch (storage_type()) {
+    case utils::kBuffer:
+      return storage_->buffer_.has_memory();
+    case utils::kTexture2D:
+    case utils::kTexture3D:
+      return storage_->image_.has_memory();
+  }
+}
+
 void vTensor::bind_allocation(const vkapi::Allocation& allocation) {
   switch (storage_type()) {
     case utils::kBuffer:
@@ -909,6 +919,18 @@ void vTensor::bind_allocation(const vkapi::Allocation& allocation) {
   }
 }
 
+void vTensor::acquire_allocation(vkapi::Allocation&& allocation) {
+  switch (storage_type()) {
+    case utils::kBuffer:
+      storage_->buffer_.acquire_allocation(std::move(allocation));
+      break;
+    case utils::kTexture2D:
+    case utils::kTexture3D:
+      storage_->image_.acquire_allocation(std::move(allocation));
+      break;
+  }
+}
+
 void vTensor::update_metadata() {
   numel_ = utils::multiply_integers(sizes_);
   strides_ = calculate_strides(sizes_, dim_order_);

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

@@ -560,6 +560,12 @@ class vTensor final {
    */
   VmaAllocationCreateInfo get_allocation_create_info() const;
 
+  /*
+   * Checks if the tensor's underlying buffer or image resource is bound to a
+   * memory allocation.
+   */
+  bool memory_is_bound() const;
+
   /*
    * Return the VkMemoryRequirements of the underlying resource
    */
@@ -570,6 +576,11 @@ class vTensor final {
    */
   void bind_allocation(const vkapi::Allocation& allocation);
 
+  /*
+   * Binds and acquires a rvalue memory allocation
+   */
+  void acquire_allocation(vkapi::Allocation&& allocation);
+
  private:
   /*
    * Assuming sizes, dim order, or axis mapping was modified, recompute all

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -356,8 +356,6 @@ ValueRef ComputeGraph::add_tensor(
     const utils::GPUMemoryLayout memory_layout,
     const int64_t shared_object_idx,
     const utils::AxisMapLayout axis_map_layout) {
-  bool allocate_memory = shared_object_idx < 0;
-
   ValueRef idx(static_cast<int>(values_.size()));
   check_no_active_value_ptrs();
   values_.emplace_back(api::vTensor(
@@ -366,10 +364,10 @@ ValueRef ComputeGraph::add_tensor(
       dtype,
       storage_type,
       memory_layout,
-      allocate_memory,
+      false,
       axis_map_layout));
 
-  if (!allocate_memory) {
+  if (shared_object_idx >= 0) {
     get_shared_object(shared_object_idx).add_user(this, idx);
   }
   return idx;
@@ -626,6 +624,17 @@ SharedObject& ComputeGraph::get_shared_object(const int64_t idx) {
   return shared_objects_.at(idx);
 }
 
+void ComputeGraph::create_dedicated_allocation_for(const ValueRef idx) {
+  vTensorPtr tensor = get_tensor(idx);
+  if (!tensor->memory_is_bound()) {
+    VmaAllocationCreateInfo alloc_create_info =
+        context()->adapter_ptr()->vma().gpuonly_resource_create_info();
+    tensor->acquire_allocation(
+        context()->adapter_ptr()->vma().create_allocation(
+            tensor->get_memory_requirements(), alloc_create_info));
+  }
+}
+
 void ComputeGraph::update_descriptor_counts(
     const vkapi::ShaderInfo& shader_info,
     bool execute) {
@@ -852,11 +861,6 @@ void ComputeGraph::prepare() {
   }
 
   execute_threshold_node_count_ = count_threshold;
-
-  for (SharedObject& shared_object : shared_objects_) {
-    shared_object.allocate(this);
-    shared_object.bind_users(this);
-  }
 }
 
 void ComputeGraph::prepare_pipelines() {
@@ -952,6 +956,18 @@ void ComputeGraph::prepack() {
   submit_current_cmd_and_wait(/*final_use=*/true);
   context_->flush();
   staging_nbytes_in_cmd_ = 0;
+
+  // Initialize allocations for intermediate tensors
+  for (SharedObject& shared_object : shared_objects_) {
+    shared_object.allocate(this);
+    shared_object.bind_users(this);
+  }
+  // Make sure all remaining tensors have allocations
+  for (int i = 0; i < values_.size(); i++) {
+    if (values_.at(i).isTensor()) {
+      create_dedicated_allocation_for(i);
+    }
+  }
 }
 
 void ComputeGraph::execute() {

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -827,6 +827,13 @@ class ComputeGraph final {
 
   SharedObject& get_shared_object(const int64_t idx);
 
+  /*
+   * Creates a dedicated memory allocation for a vTensor value, and have the
+   * tensor acquire the allocation object. If the tensor is already bound to a
+   * memory allocation, this function will be a no-op.
+   */
+  void create_dedicated_allocation_for(const ValueRef idx);
+
   //
   // Graph Preparation
   //

backends/vulkan/runtime/graph/ops/PrepackNode.cpp

@@ -97,6 +97,10 @@ void PrepackNode::encode(ComputeGraph* graph) {
   }
 
   {
+    // If the vTensor is not yet bound to a memory allocation, create a new one
+    // and aquire it.
+    graph->create_dedicated_allocation_for(packed_);
+
     vkapi::PipelineBarrier pipeline_barrier{};
     vkapi::DescriptorSet descriptor_set = context->get_descriptor_set(
         shader_, local_workgroup_size_, spec_vars_, push_constants_offset);

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

@@ -20,6 +20,7 @@ VulkanBuffer::VulkanBuffer()
       allocator_(VK_NULL_HANDLE),
       memory_{},
       owns_memory_(false),
+      memory_bundled_(false),
       is_copy_(false),
       handle_(VK_NULL_HANDLE) {}
 
@@ -33,6 +34,7 @@ VulkanBuffer::VulkanBuffer(
       allocator_(vma_allocator),
       memory_{},
       owns_memory_(allocate_memory),
+      memory_bundled_(allocate_memory),
       is_copy_(false),
       handle_(VK_NULL_HANDLE) {
   // If the buffer size is 0, allocate a buffer with a size of 1 byte. This is
@@ -77,6 +79,7 @@ VulkanBuffer::VulkanBuffer(
       allocator_(other.allocator_),
       memory_(other.memory_),
       owns_memory_(false),
+      memory_bundled_(false),
       is_copy_(true),
       handle_(other.handle_) {
   // TODO: set the offset and range appropriately
@@ -91,6 +94,7 @@ VulkanBuffer::VulkanBuffer(VulkanBuffer&& other) noexcept
       allocator_(other.allocator_),
       memory_(std::move(other.memory_)),
       owns_memory_(other.owns_memory_),
+      memory_bundled_(other.memory_bundled_),
       is_copy_(other.is_copy_),
       handle_(other.handle_) {
   other.handle_ = VK_NULL_HANDLE;
@@ -99,16 +103,19 @@ VulkanBuffer::VulkanBuffer(VulkanBuffer&& other) noexcept
 VulkanBuffer& VulkanBuffer::operator=(VulkanBuffer&& other) noexcept {
   VkBuffer tmp_buffer = handle_;
   bool tmp_owns_memory = owns_memory_;
+  bool tmp_memory_bundled = memory_bundled_;
 
   buffer_properties_ = other.buffer_properties_;
   allocator_ = other.allocator_;
   memory_ = std::move(other.memory_);
   owns_memory_ = other.owns_memory_;
+  memory_bundled_ = other.memory_bundled_;
   is_copy_ = other.is_copy_;
   handle_ = other.handle_;
 
   other.handle_ = tmp_buffer;
   other.owns_memory_ = tmp_owns_memory;
+  other.memory_bundled_ = tmp_memory_bundled;
 
   return *this;
 }
@@ -119,14 +126,22 @@ VulkanBuffer::~VulkanBuffer() {
   // ownership of the underlying resource.
   if (handle_ != VK_NULL_HANDLE && !is_copy_) {
     if (owns_memory_) {
-      vmaDestroyBuffer(allocator_, handle_, memory_.allocation);
+      if (memory_bundled_) {
+        vmaDestroyBuffer(allocator_, handle_, memory_.allocation);
+        // Prevent the underlying memory allocation from being freed; it was
+        // freed by vmaDestroyImage
+        memory_.allocation = VK_NULL_HANDLE;
+      } else {
+        vkDestroyBuffer(this->device(), handle_, nullptr);
+        // Allow underlying memory allocation to be freed by the destructor of
+        // Allocation class
+      }
     } else {
       vkDestroyBuffer(this->device(), handle_, nullptr);
+      // Prevent the underlying memory allocation from being freed since this
+      // object doesn't own it
+      memory_.allocation = VK_NULL_HANDLE;
     }
-    // Prevent the underlying memory allocation from being freed; it was either
-    // freed by vmaDestroyBuffer, or this resource does not own the underlying
-    // memory
-    memory_.allocation = VK_NULL_HANDLE;
   }
 }
 
@@ -136,6 +151,24 @@ VmaAllocationInfo VulkanBuffer::allocation_info() const {
   return info;
 }
 
+void VulkanBuffer::bind_allocation_impl(const Allocation& memory) {
+  VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!");
+  if (!is_copy_) {
+    VK_CHECK(vmaBindBufferMemory(allocator_, memory.allocation, handle_));
+  }
+}
+
+void VulkanBuffer::bind_allocation(const Allocation& memory) {
+  bind_allocation_impl(memory);
+  memory_.allocation = memory.allocation;
+}
+
+void VulkanBuffer::acquire_allocation(Allocation&& memory) {
+  bind_allocation_impl(memory);
+  memory_ = std::move(memory);
+  owns_memory_ = true;
+}
+
 VkMemoryRequirements VulkanBuffer::get_memory_requirements() const {
   VkMemoryRequirements memory_requirements;
   vkGetBufferMemoryRequirements(this->device(), handle_, &memory_requirements);

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

@@ -100,6 +100,10 @@ class VulkanBuffer final {
   Allocation memory_;
   // Indicates whether the underlying memory is owned by this resource
   bool owns_memory_;
+  // Indicates whether the allocation for the buffer was created with the buffer
+  // via vmaCreateBuffer; if this is false, the memory is owned but was bound
+  // separately via vmaBindBufferMemory
+  bool memory_bundled_;
   // Indicates whether this VulkanBuffer was copied from another VulkanBuffer,
   // thus it does not have ownership of the underlying VKBuffer
   bool is_copy_;
@@ -162,13 +166,21 @@ class VulkanBuffer final {
     return (handle_ == other.handle_) && is_copy_;
   }
 
-  inline void bind_allocation(const Allocation& memory) {
-    VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!");
-    if (!is_copy_) {
-      VK_CHECK(vmaBindBufferMemory(allocator_, memory.allocation, handle_));
-    }
-    memory_.allocation = memory.allocation;
-  }
+ private:
+  void bind_allocation_impl(const Allocation& memory);
+
+ public:
+  /*
+   * Given a memory allocation, bind it to the underlying VkImage. The lifetime
+   * of the memory allocation is assumed to be managed externally.
+   */
+  void bind_allocation(const Allocation& memory);
+
+  /*
+   * Given a rvalue memory allocation, bind it to the underlying VkImage and
+   * also acquire ownership of the memory allocation.
+   */
+  void acquire_allocation(Allocation&& memory);
 
   VkMemoryRequirements get_memory_requirements() const;