use core::array;
use crate::core_3d::{
graph::{Core3d, Node3d},
prepare_core_3d_depth_textures,
};
use bevy_app::{App, Plugin};
use bevy_asset::{load_internal_asset, weak_handle, Handle};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::{
component::Component,
entity::Entity,
prelude::{resource_exists, Without},
query::{Or, QueryState, With},
resource::Resource,
schedule::IntoScheduleConfigs as _,
system::{lifetimeless::Read, Commands, Local, Query, Res, ResMut},
world::{FromWorld, World},
};
use bevy_math::{uvec2, UVec2, Vec4Swizzles as _};
use bevy_render::batching::gpu_preprocessing::GpuPreprocessingSupport;
use bevy_render::{
experimental::occlusion_culling::{
OcclusionCulling, OcclusionCullingSubview, OcclusionCullingSubviewEntities,
},
render_graph::{Node, NodeRunError, RenderGraphApp, RenderGraphContext},
render_resource::{
binding_types::{sampler, texture_2d, texture_2d_multisampled, texture_storage_2d},
BindGroup, BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries,
CachedComputePipelineId, ComputePassDescriptor, ComputePipeline, ComputePipelineDescriptor,
Extent3d, IntoBinding, PipelineCache, PushConstantRange, Sampler, SamplerBindingType,
SamplerDescriptor, Shader, ShaderStages, SpecializedComputePipeline,
SpecializedComputePipelines, StorageTextureAccess, TextureAspect, TextureDescriptor,
TextureDimension, TextureFormat, TextureSampleType, TextureUsages, TextureView,
TextureViewDescriptor, TextureViewDimension,
},
renderer::{RenderContext, RenderDevice},
texture::TextureCache,
view::{ExtractedView, NoIndirectDrawing, ViewDepthTexture},
Render, RenderApp, RenderSet,
};
use bitflags::bitflags;
use tracing::debug;
pub const DOWNSAMPLE_DEPTH_SHADER_HANDLE: Handle<Shader> =
weak_handle!("a09a149e-5922-4fa4-9170-3c1a13065364");
pub const DEPTH_PYRAMID_MIP_COUNT: usize = 12;
pub struct MipGenerationPlugin;
impl Plugin for MipGenerationPlugin {
fn build(&self, app: &mut App) {
load_internal_asset!(
app,
DOWNSAMPLE_DEPTH_SHADER_HANDLE,
"downsample_depth.wgsl",
Shader::from_wgsl
);
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app
.init_resource::<SpecializedComputePipelines<DownsampleDepthPipeline>>()
.add_render_graph_node::<DownsampleDepthNode>(Core3d, Node3d::EarlyDownsampleDepth)
.add_render_graph_node::<DownsampleDepthNode>(Core3d, Node3d::LateDownsampleDepth)
.add_render_graph_edges(
Core3d,
(
Node3d::EarlyPrepass,
Node3d::EarlyDeferredPrepass,
Node3d::EarlyDownsampleDepth,
Node3d::LatePrepass,
Node3d::LateDeferredPrepass,
),
)
.add_render_graph_edges(
Core3d,
(
Node3d::EndMainPass,
Node3d::LateDownsampleDepth,
Node3d::EndMainPassPostProcessing,
),
)
.add_systems(
Render,
create_downsample_depth_pipelines.in_set(RenderSet::Prepare),
)
.add_systems(
Render,
(
prepare_view_depth_pyramids,
prepare_downsample_depth_view_bind_groups,
)
.chain()
.in_set(RenderSet::PrepareResources)
.run_if(resource_exists::<DownsampleDepthPipelines>)
.after(prepare_core_3d_depth_textures),
);
}
fn finish(&self, app: &mut App) {
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app.init_resource::<DepthPyramidDummyTexture>();
}
}
pub struct DownsampleDepthNode {
main_view_query: QueryState<(
Read<ViewDepthPyramid>,
Read<ViewDownsampleDepthBindGroup>,
Read<ViewDepthTexture>,
Option<Read<OcclusionCullingSubviewEntities>>,
)>,
shadow_view_query: QueryState<(
Read<ViewDepthPyramid>,
Read<ViewDownsampleDepthBindGroup>,
Read<OcclusionCullingSubview>,
)>,
}
impl FromWorld for DownsampleDepthNode {
fn from_world(world: &mut World) -> Self {
Self {
main_view_query: QueryState::new(world),
shadow_view_query: QueryState::new(world),
}
}
}
impl Node for DownsampleDepthNode {
fn update(&mut self, world: &mut World) {
self.main_view_query.update_archetypes(world);
self.shadow_view_query.update_archetypes(world);
}
fn run<'w>(
&self,
render_graph_context: &mut RenderGraphContext,
render_context: &mut RenderContext<'w>,
world: &'w World,
) -> Result<(), NodeRunError> {
let Ok((
view_depth_pyramid,
view_downsample_depth_bind_group,
view_depth_texture,
maybe_view_light_entities,
)) = self
.main_view_query
.get_manual(world, render_graph_context.view_entity())
else {
return Ok(());
};
downsample_depth(
render_graph_context,
render_context,
world,
view_depth_pyramid,
view_downsample_depth_bind_group,
uvec2(
view_depth_texture.texture.width(),
view_depth_texture.texture.height(),
),
view_depth_texture.texture.sample_count(),
)?;
if let Some(view_light_entities) = maybe_view_light_entities {
for &view_light_entity in &view_light_entities.0 {
let Ok((view_depth_pyramid, view_downsample_depth_bind_group, occlusion_culling)) =
self.shadow_view_query.get_manual(world, view_light_entity)
else {
continue;
};
downsample_depth(
render_graph_context,
render_context,
world,
view_depth_pyramid,
view_downsample_depth_bind_group,
UVec2::splat(occlusion_culling.depth_texture_size),
1,
)?;
}
}
Ok(())
}
}
fn downsample_depth<'w>(
render_graph_context: &mut RenderGraphContext,
render_context: &mut RenderContext<'w>,
world: &'w World,
view_depth_pyramid: &ViewDepthPyramid,
view_downsample_depth_bind_group: &ViewDownsampleDepthBindGroup,
view_size: UVec2,
sample_count: u32,
) -> Result<(), NodeRunError> {
let downsample_depth_pipelines = world.resource::<DownsampleDepthPipelines>();
let pipeline_cache = world.resource::<PipelineCache>();
let (Some(first_downsample_depth_pipeline_id), Some(second_downsample_depth_pipeline_id)) =
(if sample_count > 1 {
(
downsample_depth_pipelines.first_multisample.pipeline_id,
downsample_depth_pipelines.second_multisample.pipeline_id,
)
} else {
(
downsample_depth_pipelines.first.pipeline_id,
downsample_depth_pipelines.second.pipeline_id,
)
})
else {
return Ok(());
};
let (Some(first_downsample_depth_pipeline), Some(second_downsample_depth_pipeline)) = (
pipeline_cache.get_compute_pipeline(first_downsample_depth_pipeline_id),
pipeline_cache.get_compute_pipeline(second_downsample_depth_pipeline_id),
) else {
return Ok(());
};
view_depth_pyramid.downsample_depth(
&format!("{:?}", render_graph_context.label()),
render_context,
view_size,
view_downsample_depth_bind_group,
first_downsample_depth_pipeline,
second_downsample_depth_pipeline,
);
Ok(())
}
#[derive(Resource)]
pub struct DownsampleDepthPipeline {
bind_group_layout: BindGroupLayout,
pipeline_id: Option<CachedComputePipelineId>,
}
impl DownsampleDepthPipeline {
fn new(bind_group_layout: BindGroupLayout) -> DownsampleDepthPipeline {
DownsampleDepthPipeline {
bind_group_layout,
pipeline_id: None,
}
}
}
#[derive(Resource)]
pub struct DownsampleDepthPipelines {
first: DownsampleDepthPipeline,
second: DownsampleDepthPipeline,
first_multisample: DownsampleDepthPipeline,
second_multisample: DownsampleDepthPipeline,
sampler: Sampler,
}
fn create_downsample_depth_pipelines(
mut commands: Commands,
render_device: Res<RenderDevice>,
pipeline_cache: Res<PipelineCache>,
mut specialized_compute_pipelines: ResMut<SpecializedComputePipelines<DownsampleDepthPipeline>>,
gpu_preprocessing_support: Res<GpuPreprocessingSupport>,
mut has_run: Local<bool>,
) {
if *has_run {
return;
}
*has_run = true;
if !gpu_preprocessing_support.is_culling_supported() {
debug!("Downsample depth is not supported on this platform.");
return;
}
let standard_bind_group_layout =
create_downsample_depth_bind_group_layout(&render_device, false);
let multisampled_bind_group_layout =
create_downsample_depth_bind_group_layout(&render_device, true);
let sampler = render_device.create_sampler(&SamplerDescriptor {
label: Some("depth pyramid sampler"),
..SamplerDescriptor::default()
});
let mut downsample_depth_pipelines = DownsampleDepthPipelines {
first: DownsampleDepthPipeline::new(standard_bind_group_layout.clone()),
second: DownsampleDepthPipeline::new(standard_bind_group_layout.clone()),
first_multisample: DownsampleDepthPipeline::new(multisampled_bind_group_layout.clone()),
second_multisample: DownsampleDepthPipeline::new(multisampled_bind_group_layout.clone()),
sampler,
};
downsample_depth_pipelines.first.pipeline_id = Some(specialized_compute_pipelines.specialize(
&pipeline_cache,
&downsample_depth_pipelines.first,
DownsampleDepthPipelineKey::empty(),
));
downsample_depth_pipelines.second.pipeline_id = Some(specialized_compute_pipelines.specialize(
&pipeline_cache,
&downsample_depth_pipelines.second,
DownsampleDepthPipelineKey::SECOND_PHASE,
));
downsample_depth_pipelines.first_multisample.pipeline_id =
Some(specialized_compute_pipelines.specialize(
&pipeline_cache,
&downsample_depth_pipelines.first_multisample,
DownsampleDepthPipelineKey::MULTISAMPLE,
));
downsample_depth_pipelines.second_multisample.pipeline_id =
Some(specialized_compute_pipelines.specialize(
&pipeline_cache,
&downsample_depth_pipelines.second_multisample,
DownsampleDepthPipelineKey::SECOND_PHASE | DownsampleDepthPipelineKey::MULTISAMPLE,
));
commands.insert_resource(downsample_depth_pipelines);
}
fn create_downsample_depth_bind_group_layout(
render_device: &RenderDevice,
is_multisampled: bool,
) -> BindGroupLayout {
render_device.create_bind_group_layout(
if is_multisampled {
"downsample multisample depth bind group layout"
} else {
"downsample depth bind group layout"
},
&BindGroupLayoutEntries::sequential(
ShaderStages::COMPUTE,
(
if is_multisampled {
texture_2d_multisampled(TextureSampleType::Depth)
} else {
texture_2d(TextureSampleType::Depth)
},
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::ReadWrite),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
sampler(SamplerBindingType::NonFiltering),
),
),
)
}
bitflags! {
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct DownsampleDepthPipelineKey: u8 {
const MULTISAMPLE = 1;
const SECOND_PHASE = 2;
}
}
impl SpecializedComputePipeline for DownsampleDepthPipeline {
type Key = DownsampleDepthPipelineKey;
fn specialize(&self, key: Self::Key) -> ComputePipelineDescriptor {
let mut shader_defs = vec![];
if key.contains(DownsampleDepthPipelineKey::MULTISAMPLE) {
shader_defs.push("MULTISAMPLE".into());
}
let label = format!(
"downsample depth{}{} pipeline",
if key.contains(DownsampleDepthPipelineKey::MULTISAMPLE) {
" multisample"
} else {
""
},
if key.contains(DownsampleDepthPipelineKey::SECOND_PHASE) {
" second phase"
} else {
" first phase"
}
)
.into();
ComputePipelineDescriptor {
label: Some(label),
layout: vec![self.bind_group_layout.clone()],
push_constant_ranges: vec![PushConstantRange {
stages: ShaderStages::COMPUTE,
range: 0..4,
}],
shader: DOWNSAMPLE_DEPTH_SHADER_HANDLE,
shader_defs,
entry_point: if key.contains(DownsampleDepthPipelineKey::SECOND_PHASE) {
"downsample_depth_second".into()
} else {
"downsample_depth_first".into()
},
zero_initialize_workgroup_memory: false,
}
}
}
#[derive(Resource, Deref, DerefMut)]
pub struct DepthPyramidDummyTexture(TextureView);
impl FromWorld for DepthPyramidDummyTexture {
fn from_world(world: &mut World) -> Self {
let render_device = world.resource::<RenderDevice>();
DepthPyramidDummyTexture(create_depth_pyramid_dummy_texture(
render_device,
"depth pyramid dummy texture",
"depth pyramid dummy texture view",
))
}
}
pub fn create_depth_pyramid_dummy_texture(
render_device: &RenderDevice,
texture_label: &'static str,
texture_view_label: &'static str,
) -> TextureView {
render_device
.create_texture(&TextureDescriptor {
label: Some(texture_label),
size: Extent3d {
width: 1,
height: 1,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: TextureDimension::D2,
format: TextureFormat::R32Float,
usage: TextureUsages::STORAGE_BINDING,
view_formats: &[],
})
.create_view(&TextureViewDescriptor {
label: Some(texture_view_label),
format: Some(TextureFormat::R32Float),
dimension: Some(TextureViewDimension::D2),
usage: None,
aspect: TextureAspect::All,
base_mip_level: 0,
mip_level_count: Some(1),
base_array_layer: 0,
array_layer_count: Some(1),
})
}
#[derive(Component)]
pub struct ViewDepthPyramid {
pub all_mips: TextureView,
pub mips: [TextureView; DEPTH_PYRAMID_MIP_COUNT],
pub mip_count: u32,
}
impl ViewDepthPyramid {
pub fn new(
render_device: &RenderDevice,
texture_cache: &mut TextureCache,
depth_pyramid_dummy_texture: &TextureView,
size: UVec2,
texture_label: &'static str,
texture_view_label: &'static str,
) -> ViewDepthPyramid {
let depth_pyramid_size = Extent3d {
width: size.x.div_ceil(2),
height: size.y.div_ceil(2),
depth_or_array_layers: 1,
};
let depth_pyramid_mip_count = depth_pyramid_size.max_mips(TextureDimension::D2);
let depth_pyramid = texture_cache.get(
render_device,
TextureDescriptor {
label: Some(texture_label),
size: depth_pyramid_size,
mip_level_count: depth_pyramid_mip_count,
sample_count: 1,
dimension: TextureDimension::D2,
format: TextureFormat::R32Float,
usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
view_formats: &[],
},
);
let depth_pyramid_mips = array::from_fn(|i| {
if (i as u32) < depth_pyramid_mip_count {
depth_pyramid.texture.create_view(&TextureViewDescriptor {
label: Some(texture_view_label),
format: Some(TextureFormat::R32Float),
dimension: Some(TextureViewDimension::D2),
usage: None,
aspect: TextureAspect::All,
base_mip_level: i as u32,
mip_level_count: Some(1),
base_array_layer: 0,
array_layer_count: Some(1),
})
} else {
(*depth_pyramid_dummy_texture).clone()
}
});
let depth_pyramid_all_mips = depth_pyramid.default_view.clone();
Self {
all_mips: depth_pyramid_all_mips,
mips: depth_pyramid_mips,
mip_count: depth_pyramid_mip_count,
}
}
pub fn create_bind_group<'a, R>(
&'a self,
render_device: &RenderDevice,
label: &'static str,
bind_group_layout: &BindGroupLayout,
source_image: R,
sampler: &'a Sampler,
) -> BindGroup
where
R: IntoBinding<'a>,
{
render_device.create_bind_group(
label,
bind_group_layout,
&BindGroupEntries::sequential((
source_image,
&self.mips[0],
&self.mips[1],
&self.mips[2],
&self.mips[3],
&self.mips[4],
&self.mips[5],
&self.mips[6],
&self.mips[7],
&self.mips[8],
&self.mips[9],
&self.mips[10],
&self.mips[11],
sampler,
)),
)
}
pub fn downsample_depth(
&self,
label: &str,
render_context: &mut RenderContext,
view_size: UVec2,
downsample_depth_bind_group: &BindGroup,
downsample_depth_first_pipeline: &ComputePipeline,
downsample_depth_second_pipeline: &ComputePipeline,
) {
let command_encoder = render_context.command_encoder();
let mut downsample_pass = command_encoder.begin_compute_pass(&ComputePassDescriptor {
label: Some(label),
timestamp_writes: None,
});
downsample_pass.set_pipeline(downsample_depth_first_pipeline);
downsample_pass.set_push_constants(0, &self.mip_count.to_le_bytes());
downsample_pass.set_bind_group(0, downsample_depth_bind_group, &[]);
downsample_pass.dispatch_workgroups(view_size.x.div_ceil(64), view_size.y.div_ceil(64), 1);
if self.mip_count >= 7 {
downsample_pass.set_pipeline(downsample_depth_second_pipeline);
downsample_pass.dispatch_workgroups(1, 1, 1);
}
}
}
pub fn prepare_view_depth_pyramids(
mut commands: Commands,
render_device: Res<RenderDevice>,
mut texture_cache: ResMut<TextureCache>,
depth_pyramid_dummy_texture: Res<DepthPyramidDummyTexture>,
views: Query<(Entity, &ExtractedView), (With<OcclusionCulling>, Without<NoIndirectDrawing>)>,
) {
for (view_entity, view) in &views {
commands.entity(view_entity).insert(ViewDepthPyramid::new(
&render_device,
&mut texture_cache,
&depth_pyramid_dummy_texture,
view.viewport.zw(),
"view depth pyramid texture",
"view depth pyramid texture view",
));
}
}
#[derive(Component, Deref, DerefMut)]
pub struct ViewDownsampleDepthBindGroup(BindGroup);
fn prepare_downsample_depth_view_bind_groups(
mut commands: Commands,
render_device: Res<RenderDevice>,
downsample_depth_pipelines: Res<DownsampleDepthPipelines>,
view_depth_textures: Query<
(
Entity,
&ViewDepthPyramid,
Option<&ViewDepthTexture>,
Option<&OcclusionCullingSubview>,
),
Or<(With<ViewDepthTexture>, With<OcclusionCullingSubview>)>,
>,
) {
for (view_entity, view_depth_pyramid, view_depth_texture, shadow_occlusion_culling) in
&view_depth_textures
{
let is_multisampled = view_depth_texture
.is_some_and(|view_depth_texture| view_depth_texture.texture.sample_count() > 1);
commands
.entity(view_entity)
.insert(ViewDownsampleDepthBindGroup(
view_depth_pyramid.create_bind_group(
&render_device,
if is_multisampled {
"downsample multisample depth bind group"
} else {
"downsample depth bind group"
},
if is_multisampled {
&downsample_depth_pipelines
.first_multisample
.bind_group_layout
} else {
&downsample_depth_pipelines.first.bind_group_layout
},
match (view_depth_texture, shadow_occlusion_culling) {
(Some(view_depth_texture), _) => view_depth_texture.view(),
(None, Some(shadow_occlusion_culling)) => {
&shadow_occlusion_culling.depth_texture_view
}
(None, None) => panic!("Should never happen"),
},
&downsample_depth_pipelines.sampler,
),
));
}
}