collective
Browse / github.com/storytold/effectcraft / patch
av1enc: motion search with visited set and early-exit SAD, fixed-size SATD, region-only snapshots
av1enc: half-pel planes, motion search, SATD and decoded-flag snapshot in the decision code (-21 % together with #2373/#2374) ## Where the time goes At the export default, mode decision is 90-95 % of the encode: `try_tx` 50 % (inverse transform 21 %, forward 11 %), intra checks 45 % (including intra tried inside inter frames), motion search 15-20 %, SATD 9 % (24 % on key frames), half-pel planes 4-6 %, loop filter 2 %, entropy coding 1 %. Baseline: about 1 frame per second at 1080p. ## Proposal (all output-identical) - A-I4 (`pred.rs`, `Decider::new`): build the half-pel planes in one pass. - A-I5 (`decide.rs`): motion search with a visited set and a partial-SAD early exit. - A-I6: SATD with a fixed-size Hadamard (`
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- vangel76 (imported from pull request #2406, open)
- Base
991904888a20b7da72b2a4880b21c7d912154736- Size
- +166 −48 in 1 file
- Review
- pending
- Shared
- 2026-10-28
- Id
256d2acd22aa29cfce094fd3
How it went
No other operator has reported using it yet.
Diff
raw · crates/av1enc/src/decide.rs
@@ -110,7 +110,8 @@ struct TxResult { /// Saved encoder state of a block area (reconstruction, BlockDecoded flags). struct Snapshot { planes: [Vec<u16>; 3],- decoded: Box<[[[bool; 35]; 35]; 3]>,+ /// BlockDecoded flags of the block area only (the only flags coding inside it can change).+ decoded: [[[bool; 16]; 16]; 3], } pub(crate) struct Decider<'a> {@@ -211,6 +212,13 @@ impl<'a> Decider<'a> { } } + /// The BlockDecoded cells `mark_decoded` can set for blocks inside (r, c, n4): rows and+ /// columns `by + 1 .. by + 1 + step` / `bx + 1 ..` of `decoded[plane]`.+ fn decoded_region(plane: usize, r: usize, c: usize, n4: usize) -> (usize, usize, usize) {+ let s = (plane > 0) as usize;+ ((r & 15) >> s, (c & 15) >> s, (n4 >> s).max(1))+ }+ fn snapshot(&self, r: usize, c: usize, n4: usize) -> Snapshot { let planes = std::array::from_fn(|p| { let s = (p > 0) as usize;@@ -222,7 +230,14 @@ impl<'a> Decider<'a> { } v });- Snapshot { planes, decoded: self.decoded.clone() }+ let mut decoded = [[[false; 16]; 16]; 3];+ for (p, d) in decoded.iter_mut().enumerate() {+ let (by, bx, step) = Self::decoded_region(p, r, c, n4);+ for i in 0..step {+ d[i][..step].copy_from_slice(&self.decoded[p][by + i + 1][bx + 1..bx + 1 + step]);+ }+ }+ Snapshot { planes, decoded } } fn restore(&mut self, snap: &Snapshot, r: usize, c: usize, n4: usize) {@@ -234,7 +249,12 @@ impl<'a> Decider<'a> { pl.data[y * pl.stride + x0..y * pl.stride + x0 + n].copy_from_slice(&snap.planes[p][i * n..i * n + n]); } }- self.decoded.clone_from(&snap.decoded);+ for p in 0..3 {+ let (by, bx, step) = Self::decoded_region(p, r, c, n4);+ for i in 0..step {+ self.decoded[p][by + i + 1][bx + 1..bx + 1 + step].copy_from_slice(&snap.decoded[p][i][..step]);+ }+ } } fn set_mvs(&mut self, r: usize, c: usize, n4: usize, mv: [i32; 2]) {@@ -519,15 +539,22 @@ impl<'a> Decider<'a> { let comp = |d: i32| if d == 0 { 0.5 } else { 2.0 + 2.0 * ((d.abs() / 8) as f64 + 1.0).log2() }; comp(m[0] - pred_mv[0]) + comp(m[1] - pred_mv[1]) };- let cost =- |m: [i32; 2]| -> f64 { self.sad_ref(&refr.planes[0], x0 + (m[1] >> 3), y0 + (m[0] >> 3), c * 4, r * 4, n) as f64 + self.lambda_sad * mv_bits(m) };+ // `bound` is the cost to beat: the SAD stops early once the cost cannot be below it (the+ // callers only accept a strictly lower cost, so the result is unchanged)+ let cost = |m: [i32; 2], bound: f64| -> f64 {+ let lb = self.lambda_sad * mv_bits(m);+ self.sad_ref(&refr.planes[0], x0 + (m[1] >> 3), y0 + (m[0] >> 3), c * 4, r * 4, n, lb, bound) as f64 + lb+ };+ // positions already evaluated cannot beat `best` again (costs are deterministic)+ let mut seen: Vec<[i32; 2]> = Vec::with_capacity(64); let mut best = ([0, 0], f64::MAX); for &m in &cands { let m = [m[0] & !7, m[1] & !7];- if !valid(m) {+ if !valid(m) || seen.contains(&m) { continue; }- let cst = cost(m);+ seen.push(m);+ let cst = cost(m, best.1); if cst < best.1 { best = (m, cst); }@@ -548,10 +575,11 @@ impl<'a> Decider<'a> { let center = best.0; for d in [[-step, 0], [step, 0], [0, -step], [0, step]] { let m = [center[0] + d[0], center[1] + d[1]];- if !valid(m) {+ if !valid(m) || seen.contains(&m) { continue; }- let cst = cost(m);+ seen.push(m);+ let cst = cost(m, best.1); if cst < best.1 { best = (m, cst); moved = true;@@ -562,15 +590,18 @@ impl<'a> Decider<'a> { } // half- then quarter-sample refinement (plus the exact neighbour candidates) let mut buf = vec![0u16; n * n];- let sub_cost = |m: [i32; 2], buf: &mut Vec<u16>| -> f64 {+ let sub_cost = |m: [i32; 2], buf: &mut Vec<u16>, bound: f64| -> f64 { self.search_block(m, c * 4, r * 4, n, buf);- self.sad(0, c * 4, r * 4, n, buf) as f64 + self.lambda_sad * mv_bits(m)+ let lb = self.lambda_sad * mv_bits(m);+ self.sad_bounded(0, c * 4, r * 4, n, buf, lb, bound) as f64 + lb };+ let mut seen_sub: Vec<[i32; 2]> = Vec::with_capacity(24); for m in cands { if m[0] & 7 != 0 || m[1] & 7 != 0 { let m = [m[0] & !1, m[1] & !1];- if valid(m) {- let cst = sub_cost(m, &mut buf);+ if valid(m) && !seen_sub.contains(&m) {+ seen_sub.push(m);+ let cst = sub_cost(m, &mut buf, best.1); if cst < best.1 { best = (m, cst); }@@ -581,10 +612,11 @@ impl<'a> Decider<'a> { let center = best.0; for d in [[-1, -1], [-1, 0], [-1, 1], [0, -1], [0, 1], [1, -1], [1, 0], [1, 1]] { let m = [center[0] + d[0] * step, center[1] + d[1] * step];- if !valid(m) {+ if !valid(m) || seen_sub.contains(&m) { continue; }- let cst = sub_cost(m, &mut buf);+ seen_sub.push(m);+ let cst = sub_cost(m, &mut buf, best.1); if cst < best.1 { best = (m, cst); }@@ -670,7 +702,8 @@ impl<'a> Decider<'a> { } /// SAD of the source block at (x, y) (luma) against the reference block at (rx, ry).- fn sad_ref(&self, refp: &Plane, rx: i32, ry: i32, x: usize, y: usize, n: usize) -> u64 {+ /// Stops after a row once `sad + lb >= bound` (the cost can no longer win).+ fn sad_ref(&self, refp: &Plane, rx: i32, ry: i32, x: usize, y: usize, n: usize, lb: f64, bound: f64) -> u64 { let (vw, vh) = (self.g.width, self.g.height); let w = n.min(vw.saturating_sub(x)); let h = n.min(vh.saturating_sub(y));@@ -690,57 +723,78 @@ impl<'a> Decider<'a> { sad += (s[j] as i32 - b as i32).unsigned_abs() as u64; } }+ if sad as f64 + lb >= bound {+ return sad;+ } } sad } + /// SAD of a prediction against the source over the visible part of the block. Stops after a row+ /// once `sad + lb >= bound` (the cost can no longer win).+ fn sad_bounded(&self, p: usize, x: usize, y: usize, n: usize, pred: &[u16], lb: f64, bound: f64) -> u64 {+ let (vw, vh) = self.g.visible(p);+ let w = n.min(vw.saturating_sub(x));+ let h = n.min(vh.saturating_sub(y));+ let src = &self.src[p];+ let mut s = 0u64;+ for i in 0..h {+ let row = &src.row(y + i)[x..x + w];+ s += row.iter().zip(&pred[i * n..i * n + w]).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs()).sum::<u32>() as u64;+ if s as f64 + lb >= bound {+ return s;+ }+ }+ s+ }+ /// Sum of absolute Hadamard-transformed differences (4x4 or 8x8 sub-blocks) of a prediction /// against the source, scaled to about the size of the SAD. fn satd(&self, p: usize, x: usize, y: usize, n: usize, pred: &[u16]) -> u64 {+ if n >= 8 { self.satd_k::<8>(p, x, y, n, pred) } else { self.satd_k::<4>(p, x, y, n, pred) }+ }++ /// `satd` with K x K Hadamard sub-blocks (fixed sizes so the butterflies vectorise).+ fn satd_k<const K: usize>(&self, p: usize, x: usize, y: usize, n: usize, pred: &[u16]) -> u64 { let src = &self.src[p];- let k = n.min(8); let mut total = 0u64;- let mut d = [0i32; 64];- for by in (0..n).step_by(k) {- for bx in (0..n).step_by(k) {- for i in 0..k {- let row = &src.row(y + by + i)[x + bx..x + bx + k];- for j in 0..k {- d[i * k + j] = row[j] as i32 - pred[(by + i) * n + bx + j] as i32;+ for by in (0..n).step_by(K) {+ for bx in (0..n).step_by(K) {+ // d[i] holds row i of the differences+ let mut d = [[0i32; K]; K];+ for (i, di) in d.iter_mut().enumerate() {+ let row = &src.row(y + by + i)[x + bx..x + bx + K];+ let pr = &pred[(by + i) * n + bx..(by + i) * n + bx + K];+ for j in 0..K {+ di[j] = row[j] as i32 - pr[j] as i32; } }- for i in 0..k {- hadamard(&mut d[i * k..i * k + k]);+ // vertical butterflies across rows (all columns at once)+ let mut h = 1;+ while h < K {+ for i0 in (0..K).step_by(2 * h) {+ for i in i0..i0 + h {+ let (a, b) = (d[i], d[i + h]);+ for j in 0..K {+ d[i][j] = a[j] + b[j];+ d[i + h][j] = a[j] - b[j];+ }+ }+ }+ h *= 2; }- let mut col = [0i32; 8];+ // horizontal butterflies within each row let mut s = 0u64;- for j in 0..k {- for i in 0..k {- col[i] = d[i * k + j];- }- hadamard(&mut col[..k]);- s += col[..k].iter().map(|v| v.unsigned_abs() as u64).sum::<u64>();+ for di in d.iter_mut() {+ hadamard(di);+ s += di.iter().map(|v| v.unsigned_abs() as u64).sum::<u64>(); }- total += s / (k as u64 / 2);+ total += s / (K as u64 / 2); } } total } - /// SAD of a prediction against the source over the visible part of the block.- fn sad(&self, p: usize, x: usize, y: usize, n: usize, pred: &[u16]) -> u64 {- let (vw, vh) = self.g.visible(p);- let w = n.min(vw.saturating_sub(x));- let h = n.min(vh.saturating_sub(y));- let src = &self.src[p];- let mut s = 0u64;- for i in 0..h {- let row = &src.row(y + i)[x..x + w];- s += row.iter().zip(&pred[i * n..i * n + w]).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs()).sum::<u32>() as u64;- }- s- }- fn sse(&self, p: usize, x: usize, y: usize, n: usize, blk: &[u16]) -> f64 { let (vw, vh) = self.g.visible(p); let w = n.min(vw.saturating_sub(x));@@ -889,3 +943,67 @@ pub(crate) fn scan_for(tx_sz: usize) -> &'static [u16] { _ => &DEFAULT_SCAN_32X32, } }++#[cfg(test)]+mod tests {+ use super::*;++ /// The SATD as it was before the fixed-size version: runtime block size, column gather.+ fn satd_ref(src: &Plane, x: usize, y: usize, n: usize, pred: &[u16]) -> u64 {+ let k = n.min(8);+ let mut total = 0u64;+ let mut d = [0i32; 64];+ for by in (0..n).step_by(k) {+ for bx in (0..n).step_by(k) {+ for i in 0..k {+ let row = &src.row(y + by + i)[x + bx..x + bx + k];+ for j in 0..k {+ d[i * k + j] = row[j] as i32 - pred[(by + i) * n + bx + j] as i32;+ }+ }+ for i in 0..k {+ hadamard(&mut d[i * k..i * k + k]);+ }+ let mut col = [0i32; 8];+ let mut s = 0u64;+ for j in 0..k {+ for i in 0..k {+ col[i] = d[i * k + j];+ }+ hadamard(&mut col[..k]);+ s += col[..k].iter().map(|v| v.unsigned_abs() as u64).sum::<u64>();+ }+ total += s / (k as u64 / 2);+ }+ }+ total+ }++ #[test]+ fn satd_fixed_size_matches_runtime_size() {+ let g = FrameGeom::new(100, 80);+ let mut state = 0x9e37_79b9_7f4a_7c15u64;+ let mut next = move || {+ state ^= state << 13;+ state ^= state >> 7;+ state ^= state << 17;+ state+ };+ for bd in [8u32, 10] {+ let max = (1u64 << bd) - 1;+ let mut mk = |w: usize, h: usize| {+ let mut p = Plane::new(w, h);+ p.data.iter_mut().for_each(|v| *v = (next() % (max + 1)) as u16);+ p+ };+ let src = [mk(128, 128), mk(64, 64), mk(64, 64)];+ let dec = Decider::new(&g, bd, 100, true, &src, None);+ for n in [4usize, 8, 16, 32] {+ let pred: Vec<u16> = (0..n * n).map(|_| (next() % (max + 1)) as u16).collect();+ for (p, x, y) in [(0, 0, 0), (0, 32, 40), (1, 8, 4)] {+ assert_eq!(dec.satd(p, x, y, n, &pred), satd_ref(&src[p], x, y, n, &pred), "n {n} plane {p}");+ }+ }+ }+ }+}