bevy_pbr2: Port directional lights and shadows to wgsl

pipelined/bevy_pbr2/src/render/light.rs

@@ -46,27 +46,30 @@ pub struct GpuOmniLight {
     view_proj: Mat4,
     color: Vec4,
     position: Vec3,
-    shadow_bias_min_max: Vec2,
     range: f32,
     radius: f32,
+    shadow_bias_min: f32,
+    shadow_bias_max: f32,
 }
 
 #[repr(C)]
 #[derive(Copy, Clone, AsStd140, Default, Debug)]
 pub struct GpuDirectionalLight {
+    view_projection: Mat4,
     color: Vec4,
     dir_to_light: Vec3,
-    shadow_bias_min_max: Vec2,
-    view_projection: Mat4,
+    shadow_bias_min: f32,
+    shadow_bias_max: f32,
 }
 
 #[repr(C)]
 #[derive(Copy, Clone, Debug, AsStd140)]
 pub struct GpuLights {
     omni_lights: [GpuOmniLight; MAX_OMNI_LIGHTS],
     directional_lights: [GpuDirectionalLight; MAX_DIRECTIONAL_LIGHTS],
-    len: UVec4,
     ambient_color: Vec4,
+    n_omni_lights: u32,
+    n_directional_lights: u32,
 }
 
 // NOTE: this must be kept in sync with the same constants in pbr.frag
@@ -280,12 +283,8 @@ pub fn prepare_lights(
 
         let mut gpu_lights = GpuLights {
             ambient_color: ambient_color.into(),
-            len: UVec4::new(
-                omni_lights.iter().len() as u32,
-                directional_lights.iter().len() as u32,
-                0,
-                0,
-            ),
+            n_omni_lights: omni_lights.iter().len() as u32,
+            n_directional_lights: directional_lights.iter().len() as u32,
             omni_lights: [GpuOmniLight::default(); MAX_OMNI_LIGHTS],
             directional_lights: [GpuDirectionalLight::default(); MAX_DIRECTIONAL_LIGHTS],
         };
@@ -316,12 +315,13 @@ pub fn prepare_lights(
                 // premultiply color by intensity
                 // we don't use the alpha at all, so no reason to multiply only [0..3]
                 color: (light.color.as_rgba_linear() * light.intensity).into(),
-                radius: light.radius.into(),
-                position: light.transform.translation.into(),
+                radius: light.radius,
+                position: light.transform.translation,
                 range: 1.0 / (light.range * light.range),
                 // this could technically be copied to the gpu from the light's ViewUniforms
                 view_proj: projection * view_transform.compute_matrix().inverse(),
-                shadow_bias_min_max: light.shadow_bias_min_max,
+                shadow_bias_min: light.shadow_bias_min_max.x,
+                shadow_bias_max: light.shadow_bias_min_max.y,
             };
 
             let view_light_entity = commands
@@ -374,10 +374,11 @@ pub fn prepare_lights(
                 // premultiply color by intensity
                 // we don't use the alpha at all, so no reason to multiply only [0..3]
                 color: (light.color.as_rgba_linear() * intensity).into(),
-                dir_to_light: dir_to_light.into(),
+                dir_to_light,
                 // NOTE: * view is correct, it should not be view.inverse() here
                 view_projection: projection * view,
-                shadow_bias_min_max: light.shadow_bias_min_max,
+                shadow_bias_min: light.shadow_bias_min_max.x,
+                shadow_bias_max: light.shadow_bias_min_max.y,
             };
 
             let depth_texture_view =

pipelined/bevy_pbr2/src/render/pbr.wgsl

@@ -93,15 +93,27 @@ struct OmniLight {
     position: vec3<f32>;
     range: f32;
     radius: f32;
+    shadow_bias_min: f32;
+    shadow_bias_max: f32;
+};
+
+struct DirectionalLight {
+    view_projection: mat4x4<f32>;
+    color: vec4<f32>;
+    direction_to_light: vec3<f32>;
+    shadow_bias_min: f32;
+    shadow_bias_max: f32;
 };
 
 [[block]]
 struct Lights {
     // NOTE: this array size must be kept in sync with the constants defined bevy_pbr2/src/render/light.rs
     // TODO: this can be removed if we move to storage buffers for light arrays
     omni_lights: array<OmniLight, 10>;
+    directional_lights: array<DirectionalLight, 1>;
     ambient_color: vec4<f32>;
-    num_lights: u32;
+    n_omni_lights: u32;
+    n_directional_lights: u32;
 };
 
 let FLAGS_BASE_COLOR_TEXTURE_BIT: u32         = 1u;
@@ -346,7 +358,22 @@ fn omni_light(
     return ((diffuse + specular_light) * light.color.rgb) * (rangeAttenuation * NoL);
 }
 
-fn fetch_shadow(light_id: i32, homogeneous_coords: vec4<f32>) -> f32 {
+fn directional_light(light: DirectionalLight, roughness: f32, NdotV: f32, normal: vec3<f32>, view: vec3<f32>, R: vec3<f32>, F0: vec3<f32>, diffuseColor: vec3<f32>) -> vec3<f32> {
+    let incident_light = light.direction_to_light.xyz;
+
+    let half_vector = normalize(incident_light + view);
+    let NoL = saturate(dot(normal, incident_light));
+    let NoH = saturate(dot(normal, half_vector));
+    let LoH = saturate(dot(incident_light, half_vector));
+
+    let diffuse = diffuseColor * Fd_Burley(roughness, NdotV, NoL, LoH);
+    let specularIntensity = 1.0;
+    let specular_light = specular(F0, roughness, half_vector, NdotV, NoL, NoH, LoH, specularIntensity);
+
+    return (specular_light + diffuse) * light.color.rgb * NoL;
+}
+
+fn fetch_shadow(light_id: i32, homogeneous_coords: vec4<f32>, shadow_bias: f32) -> f32 {
     if (homogeneous_coords.w <= 0.0) {
         return 1.0;
     }
@@ -356,7 +383,9 @@ fn fetch_shadow(light_id: i32, homogeneous_coords: vec4<f32>) -> f32 {
     // compute texture coordinates for shadow lookup
     let light_local = homogeneous_coords.xy * flip_correction * proj_correction + vec2<f32>(0.5, 0.5);
     // do the lookup, using HW PCF and comparison
-    return textureSampleCompare(shadow_textures, shadow_textures_sampler, light_local, i32(light_id), homogeneous_coords.z * proj_correction);
+    // NOTE: Due to non-uniform control flow above, we must use the level variant of the texture
+    //       sampler to avoid use of implicit derivatives causing possible undefined behavior.
+    return textureSampleCompareLevel(shadow_textures, shadow_textures_sampler, light_local, i32(light_id), homogeneous_coords.z * proj_correction - shadow_bias);
 }
 
 struct FragmentInput {
@@ -445,10 +474,27 @@ fn fragment(in: FragmentInput) -> [[location(0)]] vec4<f32> {
 
         // accumulate color
         var light_accum: vec3<f32> = vec3<f32>(0.0);
-        for (var i: i32 = 0; i < i32(lights.num_lights); i = i + 1) {
+        let n_omni_lights = i32(lights.n_omni_lights);
+        let n_directional_lights = i32(lights.n_directional_lights);
+        for (var i: i32 = 0; i < n_omni_lights; i = i + 1) {
             let light = lights.omni_lights[i];
             let light_contrib = omni_light(in.world_position.xyz, light, roughness, NdotV, N, V, R, F0, diffuse_color);
-            let shadow = fetch_shadow(i, light.view_projection * in.world_position);
+            let dir_to_light = normalize(light.position.xyz - in.world_position.xyz);
+            let shadow_bias = max(
+                light.shadow_bias_max * (1.0 - dot(in.world_normal, dir_to_light)),
+                light.shadow_bias_min
+            );
+            let shadow = fetch_shadow(i, light.view_projection * in.world_position, shadow_bias);
+            light_accum = light_accum + light_contrib * shadow;
+        }
+        for (var i: i32 = 0; i < n_directional_lights; i = i + 1) {
+            let light = lights.directional_lights[i];
+            let light_contrib = directional_light(light, roughness, NdotV, N, V, R, F0, diffuse_color);
+            let shadow_bias = max(
+                light.shadow_bias_max * (1.0 - dot(in.world_normal, light.direction_to_light.xyz)),
+                light.shadow_bias_min
+            );
+            let shadow = fetch_shadow(n_omni_lights + i, light.view_projection * in.world_position, shadow_bias);
             light_accum = light_accum + light_contrib * shadow;
         }