Files
vxls/src/voxel/pipeline.rs
T
2026-08-31 15:25:03 +02:00

800 lines
30 KiB
Rust

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<const N: usize>
where
[(); N * N * N]:,
{
chunk_allocations: Vec<bool>,
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<const N: usize> VoxelPipeline<N>
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::<u32>() 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::<u32>() 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::<CacheChunkObject>() 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::<u32>() 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::<StructurePoolElement<N>>() 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::<ColorPoolElement<N>>() 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::<LocationPoolElement>() 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::<RequestBufferElement<N>>() 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::<u32>() 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::<Vec<_>>().as_slice(),
));
request_sort_buffer.unmap();
let usage_buffer = device.create_buffer(&wgpu::wgt::BufferDescriptor {
label: Some("Usage buffer buffer"),
size: size_of::<u32>() 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<storage, read_write> chunk_requests: array<u32>;
@compute
@workgroup_size(16)
fn main(
@builtin(global_invocation_id) global_invocation_id: vec3<u32>
)
{
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<atomic<u32>, 64>
}
@group(0) @binding(0) var<storage, read_write> structure_pool: array<u32>;
@group(0) @binding(1) var<storage, read_write> color_pool: array<u32>;
@group(0) @binding(2) var<storage, read_write> location_pool: array<u32>;
@group(0) @binding(3) var<storage, read_write> request_buffer: array<RequestElement>;
@group(0) @binding(4) var<storage, read_write> usage_buffer: array<u32>;
@group(1) @binding(0) var<storage, read_write> sort_indirection: array<u32>;
@compute
@workgroup_size(16)
fn main(
@builtin(global_invocation_id) global_invocation_id: vec3<u32>
)
{
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::<CacheChunkObject>() 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::<Vec<_>>();
self.chunk_indices = self.device.create_buffer(&wgpu::wgt::BufferDescriptor {
label: Some("Chunk index buffer"),
size: size_of::<u32>() 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::<u32>() 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<ChunkHandle>
{
// 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::<ChunkObject>() 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::<u32>() 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::<CacheChunkObject>() as u64,
NonZero::new(size_of::<CacheChunkObject>() 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()
}
}