diff --git a/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.lock b/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.lock index d370303..34f8b05 100644 --- a/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.lock +++ b/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.lock @@ -774,9 +774,11 @@ checksum = "f25c0e292a7ca6d6498557ff1df68f32c99850012b6ea401cf8daf771f22ff53" name = "draw-compute" version = "0.1.0" dependencies = [ + "bytemuck", "eframe", "egui", "egui-wgpu", + "glam 0.30.1", "shader", "wgpu", ] diff --git a/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.toml b/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.toml index 64f205f..68ee9ee 100644 --- a/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.toml +++ b/blog/2025-03-30-draw-compute-shader/draw-compute/Cargo.toml @@ -4,9 +4,11 @@ version = "0.1.0" edition = "2021" [dependencies] +bytemuck = { version = "1.22", features = ["derive"] } eframe = { version = "0.31", features = ["wgpu"] } egui = "0.31" egui-wgpu = "0.31" +glam = { version = "0.30", default-features = false, features = ["bytemuck", "libm"] } wgpu = { version = "24.0", features = ["spirv"] } shader = { path = "shader" } diff --git a/blog/2025-03-30-draw-compute-shader/draw-compute/shader/src/lib.rs b/blog/2025-03-30-draw-compute-shader/draw-compute/shader/src/lib.rs index 1ff6169..91c617d 100644 --- a/blog/2025-03-30-draw-compute-shader/draw-compute/shader/src/lib.rs +++ b/blog/2025-03-30-draw-compute-shader/draw-compute/shader/src/lib.rs @@ -1,27 +1,41 @@ #![no_std] + +use glam::Vec4Swizzles; use spirv_std::spirv; #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)] #[repr(C)] pub struct Viewport { - pub offset_x: u32, - pub offset_y: u32, - pub width: u32, - pub height: u32, + pub image_size: glam::UVec2, + pub viewport_offset: glam::UVec2, + pub viewport_size: glam::UVec2, } const BLOCK_SIZE: u32 = 16; const BLACK: glam::Vec4 = glam::vec4(0.0, 0.0, 0.0, 1.0); const WHITE: glam::Vec4 = glam::vec4(1.0, 1.0, 1.0, 1.0); +fn image_index(x: usize, y: usize, width: usize) -> usize { + y * width + x +} + #[spirv(compute(threads(1)))] pub fn main_cs( #[spirv(uniform, descriptor_set = 0, binding = 0)] viewport: &Viewport, #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] image: &mut [glam::Vec4], ) { - for x in 0..viewport.width as usize { - for y in 0..viewport.height as usize { - let image_index = y * viewport.width as usize + x; + let width = viewport.image_size.x as usize; + let height = viewport.image_size.y as usize; + for x in 0..width { + for y in 0..height { + let index = image_index(x, y, width); + if x == 0 + || x == width - 1 + || y == 0 + || y == height - 1 + { + image[index] = WHITE; + } } } } @@ -42,5 +56,13 @@ pub fn main_fs( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] image: &mut [glam::Vec4], output: &mut glam::Vec4, ) { - *output = BLACK; + let pixel_coordinate = frag_coord.xy().as_usizevec2(); + let (pixel_x, pixel_y) = (pixel_coordinate.x, pixel_coordinate.y); + let index = image_index(pixel_x, pixel_y, viewport.viewport_size.x as usize); + + *output = if index < image.len() { + image[index] + } else { + BLACK + }; } diff --git a/blog/2025-03-30-draw-compute-shader/draw-compute/src/main.rs b/blog/2025-03-30-draw-compute-shader/draw-compute/src/main.rs index 04cc5e7..08b7666 100644 --- a/blog/2025-03-30-draw-compute-shader/draw-compute/src/main.rs +++ b/blog/2025-03-30-draw-compute-shader/draw-compute/src/main.rs @@ -3,6 +3,7 @@ use eframe::wgpu::{CommandBuffer, CommandEncoder, Device, Queue, RenderPass}; use eframe::Frame; use egui::{Context, Sense}; use egui_wgpu::{CallbackResources, CallbackTrait, ScreenDescriptor}; +use shader::Viewport; struct DrawResources { device: wgpu::Device, @@ -10,6 +11,7 @@ struct DrawResources { bind_group: wgpu::BindGroup, viewport_buffer: wgpu::Buffer, image_buffer: wgpu::Buffer, + image_buffer_size: glam::UVec2, compute_pipeline: wgpu::ComputePipeline, render_pipeline: wgpu::RenderPipeline, } @@ -148,6 +150,7 @@ impl DrawResources { let bind_group_layout = Self::bind_group_layout(device); let viewport_buffer = Self::viewport_buffer(device); let image_buffer = Self::image_buffer(device, width, height); + let image_buffer_size = glam::uvec2(width as u32, height as u32); let bind_group = Self::bind_group(device, &bind_group_layout, &viewport_buffer, &image_buffer); @@ -162,6 +165,7 @@ impl DrawResources { bind_group, viewport_buffer, image_buffer, + image_buffer_size, compute_pipeline, render_pipeline, } @@ -187,26 +191,38 @@ impl CallbackTrait for DrawCallback { fn prepare( &self, _device: &Device, - _queue: &Queue, + queue: &Queue, _screen_descriptor: &ScreenDescriptor, egui_encoder: &mut CommandEncoder, callback_resources: &mut CallbackResources, ) -> Vec { + + let resources = callback_resources + .get_mut::() + .expect("missing draw resources"); + + let mut viewport = Viewport { + image_size: resources.image_buffer_size, + viewport_offset: glam::uvec2(self.draw_rect.min.x as u32, self.draw_rect.min.y as u32), + viewport_size: glam::uvec2(self.draw_rect.size().x as u32, self.draw_rect.size().y as u32), + }; + if !self.draw_resize { + queue.write_buffer(&resources.viewport_buffer, 0, bytemuck::cast_slice(&[viewport])); return vec![]; } + let draw_size = self.draw_rect.size(); + resources.resize(draw_size.x as u64, draw_size.y as u64); + + viewport.image_size = resources.image_buffer_size; + queue.write_buffer(&resources.viewport_buffer, 0, bytemuck::cast_slice(&[viewport])); + let mut compute_pass = egui_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("compute"), timestamp_writes: None, }); - let resources = callback_resources - .get_mut::() - .expect("missing draw resources"); - let draw_size = self.draw_rect.size(); - resources.resize(draw_size.x as u64, draw_size.y as u64); - compute_pass.set_pipeline(&resources.compute_pipeline); compute_pass.set_bind_group(0, &resources.bind_group, &[]);