working buffer compactor

This commit is contained in:
2026-09-17 20:52:09 +02:00
parent 7c78085b61
commit f34550e692
5 changed files with 326 additions and 26 deletions
+1 -1
View File
@@ -1,4 +1,4 @@
all: voxel.spv
all: voxel.spv compaction.spv
%.spv: %.slang
slangc $< -O3 -fvk-use-entrypoint-name -target spirv -o $@
+23 -11
View File
@@ -1,13 +1,13 @@
[[vk::binding(0, 0)]]
RWStructuredBuffer<uint32_t> count_buffer;
[[vk::binding(1, 0)]]
RWStructuredBuffer<uint32_t> index_buffer;
[[vk::binding(0, 1)]]
RWStructuredBuffer<uint32_t> reduced_buffer;
[[vk::binding(1, 1)]]
RWStructuredBuffer<uint32_t> sum_buffer;
[[vk::binding(2, 1)]]
RWStructuredBuffer<uint32_t> compaction_buffer;
groupshared uint32_t local_data[256 * 2];
static uint32_t THREAD_WIDTH = 256;
@@ -50,7 +50,7 @@ void block_sum(
var width : uint32_t = 2;
while (width <= DATA_WIDTH)
{
let dest_index = width * (thread_index + 1) - 1;
let dest_index = width * (local_thread_index + 1) - 1;
let get_index = dest_index - (width / 2);
// println!("{}, {}", get_index, dest_index);
if (dest_index < DATA_WIDTH)
@@ -61,10 +61,14 @@ void block_sum(
GroupMemoryBarrierWithGroupSync();
}
// Write to reduced buffer
reduced_buffer[workgroup_id.x] = local_data[DATA_WIDTH - 1];
// reduced_buffer[workgroup_id.x] = local_data[DATA_WIDTH - 1];
local_data[DATA_WIDTH - 1] = 0;
while (width >= 2)
{
let dest_index = width * (thread_index + 1) - 1;
let dest_index = width * (local_thread_index + 1) - 1;
let get_index = dest_index - (width / 2);
// println!("{}, {}", get_index, dest_index);
if (dest_index < DATA_WIDTH)
@@ -81,9 +85,6 @@ void block_sum(
// Dump back to sum buffer
sum_buffer[2 * thread_index] = local_data[2 * local_thread_index];
sum_buffer[2 * thread_index + 1] = local_data[2 * local_thread_index + 1];
// Write to reduced buffer
reduced_buffer[workgroup_id.x] = local_data[DATA_WIDTH - 1];
}
[numthreads(1, 1, 1)]
@@ -107,7 +108,7 @@ void linear_reduced_sum(
[numthreads(256, 1, 1)]
[shader("compute")]
void uniform_add(
void scatter(
uint32_t3 workgroup_id: SV_GroupID,
uint32_t3 local_thread_id: SV_GroupThreadID,
uint32_t3 global_thread_id: SV_DispatchThreadID)
@@ -116,9 +117,20 @@ void uniform_add(
let local_thread_index = local_thread_id.x;
// Gather
let reduced_value = reduced_buffer[global_thread_id.x];
let reduced_value = reduced_buffer[workgroup_id.x];
// Apply
sum_buffer[thread_index * 2] += reduced_value;
sum_buffer[thread_index * 2 + 1] += reduced_value;
let sum1 = sum_buffer[thread_index * 2] + reduced_value;
let sum2 = sum_buffer[thread_index * 2 + 1] + reduced_value;
// Scatter index to compact index buffer if predicate is true
if (count_buffer[thread_index * 2] != 0)
{
index_buffer[sum1] = thread_index * 2;
}
if (count_buffer[thread_index * 2 + 1] != 0)
{
index_buffer[sum2] = thread_index * 2 + 1;
}
}
Binary file not shown.
+92
View File
@@ -2,6 +2,7 @@
#![feature(float_algebraic)]
use std::collections::HashMap;
use std::collections::HashSet;
use std::ops::Div;
use std::sync::Arc;
use std::sync::mpsc::sync_channel;
@@ -68,6 +69,7 @@ use crate::voxel_cache::data::ColorBytes;
use crate::voxel_cache::data::DestinationElement;
use crate::voxel_cache::data::LocationPoolElement;
use crate::voxel_cache::data::StructurePoolElement;
use crate::voxel_cache::indirect_buffer::BufferCompactor;
use crate::voxel_cache::producer_interface::CacheProducerInterface;
use crate::voxel_cache::producer_interface::CacheRequest;
@@ -173,6 +175,96 @@ impl State
.await
.unwrap();
let mut encoder = device.create_command_encoder(&Default::default());
// --------------------------
let mut count_buffer = vec![0u32; 65536];
let mut check_buffer = vec![];
// Spread random numbers throughout the buffer
let mut picked = HashSet::new();
for i in 1..=1000
{
let mut r = rand::random_range(0..(count_buffer.len()));
while picked.contains(&r)
{
r = rand::random_range(0..(count_buffer.len()));
}
picked.insert(r);
check_buffer.push(r);
count_buffer[r] = i;
}
check_buffer.sort();
count_buffer
.iter()
.enumerate()
.filter(|(_, x)| **x != 0)
.take(100)
.for_each(|(x, y)| println!("{}, {}", x, y));
let buffer_compactor = BufferCompactor::new(&device, count_buffer.len());
let gpu_count_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: None,
contents: bytemuck::cast_slice(count_buffer.as_slice()),
usage: BufferUsages::STORAGE,
});
let gpu_index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: None,
contents: bytemuck::cast_slice(vec![0u32; count_buffer.len()].as_slice()),
usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC,
});
buffer_compactor.compact_buffer(
&device,
&mut encoder,
&gpu_count_buffer,
&gpu_index_buffer,
);
queue.submit([encoder.finish()]);
// --------------------------
wgpu::util::DownloadBuffer::read_buffer(
&device,
&queue,
&buffer_compactor.reduced_buffer.slice(..),
//&buffer_compactor.total_sum_buffer(),
|result| {
let result = result.unwrap();
let indices: Vec<u32> = result
.to_vec()
.chunks(4)
.map(|x| u32::from_ne_bytes([x[0], x[1], x[2], x[3]]))
.collect();
println!("Total sum {:?}", indices);
},
);
wgpu::util::DownloadBuffer::read_buffer(
&device,
&queue,
&gpu_index_buffer.slice(..),
|result| {
let result = result.unwrap();
let indices: Vec<u32> = result
.to_vec()
.chunks(4)
.map(|x| u32::from_ne_bytes([x[0], x[1], x[2], x[3]]))
.collect();
println!("{:?}", &indices[..1010]);
},
);
for _ in 0..16
{
device.poll(wgpu::wgt::PollType::Poll);
}
println!("Check buffer : {:?}", check_buffer);
panic!();
let size = window.inner_size();
let surface = instance.create_surface(window.clone()).unwrap();
+210 -14
View File
@@ -1,44 +1,240 @@
use wgpu::{Buffer, BufferUsages, Device, util::DeviceExt};
use wgpu::BindGroup;
use wgpu::BindGroupLayout;
use wgpu::Buffer;
use wgpu::BufferSlice;
use wgpu::BufferUsages;
use wgpu::CommandEncoder;
use wgpu::ComputePipeline;
use wgpu::Device;
use wgpu::ShaderStages;
use wgpu::util::DeviceExt;
pub struct BufferCompactor
pub struct NonZeroBufferCompactor
{
size: usize,
block_count: usize,
reduced_buffer: Buffer,
sum_buffer: Buffer,
compaction_buffer: Buffer,
count_buffer_bg_layout: BindGroupLayout,
compaction_bg: BindGroup,
reduce_sum_pipeline: ComputePipeline,
partial_sum_pipeline: ComputePipeline,
scatter_pipeline: ComputePipeline,
}
impl BufferCompactor
impl NonZeroBufferCompactor
{
const THREAD_COUNT: usize = 256;
pub fn new(device: &Device, size: usize) -> Self
{
let block_count = size.div_ceil(size);
let block_count = size.div_ceil(Self::THREAD_COUNT * 2);
let reduced_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("indirect_buffer_reduced"),
contents: bytemuck::cast_slice(vec![0; block_count].as_slice()),
usage: BufferUsages::STORAGE,
contents: bytemuck::cast_slice(vec![0; block_count + 1].as_slice()),
usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC,
});
let sum_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("indirect_buffer_sum"),
contents: bytemuck::cast_slice(vec![0; size].as_slice()),
contents: bytemuck::cast_slice(
vec![0; block_count * Self::THREAD_COUNT * 2].as_slice(),
),
usage: BufferUsages::STORAGE,
});
let compaction_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("indirect_buffer_compaction"),
contents: bytemuck::cast_slice(vec![0; size].as_slice()),
usage: BufferUsages::STORAGE,
let count_buffer_bg_layout =
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("count_buffer_bg_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 compaction_bg_layout =
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("count_buffer_bg_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 compaction_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("compaction_bg"),
layout: &compaction_bg_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: reduced_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: sum_buffer.as_entire_binding(),
},
],
});
BufferCompactor {
let compaction_pipeline_layout =
device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("compaction_pipeline_layout"),
bind_group_layouts: &[Some(&count_buffer_bg_layout), Some(&compaction_bg_layout)],
immediate_size: 0,
});
let compaction_shader_module = unsafe {
device.create_shader_module_trusted(
wgpu::ShaderModuleDescriptor {
label: Some("compaction_shader_layout"),
source: wgpu::ShaderSource::SpirV(wgpu::__macro_helpers::Cow::Borrowed(
wgpu::include_spirv_source!("../../shaders/compaction.spv"),
)),
},
wgpu::ShaderRuntimeChecks {
bounds_checks: false,
force_loop_bounding: false,
ray_query_initialization_tracking: false,
task_shader_dispatch_tracking: false,
mesh_shader_primitive_indices_clamp: false,
int_div_checks: false,
},
)
};
let reduce_sum_pipeline =
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("reduce_sum_pipeline"),
layout: Some(&compaction_pipeline_layout),
module: &compaction_shader_module,
entry_point: Some("block_sum"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
let partial_sum_pipeline =
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("partial_sum_pipeline "),
layout: Some(&compaction_pipeline_layout),
module: &compaction_shader_module,
entry_point: Some("linear_reduced_sum"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
let scatter_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("uniform_add_pipeline"),
layout: Some(&compaction_pipeline_layout),
module: &compaction_shader_module,
entry_point: Some("scatter"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
NonZeroBufferCompactor {
size,
block_count,
reduced_buffer,
sum_buffer,
compaction_buffer,
count_buffer_bg_layout,
compaction_bg,
reduce_sum_pipeline,
partial_sum_pipeline,
scatter_pipeline,
}
}
pub fn total_sum_buffer(&self) -> BufferSlice<'_>
{
self.reduced_buffer
.slice(((size_of::<u32>() * self.block_count) as u64)..)
}
pub fn copy_count_to_buffer(&self, encoder: &mut CommandEncoder, target: &Buffer)
{
encoder.copy_buffer_to_buffer(
&self.reduced_buffer,
(size_of::<u32>() * self.block_count) as u64,
target,
0,
size_of::<u32>() as u64,
);
}
pub fn compact_buffer(
&self,
device: &Device,
encoder: &mut CommandEncoder,
count_buffer: &Buffer,
index_buffer: &Buffer,
)
{
let count_buffer_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("count_buffer_bind_group"),
layout: &self.count_buffer_bg_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: count_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: index_buffer.as_entire_binding(),
},
],
});
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("buffer_compactor_compute_pass"),
timestamp_writes: None,
});
compute_pass.set_bind_group(0, &count_buffer_bind_group, &[]);
compute_pass.set_bind_group(1, &self.compaction_bg, &[]);
// Perform local running sum, then reduce
compute_pass.set_pipeline(&self.reduce_sum_pipeline);
compute_pass.dispatch_workgroups(self.block_count as u32, 1, 1);
// Partial sums are now in the reuced sum buffer. Perform linear naive runnig sum on it
compute_pass.set_pipeline(&self.partial_sum_pipeline);
compute_pass.dispatch_workgroups(1, 1, 1);
// We can now reapply the partial running sum on the subblocks
compute_pass.set_pipeline(&self.scatter_pipeline);
compute_pass.dispatch_workgroups(self.block_count as u32, 1, 1);
}
}