Screen Space Ambient Occlusion (SSAO) MVP (bevyengine#7402)

![image](https://github.com/bevyengine/bevy/assets/47158642/dbb62645-f639-4f2b-b84b-26fd915c186d)

# Objective

- Add Screen space ambient occlusion (SSAO). SSAO approximates
small-scale, local occlusion of _indirect_ diffuse light between
objects. SSAO does not apply to direct lighting, such as point or
directional lights.
- This darkens creases, e.g. on staircases, and gives nice contact
shadows where objects meet, giving entities a more "grounded" feel.
- Closes bevyengine#3632.

## Solution

- Implement the GTAO algorithm.
-
https://www.activision.com/cdn/research/Practical_Real_Time_Strategies_for_Accurate_Indirect_Occlusion_NEW%20VERSION_COLOR.pdf
-
https://blog.selfshadow.com/publications/s2016-shading-course/activision/s2016_pbs_activision_occlusion.pdf
- Source code heavily based on [Intel's
XeGTAO](https://github.com/GameTechDev/XeGTAO/blob/0d177ce06bfa642f64d8af4de1197ad1bcb862d4/Source/Rendering/Shaders/XeGTAO.hlsli).
- Add an SSAO bevy example.

## Algorithm Overview
* Run a depth and normal prepass
* Create downscaled mips of the depth texture (preprocess_depths pass)
* GTAO pass - for each pixel, take several random samples from the
depth+normal buffers, reconstruct world position, raytrace in screen
space to estimate occlusion. Rather then doing completely random samples
on a hemisphere, you choose random _slices_ of the hemisphere, and then
can analytically compute the full occlusion of that slice. Also compute
edges based on depth differences here.
* Spatial denoise pass - bilateral blur, using edge detection to not
blur over edges. This is the final SSAO result.
* Main pass - if SSAO exists, sample the SSAO texture, and set occlusion
to be the minimum of ssao/material occlusion. This then feeds into the
rest of the PBR shader as normal.

---

## Future Improvements
- Maybe remove the low quality preset for now (too noisy)
- WebGPU fallback (see below)
- Faster depth->world position (see reverted code)
- Bent normals 
- Try interleaved gradient noise or spatiotemporal blue noise
- Replace the spatial denoiser with a combined spatial+temporal denoiser
- Render at half resolution and use a bilateral upsample
- Better multibounce approximation
(https://drive.google.com/file/d/1SyagcEVplIm2KkRD3WQYSO9O0Iyi1hfy/view)

## Far-Future Performance Improvements
- F16 math (missing naga-wgsl support
https://github.com/gfx-rs/naga/issues/1884)
- Faster coordinate space conversion for normals
- Faster depth mipchain creation
(https://github.com/GPUOpen-Effects/FidelityFX-SPD) (wgpu/naga does not
currently support subgroup ops)
- Deinterleaved SSAO for better cache efficiency
(https://developer.nvidia.com/sites/default/files/akamai/gameworks/samples/DeinterleavedTexturing.pdf)

## Other Interesting Papers
- Visibility bitmask
(https://link.springer.com/article/10.1007/s00371-022-02703-y,
https://cdrinmatane.github.io/posts/cgspotlight-slides/)
- Screen space diffuse lighting
(https://github.com/Patapom/GodComplex/blob/master/Tests/TestHBIL/2018%20Mayaux%20-%20Horizon-Based%20Indirect%20Lighting%20(HBIL).pdf)

## Platform Support
* SSAO currently does not work on DirectX12 due to issues with wgpu and
naga:
  * gfx-rs/wgpu#3798
  * gfx-rs/naga#2353
* SSAO currently does not work on WebGPU because r16float is not a valid
storage texture format
https://gpuweb.github.io/gpuweb/wgsl/#storage-texel-formats. We can fix
this with a fallback to r32float.

---

## Changelog

- Added ScreenSpaceAmbientOcclusionSettings,
ScreenSpaceAmbientOcclusionQualityLevel, and
ScreenSpaceAmbientOcclusionBundle

---------

Co-authored-by: IceSentry <c.giguere42@gmail.com>
Co-authored-by: IceSentry <IceSentry@users.noreply.github.com>
Co-authored-by: Daniel Chia <danstryder@gmail.com>
Co-authored-by: Elabajaba <Elabajaba@users.noreply.github.com>
Co-authored-by: Robert Swain <robert.swain@gmail.com>
Co-authored-by: robtfm <50659922+robtfm@users.noreply.github.com>
Co-authored-by: Brandon Dyer <brandondyer64@gmail.com>
Co-authored-by: Edgar Geier <geieredgar@gmail.com>
Co-authored-by: Nicola Papale <nicopap@users.noreply.github.com>
Co-authored-by: Carter Anderson <mcanders1@gmail.com>

Cargo.toml

@@ -517,6 +517,16 @@ description = "Create a custom material to draw 3d lines"
 category = "3D Rendering"
 wasm = true
 
+[[example]]
+name = "ssao"
+path = "examples/3d/ssao.rs"
+
+[package.metadata.example.ssao]
+name = "Screen Space Ambient Occlusion"
+description = "A scene showcasing screen space ambient occlusion"
+category = "3D Rendering"
+wasm = false
+
 [[example]]
 name = "spotlight"
 path = "examples/3d/spotlight.rs"

crates/bevy_core_pipeline/src/prepass/node.rs

@@ -94,7 +94,6 @@ impl ViewNode for PrepassNode {
                     stencil_ops: None,
                 }),
             });
-
             if let Some(viewport) = camera.viewport.as_ref() {
                 render_pass.set_camera_viewport(viewport);
             }

crates/bevy_pbr/src/lib.rs

@@ -12,6 +12,7 @@ mod parallax;
 mod pbr_material;
 mod prepass;
 mod render;
+mod ssao;
 
 pub use alpha::*;
 pub use bundle::*;
@@ -23,6 +24,7 @@ pub use parallax::*;
 pub use pbr_material::*;
 pub use prepass::*;
 pub use render::*;
+pub use ssao::*;
 
 pub mod prelude {
     #[doc(hidden)]
@@ -38,6 +40,7 @@ pub mod prelude {
         material::{Material, MaterialPlugin},
         parallax::ParallaxMappingMethod,
         pbr_material::StandardMaterial,
+        ssao::ScreenSpaceAmbientOcclusionPlugin,
     };
 }
 
@@ -184,6 +187,7 @@ impl Plugin for PbrPlugin {
                 prepass_enabled: self.prepass_enabled,
                 ..Default::default()
             })
+            .add_plugin(ScreenSpaceAmbientOcclusionPlugin)
             .add_plugin(EnvironmentMapPlugin)
             .init_resource::<AmbientLight>()
             .init_resource::<GlobalVisiblePointLights>()

crates/bevy_pbr/src/material.rs

@@ -1,7 +1,7 @@
 use crate::{
     render, AlphaMode, DrawMesh, DrawPrepass, EnvironmentMapLight, MeshPipeline, MeshPipelineKey,
-    MeshUniform, PrepassPipelinePlugin, PrepassPlugin, RenderLightSystems, SetMeshBindGroup,
-    SetMeshViewBindGroup, Shadow,
+    MeshUniform, PrepassPipelinePlugin, PrepassPlugin, RenderLightSystems,
+    ScreenSpaceAmbientOcclusionSettings, SetMeshBindGroup, SetMeshViewBindGroup, Shadow,
 };
 use bevy_app::{App, Plugin};
 use bevy_asset::{AddAsset, AssetEvent, AssetServer, Assets, Handle};
@@ -387,6 +387,7 @@ pub fn queue_material_meshes<M: Material>(
         Option<&Tonemapping>,
         Option<&DebandDither>,
         Option<&EnvironmentMapLight>,
+        Option<&ScreenSpaceAmbientOcclusionSettings>,
         Option<&NormalPrepass>,
         Option<&TemporalAntiAliasSettings>,
         &mut RenderPhase<Opaque3d>,
@@ -402,6 +403,7 @@ pub fn queue_material_meshes<M: Material>(
         tonemapping,
         dither,
         environment_map,
+        ssao,
         normal_prepass,
         taa_settings,
         mut opaque_phase,
@@ -455,6 +457,10 @@ pub fn queue_material_meshes<M: Material>(
             }
         }
 
+        if ssao.is_some() {
+            view_key |= MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION;
+        }
+
         let rangefinder = view.rangefinder3d();
         for visible_entity in &visible_entities.entities {
             if let Ok((material_handle, mesh_handle, mesh_uniform)) =

crates/bevy_pbr/src/render/mesh.rs

@@ -1,9 +1,9 @@
 use crate::{
     environment_map, prepass, EnvironmentMapLight, FogMeta, GlobalLightMeta, GpuFog, GpuLights,
     GpuPointLights, LightMeta, NotShadowCaster, NotShadowReceiver, PreviousGlobalTransform,
-    ShadowSamplers, ViewClusterBindings, ViewFogUniformOffset, ViewLightsUniformOffset,
-    ViewShadowBindings, CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT, MAX_CASCADES_PER_LIGHT,
-    MAX_DIRECTIONAL_LIGHTS,
+    ScreenSpaceAmbientOcclusionTextures, ShadowSamplers, ViewClusterBindings, ViewFogUniformOffset,
+    ViewLightsUniformOffset, ViewShadowBindings, CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT,
+    MAX_CASCADES_PER_LIGHT, MAX_DIRECTIONAL_LIGHTS,
 };
 use bevy_app::Plugin;
 use bevy_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
@@ -433,22 +433,33 @@ impl FromWorld for MeshPipeline {
                     },
                     count: None,
                 },
+                // Screen space ambient occlusion texture
+                BindGroupLayoutEntry {
+                    binding: 11,
+                    visibility: ShaderStages::FRAGMENT,
+                    ty: BindingType::Texture {
+                        multisampled: false,
+                        sample_type: TextureSampleType::Float { filterable: false },
+                        view_dimension: TextureViewDimension::D2,
+                    },
+                    count: None,
+                },
             ];
 
             // EnvironmentMapLight
             let environment_map_entries =
-                environment_map::get_bind_group_layout_entries([11, 12, 13]);
+                environment_map::get_bind_group_layout_entries([12, 13, 14]);
             entries.extend_from_slice(&environment_map_entries);
 
             // Tonemapping
-            let tonemapping_lut_entries = get_lut_bind_group_layout_entries([14, 15]);
+            let tonemapping_lut_entries = get_lut_bind_group_layout_entries([15, 16]);
             entries.extend_from_slice(&tonemapping_lut_entries);
 
             if cfg!(any(not(feature = "webgl"), not(target_arch = "wasm32")))
                 || (cfg!(all(feature = "webgl", target_arch = "wasm32")) && !multisampled)
             {
                 entries.extend_from_slice(&prepass::get_bind_group_layout_entries(
-                    [16, 17, 18],
+                    [17, 18, 19],
                     multisampled,
                 ));
             }
@@ -586,8 +597,9 @@ bitflags::bitflags! {
         const MAY_DISCARD                       = (1 << 6); // Guards shader codepaths that may discard, allowing early depth tests in most cases
                                                             // See: https://www.khronos.org/opengl/wiki/Early_Fragment_Test
         const ENVIRONMENT_MAP                   = (1 << 7);
-        const DEPTH_CLAMP_ORTHO                 = (1 << 8);
-        const TAA                               = (1 << 9);
+        const SCREEN_SPACE_AMBIENT_OCCLUSION    = (1 << 8);
+        const DEPTH_CLAMP_ORTHO                 = (1 << 9);
+        const TAA                               = (1 << 10);
         const BLEND_RESERVED_BITS               = Self::BLEND_MASK_BITS << Self::BLEND_SHIFT_BITS; // ← Bitmask reserving bits for the blend state
         const BLEND_OPAQUE                      = (0 << Self::BLEND_SHIFT_BITS);                   // ← Values are just sequential within the mask, and can range from 0 to 3
         const BLEND_PREMULTIPLIED_ALPHA         = (1 << Self::BLEND_SHIFT_BITS);                   //
@@ -727,6 +739,10 @@ impl SpecializedMeshPipeline for MeshPipeline {
             bind_group_layout.push(self.mesh_layout.clone());
         };
 
+        if key.contains(MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION) {
+            shader_defs.push("SCREEN_SPACE_AMBIENT_OCCLUSION".into());
+        }
+
         let vertex_buffer_layout = layout.get_layout(&vertex_attributes)?;
 
         let (label, blend, depth_write_enabled);
@@ -974,6 +990,7 @@ pub fn queue_mesh_view_bind_groups(
         Entity,
         &ViewShadowBindings,
         &ViewClusterBindings,
+        Option<&ScreenSpaceAmbientOcclusionTextures>,
         Option<&ViewPrepassTextures>,
         Option<&EnvironmentMapLight>,
         &Tonemapping,
@@ -1003,11 +1020,17 @@ pub fn queue_mesh_view_bind_groups(
             entity,
             view_shadow_bindings,
             view_cluster_bindings,
+            ssao_textures,
             prepass_textures,
             environment_map,
             tonemapping,
         ) in &views
         {
+            let fallback_ssao = fallback_images
+                .image_for_samplecount(1)
+                .texture_view
+                .clone();
+
             let layout = if msaa.samples() > 1 {
                 &mesh_pipeline.view_layout_multisampled
             } else {
@@ -1063,18 +1086,26 @@ pub fn queue_mesh_view_bind_groups(
                     binding: 10,
                     resource: fog_binding.clone(),
                 },
+                BindGroupEntry {
+                    binding: 11,
+                    resource: BindingResource::TextureView(
+                        ssao_textures
+                            .map(|t| &t.screen_space_ambient_occlusion_texture.default_view)
+                            .unwrap_or(&fallback_ssao),
+                    ),
+                },
             ];
 
             let env_map = environment_map::get_bindings(
                 environment_map,
                 &images,
                 &fallback_cubemap,
-                [11, 12, 13],
+                [12, 13, 14],
             );
             entries.extend_from_slice(&env_map);
 
             let tonemapping_luts =
-                get_lut_bindings(&images, &tonemapping_luts, tonemapping, [14, 15]);
+                get_lut_bindings(&images, &tonemapping_luts, tonemapping, [15, 16]);
             entries.extend_from_slice(&tonemapping_luts);
 
             // When using WebGL, we can't have a depth texture with multisampling
@@ -1086,7 +1117,7 @@ pub fn queue_mesh_view_bind_groups(
                     &mut fallback_images,
                     &mut fallback_depths,
                     &msaa,
-                    [16, 17, 18],
+                    [17, 18, 19],
                 ));
             }
 

crates/bevy_pbr/src/render/mesh_view_bindings.wgsl

@@ -47,29 +47,32 @@ var<uniform> globals: Globals;
 var<uniform> fog: Fog;
 
 @group(0) @binding(11)
-var environment_map_diffuse: texture_cube<f32>;
+var screen_space_ambient_occlusion_texture: texture_2d<f32>;
+
 @group(0) @binding(12)
-var environment_map_specular: texture_cube<f32>;
+var environment_map_diffuse: texture_cube<f32>;
 @group(0) @binding(13)
+var environment_map_specular: texture_cube<f32>;
+@group(0) @binding(14)
 var environment_map_sampler: sampler;
 
-@group(0) @binding(14)
-var dt_lut_texture: texture_3d<f32>;
 @group(0) @binding(15)
+var dt_lut_texture: texture_3d<f32>;
+@group(0) @binding(16)
 var dt_lut_sampler: sampler;
 
 #ifdef MULTISAMPLED
-@group(0) @binding(16)
-var depth_prepass_texture: texture_depth_multisampled_2d;
 @group(0) @binding(17)
-var normal_prepass_texture: texture_multisampled_2d<f32>;
+var depth_prepass_texture: texture_depth_multisampled_2d;
 @group(0) @binding(18)
+var normal_prepass_texture: texture_multisampled_2d<f32>;
+@group(0) @binding(19)
 var motion_vector_prepass_texture: texture_multisampled_2d<f32>;
 #else
-@group(0) @binding(16)
-var depth_prepass_texture: texture_depth_2d;
 @group(0) @binding(17)
-var normal_prepass_texture: texture_2d<f32>;
+var depth_prepass_texture: texture_depth_2d;
 @group(0) @binding(18)
+var normal_prepass_texture: texture_2d<f32>;
+@group(0) @binding(19)
 var motion_vector_prepass_texture: texture_2d<f32>;
 #endif

crates/bevy_pbr/src/render/pbr.wgsl

@@ -13,6 +13,10 @@
 
 #import bevy_pbr::prepass_utils
 
+#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
+#import bevy_pbr::gtao_utils
+#endif
+
 struct FragmentInput {
     @builtin(front_facing) is_front: bool,
     @builtin(position) frag_coord: vec4<f32>,
@@ -88,12 +92,20 @@ fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
         pbr_input.material.metallic = metallic;
         pbr_input.material.perceptual_roughness = perceptual_roughness;
 
-        var occlusion: f32 = 1.0;
+        // TODO: Split into diffuse/specular occlusion?
+        var occlusion: vec3<f32> = vec3(1.0);
 #ifdef VERTEX_UVS
         if ((material.flags & STANDARD_MATERIAL_FLAGS_OCCLUSION_TEXTURE_BIT) != 0u) {
-            occlusion = textureSample(occlusion_texture, occlusion_sampler, uv).r;
+            occlusion = vec3(textureSample(occlusion_texture, occlusion_sampler, in.uv).r);
         }
 #endif
+#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
+        let ssao = textureLoad(screen_space_ambient_occlusion_texture, vec2<i32>(in.frag_coord.xy), 0i).r;
+        let ssao_multibounce = gtao_multibounce(ssao, pbr_input.material.base_color.rgb);
+        occlusion = min(occlusion, ssao_multibounce);
+#endif
+        pbr_input.occlusion = occlusion;
+
         pbr_input.frag_coord = in.frag_coord;
         pbr_input.world_position = in.world_position;
 

crates/bevy_pbr/src/render/pbr_ambient.wgsl

@@ -10,7 +10,7 @@ fn ambient_light(
     diffuse_color: vec3<f32>,
     specular_color: vec3<f32>,
     perceptual_roughness: f32,
-    occlusion: f32,
+    occlusion: vec3<f32>,
 ) -> vec3<f32> {
     let diffuse_ambient = EnvBRDFApprox(diffuse_color, F_AB(1.0, NdotV)) * occlusion;
     let specular_ambient = EnvBRDFApprox(specular_color, F_AB(perceptual_roughness, NdotV));

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -126,7 +126,7 @@ fn calculate_view(
 
 struct PbrInput {
     material: StandardMaterial,
-    occlusion: f32,
+    occlusion: vec3<f32>,
     frag_coord: vec4<f32>,
     world_position: vec4<f32>,
     // Normalized world normal used for shadow mapping as normal-mapping is not used for shadow
@@ -146,7 +146,7 @@ fn pbr_input_new() -> PbrInput {
     var pbr_input: PbrInput;
 
     pbr_input.material = standard_material_new();
-    pbr_input.occlusion = 1.0;
+    pbr_input.occlusion = vec3<f32>(1.0);
 
     pbr_input.frag_coord = vec4<f32>(0.0, 0.0, 0.0, 1.0);
     pbr_input.world_position = vec4<f32>(0.0, 0.0, 0.0, 1.0);

crates/bevy_pbr/src/render/utils.wgsl

@@ -1,6 +1,7 @@
 #define_import_path bevy_pbr::utils
 
 const PI: f32 = 3.141592653589793;
+const HALF_PI: f32 = 1.57079632679;
 const E: f32 = 2.718281828459045;
 
 fn hsv2rgb(hue: f32, saturation: f32, value: f32) -> vec3<f32> {