bevy_post_process/bloom/

mod.rs

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,
};
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_math::{ops, UVec2};
use bevy_render::{
    camera::ExtractedCamera,
    diagnostic::RecordDiagnostics,
    extract_component::{
        ComponentUniforms, DynamicUniformIndex, ExtractComponentPlugin, UniformComponentPlugin,
    },
    render_resource::*,
    renderer::{RenderContext, RenderDevice, ViewQuery},
    texture::{CachedTexture, 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,
};

const BLOOM_TEXTURE_FORMAT: TextureFormat = TextureFormat::Rg11b10Ufloat;

#[derive(Default)]
pub struct BloomPlugin;

impl Plugin for BloomPlugin {
    fn build(&self, app: &mut App) {
        embedded_asset!(app, "bloom.wesl");

        app.add_plugins((
            ExtractComponentPlugin::<Bloom>::default(),
            UniformComponentPlugin::<BloomUniforms>::default(),
        ));

        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
            return;
        };
        render_app
            .init_gpu_resource::<SpecializedRenderPipelines<BloomDownsamplingPipeline>>()
            .init_gpu_resource::<SpecializedRenderPipelines<BloomUpsamplingPipeline>>()
            .add_systems(
                RenderStartup,
                (
                    init_bloom_downsampling_pipeline,
                    init_bloom_upscaling_pipeline,
                ),
            )
            .add_systems(
                Render,
                (
                    prepare_downsampling_pipeline.in_set(RenderSystems::Prepare),
                    prepare_upsampling_pipeline.in_set(RenderSystems::Prepare),
                    prepare_bloom_textures.in_set(RenderSystems::PrepareResources),
                    prepare_bloom_bind_groups.in_set(RenderSystems::PrepareBindGroups),
                ),
            )
            .add_systems(
                Core3d,
                bloom.before(tonemapping).in_set(Core3dSystems::PostProcess),
            )
            .add_systems(
                Core2d,
                bloom.before(tonemapping).in_set(Core2dSystems::PostProcess),
            );
    }
}

pub fn bloom(
    view: ViewQuery<(
        &ExtractedCamera,
        &ViewTarget,
        &BloomTexture,
        &BloomBindGroups,
        &DynamicUniformIndex<BloomUniforms>,
        &Bloom,
        &UpsamplingPipelineIds,
        &BloomDownsamplingPipelineIds,
    )>,
    downsampling_pipeline_res: Res<BloomDownsamplingPipeline>,
    pipeline_cache: Res<PipelineCache>,
    uniforms: Res<ComponentUniforms<BloomUniforms>>,
    mut ctx: RenderContext,
) {
    let (
        camera,
        view_target,
        bloom_texture,
        bind_groups,
        uniform_index,
        bloom_settings,
        upsampling_pipeline_ids,
        downsampling_pipeline_ids,
    ) = view.into_inner();

    if bloom_settings.intensity == 0.0 || !camera.hdr {
        return;
    }

    let (
        Some(uniforms_binding),
        Some(downsampling_first_pipeline),
        Some(downsampling_pipeline),
        Some(upsampling_pipeline),
        Some(upsampling_final_pipeline),
    ) = (
        uniforms.binding(),
        pipeline_cache.get_render_pipeline(downsampling_pipeline_ids.first),
        pipeline_cache.get_render_pipeline(downsampling_pipeline_ids.main),
        pipeline_cache.get_render_pipeline(upsampling_pipeline_ids.id_main),
        pipeline_cache.get_render_pipeline(upsampling_pipeline_ids.id_final),
    )
    else {
        return;
    };

    let view_texture = view_target.main_texture_view();
    let view_texture_unsampled = view_target.get_unsampled_color_attachment();

    // Create the first downsampling bind group (reads from main texture)
    let downsampling_first_bind_group = ctx.render_device().create_bind_group(
        "bloom_downsampling_first_bind_group",
        &pipeline_cache.get_bind_group_layout(&downsampling_pipeline_res.bind_group_layout),
        &BindGroupEntries::sequential((
            view_texture,
            &bind_groups.sampler,
            uniforms_binding.clone(),
        )),
    );

    let diagnostics = ctx.diagnostic_recorder();
    let diagnostics = diagnostics.as_deref();
    let time_span = diagnostics.time_span(ctx.command_encoder(), "bloom");

    let command_encoder = ctx.command_encoder();
    command_encoder.push_debug_group("bloom");

    // First downsample pass
    {
        let view = &bloom_texture.view(0);
        let mut downsampling_first_pass =
            command_encoder.begin_render_pass(&RenderPassDescriptor {
                label: Some("bloom_downsampling_first_pass"),
                color_attachments: &[Some(RenderPassColorAttachment {
                    view,
                    depth_slice: None,
                    resolve_target: None,
                    ops: Operations::default(),
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
                multiview_mask: None,
            });
        downsampling_first_pass.set_pipeline(downsampling_first_pipeline);
        downsampling_first_pass.set_bind_group(
            0,
            &downsampling_first_bind_group,
            &[uniform_index.index()],
        );
        downsampling_first_pass.draw(0..3, 0..1);
    }

    // Other downsample passes
    for mip in 1..bloom_texture.mip_count {
        let view = &bloom_texture.view(mip);
        let mut downsampling_pass = command_encoder.begin_render_pass(&RenderPassDescriptor {
            label: Some("bloom_downsampling_pass"),
            color_attachments: &[Some(RenderPassColorAttachment {
                view,
                depth_slice: None,
                resolve_target: None,
                ops: Operations::default(),
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });
        downsampling_pass.set_pipeline(downsampling_pipeline);
        downsampling_pass.set_bind_group(
            0,
            &bind_groups.downsampling_bind_groups[mip as usize - 1],
            &[uniform_index.index()],
        );
        downsampling_pass.draw(0..3, 0..1);
    }

    // Upsample passes except the final one
    for mip in (1..bloom_texture.mip_count).rev() {
        let view = &bloom_texture.view(mip - 1);
        let mut upsampling_pass = command_encoder.begin_render_pass(&RenderPassDescriptor {
            label: Some("bloom_upsampling_pass"),
            color_attachments: &[Some(RenderPassColorAttachment {
                view,
                depth_slice: None,
                resolve_target: None,
                ops: Operations {
                    load: LoadOp::Load,
                    store: StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });
        upsampling_pass.set_pipeline(upsampling_pipeline);
        upsampling_pass.set_bind_group(
            0,
            &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);
    }

    // Final upsample pass
    {
        let mut upsampling_final_pass = command_encoder.begin_render_pass(&RenderPassDescriptor {
            label: Some("bloom_upsampling_final_pass"),
            color_attachments: &[Some(view_texture_unsampled)],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            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()],
        );
        if let Some(viewport) = camera.viewport.as_ref() {
            upsampling_final_pass.set_viewport(
                viewport.physical_position.x as f32,
                viewport.physical_position.y as f32,
                viewport.physical_size.x as f32,
                viewport.physical_size.y as f32,
                viewport.depth.start,
                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);
    }

    command_encoder.pop_debug_group();
    time_span.end(ctx.command_encoder());
}

#[derive(Component)]
pub struct BloomTexture {
    // First mip is half the screen resolution, successive mips are half the previous
    #[cfg(any(
        not(feature = "webgl"),
        not(target_arch = "wasm32"),
        feature = "webgpu"
    ))]
    texture: CachedTexture,
    // WebGL does not support binding specific mip levels for sampling, fallback to separate textures instead
    #[cfg(all(feature = "webgl", target_arch = "wasm32", not(feature = "webgpu")))]
    texture: Vec<CachedTexture>,
    mip_count: u32,
}

impl BloomTexture {
    #[cfg(any(
        not(feature = "webgl"),
        not(target_arch = "wasm32"),
        feature = "webgpu"
    ))]
    fn view(&self, base_mip_level: u32) -> TextureView {
        self.texture.texture.create_view(&TextureViewDescriptor {
            base_mip_level,
            mip_level_count: Some(1u32),
            ..Default::default()
        })
    }
    #[cfg(all(feature = "webgl", target_arch = "wasm32", not(feature = "webgpu")))]
    fn view(&self, base_mip_level: u32) -> TextureView {
        self.texture[base_mip_level as usize]
            .texture
            .create_view(&TextureViewDescriptor {
                base_mip_level: 0,
                mip_level_count: Some(1u32),
                ..Default::default()
            })
    }
}

fn prepare_bloom_textures(
    mut commands: Commands,
    mut texture_cache: ResMut<TextureCache>,
    render_device: Res<RenderDevice>,
    views: Query<(Entity, &ExtractedCamera, &Bloom)>,
) {
    for (entity, camera, bloom) in &views {
        if let Some(viewport) = camera.physical_viewport_size {
            // How many times we can halve the resolution minus one so we don't go unnecessarily low
            let mip_count = bloom.max_mip_dimension.ilog2().max(2) - 1;
            let mip_height_ratio = if viewport.y != 0 {
                bloom.max_mip_dimension as f32 / viewport.y as f32
            } else {
                0.
            };

            let texture_descriptor = TextureDescriptor {
                label: Some("bloom_texture"),
                size: (viewport.as_vec2() * mip_height_ratio)
                    .round()
                    .as_uvec2()
                    .max(UVec2::ONE)
                    .to_extents(),
                mip_level_count: mip_count,
                sample_count: 1,
                dimension: TextureDimension::D2,
                format: BLOOM_TEXTURE_FORMAT,
                usage: TextureUsages::RENDER_ATTACHMENT | TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            };

            #[cfg(any(
                not(feature = "webgl"),
                not(target_arch = "wasm32"),
                feature = "webgpu"
            ))]
            let texture = texture_cache.get(&render_device, texture_descriptor);
            #[cfg(all(feature = "webgl", target_arch = "wasm32", not(feature = "webgpu")))]
            let texture: Vec<CachedTexture> = (0..mip_count)
                .map(|mip| {
                    texture_cache.get(
                        &render_device,
                        TextureDescriptor {
                            size: Extent3d {
                                width: (texture_descriptor.size.width >> mip).max(1),
                                height: (texture_descriptor.size.height >> mip).max(1),
                                depth_or_array_layers: 1,
                            },
                            mip_level_count: 1,
                            ..texture_descriptor.clone()
                        },
                    )
                })
                .collect();

            commands
                .entity(entity)
                .insert(BloomTexture { texture, mip_count });
        }
    }
}

#[derive(Component)]
pub struct BloomBindGroups {
    #[cfg(any(
        not(feature = "webgl"),
        not(target_arch = "wasm32"),
        feature = "webgpu"
    ))]
    cache_key: (TextureId, BufferId),
    #[cfg(all(feature = "webgl", target_arch = "wasm32", not(feature = "webgpu")))]
    cache_key: (Vec<TextureId>, BufferId),
    downsampling_bind_groups: Box<[BindGroup]>,
    upsampling_bind_groups: Box<[BindGroup]>,
    sampler: Sampler,
}

fn prepare_bloom_bind_groups(
    mut commands: Commands,
    render_device: Res<RenderDevice>,
    downsampling_pipeline: Res<BloomDownsamplingPipeline>,
    upsampling_pipeline: Res<BloomUpsamplingPipeline>,
    views: Query<(Entity, &BloomTexture, Option<&BloomBindGroups>)>,
    uniforms: Res<ComponentUniforms<BloomUniforms>>,
    pipeline_cache: Res<PipelineCache>,
) {
    let sampler = &downsampling_pipeline.sampler;

    for (entity, bloom_texture, bloom_bind_groups) in &views {
        #[cfg(any(
            not(feature = "webgl"),
            not(target_arch = "wasm32"),
            feature = "webgpu"
        ))]
        let cache_key = (
            bloom_texture.texture.texture.id(),
            uniforms.buffer().unwrap().id(),
        );
        #[cfg(all(feature = "webgl", target_arch = "wasm32", not(feature = "webgpu")))]
        let cache_key = (
            bloom_texture
                .texture
                .iter()
                .map(|tex| tex.texture.id())
                .collect(),
            uniforms.buffer().unwrap().id(),
        );

        if let Some(b) = bloom_bind_groups
            && b.cache_key == cache_key
        {
            continue;
        }

        let bind_group_count = bloom_texture.mip_count as usize - 1;

        let mut downsampling_bind_groups = Vec::with_capacity(bind_group_count);
        for mip in 1..bloom_texture.mip_count {
            downsampling_bind_groups.push(render_device.create_bind_group(
                "bloom_downsampling_bind_group",
                &pipeline_cache.get_bind_group_layout(&downsampling_pipeline.bind_group_layout),
                &BindGroupEntries::sequential((
                    &bloom_texture.view(mip - 1),
                    sampler,
                    uniforms.binding().unwrap(),
                )),
            ));
        }

        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(
                "bloom_upsampling_bind_group",
                &pipeline_cache.get_bind_group_layout(&upsampling_pipeline.bind_group_layout),
                &BindGroupEntries::sequential((
                    &bloom_texture.view(mip),
                    sampler,
                    uniforms.binding().unwrap(),
                )),
            ));
        }

        commands.entity(entity).insert(BloomBindGroups {
            cache_key,
            downsampling_bind_groups: downsampling_bind_groups.into_boxed_slice(),
            upsampling_bind_groups: upsampling_bind_groups.into_boxed_slice(),
            sampler: sampler.clone(),
        });
    }
}

/// Calculates blend intensities of blur pyramid levels
/// during the upsampling + compositing stage.
///
/// The function assumes all pyramid levels are upsampled and
/// blended into higher frequency ones using this function to
/// calculate blend levels every time. The final (highest frequency)
/// pyramid level in not blended into anything therefore this function
/// is not applied to it. As a result, the *mip* parameter of 0 indicates
/// the second-highest frequency pyramid level (in our case that is the
/// 0th mip of the bloom texture with the original image being the
/// actual highest frequency level).
///
/// Parameters:
/// * `mip` - the index of the lower frequency pyramid level (0 - `max_mip`, where 0 indicates highest frequency mip but not the highest frequency image).
/// * `max_mip` - the index of the lowest frequency pyramid level.
///
/// This function can be visually previewed for all values of *mip* (normalized) with tweakable
/// [`Bloom`] parameters on [Desmos graphing calculator](https://www.desmos.com/calculator/ncc8xbhzzl).
fn compute_blend_factor(bloom: &Bloom, mip: f32, max_mip: f32) -> f32 {
    let mut lf_boost =
        (1.0 - ops::powf(
            1.0 - (mip / max_mip),
            1.0 / (1.0 - bloom.low_frequency_boost_curvature),
        )) * bloom.low_frequency_boost;
    let high_pass_lq = 1.0
        - (((mip / max_mip) - bloom.high_pass_frequency) / bloom.high_pass_frequency)
            .clamp(0.0, 1.0);
    lf_boost *= match bloom.composite_mode {
        BloomCompositeMode::EnergyConserving => 1.0 - bloom.intensity,
        BloomCompositeMode::Additive => 1.0,
    };

    (bloom.intensity + lf_boost) * high_pass_lq
}