rewrite lamplights to use compute shader

assets/shaders/compute/texture_write.wgsl

@@ -0,0 +1,25 @@
+// 16-aligned!
+struct LightUpdate {
+    data: vec4<u32>,
+    data2: vec4<u32>,
+    position: vec2<u32>,
+}
+
+@group(0) @binding(0) var light_tex:   texture_storage_2d<rgba32uint, write>;
+@group(0) @binding(1) var light_tex_2: texture_storage_2d<rgba32uint, write>;
+
+@group(1) @binding(0) var<storage, read> light_changes: array<LightUpdate>;
+
+@compute @workgroup_size(64,1,1)
+fn main(
+  @builtin(global_invocation_id) id: vec3<u32>,
+) {
+  let i: u32 = id.x;
+  if (i >= arrayLength(&light_changes)) {
+    return;
+  }
+  let position = light_changes[i].position;
+
+  textureStore(light_tex,   position, light_changes[i].data);
+  textureStore(light_tex_2, position, light_changes[i].data2);
+}

engine/src/gfx.rs

@@ -599,7 +599,7 @@ impl GfxContext {
     }
 
     pub fn start_frame(&mut self, sco: &SurfaceTexture) -> (Encoders, TextureView) {
-        let end = self
+        let mut end = self
             .device
             .create_command_encoder(&CommandEncoderDescriptor {
                 label: Some("End encoder"),
@@ -617,7 +617,8 @@ impl GfxContext {
         }
 
         self.render_params.upload_to_gpu(&self.queue);
-        self.lamplights.apply_changes(&self.queue);
+        self.lamplights
+            .apply_changes(&self.queue, &self.device, &mut end);
 
         (
             Encoders {

engine/src/lamplights.rs

@@ -1,14 +1,17 @@
-use crate::{Texture, TextureBuilder};
+use crate::pbuffer::PBuffer;
+use crate::{compile_shader, Texture, TextureBuilder};
+use common::FastMap;
 use geom::Vec3;
 use ordered_float::OrderedFloat;
-use wgpu::TextureFormat;
+use wgpu::{
+    BufferUsages, CommandEncoder, ComputePassDescriptor, Device, Queue, TextureFormat,
+    TextureUsages,
+};
 
 #[derive(Clone, Copy, Debug)]
 #[repr(transparent)]
 struct EncodedLight(u32);
 
-u8slice_impl!(EncodedLight);
-
 impl EncodedLight {
     fn encode(chunk_origin: Vec3, light: Vec3) -> Self {
         let diff = light - chunk_origin;
@@ -27,25 +30,34 @@ impl EncodedLight {
 
 pub type LightChunkID = (u16, u16);
 
+#[derive(Copy, Clone)]
+#[repr(C)]
 struct LightChunkUpdate {
-    x: u16,
-    y: u16,
     lights: [EncodedLight; 4],
+    lights2: [EncodedLight; 4],
+    x: u32,
+    y: u32,
+    _pad: (u32, u32), // lights (vec4) is 16 bytes aligned
 }
 
+u8slice_impl!(LightChunkUpdate);
+
 pub struct LampLights {
     pub(crate) lightdata: Texture,
     pub(crate) lightdata2: Texture,
     pending_changes: Vec<LightChunkUpdate>,
-    pending_changes2: Vec<LightChunkUpdate>,
+    changes_buffer: PBuffer,
+    buffer_layout: wgpu::BindGroupLayout,
+    texture_write_bg: wgpu::BindGroup,
+    texture_write_pipeline: wgpu::ComputePipeline,
 }
 
 impl LampLights {
     pub const LIGHTCHUNK_SIZE: u32 = 32; // in meters, side length of a light chunk, can contain at most 4 lights
-    pub const MAP_SIZE: u32 = 25 * 1024; // in meters, side length of the map
+    pub const MAP_SIZE: u32 = 50 * 512; // in meters, side length of the map
     pub const LIGHTMAP_SIZE: u32 = Self::MAP_SIZE / Self::LIGHTCHUNK_SIZE; // in light chunks
 
-    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
+    pub fn new(device: &Device, queue: &Queue) -> Self {
         let lightdata = TextureBuilder::empty(
             Self::LIGHTMAP_SIZE,
             Self::LIGHTMAP_SIZE,
@@ -55,6 +67,7 @@ impl LampLights {
         .with_label("lightdata")
         .with_sampler(Texture::nearest_sampler())
         .with_srgb(false)
+        .with_usage(TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING)
         .build(device, queue);
 
         let lightdata2 = TextureBuilder::empty(
@@ -66,17 +79,93 @@ impl LampLights {
         .with_label("lightdata2")
         .with_sampler(Texture::nearest_sampler())
         .with_srgb(false)
+        .with_usage(TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING)
         .build(device, queue);
 
+        let texture_write_module =
+            compile_shader(device, "compute/texture_write", &FastMap::default());
+
+        let textures_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+            label: Some("texture_write"),
+            entries: &[
+                wgpu::BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: wgpu::ShaderStages::COMPUTE,
+                    ty: wgpu::BindingType::StorageTexture {
+                        access: wgpu::StorageTextureAccess::WriteOnly,
+                        format: TextureFormat::Rgba32Uint,
+                        view_dimension: wgpu::TextureViewDimension::D2,
+                    },
+                    count: None,
+                },
+                wgpu::BindGroupLayoutEntry {
+                    binding: 1,
+                    visibility: wgpu::ShaderStages::COMPUTE,
+                    ty: wgpu::BindingType::StorageTexture {
+                        access: wgpu::StorageTextureAccess::WriteOnly,
+                        format: TextureFormat::Rgba32Uint,
+                        view_dimension: wgpu::TextureViewDimension::D2,
+                    },
+                    count: None,
+                },
+            ],
+        });
+
+        let buffer_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+            label: Some("buffer_layout"),
+            entries: &[wgpu::BindGroupLayoutEntry {
+                binding: 0,
+                visibility: wgpu::ShaderStages::COMPUTE,
+                ty: wgpu::BindingType::Buffer {
+                    has_dynamic_offset: false,
+                    min_binding_size: None,
+                    ty: wgpu::BufferBindingType::Storage { read_only: true },
+                },
+                count: None,
+            }],
+        });
+
+        let texture_write_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+            label: Some("texture_write"),
+            bind_group_layouts: &[&textures_layout, &buffer_layout],
+            push_constant_ranges: &[],
+        });
+
+        let texture_write_pipeline =
+            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+                label: Some("texture_write"),
+                layout: Some(&texture_write_layout),
+                module: &texture_write_module,
+                entry_point: "main",
+            });
+
+        let texture_write_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("texture_write"),
+            layout: &textures_layout,
+            entries: &[
+                wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: wgpu::BindingResource::TextureView(&lightdata.view),
+                },
+                wgpu::BindGroupEntry {
+                    binding: 1,
+                    resource: wgpu::BindingResource::TextureView(&lightdata2.view),
+                },
+            ],
+        });
+
         Self {
             lightdata,
             lightdata2,
             pending_changes: Vec::new(),
-            pending_changes2: vec![],
+            changes_buffer: PBuffer::new(BufferUsages::COPY_DST | BufferUsages::STORAGE),
+            buffer_layout,
+            texture_write_pipeline,
+            texture_write_bg,
         }
     }
 
-    pub fn reset(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
+    pub fn reset(&mut self, device: &Device, queue: &Queue) {
         *self = Self::new(device, queue);
     }
 
@@ -120,74 +209,36 @@ impl LampLights {
             }
         }
         self.pending_changes.push(LightChunkUpdate {
-            x: chunk.0,
-            y: chunk.1,
+            x: chunk.0 as u32,
+            y: chunk.1 as u32,
             lights: encoded_lights,
+            lights2: extra_lights,
+            _pad: (0, 0),
         });
-        if extra_lights[0].0 != 0 {
-            self.pending_changes2.push(LightChunkUpdate {
-                x: chunk.0,
-                y: chunk.1,
-                lights: extra_lights,
-            });
-        }
     }
 
-    pub fn apply_changes(&mut self, queue: &wgpu::Queue) {
-        for change in self.pending_changes.drain(..) {
-            // SAFETY: repr(transparent)
-            let data: [u32; 4] = unsafe { std::mem::transmute(change.lights) };
-            queue.write_texture(
-                wgpu::ImageCopyTexture {
-                    texture: &self.lightdata.texture,
-                    mip_level: 0,
-                    origin: wgpu::Origin3d {
-                        x: change.x as u32,
-                        y: change.y as u32,
-                        z: 0,
-                    },
-                    aspect: wgpu::TextureAspect::All,
-                },
-                bytemuck::cast_slice(&data),
-                wgpu::ImageDataLayout {
-                    offset: 0,
-                    bytes_per_row: Some(4 * 4),
-                    rows_per_image: Some(1),
-                },
-                wgpu::Extent3d {
-                    width: 1,
-                    height: 1,
-                    depth_or_array_layers: 1,
-                },
-            );
+    pub fn apply_changes(&mut self, queue: &Queue, device: &Device, encoder: &mut CommandEncoder) {
+        if self.pending_changes.is_empty() {
+            return;
         }
 
-        for change in self.pending_changes2.drain(..) {
-            // SAFETY: repr(transparent)
-            let data: [u32; 4] = unsafe { std::mem::transmute(change.lights) };
-            queue.write_texture(
-                wgpu::ImageCopyTexture {
-                    texture: &self.lightdata2.texture,
-                    mip_level: 0,
-                    origin: wgpu::Origin3d {
-                        x: change.x as u32,
-                        y: change.y as u32,
-                        z: 0,
-                    },
-                    aspect: wgpu::TextureAspect::All,
-                },
-                bytemuck::cast_slice(&data),
-                wgpu::ImageDataLayout {
-                    offset: 0,
-                    bytes_per_row: Some(4 * 4),
-                    rows_per_image: Some(1),
-                },
-                wgpu::Extent3d {
-                    width: 1,
-                    height: 1,
-                    depth_or_array_layers: 1,
-                },
-            );
-        }
+        self.changes_buffer
+            .write_qd(queue, device, bytemuck::cast_slice(&self.pending_changes));
+
+        let Some(buffer_bg) = self.changes_buffer.bindgroup(device, &self.buffer_layout) else {
+            self.pending_changes.clear();
+            return;
+        };
+
+        let mut compute = encoder.begin_compute_pass(&ComputePassDescriptor {
+            label: Some("lamp lights update"),
+            timestamp_writes: None,
+        });
+
+        compute.set_pipeline(&self.texture_write_pipeline);
+        compute.set_bind_group(0, &self.texture_write_bg, &[]);
+        compute.set_bind_group(1, &buffer_bg, &[]);
+        compute.dispatch_workgroups((self.pending_changes.len() as u32).div_ceil(64), 1, 1);
+        self.pending_changes.clear();
     }
 }

engine/src/pbuffer.rs

@@ -1,11 +1,12 @@
-use crate::GfxContext;
 use std::sync::Arc;
+
 use wgpu::{
-    BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor,
-    BindGroupLayoutEntry, BindingResource, BindingType, BufferBinding, BufferDescriptor,
-    BufferSize, BufferSlice, BufferUsages, Device, Queue,
+    BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindingResource, BufferBinding,
+    BufferDescriptor, BufferSize, BufferSlice, BufferUsages, Device, Queue,
 };
 
+use crate::GfxContext;
+
 /// Short for Persistent Buffer, keeps memory around to reuse it
 #[derive(Clone)]
 pub struct PBuffer {
@@ -46,12 +47,12 @@ impl PBuffer {
         );
     }
 
-    pub fn bindgroup(&self, gfx: &GfxContext, layout: &BindGroupLayout) -> Option<wgpu::BindGroup> {
+    pub fn bindgroup(&self, device: &Device, layout: &BindGroupLayout) -> Option<wgpu::BindGroup> {
         if self.len == 0 {
             return None;
         }
         let buffer = self.inner.as_ref()?;
-        Some(gfx.device.create_bind_group(&BindGroupDescriptor {
+        Some(device.create_bind_group(&BindGroupDescriptor {
             label: Some("pbuffer bg"),
             layout,
             entries: &[BindGroupEntry {
@@ -65,27 +66,6 @@ impl PBuffer {
         }))
     }
 
-    pub fn bindgroup_layout(
-        gfx: &GfxContext,
-        visibility: wgpu::ShaderStages,
-        ty: wgpu::BufferBindingType,
-    ) -> BindGroupLayout {
-        gfx.device
-            .create_bind_group_layout(&BindGroupLayoutDescriptor {
-                label: Some("pbuffer bglayout"),
-                entries: &[BindGroupLayoutEntry {
-                    binding: 0,
-                    visibility,
-                    ty: BindingType::Buffer {
-                        has_dynamic_offset: false,
-                        min_binding_size: None,
-                        ty,
-                    },
-                    count: None,
-                }],
-            })
-    }
-
     pub fn slice(&self) -> Option<BufferSlice> {
         if self.len == 0 {
             return None;

engine/src/shader.rs

@@ -36,6 +36,7 @@ fn mk_module(data: String, device: &Device) -> ShaderModule {
 }
 
 /// if type isn't provided it will be detected by looking at extension
+/// name shouldn't include "assets/shaders" or ".wgsl" extension. It will be added automatically
 pub fn compile_shader(
     device: &Device,
     name: &str,

simulation/src/map/terrain.rs

@@ -12,7 +12,7 @@ pub type TerrainChunkID = common::ChunkID<5>;
 
 pub const TERRAIN_CHUNK_RESOLUTION: usize = 32;
 
-pub(super) const CELL_SIZE: f32 = TerrainChunkID::SIZE_F32 / TERRAIN_CHUNK_RESOLUTION as f32;
+pub(super) const CELL_SIZE: f32 = TerrainChunkID::SIZE_F32 / TERRAIN_CHUNK_RESOLUTION as f32; // size is 512m
 
 const TREE_GRID_SIZE: usize = 256;