Implemented image clustering

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,

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<f32> {
+        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();

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<ClusterBuffers>,
 }
 
 #[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<u32>,
+    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::<LightCount>() 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::<LightCount>() 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<GpuLight> = self.lights.iter().map(|l| l.to_gpu()).collect();
-        queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&data));
+        let gpu_lights: Vec<GpuLight> = 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,
+        }
     }
-}
+
+    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<f32>,
+        _projection_matrix: Matrix4<f32>,
+        _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,
+        }
+    }
+}

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)]

solar_engine/src/shader.wgsl

@@ -4,6 +4,13 @@ struct VertexInput {
     @location(2) normal: vec3<f32>,
 };
 
+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<f32>,
     @location(6) model_row1: vec4<f32>,
@@ -30,6 +37,7 @@ var<uniform> light_count: LightCount;
 
 struct Globals {
     view_proj: mat4x4<f32>,
+    resolution: vec2<f32>,
 }
 
 @group(0) @binding(0)
@@ -49,6 +57,11 @@ struct GpuLight {
 @group(1) @binding(0)
 var<storage, read> all_lights: array<GpuLight>;
 
+@group(2) @binding(0)
+var<storage, read> cluster_light_indices: array<u32>;
+@group(2) @binding(1)
+var<storage, read> cluster_offsets: array<vec2<u32>>;
+
 @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<f32>, view_pos_z: f32, screen_size: vec2<f32>) -> 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<f32>) -> bool {
+    return any(vec3<bool>(v != v));
+}
+
 @fragment
 fn fs_main(input: VSOutput) -> @location(0) vec4<f32> {
     var lighting: vec3<f32> = vec3<f32>(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<f32>(input.frag_color * 2.0, 1.0);
+        return vec4<f32>(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<f32> = vec3<f32>(0.0);
 
         let light_dir = normalize(light.position - input.world_pos);

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<Shape, Geometry>,
-    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::<Globals>() 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<GpuLight> = 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];