From 28ad3f27cc543e46d41615b483c180fc059c277e Mon Sep 17 00:00:00 2001
From: pytorchbot <soumith+bot@pytorch.org>
Date: Mon, 9 Dec 2024 22:15:49 -0800
Subject: [PATCH] [ET-VK] Removing unnecessary and redundant members from
 StagingBuffer.

Pull Request resolved: https://github.com/pytorch/executorch/pull/7220

This diff removes unnecessary and redundant members from the StagingBuffer class in the Vulkan runtime API. Specifically, the `numel_` and `nbytes_` members are removed, as they can be calculated from the `dtype_` member. This simplifies the class and reduces the amount of memory used.
ghstack-source-id: 257227238
@exported-using-ghexport

Differential Revision: [D66742613](https://our.internmc.facebook.com/intern/diff/D66742613/)

Co-authored-by: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com>
---
 .../runtime/api/containers/StagingBuffer.h     | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/backends/vulkan/runtime/api/containers/StagingBuffer.h b/backends/vulkan/runtime/api/containers/StagingBuffer.h
index 6f67ae8a64..1e9f569fc4 100644
--- a/backends/vulkan/runtime/api/containers/StagingBuffer.h
+++ b/backends/vulkan/runtime/api/containers/StagingBuffer.h
@@ -23,8 +23,6 @@ class StagingBuffer final {
  private:
   Context* context_p_;
   vkapi::ScalarType dtype_;
-  size_t numel_;
-  size_t nbytes_;
   vkapi::VulkanBuffer vulkan_buffer_;
 
   void* mapped_data_;
@@ -36,10 +34,8 @@ class StagingBuffer final {
       const size_t numel)
       : context_p_(context_p),
         dtype_(dtype),
-        numel_(numel),
-        nbytes_(element_size(dtype_) * numel_),
-        vulkan_buffer_(
-            context_p_->adapter_ptr()->vma().create_staging_buffer(nbytes_)),
+        vulkan_buffer_(context_p_->adapter_ptr()->vma().create_staging_buffer(
+            element_size(dtype_) * numel)),
         mapped_data_(nullptr) {}
 
   StagingBuffer(const StagingBuffer&) = delete;
@@ -68,15 +64,15 @@ class StagingBuffer final {
   }
 
   inline size_t numel() {
-    return numel_;
+    return nbytes() / element_size(dtype_);
   }
 
   inline size_t nbytes() {
-    return nbytes_;
+    return vulkan_buffer_.mem_size();
   }
 
   inline void copy_from(const void* src, const size_t nbytes) {
-    VK_CHECK_COND(nbytes <= nbytes_);
+    VK_CHECK_COND(nbytes <= this->nbytes());
     memcpy(data(), src, nbytes);
     vmaFlushAllocation(
         vulkan_buffer_.vma_allocator(),
@@ -86,7 +82,7 @@ class StagingBuffer final {
   }
 
   inline void copy_to(void* dst, const size_t nbytes) {
-    VK_CHECK_COND(nbytes <= nbytes_);
+    VK_CHECK_COND(nbytes <= this->nbytes());
     vmaInvalidateAllocation(
         vulkan_buffer_.vma_allocator(),
         vulkan_buffer_.allocation(),
@@ -96,7 +92,7 @@ class StagingBuffer final {
   }
 
   inline void set_staging_zeros() {
-    memset(data(), 0, nbytes_);
+    memset(data(), 0, nbytes());
   }
 };