Omnilight shadow map wgsl (#15)

examples/3d/3d_scene_pipelined.rs

@@ -3,9 +3,9 @@ use bevy::{
     diagnostic::{FrameTimeDiagnosticsPlugin, LogDiagnosticsPlugin},
     ecs::prelude::*,
     input::Input,
-    math::Vec3,
+    math::{Quat, Vec3},
-    pbr2::{OmniLightBundle, PbrBundle, StandardMaterial},
+    pbr2::{OmniLight, OmniLightBundle, PbrBundle, StandardMaterial},
-    prelude::{App, Assets, KeyCode, Transform},
+    prelude::{App, Assets, BuildChildren, KeyCode, Transform},
     render2::{
         camera::PerspectiveCameraBundle,
         color::Color,
@@ -42,6 +42,32 @@ fn setup(
         }),
         ..Default::default()
     });
+
+    let mut transform = Transform::from_xyz(2.5, 2.5, 0.0);
+    transform.rotate(Quat::from_rotation_z(std::f32::consts::FRAC_PI_2));
+    commands.spawn_bundle(PbrBundle {
+        mesh: meshes.add(Mesh::from(shape::Plane { size: 5.0 })),
+        transform,
+        material: materials.add(StandardMaterial {
+            base_color: Color::INDIGO,
+            perceptual_roughness: 1.0,
+            ..Default::default()
+        }),
+        ..Default::default()
+    });
+
+    let mut transform = Transform::from_xyz(0.0, 2.5, -2.5);
+    transform.rotate(Quat::from_rotation_x(std::f32::consts::FRAC_PI_2));
+    commands.spawn_bundle(PbrBundle {
+        mesh: meshes.add(Mesh::from(shape::Plane { size: 5.0 })),
+        transform,
+        material: materials.add(StandardMaterial {
+            base_color: Color::INDIGO,
+            perceptual_roughness: 1.0,
+            ..Default::default()
+        }),
+        ..Default::default()
+    });
     // cube
     commands
         .spawn_bundle(PbrBundle {
@@ -69,14 +95,89 @@ fn setup(
             ..Default::default()
         })
         .insert(Movable);
+
     // light
-    commands.spawn_bundle(OmniLightBundle {
+    commands
-        transform: Transform::from_xyz(5.0, 8.0, 2.0),
+        .spawn_bundle(OmniLightBundle {
+            // transform: Transform::from_xyz(5.0, 8.0, 2.0),
+            transform: Transform::from_xyz(1.0, 2.0, 0.0),
+            omni_light: OmniLight {
+                color: Color::RED,
+                ..Default::default()
+            },
+            ..Default::default()
+        })
+        .with_children(|builder| {
+            builder.spawn_bundle(PbrBundle {
+                mesh: meshes.add(Mesh::from(shape::UVSphere {
+                    radius: 0.1,
+                    ..Default::default()
+                })),
+                material: materials.add(StandardMaterial {
+                    base_color: Color::RED,
+                    emissive: Color::rgba_linear(100.0, 0.0, 0.0, 0.0),
+                    ..Default::default()
+                }),
                 ..Default::default()
             });
+        });
+
+    // light
+    commands
+        .spawn_bundle(OmniLightBundle {
+            // transform: Transform::from_xyz(5.0, 8.0, 2.0),
+            transform: Transform::from_xyz(-1.0, 2.0, 0.0),
+            omni_light: OmniLight {
+                color: Color::GREEN,
+                ..Default::default()
+            },
+            ..Default::default()
+        })
+        .with_children(|builder| {
+            builder.spawn_bundle(PbrBundle {
+                mesh: meshes.add(Mesh::from(shape::UVSphere {
+                    radius: 0.1,
+                    ..Default::default()
+                })),
+                material: materials.add(StandardMaterial {
+                    base_color: Color::GREEN,
+                    emissive: Color::rgba_linear(0.0, 100.0, 0.0, 0.0),
+                    ..Default::default()
+                }),
+                ..Default::default()
+            });
+        });
+
+    // light
+    commands
+        .spawn_bundle(OmniLightBundle {
+            // transform: Transform::from_xyz(5.0, 8.0, 2.0),
+            transform: Transform::from_xyz(0.0, 4.0, 0.0),
+            omni_light: OmniLight {
+                color: Color::BLUE,
+                ..Default::default()
+            },
+            ..Default::default()
+        })
+        .with_children(|builder| {
+            builder.spawn_bundle(PbrBundle {
+                mesh: meshes.add(Mesh::from(shape::UVSphere {
+                    radius: 0.1,
+                    ..Default::default()
+                })),
+                material: materials.add(StandardMaterial {
+                    base_color: Color::BLUE,
+                    emissive: Color::rgba_linear(0.0, 0.0, 100.0, 0.0),
+                    ..Default::default()
+                }),
+                ..Default::default()
+            });
+        });
+
     // camera
     commands.spawn_bundle(PerspectiveCameraBundle {
-        transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y),
+        transform: Transform::from_xyz(-2.0, 5.0, 7.5)
+            .looking_at(Vec3::new(0.0, 2.0, 0.0), Vec3::Y),
         ..Default::default()
     });
 }

pipelined/bevy_pbr2/src/render/light.rs

@@ -1,6 +1,6 @@
 use crate::{AmbientLight, ExtractedMeshes, MeshMeta, OmniLight, PbrShaders};
 use bevy_ecs::{prelude::*, system::SystemState};
-use bevy_math::{Mat4, Vec3, Vec4};
+use bevy_math::{const_vec3, Mat4, Vec3, Vec4};
 use bevy_render2::{
     color::Color,
     core_pipeline::Transparent3dPhase,
@@ -34,10 +34,12 @@ pub struct ExtractedPointLight {
 #[derive(Copy, Clone, AsStd140, Default, Debug)]
 pub struct GpuLight {
     color: Vec4,
-    range: f32,
+    // proj: Mat4,
-    radius: f32,
     position: Vec3,
-    view_proj: Mat4,
+    inverse_square_range: f32,
+    radius: f32,
+    near: f32,
+    far: f32,
 }
 
 #[repr(C)]
@@ -54,7 +56,7 @@ pub const MAX_OMNI_LIGHTS: usize = 10;
 pub const SHADOW_SIZE: Extent3d = Extent3d {
     width: 1024,
     height: 1024,
-    depth_or_array_layers: MAX_OMNI_LIGHTS as u32,
+    depth_or_array_layers: 6 * MAX_OMNI_LIGHTS as u32,
 };
 pub const SHADOW_FORMAT: TextureFormat = TextureFormat::Depth32Float;
 
@@ -148,7 +150,7 @@ impl FromWorld for ShadowShaders {
                 topology: PrimitiveTopology::TriangleList,
                 strip_index_format: None,
                 front_face: FrontFace::Ccw,
-                cull_mode: Some(Face::Back),
+                cull_mode: None,
                 polygon_mode: PolygonMode::Fill,
                 clamp_depth: false,
                 conservative: false,
@@ -193,8 +195,52 @@ pub fn extract_lights(
     }
 }
 
+// Can't do `Vec3::Y * -1.0` because mul isn't const
+const NEGATIVE_X: Vec3 = const_vec3!([-1.0, 0.0, 0.0]);
+const NEGATIVE_Y: Vec3 = const_vec3!([0.0, -1.0, 0.0]);
+const NEGATIVE_Z: Vec3 = const_vec3!([0.0, 0.0, -1.0]);
+
+struct CubeMapFace {
+    target: Vec3,
+    up: Vec3,
+}
+
+// see https://www.khronos.org/opengl/wiki/Cubemap_Texture
+const CUBE_MAP_FACES: [CubeMapFace; 6] = [
+    // 0 	GL_TEXTURE_CUBE_MAP_POSITIVE_X
+    CubeMapFace {
+        target: NEGATIVE_X,
+        up: NEGATIVE_Y,
+    },
+    // 1 	GL_TEXTURE_CUBE_MAP_NEGATIVE_X
+    CubeMapFace {
+        target: Vec3::X,
+        up: NEGATIVE_Y,
+    },
+    // 2 	GL_TEXTURE_CUBE_MAP_POSITIVE_Y
+    CubeMapFace {
+        target: NEGATIVE_Y,
+        up: Vec3::Z,
+    },
+    // 3 	GL_TEXTURE_CUBE_MAP_NEGATIVE_Y
+    CubeMapFace {
+        target: Vec3::Y,
+        up: NEGATIVE_Z,
+    },
+    // 4 	GL_TEXTURE_CUBE_MAP_POSITIVE_Z
+    CubeMapFace {
+        target: NEGATIVE_Z,
+        up: NEGATIVE_Y,
+    },
+    // 5 	GL_TEXTURE_CUBE_MAP_NEGATIVE_Z
+    CubeMapFace {
+        target: Vec3::Z,
+        up: NEGATIVE_Y,
+    },
+];
+
 pub struct ViewLight {
-    pub depth_texture: TextureView,
+    pub depth_texture_view: TextureView,
 }
 
 pub struct ViewLights {
@@ -248,7 +294,19 @@ pub fn prepare_lights(
         };
 
         // TODO: this should select lights based on relevance to the view instead of the first ones that show up in a query
-        for (i, light) in lights.iter().enumerate().take(MAX_OMNI_LIGHTS) {
+        for (light_index, light) in lights.iter().enumerate().take(MAX_OMNI_LIGHTS) {
+            let projection =
+                Mat4::perspective_rh(std::f32::consts::FRAC_PI_2, 1.0, 0.1, light.range);
+
+            // ignore scale because we don't want to effectively scale light radius and range
+            // by applying those as a view transform to shadow map rendering of objects
+            // and ignore rotation because we want the shadow map projections to align with the axes
+            let view_translation = GlobalTransform::from_translation(light.transform.translation);
+
+            for (face_index, CubeMapFace { target, up }) in CUBE_MAP_FACES.iter().enumerate() {
+                // use the cubemap projection direction
+                let view_rotation = GlobalTransform::identity().looking_at(*target, *up);
+
                 let depth_texture_view =
                     light_depth_texture
                         .texture
@@ -259,37 +317,18 @@ pub fn prepare_lights(
                             aspect: TextureAspect::All,
                             base_mip_level: 0,
                             mip_level_count: None,
-                        base_array_layer: i as u32,
+                            base_array_layer: (light_index * 6 + face_index) as u32,
                             array_layer_count: NonZeroU32::new(1),
                         });
 
-            let view_transform = GlobalTransform::from_translation(light.transform.translation)
-                .looking_at(Vec3::default(), Vec3::Y);
-            // TODO: configure light projection based on light configuration
-            let projection =
-                Mat4::perspective_rh(std::f32::consts::FRAC_PI_2, 1.0, 0.1, light.range);
-
-            gpu_lights.lights[i] = GpuLight {
-                // 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(),
-                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(),
-            };
-
                 let view_light_entity = commands
                     .spawn()
                     .insert_bundle((
-                    ViewLight {
+                        ViewLight { depth_texture_view },
-                        depth_texture: depth_texture_view,
-                    },
                         ExtractedView {
                             width: SHADOW_SIZE.width,
                             height: SHADOW_SIZE.height,
-                        transform: view_transform.clone(),
+                            transform: view_translation * view_rotation,
                             projection,
                         },
                         RenderPhase::<ShadowPhase>::default(),
@@ -298,9 +337,36 @@ pub fn prepare_lights(
                 view_lights.push(view_light_entity);
             }
 
+            gpu_lights.lights[light_index] = GpuLight {
+                // 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(),
+                inverse_square_range: 1.0 / (light.range * light.range),
+                near: 0.1,
+                far: light.range,
+                // proj: projection,
+            };
+        }
+
+        let light_depth_texture_view =
+            light_depth_texture
+                .texture
+                .create_view(&TextureViewDescriptor {
+                    label: None,
+                    format: None,
+                    dimension: Some(TextureViewDimension::CubeArray),
+                    aspect: TextureAspect::All,
+                    base_mip_level: 0,
+                    mip_level_count: None,
+                    base_array_layer: 0 as u32,
+                    array_layer_count: None,
+                });
+
         commands.entity(entity).insert(ViewLights {
             light_depth_texture: light_depth_texture.texture,
-            light_depth_texture_view: light_depth_texture.default_view,
+            light_depth_texture_view,
             lights: view_lights,
             gpu_light_binding_index: light_meta.view_gpu_lights.push(gpu_lights),
         });
@@ -356,7 +422,7 @@ impl Node for ShadowPassNode {
                     label: Some("shadow_pass"),
                     color_attachments: &[],
                     depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
-                        view: &view_light.depth_texture,
+                        view: &view_light.depth_texture_view,
                         depth_ops: Some(Operations {
                             load: LoadOp::Clear(1.0),
                             store: true,
@@ -422,7 +488,8 @@ impl Draw for DrawShadowMesh {
             self.params.get(world);
         let view_uniform_offset = views.get(view).unwrap();
         let extracted_mesh = &extracted_meshes.into_inner().meshes[draw_key];
-        pass.set_render_pipeline(&shadow_shaders.into_inner().pipeline);
+        let shadow_shaders = shadow_shaders.into_inner();
+        pass.set_render_pipeline(&shadow_shaders.pipeline);
         pass.set_bind_group(
             0,
             light_meta

pipelined/bevy_pbr2/src/render/mod.rs

@@ -68,7 +68,7 @@ impl FromWorld for PbrShaders {
                         has_dynamic_offset: true,
                         // TODO: change this to ViewUniform::std140_size_static once crevice fixes this!
                         // Context: https://github.com/LPGhatguy/crevice/issues/29
-                        min_binding_size: BufferSize::new(1264),
+                        min_binding_size: BufferSize::new(512),
                     },
                     count: None,
                 },
@@ -79,7 +79,7 @@ impl FromWorld for PbrShaders {
                     ty: BindingType::Texture {
                         multisampled: false,
                         sample_type: TextureSampleType::Depth,
-                        view_dimension: TextureViewDimension::D2Array,
+                        view_dimension: TextureViewDimension::CubeArray,
                     },
                     count: None,
                 },

pipelined/bevy_pbr2/src/render/pbr.wgsl

@@ -89,10 +89,12 @@ struct StandardMaterial {
 
 struct OmniLight {
     color: vec4<f32>;
-    range: f32;
+    // projection: mat4x4<f32>;
-    radius: f32;
     position: vec3<f32>;
-    view_projection: mat4x4<f32>;
+    inverse_square_range: f32;
+    radius: f32;
+    near: f32;
+    far: f32;
 };
 
 [[block]]
@@ -115,7 +117,7 @@ let FLAGS_UNLIT_BIT: u32                      = 32u;
 [[group(0), binding(1)]]
 var lights: Lights;
 [[group(0), binding(2)]]
-var shadow_textures: texture_depth_2d_array;
+var shadow_textures: texture_depth_cube_array;
 [[group(0), binding(3)]]
 var shadow_textures_sampler: sampler_comparison;
 
@@ -301,7 +303,7 @@ fn omni_light(
     let light_to_frag = light.position.xyz - world_position.xyz;
     let distance_square = dot(light_to_frag, light_to_frag);
     let rangeAttenuation =
-        getDistanceAttenuation(distance_square, light.range);
+        getDistanceAttenuation(distance_square, light.inverse_square_range);
 
     // Specular.
     // Representative Point Area Lights.
@@ -346,22 +348,41 @@ fn omni_light(
     return ((diffuse + specular_light) * light.color.rgb) * (rangeAttenuation * NoL);
 }
 
-fn fetch_shadow(light_id: i32, homogeneous_coords: vec4<f32>) -> f32 {
+fn fetch_shadow(light_id: i32, frag_position: vec4<f32>) -> f32 {
-    if (homogeneous_coords.w <= 0.0) {
+    let light = lights.omni_lights[light_id];
-        return 1.0;
+
-    }
+    // because the shadow maps align with the axes and the frustum planes are at 45 degrees
-    // compensate for the Y-flip difference between the NDC and texture coordinates
+    // we can get the worldspace depth by taking the largest absolute axis
-    let flip_correction = vec2<f32>(0.5, -0.5);
+    let frag_ls = light.position.xyz - frag_position.xyz;
-    let proj_correction = 1.0 / homogeneous_coords.w;
+    let abs_position_ls = abs(frag_ls);
-    // compute texture coordinates for shadow lookup
+    let major_axis_magnitude = max(abs_position_ls.x, max(abs_position_ls.y, abs_position_ls.z));
-    let light_local = homogeneous_coords.xy * flip_correction * proj_correction + vec2<f32>(0.5, 0.5);
+
+    // do a full projection
+    // vec4 clip = light.projection * vec4(0.0, 0.0, -major_axis_magnitude, 1.0);
+    // float depth = (clip.z / clip.w);
+
+    // alternatively do only the necessary multiplications using near/far
+    let proj_r = light.far / (light.near - light.far);
+    let z = -major_axis_magnitude * proj_r + light.near * proj_r;
+    let w = major_axis_magnitude;
+    let depth = z / w;
+
+    // let shadow = texture(samplerCubeArrayShadow(t_Shadow, s_Shadow), vec4(frag_ls, i), depth - bias);
+
+    // manual depth testing
+    // float shadow = texture(samplerCubeArray(t_Shadow, s_Shadow), vec4(-frag_ls, 6 * i)).r;
+    // shadow = depth > shadow ? 0.0 : 1.0;
+    // o_Target = vec4(vec3(shadow * 20 - 19, depth * 20 - 19, 0.0), 1.0);
+    // o_Target = vec4(vec3(shadow * 20 - 19), 1.0);
+
     // do the lookup, using HW PCF and comparison
     // NOTE: Due to the non-uniform control flow above, we must use the Level variant of
     //       textureSampleCompare to avoid undefined behaviour due to some of the fragments in
     //       a quad (2x2 fragments) being processed not being sampled, and this messing with
     //       mip-mapping functionality. The shadow maps have no mipmaps so Level just samples
     //       from LOD 0.
-    return textureSampleCompareLevel(shadow_textures, shadow_textures_sampler, light_local, i32(light_id), homogeneous_coords.z * proj_correction);
+    let bias = 0.0001;
+    return textureSampleCompareLevel(shadow_textures, shadow_textures_sampler, frag_ls, i32(light_id), depth - bias);
 }
 
 struct FragmentInput {
@@ -453,7 +474,7 @@ fn fragment(in: FragmentInput) -> [[location(0)]] vec4<f32> {
         for (var i: i32 = 0; i < i32(lights.num_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 shadow = fetch_shadow(i, in.world_position);
             light_accum = light_accum + light_contrib * shadow;
         }