add temp_resources to CommandBufferMutable, BakedCommands & EncoderInFlight, remove it from ActiveSubmission

Use `CommandBufferMutable.temp_resources` to store the ray tracing related staging buffers and scratch buffers. These shouldn't have been stored in `PendingWrites.temp_resources` because they can get destroyed early (before the command buffer that uses them has been submitted) if any other command buffer gets submitted first.

Author: teoxoy

wgpu-core/src/command/mod.rs

@@ -28,6 +28,7 @@ pub use timestamp_writes::PassTimestampWrites;
 
 use self::memory_init::CommandBufferTextureMemoryActions;
 
+use crate::device::queue::TempResource;
 use crate::device::{Device, DeviceError, MissingFeatures};
 use crate::lock::{rank, Mutex};
 use crate::snatch::SnatchGuard;
@@ -432,6 +433,7 @@ impl Drop for CommandEncoder {
 pub(crate) struct BakedCommands {
     pub(crate) encoder: CommandEncoder,
     pub(crate) trackers: Tracker,
+    pub(crate) temp_resources: Vec<TempResource>,
     buffer_memory_init_actions: Vec<BufferInitTrackerAction>,
     texture_memory_actions: CommandBufferTextureMemoryActions,
 }
@@ -460,6 +462,7 @@ pub struct CommandBufferMutable {
 
     blas_actions: Vec<BlasAction>,
     tlas_actions: Vec<TlasAction>,
+    temp_resources: Vec<TempResource>,
 
     #[cfg(feature = "trace")]
     pub(crate) commands: Option<Vec<TraceCommand>>,
@@ -479,6 +482,7 @@ impl CommandBufferMutable {
         BakedCommands {
             encoder: self.encoder,
             trackers: self.trackers,
+            temp_resources: self.temp_resources,
             buffer_memory_init_actions: self.buffer_memory_init_actions,
             texture_memory_actions: self.texture_memory_actions,
         }
@@ -545,6 +549,7 @@ impl CommandBuffer {
                     pending_query_resets: QueryResetMap::new(),
                     blas_actions: Default::default(),
                     tlas_actions: Default::default(),
+                    temp_resources: Default::default(),
                     #[cfg(feature = "trace")]
                     commands: if device.trace.lock().is_some() {
                         Some(Vec::new())

wgpu-core/src/command/ray_tracing.rs

@@ -329,12 +329,9 @@ impl Global {
             }
         }
 
-        if let Some(queue) = device.get_queue() {
-            queue
-                .pending_writes
-                .lock()
-                .consume_temp(TempResource::ScratchBuffer(scratch_buffer));
-        }
+        cmd_buf_data
+            .temp_resources
+            .push(TempResource::ScratchBuffer(scratch_buffer));
 
         cmd_buf_data_guard.mark_successful();
         Ok(())
@@ -725,21 +722,15 @@ impl Global {
             }
 
             if let Some(staging_buffer) = staging_buffer {
-                if let Some(queue) = device.get_queue() {
-                    queue
-                        .pending_writes
-                        .lock()
-                        .consume_temp(TempResource::StagingBuffer(staging_buffer));
-                }
+                cmd_buf_data
+                    .temp_resources
+                    .push(TempResource::StagingBuffer(staging_buffer));
             }
         }
 
-        if let Some(queue) = device.get_queue() {
-            queue
-                .pending_writes
-                .lock()
-                .consume_temp(TempResource::ScratchBuffer(scratch_buffer));
-        }
+        cmd_buf_data
+            .temp_resources
+            .push(TempResource::ScratchBuffer(scratch_buffer));
 
         cmd_buf_data_guard.mark_successful();
         Ok(())

wgpu-core/src/device/life.rs

@@ -29,9 +29,6 @@ struct ActiveSubmission {
     /// submission has completed.
     index: SubmissionIndex,
 
-    /// Temporary resources to be freed once this queue submission has completed.
-    temp_resources: Vec<TempResource>,
-
     /// Buffers to be mapped once this submission has completed.
     mapped: Vec<Arc<Buffer>>,
 
@@ -107,11 +104,6 @@ impl ActiveSubmission {
 
     pub fn contains_blas(&self, blas: &Blas) -> bool {
         for encoder in &self.encoders {
-            // The ownership location of blas's depends on where the command encoder
-            // came from. If it is the staging command encoder on the queue, it is
-            // in the pending buffer list. If it came from a user command encoder,
-            // it is in the tracker.
-
             if encoder.trackers.blas_s.contains(blas) {
                 return true;
             }
@@ -122,11 +114,6 @@ impl ActiveSubmission {
 
     pub fn contains_tlas(&self, tlas: &Tlas) -> bool {
         for encoder in &self.encoders {
-            // The ownership location of tlas's depends on where the command encoder
-            // came from. If it is the staging command encoder on the queue, it is
-            // in the pending buffer list. If it came from a user command encoder,
-            // it is in the tracker.
-
             if encoder.trackers.tlas_s.contains(tlas) {
                 return true;
             }
@@ -203,15 +190,9 @@ impl LifetimeTracker {
     }
 
     /// Start tracking resources associated with a new queue submission.
-    pub fn track_submission(
-        &mut self,
-        index: SubmissionIndex,
-        temp_resources: impl Iterator<Item = TempResource>,
-        encoders: Vec<EncoderInFlight>,
-    ) {
+    pub fn track_submission(&mut self, index: SubmissionIndex, encoders: Vec<EncoderInFlight>) {
         self.active.push(ActiveSubmission {
             index,
-            temp_resources: temp_resources.collect(),
             mapped: Vec::new(),
             encoders,
             work_done_closures: SmallVec::new(),
@@ -332,7 +313,6 @@ impl LifetimeTracker {
                 profiling::scope!("drop command buffer trackers");
                 drop(encoder);
             }
-            drop(a.temp_resources);
             work_done_closures.extend(a.work_done_closures);
         }
         work_done_closures
@@ -347,7 +327,12 @@ impl LifetimeTracker {
             .active
             .iter_mut()
             .find(|a| a.index == last_submit_index)
-            .map(|a| &mut a.temp_resources);
+            .map(|a| {
+                // Because this resource's `last_submit_index` matches `a.index`,
+                // we know that we must have done something with the resource,
+                // so `a.encoders` should not be empty.
+                &mut a.encoders.last_mut().unwrap().temp_resources
+            });
         if let Some(resources) = resources {
             resources.push(temp_resource);
         }

wgpu-core/src/device/queue.rs

@@ -268,6 +268,7 @@ pub enum TempResource {
 pub(crate) struct EncoderInFlight {
     inner: crate::command::CommandEncoder,
     pub(crate) trackers: Tracker,
+    pub(crate) temp_resources: Vec<TempResource>,
 
     /// These are the buffers that have been tracked by `PendingWrites`.
     pub(crate) pending_buffers: FastHashMap<TrackerIndex, Arc<Buffer>>,
@@ -377,6 +378,7 @@ impl PendingWrites {
                     hal_label: None,
                 },
                 trackers: Tracker::new(),
+                temp_resources: mem::take(&mut self.temp_resources),
                 pending_buffers,
                 pending_textures,
             };
@@ -1195,6 +1197,7 @@ impl Queue {
                         active_executions.push(EncoderInFlight {
                             inner: baked.encoder,
                             trackers: baked.trackers,
+                            temp_resources: baked.temp_resources,
                             pending_buffers: FastHashMap::default(),
                             pending_textures: FastHashMap::default(),
                         });
@@ -1293,11 +1296,8 @@ impl Queue {
             profiling::scope!("cleanup");
 
             // this will register the new submission to the life time tracker
-            self.lock_life().track_submission(
-                submit_index,
-                pending_writes.temp_resources.drain(..),
-                active_executions,
-            );
+            self.lock_life()
+                .track_submission(submit_index, active_executions);
             drop(pending_writes);
 
             // This will schedule destruction of all resources that are no longer needed