diff --git a/src/main.rs b/src/main.rs index 3099845..5598674 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,65 +1,40 @@ #![feature(generic_const_exprs)] #![feature(float_algebraic)] -use core::sync; -use std::cell::RefCell; use std::collections::HashMap; -use std::fs::File; -use std::hash::Hash; -use std::rc::Rc; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::mpsc::sync_channel; use bytemuck::Pod; use bytemuck::Zeroable; use bytemuck::cast_slice; use crevice::std140::AsStd140; -use crevice::std430::AsStd430; use egui::Color32; -use egui::Label; -use egui::emath::fast_midpoint; use egui::mutex::Mutex; use egui_plot::BarChart; use glam::Mat4; use glam::Vec3; use glam::Vec4; use itertools::Itertools; -use rand::random; use rayon::iter::IndexedParallelIterator; use rayon::iter::IntoParallelRefIterator; use rayon::iter::ParallelIterator; -use wgpu::BindGroup; -use wgpu::BindGroupEntry; -use wgpu::BindGroupLayoutDescriptor; -use wgpu::BindGroupLayoutEntry; use wgpu::Buffer; use wgpu::BufferUsages; -use wgpu::DepthBiasState; use wgpu::Device; use wgpu::Extent3d; use wgpu::Features; -use wgpu::FragmentState; use wgpu::InstanceDescriptor; use wgpu::InstanceFlags; use wgpu::MemoryBudgetThresholds; -use wgpu::NoopBackendOptions; use wgpu::Operations; -use wgpu::PrimitiveState; -use wgpu::RenderPassDepthStencilAttachment; use wgpu::RenderPipeline; -use wgpu::RenderPipelineDescriptor; -use wgpu::ShaderModuleDescriptor; -use wgpu::ShaderStages; -use wgpu::StencilState; use wgpu::Texture; use wgpu::TextureUsages; use wgpu::TextureView; -use wgpu::VertexState; use wgpu::util::BufferInitDescriptor; use wgpu::util::DeviceExt; use wgpu::util::DownloadBuffer; -use wgpu::util::StagingBelt; use winit::application::ApplicationHandler; use winit::event::DeviceEvent; use winit::event::MouseScrollDelta; @@ -69,30 +44,27 @@ use winit::event_loop::ActiveEventLoop; use winit::event_loop::ControlFlow; use winit::event_loop::EventLoop; use winit::event_loop::OwnedDisplayHandle; -use winit::platform::x11::EventLoopBuilderExtX11; use winit::window::Window; use winit::window::WindowId; use crate::camera::Camera; use crate::egui_renderer::EguiRenderer; -use crate::producers::BallGenerator; use crate::producers::ChunkedProducer; -use crate::producers::Producer; use crate::producers::SineGenerator; use crate::producers::TerrainGenerator; -use crate::voxel_cache::cache::CacheNodeRequest; -use crate::voxel_cache::cache::CacheResponse; -use crate::voxel_cache::cache::ColorBytes; -use crate::voxel_cache::cache::DestinationElement; -use crate::voxel_cache::cache::LocationPoolElement; -use crate::voxel_cache::cache::VoxelCache; -use crate::voxel_cache::sparse::NTreeNodeLocator; +use crate::sparse_tree::NTreeNodeLocator; +use crate::voxel_cache::VoxelCache; +use crate::voxel_cache::data::CacheNodeRequest; +use crate::voxel_cache::data::CacheResponse; +use crate::voxel_cache::data::ColorBytes; +use crate::voxel_cache::data::DestinationElement; +use crate::voxel_cache::data::LocationPoolElement; mod camera; mod egui_renderer; mod producers; +mod sparse_tree; mod voxel_cache; -//mod tree; // struct State diff --git a/src/producers.rs b/src/producers.rs index aec64f2..96079b4 100644 --- a/src/producers.rs +++ b/src/producers.rs @@ -2,12 +2,11 @@ use std::fs::File; use std::path::Path; use glam::Vec3; -use indicatif::ProgressIterator; use itertools::Itertools; -use crate::voxel_cache::gpu::ExplicitNTreeNode; -use crate::voxel_cache::gpu::StructurePointer; -use crate::voxel_cache::sparse::Color; +use crate::sparse_tree::Color; +use crate::voxel_cache::data::ExplicitNTreeNode; +use crate::voxel_cache::data::StructurePointer; pub struct BallGenerator { diff --git a/src/voxel_cache/sparse.rs b/src/sparse_tree.rs similarity index 99% rename from src/voxel_cache/sparse.rs rename to src/sparse_tree.rs index 69b1aea..20799ff 100644 --- a/src/voxel_cache/sparse.rs +++ b/src/sparse_tree.rs @@ -5,8 +5,8 @@ use bytemuck::Pod; use bytemuck::Zeroable; use itertools::Itertools; -use crate::voxel_cache::gpu::ExplicitNTreeNode; -use crate::voxel_cache::gpu::StructurePointer; +use crate::voxel_cache::data::ExplicitNTreeNode; +use crate::voxel_cache::data::StructurePointer; #[derive(Debug, Clone, Copy, PartialEq, Pod, Zeroable)] #[repr(C)] diff --git a/src/voxel_cache.rs b/src/voxel_cache.rs index 609fb5c..d4ef929 100644 --- a/src/voxel_cache.rs +++ b/src/voxel_cache.rs @@ -1,4 +1,1042 @@ -pub mod cache; -pub mod gpu; -pub mod pipeline; -pub mod sparse; +use bytemuck::Zeroable; +use wgpu::{BindGroup, BindGroupLayout, Buffer, BufferUsages, ComputePipeline, Device, Queue, ShaderStages, CommandEncoder}; + +use crate::voxel_cache::{data::{CacheNodeRequest, CacheResponse, ColorPoolElement, LocationPoolElement, StructurePoolElement}, request_buffer::RequestBuffer, structure_table::StructureTable, usage_buffer::UsageBuffer}; + + +pub mod request_buffer; +pub mod structure_table; +pub mod usage_buffer; + +pub mod data; + +pub struct VoxelCache +{ + size: usize, + device: Device, + queue: Queue, + structure_pool: Buffer, + color_pool: Buffer, + location_pool: Buffer, + + pub structure_table: StructureTable, + + request_buffer: RequestBuffer, + + voxel_cache_bind_group: BindGroup, + structure_table_bind_group_layout: BindGroupLayout, + voxel_cache_render_bind_group_layout: BindGroupLayout, + + // Writing requests + request_write_count_buffer: Buffer, + request_write_bindgroup_layout_user: BindGroupLayout, + request_write_pipeline: ComputePipeline, + + // User response -> caching + caching_user_bindgroup_layout: BindGroupLayout, + caching_pipeline: ComputePipeline, + + // Invalidation + invalidation_pipeline: ComputePipeline, + + pub usage_buffer: UsageBuffer, +} + + +#[derive(Zeroable, bytemuck::Pod, Clone, Copy)] +#[repr(C)] +struct CachingPipelineImmediates +{ + frame_timestamp: u32, + write_pointers: u32, // Boolean +} + +unsafe fn as_raw_bytes(slice: &[T]) -> &[u8] +{ + let size = slice.len() * size_of::(); + unsafe { std::slice::from_raw_parts(slice.as_ptr() as *const u8, size) } +} + +impl VoxelCache +where + [(); N * N * N]:, +{ + pub fn new(cache_size: usize, device: Device, queue: Queue) -> Self + { + let structure_pool = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Structure pool"), + size: (size_of::>() * cache_size) as u64, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }); + + let location_pool = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Location pool"), + size: (size_of::() * cache_size) as u64, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }); + + let color_pool = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Color pool"), + size: (size_of::>() * cache_size) as u64, + usage: BufferUsages::STORAGE, + mapped_at_creation: false, + }); + + let request_buffer = RequestBuffer::new(cache_size, device.clone()); + let usage_buffer = UsageBuffer::new(cache_size, device.clone()); + + let request_write_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("total_request_count_buffer"), + size: size_of::() as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + + // Rendering bing group layouts + let voxel_cache_render_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("voxel_cache_render_bind_group_layout"), + entries: &[ + // Structure pool + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Color pool + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Location pool + wgpu::BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Request buffer + wgpu::BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Usage buffer + wgpu::BindGroupLayoutEntry { + binding: 4, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + + // Structure table stuff + // Structure table pointers + wgpu::BindGroupLayoutEntry { + binding: 5, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + + // Structure table request buffer + wgpu::BindGroupLayoutEntry { + binding: 6, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let voxel_cache_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("voxel_cache_bind_group_layout"), + entries: &[ + // Structure pool + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Color pool + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Location pool + wgpu::BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Sorted requests + wgpu::BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // LRU List + wgpu::BindGroupLayoutEntry { + binding: 4, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Usage buffer + wgpu::BindGroupLayoutEntry { + binding: 5, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // pools request count + wgpu::BindGroupLayoutEntry { + binding: 6, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // total request count + wgpu::BindGroupLayoutEntry { + binding: 7, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let voxel_cache_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("voxel_cache_bind_group"), + layout: &voxel_cache_bind_group_layout, + entries: &[ + // Structure pool + wgpu::BindGroupEntry { + binding: 0, + resource: structure_pool.as_entire_binding(), + }, + // Color pool + wgpu::BindGroupEntry { + binding: 1, + resource: color_pool.as_entire_binding(), + }, + // Location pool + wgpu::BindGroupEntry { + binding: 2, + resource: location_pool.as_entire_binding(), + }, + // Sorted requests + wgpu::BindGroupEntry { + binding: 3, + resource: request_buffer.sort_buffer().as_entire_binding(), + }, + // LRU list + wgpu::BindGroupEntry { + binding: 4, + resource: usage_buffer.sort_buffer().as_entire_binding(), + }, + // Usage buffer + wgpu::BindGroupEntry { + binding: 5, + resource: usage_buffer.usage_buffer().as_entire_binding(), + }, + // Pools request count + wgpu::BindGroupEntry { + binding: 6, + resource: request_buffer.request_count_buffer().as_entire_binding(), + }, + // Total request count + wgpu::BindGroupEntry { + binding: 7, + resource: request_write_count_buffer.as_entire_binding(), + }, + ], + }); + + let structure_table_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("structure_table_bind_group_layout"), + entries: &[ + // Pointer table + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Sorted requests + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // Requests count + wgpu::BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let children_count = N*N*N; + let wgsl_bindings = + format!(" + struct StructurePoolElement + {{ + pointers: array + }} + + struct ColorPoolElement + {{ + colors: array + }} + + struct LocationPoolElement + {{ + structure_id: u32, + structure_locator: u32 + }} + + struct SortedRequestsElement + {{ + node: u32, + child: u32 + }} + + @group(0) @binding(0) var structure_pool: array; + @group(0) @binding(1) var color_pool: array; + @group(0) @binding(2) var location_pool: array; + @group(0) @binding(3) var sorted_requests: array; + @group(0) @binding(4) var lru_list: array; + @group(0) @binding(5) var usage_buffer: array; + @group(0) @binding(6) var pools_request_count: u32; + @group(0) @binding(7) var total_request_count: u32; + + @group(1) @binding(0) var structure_table_pointers: array; + @group(1) @binding(1) var structure_table_sorted_requests: array; + @group(1) @binding(2) var structure_table_request_count: u32; + "); + + + let write_requests_bindgroup_layout_user = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("write_requests_bindgroup_layout_user"), + entries: &[ + // User target + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let write_requests_pipeline_layout = + device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("write_requests_pipeline_layout"), + bind_group_layouts: &[ + Some(&voxel_cache_bind_group_layout), + Some(&structure_table_bind_group_layout), + Some(&write_requests_bindgroup_layout_user), + ], + immediate_size: 0, + }); + + // The user bindgroup will be created on the fly + let write_requests = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("write_requests_pipeline"), + layout: Some(&write_requests_pipeline_layout), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("write_requests_pipeline_shader_module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + {wgsl_bindings} + + struct CacheNodeRequest + {{ + // Requested ressource + structure_id: u32, + structure_locator: u32, + + // Write back info + node_index: u32, + child_index: u32, + }} + + @group(2) @binding(0) var requests: array; + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + // One shader invocation per invocation on both structure table domain + // and cache domain + var index = global_invocation_id.x; + let max_requests_count = min(arrayLength(&requests), arrayLength(&lru_list)); + total_request_count = min(max_requests_count, pools_request_count + structure_table_request_count); + if(index >= total_request_count) + {{ + return; + }} + + if(index < structure_table_request_count) + {{ + var request: CacheNodeRequest; + // Index in structure table IS structure id + request.structure_id = structure_table_sorted_requests[index].node; + request.structure_locator = 0; // Root request -> locator 0 + request.node_index = request.structure_id; + request.child_index = 0xFFFFFFFF; // child index u32::MAX -> + // Structure table request + // Write request + requests[index] = request; + return; + }} + + let pool_index = index - structure_table_request_count; + var request: CacheNodeRequest; + + // Request location comes from location pool + request.structure_id = location_pool[sorted_requests[pool_index].node].structure_id; + request.structure_locator = location_pool[sorted_requests[pool_index].node].structure_locator; + + request.node_index = sorted_requests[pool_index].node; + request.child_index = sorted_requests[pool_index].child; + + // Write request + requests[index] = request; + + }} + " + ) + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + // ~~~ Caching pipeline ~~~ + // Pipeline that takes user fulling ~some~ requests + // and writes them to the cache based on the eviction list + + // Caching pipeline binds + // - Cache data + // + // - User response : + // - response structure nodes + // - response color nodes + // - response locations + // - response destinations + let caching_user_bindgroup_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("caching_user_bindgroup_layout"), + entries: &[ + // structure nodes + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // response color nodes + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // response locations + wgpu::BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + // response destinations + wgpu::BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let caching_pipeline_layout = + device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("caching_pipelien_layout"), + bind_group_layouts: &[ + Some(&voxel_cache_bind_group_layout), + Some(&structure_table_bind_group_layout), + Some(&caching_user_bindgroup_layout), + ], + immediate_size: size_of::() as u32, // Current frame timestamp + }); + + let children_count = N * N * N; + let caching_pipeline = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("caching_pipeline"), + layout: Some(&caching_pipeline_layout), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("caching_pipeline_shader_module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + + struct DestinationElement + {{ + node: u32, + child: u32 + }} + + struct CachingPipelineImmediates + {{ + frame_timestamp: u32, + write_pointers: u32 + }} + + {wgsl_bindings} + + var parameters: CachingPipelineImmediates; + + @group(2) @binding(0) var structure_nodes: array; + @group(2) @binding(1) var color_nodes: array; + @group(2) @binding(2) var locations: array; + @group(2) @binding(3) var destinations: array; + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + // Copy with indirection + let index = global_invocation_id.x; + let total = arrayLength(&structure_nodes); + let total_cache = arrayLength(&lru_list); + if(index >= total || index >= total_cache) + {{ + return; + }} + + let overwritten_element = lru_list[index]; + if(parameters.write_pointers == 0 && usage_buffer[overwritten_element] != parameters.frame_timestamp) + {{ + // Phase 1 + // Copy into cache page + for(var i = 0; i < {children_count}; i++) + {{ + structure_pool[overwritten_element].pointers[i] = select(u32(0), u32(1)<<31, structure_nodes[index].pointers[i] != 0); + }} + color_pool[overwritten_element] = color_nodes[index]; + location_pool[overwritten_element] = locations[index]; + + // Mark dirty/correct timestamp + usage_buffer[overwritten_element] = parameters.frame_timestamp + 1; + }} + + if(parameters.write_pointers != 0 && usage_buffer[overwritten_element] != parameters.frame_timestamp) + {{ + // Phase 2 + + // Point parent to new page + let new_pointer = (1 << 31) | (1 << 30) | overwritten_element; + if(destinations[index].child == 0xFFFFFFFF) + {{ + structure_table_pointers[destinations[index].node] = new_pointer; + }}else if usage_buffer[destinations[index].node] != parameters.frame_timestamp + 1 + {{ + structure_pool[destinations[index].node].pointers[destinations[index].child] = new_pointer; + }} + }} + + }} + ") + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + // ~~~ Invalidation pipeline ~~~ + // Pipeline invalidates old nodes that points to ones + // that have been replaces base on timestamp information + let invalidation_pipeline_layout = + device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("invalidation_pipeline_layout"), + bind_group_layouts: &[ + Some(&voxel_cache_bind_group_layout), + Some(&structure_table_bind_group_layout), + ], + immediate_size: size_of::() as u32, // frame_timestamp + }); + + let invalidation_pipeline = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("invalidation_pipeline"), + layout: Some(&invalidation_pipeline_layout), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("invalidation_pipeline_shader_module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + {wgsl_bindings} + + var frame_timestamp: u32; + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + var index = global_invocation_id.x; + let total_pool = arrayLength(&structure_pool); + let total_table = arrayLength(&structure_table_pointers); + + if index < total_pool + {{ + for(var i = 0; i < {children_count}; i += 1) + {{ + let structure_pointer = structure_pool[index].pointers[i]; + let pointer = structure_pointer & 0x3FFFFFFF; + let pointer_subdiv = ((structure_pointer >> 31) & 1) != 0; + let pointed_timestamp = usage_buffer[pointer]; + if(pointed_timestamp == frame_timestamp + 1) // Future + // timestamp + // -> new page + {{ + // Invalidate + structure_pool[index].pointers[i] = select(u32(0), u32(1), pointer_subdiv) << 31; + }} + }} + return; + }} + index -= total_pool; + if index < total_table + {{ + let structure_pointer = structure_table_pointers[index]; + // Is pointer pointing to something valid + let pointer_valid = ((structure_pointer >> 30) & 1) != 0; + let pointer_subdiv = ((structure_pointer >> 31) & 1) != 0; + let pointer = structure_pointer & 0x3FFFFFFF; + let pointed_timestamp = usage_buffer[pointer]; + if(pointed_timestamp == frame_timestamp + 1 && pointer_valid && pointer_subdiv) // Future + // timestamp + // -> new page + {{ + // Invalidate + structure_table_pointers[index] &= select(u32(0), u32(1), pointer_subdiv) << 31; + }} + }} + + }} + ") + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + Self { + size: cache_size, + structure_table: StructureTable::new(device.clone(), queue.clone()), + queue, + device, + structure_pool, + color_pool, + location_pool, + request_buffer, + usage_buffer, + + voxel_cache_bind_group, + structure_table_bind_group_layout, + voxel_cache_render_bind_group_layout, + + // Write requests stage + request_write_bindgroup_layout_user: write_requests_bindgroup_layout_user, + request_write_count_buffer, + request_write_pipeline: write_requests, + + // Caching operation + caching_user_bindgroup_layout, + caching_pipeline, + + // Invalidation + invalidation_pipeline, + } + } + + pub fn bind_group_layout(&self) -> BindGroupLayout + { + self.voxel_cache_render_bind_group_layout.clone() + } + + pub fn bind_group(&self) -> BindGroup + { + self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("voxel_cache_render_bind_group"), + layout: &self.voxel_cache_render_bind_group_layout, + entries: &[ + // Structure pool + wgpu::BindGroupEntry { + binding: 0, + resource: self.structure_pool.as_entire_binding(), + }, + // Color pool + wgpu::BindGroupEntry { + binding: 1, + resource: self.color_pool.as_entire_binding(), + }, + // Location pool + wgpu::BindGroupEntry { + binding: 2, + resource: self.location_pool.as_entire_binding(), + }, + // Request buffer + wgpu::BindGroupEntry { + binding: 3, + resource: self.request_buffer.request_buffer().as_entire_binding(), + }, + // Usage buffer + wgpu::BindGroupEntry { + binding: 4, + resource: self.usage_buffer.usage_buffer().as_entire_binding(), + }, + // Structure table + // Structure table pointers + wgpu::BindGroupEntry { + binding: 5, + resource: self.structure_table.pointer_table.as_entire_binding(), + }, + // Pools request count + wgpu::BindGroupEntry { + binding: 6, + resource: self.structure_table.request_buffer.request_buffer().as_entire_binding(), + }, + ], + }) + } + + fn structure_table_bind_group(&self) -> BindGroup + { + self.device.create_bind_group(&wgpu::BindGroupDescriptor + { + label: Some("structure_table_bind_group"), + layout: &self.structure_table_bind_group_layout, + entries: + &[ + wgpu::BindGroupEntry + { + binding: 0, + resource: self.structure_table.pointer_table.as_entire_binding(), + }, + wgpu::BindGroupEntry + { + binding: 1, + resource: self.structure_table.request_buffer.sort_buffer().as_entire_binding(), + }, + wgpu::BindGroupEntry + { + binding: 2, + resource: self.structure_table.request_buffer.request_count_buffer().as_entire_binding(), + }, + ], + }) + } + + pub fn cache_post_render(&mut self, encoder: &mut CommandEncoder, request_target: Buffer) + { + // Sort requests, reset request buffers to count requests + self.usage_buffer.sort_usage(encoder); + self.request_buffer.sort_requests(encoder); + + self.structure_table + .request_buffer_mut() + .sort_requests(encoder); + + // Sorted requests are now in the group + let user_target_bindgroup = self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("write_requests_user_bindgroup"), + layout: &self.request_write_bindgroup_layout_user, + entries: &[ + // user target + wgpu::BindGroupEntry { + binding: 0, + resource: request_target.as_entire_binding(), + }, + ], + }); + + let mut write_requests_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("write_requests_pass"), + timestamp_writes: None, + }); + + write_requests_pass.set_bind_group(0, Some(&self.voxel_cache_bind_group), &[]); + write_requests_pass.set_bind_group(1, Some(&self.structure_table_bind_group()), &[]); + write_requests_pass.set_bind_group(2, Some(&user_target_bindgroup), &[]); + write_requests_pass.set_pipeline(&self.request_write_pipeline); + + // Compute necessary shader invocations + let shader_invocation_count = request_target.size() / size_of::() as u64; + let workgroup_invocation_count = shader_invocation_count.div_ceil(64); + write_requests_pass.dispatch_workgroups(workgroup_invocation_count as u32, 1, 1); + drop(write_requests_pass); + + self.request_buffer.reset_requests(encoder); + self.structure_table + .request_buffer_mut() + .reset_requests(encoder); + } + + pub fn current_timestamp(&self) -> u32 + { + self.usage_buffer.timestamp() + } + + pub fn cache_insert(&mut self, encoder: &mut CommandEncoder, insertion: &CacheResponse) + { + // Each of the entry of the cache insertion, is matched with the usage list to evict + let mut cache_insertion_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("cache_insertion_pass"), + timestamp_writes: None, + }); + + let user_caching_bindgroup = self.device.create_bind_group(&wgpu::BindGroupDescriptor + { + label: Some("user_caching_bind_group"), + layout: &self.caching_user_bindgroup_layout, + entries: & + [ + wgpu::BindGroupEntry + { + binding: 0, + resource: insertion.structure_nodes.as_entire_binding(), + }, + wgpu::BindGroupEntry + { + binding: 1, + resource: insertion.color_nodes.as_entire_binding(), + }, + wgpu::BindGroupEntry + { + binding: 2, + resource: insertion.locations.as_entire_binding(), + }, + wgpu::BindGroupEntry + { + binding: 3, + resource: insertion.parents.as_entire_binding(), + } + ], + }); + + // Copy insertions to evicted lru + // Point parents to children + let structure_table_bind_group = self.structure_table_bind_group(); + + // ~~~ Write data + Dirty flagging/timestamp update ~~~ + { + cache_insertion_pass.set_pipeline(&self.caching_pipeline); + cache_insertion_pass.set_bind_group(0, Some(&self.voxel_cache_bind_group), &[]); + cache_insertion_pass.set_bind_group(1, Some(&structure_table_bind_group), &[]); + cache_insertion_pass.set_bind_group(2, Some(&user_caching_bindgroup), &[]); + cache_insertion_pass.set_immediates( + 0, + bytemuck::bytes_of(&CachingPipelineImmediates { + frame_timestamp: self.usage_buffer.timestamp(), + write_pointers: 0, // false + }), + ); + + // Compute dispatch amounts + let shader_invocation_count = + insertion.structure_nodes.size() as usize / size_of::>(); + let workgroup_invocations = shader_invocation_count.div_ceil(64); + + cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); + } + + // ~~~ Invalidate pointers ~~~ + { + cache_insertion_pass.set_pipeline(&self.invalidation_pipeline); + cache_insertion_pass.set_bind_group(2, None, &[]); + cache_insertion_pass + .set_immediates(0, bytemuck::bytes_of(&self.usage_buffer.timestamp())); + + // Compute dispatch amounts + let shader_invocation_count = self.size + self.structure_table.allocation_table.len(); + let workgroup_invocations = shader_invocation_count.div_ceil(64); + + cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); + } + + // ~~~ Write new pointers ~~~ + { + cache_insertion_pass.set_pipeline(&self.caching_pipeline); + cache_insertion_pass.set_bind_group(2, Some(&user_caching_bindgroup), &[]); + cache_insertion_pass.set_immediates( + 0, + bytemuck::bytes_of(&CachingPipelineImmediates { + frame_timestamp: self.usage_buffer.timestamp(), + write_pointers: 1, // true + }), + ); + + // Compute dispatch amounts + let shader_invocation_count = + insertion.structure_nodes.size() as usize / size_of::>(); + let workgroup_invocations = shader_invocation_count.div_ceil(64); + + cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); + } + + } + + pub fn next_frame(&mut self) + { + self.usage_buffer.next_frame(); + } + + pub fn total_request_count(&self) -> u32 + { + let (tx, rx) = std::sync::mpsc::sync_channel(1); + wgpu::util::DownloadBuffer::read_buffer( + &self.device, + &self.queue, + &self.request_write_count_buffer.slice(0..), + move |download_buffer| { + let vec = download_buffer.unwrap().to_vec(); + let value = u32::from_ne_bytes(std::array::from_fn(|i| vec[i])); + let _ = tx.try_send(value); + }, + ); + + loop + { + if let Ok(value) = rx.try_recv() + { + return value; + } + self.device.poll(wgpu::wgt::PollType::Poll).unwrap(); + } + + //panic!("Could not retrieve total request count."); + } +} diff --git a/src/voxel_cache/cache.rs b/src/voxel_cache/cache.rs deleted file mode 100644 index a2a0742..0000000 --- a/src/voxel_cache/cache.rs +++ /dev/null @@ -1,2113 +0,0 @@ -use std::num::NonZero; - -use bytemuck::Pod; -use bytemuck::Zeroable; -use bytemuck::bytes_of; -use wgpu::BindGroup; -use wgpu::BindGroupEntry; -use wgpu::BindGroupLayout; -use wgpu::Buffer; -use wgpu::BufferUsages; -use wgpu::CommandEncoder; -use wgpu::ComputePipeline; -use wgpu::Device; -use wgpu::Queue; -use wgpu::ShaderStages; -use wgpu::util::DeviceExt; -use wgpu::util::StagingBelt; - -use crate::voxel_cache::gpu::StructurePointer; -use crate::voxel_cache::sparse::Color; - -#[derive(Clone, Copy, Zeroable, Pod, Debug)] -#[repr(C)] -pub struct CacheNodeRequest -{ - // Requested ressource - pub structure_id: u32, - pub structure_locator: u32, - - // Write back info - pub node_index: u32, - pub child_index: u32, -} - -pub struct RequestBuffer -{ - pub cache_size: usize, - pub request_buffer: Buffer, - pub request_count_buffer: Buffer, - pub indirect_count_storage: Buffer, - pub element_sort_count_buffer: Buffer, - pub indirect_count: Buffer, - pub sort_buffer: Buffer, - pub reset_pipeline: ComputePipeline, - pub sort_pipeline: ComputePipeline, - pub bindgroup: BindGroup, - pub device: Device, - - pub compaction_pipeline: ComputePipeline, - pub running_sum_pipeline: ComputePipeline - -} - -pub struct RequestBufferElement -where - [(); N * N * N]:, -{ - request_count: [u32; N * N * N], -} - -impl RequestBuffer -where - [(); N * N * N]:, -{ - pub fn new(cache_size: usize, device: Device) -> Self - { - let request_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some(format!("request_buffer_{N}").as_str()), - size: (size_of::>() * cache_size) as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_DST, - mapped_at_creation: false, - }); - - let request_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Request_count_buffer"), - size: size_of::() as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - - - let indirect_count = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("indirect_count_buffer"), - usage: BufferUsages::INDIRECT | BufferUsages::COPY_DST, - size: size_of::() as u64, - mapped_at_creation: false, - }); - - let indirect_count_storage = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("indirect_count_storage_buffer"), - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - size: size_of::() as u64, - mapped_at_creation: false, - }); - - let element_sort_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("indirect_count_storage_buffer"), - usage: BufferUsages::STORAGE, - size: size_of::() as u64, - mapped_at_creation: false, - }); - - // One element number per child entry - // One number to identify node, one to identify sub child - let sort_buffer_count = cache_size * N * N * N; - let sort_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some(format!("request_sort_buffer_{N}").as_str()), - size: (size_of::<(u32, u32)>() * sort_buffer_count) as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - mapped_at_creation: true, - }); - - let init_data = (0..(cache_size as u32)) - .flat_map(|i| (0..((N * N * N) as u32)).map(move |j| (i, j))) - .flat_map(|(i, j)| [i, j]) // Because bytemuck does not like tuples ... - .collect::>(); - - sort_buffer - .get_mapped_range_mut(0..) - .unwrap() - .copy_from_slice(bytemuck::cast_slice(&init_data)); - sort_buffer.unmap(); - - let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("request_sort_pipeline_bing_group_layout"), - entries: &[ - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 4, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("request_sort_pipeline_bindgroup"), - layout: &bindgroup_layout, - entries: &[ - BindGroupEntry { - binding: 0, - resource: request_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 1, - resource: sort_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 2, - resource: request_count_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 3, - resource: indirect_count_storage.as_entire_binding(), - }, - BindGroupEntry { - binding: 4, - resource: element_sort_count_buffer.as_entire_binding(), - }, - ], - }); - - let pipeline_layouts = &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("request_buffer_pipeline_layout"), - bind_group_layouts: &[Some(&bindgroup_layout)], - immediate_size: size_of::() as u32, - }); - - let children_count = N * N * N; - let sort_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("Sort node request"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("clear_chunk_requests shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct RequestElement - {{ - children: array - }} - - struct SortElement - {{ - node: u32, - child: u32 - }} - - var both_phase: u32; - - @group(0) @binding(0) var request_buffer: array; - @group(0) @binding(1) var sort_indirection: array; - @group(0) @binding(2) var request_count: u32; - @group(0) @binding(3) var indirect_count: vec3; - @group(0) @binding(4) var element_sort_count: u32; - - fn big_fusion(index: u32, phase_size: u32) -> vec2 - {{ - // Find out in which block index this invocation pertains - let block_index = index / (phase_size / 2); - let element_index = index % (phase_size / 2); - - let offset = block_index * phase_size; - let a = offset + element_index; - let b = offset + phase_size - 1 - element_index; - return vec2(a, b); - }} - - fn small_fusion(index: u32, phase: u32, sub_phase: u32) -> vec2 - {{ - let phase_size = u32(1 << (phase - sub_phase + 1)); - - let block_index = index / (phase_size / 2); - let element_index = index % (phase_size / 2); - - let offset = block_index * phase_size; - let a = offset + element_index; - let b = a + (phase_size / 2); - return vec2(a, b); - }} - - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - - let index = global_invocation_id.x; - let total = arrayLength(&sort_indirection); - - let sub_phase = (both_phase >> 16) & 0xFFFF; - let phase = both_phase & 0xFFFF; - let phase_total_width = u32(1 << (phase + 1)); - - - // Check if phase is last - /* - if(phase == sub_phase) - {{ - let next_phase_width = phase_total_width * 2; - if(next_phase_width >= element_sort_count && phase_total_width >= element_sort_count) - {{ - // This was the final phase, stop - indirect_count = vec3(0); - }} - }} - */ - - if(phase == 0) - {{ - let a = index * 2; - let b = a + 1; - sort_indirection[a].child = a % {children_count}; - sort_indirection[b].child = b % {children_count}; - }} - - var swap_indices = vec2(0, 0); - - if(sub_phase == 0) - {{ - // Bitonic fusion - swap_indices = big_fusion(index, phase_total_width); - }}else - {{ - swap_indices = small_fusion(index, phase, sub_phase); - }} - - // Do swap - - //if(swap_indices.y >= request_count_round_up) - if(swap_indices.y >= element_sort_count) - {{ - // Suppose that swap_indices.y is -inf, dont swap - return; - }} - - let av = request_buffer[sort_indirection[swap_indices.x].node].children[sort_indirection[swap_indices.x].child]; - let bv = request_buffer[sort_indirection[swap_indices.y].node].children[sort_indirection[swap_indices.y].child]; - - if(bv > av) - {{ - let temp = sort_indirection[swap_indices.x]; - sort_indirection[swap_indices.x] = sort_indirection[swap_indices.y]; - sort_indirection[swap_indices.y] = temp; - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - let running_sum_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("running_sum_pipeline"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("running_sum shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct RequestElement - {{ - children: array - }} - - struct SortElement - {{ - node: u32, - child: u32 - }} - - var phase: u32; - - @group(0) @binding(0) var request_buffer: array; - @group(0) @binding(1) var sort_indirection: array; - @group(0) @binding(2) var request_counts: atomic; - @group(0) @binding(3) var indirect_count: u32; - @group(0) @binding(4) var node_sort_count: u32; - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - let index = global_invocation_id.x; - let len = arrayLength(&request_buffer); - if(index > len) {{ return; }} - let sindex = index * {children_count}; - - if(phase == 0) - {{ - var count = 0; - for(var i = 0; i < {children_count}; i++) - {{ - count += select(0, 1, request_buffer[index].children[i] != 0); - }} - - sort_indirection[sindex].child = select(u32(0), u32(1), count != 0); - if count != 0 - {{ - atomicAdd(&request_counts, 1); - }} - return; - }} - - if(phase == 0xFFFFFFFF) - {{ - // Double buffering bring back - sort_indirection[sindex].child = sort_indirection[sindex].node; - }} - - // Phase is not zero, running sum part - let running_sum_phase = phase - 1; - let running_sum_offset = u32((1 << running_sum_phase) * {children_count}); - - var add = u32(0); - // Double buffering - if(running_sum_phase % 2 == 0) - {{ - if(sindex >= running_sum_offset) - {{ - add = sort_indirection[sindex - running_sum_offset].child; - }} - sort_indirection[sindex].node = sort_indirection[sindex].child + add; - }} - else - {{ - if(sindex >= running_sum_offset) - {{ - add = sort_indirection[sindex - running_sum_offset].node; - }} - sort_indirection[sindex].child = sort_indirection[sindex].node + add; - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - let compaction_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("compaction_pipeline"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("compaction shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct RequestElement - {{ - children: array - }} - - struct SortElement - {{ - node: u32, - child: u32 - }} - - var phase: u32; - - @group(0) @binding(0) var request_buffer: array; - @group(0) @binding(1) var sort_indirection: array; - @group(0) @binding(2) var request_counts: atomic; - @group(0) @binding(3) var indirect_count: vec3; - @group(0) @binding(4) var element_sort_count: u32; - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - let index = global_invocation_id.x; - let sindex = index * {children_count}; - let len = arrayLength(&request_buffer); - - element_sort_count = sort_indirection[(len - 1) * {children_count}].child * {children_count}; - let invocation_count = (element_sort_count / 2) + select(u32(0), u32(1), element_sort_count % 2 != 0); - indirect_count = vec3( - (invocation_count / 64) + select(u32(0), u32(1), invocation_count % 64 != 0), - 1, 1 - ); - let destination_index = sort_indirection[sindex].child - 1; - - var count = u32(0); - for(var i = 0; i < {children_count}; i++) - {{ - count += request_buffer[index].children[i]; - }} - - if(count != 0) // Keep ? - {{ - for(var i = u32(0); i < u32({children_count}); i++) - {{ - sort_indirection[destination_index * u32({children_count}) + i].node = index; - }} - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - let reset_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("reset_node_requests"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("clear_chunk_requests shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct RequestElement - {{ - children: array - }} - - @group(0) @binding(0) var request_buffer: array; - @group(0) @binding(1) var _ignore: array; - @group(0) @binding(2) var request_count: u32; - @group(0) @binding(4) var node_sort_count: u32; - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - let index = global_invocation_id.x; - let total = arrayLength(&request_buffer); - if(index >= total) - {{ return; }} - - if(index < total) - {{ - for(var i = 0; i < {children_count}; i += 1) - {{ - /* - if(request_buffer[index].children[i] != 0) - {{ - atomicAdd(&request_count, 1); - }} - */ - request_buffer[index].children[i] = 0; - }} - request_count = 0; - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - Self { - cache_size, - request_buffer, - request_count_buffer, - indirect_count, - indirect_count_storage, - element_sort_count_buffer, - running_sum_pipeline, - sort_buffer, - sort_pipeline, - reset_pipeline, - compaction_pipeline, - bindgroup, - device, - } - } - - pub fn request_buffer(&self) -> &Buffer - { - &self.request_buffer - } - - pub fn request_count_buffer(&self) -> &Buffer - { - &self.request_count_buffer - } - - pub fn sort_buffer(&self) -> &Buffer - { - &self.sort_buffer - } - - pub fn reset_requests(&self, encoder: &mut CommandEncoder) - { - let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("request_buffer_reset_compute_pass"), - timestamp_writes: None, - }); - - compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - compute_pass.set_pipeline(&self.reset_pipeline); - - let shader_invocations = self.cache_size; // one invocation per element - let workgroup_invocations = shader_invocations.div_ceil(64); - compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - pub fn sort_requests(&self, encoder: &mut CommandEncoder) - { - let mut compaction_compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("request_buffer_compaction_compute_pass"), - timestamp_writes: None, - }); - - compaction_compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - let request_element_count = self.cache_size; // Each child slot is sorted - let workgroups_invocations = request_element_count.div_ceil(64); - - // = Perform list compaction - - // == Running sum - compaction_compute_pass.set_pipeline(&self.running_sum_pipeline); - - // Phase 0: Put ones in correct location - compaction_compute_pass.set_immediates(0, bytes_of(&0)); - compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); - // Phase _: running sum - // Running sum phase - let running_sum_steps = request_element_count.next_power_of_two().ilog2(); - for i in 1..=running_sum_steps - { - compaction_compute_pass.set_immediates(0, bytes_of(&i)); - compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); - } - - if !running_sum_steps.is_multiple_of(2) - { - // Bring back double buffer - compaction_compute_pass.set_immediates(0, bytes_of(&0xFFFFFFFF_u32)); - compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); - } - - // == Stream compaction - compaction_compute_pass.set_pipeline(&self.compaction_pipeline); - compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); - drop(compaction_compute_pass); - - // == Copy count into indirect buffer - - - // = Sort - - // == Bitonic sort - let sort_element_count = self.cache_size * N * N * N; - let phases_upper_bound = sort_element_count.next_power_of_two().ilog2(); - - // Phase 0 dispatch all - for i in 0..=phases_upper_bound - { - encoder.copy_buffer_to_buffer(&self.indirect_count_storage, 0, &self.indirect_count, 0, Some(size_of::() as u64)); - let mut sort_compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some(format!("request_buffer_sort_compute_pass_{}", 0).as_str()), - timestamp_writes: None, - }); - sort_compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - sort_compute_pass.set_pipeline(&self.sort_pipeline); - for j in 0..=i - { - sort_compute_pass.set_immediates(0, bytemuck::bytes_of( - &(i | (j << 16)) - )); - sort_compute_pass.dispatch_workgroups_indirect(&self.indirect_count, 0); - } - drop(sort_compute_pass); - } - } -} - -pub struct UsageBuffer -{ - pub cache_size: usize, - pub current_timestamp: u32, - pub usage_buffer: Buffer, - pub sort_buffer: Buffer, - pub sort_pipeline: ComputePipeline, - pub bindgroup: BindGroup, - pub device: Device, -} - -impl UsageBuffer -{ - pub fn new(cache_size: usize, device: Device) -> Self - { - let initial_timestamp = 0; - let usage_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Usage buffer"), - size: (size_of::() * cache_size) as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, - mapped_at_creation: true, - }); - let sort_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("usage_sort_buffer"), - size: (size_of::() * cache_size) as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - mapped_at_creation: true, - }); - - usage_buffer - .get_mapped_range_mut(0..) - .unwrap() - .copy_from_slice(bytemuck::cast_slice( - vec![initial_timestamp; cache_size].as_slice(), - )); - usage_buffer.unmap(); - - sort_buffer - .get_mapped_range_mut(0..) - .unwrap() - .copy_from_slice(bytemuck::cast_slice( - &(0..(cache_size as u32)).collect::>(), - )); - sort_buffer.unmap(); - - let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("usage_sort_pipeline_bing_group_layout"), - entries: &[ - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("usage_sort_pipeline_bindgroup"), - layout: &bindgroup_layout, - entries: &[ - BindGroupEntry { - binding: 0, - resource: usage_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 1, - resource: sort_buffer.as_entire_binding(), - }, - ], - }); - - let pipeline_layouts = &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("usage_buffer_pipeline_layout"), - bind_group_layouts: &[Some(&bindgroup_layout)], - immediate_size: size_of::() as u32, - }); - - let sort_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("Sort node request"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("clear_chunk_requests shader module"), - source: wgpu::ShaderSource::Wgsl( - " - - @group(0) @binding(0) var usage_buffer: array; - @group(0) @binding(1) var sort_buffer: array; - var both_phase: u32; - - fn big_fusion(index: u32, phase_size: u32) -> vec2 - {{ - // Find out in which block index this invocation pertains - let block_index = index / (phase_size / 2); - let element_index = index % (phase_size / 2); - - let offset = block_index * phase_size; - let a = offset + element_index; - let b = offset + phase_size - 1 - element_index; - return vec2(a, b); - }} - - fn small_fusion(index: u32, phase: u32, sub_phase: u32) -> vec2 - {{ - let phase_size = u32(1 << (phase - sub_phase + 1)); - - let block_index = index / (phase_size / 2); - let element_index = index % (phase_size / 2); - - let offset = block_index * phase_size; - let a = offset + element_index; - let b = a + (phase_size / 2); - return vec2(a, b); - }} - - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - - let length = arrayLength(&sort_buffer); - let index = global_invocation_id.x; - let sub_phase = (both_phase >> 16) & 0xFFFF; - let phase = both_phase & 0xFFFF; - let phase_total_width = u32(1 << (phase + 1)); - - var swap_indices = vec2(0, 0); - - if(sub_phase == 0) - {{ - // Bitonic fusion - swap_indices = big_fusion(index, phase_total_width); - }}else - {{ - swap_indices = small_fusion(index, phase, sub_phase); - }} - - // Do swap - if(swap_indices.y >= length) - {{ - // Suppose that swap_indices.y is -inf, dont swap - return; - }} - - let av = usage_buffer[sort_buffer[swap_indices.x]]; - let bv = usage_buffer[sort_buffer[swap_indices.y]]; - - if(bv < av) - {{ - let temp = sort_buffer[swap_indices.x]; - sort_buffer[swap_indices.x] = sort_buffer[swap_indices.y]; - sort_buffer[swap_indices.y] = temp; - }} - }} - " - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - Self { - cache_size, - current_timestamp: initial_timestamp, - usage_buffer, - sort_buffer, - sort_pipeline, - bindgroup, - device, - } - } - - pub fn timestamp(&self) -> u32 - { - self.current_timestamp - } - - pub fn next_frame(&mut self) - { - let (new_timestamp, _) = self.current_timestamp.overflowing_add(1); - self.current_timestamp = new_timestamp; - } - - pub fn usage_buffer(&self) -> &Buffer - { - &self.usage_buffer - } - - pub fn sort_buffer(&self) -> &Buffer - { - &self.sort_buffer - } - - pub fn sort_usage(&self, encoder: &mut CommandEncoder) - { - let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("usage_buffer_sort_compute_pass"), - timestamp_writes: None, - }); - - compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - compute_pass.set_pipeline(&self.sort_pipeline); - - // Compute required shader invocations - let sort_steps = self.cache_size.next_power_of_two().ilog2(); // Each element is sorted - let shader_invocations = self.cache_size.div_ceil(2); // bitonic sorting : - // half as many shaders - // per element - let workgroup_invocations = shader_invocations.div_ceil(64); - - for i in 0..=sort_steps - { - for j in 0..=i - { - compute_pass.set_immediates(0, bytes_of(&(j << 16 | i))); - compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - } - } -} - -// Stores root pointers to the cache -pub struct StructureTable -{ - device: Device, - queue: Queue, - allocation_table: Vec, - available_slots: usize, - pointer_table: Buffer, - request_buffer: RequestBuffer<1>, - write_staging: StagingBelt, -} - -impl StructureTable -{ - pub fn new(device: Device, queue: Queue) -> Self - { - let pointer_table = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("structure_table_pointer_table"), - size: size_of::() as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - - StructureTable { - request_buffer: RequestBuffer::new(1, device.clone()), - write_staging: StagingBelt::new(device.clone(), size_of::() as u64), - device, - queue, - allocation_table: vec![false], - available_slots: 1, - pointer_table, - } - } - - fn double_capacity(&mut self, encoder: &mut CommandEncoder) - { - self.available_slots += self.allocation_table.len(); - self.allocation_table - .extend(vec![false; self.allocation_table.len()]); - - // Copy buffers into bigger buffers - let pointer_table = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("structure_table_pointer_table"), - size: size_of::() as u64 * self.allocation_table.len() as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST, - mapped_at_creation: false, - }); - - encoder.copy_buffer_to_buffer( - &self.pointer_table, - 0, - &pointer_table, - 0, - self.pointer_table.size(), - ); - - self.pointer_table = pointer_table; - self.request_buffer = RequestBuffer::new(self.allocation_table.len(), self.device.clone()); - } - - pub fn request_buffer(&self) -> &RequestBuffer<1> - { - &self.request_buffer - } - - pub fn request_buffer_mut(&mut self) -> &mut RequestBuffer<1> - { - &mut self.request_buffer - } - - pub fn remove_structure(&mut self, structure_id: u32) - { - self.available_slots += 1; - self.allocation_table[structure_id as usize] = false; - } - - // Returns a structure Id - pub fn allocate_structure(&mut self, subdivided: bool) -> u32 - { - let mut encoder = self - .device - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("structure_table_resize_encoder"), - }); - - // Get first available element - if self.available_slots == 0 - { - self.double_capacity(&mut encoder); - } - let (first_id, _) = self - .allocation_table - .iter() - .enumerate() - .filter(|(_, allocated)| !**allocated) - .next() - .unwrap(); - self.allocation_table[first_id] = true; - self.available_slots -= 1; - - // Write empty pointer to new pointer - self.write_staging - .write_buffer( - &mut encoder, - &self.pointer_table, - size_of::() as u64 * first_id as u64, - NonZero::new(size_of::() as u64).unwrap(), - ) - .copy_from_slice(bytemuck::bytes_of( - &StructurePointer::new(subdivided, false, 0).0, - )); - self.write_staging.finish_and_recall_on_submit(&encoder); - self.queue.submit([encoder.finish()]); - - first_id as u32 - } -} - -pub struct VoxelCache -{ - size: usize, - device: Device, - queue: Queue, - structure_pool: Buffer, - color_pool: Buffer, - location_pool: Buffer, - - pub structure_table: StructureTable, - - request_buffer: RequestBuffer, - - voxel_cache_bind_group: BindGroup, - structure_table_bind_group_layout: BindGroupLayout, - voxel_cache_render_bind_group_layout: BindGroupLayout, - - // Writing requests - request_write_count_buffer: Buffer, - request_write_bindgroup_layout_user: BindGroupLayout, - request_write_pipeline: ComputePipeline, - - // User response -> caching - caching_user_bindgroup_layout: BindGroupLayout, - caching_pipeline: ComputePipeline, - - // Invalidation - invalidation_pipeline: ComputePipeline, - - pub usage_buffer: UsageBuffer, -} - -pub struct StructurePoolElement -where - [(); N * N * N]:, -{ - pointers: [StructurePointer; N * N * N], -} - -pub struct DestinationElement -{ - pub node: u32, - pub child: u32, -} - -pub struct ColorBytes(pub u8, pub u8, pub u8, pub u8); - -impl From for ColorBytes -{ - fn from(value: Color) -> Self { - Self( - (value.0 * 255.) as u8, - (value.1 * 255.) as u8, - (value.2 * 255.) as u8, - (value.3 * 255.) as u8, - ) - } -} - -pub struct ColorPoolElement -where - [(); N * N * N]:, -{ - colors: [ColorBytes; N * N * N], -} - -pub struct LocationPoolElement -{ - pub structure_id: u32, - pub structure_locator: u32, -} - -pub struct CacheResponse -{ - // Each buffer contains the same amount of elements (structure of arrays style) - - // Cache data to bring in - pub structure_nodes: Buffer, - pub color_nodes: Buffer, - pub locations: Buffer, - - // Which nodes this extends : node_index + child_index - pub parents: Buffer, -} - -#[derive(Zeroable, bytemuck::Pod, Clone, Copy)] -#[repr(C)] -struct CachingPipelineImmediates -{ - frame_timestamp: u32, - write_pointers: u32, // Boolean -} - -unsafe fn as_raw_bytes(slice: &[T]) -> &[u8] -{ - let size = slice.len() * size_of::(); - unsafe { std::slice::from_raw_parts(slice.as_ptr() as *const u8, size) } -} - -impl VoxelCache -where - [(); N * N * N]:, -{ - pub fn new(cache_size: usize, device: Device, queue: Queue) -> Self - { - let structure_pool = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Structure pool"), - size: (size_of::>() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let location_pool = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Location pool"), - size: (size_of::() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let color_pool = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Color pool"), - size: (size_of::>() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let request_buffer = RequestBuffer::new(cache_size, device.clone()); - let usage_buffer = UsageBuffer::new(cache_size, device.clone()); - - let request_write_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("total_request_count_buffer"), - size: size_of::() as u64, - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - - // Rendering bing group layouts - let voxel_cache_render_bind_group_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("voxel_cache_render_bind_group_layout"), - entries: &[ - // Structure pool - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Color pool - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Location pool - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Request buffer - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Usage buffer - wgpu::BindGroupLayoutEntry { - binding: 4, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - - // Structure table stuff - // Structure table pointers - wgpu::BindGroupLayoutEntry { - binding: 5, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - - // Structure table request buffer - wgpu::BindGroupLayoutEntry { - binding: 6, - visibility: ShaderStages::FRAGMENT, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let voxel_cache_bind_group_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("voxel_cache_bind_group_layout"), - entries: &[ - // Structure pool - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Color pool - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Location pool - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Sorted requests - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // LRU List - wgpu::BindGroupLayoutEntry { - binding: 4, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Usage buffer - wgpu::BindGroupLayoutEntry { - binding: 5, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // pools request count - wgpu::BindGroupLayoutEntry { - binding: 6, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // total request count - wgpu::BindGroupLayoutEntry { - binding: 7, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let voxel_cache_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("voxel_cache_bind_group"), - layout: &voxel_cache_bind_group_layout, - entries: &[ - // Structure pool - wgpu::BindGroupEntry { - binding: 0, - resource: structure_pool.as_entire_binding(), - }, - // Color pool - wgpu::BindGroupEntry { - binding: 1, - resource: color_pool.as_entire_binding(), - }, - // Location pool - wgpu::BindGroupEntry { - binding: 2, - resource: location_pool.as_entire_binding(), - }, - // Sorted requests - wgpu::BindGroupEntry { - binding: 3, - resource: request_buffer.sort_buffer().as_entire_binding(), - }, - // LRU list - wgpu::BindGroupEntry { - binding: 4, - resource: usage_buffer.sort_buffer().as_entire_binding(), - }, - // Usage buffer - wgpu::BindGroupEntry { - binding: 5, - resource: usage_buffer.usage_buffer().as_entire_binding(), - }, - // Pools request count - wgpu::BindGroupEntry { - binding: 6, - resource: request_buffer.request_count_buffer().as_entire_binding(), - }, - // Total request count - wgpu::BindGroupEntry { - binding: 7, - resource: request_write_count_buffer.as_entire_binding(), - }, - ], - }); - - let structure_table_bind_group_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("structure_table_bind_group_layout"), - entries: &[ - // Pointer table - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Sorted requests - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Requests count - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let children_count = N*N*N; - let wgsl_bindings = - format!(" - struct StructurePoolElement - {{ - pointers: array - }} - - struct ColorPoolElement - {{ - colors: array - }} - - struct LocationPoolElement - {{ - structure_id: u32, - structure_locator: u32 - }} - - struct SortedRequestsElement - {{ - node: u32, - child: u32 - }} - - @group(0) @binding(0) var structure_pool: array; - @group(0) @binding(1) var color_pool: array; - @group(0) @binding(2) var location_pool: array; - @group(0) @binding(3) var sorted_requests: array; - @group(0) @binding(4) var lru_list: array; - @group(0) @binding(5) var usage_buffer: array; - @group(0) @binding(6) var pools_request_count: u32; - @group(0) @binding(7) var total_request_count: u32; - - @group(1) @binding(0) var structure_table_pointers: array; - @group(1) @binding(1) var structure_table_sorted_requests: array; - @group(1) @binding(2) var structure_table_request_count: u32; - "); - - - let write_requests_bindgroup_layout_user = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("write_requests_bindgroup_layout_user"), - entries: &[ - // User target - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let write_requests_pipeline_layout = - device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("write_requests_pipeline_layout"), - bind_group_layouts: &[ - Some(&voxel_cache_bind_group_layout), - Some(&structure_table_bind_group_layout), - Some(&write_requests_bindgroup_layout_user), - ], - immediate_size: 0, - }); - - // The user bindgroup will be created on the fly - let write_requests = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("write_requests_pipeline"), - layout: Some(&write_requests_pipeline_layout), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("write_requests_pipeline_shader_module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - {wgsl_bindings} - - struct CacheNodeRequest - {{ - // Requested ressource - structure_id: u32, - structure_locator: u32, - - // Write back info - node_index: u32, - child_index: u32, - }} - - @group(2) @binding(0) var requests: array; - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - // One shader invocation per invocation on both structure table domain - // and cache domain - var index = global_invocation_id.x; - let max_requests_count = min(arrayLength(&requests), arrayLength(&lru_list)); - total_request_count = min(max_requests_count, pools_request_count + structure_table_request_count); - if(index >= total_request_count) - {{ - return; - }} - - if(index < structure_table_request_count) - {{ - var request: CacheNodeRequest; - // Index in structure table IS structure id - request.structure_id = structure_table_sorted_requests[index].node; - request.structure_locator = 0; // Root request -> locator 0 - request.node_index = request.structure_id; - request.child_index = 0xFFFFFFFF; // child index u32::MAX -> - // Structure table request - // Write request - requests[index] = request; - return; - }} - - let pool_index = index - structure_table_request_count; - var request: CacheNodeRequest; - - // Request location comes from location pool - request.structure_id = location_pool[sorted_requests[pool_index].node].structure_id; - request.structure_locator = location_pool[sorted_requests[pool_index].node].structure_locator; - - request.node_index = sorted_requests[pool_index].node; - request.child_index = sorted_requests[pool_index].child; - - // Write request - requests[index] = request; - - }} - " - ) - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - // ~~~ Caching pipeline ~~~ - // Pipeline that takes user fulling ~some~ requests - // and writes them to the cache based on the eviction list - - // Caching pipeline binds - // - Cache data - // - // - User response : - // - response structure nodes - // - response color nodes - // - response locations - // - response destinations - let caching_user_bindgroup_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("caching_user_bindgroup_layout"), - entries: &[ - // structure nodes - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // response color nodes - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // response locations - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // response destinations - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let caching_pipeline_layout = - device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("caching_pipelien_layout"), - bind_group_layouts: &[ - Some(&voxel_cache_bind_group_layout), - Some(&structure_table_bind_group_layout), - Some(&caching_user_bindgroup_layout), - ], - immediate_size: size_of::() as u32, // Current frame timestamp - }); - - let children_count = N * N * N; - let caching_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("caching_pipeline"), - layout: Some(&caching_pipeline_layout), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("caching_pipeline_shader_module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - - struct DestinationElement - {{ - node: u32, - child: u32 - }} - - struct CachingPipelineImmediates - {{ - frame_timestamp: u32, - write_pointers: u32 - }} - - {wgsl_bindings} - - var parameters: CachingPipelineImmediates; - - @group(2) @binding(0) var structure_nodes: array; - @group(2) @binding(1) var color_nodes: array; - @group(2) @binding(2) var locations: array; - @group(2) @binding(3) var destinations: array; - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - // Copy with indirection - let index = global_invocation_id.x; - let total = arrayLength(&structure_nodes); - let total_cache = arrayLength(&lru_list); - if(index >= total || index >= total_cache) - {{ - return; - }} - - let overwritten_element = lru_list[index]; - if(parameters.write_pointers == 0 && usage_buffer[overwritten_element] != parameters.frame_timestamp) - {{ - // Phase 1 - // Copy into cache page - for(var i = 0; i < {children_count}; i++) - {{ - structure_pool[overwritten_element].pointers[i] = select(u32(0), u32(1)<<31, structure_nodes[index].pointers[i] != 0); - }} - color_pool[overwritten_element] = color_nodes[index]; - location_pool[overwritten_element] = locations[index]; - - // Mark dirty/correct timestamp - usage_buffer[overwritten_element] = parameters.frame_timestamp + 1; - }} - - if(parameters.write_pointers != 0 && usage_buffer[overwritten_element] != parameters.frame_timestamp) - {{ - // Phase 2 - - // Point parent to new page - let new_pointer = (1 << 31) | (1 << 30) | overwritten_element; - if(destinations[index].child == 0xFFFFFFFF) - {{ - structure_table_pointers[destinations[index].node] = new_pointer; - }}else if usage_buffer[destinations[index].node] != parameters.frame_timestamp + 1 - {{ - structure_pool[destinations[index].node].pointers[destinations[index].child] = new_pointer; - }} - }} - - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - // ~~~ Invalidation pipeline ~~~ - // Pipeline invalidates old nodes that points to ones - // that have been replaces base on timestamp information - let invalidation_pipeline_layout = - device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("invalidation_pipeline_layout"), - bind_group_layouts: &[ - Some(&voxel_cache_bind_group_layout), - Some(&structure_table_bind_group_layout), - ], - immediate_size: size_of::() as u32, // frame_timestamp - }); - - let invalidation_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("invalidation_pipeline"), - layout: Some(&invalidation_pipeline_layout), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("invalidation_pipeline_shader_module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - {wgsl_bindings} - - var frame_timestamp: u32; - - @compute - @workgroup_size(64) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - var index = global_invocation_id.x; - let total_pool = arrayLength(&structure_pool); - let total_table = arrayLength(&structure_table_pointers); - - if index < total_pool - {{ - for(var i = 0; i < {children_count}; i += 1) - {{ - let structure_pointer = structure_pool[index].pointers[i]; - let pointer = structure_pointer & 0x3FFFFFFF; - let pointer_subdiv = ((structure_pointer >> 31) & 1) != 0; - let pointed_timestamp = usage_buffer[pointer]; - if(pointed_timestamp == frame_timestamp + 1) // Future - // timestamp - // -> new page - {{ - // Invalidate - structure_pool[index].pointers[i] = select(u32(0), u32(1), pointer_subdiv) << 31; - }} - }} - return; - }} - index -= total_pool; - if index < total_table - {{ - let structure_pointer = structure_table_pointers[index]; - // Is pointer pointing to something valid - let pointer_valid = ((structure_pointer >> 30) & 1) != 0; - let pointer_subdiv = ((structure_pointer >> 31) & 1) != 0; - let pointer = structure_pointer & 0x3FFFFFFF; - let pointed_timestamp = usage_buffer[pointer]; - if(pointed_timestamp == frame_timestamp + 1 && pointer_valid && pointer_subdiv) // Future - // timestamp - // -> new page - {{ - // Invalidate - structure_table_pointers[index] &= select(u32(0), u32(1), pointer_subdiv) << 31; - }} - }} - - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - Self { - size: cache_size, - structure_table: StructureTable::new(device.clone(), queue.clone()), - queue, - device, - structure_pool, - color_pool, - location_pool, - request_buffer, - usage_buffer, - - voxel_cache_bind_group, - structure_table_bind_group_layout, - voxel_cache_render_bind_group_layout, - - // Write requests stage - request_write_bindgroup_layout_user: write_requests_bindgroup_layout_user, - request_write_count_buffer, - request_write_pipeline: write_requests, - - // Caching operation - caching_user_bindgroup_layout, - caching_pipeline, - - // Invalidation - invalidation_pipeline, - } - } - - pub fn bind_group_layout(&self) -> BindGroupLayout - { - self.voxel_cache_render_bind_group_layout.clone() - } - - pub fn bind_group(&self) -> BindGroup - { - self.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("voxel_cache_render_bind_group"), - layout: &self.voxel_cache_render_bind_group_layout, - entries: &[ - // Structure pool - wgpu::BindGroupEntry { - binding: 0, - resource: self.structure_pool.as_entire_binding(), - }, - // Color pool - wgpu::BindGroupEntry { - binding: 1, - resource: self.color_pool.as_entire_binding(), - }, - // Location pool - wgpu::BindGroupEntry { - binding: 2, - resource: self.location_pool.as_entire_binding(), - }, - // Request buffer - wgpu::BindGroupEntry { - binding: 3, - resource: self.request_buffer.request_buffer().as_entire_binding(), - }, - // Usage buffer - wgpu::BindGroupEntry { - binding: 4, - resource: self.usage_buffer.usage_buffer().as_entire_binding(), - }, - // Structure table - // Structure table pointers - wgpu::BindGroupEntry { - binding: 5, - resource: self.structure_table.pointer_table.as_entire_binding(), - }, - // Pools request count - wgpu::BindGroupEntry { - binding: 6, - resource: self.structure_table.request_buffer.request_buffer().as_entire_binding(), - }, - ], - }) - } - - fn structure_table_bind_group(&self) -> BindGroup - { - self.device.create_bind_group(&wgpu::BindGroupDescriptor - { - label: Some("structure_table_bind_group"), - layout: &self.structure_table_bind_group_layout, - entries: - &[ - wgpu::BindGroupEntry - { - binding: 0, - resource: self.structure_table.pointer_table.as_entire_binding(), - }, - wgpu::BindGroupEntry - { - binding: 1, - resource: self.structure_table.request_buffer.sort_buffer().as_entire_binding(), - }, - wgpu::BindGroupEntry - { - binding: 2, - resource: self.structure_table.request_buffer.request_count_buffer().as_entire_binding(), - }, - ], - }) - } - - pub fn cache_post_render(&mut self, encoder: &mut CommandEncoder, request_target: Buffer) - { - // Sort requests, reset request buffers to count requests - self.usage_buffer.sort_usage(encoder); - self.request_buffer.sort_requests(encoder); - - self.structure_table - .request_buffer_mut() - .sort_requests(encoder); - - // Sorted requests are now in the group - let user_target_bindgroup = self.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("write_requests_user_bindgroup"), - layout: &self.request_write_bindgroup_layout_user, - entries: &[ - // user target - wgpu::BindGroupEntry { - binding: 0, - resource: request_target.as_entire_binding(), - }, - ], - }); - - let mut write_requests_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("write_requests_pass"), - timestamp_writes: None, - }); - - write_requests_pass.set_bind_group(0, Some(&self.voxel_cache_bind_group), &[]); - write_requests_pass.set_bind_group(1, Some(&self.structure_table_bind_group()), &[]); - write_requests_pass.set_bind_group(2, Some(&user_target_bindgroup), &[]); - write_requests_pass.set_pipeline(&self.request_write_pipeline); - - // Compute necessary shader invocations - let shader_invocation_count = request_target.size() / size_of::() as u64; - let workgroup_invocation_count = shader_invocation_count.div_ceil(64); - write_requests_pass.dispatch_workgroups(workgroup_invocation_count as u32, 1, 1); - drop(write_requests_pass); - - self.request_buffer.reset_requests(encoder); - self.structure_table - .request_buffer_mut() - .reset_requests(encoder); - } - - pub fn current_timestamp(&self) -> u32 - { - self.usage_buffer.timestamp() - } - - pub fn cache_insert(&mut self, encoder: &mut CommandEncoder, insertion: &CacheResponse) - { - // Each of the entry of the cache insertion, is matched with the usage list to evict - let mut cache_insertion_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("cache_insertion_pass"), - timestamp_writes: None, - }); - - let user_caching_bindgroup = self.device.create_bind_group(&wgpu::BindGroupDescriptor - { - label: Some("user_caching_bind_group"), - layout: &self.caching_user_bindgroup_layout, - entries: & - [ - wgpu::BindGroupEntry - { - binding: 0, - resource: insertion.structure_nodes.as_entire_binding(), - }, - wgpu::BindGroupEntry - { - binding: 1, - resource: insertion.color_nodes.as_entire_binding(), - }, - wgpu::BindGroupEntry - { - binding: 2, - resource: insertion.locations.as_entire_binding(), - }, - wgpu::BindGroupEntry - { - binding: 3, - resource: insertion.parents.as_entire_binding(), - } - ], - }); - - // Copy insertions to evicted lru - // Point parents to children - let structure_table_bind_group = self.structure_table_bind_group(); - - // ~~~ Write data + Dirty flagging/timestamp update ~~~ - { - cache_insertion_pass.set_pipeline(&self.caching_pipeline); - cache_insertion_pass.set_bind_group(0, Some(&self.voxel_cache_bind_group), &[]); - cache_insertion_pass.set_bind_group(1, Some(&structure_table_bind_group), &[]); - cache_insertion_pass.set_bind_group(2, Some(&user_caching_bindgroup), &[]); - cache_insertion_pass.set_immediates( - 0, - bytemuck::bytes_of(&CachingPipelineImmediates { - frame_timestamp: self.usage_buffer.timestamp(), - write_pointers: 0, // false - }), - ); - - // Compute dispatch amounts - let shader_invocation_count = - insertion.structure_nodes.size() as usize / size_of::>(); - let workgroup_invocations = shader_invocation_count.div_ceil(64); - - cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - // ~~~ Invalidate pointers ~~~ - { - cache_insertion_pass.set_pipeline(&self.invalidation_pipeline); - cache_insertion_pass.set_bind_group(2, None, &[]); - cache_insertion_pass - .set_immediates(0, bytemuck::bytes_of(&self.usage_buffer.timestamp())); - - // Compute dispatch amounts - let shader_invocation_count = self.size + self.structure_table.allocation_table.len(); - let workgroup_invocations = shader_invocation_count.div_ceil(64); - - cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - // ~~~ Write new pointers ~~~ - { - cache_insertion_pass.set_pipeline(&self.caching_pipeline); - cache_insertion_pass.set_bind_group(2, Some(&user_caching_bindgroup), &[]); - cache_insertion_pass.set_immediates( - 0, - bytemuck::bytes_of(&CachingPipelineImmediates { - frame_timestamp: self.usage_buffer.timestamp(), - write_pointers: 1, // true - }), - ); - - // Compute dispatch amounts - let shader_invocation_count = - insertion.structure_nodes.size() as usize / size_of::>(); - let workgroup_invocations = shader_invocation_count.div_ceil(64); - - cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - } - - pub fn next_frame(&mut self) - { - self.usage_buffer.next_frame(); - } - - pub fn total_request_count(&self) -> u32 - { - let (tx, rx) = std::sync::mpsc::sync_channel(1); - wgpu::util::DownloadBuffer::read_buffer( - &self.device, - &self.queue, - &self.request_write_count_buffer.slice(0..), - move |download_buffer| { - let vec = download_buffer.unwrap().to_vec(); - let value = u32::from_ne_bytes(std::array::from_fn(|i| vec[i])); - let _ = tx.try_send(value); - }, - ); - - loop - { - if let Ok(value) = rx.try_recv() - { - return value; - } - self.device.poll(wgpu::wgt::PollType::Poll).unwrap(); - } - - //panic!("Could not retrieve total request count."); - } -} diff --git a/src/voxel_cache/cache.wait.rs b/src/voxel_cache/cache.wait.rs deleted file mode 100644 index 899c6c4..0000000 --- a/src/voxel_cache/cache.wait.rs +++ /dev/null @@ -1,1663 +0,0 @@ -use std::num::NonZero; - -use bytemuck::Zeroable; -use wgpu::BindGroup; -use wgpu::BindGroupEntry; -use wgpu::BindGroupLayout; -use wgpu::Buffer; -use wgpu::BufferUsages; -use wgpu::CommandEncoder; -use wgpu::ComputePipeline; -use wgpu::Device; -use wgpu::Queue; -use wgpu::ShaderStages; -use wgpu::util::StagingBelt; - -use crate::voxel::gpu::StructurePointer; -use crate::voxel::sparse::Color; - -pub struct CacheNodeRequest -{ - structure_id: u32, - structure_locator: u32, - node_index: u32, - child_index: u32, -} - -pub struct RequestBuffer -{ - cache_size: usize, - request_buffer: Buffer, - request_count_buffer: Buffer, - sort_buffer: Buffer, - reset_pipeline: ComputePipeline, - sort_pipeline: ComputePipeline, - bindgroup: BindGroup, - device: Device, -} - -pub struct RequestBufferElement -where - [(); N * N * N]:, -{ - request_count: [u32; N * N * N], -} - -impl RequestBuffer -where - [(); N * N * N]:, -{ - pub fn new(cache_size: usize, device: Device) -> Self - { - let request_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Request buffer"), - size: (size_of::>() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let request_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Request_count_buffer"), - size: size_of::() as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - // One element number per child entry - // One number to identify node, one to identify sub child - let sort_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Request buffer"), - size: (size_of::<(u32, u32)>() * cache_size * N * N * N) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: true, - }); - - let init_data = (0..cache_size) - .map(|i| (0..(N * N * N)).map(move |j| (i, j))) - .flatten() - .flat_map(|(i, j)| [i, j]) // Because bytemuck does not like tuples ... - .collect::>(); - - sort_buffer - .get_mapped_range_mut(0..) - .unwrap() - .copy_from_slice(bytemuck::cast_slice(&init_data)); - sort_buffer.unmap(); - - let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("request_sort_pipeline_bing_group_layout"), - entries: &[ - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("request_sort_pipeline_bindgroup"), - layout: &bindgroup_layout, - entries: &[ - BindGroupEntry { - binding: 0, - resource: request_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 1, - resource: sort_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 2, - resource: request_count_buffer.as_entire_binding(), - }, - ], - }); - - let pipeline_layouts = &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("request_buffer_pipeline_layout"), - bind_group_layouts: &[Some(&bindgroup_layout)], - immediate_size: 0, - }); - - let children_count = N * N * N; - let sort_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("Sort node request"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("clear_chunk_requests shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct RequestElement - {{ - children: array, {children_count}> - }} - - struct SortElement - {{ - node: u32, - child: u32 - }} - - @group(0) @binding(0) var request_counts: array; - @group(0) @binding(1) var sort_indirection: array; - @group(0) @binding(2) var request_counts: u32; - - @compute - @workgroup_size(16) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - request_counts = 0; - let index = global_invocation_id.x; - let total = arrayLength(&sort_indirections); - - // Odd pass - let a = index * 2 + 1; - let b = a + 1; - - // Gather elements - let av = request_buffer[sort_indirection[a].node].children[sort_indirection[a].child]; - let bv = request_buffer[sort_indirection[b].node].children[sort_indirections[b].child]; - - if b < total && av > bv - {{ - request_buffer[sort_indirection[a].node].children[sort_indirection[a].child] = bv; - request_buffer[sort_indirection[b].node].children[sort_indirections[b].child] = av; - }} - storageBarrier(); - - // Even pass - a = index * 2; - b = a + 1; - - // Gather elements - av = request_buffer[sort_indirection[a].node].children[sort_indirection[a].child]; - bv = request_buffer[sort_indirection[b].node].children[sort_indirections[b].child]; - - if b < total && av > bv - {{ - request_buffer[sort_indirection[a].node].children[sort_indirection[a].child] = bv; - request_buffer[sort_indirection[b].node].children[sort_indirections[b].child] = av; - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - let reset_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("reset_node_requests"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("clear_chunk_requests shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct RequestElement - {{ - children: array, {children_count}> - }} - - @group(0) @binding(0) var request_buffer: array; - @group(0) @binding(1) var _ignore: array; - @group(0) @binding(2) var request_count: atomic; - - @compute - @workgroup_size(16) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - let index = global_invocation_id.x; - let total = arrayLength(&sort_indirections); - - if(index < total) - {{ - for(var i = 0; i < {children_count}; i += 1) - {{ - if(request_buffer[index].children[i] != 0) - {{ - atomicAdd(&request_count, 1); - }} - request_buffer[index].children[i] = 0; - }} - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - Self { - cache_size, - request_buffer, - request_count_buffer, - sort_buffer, - sort_pipeline, - reset_pipeline, - bindgroup, - device, - } - } - - pub fn request_buffer(&self) -> &Buffer - { - &self.request_buffer - } - - pub fn request_count_buffer(&self) -> &Buffer - { - &self.request_count_buffer - } - - pub fn sort_buffer(&self) -> &Buffer - { - &self.sort_buffer - } - - pub fn reset_requests(&self, encoder: &mut CommandEncoder) - { - let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("request_buffer_sort_compute_pass"), - timestamp_writes: None, - }); - - compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - compute_pass.set_pipeline(&self.reset_pipeline); - - let shader_invocations = self.cache_size; // one invocation per element - let workgroup_invocations = shader_invocations.div_ceil(16); - compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - pub fn sort_requests(&self, encoder: &mut CommandEncoder) - { - let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("request_buffer_sort_compute_pass"), - timestamp_writes: None, - }); - - compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - compute_pass.set_pipeline(&self.sort_pipeline); - - // Compute required shader invocations - let sort_element_count = self.cache_size * N * N * N; // Each child slot is sorted - let shader_invocations = sort_element_count.next_multiple_of(2) / 2; // odd-even sorting : - // half as many shaders - // per element - let workgroup_invocations = shader_invocations.next_multiple_of(16) / 16; - - for _ in 0..sort_element_count.div_ceil(2) - { - compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - } -} - -pub struct UsageBuffer -{ - cache_size: usize, - current_timestamp: u32, - usage_buffer: Buffer, - sort_buffer: Buffer, - sort_pipeline: ComputePipeline, - bindgroup: BindGroup, - device: Device, -} - -impl UsageBuffer -{ - pub fn new(cache_size: usize, device: Device) -> Self - { - let initial_timestamp = 0; - let usage_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Usage buffer"), - size: (size_of::() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: true, - }); - let sort_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("usage_sort_buffer"), - size: (size_of::() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: true, - }); - - usage_buffer - .get_mapped_range_mut(0..) - .unwrap() - .copy_from_slice(bytemuck::cast_slice( - vec![initial_timestamp; cache_size].as_slice(), - )); - usage_buffer.unmap(); - - sort_buffer - .get_mapped_range_mut(0..) - .unwrap() - .copy_from_slice(bytemuck::cast_slice(&(0..cache_size).collect::>())); - sort_buffer.unmap(); - - let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("usage_sort_pipeline_bing_group_layout"), - entries: &[ - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("usage_sort_pipeline_bindgroup"), - layout: &bindgroup_layout, - entries: &[ - BindGroupEntry { - binding: 0, - resource: usage_buffer.as_entire_binding(), - }, - BindGroupEntry { - binding: 1, - resource: sort_buffer.as_entire_binding(), - }, - ], - }); - - let pipeline_layouts = &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("usage_buffer_pipeline_layout"), - bind_group_layouts: &[Some(&bindgroup_layout)], - immediate_size: 0, - }); - - let sort_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("Sort node request"), - layout: Some(pipeline_layouts), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("clear_chunk_requests shader module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - @group(0) @binding(0) var usage_stamps: array; - @group(0) @binding(1) var sort_indirection: array; - - @compute - @workgroup_size(16) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - let index = global_invocation_id.x; - let total = arrayLength(&sort_indirections); - - // Odd pass - let a = index * 2 + 1; - let b = a + 1; - - // Gather elements - let av = usage_stamps[sort_indirection[a]]; - let bv = usage_stamps[sort_indirection[b]]; - - if b < total && av < bv // Reverse sort order - {{ - request_buffer[sort_indirection[a]] = bv; - request_buffer[sort_indirection[b]] = av; - }} - storageBarrier(); - - // Even pass - a = index * 2; - b = a + 1; - - // Gather elements - av = usage_stamps[sort_indirection[a]]; - bv = usage_stamps[sort_indirection[b]]; - - if b < total && av < bv // Reverse sort order - {{ - request_buffer[sort_indirection[a]] = bv; - request_buffer[sort_indirection[b]] = av; - }} - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - Self { - cache_size, - current_timestamp: initial_timestamp, - usage_buffer, - sort_buffer, - sort_pipeline, - bindgroup, - device, - } - } - - pub fn timestamp(&self) -> u32 - { - self.current_timestamp - } - - pub fn next_frame(&mut self) - { - let (new_timestamp, _) = self.current_timestamp.overflowing_add(1); - self.current_timestamp = new_timestamp; - } - - pub fn usage_buffer(&self) -> &Buffer - { - &self.usage_buffer - } - - pub fn sort_buffer(&self) -> &Buffer - { - &self.sort_buffer - } - - pub fn sort_usage(&self, encoder: &mut CommandEncoder) - { - let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("usage_buffer_sort_compute_pass"), - timestamp_writes: None, - }); - - compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); - compute_pass.set_pipeline(&self.sort_pipeline); - - // Compute required shader invocations - let sort_element_count = self.cache_size; // Each element is sorted - let shader_invocations = sort_element_count.next_multiple_of(2) / 2; // odd-even sorting : - // half as many shaders - // per element - let workgroup_invocations = shader_invocations.next_multiple_of(16) / 16; - - for _ in 0..sort_element_count.div_ceil(2) - { - compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - } -} - -// Stores root pointers to the cache -pub struct StructureTable -{ - device: Device, - queue: Queue, - allocation_table: Vec, - available_slots: usize, - pointer_table: Buffer, - request_buffer: RequestBuffer<1>, - write_staging: StagingBelt, -} - -impl StructureTable -{ - pub fn new(device: Device, queue: Queue) -> Self - { - let pointer_table = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("structure_table_pointer_table"), - size: size_of::() as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - StructureTable { - request_buffer: RequestBuffer::new(1, device.clone()), - write_staging: StagingBelt::new(device.clone(), size_of::() as u64), - device, - queue, - allocation_table: vec![false], - available_slots: 1, - pointer_table, - } - } - - fn double_capacity(&mut self, encoder: &mut CommandEncoder) - { - self.available_slots += self.allocation_table.len(); - self.allocation_table - .extend(vec![false; self.allocation_table.len()]); - - // Copy buffers into bigger buffers - let pointer_table = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("structure_table_pointer_table"), - size: size_of::() as u64 * self.allocation_table.len() as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - encoder.copy_buffer_to_buffer( - &self.pointer_table, - 0, - &pointer_table, - 0, - self.pointer_table.size(), - ); - - self.pointer_table = pointer_table; - self.request_buffer = RequestBuffer::new(self.allocation_table.len(), self.device.clone()); - } - - pub fn request_buffer(&self) -> &RequestBuffer<1> - { - &self.request_buffer - } - - pub fn request_buffer_mut(&mut self) -> &mut RequestBuffer<1> - { - &mut self.request_buffer - } - - pub fn remove_structure(&mut self, structure_id: u32) - { - self.available_slots += 1; - self.allocation_table[structure_id as usize] = false; - } - - // Returns a structure Id - pub fn allocate_structure(&mut self, subdivided: bool) -> u32 - { - let mut encoder = self - .device - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("structure_table_resize_encoder"), - }); - - // Get first available element - if self.available_slots == 0 - { - self.double_capacity(&mut encoder); - } - let (first_id, _) = self - .allocation_table - .iter() - .enumerate() - .filter(|(_, allocated)| !**allocated) - .next() - .unwrap(); - self.allocation_table[first_id] = true; - self.available_slots -= 1; - - // Write empty pointer to new pointer - self.write_staging - .write_buffer( - &mut encoder, - &self.pointer_table, - size_of::() as u64 * first_id as u64, - NonZero::new(size_of::() as u64).unwrap(), - ) - .copy_from_slice(bytemuck::bytes_of( - &StructurePointer::new(subdivided, false, 0).0, - )); - - first_id as u32 - } -} - -pub struct VoxelCache -{ - size: usize, - device: Device, - queue: Queue, - structure_pool: Buffer, - color_pool: Buffer, - location_pool: Buffer, - - structure_table: StructureTable, - - request_buffer: RequestBuffer, - - // Writing requests - request_write_count_buffer: Buffer, - request_write_bindgroup_layout_user: BindGroupLayout, - request_write_bindgroup: BindGroup, - request_write_pipeline: ComputePipeline, - - // User response -> caching - caching_user_bindgroup_layout: BindGroupLayout, - caching_pools_bindgroup: BindGroup, - caching_pipeline: ComputePipeline, - - // Invalidation - invalidation_bindgroup: BindGroup, - invalidation_pipeline: ComputePipeline, - invalidation_pipeline_bingroup_layout_dynamic: BindGroupLayout, - - usage_buffer: UsageBuffer, -} - -pub struct StructurePoolElement -where - [(); N * N * N]:, -{ - pointers: [StructurePointer; N * N * N], -} - -pub struct ColorPoolElement -where - [(); N * N * N]:, -{ - colors: [Color; N * N * N], -} - -pub struct LocationPoolElement -{ - structure_id: u32, - structure_locator: u32, -} - -pub struct CacheResponse -{ - // Each buffer contains the same amount of elements (structure of arrays style) - - // Cache data to bring in - structure_nodes: Buffer, - color_nodes: Buffer, - locations: Buffer, - - // Which nodes this extends : node_index + child_index - parents: Buffer, -} - -#[derive(Zeroable, bytemuck::Pod, Clone, Copy)] -#[repr(C)] -struct CachingPipelineImmediates -{ - frame_timestamp: u32, - write_pointers: u32, // Boolean -} - -unsafe fn as_raw_bytes(slice: &[T]) -> &[u8] -{ - let size = slice.len() * size_of::(); - unsafe { std::slice::from_raw_parts(slice.as_ptr() as *const u8, size) } -} - -impl VoxelCache -where - [(); N * N * N]:, -{ - pub fn new(cache_size: usize, device: Device, queue: Queue) -> Self - { - let structure_pool = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Structure pool"), - size: (size_of::>() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let location_pool = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Location pool"), - size: (size_of::() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let color_pool = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Color pool"), - size: (size_of::>() * cache_size) as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - let request_buffer = RequestBuffer::new(cache_size, device.clone()); - let usage_buffer = UsageBuffer::new(cache_size, device.clone()); - - let request_write_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("total_request_count_buffer"), - size: size_of::() as u64, - usage: BufferUsages::STORAGE, - mapped_at_creation: false, - }); - - // ~~~ Write requests pipeline ~~~ - // Pipeline that takes the sorted requests - // And writes the top N into a user buffer - // So it can full fill the requests - - // This pipeline binds - // - total request count buffer - // - // Cache - // - cache request sorted list - // - cache request count - // - location information to ask user for correct node - // - // Structure table - // - structure table request sorted list - // - structure table request count - // - user target buffer - let write_requests_bindgroup_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("write_requests_bindgroup_layout"), - entries: &[ - // Total request count buffer - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Cache request sorted list - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Cache request count - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Location buffer - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let write_requests_bindgroup_layout_user = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("write_requests_bindgroup_layout_user"), - entries: &[ - // Structure table request list - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Structure table request list - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // User target - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let write_requests_bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("write_requests_bindgroup"), - layout: &write_requests_bindgroup_layout, - entries: &[ - // Total request count buffer - BindGroupEntry { - binding: 0, - resource: request_write_count_buffer.as_entire_binding(), - }, - // Sorted requests - BindGroupEntry { - binding: 1, - resource: request_buffer.sort_buffer().as_entire_binding(), - }, - // Request count - BindGroupEntry { - binding: 2, - resource: request_buffer.request_count_buffer().as_entire_binding(), - }, - // Location buffer - BindGroupEntry { - binding: 3, - resource: location_pool.as_entire_binding(), - }, - ], - }); - - let write_requests_pipeline_layout = - device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("write_requests_pipeline_layout"), - bind_group_layouts: &[ - Some(&write_requests_bindgroup_layout), - Some(&write_requests_bindgroup_layout_user), - ], - immediate_size: 0, - }); - - // The user bindgroup will be created on the fly - let write_requests = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("write_requests_pipeline"), - layout: Some(&write_requests_pipeline_layout), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("write_requests_pipeline_shader_module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct SortedElement - {{ - node: u32, - children: u32 - }} - - struct LocationElement - {{ - structure_id: u32, - structure_locator: u32, - }} - - struct CacheNodeRequest - {{ - structure_id: u32, - structure_locator: u32, - child_index: u32 - }} - - @group(0) @binding(0) var total_request_count: u32; - @group(0) @binding(1) var sorted_requests: array; - @group(0) @binding(2) var sorted_request_count: u32; - @group(0) @binding(3) var locations: array; - - @group(1) @binding(0) var structure_table_requests: array; - @group(1) @binding(1) var structure_table_request_count: u32; - @group(1) @binding(2) var target: array; - - @compute - @workgroup_size(16) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - total_request_count = sorted_request_count + structure_table_request_count; - - let st_count = arrayLength(&structure_table_requests) - let cache_count = arrayLength(&target); - var index = global_invocation_id.x; - if(index >= arrayLength(&target)) - {{ - return; - }} - - if(index < structure_table_request_count) - {{ - var req: CacheNodeRequest; - req.structure_id = structure_table_requests[index].node; - req.structure_locator = 0; - req.child_index = 0xFFFFFFFF; - target[index] = req; - return; - }} - index -= structure_table_request_count; - - // Copy with indirection - let node = sorted_requests[index].node; - let child = sorted_requests[index].child; - - var req: CacheNodeRequest; - req.structure_id = locations[node].structure_id; - req.structure_locator = locations[node].structure_locator; - req.child_index = child_index; - target[index] = req; - - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - // ~~~ Caching pipeline ~~~ - // Pipeline that takes user fulling ~some~ requests - // and writes them to the cache based on the eviction list - - // Caching pipeline binds - // - Cache data : - // -- Structure pool - // -- Color pool - // -- Location pool - // - LRU list - // - Usage buffer - // - //dynamic: - // - Structure table data: - // -- Pointers - // - // - User response : - // - response structure nodes - // - response color nodes - // - response locations - // - response destinations - - let caching_pools_bindgroup_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("caching_pools_bindgroup_layout"), - entries: &[ - // Structure pool - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Color pool - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Location pool - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // LRU list - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Usage buffer - wgpu::BindGroupLayoutEntry { - binding: 4, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let caching_user_bindgroup_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("caching_user_bindgroup_layout"), - entries: &[ - // Structure table pointers - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // structure nodes - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // response color nodes - wgpu::BindGroupLayoutEntry { - binding: 2, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // response locations - wgpu::BindGroupLayoutEntry { - binding: 3, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // response destinations - wgpu::BindGroupLayoutEntry { - binding: 4, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let caching_pools_bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("caching_pools_bind_group"), - layout: &write_requests_bindgroup_layout, - entries: &[ - // Structure pool - BindGroupEntry { - binding: 0, - resource: structure_pool.as_entire_binding(), - }, - // Color pool - BindGroupEntry { - binding: 1, - resource: color_pool.as_entire_binding(), - }, - // Location pool - BindGroupEntry { - binding: 2, - resource: location_pool.as_entire_binding(), - }, - // LRU list - BindGroupEntry { - binding: 3, - resource: usage_buffer.sort_buffer().as_entire_binding(), - }, - // Usage buffer - BindGroupEntry { - binding: 4, - resource: usage_buffer.usage_buffer().as_entire_binding(), - }, - ], - }); - - let caching_pipeline_layout = - device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("caching_pipelien_layout"), - bind_group_layouts: &[ - Some(&caching_pools_bindgroup_layout), - Some(&caching_user_bindgroup_layout), - ], - immediate_size: size_of::() as u32, // Current frame timestamp - }); - - let children_count = N * N * N; - let caching_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("caching_pipeline"), - layout: Some(&caching_pipeline_layout), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("caching_pipeline_shader_module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct StructurePoolElement - {{ - pointers: array - }} - - struct ColorPoolElement - {{ - pointers: array, {children_count}> - }} - - struct LocationPoolElement - {{ - structure_id: u32, - structure_locator: u32, - }} - - struct DestinationElement - {{ - node: u32, - child: u32 - }} - - struct CachingPipelineImmediates - {{ - frame_timestamp: u32, - write_pointers: u32 - }} - - @group(0) @binding(0) var structure_pool: array; - @group(0) @binding(1) var color_pool: array; - @group(0) @binding(2) var location_pool: array; - @group(0) @binding(3) var lru_list: array; - @group(0) @binding(4) var usage_buffer: array; - - var parameters: CachingPipelineImmediates; - - @group(1) @binding(0) var structure_table_pointers: array; - @group(1) @binding(1) var structure_nodes: array; - @group(1) @binding(2) var color_nodes: array; - @group(1) @binding(3) var locations: array; - @group(1) @binding(4) var destinations: array; - - @compute - @workgroup_size(16) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - // Copy with indirection - let index = global_invocation_id.x; - let total = arrayLength(&structure_nodes); - if index < total - {{ - let overwritten_element = lru_list[index]; - if(parameters.write_pointers == 0) - {{ - // Copy into cache page - structure_pool[overwritten_element] = structure_nodes[index]; - color_pool[overwritten_element] = color_nodes[index]; - location_pool[overwritten_element] = locations[index]; - - // Mark dirty/corret timestamp - usage_buffer[overwritten_element] = frame_timestamp + 1; - }} - else - {{ - // Point parent to new page - let new_pointer = (1 << 31) | (1 << 30) | (overwritten_element & 0x3FFFFFFF); - if(destinations[index].child == 0xFFFFFFFF) - {{ - structure_table_pointers[destinations[index].node] = new_pointer - }}else - {{ - structure_pool[destinations[index].node].pointers[destinations[index].child] = new_pointer; - }} - }} - }} - - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - // ~~~ Invalidation pipeline ~~~ - // Pipeline invalidates old nodes that points to ones - // that have been replaces base on timestamp information - - // Pipeline binds - // - structure pool - // - usage buffer - //dynamic - // - structure table - let invalidation_pipeline_bingroup_layout = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("invalidation_pipeline_bingroup_layout"), - entries: &[ - // Structure pool - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - // Usage buffer - wgpu::BindGroupLayoutEntry { - binding: 1, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: true }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let invalidation_pipeline_bingroup_layout_dynamic = - device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some("invalidation_pipeline_bingroup_layout_dynamic"), - entries: &[ - // Structure table - wgpu::BindGroupLayoutEntry { - binding: 0, - visibility: ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { read_only: false }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }, - ], - }); - - let invalidation_pipeline_layout = - device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("invalidation_pipeline_layout"), - bind_group_layouts: &[ - Some(&invalidation_pipeline_bingroup_layout), - Some(&invalidation_pipeline_bingroup_layout_dynamic), - ], - immediate_size: size_of::() as u32, // frame_timestamp - }); - - let invalidation_bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("invalidation_bindgroup"), - layout: &invalidation_pipeline_bingroup_layout, - entries: &[ - wgpu::BindGroupEntry { - binding: 0, - resource: structure_pool.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 1, - resource: usage_buffer.usage_buffer().as_entire_binding(), - }, - ], - }); - - let invalidation_pipeline = - device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("invalidation_pipeline"), - layout: Some(&invalidation_pipeline_layout), - module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("caching_pipeline_shader_module"), - source: wgpu::ShaderSource::Wgsl( - format!(" - struct StructurePoolElement - {{ - pointers: array - }} - - struct LocationPoolElement - {{ - structure_id: u32, - structure_locator: u32, - }} - - @group(0) @binding(0) var structure_pool: array; - @group(0) @binding(1) var usage_buffer: array; - - @group(1) @binding(0) var structure_table: array; - - var frame_timestamp: u32; - - @compute - @workgroup_size(16) - fn main( - @builtin(global_invocation_id) global_invocation_id: vec3 - ) - {{ - // Copy with indirection - var index = global_invocation_id.x; - let total_pool = arrayLength(&structure_pool); - let total_table = arrayLength(&structure_table); - - if index < total_pool - {{ - for(var i = 0; i < {children_count}; i += 1) - {{ - let pointer = structure_pool[index].pointers[i]; - // Is pointer pointing to something valid - let pointer_valid = ((pointer >> 30) & 1) != 0; - let pointer = pointer & 0x3FFFFFFF; - let pointed_timestamp = usage_buffer[pointer]; - if(pointed_timestamp == frame_timestamp + 1) // Future - // timestamp - // -> new page - {{ - // Invalidate - structure_pool[index].pointers[i] &= !(1 << 30); - }} - }} - return; - }} - index -= table_pool; - if index < total_table - {{ - let pointer = structure_table[index]; - // Is pointer pointing to something valid - let pointer_valid = ((pointer >> 30) & 1) != 0; - let pointer = pointer & 0x3FFFFFFF; - let pointed_timestamp = usage_buffer[pointer]; - if(pointed_timestamp == frame_timestamp + 1) // Future - // timestamp - // -> new page - {{ - // Invalidate - structure_table[index] &= !(1 << 30); - }} - }} - - }} - ") - .into(), - ), - }), - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - - Self { - size: cache_size, - structure_table: StructureTable::new(device.clone(), queue.clone()), - queue, - device, - structure_pool, - color_pool, - location_pool, - request_buffer, - usage_buffer, - - // Write requests stage - request_write_bindgroup_layout_user: write_requests_bindgroup_layout_user, - request_write_count_buffer, - request_write_bindgroup: write_requests_bindgroup, - request_write_pipeline: write_requests, - - // Caching operation - caching_pools_bindgroup, - caching_user_bindgroup_layout, - caching_pipeline, - - // Invalidation - invalidation_bindgroup, - invalidation_pipeline, - invalidation_pipeline_bingroup_layout_dynamic, - } - } - - pub fn cache_post_render(&mut self, encoder: &mut CommandEncoder, request_target: Buffer) - { - // Sort requests, reset request buffers to count requests - self.usage_buffer.sort_usage(encoder); - self.request_buffer.sort_requests(encoder); - self.request_buffer.reset_requests(encoder); - - self.structure_table - .request_buffer_mut() - .reset_requests(encoder); - self.structure_table - .request_buffer_mut() - .sort_requests(encoder); - - // Sorted requests are now in the group - let user_target_bindgroup = self.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("write_requests_user_bindgroup"), - layout: &self.request_write_bindgroup_layout_user, - entries: &[ - wgpu::BindGroupEntry { - binding: 0, - resource: self - .structure_table - .request_buffer() - .sort_buffer() - .as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 1, - resource: self - .structure_table - .request_buffer - .request_count_buffer() - .as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 2, - resource: request_target.as_entire_binding(), - }, - ], - }); - - let mut write_requests_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("write_requests_pass"), - timestamp_writes: None, - }); - - write_requests_pass.set_bind_group(0, Some(&self.request_write_bindgroup), &[]); - write_requests_pass.set_bind_group(1, Some(&user_target_bindgroup), &[]); - write_requests_pass.set_pipeline(&self.request_write_pipeline); - - // Compute necessary shader invocations - let shader_invocation_count = request_target.size() / size_of::() as u64; - let workgroup_invocation_count = shader_invocation_count.div_ceil(shader_invocation_count); - write_requests_pass.dispatch_workgroups(workgroup_invocation_count as u32, 1, 1); - } - - pub fn cache_insert(&mut self, encoder: &mut CommandEncoder, insertion: &CacheResponse) - { - // Each of the entry of the cache insertion, is matched with the usage list to evict - let mut cache_insertion_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("cache_insertion_pass"), - timestamp_writes: None, - }); - - // Copy insertions to evicted lru - // Point parents to children - let user_caching_bindgroup = self.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("caching_user_bindgroup"), - layout: &self.caching_user_bindgroup_layout, - entries: &[ - wgpu::BindGroupEntry { - binding: 0, - resource: self.structure_table.pointer_table.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 1, - resource: insertion.structure_nodes.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 2, - resource: insertion.color_nodes.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 3, - resource: insertion.locations.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 4, - resource: insertion.parents.as_entire_binding(), - }, - ], - }); - - let invalidation_dynamic_bindgroup = - self.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("invalidation_dynamic_bindgroup"), - layout: &self.invalidation_pipeline_bingroup_layout_dynamic, - entries: &[wgpu::BindGroupEntry { - binding: 0, - resource: self.structure_table.pointer_table.as_entire_binding(), - }], - }); - - // ~~~ Write data + Dirty flagging/timestamp update ~~~ - { - cache_insertion_pass.set_bind_group(0, Some(&self.caching_pools_bindgroup), &[]); - cache_insertion_pass.set_bind_group(1, Some(&user_caching_bindgroup), &[]); - cache_insertion_pass.set_immediates( - 0, - bytemuck::bytes_of(&CachingPipelineImmediates { - frame_timestamp: self.usage_buffer.timestamp(), - write_pointers: 0, // false - }), - ); - - // Compute dispatch amounts - let shader_invocation_count = - insertion.structure_nodes.size() as usize / size_of::>(); - let workgroup_invocations = shader_invocation_count.div_ceil(16); - - cache_insertion_pass.set_pipeline(&self.caching_pipeline); - cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - // ~~~ Invalidate pointers ~~~ - { - cache_insertion_pass.set_bind_group(0, Some(&self.invalidation_bindgroup), &[]); - cache_insertion_pass.set_bind_group(1, Some(&invalidation_dynamic_bindgroup), &[]); - cache_insertion_pass - .set_immediates(0, bytemuck::bytes_of(&self.usage_buffer.timestamp())); - - // Compute dispatch amounts - let shader_invocation_count = self.size + self.structure_table.allocation_table.len(); - let workgroup_invocations = shader_invocation_count.div_ceil(16); - - cache_insertion_pass.set_pipeline(&self.invalidation_pipeline); - cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - - // ~~~ Write new pointers ~~~ - { - cache_insertion_pass.set_bind_group(0, Some(&self.caching_pools_bindgroup), &[]); - cache_insertion_pass.set_bind_group(1, Some(&user_caching_bindgroup), &[]); - cache_insertion_pass.set_immediates( - 0, - bytemuck::bytes_of(&CachingPipelineImmediates { - frame_timestamp: self.usage_buffer.timestamp(), - write_pointers: 1, // true - }), - ); - - // Compute dispatch amounts - let shader_invocation_count = - insertion.structure_nodes.size() as usize / size_of::>(); - let workgroup_invocations = shader_invocation_count.div_ceil(16); - - cache_insertion_pass.set_pipeline(&self.caching_pipeline); - cache_insertion_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); - } - } - - pub fn total_request_count(&self) -> u32 - { - let (tx, rx) = std::sync::mpsc::sync_channel(1); - wgpu::util::DownloadBuffer::read_buffer( - &self.device, - &self.queue, - &self.request_write_count_buffer.slice(0..), - move |download_buffer| { - let vec = download_buffer.unwrap().to_vec(); - let value = u32::from_ne_bytes(std::array::from_fn(|i| vec[i])); - tx.send(value).unwrap(); - }, - ); - - for _ in 0..16 - { - if let Ok(value) = rx.try_recv() - { - return value; - } - self.device.poll(wgpu::wgt::PollType::Poll); - } - - panic!(); - } -} diff --git a/src/voxel_cache/data.rs b/src/voxel_cache/data.rs new file mode 100644 index 0000000..ba74146 --- /dev/null +++ b/src/voxel_cache/data.rs @@ -0,0 +1,102 @@ +use bytemuck::Pod; +use bytemuck::Zeroable; +use wgpu::Buffer; + +use crate::sparse_tree::Color; + +#[derive(Clone, Copy, Zeroable, Pod, Debug)] +#[repr(C)] +pub struct CacheNodeRequest +{ + // Requested ressource + pub structure_id: u32, + pub structure_locator: u32, + + // Write back info + pub node_index: u32, + pub child_index: u32, +} + +pub struct RequestBufferElement +where + [(); N * N * N]:, +{ + request_count: [u32; N * N * N], +} + +pub struct StructurePoolElement +where + [(); N * N * N]:, +{ + pointers: [StructurePointer; N * N * N], +} + +pub struct DestinationElement +{ + pub node: u32, + pub child: u32, +} + +pub struct ColorBytes(pub u8, pub u8, pub u8, pub u8); + +impl From for ColorBytes +{ + fn from(value: Color) -> Self + { + Self( + (value.0 * 255.) as u8, + (value.1 * 255.) as u8, + (value.2 * 255.) as u8, + (value.3 * 255.) as u8, + ) + } +} + +pub struct ColorPoolElement +where + [(); N * N * N]:, +{ + colors: [ColorBytes; N * N * N], +} + +pub struct LocationPoolElement +{ + pub structure_id: u32, + pub structure_locator: u32, +} + +pub struct CacheResponse +{ + // Each buffer contains the same amount of elements (structure of arrays style) + + // Cache data to bring in + pub structure_nodes: Buffer, + pub color_nodes: Buffer, + pub locations: Buffer, + + // Which nodes this extends : node_index + child_index + pub parents: Buffer, +} + +#[derive(Clone, Copy, Pod, Zeroable)] +#[repr(transparent)] +pub struct StructurePointer(pub u32); + +impl StructurePointer +{ + pub fn new(subdivided: bool, pointer_valid: bool, pointer: u32) -> Self + { + assert!(pointer >> 30 == 0); + StructurePointer((subdivided as u32) << 31 | (pointer_valid as u32) << 30 | pointer) + } +} + +#[derive(Clone, Copy, Zeroable)] +#[repr(C)] +pub struct ExplicitNTreeNode +where + [(); N * N * N]:, +{ + pub structure: [StructurePointer; N * N * N], + pub colors: [Color; N * N * N], +} diff --git a/src/voxel_cache/gpu.rs b/src/voxel_cache/gpu.rs deleted file mode 100644 index 827267a..0000000 --- a/src/voxel_cache/gpu.rs +++ /dev/null @@ -1,27 +0,0 @@ -use bytemuck::Pod; -use bytemuck::Zeroable; - -use crate::voxel_cache::sparse::Color; - -#[derive(Clone, Copy, Pod, Zeroable)] -#[repr(transparent)] -pub struct StructurePointer(pub u32); - -impl StructurePointer -{ - pub fn new(subdivided: bool, pointer_valid: bool, pointer: u32) -> Self - { - assert!(pointer >> 30 == 0); - StructurePointer((subdivided as u32) << 31 | (pointer_valid as u32) << 30 | pointer) - } -} - -#[derive(Clone, Copy, Zeroable)] -#[repr(C)] -pub struct ExplicitNTreeNode -where - [(); N * N * N]:, -{ - pub structure: [StructurePointer; N * N * N], - pub colors: [Color; N * N * N], -} diff --git a/src/voxel_cache/pipeline.rs b/src/voxel_cache/pipeline.rs deleted file mode 100644 index 2ca03f2..0000000 --- a/src/voxel_cache/pipeline.rs +++ /dev/null @@ -1,68 +0,0 @@ -use std::num::NonZero; - -use bytemuck::cast_slice; -use crevice::std140::AsStd140; -use glam::Mat4; -use wgpu::BindGroup; -use wgpu::BindGroupDescriptor; -use wgpu::BindGroupEntry; -use wgpu::BindGroupLayout; -use wgpu::Buffer; -use wgpu::BufferUsages; -use wgpu::CommandEncoder; -use wgpu::CommandEncoderDescriptor; -use wgpu::ComputePass; -use wgpu::ComputePassDescriptor; -use wgpu::ComputePipeline; -use wgpu::Device; -use wgpu::Operations; -use wgpu::Queue; -use wgpu::RenderPipeline; -use wgpu::ShaderModuleDescriptor; -use wgpu::ShaderStages; -use wgpu::TextureFormat; -use wgpu::TextureView; -use wgpu::VertexBufferLayout; -use wgpu::util::StagingBelt; - -use crate::as_raw_bytes; -use crate::camera::Camera; -use crate::voxel_cache::cache::ColorPoolElement; -use crate::voxel_cache::cache::LocationPoolElement; -use crate::voxel_cache::cache::RequestBufferElement; -use crate::voxel_cache::cache::StructurePoolElement; -use crate::voxel_cache::gpu::StructurePointer; -use crate::voxel_cache::sparse::Color; - -// Represents a chunk to be rendered by the voxel pipeline -#[derive(Clone, Copy)] -#[repr(C)] -pub struct ChunkObject -{ - // Chunk object transform - pub transform: Mat4, - - // Chunk data - pub color: Color, - pub subdivided: bool, - - // Producer specific data - pub id: u32, -} - -#[derive(Clone, Copy)] -#[repr(C)] -pub struct CacheChunkObject -{ - // Chunk object transform - transform: Mat4, - - // Chunk data - color: Color, - - // Producer specific data - id: u32, - - // Pointer into cache - pointer: StructurePointer, -} diff --git a/src/voxel_cache/request_buffer.rs b/src/voxel_cache/request_buffer.rs index e69de29..203fef0 100644 --- a/src/voxel_cache/request_buffer.rs +++ b/src/voxel_cache/request_buffer.rs @@ -0,0 +1,644 @@ +use bytemuck::bytes_of; +use wgpu::BindGroup; +use wgpu::BindGroupEntry; +use wgpu::Buffer; +use wgpu::BufferUsages; +use wgpu::CommandEncoder; +use wgpu::ComputePipeline; +use wgpu::Device; +use wgpu::ShaderStages; + +use crate::voxel_cache::data::RequestBufferElement; + +pub struct RequestBuffer +{ + pub cache_size: usize, + pub request_buffer: Buffer, + pub request_count_buffer: Buffer, + pub indirect_count_storage: Buffer, + pub element_sort_count_buffer: Buffer, + pub indirect_count: Buffer, + pub sort_buffer: Buffer, + pub reset_pipeline: ComputePipeline, + pub sort_pipeline: ComputePipeline, + pub bindgroup: BindGroup, + pub device: Device, + + pub compaction_pipeline: ComputePipeline, + pub running_sum_pipeline: ComputePipeline, +} + +impl RequestBuffer +where + [(); N * N * N]:, +{ + pub fn new(cache_size: usize, device: Device) -> Self + { + let request_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some(format!("request_buffer_{N}").as_str()), + size: (size_of::>() * cache_size) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + let request_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Request_count_buffer"), + size: size_of::() as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + + let indirect_count = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("indirect_count_buffer"), + usage: BufferUsages::INDIRECT | BufferUsages::COPY_DST, + size: size_of::() as u64, + mapped_at_creation: false, + }); + + let indirect_count_storage = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("indirect_count_storage_buffer"), + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + size: size_of::() as u64, + mapped_at_creation: false, + }); + + let element_sort_count_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("indirect_count_storage_buffer"), + usage: BufferUsages::STORAGE, + size: size_of::() as u64, + mapped_at_creation: false, + }); + + // One element number per child entry + // One number to identify node, one to identify sub child + let sort_buffer_count = cache_size * N * N * N; + let sort_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some(format!("request_sort_buffer_{N}").as_str()), + size: (size_of::<(u32, u32)>() * sort_buffer_count) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: true, + }); + + let init_data = (0..(cache_size as u32)) + .flat_map(|i| (0..((N * N * N) as u32)).map(move |j| (i, j))) + .flat_map(|(i, j)| [i, j]) // Because bytemuck does not like tuples ... + .collect::>(); + + sort_buffer + .get_mapped_range_mut(0..) + .unwrap() + .copy_from_slice(bytemuck::cast_slice(&init_data)); + sort_buffer.unmap(); + + let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("request_sort_pipeline_bing_group_layout"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 2, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 3, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 4, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("request_sort_pipeline_bindgroup"), + layout: &bindgroup_layout, + entries: &[ + BindGroupEntry { + binding: 0, + resource: request_buffer.as_entire_binding(), + }, + BindGroupEntry { + binding: 1, + resource: sort_buffer.as_entire_binding(), + }, + BindGroupEntry { + binding: 2, + resource: request_count_buffer.as_entire_binding(), + }, + BindGroupEntry { + binding: 3, + resource: indirect_count_storage.as_entire_binding(), + }, + BindGroupEntry { + binding: 4, + resource: element_sort_count_buffer.as_entire_binding(), + }, + ], + }); + + let pipeline_layouts = &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("request_buffer_pipeline_layout"), + bind_group_layouts: &[Some(&bindgroup_layout)], + immediate_size: size_of::() as u32, + }); + + let children_count = N * N * N; + let sort_pipeline = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("Sort node request"), + layout: Some(pipeline_layouts), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("clear_chunk_requests shader module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + struct RequestElement + {{ + children: array + }} + + struct SortElement + {{ + node: u32, + child: u32 + }} + + var both_phase: u32; + + @group(0) @binding(0) var request_buffer: array; + @group(0) @binding(1) var sort_indirection: array; + @group(0) @binding(2) var request_count: u32; + @group(0) @binding(3) var indirect_count: vec3; + @group(0) @binding(4) var element_sort_count: u32; + + fn big_fusion(index: u32, phase_size: u32) -> vec2 + {{ + // Find out in which block index this invocation pertains + let block_index = index / (phase_size / 2); + let element_index = index % (phase_size / 2); + + let offset = block_index * phase_size; + let a = offset + element_index; + let b = offset + phase_size - 1 - element_index; + return vec2(a, b); + }} + + fn small_fusion(index: u32, phase: u32, sub_phase: u32) -> vec2 + {{ + let phase_size = u32(1 << (phase - sub_phase + 1)); + + let block_index = index / (phase_size / 2); + let element_index = index % (phase_size / 2); + + let offset = block_index * phase_size; + let a = offset + element_index; + let b = a + (phase_size / 2); + return vec2(a, b); + }} + + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + + let index = global_invocation_id.x; + let total = arrayLength(&sort_indirection); + + let sub_phase = (both_phase >> 16) & 0xFFFF; + let phase = both_phase & 0xFFFF; + let phase_total_width = u32(1 << (phase + 1)); + + + // Check if phase is last + /* + if(phase == sub_phase) + {{ + let next_phase_width = phase_total_width * 2; + if(next_phase_width >= element_sort_count && phase_total_width >= element_sort_count) + {{ + // This was the final phase, stop + indirect_count = vec3(0); + }} + }} + */ + + if(phase == 0) + {{ + let a = index * 2; + let b = a + 1; + sort_indirection[a].child = a % {children_count}; + sort_indirection[b].child = b % {children_count}; + }} + + var swap_indices = vec2(0, 0); + + if(sub_phase == 0) + {{ + // Bitonic fusion + swap_indices = big_fusion(index, phase_total_width); + }}else + {{ + swap_indices = small_fusion(index, phase, sub_phase); + }} + + // Do swap + + //if(swap_indices.y >= request_count_round_up) + if(swap_indices.y >= element_sort_count) + {{ + // Suppose that swap_indices.y is -inf, dont swap + return; + }} + + let av = request_buffer[sort_indirection[swap_indices.x].node].children[sort_indirection[swap_indices.x].child]; + let bv = request_buffer[sort_indirection[swap_indices.y].node].children[sort_indirection[swap_indices.y].child]; + + if(bv > av) + {{ + let temp = sort_indirection[swap_indices.x]; + sort_indirection[swap_indices.x] = sort_indirection[swap_indices.y]; + sort_indirection[swap_indices.y] = temp; + }} + }} + ") + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + let running_sum_pipeline = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("running_sum_pipeline"), + layout: Some(pipeline_layouts), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("running_sum shader module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + struct RequestElement + {{ + children: array + }} + + struct SortElement + {{ + node: u32, + child: u32 + }} + + var phase: u32; + + @group(0) @binding(0) var request_buffer: array; + @group(0) @binding(1) var sort_indirection: array; + @group(0) @binding(2) var request_counts: atomic; + @group(0) @binding(3) var indirect_count: u32; + @group(0) @binding(4) var node_sort_count: u32; + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + let index = global_invocation_id.x; + let len = arrayLength(&request_buffer); + if(index > len) {{ return; }} + let sindex = index * {children_count}; + + if(phase == 0) + {{ + var count = 0; + for(var i = 0; i < {children_count}; i++) + {{ + count += select(0, 1, request_buffer[index].children[i] != 0); + }} + + sort_indirection[sindex].child = select(u32(0), u32(1), count != 0); + if count != 0 + {{ + atomicAdd(&request_counts, 1); + }} + return; + }} + + if(phase == 0xFFFFFFFF) + {{ + // Double buffering bring back + sort_indirection[sindex].child = sort_indirection[sindex].node; + }} + + // Phase is not zero, running sum part + let running_sum_phase = phase - 1; + let running_sum_offset = u32((1 << running_sum_phase) * {children_count}); + + var add = u32(0); + // Double buffering + if(running_sum_phase % 2 == 0) + {{ + if(sindex >= running_sum_offset) + {{ + add = sort_indirection[sindex - running_sum_offset].child; + }} + sort_indirection[sindex].node = sort_indirection[sindex].child + add; + }} + else + {{ + if(sindex >= running_sum_offset) + {{ + add = sort_indirection[sindex - running_sum_offset].node; + }} + sort_indirection[sindex].child = sort_indirection[sindex].node + add; + }} + }} + ") + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + let compaction_pipeline = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("compaction_pipeline"), + layout: Some(pipeline_layouts), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("compaction shader module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + struct RequestElement + {{ + children: array + }} + + struct SortElement + {{ + node: u32, + child: u32 + }} + + var phase: u32; + + @group(0) @binding(0) var request_buffer: array; + @group(0) @binding(1) var sort_indirection: array; + @group(0) @binding(2) var request_counts: atomic; + @group(0) @binding(3) var indirect_count: vec3; + @group(0) @binding(4) var element_sort_count: u32; + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + let index = global_invocation_id.x; + let sindex = index * {children_count}; + let len = arrayLength(&request_buffer); + + element_sort_count = sort_indirection[(len - 1) * {children_count}].child * {children_count}; + let invocation_count = (element_sort_count / 2) + select(u32(0), u32(1), element_sort_count % 2 != 0); + indirect_count = vec3( + (invocation_count / 64) + select(u32(0), u32(1), invocation_count % 64 != 0), + 1, 1 + ); + let destination_index = sort_indirection[sindex].child - 1; + + var count = u32(0); + for(var i = 0; i < {children_count}; i++) + {{ + count += request_buffer[index].children[i]; + }} + + if(count != 0) // Keep ? + {{ + for(var i = u32(0); i < u32({children_count}); i++) + {{ + sort_indirection[destination_index * u32({children_count}) + i].node = index; + }} + }} + }} + ") + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + let reset_pipeline = + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("reset_node_requests"), + layout: Some(pipeline_layouts), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("clear_chunk_requests shader module"), + source: wgpu::ShaderSource::Wgsl( + format!(" + struct RequestElement + {{ + children: array + }} + + @group(0) @binding(0) var request_buffer: array; + @group(0) @binding(1) var _ignore: array; + @group(0) @binding(2) var request_count: u32; + @group(0) @binding(4) var node_sort_count: u32; + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + let index = global_invocation_id.x; + let total = arrayLength(&request_buffer); + if(index >= total) + {{ return; }} + + if(index < total) + {{ + for(var i = 0; i < {children_count}; i += 1) + {{ + /* + if(request_buffer[index].children[i] != 0) + {{ + atomicAdd(&request_count, 1); + }} + */ + request_buffer[index].children[i] = 0; + }} + request_count = 0; + }} + }} + ") + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + Self { + cache_size, + request_buffer, + request_count_buffer, + indirect_count, + indirect_count_storage, + element_sort_count_buffer, + running_sum_pipeline, + sort_buffer, + sort_pipeline, + reset_pipeline, + compaction_pipeline, + bindgroup, + device, + } + } + + pub fn request_buffer(&self) -> &Buffer + { + &self.request_buffer + } + + pub fn request_count_buffer(&self) -> &Buffer + { + &self.request_count_buffer + } + + pub fn sort_buffer(&self) -> &Buffer + { + &self.sort_buffer + } + + pub fn reset_requests(&self, encoder: &mut CommandEncoder) + { + let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("request_buffer_reset_compute_pass"), + timestamp_writes: None, + }); + + compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); + compute_pass.set_pipeline(&self.reset_pipeline); + + let shader_invocations = self.cache_size; // one invocation per element + let workgroup_invocations = shader_invocations.div_ceil(64); + compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); + } + + pub fn sort_requests(&self, encoder: &mut CommandEncoder) + { + let mut compaction_compute_pass = + encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("request_buffer_compaction_compute_pass"), + timestamp_writes: None, + }); + + compaction_compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); + let request_element_count = self.cache_size; // Each child slot is sorted + let workgroups_invocations = request_element_count.div_ceil(64); + + // = Perform list compaction + + // == Running sum + compaction_compute_pass.set_pipeline(&self.running_sum_pipeline); + + // Phase 0: Put ones in correct location + compaction_compute_pass.set_immediates(0, bytes_of(&0)); + compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); + // Phase _: running sum + // Running sum phase + let running_sum_steps = request_element_count.next_power_of_two().ilog2(); + for i in 1..=running_sum_steps + { + compaction_compute_pass.set_immediates(0, bytes_of(&i)); + compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); + } + + if !running_sum_steps.is_multiple_of(2) + { + // Bring back double buffer + compaction_compute_pass.set_immediates(0, bytes_of(&0xFFFFFFFF_u32)); + compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); + } + + // == Stream compaction + compaction_compute_pass.set_pipeline(&self.compaction_pipeline); + compaction_compute_pass.dispatch_workgroups(workgroups_invocations as u32, 1, 1); + drop(compaction_compute_pass); + + // == Copy count into indirect buffer + + // = Sort + + // == Bitonic sort + let sort_element_count = self.cache_size * N * N * N; + let phases_upper_bound = sort_element_count.next_power_of_two().ilog2(); + + // Phase 0 dispatch all + for i in 0..=phases_upper_bound + { + encoder.copy_buffer_to_buffer( + &self.indirect_count_storage, + 0, + &self.indirect_count, + 0, + Some(size_of::() as u64), + ); + let mut sort_compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some(format!("request_buffer_sort_compute_pass_{}", 0).as_str()), + timestamp_writes: None, + }); + sort_compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); + sort_compute_pass.set_pipeline(&self.sort_pipeline); + for j in 0..=i + { + sort_compute_pass.set_immediates(0, bytemuck::bytes_of(&(i | (j << 16)))); + sort_compute_pass.dispatch_workgroups_indirect(&self.indirect_count, 0); + } + drop(sort_compute_pass); + } + } +} diff --git a/src/voxel_cache/structure_table.rs b/src/voxel_cache/structure_table.rs index e69de29..c66179b 100644 --- a/src/voxel_cache/structure_table.rs +++ b/src/voxel_cache/structure_table.rs @@ -0,0 +1,129 @@ +use std::num::NonZero; + +use wgpu::Buffer; +use wgpu::BufferUsages; +use wgpu::CommandEncoder; +use wgpu::Device; +use wgpu::Queue; +use wgpu::util::StagingBelt; + +use crate::voxel_cache::data::StructurePointer; +use crate::voxel_cache::request_buffer::RequestBuffer; + +// Stores root pointers to the cache +pub struct StructureTable +{ + pub(crate) device: Device, + pub(crate) queue: Queue, + pub(crate) allocation_table: Vec, + pub(crate) available_slots: usize, + pub(crate) pointer_table: Buffer, + pub(crate) request_buffer: RequestBuffer<1>, + pub(crate) write_staging: StagingBelt, +} + +impl StructureTable +{ + pub fn new(device: Device, queue: Queue) -> Self + { + let pointer_table = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("structure_table_pointer_table"), + size: size_of::() as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + + StructureTable { + request_buffer: RequestBuffer::new(1, device.clone()), + write_staging: StagingBelt::new(device.clone(), size_of::() as u64), + device, + queue, + allocation_table: vec![false], + available_slots: 1, + pointer_table, + } + } + + fn double_capacity(&mut self, encoder: &mut CommandEncoder) + { + self.available_slots += self.allocation_table.len(); + self.allocation_table + .extend(vec![false; self.allocation_table.len()]); + + // Copy buffers into bigger buffers + let pointer_table = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("structure_table_pointer_table"), + size: size_of::() as u64 * self.allocation_table.len() as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + encoder.copy_buffer_to_buffer( + &self.pointer_table, + 0, + &pointer_table, + 0, + self.pointer_table.size(), + ); + + self.pointer_table = pointer_table; + self.request_buffer = RequestBuffer::new(self.allocation_table.len(), self.device.clone()); + } + + pub fn request_buffer(&self) -> &RequestBuffer<1> + { + &self.request_buffer + } + + pub fn request_buffer_mut(&mut self) -> &mut RequestBuffer<1> + { + &mut self.request_buffer + } + + pub fn remove_structure(&mut self, structure_id: u32) + { + self.available_slots += 1; + self.allocation_table[structure_id as usize] = false; + } + + // Returns a structure Id + pub fn allocate_structure(&mut self, subdivided: bool) -> u32 + { + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("structure_table_resize_encoder"), + }); + + // Get first available element + if self.available_slots == 0 + { + self.double_capacity(&mut encoder); + } + let (first_id, _) = self + .allocation_table + .iter() + .enumerate() + .filter(|(_, allocated)| !**allocated) + .next() + .unwrap(); + self.allocation_table[first_id] = true; + self.available_slots -= 1; + + // Write empty pointer to new pointer + self.write_staging + .write_buffer( + &mut encoder, + &self.pointer_table, + size_of::() as u64 * first_id as u64, + NonZero::new(size_of::() as u64).unwrap(), + ) + .copy_from_slice(bytemuck::bytes_of( + &StructurePointer::new(subdivided, false, 0).0, + )); + self.write_staging.finish_and_recall_on_submit(&encoder); + self.queue.submit([encoder.finish()]); + + first_id as u32 + } +} diff --git a/src/voxel_cache/usage_buffer.rs b/src/voxel_cache/usage_buffer.rs index e69de29..64b1cac 100644 --- a/src/voxel_cache/usage_buffer.rs +++ b/src/voxel_cache/usage_buffer.rs @@ -0,0 +1,249 @@ +use bytemuck::bytes_of; +use wgpu::BindGroup; +use wgpu::BindGroupEntry; +use wgpu::Buffer; +use wgpu::BufferUsages; +use wgpu::CommandEncoder; +use wgpu::ComputePipeline; +use wgpu::Device; +use wgpu::ShaderStages; + +pub struct UsageBuffer +{ + pub cache_size: usize, + pub current_timestamp: u32, + pub usage_buffer: Buffer, + pub sort_buffer: Buffer, + pub sort_pipeline: ComputePipeline, + pub bindgroup: BindGroup, + pub device: Device, +} + +impl UsageBuffer +{ + pub fn new(cache_size: usize, device: Device) -> Self + { + let initial_timestamp = 0; + let usage_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Usage buffer"), + size: (size_of::() * cache_size) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: true, + }); + let sort_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("usage_sort_buffer"), + size: (size_of::() * cache_size) as u64, + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: true, + }); + + usage_buffer + .get_mapped_range_mut(0..) + .unwrap() + .copy_from_slice(bytemuck::cast_slice( + vec![initial_timestamp; cache_size].as_slice(), + )); + usage_buffer.unmap(); + + sort_buffer + .get_mapped_range_mut(0..) + .unwrap() + .copy_from_slice(bytemuck::cast_slice( + &(0..(cache_size as u32)).collect::>(), + )); + sort_buffer.unmap(); + + let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("usage_sort_pipeline_bing_group_layout"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let bindgroup = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("usage_sort_pipeline_bindgroup"), + layout: &bindgroup_layout, + entries: &[ + BindGroupEntry { + binding: 0, + resource: usage_buffer.as_entire_binding(), + }, + BindGroupEntry { + binding: 1, + resource: sort_buffer.as_entire_binding(), + }, + ], + }); + + let pipeline_layouts = &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("usage_buffer_pipeline_layout"), + bind_group_layouts: &[Some(&bindgroup_layout)], + immediate_size: size_of::() as u32, + }); + + let sort_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("Sort node request"), + layout: Some(pipeline_layouts), + module: &device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("clear_chunk_requests shader module"), + source: wgpu::ShaderSource::Wgsl( + " + + @group(0) @binding(0) var usage_buffer: array; + @group(0) @binding(1) var sort_buffer: array; + var both_phase: u32; + + fn big_fusion(index: u32, phase_size: u32) -> vec2 + {{ + // Find out in which block index this invocation pertains + let block_index = index / (phase_size / 2); + let element_index = index % (phase_size / 2); + + let offset = block_index * phase_size; + let a = offset + element_index; + let b = offset + phase_size - 1 - element_index; + return vec2(a, b); + }} + + fn small_fusion(index: u32, phase: u32, sub_phase: u32) -> vec2 + {{ + let phase_size = u32(1 << (phase - sub_phase + 1)); + + let block_index = index / (phase_size / 2); + let element_index = index % (phase_size / 2); + + let offset = block_index * phase_size; + let a = offset + element_index; + let b = a + (phase_size / 2); + return vec2(a, b); + }} + + + @compute + @workgroup_size(64) + fn main( + @builtin(global_invocation_id) global_invocation_id: vec3 + ) + {{ + + let length = arrayLength(&sort_buffer); + let index = global_invocation_id.x; + let sub_phase = (both_phase >> 16) & 0xFFFF; + let phase = both_phase & 0xFFFF; + let phase_total_width = u32(1 << (phase + 1)); + + var swap_indices = vec2(0, 0); + + if(sub_phase == 0) + {{ + // Bitonic fusion + swap_indices = big_fusion(index, phase_total_width); + }}else + {{ + swap_indices = small_fusion(index, phase, sub_phase); + }} + + // Do swap + if(swap_indices.y >= length) + {{ + // Suppose that swap_indices.y is -inf, dont swap + return; + }} + + let av = usage_buffer[sort_buffer[swap_indices.x]]; + let bv = usage_buffer[sort_buffer[swap_indices.y]]; + + if(bv < av) + {{ + let temp = sort_buffer[swap_indices.x]; + sort_buffer[swap_indices.x] = sort_buffer[swap_indices.y]; + sort_buffer[swap_indices.y] = temp; + }} + }} + " + .into(), + ), + }), + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + + Self { + cache_size, + current_timestamp: initial_timestamp, + usage_buffer, + sort_buffer, + sort_pipeline, + bindgroup, + device, + } + } + + pub fn timestamp(&self) -> u32 + { + self.current_timestamp + } + + pub fn next_frame(&mut self) + { + let (new_timestamp, _) = self.current_timestamp.overflowing_add(1); + self.current_timestamp = new_timestamp; + } + + pub fn usage_buffer(&self) -> &Buffer + { + &self.usage_buffer + } + + pub fn sort_buffer(&self) -> &Buffer + { + &self.sort_buffer + } + + pub fn sort_usage(&self, encoder: &mut CommandEncoder) + { + let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("usage_buffer_sort_compute_pass"), + timestamp_writes: None, + }); + + compute_pass.set_bind_group(0, Some(&self.bindgroup), &[]); + compute_pass.set_pipeline(&self.sort_pipeline); + + // Compute required shader invocations + let sort_steps = self.cache_size.next_power_of_two().ilog2(); // Each element is sorted + let shader_invocations = self.cache_size.div_ceil(2); // bitonic sorting : + // half as many shaders + // per element + let workgroup_invocations = shader_invocations.div_ceil(64); + + for i in 0..=sort_steps + { + for j in 0..=i + { + compute_pass.set_immediates(0, bytes_of(&(j << 16 | i))); + compute_pass.dispatch_workgroups(workgroup_invocations as u32, 1, 1); + } + } + } +}