[Merged by Bors] - bevy_pbr: Use a special first depth slice for clustered forward

crates/bevy_pbr/src/light.rs

@@ -210,6 +210,9 @@ pub struct Clusters {
     pub(crate) tile_size: UVec2,
     /// Number of clusters in x / y / z in the view frustum
     pub(crate) axis_slices: UVec3,
+    /// Distance to the far plane of the first depth slice. The first depth slice is special
+    /// and explicitly-configured to avoid having unnecessarily many slices close to the camera.
+    pub(crate) near: f32,
     aabbs: Vec<Aabb>,
     pub(crate) lights: Vec<VisiblePointLights>,
 }
@@ -219,6 +222,7 @@ impl Clusters {
         let mut clusters = Self {
             tile_size,
             axis_slices: Default::default(),
+            near: 5.0,
             aabbs: Default::default(),
             lights: Default::default(),
         };
@@ -246,6 +250,8 @@ impl Clusters {
             (screen_size.y + 1) / tile_size.y,
             z_slices,
         );
+        // NOTE: Maximum 4096 clusters due to uniform buffer size constraints
+        assert!(self.axis_slices.x * self.axis_slices.y * self.axis_slices.z <= 4096);
     }
 }
 
@@ -320,11 +326,15 @@ fn compute_aabb_for_cluster(
         let p_max = screen_to_view(screen_size, inverse_projection, p_max, 1.0);
 
         let z_far_over_z_near = -z_far / -z_near;
-        let cluster_near = -z_near * z_far_over_z_near.powf(ijk.z / cluster_dimensions.z as f32);
+        let cluster_near = if ijk.z == 0.0 {
+            0.0
+        } else {
+            -z_near * z_far_over_z_near.powf((ijk.z - 1.0) / (cluster_dimensions.z - 1) as f32)
+        };
         // NOTE: This could be simplified to:
         // cluster_far = cluster_near * z_far_over_z_near;
         let cluster_far =
-            -z_near * z_far_over_z_near.powf((ijk.z + 1.0) / cluster_dimensions.z as f32);
+            -z_near * z_far_over_z_near.powf(ijk.z / (cluster_dimensions.z - 1) as f32);
 
         // Calculate the four intersection points of the min and max points with the cluster near and far planes
         let p_min_near = line_intersection_to_z_plane(Vec3::ZERO, p_min.xyz(), cluster_near);
@@ -387,7 +397,7 @@ pub fn update_clusters(windows: Res<Windows>, mut views: Query<(&Camera, &mut Cl
             for x in 0..clusters.axis_slices.x {
                 for z in 0..clusters.axis_slices.z {
                     aabbs.push(compute_aabb_for_cluster(
-                        camera.near,
+                        clusters.near,
                         camera.far,
                         tile_size,
                         screen_size,
@@ -428,13 +438,19 @@ impl VisiblePointLights {
     }
 }
 
-fn view_z_to_z_slice(cluster_factors: Vec2, view_z: f32, is_orthographic: bool) -> u32 {
+fn view_z_to_z_slice(
+    cluster_factors: Vec2,
+    z_slices: f32,
+    view_z: f32,
+    is_orthographic: bool,
+) -> u32 {
     if is_orthographic {
         // NOTE: view_z is correct in the orthographic case
         ((view_z - cluster_factors.x) * cluster_factors.y).floor() as u32
     } else {
         // NOTE: had to use -view_z to make it positive else log(negative) is nan
-        ((-view_z).ln() * cluster_factors.x - cluster_factors.y).floor() as u32
+        ((-view_z).ln() * cluster_factors.x - cluster_factors.y + 1.0).clamp(0.0, z_slices - 1.0)
+            as u32
     }
 }
 
@@ -449,7 +465,12 @@ fn ndc_position_to_cluster(
     let frag_coord =
         (ndc_p.xy() * Vec2::new(0.5, -0.5) + Vec2::splat(0.5)).clamp(Vec2::ZERO, Vec2::ONE);
     let xy = (frag_coord * cluster_dimensions_f32.xy()).floor();
-    let z_slice = view_z_to_z_slice(cluster_factors, view_z, is_orthographic);
+    let z_slice = view_z_to_z_slice(
+        cluster_factors,
+        cluster_dimensions.z as f32,
+        view_z,
+        is_orthographic,
+    );
     xy.as_uvec2()
         .extend(z_slice)
         .clamp(UVec3::ZERO, cluster_dimensions - UVec3::ONE)
@@ -474,7 +495,8 @@ pub fn assign_lights_to_clusters(
         let cluster_count = clusters.aabbs.len();
         let is_orthographic = camera.projection_matrix.w_axis.w == 1.0;
         let cluster_factors = calculate_cluster_factors(
-            camera.near,
+            // NOTE: Using the special cluster near value
+            clusters.near,
             camera.far,
             clusters.axis_slices.z as f32,
             is_orthographic,

crates/bevy_pbr/src/render/light.rs

@@ -129,6 +129,7 @@ pub struct GpuLights {
     // TODO: this comes first to work around a WGSL alignment issue. We need to solve this issue before releasing the renderer rework
     directional_lights: [GpuDirectionalLight; MAX_DIRECTIONAL_LIGHTS],
     ambient_color: Vec4,
+    // xyz are x/y/z cluster dimensions and w is the number of clusters
     cluster_dimensions: UVec4,
     // xy are vec2<f32>(cluster_dimensions.xy) / vec2<f32>(view.width, view.height)
     // z is cluster_dimensions.z / log(far / near)
@@ -323,6 +324,8 @@ impl SpecializedPipeline for ShadowPipeline {
 
 #[derive(Component)]
 pub struct ExtractedClusterConfig {
+    /// Special near value for cluster calculations
+    near: f32,
     /// Number of clusters in x / y / z in the view frustum
     axis_slices: UVec3,
 }
@@ -339,6 +342,7 @@ pub fn extract_clusters(mut commands: Commands, views: Query<(Entity, &Clusters)
                 data: clusters.lights.clone(),
             },
             ExtractedClusterConfig {
+                near: clusters.near,
                 axis_slices: clusters.axis_slices,
             },
         ));
@@ -551,7 +555,7 @@ pub fn calculate_cluster_factors(
     if is_orthographic {
         Vec2::new(-near, z_slices / (-far - -near))
     } else {
-        let z_slices_of_ln_zfar_over_znear = z_slices / (far / near).ln();
+        let z_slices_of_ln_zfar_over_znear = (z_slices - 1.0) / (far / near).ln();
         Vec2::new(
             z_slices_of_ln_zfar_over_znear,
             near.ln() * z_slices_of_ln_zfar_over_znear,
@@ -680,12 +684,13 @@ pub fn prepare_lights(
 
         let is_orthographic = extracted_view.projection.w_axis.w == 1.0;
         let cluster_factors_zw = calculate_cluster_factors(
-            extracted_view.near,
+            clusters.near,
             extracted_view.far,
             clusters.axis_slices.z as f32,
             is_orthographic,
         );
 
+        let n_clusters = clusters.axis_slices.x * clusters.axis_slices.y * clusters.axis_slices.z;
         let mut gpu_lights = GpuLights {
             directional_lights: [GpuDirectionalLight::default(); MAX_DIRECTIONAL_LIGHTS],
             ambient_color: Vec4::from_slice(&ambient_light.color.as_linear_rgba_f32())
@@ -696,7 +701,7 @@ pub fn prepare_lights(
                 cluster_factors_zw.x,
                 cluster_factors_zw.y,
             ),
-            cluster_dimensions: clusters.axis_slices.extend(0),
+            cluster_dimensions: clusters.axis_slices.extend(n_clusters),
             n_directional_lights: directional_lights.iter().len() as u32,
         };
 

crates/bevy_pbr/src/render/mesh_view_bind_group.wgsl

@@ -38,7 +38,7 @@ struct Lights {
     // NOTE: this array size must be kept in sync with the constants defined bevy_pbr2/src/render/light.rs
     directional_lights: array<DirectionalLight, 1u>;
     ambient_color: vec4<f32>;
-    // x/y/z dimensions
+    // x/y/z dimensions and n_clusters in w
     cluster_dimensions: vec4<u32>;
     // xy are vec2<f32>(cluster_dimensions.xy) / vec2<f32>(view.width, view.height)
     //

crates/bevy_pbr/src/render/pbr.wgsl

@@ -244,14 +244,22 @@ fn view_z_to_z_slice(view_z: f32, is_orthographic: bool) -> u32 {
         return u32(floor((view_z - lights.cluster_factors.z) * lights.cluster_factors.w));
     } else {
         // NOTE: had to use -view_z to make it positive else log(negative) is nan
-        return u32(floor(log(-view_z) * lights.cluster_factors.z - lights.cluster_factors.w));
+        return min(
+            u32(log(-view_z) * lights.cluster_factors.z - lights.cluster_factors.w + 1.0),
+            lights.cluster_dimensions.z - 1u
+        );
     }
 }
 
 fn fragment_cluster_index(frag_coord: vec2<f32>, view_z: f32, is_orthographic: bool) -> u32 {
     let xy = vec2<u32>(floor(frag_coord * lights.cluster_factors.xy));
     let z_slice = view_z_to_z_slice(view_z, is_orthographic);
-    return (xy.y * lights.cluster_dimensions.x + xy.x) * lights.cluster_dimensions.z + z_slice;
+    // NOTE: Restricting cluster index to avoid undefined behavior when accessing uniform buffer
+    // arrays based on the cluster index.
+    return min(
+        (xy.y * lights.cluster_dimensions.x + xy.x) * lights.cluster_dimensions.z + z_slice,
+        lights.cluster_dimensions.w - 1u
+    );
 }
 
 struct ClusterOffsetAndCount {