Reorganize layout, pass a new buffer

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 {

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<wgpu::Buffer>,
+    accum_image: Option<wgpu::Buffer>,
+    bind_group: Option<wgpu::BindGroup>,
 }
 
 impl AccumulatePipeline {
-    pub fn new(device: &Device, module: &ShaderModule) -> Self {
+    const BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
-        let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+        wgpu::BindGroupLayoutDescriptor {
             label: Some("accumulate"),
-            entries: &[wgpu::BindGroupLayoutEntry {
+            entries: &[
+                // transforms
+                wgpu::BindGroupLayoutEntry {
                     binding: 0,
-                count: None,
                     visibility: wgpu::ShaderStages::COMPUTE,
                     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: false },
                     },
-            }],
+                    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::<IfsConstants>() 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::<f32>() 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,
+            device: device.clone(),
-            bind_group,
+            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<wgpu::Buffer>,
+    bind_group: Option<wgpu::BindGroup>,
 }
 
 impl RenderPipeline {
-    pub fn new(device: &Device, module: &ShaderModule, format: wgpu::TextureFormat) -> Self {
+    const BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
-        let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+        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::<IfsConstants>() 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 {
+    pub fn set_accum_image(&mut self, accum_image: &wgpu::Buffer) {
-    bind_group: wgpu::BindGroup,
+        self.accum_image = Some(accum_image.clone());
+        self.bind_group.take();
     }
 
-impl RenderPass {
+    fn fetch_bind_group(&mut self) -> &wgpu::BindGroup {
-    pub fn new(device: &Device, pipeline: &RenderPipeline, accum_buffer: &wgpu::Buffer) -> Self {
+        self.bind_group
-        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            .get_or_insert_with(|| {
+                self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                     label: Some("render"),
-            layout: &pipeline.pipeline.get_bind_group_layout(0),
+                    layout: &self.bind_group_layout,
                     entries: &[wgpu::BindGroupEntry {
                         binding: 0,
-                resource: accum_buffer.as_entire_binding(),
+                        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"),
+            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::<f32>() 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<Flare>, window: Arc<Window>) -> 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 mut accumulate_pipeline = AccumulatePipeline::new(&flare.device, &flare.module);
-        let accumulate_pass = AccumulatePass::new(
+        let mut render_pipeline =
-            &flare.device,
-            &accumulate_pipeline,
-            (window_size.width, window_size.height),
-        );
-        let 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) {
+    pub fn render(&mut self, run_accumulate: bool) {
-        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) {
         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);