Refactor to allow running examples

This commit is contained in:
Bradlee Speice 2025-04-18 15:07:35 -04:00
parent 8aeb3cc32a
commit a976ddaf30
7 changed files with 413 additions and 179 deletions

View File

@ -1,15 +1,25 @@
#![no_std]
use glam::Vec4Swizzles;
use spirv_std::spirv;
pub trait DrawSettings: Copy + Sized + bytemuck::Pod + bytemuck::Zeroable {}
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct Viewport {
pub offset: glam::UVec2,
pub size: glam::UVec2,
pub image: glam::UVec2,
pub struct DrawSized {
pub image_size: glam::UVec2,
pub viewport_size: glam::UVec2,
}
impl DrawSettings for DrawSized {}
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct DrawRect {
pub image_size: glam::UVec2,
pub viewport_size: glam::UVec2,
pub viewport_offset: glam::UVec2,
}
impl DrawSettings for DrawRect {}
const BLOCK_SIZE: usize = 16;
const BLACK: glam::Vec4 = glam::vec4(0.0, 0.0, 0.0, 1.0);
@ -19,13 +29,37 @@ fn image_index(x: usize, y: usize, width: usize) -> usize {
y * width + x
}
fn in_bounds(image_coordinate: glam::Vec2, image_size: glam::Vec2) -> bool {
image_coordinate.cmpge(glam::Vec2::ZERO).all() && image_coordinate.cmplt(image_size).all()
}
#[spirv(compute(threads(1)))]
pub fn main_cs(
#[spirv(uniform, descriptor_set = 0, binding = 0)] viewport: &Viewport,
pub fn main_cs_bounding(
#[spirv(uniform, descriptor_set = 0, binding = 0)] viewport: &DrawSized,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] image: &mut [glam::Vec4],
) {
let width = viewport.image.x as usize;
let height = viewport.image.y as usize;
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 {
image[image_index(x, y, width)] =
if x == 0 || x == width - 1 || y == 0 || y == height - 1 {
WHITE
} else {
BLACK
}
}
}
}
#[spirv(compute(threads(1)))]
pub fn main_cs_blocks(
#[spirv(uniform, descriptor_set = 0, binding = 0)] viewport: &DrawSized,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] image: &mut [glam::Vec4],
) {
let width = viewport.image_size.x as usize;
let height = viewport.image_size.y as usize;
for x in 0..width {
let x_even = x / BLOCK_SIZE % 2 == 0;
for y in 0..height {
@ -48,24 +82,50 @@ pub fn main_vs(
}
#[spirv(fragment)]
pub fn main_fs(
pub fn main_fs_size(
#[spirv(frag_coord)] frag_coord: glam::Vec4,
#[spirv(uniform, descriptor_set = 0, binding = 0)] viewport: &Viewport,
#[spirv(uniform, descriptor_set = 0, binding = 0)] draw_settings: &DrawSized,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] image: &mut [glam::Vec4],
output: &mut glam::Vec4,
) {
let vp_size = viewport.size.as_vec2();
let img_size = viewport.image.as_vec2();
let vp_size = draw_settings.viewport_size.as_vec2();
let img_size = draw_settings.image_size.as_vec2();
let scale = (vp_size / img_size).min_element();
let img_offset = (vp_size / scale - img_size) / 2.0;
let img_coord = (frag_coord.xy() - viewport.offset.as_vec2()) / scale - img_offset;
let img_coord = frag_coord.xy() / scale - img_offset;
*output = if img_coord.cmpge(glam::Vec2::ZERO).all() && img_coord.cmple(img_size).all() {
*output = if in_bounds(img_coord, img_size) {
image[image_index(
img_coord.x as usize,
img_coord.y as usize,
viewport.image.x as usize,
img_size.x as usize,
)]
} else {
BLACK
}
}
#[spirv(fragment)]
pub fn main_fs_rect(
#[spirv(frag_coord)] frag_coord: glam::Vec4,
#[spirv(uniform, descriptor_set = 0, binding = 0)] draw_settings: &DrawRect,
#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] image: &mut [glam::Vec4],
output: &mut glam::Vec4,
) {
let vp_size = draw_settings.viewport_size.as_vec2();
let img_size = draw_settings.image_size.as_vec2();
let scale = (vp_size / img_size).min_element();
let img_offset = (vp_size / scale - img_size) / 2.0;
let img_coord =
(frag_coord.xy() - draw_settings.viewport_offset.as_vec2()) / scale - img_offset;
*output = if in_bounds(img_coord, img_size) {
image[image_index(
img_coord.x as usize,
img_coord.y as usize,
img_size.x as usize,
)]
} else {
BLACK

View File

@ -0,0 +1,16 @@
use draw_compute::draw_shaders::ShaderBounding;
use draw_compute::ComputeDraw;
fn main() {
let native_options = eframe::NativeOptions {
renderer: eframe::Renderer::Wgpu,
..Default::default()
};
eframe::run_native(
"Compute Draw",
native_options,
Box::new(|_cc| Ok(Box::new(ComputeDraw::<ShaderBounding>::new()))),
)
.unwrap()
}

View File

@ -0,0 +1,15 @@
use draw_compute::draw_shaders::ShaderOffset;
use draw_compute::ComputeDraw;
fn main() {
let native_options = eframe::NativeOptions {
renderer: eframe::Renderer::Wgpu,
..Default::default()
};
eframe::run_native(
"Compute Draw",
native_options,
Box::new(|_cc| Ok(Box::new(ComputeDraw::<ShaderOffset>::new()))),
).unwrap()
}

View File

@ -0,0 +1,15 @@
use draw_compute::draw_shaders::ShaderOffsetBlocks;
use draw_compute::ComputeDraw;
fn main() {
let native_options = eframe::NativeOptions {
renderer: eframe::Renderer::Wgpu,
..Default::default()
};
eframe::run_native(
"Compute Draw",
native_options,
Box::new(|_cc| Ok(Box::new(ComputeDraw::<ShaderOffsetBlocks>::new()))),
).unwrap()
}

View File

@ -1,26 +1,34 @@
use eframe::epaint::PaintCallbackInfo;
use eframe::wgpu::{CommandBuffer, CommandEncoder, Device, Queue, RenderPass};
use eframe::Frame;
use egui::{Context, Sense};
use egui_wgpu::{CallbackResources, CallbackTrait, ScreenDescriptor};
use shader::DrawSettings;
use std::marker::PhantomData;
struct DrawResources {
device: wgpu::Device,
bind_group_layout: wgpu::BindGroupLayout,
bind_group: wgpu::BindGroup,
viewport_buffer: wgpu::Buffer,
image_buffer: wgpu::Buffer,
image_size: glam::UVec2,
compute_pipeline: wgpu::ComputePipeline,
render_pipeline: wgpu::RenderPipeline,
pub trait ShaderSettings: Send + Sync {
type DrawSettings: DrawSettings;
fn compute_shader() -> &'static str;
fn fragment_shader() -> &'static str;
fn new(interact_rect: egui::Rect) -> Self;
fn write_buffer(&self, queue: &wgpu::Queue, buffer: &wgpu::Buffer, image_size: glam::UVec2);
}
impl DrawResources {
pub struct DrawResources<S: ShaderSettings> {
bind_group_layout: wgpu::BindGroupLayout,
pub bind_group: wgpu::BindGroup,
pub viewport_buffer: wgpu::Buffer,
image_buffer: wgpu::Buffer,
pub image_size: glam::UVec2,
pub compute_pipeline: wgpu::ComputePipeline,
pub render_pipeline: wgpu::RenderPipeline,
settings: PhantomData<S>,
}
impl<S: ShaderSettings> DrawResources<S> {
fn bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("compute_draw"),
entries: &[
// viewport
// draw_settings
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE | wgpu::ShaderStages::FRAGMENT,
@ -71,7 +79,7 @@ impl DrawResources {
fn viewport_buffer(device: &wgpu::Device) -> wgpu::Buffer {
device.create_buffer(&wgpu::BufferDescriptor {
label: Some("viewport"),
size: size_of::<shader::Viewport>() as u64,
size: size_of::<S::DrawSettings>() as u64,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::UNIFORM,
mapped_at_creation: false,
})
@ -105,7 +113,7 @@ impl DrawResources {
label: Some("compute"),
layout: Some(&pipeline_layout),
module: &module,
entry_point: Some("main_cs"),
entry_point: Some(S::compute_shader()),
compilation_options: Default::default(),
cache: None,
})
@ -136,7 +144,7 @@ impl DrawResources {
multisample: Default::default(),
fragment: Some(wgpu::FragmentState {
module,
entry_point: Some("main_fs"),
entry_point: Some(S::fragment_shader()),
compilation_options: Default::default(),
targets: &[Some((*format).into())],
}),
@ -145,7 +153,12 @@ impl DrawResources {
})
}
fn new(device: &wgpu::Device, format: &wgpu::TextureFormat, width: u64, height: u64) -> Self {
pub fn new(
device: &wgpu::Device,
format: &wgpu::TextureFormat,
width: u64,
height: u64,
) -> Self {
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);
@ -159,7 +172,6 @@ impl DrawResources {
let render_pipeline = Self::render_pipeline(device, &module, &bind_group_layout, format);
Self {
device: device.clone(),
bind_group_layout,
bind_group,
viewport_buffer,
@ -167,156 +179,18 @@ impl DrawResources {
image_size,
compute_pipeline,
render_pipeline,
settings: PhantomData,
}
}
fn resize(&mut self, width: u64, height: u64) {
self.image_buffer = Self::image_buffer(&self.device, width, height);
pub fn resize(&mut self, device: &wgpu::Device, width: u64, height: u64) {
self.image_buffer = Self::image_buffer(device, width, height);
self.image_size = glam::uvec2(width as u32, height as u32);
self.bind_group = Self::bind_group(
&self.device,
device,
&self.bind_group_layout,
&self.viewport_buffer,
&self.image_buffer,
);
}
}
struct DrawCallback {
draw_rect: egui::Rect,
draw_resize: bool,
}
impl CallbackTrait for DrawCallback {
fn prepare(
&self,
_device: &Device,
queue: &Queue,
_screen_descriptor: &ScreenDescriptor,
egui_encoder: &mut CommandEncoder,
callback_resources: &mut CallbackResources,
) -> Vec<CommandBuffer> {
let resources = callback_resources
.get_mut::<DrawResources>()
.expect("missing draw resources");
if self.draw_resize {
resources.resize(
self.draw_rect.size().x as u64,
self.draw_rect.size().y as u64,
);
}
let viewport = shader::Viewport {
image: resources.image_size,
offset: glam::uvec2(self.draw_rect.min.x as u32, self.draw_rect.min.y as u32),
size: glam::uvec2(
self.draw_rect.size().x as u32,
self.draw_rect.size().y as u32,
),
};
queue.write_buffer(
&resources.viewport_buffer,
0,
bytemuck::cast_slice(&[viewport]),
);
if self.draw_resize {
let mut compute_pass = egui_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("compute"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&resources.compute_pipeline);
compute_pass.set_bind_group(0, &resources.bind_group, &[]);
compute_pass.dispatch_workgroups(1, 1, 1);
}
vec![]
}
fn paint(
&self,
_info: PaintCallbackInfo,
render_pass: &mut RenderPass<'static>,
callback_resources: &CallbackResources,
) {
let resources = callback_resources.get::<DrawResources>().unwrap();
render_pass.set_pipeline(&resources.render_pipeline);
render_pass.set_bind_group(0, &resources.bind_group, &[]);
render_pass.draw(0..3, 0..1);
}
}
#[derive(Copy, Clone)]
struct ComputeDraw {
initial_draw: bool,
}
impl eframe::App for ComputeDraw {
fn update(&mut self, ctx: &Context, frame: &mut Frame) {
let initial_draw = self.initial_draw;
self.initial_draw = false;
if initial_draw {
let wgpu_render_state = frame.wgpu_render_state().expect("missing WGPU state");
let device = wgpu_render_state.device.clone();
let format = wgpu_render_state.target_format.clone();
let callback_resources = &mut wgpu_render_state
.renderer
.as_ref()
.write()
.callback_resources;
// Guess an initial size for initializing GPU resources, it will be adjusted later
callback_resources.insert(DrawResources::new(&device, &format, 800, 600));
}
egui::TopBottomPanel::bottom("bottom").show(ctx, |ui| {
let wgpu_render_state = frame.wgpu_render_state().expect("missing WGPU state");
let image_size = wgpu_render_state
.renderer
.as_ref()
.read()
.callback_resources
.get::<DrawResources>()
.unwrap()
.image_size;
ui.label(format!("Viewport: image={image_size}"))
});
egui::CentralPanel::default().show(ctx, |ui| {
egui::Frame::canvas(ui.style()).show(ui, |ui| {
let interact_rect = ui.available_rect_before_wrap();
let (response, painter) = ui.allocate_painter(interact_rect.size(), Sense::click());
let callback = DrawCallback {
draw_rect: interact_rect,
draw_resize: initial_draw || response.clicked(),
};
painter.add(egui_wgpu::Callback::new_paint_callback(
interact_rect,
callback,
))
});
});
}
}
fn main() {
let native_options = eframe::NativeOptions {
renderer: eframe::Renderer::Wgpu,
..Default::default()
};
eframe::run_native(
"Compute Draw",
native_options,
Box::new(|_cc| Ok(Box::new(ComputeDraw { initial_draw: true }))),
)
.unwrap()
}

View File

@ -0,0 +1,96 @@
use crate::draw_resources::ShaderSettings;
use glam::UVec2;
use shader::{DrawRect, DrawSized};
use wgpu::{Buffer, Queue};
pub struct ShaderBounding {
draw_size: egui::Vec2,
}
impl ShaderSettings for ShaderBounding {
type DrawSettings = DrawSized;
fn compute_shader() -> &'static str {
"main_cs_bounding"
}
fn fragment_shader() -> &'static str {
"main_fs_size"
}
fn new(interact_rect: egui::Rect) -> Self {
Self {
draw_size: interact_rect.size()
}
}
fn write_buffer(&self, queue: &Queue, buffer: &Buffer, image_size: glam::UVec2) {
let draw_settings = DrawSized {
image_size,
viewport_size: glam::uvec2(self.draw_size.x as u32, self.draw_size.y as u32),
};
queue.write_buffer(&buffer, 0, bytemuck::cast_slice(&[draw_settings]));
}
}
pub struct ShaderOffset {
draw_rect: egui::Rect,
}
impl ShaderSettings for ShaderOffset {
type DrawSettings = DrawRect;
fn compute_shader() -> &'static str {
"main_cs_bounding"
}
fn fragment_shader() -> &'static str {
"main_fs_rect"
}
fn new(interact_rect: egui::Rect) -> Self {
Self { draw_rect: interact_rect }
}
fn write_buffer(&self, queue: &Queue, buffer: &Buffer, image_size: UVec2) {
let viewport_size = glam::uvec2(self.draw_rect.size().x as u32, self.draw_rect.size().y as u32);
let viewport_offset = glam::uvec2(self.draw_rect.min.x as u32, self.draw_rect.min.y as u32);
let draw_settings = DrawRect {
image_size,
viewport_size,
viewport_offset
};
queue.write_buffer(&buffer, 0, bytemuck::cast_slice(&[draw_settings]));
}
}
pub struct ShaderOffsetBlocks {
draw_rect: egui::Rect,
}
impl ShaderSettings for ShaderOffsetBlocks {
type DrawSettings = DrawRect;
fn compute_shader() -> &'static str {
"main_cs_blocks"
}
fn fragment_shader() -> &'static str {
"main_fs_rect"
}
fn new(interact_rect: egui::Rect) -> Self {
Self { draw_rect: interact_rect }
}
fn write_buffer(&self, queue: &Queue, buffer: &Buffer, image_size: UVec2) {
let viewport_size = glam::uvec2(self.draw_rect.size().x as u32, self.draw_rect.size().y as u32);
let viewport_offset = glam::uvec2(self.draw_rect.min.x as u32, self.draw_rect.min.y as u32);
let draw_settings = DrawRect {
image_size,
viewport_size,
viewport_offset
};
queue.write_buffer(&buffer, 0, bytemuck::cast_slice(&[draw_settings]));
}
}

View File

@ -0,0 +1,158 @@
mod draw_resources;
pub mod draw_shaders;
use std::marker::PhantomData;
use eframe::Frame;
use egui::{Context, Sense};
use crate::draw_resources::{DrawResources, ShaderSettings};
use eframe::epaint::PaintCallbackInfo;
use egui_wgpu::{CallbackResources, CallbackTrait, ScreenDescriptor};
use wgpu::{CommandBuffer, CommandEncoder, Device, Queue, RenderPass};
struct DrawCallback<S> {
interact_rect: egui::Rect,
interact_resize: bool,
settings: S,
}
impl<S: ShaderSettings> DrawCallback<S> {
fn new(interact_rect: egui::Rect, interact_resize: bool) -> Self {
Self {
interact_rect,
interact_resize,
settings: S::new(interact_rect),
}
}
}
impl<S: ShaderSettings + Send + Sync + 'static> CallbackTrait for DrawCallback<S> {
fn prepare(
&self,
device: &Device,
queue: &Queue,
_screen_descriptor: &ScreenDescriptor,
egui_encoder: &mut CommandEncoder,
callback_resources: &mut CallbackResources,
) -> Vec<CommandBuffer> {
let resources = callback_resources
.get_mut::<DrawResources<S>>()
.expect("missing draw resources");
if self.interact_resize {
resources.resize(
device,
self.interact_rect.width() as u64,
self.interact_rect.height() as u64,
);
}
self.settings.write_buffer(queue, &resources.viewport_buffer, resources.image_size);
if self.interact_resize {
let mut compute_pass = egui_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("compute"),
timestamp_writes: None,
});
compute_pass.set_pipeline(&resources.compute_pipeline);
compute_pass.set_bind_group(0, &resources.bind_group, &[]);
compute_pass.dispatch_workgroups(1, 1, 1);
}
vec![]
}
fn paint(
&self,
_info: PaintCallbackInfo,
render_pass: &mut RenderPass<'static>,
callback_resources: &CallbackResources,
) {
let resources = callback_resources.get::<DrawResources<S>>().unwrap();
render_pass.set_pipeline(&resources.render_pipeline);
render_pass.set_bind_group(0, &resources.bind_group, &[]);
render_pass.draw(0..3, 0..1);
}
}
#[derive(Copy, Clone)]
pub struct ComputeDraw<S> {
initial_draw: bool,
settings: PhantomData<S>,
}
impl <S: ShaderSettings> ComputeDraw<S> {
pub fn new() -> Self {
ComputeDraw {
initial_draw: true,
settings: PhantomData,
}
}
}
impl<S: ShaderSettings + 'static> eframe::App for ComputeDraw<S> {
fn update(&mut self, ctx: &Context, frame: &mut Frame) {
let initial_draw = self.initial_draw;
self.initial_draw = false;
if initial_draw {
let wgpu_render_state = frame.wgpu_render_state().expect("missing WGPU state");
let device = wgpu_render_state.device.clone();
let format = wgpu_render_state.target_format.clone();
let callback_resources = &mut wgpu_render_state
.renderer
.as_ref()
.write()
.callback_resources;
// Guess an initial size for initializing GPU resources, it will be adjusted later
callback_resources.insert(DrawResources::<S>::new(&device, &format, 800, 600));
}
/*
egui::TopBottomPanel::bottom("bottom").show(ctx, |ui| {
let wgpu_render_state = frame.wgpu_render_state().expect("missing WGPU state");
let image_size = wgpu_render_state
.renderer
.as_ref()
.read()
.callback_resources
.get::<DrawResources<DrawSimple>>()
.unwrap()
.image_size;
ui.label(format!("Viewport: image={image_size}"))
});
*/
egui::CentralPanel::default().show(ctx, |ui| {
egui::Frame::canvas(ui.style()).show(ui, |ui| {
let interact_rect = ui.available_rect_before_wrap();
let (response, painter) = ui.allocate_painter(interact_rect.size(), Sense::click());
painter.add(egui_wgpu::Callback::new_paint_callback(
interact_rect,
DrawCallback::<S>::new(interact_rect, initial_draw || response.clicked()),
))
});
});
}
}
/*
fn main() {
let native_options = eframe::NativeOptions {
renderer: eframe::Renderer::Wgpu,
..Default::default()
};
eframe::run_native(
"Compute Draw",
native_options,
Box::new(|_cc| Ok(Box::new(ComputeDraw { initial_draw: true }))),
)
.unwrap()
}
*/