flare/ifs-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reorganize layout, pass a new buffer

Browse Source
This commit is contained in:
Bradlee Speice 2025-01-28 22:02:29 -05:00
parent 203435ca72
commit 6d02b3dcc1
2 changed files with 244 additions and 130 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
crates/flare-shader/src/lib.rs
View File

@ -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 {

341
crates/flare/src/main.rs
View File

@ -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 {
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 {
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,
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<wgpu::Buffer>,
bind_group: Option<wgpu::BindGroup>,
}
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::<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 {
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: &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 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);