Skip to main content

pito_codecs/vp9/enc/rc/
analysis.rs

1use super::FrameKind;
2use rayon::prelude::*;
3
4#[derive(Clone, Copy, Debug, Default, PartialEq)]
5pub struct FrameAnalysis {
6    pub intra_cost: u64,
7    pub inter_cost: u64,
8    pub has_prev: bool,
9    pub blocks: u32,
10}
11
12impl FrameAnalysis {
13    pub fn cost(&self, kind: FrameKind) -> u64 {
14        let c = if kind == FrameKind::Key || !self.has_prev {
15            self.intra_cost
16        } else {
17            self.inter_cost
18        };
19        c.max(1)
20    }
21
22    pub fn inter_ratio(&self) -> f64 {
23        if !self.has_prev {
24            return 1.0;
25        }
26        self.inter_cost as f64 / self.intra_cost.max(1) as f64
27    }
28
29    pub fn scene_cut(&self) -> bool {
30        self.has_prev && self.inter_ratio() > 0.72
31    }
32}
33
34pub struct Analyser {
35    width: usize,
36    height: usize,
37    pw: usize,
38    ph: usize,
39    bw: usize,
40    bh: usize,
41    cur: Vec<u8>,
42    prev: Vec<u8>,
43    has_prev: bool,
44    mvs: Vec<(i8, i8)>,
45    prev_mvs: Vec<(i8, i8)>,
46    activity: Vec<u32>,
47    rows: Vec<(u64, u64)>,
48    global: (i8, i8),
49}
50
51const RANGE: i32 = 48;
52
53impl Analyser {
54    pub fn new(width: usize, height: usize) -> Analyser {
55        let pw = width.div_ceil(2).max(1).next_multiple_of(8);
56        let ph = height.div_ceil(2).max(1).next_multiple_of(8);
57        let (bw, bh) = (pw / 8, ph / 8);
58        Analyser {
59            width,
60            height,
61            pw,
62            ph,
63            bw,
64            bh,
65            cur: vec![0; pw * ph],
66            prev: vec![0; pw * ph],
67            has_prev: false,
68            mvs: vec![(0, 0); bw * bh],
69            prev_mvs: vec![(0, 0); bw * bh],
70            activity: vec![0; bw * bh],
71            rows: vec![(0, 0); bh],
72            global: (0, 0),
73        }
74    }
75
76    pub fn blocks(&self) -> (usize, usize) {
77        (self.bw, self.bh)
78    }
79
80    pub fn activity(&self) -> &[u32] {
81        &self.activity
82    }
83
84    pub fn reset(&mut self) {
85        self.has_prev = false;
86        self.prev_mvs.fill((0, 0));
87        self.global = (0, 0);
88    }
89
90    pub fn run(&mut self, y: &[u8], stride: usize) -> FrameAnalysis {
91        assert!(stride >= self.width && y.len() >= stride * (self.height - 1) + self.width);
92        self.downsample(y, stride);
93        let (pw, ph, bw) = (self.pw, self.ph, self.bw);
94        let cur = &self.cur[..];
95        let prev = self.has_prev.then_some(&self.prev[..]);
96        let prev_mvs = &self.prev_mvs[..];
97        let global = self.global;
98        self.activity
99            .par_chunks_mut(bw)
100            .zip(self.mvs.par_chunks_mut(bw))
101            .zip(self.rows.par_iter_mut())
102            .enumerate()
103            .for_each(|(by, ((act, mvs), sums))| {
104                let mut intra_sum = 0u64;
105                let mut best_sum = 0u64;
106                for bx in 0..bw {
107                    let (x0, y0) = (bx * 8, by * 8);
108                    let src = &cur[y0 * pw + x0..];
109                    let intra = satd8(src, pw, None);
110                    act[bx] = intra;
111                    intra_sum += intra as u64;
112                    let Some(prev) = prev else {
113                        best_sum += intra as u64;
114                        continue;
115                    };
116                    let left = if bx > 0 { mvs[bx - 1] } else { (0, 0) };
117                    let mv = search(
118                        src,
119                        prev,
120                        pw,
121                        ph,
122                        x0,
123                        y0,
124                        [prev_mvs[by * bw + bx], left, global],
125                    );
126                    mvs[bx] = mv;
127                    let rx = (x0 as i32 + mv.0 as i32) as usize;
128                    let ry = (y0 as i32 + mv.1 as i32) as usize;
129                    let inter = satd8(src, pw, Some(&prev[ry * pw + rx..]));
130                    best_sum += inter.min(intra) as u64;
131                }
132                *sums = (intra_sum, best_sum);
133            });
134        let (intra, best) = self
135            .rows
136            .iter()
137            .fold((0u64, 0u64), |a, r| (a.0 + r.0, a.1 + r.1));
138        let out = FrameAnalysis {
139            intra_cost: intra,
140            inter_cost: best,
141            has_prev: self.has_prev,
142            blocks: (self.bw * self.bh) as u32,
143        };
144        std::mem::swap(&mut self.cur, &mut self.prev);
145        if self.has_prev {
146            std::mem::swap(&mut self.mvs, &mut self.prev_mvs);
147            self.global = dominant(&self.prev_mvs);
148        }
149        self.has_prev = true;
150        out
151    }
152
153    fn downsample(&mut self, y: &[u8], stride: usize) {
154        let (w, h, pw) = (self.width, self.height, self.pw);
155        let full = (w / 2).min(pw);
156        self.cur
157            .par_chunks_mut(pw)
158            .enumerate()
159            .for_each(|(r, out)| {
160                let r0 = &y[(2 * r).min(h - 1) * stride..][..w];
161                let r1 = &y[(2 * r + 1).min(h - 1) * stride..][..w];
162                for ((o, a), b) in out[..full]
163                    .iter_mut()
164                    .zip(r0.chunks_exact(2))
165                    .zip(r1.chunks_exact(2))
166                {
167                    *o = ((a[0] as u16 + a[1] as u16 + b[0] as u16 + b[1] as u16 + 2) >> 2) as u8;
168                }
169                for (x, o) in out.iter_mut().enumerate().skip(full) {
170                    let x0 = (2 * x).min(w - 1);
171                    let x1 = (2 * x + 1).min(w - 1);
172                    let s = r0[x0] as u16 + r0[x1] as u16 + r1[x0] as u16 + r1[x1] as u16;
173                    *o = ((s + 2) >> 2) as u8;
174                }
175            });
176    }
177}
178
179fn dominant(mvs: &[(i8, i8)]) -> (i8, i8) {
180    let mut xs = [0u32; 256];
181    let mut ys = [0u32; 256];
182    for &(x, y) in mvs {
183        xs[x as u8 as usize] += 1;
184        ys[y as u8 as usize] += 1;
185    }
186    let pick = |h: &[u32; 256]| {
187        let i = (0..256).max_by_key(|&i| h[i]).unwrap_or(0);
188        i as u8 as i8
189    };
190    (pick(&xs), pick(&ys))
191}
192
193fn search(
194    src: &[u8],
195    prev: &[u8],
196    pw: usize,
197    ph: usize,
198    x0: usize,
199    y0: usize,
200    cands: [(i8, i8); 3],
201) -> (i8, i8) {
202    let (xmin, xmax) = (
203        (-(x0 as i32)).max(-RANGE),
204        ((pw - 8 - x0) as i32).min(RANGE),
205    );
206    let (ymin, ymax) = (
207        (-(y0 as i32)).max(-RANGE),
208        ((ph - 8 - y0) as i32).min(RANGE),
209    );
210    let sad_at = |mx: i32, my: i32| {
211        let o = (y0 as i32 + my) as usize * pw + (x0 as i32 + mx) as usize;
212        sad8(src, &prev[o..], pw)
213    };
214    let mut best = (0i32, 0i32);
215    let mut best_sad = sad_at(0, 0);
216    for c in cands {
217        let (mx, my) = (
218            (c.0 as i32).clamp(xmin, xmax),
219            (c.1 as i32).clamp(ymin, ymax),
220        );
221        if (mx, my) != best {
222            let s = sad_at(mx, my);
223            if s < best_sad {
224                best_sad = s;
225                best = (mx, my);
226            }
227        }
228    }
229    let mut step = 4;
230    let mut guard = 0;
231    while step > 0 && guard < 24 {
232        guard += 1;
233        let mut moved = false;
234        for (dx, dy) in [
235            (-step, 0),
236            (step, 0),
237            (0, -step),
238            (0, step),
239            (-step, -step),
240            (step, -step),
241            (-step, step),
242            (step, step),
243        ] {
244            let (mx, my) = (best.0 + dx, best.1 + dy);
245            if mx < xmin || mx > xmax || my < ymin || my > ymax {
246                continue;
247            }
248            let s = sad_at(mx, my);
249            if s < best_sad {
250                best_sad = s;
251                best = (mx, my);
252                moved = true;
253            }
254        }
255        if !moved {
256            step /= 2;
257        }
258    }
259    (best.0 as i8, best.1 as i8)
260}
261
262#[cfg(target_arch = "x86_64")]
263fn sad8(a: &[u8], b: &[u8], stride: usize) -> u32 {
264    use std::arch::x86_64::*;
265    assert!(a.len() >= 7 * stride + 8 && b.len() >= 7 * stride + 8);
266    unsafe {
267        let mut acc = _mm_setzero_si128();
268        for r in (0..8).step_by(2) {
269            let pa = a.as_ptr().add(r * stride);
270            let pb = b.as_ptr().add(r * stride);
271            let va = _mm_unpacklo_epi64(
272                _mm_loadl_epi64(pa as *const __m128i),
273                _mm_loadl_epi64(pa.add(stride) as *const __m128i),
274            );
275            let vb = _mm_unpacklo_epi64(
276                _mm_loadl_epi64(pb as *const __m128i),
277                _mm_loadl_epi64(pb.add(stride) as *const __m128i),
278            );
279            acc = _mm_add_epi64(acc, _mm_sad_epu8(va, vb));
280        }
281        (_mm_cvtsi128_si32(acc) + _mm_extract_epi16(acc, 4)) as u32
282    }
283}
284
285#[cfg(not(target_arch = "x86_64"))]
286fn sad8(a: &[u8], b: &[u8], stride: usize) -> u32 {
287    let mut s = 0u32;
288    for r in 0..8 {
289        for (x, y) in a[r * stride..][..8].iter().zip(&b[r * stride..][..8]) {
290            s += x.abs_diff(*y) as u32;
291        }
292    }
293    s
294}
295
296fn satd8(src: &[u8], stride: usize, pred: Option<&[u8]>) -> u32 {
297    let mut d = [[0i32; 8]; 8];
298    for (r, row) in d.iter_mut().enumerate() {
299        let s = &src[r * stride..][..8];
300        match pred {
301            Some(p) => {
302                let p = &p[r * stride..][..8];
303                for x in 0..8 {
304                    row[x] = s[x] as i32 - p[x] as i32;
305                }
306            }
307            None => {
308                for x in 0..8 {
309                    row[x] = s[x] as i32;
310                }
311            }
312        }
313        hadamard8(row);
314    }
315    let mut sum = 0u32;
316    let mut dc = 0u32;
317    for x in 0..8 {
318        let mut col = [
319            d[0][x], d[1][x], d[2][x], d[3][x], d[4][x], d[5][x], d[6][x], d[7][x],
320        ];
321        hadamard8(&mut col);
322        if x == 0 {
323            dc = col[0].unsigned_abs();
324        }
325        sum += col.iter().map(|v| v.unsigned_abs()).sum::<u32>();
326    }
327    if pred.is_none() {
328        (sum - dc) >> 2
329    } else {
330        sum >> 2
331    }
332}
333
334fn hadamard8(v: &mut [i32; 8]) {
335    let mut h = 4;
336    while h > 0 {
337        for i in 0..8 {
338            if i & h == 0 {
339                let (a, b) = (v[i], v[i + h]);
340                v[i] = a + b;
341                v[i + h] = a - b;
342            }
343        }
344        h >>= 1;
345    }
346}