bevy_pbr/
cluster.rs

1use core::num::NonZero;
2
3use bevy_camera::Camera;
4use bevy_ecs::{entity::EntityHashMap, prelude::*};
5use bevy_light::cluster::{ClusterableObjectCounts, Clusters, GlobalClusterSettings};
6use bevy_math::{uvec4, UVec3, UVec4, Vec4};
7use bevy_render::{
8    render_resource::{
9        BindingResource, BufferBindingType, ShaderSize, ShaderType, StorageBuffer, UniformBuffer,
10    },
11    renderer::{RenderAdapter, RenderDevice, RenderQueue},
12    sync_world::RenderEntity,
13    Extract,
14};
15use tracing::warn;
16
17use crate::MeshPipeline;
18
19// NOTE: this must be kept in sync with the same constants in
20// `mesh_view_types.wgsl`.
21pub const MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS: usize = 204;
22// Make sure that the clusterable object buffer doesn't overflow the maximum
23// size of a UBO on WebGL 2.
24const _: () =
25    assert!(size_of::<GpuClusterableObject>() * MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS <= 16384);
26
27// NOTE: Clustered-forward rendering requires 3 storage buffer bindings so check that
28// at least that many are supported using this constant and SupportedBindingType::from_device()
29pub const CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT: u32 = 3;
30
31// this must match CLUSTER_COUNT_SIZE in pbr.wgsl
32// and must be large enough to contain MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS
33const CLUSTER_COUNT_SIZE: u32 = 9;
34
35const CLUSTER_OFFSET_MASK: u32 = (1 << (32 - (CLUSTER_COUNT_SIZE * 2))) - 1;
36const CLUSTER_COUNT_MASK: u32 = (1 << CLUSTER_COUNT_SIZE) - 1;
37
38pub(crate) fn make_global_cluster_settings(world: &World) -> GlobalClusterSettings {
39    let device = world.resource::<RenderDevice>();
40    let adapter = world.resource::<RenderAdapter>();
41    let clustered_decals_are_usable =
42        crate::decal::clustered::clustered_decals_are_usable(device, adapter);
43    let supports_storage_buffers = matches!(
44        device.get_supported_read_only_binding_type(CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT),
45        BufferBindingType::Storage { .. }
46    );
47    GlobalClusterSettings {
48        supports_storage_buffers,
49        clustered_decals_are_usable,
50        max_uniform_buffer_clusterable_objects: MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS,
51        view_cluster_bindings_max_indices: ViewClusterBindings::MAX_INDICES,
52    }
53}
54
55#[derive(Copy, Clone, ShaderType, Default, Debug)]
56pub struct GpuClusterableObject {
57    // For point lights: the lower-right 2x2 values of the projection matrix [2][2] [2][3] [3][2] [3][3]
58    // For spot lights: 2 components of the direction (x,z), spot_scale and spot_offset
59    pub(crate) light_custom_data: Vec4,
60    pub(crate) color_inverse_square_range: Vec4,
61    pub(crate) position_radius: Vec4,
62    pub(crate) flags: u32,
63    pub(crate) shadow_depth_bias: f32,
64    pub(crate) shadow_normal_bias: f32,
65    pub(crate) spot_light_tan_angle: f32,
66    pub(crate) soft_shadow_size: f32,
67    pub(crate) shadow_map_near_z: f32,
68    pub(crate) decal_index: u32,
69    pub(crate) pad: f32,
70}
71
72#[derive(Resource)]
73pub struct GlobalClusterableObjectMeta {
74    pub gpu_clusterable_objects: GpuClusterableObjects,
75    pub entity_to_index: EntityHashMap<usize>,
76}
77
78pub enum GpuClusterableObjects {
79    Uniform(UniformBuffer<GpuClusterableObjectsUniform>),
80    Storage(StorageBuffer<GpuClusterableObjectsStorage>),
81}
82
83#[derive(ShaderType)]
84pub struct GpuClusterableObjectsUniform {
85    data: Box<[GpuClusterableObject; MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS]>,
86}
87
88#[derive(ShaderType, Default)]
89pub struct GpuClusterableObjectsStorage {
90    #[size(runtime)]
91    data: Vec<GpuClusterableObject>,
92}
93
94#[derive(Component)]
95pub struct ExtractedClusterConfig {
96    /// Special near value for cluster calculations
97    pub(crate) near: f32,
98    pub(crate) far: f32,
99    /// Number of clusters in `X` / `Y` / `Z` in the view frustum
100    pub(crate) dimensions: UVec3,
101}
102
103enum ExtractedClusterableObjectElement {
104    ClusterHeader(ClusterableObjectCounts),
105    ClusterableObjectEntity(Entity),
106}
107
108#[derive(Component)]
109pub struct ExtractedClusterableObjects {
110    data: Vec<ExtractedClusterableObjectElement>,
111}
112
113#[derive(ShaderType)]
114struct GpuClusterOffsetsAndCountsUniform {
115    data: Box<[UVec4; ViewClusterBindings::MAX_UNIFORM_ITEMS]>,
116}
117
118#[derive(ShaderType, Default)]
119struct GpuClusterableObjectIndexListsStorage {
120    #[size(runtime)]
121    data: Vec<u32>,
122}
123
124#[derive(ShaderType, Default)]
125struct GpuClusterOffsetsAndCountsStorage {
126    /// The starting offset, followed by the number of point lights, spot
127    /// lights, reflection probes, and irradiance volumes in each cluster, in
128    /// that order. The remaining fields are filled with zeroes.
129    #[size(runtime)]
130    data: Vec<[UVec4; 2]>,
131}
132
133enum ViewClusterBuffers {
134    Uniform {
135        // NOTE: UVec4 is because all arrays in Std140 layout have 16-byte alignment
136        clusterable_object_index_lists: UniformBuffer<GpuClusterableObjectIndexListsUniform>,
137        // NOTE: UVec4 is because all arrays in Std140 layout have 16-byte alignment
138        cluster_offsets_and_counts: UniformBuffer<GpuClusterOffsetsAndCountsUniform>,
139    },
140    Storage {
141        clusterable_object_index_lists: StorageBuffer<GpuClusterableObjectIndexListsStorage>,
142        cluster_offsets_and_counts: StorageBuffer<GpuClusterOffsetsAndCountsStorage>,
143    },
144}
145
146#[derive(Component)]
147pub struct ViewClusterBindings {
148    n_indices: usize,
149    n_offsets: usize,
150    buffers: ViewClusterBuffers,
151}
152
153pub fn init_global_clusterable_object_meta(
154    mut commands: Commands,
155    render_device: Res<RenderDevice>,
156) {
157    commands.insert_resource(GlobalClusterableObjectMeta::new(
158        render_device.get_supported_read_only_binding_type(CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT),
159    ));
160}
161
162impl GlobalClusterableObjectMeta {
163    pub fn new(buffer_binding_type: BufferBindingType) -> Self {
164        Self {
165            gpu_clusterable_objects: GpuClusterableObjects::new(buffer_binding_type),
166            entity_to_index: EntityHashMap::default(),
167        }
168    }
169}
170
171impl GpuClusterableObjects {
172    fn new(buffer_binding_type: BufferBindingType) -> Self {
173        match buffer_binding_type {
174            BufferBindingType::Storage { .. } => Self::storage(),
175            BufferBindingType::Uniform => Self::uniform(),
176        }
177    }
178
179    fn uniform() -> Self {
180        Self::Uniform(UniformBuffer::default())
181    }
182
183    fn storage() -> Self {
184        Self::Storage(StorageBuffer::default())
185    }
186
187    pub(crate) fn set(&mut self, mut clusterable_objects: Vec<GpuClusterableObject>) {
188        match self {
189            GpuClusterableObjects::Uniform(buffer) => {
190                let len = clusterable_objects
191                    .len()
192                    .min(MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS);
193                let src = &clusterable_objects[..len];
194                let dst = &mut buffer.get_mut().data[..len];
195                dst.copy_from_slice(src);
196            }
197            GpuClusterableObjects::Storage(buffer) => {
198                buffer.get_mut().data.clear();
199                buffer.get_mut().data.append(&mut clusterable_objects);
200            }
201        }
202    }
203
204    pub(crate) fn write_buffer(
205        &mut self,
206        render_device: &RenderDevice,
207        render_queue: &RenderQueue,
208    ) {
209        match self {
210            GpuClusterableObjects::Uniform(buffer) => {
211                buffer.write_buffer(render_device, render_queue);
212            }
213            GpuClusterableObjects::Storage(buffer) => {
214                buffer.write_buffer(render_device, render_queue);
215            }
216        }
217    }
218
219    pub fn binding(&self) -> Option<BindingResource<'_>> {
220        match self {
221            GpuClusterableObjects::Uniform(buffer) => buffer.binding(),
222            GpuClusterableObjects::Storage(buffer) => buffer.binding(),
223        }
224    }
225
226    pub fn min_size(buffer_binding_type: BufferBindingType) -> NonZero<u64> {
227        match buffer_binding_type {
228            BufferBindingType::Storage { .. } => GpuClusterableObjectsStorage::min_size(),
229            BufferBindingType::Uniform => GpuClusterableObjectsUniform::min_size(),
230        }
231    }
232}
233
234impl Default for GpuClusterableObjectsUniform {
235    fn default() -> Self {
236        Self {
237            data: Box::new(
238                [GpuClusterableObject::default(); MAX_UNIFORM_BUFFER_CLUSTERABLE_OBJECTS],
239            ),
240        }
241    }
242}
243
244/// Extracts clusters from the main world from the render world.
245pub fn extract_clusters(
246    mut commands: Commands,
247    views: Extract<Query<(RenderEntity, &Clusters, &Camera)>>,
248    mapper: Extract<Query<RenderEntity>>,
249) {
250    for (entity, clusters, camera) in &views {
251        let mut entity_commands = commands
252            .get_entity(entity)
253            .expect("Clusters entity wasn't synced.");
254        if !camera.is_active {
255            entity_commands.remove::<(ExtractedClusterableObjects, ExtractedClusterConfig)>();
256            continue;
257        }
258
259        let entity_count: usize = clusters
260            .clusterable_objects
261            .iter()
262            .map(|l| l.entities.len())
263            .sum();
264        let mut data = Vec::with_capacity(clusters.clusterable_objects.len() + entity_count);
265        for cluster_objects in &clusters.clusterable_objects {
266            data.push(ExtractedClusterableObjectElement::ClusterHeader(
267                cluster_objects.counts,
268            ));
269            for clusterable_entity in &cluster_objects.entities {
270                if let Ok(entity) = mapper.get(*clusterable_entity) {
271                    data.push(ExtractedClusterableObjectElement::ClusterableObjectEntity(
272                        entity,
273                    ));
274                }
275            }
276        }
277
278        entity_commands.insert((
279            ExtractedClusterableObjects { data },
280            ExtractedClusterConfig {
281                near: clusters.near,
282                far: clusters.far,
283                dimensions: clusters.dimensions,
284            },
285        ));
286    }
287}
288
289pub fn prepare_clusters(
290    mut commands: Commands,
291    render_device: Res<RenderDevice>,
292    render_queue: Res<RenderQueue>,
293    mesh_pipeline: Res<MeshPipeline>,
294    global_clusterable_object_meta: Res<GlobalClusterableObjectMeta>,
295    views: Query<(Entity, &ExtractedClusterableObjects)>,
296) {
297    let render_device = render_device.into_inner();
298    let supports_storage_buffers = matches!(
299        mesh_pipeline.clustered_forward_buffer_binding_type,
300        BufferBindingType::Storage { .. }
301    );
302    for (entity, extracted_clusters) in &views {
303        let mut view_clusters_bindings =
304            ViewClusterBindings::new(mesh_pipeline.clustered_forward_buffer_binding_type);
305        view_clusters_bindings.clear();
306
307        for record in &extracted_clusters.data {
308            match record {
309                ExtractedClusterableObjectElement::ClusterHeader(counts) => {
310                    let offset = view_clusters_bindings.n_indices();
311                    view_clusters_bindings.push_offset_and_counts(offset, counts);
312                }
313                ExtractedClusterableObjectElement::ClusterableObjectEntity(entity) => {
314                    if let Some(clusterable_object_index) =
315                        global_clusterable_object_meta.entity_to_index.get(entity)
316                    {
317                        if view_clusters_bindings.n_indices() >= ViewClusterBindings::MAX_INDICES
318                            && !supports_storage_buffers
319                        {
320                            warn!(
321                                "Clusterable object index lists are full! The clusterable \
322                                 objects in the view are present in too many clusters."
323                            );
324                            break;
325                        }
326                        view_clusters_bindings.push_index(*clusterable_object_index);
327                    }
328                }
329            }
330        }
331
332        view_clusters_bindings.write_buffers(render_device, &render_queue);
333
334        commands.entity(entity).insert(view_clusters_bindings);
335    }
336}
337
338impl ViewClusterBindings {
339    pub const MAX_OFFSETS: usize = 16384 / 4;
340    const MAX_UNIFORM_ITEMS: usize = Self::MAX_OFFSETS / 4;
341    pub const MAX_INDICES: usize = 16384;
342
343    pub fn new(buffer_binding_type: BufferBindingType) -> Self {
344        Self {
345            n_indices: 0,
346            n_offsets: 0,
347            buffers: ViewClusterBuffers::new(buffer_binding_type),
348        }
349    }
350
351    pub fn clear(&mut self) {
352        match &mut self.buffers {
353            ViewClusterBuffers::Uniform {
354                clusterable_object_index_lists,
355                cluster_offsets_and_counts,
356            } => {
357                *clusterable_object_index_lists.get_mut().data =
358                    [UVec4::ZERO; Self::MAX_UNIFORM_ITEMS];
359                *cluster_offsets_and_counts.get_mut().data = [UVec4::ZERO; Self::MAX_UNIFORM_ITEMS];
360            }
361            ViewClusterBuffers::Storage {
362                clusterable_object_index_lists,
363                cluster_offsets_and_counts,
364                ..
365            } => {
366                clusterable_object_index_lists.get_mut().data.clear();
367                cluster_offsets_and_counts.get_mut().data.clear();
368            }
369        }
370    }
371
372    fn push_offset_and_counts(&mut self, offset: usize, counts: &ClusterableObjectCounts) {
373        match &mut self.buffers {
374            ViewClusterBuffers::Uniform {
375                cluster_offsets_and_counts,
376                ..
377            } => {
378                let array_index = self.n_offsets >> 2; // >> 2 is equivalent to / 4
379                if array_index >= Self::MAX_UNIFORM_ITEMS {
380                    warn!("cluster offset and count out of bounds!");
381                    return;
382                }
383                let component = self.n_offsets & ((1 << 2) - 1);
384                let packed =
385                    pack_offset_and_counts(offset, counts.point_lights, counts.spot_lights);
386
387                cluster_offsets_and_counts.get_mut().data[array_index][component] = packed;
388            }
389            ViewClusterBuffers::Storage {
390                cluster_offsets_and_counts,
391                ..
392            } => {
393                cluster_offsets_and_counts.get_mut().data.push([
394                    uvec4(
395                        offset as u32,
396                        counts.point_lights,
397                        counts.spot_lights,
398                        counts.reflection_probes,
399                    ),
400                    uvec4(counts.irradiance_volumes, counts.decals, 0, 0),
401                ]);
402            }
403        }
404
405        self.n_offsets += 1;
406    }
407
408    pub fn n_indices(&self) -> usize {
409        self.n_indices
410    }
411
412    pub fn push_index(&mut self, index: usize) {
413        match &mut self.buffers {
414            ViewClusterBuffers::Uniform {
415                clusterable_object_index_lists,
416                ..
417            } => {
418                let array_index = self.n_indices >> 4; // >> 4 is equivalent to / 16
419                let component = (self.n_indices >> 2) & ((1 << 2) - 1);
420                let sub_index = self.n_indices & ((1 << 2) - 1);
421                let index = index as u32;
422
423                clusterable_object_index_lists.get_mut().data[array_index][component] |=
424                    index << (8 * sub_index);
425            }
426            ViewClusterBuffers::Storage {
427                clusterable_object_index_lists,
428                ..
429            } => {
430                clusterable_object_index_lists
431                    .get_mut()
432                    .data
433                    .push(index as u32);
434            }
435        }
436
437        self.n_indices += 1;
438    }
439
440    pub fn write_buffers(&mut self, render_device: &RenderDevice, render_queue: &RenderQueue) {
441        match &mut self.buffers {
442            ViewClusterBuffers::Uniform {
443                clusterable_object_index_lists,
444                cluster_offsets_and_counts,
445            } => {
446                clusterable_object_index_lists.write_buffer(render_device, render_queue);
447                cluster_offsets_and_counts.write_buffer(render_device, render_queue);
448            }
449            ViewClusterBuffers::Storage {
450                clusterable_object_index_lists,
451                cluster_offsets_and_counts,
452            } => {
453                clusterable_object_index_lists.write_buffer(render_device, render_queue);
454                cluster_offsets_and_counts.write_buffer(render_device, render_queue);
455            }
456        }
457    }
458
459    pub fn clusterable_object_index_lists_binding(&self) -> Option<BindingResource<'_>> {
460        match &self.buffers {
461            ViewClusterBuffers::Uniform {
462                clusterable_object_index_lists,
463                ..
464            } => clusterable_object_index_lists.binding(),
465            ViewClusterBuffers::Storage {
466                clusterable_object_index_lists,
467                ..
468            } => clusterable_object_index_lists.binding(),
469        }
470    }
471
472    pub fn offsets_and_counts_binding(&self) -> Option<BindingResource<'_>> {
473        match &self.buffers {
474            ViewClusterBuffers::Uniform {
475                cluster_offsets_and_counts,
476                ..
477            } => cluster_offsets_and_counts.binding(),
478            ViewClusterBuffers::Storage {
479                cluster_offsets_and_counts,
480                ..
481            } => cluster_offsets_and_counts.binding(),
482        }
483    }
484
485    pub fn min_size_clusterable_object_index_lists(
486        buffer_binding_type: BufferBindingType,
487    ) -> NonZero<u64> {
488        match buffer_binding_type {
489            BufferBindingType::Storage { .. } => GpuClusterableObjectIndexListsStorage::min_size(),
490            BufferBindingType::Uniform => GpuClusterableObjectIndexListsUniform::min_size(),
491        }
492    }
493
494    pub fn min_size_cluster_offsets_and_counts(
495        buffer_binding_type: BufferBindingType,
496    ) -> NonZero<u64> {
497        match buffer_binding_type {
498            BufferBindingType::Storage { .. } => GpuClusterOffsetsAndCountsStorage::min_size(),
499            BufferBindingType::Uniform => GpuClusterOffsetsAndCountsUniform::min_size(),
500        }
501    }
502}
503
504impl ViewClusterBuffers {
505    fn new(buffer_binding_type: BufferBindingType) -> Self {
506        match buffer_binding_type {
507            BufferBindingType::Storage { .. } => Self::storage(),
508            BufferBindingType::Uniform => Self::uniform(),
509        }
510    }
511
512    fn uniform() -> Self {
513        ViewClusterBuffers::Uniform {
514            clusterable_object_index_lists: UniformBuffer::default(),
515            cluster_offsets_and_counts: UniformBuffer::default(),
516        }
517    }
518
519    fn storage() -> Self {
520        ViewClusterBuffers::Storage {
521            clusterable_object_index_lists: StorageBuffer::default(),
522            cluster_offsets_and_counts: StorageBuffer::default(),
523        }
524    }
525}
526
527// Compresses the offset and counts of point and spot lights so that they fit in
528// a UBO.
529//
530// This function is only used if storage buffers are unavailable on this
531// platform: typically, on WebGL 2.
532//
533// NOTE: With uniform buffer max binding size as 16384 bytes
534// that means we can fit 204 clusterable objects in one uniform
535// buffer, which means the count can be at most 204 so it
536// needs 9 bits.
537// The array of indices can also use u8 and that means the
538// offset in to the array of indices needs to be able to address
539// 16384 values. log2(16384) = 14 bits.
540// We use 32 bits to store the offset and counts so
541// we pack the offset into the upper 14 bits of a u32,
542// the point light count into bits 9-17, and the spot light count into bits 0-8.
543//  [ 31     ..     18 | 17      ..      9 | 8       ..     0 ]
544//  [      offset      | point light count | spot light count ]
545//
546// NOTE: This assumes CPU and GPU endianness are the same which is true
547// for all common and tested x86/ARM CPUs and AMD/NVIDIA/Intel/Apple/etc GPUs
548//
549// NOTE: On platforms that use this function, we don't cluster light probes, so
550// the number of light probes is irrelevant.
551fn pack_offset_and_counts(offset: usize, point_count: u32, spot_count: u32) -> u32 {
552    ((offset as u32 & CLUSTER_OFFSET_MASK) << (CLUSTER_COUNT_SIZE * 2))
553        | ((point_count & CLUSTER_COUNT_MASK) << CLUSTER_COUNT_SIZE)
554        | (spot_count & CLUSTER_COUNT_MASK)
555}
556
557#[derive(ShaderType)]
558struct GpuClusterableObjectIndexListsUniform {
559    data: Box<[UVec4; ViewClusterBindings::MAX_UNIFORM_ITEMS]>,
560}
561
562// NOTE: Assert at compile time that GpuClusterableObjectIndexListsUniform
563// fits within the maximum uniform buffer binding size
564const _: () = assert!(GpuClusterableObjectIndexListsUniform::SHADER_SIZE.get() <= 16384);
565
566impl Default for GpuClusterableObjectIndexListsUniform {
567    fn default() -> Self {
568        Self {
569            data: Box::new([UVec4::ZERO; ViewClusterBindings::MAX_UNIFORM_ITEMS]),
570        }
571    }
572}
573
574impl Default for GpuClusterOffsetsAndCountsUniform {
575    fn default() -> Self {
576        Self {
577            data: Box::new([UVec4::ZERO; ViewClusterBindings::MAX_UNIFORM_ITEMS]),
578        }
579    }
580}