codecraft 0.2.0

//! The cluster grid: which lights reach which part of the frustum (forward+).
//! See `clustered.wgsl` for the culling geometry.

use glam::Mat4;

use crate::sceneobjects::lights::{GpuLight, MAX_LIGHTS};

/// Cells across, down, and into the frame; 16x9 keeps a cell roughly square, 24 slices is the usual.
pub const GRID: [u32; 3] = [16, 9, 24];

/// How many cells there are.
pub const CELLS: u32 = GRID[0] * GRID[1] * GRID[2];

/// How many lights one cell can hold; a cell that wants more drops the rest.
pub const MAX_LIGHTS_PER_CELL: u32 = 64;

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct ClusterUniform {
    inv_proj: [f32; 16],
    /// x, y: the frame in pixels. z, w: near and far.
    screen: [f32; 4],
    grid: [u32; 4],
    /// x, y: scale and bias taking `log(view depth)` to a slice.
    slice: [f32; 4],
}

/// Scale and bias mapping a view-space depth to a slice: `slice = log(depth) * scale + bias`.
pub fn slice_scale_bias(near: f32, far: f32) -> (f32, f32) {
    let slices = GRID[2] as f32;
    let scale = slices / (far / near).ln();
    (scale, -near.ln() * scale)
}

/// The grid, the lights in it, and the pass that fills it.
pub struct Clusters {
    pipeline: wgpu::ComputePipeline,
    uniform: wgpu::Buffer,
    view: wgpu::Buffer,
    pub lights: wgpu::Buffer,
    pub counts: wgpu::Buffer,
    pub indices: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    count: u32,
}

impl Clusters {
    pub fn new(device: &wgpu::Device) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("cluster culling shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("clustered.wgsl").into()),
        });

        let uniform = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster uniform"),
            size: std::mem::size_of::<ClusterUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let view = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster view matrix"),
            size: std::mem::size_of::<[f32; 16]>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let lights = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("local lights"),
            size: (MAX_LIGHTS * std::mem::size_of::<GpuLight>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let counts = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster light counts"),
            size: (CELLS as usize * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        let indices = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster light indices"),
            size: ((CELLS * MAX_LIGHTS_PER_CELL) as usize * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });

        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("cluster bind group layout"),
            entries: &[
                uniform_entry(0),
                storage_entry(1, true),
                storage_entry(2, false),
                storage_entry(3, false),
                uniform_entry(4),
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("cluster bind group"),
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: lights.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: counts.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: indices.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: view.as_entire_binding(),
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("cluster pipeline layout"),
            bind_group_layouts: &[Some(&layout)],
            immediate_size: 0,
        });

        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("cluster culling pipeline"),
            layout: Some(&pipeline_layout),
            module: &shader,
            entry_point: Some("cull_cs"),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            cache: None,
        });

        Self {
            pipeline,
            uniform,
            view,
            lights,
            counts,
            indices,
            bind_group,
            count: 0,
        }
    }

    /// Send this frame's lights and the grid they will be sorted into.
    #[allow(clippy::too_many_arguments)]
    pub fn prepare(
        &mut self,
        queue: &wgpu::Queue,
        lights: &[GpuLight],
        view: Mat4,
        proj: Mat4,
        width: u32,
        height: u32,
        near: f32,
        far: f32,
    ) {
        self.count = lights.len() as u32;
        if !lights.is_empty() {
            queue.write_buffer(&self.lights, 0, bytemuck::cast_slice(lights));
        }
        queue.write_buffer(&self.view, 0, bytemuck::cast_slice(&view.to_cols_array()));

        let (scale, bias) = slice_scale_bias(near, far);
        let uniform = ClusterUniform {
            inv_proj: proj.inverse().to_cols_array(),
            screen: [width.max(1) as f32, height.max(1) as f32, near, far],
            grid: [GRID[0], GRID[1], GRID[2], self.count],
            slice: [scale, bias, 0.0, 0.0],
        };
        queue.write_buffer(&self.uniform, 0, bytemuck::bytes_of(&uniform));
    }

    /// Sort this frame's lights into the grid.
    pub fn cull(&self, encoder: &mut wgpu::CommandEncoder) {
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("cluster culling pass"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &self.bind_group, &[]);
        // One invocation per cell, in groups of sixty-four.
        pass.dispatch_workgroups(CELLS.div_ceil(64), 1, 1);
    }
}

fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Uniform,
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Storage { read_only },
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_grid_is_a_few_thousand_cells() {
        assert_eq!(CELLS, 16 * 9 * 24);
        assert_eq!(CELLS, 3456);
    }

    #[test]
    fn a_depth_maps_to_the_slice_it_belongs_in() {
        let (near, far) = (0.1, 100.0);
        let (scale, bias) = slice_scale_bias(near, far);
        let slice = |depth: f32| (depth.ln() * scale + bias).floor() as i32;

        // The far plane lands one past the last slice (the shader clamps it), and rounding can leave it a hair short.
        assert_eq!(slice(near), 0);
        assert!(
            (GRID[2] as i32 - 1..=GRID[2] as i32).contains(&slice(far)),
            "the far plane should land at the back of the grid, got {}",
            slice(far),
        );
        assert_eq!(slice(far).min(GRID[2] as i32 - 1), GRID[2] as i32 - 1);

        let boundary = |k: u32| near * (far / near).powf(k as f32 / GRID[2] as f32);
        assert_eq!(slice(boundary(5) + 1.0e-4), 5);
        assert!(
            boundary(20) - boundary(19) > boundary(2) - boundary(1),
            "a far slice covers more depth than a near one",
        );
    }
}