debug view

7 files changed, 259 insertions, 74 deletions

diff --git a/flowy/src/main.rs b/flowy/src/main.rs
index a0a4024..6620ed8 100644
--- a/flowy/src/main.rs
+++ b/flowy/src/main.rs
@@ -13,7 +13,7 @@ use wgpu::{
     MaintainBase, PowerPreference, RequestAdapterOptions,
 };
 
-use crate::motion::RoundParams;
+use crate::motion::{debug::MotionDebugger, RoundParams};
 
 fn main() {
     env_logger::init_from_env("LOG");
@@ -38,7 +38,7 @@ fn main() {
         .request_device(
             &DeviceDescriptor {
                 label: None,
-                features: Features::BGRA8UNORM_STORAGE,
+                features: Features::empty(),
                 limits: Limits::default(),
             },
             None,
@@ -46,7 +46,9 @@ fn main() {
         .block_on()
         .unwrap();
 
-    let (width, height) = (1920, 1080);
+    let (mut framework, fparams) = Framework::init();
+
+    let (width, height) = (fparams.width, fparams.height);
     let bsize = 8;
     let params = Params {
         width,
@@ -61,24 +63,30 @@ fn main() {
         block_width: bsize,
         block_height: bsize,
         blocks: (width / bsize) * (height / bsize),
+        init_debug: true,
     };
 
     let bufs = CommonBuffers::create(&device, &params);
     let menc = MotionEncoder::create(&device, &params, &bufs);
     let mdec = MotionDecoder::create(&device, &params, &bufs);
+    let mdeb = MotionDebugger::create(&device, &params, &bufs);
 
     let mut buffer = vec![0u8; width * height * 4];
 
-    let (mut framework, fparams) = Framework::init();
-
     menc.write_uniforms(&queue);
     mdec.write_uniforms(&queue);
+    mdeb.write_uniforms(&queue);
 
     info!("ready (setup took {:?})", t.elapsed());
-    
+
     let mut i = 0;
     loop {
-        let rp = RoundParams { swap: i % 2 };
+        let rp = RoundParams {
+            swap: i % 2,
+            debug: fparams.debug == 2,
+            preview: fparams.debug > 0,
+        };
+        eprintln!("{params:?} {rp:?}");
         debug!("read");
         stdin().read_exact(&mut buffer).unwrap();
         framework.next_frame_manual();
@@ -93,8 +101,13 @@ fn main() {
         if i > 1 {
             menc.pass(&mut encoder, &params, &rp);
             mdec.pass(&mut encoder, &params, &rp);
+            if rp.debug {
+                mdeb.pass(&mut encoder, &params, &rp);
+            }
+        }
+        if rp.preview {
+            bufs.prepare_texture_download(&mut encoder, &params, &rp);
         }
-        bufs.prepare_texture_download(&mut encoder, &params, &rp);
 
         queue.submit(Some(encoder.finish()));
         device.poll(MaintainBase::Wait);
diff --git a/flowy/src/motion/debug.rs b/flowy/src/motion/debug.rs
index 8b13789..9029b87 100644
--- a/flowy/src/motion/debug.rs
+++ b/flowy/src/motion/debug.rs
@@ -1 +1,175 @@
+use super::{CommonBuffers, Params, RoundParams};
+use bytemuck::{Pod, Zeroable};
+use std::mem::size_of;
+use wgpu::{
+    include_wgsl, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
+    BindGroupLayoutEntry, BindingType, Buffer, BufferDescriptor, BufferUsages, CommandEncoder,
+    ComputePipeline, ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue,
+    ShaderStages, TextureSampleType, TextureViewDimension,
+};
 
+pub struct MotionDebugger {
+    pipeline: ComputePipeline,
+    bind_groups: [BindGroup; 2],
+
+    uniform_buffer: Buffer,
+    uniform: DebuggerUniform,
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable, Default)]
+pub struct DebuggerUniform {
+    block_size: [i32; 2],
+    output_stride: i32,
+    _pad: u32,
+}
+
+impl MotionDebugger {
+    pub fn create(device: &Device, params: &Params, bufs: &CommonBuffers) -> Self {
+        let uniform_buffer = device.create_buffer(&BufferDescriptor {
+            label: Some("encoder uniforms"),
+            size: size_of::<DebuggerUniform>() as u64,
+            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+        let uniform = DebuggerUniform {
+            block_size: [params.block_width as i32, params.block_height as i32],
+            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: true },
+                    },
+                    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,
+                },
+                BindGroupLayoutEntry {
+                    binding: 4,
+                    count: None,
+                    ty: BindingType::StorageTexture {
+                        access: wgpu::StorageTextureAccess::WriteOnly,
+                        format: wgpu::TextureFormat::Rgba8Unorm,
+                        view_dimension: TextureViewDimension::D2,
+                    },
+                    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[i].create_view(&Default::default()),
+                        ),
+                    },
+                    BindGroupEntry {
+                        binding: 3,
+                        resource: wgpu::BindingResource::TextureView(
+                            &bufs.textures[1 - i].create_view(&Default::default()),
+                        ),
+                    },
+                    BindGroupEntry {
+                        binding: 4,
+                        resource: wgpu::BindingResource::TextureView(
+                            &bufs
+                                .debug_output
+                                .as_ref()
+                                .unwrap()
+                                .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!("debug.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.blocks_x as u32, params.blocks_y as u32, 1);
+    }
+}
diff --git a/flowy/src/motion/debug.wgsl b/flowy/src/motion/debug.wgsl
index e69de29..3e740f6 100644
--- a/flowy/src/motion/debug.wgsl
+++ b/flowy/src/motion/debug.wgsl
@@ -0,0 +1,38 @@
+
+
+struct Params {
+    block_size: vec2<i32>,
+    offsets_stride: u32
+}
+
+struct BlockOffset {
+    score: f32,
+    offset: vec2<i32>,
+    tint: vec3<f32>,
+}
+
+@group(0) @binding(0) var<uniform> params: Params;
+@group(0) @binding(1) var<storage, read> offsets: array<BlockOffset>;
+@group(0) @binding(2) var next: texture_2d<f32>;
+@group(0) @binding(3) var prev: texture_2d<f32>;
+@group(0) @binding(4) var out: texture_storage_2d<rgba8unorm, write>;
+
+@compute @workgroup_size(1)fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let uv = vec2<i32>(global_id.xy) * params.block_size;
+
+    let bl = offsets[global_id.x + global_id.y * params.offsets_stride];
+
+    for (var x = 0; x < params.block_size.x; x++) {
+        for (var y = 0; y < params.block_size.y; y++) {
+            let base = uv + vec2(x, y);
+            let col = vec4(colormap_vec(vec2<f32>(bl.offset) * 0.05),1.);
+            textureStore(out, base, 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/dec.rs b/flowy/src/motion/dec.rs
index 518fbd2..e62347c 100644
--- a/flowy/src/motion/dec.rs
+++ b/flowy/src/motion/dec.rs
@@ -66,7 +66,7 @@ impl MotionDecoder {
                     count: None,
                     ty: BindingType::StorageTexture {
                         access: wgpu::StorageTextureAccess::WriteOnly,
-                        format: wgpu::TextureFormat::Bgra8Unorm,
+                        format: wgpu::TextureFormat::Rgba8Unorm,
                         view_dimension: TextureViewDimension::D2,
                     },
                     visibility: ShaderStages::COMPUTE,
@@ -153,35 +153,3 @@ impl MotionDecoder {
         cpass.dispatch_workgroups(params.blocks_x as u32, params.blocks_y as u32, 1);
     }
 }
-
-/*
-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/dec.wgsl b/flowy/src/motion/dec.wgsl
index 2fa3f3b..f4db974 100644
--- a/flowy/src/motion/dec.wgsl
+++ b/flowy/src/motion/dec.wgsl
@@ -12,7 +12,7 @@ struct BlockOffset {
 
 @group(0) @binding(0) var<uniform> params: Params;
 @group(0) @binding(1) var<storage, read> offsets: array<BlockOffset>;
-@group(0) @binding(2) var next: texture_storage_2d<bgra8unorm, write>;
+@group(0) @binding(2) var next: texture_storage_2d<rgba8unorm, write>;
 @group(0) @binding(3) var prev: texture_2d<f32>;
 
 @compute @workgroup_size(1) fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
diff --git a/flowy/src/motion/enc.wgsl b/flowy/src/motion/enc.wgsl
index 46fe363..6ba3081 100644
--- a/flowy/src/motion/enc.wgsl
+++ b/flowy/src/motion/enc.wgsl
@@ -66,31 +66,19 @@ fn apply_tint(uv: vec2<i32>) {
 
     let tint = (average_ncol - average_pcol) / f32(params.block_size.x * params.block_size.y);
 
-    // var err = 0.;
-    // for (var x = 0; x < params.block_size.x; x++) {
-    //     for (var y = 0; y < params.block_size.y; y++) {
-    //         let base = uv + vec2(x, y);
-    //         let pcol = textureLoad(prev, base + best_offset, 0).rgb + tint;
-    //         let ncol = textureLoad(next, base, 0).rgb;
-    //         err += distance(pcol, ncol);
-    //     }
-    // }
-    // if err < best_error {
-    best_tint = tint;
-    // }
+    var err = 0.;
+    for (var x = 0; x < params.block_size.x; x++) {
+        for (var y = 0; y < params.block_size.y; y++) {
+            let base = uv + vec2(x, y);
+            let pcol = textureLoad(prev, base + best_offset, 0).rgb + tint;
+            let ncol = textureLoad(next, base, 0).rgb;
+            err += distance(pcol, ncol);
+        }
+    }
+    if err < best_error {
+        best_tint = tint;
+    }
 }
-// fn search_offset(uv: vec2<i32>) -> SearchRes {
-//     var best_err = 100000000.;
-//     var best_offset = vec2(0);
-//     // TODO: better ordering
-//     for (var ox = -params.search_radius; ox <= params.search_radius; ox++) {
-//         for (var oy = -params.search_radius; oy <= params.search_radius; oy++) {
-//             let offset = vec2(ox, oy);
-//         }
-//     }
-//     return SearchRes(best_offset, best_err);
-// }
-
 
 fn test_offset(uv: vec2<i32>, offset: vec2<i32>) {
     var err = 0.;
@@ -108,8 +96,3 @@ fn test_offset(uv: vec2<i32>, offset: vec2<i32>) {
     }
 }
 
-
-// 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
index dfdf7a6..cd53550 100644
--- a/flowy/src/motion/mod.rs
+++ b/flowy/src/motion/mod.rs
@@ -8,6 +8,7 @@ use wgpu::{
     Texture, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
 };
 
+#[derive(Debug)]
 pub struct Params {
     pub width: usize,
     pub height: usize,
@@ -17,10 +18,14 @@ pub struct Params {
     pub blocks_x: usize,
     pub blocks_y: usize,
     pub blocks: usize,
+    pub init_debug: bool,
 }
 
+#[derive(Debug)]
 pub struct RoundParams {
     pub swap: usize,
+    pub debug: bool,
+    pub preview: bool,
 }
 
 pub struct CommonBuffers {
@@ -48,7 +53,7 @@ impl CommonBuffers {
                 mip_level_count: 1,
                 sample_count: 1,
                 dimension: TextureDimension::D2,
-                format: TextureFormat::Bgra8Unorm,
+                format: TextureFormat::Rgba8Unorm,
                 usage: TextureUsages::TEXTURE_BINDING
                     | TextureUsages::STORAGE_BINDING
                     | TextureUsages::COPY_DST
@@ -63,7 +68,7 @@ impl CommonBuffers {
             mip_level_count: 1,
             sample_count: 1,
             dimension: TextureDimension::D2,
-            format: TextureFormat::Bgra8Unorm,
+            format: TextureFormat::Rgba8Unorm,
             usage: TextureUsages::STORAGE_BINDING | TextureUsages::COPY_SRC,
             view_formats: &[],
         }));
@@ -123,7 +128,11 @@ impl CommonBuffers {
     ) {
         encoder.copy_texture_to_buffer(
             wgpu::ImageCopyTexture {
-                texture: &self.textures[rp.swap],
+                texture: if rp.debug {
+                    self.debug_output.as_ref().unwrap()
+                } else {
+                    &self.textures[rp.swap]
+                },
                 mip_level: 0,
                 origin: wgpu::Origin3d::ZERO,
                 aspect: wgpu::TextureAspect::All,