From 5559a40d956a0c879d607cce3af4e15251ccfbb6 Mon Sep 17 00:00:00 2001 From: Paris DOUADY Date: Thu, 8 Feb 2024 17:33:22 +0100 Subject: [PATCH] rewrite lamplights to use compute shader --- assets/shaders/compute/texture_write.wgsl | 25 +++ engine/src/gfx.rs | 5 +- engine/src/lamplights.rs | 199 ++++++++++++++-------- engine/src/pbuffer.rs | 34 +--- engine/src/shader.rs | 1 + simulation/src/map/terrain.rs | 2 +- 6 files changed, 162 insertions(+), 104 deletions(-) create mode 100644 assets/shaders/compute/texture_write.wgsl diff --git a/assets/shaders/compute/texture_write.wgsl b/assets/shaders/compute/texture_write.wgsl new file mode 100644 index 00000000..f132f2ca --- /dev/null +++ b/assets/shaders/compute/texture_write.wgsl @@ -0,0 +1,25 @@ +// 16-aligned! +struct LightUpdate { + data: vec4, + data2: vec4, + position: vec2, +} + +@group(0) @binding(0) var light_tex: texture_storage_2d; +@group(0) @binding(1) var light_tex_2: texture_storage_2d; + +@group(1) @binding(0) var light_changes: array; + +@compute @workgroup_size(64,1,1) +fn main( + @builtin(global_invocation_id) id: vec3, +) { + 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); +} \ No newline at end of file diff --git a/engine/src/gfx.rs b/engine/src/gfx.rs index c80f6d04..cf967221 100644 --- a/engine/src/gfx.rs +++ b/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 { diff --git a/engine/src/lamplights.rs b/engine/src/lamplights.rs index 3e675f7c..0f77acb4 100644 --- a/engine/src/lamplights.rs +++ b/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, - pending_changes2: Vec, + 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(); } } diff --git a/engine/src/pbuffer.rs b/engine/src/pbuffer.rs index 2ba05253..8bcc2b1e 100644 --- a/engine/src/pbuffer.rs +++ b/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 { + pub fn bindgroup(&self, device: &Device, layout: &BindGroupLayout) -> Option { 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 { if self.len == 0 { return None; diff --git a/engine/src/shader.rs b/engine/src/shader.rs index a47c3565..b0c39521 100644 --- a/engine/src/shader.rs +++ b/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, diff --git a/simulation/src/map/terrain.rs b/simulation/src/map/terrain.rs index 1e22da32..139b6948 100644 --- a/simulation/src/map/terrain.rs +++ b/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;