From 4371bb274946f257272b1caed21843e3e481b449 Mon Sep 17 00:00:00 2001 From: Verox001 Date: Wed, 7 May 2025 19:45:35 +0200 Subject: [PATCH] Implemented image clustering --- simulator/src/main.rs | 2 +- solar_engine/src/camera.rs | 4 + solar_engine/src/light.rs | 212 ++++++++++++++++++++++++++++++++--- solar_engine/src/render.rs | 4 +- solar_engine/src/shader.wgsl | 49 +++++++- solar_engine/src/state.rs | 107 ++++++++++++------ 6 files changed, 322 insertions(+), 56 deletions(-) diff --git a/simulator/src/main.rs b/simulator/src/main.rs index 10a5719..06f14bf 100644 --- a/simulator/src/main.rs +++ b/simulator/src/main.rs @@ -57,7 +57,7 @@ pub async fn run() { direction: Vector3::new(0.0, 0.0, 0.0), color: Vector3::from([1.0, 1.0, 0.8]), intensity: 5.0, - range: 1.0, + range: 100000000.0, inner_cutoff: 0.0, light_type: LightType::Point, outer_cutoff: 0.0, diff --git a/solar_engine/src/camera.rs b/solar_engine/src/camera.rs index 823a300..4ff3e5a 100644 --- a/solar_engine/src/camera.rs +++ b/solar_engine/src/camera.rs @@ -29,6 +29,10 @@ impl Camera { let proj = perspective(self.fov_y, self.aspect, self.znear, self.zfar); proj * view } + + pub fn build_view_matrix(&self) -> Matrix4 { + Matrix4::look_at_rh(self.eye, self.target, self.up) + } pub fn rotate_yaw_pitch(&mut self, yaw: f32, pitch: f32) { let dir = (self.target - self.eye).normalize(); diff --git a/solar_engine/src/light.rs b/solar_engine/src/light.rs index 16762e7..3e56e05 100644 --- a/solar_engine/src/light.rs +++ b/solar_engine/src/light.rs @@ -1,5 +1,6 @@ use bytemuck::{Pod, Zeroable}; -use cgmath::Vector3; +use cgmath::{EuclideanSpace, Matrix4, Point3, Transform, Vector3}; +use wgpu::util::DeviceExt; #[repr(u32)] #[derive(Clone, Copy, Debug)] @@ -9,6 +10,12 @@ pub enum LightType { Spot = 2 } +pub const CLUSTER_COUNT_X: usize = 16; +pub const CLUSTER_COUNT_Y: usize = 9; +pub const CLUSTER_COUNT_Z: usize = 24; +pub const MAX_LIGHTS_PER_CLUSTER: usize = 32; +pub const TOTAL_CLUSTERS: usize = CLUSTER_COUNT_X * CLUSTER_COUNT_Y * CLUSTER_COUNT_Z; + #[repr(C)] #[derive(Clone, Copy, Pod, Zeroable, Debug)] pub struct GpuLight { @@ -53,7 +60,8 @@ pub struct LightManager { pub buffer: wgpu::Buffer, pub bind_group: wgpu::BindGroup, pub count_buffer: wgpu::Buffer, - pub layout: wgpu::BindGroupLayout + pub layout: wgpu::BindGroupLayout, + pub cluster_buffers: Option, } #[repr(C)] @@ -62,6 +70,19 @@ pub struct LightCount { pub count: u32, } +pub struct ClusterBuffers { + pub light_indices: wgpu::Buffer, + pub offsets: wgpu::Buffer, + pub bind_group: wgpu::BindGroup, + pub layout: wgpu::BindGroupLayout, +} + +#[derive(Debug, Clone)] +pub struct ClusterAssignment { + pub cluster_light_indices: Vec, + pub cluster_offsets: Vec<(u32, u32)>, +} + impl LightManager { pub fn new(device: &wgpu::Device, max_lights: usize) -> Self { let buffer = device.create_buffer(&wgpu::BufferDescriptor { @@ -71,10 +92,16 @@ impl LightManager { mapped_at_creation: false, }); + let count_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Light Count Buffer"), + size: size_of::() as wgpu::BufferAddress, + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("Light Bind Group Layout"), entries: &[ - // Binding 0: Storage buffer for lights wgpu::BindGroupLayoutEntry { binding: 0, visibility: wgpu::ShaderStages::VERTEX_FRAGMENT, @@ -85,7 +112,6 @@ impl LightManager { }, count: None, }, - // Binding 1: Uniform buffer for light count wgpu::BindGroupLayoutEntry { binding: 1, visibility: wgpu::ShaderStages::FRAGMENT, @@ -99,13 +125,6 @@ impl LightManager { ], }); - let count_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Light Count Buffer"), - size: size_of::() as wgpu::BufferAddress, - usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, - mapped_at_creation: false, - }); - let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { label: Some("Light Bind Group"), layout: &layout, @@ -124,9 +143,10 @@ impl LightManager { Self { lights: Vec::new(), buffer, + count_buffer, bind_group, layout, - count_buffer, + cluster_buffers: None, } } @@ -134,9 +154,13 @@ impl LightManager { self.lights.push(light); } + pub fn clear(&mut self) { + self.lights.clear(); + } + pub fn update_gpu(&self, queue: &wgpu::Queue) { - let data: Vec = self.lights.iter().map(|l| l.to_gpu()).collect(); - queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&data)); + let gpu_lights: Vec = self.lights.iter().map(|l| l.to_gpu()).collect(); + queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&gpu_lights)); let count = LightCount { count: self.lights.len() as u32, @@ -151,8 +175,160 @@ impl LightManager { pub fn layout(&self) -> &wgpu::BindGroupLayout { &self.layout } - - pub fn clear(&mut self) { - self.lights.clear(); + + pub fn create_cluster_buffers( + &self, + device: &wgpu::Device, + assignment: &ClusterAssignment, + ) -> ClusterBuffers { + let cluster_light_indices = if assignment.cluster_light_indices.is_empty() { + vec![0u32] + } else { + assignment.cluster_light_indices.clone() + }; + + let offset_pairs: Vec<[u32; 2]> = if assignment.cluster_offsets.is_empty() { + vec![[0, 0]] + } else { + assignment.cluster_offsets + .iter() + .map(|&(o, c)| [o, c]) + .collect() + }; + + let light_index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Cluster Light Indices"), + contents: bytemuck::cast_slice(&cluster_light_indices), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + }); + + let offset_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Cluster Offsets"), + contents: bytemuck::cast_slice(&offset_pairs), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + }); + + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("Cluster Bind Group Layout"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("Cluster Bind Group"), + layout: &layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: light_index_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: offset_buffer.as_entire_binding(), + }, + ], + }); + + ClusterBuffers { + light_indices: light_index_buffer, + offsets: offset_buffer, + layout, + bind_group, + } } -} \ No newline at end of file + + pub fn update_cluster_buffers( + &mut self, + device: &wgpu::Device, + queue: &wgpu::Queue, + assignment: &ClusterAssignment, + ) { + if let Some(buffers) = &self.cluster_buffers { + queue.write_buffer(&buffers.light_indices, 0, bytemuck::cast_slice(&assignment.cluster_light_indices)); + + let offset_pairs: Vec<[u32; 2]> = assignment + .cluster_offsets + .iter() + .map(|&(o, c)| [o, c]) + .collect(); + + queue.write_buffer(&buffers.offsets, 0, bytemuck::cast_slice(&offset_pairs)); + } else { + let buffers = self.create_cluster_buffers(device, assignment); + self.cluster_buffers = Some(buffers); + } + } + + pub fn compute_cluster_assignments( + &self, + view_matrix: Matrix4, + _projection_matrix: Matrix4, + _screen_width: f32, + _screen_height: f32, + ) -> ClusterAssignment { + let mut cluster_light_lists = vec![Vec::new(); TOTAL_CLUSTERS]; + + let log_near = 0.1f32.log2(); + let log_far = 1000.0f32.log2(); + let log_range = log_far - log_near; + + for (i, light) in self.lights.iter().enumerate() { + let pos_view = view_matrix.transform_point(Point3::from_vec(light.position)); + let radius = light.range; + + let z_bounds = [ + (-pos_view.z - radius).max(0.1).log2(), + (-pos_view.z + radius).max(0.1).log2(), + ]; + let z_start = ((z_bounds[0].min(z_bounds[1]) - log_near) / log_range * CLUSTER_COUNT_Z as f32).floor() as usize; + let z_end = ((z_bounds[0].max(z_bounds[1]) - log_near) / log_range * CLUSTER_COUNT_Z as f32).ceil() as usize; + + for z in z_start.min(CLUSTER_COUNT_Z)..z_end.min(CLUSTER_COUNT_Z) { + for y in 0..CLUSTER_COUNT_Y { + for x in 0..CLUSTER_COUNT_X { + let cluster = x + y * CLUSTER_COUNT_X + z * CLUSTER_COUNT_X * CLUSTER_COUNT_Y; + + if cluster_light_lists[cluster].len() < MAX_LIGHTS_PER_CLUSTER { + cluster_light_lists[cluster].push(i as u32); + } + } + } + } + } + + let mut cluster_light_indices = Vec::new(); + let mut cluster_offsets = Vec::with_capacity(TOTAL_CLUSTERS); + + for lights in cluster_light_lists { + let offset = cluster_light_indices.len() as u32; + let count = lights.len() as u32; + cluster_light_indices.extend(lights); + cluster_offsets.push((offset, count)); + } + + ClusterAssignment { + cluster_light_indices, + cluster_offsets, + } + } +} diff --git a/solar_engine/src/render.rs b/solar_engine/src/render.rs index da1f440..9cadbe8 100644 --- a/solar_engine/src/render.rs +++ b/solar_engine/src/render.rs @@ -1,7 +1,9 @@ #[repr(C)] -#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)] +#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)] pub struct Globals { pub view_proj: [[f32; 4]; 4], + pub resolution: [f32; 2], + pub _padding: [f32; 2], } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] diff --git a/solar_engine/src/shader.wgsl b/solar_engine/src/shader.wgsl index 0da34f3..46930e0 100644 --- a/solar_engine/src/shader.wgsl +++ b/solar_engine/src/shader.wgsl @@ -4,6 +4,13 @@ struct VertexInput { @location(2) normal: vec3, }; +const CLUSTER_COUNT_X: u32 = 16u; +const CLUSTER_COUNT_Y: u32 = 9u; +const CLUSTER_COUNT_Z: u32 = 24u; + +const NEAR_PLANE: f32 = 0.1; +const FAR_PLANE: f32 = 1000.0; + struct InstanceInput { @location(5) model_row0: vec4, @location(6) model_row1: vec4, @@ -30,6 +37,7 @@ var light_count: LightCount; struct Globals { view_proj: mat4x4, + resolution: vec2, } @group(0) @binding(0) @@ -49,6 +57,11 @@ struct GpuLight { @group(1) @binding(0) var all_lights: array; +@group(2) @binding(0) +var cluster_light_indices: array; +@group(2) @binding(1) +var cluster_offsets: array>; + @vertex fn vs_main(vertex: VertexInput, instance: InstanceInput) -> VSOutput { var out: VSOutput; @@ -76,17 +89,47 @@ fn vs_main(vertex: VertexInput, instance: InstanceInput) -> VSOutput { return out; } +fn compute_cluster_id(frag_coord: vec4, view_pos_z: f32, screen_size: vec2) -> u32 { + let x_frac = frag_coord.x / screen_size.x; + let y_frac = frag_coord.y / screen_size.y; + + let x = clamp(u32(x_frac * f32(CLUSTER_COUNT_X)), 0u, CLUSTER_COUNT_X - 1u); + let y = clamp(u32(y_frac * f32(CLUSTER_COUNT_Y)), 0u, CLUSTER_COUNT_Y - 1u); + + // Z: logarithmic depth + let depth = -view_pos_z; // view-space z is negative + let depth_clamped = clamp(depth, NEAR_PLANE, FAR_PLANE); + let log_depth = log2(depth_clamped); + let z = clamp(u32((log_depth / log2(FAR_PLANE / NEAR_PLANE)) * f32(CLUSTER_COUNT_Z)), 0u, CLUSTER_COUNT_Z - 1u); + + return x + y * CLUSTER_COUNT_X + z * CLUSTER_COUNT_X * CLUSTER_COUNT_Y; +} + +fn is_nan_f32(x: f32) -> bool { + return x != x; +} + +fn is_nan_vec3(v: vec3) -> bool { + return any(vec3(v != v)); +} + @fragment fn fs_main(input: VSOutput) -> @location(0) vec4 { var lighting: vec3 = vec3(0.0); let always_lit = (input.flags & 0x1u) != 0u; + let cluster_id = compute_cluster_id(input.position, input.world_pos.z, globals.resolution); + let offset_info = cluster_offsets[cluster_id]; + let offset = offset_info.x; + let count = offset_info.y; + if (always_lit) { - return vec4(input.frag_color * 2.0, 1.0); + return vec4(input.frag_color, 2.0); } - for (var i = 0u; i < light_count.count; i = i + 1u) { - let light = all_lights[i]; + for (var i = 0u; i < count; i = i + 1u) { + let light_index = cluster_light_indices[offset + i]; + let light = all_lights[light_index]; var light_contrib: vec3 = vec3(0.0); let light_dir = normalize(light.position - input.world_pos); diff --git a/solar_engine/src/state.rs b/solar_engine/src/state.rs index 662f6e1..c30a257 100644 --- a/solar_engine/src/state.rs +++ b/solar_engine/src/state.rs @@ -5,11 +5,11 @@ use cgmath::{perspective, Deg, Matrix4, Point3, Vector3}; use log::info; use pollster::FutureExt; use wgpu::util::DeviceExt; -use wgpu::{Adapter, Device, Instance, PresentMode, Queue, Surface, SurfaceCapabilities, SurfaceConfiguration}; +use wgpu::{Adapter, BindGroup, BindGroupLayout, Device, Instance, PresentMode, Queue, Surface, SurfaceCapabilities, SurfaceConfiguration}; use winit::dpi::PhysicalSize; use winit::window::{Window}; use crate::camera::Camera; -use crate::light::{GpuLight, LightManager}; +use crate::light::{ClusterBuffers, GpuLight, LightManager}; use crate::render::{create_circle_vertices, create_sphere_vertices, Geometry, Globals, InstanceRaw, RenderInstance, SampleCount, Shape, Vertex}; pub struct State<'a> { @@ -28,7 +28,7 @@ pub struct State<'a> { instance_buffer: wgpu::Buffer, geometries: HashMap, - global_bind_group: wgpu::BindGroup, + global_bind_group: BindGroup, global_buffer: wgpu::Buffer, camera: Camera, @@ -58,37 +58,15 @@ impl<'a> State<'a> { let globals = Globals { view_proj: view_proj.into(), + resolution: [config.width as f32, config.height as f32], + _padding: [0.0, 0.0], }; - let global_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("Global Uniform Buffer"), - contents: bytemuck::cast_slice(&[globals]), - usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, - }); - let global_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("Global Bind Group Layout"), - entries: &[wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: wgpu::ShaderStages::VERTEX, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Uniform, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }], - }); - let global_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { - layout: &global_bind_group_layout, - entries: &[wgpu::BindGroupEntry { - binding: 0, - resource: global_buffer.as_entire_binding(), - }], - label: Some("Global Bind Group"), - }); + let (global_buffer, global_bind_group_layout, global_bind_group) = Self::create_global_buffer(&device); + queue.write_buffer(&global_buffer, 0, bytemuck::cast_slice(&[globals])); - let light_manager = LightManager::new(&device, 10); + let mut light_manager = LightManager::new(&device, 10); - let render_pipeline = Self::create_render_pipeline(&device, &config, sample_count.0, &global_bind_group_layout, &light_manager); + let render_pipeline = Self::create_render_pipeline(&queue, &device, &config, sample_count.0, &global_bind_group_layout, &mut light_manager, &camera); let geometries = Self::create_geometries(&device); let instances = vec![]; @@ -120,6 +98,40 @@ impl<'a> State<'a> { light_manager } } + + fn create_global_buffer(device: &wgpu::Device) -> (wgpu::Buffer, BindGroupLayout, BindGroup) { + let global_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Global Buffer"), + size: size_of::() as u64, + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + let global_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("Global Bind Group Layout"), + entries: &[wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::VERTEX_FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }], + }); + + let global_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + layout: &global_bind_group_layout, + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: global_buffer.as_entire_binding(), + }], + label: Some("Global Bind Group"), + }); + + (global_buffer, global_bind_group_layout, global_bind_group) + } fn update_lights(&mut self) { let light_data: Vec = self.light_manager.lights.iter().map(|l| l.to_gpu()).collect(); @@ -206,16 +218,29 @@ impl<'a> State<'a> { }) } - fn create_render_pipeline(device: &Device, config: &SurfaceConfiguration, sample_count: u32, global_bind_group_layout: &wgpu::BindGroupLayout, light_manager: &LightManager) -> wgpu::RenderPipeline { + fn create_render_pipeline(queue: &Queue, device: &Device, config: &SurfaceConfiguration, sample_count: u32, global_bind_group_layout: &BindGroupLayout, light_manager: &mut LightManager, camera: &Camera) -> wgpu::RenderPipeline { let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("Shader"), source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()), }); + let cluster_assignment = light_manager.compute_cluster_assignments( + camera.build_view_matrix(), + camera.build_view_projection_matrix(), + config.width as f32, + config.height as f32, + ); + + let cluster_buffers = light_manager.create_cluster_buffers(&device, &cluster_assignment); + let render_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("Render Pipeline Layout"), - bind_group_layouts: &[&global_bind_group_layout, &light_manager.layout], + bind_group_layouts: &[ + &global_bind_group_layout, + &light_manager.layout, + &cluster_buffers.layout, + ], push_constant_ranges: &[], }); @@ -326,6 +351,8 @@ impl<'a> State<'a> { let new_globals = Globals { view_proj: view_proj.into(), + resolution: [self.config.width as f32, self.config.height as f32], + _padding: [0.0, 0.0], }; self.queue.write_buffer(&self.global_buffer, 0, bytemuck::cast_slice(&[new_globals])); self.depth_texture = Self::create_depth_texture(&self.device, self.config.width, self.config.height, self.sample_count.get()); @@ -355,9 +382,19 @@ impl<'a> State<'a> { let view_proj = self.camera.build_view_projection_matrix(); let globals = Globals { view_proj: view_proj.into(), + resolution: [self.config.width as f32, self.config.height as f32], + _padding: [0.0, 0.0], }; self.queue.write_buffer(&self.global_buffer, 0, bytemuck::cast_slice(&[globals])); + let assignment = self.light_manager.compute_cluster_assignments( + self.camera.build_view_matrix(), + self.camera.build_view_projection_matrix(), + self.config.width as f32, + self.config.height as f32, + ); + self.light_manager.update_cluster_buffers(&self.device, &self.queue, &assignment); + let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("Render Encoder"), }); @@ -397,6 +434,10 @@ impl<'a> State<'a> { self.light_manager.update_gpu(&self.queue); render_pass.set_bind_group(1, &self.light_manager.bind_group, &[]); + if let Some(clusters) = &self.light_manager.cluster_buffers { + render_pass.set_bind_group(2, &clusters.bind_group, &[]); + } + for shape in self.geometries.keys().copied() { let geometry = &self.geometries[&shape];