Add MAY_DISCARD shader def, enabling early depth tests for most cases (#6697)

# Objective

- Right now we can't really benefit from [early depth
testing](https://www.khronos.org/opengl/wiki/Early_Fragment_Test) in our
PBR shader because it includes codepaths with `discard`, even for
situations where they are not necessary.

## Solution

- This PR introduces a new `MeshPipelineKey` and shader def,
`MAY_DISCARD`;
- All possible material/mesh options that that may result in `discard`s
being needed must set `MAY_DISCARD` ahead of time:
- Right now, this is only `AlphaMode::Mask(f32)`, but in the future
might include other options/effects; (e.g. one effect I'm personally
interested in is bayer dither pseudo-transparency for LOD transitions of
opaque meshes)
- Shader codepaths that can `discard` are guarded by an `#ifdef
MAY_DISCARD` preprocessor directive:
  - Right now, this is just one branch in `alpha_discard()`;
- If `MAY_DISCARD` is _not_ set, the `@early_depth_test` attribute is
added to the PBR fragment shader. This is a not yet documented, possibly
non-standard WGSL extension I found browsing Naga's source code. [I
opened a PR to document it
there](gfx-rs/naga#2132). My understanding is
that for backends where this attribute is supported, it will force an
explicit opt-in to early depth test. (e.g. via
`layout(early_fragment_tests) in;` in GLSL)

## Caveats

- I included `@early_depth_test` for the sake of us being explicit, and
avoiding the need for the driver to be “smart” about enabling this
feature. That way, if we make a mistake and include a `discard`
unguarded by `MAY_DISCARD`, it will either produce errors or noticeable
visual artifacts so that we'll catch early, instead of causing a
performance regression.
- I'm not sure explicit early depth test is supported on the naga Metal
backend, which is what I'm currently using, so I can't really test the
explicit early depth test enable, I would like others with Vulkan/GL
hardware to test it if possible;
- I would like some guidance on how to measure/verify the performance
benefits of this;
- If I understand it correctly, this, or _something like this_ is needed
to fully reap the performance gains enabled by #6284;
- This will _most definitely_ conflict with #6284 and #6644. I can fix
the conflicts as needed, depending on whether/the order they end up
being merging in.

---

## Changelog

### Changed

- Early depth tests are now enabled whenever possible for meshes using
`StandardMaterial`, reducing the number of fragments evaluated for
scenes with lots of occlusions.

Author: coreh

crates/bevy_pbr/src/material.rs

@@ -472,6 +472,9 @@ pub fn queue_material_meshes<M: Material>(
                         AlphaMode::Multiply => {
                             mesh_key |= MeshPipelineKey::BLEND_MULTIPLY;
                         }
+                        AlphaMode::Mask(_) => {
+                            mesh_key |= MeshPipelineKey::MAY_DISCARD;
+                        }
                         _ => (),
                     }
 

crates/bevy_pbr/src/prepass/mod.rs

@@ -364,8 +364,8 @@ where
             shader_defs.push("DEPTH_PREPASS".into());
         }
 
-        if key.mesh_key.contains(MeshPipelineKey::ALPHA_MASK) {
-            shader_defs.push("ALPHA_MASK".into());
+        if key.mesh_key.contains(MeshPipelineKey::MAY_DISCARD) {
+            shader_defs.push("MAY_DISCARD".into());
         }
 
         let blend_key = key
@@ -467,9 +467,7 @@ where
         // is enabled or the material uses alpha cutoff values and doesn't rely on the standard
         // prepass shader
         let fragment_required = !targets.is_empty()
-            || ((key.mesh_key.contains(MeshPipelineKey::ALPHA_MASK)
-                || blend_key == MeshPipelineKey::BLEND_PREMULTIPLIED_ALPHA
-                || blend_key == MeshPipelineKey::BLEND_ALPHA)
+            || (key.mesh_key.contains(MeshPipelineKey::MAY_DISCARD)
                 && self.material_fragment_shader.is_some());
 
         let fragment = fragment_required.then(|| {
@@ -967,7 +965,7 @@ pub fn queue_prepass_material_meshes<M: Material>(
             let alpha_mode = material.properties.alpha_mode;
             match alpha_mode {
                 AlphaMode::Opaque => {}
-                AlphaMode::Mask(_) => mesh_key |= MeshPipelineKey::ALPHA_MASK,
+                AlphaMode::Mask(_) => mesh_key |= MeshPipelineKey::MAY_DISCARD,
                 AlphaMode::Blend
                 | AlphaMode::Premultiplied
                 | AlphaMode::Add

crates/bevy_pbr/src/render/light.rs

@@ -1608,11 +1608,11 @@ pub fn queue_shadows<M: Material>(
                         }
                         let alpha_mode = material.properties.alpha_mode;
                         match alpha_mode {
-                            AlphaMode::Mask(_) => {
-                                mesh_key |= MeshPipelineKey::ALPHA_MASK;
-                            }
-                            AlphaMode::Blend | AlphaMode::Premultiplied | AlphaMode::Add => {
-                                mesh_key |= MeshPipelineKey::BLEND_PREMULTIPLIED_ALPHA;
+                            AlphaMode::Mask(_)
+                            | AlphaMode::Blend
+                            | AlphaMode::Premultiplied
+                            | AlphaMode::Add => {
+                                mesh_key |= MeshPipelineKey::MAY_DISCARD;
                             }
                             _ => {}
                         }

crates/bevy_pbr/src/render/mesh.rs

@@ -582,7 +582,8 @@ bitflags::bitflags! {
         const DEPTH_PREPASS                     = (1 << 3);
         const NORMAL_PREPASS                    = (1 << 4);
         const MOTION_VECTOR_PREPASS             = (1 << 5);
-        const ALPHA_MASK                        = (1 << 6);
+        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 BLEND_RESERVED_BITS               = Self::BLEND_MASK_BITS << Self::BLEND_SHIFT_BITS; // ← Bitmask reserving bits for the blend state
@@ -795,6 +796,10 @@ impl SpecializedMeshPipeline for MeshPipeline {
             }
         }
 
+        if key.contains(MeshPipelineKey::MAY_DISCARD) {
+            shader_defs.push("MAY_DISCARD".into());
+        }
+
         if key.contains(MeshPipelineKey::ENVIRONMENT_MAP) {
             shader_defs.push("ENVIRONMENT_MAP".into());
         }

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -14,16 +14,20 @@ fn alpha_discard(material: StandardMaterial, output_color: vec4<f32>) -> vec4<f3
     if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE {
         // NOTE: If rendering as opaque, alpha should be ignored so set to 1.0
         color.a = 1.0;
-    } else if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_MASK {
+    }
+
+#ifdef MAY_DISCARD
+    else if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_MASK {
         if color.a >= material.alpha_cutoff {
             // NOTE: If rendering as masked alpha and >= the cutoff, render as fully opaque
             color.a = 1.0;
         } else {
-            // NOTE: output_color.a < in.material.alpha_cutoff should not is not rendered
-            // NOTE: This and any other discards mean that early-z testing cannot be done!
+            // NOTE: output_color.a < in.material.alpha_cutoff should not be rendered
             discard;
         }
     }
+#endif
+
     return color;
 }
 

crates/bevy_pbr/src/render/pbr_prepass.wgsl

@@ -30,19 +30,7 @@ const PREMULTIPLIED_ALPHA_CUTOFF = 0.05;
 // We can use a simplified version of alpha_discard() here since we only need to handle the alpha_cutoff
 fn prepass_alpha_discard(in: FragmentInput) {
 
-// This is a workaround since the preprocessor does not support
-// #if defined(ALPHA_MASK) || defined(BLEND_PREMULTIPLIED_ALPHA)
-#ifndef ALPHA_MASK
-#ifndef BLEND_PREMULTIPLIED_ALPHA
-#ifndef BLEND_ALPHA
-
-#define EMPTY_PREPASS_ALPHA_DISCARD
-
-#endif // BLEND_ALPHA
-#endif // BLEND_PREMULTIPLIED_ALPHA not defined
-#endif // ALPHA_MASK not defined
-
-#ifndef EMPTY_PREPASS_ALPHA_DISCARD
+#ifdef MAY_DISCARD
     var output_color: vec4<f32> = material.base_color;
 
 #ifdef VERTEX_UVS
@@ -51,22 +39,22 @@ fn prepass_alpha_discard(in: FragmentInput) {
     }
 #endif // VERTEX_UVS
 
-#ifdef ALPHA_MASK
-    if ((material.flags & STANDARD_MATERIAL_FLAGS_ALPHA_MODE_MASK) != 0u) && output_color.a < material.alpha_cutoff {
-        discard;
-    }
-#else // BLEND_PREMULTIPLIED_ALPHA || BLEND_ALPHA
     let alpha_mode = material.flags & STANDARD_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
-    if (alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND || alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_ADD)
-        && output_color.a < PREMULTIPLIED_ALPHA_CUTOFF {
-        discard;
-    } else if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_PREMULTIPLIED
-        && all(output_color < vec4(PREMULTIPLIED_ALPHA_CUTOFF)) {
-        discard;
+    if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_MASK {
+        if output_color.a < material.alpha_cutoff {
+            discard;
+        }
+    } else if (alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND || alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_ADD) {
+        if output_color.a < PREMULTIPLIED_ALPHA_CUTOFF {
+            discard;
+        }
+    } else if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_PREMULTIPLIED {
+        if all(output_color < vec4(PREMULTIPLIED_ALPHA_CUTOFF)) {
+            discard;
+        }
     }
-#endif // !ALPHA_MASK
 
-#endif // EMPTY_PREPASS_ALPHA_DISCARD not defined
+#endif // MAY_DISCARD
 }
 
 #ifdef PREPASS_FRAGMENT