aboutsummaryrefslogtreecommitdiff
path: root/flowy/src/motion
diff options
context:
space:
mode:
Diffstat (limited to 'flowy/src/motion')
-rw-r--r--flowy/src/motion/dec.rs33
-rw-r--r--flowy/src/motion/enc.rs156
-rw-r--r--flowy/src/motion/enc.wgsl49
-rw-r--r--flowy/src/motion/mod.rs150
4 files changed, 388 insertions, 0 deletions
diff --git a/flowy/src/motion/dec.rs b/flowy/src/motion/dec.rs
new file mode 100644
index 0000000..84aeff7
--- /dev/null
+++ b/flowy/src/motion/dec.rs
@@ -0,0 +1,33 @@
+pub struct MotionDecoder {}
+
+/*
+BindGroupLayoutEntry {
+ binding: 0,
+ count: None,
+ ty: wgpu::BindingType::StorageTexture {
+ access: wgpu::StorageTextureAccess::WriteOnly,
+ format: TextureFormat::Bgra8Unorm,
+ view_dimension: wgpu::TextureViewDimension::D2,
+ },
+ visibility: ShaderStages::COMPUTE,
+ },
+ BindGroupLayoutEntry {
+ binding: 1,
+ count: None,
+ ty: BindingType::Texture {
+ sample_type: TextureSampleType::Float { filterable: false },
+ view_dimension: TextureViewDimension::D2,
+ multisampled: false,
+ },
+ visibility: ShaderStages::COMPUTE,
+ },
+ BindGroupLayoutEntry {
+ binding: 2,
+ count: None,
+ ty: BindingType::Texture {
+ sample_type: TextureSampleType::Float { filterable: false },
+ view_dimension: TextureViewDimension::D2,
+ multisampled: false,
+ },
+ visibility: ShaderStages::COMPUTE,
+ }, */
diff --git a/flowy/src/motion/enc.rs b/flowy/src/motion/enc.rs
new file mode 100644
index 0000000..aa45d7b
--- /dev/null
+++ b/flowy/src/motion/enc.rs
@@ -0,0 +1,156 @@
+use std::mem::size_of;
+
+use super::{CommonBuffers, Params, RoundParams};
+use bytemuck::{Pod, Zeroable};
+use wgpu::{
+ include_wgsl, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
+ BindGroupLayoutEntry, BindingType, Buffer, BufferDescriptor, BufferUsages, CommandEncoder,
+ ComputePipeline, ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue,
+ ShaderStages, TextureFormat, TextureSampleType, TextureViewDimension,
+};
+
+pub struct MotionEncoder {
+ pipeline: ComputePipeline,
+ bind_groups: [BindGroup; 2],
+
+ uniform_buffer: Buffer,
+ uniform: EncoderUniform,
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable, Default)]
+pub struct EncoderUniform {
+ block_size: [i32; 2],
+ output_stride: i32,
+ _pad: u32,
+}
+
+impl MotionEncoder {
+ pub fn create(device: &Device, params: &Params, bufs: &CommonBuffers) -> Self {
+ let uniform_buffer = device.create_buffer(&BufferDescriptor {
+ label: Some("encoder uniforms"),
+ size: size_of::<EncoderUniform>() as u64,
+ usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
+ mapped_at_creation: false,
+ });
+ let uniform = EncoderUniform {
+ block_size: [8, 8],
+ output_stride: (params.width / params.block_width) as i32,
+ ..Default::default()
+ };
+
+ let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+ label: None,
+ entries: &[
+ BindGroupLayoutEntry {
+ binding: 0,
+ count: None,
+ ty: BindingType::Buffer {
+ ty: wgpu::BufferBindingType::Uniform,
+ has_dynamic_offset: false,
+ min_binding_size: None,
+ },
+ visibility: ShaderStages::COMPUTE,
+ },
+ BindGroupLayoutEntry {
+ binding: 1,
+ count: None,
+ ty: BindingType::Buffer {
+ has_dynamic_offset: false,
+ min_binding_size: None,
+ ty: wgpu::BufferBindingType::Storage { read_only: false },
+ },
+ visibility: ShaderStages::COMPUTE,
+ },
+ BindGroupLayoutEntry {
+ binding: 2,
+ count: None,
+ ty: BindingType::Texture {
+ sample_type: TextureSampleType::Float { filterable: false },
+ view_dimension: TextureViewDimension::D2,
+ multisampled: false,
+ },
+ visibility: ShaderStages::COMPUTE,
+ },
+ BindGroupLayoutEntry {
+ binding: 3,
+ count: None,
+ ty: BindingType::Texture {
+ sample_type: TextureSampleType::Float { filterable: false },
+ view_dimension: TextureViewDimension::D2,
+ multisampled: false,
+ },
+ visibility: ShaderStages::COMPUTE,
+ },
+ ],
+ });
+
+ let bind_groups = [0, 1].map(|i| {
+ device.create_bind_group(&BindGroupDescriptor {
+ label: None,
+ layout: &bind_group_layout,
+ entries: &[
+ BindGroupEntry {
+ binding: 0,
+ resource: uniform_buffer.as_entire_binding(),
+ },
+ BindGroupEntry {
+ binding: 1,
+ resource: bufs.offsets.as_entire_binding(),
+ },
+ BindGroupEntry {
+ binding: 2,
+ resource: wgpu::BindingResource::TextureView(
+ &bufs.textures[1 - i].create_view(&Default::default()),
+ ),
+ },
+ BindGroupEntry {
+ binding: 3,
+ resource: wgpu::BindingResource::TextureView(
+ &bufs.textures[1 - i].create_view(&Default::default()),
+ ),
+ },
+ ],
+ })
+ });
+
+ let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
+ label: None,
+ bind_group_layouts: &[&bind_group_layout],
+ push_constant_ranges: &[],
+ });
+
+ let module = device.create_shader_module(include_wgsl!("enc.wgsl"));
+ let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
+ label: None,
+ layout: Some(&pipeline_layout),
+ module: &module,
+ entry_point: "main",
+ });
+
+ Self {
+ bind_groups,
+ uniform,
+ uniform_buffer,
+ pipeline,
+ }
+ }
+
+ pub fn write_uniforms(&self, queue: Queue) {
+ queue.write_buffer(
+ &self.uniform_buffer,
+ 0,
+ bytemuck::cast_slice(&[self.uniform]),
+ )
+ }
+
+ pub fn pass(&self, encoder: &mut CommandEncoder, params: &Params, rp: &RoundParams) {
+ let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+ label: None,
+ timestamp_writes: None,
+ });
+ cpass.set_pipeline(&self.pipeline);
+ cpass.set_bind_group(0, &self.bind_groups[rp.swap], &[]);
+ cpass.dispatch_workgroups(params.extent.width, params.extent.height, 1);
+ }
+}
diff --git a/flowy/src/motion/enc.wgsl b/flowy/src/motion/enc.wgsl
new file mode 100644
index 0000000..cc185e9
--- /dev/null
+++ b/flowy/src/motion/enc.wgsl
@@ -0,0 +1,49 @@
+
+struct Params {
+ block_size: vec2<i32>,
+ output_stride: i32
+}
+
+struct BlockOffset {
+ score: f32,
+ offset: vec2<i32>,
+}
+
+@group(0) @binding(0) var<uniform> params: Params;
+@group(0) @binding(1) var<storage, read_write> output: BlockOffset;
+@group(0) @binding(2) var prev: texture_2d<f32>;
+@group(0) @binding(3) var next: texture_2d<f32>;
+
+@compute @workgroup_size(1) fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+ let uv = vec2<i32>(global_id.xy) * params.block_size;
+
+ // output[global_id.x + global_id.y * params.output_stride] = BlockOffset(0., uv);
+
+ // let orig = textureLoad(texa, uv, 0);
+
+ // var best_sim = 1000.;
+ // var best_coord = vec2(0,0);
+ // for (var x = -SDIST; x <= SDIST; x++) {
+ // for (var y = -SDIST; y <= SDIST; y++) {
+ // let ov = uv + vec2(x,y);
+ // let samp = textureLoad(texb, ov, 0);
+ // let d = distance(samp.rgb, orig.rgb)
+ // + distance(sampnn.rgb, orignn.rgb)
+ // + distance(sampnp.rgb, orignp.rgb)
+ // + distance(samppn.rgb, origpn.rgb)
+ // + distance(samppp.rgb, origpp.rgb);
+ // if d < best_sim {
+ // best_sim = d;
+ // best_coord = vec2(x,y);
+ // }
+ // }}
+
+ // let col = vec4(colormap_vec(vec2<f32>(best_coord) / f32(SDIST)), 1.);
+
+ // textureStore(output, global_id.xy, col);
+}
+
+// fn colormap_vec(v: vec2<f32>) -> vec3<f32> {
+// return vec3(v.y, v.x - 0.5 * v.y, -v.x - 0.5 * v.y);
+// }
+
diff --git a/flowy/src/motion/mod.rs b/flowy/src/motion/mod.rs
new file mode 100644
index 0000000..acb1396
--- /dev/null
+++ b/flowy/src/motion/mod.rs
@@ -0,0 +1,150 @@
+pub mod dec;
+pub mod enc;
+
+use std::mem::size_of;
+use wgpu::{
+ Buffer, BufferUsages, CommandEncoder, Device, Extent3d, ImageCopyTexture, Origin3d, Queue,
+ Texture, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
+};
+
+pub struct Params {
+ pub width: usize,
+ pub height: usize,
+ pub extent: Extent3d,
+ pub block_width: usize,
+ pub block_height: usize,
+ pub blocks: usize,
+}
+
+pub struct RoundParams {
+ pub swap: usize,
+}
+
+pub struct CommonBuffers {
+ textures: [Texture; 2],
+ offsets: Buffer,
+ offsets_download: Option<Buffer>,
+ debug_output: Option<Texture>,
+ texture_download: Option<Buffer>,
+}
+
+#[repr(C)]
+pub struct BlockOffset {
+ score: f32,
+ _pad: u32,
+ offset: [f32; 2],
+}
+
+impl CommonBuffers {
+ pub fn create(device: &Device, params: &Params) -> Self {
+ let textures = [(), ()].map(|_| {
+ device.create_texture(&TextureDescriptor {
+ label: None,
+ size: params.extent,
+ mip_level_count: 1,
+ sample_count: 1,
+ dimension: TextureDimension::D2,
+ format: TextureFormat::Bgra8Unorm,
+ usage: TextureUsages::TEXTURE_BINDING
+ | TextureUsages::COPY_DST
+ | TextureUsages::COPY_SRC,
+ view_formats: &[],
+ })
+ });
+
+ let debug_output = Some(device.create_texture(&TextureDescriptor {
+ label: None,
+ size: params.extent,
+ mip_level_count: 1,
+ sample_count: 1,
+ dimension: TextureDimension::D2,
+ format: TextureFormat::Bgra8Unorm,
+ usage: TextureUsages::STORAGE_BINDING | TextureUsages::COPY_SRC,
+ view_formats: &[],
+ }));
+
+ let texture_download = Some(device.create_buffer(&wgpu::BufferDescriptor {
+ label: None,
+ size: (params.width * params.height * 4) as u64,
+ usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
+ mapped_at_creation: false,
+ }));
+
+ let offsets = device.create_buffer(&wgpu::BufferDescriptor {
+ label: None,
+ size: (params.blocks * size_of::<BlockOffset>()) as u64,
+ usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
+ mapped_at_creation: false,
+ });
+ let offsets_download = Some(device.create_buffer(&wgpu::BufferDescriptor {
+ label: None,
+ size: (params.blocks * size_of::<BlockOffset>()) as u64,
+ usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
+ mapped_at_creation: false,
+ }));
+
+ Self {
+ debug_output,
+ textures,
+ offsets_download,
+ offsets,
+ texture_download,
+ }
+ }
+
+ pub fn upload(&self, queue: &Queue, params: &Params, rp: &RoundParams, buffer: &[u8]) {
+ queue.write_texture(
+ ImageCopyTexture {
+ aspect: wgpu::TextureAspect::All,
+ mip_level: 0,
+ origin: Origin3d::ZERO,
+ texture: &self.textures[rp.swap],
+ },
+ buffer,
+ wgpu::ImageDataLayout {
+ offset: 0,
+ bytes_per_row: Some(params.extent.width * 4),
+ rows_per_image: Some(params.extent.height),
+ },
+ params.extent,
+ );
+ }
+
+ pub fn prepare_texture_download(
+ &self,
+ encoder: &mut CommandEncoder,
+ params: &Params,
+ rp: &RoundParams,
+ ) {
+ encoder.copy_texture_to_buffer(
+ wgpu::ImageCopyTexture {
+ texture: &self.textures[rp.swap],
+ mip_level: 0,
+ origin: wgpu::Origin3d::ZERO,
+ aspect: wgpu::TextureAspect::All,
+ },
+ wgpu::ImageCopyBuffer {
+ buffer: self.texture_download.as_ref().unwrap(),
+ layout: wgpu::ImageDataLayout {
+ offset: 0,
+ bytes_per_row: Some(params.extent.width * 4),
+ rows_per_image: Some(params.extent.height),
+ },
+ },
+ params.extent,
+ );
+ }
+
+ pub fn download(&self, device: &Device, buffer: &mut [u8]) {
+ let buffer_slice = self.texture_download.as_ref().unwrap().slice(..);
+ let (sender, receiver) = oneshot::channel();
+ buffer_slice.map_async(wgpu::MapMode::Read, move |r| sender.send(r).unwrap());
+ device.poll(wgpu::Maintain::Wait);
+ receiver.recv().unwrap().unwrap();
+ {
+ let view = buffer_slice.get_mapped_range();
+ buffer.copy_from_slice(&view[..]);
+ }
+ self.texture_download.as_ref().unwrap().unmap();
+ }
+}