diff --git a/crates/hevc/src/intra.rs b/crates/hevc/src/intra.rs --- a/crates/hevc/src/intra.rs +++ b/crates/hevc/src/intra.rs @@ -102,91 +102,121 @@ pub fn filter(r: &mut Refs, mode: u32, n: usize, strong_enabled: bool, is_luma: /// Predict an n x n block with `mode` into `out` (stride `os`). `edge_filters`: DC / horizontal / /// vertical boundary smoothing (luma blocks smaller than 32x32). pub fn predict(r: &Refs, mode: u32, n: usize, edge_filters: bool, bit_depth: u32, out: &mut [S], os: usize) { - let max = (1i32 << bit_depth) - 1; + match n { + 4 => predict_n::(r, mode, edge_filters, bit_depth, out, os), + 8 => predict_n::(r, mode, edge_filters, bit_depth, out, os), + 16 => predict_n::(r, mode, edge_filters, bit_depth, out, os), + _ => predict_n::(r, mode, edge_filters, bit_depth, out, os), + } +} + +#[inline(always)] +fn predict_n(r: &Refs, mode: u32, edge_filters: bool, bit_depth: u32, out: &mut [S], os: usize) { match mode { - 0 => { - // planar - let log2 = n.trailing_zeros(); - let tr = r.top[1 + n] as i32; - let bl = r.left[1 + n] as i32; - for y in 0..n { - let l = r.left[1 + y] as i32; - for x in 0..n { - let t = r.top[1 + x] as i32; - let v = ((n - 1 - x) as i32 * l + (x as i32 + 1) * tr + (n - 1 - y) as i32 * t + (y as i32 + 1) * bl + n as i32) >> (log2 + 1); - out[y * os + x] = S::from_i32(v); - } - } - } + 0 => planar::(r, out, os), 1 => { - let log2 = n.trailing_zeros(); - let sum: u32 = r.top[1..=n].iter().map(|&v| v as u32).sum::() + r.left[1..=n].iter().map(|&v| v as u32).sum::(); - let dc = ((sum + n as u32) >> (log2 + 1)) as u16; + let log2 = N.trailing_zeros(); + let sum: u32 = r.top[1..=N].iter().map(|&v| v as u32).sum::() + r.left[1..=N].iter().map(|&v| v as u32).sum::(); + let dc = ((sum + N as u32) >> (log2 + 1)) as u16; let dcs = S::from_u16(dc); - for y in 0..n { - out[y * os..y * os + n].fill(dcs); + for y in 0..N { + out[y * os..y * os + N].fill(dcs); } - if edge_filters && n < 32 { + if edge_filters && N < 32 { let d = dc as u32; out[0] = S::from_i32(((r.left[1] as u32 + 2 * d + r.top[1] as u32 + 2) >> 2) as i32); - for x in 1..n { + for x in 1..N { out[x] = S::from_i32(((r.top[1 + x] as u32 + 3 * d + 2) >> 2) as i32); } - for y in 1..n { + for y in 1..N { out[y * os] = S::from_i32(((r.left[1 + y] as u32 + 3 * d + 2) >> 2) as i32); } } } - _ => { - let angle = INTRA_PRED_ANGLE[mode as usize] as i32; - // reference array ref[-n..=2n] stored at offset n - let mut refa = [0i32; 3 * 64 + 1]; - let base = n as i32; - let (main, side) = if mode >= 18 { (&r.top, &r.left) } else { (&r.left, &r.top) }; - for x in 0..=n { - refa[(base + x as i32) as usize] = main[x] as i32; // ref[x] = p[-1+x][-1] - } - if angle < 0 { - let lim = (n as i32 * angle) >> 5; - if lim < -1 { - let inv = INV_ANGLE[(mode - 11) as usize] as i32; - for x in lim..=-1 { - let idx = -1 + ((x * inv + 128) >> 8); - refa[(base + x) as usize] = side[(idx + 1) as usize] as i32; - } - } - } else { - for x in n + 1..=2 * n { - refa[(base + x as i32) as usize] = main[x] as i32; - } + _ => angular::(r, mode, edge_filters, bit_depth, out, os), + } +} + +fn planar(r: &Refs, out: &mut [S], os: usize) { + let log2 = N.trailing_zeros(); + let tr = r.top[1 + N] as i32; + let bl = r.left[1 + N] as i32; + let top = &r.top[1..1 + N]; + for y in 0..N { + let l = r.left[1 + y] as i32; + let (wt, wb) = ((N - 1 - y) as i32, (y + 1) as i32); + for (x, (o, &t)) in out[y * os..y * os + N].iter_mut().zip(top).enumerate() { + let v = (N - 1 - x) as i32 * l + (x as i32 + 1) * tr + wt * t as i32 + wb * bl + N as i32; + *o = S::from_i32(v >> (log2 + 1)); + } + } +} + +/// Angular modes 2..34 (8.4.4.2.6). Horizontal modes (< 18) are predicted as the vertical mode +/// on the transposed references into a tile, then transposed into `out`. +/// Row `j` of a vertical angular prediction from `ref[x] = refs[base + x]`. +#[inline(always)] +fn angular_row(refs: &[u16], base: i32, angle: i32, j: usize) -> [u16; N] { + let pos = (j as i32 + 1) * angle; + let fact = (pos & 31) as u32; + let p = &refs[(base + (pos >> 5) + 1) as usize..][..N + 1]; + let mut v = [0u16; N]; + if fact == 0 { + v.copy_from_slice(&p[..N]); + } else { + for (v, (&a, &b)) in v.iter_mut().zip(p.iter().zip(&p[1..])) { + *v = (((32 - fact) * a as u32 + fact * b as u32 + 16) >> 5) as u16; + } + } + v +} + +fn angular(r: &Refs, mode: u32, edge_filters: bool, bit_depth: u32, out: &mut [S], os: usize) { + let angle = INTRA_PRED_ANGLE[mode as usize] as i32; + let (main, side) = if mode >= 18 { (&r.top, &r.left) } else { (&r.left, &r.top) }; + // ref[x] = refs[base + x] for x = -N..=2N; ref[x] = main[x] for x >= 0 + let mut ext = [0u16; 3 * 32 + 1]; + let (refs, base): (&[u16], i32) = if angle < 0 { + ext[N..=2 * N].copy_from_slice(&main[..=N]); + let lim = (N as i32 * angle) >> 5; + if lim < -1 { + let inv = INV_ANGLE[(mode - 11) as usize] as i32; + for x in lim..=-1 { + ext[(N as i32 + x) as usize] = side[((x * inv + 128) >> 8) as usize]; } - let horizontal = mode < 18; - for j in 0..n { - let pos = (j as i32 + 1) * angle; - let idx = pos >> 5; - let fact = pos & 31; - for i in 0..n { - let k = (base + i as i32 + idx + 1) as usize; - let v = if fact != 0 { ((32 - fact) * refa[k] + fact * refa[k + 1] + 16) >> 5 } else { refa[k] }; - if horizontal { - out[i * os + j] = S::from_i32(v); - } else { - out[j * os + i] = S::from_i32(v); - } - } + } + (&ext, N as i32) + } else { + (main, 0) + }; + if mode >= 18 { + for j in 0..N { + for (o, v) in out[j * os..j * os + N].iter_mut().zip(angular_row::(refs, base, angle, j)) { + *o = S::from_u16(v); } - if edge_filters && n < 32 { - if mode == 26 { - for y in 0..n { - let v = r.top[1] as i32 + ((r.left[1 + y] as i32 - r.left[0] as i32) >> 1); - out[y * os] = S::from_i32(v.clamp(0, max)); - } - } else if mode == 10 { - for x in 0..n { - let v = r.left[1] as i32 + ((r.top[1 + x] as i32 - r.top[0] as i32) >> 1); - out[x] = S::from_i32(v.clamp(0, max)); - } - } + } + } else { + let mut tile = [[0u16; N]; N]; + for (j, t) in tile.iter_mut().enumerate() { + *t = angular_row::(refs, base, angle, j); + } + for y in 0..N { + for (o, t) in out[y * os..y * os + N].iter_mut().zip(&tile) { + *o = S::from_u16(t[y]); + } + } + } + if edge_filters && N < 32 { + let max = (1i32 << bit_depth) - 1; + if mode == 26 { + for y in 0..N { + let v = r.top[1] as i32 + ((r.left[1 + y] as i32 - r.left[0] as i32) >> 1); + out[y * os] = S::from_i32(v.clamp(0, max)); + } + } else if mode == 10 { + for x in 0..N { + let v = r.left[1] as i32 + ((r.top[1 + x] as i32 - r.top[0] as i32) >> 1); + out[x] = S::from_i32(v.clamp(0, max)); } } } @@ -238,6 +268,80 @@ mod tests { assert_eq!(out[4 + 1], r.top[1 + 1 + 1 + 1]); } + /// The previous per-sample implementation of planar and angular prediction (8.4.4.2.4, 8.4.4.2.6). + fn reference(r: &Refs, mode: u32, n: usize, out: &mut [u16]) { + if mode == 0 { + let log2 = n.trailing_zeros(); + let (tr, bl) = (r.top[1 + n] as i32, r.left[1 + n] as i32); + for y in 0..n { + for x in 0..n { + let (l, t) = (r.left[1 + y] as i32, r.top[1 + x] as i32); + out[y * n + x] = + (((n - 1 - x) as i32 * l + (x as i32 + 1) * tr + (n - 1 - y) as i32 * t + (y as i32 + 1) * bl + n as i32) >> (log2 + 1)) as u16; + } + } + return; + } + let angle = INTRA_PRED_ANGLE[mode as usize] as i32; + let mut refa = [0i32; 3 * 64 + 1]; + let base = n as i32; + let (main, side) = if mode >= 18 { (&r.top, &r.left) } else { (&r.left, &r.top) }; + for x in 0..=n { + refa[(base + x as i32) as usize] = main[x] as i32; + } + if angle < 0 { + let lim = (n as i32 * angle) >> 5; + if lim < -1 { + let inv = INV_ANGLE[(mode - 11) as usize] as i32; + for x in lim..=-1 { + refa[(base + x) as usize] = side[((x * inv + 128) >> 8) as usize] as i32; + } + } + } else { + for x in n + 1..=2 * n { + refa[(base + x as i32) as usize] = main[x] as i32; + } + } + for j in 0..n { + let pos = (j as i32 + 1) * angle; + let (idx, fact) = (pos >> 5, pos & 31); + for i in 0..n { + let k = (base + i as i32 + idx + 1) as usize; + let v = if fact != 0 { ((32 - fact) * refa[k] + fact * refa[k + 1] + 16) >> 5 } else { refa[k] }; + if mode < 18 { + out[i * n + j] = v as u16; + } else { + out[j * n + i] = v as u16; + } + } + } + } + + #[test] + fn planar_and_angular_match_reference() { + let mut rng = 0x9e37_79b9_7f4a_7c15u64; + let mut next = move || { + rng ^= rng << 13; + rng ^= rng >> 7; + rng ^= rng << 17; + rng + }; + for bit_depth in [8u32, 10, 12] { + for n in [4usize, 8, 16, 32] { + for _ in 0..4 { + let v: Vec = (0..257).map(|_| (next() % (1 << bit_depth)) as u16).collect(); + let r = refs(|i| v[(i + 128) as usize]); + for mode in (0..35).filter(|&m| m != 1) { + let (mut a, mut b) = (vec![0u16; n * n], vec![0u16; n * n]); + predict(&r, mode, n, false, bit_depth, &mut a, n); + reference(&r, mode, n, &mut b); + assert_eq!(a, b, "n {n} mode {mode} bit depth {bit_depth}"); + } + } + } + } + } + #[test] fn substitution() { let mut r = Refs::default();