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::ColorPoolElement; use crate::voxel::cache::LocationPoolElement; use crate::voxel::cache::RequestBufferElement; use crate::voxel::cache::StructurePoolElement; use crate::voxel::gpu::StructurePointer; use crate::voxel::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, } #[derive(Clone, Copy)] pub struct ChunkHandle(usize); pub struct CacheRequest { // Records how many rays requested a // resource count: u32, } pub struct VoxelPipeline where [(); N * N * N]:, { chunk_allocations: Vec, chunk_indices: Buffer, chunk_staging: StagingBelt, chunk_objects: Buffer, chunk_requests: Buffer, chunk_objects_bind_group_layout: BindGroupLayout, ray_bind_group_layout: BindGroupLayout, chunk_objects_bind_group: BindGroup, chunk_requests_bind_group: BindGroup, render_pipeline: RenderPipeline, // Cache pools structure_pool: Buffer, color_pool: Buffer, location_pool: Buffer, // Cache interaction request_buffer: Buffer, usage_buffer: Buffer, cache_bind_group: BindGroup, device: Device, queue: Queue, // Cache keeping shaders clear_chunk_request: ComputePipeline, sort_requests: ComputePipeline, } #[derive(AsStd140)] struct RenderPipelineImmediate { view_proj: Mat4, } impl VoxelPipeline where [(); N * N * N]:, { pub fn new( cache_size: usize, device: Device, queue: Queue, surface_format: TextureFormat, ) -> Self { let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("Main shader module"), source: wgpu::ShaderSource::Wgsl( std::fs::read_to_string("shaders/voxel.wgsl") .unwrap() .into(), ), }); let cache_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("Cache bind group layout"), entries: &[ // Location 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, }, ], }); let chunk_objects_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("Chunk objects bg"), entries: &[wgpu::BindGroupLayoutEntry { binding: 0, visibility: ShaderStages::VERTEX, ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, }, count: None, }], }); let chunk_requests_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("Ray bind group layout"), entries: &[wgpu::BindGroupLayoutEntry { binding: 0, visibility: ShaderStages::FRAGMENT | ShaderStages::COMPUTE, ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: false }, has_dynamic_offset: false, min_binding_size: None, }, count: None, }], }); let request_buffer_sort_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("Ray bind 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, }], }); let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("Voxel pipeline layout"), bind_group_layouts: &[ Some(&chunk_objects_bind_group_layout), Some(&chunk_requests_bind_group_layout), Some(&cache_bind_group_layout), ], immediate_size: RenderPipelineImmediate::std140_size_static() as u32, }); let chunk_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { label: Some("Render pipeline"), layout: Some(&pipeline_layout), vertex: wgpu::VertexState { module: &shader_module, entry_point: Some("chunk"), compilation_options: Default::default(), buffers: &[Some(VertexBufferLayout { array_stride: size_of::() as u64, step_mode: wgpu::VertexStepMode::Instance, attributes: &[wgpu::VertexAttribute { format: wgpu::VertexFormat::Uint32, offset: 0, shader_location: 0, }], })], }, primitive: wgpu::PrimitiveState { topology: wgpu::PrimitiveTopology::TriangleList, strip_index_format: None, front_face: wgpu::FrontFace::Ccw, cull_mode: None, unclipped_depth: false, polygon_mode: wgpu::PolygonMode::Fill, conservative: false, }, depth_stencil: Some(wgpu::DepthStencilState { format: wgpu::TextureFormat::Depth24PlusStencil8, depth_write_enabled: Some(true), depth_compare: Some(wgpu::CompareFunction::LessEqual), stencil: wgpu::StencilState::default(), bias: wgpu::DepthBiasState::default(), }), multisample: wgpu::MultisampleState::default(), fragment: Some(wgpu::FragmentState { module: &shader_module, entry_point: Some("fragment"), compilation_options: wgpu::PipelineCompilationOptions::default(), targets: &[Some(wgpu::ColorTargetState { format: surface_format, blend: None, write_mask: wgpu::ColorWrites::default(), })], }), multiview_mask: None, cache: None, }); let chunk_indices = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Chunk index buffer"), size: size_of::() as u64, usage: BufferUsages::VERTEX | BufferUsages::COPY_DST, mapped_at_creation: true, }); chunk_indices.unmap(); let chunk_objects = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Chunk buffer"), size: size_of::() as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let chunk_requests = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Chunk request buffer"), size: size_of::() as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); // Pools let structure_pool = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Structure pool"), size: size_of::>() as u64 * cache_size as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let color_pool = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Color pool"), size: size_of::>() as u64 * cache_size as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let location_pool = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Locatino pool"), size: size_of::() as u64 * cache_size as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let request_buffer = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Request buffer"), size: size_of::>() as u64 * cache_size as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let request_sort_buffer = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Request sort buffer"), size: size_of::() as u64 * cache_size as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: true, }); request_sort_buffer .get_mapped_range_mut(0..) .unwrap() .copy_from_slice(cast_slice( (0..cache_size as u32).collect::>().as_slice(), )); request_sort_buffer.unmap(); let usage_buffer = device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Usage buffer buffer"), size: size_of::() as u64 * cache_size as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let cache_bind_group = device.create_bind_group(&BindGroupDescriptor { label: Some("Cache bind group"), layout: &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(), }, // Request buffer wgpu::BindGroupEntry { binding: 3, resource: request_buffer.as_entire_binding(), }, // Usage buffer wgpu::BindGroupEntry { binding: 4, resource: structure_pool.as_entire_binding(), }, ], }); // Cache keeping shaders let clear_chunk_request = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some("Clean chunk request"), layout: Some( &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("clear chunk requests layout"), bind_group_layouts: &[Some(&chunk_requests_bind_group_layout)], immediate_size: 0, }), ), module: &device.create_shader_module(ShaderModuleDescriptor { label: Some("clear_chunk_requests shader module"), source: wgpu::ShaderSource::Wgsl( " @group(0) @binding(0) var chunk_requests: array; @compute @workgroup_size(16) fn main( @builtin(global_invocation_id) global_invocation_id: vec3 ) { let index = global_invocation_id.x; let total = arrayLength(&chunk_requests); if(index < total) { chunk_requests[index] = 0; } } " .into(), ), }), entry_point: Some("main"), compilation_options: Default::default(), cache: None, }); // Sorting requests let sort_requests_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { label: Some("Sort request bing group"), layout: &request_buffer_sort_bind_group_layout, entries: &[BindGroupEntry { binding: 0, resource: request_sort_buffer.as_entire_binding(), }], }); let sort_requests = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some("Sort node request"), layout: Some( &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("Sort node requests"), bind_group_layouts: &[Some(&cache_bind_group_layout), Some(&request_buffer_sort_bind_group_layout)], immediate_size: 0, }), ), module: &device.create_shader_module(ShaderModuleDescriptor { label: Some("clear_chunk_requests shader module"), source: wgpu::ShaderSource::Wgsl( " struct RequestElement { children: array, 64> } @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 request_buffer: array; @group(0) @binding(4) var usage_buffer: array; @group(1) @binding(0) 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(&chunk_requests); // Odd pass let a = index * 2 + 1; let b = a + 1; // Gather elements let av = request_buffer[sort_indirection[a]]; let bv = request_buffer[sort_indirection[b]]; if b < total && av > bv { request_buffer[sort_indirection[a]] = bv; request_buffer[sort_indirection[b]] = av; } storageBarrier(); // Even pass a = index * 2; b = a + 1; // Gather elements av = request_buffer[sort_indirection[a]]; bv = request_buffer[sort_indirection[b]]; if b < total && av > bv { request_buffer[sort_indirection[a]] = bv; request_buffer[sort_indirection[b]] = av; } } " .into(), ), }), entry_point: Some("main"), compilation_options: Default::default(), cache: None, }); VoxelPipeline { // Only one slot, no allocated chunks at the beginning chunk_allocations: vec![false], chunk_objects_bind_group: device.create_bind_group(&wgpu::BindGroupDescriptor { label: Some("Chunk objects bind group"), layout: &chunk_objects_bind_group_layout, entries: &[wgpu::BindGroupEntry { binding: 0, resource: chunk_objects.as_entire_binding(), }], }), chunk_requests_bind_group: device.create_bind_group(&wgpu::BindGroupDescriptor { label: Some("Ray bind group"), layout: &chunk_requests_bind_group_layout, entries: &[wgpu::BindGroupEntry { binding: 0, resource: chunk_requests.as_entire_binding(), }], }), chunk_objects, chunk_requests, chunk_indices, structure_pool, color_pool, location_pool, request_buffer, usage_buffer, cache_bind_group, chunk_objects_bind_group_layout, ray_bind_group_layout: chunk_requests_bind_group_layout, chunk_staging: StagingBelt::new(device.clone(), size_of::() as u64), render_pipeline: chunk_pipeline, device, queue, clear_chunk_request, sort_requests, } } pub fn render( &mut self, encoder: &mut CommandEncoder, texture_view: &TextureView, depth_buffer_view: &TextureView, camera: &Camera, ) { let mut renderpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { label: None, color_attachments: &[Some(wgpu::RenderPassColorAttachment { view: texture_view, depth_slice: None, resolve_target: None, ops: wgpu::Operations { load: wgpu::LoadOp::Clear(wgpu::Color::BLACK), store: wgpu::StoreOp::Store, }, })], depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment { view: depth_buffer_view, depth_ops: Some(Operations { load: wgpu::LoadOp::Clear(1.), store: wgpu::StoreOp::Discard, }), stencil_ops: None, }), timestamp_writes: None, occlusion_query_set: None, multiview_mask: None, }); renderpass.set_pipeline(&self.render_pipeline); renderpass.set_bind_group(0, Some(&self.chunk_objects_bind_group), &[]); renderpass.set_bind_group(1, Some(&self.chunk_requests_bind_group), &[]); renderpass.set_bind_group(2, Some(&self.cache_bind_group), &[]); renderpass.set_vertex_buffer(0, self.chunk_indices.slice(0..)); renderpass.set_immediates( 0, RenderPipelineImmediate { view_proj: camera.view_proj(), } .as_std140() .as_bytes(), ); renderpass.draw( 0..36, 0..(self.chunk_allocations.iter().filter(|x| **x).count() as u32), ); // End the renderpass. drop(renderpass); let mut compute_pass = encoder.begin_compute_pass(&ComputePassDescriptor { label: Some("cache keeping pass"), timestamp_writes: None, }); compute_pass.set_bind_group(0, Some(&self.chunk_requests_bind_group), &[]); compute_pass.set_pipeline(&self.clear_chunk_request); compute_pass.dispatch_workgroups( self.chunk_allocations.len().next_multiple_of(16) as u32 / 16, 1, 1, ); drop(compute_pass) } fn update_indices(&mut self) { let indices = self .chunk_allocations .iter() .enumerate() .filter(|(_, b)| **b) .map(|(i, _)| i as u32) .collect::>(); self.chunk_indices = self.device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Chunk index buffer"), size: size_of::() as u64 * indices.len() as u64, usage: BufferUsages::VERTEX | BufferUsages::COPY_DST, mapped_at_creation: true, }); // Copy chunk indices into new buffer self.chunk_indices .get_mapped_range_mut(0..(size_of::() as u64 * indices.len() as u64)) .unwrap() .copy_from_slice(cast_slice(&indices)); self.chunk_indices.unmap(); } pub fn remove_chunks(&mut self, handles: &[ChunkHandle]) { for handle in handles.iter() { self.chunk_allocations[handle.0] = false; } self.update_indices(); } pub fn push_new_chunks(&mut self, objects: &[ChunkObject]) -> Vec { // Find room for new chunks let mut destinations = vec![0; objects.len()]; let mut encoder = self .device .create_command_encoder(&CommandEncoderDescriptor { label: Some("Chunk buffer writes"), }); // count available space let space = self.chunk_allocations.iter().filter(|x| !*x).count(); if space < objects.len() { // Allocate more space // Get first bigger power of two let necessary_space = (objects.len() + self.chunk_allocations.len()).next_power_of_two(); dbg!(necessary_space); self.chunk_allocations .extend(vec![false; necessary_space - self.chunk_allocations.len()]); // Make buffer bigger let old_buffer = self.chunk_objects.clone(); self.chunk_objects = self.device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Chunk buffer"), size: size_of::() as u64 * self.chunk_allocations.len() as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); let old_chunk_requests = self.chunk_requests.clone(); self.chunk_requests = self.device.create_buffer(&wgpu::wgt::BufferDescriptor { label: Some("Chunk request buffer"), size: size_of::() as u64 * self.chunk_allocations.len() as u64, usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, mapped_at_creation: false, }); self.chunk_objects_bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { label: Some("Chunk objects bind group"), layout: &self.chunk_objects_bind_group_layout, entries: &[wgpu::BindGroupEntry { binding: 0, resource: self.chunk_objects.as_entire_binding(), }], }); self.chunk_requests_bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { label: Some("Ray bind group"), layout: &self.ray_bind_group_layout, entries: &[wgpu::BindGroupEntry { binding: 0, resource: self.chunk_requests.as_entire_binding(), }], }); encoder.copy_buffer_to_buffer( &old_chunk_requests, 0, &self.chunk_requests, 0, old_chunk_requests.size(), ); encoder.copy_buffer_to_buffer( &old_buffer, 0, &self.chunk_objects, 0, old_buffer.size(), ); } let mut dest_ptr = 0; for (i, allocated) in self .chunk_allocations .iter_mut() .enumerate() .filter(|(_, allocated)| !**allocated) .take(destinations.len()) { if !*allocated { *allocated = true; destinations[dest_ptr] = i; dest_ptr += 1; } } // Write each new chunk for (destination, object) in destinations.iter().zip(objects.iter()) { let cache_object = CacheChunkObject { transform: object.transform, id: object.id, color: object.color, pointer: StructurePointer::new(object.subdivided, false, 0), }; let mut view = self.chunk_staging.write_buffer( &mut encoder, &self.chunk_objects, *destination as u64 * size_of::() as u64, NonZero::new(size_of::() as u64).unwrap(), ); let temp_slice = [cache_object]; view.copy_from_slice(unsafe { as_raw_bytes(&temp_slice) }); } self.update_indices(); self.chunk_staging.finish_and_recall_on_submit(&encoder); self.queue.submit([encoder.finish()]); destinations.iter().map(|d| ChunkHandle(*d)).collect() } }