Add lens dirt support to Bloom

crates/bevy_post_process/src/bloom/bloom.wgsl

@@ -6,17 +6,28 @@
 // * [COD] - Next Generation Post Processing in Call of Duty - http://www.iryoku.com/next-generation-post-processing-in-call-of-duty-advanced-warfare
 // * [PBB] - Physically Based Bloom - https://learnopengl.com/Guest-Articles/2022/Phys.-Based-Bloom
 
+#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
+
 struct BloomUniforms {
     threshold_precomputations: vec4<f32>,
     viewport: vec4<f32>,
     scale: vec2<f32>,
     aspect: f32,
+    lens_dirt_intensity: f32,
+    lens_dirt_tint: vec3<f32>,
+    padding: u32,
 };
 
 @group(0) @binding(0) var input_texture: texture_2d<f32>;
 @group(0) @binding(1) var s: sampler;
 
 @group(0) @binding(2) var<uniform> uniforms: BloomUniforms;
+@group(0) @binding(3) var<storage, read> blend_factor: f32;
+
+#ifdef LENS_DIRT
+@group(0) @binding(4) var dirt_texture: texture_2d<f32>;
+@group(0) @binding(5) var dirt_sampler: sampler;
+#endif
 
 #ifdef FIRST_DOWNSAMPLE
 // https://catlikecoding.com/unity/tutorials/advanced-rendering/bloom/#3.4
@@ -155,9 +166,10 @@ fn sample_input_3x3_tent(uv: vec2<f32>) -> vec3<f32> {
 
 #ifdef FIRST_DOWNSAMPLE
 @fragment
-fn downsample_first(@location(0) output_uv: vec2<f32>) -> @location(0) vec4<f32> {
-    let sample_uv = uniforms.viewport.xy + output_uv * uniforms.viewport.zw;
+fn downsample_first(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    let sample_uv = uniforms.viewport.xy + in.uv * uniforms.viewport.zw;
     var sample = sample_input_13_tap(sample_uv);
+
     // Lower bound of 0.0001 is to avoid propagating multiplying by 0.0 through the
     // downscaling and upscaling which would result in black boxes.
     // The upper bound is to prevent NaNs.
@@ -172,12 +184,28 @@ fn downsample_first(@location(0) output_uv: vec2<f32>) -> @location(0) vec4<f32>
 }
 #endif
 
+#ifdef LENS_DIRT
+@fragment
+fn upsample_final(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    let bloom = sample_input_3x3_tent(in.uv);
+    let bloom_intensity = max(bloom.r, max(bloom.g, bloom.b));
+
+    let dirt = textureSample(dirt_texture, dirt_sampler, in.uv).r;
+    let amount = clamp(dirt * uniforms.lens_dirt_intensity * bloom_intensity, 0.0, 1.0);
+
+    let result_bloom = mix(bloom.rgb, bloom.rgb * uniforms.lens_dirt_tint.rgb, amount);
+    let alpha = mix(blend_factor, 1.0, amount);
+
+    return vec4<f32>(result_bloom, alpha);
+}
+#endif
+
 @fragment
-fn downsample(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
-    return vec4<f32>(sample_input_13_tap(uv), 1.0);
+fn downsample(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    return vec4<f32>(sample_input_13_tap(in.uv), 1.0);
 }
 
 @fragment
-fn upsample(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
-    return vec4<f32>(sample_input_3x3_tent(uv), 1.0);
+fn upsample(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    return vec4<f32>(sample_input_3x3_tent(in.uv), blend_factor);
 }

crates/bevy_post_process/src/bloom/downsampling_pipeline.rs

@@ -1,13 +1,13 @@
-use bevy_core_pipeline::FullscreenShader;
+use super::{settings::BloomUniforms, Bloom, BLOOM_TEXTURE_FORMAT};
 
-use super::{Bloom, BLOOM_TEXTURE_FORMAT};
 use bevy_asset::{load_embedded_asset, AssetServer, Handle};
+use bevy_core_pipeline::FullscreenShader;
 use bevy_ecs::{
     prelude::{Component, Entity},
     resource::Resource,
     system::{Commands, Query, Res, ResMut},
 };
-use bevy_math::{Vec2, Vec4};
+use bevy_math::Vec2;
 use bevy_render::{
     render_resource::{
         binding_types::{sampler, texture_2d, uniform_buffer},
@@ -42,17 +42,6 @@ pub struct BloomDownsamplingPipelineKeys {
     uniform_scale: bool,
 }
 
-/// The uniform struct extracted from [`Bloom`] attached to a Camera.
-/// Will be available for use in the Bloom shader.
-#[derive(Component, ShaderType, Clone)]
-pub struct BloomUniforms {
-    // Precomputed values used when thresholding, see https://catlikecoding.com/unity/tutorials/advanced-rendering/bloom/#3.4
-    pub threshold_precomputations: Vec4,
-    pub viewport: Vec4,
-    pub scale: Vec2,
-    pub aspect: f32,
-}
-
 pub fn init_bloom_downsampling_pipeline(
     mut commands: Commands,
     render_device: Res<RenderDevice>,

crates/bevy_post_process/src/bloom/mod.rs

@@ -2,41 +2,42 @@ mod downsampling_pipeline;
 mod settings;
 mod upsampling_pipeline;
 
-use bevy_image::ToExtents;
-pub use settings::{Bloom, BloomCompositeMode, BloomPrefilter};
-
 use crate::bloom::{
-    downsampling_pipeline::init_bloom_downsampling_pipeline,
-    upsampling_pipeline::init_bloom_upscaling_pipeline,
+    downsampling_pipeline::{
+        init_bloom_downsampling_pipeline, prepare_downsampling_pipeline, BloomDownsamplingPipeline,
+        BloomDownsamplingPipelineIds,
+    },
+    settings::BloomUniforms,
+    upsampling_pipeline::{
+        init_bloom_upscaling_pipeline, prepare_upsampling_pipeline, BloomUpsamplingPipeline,
+        UpsamplingPipelineIds,
+    },
 };
+pub use settings::{Bloom, BloomCompositeMode, BloomPrefilter, LensDirt};
+
 use bevy_app::{App, Plugin};
 use bevy_asset::embedded_asset;
-use bevy_color::{Gray, LinearRgba};
 use bevy_core_pipeline::{
     schedule::{Core2d, Core2dSystems, Core3d, Core3dSystems},
     tonemapping::tonemapping,
 };
 use bevy_ecs::prelude::*;
+use bevy_image::ToExtents;
 use bevy_math::{ops, UVec2};
 use bevy_render::{
     camera::ExtractedCamera,
     diagnostic::RecordDiagnostics,
     extract_component::{
         ComponentUniforms, DynamicUniformIndex, ExtractComponentPlugin, UniformComponentPlugin,
     },
+    render_asset::RenderAssets,
     render_resource::*,
     renderer::{RenderContext, RenderDevice, ViewQuery},
-    texture::{CachedTexture, TextureCache},
+    texture::{CachedTexture, FallbackImage, GpuImage, TextureCache},
     view::ViewTarget,
     GpuResourceAppExt, Render, RenderApp, RenderStartup, RenderSystems,
 };
-use downsampling_pipeline::{
-    prepare_downsampling_pipeline, BloomDownsamplingPipeline, BloomDownsamplingPipelineIds,
-    BloomUniforms,
-};
-use upsampling_pipeline::{
-    prepare_upsampling_pipeline, BloomUpsamplingPipeline, UpsamplingPipelineIds,
-};
+use core::num::NonZero;
 
 const BLOOM_TEXTURE_FORMAT: TextureFormat = TextureFormat::Rg11b10Ufloat;
 
@@ -97,8 +98,11 @@ pub fn bloom(
         &BloomDownsamplingPipelineIds,
     )>,
     downsampling_pipeline_res: Res<BloomDownsamplingPipeline>,
+    upsampling_pipeline_res: Res<BloomUpsamplingPipeline>,
     pipeline_cache: Res<PipelineCache>,
     uniforms: Res<ComponentUniforms<BloomUniforms>>,
+    gpu_images: Res<RenderAssets<GpuImage>>,
+    fallback: Res<FallbackImage>,
     mut ctx: RenderContext,
 ) {
     let (
@@ -147,6 +151,40 @@ pub fn bloom(
         )),
     );
 
+    // Use dirt pipeline if ready, otherwise standard final pipeline.
+    let (final_pipeline, use_dirt) = match upsampling_pipeline_ids.id_final_dirt {
+        Some(dirt_id) => match pipeline_cache.get_render_pipeline(dirt_id) {
+            Some(dirt_pipeline) => (dirt_pipeline, true),
+            None => (upsampling_final_pipeline, false),
+        },
+        None => (upsampling_final_pipeline, false),
+    };
+
+    // Choose the appropriate final bind group.
+    let final_bind_group = match (use_dirt, &bloom_settings.lens_dirt.texture) {
+        (true, Some(lens_dirt_texture)) => {
+            let dirt_image = gpu_images.get(lens_dirt_texture).unwrap_or(&fallback.d2);
+            ctx.render_device().create_bind_group(
+                "bloom_final_dirt_bind_group",
+                &pipeline_cache
+                    .get_bind_group_layout(&upsampling_pipeline_res.dirt_bind_group_layout),
+                &BindGroupEntries::sequential((
+                    &bloom_texture.view(0),
+                    &bind_groups.sampler,
+                    uniforms_binding.clone(),
+                    BufferBinding {
+                        buffer: &bind_groups.blend_factor_buffers[0],
+                        offset: 0,
+                        size: NonZero::<u64>::new(4),
+                    },
+                    &dirt_image.texture_view,
+                    &bind_groups.sampler,
+                )),
+            )
+        }
+        _ => bind_groups.upsampling_bind_groups[(bloom_texture.mip_count - 1) as usize].clone(),
+    };
+
     let diagnostics = ctx.diagnostic_recorder();
     let diagnostics = diagnostics.as_deref();
     let time_span = diagnostics.time_span(ctx.command_encoder(), "bloom");
@@ -230,12 +268,6 @@ pub fn bloom(
             &bind_groups.upsampling_bind_groups[(bloom_texture.mip_count - mip - 1) as usize],
             &[uniform_index.index()],
         );
-        let blend = compute_blend_factor(
-            bloom_settings,
-            mip as f32,
-            (bloom_texture.mip_count - 1) as f32,
-        );
-        upsampling_pass.set_blend_constant(LinearRgba::gray(blend).into());
         upsampling_pass.draw(0..3, 0..1);
     }
 
@@ -249,12 +281,8 @@ pub fn bloom(
             occlusion_query_set: None,
             multiview_mask: None,
         });
-        upsampling_final_pass.set_pipeline(upsampling_final_pipeline);
-        upsampling_final_pass.set_bind_group(
-            0,
-            &bind_groups.upsampling_bind_groups[(bloom_texture.mip_count - 1) as usize],
-            &[uniform_index.index()],
-        );
+        upsampling_final_pass.set_pipeline(final_pipeline);
+        upsampling_final_pass.set_bind_group(0, &final_bind_group, &[uniform_index.index()]);
         if let Some(viewport) = camera.viewport.as_ref() {
             upsampling_final_pass.set_viewport(
                 viewport.physical_position.x as f32,
@@ -265,8 +293,6 @@ pub fn bloom(
                 viewport.depth.end,
             );
         }
-        let blend = compute_blend_factor(bloom_settings, 0.0, (bloom_texture.mip_count - 1) as f32);
-        upsampling_final_pass.set_blend_constant(LinearRgba::gray(blend).into());
         upsampling_final_pass.draw(0..3, 0..1);
     }
 
@@ -388,6 +414,7 @@ pub struct BloomBindGroups {
     cache_key: (Vec<TextureId>, BufferId),
     downsampling_bind_groups: Box<[BindGroup]>,
     upsampling_bind_groups: Box<[BindGroup]>,
+    blend_factor_buffers: Box<[Buffer]>,
     sampler: Sampler,
 }
 
@@ -396,13 +423,13 @@ fn prepare_bloom_bind_groups(
     render_device: Res<RenderDevice>,
     downsampling_pipeline: Res<BloomDownsamplingPipeline>,
     upsampling_pipeline: Res<BloomUpsamplingPipeline>,
-    views: Query<(Entity, &BloomTexture, Option<&BloomBindGroups>)>,
+    views: Query<(Entity, &BloomTexture, &Bloom, Option<&BloomBindGroups>)>,
     uniforms: Res<ComponentUniforms<BloomUniforms>>,
     pipeline_cache: Res<PipelineCache>,
 ) {
     let sampler = &downsampling_pipeline.sampler;
 
-    for (entity, bloom_texture, bloom_bind_groups) in &views {
+    for (entity, bloom_texture, bloom, bloom_bind_groups) in &views {
         #[cfg(any(
             not(feature = "webgl"),
             not(target_arch = "wasm32"),
@@ -443,6 +470,18 @@ fn prepare_bloom_bind_groups(
             ));
         }
 
+        let mip_count = bloom_texture.mip_count;
+        let blend_factor_buffers: Box<[Buffer]> = (0..mip_count)
+            .map(|mip| {
+                let factor = compute_blend_factor(bloom, mip as f32, (mip_count - 1) as f32);
+                render_device.create_buffer_with_data(&BufferInitDescriptor {
+                    label: Some("bloom_blend_factor"),
+                    contents: &factor.to_ne_bytes(),
+                    usage: BufferUsages::STORAGE,
+                })
+            })
+            .collect();
+
         let mut upsampling_bind_groups = Vec::with_capacity(bind_group_count);
         for mip in (0..bloom_texture.mip_count).rev() {
             upsampling_bind_groups.push(render_device.create_bind_group(
@@ -452,6 +491,11 @@ fn prepare_bloom_bind_groups(
                     &bloom_texture.view(mip),
                     sampler,
                     uniforms.binding().unwrap(),
+                    BufferBinding {
+                        buffer: &blend_factor_buffers[mip as usize],
+                        offset: 0,
+                        size: NonZero::<u64>::new(4),
+                    },
                 )),
             ));
         }
@@ -460,6 +504,7 @@ fn prepare_bloom_bind_groups(
             cache_key,
             downsampling_bind_groups: downsampling_bind_groups.into_boxed_slice(),
             upsampling_bind_groups: upsampling_bind_groups.into_boxed_slice(),
+            blend_factor_buffers,
             sampler: sampler.clone(),
         });
     }