1use core::array;
9
10use crate::core_3d::{
11 graph::{Core3d, Node3d},
12 prepare_core_3d_depth_textures,
13};
14use bevy_app::{App, Plugin};
15use bevy_asset::{embedded_asset, load_embedded_asset, Handle};
16use bevy_derive::{Deref, DerefMut};
17use bevy_ecs::{
18 component::Component,
19 entity::Entity,
20 prelude::{resource_exists, Without},
21 query::{Or, QueryState, With},
22 resource::Resource,
23 schedule::IntoScheduleConfigs as _,
24 system::{lifetimeless::Read, Commands, Local, Query, Res, ResMut},
25 world::{FromWorld, World},
26};
27use bevy_math::{uvec2, UVec2, Vec4Swizzles as _};
28use bevy_render::{
29 batching::gpu_preprocessing::GpuPreprocessingSupport,
30 render_resource::BindGroupLayoutDescriptor, RenderStartup,
31};
32use bevy_render::{
33 experimental::occlusion_culling::{
34 OcclusionCulling, OcclusionCullingSubview, OcclusionCullingSubviewEntities,
35 },
36 render_graph::{Node, NodeRunError, RenderGraphContext, RenderGraphExt},
37 render_resource::{
38 binding_types::{sampler, texture_2d, texture_2d_multisampled, texture_storage_2d},
39 BindGroup, BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries,
40 CachedComputePipelineId, ComputePassDescriptor, ComputePipeline, ComputePipelineDescriptor,
41 Extent3d, IntoBinding, PipelineCache, PushConstantRange, Sampler, SamplerBindingType,
42 SamplerDescriptor, ShaderStages, SpecializedComputePipeline, SpecializedComputePipelines,
43 StorageTextureAccess, TextureAspect, TextureDescriptor, TextureDimension, TextureFormat,
44 TextureSampleType, TextureUsages, TextureView, TextureViewDescriptor, TextureViewDimension,
45 },
46 renderer::{RenderContext, RenderDevice},
47 texture::TextureCache,
48 view::{ExtractedView, NoIndirectDrawing, ViewDepthTexture},
49 Render, RenderApp, RenderSystems,
50};
51use bevy_shader::Shader;
52use bevy_utils::default;
53use bitflags::bitflags;
54use tracing::debug;
55
56#[derive(Resource, Deref)]
58pub struct DownsampleDepthShader(Handle<Shader>);
59
60pub const DEPTH_PYRAMID_MIP_COUNT: usize = 12;
65
66pub struct MipGenerationPlugin;
72
73impl Plugin for MipGenerationPlugin {
74 fn build(&self, app: &mut App) {
75 embedded_asset!(app, "downsample_depth.wgsl");
76
77 let downsample_depth_shader = load_embedded_asset!(app, "downsample_depth.wgsl");
78
79 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
80 return;
81 };
82
83 render_app
84 .insert_resource(DownsampleDepthShader(downsample_depth_shader))
85 .init_resource::<SpecializedComputePipelines<DownsampleDepthPipeline>>()
86 .add_render_graph_node::<DownsampleDepthNode>(Core3d, Node3d::EarlyDownsampleDepth)
87 .add_render_graph_node::<DownsampleDepthNode>(Core3d, Node3d::LateDownsampleDepth)
88 .add_render_graph_edges(
89 Core3d,
90 (
91 Node3d::EarlyPrepass,
92 Node3d::EarlyDeferredPrepass,
93 Node3d::EarlyDownsampleDepth,
94 Node3d::LatePrepass,
95 Node3d::LateDeferredPrepass,
96 ),
97 )
98 .add_render_graph_edges(
99 Core3d,
100 (
101 Node3d::StartMainPassPostProcessing,
102 Node3d::LateDownsampleDepth,
103 Node3d::EndMainPassPostProcessing,
104 ),
105 )
106 .add_systems(RenderStartup, init_depth_pyramid_dummy_texture)
107 .add_systems(
108 Render,
109 create_downsample_depth_pipelines.in_set(RenderSystems::Prepare),
110 )
111 .add_systems(
112 Render,
113 (
114 prepare_view_depth_pyramids,
115 prepare_downsample_depth_view_bind_groups,
116 )
117 .chain()
118 .in_set(RenderSystems::PrepareResources)
119 .run_if(resource_exists::<DownsampleDepthPipelines>)
120 .after(prepare_core_3d_depth_textures),
121 );
122 }
123}
124
125pub struct DownsampleDepthNode {
142 main_view_query: QueryState<(
145 Read<ViewDepthPyramid>,
146 Read<ViewDownsampleDepthBindGroup>,
147 Read<ViewDepthTexture>,
148 Option<Read<OcclusionCullingSubviewEntities>>,
149 )>,
150 shadow_view_query: QueryState<(
152 Read<ViewDepthPyramid>,
153 Read<ViewDownsampleDepthBindGroup>,
154 Read<OcclusionCullingSubview>,
155 )>,
156}
157
158impl FromWorld for DownsampleDepthNode {
159 fn from_world(world: &mut World) -> Self {
160 Self {
161 main_view_query: QueryState::new(world),
162 shadow_view_query: QueryState::new(world),
163 }
164 }
165}
166
167impl Node for DownsampleDepthNode {
168 fn update(&mut self, world: &mut World) {
169 self.main_view_query.update_archetypes(world);
170 self.shadow_view_query.update_archetypes(world);
171 }
172
173 fn run<'w>(
174 &self,
175 render_graph_context: &mut RenderGraphContext,
176 render_context: &mut RenderContext<'w>,
177 world: &'w World,
178 ) -> Result<(), NodeRunError> {
179 let Ok((
180 view_depth_pyramid,
181 view_downsample_depth_bind_group,
182 view_depth_texture,
183 maybe_view_light_entities,
184 )) = self
185 .main_view_query
186 .get_manual(world, render_graph_context.view_entity())
187 else {
188 return Ok(());
189 };
190
191 downsample_depth(
193 render_graph_context,
194 render_context,
195 world,
196 view_depth_pyramid,
197 view_downsample_depth_bind_group,
198 uvec2(
199 view_depth_texture.texture.width(),
200 view_depth_texture.texture.height(),
201 ),
202 view_depth_texture.texture.sample_count(),
203 )?;
204
205 if let Some(view_light_entities) = maybe_view_light_entities {
207 for &view_light_entity in &view_light_entities.0 {
208 let Ok((view_depth_pyramid, view_downsample_depth_bind_group, occlusion_culling)) =
209 self.shadow_view_query.get_manual(world, view_light_entity)
210 else {
211 continue;
212 };
213 downsample_depth(
214 render_graph_context,
215 render_context,
216 world,
217 view_depth_pyramid,
218 view_downsample_depth_bind_group,
219 UVec2::splat(occlusion_culling.depth_texture_size),
220 1,
221 )?;
222 }
223 }
224
225 Ok(())
226 }
227}
228
229fn downsample_depth<'w>(
232 render_graph_context: &mut RenderGraphContext,
233 render_context: &mut RenderContext<'w>,
234 world: &'w World,
235 view_depth_pyramid: &ViewDepthPyramid,
236 view_downsample_depth_bind_group: &ViewDownsampleDepthBindGroup,
237 view_size: UVec2,
238 sample_count: u32,
239) -> Result<(), NodeRunError> {
240 let downsample_depth_pipelines = world.resource::<DownsampleDepthPipelines>();
241 let pipeline_cache = world.resource::<PipelineCache>();
242
243 let (Some(first_downsample_depth_pipeline_id), Some(second_downsample_depth_pipeline_id)) =
250 (if sample_count > 1 {
251 (
252 downsample_depth_pipelines.first_multisample.pipeline_id,
253 downsample_depth_pipelines.second_multisample.pipeline_id,
254 )
255 } else {
256 (
257 downsample_depth_pipelines.first.pipeline_id,
258 downsample_depth_pipelines.second.pipeline_id,
259 )
260 })
261 else {
262 return Ok(());
263 };
264
265 let (Some(first_downsample_depth_pipeline), Some(second_downsample_depth_pipeline)) = (
267 pipeline_cache.get_compute_pipeline(first_downsample_depth_pipeline_id),
268 pipeline_cache.get_compute_pipeline(second_downsample_depth_pipeline_id),
269 ) else {
270 return Ok(());
271 };
272
273 view_depth_pyramid.downsample_depth(
275 &format!("{:?}", render_graph_context.label()),
276 render_context,
277 view_size,
278 view_downsample_depth_bind_group,
279 first_downsample_depth_pipeline,
280 second_downsample_depth_pipeline,
281 );
282 Ok(())
283}
284
285#[derive(Resource)]
287pub struct DownsampleDepthPipeline {
288 bind_group_layout: BindGroupLayoutDescriptor,
290 pipeline_id: Option<CachedComputePipelineId>,
292 shader: Handle<Shader>,
294}
295
296impl DownsampleDepthPipeline {
297 fn new(
303 bind_group_layout: BindGroupLayoutDescriptor,
304 shader: Handle<Shader>,
305 ) -> DownsampleDepthPipeline {
306 DownsampleDepthPipeline {
307 bind_group_layout,
308 pipeline_id: None,
309 shader,
310 }
311 }
312}
313
314#[derive(Resource)]
316pub struct DownsampleDepthPipelines {
317 first: DownsampleDepthPipeline,
320 second: DownsampleDepthPipeline,
323 first_multisample: DownsampleDepthPipeline,
325 second_multisample: DownsampleDepthPipeline,
327 sampler: Sampler,
330}
331
332fn create_downsample_depth_pipelines(
335 mut commands: Commands,
336 render_device: Res<RenderDevice>,
337 pipeline_cache: Res<PipelineCache>,
338 mut specialized_compute_pipelines: ResMut<SpecializedComputePipelines<DownsampleDepthPipeline>>,
339 gpu_preprocessing_support: Res<GpuPreprocessingSupport>,
340 downsample_depth_shader: Res<DownsampleDepthShader>,
341 mut has_run: Local<bool>,
342) {
343 if *has_run {
348 return;
349 }
350 *has_run = true;
351
352 if !gpu_preprocessing_support.is_culling_supported() {
353 debug!("Downsample depth is not supported on this platform.");
354 return;
355 }
356
357 let standard_bind_group_layout = create_downsample_depth_bind_group_layout(false);
362 let multisampled_bind_group_layout = create_downsample_depth_bind_group_layout(true);
363
364 let sampler = render_device.create_sampler(&SamplerDescriptor {
366 label: Some("depth pyramid sampler"),
367 ..SamplerDescriptor::default()
368 });
369
370 let mut downsample_depth_pipelines = DownsampleDepthPipelines {
372 first: DownsampleDepthPipeline::new(
373 standard_bind_group_layout.clone(),
374 downsample_depth_shader.0.clone(),
375 ),
376 second: DownsampleDepthPipeline::new(
377 standard_bind_group_layout.clone(),
378 downsample_depth_shader.0.clone(),
379 ),
380 first_multisample: DownsampleDepthPipeline::new(
381 multisampled_bind_group_layout.clone(),
382 downsample_depth_shader.0.clone(),
383 ),
384 second_multisample: DownsampleDepthPipeline::new(
385 multisampled_bind_group_layout.clone(),
386 downsample_depth_shader.0.clone(),
387 ),
388 sampler,
389 };
390
391 downsample_depth_pipelines.first.pipeline_id = Some(specialized_compute_pipelines.specialize(
394 &pipeline_cache,
395 &downsample_depth_pipelines.first,
396 DownsampleDepthPipelineKey::empty(),
397 ));
398 downsample_depth_pipelines.second.pipeline_id = Some(specialized_compute_pipelines.specialize(
399 &pipeline_cache,
400 &downsample_depth_pipelines.second,
401 DownsampleDepthPipelineKey::SECOND_PHASE,
402 ));
403 downsample_depth_pipelines.first_multisample.pipeline_id =
404 Some(specialized_compute_pipelines.specialize(
405 &pipeline_cache,
406 &downsample_depth_pipelines.first_multisample,
407 DownsampleDepthPipelineKey::MULTISAMPLE,
408 ));
409 downsample_depth_pipelines.second_multisample.pipeline_id =
410 Some(specialized_compute_pipelines.specialize(
411 &pipeline_cache,
412 &downsample_depth_pipelines.second_multisample,
413 DownsampleDepthPipelineKey::SECOND_PHASE | DownsampleDepthPipelineKey::MULTISAMPLE,
414 ));
415
416 commands.insert_resource(downsample_depth_pipelines);
417}
418
419fn create_downsample_depth_bind_group_layout(is_multisampled: bool) -> BindGroupLayoutDescriptor {
421 BindGroupLayoutDescriptor::new(
422 if is_multisampled {
423 "downsample multisample depth bind group layout"
424 } else {
425 "downsample depth bind group layout"
426 },
427 &BindGroupLayoutEntries::sequential(
428 ShaderStages::COMPUTE,
429 (
430 if is_multisampled {
434 texture_2d_multisampled(TextureSampleType::Depth)
435 } else {
436 texture_2d(TextureSampleType::Depth)
437 },
438 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
440 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
441 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
442 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
443 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
444 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::ReadWrite),
445 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
446 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
447 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
448 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
449 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
450 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
451 sampler(SamplerBindingType::NonFiltering),
452 ),
453 ),
454 )
455}
456
457bitflags! {
458 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
464 pub struct DownsampleDepthPipelineKey: u8 {
465 const MULTISAMPLE = 1;
467 const SECOND_PHASE = 2;
470 }
471}
472
473impl SpecializedComputePipeline for DownsampleDepthPipeline {
474 type Key = DownsampleDepthPipelineKey;
475
476 fn specialize(&self, key: Self::Key) -> ComputePipelineDescriptor {
477 let mut shader_defs = vec![];
478 if key.contains(DownsampleDepthPipelineKey::MULTISAMPLE) {
479 shader_defs.push("MULTISAMPLE".into());
480 }
481
482 let label = format!(
483 "downsample depth{}{} pipeline",
484 if key.contains(DownsampleDepthPipelineKey::MULTISAMPLE) {
485 " multisample"
486 } else {
487 ""
488 },
489 if key.contains(DownsampleDepthPipelineKey::SECOND_PHASE) {
490 " second phase"
491 } else {
492 " first phase"
493 }
494 )
495 .into();
496
497 ComputePipelineDescriptor {
498 label: Some(label),
499 layout: vec![self.bind_group_layout.clone()],
500 push_constant_ranges: vec![PushConstantRange {
501 stages: ShaderStages::COMPUTE,
502 range: 0..4,
503 }],
504 shader: self.shader.clone(),
505 shader_defs,
506 entry_point: Some(if key.contains(DownsampleDepthPipelineKey::SECOND_PHASE) {
507 "downsample_depth_second".into()
508 } else {
509 "downsample_depth_first".into()
510 }),
511 ..default()
512 }
513 }
514}
515
516#[derive(Resource, Deref, DerefMut)]
519pub struct DepthPyramidDummyTexture(TextureView);
520
521pub fn init_depth_pyramid_dummy_texture(mut commands: Commands, render_device: Res<RenderDevice>) {
522 commands.insert_resource(DepthPyramidDummyTexture(
523 create_depth_pyramid_dummy_texture(
524 &render_device,
525 "depth pyramid dummy texture",
526 "depth pyramid dummy texture view",
527 ),
528 ));
529}
530
531pub fn create_depth_pyramid_dummy_texture(
534 render_device: &RenderDevice,
535 texture_label: &'static str,
536 texture_view_label: &'static str,
537) -> TextureView {
538 render_device
539 .create_texture(&TextureDescriptor {
540 label: Some(texture_label),
541 size: Extent3d::default(),
542 mip_level_count: 1,
543 sample_count: 1,
544 dimension: TextureDimension::D2,
545 format: TextureFormat::R32Float,
546 usage: TextureUsages::STORAGE_BINDING,
547 view_formats: &[],
548 })
549 .create_view(&TextureViewDescriptor {
550 label: Some(texture_view_label),
551 format: Some(TextureFormat::R32Float),
552 dimension: Some(TextureViewDimension::D2),
553 usage: None,
554 aspect: TextureAspect::All,
555 base_mip_level: 0,
556 mip_level_count: Some(1),
557 base_array_layer: 0,
558 array_layer_count: Some(1),
559 })
560}
561
562#[derive(Component)]
567pub struct ViewDepthPyramid {
568 pub all_mips: TextureView,
570 pub mips: [TextureView; DEPTH_PYRAMID_MIP_COUNT],
572 pub mip_count: u32,
577}
578
579impl ViewDepthPyramid {
580 pub fn new(
582 render_device: &RenderDevice,
583 texture_cache: &mut TextureCache,
584 depth_pyramid_dummy_texture: &TextureView,
585 size: UVec2,
586 texture_label: &'static str,
587 texture_view_label: &'static str,
588 ) -> ViewDepthPyramid {
589 let depth_pyramid_size = Extent3d {
591 width: size.x.div_ceil(2),
592 height: size.y.div_ceil(2),
593 depth_or_array_layers: 1,
594 };
595
596 let depth_pyramid_mip_count = depth_pyramid_size.max_mips(TextureDimension::D2);
598
599 let depth_pyramid = texture_cache.get(
601 render_device,
602 TextureDescriptor {
603 label: Some(texture_label),
604 size: depth_pyramid_size,
605 mip_level_count: depth_pyramid_mip_count,
606 sample_count: 1,
607 dimension: TextureDimension::D2,
608 format: TextureFormat::R32Float,
609 usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
610 view_formats: &[],
611 },
612 );
613
614 let depth_pyramid_mips = array::from_fn(|i| {
616 if (i as u32) < depth_pyramid_mip_count {
617 depth_pyramid.texture.create_view(&TextureViewDescriptor {
618 label: Some(texture_view_label),
619 format: Some(TextureFormat::R32Float),
620 dimension: Some(TextureViewDimension::D2),
621 usage: None,
622 aspect: TextureAspect::All,
623 base_mip_level: i as u32,
624 mip_level_count: Some(1),
625 base_array_layer: 0,
626 array_layer_count: Some(1),
627 })
628 } else {
629 (*depth_pyramid_dummy_texture).clone()
630 }
631 });
632
633 let depth_pyramid_all_mips = depth_pyramid.default_view.clone();
635
636 Self {
637 all_mips: depth_pyramid_all_mips,
638 mips: depth_pyramid_mips,
639 mip_count: depth_pyramid_mip_count,
640 }
641 }
642
643 pub fn create_bind_group<'a, R>(
646 &'a self,
647 render_device: &RenderDevice,
648 label: &'static str,
649 bind_group_layout: &BindGroupLayout,
650 source_image: R,
651 sampler: &'a Sampler,
652 ) -> BindGroup
653 where
654 R: IntoBinding<'a>,
655 {
656 render_device.create_bind_group(
657 label,
658 bind_group_layout,
659 &BindGroupEntries::sequential((
660 source_image,
661 &self.mips[0],
662 &self.mips[1],
663 &self.mips[2],
664 &self.mips[3],
665 &self.mips[4],
666 &self.mips[5],
667 &self.mips[6],
668 &self.mips[7],
669 &self.mips[8],
670 &self.mips[9],
671 &self.mips[10],
672 &self.mips[11],
673 sampler,
674 )),
675 )
676 }
677
678 pub fn downsample_depth(
682 &self,
683 label: &str,
684 render_context: &mut RenderContext,
685 view_size: UVec2,
686 downsample_depth_bind_group: &BindGroup,
687 downsample_depth_first_pipeline: &ComputePipeline,
688 downsample_depth_second_pipeline: &ComputePipeline,
689 ) {
690 let command_encoder = render_context.command_encoder();
691 let mut downsample_pass = command_encoder.begin_compute_pass(&ComputePassDescriptor {
692 label: Some(label),
693 timestamp_writes: None,
694 });
695 downsample_pass.set_pipeline(downsample_depth_first_pipeline);
696 downsample_pass.set_push_constants(0, &self.mip_count.to_le_bytes());
698 downsample_pass.set_bind_group(0, downsample_depth_bind_group, &[]);
699 downsample_pass.dispatch_workgroups(view_size.x.div_ceil(64), view_size.y.div_ceil(64), 1);
700
701 if self.mip_count >= 7 {
702 downsample_pass.set_pipeline(downsample_depth_second_pipeline);
703 downsample_pass.dispatch_workgroups(1, 1, 1);
704 }
705 }
706}
707
708pub fn prepare_view_depth_pyramids(
710 mut commands: Commands,
711 render_device: Res<RenderDevice>,
712 mut texture_cache: ResMut<TextureCache>,
713 depth_pyramid_dummy_texture: Res<DepthPyramidDummyTexture>,
714 views: Query<(Entity, &ExtractedView), (With<OcclusionCulling>, Without<NoIndirectDrawing>)>,
715) {
716 for (view_entity, view) in &views {
717 commands.entity(view_entity).insert(ViewDepthPyramid::new(
718 &render_device,
719 &mut texture_cache,
720 &depth_pyramid_dummy_texture,
721 view.viewport.zw(),
722 "view depth pyramid texture",
723 "view depth pyramid texture view",
724 ));
725 }
726}
727
728#[derive(Component, Deref, DerefMut)]
733pub struct ViewDownsampleDepthBindGroup(BindGroup);
734
735fn prepare_downsample_depth_view_bind_groups(
738 mut commands: Commands,
739 render_device: Res<RenderDevice>,
740 downsample_depth_pipelines: Res<DownsampleDepthPipelines>,
741 pipeline_cache: Res<PipelineCache>,
742 view_depth_textures: Query<
743 (
744 Entity,
745 &ViewDepthPyramid,
746 Option<&ViewDepthTexture>,
747 Option<&OcclusionCullingSubview>,
748 ),
749 Or<(With<ViewDepthTexture>, With<OcclusionCullingSubview>)>,
750 >,
751) {
752 for (view_entity, view_depth_pyramid, view_depth_texture, shadow_occlusion_culling) in
753 &view_depth_textures
754 {
755 let is_multisampled = view_depth_texture
756 .is_some_and(|view_depth_texture| view_depth_texture.texture.sample_count() > 1);
757 commands
758 .entity(view_entity)
759 .insert(ViewDownsampleDepthBindGroup(
760 view_depth_pyramid.create_bind_group(
761 &render_device,
762 if is_multisampled {
763 "downsample multisample depth bind group"
764 } else {
765 "downsample depth bind group"
766 },
767 &pipeline_cache.get_bind_group_layout(if is_multisampled {
768 &downsample_depth_pipelines
769 .first_multisample
770 .bind_group_layout
771 } else {
772 &downsample_depth_pipelines.first.bind_group_layout
773 }),
774 match (view_depth_texture, shadow_occlusion_culling) {
775 (Some(view_depth_texture), _) => view_depth_texture.view(),
776 (None, Some(shadow_occlusion_culling)) => {
777 &shadow_occlusion_culling.depth_texture_view
778 }
779 (None, None) => panic!("Should never happen"),
780 },
781 &downsample_depth_pipelines.sampler,
782 ),
783 ));
784 }
785}