Something went wrong. Try again.
personal knowledge base
djot markdown knowledge wiki
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417use super::SimulationParameters;use super::{Edge, Point, Vector2f};
pub struct Simulation<'a> { device: &'a wgpu::Device, queue: &'a wgpu::Queue, bindings: Vec<Binding>, points_count: u32, edges_count: u32, points_mapping: wgpu::Buffer, pipeline: wgpu::ComputePipeline,}
struct Buffer { buffer: wgpu::Buffer, visibility: wgpu::ShaderStages, type_: wgpu::BindingType,}
impl Buffer { fn new( device: &wgpu::Device, label: Option<&str>, size: u64, usage: wgpu::BufferUsages, visibility: wgpu::ShaderStages, type_: wgpu::BindingType, ) -> Self { let buffer = device.create_buffer(&wgpu::BufferDescriptor { label, size, usage, mapped_at_creation: false, }); Self { buffer, visibility, type_, } }}
struct BindingBuilder<'a> { device: &'a wgpu::Device, queue: &'a wgpu::Queue, buffers: Vec<Buffer>, label: Option<&'a str>, id: u32,}
impl<'a> BindingBuilder<'a> { pub fn new( device: &'a wgpu::Device, queue: &'a wgpu::Queue, label: Option<&'a str>, id: u32, ) -> Self { Self { device, queue, buffers: Vec::new(), label, id, } }
pub fn add_buffer( &mut self, label: Option<&str>, size: u64, usage: wgpu::BufferUsages, visibility: wgpu::ShaderStages, type_: wgpu::BindingType, data: Option<&[u8]>, ) -> &mut Self { let buffer = Buffer::new(self.device, label, size, usage, visibility, type_); if let Some(data) = data { self.queue.write_buffer(&buffer.buffer, 0, data); } self.buffers.push(buffer); self }
fn layout(&self) -> wgpu::BindGroupLayout { self.device .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: self.label, entries: &self .buffers .iter() .enumerate() .map(|(i, buffer)| wgpu::BindGroupLayoutEntry { binding: i as u32, visibility: buffer.visibility, ty: buffer.type_.clone(), count: None, }) .collect::<Vec<_>>(), }) }
fn bind_group(&self, layout: &wgpu::BindGroupLayout) -> wgpu::BindGroup { self.device.create_bind_group(&wgpu::BindGroupDescriptor { label: self.label, layout, entries: &self .buffers .iter() .enumerate() .map(|(i, buffer)| wgpu::BindGroupEntry { binding: i as u32, resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { buffer: &buffer.buffer, offset: 0, size: None, }), }) .collect::<Vec<_>>(), }) }
fn build(&self) -> Binding { let layout = self.layout(); Binding { id: self.id, bind_group: self.bind_group(&layout), layout, } }}
struct Binding { id: u32, bind_group: wgpu::BindGroup, layout: wgpu::BindGroupLayout,}
impl Binding { fn layout(&self) -> &wgpu::BindGroupLayout { &self.layout } fn set(&self, pass: &mut wgpu::ComputePass) { pass.set_bind_group(self.id, &self.bind_group, &[]); }}
impl<'a> Simulation<'a> { fn create_binding_0( device: &'a wgpu::Device, queue: &'a wgpu::Queue, points: &[Point], ) -> Binding { BindingBuilder::new(device, queue, Some("Simulation Bind Group 0"), 0) .add_buffer( Some("Points Buffer"), (points.len() * std::mem::size_of::<Point>()) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: false }, has_dynamic_offset: false, min_binding_size: None, }, Some(bytemuck::cast_slice(points)), ) .build() }
fn create_binding_1( device: &'a wgpu::Device, queue: &'a wgpu::Queue, params: SimulationParameters, ) -> Binding { BindingBuilder::new(device, queue, Some("Simulation Bind Group 1"), 1) .add_buffer( Some("Params Buffer"), std::mem::size_of::<SimulationParameters>() as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::UNIFORM, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None, }, Some(bytemuck::cast_slice(&[params])), ) .build() }
fn create_binging_2( device: &'a wgpu::Device, queue: &'a wgpu::Queue, points: &[Point], edges: &[Edge], ) -> Binding { let mut edges_start_p1 = Vec::with_capacity(points.len()); let mut edges_start_p2 = Vec::with_capacity(points.len()); let mut edges_p1 = Vec::with_capacity(edges.len()); let mut edges_p2 = Vec::with_capacity(edges.len()); for i in 0..points.len() { edges_start_p1.push(edges_p1.len() as u32); edges_start_p2.push(edges_p2.len() as u32);
for j in 0..edges.len() { if edges[j].start == i as u32 { edges_p1.push(edges[j].end); } else if edges[j].end == i as u32 { edges_p2.push(edges[j].start); } } } BindingBuilder::new(device, queue, Some("Simulation Bind Group 2"), 2) .add_buffer( Some("Edges Start P1 Buffer"), (edges_start_p1.len() * std::mem::size_of::<u32>()) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, }, Some(bytemuck::cast_slice(&edges_start_p1)), ) .add_buffer( Some("Edges Start P2 Buffer"), (edges_start_p2.len() * std::mem::size_of::<u32>()) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, }, Some(bytemuck::cast_slice(&edges_start_p2)), ) .add_buffer( Some("Edges P1 Buffer"), (edges_p1.len() * std::mem::size_of::<u32>()) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, }, Some(bytemuck::cast_slice(&edges_p1)), ) .add_buffer( Some("Edges P2 Buffer"), (edges_p2.len() * std::mem::size_of::<u32>()) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, }, Some(bytemuck::cast_slice(&edges_p2)), ) .build() }
fn create_binding3( device: &'a wgpu::Device, queue: &'a wgpu::Queue, points_count: u32, ) -> Binding { BindingBuilder::new(device, queue, Some("Simulation Bind Group 3"), 3) .add_buffer( Some("Forces Buffer"), (points_count * std::mem::size_of::<Vector2f>() as u32) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: false }, has_dynamic_offset: false, min_binding_size: None, }, None, ) .add_buffer( Some("Velocities Buffer"), (points_count * std::mem::size_of::<Vector2f>() as u32) as u64, wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::STORAGE, wgpu::ShaderStages::COMPUTE, wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: false }, has_dynamic_offset: false, min_binding_size: None, }, None, ) .build() }
pub async fn new( device: &'a wgpu::Device, queue: &'a wgpu::Queue, points: &[Point], edges: &[Edge], params: SimulationParameters, ) -> Self { let points_mapping = device.create_buffer(&wgpu::BufferDescriptor { label: Some("Points Mapping Buffer"), size: (points.len() * std::mem::size_of::<Point>()) as u64, usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false, });
assert!(points.len() + 63 / 64 < u16::MAX as usize); let simulation_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("Simulation Shader"), source: wgpu::ShaderSource::Wgsl(include_str!("simulation.wgsl").into()), }); let mut bindings = vec![]; bindings.push(Self::create_binding_0(device, queue, points)); bindings.push(Self::create_binding_1(device, queue, params)); bindings.push(Self::create_binging_2(device, queue, points, edges)); bindings.push(Self::create_binding3(device, queue, points.len() as u32));
let layouts = bindings .iter() .map(|binding| binding.layout()) .collect::<Vec<_>>(); let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("Simulation Pipeline Layout"), bind_group_layouts: &layouts, push_constant_ranges: &[], }); let step = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some("Simulation Step"), layout: Some(&pipeline_layout), module: &simulation_shader, entry_point: Some("step"), cache: None, compilation_options: wgpu::PipelineCompilationOptions::default(), });
device.poll(wgpu::Maintain::Wait).panic_on_timeout();
Self { device, queue, points_count: points.len() as u32, edges_count: edges.len() as u32, points_mapping, bindings, pipeline: step, } }}
impl<'a> super::Simulation for Simulation<'a> { fn step(&self) { let mut encoder = self .device .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("Simulation Step Encoder"), }); { let desc = wgpu::ComputePassDescriptor { label: Some("Simulation Step Pass"), timestamp_writes: None, }; let mut pass = encoder.begin_compute_pass(&desc); pass.set_pipeline(&self.pipeline); for binding in &self.bindings { binding.set(&mut pass); } pass.dispatch_workgroups(self.points_count + 63 / 64, 64, 1); } // submit the encoder self.queue.submit(Some(encoder.finish())); }
async fn get_points(&self) -> Box<[Point]> { // copy the points buffer to the mapping buffer let mut encoder = self .device .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("Copy Points Buffer"), }); { encoder.copy_buffer_to_buffer( &self.points_mapping, 0, &self.points_mapping, 0, (self.points_count * std::mem::size_of::<Point>() as u32) as u64, ); } // submit the command self.queue.submit(Some(encoder.finish())); // map the points buffer self.points_mapping .slice(..) .map_async(wgpu::MapMode::Read, |r| r.unwrap()); self.device.poll(wgpu::Maintain::Wait).panic_on_timeout();
(bytemuck::cast_slice::<u8, Point>(&self.points_mapping.slice(..).get_mapped_range())) .to_vec() .into_boxed_slice() }}