[ET-VK] Persistently map staging buffers

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

@@ -14,6 +14,8 @@
 
 #include <executorch/backends/vulkan/runtime/vk_api/memory/Buffer.h>
 
+#include <cstring>
+
 namespace vkcompute {
 namespace api {
 
@@ -55,13 +57,41 @@ class StagingBuffer final {
     return vulkan_buffer_;
   }
 
+  inline void* data() {
+    return vulkan_buffer_.allocation_info().pMappedData;
+  }
+
   inline size_t numel() {
     return numel_;
   }
 
   inline size_t nbytes() {
     return nbytes_;
   }
+
+  inline void copy_from(const void* src, const size_t nbytes) {
+    VK_CHECK_COND(nbytes <= nbytes_);
+    memcpy(data(), src, nbytes);
+    vmaFlushAllocation(
+        vulkan_buffer_.vma_allocator(),
+        vulkan_buffer_.allocation(),
+        0u,
+        VK_WHOLE_SIZE);
+  }
+
+  inline void copy_to(void* dst, const size_t nbytes) {
+    VK_CHECK_COND(nbytes <= nbytes_);
+    vmaInvalidateAllocation(
+        vulkan_buffer_.vma_allocator(),
+        vulkan_buffer_.allocation(),
+        0u,
+        VK_WHOLE_SIZE);
+    memcpy(dst, data(), nbytes);
+  }
+
+  inline void set_staging_zeros() {
+    memset(data(), 0, nbytes_);
+  }
 };
 
 } // namespace api

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -401,7 +401,7 @@ void ComputeGraph::copy_into_staging(
     const size_t numel) {
   StagingPtr staging = get_staging(idx);
   size_t nbytes = numel * vkapi::element_size(staging->dtype());
-  copy_ptr_to_staging(data, *staging, nbytes);
+  staging->copy_from(data, nbytes);
 }
 
 void ComputeGraph::copy_from_staging(
@@ -410,7 +410,7 @@ void ComputeGraph::copy_from_staging(
     const size_t numel) {
   StagingPtr staging = get_staging(idx);
   size_t nbytes = numel * vkapi::element_size(staging->dtype());
-  copy_staging_to_ptr(*staging, data, nbytes);
+  staging->copy_to(data, nbytes);
 }
 
 void ComputeGraph::prepare() {

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

@@ -53,16 +53,15 @@ api::StagingBuffer PrepackNode::create_staging_buffer(ComputeGraph* graph) {
   if (graph->val_is_none(tref_)) {
     size_t numel = utils::multiply_integers(packed->sizes());
     api::StagingBuffer staging(graph->context(), packed->dtype(), numel);
-    size_t nbytes = numel * vkapi::element_size(packed->dtype());
-    set_staging_zeros(staging, nbytes);
+    staging.set_staging_zeros();
     return staging;
   }
 
   TensorRefPtr tref = graph->get_tref(tref_);
   size_t numel = utils::multiply_integers(tref->sizes);
   api::StagingBuffer staging(graph->context(), tref->dtype, numel);
   size_t nbytes = numel * vkapi::element_size(tref->dtype);
-  copy_ptr_to_staging(tref->data, staging, nbytes);
+  staging.copy_from(tref->data, nbytes);
   return staging;
 }
 

backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp

@@ -13,88 +13,8 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h>
 
-#include <cstring>
-
 namespace vkcompute {
 
-template <typename T>
-void memcpy_to_mapping_impl(
-    const void* src,
-    vkapi::MemoryMap& dst_mapping,
-    const size_t nbytes) {
-  T* data_ptr = dst_mapping.template data<T>();
-  memcpy(data_ptr, reinterpret_cast<const T*>(src), nbytes);
-}
-
-template <typename T>
-void memcpy_from_mapping_impl(
-    vkapi::MemoryMap& src_mapping,
-    void* dst,
-    const size_t nbytes) {
-  T* data_ptr = src_mapping.template data<T>();
-  memcpy(reinterpret_cast<T*>(dst), data_ptr, nbytes);
-}
-
-void memcpy_to_mapping(
-    const void* src,
-    vkapi::MemoryMap& dst_mapping,
-    const size_t nbytes,
-    const vkapi::ScalarType dtype) {
-#define DTYPE_CASE(ctype, vkformat, name)                    \
-  case vkapi::ScalarType::name:                              \
-    memcpy_to_mapping_impl<ctype>(src, dst_mapping, nbytes); \
-    break;
-
-  switch (dtype) {
-    VK_FORALL_SCALAR_TYPES(DTYPE_CASE)
-    default:
-      VK_THROW("Unrecognized dtype!");
-  }
-#undef DTYPE_CASE
-}
-
-void memcpy_from_mapping(
-    vkapi::MemoryMap& src_mapping,
-    void* dst,
-    const size_t nbytes,
-    const vkapi::ScalarType dtype) {
-#define DTYPE_CASE(ctype, vkformat, name)                      \
-  case vkapi::ScalarType::name:                                \
-    memcpy_from_mapping_impl<ctype>(src_mapping, dst, nbytes); \
-    break;
-
-  switch (dtype) {
-    VK_FORALL_SCALAR_TYPES(DTYPE_CASE)
-    default:
-      VK_THROW("Unrecognized dtype!");
-  }
-#undef DTYPE_CASE
-}
-
-void copy_ptr_to_staging(
-    const void* src,
-    api::StagingBuffer& staging,
-    const size_t nbytes) {
-  vkapi::MemoryMap mapping(staging.buffer(), vkapi::MemoryAccessType::WRITE);
-  mapping.invalidate();
-  memcpy_to_mapping(src, mapping, nbytes, staging.dtype());
-}
-
-void copy_staging_to_ptr(
-    api::StagingBuffer& staging,
-    void* dst,
-    const size_t nbytes) {
-  vkapi::MemoryMap mapping(staging.buffer(), vkapi::MemoryAccessType::READ);
-  mapping.invalidate();
-  memcpy_from_mapping(mapping, dst, nbytes, staging.dtype());
-}
-
-void set_staging_zeros(api::StagingBuffer& staging, const size_t nbytes) {
-  vkapi::MemoryMap mapping(staging.buffer(), vkapi::MemoryAccessType::WRITE);
-  uint8_t* data_ptr = mapping.template data<uint8_t>();
-  memset(data_ptr, 0, staging.nbytes());
-}
-
 vkapi::ShaderInfo get_nchw_to_tensor_shader(
     const api::vTensor& v_dst,
     const bool int8_buffer_enabled) {

backends/vulkan/runtime/graph/ops/utils/StagingUtils.h

@@ -12,25 +12,6 @@
 
 namespace vkcompute {
 
-//
-// Functions to copy data into and out of a staging buffer
-//
-
-void copy_ptr_to_staging(
-    const void* src,
-    api::StagingBuffer& staging,
-    const size_t nbytes);
-void copy_staging_to_ptr(
-    api::StagingBuffer& staging,
-    void* dst,
-    const size_t nbytes);
-
-void set_staging_zeros(api::StagingBuffer& staging, const size_t nbytes);
-
-//
-// Functions to get shaders
-//
-
 vkapi::ShaderInfo get_nchw_to_tensor_shader(
     const api::vTensor& v_dst,
     bool int8_buffer_enabled = true);

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

@@ -30,6 +30,7 @@ Allocation::Allocation()
       create_info{},
       allocator(VK_NULL_HANDLE),
       allocation(VK_NULL_HANDLE),
+      allocation_info({}),
       is_copy_(false) {}
 
 Allocation::Allocation(
@@ -40,6 +41,7 @@ Allocation::Allocation(
       create_info(create_info),
       allocator(vma_allocator),
       allocation(VK_NULL_HANDLE),
+      allocation_info({}),
       is_copy_(false) {
   VK_CHECK(vmaAllocateMemory(
       allocator, &memory_requirements, &create_info, &allocation, nullptr));
@@ -50,15 +52,18 @@ Allocation::Allocation(const Allocation& other) noexcept
       create_info(other.create_info),
       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),
       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 {
@@ -68,9 +73,11 @@ Allocation& Allocation::operator=(Allocation&& other) noexcept {
   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

@@ -62,6 +62,8 @@ struct Allocation final {
   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

@@ -142,7 +142,8 @@ VulkanBuffer Allocator::create_staging_buffer(const VkDeviceSize size) {
   // Staging buffers are accessed by both the CPU and GPU, so set the
   // appropriate flags to indicate that the host device will be accessing
   // the data from this buffer.
-  alloc_create_info.flags |= VMA_ALLOCATION_CREATE_HOST_ACCESS_RANDOM_BIT;
+  alloc_create_info.flags |= VMA_ALLOCATION_CREATE_HOST_ACCESS_RANDOM_BIT |
+      VMA_ALLOCATION_CREATE_MAPPED_BIT;
   alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_HOST;
   alloc_create_info.requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
   alloc_create_info.preferredFlags =

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

@@ -67,7 +67,7 @@ VulkanBuffer::VulkanBuffer(
         &allocation_create_info,
         &handle_,
         &(memory_.allocation),
-        nullptr));
+        &(memory_.allocation_info)));
   } else {
     VmaAllocatorInfo allocator_info{};
     vmaGetAllocatorInfo(allocator_, &allocator_info);

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

@@ -114,6 +114,10 @@ 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);
   }

backends/vulkan/test/utils/test_utils.cpp

@@ -326,15 +326,15 @@ void record_reference_matmul(
 void fill_vtensor(api::vTensor& vten, std::vector<float>& data) {
   api::StagingBuffer staging_buffer(api::context(), vten.dtype(), data.size());
 
-#define CASE(ctype, name)                                                     \
-  case vkapi::ScalarType::name: {                                             \
-    std::vector<ctype> data_converted;                                        \
-    data_converted.resize(data.size());                                       \
-    for (int i = 0; i < data.size(); ++i) {                                   \
-      data_converted[i] = ctype(data[i]);                                     \
-    }                                                                         \
-    copy_ptr_to_staging(                                                      \
-        data_converted.data(), staging_buffer, vten.staging_buffer_nbytes()); \
+#define CASE(ctype, name)                                     \
+  case vkapi::ScalarType::name: {                             \
+    std::vector<ctype> data_converted;                        \
+    data_converted.resize(data.size());                       \
+    for (int i = 0; i < data.size(); ++i) {                   \
+      data_converted[i] = ctype(data[i]);                     \
+    }                                                         \
+    staging_buffer.copy_from(                                 \
+        data_converted.data(), vten.staging_buffer_nbytes()); \
   } break;
 
   switch (vten.dtype()) {
@@ -424,14 +424,14 @@ void extract_vtensor(api::vTensor& vten, std::vector<float>& data) {
   api::context()->submit_cmd_to_gpu(fence.get_submit_handle());
   fence.wait();
 
-#define CASE(ctype, name)                                                     \
-  case vkapi::ScalarType::name: {                                             \
-    std::vector<ctype> data_converted(data.size());                           \
-    copy_staging_to_ptr(                                                      \
-        staging_buffer, data_converted.data(), vten.staging_buffer_nbytes()); \
-    for (int i = 0; i < data.size(); ++i) {                                   \
-      data[i] = float(data_converted[i]);                                     \
-    }                                                                         \
+#define CASE(ctype, name)                                     \
+  case vkapi::ScalarType::name: {                             \
+    std::vector<ctype> data_converted(data.size());           \
+    staging_buffer.copy_to(                                   \
+        data_converted.data(), vten.staging_buffer_nbytes()); \
+    for (int i = 0; i < data.size(); ++i) {                   \
+      data[i] = float(data_converted[i]);                     \
+    }                                                         \
   } break;
 
   switch (vten.dtype()) {

backends/vulkan/test/utils/test_utils.h

@@ -132,7 +132,7 @@ fill_staging(api::StagingBuffer& staging, float val, int numel = -1) {
   }
   std::vector<float> data(numel);
   std::fill(data.begin(), data.end(), val);
-  copy_ptr_to_staging(data.data(), staging, sizeof(float) * numel);
+  staging.copy_from(data.data(), sizeof(float) * numel);
 }
 
 void fill_vtensor(api::vTensor& vten, std::vector<float>& data);
@@ -169,7 +169,7 @@ check_staging_buffer(api::StagingBuffer& staging, float val, int numel = -1) {
     numel = staging.numel();
   }
   std::vector<float> data(numel);
-  copy_staging_to_ptr(staging, data.data(), sizeof(float) * numel);
+  staging.copy_to(data.data(), sizeof(float) * numel);
 
   for (size_t i = 0; i < data.size(); ++i) {
     CHECK_VALUE(data, i, val);