Rework of wgpu_hal vulkan buffer importing (#7824)

Co-authored-by: Connor Fitzgerald <connorwadefitzgerald@gmail.com>

wgpu-hal/src/vulkan/device.rs

@@ -808,17 +808,6 @@ impl super::Device {
         })
     }
 
-    /// # Safety
-    ///
-    /// - `vk_buffer`'s memory must be managed by the caller
-    /// - Externally imported buffers can't be mapped by `wgpu`
-    pub unsafe fn buffer_from_raw(vk_buffer: vk::Buffer) -> super::Buffer {
-        super::Buffer {
-            raw: vk_buffer,
-            block: None,
-        }
-    }
-
     fn create_shader_module_impl(
         &self,
         spv: &[u32],
@@ -1153,7 +1142,7 @@ impl crate::Device for super::Device {
 
         Ok(super::Buffer {
             raw,
-            block: Some(Mutex::new(block)),
+            block: Some(Mutex::new(super::BufferMemoryBacking::Managed(block))),
         })
     }
     unsafe fn destroy_buffer(&self, buffer: super::Buffer) {
@@ -1161,7 +1150,14 @@ impl crate::Device for super::Device {
         if let Some(block) = buffer.block {
             let block = block.into_inner();
             self.counters.buffer_memory.sub(block.size() as isize);
-            unsafe { self.mem_allocator.lock().dealloc(&*self.shared, block) };
+            match block {
+                super::BufferMemoryBacking::Managed(block) => unsafe {
+                    self.mem_allocator.lock().dealloc(&*self.shared, block)
+                },
+                super::BufferMemoryBacking::VulkanMemory { memory, .. } => unsafe {
+                    self.shared.raw.free_memory(memory, None);
+                },
+            }
         }
 
         self.counters.buffers.sub(1);
@@ -1179,18 +1175,27 @@ impl crate::Device for super::Device {
         if let Some(ref block) = buffer.block {
             let size = range.end - range.start;
             let mut block = block.lock();
-            let ptr = unsafe { block.map(&*self.shared, range.start, size as usize)? };
-            let is_coherent = block
-                .props()
-                .contains(gpu_alloc::MemoryPropertyFlags::HOST_COHERENT);
-            Ok(crate::BufferMapping { ptr, is_coherent })
+            if let super::BufferMemoryBacking::Managed(ref mut block) = *block {
+                let ptr = unsafe { block.map(&*self.shared, range.start, size as usize)? };
+                let is_coherent = block
+                    .props()
+                    .contains(gpu_alloc::MemoryPropertyFlags::HOST_COHERENT);
+                Ok(crate::BufferMapping { ptr, is_coherent })
+            } else {
+                crate::hal_usage_error("tried to map externally created buffer")
+            }
         } else {
             crate::hal_usage_error("tried to map external buffer")
         }
     }
     unsafe fn unmap_buffer(&self, buffer: &super::Buffer) {
         if let Some(ref block) = buffer.block {
-            unsafe { block.lock().unmap(&*self.shared) };
+            match &mut *block.lock() {
+                super::BufferMemoryBacking::Managed(block) => unsafe { block.unmap(&*self.shared) },
+                super::BufferMemoryBacking::VulkanMemory { .. } => {
+                    crate::hal_usage_error("tried to unmap externally created buffer")
+                }
+            };
         } else {
             crate::hal_usage_error("tried to unmap external buffer")
         }

wgpu-hal/src/vulkan/mod.rs

@@ -776,11 +776,70 @@ impl Drop for Queue {
         unsafe { self.relay_semaphores.lock().destroy(&self.device.raw) };
     }
 }
-
+#[derive(Debug)]
+enum BufferMemoryBacking {
+    Managed(gpu_alloc::MemoryBlock<vk::DeviceMemory>),
+    VulkanMemory {
+        memory: vk::DeviceMemory,
+        offset: u64,
+        size: u64,
+    },
+}
+impl BufferMemoryBacking {
+    fn memory(&self) -> &vk::DeviceMemory {
+        match self {
+            Self::Managed(m) => m.memory(),
+            Self::VulkanMemory { memory, .. } => memory,
+        }
+    }
+    fn offset(&self) -> u64 {
+        match self {
+            Self::Managed(m) => m.offset(),
+            Self::VulkanMemory { offset, .. } => *offset,
+        }
+    }
+    fn size(&self) -> u64 {
+        match self {
+            Self::Managed(m) => m.size(),
+            Self::VulkanMemory { size, .. } => *size,
+        }
+    }
+}
 #[derive(Debug)]
 pub struct Buffer {
     raw: vk::Buffer,
-    block: Option<Mutex<gpu_alloc::MemoryBlock<vk::DeviceMemory>>>,
+    block: Option<Mutex<BufferMemoryBacking>>,
+}
+impl Buffer {
+    /// # Safety
+    ///
+    /// - `vk_buffer`'s memory must be managed by the caller
+    /// - Externally imported buffers can't be mapped by `wgpu`
+    pub unsafe fn from_raw(vk_buffer: vk::Buffer) -> Self {
+        Self {
+            raw: vk_buffer,
+            block: None,
+        }
+    }
+    /// # Safety
+    /// - We will use this buffer and the buffer's backing memory range as if we have exclusive ownership over it, until the wgpu resource is dropped and the wgpu-hal object is cleaned up
+    /// - Externally imported buffers can't be mapped by `wgpu`
+    /// - `offset` and `size` must be valid with the allocation of `memory`
+    pub unsafe fn from_raw_managed(
+        vk_buffer: vk::Buffer,
+        memory: vk::DeviceMemory,
+        offset: u64,
+        size: u64,
+    ) -> Self {
+        Self {
+            raw: vk_buffer,
+            block: Some(Mutex::new(BufferMemoryBacking::VulkanMemory {
+                memory,
+                offset,
+                size,
+            })),
+        }
+    }
 }
 
 impl crate::DynBuffer for Buffer {}