WIP: export buffer handles

wgpu-core/src/conv.rs

@@ -93,6 +93,10 @@ pub fn map_buffer_usage(usage: wgt::BufferUsages) -> hal::BufferUses {
         hal::BufferUses::QUERY_RESOLVE,
         usage.contains(wgt::BufferUsages::QUERY_RESOLVE),
     );
+    u.set(
+        hal::BufferUses::EXPORT_BUFFER_HANDLE,
+        usage.contains(wgt::BufferUsages::BUFFER_HANDLE)
+    );
     u
 }
 

wgpu-hal/src/dx12/device.rs

@@ -15,7 +15,7 @@ use windows::{
         System::Threading,
     },
 };
-
+use wgt::BufferHandle;
 use super::{conv, descriptor, D3D12Lib};
 use crate::{
     auxil::{self, dxgi::result::HResult},
@@ -1847,4 +1847,8 @@ impl crate::Device for super::Device {
             total_reserved_bytes: upstream.total_reserved_bytes,
         })
     }
+
+    unsafe fn get_buffer_handle(&self, _buffer: &<Self::A as crate::Api>::Buffer) -> Result<wgt::BufferHandle, crate::DeviceError> {
+        todo!("buffer handle not yet supported");
+    }
 }

wgpu-hal/src/empty.rs

@@ -1,6 +1,7 @@
 #![allow(unused_variables)]
 
 use std::ops::Range;
+use wgt::BufferHandle;
 
 #[derive(Clone, Debug)]
 pub struct Api;
@@ -306,6 +307,10 @@ impl crate::Device for Context {
     fn get_internal_counters(&self) -> wgt::HalCounters {
         Default::default()
     }
+
+    unsafe fn get_buffer_handle(&self, _buffer: &<Self::A as crate::Api>::Buffer) -> Result<wgt::BufferHandle, crate::DeviceError> {
+        unimplemented!("buffer handle not supported");
+    }
 }
 
 impl crate::CommandEncoder for Encoder {

wgpu-hal/src/gles/device.rs

@@ -10,6 +10,7 @@ use std::{
 
 use arrayvec::ArrayVec;
 use std::sync::atomic::Ordering;
+use wgt::BufferHandle;
 
 type ShaderStage<'a> = (
     naga::ShaderStage,
@@ -1625,6 +1626,10 @@ impl crate::Device for super::Device {
     fn get_internal_counters(&self) -> wgt::HalCounters {
         self.counters.clone()
     }
+
+    unsafe fn get_buffer_handle(&self, _buffer: &<Self::A as crate::Api>::Buffer) -> Result<wgt::BufferHandle, crate::DeviceError> {
+        unimplemented!("buffer handle not supported");
+    }
 }
 
 #[cfg(send_sync)]

wgpu-hal/src/lib.rs

@@ -956,6 +956,8 @@ pub trait Device: WasmNotSendSync {
     fn generate_allocator_report(&self) -> Option<wgt::AllocatorReport> {
         None
     }
+
+    unsafe fn get_buffer_handle(&self, _buffer: &<Self::A as Api>::Buffer) -> Result<wgt::BufferHandle, crate::DeviceError>;
 }
 
 pub trait Queue: WasmNotSendSync {
@@ -1569,6 +1571,7 @@ bitflags::bitflags! {
         const ACCELERATION_STRUCTURE_SCRATCH = 1 << 11;
         const BOTTOM_LEVEL_ACCELERATION_STRUCTURE_INPUT = 1 << 12;
         const TOP_LEVEL_ACCELERATION_STRUCTURE_INPUT = 1 << 13;
+        const EXPORT_BUFFER_HANDLE = 1 << 14;
         /// The combination of states that a buffer may be in _at the same time_.
         const INCLUSIVE = Self::MAP_READ.bits() | Self::COPY_SRC.bits() |
             Self::INDEX.bits() | Self::VERTEX.bits() | Self::UNIFORM.bits() |

wgpu-hal/src/metal/device.rs

@@ -5,7 +5,7 @@ use std::{
     sync::{atomic, Arc},
     thread, time,
 };
-
+use wgt::BufferHandle;
 use super::conv;
 use crate::auxil::map_naga_stage;
 
@@ -1423,4 +1423,8 @@ impl crate::Device for super::Device {
     fn get_internal_counters(&self) -> wgt::HalCounters {
         self.counters.clone()
     }
+
+    fn get_buffer_handle(&self, _buffer: &<Self::A as crate::Api>::Buffer) -> Result<wgt::BufferHandle, crate::DeviceError> {
+        unimplemented!("buffer handle not supported");
+    }
 }

wgpu-hal/src/vulkan/adapter.rs

@@ -1,4 +1,4 @@
-use super::conv;
+use super::{BufferHandleFunctions, conv, PreferredBufferHandle};
 
 use ash::{amd, ext, google, khr, vk};
 use parking_lot::Mutex;
@@ -111,6 +111,8 @@ pub struct PhysicalDeviceFeatures {
 
     /// Features provided by `VK_EXT_subgroup_size_control`, promoted to Vulkan 1.3.
     subgroup_size_control: Option<vk::PhysicalDeviceSubgroupSizeControlFeatures<'static>>,
+
+    buffer_handle: Option<vk::PhysicalDeviceExternalBufferInfo<'static>>
 }
 
 impl PhysicalDeviceFeatures {
@@ -451,6 +453,20 @@ impl PhysicalDeviceFeatures {
             } else {
                 None
             },
+            buffer_handle: if device_api_version >= vk::API_VERSION_1_1 || (enabled_extensions.contains(&khr::external_memory_capabilities::NAME)) {
+                Some(
+                    vk::PhysicalDeviceExternalBufferInfo::default()
+                        .usage(conv::map_buffer_usage(crate::BufferUses::COPY_SRC | crate::BufferUses::COPY_DST))
+                        .handle_type(
+                            #[cfg(target_os = "windows")]
+                            vk::ExternalMemoryHandleTypeFlags::OPAQUE_WIN32_KHR,
+                            #[cfg(not(target_os = "windows"))]
+                            vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD,
+                        )
+                )
+            } else {
+                None
+            }
         }
     }
 
@@ -470,6 +486,7 @@ impl PhysicalDeviceFeatures {
     ) -> (wgt::Features, wgt::DownlevelFlags) {
         use crate::auxil::db;
         use wgt::{DownlevelFlags as Df, Features as F};
+        println!("{:?}", caps.supported_extensions);
         let mut features = F::empty()
             | F::SPIRV_SHADER_PASSTHROUGH
             | F::MAPPABLE_PRIMARY_BUFFERS
@@ -733,6 +750,12 @@ impl PhysicalDeviceFeatures {
 
         features.set(F::RAY_QUERY, caps.supports_extension(khr::ray_query::NAME));
 
+        features.set(
+            F::BUFFER_HANDLE,
+            caps.supports_extension(khr::external_memory::NAME)
+                && (caps.supports_extension(khr::external_memory_win32::NAME) || caps.supports_extension(khr::external_memory_fd::NAME)),
+        );
+
         let rg11b10ufloat_renderable = supports_format(
             instance,
             phd,
@@ -842,6 +865,8 @@ pub struct PhysicalDeviceProperties {
     /// `VK_EXT_subgroup_size_control` extension, promoted to Vulkan 1.3.
     subgroup_size_control: Option<vk::PhysicalDeviceSubgroupSizeControlProperties<'static>>,
 
+    prefered_buffer_handle: super::PreferredBufferHandle,
+
     /// The device API version.
     ///
     /// Which is the version of Vulkan supported for device-level functionality.
@@ -988,6 +1013,11 @@ impl PhysicalDeviceProperties {
             extensions.push(khr::ray_query::NAME);
         }
 
+        if requested_features.contains(wgt::Features::BUFFER_HANDLE) {
+            extensions.push(khr::external_memory::NAME);
+            extensions.push(khr::external_memory_capabilities::NAME);
+        }
+
         // Require `VK_EXT_conservative_rasterization` if the associated feature was requested
         if requested_features.contains(wgt::Features::CONSERVATIVE_RASTERIZATION) {
             extensions.push(ext::conservative_rasterization::NAME);
@@ -1193,6 +1223,13 @@ impl super::InstanceShared {
                         .retain(|&x| x.extension_name_as_c_str() != Ok(ext::robustness2::NAME));
                 }
             };
+            if capabilities.supports_extension(khr::external_memory::NAME) {
+                if capabilities.supports_extension(khr::external_memory_win32::NAME) {
+                    capabilities.prefered_buffer_handle = super::PreferredBufferHandle::Win32;
+                } else if capabilities.supports_extension(khr::external_memory_fd::NAME) {
+                    capabilities.prefered_buffer_handle = super::PreferredBufferHandle::Fd;
+                }
+            }
             capabilities
         };
 
@@ -1504,6 +1541,7 @@ impl super::Instance {
                 }),
             image_format_list: phd_capabilities.device_api_version >= vk::API_VERSION_1_2
                 || phd_capabilities.supports_extension(khr::image_format_list::NAME),
+            preferred_buffer_handle: phd_capabilities.prefered_buffer_handle.clone(),
         };
         let capabilities = crate::Capabilities {
             limits: phd_capabilities.to_wgpu_limits(),
@@ -1680,6 +1718,12 @@ impl super::Adapter {
             None
         };
 
+        let buffer_handle_fns = match self.private_caps.preferred_buffer_handle {
+            PreferredBufferHandle::None => BufferHandleFunctions::None,
+            PreferredBufferHandle::Win32 => BufferHandleFunctions::Win32(khr::external_memory_win32::Device::new(&self.instance.raw, &raw_device)),
+            PreferredBufferHandle::Fd => BufferHandleFunctions::Fd(khr::external_memory_fd::Device::new(&self.instance.raw, &raw_device))
+        };
+
         let naga_options = {
             use naga::back::spv;
 
@@ -1840,6 +1884,7 @@ impl super::Adapter {
                 draw_indirect_count: indirect_count_fn,
                 timeline_semaphore: timeline_semaphore_fn,
                 ray_tracing: ray_tracing_fns,
+                buffer_handle: buffer_handle_fns,
             },
             pipeline_cache_validation_key,
             vendor_id: self.phd_capabilities.properties.vendor_id,

wgpu-hal/src/vulkan/conv.rs

@@ -1,4 +1,5 @@
 use ash::vk;
+use crate::vulkan::PreferredBufferHandle;
 
 impl super::PrivateCapabilities {
     pub fn map_texture_format(&self, format: wgt::TextureFormat) -> vk::Format {
@@ -964,3 +965,11 @@ pub fn map_acceleration_structure_usage_to_barrier(
 
     (stages, access)
 }
+
+pub(super) fn map_preferred_buffer_handle(preferred_buffer_handle: &super::PreferredBufferHandle) -> Option<vk::ExternalMemoryHandleTypeFlags> {
+    match preferred_buffer_handle {
+        PreferredBufferHandle::None => None,
+        PreferredBufferHandle::Win32 => Some(vk::ExternalMemoryHandleTypeFlags::OPAQUE_WIN32),
+        PreferredBufferHandle::Fd => Some(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD),
+    }
+}

wgpu-hal/src/vulkan/device.rs

@@ -1,4 +1,4 @@
-use super::conv;
+use super::{BufferHandleFunctions, conv};
 
 use arrayvec::ArrayVec;
 use ash::{khr, vk};
@@ -13,6 +13,8 @@ use std::{
     ptr,
     sync::Arc,
 };
+use std::ffi::c_void;
+use bitflags::Flags;
 
 impl super::DeviceShared {
     pub(super) unsafe fn set_object_name(&self, object: impl vk::Handle, name: &str) {
@@ -315,6 +317,19 @@ impl gpu_alloc::MemoryDevice<vk::DeviceMemory> for super::DeviceShared {
             info = info.push_next(&mut info_flags);
         }
 
+        let mut export_mem_info;
+        let mut export_mem_info_win;
+
+        if self.private_caps.preferred_buffer_handle == super::PreferredBufferHandle::Win32  {
+            export_mem_info_win = vk::ExportMemoryWin32HandleInfoKHR::default();
+            info = info.push_next(&mut export_mem_info_win);
+        }
+
+        if self.private_caps.preferred_buffer_handle != super::PreferredBufferHandle::None {
+            export_mem_info = vk::ExportMemoryAllocateInfo::default().handle_types(conv::map_preferred_buffer_handle(&self.private_caps.preferred_buffer_handle).unwrap());
+            info = info.push_next(&mut export_mem_info);
+        }
+
         match unsafe { self.raw.allocate_memory(&info, None) } {
             Ok(memory) => {
                 self.memory_allocations_counter.add(1);
@@ -914,16 +929,32 @@ impl crate::Device for super::Device {
             req.alignment
         } - 1;
 
-        let block = unsafe {
-            self.mem_allocator.lock().alloc(
-                &*self.shared,
-                gpu_alloc::Request {
-                    size: req.size,
-                    align_mask: alignment_mask,
-                    usage: alloc_usage,
-                    memory_types: req.memory_type_bits & self.valid_ash_memory_types,
-                },
-            )?
+        let block = if desc.usage.contains(crate::BufferUses::EXPORT_BUFFER_HANDLE) {
+            unsafe {
+                self.mem_allocator.lock().alloc_with_dedicated(
+                    &*self.shared,
+                    gpu_alloc::Request {
+                        size: req.size,
+                        align_mask: alignment_mask,
+                        usage: alloc_usage,
+                        memory_types: req.memory_type_bits & self.valid_ash_memory_types,
+                    },
+                    // it seems exporting buffer handles does not take offset
+                    gpu_alloc::Dedicated::Required,
+                )?
+            }
+        } else {
+            unsafe {
+                self.mem_allocator.lock().alloc(
+                    &*self.shared,
+                    gpu_alloc::Request {
+                        size: req.size,
+                        align_mask: alignment_mask,
+                        usage: alloc_usage,
+                        memory_types: req.memory_type_bits & self.valid_ash_memory_types,
+                    },
+                )?
+            }
         };
 
         unsafe {
@@ -2463,6 +2494,29 @@ impl crate::Device for super::Device {
 
         self.counters.clone()
     }
+
+    unsafe fn get_buffer_handle(&self, buffer: &super::Buffer) -> Result<wgt::BufferHandle, crate::DeviceError> {
+        Ok(match &self.shared.extension_fns.buffer_handle {
+            BufferHandleFunctions::None => unreachable!("no extension enabled"),
+            BufferHandleFunctions::Win32(win_32) => unsafe {
+                wgt::BufferHandle::Win32(
+                    win_32.get_memory_win32_handle(
+                        &vk::MemoryGetWin32HandleInfoKHR::default()
+                            .handle_type(vk::ExternalMemoryHandleTypeFlags::OPAQUE_WIN32_KHR)
+                            .memory(*buffer.block.as_ref().expect("cannot export imported buffer").lock().memory()),
+                    ).map_err(super::map_host_device_oom_err)? as *mut c_void
+                ) },
+            BufferHandleFunctions::Fd(fd) => unsafe {
+                wgt::BufferHandle::Fd(
+                    fd.get_memory_fd(
+                        &vk::MemoryGetFdInfoKHR::default()
+                            .handle_type(vk::ExternalMemoryHandleTypeFlags::OPAQUE_FD)
+                            .memory(*buffer.block.as_ref().expect("cannot export imported buffer").lock().memory()),
+                    ).map_err(super::map_host_device_oom_err)?
+                ) },
+            }
+        )
+    }
 }
 
 impl super::DeviceShared {

wgpu-hal/src/vulkan/instance.rs

@@ -366,6 +366,17 @@ impl super::Instance {
                 None
             };
 
+        let get_external_memory_capability =
+            if extensions.contains(&khr::external_memory_capabilities::NAME) {
+                log::debug!("Enabling device properties2");
+                Some(khr::external_memory_capabilities::Instance::new(
+                    &entry,
+                    &raw_instance,
+                ))
+            } else {
+                None
+            };
+
         let drop_guard = crate::DropGuard::from_option(drop_callback);
 
         Ok(Self {
@@ -376,6 +387,7 @@ impl super::Instance {
                 flags,
                 debug_utils,
                 get_physical_device_properties,
+                get_external_memory_capability,
                 entry,
                 has_nv_optimus,
                 instance_api_version,