Merge pull request bevyengine#17 from ricky26/meshlet

Optimise extract_meshlet_meshes and prepare_meshlet_per_frame_resources

crates/bevy_pbr/src/meshlet/gpu_scene.rs

@@ -25,6 +25,7 @@ use std::{
  ops::{DerefMut, Range},
  sync::Arc,
 };
+use encase::internal::WriteInto;
 
 /// Create and queue for uploading to the GPU [`MeshUniform`] components for
 /// [`MeshletMesh`] entities, as well as queuing uploads for any new meshlet mesh
@@ -128,6 +129,24 @@ pub fn queue_material_meshlet_meshes<M: Material>(
  }
 }
 
+fn upload_storage_buffer<T: ShaderSize + bytemuck::Pod>(
+ buffer: &mut StorageBuffer<Vec<T>>,
+ render_device: &RenderDevice,
+ render_queue: &RenderQueue,
+) where Vec<T>: WriteInto {
+ let inner = buffer.buffer();
+ let capacity = inner.as_deref().map_or(0, |b| b.size());
+ let size = buffer.get().size().get() as BufferAddress;
+
+ if capacity >= size {
+ let inner = inner.unwrap();
+ let bytes = bytemuck::cast_slice(buffer.get().as_slice());
+ render_queue.write_buffer(inner, 0, bytes);
+ } else {
+ buffer.write_buffer(render_device, render_queue);
+ }
+}
+
 pub fn prepare_meshlet_per_frame_resources(
  mut gpu_scene: ResMut<MeshletGpuScene>,
  views: Query<(Entity, &ExtractedView, AnyOf<(&Camera3d, &ShadowView)>)>,
@@ -144,21 +163,15 @@ pub fn prepare_meshlet_per_frame_resources(
  return;
  }
 
+ let gpu_scene = gpu_scene.as_mut();
+
  gpu_scene
  .instance_uniforms
  .write_buffer(&render_device, &render_queue);
- gpu_scene
- .instance_material_ids
- .write_buffer(&render_device, &render_queue);
- gpu_scene
- .thread_instance_ids
- .write_buffer(&render_device, &render_queue);
- gpu_scene
- .thread_meshlet_ids
- .write_buffer(&render_device, &render_queue);
- gpu_scene
- .previous_thread_ids
- .write_buffer(&render_device, &render_queue);
+ upload_storage_buffer(&mut gpu_scene.instance_material_ids, &*render_device, &*render_queue);
+ upload_storage_buffer(&mut gpu_scene.thread_instance_ids, &*render_device, &*render_queue);
+ upload_storage_buffer(&mut gpu_scene.thread_meshlet_ids, &*render_device, &*render_queue);
+ upload_storage_buffer(&mut gpu_scene.previous_thread_ids, &*render_device, &*render_queue);
 
  let needed_buffer_size = 4 * gpu_scene.scene_index_count;
  let visibility_buffer_draw_index_buffer =
@@ -745,18 +758,16 @@ impl MeshletGpuScene {
  .previous_thread_id_starts
  .entry((instance, handle.id()))
  .or_insert((0, true));
+ let previous_thread_ids = if previous_thread_id_start.1 {
+ 0..(meshlets_slice.len() as u32)
+ } else {
+ let start = previous_thread_id_start.0;
+ start..(meshlets_slice.len() as u32 + start)
+ };
 
- for (i, meshlet_index) in meshlets_slice.into_iter().enumerate() {
- self.thread_instance_ids.get_mut().push(instance_index);
- self.thread_meshlet_ids.get_mut().push(meshlet_index);
- self.previous_thread_ids
- .get_mut()
- .push(if previous_thread_id_start.1 {
- 0
- } else {
- previous_thread_id_start.0 + i as u32
- });
- }
+ self.thread_instance_ids.get_mut().extend(std::iter::repeat(instance_index).take(meshlets_slice.len()));
+ self.thread_meshlet_ids.get_mut().extend(meshlets_slice);
+ self.previous_thread_ids.get_mut().extend(previous_thread_ids);
 
  *previous_thread_id_start = (current_thread_id_start, true);
  }