Implement per-pixel linked list for OIT

crates/bevy_core_pipeline/src/core_3d/main_transparent_pass_3d_node.rs

@@ -1,12 +1,15 @@
-use crate::core_3d::Transparent3d;
+use crate::{
+    core_3d::Transparent3d,
+    oit::{resolve::OitResolvePipelineId, OrderIndependentTransparencySettings},
+};
 use bevy_camera::{MainPassResolutionOverride, Viewport};
 use bevy_ecs::{prelude::*, query::QueryItem};
 use bevy_render::{
     camera::ExtractedCamera,
     diagnostic::RecordDiagnostics,
     render_graph::{NodeRunError, RenderGraphContext, ViewNode},
     render_phase::ViewSortedRenderPhases,
-    render_resource::{RenderPassDescriptor, StoreOp},
+    render_resource::{PipelineCache, RenderPassDescriptor, StoreOp},
     renderer::RenderContext,
     view::{ExtractedView, ViewDepthTexture, ViewTarget},
 };
@@ -26,12 +29,14 @@ impl ViewNode for MainTransparentPass3dNode {
         &'static ViewTarget,
         &'static ViewDepthTexture,
         Option<&'static MainPassResolutionOverride>,
+        Has<OrderIndependentTransparencySettings>,
+        Option<&'static OitResolvePipelineId>,
     );
     fn run(
         &self,
         graph: &mut RenderGraphContext,
         render_context: &mut RenderContext,
-        (camera, view, target, depth, resolution_override): QueryItem<Self::ViewQuery>,
+        (camera, view, target, depth, resolution_override, has_oit, oit_resolve_pipeline_id): QueryItem<Self::ViewQuery>,
         world: &World,
     ) -> Result<(), NodeRunError> {
         let view_entity = graph.view_entity();
@@ -47,6 +52,22 @@ impl ViewNode for MainTransparentPass3dNode {
         };
 
         if !transparent_phase.items.is_empty() {
+            if has_oit {
+                // We can't run transparent phase if OitResolvePipelineId is not ready
+                // Otherwise we will write to `oit_atomic_counter` and `oit_heads` buffer without resetting them
+                // which causes corrupted linked list(can have circular references) on the next pass
+                let Some(oit_resolve_pipeline_id) = oit_resolve_pipeline_id else {
+                    return Ok(());
+                };
+                let pipeline_cache = world.resource::<PipelineCache>();
+                if pipeline_cache
+                    .get_render_pipeline(oit_resolve_pipeline_id.0)
+                    .is_none()
+                {
+                    return Ok(());
+                }
+            }
+
             // Run the transparent pass, sorted back-to-front
             // NOTE: Scoped to drop the mutable borrow of render_context
             #[cfg(feature = "trace")]

crates/bevy_core_pipeline/src/oit/mod.rs

@@ -1,10 +1,9 @@
 //! Order Independent Transparency (OIT) for 3d rendering. See [`OrderIndependentTransparencyPlugin`] for more details.
 
 use bevy_app::prelude::*;
-use bevy_camera::{Camera3d, RenderTarget};
-use bevy_ecs::{component::*, lifecycle::ComponentHook, prelude::*};
+use bevy_camera::Camera3d;
+use bevy_ecs::{component::*, prelude::*};
 use bevy_math::UVec2;
-use bevy_platform::collections::HashSet;
 use bevy_platform::time::Instant;
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use bevy_render::{
@@ -17,12 +16,11 @@ use bevy_render::{
     Render, RenderApp, RenderStartup, RenderSystems,
 };
 use bevy_shader::load_shader_library;
-use bevy_window::PrimaryWindow;
 use resolve::{
     node::{OitResolveNode, OitResolvePass},
     OitResolvePlugin,
 };
-use tracing::{trace, warn};
+use tracing::trace;
 
 use crate::core_3d::graph::{Core3d, Node3d};
 
@@ -35,13 +33,16 @@ pub mod resolve;
 // depth peeling, stochastic transparency, ray tracing etc.
 // This should probably be done by adding an enum to this component.
 // We use the same struct to pass on the settings to the drawing shader.
-#[derive(Clone, Copy, ExtractComponent, Reflect, ShaderType)]
+#[derive(Clone, Copy, ExtractComponent, Reflect, ShaderType, Component)]
 #[reflect(Clone, Default)]
 pub struct OrderIndependentTransparencySettings {
-    /// Controls how many layers will be used to compute the blending.
-    /// The more layers you use the more memory it will use but it will also give better results.
-    /// 8 is generally recommended, going above 32 is probably not worth it in the vast majority of cases
-    pub layer_count: i32,
+    /// Controls how many fragments will be exactly sorted.
+    /// 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 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.
     /// This can be tweaked to optimize quality / layers count. Higher values will
     /// allow lower number of layers and a better performance, compromising quality.
@@ -51,32 +52,13 @@ pub struct OrderIndependentTransparencySettings {
 impl Default for OrderIndependentTransparencySettings {
     fn default() -> Self {
         Self {
-            layer_count: 8,
+            sorted_fragment_max_count: 8,
+            fragments_per_pixel_average: 4.0,
             alpha_threshold: 0.0,
         }
     }
 }
 
-// OrderIndependentTransparencySettings is also a Component. We explicitly implement the trait so
-// we can hook on_add to issue a warning in case `layer_count` is seemingly too high.
-impl Component for OrderIndependentTransparencySettings {
-    const STORAGE_TYPE: StorageType = StorageType::SparseSet;
-    type Mutability = Mutable;
-
-    fn on_add() -> Option<ComponentHook> {
-        Some(|world, context| {
-            if let Some(value) = world.get::<OrderIndependentTransparencySettings>(context.entity)
-                && value.layer_count > 32
-            {
-                warn!("{}OrderIndependentTransparencySettings layer_count set to {} might be too high.",
-                        context.caller.map(|location|format!("{location}: ")).unwrap_or_default(),
-                        value.layer_count
-                    );
-            }
-        })
-    }
-}
-
 /// A plugin that adds support for Order Independent Transparency (OIT).
 /// This can correctly render some scenes that would otherwise have artifacts due to alpha blending, but uses more memory.
 ///
@@ -88,8 +70,8 @@ impl Component for OrderIndependentTransparencySettings {
 /// # Implementation details
 /// This implementation uses 2 passes.
 ///
-/// The first pass writes the depth and color of all the fragments to a big buffer.
-/// The buffer contains N layers for each pixel, where N can be set with [`OrderIndependentTransparencySettings::layer_count`].
+/// The first pass constructs a linked list which stores depth and color of all fragments in a big buffer.
+/// The linked list capacity can be set with [`OrderIndependentTransparencySettings::fragments_per_pixel_average`].
 /// This pass is essentially a forward pass.
 ///
 /// The second pass is a single fullscreen triangle pass that sorts all the fragments then blends them together
@@ -103,8 +85,7 @@ impl Plugin for OrderIndependentTransparencyPlugin {
             ExtractComponentPlugin::<OrderIndependentTransparencySettings>::default(),
             OitResolvePlugin,
         ))
-        .add_systems(Update, check_msaa)
-        .add_systems(Last, configure_depth_texture_usages);
+        .add_systems(Update, check_msaa);
 
         let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
             return;
@@ -114,7 +95,10 @@ impl Plugin for OrderIndependentTransparencyPlugin {
             .add_systems(RenderStartup, init_oit_buffers)
             .add_systems(
                 Render,
-                prepare_oit_buffers.in_set(RenderSystems::PrepareResources),
+                (
+                    configure_camera_depth_usages.in_set(RenderSystems::ManageViews),
+                    prepare_oit_buffers.in_set(RenderSystems::PrepareResources),
+                ),
             );
 
         render_app
@@ -130,34 +114,17 @@ impl Plugin for OrderIndependentTransparencyPlugin {
     }
 }
 
-// WARN This should only happen for cameras with the [`OrderIndependentTransparencySettings`] component
-// but when multiple cameras are present on the same window
-// bevy reuses the same depth texture so we need to set this on all cameras with the same render target.
-fn configure_depth_texture_usages(
-    p: Query<Entity, With<PrimaryWindow>>,
-    cameras: Query<(&RenderTarget, Has<OrderIndependentTransparencySettings>)>,
-    mut new_cameras: Query<(&mut Camera3d, &RenderTarget), Added<Camera3d>>,
+fn configure_camera_depth_usages(
+    mut cameras: Query<
+        &mut Camera3d,
+        (
+            Changed<Camera3d>,
+            With<OrderIndependentTransparencySettings>,
+        ),
+    >,
 ) {
-    if new_cameras.is_empty() {
-        return;
-    }
-
-    // Find all the render target that potentially uses OIT
-    let primary_window = p.single().ok();
-    let mut render_target_has_oit = <HashSet<_>>::default();
-    for (render_target, has_oit) in &cameras {
-        if has_oit {
-            render_target_has_oit.insert(render_target.normalize(primary_window));
-        }
-    }
-
-    // Update the depth texture usage for cameras with a render target that has OIT
-    for (mut camera_3d, render_target) in &mut new_cameras {
-        if render_target_has_oit.contains(&render_target.normalize(primary_window)) {
-            let mut usages = TextureUsages::from(camera_3d.depth_texture_usages);
-            usages |= TextureUsages::RENDER_ATTACHMENT | TextureUsages::TEXTURE_BINDING;
-            camera_3d.depth_texture_usages = usages.into();
-        }
+    for mut camera in &mut cameras {
+        camera.depth_texture_usages.0 |= TextureUsages::TEXTURE_BINDING.bits();
     }
 }
 
@@ -169,18 +136,25 @@ fn check_msaa(cameras: Query<&Msaa, With<OrderIndependentTransparencySettings>>)
     }
 }
 
+#[derive(Clone, Copy, ShaderType)]
+pub struct OitFragmentNode {
+    pub color: u32,
+    pub depth_alpha: u32,
+    pub next: u32,
+}
+
 /// Holds the buffers that contain the data of all OIT layers.
 /// We use one big buffer for the entire app. Each camera will reuse it so it will
 /// always be the size of the biggest OIT enabled camera.
 #[derive(Resource)]
 pub struct OitBuffers {
-    /// The OIT layers containing depth and color for each fragments.
+    pub settings: DynamicUniformBuffer<OrderIndependentTransparencySettings>,
+    /// The OIT buffers containing color, depth and linked next node for each fragments.
     /// This is essentially used as a 3d array where xy is the screen coordinate and z is
     /// the list of fragments rendered with OIT.
-    pub layers: BufferVec<UVec2>,
-    /// Buffer containing the index of the last layer that was written for each fragment.
-    pub layer_ids: BufferVec<i32>,
-    pub settings: DynamicUniformBuffer<OrderIndependentTransparencySettings>,
+    pub nodes: BufferVec<OitFragmentNode>,
+    pub heads: BufferVec<u32>,
+    pub atomic_counter: BufferVec<u32>,
 }
 
 pub fn init_oit_buffers(
@@ -190,22 +164,27 @@ pub fn init_oit_buffers(
 ) {
     // initialize buffers with something so there's a valid binding
 
-    let mut layers = BufferVec::new(BufferUsages::COPY_DST | BufferUsages::STORAGE);
-    layers.set_label(Some("oit_layers"));
-    layers.reserve(1, &render_device);
-    layers.write_buffer(&render_device, &render_queue);
+    let mut nodes = BufferVec::new(BufferUsages::COPY_DST | BufferUsages::STORAGE);
+    nodes.set_label(Some("oit_nodes"));
+    nodes.reserve(1, &render_device);
+
+    let mut heads = BufferVec::new(BufferUsages::COPY_DST | BufferUsages::STORAGE);
+    heads.set_label(Some("oit_heads"));
+    heads.push(u32::MAX);
+    heads.write_buffer(&render_device, &render_queue);
 
-    let mut layer_ids = BufferVec::new(BufferUsages::COPY_DST | BufferUsages::STORAGE);
-    layer_ids.set_label(Some("oit_layer_ids"));
-    layer_ids.reserve(1, &render_device);
-    layer_ids.write_buffer(&render_device, &render_queue);
+    let mut atomic_counter = BufferVec::new(BufferUsages::COPY_DST | BufferUsages::STORAGE);
+    atomic_counter.set_label(Some("oit_atomic_counter"));
+    atomic_counter.push(0);
+    atomic_counter.write_buffer(&render_device, &render_queue);
 
     let mut settings = DynamicUniformBuffer::default();
     settings.set_label(Some("oit_settings"));
 
     commands.insert_resource(OitBuffers {
-        layers,
-        layer_ids,
+        nodes,
+        heads,
+        atomic_counter,
         settings,
     });
 }
@@ -233,52 +212,43 @@ pub fn prepare_oit_buffers(
     mut buffers: ResMut<OitBuffers>,
 ) {
     // Get the max buffer size for any OIT enabled camera
-    let mut max_layer_ids_size = usize::MIN;
-    let mut max_layers_size = usize::MIN;
+    let mut max_size = UVec2::new(0, 0);
+    let mut fragments_per_pixel_average = 0f32;
     for (camera, settings) in &cameras {
         let Some(size) = camera.physical_target_size else {
             continue;
         };
-
-        let layer_count = settings.layer_count as usize;
-        let size = (size.x * size.y) as usize;
-        max_layer_ids_size = max_layer_ids_size.max(size);
-        max_layers_size = max_layers_size.max(size * layer_count);
+        max_size = max_size.max(size);
+        fragments_per_pixel_average =
+            fragments_per_pixel_average.max(settings.fragments_per_pixel_average);
     }
 
-    // Create or update the layers buffer based on the max size
-    if buffers.layers.capacity() < max_layers_size {
+    // Create or update the heads buffer based on the max size
+    let heads_size = (max_size.x * max_size.y) as usize;
+    if buffers.heads.capacity() < heads_size {
         let start = Instant::now();
-        buffers.layers.reserve(max_layers_size, &render_device);
-        let remaining = max_layers_size - buffers.layers.capacity();
-        for _ in 0..remaining {
-            buffers.layers.push(UVec2::ZERO);
+        buffers.heads.clear();
+        buffers.heads.reserve(heads_size, &render_device);
+        for _ in 0..heads_size {
+            buffers.heads.push(u32::MAX);
         }
-        buffers.layers.write_buffer(&render_device, &render_queue);
+        buffers.heads.write_buffer(&render_device, &render_queue);
         trace!(
-            "OIT layers buffer updated in {:.01}ms with total size {} MiB",
+            "OIT heads buffer updated in {:.01}ms with total size {} MiB",
             start.elapsed().as_millis(),
-            buffers.layers.capacity() * size_of::<UVec2>() / 1024 / 1024,
+            buffers.heads.capacity() * size_of::<u32>() / 1024 / 1024,
         );
     }
 
-    // Create or update the layer_ids buffer based on the max size
-    if buffers.layer_ids.capacity() < max_layer_ids_size {
+    // Create or update the nodes buffer based on the max size
+    let nodes_size = ((max_size.x * max_size.y) as f32 * fragments_per_pixel_average) as usize;
+    if buffers.nodes.capacity() < nodes_size {
         let start = Instant::now();
-        buffers
-            .layer_ids
-            .reserve(max_layer_ids_size, &render_device);
-        let remaining = max_layer_ids_size - buffers.layer_ids.capacity();
-        for _ in 0..remaining {
-            buffers.layer_ids.push(0);
-        }
-        buffers
-            .layer_ids
-            .write_buffer(&render_device, &render_queue);
+        buffers.nodes.reserve(nodes_size, &render_device);
         trace!(
-            "OIT layer ids buffer updated in {:.01}ms with total size {} MiB",
+            "OIT nodes buffer updated in {:.01}ms with total size {} MiB",
             start.elapsed().as_millis(),
-            buffers.layer_ids.capacity() * size_of::<UVec2>() / 1024 / 1024,
+            buffers.nodes.capacity() * size_of::<OitFragmentNode>() / 1024 / 1024,
         );
     }