Adding alpha_threshold to OrderIndependentTransparencySettings for user-level optimization (bevyengine#16090)

# Objective

Order independent transparency can filter fragment writes based on the
alpha value and it is currently hard-coded to anything higher than 0.0.
By making that value configurable, users can optimize fragment writes,
potentially reducing the number of layers needed and improving
performance in favor of some transparency quality.

## Solution

This PR adds `alpha_threshold` to the
OrderIndependentTransparencySettings component and uses the struct to
configure a corresponding shader uniform. This uniform is then used
instead of the hard-coded value.

To configure OIT with a custom alpha threshold, use:

```rust
fn setup(mut commands: Commands) {
    commands.spawn((
        Camera3d::default(),
        OrderIndependentTransparencySettings {
            layer_count: 8,
            alpha_threshold: 0.2,
        },
    ));
}
```

## Testing

I tested this change using the included OIT example, as well as with two
additional projects.

## Migration Guide

If you previously explicitly initialized
OrderIndependentTransparencySettings with your own `layer_count`, you
will now have to add either a `..default()` statement or an explicit
`alpha_threshold` value:

```rust
fn setup(mut commands: Commands) {
    commands.spawn((
        Camera3d::default(),
        OrderIndependentTransparencySettings {
            layer_count: 16,
            ..default()
        },
    ));
}
```

---------

Co-authored-by: JMS55 <[email protected]>

Author: hexroll

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};
@@ -2198,7 +2198,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

@@ -377,7 +377,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),
+                ),
             ));
         }
     }
@@ -694,16 +697,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

@@ -109,5 +109,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,
+};