From fe8238421fd4af7e650e37fac8c978d88c1a9ca1 Mon Sep 17 00:00:00 2001
From: Julian Orth <ju.orth@gmail.com>
Date: Fri, 6 Sep 2024 19:47:34 +0200
Subject: [PATCH] vulkan: only use coherency functions for non-coherent memory

---
 src/gfx_apis/vulkan/allocator.rs |  7 ++++++-
 src/gfx_apis/vulkan/staging.rs   | 24 ++++++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/gfx_apis/vulkan/allocator.rs b/src/gfx_apis/vulkan/allocator.rs
index 83f661ba..4f2568e4 100644
--- a/src/gfx_apis/vulkan/allocator.rs
+++ b/src/gfx_apis/vulkan/allocator.rs
@@ -4,7 +4,7 @@ use {
         utils::{numcell::NumCell, ptr_ext::MutPtrExt},
     },
     ash::vk::{DeviceMemory, DeviceSize, MemoryRequirements},
-    gpu_alloc::{Config, GpuAllocator, MemoryBlock, Request, UsageFlags},
+    gpu_alloc::{Config, GpuAllocator, MemoryBlock, MemoryPropertyFlags, Request, UsageFlags},
     gpu_alloc_ash::AshMemoryDevice,
     std::{
         cell::{Cell, UnsafeCell},
@@ -25,6 +25,7 @@ pub struct VulkanAllocation {
     pub(super) offset: DeviceSize,
     pub(super) mem: Option<*mut u8>,
     pub(super) size: DeviceSize,
+    pub(super) coherency_mask: Option<u64>,
     block: Cell<Option<MemoryBlock<DeviceMemory>>>,
 }
 
@@ -111,6 +112,10 @@ impl VulkanAllocator {
             offset: block.offset(),
             mem: ptr,
             size: block.size(),
+            coherency_mask: match block.props().contains(MemoryPropertyFlags::HOST_COHERENT) {
+                true => None,
+                false => Some(self.non_coherent_atom_mask),
+            },
             block: Cell::new(Some(block)),
         })
     }
diff --git a/src/gfx_apis/vulkan/staging.rs b/src/gfx_apis/vulkan/staging.rs
index b8ccc2ef..d09d705a 100644
--- a/src/gfx_apis/vulkan/staging.rs
+++ b/src/gfx_apis/vulkan/staging.rs
@@ -72,27 +72,31 @@ impl VulkanStagingBuffer {
         F: FnOnce(*mut u8, usize) -> T,
     {
         let t = f(self.allocation.mem.unwrap(), self.size as usize);
-        let range = self.range();
-        let res = unsafe { self.device.device.flush_mapped_memory_ranges(&[range]) };
-        res.map_err(VulkanError::FlushMemory).map(|_| t)
+        if let Some(mask) = self.allocation.coherency_mask {
+            let range = self.incoherent_range(mask);
+            let res = unsafe { self.device.device.flush_mapped_memory_ranges(&[range]) };
+            res.map_err(VulkanError::FlushMemory)?;
+        }
+        Ok(t)
     }
 
     pub fn download<T, F>(&self, f: F) -> Result<T, VulkanError>
     where
         F: FnOnce(*const u8, usize) -> T,
     {
-        let range = self.range();
-        let res = unsafe { self.device.device.invalidate_mapped_memory_ranges(&[range]) };
-        res.map_err(VulkanError::FlushMemory)?;
+        if let Some(mask) = self.allocation.coherency_mask {
+            let range = self.incoherent_range(mask);
+            let res = unsafe { self.device.device.invalidate_mapped_memory_ranges(&[range]) };
+            res.map_err(VulkanError::FlushMemory)?;
+        }
         Ok(f(self.allocation.mem.unwrap(), self.size as usize))
     }
 
-    fn range(&self) -> MappedMemoryRange {
-        let atom_mask = self.allocation.allocator.non_coherent_atom_mask;
+    fn incoherent_range(&self, mask: u64) -> MappedMemoryRange {
         MappedMemoryRange::default()
             .memory(self.allocation.memory)
-            .offset(self.allocation.offset & !atom_mask)
-            .size((self.allocation.size + atom_mask) & !atom_mask)
+            .offset(self.allocation.offset & !mask)
+            .size((self.allocation.size + mask) & !mask)
     }
 }