Speed up device-to-host staging reads (pytorch#22023)

Differential Revision: D116964600

Pull Request resolved: pytorch#22023

Author: SS-JIA

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]);

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]);

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);
 }

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)