1 files changed, 73 insertions, 46 deletions

diff --git a/flowy/src/motion/enc.wgsl b/flowy/src/motion/enc.wgsl
index b94bac9..46fe363 100644
--- a/flowy/src/motion/enc.wgsl
+++ b/flowy/src/motion/enc.wgsl
@@ -17,71 +17,98 @@ struct BlockOffset {
 @group(0) @binding(2) var next: texture_2d<f32>;
 @group(0) @binding(3) var prev: texture_2d<f32>;
 
-@compute @workgroup_size(1) fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+var<private> best_offset: vec2<i32> = vec2(0);
+var<private> best_error: f32 = 100000.;
+var<private> best_tint: vec3<f32> = vec3(0.);
+
+@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 _f = search_offset(uv);
-    let best_err = _f.error;
-    let best_offset = _f.offset;
+    test_offset(uv, vec2(0, 0));
+
+    do_dist(uv, 32);
+    do_dist(uv, 16);
+    do_dist(uv, 10);
+    do_dist(uv, 8);
+    do_dist(uv, 6);
+    do_dist(uv, 3);
+    do_dist(uv, 4);
+    do_dist(uv, 2);
+    do_dist(uv, 1);
+
+    apply_tint(uv);
 
-    var best_tint = vec3(0.);
+    offsets[global_id.x + global_id.y * params.offsets_stride] = BlockOffset(best_error, best_offset, best_tint);
+}
+
+fn do_dist(uv: vec2<i32>, n: i32) {
+    test_offset(uv, vec2(0, n));
+    test_offset(uv, vec2(n, n));
+    test_offset(uv, vec2(n, 0));
+    test_offset(uv, vec2(n, -n));
+    test_offset(uv, vec2(0, -n));
+    test_offset(uv, vec2(-n, -n));
+    test_offset(uv, vec2(-n, 0));
+    test_offset(uv, vec2(-n, n));
+}
+
+fn apply_tint(uv: vec2<i32>) {
 
     var average_pcol = vec3(0.);
     var average_ncol = vec3(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);
-        average_pcol += textureLoad(prev, base+best_offset, 0).rgb;
-        average_ncol += textureLoad(next, base, 0).rgb;
-    }}
+        for (var y = 0; y < params.block_size.y; y++) {
+            let base = uv + vec2(x, y);
+            average_pcol += textureLoad(prev, base + best_offset, 0).rgb;
+            average_ncol += textureLoad(next, base, 0).rgb;
+        }
+    }
 
     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_err {
-        best_tint = tint;
-    }
-
-    offsets[global_id.x + global_id.y * params.offsets_stride] = BlockOffset(best_err, best_offset, best_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);
+// }
 
-struct SearchRes {offset: vec2<i32>, error: f32}
-
-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);
 
-        var err = 0.;
-        for (var x = 0; x < params.block_size.x; x++) {
+fn test_offset(uv: vec2<i32>, offset: vec2<i32>) {
+    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+offset, 0).rgb;
+            let base = uv + vec2(x, y);
+            let pcol = textureLoad(prev, base + offset, 0).rgb;
             let ncol = textureLoad(next, base, 0).rgb;
             err += distance(pcol, ncol);
-        }}
-        
-        if err < best_err {
-            best_err = err;
-            best_offset = offset;
-            if err < params.skip_threshold {
-                return SearchRes(offset, err);
-            }
         }
-    }}
-    return SearchRes(best_offset,best_err);
+    }
+    if err < best_error {
+        best_error = err;
+        best_offset = offset;
+    }
 }
 
+
 // fn colormap_vec(v: vec2<f32>) -> vec3<f32> {
 //     return vec3(v.y, v.x - 0.5 * v.y, -v.x - 0.5 * v.y);
 // }