Add depth and normal prepass (bevyengine#6284)

# Objective

- Add a configurable prepass
- A depth prepass is useful for various shader effects and to reduce overdraw. It can be expansive depending on the scene so it's important to be able to disable it if you don't need any effects that uses it or don't suffer from excessive overdraw.
- The goal is to eventually use it for things like TAA, Ambient Occlusion, SSR and various other techniques that can benefit from having a prepass.

## Solution

The prepass node is inserted before the main pass. It runs for each `Camera3d` with a prepass component (`DepthPrepass`, `NormalPrepass`). The presence of one of those components is used to determine which textures are generated in the prepass. When any prepass is enabled, the depth buffer generated will be used by the main pass to reduce overdraw.

The prepass runs for each `Material` created with the `MaterialPlugin::prepass_enabled` option set to `true`. You can overload the shader used by the prepass by using `Material::prepass_vertex_shader()` and/or `Material::prepass_fragment_shader()`. It will also use the `Material::specialize()` for more advanced use cases. It is enabled by default on all materials.

The prepass works on opaque materials and materials using an alpha mask. Transparent materials are ignored.

The `StandardMaterial` overloads the prepass fragment shader to support alpha mask and normal maps.

---

## Changelog

- Add a new `PrepassNode` that runs before the main pass
- Add a `PrepassPlugin` to extract/prepare/queue the necessary data
- Add a `DepthPrepass` and `NormalPrepass` component to control which textures will be created by the prepass and available in later passes.
- Add a new `prepass_enabled` flag to the `MaterialPlugin` that will control if a material uses the prepass or not.
- Add a new `prepass_enabled` flag to the `PbrPlugin` to control if the StandardMaterial uses the prepass. Currently defaults to false.
- Add `Material::prepass_vertex_shader()` and `Material::prepass_fragment_shader()` to control the prepass from the `Material`

## Notes

In bevy's sample 3d scene, the performance is actually worse when enabling the prepass, but on more complex scenes the performance is generally better. I would like more testing on this, but @DGriffin91 has reported a very noticeable improvements in some scenes.

The prepass is also used by @JMS55 for TAA and GTAO

discord thread: <https://discord.com/channels/691052431525675048/1011624228627419187>

This PR was built on top of the work of multiple people

Co-Authored-By: @superdump 
Co-Authored-By: @robtfm 
Co-Authored-By: @JMS55 

Co-authored-by: Charles <IceSentry@users.noreply.github.com>
Co-authored-by: JMS55 <47158642+JMS55@users.noreply.github.com>

Cargo.toml

@@ -1220,6 +1220,17 @@ description = "A shader and a material that uses it"
 category = "Shaders"
 wasm = true
 
+[[example]]
+name = "shader_prepass"
+path = "examples/shader/shader_prepass.rs"
+
+[package.metadata.example.shader_prepass]
+name = "Material Prepass"
+description = "A shader that uses the depth texture generated in a prepass"
+category = "Shaders"
+wasm = false
+
+
 [[example]]
 name = "shader_material_screenspace_texture"
 path = "examples/shader/shader_material_screenspace_texture.rs"

assets/shaders/show_prepass.wgsl

@@ -0,0 +1,26 @@
+#import bevy_pbr::mesh_types
+#import bevy_pbr::mesh_view_bindings
+#import bevy_pbr::utils
+
+@group(1) @binding(0)
+var<uniform> show_depth: f32;
+@group(1) @binding(1)
+var<uniform> show_normal: f32;
+
+@fragment
+fn fragment(
+    @builtin(position) frag_coord: vec4<f32>,
+    @builtin(sample_index) sample_index: u32,
+    #import bevy_pbr::mesh_vertex_output
+) -> @location(0) vec4<f32> {
+    if show_depth == 1.0 {
+        let depth = prepass_depth(frag_coord, sample_index);
+        return vec4(depth, depth, depth, 1.0);
+    } else if show_normal == 1.0 {
+        let normal = prepass_normal(frag_coord, sample_index);
+        return vec4(normal, 1.0);
+    } else {
+        // transparent
+        return vec4(0.0);
+    }
+}

crates/bevy_core_pipeline/src/core_3d/main_pass_3d_node.rs

@@ -1,6 +1,7 @@
 use crate::{
     clear_color::{ClearColor, ClearColorConfig},
     core_3d::{AlphaMask3d, Camera3d, Opaque3d, Transparent3d},
+    prepass::{DepthPrepass, NormalPrepass},
 };
 use bevy_ecs::prelude::*;
 use bevy_render::{
@@ -14,6 +15,8 @@ use bevy_render::{
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::info_span;
 
+use super::Camera3dDepthLoadOp;
+
 pub struct MainPass3dNode {
     query: QueryState<
         (
@@ -24,6 +27,8 @@ pub struct MainPass3dNode {
             &'static Camera3d,
             &'static ViewTarget,
             &'static ViewDepthTexture,
+            Option<&'static DepthPrepass>,
+            Option<&'static NormalPrepass>,
         ),
         With<ExtractedView>,
     >,
@@ -55,13 +60,20 @@ impl Node for MainPass3dNode {
         world: &World,
     ) -> Result<(), NodeRunError> {
         let view_entity = graph.get_input_entity(Self::IN_VIEW)?;
-        let (camera, opaque_phase, alpha_mask_phase, transparent_phase, camera_3d, target, depth) =
-            match self.query.get_manual(world, view_entity) {
-                Ok(query) => query,
-                Err(_) => {
-                    return Ok(());
-                } // No window
-            };
+        let Ok((
+            camera,
+            opaque_phase,
+            alpha_mask_phase,
+            transparent_phase,
+            camera_3d,
+            target,
+            depth,
+            depth_prepass,
+            normal_prepass,
+        )) = self.query.get_manual(world, view_entity) else {
+            // No window
+            return Ok(());
+        };
 
         // Always run opaque pass to ensure screen is cleared
         {
@@ -88,8 +100,15 @@ impl Node for MainPass3dNode {
                     view: &depth.view,
                     // NOTE: The opaque main pass loads the depth buffer and possibly overwrites it
                     depth_ops: Some(Operations {
-                        // NOTE: 0.0 is the far plane due to bevy's use of reverse-z projections.
-                        load: camera_3d.depth_load_op.clone().into(),
+                        load: if depth_prepass.is_some() || normal_prepass.is_some() {
+                            // if any prepass runs, it will generate a depth buffer so we should use it,
+                            // even if only the normal_prepass is used.
+                            Camera3dDepthLoadOp::Load
+                        } else {
+                            // NOTE: 0.0 is the far plane due to bevy's use of reverse-z projections.
+                            camera_3d.depth_load_op.clone()
+                        }
+                        .into(),
                         store: true,
                     }),
                     stencil_ops: None,

crates/bevy_core_pipeline/src/core_3d/mod.rs

@@ -7,6 +7,7 @@ pub mod graph {
         pub const VIEW_ENTITY: &str = "view_entity";
     }
     pub mod node {
+        pub const PREPASS: &str = "prepass";
         pub const MAIN_PASS: &str = "main_pass";
         pub const BLOOM: &str = "bloom";
         pub const TONEMAPPING: &str = "tonemapping";
@@ -43,7 +44,11 @@ use bevy_render::{
 };
 use bevy_utils::{FloatOrd, HashMap};
 
-use crate::{tonemapping::TonemappingNode, upscaling::UpscalingNode};
+use crate::{
+    prepass::{node::PrepassNode, DepthPrepass},
+    tonemapping::TonemappingNode,
+    upscaling::UpscalingNode,
+};
 
 pub struct Core3dPlugin;
 
@@ -68,20 +73,29 @@ impl Plugin for Core3dPlugin {
             .add_system_to_stage(RenderStage::PhaseSort, sort_phase_system::<AlphaMask3d>)
             .add_system_to_stage(RenderStage::PhaseSort, sort_phase_system::<Transparent3d>);
 
+        let prepass_node = PrepassNode::new(&mut render_app.world);
         let pass_node_3d = MainPass3dNode::new(&mut render_app.world);
         let tonemapping = TonemappingNode::new(&mut render_app.world);
         let upscaling = UpscalingNode::new(&mut render_app.world);
         let mut graph = render_app.world.resource_mut::<RenderGraph>();
 
         let mut draw_3d_graph = RenderGraph::default();
+        draw_3d_graph.add_node(graph::node::PREPASS, prepass_node);
         draw_3d_graph.add_node(graph::node::MAIN_PASS, pass_node_3d);
         draw_3d_graph.add_node(graph::node::TONEMAPPING, tonemapping);
         draw_3d_graph.add_node(graph::node::END_MAIN_PASS_POST_PROCESSING, EmptyNode);
         draw_3d_graph.add_node(graph::node::UPSCALING, upscaling);
+
         let input_node_id = draw_3d_graph.set_input(vec![SlotInfo::new(
             graph::input::VIEW_ENTITY,
             SlotType::Entity,
         )]);
+        draw_3d_graph.add_slot_edge(
+            input_node_id,
+            graph::input::VIEW_ENTITY,
+            graph::node::PREPASS,
+            PrepassNode::IN_VIEW,
+        );
         draw_3d_graph.add_slot_edge(
             input_node_id,
             graph::input::VIEW_ENTITY,
@@ -100,6 +114,7 @@ impl Plugin for Core3dPlugin {
             graph::node::UPSCALING,
             UpscalingNode::IN_VIEW,
         );
+        draw_3d_graph.add_node_edge(graph::node::PREPASS, graph::node::MAIN_PASS);
         draw_3d_graph.add_node_edge(graph::node::MAIN_PASS, graph::node::TONEMAPPING);
         draw_3d_graph.add_node_edge(
             graph::node::TONEMAPPING,
@@ -253,7 +268,7 @@ pub fn prepare_core_3d_depth_textures(
     msaa: Res<Msaa>,
     render_device: Res<RenderDevice>,
     views_3d: Query<
-        (Entity, &ExtractedCamera),
+        (Entity, &ExtractedCamera, Option<&DepthPrepass>),
         (
             With<RenderPhase<Opaque3d>>,
             With<RenderPhase<AlphaMask3d>>,
@@ -262,34 +277,46 @@ pub fn prepare_core_3d_depth_textures(
     >,
 ) {
     let mut textures = HashMap::default();
-    for (entity, camera) in &views_3d {
-        if let Some(physical_target_size) = camera.physical_target_size {
-            let cached_texture = textures
-                .entry(camera.target.clone())
-                .or_insert_with(|| {
-                    texture_cache.get(
-                        &render_device,
-                        TextureDescriptor {
-                            label: Some("view_depth_texture"),
-                            size: Extent3d {
-                                depth_or_array_layers: 1,
-                                width: physical_target_size.x,
-                                height: physical_target_size.y,
-                            },
-                            mip_level_count: 1,
-                            sample_count: msaa.samples,
-                            dimension: TextureDimension::D2,
-                            format: TextureFormat::Depth32Float, /* PERF: vulkan docs recommend using 24
-                                                                  * bit depth for better performance */
-                            usage: TextureUsages::RENDER_ATTACHMENT,
-                        },
-                    )
-                })
-                .clone();
-            commands.entity(entity).insert(ViewDepthTexture {
-                texture: cached_texture.texture,
-                view: cached_texture.default_view,
-            });
-        }
+    for (entity, camera, depth_prepass) in &views_3d {
+        let Some(physical_target_size) = camera.physical_target_size else {
+            continue;
+        };
+
+        let cached_texture = textures
+            .entry(camera.target.clone())
+            .or_insert_with(|| {
+                // Default usage required to write to the depth texture
+                let mut usage = TextureUsages::RENDER_ATTACHMENT;
+                if depth_prepass.is_some() {
+                    // Required to read the output of the prepass
+                    usage |= TextureUsages::COPY_SRC;
+                }
+
+                // The size of the depth texture
+                let size = Extent3d {
+                    depth_or_array_layers: 1,
+                    width: physical_target_size.x,
+                    height: physical_target_size.y,
+                };
+
+                let descriptor = TextureDescriptor {
+                    label: Some("view_depth_texture"),
+                    size,
+                    mip_level_count: 1,
+                    sample_count: msaa.samples,
+                    dimension: TextureDimension::D2,
+                    // PERF: vulkan docs recommend using 24 bit depth for better performance
+                    format: TextureFormat::Depth32Float,
+                    usage,
+                };
+
+                texture_cache.get(&render_device, descriptor)
+            })
+            .clone();
+
+        commands.entity(entity).insert(ViewDepthTexture {
+            texture: cached_texture.texture,
+            view: cached_texture.default_view,
+        });
     }
 }

crates/bevy_core_pipeline/src/lib.rs

@@ -4,6 +4,7 @@ pub mod core_2d;
 pub mod core_3d;
 pub mod fullscreen_vertex_shader;
 pub mod fxaa;
+pub mod prepass;
 pub mod tonemapping;
 pub mod upscaling;
 
@@ -23,6 +24,7 @@ use crate::{
     core_3d::Core3dPlugin,
     fullscreen_vertex_shader::FULLSCREEN_SHADER_HANDLE,
     fxaa::FxaaPlugin,
+    prepass::{DepthPrepass, NormalPrepass},
     tonemapping::TonemappingPlugin,
     upscaling::UpscalingPlugin,
 };
@@ -44,6 +46,8 @@ impl Plugin for CorePipelinePlugin {
 
         app.register_type::<ClearColor>()
             .register_type::<ClearColorConfig>()
+            .register_type::<DepthPrepass>()
+            .register_type::<NormalPrepass>()
             .init_resource::<ClearColor>()
             .add_plugin(ExtractResourcePlugin::<ClearColor>::default())
             .add_plugin(Core2dPlugin)