Merge pull request #19 from rodolphito/meshlet

Add view instance visibility buffers to gpu scene

crates/bevy_pbr/src/meshlet/cull_meshlets.wgsl

@@ -5,6 +5,7 @@
     meshlet_instance_uniforms,
     meshlet_occlusion,
     view,
+    should_cull_instance,
     get_meshlet_previous_occlusion,
 }
 #ifdef MESHLET_SECOND_CULLING_PASS
@@ -21,9 +22,12 @@
 fn cull_meshlets(@builtin(global_invocation_id) thread_id: vec3<u32>) {
     // Fetch the instanced meshlet data
     if thread_id.x >= arrayLength(&meshlet_thread_meshlet_ids) { return; }
+    let instance_id = meshlet_thread_instance_ids[thread_id.x];
+    if should_cull_instance(instance_id) {
+        return;
+    }
     let meshlet_id = meshlet_thread_meshlet_ids[thread_id.x];
     let bounding_sphere = meshlet_bounding_spheres[meshlet_id];
-    let instance_id = meshlet_thread_instance_ids[thread_id.x];
     let instance_uniform = meshlet_instance_uniforms[instance_id];
     let model = affine_to_square(instance_uniform.model);
     let model_scale = max(length(model[0]), max(length(model[1]), length(model[2])));

crates/bevy_pbr/src/meshlet/gpu_scene.rs

@@ -17,13 +17,14 @@ use bevy_render::{
     render_resource::{binding_types::*, *},
     renderer::{RenderDevice, RenderQueue},
     texture::{CachedTexture, TextureCache},
-    view::{ExtractedView, ViewDepthTexture, ViewUniform, ViewUniforms},
+    view::{ExtractedView, RenderLayers, ViewDepthTexture, ViewUniform, ViewUniforms},
     MainWorld,
 };
 use bevy_transform::components::GlobalTransform;
 use bevy_utils::{default, EntityHashMap, HashMap, HashSet};
 use encase::internal::WriteInto;
 use std::{
+    iter,
     ops::{DerefMut, Range},
     sync::Arc,
 };
@@ -42,6 +43,7 @@ pub fn extract_meshlet_meshes(
             &Handle<MeshletMesh>,
             &GlobalTransform,
             Option<&PreviousGlobalTransform>,
+            Option<&RenderLayers>,
             Has<NotShadowReceiver>,
             Has<NotShadowCaster>,
         )>,
@@ -71,10 +73,17 @@ pub fn extract_meshlet_meshes(
         }
     }
 
-    // TODO: Handle not_shadow_caster
     for (
         instance_index,
-        (instance, handle, transform, previous_transform, not_shadow_receiver, _not_shadow_caster),
+        (
+            instance,
+            handle,
+            transform,
+            previous_transform,
+            render_layers,
+            not_shadow_receiver,
+            not_shadow_caster,
+        ),
     ) in query.iter().enumerate()
     {
         if asset_server.is_managed(handle.id())
@@ -83,7 +92,14 @@ pub fn extract_meshlet_meshes(
             continue;
         }
 
-        gpu_scene.queue_meshlet_mesh_upload(instance, handle, &mut assets, instance_index as u32);
+        gpu_scene.queue_meshlet_mesh_upload(
+            instance,
+            render_layers.cloned().unwrap_or(default()),
+            not_shadow_caster,
+            handle,
+            &mut assets,
+            instance_index as u32,
+        );
 
         let transform = transform.affine();
         let previous_transform = previous_transform.map(|t| t.0).unwrap_or(transform);
@@ -139,7 +155,7 @@ pub fn queue_material_meshlet_meshes<M: Material>(
     // TODO: Ideally we could parallelize this system, both between different materials, and the loop over instances
     let gpu_scene = gpu_scene.deref_mut();
 
-    for (i, instance) in gpu_scene.instances.iter().enumerate() {
+    for (i, (instance, _, _)) in gpu_scene.instances.iter().enumerate() {
         if let Some(material_asset_id) = render_material_instances.get(instance) {
             let material_asset_id = material_asset_id.untyped();
             if let Some(material_id) = gpu_scene.material_id_lookup.get(&material_asset_id) {
@@ -173,7 +189,12 @@ fn upload_storage_buffer<T: ShaderSize + bytemuck::Pod>(
 
 pub fn prepare_meshlet_per_frame_resources(
     mut gpu_scene: ResMut<MeshletGpuScene>,
-    views: Query<(Entity, &ExtractedView, AnyOf<(&Camera3d, &ShadowView)>)>,
+    views: Query<(
+        Entity,
+        &ExtractedView,
+        Option<&RenderLayers>,
+        AnyOf<(&Camera3d, &ShadowView)>,
+    )>,
     mut texture_cache: ResMut<TextureCache>,
     render_queue: Res<RenderQueue>,
     render_device: Res<RenderDevice>,
@@ -230,7 +251,38 @@ pub fn prepare_meshlet_per_frame_resources(
         };
 
     let needed_buffer_size = gpu_scene.scene_meshlet_count.div_ceil(32) as u64 * 4;
-    for (view_entity, view, (_, shadow_view)) in &views {
+    for (view_entity, view, render_layers, (_, shadow_view)) in &views {
+        let instance_visibility = gpu_scene
+            .view_instance_visibility
+            .entry(view_entity)
+            .or_insert_with(|| {
+                let mut buffer = StorageBuffer::default();
+                buffer.set_label(Some("meshlet_view_instance_visibility"));
+                buffer
+            });
+        for (instance_index, (_, layers, not_shadow_caster)) in
+            gpu_scene.instances.iter().enumerate()
+        {
+            // if either the layers don't match the view's layers
+            // or this is a shadow view and the instance is not a shadow caster
+            if !render_layers.unwrap_or(&default()).intersects(layers)
+                || (shadow_view.is_some() && *not_shadow_caster)
+            {
+                // hide the instance for this view
+                let vec = instance_visibility.get_mut();
+                let index = instance_index / 32;
+                let bit = instance_index - index * 32;
+                if vec.len() <= index {
+                    vec.extend(iter::repeat(0).take(index - vec.len() + 1));
+                }
+                vec[index] |= 1 << bit;
+            }
+        }
+
+        upload_storage_buffer(instance_visibility, &render_device, &render_queue);
+
+        let instance_visibility = instance_visibility.buffer().unwrap().clone();
+
         let create_occlusion_buffer = || {
             render_device.create_buffer(&BufferDescriptor {
                 label: Some("meshlet_occlusion_buffer"),
@@ -365,6 +417,7 @@ pub fn prepare_meshlet_per_frame_resources(
             previous_occlusion_buffer,
             occlusion_buffer,
             occlusion_buffer_needs_clearing,
+            instance_visibility,
             visibility_buffer: not_shadow_view
                 .then(|| texture_cache.get(&render_device, visibility_buffer)),
             visibility_buffer_draw_command_buffer_first,
@@ -401,6 +454,9 @@ pub fn prepare_meshlet_view_bind_groups(
             gpu_scene.meshlet_bounding_spheres.binding(),
             gpu_scene.thread_instance_ids.binding().unwrap(),
             gpu_scene.instance_uniforms.binding().unwrap(),
+            gpu_scene.view_instance_visibility[&view_entity]
+                .binding()
+                .unwrap(),
             view_resources.occlusion_buffer.as_entire_binding(),
             gpu_scene.previous_thread_ids.binding().unwrap(),
             view_resources.previous_occlusion_buffer.as_entire_binding(),
@@ -555,8 +611,9 @@ pub struct MeshletGpuScene {
     next_material_id: u32,
     material_id_lookup: HashMap<UntypedAssetId, u32>,
     material_ids_present_in_scene: HashSet<u32>,
-    instances: Vec<Entity>,
+    instances: Vec<(Entity, RenderLayers, bool)>,
     instance_uniforms: StorageBuffer<Vec<MeshUniform>>,
+    view_instance_visibility: EntityHashMap<Entity, StorageBuffer<Vec<u32>>>,
     instance_material_ids: StorageBuffer<Vec<u32>>,
     thread_instance_ids: StorageBuffer<Vec<u32>>,
     thread_meshlet_ids: StorageBuffer<Vec<u32>>,
@@ -600,6 +657,7 @@ impl FromWorld for MeshletGpuScene {
                 buffer.set_label(Some("meshlet_instance_uniforms"));
                 buffer
             },
+            view_instance_visibility: EntityHashMap::default(),
             instance_material_ids: {
                 let mut buffer = StorageBuffer::default();
                 buffer.set_label(Some("meshlet_instance_material_ids"));
@@ -634,6 +692,7 @@ impl FromWorld for MeshletGpuScene {
                         storage_buffer_read_only_sized(false, None),
                         storage_buffer_read_only_sized(false, None),
                         storage_buffer_read_only_sized(false, None),
+                        storage_buffer_read_only_sized(false, None),
                         storage_buffer_sized(false, None),
                         storage_buffer_read_only_sized(false, None),
                         storage_buffer_read_only_sized(false, None),
@@ -725,6 +784,9 @@ impl MeshletGpuScene {
         self.material_id_lookup.clear();
         self.material_ids_present_in_scene.clear();
         self.instances.clear();
+        self.view_instance_visibility
+            .values_mut()
+            .for_each(|b| b.get_mut().clear());
         self.instance_uniforms.get_mut().clear();
         self.instance_material_ids.get_mut().clear();
         self.thread_instance_ids.get_mut().clear();
@@ -740,6 +802,8 @@ impl MeshletGpuScene {
     fn queue_meshlet_mesh_upload(
         &mut self,
         instance: Entity,
+        render_layers: RenderLayers,
+        not_shadow_caster: bool,
         handle: &Handle<MeshletMesh>,
         assets: &mut Assets<MeshletMesh>,
         instance_index: u32,
@@ -779,7 +843,8 @@ impl MeshletGpuScene {
             )
         };
 
-        self.instances.push(instance);
+        self.instances
+            .push((instance, render_layers, not_shadow_caster));
         self.instance_material_ids.get_mut().push(0);
 
         let (buffer_slices, index_count) = self
@@ -864,6 +929,7 @@ pub struct MeshletViewResources {
     previous_occlusion_buffer: Buffer,
     pub occlusion_buffer: Buffer,
     pub occlusion_buffer_needs_clearing: bool,
+    pub instance_visibility: Buffer,
     pub visibility_buffer: Option<CachedTexture>,
     pub visibility_buffer_draw_command_buffer_first: Buffer,
     pub visibility_buffer_draw_command_buffer_second: Buffer,

crates/bevy_pbr/src/meshlet/meshlet_bindings.wgsl

@@ -50,11 +50,18 @@ struct DrawIndexedIndirect {
 @group(0) @binding(1) var<storage, read> meshlet_bounding_spheres: array<MeshletBoundingSphere>;
 @group(0) @binding(2) var<storage, read> meshlet_thread_instance_ids: array<u32>;
 @group(0) @binding(3) var<storage, read> meshlet_instance_uniforms: array<Mesh>;
-@group(0) @binding(4) var<storage, read_write> meshlet_occlusion: array<atomic<u32>>; // packed bool's
-@group(0) @binding(5) var<storage, read> meshlet_previous_thread_ids: array<u32>;
-@group(0) @binding(6) var<storage, read> meshlet_previous_occlusion: array<u32>; // packed bool's
-@group(0) @binding(7) var<uniform> view: View;
-@group(0) @binding(8) var depth_pyramid: texture_2d<f32>;
+@group(0) @binding(4) var<storage, read> meshlet_view_instance_visibility: array<u32>;
+@group(0) @binding(5) var<storage, read_write> meshlet_occlusion: array<atomic<u32>>; // packed bool's
+@group(0) @binding(6) var<storage, read> meshlet_previous_thread_ids: array<u32>;
+@group(0) @binding(7) var<storage, read> meshlet_previous_occlusion: array<u32>; // packed bool's
+@group(0) @binding(8) var<uniform> view: View;
+@group(0) @binding(9) var depth_pyramid: texture_2d<f32>;
+
+fn should_cull_instance(instance_id: u32) -> bool {
+    let bit_offset = instance_id % 32u;
+    let packed_visibility = meshlet_view_instance_visibility[instance_id / 32u];
+    return bool(extractBits(packed_visibility, bit_offset, 1u));
+}
 
 fn get_meshlet_previous_occlusion(thread_id: u32) -> bool {
     let previous_thread_id = meshlet_previous_thread_ids[thread_id];

examples/3d/meshlet.rs

@@ -58,7 +58,7 @@ fn setup(
         },
         cascade_shadow_config: CascadeShadowConfigBuilder {
             num_cascades: 1,
-            maximum_distance: 15.0,
+            maximum_distance: 5.0,
             ..default()
         }
         .build(),