From aaa5ea4be0c8603f5778f0125a03cd1cc1024bb0 Mon Sep 17 00:00:00 2001 From: mars Date: Sat, 7 May 2022 21:50:27 -0600 Subject: [PATCH] Mesh compute skinning! (but no rasterizing yet) --- shaders/mesh_skinning.wgsl | 42 +++++++ src/main.rs | 2 + src/pass/mesh.rs | 243 ++++++++++++++++++++++++++++++++++++- 3 files changed, 283 insertions(+), 4 deletions(-) create mode 100644 shaders/mesh_skinning.wgsl diff --git a/shaders/mesh_skinning.wgsl b/shaders/mesh_skinning.wgsl new file mode 100644 index 0000000..b091650 --- /dev/null +++ b/shaders/mesh_skinning.wgsl @@ -0,0 +1,42 @@ +#include oct_encoding.wgsl +#include skin.wgsl + +struct SkinningUniform { + transform: mat4x4; + src_offset: u32; + dst_offset: u32; + count: u32; +}; + +[[group(0), binding(0)]] +var skinning_ubo: SkinningUniform; + +[[group(0), binding(1)]] +var dst_vertices: SkinnedVertexArray; + +[[group(0), binding(2)]] +var src_vertices: SkinnedVertexArray; + +[[stage(compute), workgroup_size(64)]] +fn cs_main( + [[builtin(global_invocation_id)]] global_invocation_id: vec3, +) { + let vertex_index = global_invocation_id.x; + if (vertex_index >= skinning_ubo.count) { + return; + } + + let src_index = skinning_ubo.src_offset + vertex_index; + let dst_index = skinning_ubo.dst_offset + vertex_index; + + let ptf = src_vertices.data[src_index].ptf; + let position = ptf.xyz; + let tan_frame = bitcast(ptf.w); + + let transform = skinning_ubo.transform; + let position = transform * vec4(position, 1.0); + let tan_frame = tan_frame_transform(tan_frame, transform); + + let dst_ptf = vec4(position.xyz, bitcast(tan_frame)); + dst_vertices.data[dst_index].ptf = dst_ptf; +} diff --git a/src/main.rs b/src/main.rs index 124ca1f..2e00356 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,10 +26,12 @@ fn main() { let shader_watcher = shader::ShaderWatcher::new(shader_store.to_owned(), shaders_dir).unwrap(); let mesh_forward = shader_watcher.add_file("mesh_forward.wgsl").unwrap(); + let mesh_skinning = shader_watcher.add_file("mesh_skinning.wgsl").unwrap(); let mesh_shaders = pass::mesh::ShaderInfo { store: shader_store.clone(), forward: mesh_forward, + skinning: mesh_skinning, }; let mesh_pass = pass::mesh::MeshPass::new( diff --git a/src/pass/mesh.rs b/src/pass/mesh.rs index 82276be..949e44d 100644 --- a/src/pass/mesh.rs +++ b/src/pass/mesh.rs @@ -1,12 +1,14 @@ use super::*; +use crate::gpu::GpuVec; use crate::mesh::*; +use crate::shader::{ShaderHandle, ShaderStore}; use crate::viewport::ViewportInfo; use crate::RenderLayouts; -use crate::shader::{ShaderStore, ShaderHandle}; pub struct ShaderInfo { pub store: Arc, pub forward: ShaderHandle, + pub skinning: ShaderHandle, } #[repr(C)] @@ -18,7 +20,7 @@ pub struct Vertex { impl Attribute for Vertex { fn get_usages() -> wgpu::BufferUsages { - wgpu::BufferUsages::VERTEX + wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::VERTEX } } @@ -37,7 +39,38 @@ impl Vertex { pub type Index = u32; -pub struct FrameData {} +#[repr(C)] +#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)] +struct SkinningUniform { + transform: [[f32; 4]; 4], + src_offset: u32, + dst_offset: u32, + count: u32, + _padding: u32, +} + +#[derive(Debug)] +struct MeshCommand { + vertex_offset: usize, + vertex_count: usize, + skinned_offset: usize, + index_offset: usize, + index_count: usize, + skinning_index: usize, +} + +#[derive(Debug)] +struct MeshGroupCommands { + binding_indices: MeshLayoutBindingIndices, + bind_group: wgpu::BindGroup, + meshes: Vec, +} + +pub struct FrameData { + skinned_vertices: GpuVec, + skinning_uniforms: GpuVec, + groups: Vec, +} pub struct MeshPass { device: Arc, @@ -49,6 +82,8 @@ pub struct MeshPass { index_attr_id: AttrId, mesh_layout_id: MeshLayoutId, example_mesh: MeshHandle, + skinning_bind_group_layout: wgpu::BindGroupLayout, + skinning_pipeline: wgpu::ComputePipeline, depth_pipeline: wgpu::RenderPipeline, opaque_pipeline: wgpu::RenderPipeline, target_info: ViewportInfo, @@ -177,6 +212,61 @@ impl MeshPass { drop(shader); + let skinning_bind_group_layout = + device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("Skinning 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: true, + min_binding_size: None, // TODO ??? + }, + 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::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + + let skinning_pipeline_layout = + device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("Skinning Pipeline Layout"), + bind_group_layouts: &[&skinning_bind_group_layout], + push_constant_ranges: &[], + }); + + let skinning_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("Skinning Pipeline"), + layout: Some(&skinning_pipeline_layout), + module: shader_info + .store + .get(&shader_info.skinning) + .unwrap() + .as_ref(), + entry_point: "cs_main", + }); + Self { device, layouts, @@ -187,6 +277,8 @@ impl MeshPass { vertex_attr_id, mesh_layout_id, example_mesh, + skinning_bind_group_layout, + skinning_pipeline, depth_pipeline, opaque_pipeline, target_info, @@ -198,16 +290,133 @@ impl RenderPass for MeshPass { type FrameData = FrameData; fn create_frame_data(&self) -> FrameData { - FrameData {} + FrameData { + skinned_vertices: GpuVec::new( + self.device.clone(), + Vertex::get_usages(), + 1024 * 128, + Some("Skinned Vertices".to_string()), + ), + skinning_uniforms: GpuVec::new( + self.device.clone(), + wgpu::BufferUsages::STORAGE, + 1024 * 128, + Some("Skinning Uniforms".to_string()), + ), + groups: Default::default(), + } } fn begin_frame(&self, data: &mut FrameData, phases: &mut Vec, queue: &wgpu::Queue) { println!("MeshPass::begin_frame()"); phases.push(Phase::Upload); + phases.push(Phase::Skinning); phases.push(Phase::Depth); phases.push(Phase::Opaque); phases.push(Phase::Transparent); + + data.groups.clear(); + data.skinning_uniforms.clear(); + + let meshes = &[&self.example_mesh, &self.example_mesh]; + + let mesh_bindings = self + .mesh_pool + .iter_meshes(self.mesh_layout_id, meshes.iter(), |v| v) + .unwrap(); + + let mut skinned_cursor = 0; + + for MeshLayoutInstances { + bindings, + instances, + } in mesh_bindings.iter() + { + let pools = self.mesh_pool.get_bindings(bindings.clone()); + let vertices_pool = pools.get(self.vertex_attr_id).unwrap(); + + // TODO defer bind group creation into separate Vec after GpuVecs have been written + let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &self.skinning_bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { + buffer: data.skinning_uniforms.as_ref(), + offset: 0, + // TODO ugly! + size: Some( + std::num::NonZeroU64::new( + std::mem::size_of::() as u64 + ) + .unwrap(), + ), + }), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { + buffer: data.skinned_vertices.as_ref(), + offset: 0, + size: None, + }), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { + buffer: vertices_pool.get_buffer(), + offset: 0, + size: None, + }), + }, + ], + }); + + let mut group = MeshGroupCommands { + binding_indices: bindings.clone(), + meshes: Default::default(), + bind_group, + }; + + for (_mesh, infos) in instances { + let vertices = infos.iter().find(|i| i.0 == self.vertex_attr_id).unwrap().1; + let indices = infos.iter().find(|i| i.0 == self.index_attr_id).unwrap().1; + + group.meshes.push(MeshCommand { + vertex_offset: vertices.offset, + vertex_count: vertices.count, + skinned_offset: skinned_cursor, + index_offset: indices.offset, + index_count: indices.count, + skinning_index: data.skinning_uniforms.len(), + }); + + data.skinning_uniforms.push(SkinningUniform { + transform: [ + [1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.], + ], + src_offset: vertices.offset as u32, + dst_offset: skinned_cursor as u32, + count: vertices.count as u32, + _padding: 0, + }); + + skinned_cursor += vertices.count; + } + + data.groups.push(group); + } + + println!("commands: {:#?}", data.groups); + + data.skinned_vertices.reserve(skinned_cursor); + data.skinned_vertices.write(&queue); + data.skinning_uniforms.write(&queue); } fn record_commands(&self, data: PhaseData<&FrameData>, cmds: &mut wgpu::CommandEncoder) { @@ -217,6 +426,32 @@ impl RenderPass for MeshPass { } } + fn record_compute<'a>( + &'a self, + data: PhaseData<&'a FrameData>, + cmds: &mut wgpu::ComputePass<'a>, + ) { + cmds.set_pipeline(&self.skinning_pipeline); + for group in data.frame_data.groups.iter() { + for mesh in group.meshes.iter() { + let ubo_offset = data + .frame_data + .skinning_uniforms + .buf_offset(mesh.skinning_index); + cmds.set_bind_group(0, &group.bind_group, &[ubo_offset as u32]); + + // TODO use div_ceil instead + let workgroup_num = if mesh.vertex_count % 64 == 0 { + mesh.vertex_count / 64 + } else { + mesh.vertex_count / 64 + 1 + }; + + cmds.dispatch(workgroup_num as u32, 1, 1); + } + } + } + fn record_render(&self, data: PhaseData<&FrameData>) -> Option { println!("MeshPass::record_render(phase: {:?})", data.phase);