Meshlet texture atomics (#17765)

* Use texture atomics rather than buffer atomics for the visbuffer
(haven't tested perf on a raster-heavy scene yet)
* Unfortunately to clear the visbuffer we now need a compute pass to
clear it. Using wgpu's clear_texture function internally uses a buffer
-> image copy that's insanely expensive. Ideally it should be using
vkCmdClearColorImage, which I've opened an issue for
gfx-rs/wgpu#7090. For now we'll have to stick
with a custom compute pass and all the extra code that brings.
* Faster resolve depth pass by discarding 0 depth pixels instead of
redundantly writing zero (2x faster for big depth textures like shadow
views)

Author: JMS55

crates/bevy_core_pipeline/src/experimental/mip_generation/downsample_depth.wgsl

@@ -1,8 +1,8 @@
 #ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
-@group(0) @binding(0) var<storage, read> mip_0: array<u64>; // Per pixel
+@group(0) @binding(0) var mip_0: texture_storage_2d<r64uint, read>;
 #else
 #ifdef MESHLET
-@group(0) @binding(0) var<storage, read> mip_0: array<u32>; // Per pixel
+@group(0) @binding(0) var mip_0: texture_storage_2d<r32uint, read>;
 #else   // MESHLET
 #ifdef MULTISAMPLE
 @group(0) @binding(0) var mip_0: texture_depth_multisampled_2d;
@@ -24,7 +24,7 @@
 @group(0) @binding(11) var mip_11: texture_storage_2d<r32float, write>;
 @group(0) @binding(12) var mip_12: texture_storage_2d<r32float, write>;
 @group(0) @binding(13) var samplr: sampler;
-struct Constants { max_mip_level: u32, view_width: u32 }
+struct Constants { max_mip_level: u32 }
 var<push_constant> constants: Constants;
 
 /// Generates a hierarchical depth buffer.
@@ -39,7 +39,6 @@ var<workgroup> intermediate_memory: array<array<f32, 16>, 16>;
 @compute
 @workgroup_size(256, 1, 1)
 fn downsample_depth_first(
-    @builtin(num_workgroups) num_workgroups: vec3u,
     @builtin(workgroup_id) workgroup_id: vec3u,
     @builtin(local_invocation_index) local_invocation_index: u32,
 ) {
@@ -309,12 +308,13 @@ fn reduce_load_mip_6(tex: vec2u) -> f32 {
 }
 
 fn load_mip_0(x: u32, y: u32) -> f32 {
-    let i = y * constants.view_width + x;
 #ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
-    return bitcast<f32>(u32(mip_0[i] >> 32u));
+    let visibility = textureLoad(mip_0, vec2(x, y)).r;
+    return bitcast<f32>(u32(visibility >> 32u));
 #else   // MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
 #ifdef MESHLET
-    return bitcast<f32>(mip_0[i]);
+    let visibility = textureLoad(mip_0, vec2(x, y)).r;
+    return bitcast<f32>(visibility);
 #else   // MESHLET
     // Downsample the top level.
 #ifdef MULTISAMPLE

crates/bevy_core_pipeline/src/experimental/mip_generation/mod.rs

@@ -427,7 +427,7 @@ impl SpecializedComputePipeline for DownsampleDepthPipeline {
             layout: vec![self.bind_group_layout.clone()],
             push_constant_ranges: vec![PushConstantRange {
                 stages: ShaderStages::COMPUTE,
-                range: 0..8,
+                range: 0..4,
             }],
             shader: DOWNSAMPLE_DEPTH_SHADER_HANDLE,
             shader_defs,
@@ -627,9 +627,8 @@ impl ViewDepthPyramid {
             timestamp_writes: None,
         });
         downsample_pass.set_pipeline(downsample_depth_first_pipeline);
-        // Pass the mip count and the texture width as push constants, for
-        // simplicity.
-        downsample_pass.set_push_constants(0, bytemuck::cast_slice(&[self.mip_count, view_size.x]));
+        // Pass the mip count as a push constant, for simplicity.
+        downsample_pass.set_push_constants(0, &self.mip_count.to_le_bytes());
         downsample_pass.set_bind_group(0, downsample_depth_bind_group, &[]);
         downsample_pass.dispatch_workgroups(view_size.x.div_ceil(64), view_size.y.div_ceil(64), 1);
 

crates/bevy_pbr/src/meshlet/clear_visibility_buffer.wgsl

@@ -0,0 +1,18 @@
+#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
+@group(0) @binding(0) var meshlet_visibility_buffer: texture_storage_2d<r64uint, write>;
+#else
+@group(0) @binding(0) var meshlet_visibility_buffer: texture_storage_2d<r32uint, write>;
+#endif
+var<push_constant> view_size: vec2<u32>;
+
+@compute
+@workgroup_size(16, 16, 1)
+fn clear_visibility_buffer(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    if any(global_id.xy >= view_size) { return; }
+
+#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
+    textureStore(meshlet_visibility_buffer, global_id.xy, vec4(0lu));
+#else
+    textureStore(meshlet_visibility_buffer, global_id.xy, vec4(0u));
+#endif
+}

crates/bevy_pbr/src/meshlet/meshlet_bindings.wgsl

@@ -100,9 +100,9 @@ fn cluster_is_second_pass_candidate(cluster_id: u32) -> bool {
 @group(0) @binding(6) var<storage, read> meshlet_raster_clusters: array<u32>; // Single object shared between all workgroups
 @group(0) @binding(7) var<storage, read> meshlet_software_raster_cluster_count: u32;
 #ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
-@group(0) @binding(8) var<storage, read_write> meshlet_visibility_buffer: array<atomic<u64>>; // Per pixel
+@group(0) @binding(8) var meshlet_visibility_buffer: texture_storage_2d<r64uint, atomic>;
 #else
-@group(0) @binding(8) var<storage, read_write> meshlet_visibility_buffer: array<atomic<u32>>; // Per pixel
+@group(0) @binding(8) var meshlet_visibility_buffer: texture_storage_2d<r32uint, atomic>;
 #endif
 @group(0) @binding(9) var<uniform> view: View;
 
@@ -149,7 +149,7 @@ fn get_meshlet_vertex_position(meshlet: ptr<function, Meshlet>, vertex_id: u32)
 #endif
 
 #ifdef MESHLET_MESH_MATERIAL_PASS
-@group(1) @binding(0) var<storage, read> meshlet_visibility_buffer: array<u64>; // Per pixel
+@group(1) @binding(0) var meshlet_visibility_buffer: texture_storage_2d<r64uint, read>;
 @group(1) @binding(1) var<storage, read> meshlet_cluster_meshlet_ids: array<u32>; // Per cluster
 @group(1) @binding(2) var<storage, read> meshlets: array<Meshlet>; // Per meshlet
 @group(1) @binding(3) var<storage, read> meshlet_indices: array<u32>; // Many per meshlet

crates/bevy_pbr/src/meshlet/mod.rs

@@ -106,9 +106,9 @@ const MESHLET_MESH_MATERIAL_SHADER_HANDLE: Handle<Shader> =
 /// * Requires preprocessing meshes. See [`MeshletMesh`] for details.
 /// * Limitations on the kinds of materials you can use. See [`MeshletMesh`] for details.
 ///
-/// This plugin requires a fairly recent GPU that supports [`WgpuFeatures::SHADER_INT64_ATOMIC_MIN_MAX`].
+/// This plugin requires a fairly recent GPU that supports [`WgpuFeatures::TEXTURE_INT64_ATOMIC`].
 ///
-/// This plugin currently works only on the Vulkan backend.
+/// This plugin currently works only on the Vulkan and Metal backends.
 ///
 /// This plugin is not compatible with [`Msaa`]. Any camera rendering a [`MeshletMesh`] must have
 /// [`Msaa`] set to [`Msaa::Off`].
@@ -133,7 +133,8 @@ pub struct MeshletPlugin {
 impl MeshletPlugin {
     /// [`WgpuFeatures`] required for this plugin to function.
     pub fn required_wgpu_features() -> WgpuFeatures {
-        WgpuFeatures::SHADER_INT64_ATOMIC_MIN_MAX
+        WgpuFeatures::TEXTURE_INT64_ATOMIC
+            | WgpuFeatures::TEXTURE_ATOMIC
             | WgpuFeatures::SHADER_INT64
             | WgpuFeatures::SUBGROUP
             | WgpuFeatures::DEPTH_CLIP_CONTROL
@@ -151,6 +152,12 @@ impl Plugin for MeshletPlugin {
             std::process::exit(1);
         }
 
+        load_internal_asset!(
+            app,
+            MESHLET_CLEAR_VISIBILITY_BUFFER_SHADER_HANDLE,
+            "clear_visibility_buffer.wgsl",
+            Shader::from_wgsl
+        );
         load_internal_asset!(
             app,
             MESHLET_BINDINGS_SHADER_HANDLE,