Merge remote-tracking branch 'pcwalton/shared-memory-mesh-collection' into shared-memory-mesh-collection

Author: pcwalton

crates/bevy_pbr/src/render/mesh.rs

@@ -1413,7 +1413,7 @@ impl RenderMeshInstanceGpuPrepared {
         entity: MainEntity,
         render_mesh_instances: &mut MainEntityHashMap<RenderMeshInstanceGpu>,
         current_input_buffer: &mut InstanceInputUniformBuffer<MeshInputUniform>,
-        previous_input_buffer: &mut InstanceInputUniformBuffer<MeshInputUniform>,
+        previous_input_buffer: &InstanceInputUniformBuffer<MeshInputUniform>,
     ) -> Option<u32> {
         // Did the last frame contain this entity as well?
         let current_uniform_index;
@@ -1431,7 +1431,7 @@ impl RenderMeshInstanceGpuPrepared {
                 // shader will need it to compute motion vectors.
                 let previous_mesh_input_uniform =
                     current_input_buffer.get_unchecked(current_uniform_index);
-                let previous_input_index = previous_input_buffer.add(previous_mesh_input_uniform);
+                let previous_input_index = previous_input_buffer.push(previous_mesh_input_uniform);
                 self.mesh_input_uniform.previous_input_index = previous_input_index;
 
                 // Write in the new mesh input uniform.
@@ -1986,6 +1986,9 @@ pub fn collect_meshes_for_gpu_building(
         mesh_culling_data_buffer.grow(current_input_buffer.len() as u32);
     }
 
+    // Pre-allocate the previous input buffer for concurrent pushes.
+    previous_input_buffer.reserve(current_input_buffer.len() as u32);
+
     // Channels used by parallel workers to send data to the single consumer.
     let (prepared_tx, prepared_rx) = mpsc::channel();
     let (reextract_tx, reextract_rx) = mpsc::channel();
@@ -2001,6 +2004,7 @@ pub fn collect_meshes_for_gpu_building(
         let frame_count = *frame_count;
         let render_mesh_instances = &*render_mesh_instances;
         let current_input_buffer = &*current_input_buffer;
+        let previous_input_buffer = &*previous_input_buffer;
         let mesh_culling_data_buffer = &*mesh_culling_data_buffer;
 
         // Spawn workers on the taskpool to prepare and update meshes in parallel.
@@ -2069,7 +2073,7 @@ pub fn collect_meshes_for_gpu_building(
                                         skin_uniforms,
                                         frame_count,
                                     ) {
-                                        Some(prepared) => {
+                                        Some(mut prepared) => {
                                             if let Some(render_mesh_instance) =
                                                 render_mesh_instances.get(&entity)
                                                 && prepared.render_layers
@@ -2090,6 +2094,15 @@ pub fn collect_meshes_for_gpu_building(
                                                 let current_uniform_index = render_mesh_instance
                                                     .gpu_specific
                                                     .current_uniform_index();
+
+                                                let previous_mesh_input_uniform =
+                                                    current_input_buffer
+                                                        .get_unchecked(current_uniform_index);
+                                                let previous_input_index = previous_input_buffer
+                                                    .push(previous_mesh_input_uniform);
+                                                prepared.mesh_input_uniform.previous_input_index =
+                                                    previous_input_index;
+
                                                 current_input_buffer.set(
                                                     current_uniform_index,
                                                     prepared.mesh_input_uniform,
@@ -2153,6 +2166,7 @@ pub fn collect_meshes_for_gpu_building(
         let entity = batch;
         meshes_to_reextract_next_frame.insert(entity);
     }
+    previous_input_buffer.truncate();
     // Buffers can't be empty. Make sure there's something in the previous input buffer.
     previous_input_buffer.ensure_nonempty();
 }

crates/bevy_render/src/batching/gpu_preprocessing.rs

@@ -1,6 +1,11 @@
 //! Batching functionality when GPU preprocessing is in use.
 
-use core::{any::TypeId, marker::PhantomData, mem};
+use core::{
+    any::TypeId,
+    marker::PhantomData,
+    mem,
+    sync::atomic::{AtomicU32, Ordering},
+};
 
 use bevy_app::{App, Plugin};
 use bevy_derive::{Deref, DerefMut};
@@ -285,6 +290,9 @@ where
     /// When adding data, we preferentially overwrite these slots first before
     /// growing the buffer itself.
     free_uniform_indices: Vec<u32>,
+
+    /// The number of elements pushed since the last [`Self::reserve`].
+    atomic_len: AtomicU32,
 }
 
 impl<BDI> InstanceInputUniformBuffer<BDI>
@@ -296,13 +304,15 @@ where
         InstanceInputUniformBuffer {
             buffer: AtomicRawBufferVec::new(BufferUsages::STORAGE),
             free_uniform_indices: vec![],
+            atomic_len: AtomicU32::new(0),
         }
     }
 
     /// Clears the buffer and entity list out.
     pub fn clear(&mut self) {
         self.buffer.clear();
         self.free_uniform_indices.clear();
+        *self.atomic_len.get_mut() = 0;
     }
 
     /// Returns the [`RawBufferVec`] corresponding to this input uniform buffer.
@@ -359,6 +369,31 @@ where
         self.buffer.set(uniform_index, element);
     }
 
+    /// Pre-allocates capacity for concurrent [`Self::push`] calls.
+    pub fn reserve(&mut self, capacity: u32) {
+        self.buffer.grow(capacity);
+        *self.atomic_len.get_mut() = 0;
+    }
+
+    /// Appends a value and returns its index. Thread-safe.
+    ///
+    /// [`Self::reserve`] must have been called first with sufficient capacity.
+    pub fn push(&self, value: BDI) -> u32 {
+        let index = self.atomic_len.fetch_add(1, Ordering::Relaxed);
+        debug_assert!(
+            (index as usize) < self.buffer.len() as usize,
+            "push exceeded pre-allocated capacity"
+        );
+        self.buffer.set(index, value);
+        index
+    }
+
+    /// Trims the buffer to the number of elements pushed since the last
+    /// [`Self::reserve`].
+    pub fn truncate(&mut self) {
+        self.buffer.truncate(*self.atomic_len.get_mut());
+    }
+
     // Ensures that the buffers are nonempty, which the GPU requires before an
     // upload can take place.
     pub fn ensure_nonempty(&mut self) {

crates/bevy_render/src/render_resource/buffer_vec.rs

@@ -407,7 +407,7 @@ where
         self.buffer.as_ref()
     }
 
-    /// Grows to the buffer by adding default values so that it's at least the
+    /// Grows the buffer by adding default values so that it's at least the
     /// given size.
     ///
     /// If the buffer is already large enough, this method does nothing.
@@ -416,6 +416,13 @@ where
             self.values.resize_with(new_len as usize, T::Blob::default);
         }
     }
+
+    /// Truncates the buffer to the given length.
+    ///
+    /// If the buffer is already shorter, this method does nothing.
+    pub fn truncate(&mut self, len: u32) {
+        self.values.truncate(len as usize);
+    }
 }
 
 /// Like [`RawBufferVec`], but doesn't require that the data type `T` be