introduce DeviceBufferTracker which holds weak references to buffers

wgpu-core/src/command/mod.rs

@@ -430,10 +430,11 @@ impl<A: HalApi> CommandBuffer<A> {
  ) {
  profiling::scope!("insert_barriers_from_device_tracker");
 
- base.buffers.set_from_tracker(&head.buffers);
- base.textures.set_from_tracker(&head.textures);
+ let buffer_barriers = base
+ .buffers
+ .set_from_tracker_and_drain_transitions(&head.buffers, snatch_guard);
 
- let buffer_barriers = base.buffers.drain_transitions(snatch_guard);
+ base.textures.set_from_tracker(&head.textures);
  let (transitions, textures) = base.textures.drain_transitions(snatch_guard);
  let texture_barriers = transitions
  .into_iter()

wgpu-core/src/device/resource.rs

@@ -3549,7 +3549,9 @@ impl<A: HalApi> Device<A> {
  // During these iterations, we discard all errors. We don't care!
  let trackers = self.trackers.lock();
  for buffer in trackers.buffers.used_resources() {
- let _ = buffer.destroy();
+ if let Some(buffer) = Weak::upgrade(&buffer) {
+ let _ = buffer.destroy();
+ }
  }
  for texture in trackers.textures.used_resources() {
  let _ = texture.destroy();

wgpu-core/src/track/buffer.rs

@@ -5,7 +5,7 @@
  * one subresource, they have no selector.
 !*/
 
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use super::{PendingTransition, TrackerIndex};
 use crate::{
@@ -231,7 +231,7 @@ impl<A: HalApi> BufferUsageScope<A> {
  }
 }
 
-/// Stores all buffer state within a command buffer or device.
+/// Stores all buffer state within a command buffer.
 pub(crate) struct BufferTracker<A: HalApi> {
  start: Vec<BufferUses>,
  end: Vec<BufferUses>,
@@ -294,38 +294,6 @@ impl<A: HalApi> BufferTracker<A> {
  buffer_barriers
  }
 
- /// Inserts a single buffer and its state into the resource tracker.
- ///
- /// If the resource already exists in the tracker, this will panic.
- ///
- /// If the ID is higher than the length of internal vectors,
- /// the vectors will be extended. A call to set_size is not needed.
- pub fn insert_single(&mut self, resource: &Arc<Buffer<A>>, state: BufferUses) {
- let index = resource.tracker_index().as_usize();
-
- self.allow_index(index);
-
- self.tracker_assert_in_bounds(index);
-
- unsafe {
- let currently_owned = self.metadata.contains_unchecked(index);
-
- if currently_owned {
- panic!("Tried to insert buffer already tracked");
- }
-
- insert(
- Some(&mut self.start),
- &mut self.end,
- &mut self.metadata,
- index,
- BufferStateProvider::Direct { state },
- None,
- ResourceMetadataProvider::Direct { resource },
- )
- }
- }
-
  /// Sets the state of a single buffer.
  ///
  /// If a transition is needed to get the buffer into the given state, that transition
@@ -494,6 +462,131 @@ impl<A: HalApi> BufferTracker<A> {
  }
 }
 
+/// Stores all buffer state within a device.
+pub(crate) struct DeviceBufferTracker<A: HalApi> {
+ current_states: Vec<BufferUses>,
+ metadata: ResourceMetadata<Weak<Buffer<A>>>,
+ temp: Vec<PendingTransition<BufferUses>>,
+}
+
+impl<A: HalApi> DeviceBufferTracker<A> {
+ pub fn new() -> Self {
+ Self {
+ current_states: Vec::new(),
+ metadata: ResourceMetadata::new(),
+ temp: Vec::new(),
+ }
+ }
+
+ fn tracker_assert_in_bounds(&self, index: usize) {
+ strict_assert!(index < self.current_states.len());
+ self.metadata.tracker_assert_in_bounds(index);
+ }
+
+ /// Extend the vectors to let the given index be valid.
+ fn allow_index(&mut self, index: usize) {
+ if index >= self.current_states.len() {
+ self.current_states.resize(index + 1, BufferUses::empty());
+ self.metadata.set_size(index + 1);
+ }
+ }
+
+ /// Returns a list of all buffers tracked.
+ pub fn used_resources(&self) -> impl Iterator<Item = Weak<Buffer<A>>> + '_ {
+ self.metadata.owned_resources()
+ }
+
+ /// Inserts a single buffer and its state into the resource tracker.
+ ///
+ /// If the resource already exists in the tracker, it will be overwritten.
+ pub fn insert_single(&mut self, buffer: &Arc<Buffer<A>>, state: BufferUses) {
+ let index = buffer.tracker_index().as_usize();
+
+ self.allow_index(index);
+
+ self.tracker_assert_in_bounds(index);
+
+ unsafe {
+ insert(
+ None,
+ &mut self.current_states,
+ &mut self.metadata,
+ index,
+ BufferStateProvider::Direct { state },
+ None,
+ ResourceMetadataProvider::Direct {
+ resource: &Arc::downgrade(buffer),
+ },
+ )
+ }
+ }
+
+ /// Sets the state of a single buffer.
+ ///
+ /// If a transition is needed to get the buffer into the given state, that transition
+ /// is returned. No more than one transition is needed.
+ pub fn set_single(
+ &mut self,
+ buffer: &Arc<Buffer<A>>,
+ state: BufferUses,
+ ) -> Option<PendingTransition<BufferUses>> {
+ let index: usize = buffer.tracker_index().as_usize();
+
+ self.tracker_assert_in_bounds(index);
+
+ let start_state_provider = BufferStateProvider::Direct { state };
+
+ unsafe {
+ barrier(
+ &mut self.current_states,
+ index,
+ start_state_provider.clone(),
+ &mut self.temp,
+ )
+ };
+ unsafe { update(&mut self.current_states, index, start_state_provider) };
+
+ strict_assert!(self.temp.len() <= 1);
+
+ self.temp.pop()
+ }
+
+ /// Sets the given state for all buffers in the given tracker.
+ ///
+ /// If a transition is needed to get the buffers into the needed state,
+ /// those transitions are returned.
+ pub fn set_from_tracker_and_drain_transitions<'a, 'b: 'a>(
+ &'a mut self,
+ tracker: &'a BufferTracker<A>,
+ snatch_guard: &'b SnatchGuard<'b>,
+ ) -> impl Iterator<Item = BufferBarrier<'a, A>> {
+ for index in tracker.metadata.owned_indices() {
+ self.tracker_assert_in_bounds(index);
+
+ let start_state_provider = BufferStateProvider::Indirect {
+ state: &tracker.start,
+ };
+ let end_state_provider = BufferStateProvider::Indirect {
+ state: &tracker.end,
+ };
+ unsafe {
+ barrier(
+ &mut self.current_states,
+ index,
+ start_state_provider,
+ &mut self.temp,
+ )
+ };
+ unsafe { update(&mut self.current_states, index, end_state_provider) };
+ }
+
+ self.temp.drain(..).map(|pending| {
+ let buf = unsafe { tracker.metadata.get_resource_unchecked(pending.id as _) };
+ pending.into_hal(buf, snatch_guard)
+ })
+ }
+}
+
 /// Source of Buffer State.
 #[derive(Debug, Clone)]
 enum BufferStateProvider<'a> {
@@ -619,14 +712,14 @@ unsafe fn insert_or_barrier_update<A: HalApi>(
 }
 
 #[inline(always)]
-unsafe fn insert<A: HalApi>(
+unsafe fn insert<T: Clone>(
  start_states: Option<&mut [BufferUses]>,
  current_states: &mut [BufferUses],
- resource_metadata: &mut ResourceMetadata<Arc<Buffer<A>>>,
+ resource_metadata: &mut ResourceMetadata<T>,
  index: usize,
  start_state_provider: BufferStateProvider<'_>,
  end_state_provider: Option<BufferStateProvider<'_>>,
- metadata_provider: ResourceMetadataProvider<'_, Arc<Buffer<A>>>,
+ metadata_provider: ResourceMetadataProvider<'_, T>,
 ) {
  let new_start_state = unsafe { start_state_provider.get_state(index) };
  let new_end_state =

wgpu-core/src/track/mod.rs

@@ -113,7 +113,9 @@ use crate::{
 use std::{fmt, ops, sync::Arc};
 use thiserror::Error;
 
-pub(crate) use buffer::{BufferBindGroupState, BufferTracker, BufferUsageScope};
+pub(crate) use buffer::{
+ BufferBindGroupState, BufferTracker, BufferUsageScope, DeviceBufferTracker,
+};
 use metadata::{ResourceMetadata, ResourceMetadataProvider};
 pub(crate) use stateless::{StatelessBindGroupState, StatelessTracker};
 pub(crate) use texture::{
@@ -600,14 +602,14 @@ impl<'a, A: HalApi> UsageScope<'a, A> {
 
 /// A tracker used by Device.
 pub(crate) struct DeviceTracker<A: HalApi> {
- pub buffers: BufferTracker<A>,
+ pub buffers: DeviceBufferTracker<A>,
  pub textures: TextureTracker<A>,
 }
 
 impl<A: HalApi> DeviceTracker<A> {
  pub fn new() -> Self {
  Self {
- buffers: BufferTracker::new(),
+ buffers: DeviceBufferTracker::new(),
  textures: TextureTracker::new(),
  }
  }