address feedback

Author: atlv24

crates/bevy_pbr/src/meshlet/asset.rs

@@ -169,6 +169,11 @@ impl AssetSaver for MeshletMeshSaver {
             .write_all(&MESHLET_MESH_ASSET_VERSION.to_le_bytes())
             .await?;
 
+        writer.write_all(bytemuck::bytes_of(&asset.aabb)).await?;
+        writer
+            .write_all(bytemuck::bytes_of(&asset.bvh_depth))
+            .await?;
+
         // Compress and write asset data
         let mut writer = FrameEncoder::new(AsyncWriteSyncAdapter(writer));
         write_slice(&asset.vertex_positions, &mut writer)?;
@@ -178,8 +183,6 @@ impl AssetSaver for MeshletMeshSaver {
         write_slice(&asset.bvh, &mut writer)?;
         write_slice(&asset.meshlets, &mut writer)?;
         write_slice(&asset.meshlet_cull_data, &mut writer)?;
-        writer.write_all(bytemuck::bytes_of(&asset.aabb))?;
-        writer.write_all(bytemuck::bytes_of(&asset.bvh_depth))?;
         writer.finish()?;
 
         Ok(())
@@ -212,6 +215,13 @@ impl AssetLoader for MeshletMeshLoader {
             return Err(MeshletMeshSaveOrLoadError::WrongVersion { found: version });
         }
 
+        let mut bytes = [0u8; 24];
+        reader.read_exact(&mut bytes).await?;
+        let aabb = bytemuck::cast(bytes);
+        let mut bytes = [0u8; 4];
+        reader.read_exact(&mut bytes).await?;
+        let bvh_depth = u32::from_le_bytes(bytes);
+
         // Load and decompress asset data
         let reader = &mut FrameDecoder::new(AsyncReadSyncAdapter(reader));
         let vertex_positions = read_slice(reader)?;
@@ -221,12 +231,6 @@ impl AssetLoader for MeshletMeshLoader {
         let bvh = read_slice(reader)?;
         let meshlets = read_slice(reader)?;
         let meshlet_cull_data = read_slice(reader)?;
-        let mut bytes = [0u8; 24];
-        reader.read_exact(&mut bytes)?;
-        let aabb = bytemuck::cast(bytes);
-        let mut bytes = [0u8; 4];
-        reader.read_exact(&mut bytes)?;
-        let bvh_depth = u32::from_le_bytes(bytes);
 
         Ok(MeshletMesh {
             vertex_positions,

crates/bevy_pbr/src/meshlet/cull_bvh.wgsl

@@ -25,7 +25,7 @@
 }
 
 @compute
-@workgroup_size(128, 1, 1) // 8 threads per node
+@workgroup_size(128, 1, 1) // 8 threads per node, 16 nodes per workgroup
 fn cull_bvh(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
     // Calculate the queue ID for this thread
     let dispatch_id = global_invocation_id.x;

crates/bevy_pbr/src/meshlet/mod.rs

@@ -147,10 +147,10 @@ impl Plugin for MeshletPlugin {
             std::process::exit(1);
         }
 
-        embedded_asset!(app, "clear_visibility_buffer.wgsl");
         load_shader_library!(app, "meshlet_bindings.wgsl");
         load_shader_library!(app, "visibility_buffer_resolve.wgsl");
         load_shader_library!(app, "meshlet_cull_shared.wgsl");
+        embedded_asset!(app, "clear_visibility_buffer.wgsl");
         embedded_asset!(app, "cull_instances.wgsl");
         embedded_asset!(app, "cull_bvh.wgsl");
         embedded_asset!(app, "cull_clusters.wgsl");

crates/bevy_pbr/src/meshlet/persistent_buffer_impls.rs

@@ -1,4 +1,4 @@
-use crate::meshlet::asset::{BvhNode, MeshletAabbErrorOffset, MeshletCullData};
+use crate::meshlet::asset::{BvhNode, MeshletCullData};
 
 use super::{asset::Meshlet, persistent_buffer::PersistentGpuBufferable};
 use alloc::sync::Arc;
@@ -18,27 +18,31 @@ impl PersistentGpuBufferable for Arc<[BvhNode]> {
         buffer_slice: &mut [u8],
         buffer_offset: BufferAddress,
     ) {
+        const SIZE: usize = size_of::<BvhNode>();
+        for (i, &node) in self.iter().enumerate() {
+            let bytes: [u8; SIZE] =
+                bytemuck::cast(node.offset_aabbs(base_meshlet_index, buffer_offset));
+            buffer_slice[i * SIZE..(i + 1) * SIZE].copy_from_slice(&bytes);
+        }
+    }
+}
+
+impl BvhNode {
+    fn offset_aabbs(mut self, base_meshlet_index: u32, buffer_offset: BufferAddress) -> Self {
         let size = size_of::<BvhNode>();
         let base_bvh_node_index = (buffer_offset / size as u64) as u32;
-        for (i, &node) in self.iter().enumerate() {
-            let bytes = bytemuck::cast::<_, [u8; size_of::<BvhNode>()]>(BvhNode {
-                aabbs: core::array::from_fn(|i| {
-                    let aabb = node.aabbs[i];
-                    MeshletAabbErrorOffset {
-                        child_offset: aabb.child_offset
-                            + if node.child_counts[i] == u8::MAX {
-                                base_bvh_node_index
-                            } else {
-                                base_meshlet_index
-                            },
-                        ..aabb
-                    }
-                }),
-                ..node
-            });
-            let i = i * size;
-            buffer_slice[i..(i + size)].clone_from_slice(&bytes);
+        for i in 0..self.aabbs.len() {
+            self.aabbs[i].child_offset += if self.child_is_bvh_node(i) {
+                base_bvh_node_index
+            } else {
+                base_meshlet_index
+            };
         }
+        self
+    }
+
+    fn child_is_bvh_node(&self, i: usize) -> bool {
+        self.child_counts[i] == u8::MAX
     }
 }
 

crates/bevy_pbr/src/meshlet/resource_manager.rs

@@ -73,6 +73,7 @@ impl ResourceManager {
     pub fn new(cluster_buffer_slots: u32, render_device: &RenderDevice) -> Self {
         let needs_dispatch_remap =
             cluster_buffer_slots > render_device.limits().max_compute_workgroups_per_dimension;
+        // The IDs are a (u32, u32) of instance and index.
         let cull_queue_size = 2 * cluster_buffer_slots as u64 * size_of::<u32>() as u64;
 
         Self {