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