Improve motion tracking withi NL4D and confidence variance

Cargo.toml

@@ -20,3 +20,8 @@ strum_macros = "0.28"
 bon = "3"
 cubecl = { version = "0.10", features = ["std"] }
 av-denoise-core = { path = "av-denoise-core", version = "0.4.0-alpha7", default-features = false }
+
+[profile.test]
+opt-level = 3
+debug-assertions = true
+overflow-checks = true

av-denoise-core/Cargo.toml

@@ -71,3 +71,7 @@ harness = false
 [[bench]]
 name = "reseed"
 harness = false
+
+[[bench]]
+name = "mc_accuracy"
+harness = false

av-denoise-core/benches/bench_kernels.rs

@@ -29,6 +29,7 @@ use kernels::mc_chain_compose::ChainComposeBench;
 use kernels::mc_confidence::McConfidenceBench;
 use kernels::mc_downscale::DownscaleBench;
 use kernels::mc_warp::WarpBench;
+use kernels::mv_regularise::MvRegulariseBench;
 use kernels::noise_partial::NoisePartialBench;
 use kernels::pack_wire::PackWireBench;
 use kernels::temporal_noise_stats::TemporalNoiseStatsBench;
@@ -220,11 +221,23 @@ fn run_all<R: Runtime>(backend: &str, device: &R::Device) {
     run(ChainComposeBench {
         client: client.clone(),
     });
+    run(MvRegulariseBench {
+        client: client.clone(),
+    });
+    for &(ch, ch_name) in CHANNELS {
+        run(CollabFusedBench {
+            client: client.clone(),
+            ch,
+            ch_name,
+            split_mv: false,
+        });
+    }
     for &(ch, ch_name) in CHANNELS {
         run(CollabFusedBench {
             client: client.clone(),
             ch,
             ch_name,
+            split_mv: true,
         });
     }
     for &(ch, ch_name) in CHANNELS {

av-denoise-core/benches/kernels/collab_fused.rs

@@ -14,13 +14,13 @@ use super::nl4d_geometry::{
     CONFIDENCE_VARIANCE,
     K_MAX,
     LAMBDA_HT,
+    MISMATCH_SCALE2,
     N_FRAMES,
     NEIGHBOUR_SLOTS,
     RADIUS,
     REFINE,
     SIGMA,
     SPATIAL_RADIUS,
-    THSAD,
     conf_stride,
     mv_stride,
 };
@@ -39,12 +39,37 @@ use super::{H, W, block_sync, make_padded_frame, shapes_with_ch, stored_channels
 /// block skips its comparisons entirely, so leaving it always open here
 /// measures the worst case. A bench that gates freely would report a
 /// time well under the real one.
+///
+/// `split_mv` picks which of the two motion fields the arm runs on, and
+/// the two bracket the real cost of the covering-block search.
+///
+/// `false` gives a zeroed field. Every block covering a patch then
+/// predicts the same position, all four rectangles coincide, three of
+/// them are dropped by the duplicate check and no extra pixel
+/// comparison runs. That arm measures the duplicate check on its own.
+///
+/// `true` gives each block a vector from its own grid parity, spaced
+/// eight pixels apart, which is further than the refine window is wide.
+/// The four rectangles covering a patch are then disjoint, nothing
+/// deduplicates, and the neighbour search scores four times the
+/// positions. That arm is the worst case, and a real motion field lands
+/// between the two.
 pub struct CollabFusedBench<R: Runtime> {
     pub client: ComputeClient<R>,
     pub ch: u32,
     pub ch_name: &'static str,
+    pub split_mv: bool,
 }
 
+/// How far apart two neighbouring blocks' vectors sit in the split
+/// field, in pixels.
+///
+/// `REFINE` is the rectangle's half-width, so two rectangles stay
+/// disjoint once their centres are more than `2 * REFINE` apart. Eight
+/// clears that with room and keeps every predicted position well inside
+/// a 1080p frame.
+const SPLIT_MV_SPACING: i32 = 8;
+
 #[derive(Clone)]
 pub struct CollabFusedInput {
     pub ring: Handle,
@@ -83,7 +108,18 @@ impl<R: Runtime> Benchmark for CollabFusedBench<R> {
         let mv_stride = mv_stride(blocks_x, blocks_y, align);
         let conf_stride = conf_stride(blocks_x, blocks_y, align);
 
-        let mv_data = vec![0i32; (2 * RADIUS * mv_stride) as usize];
+        let mut mv_data = vec![0i32; (2 * RADIUS * mv_stride) as usize];
+        if self.split_mv {
+            for t in 0..2 * RADIUS {
+                for by in 0..blocks_y {
+                    for bx in 0..blocks_x {
+                        let base = (t * mv_stride + (by * blocks_x + bx) * 2) as usize;
+                        mv_data[base] = (bx % 2) as i32 * SPLIT_MV_SPACING;
+                        mv_data[base + 1] = (by % 2) as i32 * SPLIT_MV_SPACING;
+                    }
+                }
+            }
+        }
         let mv_field = self.client.create_from_slice(i32::as_bytes(&mv_data));
         let conf_data = vec![1.0f32; (2 * RADIUS * conf_stride) as usize];
         let confidence = self.client.create_from_slice(f32::as_bytes(&conf_data));
@@ -159,7 +195,7 @@ impl<R: Runtime> Benchmark for CollabFusedBench<R> {
                 CENTRE_SLOT,
                 0.0f32,
                 0.0f32,
-                THSAD,
+                MISMATCH_SCALE2,
                 LAMBDA_HT,
                 weight_scale(SIGMA, &dct_noise_profile(0.0)),
                 cross_frame_accum_scale(SPATIAL_RADIUS, RADIUS),
@@ -185,7 +221,8 @@ impl<R: Runtime> Benchmark for CollabFusedBench<R> {
     }
 
     fn name(&self) -> String {
-        format!("collab_fused_1080p_{}", self.ch_name)
+        let field = if self.split_mv { "_split_mv" } else { "" };
+        format!("collab_fused_1080p_{}{field}", self.ch_name)
     }
 
     fn sync(&self) {

av-denoise-core/benches/kernels/mod.rs

@@ -25,6 +25,7 @@ pub mod mc_chain_compose;
 pub mod mc_confidence;
 pub mod mc_downscale;
 pub mod mc_warp;
+pub mod mv_regularise;
 pub mod nl4d_geometry;
 pub mod noise_partial;
 pub mod pack_wire;

av-denoise-core/benches/kernels/mv_regularise.rs

@@ -0,0 +1,92 @@
+use av_denoise_core::nl4d::kernels::nl4d_mv_regularise;
+use cubecl::benchmark::Benchmark;
+use cubecl::prelude::*;
+use cubecl::server::Handle;
+
+use super::{H, W, block_sync, make_synthetic_frame, shapes_with_ch};
+
+const BLKSIZE: u32 = 16;
+const STEP: u32 = 8;
+const THSAD_PIXEL: f32 = 0.02;
+const FIELD_LAMBDA: f32 = 1.0;
+
+/// The nl4d field regularisation pass over one neighbour at 1080p.
+pub struct MvRegulariseBench<R: Runtime> {
+    pub client: ComputeClient<R>,
+}
+
+#[derive(Clone)]
+pub struct RegulariseInput {
+    pub centre: Handle,
+    pub neighbour: Handle,
+    pub mv_in: Handle,
+    pub mv_out: Handle,
+    pub confidence: Handle,
+}
+
+impl<R: Runtime> Benchmark for MvRegulariseBench<R> {
+    type Input = RegulariseInput;
+    type Output = ();
+
+    fn prepare(&self) -> Self::Input {
+        let blocks = (W.div_ceil(STEP) * H.div_ceil(STEP)) as usize;
+        let centre = self
+            .client
+            .create_from_slice(f32::as_bytes(&make_synthetic_frame(W, H, 1)));
+        let neighbour = self
+            .client
+            .create_from_slice(f32::as_bytes(&make_synthetic_frame(W, H, 1)));
+        let field: Vec<i32> = (0..2 * blocks).map(|i| (i % 7) as i32 - 3).collect();
+        let mv_in = self.client.create_from_slice(i32::as_bytes(&field));
+        let mv_out = self.client.empty(2 * blocks * size_of::<i32>());
+        let confidence = self.client.empty(blocks * size_of::<f32>());
+        RegulariseInput {
+            centre,
+            neighbour,
+            mv_in,
+            mv_out,
+            confidence,
+        }
+    }
+
+    fn execute(&self, args: Self::Input) -> Result<(), String> {
+        let blocks_x = W.div_ceil(STEP);
+        let blocks_y = H.div_ceil(STEP);
+        let blocks = (blocks_x * blocks_y) as usize;
+        let block_area = (BLKSIZE * BLKSIZE) as f32;
+        unsafe {
+            nl4d_mv_regularise::launch_unchecked::<R>(
+                &self.client,
+                CubeCount::new_2d(blocks_x, blocks_y),
+                CubeDim::new_2d(8, 8),
+                ArrayArg::from_raw_parts(args.centre.clone(), (W * H) as usize),
+                ArrayArg::from_raw_parts(args.neighbour.clone(), (W * H) as usize),
+                ArrayArg::from_raw_parts(args.mv_in.clone(), 2 * blocks),
+                ArrayArg::from_raw_parts(args.mv_out.clone(), 2 * blocks),
+                ArrayArg::from_raw_parts(args.confidence.clone(), blocks),
+                FIELD_LAMBDA * block_area * THSAD_PIXEL,
+                0.0,
+                block_area * THSAD_PIXEL,
+                W,
+                H,
+                BLKSIZE,
+                STEP,
+                blocks_x,
+                blocks_y,
+            );
+        }
+        Ok(())
+    }
+
+    fn name(&self) -> String {
+        "nl4d_mv_regularise_1080p_luma".to_string()
+    }
+
+    fn sync(&self) {
+        block_sync(&self.client);
+    }
+
+    fn shapes(&self) -> Vec<Vec<usize>> {
+        shapes_with_ch(1)
+    }
+}

av-denoise-core/benches/kernels/nl4d_geometry.rs

@@ -14,7 +14,7 @@ pub const SPATIAL_RADIUS: u32 = 9;
 /// `collab::MAX_K`, the group size the filter runs at.
 pub const K_MAX: u32 = 8;
 /// `Nl4dParams::default().lambda_ht`.
-pub const LAMBDA_HT: f32 = 5.3;
+pub const LAMBDA_HT: f32 = 5.2;
 /// `Nl4dParams::default().confidence_variance`, the `use_member_sigma`
 /// flag `collab_fused` compiles against.
 pub const CONFIDENCE_VARIANCE: bool = true;
@@ -26,10 +26,9 @@ pub const BLK_STEP: u32 = 8;
 /// (`MotionCompensationMode::Mvtools`'s `blksize`), distinct from
 /// [`BLK_STEP`] above.
 pub const BLKSIZE: u32 = 16;
-/// `thsad(BLKSIZE, 1.0)` in normalised SAD units (block_area *
-/// THSAD_PIXEL, see `crate::nlmeans::motion::thsad`), hand-computed
-/// here since that function is crate-private.
-pub const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02;
+/// `Nl4dParams::default().mismatch_scale` squared, the kernel's
+/// `mismatch_scale2` argument.
+pub const MISMATCH_SCALE2: f32 = 1.0;
 
 /// Frames in the ring a pass reads.
 pub const N_FRAMES: u32 = 2 * RADIUS + 1;