Deduplicate BindGroupLayout by value

wgpu-core/src/binding_model.rs

@@ -5,7 +5,7 @@
 use crate::{
     id::{BindGroupLayoutId, BufferId, DeviceId, SamplerId, TextureViewId},
     track::{TrackerSet, DUMMY_SELECTOR},
-    FastHashMap, LifeGuard, RefCount, Stored, MAX_BIND_GROUPS,
+    FastHashMap, LifeGuard, MultiRefCount, RefCount, Stored, MAX_BIND_GROUPS,
 };
 
 use arrayvec::ArrayVec;
@@ -54,7 +54,7 @@ pub(crate) type BindEntryMap = FastHashMap<u32, wgt::BindGroupLayoutEntry>;
 pub struct BindGroupLayout<B: hal::Backend> {
     pub(crate) raw: B::DescriptorSetLayout,
     pub(crate) device_id: Stored<DeviceId>,
-    pub(crate) life_guard: LifeGuard,
+    pub(crate) multi_ref_count: MultiRefCount,
     pub(crate) entries: BindEntryMap,
     pub(crate) desc_counts: DescriptorCounts,
     pub(crate) dynamic_count: usize,

wgpu-core/src/device/life.rs

@@ -252,7 +252,7 @@ impl<B: hal::Backend> LifetimeTracker<B> {
         });
     }
 
-    pub fn map(&mut self, buffer: id::BufferId, ref_count: RefCount) {
+    pub(crate) fn map(&mut self, buffer: id::BufferId, ref_count: RefCount) {
         self.mapped.push(Stored {
             value: buffer,
             ref_count,

wgpu-core/src/device/mod.rs

@@ -8,7 +8,8 @@ use crate::{
     hub::{GfxBackend, Global, GlobalIdentityHandlerFactory, Hub, Input, Token},
     id, pipeline, resource, span, swap_chain,
     track::{BufferState, TextureState, TrackerSet},
-    validation, FastHashMap, LifeGuard, PrivateFeatures, Stored, SubmissionIndex, MAX_BIND_GROUPS,
+    validation, FastHashMap, LifeGuard, MultiRefCount, PrivateFeatures, Stored, SubmissionIndex,
+    MAX_BIND_GROUPS,
 };
 
 use arrayvec::ArrayVec;
@@ -1197,23 +1198,22 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             }
         }
 
-        // TODO: deduplicate the bind group layouts at some level.
-        // We can't do it right here, because in the remote scenario
-        // the client need to know if the same ID can be used, or not.
-        if false {
+        let (device_guard, mut token) = hub.devices.read(&mut token);
+        let device = &device_guard[device_id];
+
+        // If there is an equivalent BGL, just bump the refcount and return it.
+        {
             let (bgl_guard, _) = hub.bind_group_layouts.read(&mut token);
             let bind_group_layout_id = bgl_guard
                 .iter(device_id.backend())
                 .find(|(_, bgl)| bgl.entries == entry_map);
 
-            if let Some((id, _)) = bind_group_layout_id {
+            if let Some((id, value)) = bind_group_layout_id {
+                value.multi_ref_count.inc();
                 return Ok(id);
             }
         }
 
-        let (device_guard, mut token) = hub.devices.read(&mut token);
-        let device = &device_guard[device_id];
-
         // Validate the count parameter
         for binding in desc
             .bindings
@@ -1275,7 +1275,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                 value: device_id,
                 ref_count: device.life_guard.add_ref(),
             },
-            life_guard: LifeGuard::new(),
+            multi_ref_count: MultiRefCount::new(),
             entries: entry_map,
             desc_counts: raw_bindings.iter().cloned().collect(),
             dynamic_count: desc
@@ -1309,12 +1309,12 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let hub = B::hub(self);
         let mut token = Token::root();
         let (device_id, ref_count) = {
-            let (mut bind_group_layout_guard, _) = hub.bind_group_layouts.write(&mut token);
-            let layout = &mut bind_group_layout_guard[bind_group_layout_id];
-            (
-                layout.device_id.value,
-                layout.life_guard.ref_count.take().unwrap(),
-            )
+            let (bind_group_layout_guard, _) = hub.bind_group_layouts.read(&mut token);
+            let layout = &bind_group_layout_guard[bind_group_layout_id];
+            match layout.multi_ref_count.dec() {
+                Some(last) => (layout.device_id.value, last),
+                None => return,
+            }
         };
 
         let (device_guard, mut token) = hub.devices.read(&mut token);
@@ -1378,7 +1378,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                     .iter()
                     .map(|&id| Stored {
                         value: id,
-                        ref_count: bind_group_layout_guard[id].life_guard.add_ref(),
+                        ref_count: bind_group_layout_guard[id].multi_ref_count.add_ref(),
                     })
                     .collect()
             },

wgpu-core/src/lib.rs

@@ -62,9 +62,9 @@ type Epoch = u32;
 
 pub type RawString = *const c_char;
 
-//TODO: make it private. Currently used for swapchain creation impl.
+/// Reference count object that is 1:1 with each reference.
 #[derive(Debug)]
-pub struct RefCount(ptr::NonNull<AtomicUsize>);
+struct RefCount(ptr::NonNull<AtomicUsize>);
 
 unsafe impl Send for RefCount {}
 unsafe impl Sync for RefCount {}
@@ -89,7 +89,7 @@ impl RefCount {
     /// logic. To use this safely from outside of `Drop::drop`, the calling function must move
     /// `Self` into a `ManuallyDrop`.
     unsafe fn rich_drop_inner(&mut self) -> bool {
-        if self.0.as_ref().fetch_sub(1, Ordering::Relaxed) == 1 {
+        if self.0.as_ref().fetch_sub(1, Ordering::AcqRel) == 1 {
             let _ = Box::from_raw(self.0.as_ptr());
             true
         } else {
@@ -100,7 +100,7 @@ impl RefCount {
 
 impl Clone for RefCount {
     fn clone(&self) -> Self {
-        let old_size = unsafe { self.0.as_ref() }.fetch_add(1, Ordering::Relaxed);
+        let old_size = unsafe { self.0.as_ref() }.fetch_add(1, Ordering::Release);
         assert!(old_size < Self::MAX);
         RefCount(self.0)
     }
@@ -136,6 +136,46 @@ fn loom() {
     });
 }
 
+/// Reference count object that tracks multiple references.
+#[derive(Debug)]
+struct MultiRefCount(ptr::NonNull<AtomicUsize>);
+
+unsafe impl Send for MultiRefCount {}
+unsafe impl Sync for MultiRefCount {}
+
+impl MultiRefCount {
+    fn new() -> Self {
+        let bx = Box::new(AtomicUsize::new(1));
+        MultiRefCount(unsafe { ptr::NonNull::new_unchecked(Box::into_raw(bx)) })
+    }
+
+    fn inc(&self) {
+        unsafe { self.0.as_ref() }.fetch_add(1, Ordering::Release);
+    }
+
+    fn add_ref(&self) -> RefCount {
+        self.inc();
+        RefCount(self.0)
+    }
+
+    fn dec(&self) -> Option<RefCount> {
+        match unsafe { self.0.as_ref() }.fetch_sub(1, Ordering::AcqRel) {
+            0 => unreachable!(),
+            1 => Some(self.add_ref()),
+            _ => None,
+        }
+    }
+}
+
+impl Drop for MultiRefCount {
+    fn drop(&mut self) {
+        // We don't do anything here. We rely on the fact that
+        // `dec` was called before `MultiRefCount` got dropped,
+        // which spawned `RefCount`, which upon deletion would
+        // destroy the Box.
+    }
+}
+
 #[derive(Debug)]
 struct LifeGuard {
     ref_count: Option<RefCount>,

wgpu-core/src/resource.rs

@@ -73,7 +73,7 @@ pub enum BufferMapAsyncStatus {
 }
 
 #[derive(Debug)]
-pub enum BufferMapState<B: hal::Backend> {
+pub(crate) enum BufferMapState<B: hal::Backend> {
     /// Mapped at creation.
     Init {
         ptr: NonNull<u8>,
@@ -118,7 +118,7 @@ impl BufferMapOperation {
 }
 
 #[derive(Debug)]
-pub struct BufferPendingMapping {
+pub(crate) struct BufferPendingMapping {
     pub sub_range: hal::buffer::SubRange,
     pub op: BufferMapOperation,
     // hold the parent alive while the mapping is active

wgpu-core/src/track/mod.rs

@@ -195,7 +195,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     }
 
     /// Remove an id from the tracked map.
-    pub fn remove(&mut self, id: S::Id) -> bool {
+    pub(crate) fn remove(&mut self, id: S::Id) -> bool {
         let (index, epoch, backend) = id.unzip();
         debug_assert_eq!(backend, self.backend);
         match self.map.remove(&index) {
@@ -208,7 +208,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     }
 
     /// Removes the resource from the tracker if we are holding the last reference.
-    pub fn remove_abandoned(&mut self, id: S::Id) -> bool {
+    pub(crate) fn remove_abandoned(&mut self, id: S::Id) -> bool {
         let (index, epoch, backend) = id.unzip();
         debug_assert_eq!(backend, self.backend);
         match self.map.entry(index) {
@@ -226,7 +226,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     }
 
     /// Try to optimize the internal representation.
-    pub fn optimize(&mut self) {
+    pub(crate) fn optimize(&mut self) {
         for resource in self.map.values_mut() {
             resource.state.optimize();
         }
@@ -248,7 +248,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     /// Initialize a resource to be used.
     ///
     /// Returns false if the resource is already registered.
-    pub fn init(&mut self, id: S::Id, ref_count: RefCount, state: S) -> Result<(), &S> {
+    pub(crate) fn init(&mut self, id: S::Id, ref_count: RefCount, state: S) -> Result<(), &S> {
         let (index, epoch, backend) = id.unzip();
         debug_assert_eq!(backend, self.backend);
         match self.map.entry(index) {
@@ -302,7 +302,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     /// Extend the usage of a specified resource.
     ///
     /// Returns conflicting transition as an error.
-    pub fn change_extend(
+    pub(crate) fn change_extend(
         &mut self,
         id: S::Id,
         ref_count: &RefCount,
@@ -315,7 +315,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     }
 
     /// Replace the usage of a specified resource.
-    pub fn change_replace(
+    pub(crate) fn change_replace(
         &mut self,
         id: S::Id,
         ref_count: &RefCount,
@@ -332,7 +332,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     /// Turn the tracking from the "expand" mode into the "replace" one,
     /// installing the selected usage as the "first".
     /// This is a special operation only used by the render pass attachments.
-    pub fn prepend(
+    pub(crate) fn prepend(
         &mut self,
         id: S::Id,
         ref_count: &RefCount,
@@ -346,7 +346,7 @@ impl<S: ResourceState> ResourceTracker<S> {
 
     /// Merge another tracker into `self` by extending the current states
     /// without any transitions.
-    pub fn merge_extend(&mut self, other: &Self) -> Result<(), PendingTransition<S>> {
+    pub(crate) fn merge_extend(&mut self, other: &Self) -> Result<(), PendingTransition<S>> {
         debug_assert_eq!(self.backend, other.backend);
         for (&index, new) in other.map.iter() {
             match self.map.entry(index) {
@@ -365,7 +365,7 @@ impl<S: ResourceState> ResourceTracker<S> {
 
     /// Merge another tracker, adding it's transitions to `self`.
     /// Transitions the current usage to the new one.
-    pub fn merge_replace<'a>(&'a mut self, other: &'a Self) -> Drain<PendingTransition<S>> {
+    pub(crate) fn merge_replace<'a>(&'a mut self, other: &'a Self) -> Drain<PendingTransition<S>> {
         for (&index, new) in other.map.iter() {
             match self.map.entry(index) {
                 Entry::Vacant(e) => {
@@ -389,7 +389,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     /// the last read-only usage, if possible.
     ///
     /// Returns the old usage as an error if there is a conflict.
-    pub fn use_extend<'a, T: 'a + Borrow<RefCount>>(
+    pub(crate) fn use_extend<'a, T: 'a + Borrow<RefCount>>(
         &mut self,
         storage: &'a Storage<T, S::Id>,
         id: S::Id,
@@ -406,7 +406,7 @@ impl<S: ResourceState> ResourceTracker<S> {
     /// Combines storage access by 'Id' with the transition that replaces
     /// the last usage with a new one, returning an iterator over these
     /// transitions.
-    pub fn use_replace<'a, T: 'a + Borrow<RefCount>>(
+    pub(crate) fn use_replace<'a, T: 'a + Borrow<RefCount>>(
         &mut self,
         storage: &'a Storage<T, S::Id>,
         id: S::Id,