pito_codecs/vp9/enc/tx/
mod.rs

1#[cfg(target_arch = "x86_64")]
2mod avx2;
3mod lane;
4mod net;
5mod quant;
6mod tables;
7mod xform;
8
9#[cfg(test)]
10mod tests;
11
12use lane::{Isa, Scalar};
13use quant::Scan;
14
15pub use quant::{MAX_LEVEL, bit_cost, token};
16pub use tables::{COEFBAND_4X4, COEFBAND_8X8PLUS, ENERGY_CLASS, bands, neighbors, scan};
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
19pub struct Quant {
20    pub dc: i32,
21    pub ac: i32,
22}
23
24pub struct TokenCosts<'a> {
25    pub more: &'a [[[u32; 2]; 6]; 6],
26    pub token: &'a [[[u32; 11]; 6]; 6],
27}
28
29pub struct Block<'a> {
30    pub src: &'a [u8],
31    pub src_stride: usize,
32    pub dst: &'a mut [u8],
33    pub dst_stride: usize,
34    pub tx_size: u8,
35    pub tx_type: u8,
36    pub lossless: bool,
37    pub ctx: usize,
38    pub is_inter: bool,
39    pub chroma: bool,
40}
41
42#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
43pub struct Out {
44    pub eob: usize,
45    pub dist: u64,
46    pub rate: u32,
47}
48
49#[repr(C, align(32))]
50struct Buf<const N2: usize>([i16; N2]);
51
52fn scan_for(tx_size: u8, tx_type: u8) -> Scan<'static> {
53    let set = tables::scan_set(tx_size, tx_type);
54    Scan {
55        scan: scan(tx_size, tx_type),
56        nb: &set.neighbors,
57        bands: bands(tx_size),
58        iscan_t: &set.iscan_t,
59        log2n: 2 + tx_size as usize,
60    }
61}
62
63fn effective_type(tx_size: u8, tx_type: u8, lossless: bool) -> u8 {
64    if lossless || tx_size == 3 {
65        0
66    } else {
67        tx_type & 3
68    }
69}
70
71pub fn encode(
72    b: Block,
73    q: Quant,
74    costs: &TokenCosts,
75    lambda: u32,
76    levels: &mut [i16; 1024],
77) -> Out {
78    #[cfg(target_arch = "x86_64")]
79    if crate::media::isa::avx2_fma() {
80        return unsafe { encode_avx2(b, q, costs, lambda, levels) };
81    }
82    encode_with::<Scalar>(b, q, costs, lambda, levels)
83}
84
85#[cfg(target_arch = "x86_64")]
86#[target_feature(enable = "avx2,fma")]
87unsafe fn encode_avx2(
88    b: Block,
89    q: Quant,
90    costs: &TokenCosts,
91    lambda: u32,
92    levels: &mut [i16; 1024],
93) -> Out {
94    encode_with::<avx2::Avx2>(b, q, costs, lambda, levels)
95}
96
97#[allow(clippy::too_many_arguments)]
98pub fn reconstruct(
99    levels: &[i16],
100    eob: usize,
101    tx_size: u8,
102    tx_type: u8,
103    lossless: bool,
104    q: Quant,
105    dst: &mut [u8],
106    dst_stride: usize,
107) {
108    #[cfg(target_arch = "x86_64")]
109    if crate::media::isa::avx2_fma() {
110        return unsafe {
111            reconstruct_avx2(levels, eob, tx_size, tx_type, lossless, q, dst, dst_stride)
112        };
113    }
114    reconstruct_with::<Scalar>(levels, eob, tx_size, tx_type, lossless, q, dst, dst_stride)
115}
116
117#[cfg(target_arch = "x86_64")]
118#[target_feature(enable = "avx2,fma")]
119#[allow(clippy::too_many_arguments)]
120unsafe fn reconstruct_avx2(
121    levels: &[i16],
122    eob: usize,
123    tx_size: u8,
124    tx_type: u8,
125    lossless: bool,
126    q: Quant,
127    dst: &mut [u8],
128    dst_stride: usize,
129) {
130    reconstruct_with::<avx2::Avx2>(levels, eob, tx_size, tx_type, lossless, q, dst, dst_stride)
131}
132
133#[inline(always)]
134#[allow(clippy::too_many_arguments)]
135fn reconstruct_with<I: Isa>(
136    levels: &[i16],
137    eob: usize,
138    tx_size: u8,
139    tx_type: u8,
140    lossless: bool,
141    q: Quant,
142    dst: &mut [u8],
143    dst_stride: usize,
144) {
145    if eob == 0 {
146        return;
147    }
148    let tx_size = if lossless { 0 } else { tx_size.min(3) };
149    let tx_type = effective_type(tx_size, tx_type, lossless);
150    let s = scan_for(tx_size, tx_type);
151    let r = Recon {
152        tx_size,
153        tx_type,
154        lossless,
155        dst,
156        stride: dst_stride,
157    };
158    match tx_size {
159        0 => recon_n::<I, 16>(levels, eob, &s, q, r),
160        1 => recon_n::<I, 64>(levels, eob, &s, q, r),
161        2 => recon_n::<I, 256>(levels, eob, &s, q, r),
162        _ => recon_n::<I, 1024>(levels, eob, &s, q, r),
163    }
164}
165
166struct Recon<'a> {
167    tx_size: u8,
168    tx_type: u8,
169    lossless: bool,
170    dst: &'a mut [u8],
171    stride: usize,
172}
173
174#[inline(always)]
175fn recon_n<I: Isa, const N2: usize>(levels: &[i16], eob: usize, s: &Scan, q: Quant, r: Recon) {
176    let mut dq = Buf([0; N2]);
177    let rows = quant::dequantise(levels, eob, s, q, &mut dq.0);
178    apply::<I>(&mut dq.0, eob, rows, r);
179}
180
181#[inline(always)]
182fn apply<I: Isa>(dq: &mut [i16], eob: usize, rows: usize, r: Recon) {
183    if eob == 0 {
184        return;
185    }
186    if r.lossless {
187        xform::inverse_wht_add(dq, r.dst, r.stride);
188    } else if eob == 1 && r.tx_type == 0 {
189        xform::dc_add(dq[0] as i32, r.tx_size, r.dst, r.stride);
190    } else {
191        xform::inverse_add::<I>(dq, r.tx_size, r.tx_type, rows, r.dst, r.stride);
192    }
193}
194
195#[inline(always)]
196fn encode_with<I: Isa>(
197    b: Block,
198    q: Quant,
199    costs: &TokenCosts,
200    lambda: u32,
201    levels: &mut [i16; 1024],
202) -> Out {
203    let tx_size = if b.lossless { 0 } else { b.tx_size.min(3) };
204    match tx_size {
205        0 => encode_n::<I, 4, 16>(b, q, costs, lambda, levels),
206        1 => encode_n::<I, 8, 64>(b, q, costs, lambda, levels),
207        2 => encode_n::<I, 16, 256>(b, q, costs, lambda, levels),
208        _ => encode_n::<I, 32, 1024>(b, q, costs, lambda, levels),
209    }
210}
211
212#[inline(always)]
213fn encode_n<I: Isa, const N: usize, const N2: usize>(
214    b: Block,
215    q: Quant,
216    costs: &TokenCosts,
217    lambda: u32,
218    levels: &mut [i16; 1024],
219) -> Out {
220    let lossless = b.lossless;
221    let tx_size = N.trailing_zeros() as u8 - 2;
222    let tx_type = effective_type(tx_size, b.tx_type, lossless);
223    let s = scan_for(tx_size, tx_type);
224    let mut co = Buf([0; N2]);
225    for i in 0..N {
226        let src = &b.src[i * b.src_stride..][..N];
227        let pred = &b.dst[i * b.dst_stride..][..N];
228        for ((d, &x), &p) in co.0[i * N..][..N].iter_mut().zip(src).zip(pred) {
229            *d = x as i16 - p as i16;
230        }
231    }
232    let mut dq = Buf([0; N2]);
233    let (eob, rows, rate) = if lossless {
234        xform::forward_wht(&mut co.0);
235        let mut eob = 0;
236        for (c, &pos) in s.scan.iter().enumerate() {
237            let v = co.0[s.tpos(pos as usize)];
238            levels[c] = v;
239            if v != 0 {
240                eob = c + 1;
241            }
242        }
243        for (d, &c) in dq.0.iter_mut().zip(&co.0) {
244            *d = c.wrapping_mul(4);
245        }
246        (eob, 4, quant::rate::<N2>(levels, eob, &s, b.ctx, costs))
247    } else {
248        xform::forward::<I>(&mut co.0, tx_size, tx_type);
249        if lambda == 0 {
250            let deadzone = if b.is_inter { 14 } else { 21 };
251            let (eob, rows) = quant::quantise::<N2>(&co.0, &mut dq.0, &s, q, deadzone, levels);
252            (eob, rows, quant::rate::<N2>(levels, eob, &s, b.ctx, costs))
253        } else {
254            let (eob, rate) = quant::trellis::<N2>(&co.0, &s, q, b.ctx, costs, lambda, levels);
255            let rows = quant::dequantise(levels, eob, &s, q, &mut dq.0);
256            (eob, rows, rate)
257        }
258    };
259    let r = Recon {
260        tx_size,
261        tx_type,
262        lossless,
263        dst: &mut *b.dst,
264        stride: b.dst_stride,
265    };
266    apply::<I>(&mut dq.0, eob, rows, r);
267    let mut dist = 0u32;
268    for i in 0..N {
269        let src = &b.src[i * b.src_stride..][..N];
270        let rec = &b.dst[i * b.dst_stride..][..N];
271        for (&x, &r) in src.iter().zip(rec) {
272            let d = x as i32 - r as i32;
273            dist += (d * d) as u32;
274        }
275    }
276    Out {
277        eob,
278        dist: dist as u64,
279        rate,
280    }
281}
pito_codecs/vp9/enc/tx/mod.rs

pito_codecs/vp9/enc/tx/
mod.rs
