From 5248831232fa22a1f3d6515f7f6c7bee8994faf2 Mon Sep 17 00:00:00 2001 From: metamuffin Date: Mon, 21 Nov 2022 16:59:50 +0100 Subject: unify repos --- vgcodec/src/app.rs | 37 +++++++++++++++ vgcodec/src/diff.rs | 120 +++++++++++++++++++++++++++++++++++++++++++++++++ vgcodec/src/diff.wgsl | 17 +++++++ vgcodec/src/export.rs | 80 +++++++++++++++++++++++++++++++++ vgcodec/src/main.rs | 76 +++++++++++++++++++++++++++++++ vgcodec/src/paint.rs | 120 +++++++++++++++++++++++++++++++++++++++++++++++++ vgcodec/src/paint.wgsl | 19 ++++++++ 7 files changed, 469 insertions(+) create mode 100644 vgcodec/src/app.rs create mode 100644 vgcodec/src/diff.rs create mode 100644 vgcodec/src/diff.wgsl create mode 100644 vgcodec/src/export.rs create mode 100644 vgcodec/src/main.rs create mode 100644 vgcodec/src/paint.rs create mode 100644 vgcodec/src/paint.wgsl (limited to 'vgcodec/src') diff --git a/vgcodec/src/app.rs b/vgcodec/src/app.rs new file mode 100644 index 0000000..0c063a5 --- /dev/null +++ b/vgcodec/src/app.rs @@ -0,0 +1,37 @@ +use std::sync::Arc; + +use wgpu::{Adapter, Device, Instance, Queue}; + +pub struct App { + pub instance: Instance, + pub device: Device, + pub adapter: Adapter, + pub queue: Queue, +} + +impl App { + pub async fn new() -> Arc { + let instance = wgpu::Instance::new(wgpu::Backends::all()); + let adapter = instance + .request_adapter(&wgpu::RequestAdapterOptions::default()) + .await + .unwrap(); + let (device, queue) = adapter + .request_device( + &wgpu::DeviceDescriptor { + label: None, + features: wgpu::Features::empty(), + limits: wgpu::Limits::downlevel_defaults(), + }, + None, + ) + .await + .unwrap(); + Arc::new(Self { + adapter, + device, + instance, + queue, + }) + } +} diff --git a/vgcodec/src/diff.rs b/vgcodec/src/diff.rs new file mode 100644 index 0000000..44efabe --- /dev/null +++ b/vgcodec/src/diff.rs @@ -0,0 +1,120 @@ +use std::{borrow::Cow, sync::Arc}; +use wgpu::{util::DeviceExt, BindGroup, Buffer, ComputePipeline, Extent3d, Texture}; + +use crate::app::App; + +pub struct Differ { + app: Arc, + size: Extent3d, + + pipeline: ComputePipeline, + bind_group: BindGroup, + + counter_buffer_size: u64, + counter_buffer: Buffer, + counter_staging_buffer: Buffer, +} + +impl Differ { + pub fn new(app: &Arc, extent: Extent3d, tex_a: &Texture, tex_b: &Texture) -> Self { + let App { device, .. } = app.as_ref(); + let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: None, + source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("diff.wgsl"))), + }); + + let counter_buffer_size = std::mem::size_of::() as wgpu::BufferAddress; + let counter_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: counter_buffer_size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let counter_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&[0u32]), + usage: wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_DST + | wgpu::BufferUsages::COPY_SRC, + }); + + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &cs_module, + entry_point: "main", + }); + + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView( + &tex_a.create_view(&wgpu::TextureViewDescriptor::default()), + ), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::TextureView( + &tex_b.create_view(&wgpu::TextureViewDescriptor::default()), + ), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: counter_buffer.as_entire_binding(), + }, + ], + }); + Self { + app: app.clone(), + size: extent, + pipeline, + counter_buffer_size, + counter_buffer, + bind_group, + counter_staging_buffer, + } + } + + pub async fn run(&mut self) -> u32 { + let App { device, queue, .. } = self.app.as_ref(); + + let mut encoder = + device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + { + let mut cpass = + encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None }); + cpass.set_pipeline(&self.pipeline); + cpass.set_bind_group(0, &self.bind_group, &[]); + cpass.dispatch_workgroups(self.size.width, self.size.height, 1); + } + + encoder.copy_buffer_to_buffer( + &self.counter_buffer, + 0, + &self.counter_staging_buffer, + 0, + self.counter_buffer_size, + ); + + queue.submit(Some(encoder.finish())); + + let buffer_slice = self.counter_staging_buffer.slice(..); + let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel(); + buffer_slice.map_async(wgpu::MapMode::Read, move |v| { + sender.send(v.unwrap()).unwrap() + }); + device.poll(wgpu::Maintain::Wait); + receiver.receive().await; + + let data = buffer_slice.get_mapped_range(); + let result: u32 = bytemuck::cast_slice(&data).to_vec()[0]; + + drop(data); + self.counter_staging_buffer.unmap(); + + result + } +} diff --git a/vgcodec/src/diff.wgsl b/vgcodec/src/diff.wgsl new file mode 100644 index 0000000..bb5a3e1 --- /dev/null +++ b/vgcodec/src/diff.wgsl @@ -0,0 +1,17 @@ +@group(0) @binding(0) +var tex_a: texture_2d; +@group(0) @binding(1) +var tex_b: texture_2d; + +@group(0) @binding(2) +var exp: atomic; + +@compute @workgroup_size(1) +fn main(@builtin(global_invocation_id) global_id: vec3) { + var col_a = textureLoad(tex_a, vec2(i32(global_id.x), i32(global_id.y)), 0); + var col_b = textureLoad(tex_b, vec2(i32(global_id.x), i32(global_id.y)), 0); + var diff = col_a - col_b; + var diffsum = diff.r + diff.g + diff.b; + atomicAdd(&exp, u32(diffsum)); +} + diff --git a/vgcodec/src/export.rs b/vgcodec/src/export.rs new file mode 100644 index 0000000..7b02085 --- /dev/null +++ b/vgcodec/src/export.rs @@ -0,0 +1,80 @@ +use crate::app::App; +use image::RgbaImage; +use std::{num::NonZeroU32, sync::Arc}; +use wgpu::{ + Buffer, Extent3d, ImageCopyBuffer, ImageCopyTexture, ImageDataLayout, Origin3d, Texture, +}; + +pub struct Exporter { + app: Arc, + size: Extent3d, + + padded_bytes_per_row: u32, + export_buffer: Buffer, +} + +impl Exporter { + pub fn new(app: &Arc, size: Extent3d) -> Self { + let App { device, .. } = app.as_ref(); + + let bytes_per_pixel = std::mem::size_of::() as u32; + let unpadded_bytes_per_row = size.width * bytes_per_pixel; + let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT; + let padded_bytes_per_row_padding = (align - unpadded_bytes_per_row % align) % align; + let padded_bytes_per_row = unpadded_bytes_per_row + padded_bytes_per_row_padding; + + let export_buffer_size = (padded_bytes_per_row * size.height) as u64; + let export_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: export_buffer_size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + Self { + padded_bytes_per_row, + app: app.clone(), + size, + export_buffer, + } + } + pub async fn run(&self, texture: &Texture, save_path: &str) { + let App { device, queue, .. } = self.app.as_ref(); + + let mut encoder = + device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + + encoder.copy_texture_to_buffer( + ImageCopyTexture { + texture, + aspect: wgpu::TextureAspect::All, + mip_level: 0, + origin: Origin3d::ZERO, + }, + ImageCopyBuffer { + buffer: &self.export_buffer, + layout: ImageDataLayout { + offset: 0, + bytes_per_row: Some(NonZeroU32::new(self.padded_bytes_per_row).unwrap()), + rows_per_image: None, + }, + }, + self.size, + ); + + queue.submit(Some(encoder.finish())); + + let buffer_slice = self.export_buffer.slice(..); + let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel(); + buffer_slice.map_async(wgpu::MapMode::Read, move |v| { + sender.send(v.unwrap()).unwrap() + }); + device.poll(wgpu::Maintain::Wait); + receiver.receive().await; + + let data = buffer_slice.get_mapped_range(); + let result: Vec = bytemuck::cast_slice(&data).to_vec(); + let image = RgbaImage::from_raw(self.size.width, self.size.height, result).unwrap(); + image.save(save_path).unwrap(); + } +} diff --git a/vgcodec/src/main.rs b/vgcodec/src/main.rs new file mode 100644 index 0000000..9450894 --- /dev/null +++ b/vgcodec/src/main.rs @@ -0,0 +1,76 @@ +pub mod app; +pub mod diff; +pub mod export; +pub mod paint; + +use app::App; +use diff::Differ; +use image::EncodableLayout; +use wgpu::{Extent3d, Queue, Texture, TextureUsages}; + +use crate::export::Exporter; + +fn main() { + env_logger::init_from_env("LOG"); + pollster::block_on(run()); +} + +async fn run() { + let app = app::App::new().await; + + let App { device, queue, .. } = app.as_ref(); + + let img_target = image::open("a/a.png").unwrap().into_rgba8(); + let img_initial = image::open("a/initial.png").unwrap().into_rgba8(); + + let size = wgpu::Extent3d { + width: img_initial.width(), + height: img_initial.height(), + depth_or_array_layers: 1, + }; + + let create_texture = || { + device.create_texture(&wgpu::TextureDescriptor { + size, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: wgpu::TextureFormat::Rgba8Unorm, + usage: TextureUsages::COPY_DST + | TextureUsages::TEXTURE_BINDING + | TextureUsages::COPY_SRC, + label: None, + }) + }; + + let tex_target = create_texture(); + let tex_approx = create_texture(); + + write_texture(queue, &tex_target, img_target.as_bytes(), size); + write_texture(queue, &tex_approx, img_initial.as_bytes(), size); + + let mut differ = Differ::new(&app, size, &tex_approx, &tex_target); + let exporter = Exporter::new(&app, size); + + println!("{}", differ.run().await); + + exporter.run(&tex_approx, "a/approx.png").await; +} + +pub fn write_texture(queue: &Queue, target: &Texture, data: &[u8], size: Extent3d) { + queue.write_texture( + wgpu::ImageCopyTexture { + texture: &target, + mip_level: 0, + origin: wgpu::Origin3d::ZERO, + aspect: wgpu::TextureAspect::All, + }, + &data, + wgpu::ImageDataLayout { + offset: 0, + bytes_per_row: Some(std::num::NonZeroU32::try_from((size.width * 4) as u32).unwrap()), + rows_per_image: None, + }, + size, + ); +} diff --git a/vgcodec/src/paint.rs b/vgcodec/src/paint.rs new file mode 100644 index 0000000..3da0647 --- /dev/null +++ b/vgcodec/src/paint.rs @@ -0,0 +1,120 @@ +use std::{borrow::Cow, sync::Arc}; +use wgpu::{util::DeviceExt, BindGroup, Buffer, ComputePipeline, Extent3d, Texture}; + +use crate::app::App; + +pub struct Painter { + app: Arc, + size: Extent3d, + + pipeline: ComputePipeline, + bind_group: BindGroup, + + uniform_buffer_size: u64, + uniform_buffer: Buffer, + uniform_staging_buffer: Buffer, +} + +impl Painter { + pub fn new(app: &Arc, extent: Extent3d, tex_a: &Texture, tex_b: &Texture) -> Self { + let App { device, .. } = app.as_ref(); + let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: None, + source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("paint.wgsl"))), + }); + + let uniform_buffer_size = std::mem::size_of::() as wgpu::BufferAddress; + let uniform_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: uniform_buffer_size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&[0u32]), + usage: wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_DST + | wgpu::BufferUsages::COPY_SRC, + }); + + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &cs_module, + entry_point: "main", + }); + + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView( + &tex_a.create_view(&wgpu::TextureViewDescriptor::default()), + ), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: wgpu::BindingResource::TextureView( + &tex_b.create_view(&wgpu::TextureViewDescriptor::default()), + ), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: uniform_buffer.as_entire_binding(), + }, + ], + }); + Self { + app: app.clone(), + size: extent, + pipeline, + uniform_buffer_size, + uniform_buffer, + bind_group, + uniform_staging_buffer, + } + } + + pub async fn run(&mut self) -> u32 { + let App { device, queue, .. } = self.app.as_ref(); + + let mut encoder = + device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + { + let mut cpass = + encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None }); + cpass.set_pipeline(&self.pipeline); + cpass.set_bind_group(0, &self.bind_group, &[]); + cpass.dispatch_workgroups(self.size.width, self.size.height, 1); + } + + encoder.copy_buffer_to_buffer( + &self.uniform_buffer, + 0, + &self.uniform_staging_buffer, + 0, + self.uniform_buffer_size, + ); + + queue.submit(Some(encoder.finish())); + + let buffer_slice = self.uniform_staging_buffer.slice(..); + let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel(); + buffer_slice.map_async(wgpu::MapMode::Read, move |v| { + sender.send(v.unwrap()).unwrap() + }); + device.poll(wgpu::Maintain::Wait); + receiver.receive().await; + + let data = buffer_slice.get_mapped_range(); + let result: u32 = bytemuck::cast_slice(&data).to_vec()[0]; + + drop(data); + self.uniform_staging_buffer.unmap(); + + result + } +} diff --git a/vgcodec/src/paint.wgsl b/vgcodec/src/paint.wgsl new file mode 100644 index 0000000..02acd8e --- /dev/null +++ b/vgcodec/src/paint.wgsl @@ -0,0 +1,19 @@ +struct Uniforms { + x: f32, + y: f32, + radius: f32, + r: f32, + g: f32, + b: f32 +}; + +@group(0) @binding(0) var uniforms: Uniforms; +@group(0) @binding(1) var tex: texture_storage_2d; + +@compute @workgroup_size(1) +fn main(@builtin(global_invocation_id) global_id: vec3) { + let coords = vec2(global_id.xy); + if distance(vec2(coords), vec2(uniforms.x, uniforms.y)) < uniforms.radius { + textureStore(tex, coords.xy, vec4(vec3(uniforms.r, uniforms.g, uniforms.b), 1.0)); + } +} -- cgit v1.2.3-70-g09d2