From ed95dfe9b4d2e9dc78470ca11c75b61313cb05a2 Mon Sep 17 00:00:00 2001
From: robtfm <50659922+robtfm@users.noreply.github.com>
Date: Sat, 23 Mar 2024 11:42:08 +0000
Subject: [PATCH] Pool tracker vecs (#5414)

* pool tracker vecs

* pool

* ci

* move pool to device

* use pool ref, cleanup and comment

* suspect all the future suspects (#5413)

* suspect all the future suspects

* changelog

* changelog

* review feedback

---------

Co-authored-by: Andreas Reich <r_andreas2@web.de>
---
 CHANGELOG.md                     |  1 +
 wgpu-core/src/command/compute.rs |  8 ++---
 wgpu-core/src/command/render.rs  | 12 ++++----
 wgpu-core/src/device/queue.rs    |  4 +--
 wgpu-core/src/device/resource.rs | 11 ++++++-
 wgpu-core/src/track/buffer.rs    | 12 +++++---
 wgpu-core/src/track/metadata.rs  |  5 ++++
 wgpu-core/src/track/mod.rs       | 50 +++++++++++++++++++++++---------
 wgpu-core/src/track/texture.rs   | 13 +++++++--
 9 files changed, 83 insertions(+), 33 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 51c0da98b1..25cdb8fed1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -134,6 +134,7 @@ Bottom level categories:
 - Fix linking when targeting android. By @ashdnazg in [#5326](https://github.com/gfx-rs/wgpu/pull/5326).
 - fix resource leak for buffers/textures dropped while having pending writes. By @robtfm in [#5413](https://github.com/gfx-rs/wgpu/pull/5413)
 - Failing to set the device lost closure will call the closure before returning. By @bradwerth in [#5358](https://github.com/gfx-rs/wgpu/pull/5358).
+- Use memory pooling for UsageScopes to avoid frequent large allocations. by @robtfm in [#5414](https://github.com/gfx-rs/wgpu/pull/5414)
 
 #### Naga
 - In spv-in, remove unnecessary "gl_PerVertex" name check so unused builtins will always be skipped. By @Imberflur in [#5227](https://github.com/gfx-rs/wgpu/pull/5227).
diff --git a/wgpu-core/src/command/compute.rs b/wgpu-core/src/command/compute.rs
index c2fd3ab397..941197b33a 100644
--- a/wgpu-core/src/command/compute.rs
+++ b/wgpu-core/src/command/compute.rs
@@ -272,14 +272,14 @@ where
     }
 }
 
-struct State<A: HalApi> {
+struct State<'a, A: HalApi> {
     binder: Binder<A>,
     pipeline: Option<id::ComputePipelineId>,
-    scope: UsageScope<A>,
+    scope: UsageScope<'a, A>,
     debug_scope_depth: u32,
 }
 
-impl<A: HalApi> State<A> {
+impl<'a, A: HalApi> State<'a, A> {
     fn is_ready(&self) -> Result<(), DispatchError> {
         let bind_mask = self.binder.invalid_mask();
         if bind_mask != 0 {
@@ -407,7 +407,7 @@ impl Global {
         let mut state = State {
             binder: Binder::new(),
             pipeline: None,
-            scope: UsageScope::new(&device.tracker_indices),
+            scope: device.new_usage_scope(),
             debug_scope_depth: 0,
         };
         let mut temp_offsets = Vec::new();
diff --git a/wgpu-core/src/command/render.rs b/wgpu-core/src/command/render.rs
index 9141ddb021..8b22b4275a 100644
--- a/wgpu-core/src/command/render.rs
+++ b/wgpu-core/src/command/render.rs
@@ -739,9 +739,9 @@ impl<A: HalApi> TextureView<A> {
 const MAX_TOTAL_ATTACHMENTS: usize = hal::MAX_COLOR_ATTACHMENTS + hal::MAX_COLOR_ATTACHMENTS + 1;
 type AttachmentDataVec<T> = ArrayVec<T, MAX_TOTAL_ATTACHMENTS>;
 
-struct RenderPassInfo<'a, A: HalApi> {
+struct RenderPassInfo<'a, 'd, A: HalApi> {
     context: RenderPassContext,
-    usage_scope: UsageScope<A>,
+    usage_scope: UsageScope<'d, A>,
     /// All render attachments, including depth/stencil
     render_attachments: AttachmentDataVec<RenderAttachment<'a, A>>,
     is_depth_read_only: bool,
@@ -754,7 +754,7 @@ struct RenderPassInfo<'a, A: HalApi> {
     multiview: Option<NonZeroU32>,
 }
 
-impl<'a, A: HalApi> RenderPassInfo<'a, A> {
+impl<'a, 'd, A: HalApi> RenderPassInfo<'a, 'd, A> {
     fn add_pass_texture_init_actions<V>(
         channel: &PassChannel<V>,
         texture_memory_actions: &mut CommandBufferTextureMemoryActions<A>,
@@ -790,7 +790,7 @@ impl<'a, A: HalApi> RenderPassInfo<'a, A> {
     }
 
     fn start(
-        device: &Device<A>,
+        device: &'d Device<A>,
         label: Option<&str>,
         color_attachments: &[Option<RenderPassColorAttachment>],
         depth_stencil_attachment: Option<&RenderPassDepthStencilAttachment>,
@@ -1214,7 +1214,7 @@ impl<'a, A: HalApi> RenderPassInfo<'a, A> {
 
         Ok(Self {
             context,
-            usage_scope: UsageScope::new(&device.tracker_indices),
+            usage_scope: device.new_usage_scope(),
             render_attachments,
             is_depth_read_only,
             is_stencil_read_only,
@@ -1230,7 +1230,7 @@ impl<'a, A: HalApi> RenderPassInfo<'a, A> {
         mut self,
         raw: &mut A::CommandEncoder,
         snatch_guard: &SnatchGuard,
-    ) -> Result<(UsageScope<A>, SurfacesInDiscardState<A>), RenderPassErrorInner> {
+    ) -> Result<(UsageScope<'d, A>, SurfacesInDiscardState<A>), RenderPassErrorInner> {
         profiling::scope!("RenderPassInfo::finish");
         unsafe {
             raw.end_render_pass();
diff --git a/wgpu-core/src/device/queue.rs b/wgpu-core/src/device/queue.rs
index 6ebb9eb09b..cb558ad6f0 100644
--- a/wgpu-core/src/device/queue.rs
+++ b/wgpu-core/src/device/queue.rs
@@ -1155,7 +1155,7 @@ impl Global {
                 + 1;
             let mut active_executions = Vec::new();
 
-            let mut used_surface_textures = track::TextureUsageScope::new();
+            let mut used_surface_textures = track::TextureUsageScope::default();
 
             let snatch_guard = device.snatchable_lock.read();
 
@@ -1435,7 +1435,7 @@ impl Global {
                                 baked.encoder.end_encoding().unwrap()
                             };
                             baked.list.push(present);
-                            used_surface_textures = track::TextureUsageScope::new();
+                            used_surface_textures = track::TextureUsageScope::default();
                         }
 
                         // done
diff --git a/wgpu-core/src/device/resource.rs b/wgpu-core/src/device/resource.rs
index 28ba0eafb1..04ca94ca96 100644
--- a/wgpu-core/src/device/resource.rs
+++ b/wgpu-core/src/device/resource.rs
@@ -28,7 +28,10 @@ use crate::{
     resource_log,
     snatch::{SnatchGuard, SnatchLock, Snatchable},
     storage::Storage,
-    track::{BindGroupStates, TextureSelector, Tracker, TrackerIndexAllocators},
+    track::{
+        BindGroupStates, TextureSelector, Tracker, TrackerIndexAllocators, UsageScope,
+        UsageScopePool,
+    },
     validation::{
         self, check_buffer_usage, check_texture_usage, validate_color_attachment_bytes_per_sample,
     },
@@ -135,6 +138,7 @@ pub struct Device<A: HalApi> {
     pub(crate) deferred_destroy: Mutex<Vec<DeferredDestroy<A>>>,
     #[cfg(feature = "trace")]
     pub(crate) trace: Mutex<Option<trace::Trace>>,
+    pub(crate) usage_scopes: UsageScopePool<A>,
 }
 
 pub(crate) enum DeferredDestroy<A: HalApi> {
@@ -296,6 +300,7 @@ impl<A: HalApi> Device<A> {
             instance_flags,
             pending_writes: Mutex::new(Some(pending_writes)),
             deferred_destroy: Mutex::new(Vec::new()),
+            usage_scopes: Default::default(),
         })
     }
 
@@ -3568,6 +3573,10 @@ impl<A: HalApi> Device<A> {
             let _ = texture.destroy();
         }
     }
+
+    pub(crate) fn new_usage_scope(&self) -> UsageScope<'_, A> {
+        UsageScope::new_pooled(&self.usage_scopes, &self.tracker_indices)
+    }
 }
 
 impl<A: HalApi> Device<A> {
diff --git a/wgpu-core/src/track/buffer.rs b/wgpu-core/src/track/buffer.rs
index a30ac2a225..6cf1fdda6f 100644
--- a/wgpu-core/src/track/buffer.rs
+++ b/wgpu-core/src/track/buffer.rs
@@ -108,23 +108,27 @@ impl<A: HalApi> BufferBindGroupState<A> {
 #[derive(Debug)]
 pub(crate) struct BufferUsageScope<A: HalApi> {
     state: Vec<BufferUses>,
-
     metadata: ResourceMetadata<Buffer<A>>,
 }
 
-impl<A: HalApi> BufferUsageScope<A> {
-    pub fn new() -> Self {
+impl<A: HalApi> Default for BufferUsageScope<A> {
+    fn default() -> Self {
         Self {
             state: Vec::new(),
-
             metadata: ResourceMetadata::new(),
         }
     }
+}
 
+impl<A: HalApi> BufferUsageScope<A> {
     fn tracker_assert_in_bounds(&self, index: usize) {
         strict_assert!(index < self.state.len());
         self.metadata.tracker_assert_in_bounds(index);
     }
+    pub fn clear(&mut self) {
+        self.state.clear();
+        self.metadata.clear();
+    }
 
     /// Sets the size of all the vectors inside the tracker.
     ///
diff --git a/wgpu-core/src/track/metadata.rs b/wgpu-core/src/track/metadata.rs
index 744783a7fa..3e71e0e084 100644
--- a/wgpu-core/src/track/metadata.rs
+++ b/wgpu-core/src/track/metadata.rs
@@ -39,6 +39,11 @@ impl<T: Resource> ResourceMetadata<T> {
         resize_bitvec(&mut self.owned, size);
     }
 
+    pub(super) fn clear(&mut self) {
+        self.resources.clear();
+        self.owned.clear();
+    }
+
     /// Ensures a given index is in bounds for all arrays and does
     /// sanity checks of the presence of a refcount.
     ///
diff --git a/wgpu-core/src/track/mod.rs b/wgpu-core/src/track/mod.rs
index 9ca37ebadc..374dfe7493 100644
--- a/wgpu-core/src/track/mod.rs
+++ b/wgpu-core/src/track/mod.rs
@@ -480,8 +480,8 @@ impl<A: HalApi> RenderBundleScope<A> {
     /// Create the render bundle scope and pull the maximum IDs from the hubs.
     pub fn new() -> Self {
         Self {
-            buffers: RwLock::new(BufferUsageScope::new()),
-            textures: RwLock::new(TextureUsageScope::new()),
+            buffers: RwLock::new(BufferUsageScope::default()),
+            textures: RwLock::new(TextureUsageScope::default()),
             bind_groups: RwLock::new(StatelessTracker::new()),
             render_pipelines: RwLock::new(StatelessTracker::new()),
             query_sets: RwLock::new(StatelessTracker::new()),
@@ -512,28 +512,52 @@ impl<A: HalApi> RenderBundleScope<A> {
     }
 }
 
+/// A pool for storing the memory used by [`UsageScope`]s. We take and store this memory when the
+/// scope is dropped to avoid reallocating. The memory required only grows and allocation cost is
+/// significant when a large number of resources have been used.
+pub(crate) type UsageScopePool<A> = Mutex<Vec<(BufferUsageScope<A>, TextureUsageScope<A>)>>;
+
 /// A usage scope tracker. Only needs to store stateful resources as stateless
 /// resources cannot possibly have a usage conflict.
 #[derive(Debug)]
-pub(crate) struct UsageScope<A: HalApi> {
+pub(crate) struct UsageScope<'a, A: HalApi> {
+    pub pool: &'a UsageScopePool<A>,
     pub buffers: BufferUsageScope<A>,
     pub textures: TextureUsageScope<A>,
 }
 
-impl<A: HalApi> UsageScope<A> {
-    /// Create the render bundle scope and pull the maximum IDs from the hubs.
-    pub fn new(tracker_indices: &TrackerIndexAllocators) -> Self {
-        let mut value = Self {
-            buffers: BufferUsageScope::new(),
-            textures: TextureUsageScope::new(),
-        };
+impl<'a, A: HalApi> Drop for UsageScope<'a, A> {
+    fn drop(&mut self) {
+        // clear vecs and push into pool
+        self.buffers.clear();
+        self.textures.clear();
+        self.pool.lock().push((
+            std::mem::take(&mut self.buffers),
+            std::mem::take(&mut self.textures),
+        ));
+    }
+}
 
-        value.buffers.set_size(tracker_indices.buffers.size());
-        value.textures.set_size(tracker_indices.textures.size());
+impl<A: HalApi> UsageScope<'static, A> {
+    pub fn new_pooled<'d>(
+        pool: &'d UsageScopePool<A>,
+        tracker_indices: &TrackerIndexAllocators,
+    ) -> UsageScope<'d, A> {
+        let pooled = pool.lock().pop().unwrap_or_default();
+
+        let mut scope = UsageScope::<'d, A> {
+            pool,
+            buffers: pooled.0,
+            textures: pooled.1,
+        };
 
-        value
+        scope.buffers.set_size(tracker_indices.buffers.size());
+        scope.textures.set_size(tracker_indices.textures.size());
+        scope
     }
+}
 
+impl<'a, A: HalApi> UsageScope<'a, A> {
     /// Merge the inner contents of a bind group into the usage scope.
     ///
     /// Only stateful things are merged in here, all other resources are owned
diff --git a/wgpu-core/src/track/texture.rs b/wgpu-core/src/track/texture.rs
index e7c4707c93..3cf95ff38a 100644
--- a/wgpu-core/src/track/texture.rs
+++ b/wgpu-core/src/track/texture.rs
@@ -210,6 +210,7 @@ pub(crate) struct TextureStateSet {
     simple: Vec<TextureUses>,
     complex: FastHashMap<usize, ComplexTextureState>,
 }
+
 impl TextureStateSet {
     fn new() -> Self {
         Self {
@@ -235,15 +236,16 @@ pub(crate) struct TextureUsageScope<A: HalApi> {
     metadata: ResourceMetadata<Texture<A>>,
 }
 
-impl<A: HalApi> TextureUsageScope<A> {
-    pub fn new() -> Self {
+impl<A: HalApi> Default for TextureUsageScope<A> {
+    fn default() -> Self {
         Self {
             set: TextureStateSet::new(),
-
             metadata: ResourceMetadata::new(),
         }
     }
+}
 
+impl<A: HalApi> TextureUsageScope<A> {
     fn tracker_assert_in_bounds(&self, index: usize) {
         self.metadata.tracker_assert_in_bounds(index);
 
@@ -258,6 +260,11 @@ impl<A: HalApi> TextureUsageScope<A> {
         });
     }
 
+    pub fn clear(&mut self) {
+        self.set.clear();
+        self.metadata.clear();
+    }
+
     /// Sets the size of all the vectors inside the tracker.
     ///
     /// Must be called with the highest possible Texture ID before