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::{load_internal_asset, weak_handle, 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::batching::gpu_preprocessing::GpuPreprocessingSupport;
29use bevy_render::{
30 experimental::occlusion_culling::{
31 OcclusionCulling, OcclusionCullingSubview, OcclusionCullingSubviewEntities,
32 },
33 render_graph::{Node, NodeRunError, RenderGraphApp, RenderGraphContext},
34 render_resource::{
35 binding_types::{sampler, texture_2d, texture_2d_multisampled, texture_storage_2d},
36 BindGroup, BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries,
37 CachedComputePipelineId, ComputePassDescriptor, ComputePipeline, ComputePipelineDescriptor,
38 Extent3d, IntoBinding, PipelineCache, PushConstantRange, Sampler, SamplerBindingType,
39 SamplerDescriptor, Shader, ShaderStages, SpecializedComputePipeline,
40 SpecializedComputePipelines, StorageTextureAccess, TextureAspect, TextureDescriptor,
41 TextureDimension, TextureFormat, TextureSampleType, TextureUsages, TextureView,
42 TextureViewDescriptor, TextureViewDimension,
43 },
44 renderer::{RenderContext, RenderDevice},
45 texture::TextureCache,
46 view::{ExtractedView, NoIndirectDrawing, ViewDepthTexture},
47 Render, RenderApp, RenderSet,
48};
49use bitflags::bitflags;
50use tracing::debug;
51
52pub const DOWNSAMPLE_DEPTH_SHADER_HANDLE: Handle<Shader> =
54 weak_handle!("a09a149e-5922-4fa4-9170-3c1a13065364");
55
56pub const DEPTH_PYRAMID_MIP_COUNT: usize = 12;
61
62pub struct MipGenerationPlugin;
68
69impl Plugin for MipGenerationPlugin {
70 fn build(&self, app: &mut App) {
71 load_internal_asset!(
72 app,
73 DOWNSAMPLE_DEPTH_SHADER_HANDLE,
74 "downsample_depth.wgsl",
75 Shader::from_wgsl
76 );
77
78 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
79 return;
80 };
81
82 render_app
83 .init_resource::<SpecializedComputePipelines<DownsampleDepthPipeline>>()
84 .add_render_graph_node::<DownsampleDepthNode>(Core3d, Node3d::EarlyDownsampleDepth)
85 .add_render_graph_node::<DownsampleDepthNode>(Core3d, Node3d::LateDownsampleDepth)
86 .add_render_graph_edges(
87 Core3d,
88 (
89 Node3d::EarlyPrepass,
90 Node3d::EarlyDeferredPrepass,
91 Node3d::EarlyDownsampleDepth,
92 Node3d::LatePrepass,
93 Node3d::LateDeferredPrepass,
94 ),
95 )
96 .add_render_graph_edges(
97 Core3d,
98 (
99 Node3d::EndMainPass,
100 Node3d::LateDownsampleDepth,
101 Node3d::EndMainPassPostProcessing,
102 ),
103 )
104 .add_systems(
105 Render,
106 create_downsample_depth_pipelines.in_set(RenderSet::Prepare),
107 )
108 .add_systems(
109 Render,
110 (
111 prepare_view_depth_pyramids,
112 prepare_downsample_depth_view_bind_groups,
113 )
114 .chain()
115 .in_set(RenderSet::PrepareResources)
116 .run_if(resource_exists::<DownsampleDepthPipelines>)
117 .after(prepare_core_3d_depth_textures),
118 );
119 }
120
121 fn finish(&self, app: &mut App) {
122 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
123 return;
124 };
125 render_app.init_resource::<DepthPyramidDummyTexture>();
126 }
127}
128
129pub struct DownsampleDepthNode {
146 main_view_query: QueryState<(
149 Read<ViewDepthPyramid>,
150 Read<ViewDownsampleDepthBindGroup>,
151 Read<ViewDepthTexture>,
152 Option<Read<OcclusionCullingSubviewEntities>>,
153 )>,
154 shadow_view_query: QueryState<(
156 Read<ViewDepthPyramid>,
157 Read<ViewDownsampleDepthBindGroup>,
158 Read<OcclusionCullingSubview>,
159 )>,
160}
161
162impl FromWorld for DownsampleDepthNode {
163 fn from_world(world: &mut World) -> Self {
164 Self {
165 main_view_query: QueryState::new(world),
166 shadow_view_query: QueryState::new(world),
167 }
168 }
169}
170
171impl Node for DownsampleDepthNode {
172 fn update(&mut self, world: &mut World) {
173 self.main_view_query.update_archetypes(world);
174 self.shadow_view_query.update_archetypes(world);
175 }
176
177 fn run<'w>(
178 &self,
179 render_graph_context: &mut RenderGraphContext,
180 render_context: &mut RenderContext<'w>,
181 world: &'w World,
182 ) -> Result<(), NodeRunError> {
183 let Ok((
184 view_depth_pyramid,
185 view_downsample_depth_bind_group,
186 view_depth_texture,
187 maybe_view_light_entities,
188 )) = self
189 .main_view_query
190 .get_manual(world, render_graph_context.view_entity())
191 else {
192 return Ok(());
193 };
194
195 downsample_depth(
197 render_graph_context,
198 render_context,
199 world,
200 view_depth_pyramid,
201 view_downsample_depth_bind_group,
202 uvec2(
203 view_depth_texture.texture.width(),
204 view_depth_texture.texture.height(),
205 ),
206 view_depth_texture.texture.sample_count(),
207 )?;
208
209 if let Some(view_light_entities) = maybe_view_light_entities {
211 for &view_light_entity in &view_light_entities.0 {
212 let Ok((view_depth_pyramid, view_downsample_depth_bind_group, occlusion_culling)) =
213 self.shadow_view_query.get_manual(world, view_light_entity)
214 else {
215 continue;
216 };
217 downsample_depth(
218 render_graph_context,
219 render_context,
220 world,
221 view_depth_pyramid,
222 view_downsample_depth_bind_group,
223 UVec2::splat(occlusion_culling.depth_texture_size),
224 1,
225 )?;
226 }
227 }
228
229 Ok(())
230 }
231}
232
233fn downsample_depth<'w>(
236 render_graph_context: &mut RenderGraphContext,
237 render_context: &mut RenderContext<'w>,
238 world: &'w World,
239 view_depth_pyramid: &ViewDepthPyramid,
240 view_downsample_depth_bind_group: &ViewDownsampleDepthBindGroup,
241 view_size: UVec2,
242 sample_count: u32,
243) -> Result<(), NodeRunError> {
244 let downsample_depth_pipelines = world.resource::<DownsampleDepthPipelines>();
245 let pipeline_cache = world.resource::<PipelineCache>();
246
247 let (Some(first_downsample_depth_pipeline_id), Some(second_downsample_depth_pipeline_id)) =
254 (if sample_count > 1 {
255 (
256 downsample_depth_pipelines.first_multisample.pipeline_id,
257 downsample_depth_pipelines.second_multisample.pipeline_id,
258 )
259 } else {
260 (
261 downsample_depth_pipelines.first.pipeline_id,
262 downsample_depth_pipelines.second.pipeline_id,
263 )
264 })
265 else {
266 return Ok(());
267 };
268
269 let (Some(first_downsample_depth_pipeline), Some(second_downsample_depth_pipeline)) = (
271 pipeline_cache.get_compute_pipeline(first_downsample_depth_pipeline_id),
272 pipeline_cache.get_compute_pipeline(second_downsample_depth_pipeline_id),
273 ) else {
274 return Ok(());
275 };
276
277 view_depth_pyramid.downsample_depth(
279 &format!("{:?}", render_graph_context.label()),
280 render_context,
281 view_size,
282 view_downsample_depth_bind_group,
283 first_downsample_depth_pipeline,
284 second_downsample_depth_pipeline,
285 );
286 Ok(())
287}
288
289#[derive(Resource)]
291pub struct DownsampleDepthPipeline {
292 bind_group_layout: BindGroupLayout,
294 pipeline_id: Option<CachedComputePipelineId>,
296}
297
298impl DownsampleDepthPipeline {
299 fn new(bind_group_layout: BindGroupLayout) -> DownsampleDepthPipeline {
304 DownsampleDepthPipeline {
305 bind_group_layout,
306 pipeline_id: None,
307 }
308 }
309}
310
311#[derive(Resource)]
313pub struct DownsampleDepthPipelines {
314 first: DownsampleDepthPipeline,
317 second: DownsampleDepthPipeline,
320 first_multisample: DownsampleDepthPipeline,
322 second_multisample: DownsampleDepthPipeline,
324 sampler: Sampler,
327}
328
329fn create_downsample_depth_pipelines(
332 mut commands: Commands,
333 render_device: Res<RenderDevice>,
334 pipeline_cache: Res<PipelineCache>,
335 mut specialized_compute_pipelines: ResMut<SpecializedComputePipelines<DownsampleDepthPipeline>>,
336 gpu_preprocessing_support: Res<GpuPreprocessingSupport>,
337 mut has_run: Local<bool>,
338) {
339 if *has_run {
344 return;
345 }
346 *has_run = true;
347
348 if !gpu_preprocessing_support.is_culling_supported() {
349 debug!("Downsample depth is not supported on this platform.");
350 return;
351 }
352
353 let standard_bind_group_layout =
358 create_downsample_depth_bind_group_layout(&render_device, false);
359 let multisampled_bind_group_layout =
360 create_downsample_depth_bind_group_layout(&render_device, true);
361
362 let sampler = render_device.create_sampler(&SamplerDescriptor {
364 label: Some("depth pyramid sampler"),
365 ..SamplerDescriptor::default()
366 });
367
368 let mut downsample_depth_pipelines = DownsampleDepthPipelines {
370 first: DownsampleDepthPipeline::new(standard_bind_group_layout.clone()),
371 second: DownsampleDepthPipeline::new(standard_bind_group_layout.clone()),
372 first_multisample: DownsampleDepthPipeline::new(multisampled_bind_group_layout.clone()),
373 second_multisample: DownsampleDepthPipeline::new(multisampled_bind_group_layout.clone()),
374 sampler,
375 };
376
377 downsample_depth_pipelines.first.pipeline_id = Some(specialized_compute_pipelines.specialize(
380 &pipeline_cache,
381 &downsample_depth_pipelines.first,
382 DownsampleDepthPipelineKey::empty(),
383 ));
384 downsample_depth_pipelines.second.pipeline_id = Some(specialized_compute_pipelines.specialize(
385 &pipeline_cache,
386 &downsample_depth_pipelines.second,
387 DownsampleDepthPipelineKey::SECOND_PHASE,
388 ));
389 downsample_depth_pipelines.first_multisample.pipeline_id =
390 Some(specialized_compute_pipelines.specialize(
391 &pipeline_cache,
392 &downsample_depth_pipelines.first_multisample,
393 DownsampleDepthPipelineKey::MULTISAMPLE,
394 ));
395 downsample_depth_pipelines.second_multisample.pipeline_id =
396 Some(specialized_compute_pipelines.specialize(
397 &pipeline_cache,
398 &downsample_depth_pipelines.second_multisample,
399 DownsampleDepthPipelineKey::SECOND_PHASE | DownsampleDepthPipelineKey::MULTISAMPLE,
400 ));
401
402 commands.insert_resource(downsample_depth_pipelines);
403}
404
405fn create_downsample_depth_bind_group_layout(
407 render_device: &RenderDevice,
408 is_multisampled: bool,
409) -> BindGroupLayout {
410 render_device.create_bind_group_layout(
411 if is_multisampled {
412 "downsample multisample depth bind group layout"
413 } else {
414 "downsample depth bind group layout"
415 },
416 &BindGroupLayoutEntries::sequential(
417 ShaderStages::COMPUTE,
418 (
419 if is_multisampled {
423 texture_2d_multisampled(TextureSampleType::Depth)
424 } else {
425 texture_2d(TextureSampleType::Depth)
426 },
427 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
429 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
430 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
431 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
432 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
433 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::ReadWrite),
434 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
435 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
436 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
437 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
438 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
439 texture_storage_2d(TextureFormat::R32Float, StorageTextureAccess::WriteOnly),
440 sampler(SamplerBindingType::NonFiltering),
441 ),
442 ),
443 )
444}
445
446bitflags! {
447 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
453 pub struct DownsampleDepthPipelineKey: u8 {
454 const MULTISAMPLE = 1;
456 const SECOND_PHASE = 2;
459 }
460}
461
462impl SpecializedComputePipeline for DownsampleDepthPipeline {
463 type Key = DownsampleDepthPipelineKey;
464
465 fn specialize(&self, key: Self::Key) -> ComputePipelineDescriptor {
466 let mut shader_defs = vec![];
467 if key.contains(DownsampleDepthPipelineKey::MULTISAMPLE) {
468 shader_defs.push("MULTISAMPLE".into());
469 }
470
471 let label = format!(
472 "downsample depth{}{} pipeline",
473 if key.contains(DownsampleDepthPipelineKey::MULTISAMPLE) {
474 " multisample"
475 } else {
476 ""
477 },
478 if key.contains(DownsampleDepthPipelineKey::SECOND_PHASE) {
479 " second phase"
480 } else {
481 " first phase"
482 }
483 )
484 .into();
485
486 ComputePipelineDescriptor {
487 label: Some(label),
488 layout: vec![self.bind_group_layout.clone()],
489 push_constant_ranges: vec![PushConstantRange {
490 stages: ShaderStages::COMPUTE,
491 range: 0..4,
492 }],
493 shader: DOWNSAMPLE_DEPTH_SHADER_HANDLE,
494 shader_defs,
495 entry_point: if key.contains(DownsampleDepthPipelineKey::SECOND_PHASE) {
496 "downsample_depth_second".into()
497 } else {
498 "downsample_depth_first".into()
499 },
500 zero_initialize_workgroup_memory: false,
501 }
502 }
503}
504
505#[derive(Resource, Deref, DerefMut)]
508pub struct DepthPyramidDummyTexture(TextureView);
509
510impl FromWorld for DepthPyramidDummyTexture {
511 fn from_world(world: &mut World) -> Self {
512 let render_device = world.resource::<RenderDevice>();
513
514 DepthPyramidDummyTexture(create_depth_pyramid_dummy_texture(
515 render_device,
516 "depth pyramid dummy texture",
517 "depth pyramid dummy texture view",
518 ))
519 }
520}
521
522pub fn create_depth_pyramid_dummy_texture(
525 render_device: &RenderDevice,
526 texture_label: &'static str,
527 texture_view_label: &'static str,
528) -> TextureView {
529 render_device
530 .create_texture(&TextureDescriptor {
531 label: Some(texture_label),
532 size: Extent3d {
533 width: 1,
534 height: 1,
535 depth_or_array_layers: 1,
536 },
537 mip_level_count: 1,
538 sample_count: 1,
539 dimension: TextureDimension::D2,
540 format: TextureFormat::R32Float,
541 usage: TextureUsages::STORAGE_BINDING,
542 view_formats: &[],
543 })
544 .create_view(&TextureViewDescriptor {
545 label: Some(texture_view_label),
546 format: Some(TextureFormat::R32Float),
547 dimension: Some(TextureViewDimension::D2),
548 usage: None,
549 aspect: TextureAspect::All,
550 base_mip_level: 0,
551 mip_level_count: Some(1),
552 base_array_layer: 0,
553 array_layer_count: Some(1),
554 })
555}
556
557#[derive(Component)]
562pub struct ViewDepthPyramid {
563 pub all_mips: TextureView,
565 pub mips: [TextureView; DEPTH_PYRAMID_MIP_COUNT],
567 pub mip_count: u32,
572}
573
574impl ViewDepthPyramid {
575 pub fn new(
577 render_device: &RenderDevice,
578 texture_cache: &mut TextureCache,
579 depth_pyramid_dummy_texture: &TextureView,
580 size: UVec2,
581 texture_label: &'static str,
582 texture_view_label: &'static str,
583 ) -> ViewDepthPyramid {
584 let depth_pyramid_size = Extent3d {
586 width: size.x.div_ceil(2),
587 height: size.y.div_ceil(2),
588 depth_or_array_layers: 1,
589 };
590
591 let depth_pyramid_mip_count = depth_pyramid_size.max_mips(TextureDimension::D2);
593
594 let depth_pyramid = texture_cache.get(
596 render_device,
597 TextureDescriptor {
598 label: Some(texture_label),
599 size: depth_pyramid_size,
600 mip_level_count: depth_pyramid_mip_count,
601 sample_count: 1,
602 dimension: TextureDimension::D2,
603 format: TextureFormat::R32Float,
604 usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
605 view_formats: &[],
606 },
607 );
608
609 let depth_pyramid_mips = array::from_fn(|i| {
611 if (i as u32) < depth_pyramid_mip_count {
612 depth_pyramid.texture.create_view(&TextureViewDescriptor {
613 label: Some(texture_view_label),
614 format: Some(TextureFormat::R32Float),
615 dimension: Some(TextureViewDimension::D2),
616 usage: None,
617 aspect: TextureAspect::All,
618 base_mip_level: i as u32,
619 mip_level_count: Some(1),
620 base_array_layer: 0,
621 array_layer_count: Some(1),
622 })
623 } else {
624 (*depth_pyramid_dummy_texture).clone()
625 }
626 });
627
628 let depth_pyramid_all_mips = depth_pyramid.default_view.clone();
630
631 Self {
632 all_mips: depth_pyramid_all_mips,
633 mips: depth_pyramid_mips,
634 mip_count: depth_pyramid_mip_count,
635 }
636 }
637
638 pub fn create_bind_group<'a, R>(
641 &'a self,
642 render_device: &RenderDevice,
643 label: &'static str,
644 bind_group_layout: &BindGroupLayout,
645 source_image: R,
646 sampler: &'a Sampler,
647 ) -> BindGroup
648 where
649 R: IntoBinding<'a>,
650 {
651 render_device.create_bind_group(
652 label,
653 bind_group_layout,
654 &BindGroupEntries::sequential((
655 source_image,
656 &self.mips[0],
657 &self.mips[1],
658 &self.mips[2],
659 &self.mips[3],
660 &self.mips[4],
661 &self.mips[5],
662 &self.mips[6],
663 &self.mips[7],
664 &self.mips[8],
665 &self.mips[9],
666 &self.mips[10],
667 &self.mips[11],
668 sampler,
669 )),
670 )
671 }
672
673 pub fn downsample_depth(
677 &self,
678 label: &str,
679 render_context: &mut RenderContext,
680 view_size: UVec2,
681 downsample_depth_bind_group: &BindGroup,
682 downsample_depth_first_pipeline: &ComputePipeline,
683 downsample_depth_second_pipeline: &ComputePipeline,
684 ) {
685 let command_encoder = render_context.command_encoder();
686 let mut downsample_pass = command_encoder.begin_compute_pass(&ComputePassDescriptor {
687 label: Some(label),
688 timestamp_writes: None,
689 });
690 downsample_pass.set_pipeline(downsample_depth_first_pipeline);
691 downsample_pass.set_push_constants(0, &self.mip_count.to_le_bytes());
693 downsample_pass.set_bind_group(0, downsample_depth_bind_group, &[]);
694 downsample_pass.dispatch_workgroups(view_size.x.div_ceil(64), view_size.y.div_ceil(64), 1);
695
696 if self.mip_count >= 7 {
697 downsample_pass.set_pipeline(downsample_depth_second_pipeline);
698 downsample_pass.dispatch_workgroups(1, 1, 1);
699 }
700 }
701}
702
703pub fn prepare_view_depth_pyramids(
705 mut commands: Commands,
706 render_device: Res<RenderDevice>,
707 mut texture_cache: ResMut<TextureCache>,
708 depth_pyramid_dummy_texture: Res<DepthPyramidDummyTexture>,
709 views: Query<(Entity, &ExtractedView), (With<OcclusionCulling>, Without<NoIndirectDrawing>)>,
710) {
711 for (view_entity, view) in &views {
712 commands.entity(view_entity).insert(ViewDepthPyramid::new(
713 &render_device,
714 &mut texture_cache,
715 &depth_pyramid_dummy_texture,
716 view.viewport.zw(),
717 "view depth pyramid texture",
718 "view depth pyramid texture view",
719 ));
720 }
721}
722
723#[derive(Component, Deref, DerefMut)]
728pub struct ViewDownsampleDepthBindGroup(BindGroup);
729
730fn prepare_downsample_depth_view_bind_groups(
733 mut commands: Commands,
734 render_device: Res<RenderDevice>,
735 downsample_depth_pipelines: Res<DownsampleDepthPipelines>,
736 view_depth_textures: Query<
737 (
738 Entity,
739 &ViewDepthPyramid,
740 Option<&ViewDepthTexture>,
741 Option<&OcclusionCullingSubview>,
742 ),
743 Or<(With<ViewDepthTexture>, With<OcclusionCullingSubview>)>,
744 >,
745) {
746 for (view_entity, view_depth_pyramid, view_depth_texture, shadow_occlusion_culling) in
747 &view_depth_textures
748 {
749 let is_multisampled = view_depth_texture
750 .is_some_and(|view_depth_texture| view_depth_texture.texture.sample_count() > 1);
751 commands
752 .entity(view_entity)
753 .insert(ViewDownsampleDepthBindGroup(
754 view_depth_pyramid.create_bind_group(
755 &render_device,
756 if is_multisampled {
757 "downsample multisample depth bind group"
758 } else {
759 "downsample depth bind group"
760 },
761 if is_multisampled {
762 &downsample_depth_pipelines
763 .first_multisample
764 .bind_group_layout
765 } else {
766 &downsample_depth_pipelines.first.bind_group_layout
767 },
768 match (view_depth_texture, shadow_occlusion_culling) {
769 (Some(view_depth_texture), _) => view_depth_texture.view(),
770 (None, Some(shadow_occlusion_culling)) => {
771 &shadow_occlusion_culling.depth_texture_view
772 }
773 (None, None) => panic!("Should never happen"),
774 },
775 &downsample_depth_pipelines.sampler,
776 ),
777 ));
778 }
779}