Fix AABB projection for meshlet occlusion culling

crates/bevy_pbr/src/meshlet/cull_meshlets.wgsl

@@ -9,7 +9,7 @@
     get_meshlet_previous_occlusion,
 }
 #ifdef MESHLET_SECOND_CULLING_PASS
-#import bevy_pbr::meshlet_bindings::depth_pyramid
+#import bevy_pbr::meshlet_bindings::{depth_pyramid, aabb_uv_to_viewport}
 #endif
 #import bevy_render::maths::affine3_to_square
 
@@ -55,10 +55,8 @@ fn cull_meshlets(@builtin(global_invocation_id) cluster_id: vec3<u32>) {
         let bounding_sphere_center_view_space = (view.inverse_view * vec4(bounding_sphere_center.xyz, 1.0)).xyz;
         let aabb = project_view_space_sphere_to_screen_space_aabb(bounding_sphere_center_view_space, bounding_sphere_radius);
 
-        // Halve the AABB size because the first depth mip resampling pass cut the full screen resolution into a power of two conservatively
-        let depth_pyramid_size_mip_0 = vec2<f32>(textureDimensions(depth_pyramid, 0)) * 0.5;
-        let width = (aabb.z - aabb.x) * depth_pyramid_size_mip_0.x;
-        let height = (aabb.w - aabb.y) * depth_pyramid_size_mip_0.y;
+        let width = (aabb.z - aabb.x) * aabb_uv_to_viewport.x;
+        let height = (aabb.w - aabb.y) * aabb_uv_to_viewport.y;
         let depth_level = max(0, i32(ceil(log2(max(width, height))))); // TODO: Naga doesn't like this being a u32
         let depth_pyramid_size = vec2<f32>(textureDimensions(depth_pyramid, depth_level));
         let aabb_top_left = vec2<u32>(aabb.xy * depth_pyramid_size);

crates/bevy_pbr/src/meshlet/gpu_scene.rs

@@ -352,7 +352,7 @@ pub fn prepare_meshlet_per_frame_resources(
             });
 
         let depth_size = Extent3d {
-            // If not a power of 2, round down to the nearest power of 2 to ensure depth is conservative
+            // Round down to the nearest power of 2 to ensure depth is conservative
             width: previous_power_of_2(view.viewport.z),
             height: previous_power_of_2(view.viewport.w),
             depth_or_array_layers: 1,

crates/bevy_pbr/src/meshlet/material_draw_prepare.rs

@@ -134,8 +134,6 @@ pub fn prepare_material_meshlet_meshes_main_opaque_pass<M: Material>(
             view_key |= MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION;
         }
 
-        // TODO: Lightmaps
-
         view_key |= MeshPipelineKey::from_primitive_topology(PrimitiveTopology::TriangleList);
 
         for material_id in render_material_instances.values() {

crates/bevy_pbr/src/meshlet/meshlet_bindings.wgsl

@@ -55,6 +55,7 @@ struct DrawIndirectArgs {
 @group(0) @binding(7) var<storage, read> meshlet_previous_occlusion: array<u32>; // 1 bit per cluster (instance of a meshlet), packed as a bitmask
 @group(0) @binding(8) var<uniform> view: View;
 @group(0) @binding(9) var depth_pyramid: texture_2d<f32>; // Generated from the first raster pass (unused in the first pass but still bound)
+var<push_constant> aabb_uv_to_viewport: vec2<f32>; // Multiplier to convert UV-space AABB to view-space (unused in the first pass but still bound)
 
 fn should_cull_instance(instance_id: u32) -> bool {
     let bit_offset = instance_id % 32u;

crates/bevy_pbr/src/meshlet/mod.rs

@@ -75,6 +75,8 @@ use bevy_ecs::{
 use bevy_render::{
     render_graph::{RenderGraphApp, ViewNodeRunner},
     render_resource::{Shader, TextureUsages},
+    renderer::RenderDevice,
+    settings::WgpuFeatures,
     view::{prepare_view_targets, InheritedVisibility, Msaa, ViewVisibility, Visibility},
     ExtractSchedule, Render, RenderApp, RenderSet,
 };
@@ -102,7 +104,7 @@ const MESHLET_MESH_MATERIAL_SHADER_HANDLE: Handle<Shader> =
 ///
 /// This plugin is not compatible with [`Msaa`], and adding this plugin will disable it.
 ///
-/// This plugin does not work on the WebGL2 backend.
+/// This plugin does not work on WASM.
 ///
 /// ![A render of the Stanford dragon as a `MeshletMesh`](https://raw.githubusercontent.com/bevyengine/bevy/meshlet/crates/bevy_pbr/src/meshlet/meshlet_preview.png)
 pub struct MeshletPlugin;
@@ -168,6 +170,15 @@ impl Plugin for MeshletPlugin {
             return;
         };
 
+        if !render_app
+            .world()
+            .resource::<RenderDevice>()
+            .features()
+            .contains(WgpuFeatures::PUSH_CONSTANTS)
+        {
+            return;
+        }
+
         render_app
             .add_render_graph_node::<MeshletVisibilityBufferRasterPassNode>(
                 Core3d,

crates/bevy_pbr/src/meshlet/pipelines.rs

@@ -46,7 +46,10 @@ impl FromWorld for MeshletPipelines {
             cull_first: pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
                 label: Some("meshlet_culling_first_pipeline".into()),
                 layout: vec![cull_layout.clone()],
-                push_constant_ranges: vec![],
+                push_constant_ranges: vec![PushConstantRange {
+                    stages: ShaderStages::COMPUTE,
+                    range: 0..8,
+                }],
                 shader: MESHLET_CULLING_SHADER_HANDLE,
                 shader_defs: vec!["MESHLET_CULLING_PASS".into()],
                 entry_point: "cull_meshlets".into(),
@@ -55,7 +58,10 @@ impl FromWorld for MeshletPipelines {
             cull_second: pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
                 label: Some("meshlet_culling_second_pipeline".into()),
                 layout: vec![cull_layout],
-                push_constant_ranges: vec![],
+                push_constant_ranges: vec![PushConstantRange {
+                    stages: ShaderStages::COMPUTE,
+                    range: 0..8,
+                }],
                 shader: MESHLET_CULLING_SHADER_HANDLE,
                 shader_defs: vec![
                     "MESHLET_CULLING_PASS".into(),

crates/bevy_pbr/src/meshlet/visibility_buffer_raster_node.rs

@@ -8,12 +8,13 @@ use bevy_ecs::{
     query::QueryState,
     world::{FromWorld, World},
 };
+use bevy_math::Vec2;
 use bevy_render::{
     camera::ExtractedCamera,
     render_graph::{Node, NodeRunError, RenderGraphContext},
     render_resource::*,
     renderer::RenderContext,
-    view::{ViewDepthTexture, ViewUniformOffset},
+    view::{ExtractedView, ViewDepthTexture, ViewUniformOffset},
 };
 
 /// Rasterize meshlets into a depth buffer, and optional visibility buffer + material depth buffer for shading passes.
@@ -29,6 +30,7 @@ pub struct MeshletVisibilityBufferRasterPassNode {
     view_light_query: QueryState<(
         &'static ShadowView,
         &'static LightEntity,
+        &'static ExtractedView,
         &'static ViewUniformOffset,
         &'static MeshletViewBindGroups,
         &'static MeshletViewResources,
@@ -88,6 +90,9 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
             .cbrt()
             .ceil() as u32;
 
+        let aabb_uv_to_viewport =
+            calculate_aabb_uv_to_viewport(meshlet_view_resources, view_depth.texture.size());
+
         render_context
             .command_encoder()
             .push_debug_group("meshlet_visibility_buffer_raster_pass");
@@ -105,6 +110,7 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
             view_offset,
             culling_first_pipeline,
             culling_workgroups,
+            aabb_uv_to_viewport,
         );
         write_index_buffer_pass(
             "meshlet_write_index_buffer_first_pass",
@@ -142,6 +148,7 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
             view_offset,
             culling_second_pipeline,
             culling_workgroups,
+            aabb_uv_to_viewport,
         );
         write_index_buffer_pass(
             "meshlet_write_index_buffer_second_pass",
@@ -174,6 +181,7 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
             let Ok((
                 shadow_view,
                 light_type,
+                extracted_view,
                 view_offset,
                 meshlet_view_bind_groups,
                 meshlet_view_resources,
@@ -182,6 +190,15 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
                 continue;
             };
 
+            let aabb_uv_to_viewport = calculate_aabb_uv_to_viewport(
+                meshlet_view_resources,
+                Extent3d {
+                    width: extracted_view.viewport.z,
+                    height: extracted_view.viewport.w,
+                    depth_or_array_layers: 1,
+                },
+            );
+
             let shadow_visibility_buffer_pipeline = match light_type {
                 LightEntity::Directional { .. } => visibility_buffer_raster_depth_only_clamp_ortho,
                 _ => visibility_buffer_raster_depth_only_pipeline,
@@ -205,6 +222,7 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
                 view_offset,
                 culling_first_pipeline,
                 culling_workgroups,
+                aabb_uv_to_viewport,
             );
             write_index_buffer_pass(
                 "meshlet_write_index_buffer_first_pass",
@@ -242,6 +260,7 @@ impl Node for MeshletVisibilityBufferRasterPassNode {
                 view_offset,
                 culling_second_pipeline,
                 culling_workgroups,
+                aabb_uv_to_viewport,
             );
             write_index_buffer_pass(
                 "meshlet_write_index_buffer_second_pass",
@@ -275,6 +294,7 @@ fn cull_pass(
     view_offset: &ViewUniformOffset,
     culling_pipeline: &ComputePipeline,
     culling_workgroups: u32,
+    aabb_uv_to_viewport: Vec2,
 ) {
     let command_encoder = render_context.command_encoder();
     let mut cull_pass = command_encoder.begin_compute_pass(&ComputePassDescriptor {
@@ -283,6 +303,7 @@ fn cull_pass(
     });
     cull_pass.set_bind_group(0, &meshlet_view_bind_groups.culling, &[view_offset.offset]);
     cull_pass.set_pipeline(culling_pipeline);
+    cull_pass.set_push_constants(0, bytemuck::cast_slice(&aabb_uv_to_viewport.to_array()));
     cull_pass.dispatch_workgroups(culling_workgroups, 1, 1);
 }
 
@@ -446,3 +467,20 @@ fn copy_material_depth_pass(
         copy_pass.draw(0..3, 0..1);
     }
 }
+
+/// Meshlet AABBs are calculated against the original view, and need to be scaled to match the depth pyramid.
+///
+/// Calculates the depth pyramid size scaled by how much the first depth mip resampling cut the view depth resolution by.
+/// E.g. `view_depth = 200, depth_pyramid = 100, aabb_uv_to_viewport = 100 * (100 / 200) = 50`.
+fn calculate_aabb_uv_to_viewport(
+    meshlet_view_resources: &MeshletViewResources,
+    view_depth_size: Extent3d,
+) -> Vec2 {
+    let depth_pyramid_mip_0_size = meshlet_view_resources.depth_pyramid.texture.size();
+    let depth_pyramid_mip_0_size = Vec2::new(
+        depth_pyramid_mip_0_size.width as f32,
+        depth_pyramid_mip_0_size.height as f32,
+    );
+    let view_depth_size = Vec2::new(view_depth_size.width as f32, view_depth_size.height as f32);
+    (depth_pyramid_mip_0_size) * (depth_pyramid_mip_0_size / view_depth_size)
+}