Initial compute/vertex/fragment render system

crates/flare-shader/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "flare-shader"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lib]
+crate-type = ["dylib", "lib"]
+
+[dependencies]
+bytemuck.workspace = true
+glam.workspace = true
+spirv-std.workspace = true
+
+[lints]
+workspace = true

crates/flare-shader/src/lib.rs

@@ -0,0 +1,141 @@
+#![cfg_attr(target_arch = "spirv", no_std)]
+
+use glam::{UVec2, Vec4, Vec4Swizzles, uvec2, vec2, vec4};
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+#[repr(C)]
+pub struct IfsConstants {
+    accum_width: u32,
+    accum_height: u32,
+    viewport_width: u32,
+    viewport_height: u32,
+    background_color: Vec4,
+}
+
+impl IfsConstants {
+    pub fn new(
+        accum_width: u32,
+        accum_height: u32,
+        viewport_width: u32,
+        viewport_height: u32,
+        background_color: Vec4,
+    ) -> Self {
+        Self {
+            accum_width,
+            accum_height,
+            viewport_width,
+            viewport_height,
+            background_color,
+        }
+    }
+
+    pub fn viewport_dimensions(&self) -> UVec2 {
+        uvec2(self.viewport_width, self.viewport_height)
+    }
+
+    pub fn with_accumulate(&mut self, width: u32, height: u32) {
+        self.accum_width = width;
+        self.accum_height = height;
+    }
+
+    pub fn with_viewport(&mut self, width: u32, height: u32) {
+        self.viewport_width = width;
+        self.viewport_height = height;
+    }
+}
+
+pub trait Color {
+    type Element;
+    const BLACK: Self;
+
+    const WHITE: Self;
+}
+
+impl Color for Vec4 {
+    type Element = f32;
+
+    const BLACK: Self = vec4(0., 0., 0., 1.);
+    const WHITE: Self = vec4(1., 1., 1., 1.);
+}
+
+pub(crate) fn image_index(pixel_x: usize, pixel_y: usize, image_width: u32) -> usize {
+    pixel_x + pixel_y * image_width as usize
+}
+
+#[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],
+) {
+    let block_size = 64;
+    for i_width in 0..constants.accum_width as usize {
+        for i_height in 0..constants.accum_height as usize {
+            let color = if (i_width / block_size % 2 == 1) != (i_height / block_size % 2 == 1) {
+                Vec4::BLACK
+            } else {
+                Vec4::WHITE
+            };
+
+            accum_image[image_index(i_width, i_height, constants.accum_width)] = color;
+        }
+    }
+}
+
+#[spirv(vertex)]
+pub fn main_vs(
+    #[spirv(vertex_index)] vert_id: u32,
+    #[spirv(position, invariant)] position: &mut Vec4,
+) {
+    // Create a "quad" that fills the viewport for the fragment shader.
+    // The `draw` call issued by the main application will be for three vertex ID's (0, 1, 2).
+    // This code maps them to the points (-1, -1), (3, -1), and (-1, 3) respectively.
+    // Because the interior of that triangle covers the entire viewport,
+    // the GPU clips to the viewport and invokes the fragment shader for each pixel.
+    // https://stackoverflow.com/a/59739538
+    // https://www.saschawillems.de/blog/2016/08/13/vulkan-tutorial-on-rendering-a-fullscreen-quad-without-buffers/
+    let output_uv = vec2(((vert_id << 1) & 2) as f32, (vert_id & 2) as f32);
+    *position = (output_uv * 2.0 - 1.0, 0.0, 1.0).into();
+}
+
+#[spirv(fragment)]
+pub fn main_fs(
+    #[spirv(frag_coord)] frag_coord: Vec4,
+    #[spirv(push_constant)] ifs_constants: &IfsConstants,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] accum_image: &[Vec4],
+    output: &mut Vec4,
+) {
+    // Bootleg texture sampling; map from viewport image pixel coordinates to accumulator image
+    // pixel coordinates
+    let viewport_coordinate = frag_coord.xy().as_uvec2();
+
+    let a_width = ifs_constants.accum_width as f32;
+    let a_height = ifs_constants.accum_height as f32;
+    let v_width = ifs_constants.viewport_width as f32;
+    let v_height = ifs_constants.viewport_height as f32;
+
+    // Scale both width and height by the same factor; preserves aspect ratio
+    let scale_width = a_width / v_width;
+    let scale_height = a_height / v_height;
+    let scale = scale_width.max(scale_height);
+
+    // Re-center the image in the viewport after scale
+    let offset_x = (v_width * scale - a_width) / 2.0;
+    let offset_y = (v_height * scale - a_height) / 2.0;
+
+    let accum_coordinate = viewport_coordinate.as_vec2() * scale - vec2(offset_x, offset_y);
+
+    if accum_coordinate.x < 0.0
+        || accum_coordinate.x >= ifs_constants.accum_width as f32
+        || accum_coordinate.y < 0.0
+        || accum_coordinate.y >= ifs_constants.accum_height as f32
+    {
+        *output = ifs_constants.background_color;
+    } else {
+        *output = accum_image[image_index(
+            accum_coordinate.x as usize,
+            accum_coordinate.y as usize,
+            ifs_constants.accum_width,
+        )];
+    }
+}

crates/flare/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "flare"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+bytemuck.workspace = true
+env_logger.workspace = true
+flare-shader = { path = "../flare-shader" }
+futures-executor.workspace = true
+glam.workspace = true
+wgpu.workspace = true
+winit.workspace = true
+
+[build-dependencies]
+spirv-builder.workspace = true
+
+[lints]
+workspace = true

crates/flare/build.rs

@@ -0,0 +1,18 @@
+use spirv_builder::{MetadataPrintout, SpirvBuilder};
+use std::path::PathBuf;
+use std::{env, fs};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let result = SpirvBuilder::new("../flare-shader", "spirv-unknown-vulkan1.1")
+        .print_metadata(MetadataPrintout::Full)
+        .release(false)
+        .build()?;
+
+    // Copy the SPIR-V to this crate's output directory
+    let shader_path = result.module.unwrap_single();
+    let output_path = PathBuf::from(env::var("OUT_DIR")?).join("flare.spv");
+    fs::copy(shader_path, &output_path)?;
+
+    println!("Generated shader; path={}", output_path.display());
+    Ok(())
+}

crates/flare/src/main.rs

@@ -0,0 +1,421 @@
+use flare_shader::{Color, IfsConstants};
+use futures_executor::block_on;
+use glam::Vec4;
+use std::sync::Arc;
+use wgpu::{
+    Adapter, Device, Features, Instance, Queue, ShaderModule, Surface, SurfaceConfiguration,
+};
+use winit::event::MouseButton;
+use winit::{
+    application::ApplicationHandler,
+    dpi::LogicalSize,
+    event::WindowEvent,
+    event_loop::{ActiveEventLoop, EventLoop},
+    window::{Window, WindowAttributes, WindowId},
+};
+
+struct AccumulatePipeline {
+    pipeline: wgpu::ComputePipeline,
+}
+
+impl AccumulatePipeline {
+    pub fn new(device: &Device, module: &ShaderModule) -> Self {
+        let bindgroup_layout = device.create_bind_group_layout(&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 },
+                },
+            }],
+        });
+        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+            label: Some("accumulate"),
+            bind_group_layouts: &[&bindgroup_layout],
+            push_constant_ranges: &[wgpu::PushConstantRange {
+                stages: wgpu::ShaderStages::COMPUTE,
+                range: 0..size_of::<IfsConstants>() as u32,
+            }],
+        });
+        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+            label: Some("accumulate"),
+            layout: Some(&pipeline_layout),
+            module,
+            entry_point: Some("main_cs"),
+            compilation_options: Default::default(),
+            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,
+        }
+    }
+}
+
+struct RenderPipeline {
+    pipeline: wgpu::RenderPipeline,
+}
+
+impl RenderPipeline {
+    pub fn new(device: &Device, module: &ShaderModule, format: wgpu::TextureFormat) -> Self {
+        let bindgroup_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+            label: Some("render"),
+            entries: &[wgpu::BindGroupLayoutEntry {
+                binding: 0,
+                count: None,
+                visibility: wgpu::ShaderStages::FRAGMENT,
+                ty: wgpu::BindingType::Buffer {
+                    has_dynamic_offset: false,
+                    min_binding_size: None,
+                    ty: wgpu::BufferBindingType::Storage { read_only: true },
+                },
+            }],
+        });
+        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+            label: Some("render"),
+            bind_group_layouts: &[&bindgroup_layout],
+            push_constant_ranges: &[wgpu::PushConstantRange {
+                stages: wgpu::ShaderStages::FRAGMENT,
+                range: 0..size_of::<IfsConstants>() as u32,
+            }],
+        });
+        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("render"),
+            layout: Some(&pipeline_layout),
+            vertex: wgpu::VertexState {
+                module,
+                entry_point: Some("main_vs"),
+                compilation_options: Default::default(),
+                buffers: &[],
+            },
+            primitive: Default::default(),
+            depth_stencil: None,
+            multisample: Default::default(),
+            fragment: Some(wgpu::FragmentState {
+                module,
+                entry_point: Some("main_fs"),
+                compilation_options: Default::default(),
+                targets: &[Some(wgpu::ColorTargetState {
+                    format,
+                    blend: None,
+                    write_mask: Default::default(),
+                })],
+            }),
+            multiview: None,
+            cache: None,
+        });
+        Self { pipeline }
+    }
+}
+
+struct RenderPass {
+    bind_group: wgpu::BindGroup,
+}
+
+impl RenderPass {
+    pub fn new(device: &Device, pipeline: &RenderPipeline, accum_buffer: &wgpu::Buffer) -> Self {
+        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("render"),
+            layout: &pipeline.pipeline.get_bind_group_layout(0),
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: accum_buffer.as_entire_binding(),
+            }],
+        });
+        Self { bind_group }
+    }
+}
+
+struct Flare {
+    instance: Instance,
+    adapter: Adapter,
+    device: Device,
+    queue: Queue,
+    module: ShaderModule,
+}
+
+impl Flare {
+    pub fn new() -> Self {
+        let backends = wgpu::Backends::from_env().unwrap_or_default();
+        let instance = Instance::new(&wgpu::InstanceDescriptor {
+            backends,
+            ..Default::default()
+        });
+
+        let adapter = instance.request_adapter(&wgpu::RequestAdapterOptions {
+            power_preference: wgpu::PowerPreference::HighPerformance,
+            force_fallback_adapter: false,
+            compatible_surface: None,
+        });
+        let adapter = block_on(adapter).expect("Failed to find GPU adapter");
+
+        let required_limits = wgpu::Limits {
+            max_push_constant_size: size_of::<IfsConstants>() as u32,
+            ..Default::default()
+        };
+
+        let device = adapter.request_device(
+            &wgpu::DeviceDescriptor {
+                label: Some("flare"),
+                required_features: Features::TIMESTAMP_QUERY | Features::PUSH_CONSTANTS,
+                required_limits,
+                memory_hints: Default::default(),
+            },
+            None,
+        );
+        let (device, queue) = block_on(device).expect("Failed to find GPU device");
+        let module = device
+            .create_shader_module(wgpu::include_spirv!(concat!(env!("OUT_DIR"), "/flare.spv")));
+
+        Flare {
+            instance,
+            adapter,
+            device,
+            queue,
+            module,
+        }
+    }
+}
+
+struct FlareRender<'window> {
+    flare: Arc<Flare>,
+    surface: Surface<'window>,
+    surface_configuration: SurfaceConfiguration,
+    accumulate_pipeline: AccumulatePipeline,
+    accumulate_pass: AccumulatePass,
+    render_pipeline: RenderPipeline,
+    render_pass: RenderPass,
+    ifs_constants: IfsConstants,
+}
+
+impl FlareRender<'_> {
+    pub fn new(flare: Arc<Flare>, window: Arc<Window>) -> Self {
+        let window_size = window.inner_size();
+
+        let surface = flare
+            .instance
+            .create_surface(window.clone())
+            .expect("Unable to create surface");
+        let mut surface_configuration = surface
+            .get_default_config(&flare.adapter, window_size.width, window_size.height)
+            .expect("Unable to get surface config");
+        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 =
+            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,
+            window_size.width,
+            window_size.height,
+            Vec4::BLACK,
+        );
+
+        Self {
+            flare,
+            surface,
+            surface_configuration,
+            accumulate_pipeline,
+            accumulate_pass,
+            render_pipeline,
+            render_pass,
+            ifs_constants,
+        }
+    }
+
+    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) {
+        let output = self
+            .surface
+            .get_current_texture()
+            .expect("Failed to get current texture");
+        let output_view = output.texture.create_view(&Default::default());
+        let mut encoder = self
+            .flare
+            .device
+            .create_command_encoder(&Default::default());
+
+        if run_accumulate {
+            self.begin_accumulate(&mut encoder);
+        }
+        self.begin_render(&mut encoder, &output_view);
+
+        self.flare.queue.submit(Some(encoder.finish()));
+        output.present();
+    }
+
+    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);
+    }
+
+    pub fn resize_viewport(&mut self, width: u32, height: u32) {
+        self.surface_configuration.width = width;
+        self.surface_configuration.height = height;
+        self.surface
+            .configure(&self.flare.device, &self.surface_configuration);
+        self.ifs_constants.with_viewport(width, height);
+    }
+}
+
+struct Application<'window> {
+    flare: Arc<Flare>,
+    window: Option<Arc<Window>>,
+    flare_render: Option<FlareRender<'window>>,
+}
+
+impl Application<'_> {
+    pub fn new() -> Self {
+        Self {
+            flare: Arc::new(Flare::new()),
+            window: None,
+            flare_render: None,
+        }
+    }
+}
+
+impl ApplicationHandler for Application<'_> {
+    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
+        let attributes = WindowAttributes::default()
+            .with_title("Flare")
+            .with_inner_size(LogicalSize::new(1024, 768));
+
+        let window = Arc::new(
+            event_loop
+                .create_window(attributes)
+                .expect("Failed to create window"),
+        );
+
+        self.window = Some(window.clone());
+
+        let flare_render = FlareRender::new(self.flare.clone(), window);
+        flare_render.render(true);
+
+        self.flare_render = Some(flare_render);
+    }
+
+    fn window_event(
+        &mut self,
+        event_loop: &ActiveEventLoop,
+        _window_id: WindowId,
+        event: WindowEvent,
+    ) {
+        match event {
+            WindowEvent::Resized(size) => {
+                let flare_render = self.flare_render.as_mut().unwrap();
+                flare_render.resize_viewport(size.width, size.height);
+                flare_render.render(false);
+            }
+            WindowEvent::MouseInput {
+                button: MouseButton::Left,
+                ..
+            } => {
+                let flare_render = self.flare_render.as_mut().unwrap();
+                flare_render.resize_accumulate();
+                flare_render.render(true);
+            }
+            WindowEvent::CloseRequested => event_loop.exit(),
+            _ => (),
+        }
+    }
+}
+
+pub fn main() -> anyhow::Result<()> {
+    env_logger::init();
+    let _module = wgpu::include_spirv!(concat!(env!("OUT_DIR"), "/flare.spv"));
+
+    let event_loop = EventLoop::new()?;
+    let mut application = Application::new();
+    event_loop.run_app(&mut application)?;
+
+    Ok(())
+}