use crate::ecs::sdf::brick_map::{BRICK_CORNER_COUNT, BRICK_SIZE, BRICKS_PER_ATLAS_DIM};
use crate::ecs::sdf::primitives::{CsgOperation, SdfEdit, SdfPrimitive};
use crate::ecs::sdf::resources::TerrainConfig;
const MAX_EDITS: usize = 1024;
const MAX_DIRTY_BRICKS_PER_FRAME: usize = 4096;
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct ComputeUniforms {
brick_count: u32,
terrain_enabled: u32,
terrain_base_height: f32,
terrain_material_id: u32,
edit_count: u32,
smoothness_scale: f32,
terrain_seed: u32,
terrain_octaves: u32,
terrain_frequency: f32,
terrain_amplitude: f32,
terrain_gain: f32,
_pad0: u32,
_pad1: u32,
_pad2: u32,
_pad3: u32,
_pad4: u32,
}
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct GpuSdfEdit {
inverse_transform: [[f32; 4]; 4],
primitive_type: u32,
operation_type: u32,
material_id: u32,
smoothness: f32,
uniform_scale: f32,
param0: f32,
param1: f32,
param2: f32,
param3: f32,
_pad0: u32,
_pad1: u32,
_pad2: u32,
}
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct GpuDirtyBrick {
pub world_origin: [f32; 3],
pub voxel_size: f32,
pub atlas_slot: u32,
pub material_atlas_x: u32,
pub material_atlas_y: u32,
pub material_atlas_z: u32,
}
pub struct SdfComputePipelines {
corners_pipeline: wgpu::ComputePipeline,
materials_pipeline: wgpu::ComputePipeline,
bind_group_layout: wgpu::BindGroupLayout,
uniforms_buffer: wgpu::Buffer,
edits_buffer: wgpu::Buffer,
dirty_bricks_buffer: wgpu::Buffer,
surface_flags_buffer: wgpu::Buffer,
bind_group: Option<wgpu::BindGroup>,
last_brick_count: u32,
}
impl SdfComputePipelines {
pub fn new(device: &wgpu::Device) -> Self {
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("SDF Evaluate Compute Shader"),
source: wgpu::ShaderSource::Wgsl(
include_str!("../../shaders/sdf_evaluate.wgsl").into(),
),
});
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("SDF Compute Bind Group Layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 3,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::StorageTexture {
access: wgpu::StorageTextureAccess::WriteOnly,
format: wgpu::TextureFormat::R32Float,
view_dimension: wgpu::TextureViewDimension::D3,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 4,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::StorageTexture {
access: wgpu::StorageTextureAccess::WriteOnly,
format: wgpu::TextureFormat::R32Uint,
view_dimension: wgpu::TextureViewDimension::D3,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 5,
visibility: wgpu::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("SDF Compute Pipeline Layout"),
bind_group_layouts: &[Some(&bind_group_layout)],
immediate_size: 0,
});
let corners_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("SDF Evaluate Corners Pipeline"),
layout: Some(&pipeline_layout),
module: &shader,
entry_point: Some("evaluate_corners"),
compilation_options: Default::default(),
cache: None,
});
let materials_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("SDF Evaluate Materials Pipeline"),
layout: Some(&pipeline_layout),
module: &shader,
entry_point: Some("evaluate_materials"),
compilation_options: Default::default(),
cache: None,
});
let uniforms_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("SDF Compute Uniforms"),
size: std::mem::size_of::<ComputeUniforms>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let edits_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("SDF Compute Edits"),
size: (std::mem::size_of::<GpuSdfEdit>() * MAX_EDITS) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let dirty_bricks_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("SDF Compute Dirty Bricks"),
size: (std::mem::size_of::<GpuDirtyBrick>() * MAX_DIRTY_BRICKS_PER_FRAME) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let surface_flags_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("SDF Compute Surface Flags"),
size: (std::mem::size_of::<u32>() * MAX_DIRTY_BRICKS_PER_FRAME) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
Self {
corners_pipeline,
materials_pipeline,
bind_group_layout,
uniforms_buffer,
edits_buffer,
dirty_bricks_buffer,
surface_flags_buffer,
bind_group: None,
last_brick_count: 0,
}
}
pub fn prepare(
&mut self,
queue: &wgpu::Queue,
edits: &[SdfEdit],
dirty_bricks: &[GpuDirtyBrick],
terrain: &TerrainConfig,
smoothness_scale: f32,
) {
let brick_count = dirty_bricks.len().min(MAX_DIRTY_BRICKS_PER_FRAME);
self.last_brick_count = brick_count as u32;
if brick_count == 0 {
return;
}
let edit_count = edits.len().min(MAX_EDITS);
let uniforms = ComputeUniforms {
brick_count: brick_count as u32,
terrain_enabled: if terrain.enabled { 1 } else { 0 },
terrain_base_height: terrain.base_height,
terrain_material_id: terrain.material_id,
edit_count: edit_count as u32,
smoothness_scale,
terrain_seed: terrain.seed,
terrain_octaves: terrain.octaves,
terrain_frequency: terrain.frequency,
terrain_amplitude: terrain.amplitude,
terrain_gain: terrain.gain,
_pad0: 0,
_pad1: 0,
_pad2: 0,
_pad3: 0,
_pad4: 0,
};
queue.write_buffer(&self.uniforms_buffer, 0, bytemuck::cast_slice(&[uniforms]));
let gpu_edits: Vec<GpuSdfEdit> =
edits.iter().take(edit_count).map(serialize_edit).collect();
if !gpu_edits.is_empty() {
queue.write_buffer(&self.edits_buffer, 0, bytemuck::cast_slice(&gpu_edits));
}
queue.write_buffer(
&self.dirty_bricks_buffer,
0,
bytemuck::cast_slice(&dirty_bricks[..brick_count]),
);
let zero_flags = vec![0u32; brick_count];
queue.write_buffer(
&self.surface_flags_buffer,
0,
bytemuck::cast_slice(&zero_flags),
);
}
pub fn create_bind_group(
&mut self,
device: &wgpu::Device,
brick_atlas_view: &wgpu::TextureView,
material_atlas_view: &wgpu::TextureView,
) {
self.bind_group = Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("SDF Compute Bind Group"),
layout: &self.bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: self.uniforms_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: self.edits_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: self.dirty_bricks_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: wgpu::BindingResource::TextureView(brick_atlas_view),
},
wgpu::BindGroupEntry {
binding: 4,
resource: wgpu::BindingResource::TextureView(material_atlas_view),
},
wgpu::BindGroupEntry {
binding: 5,
resource: self.surface_flags_buffer.as_entire_binding(),
},
],
}));
}
pub fn dispatch(&self, encoder: &mut wgpu::CommandEncoder) {
if self.last_brick_count == 0 {
return;
}
let bind_group = match &self.bind_group {
Some(bind_group) => bind_group,
None => return,
};
let total_corners = self.last_brick_count
* (BRICK_CORNER_COUNT * BRICK_CORNER_COUNT * BRICK_CORNER_COUNT) as u32;
let corner_workgroups = total_corners.div_ceil(128);
{
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("SDF Evaluate Corners"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&self.corners_pipeline);
compute_pass.set_bind_group(0, bind_group, &[]);
compute_pass.dispatch_workgroups(corner_workgroups, 1, 1);
}
let total_voxels = self.last_brick_count * (BRICK_SIZE * BRICK_SIZE * BRICK_SIZE) as u32;
let material_workgroups = total_voxels.div_ceil(128);
{
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("SDF Evaluate Materials"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&self.materials_pipeline);
compute_pass.set_bind_group(0, bind_group, &[]);
compute_pass.dispatch_workgroups(material_workgroups, 1, 1);
}
}
pub fn reset_frame(&mut self) {
self.last_brick_count = 0;
}
pub fn has_pending_work(&self) -> bool {
self.last_brick_count > 0
}
}
fn serialize_edit(edit: &SdfEdit) -> GpuSdfEdit {
let inverse_transform: [[f32; 4]; 4] = (*edit.inverse_transform()).into();
let (primitive_type, param0, param1, param2, param3) = match *edit.primitive() {
SdfPrimitive::Sphere { radius } => (0u32, radius, 0.0, 0.0, 0.0),
SdfPrimitive::Box { half_extents } => {
(1u32, half_extents.x, half_extents.y, half_extents.z, 0.0)
}
SdfPrimitive::Cylinder {
radius,
half_height,
} => (2u32, radius, half_height, 0.0, 0.0),
SdfPrimitive::Torus {
major_radius,
minor_radius,
} => (3u32, major_radius, minor_radius, 0.0, 0.0),
SdfPrimitive::Capsule {
radius,
half_height,
} => (4u32, radius, half_height, 0.0, 0.0),
SdfPrimitive::Plane { normal, offset } => (5u32, normal.x, normal.y, normal.z, offset),
};
let operation_type = match edit.operation() {
CsgOperation::Union => 0u32,
CsgOperation::Subtraction => 1u32,
CsgOperation::Intersection => 2u32,
CsgOperation::SmoothUnion { .. } => 3u32,
CsgOperation::SmoothSubtraction { .. } => 4u32,
CsgOperation::SmoothIntersection { .. } => 5u32,
};
GpuSdfEdit {
inverse_transform,
primitive_type,
operation_type,
material_id: edit.material_id(),
smoothness: edit.smoothness(),
uniform_scale: edit.uniform_scale(),
param0,
param1,
param2,
param3,
_pad0: 0,
_pad1: 0,
_pad2: 0,
}
}
pub fn prepare_dirty_brick_metadata(
atlas_slot: u32,
world_origin: nalgebra_glm::Vec3,
voxel_size: f32,
) -> GpuDirtyBrick {
let bricks_per_slice = BRICKS_PER_ATLAS_DIM * BRICKS_PER_ATLAS_DIM;
let brick_z = atlas_slot as usize / bricks_per_slice;
let remainder = atlas_slot as usize % bricks_per_slice;
let brick_y = remainder / BRICKS_PER_ATLAS_DIM;
let brick_x = remainder % BRICKS_PER_ATLAS_DIM;
GpuDirtyBrick {
world_origin: [world_origin.x, world_origin.y, world_origin.z],
voxel_size,
atlas_slot,
material_atlas_x: (brick_x * BRICK_SIZE) as u32,
material_atlas_y: (brick_y * BRICK_SIZE) as u32,
material_atlas_z: (brick_z * BRICK_SIZE) as u32,
}
}