make use of depth prepass to filter out fragments

beicause · beicause · commit 9d0dbc3bdff4 · 2025-11-16T18:54:15.000+08:00
diff --git a/crates/bevy_core_pipeline/src/oit/mod.rs b/crates/bevy_core_pipeline/src/oit/mod.rs
@@ -3,7 +3,7 @@
 use bevy_app::prelude::*;
 use bevy_camera::Camera3d;
 use bevy_ecs::{component::*, prelude::*};
-use bevy_math::{UVec2, UVec3};
+use bevy_math::UVec2;
 use bevy_platform::time::Instant;
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use bevy_render::{
@@ -40,7 +40,7 @@ pub struct OrderIndependentTransparencySettings {
     /// If the scene has more fragments than this, they will be merged approximately.
     /// More sorted fragments is more accurate but will be slower.
     pub sorted_fragment_max_count: u32,
-    /// The average fragments per pixel stored in the buffer. This should be bigger enough to make oit succeed.
+    /// The average fragments per pixel stored in the buffer. This should be bigger enough otherwise the fragments will be discarded.
     /// Higher values increase memory usage.
     pub fragments_per_pixel_average: f32,
     /// Threshold for which fragments will be added to the blending layers.
@@ -239,13 +239,15 @@ pub fn prepare_oit_buffers(
     if buffers.headers.capacity() < headers_size {
         let start = Instant::now();
         buffers.headers.clear();
+        buffers.headers.reserve(headers_size, &render_device);
         for _ in 0..headers_size {
             buffers.headers.push(u32::MAX);
         }
         buffers.headers.write_buffer(&render_device, &render_queue);
         trace!(
-            "OIT headers buffer updated in {:.01}ms with total size {} MiB",
+            "OIT headers buffer updated in {:.01}ms with capacity {}, total size {} MiB",
             start.elapsed().as_millis(),
+            buffers.headers.capacity(),
             buffers.headers.capacity() * size_of::<u32>() / 1024 / 1024,
         );
     }
@@ -255,14 +257,16 @@ pub fn prepare_oit_buffers(
     if buffers.nodes.capacity() < nodes_size {
         let start = Instant::now();
         buffers.nodes.clear();
+        buffers.nodes.reserve(nodes_size, &render_device);
         for _ in 0..nodes_size {
             buffers.nodes.push(OitFragmentNode::default());
         }
         buffers.nodes.write_buffer(&render_device, &render_queue);
         trace!(
-            "OIT nodes buffer updated in {:.01}ms with total size {} MiB",
+            "OIT nodes buffer updated in {:.01}ms with capacity {}, total size {} MiB",
             start.elapsed().as_millis(),
-            buffers.nodes.capacity() * size_of::<UVec3>() / 1024 / 1024,
+            buffers.nodes.capacity(),
+            buffers.nodes.capacity() * size_of::<OitFragmentNode>() / 1024 / 1024,
         );
     }
 
diff --git a/crates/bevy_core_pipeline/src/oit/oit_draw.wgsl b/crates/bevy_core_pipeline/src/oit/oit_draw.wgsl
@@ -2,24 +2,30 @@
 
 #import bevy_pbr::mesh_view_bindings::{view, oit_nodes, oit_headers, oit_atomic_counter, oit_settings}
 #import bevy_pbr::mesh_view_types::OitFragmentNode
+#import bevy_pbr::prepass_utils
 
 #ifdef OIT_ENABLED
 // Add the fragment to the oit buffer
 fn oit_draw(position: vec4f, color: vec4f) {
+#ifdef DEPTH_PREPASS
+    if position.z < prepass_utils::prepass_depth(position, 0u) {
+        return;
+    }
+#endif
     // Don't add fully transparent fragments to the list
     // because we don't want to have to sort them in the resolve pass
     if color.a < oit_settings.alpha_threshold {
         return;
     }
     // get the index of the current fragment relative to the screen size
     let screen_index = u32(floor(position.x) + floor(position.y) * view.viewport.z);
-    // get the size of the buffer.
-    // It's always the size of the screen
+    // get the size of oit_nodes. It's screen_size * fragments_per_pixel_average
     let buffer_size = u32(view.viewport.z * view.viewport.w * oit_settings.fragments_per_pixel_average);
 
     var new_node_index = atomicAdd(&oit_atomic_counter, 1);
     // exit early if we've reached the maximum amount of fragments nodes
     if new_node_index >= buffer_size {
+        // TODO for tail blending we should return the color here
         return;
     }
 
diff --git a/crates/bevy_core_pipeline/src/oit/resolve/mod.rs b/crates/bevy_core_pipeline/src/oit/resolve/mod.rs
@@ -1,7 +1,7 @@
 use core::num::NonZero;
 
 use super::OitBuffers;
-use crate::{oit::OrderIndependentTransparencySettings, FullscreenShader};
+use crate::{oit::OrderIndependentTransparencySettings, prepass::DepthPrepass, FullscreenShader};
 use bevy_app::Plugin;
 use bevy_asset::{embedded_asset, load_embedded_asset, AssetServer};
 use bevy_derive::Deref;
@@ -144,6 +144,7 @@ pub struct OitResolvePipelineId(pub CachedRenderPipelineId);
 pub struct OitResolvePipelineKey {
     hdr: bool,
     sorted_fragment_max_count: u32,
+    depth_prepass: bool,
 }
 
 pub fn queue_oit_resolve_pipeline(
@@ -155,6 +156,7 @@ pub fn queue_oit_resolve_pipeline(
             Entity,
             &ExtractedView,
             &OrderIndependentTransparencySettings,
+            Has<DepthPrepass>,
         ),
         With<OrderIndependentTransparencySettings>,
     >,
@@ -165,11 +167,12 @@ pub fn queue_oit_resolve_pipeline(
     mut cached_pipeline_id: Local<EntityHashMap<(OitResolvePipelineKey, CachedRenderPipelineId)>>,
 ) {
     let mut current_view_entities = EntityHashSet::default();
-    for (e, view, oit_settings) in &views {
+    for (e, view, oit_settings, depth_prepass) in &views {
         current_view_entities.insert(e);
         let key = OitResolvePipelineKey {
             hdr: view.hdr,
             sorted_fragment_max_count: oit_settings.sorted_fragment_max_count,
+            depth_prepass,
         };
 
         if let Some((cached_key, id)) = cached_pipeline_id.get(&e)
@@ -210,19 +213,23 @@ fn specialize_oit_resolve_pipeline(
     } else {
         TextureFormat::bevy_default()
     };
+    let mut layout = vec![resolve_pipeline.view_bind_group_layout.clone()];
+    let mut shader_defs = vec![ShaderDefVal::UInt(
+        "SORTED_FRAGMENT_MAX_COUNT".into(),
+        key.sorted_fragment_max_count,
+    )];
+    if key.depth_prepass {
+        shader_defs.push(ShaderDefVal::Bool("DEPTH_PREPASS".into(), true));
+    } else {
+        layout.push(resolve_pipeline.oit_depth_bind_group_layout.clone());
+    }
 
     RenderPipelineDescriptor {
         label: Some("oit_resolve_pipeline".into()),
-        layout: vec![
-            resolve_pipeline.view_bind_group_layout.clone(),
-            resolve_pipeline.oit_depth_bind_group_layout.clone(),
-        ],
+        layout,
         fragment: Some(FragmentState {
             shader: load_embedded_asset!(asset_server, "oit_resolve.wgsl"),
-            shader_defs: vec![ShaderDefVal::UInt(
-                "SORTED_FRAGMENT_MAX_COUNT".into(),
-                key.sorted_fragment_max_count,
-            )],
+            shader_defs,
             targets: vec![Some(ColorTargetState {
                 format,
                 blend: Some(BlendState {
diff --git a/crates/bevy_core_pipeline/src/oit/resolve/node.rs b/crates/bevy_core_pipeline/src/oit/resolve/node.rs
@@ -9,6 +9,8 @@ use bevy_render::{
     view::{ViewDepthTexture, ViewTarget, ViewUniformOffset},
 };
 
+use crate::prepass::DepthPrepass;
+
 use super::{OitResolveBindGroup, OitResolvePipeline, OitResolvePipelineId};
 
 /// Render label for the OIT resolve pass.
@@ -26,15 +28,22 @@ impl ViewNode for OitResolveNode {
         &'static OitResolvePipelineId,
         &'static ViewDepthTexture,
         Option<&'static MainPassResolutionOverride>,
+        Has<DepthPrepass>,
     );
 
     fn run(
         &self,
         _graph: &mut RenderGraphContext,
         render_context: &mut RenderContext,
-        (camera, view_target, view_uniform, oit_resolve_pipeline_id, depth, resolution_override): QueryItem<
-            Self::ViewQuery,
-        >,
+        (
+            camera,
+            view_target,
+            view_uniform,
+            oit_resolve_pipeline_id,
+            depth,
+            resolution_override,
+            depth_prepass,
+        ): QueryItem<Self::ViewQuery>,
         world: &World,
     ) -> Result<(), NodeRunError> {
         let Some(resolve_pipeline) = world.get_resource::<OitResolvePipeline>() else {
@@ -53,12 +62,18 @@ impl ViewNode for OitResolveNode {
 
             let diagnostics = render_context.diagnostic_recorder();
 
-            let depth_bind_group = render_context.render_device().create_bind_group(
-                "oit_resolve_depth_bind_group",
-                &pipeline_cache
-                    .get_bind_group_layout(&resolve_pipeline.oit_depth_bind_group_layout),
-                &BindGroupEntries::single(depth.view()),
-            );
+            let depth_bind_group = if !depth_prepass {
+                Some(
+                    render_context.render_device().create_bind_group(
+                        "oit_resolve_depth_bind_group",
+                        &pipeline_cache
+                            .get_bind_group_layout(&resolve_pipeline.oit_depth_bind_group_layout),
+                        &BindGroupEntries::single(depth.view()),
+                    ),
+                )
+            } else {
+                None
+            };
 
             let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
                 label: Some("oit_resolve"),
@@ -77,8 +92,9 @@ impl ViewNode for OitResolveNode {
 
             render_pass.set_render_pipeline(pipeline);
             render_pass.set_bind_group(0, bind_group, &[view_uniform.offset]);
-            render_pass.set_bind_group(1, &depth_bind_group, &[]);
-
+            if let Some(depth_bind_group) = &depth_bind_group {
+                render_pass.set_bind_group(1, depth_bind_group, &[]);
+            }
             render_pass.draw(0..3, 0..1);
 
             pass_span.end(&mut render_pass);
diff --git a/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl b/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
@@ -6,7 +6,9 @@
 @group(0) @binding(2) var<storage, read_write> headers: array<u32>;
 @group(0) @binding(3) var<storage, read_write> atomic_counter: u32;
 
+#ifndef DEPTH_PREPASS
 @group(1) @binding(0) var depth: texture_depth_2d;
+#endif
 
 struct OitFragment {
     color: vec3<f32>,
@@ -35,11 +37,15 @@ fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
         }
         return vec4(0.0);
     } else {
-        // Load depth for manual depth testing.
+#ifndef DEPTH_PREPASS
+        // If depth prepass is disabled, load depth for manual depth testing.
         // This is necessary because early z doesn't seem to trigger in the transparent pass.
         // This should be done during the draw pass so those fragments simply don't exist in the list,
         // but this requires a bigger refactor
         let d = textureLoad(depth, vec2<i32>(in.position.xy), 0);
+#else
+        let d = 0.0;
+#endif
         let color = resolve(header, d);
         headers[screen_index] = LINKED_LIST_END_SENTINEL;
         return color;
@@ -62,17 +68,19 @@ fn resolve(header: u32, opaque_depth: f32) -> vec4<f32> {
         let depth_alpha = bevy_core_pipeline::oit::unpack_24bit_depth_8bit_alpha(fragment_node.depth_alpha);
         current_node = fragment_node.next;
 
+#ifndef DEPTH_PREPASS
         // depth testing
         if depth_alpha.x < opaque_depth {
             continue;
         }
+#endif
 
         if sorted_frag_count < SORTED_FRAGMENT_MAX_COUNT {
             // There is still room in the sorted list.
             // Insert the fragment so that the list stay sorted.
             var i = sorted_frag_count;
             for(; i > 0; i -= 1) {
-                // short-circuit can't be used in for-loop, https://github.com/gfx-rs/wgpu/issues/4394
+                // short-circuit can't be used in for(;;;), https://github.com/gfx-rs/wgpu/issues/4394
                 if depth_alpha.x < fragment_list[i - 1].depth {
                     fragment_list[i] = fragment_list[i - 1];
                 } else {
@@ -91,7 +99,7 @@ fn resolve(header: u32, opaque_depth: f32) -> vec4<f32> {
             final_color = blend(vec4f(fragment_list[0].color * fragment_list[0].alpha, fragment_list[0].alpha), final_color);
             var i = 0u;
             for(; i < SORTED_FRAGMENT_MAX_COUNT - 1; i += 1) {
-                // short-circuit can't be used in for-loop, https://github.com/gfx-rs/wgpu/issues/4394
+                // short-circuit can't be used in for(;;;), https://github.com/gfx-rs/wgpu/issues/4394
                 if fragment_list[i + 1].depth < depth_alpha.x {
                     fragment_list[i] = fragment_list[i + 1];
                 } else {
diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs
@@ -2379,6 +2379,7 @@ impl SpecializedMeshPipeline for MeshPipeline {
         let (mut is_opaque, mut alpha_to_coverage_enabled) = (false, false);
         if key.contains(MeshPipelineKey::OIT_ENABLED) && pass == MeshPipelineKey::BLEND_ALPHA {
             label = "oit_mesh_pipeline".into();
+            // TODO tail blending would need alpha blending
             blend = None;
             shader_defs.push("OIT_ENABLED".into());
             // TODO it should be possible to use this to combine MSAA and OIT
diff --git a/examples/3d/order_independent_transparency.rs b/examples/3d/order_independent_transparency.rs
@@ -6,7 +6,7 @@
 use bevy::{
     camera::visibility::RenderLayers,
     color::palettes::css::{BLUE, GREEN, RED, YELLOW},
-    core_pipeline::oit::OrderIndependentTransparencySettings,
+    core_pipeline::{oit::OrderIndependentTransparencySettings, prepass::DepthPrepass},
     prelude::*,
 };
 
@@ -33,6 +33,8 @@ fn setup(
         RenderLayers::layer(1),
         // Msaa currently doesn't work with OIT
         Msaa::Off,
+        // Optional: depth prepass can help OIT filter out fragments occluded by opaque objects
+        DepthPrepass,
     ));
 
     // light
