Add second SSAO denoise pass

crates/bevy_pbr/src/ssao/gtao.wgsl

@@ -27,7 +27,7 @@ fn load_noise(pixel_coordinates: vec2<i32>) -> vec2<f32> {
     var index = textureLoad(hilbert_index, pixel_coordinates % 64, 0).r;
 
 #ifdef TEMPORAL_NOISE
-    index += 72u * (globals.frame_count % 64u);
+    index += 288u * (globals.frame_count % 64u);
 #endif
 
     // R2 sequence - http://extremelearning.com.au/unreasonable-effectiveness-of-quasirandom-sequences
@@ -158,7 +158,7 @@ fn gtao(@builtin(global_invocation_id) global_id: vec3<u32>) {
         }
     }
     visibility /= slice_count;
-    visibility = pow(visibility, 2.2);
+    visibility *= visibility * visibility;
     visibility = clamp(visibility, 0.03, 1.0);
 
     textureStore(ambient_occlusion, pixel_coordinates, vec4<f32>(visibility, 0.0, 0.0, 0.0));

crates/bevy_pbr/src/ssao/mod.rs

@@ -264,6 +264,11 @@ impl Node for SSAONode {
                 &bind_groups.common_bind_group,
                 &[view_uniform_offset.offset],
             );
+
+            denoise_pass.set_bind_group(0, &bind_groups.denoise1_bind_group, &[]);
+            denoise_pass.dispatch_workgroups((camera_size.x + 7) / 8, (camera_size.y + 7) / 8, 1);
+
+            denoise_pass.set_bind_group(0, &bind_groups.denoise2_bind_group, &[]);
             denoise_pass.dispatch_workgroups((camera_size.x + 7) / 8, (camera_size.y + 7) / 8, 1);
         }
 
@@ -606,8 +611,9 @@ fn extract_ssao_settings(
 #[derive(Component)]
 pub struct ScreenSpaceAmbientOcclusionTextures {
     preprocessed_depth_texture: CachedTexture,
-    ssao_noisy_texture: CachedTexture, // Pre-denoised texture
-    pub screen_space_ambient_occlusion_texture: CachedTexture, // Denoised texture
+    ssao_noisy1_texture: CachedTexture, // Pre-denoised texture
+    ssao_noisy2_texture: CachedTexture, // Denoised first pass texture
+    pub screen_space_ambient_occlusion_texture: CachedTexture, // Denoised second pass texture
     depth_differences_texture: CachedTexture,
 }
 
@@ -618,7 +624,8 @@ fn prepare_ssao_textures(
     views: Query<(Entity, &ExtractedCamera), With<ScreenSpaceAmbientOcclusionSettings>>,
 ) {
     let mut preprocessed_depth_textures = HashMap::default();
-    let mut ssao_noisy_textures = HashMap::default();
+    let mut ssao_noisy1_textures = HashMap::default();
+    let mut ssao_noisy2_textures = HashMap::default();
     let mut ssao_textures = HashMap::default();
     let mut depth_differences_textures = HashMap::default();
     for (entity, camera) in &views {
@@ -644,15 +651,29 @@ fn prepare_ssao_textures(
                 .clone();
 
             let texture_descriptor = TextureDescriptor {
-                label: Some("ssao_noisy_texture"),
+                label: Some("ssao_noisy1_texture"),
+                size,
+                mip_level_count: 1,
+                sample_count: 1,
+                dimension: TextureDimension::D2,
+                format: TextureFormat::R32Float,
+                usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
+            };
+            let ssao_noisy1_texture = ssao_noisy1_textures
+                .entry(camera.target.clone())
+                .or_insert_with(|| texture_cache.get(&render_device, texture_descriptor.clone()))
+                .clone();
+
+            let texture_descriptor = TextureDescriptor {
+                label: Some("ssao_noisy2_texture"),
                 size,
                 mip_level_count: 1,
                 sample_count: 1,
                 dimension: TextureDimension::D2,
                 format: TextureFormat::R32Float,
                 usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
             };
-            let ssao_noisy_texture = ssao_noisy_textures
+            let ssao_noisy2_texture = ssao_noisy2_textures
                 .entry(camera.target.clone())
                 .or_insert_with(|| texture_cache.get(&render_device, texture_descriptor.clone()))
                 .clone();
@@ -689,7 +710,8 @@ fn prepare_ssao_textures(
                 .entity(entity)
                 .insert(ScreenSpaceAmbientOcclusionTextures {
                     preprocessed_depth_texture,
-                    ssao_noisy_texture,
+                    ssao_noisy1_texture,
+                    ssao_noisy2_texture,
                     screen_space_ambient_occlusion_texture: ssao_texture,
                     depth_differences_texture,
                 });
@@ -730,7 +752,8 @@ struct SSAOBindGroups {
     common_bind_group: BindGroup,
     preprocess_depth_bind_group: BindGroup,
     gtao_bind_group: BindGroup,
-    denoise_bind_group: BindGroup,
+    denoise1_bind_group: BindGroup,
+    denoise2_bind_group: BindGroup,
 }
 
 fn queue_ssao_bind_groups(
@@ -875,7 +898,7 @@ fn queue_ssao_bind_groups(
                 BindGroupEntry {
                     binding: 3,
                     resource: BindingResource::TextureView(
-                        &ssao_textures.ssao_noisy_texture.default_view,
+                        &ssao_textures.ssao_noisy1_texture.default_view,
                     ),
                 },
                 BindGroupEntry {
@@ -891,14 +914,39 @@ fn queue_ssao_bind_groups(
             ],
         });
 
-        let denoise_bind_group = render_device.create_bind_group(&BindGroupDescriptor {
-            label: Some("ssao_denoise_bind_group"),
+        let denoise1_bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+            label: Some("ssao_denoise1_bind_group"),
+            layout: &pipelines.denoise_bind_group_layout,
+            entries: &[
+                BindGroupEntry {
+                    binding: 0,
+                    resource: BindingResource::TextureView(
+                        &ssao_textures.ssao_noisy1_texture.default_view,
+                    ),
+                },
+                BindGroupEntry {
+                    binding: 1,
+                    resource: BindingResource::TextureView(
+                        &ssao_textures.depth_differences_texture.default_view,
+                    ),
+                },
+                BindGroupEntry {
+                    binding: 2,
+                    resource: BindingResource::TextureView(
+                        &ssao_textures.ssao_noisy2_texture.default_view,
+                    ),
+                },
+            ],
+        });
+
+        let denoise2_bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+            label: Some("ssao_denoise2_bind_group"),
             layout: &pipelines.denoise_bind_group_layout,
             entries: &[
                 BindGroupEntry {
                     binding: 0,
                     resource: BindingResource::TextureView(
-                        &ssao_textures.ssao_noisy_texture.default_view,
+                        &ssao_textures.ssao_noisy2_texture.default_view,
                     ),
                 },
                 BindGroupEntry {
@@ -922,7 +970,8 @@ fn queue_ssao_bind_groups(
             common_bind_group,
             preprocess_depth_bind_group,
             gtao_bind_group,
-            denoise_bind_group,
+            denoise1_bind_group,
+            denoise2_bind_group,
         });
     }
 }

crates/bevy_pbr/src/ssao/preprocess_depth.wgsl

@@ -1,6 +1,6 @@
 // Inputs a depth texture and outputs a MIP-chain of depths.
 //
-// Because GTAO's performance is bound by texture reads, this increases
+// Because SSAO's performance is bound by texture reads, this increases
 // performance over using the full resolution depth for every sample.
 
 // Reference: https://research.nvidia.com/sites/default/files/pubs/2012-06_Scalable-Ambient-Obscurance/McGuire12SAO.pdf, section 2.2