render: provide shadow texture atlas frame for each shadow casting light

lyra-game/src/render/graph/passes/meshes.rs

@@ -14,8 +14,7 @@ use crate::render::{
 };
 
 use super::{
-    BasePassSlots, LightBasePassSlots, LightCullComputePassSlots, MeshBufferStorage, RenderAssets,
+    BasePassSlots, LightBasePassSlots, LightCullComputePassSlots, MeshBufferStorage, RenderAssets, RenderMeshes, ShadowMapsPassSlots
-    RenderMeshes, ShadowMapsPassSlots,
 };
 
 #[derive(Debug, Hash, Clone, Default, PartialEq, RenderGraphLabel)]
@@ -160,7 +159,7 @@ impl Node for MeshPass {
                         binding: 3,
                         visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                         ty: wgpu::BindingType::Buffer {
-                            ty: wgpu::BufferBindingType::Uniform,
+                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                             has_dynamic_offset: false,
                             min_binding_size: None,
                         },

lyra-game/src/render/graph/passes/shadows.rs

@@ -1,4 +1,4 @@
-use std::{collections::VecDeque, mem, num::NonZeroU64, ops::Deref, rc::Rc, sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}};
+use std::{collections::VecDeque, mem, num::NonZeroU64, rc::Rc, sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}};
 
 use lyra_ecs::{
     query::{filter::Has, Entities}, AtomicRef, Component, Entity, ResourceData
@@ -29,15 +29,20 @@ pub enum ShadowMapsPassSlots {
 #[derive(Debug, Clone, Hash, PartialEq, RenderGraphLabel)]
 pub struct ShadowMapsPassLabel;
 
+#[derive(Clone, Copy)]
 struct LightDepthMap {
-    light_projection_buffer: Arc<wgpu::Buffer>,
+    //light_projection_buffer: Arc<wgpu::Buffer>,
-    bindgroup: wgpu::BindGroup,
+    //bindgroup: wgpu::BindGroup,
     atlas_index: u64,
+    uniform_index: u64,
 }
 
 pub struct ShadowMapsPass {
     bgl: Arc<wgpu::BindGroupLayout>,
     atlas_size_buffer: Arc<wgpu::Buffer>,
+    light_uniforms_buffer: Arc<wgpu::Buffer>,
+    light_uniforms_index: u64,
+    uniforms_bg: Arc<wgpu::BindGroup>,
     /// depth maps for a light owned by an entity.
     depth_maps: FxHashMap<Entity, LightDepthMap>,
 
@@ -63,8 +68,8 @@ impl ShadowMapsPass {
                         binding: 0,
                         visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                         ty: wgpu::BindingType::Buffer {
-                            ty: wgpu::BufferBindingType::Uniform,
+                            ty: wgpu::BufferBindingType::Storage { read_only: true },
-                            has_dynamic_offset: false,
+                            has_dynamic_offset: true,
                             min_binding_size: Some(
                                 NonZeroU64::new(mem::size_of::<LightShadowUniform>() as _).unwrap(),
                             ),
@@ -101,8 +106,32 @@ impl ShadowMapsPass {
             ..Default::default()
         });
 
+        let uniforms_buffer =
+            device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some("buffer_shadow_maps_light"),
+                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
+                size: device.limits().max_storage_buffer_binding_size as u64,
+                mapped_at_creation: false,
+            });
+
+        let uniforms_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("bind_group_shadows"),
+            layout: &bgl,
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
+                    buffer: &uniforms_buffer,
+                    offset: 0,
+                    size: Some(NonZeroU64::new(mem::size_of::<LightShadowUniform>() as _).unwrap()),
+                }),
+            }],
+        });
+
         Self {
             bgl,
+            light_uniforms_buffer: Arc::new(uniforms_buffer),
+            light_uniforms_index: 0,
+            uniforms_bg: Arc::new(uniforms_bg),
             atlas_size_buffer: Arc::new(atlas_size_buffer),
             depth_maps: Default::default(),
             transform_buffers: None,
@@ -116,7 +145,7 @@ impl ShadowMapsPass {
     }
 
     /// Create a depth map and return the id of the depth map in the texture atlas.
-    fn create_depth_map(&mut self, device: &wgpu::Device, entity: Entity, light_pos: Transform) -> u64 {
+    fn create_depth_map(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, entity: Entity, light_pos: Transform) -> LightDepthMap {
         const NEAR_PLANE: f32 = 0.1;
         const FAR_PLANE: f32 = 45.0;
 
@@ -137,38 +166,23 @@ impl ShadowMapsPass {
             atlas_frame,
         };
 
-        let light_projection_buffer =
+        let uniform_index = self.light_uniforms_index;
-            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+        self.light_uniforms_index += 1;
-                label: Some("buffer_shadow_maps_light"),
-                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
-                contents: bytemuck::bytes_of(&uniform),
-            });
 
-        let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
+        //self.light_uniforms_buffer
-            label: Some("shadow_maps_bind_group"),
+        let offset = uniform_index_offset(&device.limits(), uniform_index);
-            layout: &self.bgl,
+        queue.write_buffer(&self.light_uniforms_buffer, offset as u64, bytemuck::bytes_of(&uniform));
-            entries: &[
-                wgpu::BindGroupEntry {
-                    binding: 0,
-                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
-                        buffer: &light_projection_buffer,
-                        offset: 0,
-                        size: None,
-                    }),
-                }
-            ],
-        });
 
+        let v = LightDepthMap {
+            atlas_index,
+            uniform_index,
+        };
         self.depth_maps.insert(
             entity,
-            LightDepthMap {
+            v,
-                light_projection_buffer: Arc::new(light_projection_buffer),
-                bindgroup: bg,
-                atlas_index
-            },
         );
 
-        atlas_index
+        v
     }
 
     fn transform_buffers(&self) -> AtomicRef<TransformBuffers> {
@@ -211,10 +225,10 @@ impl Node for ShadowMapsPass {
             Some(SlotValue::Sampler(self.atlas_sampler.clone())),
         );
 
-        node.add_sampler_slot(
+        node.add_buffer_slot(
             ShadowMapsPassSlots::ShadowLightUniformsBuffer,
             SlotAttribute::Output,
-            Some(SlotValue::Lazy),
+            Some(SlotValue::Buffer(self.light_uniforms_buffer.clone())),
         );
 
         node.add_buffer_slot(
@@ -230,7 +244,7 @@ impl Node for ShadowMapsPass {
         &mut self,
         graph: &mut crate::render::graph::RenderGraph,
         world: &mut lyra_ecs::World,
-        _: &mut crate::render::graph::RenderGraphContext,
+        context: &mut crate::render::graph::RenderGraphContext,
     ) {
         self.render_meshes = world.try_get_resource_data::<RenderMeshes>();
         self.transform_buffers = world.try_get_resource_data::<TransformBuffers>();
@@ -245,7 +259,7 @@ impl Node for ShadowMapsPass {
             if !self.depth_maps.contains_key(&entity) {
                
                 // TODO: dont pack the textures as they're added
-                let atlas_index = self.create_depth_map(graph.device(), entity, *pos);
+                let atlas_index = self.create_depth_map(graph.device(), &context.queue, entity, *pos);
                 index_components_queue.push_back((entity, atlas_index));
 
                 debug!("Created depth map for {:?} light entity", entity);
@@ -253,17 +267,13 @@ impl Node for ShadowMapsPass {
         }
         
         // now consume from the queue adding the components to the entities
-        while let Some((entity, atlas_id)) = index_components_queue.pop_front() {
+        while let Some((entity, depth)) = index_components_queue.pop_front() {
-            world.insert(entity, LightShadowMapId(atlas_id));
+            world.insert(entity, LightShadowMapId {
+                atlas_index: depth.atlas_index,
+                uniform_index: depth.uniform_index,
+            });
         }
 
-        // update the light projection buffer slot
-        let (_, dir_depth_map) = self.depth_maps.iter().next().unwrap();
-        let val = graph
-            .slot_value_mut(ShadowMapsPassSlots::ShadowLightUniformsBuffer)
-            .unwrap();
-        *val = SlotValue::Buffer(dir_depth_map.light_projection_buffer.clone());
-
         if self.pipeline.is_none() {
             let shader = Rc::new(Shader {
                 label: Some("shader_shadows".into()),
@@ -348,6 +358,7 @@ impl Node for ShadowMapsPass {
             });
             pass.set_pipeline(&pipeline);
             let viewport = atlas.texture_viewport(dir_depth_map.atlas_index);
+            debug!("Rendering shadow map to viewport: {viewport:?}, uniform index: {}", dir_depth_map.uniform_index);
             // only render to the light's map in the atlas
             pass.set_viewport(viewport.offset.x as _, viewport.offset.y as _, viewport.size.x as _, viewport.size.y as _, 0.0, 1.0);
             // only clear the light map in the atlas
@@ -362,7 +373,9 @@ impl Node for ShadowMapsPass {
                 }
                 let buffers = buffers.unwrap();
 
-                pass.set_bind_group(0, &dir_depth_map.bindgroup, &[]);
+                let uniform_index = uniform_index_offset(&context.device.limits(), dir_depth_map.uniform_index);
+                //debug!("Uniform offset: {uniform_index}");
+                pass.set_bind_group(0, &self.uniforms_bg, &[uniform_index]);
 
                 // Get the bindgroup for job's transform and bind to it using an offset.
                 let bindgroup = transforms.bind_group(job.transform_id);
@@ -395,7 +408,7 @@ impl Node for ShadowMapsPass {
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
-struct LightShadowUniform {
+pub struct LightShadowUniform {
     space_mat: glam::Mat4,
     atlas_frame: AtlasViewport, // 2xUVec2 (4xf32), so no padding needed
 }
@@ -405,13 +418,18 @@ struct LightShadowUniform {
 /// An entity owns a light. If that light casts shadows, this will contain the ID of the shadow
 /// map inside of the [`TextureAtlas`].
 #[derive(Debug, Default, Copy, Clone, Component)]
-pub struct LightShadowMapId(u64);
+pub struct LightShadowMapId {
+    atlas_index: u64,
+    uniform_index: u64,
+}
 
-impl Deref for LightShadowMapId {
+impl LightShadowMapId {
-    type Target = u64;
+    pub fn atlas_index(&self) -> u64 {
+        self.atlas_index
+    }
 
-    fn deref(&self) -> &Self::Target {
+    pub fn uniform_index(&self) -> u64 {
-        &self.0
+        self.uniform_index
     }
 }
 
@@ -428,3 +446,8 @@ impl LightShadowMapAtlas {
         self.0.write().unwrap()
     }
 }
+
+fn uniform_index_offset(limits: &wgpu::Limits, uniform_idx: u64) -> u32 {
+    let t = uniform_idx as u32 % (limits.max_storage_buffer_binding_size / mem::size_of::<LightShadowUniform>() as u32);
+    t * limits.min_uniform_buffer_offset_alignment
+}

lyra-game/src/render/light/mod.rs

@@ -12,6 +12,8 @@ use crate::math::Transform;
 
 use self::directional::DirectionalLight;
 
+use super::graph::LightShadowMapId;
+
 const MAX_LIGHT_COUNT: usize = 16;
 
 /// A struct that stores a list of lights in a wgpu::Buffer.
@@ -166,18 +168,21 @@ impl LightUniformBuffers {
         let _ = world_tick;
         let mut lights = vec![];
 
-        for (point_light, transform) in world.view_iter::<(&PointLight, &Transform)>() {
+        for (point_light, transform, shadow_map_id) in world.view_iter::<(&PointLight, &Transform, Option<&LightShadowMapId>)>() {
-            let uniform = LightUniform::from_point_light_bundle(&point_light, &transform);
+            let shadow_map_id = shadow_map_id.map(|m| m.clone());
+            let uniform = LightUniform::from_point_light_bundle(&point_light, &transform, shadow_map_id);
             lights.push(uniform);
         }
 
-        for (spot_light, transform) in world.view_iter::<(&SpotLight, &Transform)>() {
+        for (spot_light, transform, shadow_map_id) in world.view_iter::<(&SpotLight, &Transform, Option<&LightShadowMapId>)>() {
-            let uniform = LightUniform::from_spot_light_bundle(&spot_light, &transform);
+            let shadow_map_id = shadow_map_id.map(|m| m.clone());
+            let uniform = LightUniform::from_spot_light_bundle(&spot_light, &transform, shadow_map_id);
             lights.push(uniform);
         }
 
-        for (dir_light, transform) in world.view_iter::<(&DirectionalLight, &Transform)>() {
+        for (dir_light, transform, shadow_map_id) in world.view_iter::<(&DirectionalLight, &Transform, Option<&LightShadowMapId>)>() {
-            let uniform = LightUniform::from_directional_bundle(&dir_light, &transform);
+            let shadow_map_id = shadow_map_id.map(|m| m.clone());
+            let uniform = LightUniform::from_directional_bundle(&dir_light, &transform, shadow_map_id);
             lights.push(uniform);
         }
 
@@ -216,10 +221,11 @@ pub(crate) struct LightUniform {
 
     pub spot_cutoff_rad: f32,
     pub spot_outer_cutoff_rad: f32,
+    pub light_shadow_uniform_index: i32,
 }
 
 impl LightUniform {
-    pub fn from_point_light_bundle(light: &PointLight, transform: &Transform) -> Self {
+    pub fn from_point_light_bundle(light: &PointLight, transform: &Transform, map_id: Option<LightShadowMapId>) -> Self {
         Self {
             light_type: LightType::Point as u32,
             enabled: light.enabled as u32,
@@ -233,11 +239,12 @@ impl LightUniform {
 
             spot_cutoff_rad: 0.0,
             spot_outer_cutoff_rad: 0.0,
+            light_shadow_uniform_index:  map_id.map(|m| m.uniform_index() as i32).unwrap_or(-1),
 
         }
     }
 
-    pub fn from_directional_bundle(light: &DirectionalLight, transform: &Transform) -> Self {
+    pub fn from_directional_bundle(light: &DirectionalLight, transform: &Transform, map_id: Option<LightShadowMapId>) -> Self {
         Self {
             light_type: LightType::Directional as u32,
             enabled: light.enabled as u32,
@@ -251,11 +258,12 @@ impl LightUniform {
 
             spot_cutoff_rad: 0.0,
             spot_outer_cutoff_rad: 0.0,
+            light_shadow_uniform_index:  map_id.map(|m| m.uniform_index() as i32).unwrap_or(-1),
         }
     }
 
     // Create the SpotLightUniform from an ECS bundle
-    pub fn from_spot_light_bundle(light: &SpotLight, transform: &Transform) -> Self {
+    pub fn from_spot_light_bundle(light: &SpotLight, transform: &Transform, map_id: Option<LightShadowMapId>) -> Self {
         Self {
             light_type: LightType::Spotlight as u32,
             enabled: light.enabled as u32,
@@ -269,6 +277,7 @@ impl LightUniform {
 
             spot_cutoff_rad: light.cutoff.to_radians(),
             spot_outer_cutoff_rad: light.outer_cutoff.to_radians(),
+            light_shadow_uniform_index:  map_id.map(|m| m.uniform_index() as i32).unwrap_or(-1),
         }
     }
 }

lyra-game/src/render/shaders/base.wgsl

@@ -54,6 +54,7 @@ struct Light {
 
     spot_cutoff: f32,
     spot_outer_cutoff: f32,
+    light_shadow_uniform_index: i32,
 };
 
 struct Lights {
@@ -85,9 +86,7 @@ fn vs_main(
     // the normal mat is actually only a mat3x3, but there's a bug in wgpu: https://github.com/gfx-rs/wgpu-rs/issues/36
     let normal_mat4 = u_model_transform_data.normal_matrix;
     let normal_mat = mat3x3(normal_mat4[0].xyz, normal_mat4[1].xyz, normal_mat4[2].xyz);
-    out.world_normal = normalize(normal_mat * model.normal, );
+    out.world_normal = normalize(normal_mat * model.normal);
-    
-    out.frag_pos_light_space = u_light_shadow.light_space_matrix * world_position;
 
     return out;
 }
@@ -114,10 +113,15 @@ struct LightShadowMapUniform {
     atlas_frame: TextureAtlasFrame,
 }
 
+struct LightShadowMapUniformAligned {
+    @size(256)
+    inner: LightShadowMapUniform
+}
+
 @group(4) @binding(0)
 var<storage, read_write> u_light_indices: array<u32>;
 @group(4) @binding(1)
-var t_light_grid: texture_storage_2d<rg32uint, read_write>; // vec2<u32>
+var t_light_grid: texture_storage_2d<rg32uint, read_write>; // rg32uint = vec2<u32>
 
 @group(5) @binding(0)
 var t_shadow_maps_atlas: texture_depth_2d;
@@ -126,7 +130,7 @@ var s_shadow_maps_atlas: sampler;
 @group(5) @binding(2)
 var<uniform> u_shadow_maps_atlas_size: vec2<u32>;
 @group(5) @binding(3)
-var<uniform> u_light_shadow: LightShadowMapUniform;
+var<storage, read> u_light_shadow: array<LightShadowMapUniformAligned>;
 
 @fragment
 fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
@@ -148,13 +152,18 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     let light_offset = tile.x;
     let light_count = tile.y;
 
+    let atlas_dimensions: vec2<i32> = textureDimensions(t_shadow_maps_atlas);
+
     for (var i = 0u; i < light_count; i++) {
         let light_index = u_light_indices[light_offset + i];
         let light: Light = u_lights.data[light_index];
 
         if (light.light_ty == LIGHT_TY_DIRECTIONAL) {
             let light_dir = normalize(-light.direction);
-            let shadow = calc_shadow(in.world_normal, light_dir, in.frag_pos_light_space, u_light_shadow.atlas_frame);
+            let shadow_u: LightShadowMapUniform = u_light_shadow[light.light_shadow_uniform_index].inner;
+            let frag_pos_light_space = shadow_u.light_space_matrix * vec4<f32>(in.world_position, 1.0);
+
+            let shadow = calc_shadow(in.world_normal, light_dir, frag_pos_light_space, atlas_dimensions, shadow_u.atlas_frame);
             light_res += blinn_phong_dir_light(in.world_position, in.world_normal, light, u_material, specular_color, shadow);
         } else if (light.light_ty == LIGHT_TY_POINT) {
             light_res += blinn_phong_point_light(in.world_position, in.world_normal, light, u_material, specular_color);
@@ -167,7 +176,7 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     return vec4<f32>(light_object_res, object_color.a);
 }
 
-fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: vec4<f32>, atlas_region: TextureAtlasFrame) -> f32 {
+fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: vec4<f32>, atlas_dimensions: vec2<i32>, atlas_region: TextureAtlasFrame) -> f32 {
     var proj_coords = frag_pos_light_space.xyz / frag_pos_light_space.w;
     // for some reason the y component is clipped after transforming
     proj_coords.y = -proj_coords.y;
@@ -181,8 +190,8 @@ fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: ve
     let xy_remapped = proj_coords.xy * 0.5 + 0.5;
 
     // no need to get the y since the maps are square
-    let atlas_start = f32(atlas_region.offset.x) / f32(u_shadow_maps_atlas_size.x);
+    let atlas_start = f32(atlas_region.offset.x) / f32(atlas_dimensions.x);
-    let atlas_end = f32(atlas_region.offset.x + atlas_region.size.x) / f32(u_shadow_maps_atlas_size.x);
+    let atlas_end = f32(atlas_region.offset.x + atlas_region.size.x) / f32(atlas_dimensions.x);
     // lerp the tex coords to the shadow map for this light.
     proj_coords.x = mix(atlas_start, atlas_end, xy_remapped.x);
     proj_coords.y = mix(atlas_start, atlas_end, xy_remapped.y);
@@ -195,7 +204,7 @@ fn calc_shadow(normal: vec3<f32>, light_dir: vec3<f32>, frag_pos_light_space: ve
 
     // must manually apply offset to the texture coords since `textureSampleLevel` requires a
     // const value.
-    let offset_coords = proj_coords.xy + (vec2<f32>(atlas_region.offset) / vec2<f32>(u_shadow_maps_atlas_size));
+    let offset_coords = proj_coords.xy + (vec2<f32>(atlas_region.offset) / vec2<f32>(atlas_dimensions));
     let closest_depth = textureSampleLevel(t_shadow_maps_atlas, s_shadow_maps_atlas, offset_coords, 0.0);
     let current_depth = proj_coords.z;
 

lyra-game/src/render/shaders/shadows.wgsl

@@ -14,7 +14,7 @@ struct LightShadowMapUniform {
 }
 
 @group(0) @binding(0)
-var<uniform> u_light_shadow: LightShadowMapUniform;
+var<storage, read> u_light_shadow: LightShadowMapUniform;
 
 @group(1) @binding(0)
 var<uniform> u_model_transform_data: TransformData;