Adding alpha_threshold to OrderIndependentTransparencySettings for user-level optimization

crates/bevy_core_pipeline/src/oit/mod.rs

@@ -2,18 +2,24 @@
 
 use bevy_app::prelude::*;
 use bevy_asset::{load_internal_asset, Handle};
-use bevy_ecs::prelude::*;
+use bevy_ecs::{component::*, prelude::*};
 use bevy_math::UVec2;
+use bevy_reflect::Reflect;
 use bevy_render::{
     camera::{Camera, ExtractedCamera},
     extract_component::{ExtractComponent, ExtractComponentPlugin},
     render_graph::{RenderGraphApp, ViewNodeRunner},
-    render_resource::{BufferUsages, BufferVec, DynamicUniformBuffer, Shader, TextureUsages},
+    render_resource::{
+        BufferUsages, BufferVec, DynamicUniformBuffer, Shader, ShaderType, TextureUsages,
+    },
     renderer::{RenderDevice, RenderQueue},
     view::Msaa,
     Render, RenderApp, RenderSet,
 };
-use bevy_utils::{tracing::trace, HashSet, Instant};
+use bevy_utils::{
+    tracing::{trace, warn},
+    HashSet, Instant,
+};
 use bevy_window::PrimaryWindow;
 use resolve::{
     node::{OitResolveNode, OitResolvePass},
@@ -36,17 +42,41 @@ pub const OIT_DRAW_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(404252
 // TODO consider supporting multiple OIT techniques like WBOIT, Moment Based OIT,
 // depth peeling, stochastic transparency, ray tracing etc.
 // This should probably be done by adding an enum to this component.
-#[derive(Component, Clone, Copy, ExtractComponent)]
+// We use the same struct to pass on the settings to the drawing shader.
+#[derive(Clone, Copy, ExtractComponent, Reflect, ShaderType)]
 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 16 is probably not worth it in the vast majority of cases
-    pub layer_count: u8,
+    /// 8 is generally recommended, going above 32 is probably not worth it in the vast majority of cases
+    pub layer_count: i32,
+    /// 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.
+    pub alpha_threshold: f32,
 }
 
 impl Default for OrderIndependentTransparencySettings {
     fn default() -> Self {
-        Self { layer_count: 8 }
+        Self {
+            layer_count: 8,
+            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;
+
+    fn register_component_hooks(hooks: &mut ComponentHooks) {
+        hooks.on_add(|world, entity, _| {
+            if let Some(value) = world.get::<OrderIndependentTransparencySettings>(entity) {
+                if value.layer_count > 32 {
+                    warn!("OrderIndependentTransparencySettings layer_count set to {} might be too high.", value.layer_count);
+                }
+            }
+        });
     }
 }
 
@@ -82,7 +112,8 @@ impl Plugin for OrderIndependentTransparencyPlugin {
             OitResolvePlugin,
         ))
         .add_systems(Update, check_msaa)
-        .add_systems(Last, configure_depth_texture_usages);
+        .add_systems(Last, configure_depth_texture_usages)
+        .register_type::<OrderIndependentTransparencySettings>();
 
         let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
             return;
@@ -164,7 +195,7 @@ pub struct OitBuffers {
     pub layers: BufferVec<UVec2>,
     /// Buffer containing the index of the last layer that was written for each fragment.
     pub layer_ids: BufferVec<i32>,
-    pub layers_count_uniforms: DynamicUniformBuffer<i32>,
+    pub settings: DynamicUniformBuffer<OrderIndependentTransparencySettings>,
 }
 
 impl FromWorld for OitBuffers {
@@ -184,19 +215,19 @@ impl FromWorld for OitBuffers {
         layer_ids.reserve(1, render_device);
         layer_ids.write_buffer(render_device, render_queue);
 
-        let mut layers_count_uniforms = DynamicUniformBuffer::default();
-        layers_count_uniforms.set_label(Some("oit_layers_count"));
+        let mut settings = DynamicUniformBuffer::default();
+        settings.set_label(Some("oit_settings"));
 
         Self {
             layers,
             layer_ids,
-            layers_count_uniforms,
+            settings,
         }
     }
 }
 
 #[derive(Component)]
-pub struct OitLayersCountOffset {
+pub struct OrderIndependentTransparencySettingsOffset {
     pub offset: u32,
 }
 
@@ -268,16 +299,16 @@ pub fn prepare_oit_buffers(
         );
     }
 
-    if let Some(mut writer) = buffers.layers_count_uniforms.get_writer(
+    if let Some(mut writer) = buffers.settings.get_writer(
         camera_oit_uniforms.iter().len(),
         &render_device,
         &render_queue,
     ) {
         for (entity, settings) in &camera_oit_uniforms {
-            let offset = writer.write(&(settings.layer_count as i32));
+            let offset = writer.write(settings);
             commands
                 .entity(entity)
-                .insert(OitLayersCountOffset { offset });
+                .insert(OrderIndependentTransparencySettingsOffset { offset });
         }
     }
 }

crates/bevy_core_pipeline/src/oit/oit_draw.wgsl

@@ -1,14 +1,13 @@
 #define_import_path bevy_core_pipeline::oit
 
-#import bevy_pbr::mesh_view_bindings::{view, oit_layers, oit_layer_ids, oit_layers_count}
+#import bevy_pbr::mesh_view_bindings::{view, oit_layers, oit_layer_ids, oit_settings}
 
 #ifdef OIT_ENABLED
 // Add the fragment to the oit buffer
 fn oit_draw(position: vec4f, color: vec4f) {
     // Don't add fully transparent fragments to the list
     // because we don't want to have to sort them in the resolve pass
-    // TODO should this be comparing with < espilon ?
-    if color.a == 0.0 {
+    if color.a < oit_settings.alpha_threshold {
         return;
     }
     // get the index of the current fragment relative to the screen size
@@ -20,10 +19,10 @@ fn oit_draw(position: vec4f, color: vec4f) {
     // gets the layer index of the current fragment
     var layer_id = atomicAdd(&oit_layer_ids[screen_index], 1);
     // exit early if we've reached the maximum amount of fragments per layer
-    if layer_id >= oit_layers_count {
+    if layer_id >= oit_settings.layers_count {
         // force to store the oit_layers_count to make sure we don't
         // accidentally increase the index above the maximum value
-        atomicStore(&oit_layer_ids[screen_index], oit_layers_count);
+        atomicStore(&oit_layer_ids[screen_index], oit_settings.layers_count);
         // TODO for tail blending we should return the color here
         return;
     }

crates/bevy_core_pipeline/src/oit/resolve/mod.rs

@@ -121,7 +121,7 @@ pub struct OitResolvePipelineId(pub CachedRenderPipelineId);
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct OitResolvePipelineKey {
     hdr: bool,
-    layer_count: u8,
+    layer_count: i32,
 }
 
 #[allow(clippy::too_many_arguments)]

crates/bevy_pbr/src/render/mesh.rs

@@ -5,7 +5,7 @@ use bevy_asset::{load_internal_asset, AssetId};
 use bevy_core_pipeline::{
     core_3d::{AlphaMask3d, Opaque3d, Transmissive3d, Transparent3d, CORE_3D_DEPTH_FORMAT},
     deferred::{AlphaMask3dDeferred, Opaque3dDeferred},
-    oit::{prepare_oit_buffers, OitLayersCountOffset},
+    oit::{prepare_oit_buffers, OrderIndependentTransparencySettingsOffset},
     prepass::MotionVectorPrepass,
 };
 use bevy_derive::{Deref, DerefMut};
@@ -2195,7 +2195,7 @@ impl<P: PhaseItem, const I: usize> RenderCommand<P> for SetMeshViewBindGroup<I>
         Read<ViewScreenSpaceReflectionsUniformOffset>,
         Read<ViewEnvironmentMapUniformOffset>,
         Read<MeshViewBindGroup>,
-        Option<Read<OitLayersCountOffset>>,
+        Option<Read<OrderIndependentTransparencySettingsOffset>>,
     );
     type ItemQuery = ();
 

crates/bevy_pbr/src/render/mesh_view_bindings.rs

@@ -375,7 +375,10 @@ fn layout_entries(
                 // oit_layer_ids,
                 (32, storage_buffer_sized(false, None)),
                 // oit_layer_count
-                (33, uniform_buffer::<i32>(true)),
+                (
+                    33,
+                    uniform_buffer::<OrderIndependentTransparencySettings>(true),
+                ),
             ));
         }
     }
@@ -690,16 +693,16 @@ pub fn prepare_mesh_view_bind_groups(
                 if let (
                     Some(oit_layers_binding),
                     Some(oit_layer_ids_binding),
-                    Some(oit_layers_count_uniforms_binding),
+                    Some(oit_settings_binding),
                 ) = (
                     oit_buffers.layers.binding(),
                     oit_buffers.layer_ids.binding(),
-                    oit_buffers.layers_count_uniforms.binding(),
+                    oit_buffers.settings.binding(),
                 ) {
                     entries = entries.extend_with_indices((
                         (31, oit_layers_binding.clone()),
                         (32, oit_layer_ids_binding.clone()),
-                        (33, oit_layers_count_uniforms_binding.clone()),
+                        (33, oit_settings_binding.clone()),
                     ));
                 }
             }

crates/bevy_pbr/src/render/mesh_view_bindings.wgsl

@@ -105,5 +105,5 @@ const VISIBILITY_RANGE_UNIFORM_BUFFER_SIZE: u32 = 64u;
 #ifdef OIT_ENABLED
 @group(0) @binding(31) var<storage, read_write> oit_layers: array<vec2<u32>>;
 @group(0) @binding(32) var<storage, read_write> oit_layer_ids: array<atomic<i32>>;
-@group(0) @binding(33) var<uniform> oit_layers_count: i32;
+@group(0) @binding(33) var<uniform> oit_settings: types::OrderIndependentTransparencySettings;
 #endif OIT_ENABLED

crates/bevy_pbr/src/render/mesh_view_types.wgsl

@@ -158,4 +158,10 @@ struct ScreenSpaceReflectionsSettings {
 struct EnvironmentMapUniform {
     // Transformation matrix for the environment cubemaps in world space.
     transform: mat4x4<f32>,
-};
+};
+
+// Shader version of the order independent transparency settings component.
+struct OrderIndependentTransparencySettings {
+  layers_count: i32,
+  alpha_threshold: f32,
+};