use glam::Vec3;
use thiserror::Error;
use wgpu::util::DeviceExt;
use super::transpiler::WgslShader;
use crate::types::SdfNode;
#[derive(Error, Debug)]
pub enum GpuError {
#[error("Failed to create GPU adapter")]
NoAdapter,
#[error("Failed to create GPU device: {0}")]
DeviceCreation(String),
#[error("Shader compilation error: {0}")]
ShaderCompilation(String),
#[error("Buffer mapping error: {0}")]
BufferMapping(String),
}
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct GpuInputPoint {
x: f32,
y: f32,
z: f32,
_pad: f32,
}
impl From<Vec3> for GpuInputPoint {
fn from(v: Vec3) -> Self {
Self {
x: v.x,
y: v.y,
z: v.z,
_pad: 0.0,
}
}
}
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct GpuOutputDistance {
distance: f32,
_pad1: f32,
_pad2: f32,
_pad3: f32,
}
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
struct GpuOutputFull {
dist: f32,
nx: f32,
ny: f32,
nz: f32,
}
const PARAM_BUFFER_SIZE: u64 = 1024 * 16;
pub struct GpuEvaluator {
device: wgpu::Device,
queue: wgpu::Queue,
pipeline: wgpu::ComputePipeline,
bind_group_layout: wgpu::BindGroupLayout,
param_buffer: Option<wgpu::Buffer>,
full_pipeline: Option<wgpu::ComputePipeline>,
dynamic_bind_group_layout: Option<wgpu::BindGroupLayout>,
workgroup_size: u32,
}
impl GpuEvaluator {
pub fn new(node: &SdfNode) -> Result<Self, GpuError> {
let shader = WgslShader::transpile(node, super::transpiler::TranspileMode::Hardcoded);
Self::from_shader(&shader)
}
pub fn from_shader(shader: &WgslShader) -> Result<Self, GpuError> {
let compute_shader = shader.to_compute_shader();
let mut evaluator = Self::from_wgsl(&compute_shader)?;
evaluator.workgroup_size = shader.workgroup_size;
Ok(evaluator)
}
pub fn from_wgsl(wgsl_source: &str) -> Result<Self, GpuError> {
Self::from_source(wgpu::ShaderSource::Wgsl(wgsl_source.into()))
}
#[cfg(feature = "glsl")]
pub fn from_glsl_compute(glsl_source: &str) -> Result<Self, GpuError> {
Self::from_source(wgpu::ShaderSource::Glsl {
shader: glsl_source.into(),
stage: wgpu::naga::ShaderStage::Compute,
defines: wgpu::naga::FastHashMap::default(),
})
}
fn from_source(source: wgpu::ShaderSource<'_>) -> Result<Self, GpuError> {
let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
backends: wgpu::Backends::all(),
..Default::default()
});
let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
compatible_surface: None,
force_fallback_adapter: false,
}))
.ok_or(GpuError::NoAdapter)?;
let (device, queue) = pollster::block_on(adapter.request_device(
&wgpu::DeviceDescriptor {
label: Some("ALICE-SDF Device"),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::default(),
memory_hints: wgpu::MemoryHints::Performance,
},
None,
))
.map_err(|e| GpuError::DeviceCreation(e.to_string()))?;
let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("ALICE-SDF Shader"),
source,
});
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("ALICE-SDF Bind Group Layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
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: 1,
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,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("ALICE-SDF Pipeline Layout"),
bind_group_layouts: &[&bind_group_layout],
push_constant_ranges: &[],
});
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("ALICE-SDF Compute Pipeline"),
layout: Some(&pipeline_layout),
module: &shader_module,
entry_point: Some("main"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
Ok(Self {
device,
queue,
pipeline,
bind_group_layout,
param_buffer: None,
full_pipeline: None,
dynamic_bind_group_layout: None,
workgroup_size: 256,
})
}
pub fn eval_batch(&self, points: &[Vec3]) -> Result<Vec<f32>, GpuError> {
if points.is_empty() {
return Ok(Vec::new());
}
let point_count = points.len();
let input_data: Vec<GpuInputPoint> = points.iter().map(|&p| p.into()).collect();
let input_buffer = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Input Points Buffer"),
contents: bytemuck::cast_slice(&input_data),
usage: wgpu::BufferUsages::STORAGE,
});
let output_buffer_size = (point_count * std::mem::size_of::<GpuOutputDistance>()) as u64;
let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Output Distances Buffer"),
size: output_buffer_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Staging Buffer"),
size: output_buffer_size,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let count_buffer = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Point Count Buffer"),
contents: bytemuck::cast_slice(&[point_count as u32]),
usage: wgpu::BufferUsages::UNIFORM,
});
let bind_group = self.distance_bind_group(
"ALICE-SDF Bind Group",
&input_buffer,
&output_buffer,
&count_buffer,
);
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("ALICE-SDF Command Encoder"),
});
{
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("ALICE-SDF Compute Pass"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&self.pipeline);
compute_pass.set_bind_group(0, &bind_group, &[]);
let wg = self.workgroup_size;
let workgroup_count = (point_count as u32).div_ceil(wg);
compute_pass.dispatch_workgroups(workgroup_count, 1, 1);
}
encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging_buffer, 0, output_buffer_size);
self.queue.submit(std::iter::once(encoder.finish()));
let buffer_slice = staging_buffer.slice(..);
let (sender, receiver) = std::sync::mpsc::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = sender.send(result);
});
self.device.poll(wgpu::Maintain::Wait);
receiver
.recv()
.map_err(|_| GpuError::BufferMapping("GPU channel closed unexpectedly".to_string()))?
.map_err(|e| GpuError::BufferMapping(format!("{:?}", e)))?;
let data = buffer_slice.get_mapped_range();
let output: &[GpuOutputDistance] = bytemuck::cast_slice(&data);
let distances: Vec<f32> = output.iter().map(|d| d.distance).collect();
drop(data);
staging_buffer.unmap();
Ok(distances)
}
fn distance_bind_group(
&self,
label: &str,
input: &wgpu::Buffer,
output: &wgpu::Buffer,
count: &wgpu::Buffer,
) -> wgpu::BindGroup {
let mut entries = vec![
wgpu::BindGroupEntry {
binding: 0,
resource: input.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: output.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: count.as_entire_binding(),
},
];
let layout = match (&self.dynamic_bind_group_layout, &self.param_buffer) {
(Some(layout), Some(params)) => {
entries.push(wgpu::BindGroupEntry {
binding: 3,
resource: params.as_entire_binding(),
});
layout
}
_ => &self.bind_group_layout,
};
self.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(label),
layout,
entries: &entries,
})
}
pub fn device_info(&self) -> String {
"WebGPU Device (via wgpu)".to_string()
}
pub fn create_buffer_pool(&self, initial_capacity: usize) -> GpuBufferPool {
GpuBufferPool::new(&self.device, initial_capacity)
}
pub fn eval_batch_pooled(
&self,
points: &[Vec3],
pool: &mut GpuBufferPool,
) -> Result<Vec<f32>, GpuError> {
if points.is_empty() {
return Ok(Vec::new());
}
let point_count = points.len();
pool.ensure_capacity(&self.device, point_count);
let input_data: Vec<GpuInputPoint> = points.iter().map(|&p| p.into()).collect();
self.queue
.write_buffer(&pool.input_buffer, 0, bytemuck::cast_slice(&input_data));
self.queue.write_buffer(
&pool.count_buffer,
0,
bytemuck::cast_slice(&[point_count as u32]),
);
let bind_group = self.distance_bind_group(
"ALICE-SDF Pooled Bind Group",
&pool.input_buffer,
&pool.output_buffer,
&pool.count_buffer,
);
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("ALICE-SDF Pooled Encoder"),
});
{
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("ALICE-SDF Pooled Pass"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&self.pipeline);
compute_pass.set_bind_group(0, &bind_group, &[]);
let workgroup_count = (point_count as u32).div_ceil(self.workgroup_size);
compute_pass.dispatch_workgroups(workgroup_count, 1, 1);
}
let read_size = (point_count * std::mem::size_of::<GpuOutputDistance>()) as u64;
encoder.copy_buffer_to_buffer(&pool.output_buffer, 0, &pool.staging_buffer, 0, read_size);
self.queue.submit(std::iter::once(encoder.finish()));
let buffer_slice = pool.staging_buffer.slice(..read_size);
let (sender, receiver) = std::sync::mpsc::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = sender.send(result);
});
self.device.poll(wgpu::Maintain::Wait);
receiver
.recv()
.map_err(|_| GpuError::BufferMapping("GPU channel closed unexpectedly".to_string()))?
.map_err(|e| GpuError::BufferMapping(format!("{:?}", e)))?;
let data = buffer_slice.get_mapped_range();
let output: &[GpuOutputDistance] = bytemuck::cast_slice(&data);
let distances: Vec<f32> = output[..point_count].iter().map(|d| d.distance).collect();
drop(data);
pool.staging_buffer.unmap();
Ok(distances)
}
pub fn eval_batch_auto(
&self,
points: &[Vec3],
pool: &mut GpuBufferPool,
) -> Result<Vec<f32>, GpuError> {
if points.is_empty() {
return Ok(Vec::new());
}
if points.len() <= AUTO_TUNE_THRESHOLD {
return self.eval_batch_pooled(points, pool);
}
let chunk_size = optimal_chunk_size(points.len());
let mut all_distances = Vec::with_capacity(points.len());
for chunk in points.chunks(chunk_size) {
let distances = self.eval_batch_pooled(chunk, pool)?;
all_distances.extend_from_slice(&distances);
}
Ok(all_distances)
}
pub async fn new_async(node: &SdfNode) -> Result<Self, GpuError> {
let shader = WgslShader::transpile(node, super::transpiler::TranspileMode::Hardcoded);
Self::from_shader_async(&shader).await
}
pub async fn from_shader_async(shader: &WgslShader) -> Result<Self, GpuError> {
let compute_shader = shader.to_compute_shader();
let mut evaluator = Self::from_wgsl_async(&compute_shader).await?;
evaluator.workgroup_size = shader.workgroup_size;
Ok(evaluator)
}
pub async fn from_wgsl_async(wgsl_source: &str) -> Result<Self, GpuError> {
let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
backends: wgpu::Backends::all(),
..Default::default()
});
let adapter = instance
.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
compatible_surface: None,
force_fallback_adapter: false,
})
.await
.ok_or(GpuError::NoAdapter)?;
let (device, queue) = adapter
.request_device(
&wgpu::DeviceDescriptor {
label: Some("ALICE-SDF Device"),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::default(),
memory_hints: wgpu::MemoryHints::Performance,
},
None,
)
.await
.map_err(|e| GpuError::DeviceCreation(e.to_string()))?;
let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("ALICE-SDF Shader"),
source: wgpu::ShaderSource::Wgsl(wgsl_source.into()),
});
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("ALICE-SDF Bind Group Layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
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: 1,
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,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("ALICE-SDF Pipeline Layout"),
bind_group_layouts: &[&bind_group_layout],
push_constant_ranges: &[],
});
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("ALICE-SDF Compute Pipeline"),
layout: Some(&pipeline_layout),
module: &shader_module,
entry_point: Some("main"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
Ok(Self {
device,
queue,
pipeline,
bind_group_layout,
param_buffer: None,
full_pipeline: None,
dynamic_bind_group_layout: None,
workgroup_size: 256,
})
}
pub async fn eval_batch_async(&self, points: &[Vec3]) -> Result<Vec<f32>, GpuError> {
if points.is_empty() {
return Ok(Vec::new());
}
let point_count = points.len();
let input_data: Vec<GpuInputPoint> = points.iter().map(|&p| p.into()).collect();
let input_buffer = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Input Points Buffer"),
contents: bytemuck::cast_slice(&input_data),
usage: wgpu::BufferUsages::STORAGE,
});
let output_buffer_size = (point_count * std::mem::size_of::<GpuOutputDistance>()) as u64;
let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Output Distances Buffer"),
size: output_buffer_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Staging Buffer"),
size: output_buffer_size,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let count_buffer = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Point Count Buffer"),
contents: bytemuck::cast_slice(&[point_count as u32]),
usage: wgpu::BufferUsages::UNIFORM,
});
let bind_group = self.distance_bind_group(
"ALICE-SDF Bind Group",
&input_buffer,
&output_buffer,
&count_buffer,
);
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("ALICE-SDF Command Encoder"),
});
{
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("ALICE-SDF Compute Pass"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&self.pipeline);
compute_pass.set_bind_group(0, &bind_group, &[]);
let workgroup_count = (point_count as u32).div_ceil(self.workgroup_size);
compute_pass.dispatch_workgroups(workgroup_count, 1, 1);
}
encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging_buffer, 0, output_buffer_size);
self.queue.submit(std::iter::once(encoder.finish()));
let buffer_slice = staging_buffer.slice(..);
let (sender, receiver) = futures_channel::oneshot::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = sender.send(result);
});
self.device.poll(wgpu::Maintain::Wait);
receiver
.await
.map_err(|_| GpuError::BufferMapping("Channel closed".to_string()))?
.map_err(|e| GpuError::BufferMapping(format!("{:?}", e)))?;
let data = buffer_slice.get_mapped_range();
let output: &[GpuOutputDistance] = bytemuck::cast_slice(&data);
let distances: Vec<f32> = output.iter().map(|d| d.distance).collect();
drop(data);
staging_buffer.unmap();
Ok(distances)
}
pub fn new_dynamic(node: &SdfNode) -> Result<Self, GpuError> {
use super::transpiler::TranspileMode;
let shader = WgslShader::transpile(node, TranspileMode::Dynamic);
let compute_source = shader.to_compute_shader();
let full_source = shader.to_compute_shader_with_normals();
let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
backends: wgpu::Backends::all(),
..Default::default()
});
let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
compatible_surface: None,
force_fallback_adapter: false,
}))
.ok_or(GpuError::NoAdapter)?;
let (device, queue) = pollster::block_on(adapter.request_device(
&wgpu::DeviceDescriptor {
label: Some("ALICE-SDF Device (Dynamic)"),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::default(),
memory_hints: wgpu::MemoryHints::Performance,
},
None,
))
.map_err(|e| GpuError::DeviceCreation(e.to_string()))?;
let dynamic_bind_group_layout =
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("ALICE-SDF Dynamic Bind Group Layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
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: 1,
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,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
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: 3,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("ALICE-SDF Dynamic Pipeline Layout"),
bind_group_layouts: &[&dynamic_bind_group_layout],
push_constant_ranges: &[],
});
let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("ALICE-SDF Dynamic Shader"),
source: wgpu::ShaderSource::Wgsl(compute_source.into()),
});
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("ALICE-SDF Dynamic Pipeline"),
layout: Some(&pipeline_layout),
module: &shader_module,
entry_point: Some("main"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
let full_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("ALICE-SDF Full Shader"),
source: wgpu::ShaderSource::Wgsl(full_source.into()),
});
let full_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("ALICE-SDF Full Pipeline"),
layout: Some(&pipeline_layout),
module: &full_module,
entry_point: Some("main"),
compilation_options: wgpu::PipelineCompilationOptions::default(),
cache: None,
});
let mut param_data = shader.param_layout.clone();
while param_data.len() % 4 != 0 {
param_data.push(0.0);
}
param_data.resize(PARAM_BUFFER_SIZE as usize / 4, 0.0);
let param_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("SDF Params Buffer"),
contents: bytemuck::cast_slice(¶m_data),
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
});
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("ALICE-SDF Legacy Bind Group Layout"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
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: 1,
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,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
Ok(Self {
device,
queue,
pipeline,
bind_group_layout,
param_buffer: Some(param_buffer),
full_pipeline: Some(full_pipeline),
dynamic_bind_group_layout: Some(dynamic_bind_group_layout),
workgroup_size: shader.workgroup_size,
})
}
pub fn update_params(&self, node: &SdfNode) {
let params = WgslShader::extract_params(node);
let mut padded = params;
while padded.len() % 4 != 0 {
padded.push(0.0);
}
padded.resize(PARAM_BUFFER_SIZE as usize / 4, 0.0);
self.queue.write_buffer(
self.param_buffer
.as_ref()
.expect("update_params requires Dynamic mode (use new_dynamic)"),
0,
bytemuck::cast_slice(&padded),
);
}
pub fn eval_batch_full(&self, points: &[Vec3]) -> Result<Vec<(f32, Vec3)>, GpuError> {
if points.is_empty() {
return Ok(Vec::new());
}
let full_pipeline = self
.full_pipeline
.as_ref()
.expect("eval_batch_full requires Dynamic mode (use new_dynamic)");
let dyn_layout = self
.dynamic_bind_group_layout
.as_ref()
.expect("eval_batch_full requires Dynamic mode (use new_dynamic)");
let param_buf = self
.param_buffer
.as_ref()
.expect("eval_batch_full requires Dynamic mode (use new_dynamic)");
let point_count = points.len();
let input_data: Vec<GpuInputPoint> = points.iter().map(|&p| p.into()).collect();
let input_buffer = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Input Points Buffer (Full)"),
contents: bytemuck::cast_slice(&input_data),
usage: wgpu::BufferUsages::STORAGE,
});
let output_buffer_size = (point_count * std::mem::size_of::<GpuOutputFull>()) as u64;
let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Output Full Buffer"),
size: output_buffer_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Staging Buffer (Full)"),
size: output_buffer_size,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let count_buffer = self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Point Count Buffer"),
contents: bytemuck::cast_slice(&[point_count as u32]),
usage: wgpu::BufferUsages::UNIFORM,
});
let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("ALICE-SDF Full Bind Group"),
layout: dyn_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: input_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: output_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: count_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: param_buf.as_entire_binding(),
},
],
});
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("ALICE-SDF Full Encoder"),
});
{
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("ALICE-SDF Full Pass"),
timestamp_writes: None,
});
compute_pass.set_pipeline(full_pipeline);
compute_pass.set_bind_group(0, &bind_group, &[]);
let workgroup_count = (point_count as u32).div_ceil(self.workgroup_size);
compute_pass.dispatch_workgroups(workgroup_count, 1, 1);
}
encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging_buffer, 0, output_buffer_size);
self.queue.submit(std::iter::once(encoder.finish()));
let buffer_slice = staging_buffer.slice(..);
let (sender, receiver) = std::sync::mpsc::channel();
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let _ = sender.send(result);
});
self.device.poll(wgpu::Maintain::Wait);
receiver
.recv()
.map_err(|_| GpuError::BufferMapping("GPU channel closed unexpectedly".to_string()))?
.map_err(|e| GpuError::BufferMapping(format!("{:?}", e)))?;
let data = buffer_slice.get_mapped_range();
let output: &[GpuOutputFull] = bytemuck::cast_slice(&data);
let results: Vec<(f32, Vec3)> = output
.iter()
.map(|o| (o.dist, Vec3::new(o.nx, o.ny, o.nz)))
.collect();
drop(data);
staging_buffer.unmap();
Ok(results)
}
pub const fn eval_batch_submit(&self, points: Vec<Vec3>) -> GpuEvalFuture<'_> {
GpuEvalFuture::new(self, points)
}
}
const AUTO_TUNE_THRESHOLD: usize = 262144;
const fn optimal_chunk_size(total_points: usize) -> usize {
if total_points <= 262144 {
total_points } else if total_points <= 1_048_576 {
262144 } else {
524288 }
}
pub struct GpuBufferPool {
pub capacity: usize,
input_buffer: wgpu::Buffer,
output_buffer: wgpu::Buffer,
staging_buffer: wgpu::Buffer,
count_buffer: wgpu::Buffer,
}
impl GpuBufferPool {
fn new(device: &wgpu::Device, capacity: usize) -> Self {
let cap = capacity.max(256); let input_size = (cap * std::mem::size_of::<GpuInputPoint>()) as u64;
let output_size = (cap * std::mem::size_of::<GpuOutputDistance>()) as u64;
let input_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Pool Input Buffer"),
size: input_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Pool Output Buffer"),
size: output_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Pool Staging Buffer"),
size: output_size,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let count_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Pool Count Buffer"),
size: 4,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
Self {
capacity: cap,
input_buffer,
output_buffer,
staging_buffer,
count_buffer,
}
}
fn ensure_capacity(&mut self, device: &wgpu::Device, required: usize) {
if required <= self.capacity {
return;
}
let new_cap = (self.capacity * 3 / 2).max(required);
*self = Self::new(device, new_cap);
}
}
impl std::fmt::Debug for GpuBufferPool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GpuBufferPool")
.field("capacity", &self.capacity)
.field(
"input_bytes",
&(self.capacity * std::mem::size_of::<GpuInputPoint>()),
)
.field(
"output_bytes",
&(self.capacity * std::mem::size_of::<GpuOutputDistance>()),
)
.finish_non_exhaustive()
}
}
#[must_use = "the batch is evaluated only when the handle is consumed by `wait` or `resolve`"]
pub struct GpuEvalFuture<'a> {
evaluator: &'a GpuEvaluator,
points: Vec<Vec3>,
}
impl<'a> GpuEvalFuture<'a> {
const fn new(evaluator: &'a GpuEvaluator, points: Vec<Vec3>) -> Self {
Self { evaluator, points }
}
pub fn wait(self) -> Result<Vec<f32>, GpuError> {
self.evaluator.eval_batch(&self.points)
}
pub async fn resolve(self) -> Result<Vec<f32>, GpuError> {
self.evaluator.eval_batch_async(&self.points).await
}
}
impl std::fmt::Debug for GpuEvalFuture<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GpuEvalFuture")
.field("points", &self.points.len())
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for GpuEvaluator {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GpuEvaluator")
.field("device", &self.device_info())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::eval::eval;
fn has_gpu() -> bool {
let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))
.is_some()
}
#[test]
fn test_gpu_eval_sphere() {
if !has_gpu() {
eprintln!("Skipping GPU test: no GPU available");
return;
}
let sphere = SdfNode::Sphere { radius: 1.0 };
let gpu = GpuEvaluator::new(&sphere).unwrap();
let points = vec![
Vec3::ZERO,
Vec3::new(1.0, 0.0, 0.0),
Vec3::new(2.0, 0.0, 0.0),
];
let gpu_results = gpu.eval_batch(&points).unwrap();
for (i, p) in points.iter().enumerate() {
let cpu_result = eval(&sphere, *p);
assert!(
(gpu_results[i] - cpu_result).abs() < 0.001,
"Mismatch at {:?}: gpu={}, cpu={}",
p,
gpu_results[i],
cpu_result
);
}
}
#[test]
fn test_gpu_eval_complex() {
if !has_gpu() {
eprintln!("Skipping GPU test: no GPU available");
return;
}
let shape = SdfNode::Sphere { radius: 1.0 }.union(
SdfNode::Box3d {
half_extents: Vec3::splat(0.5),
}
.translate(2.0, 0.0, 0.0),
);
let gpu = GpuEvaluator::new(&shape).unwrap();
let points: Vec<Vec3> = (0..100)
.map(|i| Vec3::new((i as f32).mul_add(0.05, -2.5), 0.0, 0.0))
.collect();
let gpu_results = gpu.eval_batch(&points).unwrap();
for (i, p) in points.iter().enumerate() {
let cpu_result = eval(&shape, *p);
assert!(
(gpu_results[i] - cpu_result).abs() < 0.01,
"Mismatch at {:?}: gpu={}, cpu={}",
p,
gpu_results[i],
cpu_result
);
}
}
#[test]
fn test_gpu_eval_large_batch() {
if !has_gpu() {
eprintln!("Skipping GPU test: no GPU available");
return;
}
let sphere = SdfNode::Sphere { radius: 1.0 };
let gpu = GpuEvaluator::new(&sphere).unwrap();
let points: Vec<Vec3> = (0..10000)
.map(|i| {
let t = i as f32 / 10000.0;
Vec3::new(t * 4.0 - 2.0, (t * 10.0).sin(), (t * 10.0).cos())
})
.collect();
let gpu_results = gpu.eval_batch(&points).unwrap();
assert_eq!(gpu_results.len(), 10000);
for i in [0, 1000, 5000, 9999] {
let cpu_result = eval(&sphere, points[i]);
assert!(
(gpu_results[i] - cpu_result).abs() < 0.01,
"Mismatch at index {}: gpu={}, cpu={}",
i,
gpu_results[i],
cpu_result
);
}
}
}