diff --git a/crates/flare-shader/src/lib.rs b/crates/flare-shader/src/lib.rs index 51b323b..4750f7b 100644 --- a/crates/flare-shader/src/lib.rs +++ b/crates/flare-shader/src/lib.rs @@ -63,10 +63,21 @@ pub(crate) fn image_index(pixel_x: usize, pixel_y: usize, image_width: u32) -> u pixel_x + pixel_y * image_width as usize } +#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)] +#[repr(C)] +pub struct Transform { + variation_offset: u32, + variation_count: u32, + weight: f32, + color: f32, + color_speed: f32, +} + #[spirv(compute(threads(1)))] pub fn main_cs( #[spirv(push_constant)] constants: &IfsConstants, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] accum_image: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] transforms: &[Transform], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] accum_image: &mut [Vec4], ) { let block_size = 64; for i_width in 0..constants.accum_width as usize { diff --git a/crates/flare/src/main.rs b/crates/flare/src/main.rs index 58a9b47..89ecfa9 100644 --- a/crates/flare/src/main.rs +++ b/crates/flare/src/main.rs @@ -1,9 +1,12 @@ -use flare_shader::{Color, IfsConstants}; +use flare_shader::{Color, IfsConstants, Transform}; use futures_executor::block_on; use glam::Vec4; use std::sync::Arc; +use bytemuck::Zeroable; +use wgpu::util::DeviceExt; use wgpu::{ - Adapter, Device, Features, Instance, Queue, ShaderModule, Surface, SurfaceConfiguration, + Adapter, BindGroupLayout, Device, Features, Instance, Queue, ShaderModule, Surface, + SurfaceConfiguration, }; use winit::event::MouseButton; use winit::{ @@ -15,27 +18,74 @@ use winit::{ }; struct AccumulatePipeline { + device: Device, + bind_group_layout: BindGroupLayout, + pipeline_layout: wgpu::PipelineLayout, pipeline: wgpu::ComputePipeline, + + transforms: Option, + accum_image: Option, + bind_group: Option, } impl AccumulatePipeline { - pub fn new(device: &Device, module: &ShaderModule) -> Self { - let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + const BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> = + wgpu::BindGroupLayoutDescriptor { label: Some("accumulate"), - entries: &[wgpu::BindGroupLayoutEntry { - binding: 0, - count: None, - visibility: wgpu::ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - has_dynamic_offset: false, - min_binding_size: None, - ty: wgpu::BufferBindingType::Storage { read_only: false }, + entries: &[ + // transforms + 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, }, - }], - }); + // accum_image + 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, + }, + ], + }; + + fn create_bind_group( + device: &Device, + layout: &wgpu::BindGroupLayout, + transforms: &wgpu::Buffer, + accum_image: &wgpu::Buffer, + ) -> wgpu::BindGroup { + device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("accumulate"), + layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: transforms.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: accum_image.as_entire_binding(), + }, + ], + }) + } + + pub fn new(device: &Device, module: &ShaderModule) -> Self { + let bind_group_layout = + device.create_bind_group_layout(&Self::BIND_GROUP_LAYOUT_DESCRIPTOR); let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("accumulate"), - bind_group_layouts: &[&bindgroup_layout], + bind_group_layouts: &[&bind_group_layout], push_constant_ranges: &[wgpu::PushConstantRange { stages: wgpu::ShaderStages::COMPUTE, range: 0..size_of::() as u32, @@ -50,62 +100,110 @@ impl AccumulatePipeline { cache: None, }); - Self { pipeline } - } -} - -struct AccumulatePass { - accum_buffer: wgpu::Buffer, - bind_group: wgpu::BindGroup, -} - -impl AccumulatePass { - pub fn new(device: &Device, pipeline: &AccumulatePipeline, dimensions: (u32, u32)) -> Self { - let pixels = dimensions.0 * dimensions.1; - let elements = pixels * 4; - let accum_buffer = device.create_buffer(&wgpu::BufferDescriptor { - label: Some("accumulate"), - size: elements as u64 * size_of::() as u64, - usage: wgpu::BufferUsages::STORAGE, - mapped_at_creation: false, - }); - let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("accumulate"), - layout: &pipeline.pipeline.get_bind_group_layout(0), - entries: &[wgpu::BindGroupEntry { - binding: 0, - resource: accum_buffer.as_entire_binding(), - }], - }); Self { - accum_buffer, - bind_group, + device: device.clone(), + bind_group_layout, + pipeline_layout, + pipeline, + transforms: None, + accum_image: None, + bind_group: None, } } + + pub fn set_transforms(&mut self, transforms: &[Transform]) { + // Should be smarter about allocation in the future + self.transforms = Some( + self.device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("accumulate/transforms"), + contents: bytemuck::cast_slice(transforms), + usage: wgpu::BufferUsages::STORAGE, + }), + ); + + // Setting a new buffer invalidates the existing bindings + self.bind_group.take(); + } + + pub fn set_accum_image(&mut self, accum_image: &wgpu::Buffer) { + self.accum_image = Some(accum_image.clone()); + + // Setting a new buffer invalidates the existing bindings + self.bind_group.take(); + } + + fn fetch_bind_group(&mut self) -> &wgpu::BindGroup { + self.bind_group + .get_or_insert_with(|| { + self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("accumulate"), + layout: &self.bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: self + .transforms + .as_ref() + .expect("transforms missing") + .as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: self + .accum_image + .as_ref() + .expect("accum_image missing") + .as_entire_binding(), + }, + ], + }) + }) + } + + pub fn run(&mut self, encoder: &mut wgpu::CommandEncoder, constants: &IfsConstants) { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("accumulate"), + timestamp_writes: None, + }); + pass.set_pipeline(&self.pipeline); + pass.set_push_constants(0, bytemuck::cast_slice(&[*constants])); + pass.set_bind_group(0, self.fetch_bind_group(), &[]); + pass.dispatch_workgroups(1, 1, 1); + } } struct RenderPipeline { + device: Device, + bind_group_layout: wgpu::BindGroupLayout, pipeline: wgpu::RenderPipeline, + + accum_image: Option, + bind_group: Option, } impl RenderPipeline { - pub fn new(device: &Device, module: &ShaderModule, format: wgpu::TextureFormat) -> Self { - let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + const BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> = + wgpu::BindGroupLayoutDescriptor { label: Some("render"), entries: &[wgpu::BindGroupLayoutEntry { binding: 0, - count: None, visibility: wgpu::ShaderStages::FRAGMENT, ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, has_dynamic_offset: false, min_binding_size: None, - ty: wgpu::BufferBindingType::Storage { read_only: true }, }, + count: None, }], - }); + }; + + pub fn new(device: &Device, module: &ShaderModule, format: wgpu::TextureFormat) -> Self { + let bind_group_layout = + device.create_bind_group_layout(&Self::BIND_GROUP_LAYOUT_DESCRIPTOR); let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("render"), - bind_group_layouts: &[&bindgroup_layout], + bind_group_layouts: &[&bind_group_layout], push_constant_ranges: &[wgpu::PushConstantRange { stages: wgpu::ShaderStages::FRAGMENT, range: 0..size_of::() as u32, @@ -136,25 +234,63 @@ impl RenderPipeline { multiview: None, cache: None, }); - Self { pipeline } + Self { + device: device.clone(), + bind_group_layout, + pipeline, + accum_image: None, + bind_group: None, + } } -} -struct RenderPass { - bind_group: wgpu::BindGroup, -} + pub fn set_accum_image(&mut self, accum_image: &wgpu::Buffer) { + self.accum_image = Some(accum_image.clone()); + self.bind_group.take(); + } -impl RenderPass { - pub fn new(device: &Device, pipeline: &RenderPipeline, accum_buffer: &wgpu::Buffer) -> Self { - let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + fn fetch_bind_group(&mut self) -> &wgpu::BindGroup { + self.bind_group + .get_or_insert_with(|| { + self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("render"), + layout: &self.bind_group_layout, + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: self + .accum_image + .as_ref() + .expect("accum_image missing") + .as_entire_binding(), + }], + }) + }) + } + + pub fn run( + &mut self, + encoder: &mut wgpu::CommandEncoder, + output_view: &wgpu::TextureView, + constants: &IfsConstants, + ) { + let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { label: Some("render"), - layout: &pipeline.pipeline.get_bind_group_layout(0), - entries: &[wgpu::BindGroupEntry { - binding: 0, - resource: accum_buffer.as_entire_binding(), - }], + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: output_view, + resolve_target: None, + ops: Default::default(), + })], + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, }); - Self { bind_group } + pass.set_pipeline(&self.pipeline); + pass.set_push_constants( + wgpu::ShaderStages::FRAGMENT, + 0, + bytemuck::cast_slice(&[*constants]), + ); + pass.set_bind_group(0, self.fetch_bind_group(), &[]); + pass.draw(0..3, 0..1); } } @@ -214,13 +350,24 @@ struct FlareRender<'window> { surface: Surface<'window>, surface_configuration: SurfaceConfiguration, accumulate_pipeline: AccumulatePipeline, - accumulate_pass: AccumulatePass, render_pipeline: RenderPipeline, - render_pass: RenderPass, ifs_constants: IfsConstants, + accum_image: wgpu::Buffer, } impl FlareRender<'_> { + fn create_accum_image(device: &Device, width: u32, height: u32) -> wgpu::Buffer { + let pixels = (width * height) as u64; + let pixel_size = 4u64 * size_of::() as u64; + let size = pixels * pixel_size; + device.create_buffer(&wgpu::BufferDescriptor { + label: Some("accum_image"), + size, + usage: wgpu::BufferUsages::STORAGE, + mapped_at_creation: false, + }) + } + pub fn new(flare: Arc, window: Arc) -> Self { let window_size = window.inner_size(); @@ -234,19 +381,10 @@ impl FlareRender<'_> { surface_configuration.present_mode = wgpu::PresentMode::AutoVsync; surface.configure(&flare.device, &surface_configuration); - let accumulate_pipeline = AccumulatePipeline::new(&flare.device, &flare.module); - let accumulate_pass = AccumulatePass::new( - &flare.device, - &accumulate_pipeline, - (window_size.width, window_size.height), - ); - let render_pipeline = + let mut accumulate_pipeline = AccumulatePipeline::new(&flare.device, &flare.module); + let mut render_pipeline = RenderPipeline::new(&flare.device, &flare.module, surface_configuration.format); - let render_pass = RenderPass::new( - &flare.device, - &render_pipeline, - &accumulate_pass.accum_buffer, - ); + let ifs_constants = IfsConstants::new( window_size.width, window_size.height, @@ -255,54 +393,24 @@ impl FlareRender<'_> { Vec4::BLACK, ); + let accum_image = Self::create_accum_image(&flare.device, window_size.width, window_size.height); + + accumulate_pipeline.set_transforms(&[Transform::zeroed()]); + accumulate_pipeline.set_accum_image(&accum_image); + render_pipeline.set_accum_image(&accum_image); + Self { flare, surface, surface_configuration, accumulate_pipeline, - accumulate_pass, render_pipeline, - render_pass, ifs_constants, + accum_image, } } - fn begin_accumulate(&self, encoder: &mut wgpu::CommandEncoder) { - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("accumulate"), - timestamp_writes: None, - }); - - pass.set_pipeline(&self.accumulate_pipeline.pipeline); - pass.set_push_constants(0, bytemuck::cast_slice(&[self.ifs_constants])); - pass.set_bind_group(0, &self.accumulate_pass.bind_group, &[]); - pass.dispatch_workgroups(1, 1, 1); - } - - fn begin_render(&self, encoder: &mut wgpu::CommandEncoder, output_view: &wgpu::TextureView) { - let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { - label: Some("render"), - color_attachments: &[Some(wgpu::RenderPassColorAttachment { - view: output_view, - resolve_target: None, - ops: Default::default(), - })], - depth_stencil_attachment: None, - timestamp_writes: None, - occlusion_query_set: None, - }); - - pass.set_pipeline(&self.render_pipeline.pipeline); - pass.set_push_constants( - wgpu::ShaderStages::FRAGMENT, - 0, - bytemuck::cast_slice(&[self.ifs_constants]), - ); - pass.set_bind_group(0, &self.render_pass.bind_group, &[]); - pass.draw(0..3, 0..1); - } - - pub fn render(&self, run_accumulate: bool) { + pub fn render(&mut self, run_accumulate: bool) { let output = self .surface .get_current_texture() @@ -314,9 +422,9 @@ impl FlareRender<'_> { .create_command_encoder(&Default::default()); if run_accumulate { - self.begin_accumulate(&mut encoder); + self.accumulate_pipeline.run(&mut encoder, &self.ifs_constants); } - self.begin_render(&mut encoder, &output_view); + self.render_pipeline.run(&mut encoder, &output_view, &self.ifs_constants); self.flare.queue.submit(Some(encoder.finish())); output.present(); @@ -324,18 +432,13 @@ impl FlareRender<'_> { pub fn resize_accumulate(&mut self) { let vp_dimensions = self.ifs_constants.viewport_dimensions(); - self.accumulate_pass = AccumulatePass::new( - &self.flare.device, - &self.accumulate_pipeline, - (vp_dimensions.x, vp_dimensions.y), - ); - self.render_pass = RenderPass::new( - &self.flare.device, - &self.render_pipeline, - &self.accumulate_pass.accum_buffer, - ); self.ifs_constants .with_accumulate(vp_dimensions.x, vp_dimensions.y); + + let accum_image = Self::create_accum_image(&self.flare.device, vp_dimensions.x, vp_dimensions.y); + self.accumulate_pipeline.set_accum_image(&accum_image); + self.render_pipeline.set_accum_image(&accum_image); + self.accum_image = accum_image; } pub fn resize_viewport(&mut self, width: u32, height: u32) { @@ -377,7 +480,7 @@ impl ApplicationHandler for Application<'_> { self.window = Some(window.clone()); - let flare_render = FlareRender::new(self.flare.clone(), window); + let mut flare_render = FlareRender::new(self.flare.clone(), window); flare_render.render(true); self.flare_render = Some(flare_render);