Add support for importing external buffers (#3355)

CHANGELOG.md

@@ -41,6 +41,7 @@ Bottom level categories:
 ## Unreleased
 
 ### Changes
+- Added support for importing external buffers using `buffer_from_raw` (Dx12, Metal, Vulkan) and `create_buffer_from_hal`. By @AdrianEddy in [#3355](https://github.com/gfx-rs/wgpu/pull/3355)
 
 #### Misc Breaking Changes
 

wgpu-core/src/device/global.rs

@@ -684,6 +684,63 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         (id, Some(error))
     }
 
+    /// # Safety
+    ///
+    /// - `hal_buffer` must be created from `device_id` corresponding raw handle.
+    /// - `hal_buffer` must be created respecting `desc`
+    /// - `hal_buffer` must be initialized
+    pub unsafe fn create_buffer_from_hal<A: HalApi>(
+        &self,
+        hal_buffer: A::Buffer,
+        device_id: DeviceId,
+        desc: &resource::BufferDescriptor,
+        id_in: Input<G, id::BufferId>,
+    ) -> (id::BufferId, Option<resource::CreateBufferError>) {
+        profiling::scope!("Device::create_buffer");
+
+        let hub = A::hub(self);
+        let mut token = Token::root();
+        let fid = hub.buffers.prepare(id_in);
+
+        let (device_guard, mut token) = hub.devices.read(&mut token);
+        let error = loop {
+            let device = match device_guard.get(device_id) {
+                Ok(device) => device,
+                Err(_) => break DeviceError::Invalid.into(),
+            };
+
+            // NB: Any change done through the raw buffer handle will not be
+            // recorded in the replay
+            #[cfg(feature = "trace")]
+            if let Some(ref trace) = device.trace {
+                trace
+                    .lock()
+                    .add(trace::Action::CreateBuffer(fid.id(), desc.clone()));
+            }
+
+            let mut buffer = device.create_buffer_from_hal(hal_buffer, device_id, desc);
+
+            // Assume external buffers are initialized
+            buffer.initialization_status = crate::init_tracker::BufferInitTracker::new(0);
+
+            let ref_count = buffer.life_guard.add_ref();
+
+            let id = fid.assign(buffer, &mut token);
+            log::info!("Created buffer {:?} with {:?}", id, desc);
+
+            device
+                .trackers
+                .lock()
+                .buffers
+                .insert_single(id, ref_count, hal::BufferUses::empty());
+
+            return (id.0, None);
+        };
+
+        let id = fid.assign_error(desc.label.borrow_or_default(), &mut token);
+        (id, Some(error))
+    }
+
     pub fn texture_label<A: HalApi>(&self, id: id::TextureId) -> String {
         A::hub(self).textures.label_for_resource(id)
     }

wgpu-core/src/device/resource.rs

@@ -501,6 +501,29 @@ impl<A: HalApi> Device<A> {
         }
     }
 
+    pub fn create_buffer_from_hal(
+        &self,
+        hal_buffer: A::Buffer,
+        self_id: id::DeviceId,
+        desc: &resource::BufferDescriptor,
+    ) -> Buffer<A> {
+        debug_assert_eq!(self_id.backend(), A::VARIANT);
+
+        Buffer {
+            raw: Some(hal_buffer),
+            device_id: Stored {
+                value: id::Valid(self_id),
+                ref_count: self.life_guard.add_ref(),
+            },
+            usage: desc.usage,
+            size: desc.size,
+            initialization_status: BufferInitTracker::new(0),
+            sync_mapped_writes: None,
+            map_state: resource::BufferMapState::Idle,
+            life_guard: LifeGuard::new(desc.label.borrow_or_default()),
+        }
+    }
+
     pub(super) fn create_texture(
         &self,
         self_id: id::DeviceId,

wgpu-hal/src/dx12/device.rs

@@ -299,6 +299,17 @@ impl super::Device {
             allocation: None,
         }
     }
+
+    pub unsafe fn buffer_from_raw(
+        resource: d3d12::Resource,
+        size: wgt::BufferAddress,
+    ) -> super::Buffer {
+        super::Buffer {
+            resource,
+            size,
+            allocation: None,
+        }
+    }
 }
 
 impl crate::Device<super::Api> for super::Device {

wgpu-hal/src/metal/device.rs

@@ -263,6 +263,10 @@ impl super::Device {
         }
     }
 
+    pub unsafe fn buffer_from_raw(raw: metal::Buffer, size: wgt::BufferAddress) -> super::Buffer {
+        super::Buffer { raw, size }
+    }
+
     pub fn raw_device(&self) -> &Mutex<metal::Device> {
         &self.shared.device
     }

wgpu-hal/src/vulkan/device.rs

@@ -280,16 +280,16 @@ impl super::DeviceShared {
         &self,
         buffer: &'a super::Buffer,
         ranges: I,
-    ) -> impl 'a + Iterator<Item = vk::MappedMemoryRange> {
-        let block = buffer.block.lock();
+    ) -> Option<impl 'a + Iterator<Item = vk::MappedMemoryRange>> {
+        let block = buffer.block.as_ref()?.lock();
         let mask = self.private_caps.non_coherent_map_mask;
-        ranges.map(move |range| {
+        Some(ranges.map(move |range| {
             vk::MappedMemoryRange::builder()
                 .memory(*block.memory())
                 .offset((block.offset() + range.start) & !mask)
                 .size((range.end - range.start + mask) & !mask)
                 .build()
-        })
+        }))
     }
 
     unsafe fn free_resources(&self) {
@@ -680,6 +680,17 @@ 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],
@@ -868,16 +879,18 @@ impl crate::Device<super::Api> for super::Device {
 
         Ok(super::Buffer {
             raw,
-            block: Mutex::new(block),
+            block: Some(Mutex::new(block)),
         })
     }
     unsafe fn destroy_buffer(&self, buffer: super::Buffer) {
         unsafe { self.shared.raw.destroy_buffer(buffer.raw, None) };
-        unsafe {
-            self.mem_allocator
-                .lock()
-                .dealloc(&*self.shared, buffer.block.into_inner())
-        };
+        if let Some(block) = buffer.block {
+            unsafe {
+                self.mem_allocator
+                    .lock()
+                    .dealloc(&*self.shared, block.into_inner())
+            };
+        }
     }
 
     unsafe fn map_buffer(
@@ -885,48 +898,56 @@ impl crate::Device<super::Api> for super::Device {
         buffer: &super::Buffer,
         range: crate::MemoryRange,
     ) -> Result<crate::BufferMapping, crate::DeviceError> {
-        let size = range.end - range.start;
-        let mut block = buffer.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 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 })
+        } else {
+            Err(crate::DeviceError::OutOfMemory)
+        }
     }
     unsafe fn unmap_buffer(&self, buffer: &super::Buffer) -> Result<(), crate::DeviceError> {
-        unsafe { buffer.block.lock().unmap(&*self.shared) };
-        Ok(())
+        if let Some(ref block) = buffer.block {
+            unsafe { block.lock().unmap(&*self.shared) };
+            Ok(())
+        } else {
+            Err(crate::DeviceError::OutOfMemory)
+        }
     }
 
     unsafe fn flush_mapped_ranges<I>(&self, buffer: &super::Buffer, ranges: I)
     where
         I: Iterator<Item = crate::MemoryRange>,
     {
-        let vk_ranges = self.shared.make_memory_ranges(buffer, ranges);
-
-        unsafe {
-            self.shared
-                .raw
-                .flush_mapped_memory_ranges(
-                    &smallvec::SmallVec::<[vk::MappedMemoryRange; 32]>::from_iter(vk_ranges),
-                )
+        if let Some(vk_ranges) = self.shared.make_memory_ranges(buffer, ranges) {
+            unsafe {
+                self.shared
+                    .raw
+                    .flush_mapped_memory_ranges(
+                        &smallvec::SmallVec::<[vk::MappedMemoryRange; 32]>::from_iter(vk_ranges),
+                    )
+            }
+            .unwrap();
         }
-        .unwrap();
     }
     unsafe fn invalidate_mapped_ranges<I>(&self, buffer: &super::Buffer, ranges: I)
     where
         I: Iterator<Item = crate::MemoryRange>,
     {
-        let vk_ranges = self.shared.make_memory_ranges(buffer, ranges);
-
-        unsafe {
-            self.shared
-                .raw
-                .invalidate_mapped_memory_ranges(
-                    &smallvec::SmallVec::<[vk::MappedMemoryRange; 32]>::from_iter(vk_ranges),
-                )
+        if let Some(vk_ranges) = self.shared.make_memory_ranges(buffer, ranges) {
+            unsafe {
+                self.shared
+                    .raw
+                    .invalidate_mapped_memory_ranges(&smallvec::SmallVec::<
+                        [vk::MappedMemoryRange; 32],
+                    >::from_iter(vk_ranges))
+            }
+            .unwrap();
         }
-        .unwrap();
     }
 
     unsafe fn create_texture(

wgpu-hal/src/vulkan/mod.rs

@@ -305,7 +305,7 @@ pub struct Queue {
 #[derive(Debug)]
 pub struct Buffer {
     raw: vk::Buffer,
-    block: Mutex<gpu_alloc::MemoryBlock<vk::DeviceMemory>>,
+    block: Option<Mutex<gpu_alloc::MemoryBlock<vk::DeviceMemory>>>,
 }
 
 #[derive(Debug)]

wgpu/src/backend/direct.rs

@@ -1,7 +1,7 @@
 use crate::{
     context::{ObjectId, Unused},
     AdapterInfo, BindGroupDescriptor, BindGroupLayoutDescriptor, BindingResource, BufferBinding,
-    CommandEncoderDescriptor, ComputePassDescriptor, ComputePipelineDescriptor,
+    BufferDescriptor, CommandEncoderDescriptor, ComputePassDescriptor, ComputePipelineDescriptor,
     DownlevelCapabilities, Features, Label, Limits, LoadOp, MapMode, Operations,
     PipelineLayoutDescriptor, RenderBundleEncoderDescriptor, RenderPipelineDescriptor,
     SamplerDescriptor, ShaderModuleDescriptor, ShaderModuleDescriptorSpirV, ShaderSource,
@@ -154,6 +154,38 @@ impl Context {
         }
     }
 
+    pub unsafe fn create_buffer_from_hal<A: wgc::hal_api::HalApi>(
+        &self,
+        hal_buffer: A::Buffer,
+        device: &Device,
+        desc: &BufferDescriptor,
+    ) -> (wgc::id::BufferId, Buffer) {
+        let global = &self.0;
+        let (id, error) = unsafe {
+            global.create_buffer_from_hal::<A>(
+                hal_buffer,
+                device.id,
+                &desc.map_label(|l| l.map(Borrowed)),
+                (),
+            )
+        };
+        if let Some(cause) = error {
+            self.handle_error(
+                &device.error_sink,
+                cause,
+                LABEL,
+                desc.label,
+                "Device::create_buffer_from_hal",
+            );
+        }
+        (
+            id,
+            Buffer {
+                error_sink: Arc::clone(&device.error_sink),
+            },
+        )
+    }
+
     pub unsafe fn device_as_hal<A: wgc::hal_api::HalApi, F: FnOnce(Option<&A::Device>) -> R, R>(
         &self,
         device: &Device,

wgpu/src/lib.rs

@@ -2507,6 +2507,50 @@ impl Device {
         }
     }
 
+    /// Creates a [`Buffer`] from a wgpu-hal Buffer.
+    ///
+    /// # Safety
+    ///
+    /// - `hal_buffer` must be created from this device internal handle
+    /// - `hal_buffer` must be created respecting `desc`
+    /// - `hal_buffer` must be initialized
+    #[cfg(any(
+        not(target_arch = "wasm32"),
+        target_os = "emscripten",
+        feature = "webgl"
+    ))]
+    pub unsafe fn create_buffer_from_hal<A: wgc::hal_api::HalApi>(
+        &self,
+        hal_buffer: A::Buffer,
+        desc: &BufferDescriptor,
+    ) -> Buffer {
+        let mut map_context = MapContext::new(desc.size);
+        if desc.mapped_at_creation {
+            map_context.initial_range = 0..desc.size;
+        }
+
+        let (id, buffer) = unsafe {
+            self.context
+                .as_any()
+                .downcast_ref::<crate::backend::Context>()
+                .unwrap()
+                .create_buffer_from_hal::<A>(
+                    hal_buffer,
+                    self.data.as_ref().downcast_ref().unwrap(),
+                    desc,
+                )
+        };
+
+        Buffer {
+            context: Arc::clone(&self.context),
+            id: ObjectId::from(id),
+            data: Box::new(buffer),
+            map_context: Mutex::new(map_context),
+            size: desc.size,
+            usage: desc.usage,
+        }
+    }
+
     /// Creates a new [`Sampler`].
     ///
     /// `desc` specifies the behavior of the sampler.