From 05de5917bdde156baae08ffd147a6d6e6d355eba Mon Sep 17 00:00:00 2001
From: Stephen Jia <ssjia@meta.com>
Date: Fri, 21 Aug 2026 12:51:25 -0700
Subject: [PATCH] Speed up device-to-host staging reads (#22023)

Summary:

Prefer host-cached memory for device-to-host staging while preserving coherent memory preference for host-to-device staging. Invalidate only bytes consumed by readback, avoiding cache maintenance over padded staging capacity. This fixes slow Mali output copies while retaining Adreno memory selection.

Authored with Codex.

Differential Revision: D116964600
---
 .../runtime/api/containers/StagingBuffer.cpp  |  2 +-
 .../runtime/api/containers/StagingBuffer.h    |  4 +--
 .../runtime/vk_api/memory/Allocator.cpp       |  4 +--
 .../vulkan/test/vulkan_compute_api_test.cpp   | 33 +++++++++++++++++++
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/backends/vulkan/runtime/api/containers/StagingBuffer.cpp b/backends/vulkan/runtime/api/containers/StagingBuffer.cpp
index 53ec9c17eae..c1cd89e8315 100644
--- a/backends/vulkan/runtime/api/containers/StagingBuffer.cpp
+++ b/backends/vulkan/runtime/api/containers/StagingBuffer.cpp
@@ -178,7 +178,7 @@ void StagingBuffer::cast_float_to_half_and_copy_to(
       vulkan_buffer_.vma_allocator(),
       vulkan_buffer_.allocation(),
       0u,
-      VK_WHOLE_SIZE);
+      numel * sizeof(float));
   const float* src = reinterpret_cast<const float*>(data());
   for (size_t i = 0; i < numel; ++i) {
     dst[i] = float_to_half(src[i]);
diff --git a/backends/vulkan/runtime/api/containers/StagingBuffer.h b/backends/vulkan/runtime/api/containers/StagingBuffer.h
index 19060804693..f924fc07482 100644
--- a/backends/vulkan/runtime/api/containers/StagingBuffer.h
+++ b/backends/vulkan/runtime/api/containers/StagingBuffer.h
@@ -107,7 +107,7 @@ class StagingBuffer final {
         vulkan_buffer_.vma_allocator(),
         vulkan_buffer_.allocation(),
         0u,
-        VK_WHOLE_SIZE);
+        nbytes);
     memcpy(dst, data(), nbytes);
   }
 
@@ -118,7 +118,7 @@ class StagingBuffer final {
         vulkan_buffer_.vma_allocator(),
         vulkan_buffer_.allocation(),
         0u,
-        VK_WHOLE_SIZE);
+        numel * sizeof(SRC_T));
     const SRC_T* src = reinterpret_cast<const SRC_T*>(data());
     for (size_t i = 0; i < numel; ++i) {
       dst[i] = static_cast<DST_T>(src[i]);
diff --git a/backends/vulkan/runtime/vk_api/memory/Allocator.cpp b/backends/vulkan/runtime/vk_api/memory/Allocator.cpp
index f36b2b0c09e..b228e6de8c7 100644
--- a/backends/vulkan/runtime/vk_api/memory/Allocator.cpp
+++ b/backends/vulkan/runtime/vk_api/memory/Allocator.cpp
@@ -177,13 +177,13 @@ VulkanBuffer Allocator::create_staging_buffer(
   if (direction == CopyDirection::HOST_TO_DEVICE) {
     alloc_create_info.flags |=
         VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT;
+    alloc_create_info.preferredFlags = VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
   } else {
     alloc_create_info.flags |= allocation_strategy_device_to_host_;
+    alloc_create_info.preferredFlags = VK_MEMORY_PROPERTY_HOST_CACHED_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 =
-      VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_CACHED_BIT;
 
   return VulkanBuffer(allocator_, size, alloc_create_info, buffer_usage);
 }
diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp
index 590ecf186e2..d1e8e363afa 100644
--- a/backends/vulkan/test/vulkan_compute_api_test.cpp
+++ b/backends/vulkan/test/vulkan_compute_api_test.cpp
@@ -108,6 +108,39 @@ TEST_F(VulkanComputeAPITest, print_adapter) {
   std::cout << *(context()->adapter_ptr()) << std::endl;
 }
 
+TEST_F(VulkanComputeAPITest, device_to_host_staging_prefers_cached_memory) {
+  vkapi::VulkanBuffer staging_buffer =
+      context()->adapter_ptr()->vma().create_staging_buffer(
+          4096, vkapi::CopyDirection::DEVICE_TO_HOST);
+  const VmaAllocator allocator = staging_buffer.vma_allocator();
+  ASSERT_NE(allocator, VK_NULL_HANDLE);
+
+  const VkPhysicalDeviceMemoryProperties* memory_properties = nullptr;
+  vmaGetMemoryProperties(allocator, &memory_properties);
+  ASSERT_NE(memory_properties, nullptr);
+
+  bool has_host_cached_memory = false;
+  for (uint32_t i = 0; i < memory_properties->memoryTypeCount; ++i) {
+    const VkMemoryPropertyFlags flags =
+        memory_properties->memoryTypes[i].propertyFlags;
+    has_host_cached_memory = has_host_cached_memory ||
+        ((flags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) &&
+         (flags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT));
+  }
+  if (!has_host_cached_memory) {
+    GTEST_SKIP() << "Device does not expose host-cached visible memory";
+  }
+
+  const VmaAllocation allocation = staging_buffer.allocation();
+  ASSERT_NE(allocation, VK_NULL_HANDLE);
+
+  VkMemoryPropertyFlags selected_flags = 0;
+  vmaGetAllocationMemoryProperties(allocator, allocation, &selected_flags);
+
+  EXPECT_TRUE(selected_flags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT)
+      << "Selected memory with flags " << selected_flags;
+}
+
 #if defined(VK_KHR_pipeline_executable_properties) && \
     defined(ETVK_INSPECT_PIPELINES)