pixelflux/encoders/vaapi/
av1.rs

1/*
2 * This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
5 */
6
7//! The AV1 arm: the sequence and picture parameters, the packed sequence header and frame
8//! header OBUs, and the tile group behind them. Main profile, one tile group of as many tiles
9//! as the picture width needs, and the eight buffer slots addressed by timestamp: every frame
10//! refreshes its own slot and names the slot of the newest frame the client still has as all
11//! seven of its references, which is also where its probability contexts load from. The
12//! frame header is written with the bit positions of the quantizer, loop filter, and CDEF
13//! fields, which a constant-rate driver rewrites in place and a constant-quantizer session
14//! fills itself, and with OBU sizes of the fixed width it asks for.
15
16use va_sys::*;
17
18use super::super::bits::BitWriter;
19use super::super::codec::av1_level;
20use super::super::reference::REFERENCE_FRAMES;
21use super::{Buffers, Frame, Negotiated};
22
23const OBU_SEQUENCE_HEADER: u64 = 1;
24const OBU_FRAME_HEADER: u64 = 3;
25const ORDER_HINT_BITS: u32 = 8;
26const MAX_TILE_WIDTH: u32 = 4096;
27const MAX_TILE_AREA: u32 = 4096 * 2304;
28const MAX_TILE_COLS: u32 = 64;
29const MAX_TILE_ROWS: u32 = 64;
30/// The quantizer of a constant-rate session's headers, which the driver rewrites.
31const RATE_CONTROLLED_QINDEX: u32 = 128;
32
33/// The tile layout of a picture: uniform tiles of `cols` x `rows`, in superblocks.
34#[derive(Clone, Copy, Default)]
35struct Tiles {
36    sb_cols: u32,
37    sb_rows: u32,
38    cols: u32,
39    rows: u32,
40    cols_log2: u32,
41    rows_log2: u32,
42    min_log2_cols: u32,
43    max_log2_cols: u32,
44    min_log2_rows: u32,
45    max_log2_rows: u32,
46}
47
48pub(super) struct Arm {
49    ext2: VAConfigAttribValEncAV1Ext2,
50    tx_mode_select: bool,
51    obu_size_bytes: u32,
52    level_idx: u32,
53    tiles: Tiles,
54    sequence: VAEncSequenceParameterBufferAV1,
55    sequence_header: Vec<u8>,
56    /// The bit `base_q_idx` stands at in the last frame header OBU written.
57    qindex_offset: u32,
58}
59
60/// The in-loop filters of a frame: the deblocking level of every plane and edge, and the one
61/// CDEF strength (primary in the high four bits, secondary in the low two) with its damping.
62/// A constant-rate driver chooses its own and rewrites the header; a constant-quantizer
63/// session has to name them, and all zero leaves the picture unfiltered.
64#[derive(Clone, Copy, Default)]
65struct Filters {
66    level: u8,
67    cdef_strength: u8,
68    cdef_damping_minus_3: u8,
69}
70
71impl Filters {
72    /// The filters of a constant-quantizer frame at `qindex`, measured on Intel Arc on a
73    /// text desktop and a gradient wallpaper: deblocking rises with the quantizer, CDEF holds
74    /// primary 8 and secondary 2, and its damping widens as the steps coarsen.
75    fn at(qindex: u32) -> Self {
76        Self {
77            level: (qindex / 8).min(63) as u8,
78            cdef_strength: (8 << 2) | 2,
79            cdef_damping_minus_3: (qindex >> 6).min(3) as u8,
80        }
81    }
82}
83
84/// How many bits `value` needs, at least one.
85fn bit_width(value: u32) -> u32 {
86    (32 - value.leading_zeros()).max(1)
87}
88
89/// The smallest k with `block << k >= target`.
90fn tile_log2(block: u32, target: u32) -> u32 {
91    let mut k = 0;
92    while (block << k) < target {
93        k += 1;
94    }
95    k
96}
97
98/// An OBU: the header byte, the size in `size_bytes` of LEB128, the payload, and its
99/// trailing bits.
100fn obu(kind: u64, payload: &BitWriter, size_bytes: u32) -> Vec<u8> {
101    let mut body = payload.as_bytes().to_vec();
102    if payload.aligned() {
103        body.push(0x80);
104    } else {
105        let last = body.len() - 1;
106        body[last] |= 0x80 >> (payload.len() % 8);
107    }
108    let mut w = BitWriter::new();
109    w.u(1, 0);
110    w.u(4, kind);
111    w.u(1, 0);
112    w.u(1, 1);
113    w.u(1, 0);
114    let mut size = body.len() as u64;
115    for i in 0..size_bytes {
116        let byte = size & 0x7f;
117        size >>= 7;
118        w.u(8, byte | if i + 1 < size_bytes { 0x80 } else { 0 });
119    }
120    w.bytes(&body);
121    w.into_bytes()
122}
123
124impl Arm {
125    pub(super) fn new(_attr: Option<u32>, _ext1: Option<u32>, ext2: u32) -> Self {
126        let ext2 = VAConfigAttribValEncAV1Ext2 { value: ext2 };
127        let tx_mode_support = unsafe { ext2.bits.tx_mode_support() };
128        Self {
129            ext2,
130            tx_mode_select: tx_mode_support & 0x04 != 0,
131            obu_size_bytes: unsafe { ext2.bits.obu_size_bytes_minus1() } + 1,
132            level_idx: 0,
133            tiles: Tiles::default(),
134            sequence: unsafe { std::mem::zeroed() },
135            sequence_header: Vec::new(),
136            qindex_offset: 0,
137        }
138    }
139
140    /// Surfaces align to the 64x64 superblock.
141    pub(super) fn alignment(&self) -> (u32, u32) {
142        (64, 64)
143    }
144
145    /// The `seq_level_idx` the sequence names.
146    pub(super) fn level(&self) -> u32 {
147        self.level_idx
148    }
149
150    /// The uniform tiling of a `width` x `height` picture: as few tile columns as the
151    /// codec's widest tile allows and one tile row, or more rows where the area limit asks.
152    fn tiles(width: u32, height: u32) -> Result<Tiles, String> {
153        let mi_cols = 2 * ((width + 7) >> 3);
154        let mi_rows = 2 * ((height + 7) >> 3);
155        let sb_cols = (mi_cols + 15) >> 4;
156        let sb_rows = (mi_rows + 15) >> 4;
157        let sb_size = 6;
158        let max_tile_width_sb = MAX_TILE_WIDTH >> sb_size;
159        let max_tile_area_sb = MAX_TILE_AREA >> (2 * sb_size);
160        let min_log2_cols = tile_log2(max_tile_width_sb, sb_cols);
161        let max_log2_cols = tile_log2(1, sb_cols.min(MAX_TILE_COLS));
162        let max_log2_rows = tile_log2(1, sb_rows.min(MAX_TILE_ROWS));
163        let min_log2_tiles = min_log2_cols.max(tile_log2(max_tile_area_sb, sb_rows * sb_cols));
164        let cols = sb_cols.div_ceil(max_tile_width_sb).max(1);
165        let cols_log2 = tile_log2(1, cols);
166        let tile_width_sb = (sb_cols + (1 << cols_log2) - 1) >> cols_log2;
167        for rows in 1..=sb_rows.min(MAX_TILE_ROWS) {
168            let rows_log2 = tile_log2(1, rows);
169            let tile_height_sb = (sb_rows + (1 << rows_log2) - 1) >> rows_log2;
170            if sb_cols.div_ceil(tile_width_sb) == cols
171                && sb_rows.div_ceil(tile_height_sb) == rows
172                && tile_height_sb <= max_tile_area_sb / tile_width_sb
173            {
174                return Ok(Tiles {
175                    sb_cols,
176                    sb_rows,
177                    cols,
178                    rows,
179                    cols_log2,
180                    rows_log2,
181                    min_log2_cols,
182                    max_log2_cols,
183                    min_log2_rows: min_log2_tiles.saturating_sub(cols_log2),
184                    max_log2_rows,
185                });
186            }
187        }
188        Err(format!("no uniform AV1 tiling of {width}x{height}"))
189    }
190
191    /// The stream's sequence: its level from the ladder, its tiling, and the sequence header
192    /// OBU the key frames carry.
193    pub(super) fn configure(
194        &mut self,
195        n: &Negotiated,
196        surface_width: u32,
197        surface_height: u32,
198    ) -> Result<(), String> {
199        self.tiles = Self::tiles(n.width, n.height)?;
200        let tiles = self.tiles.cols * self.tiles.rows;
201        let max_tiles = unsafe { self.ext2.bits.max_tile_num_minus1() } + 1;
202        if max_tiles > 1 && tiles > max_tiles {
203            return Err(format!(
204                "this VA-API driver encodes at most {max_tiles} AV1 tiles, and {}x{} needs {tiles}",
205                n.width, n.height
206            ));
207        }
208        let _ = (surface_width, surface_height);
209        self.level_idx = av1_level(
210            n.width,
211            n.height,
212            n.fps.ceil(),
213            u64::from(n.bits_per_second),
214        );
215        let cbr = n.rc_mode != VA_RC_CQP;
216
217        let mut w = BitWriter::new();
218        w.u(3, 0);
219        w.flag(false);
220        w.flag(false);
221        w.flag(false);
222        w.flag(false);
223        w.u(5, 0);
224        w.u(12, 0);
225        w.u(5, u64::from(self.level_idx));
226        if self.level_idx > 7 {
227            w.flag(false);
228        }
229        let (wbits, hbits) = (bit_width(n.width - 1), bit_width(n.height - 1));
230        w.u(4, u64::from(wbits - 1));
231        w.u(4, u64::from(hbits - 1));
232        w.u(wbits, u64::from(n.width - 1));
233        w.u(hbits, u64::from(n.height - 1));
234        w.flag(false);
235        w.flag(false);
236        w.flag(false);
237        w.flag(false);
238        w.flag(false);
239        w.flag(false);
240        w.flag(false);
241        w.flag(false);
242        w.flag(true);
243        w.flag(false);
244        w.flag(false);
245        w.flag(false);
246        w.u(1, 0);
247        w.u(3, u64::from(ORDER_HINT_BITS - 1));
248        w.flag(false);
249        w.flag(true);
250        w.flag(false);
251        w.flag(n.bit_depth == 10);
252        w.flag(false);
253        w.flag(true);
254        w.u(8, 1);
255        w.u(8, 1);
256        w.u(8, 1);
257        w.flag(false);
258        w.u(2, 0);
259        w.flag(false);
260        w.flag(false);
261        self.sequence_header = obu(OBU_SEQUENCE_HEADER, &w, self.obu_size_bytes);
262
263        let s = &mut self.sequence;
264        *s = unsafe { std::mem::zeroed() };
265        s.seq_profile = 0;
266        s.seq_level_idx = self.level_idx as u8;
267        s.seq_tier = 0;
268        s.intra_period = i32::MAX as u32;
269        s.ip_period = 1;
270        s.order_hint_bits_minus_1 = (ORDER_HINT_BITS - 1) as u8;
271        unsafe {
272            let f = &mut s.seq_fields.bits;
273            f.set_enable_order_hint(1);
274            f.set_bit_depth_minus8(n.bit_depth - 8);
275            f.set_enable_cdef(1);
276        }
277        if cbr {
278            s.bits_per_second = n.bits_per_second;
279        }
280        Ok(())
281    }
282
283    pub(super) fn sequence(&mut self, n: &Negotiated, out: &mut Buffers) {
284        out.push(VAEncSequenceParameterBufferType, &self.sequence);
285        if n.packed & VA_ENC_PACKED_HEADER_SEQUENCE != 0 {
286            let bits = 8 * self.sequence_header.len() as u32;
287            out.push_packed(VAEncPackedHeaderSequence, &self.sequence_header, bits);
288        }
289    }
290
291    /// The `base_q_idx` the driver's rate control wrote into the frame header of the coded
292    /// temporal unit `tu`, at the offset the session named for it; None where the unit carries
293    /// no frame header laid out as the session wrote it.
294    pub(super) fn coded_qindex(&self, tu: &[u8]) -> Option<u32> {
295        let mut pos = 0usize;
296        while pos < tu.len() {
297            let header = tu[pos];
298            if header & 0x06 != 0x02 {
299                return None;
300            }
301            let (mut size, mut shift, mut i) = (0usize, 0u32, pos + 1);
302            loop {
303                let byte = *tu.get(i)?;
304                i += 1;
305                size |= ((byte & 0x7f) as usize) << shift;
306                shift += 7;
307                if byte & 0x80 == 0 {
308                    break;
309                }
310                if shift > 28 {
311                    return None;
312                }
313            }
314            if matches!((header >> 3) & 0x0f, 3 | 6) {
315                if self.qindex_offset == 0 || i - pos != 1 + self.obu_size_bytes as usize {
316                    return None;
317                }
318                let bit = pos * 8 + self.qindex_offset as usize;
319                let (byte, lead) = (bit / 8, bit % 8);
320                let pair = (u32::from(*tu.get(byte)?) << 8) | u32::from(*tu.get(byte + 1)?);
321                return Some((pair >> (8 - lead)) & 0xff);
322            }
323            pos = i + size;
324        }
325        None
326    }
327
328    /// The picture: its slot, references, order hint, the frame header OBU with the
329    /// offsets a driver patches, and the one tile group.
330    pub(super) fn picture(
331        &mut self,
332        n: &Negotiated,
333        frame: &Frame,
334        out: &mut Buffers,
335    ) -> Result<(), String> {
336        let slot_of = |pts: u64| (pts % u64::from(REFERENCE_FRAMES)) as u8;
337        let cbr = n.rc_mode != VA_RC_CQP;
338        let qindex = if cbr {
339            RATE_CONTROLLED_QINDEX
340        } else {
341            frame.qp.clamp(1, 255)
342        };
343        let order_hint = ((frame.pts - frame.key_pts) & ((1 << ORDER_HINT_BITS) - 1)) as u32;
344        let slot = slot_of(frame.pts);
345        let reference_slot = match (frame.key, frame.reference) {
346            (true, _) => None,
347            (false, Some((pts, _))) => Some(slot_of(pts)),
348            (false, None) => return Err("a predicted AV1 picture without a reference".into()),
349        };
350        let t = self.tiles;
351
352        let render_size = |w: &mut BitWriter| {
353            w.flag(true);
354            w.u(16, u64::from(n.width - 1));
355            w.u(16, u64::from(n.height - 1));
356        };
357        let mut w = BitWriter::new();
358        w.flag(false);
359        w.u(2, if frame.key { 0 } else { 1 });
360        w.flag(true);
361        if !frame.key {
362            w.flag(false);
363        }
364        w.flag(false);
365        w.flag(false);
366        w.u(ORDER_HINT_BITS, u64::from(order_hint));
367        if !frame.key {
368            w.u(3, 0);
369            w.u(8, 1 << slot);
370            w.flag(false);
371            for _ in 0..7 {
372                w.u(3, u64::from(reference_slot.unwrap()));
373            }
374            render_size(&mut w);
375            w.flag(false);
376            w.flag(false);
377            w.u(2, 0);
378            w.flag(false);
379        } else {
380            render_size(&mut w);
381        }
382        w.flag(false);
383        w.flag(true);
384        for log2 in t.min_log2_cols..t.max_log2_cols {
385            let more = log2 < t.cols_log2;
386            w.flag(more);
387            if !more {
388                break;
389            }
390        }
391        for log2 in t.min_log2_rows..t.max_log2_rows {
392            let more = log2 < t.rows_log2;
393            w.flag(more);
394            if !more {
395                break;
396            }
397        }
398        if t.cols_log2 > 0 || t.rows_log2 > 0 {
399            w.u(t.cols_log2 + t.rows_log2, 0);
400            w.u(
401                2,
402                u64::from(unsafe { self.ext2.bits.tile_size_bytes_minus1() }),
403            );
404        }
405        let header_prefix_bits = 8 * (1 + self.obu_size_bytes);
406        let qindex_offset = header_prefix_bits + w.len() as u32;
407        w.u(8, u64::from(qindex));
408        w.flag(false);
409        w.flag(false);
410        w.flag(false);
411        w.flag(false);
412        w.flag(false);
413        w.flag(false);
414        let filters = if cbr {
415            Filters::default()
416        } else {
417            Filters::at(qindex)
418        };
419        let loopfilter_offset = header_prefix_bits + w.len() as u32;
420        w.u(6, u64::from(filters.level));
421        w.u(6, u64::from(filters.level));
422        if filters.level > 0 {
423            w.u(6, u64::from(filters.level));
424            w.u(6, u64::from(filters.level));
425        }
426        w.u(3, 0);
427        w.flag(false);
428        let cdef_offset = header_prefix_bits + w.len() as u32;
429        w.u(2, u64::from(filters.cdef_damping_minus_3));
430        w.u(2, 0);
431        w.u(6, u64::from(filters.cdef_strength));
432        w.u(6, u64::from(filters.cdef_strength));
433        let cdef_size = header_prefix_bits + w.len() as u32 - cdef_offset;
434        w.flag(self.tx_mode_select);
435        if !frame.key {
436            w.flag(false);
437            for _ in 0..7 {
438                w.flag(false);
439            }
440        }
441        w.flag(false);
442        let frame_header = obu(OBU_FRAME_HEADER, &w, self.obu_size_bytes);
443        let frame_header_bits = 8 * frame_header.len() as u32;
444
445        let mut pic: VAEncPictureParameterBufferAV1 = unsafe { std::mem::zeroed() };
446        pic.frame_width_minus_1 = (n.width - 1) as u16;
447        pic.frame_height_minus_1 = (n.height - 1) as u16;
448        pic.reconstructed_frame = frame.recon;
449        pic.coded_buf = frame.coded;
450        pic.reference_frames = [VA_INVALID_SURFACE; 8];
451        for &(pts, surface, _) in frame.held {
452            if !frame.key && pts >= frame.key_pts {
453                pic.reference_frames[slot_of(pts) as usize] = surface;
454            }
455        }
456        pic.ref_frame_idx = [reference_slot.unwrap_or(0); 7];
457        pic.primary_ref_frame = if frame.key { 7 } else { 0 };
458        pic.order_hint = order_hint as u8;
459        pic.refresh_frame_flags = if frame.key { 0xff } else { 1 << slot };
460        unsafe {
461            pic.ref_frame_ctrl_l0
462                .fields
463                .set_search_idx0(if frame.key { 0 } else { 1 });
464            let f = &mut pic.picture_flags.bits;
465            f.set_frame_type(if frame.key { 0 } else { 1 });
466            f.set_error_resilient_mode(u32::from(frame.key));
467            f.set_enable_frame_obu(0);
468            let m = &mut pic.mode_control_flags.bits;
469            m.set_reference_mode(0);
470            m.set_tx_mode(if self.tx_mode_select { 2 } else { 1 });
471            pic.tile_group_obu_hdr_info.bits.set_obu_has_size_field(1);
472        }
473        pic.base_qindex = qindex as u8;
474        pic.filter_level = [filters.level; 2];
475        pic.filter_level_u = filters.level;
476        pic.filter_level_v = filters.level;
477        pic.cdef_damping_minus_3 = filters.cdef_damping_minus_3;
478        pic.cdef_y_strengths[0] = filters.cdef_strength;
479        pic.cdef_uv_strengths[0] = filters.cdef_strength;
480        pic.tile_cols = t.cols as u8;
481        pic.tile_rows = t.rows as u8;
482        let tile_width_sb = (t.sb_cols + (1 << t.cols_log2) - 1) >> t.cols_log2;
483        let tile_height_sb = (t.sb_rows + (1 << t.rows_log2) - 1) >> t.rows_log2;
484        for i in 0..t.cols as usize {
485            pic.width_in_sbs_minus_1[i] = (if i + 1 == t.cols as usize {
486                t.sb_cols - (t.cols - 1) * tile_width_sb
487            } else {
488                tile_width_sb
489            } - 1) as u16;
490        }
491        for i in 0..t.rows as usize {
492            pic.height_in_sbs_minus_1[i] = (if i + 1 == t.rows as usize {
493                t.sb_rows - (t.rows - 1) * tile_height_sb
494            } else {
495                tile_height_sb
496            } - 1) as u16;
497        }
498        pic.num_tile_groups_minus1 = 0;
499        if cbr {
500            pic.min_base_qindex = n.min_qp.clamp(1, 255) as u8;
501            pic.max_base_qindex = if n.max_qp > 0 {
502                n.max_qp.clamp(1, 255)
503            } else {
504                255
505            } as u8;
506            pic.bit_offset_qindex = qindex_offset;
507            self.qindex_offset = qindex_offset;
508            pic.bit_offset_loopfilter_params = loopfilter_offset;
509            pic.bit_offset_cdef_params = cdef_offset;
510            pic.size_in_bits_cdef_params = cdef_size;
511            pic.size_in_bits_frame_hdr_obu = frame_header_bits;
512            pic.byte_offset_frame_hdr_obu_size = if frame.key {
513                self.sequence_header.len() as u32
514            } else {
515                0
516            } + 1;
517        }
518        out.push(VAEncPictureParameterBufferType, &pic);
519        if n.packed & VA_ENC_PACKED_HEADER_PICTURE != 0 {
520            out.push_packed(VAEncPackedHeaderPicture, &frame_header, frame_header_bits);
521        }
522        let mut group: VAEncTileGroupBufferAV1 = unsafe { std::mem::zeroed() };
523        group.tg_start = 0;
524        group.tg_end = (t.cols * t.rows - 1) as u8;
525        out.push(VAEncSliceParameterBufferType, &group);
526        Ok(())
527    }
528}
529
530#[cfg(test)]
531mod tests {
532    use super::*;
533
534    /// A picture up to the widest tile is one tile; a wider one takes the fewest columns the
535    /// codec allows, and every layout the checks compute is uniform.
536    #[test]
537    fn tiling_is_uniform_and_minimal() {
538        let hd = Arm::tiles(1920, 1080).unwrap();
539        assert_eq!((hd.cols, hd.rows, hd.cols_log2, hd.rows_log2), (1, 1, 0, 0));
540        assert_eq!((hd.sb_cols, hd.sb_rows), (30, 17));
541        let uhd = Arm::tiles(3840, 2160).unwrap();
542        assert_eq!((uhd.cols, uhd.rows), (1, 1));
543        let wide = Arm::tiles(5120, 1440).unwrap();
544        assert_eq!((wide.cols, wide.cols_log2, wide.min_log2_cols), (2, 1, 1));
545        assert_eq!(
546            Arm::tiles(64, 64).unwrap().max_log2_cols,
547            0,
548            "a single superblock offers no tile columns to choose"
549        );
550    }
551
552    /// An OBU's size field takes the width the driver asked for, and the payload ends in
553    /// its trailing bits whether or not it stood on a byte boundary.
554    #[test]
555    fn obus_carry_fixed_width_sizes_and_trailing_bits() {
556        let mut w = BitWriter::new();
557        w.u(3, 5);
558        let one = obu(OBU_SEQUENCE_HEADER, &w, 1);
559        assert_eq!(one, [0x0a, 0x01, 0b1011_0000]);
560        let four = obu(OBU_SEQUENCE_HEADER, &w, 4);
561        assert_eq!(four, [0x0a, 0x81, 0x80, 0x80, 0x00, 0b1011_0000]);
562        let mut w = BitWriter::new();
563        w.u(8, 0xff);
564        assert_eq!(
565            obu(OBU_FRAME_HEADER, &w, 1),
566            [0x1a, 0x02, 0xff, 0x80],
567            "an aligned payload takes a trailing byte"
568        );
569    }
570}
pixelflux/encoders/vaapi/av1.rs

pixelflux/encoders/vaapi/
av1.rs
