pixelflux/encoders/vaapi/
h264.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 H.264 arm: the sequence, picture, and slice parameters, the packed SPS, PPS, and slice
8//! headers, and the reference lists behind them. The stream is High profile with CABAC and the
9//! 8x8 transform where the driver serves High, a single reference per slice, and a decoded
10//! picture buffer of the frames the level admits: a frame predicts from the newest one the
11//! client still has, which the slice header's reference list modification names whenever that
12//! is not the previous frame.
13
14use va_sys::*;
15
16use super::super::bits::{BitWriter, nal_unit};
17use super::super::codec::h264_level;
18use super::{Buffers, Frame, Negotiated};
19
20/// `log2_max_frame_num_minus4`: sixteen bits of `frame_num`, so the wrap that costs a key
21/// frame comes every 65536 frames.
22const LOG2_MAX_FRAME_NUM_MINUS4: u32 = 12;
23/// The quantizer the PPS names; every slice carries its own delta from it.
24const PIC_INIT_QP: u32 = 26;
25/// The lowest quantizer a picture is coded at: below it radeonsi's VCN 1 codes sharp-edged intra
26/// blocks that decode wrong.
27pub(super) const MIN_QP: u32 = 7;
28
29pub(super) struct Arm {
30    profile: VAProfile,
31    level_idc: u32,
32    mb_width: u32,
33    mb_height: u32,
34    /// Slices per picture and macroblock rows per slice.
35    slices: (u32, u32),
36    cabac: bool,
37    transform_8x8: bool,
38    dpb: u32,
39    sequence: VAEncSequenceParameterBufferH264,
40    sps: Vec<u8>,
41    pps: Vec<u8>,
42    idr_pic_id: u16,
43}
44
45impl Arm {
46    pub(super) fn new(profile: VAProfile) -> Self {
47        Self {
48            profile,
49            level_idc: 0,
50            mb_width: 0,
51            mb_height: 0,
52            slices: (1, 1),
53            cabac: profile != VAProfileH264ConstrainedBaseline,
54            transform_8x8: profile == VAProfileH264High,
55            dpb: 1,
56            sequence: unsafe { std::mem::zeroed() },
57            sps: Vec::new(),
58            pps: Vec::new(),
59            idr_pic_id: 0,
60        }
61    }
62
63    /// How many values `frame_num` takes before it wraps.
64    pub(super) fn frame_num_range(&self) -> u32 {
65        1 << (4 + LOG2_MAX_FRAME_NUM_MINUS4)
66    }
67
68    pub(super) fn slices(&self) -> (u32, u32) {
69        self.slices
70    }
71
72    /// The `level_idc` the sequence names.
73    pub(super) fn level(&self) -> u32 {
74        self.level_idc
75    }
76
77    /// The stream's sequence: the level the ladder names at this geometry and rate, never
78    /// below the one the decoded picture buffer was sized for, the picture size in macroblocks
79    /// with the cropping that trims it, and the SPS and PPS the key frames carry.
80    pub(super) fn configure(
81        &mut self,
82        n: &Negotiated,
83        surface_width: u32,
84        surface_height: u32,
85        slices: (u32, u32),
86    ) {
87        self.mb_width = surface_width / 16;
88        self.mb_height = surface_height / 16;
89        self.slices = slices;
90        self.dpb = n.dpb;
91        self.level_idc = h264_level(
92            n.width,
93            n.height,
94            n.fps.ceil(),
95            u64::from(n.bits_per_second),
96        )
97        .max(n.dpb_level);
98        let profile_idc = if self.profile == VAProfileH264High {
99            100
100        } else if self.profile == VAProfileH264Main {
101            77
102        } else {
103            66
104        };
105        let crop = (self.mb_width * 16 != n.width) || (self.mb_height * 16 != n.height);
106
107        let mut w = BitWriter::new();
108        w.u(8, profile_idc);
109        let (set1, set3, set4, set5) = (
110            profile_idc != 100,
111            false,
112            profile_idc != 66,
113            profile_idc != 66,
114        );
115        w.flag(false);
116        w.flag(set1);
117        w.flag(false);
118        w.flag(set3);
119        w.flag(set4);
120        w.flag(set5);
121        w.u(2, 0);
122        w.u(8, u64::from(self.level_idc));
123        w.ue(0);
124        if profile_idc == 100 {
125            w.ue(1);
126            w.ue(0);
127            w.ue(0);
128            w.flag(false);
129            w.flag(false);
130        }
131        w.ue(LOG2_MAX_FRAME_NUM_MINUS4);
132        w.ue(2);
133        w.ue(self.dpb);
134        w.flag(false);
135        w.ue(self.mb_width - 1);
136        w.ue(self.mb_height - 1);
137        w.flag(true);
138        w.flag(true);
139        w.flag(crop);
140        if crop {
141            w.ue(0);
142            w.ue((self.mb_width * 16 - n.width) / 2);
143            w.ue(0);
144            w.ue((self.mb_height * 16 - n.height) / 2);
145        }
146        w.flag(true);
147        w.flag(false);
148        w.flag(false);
149        w.flag(true);
150        w.u(3, 5);
151        w.flag(false);
152        w.flag(true);
153        w.u(8, 1);
154        w.u(8, 1);
155        w.u(8, 1);
156        w.flag(false);
157        w.flag(true);
158        w.u(32, u64::from(n.fps.den));
159        w.u(32, 2 * u64::from(n.fps.num));
160        w.flag(true);
161        w.flag(false);
162        w.flag(false);
163        w.flag(false);
164        w.flag(true);
165        w.flag(true);
166        w.ue(0);
167        w.ue(0);
168        w.ue(15);
169        w.ue(15);
170        w.ue(0);
171        w.ue(self.dpb);
172        w.trailing_bits();
173        self.sps = nal_unit(&[0x67], w.as_bytes());
174
175        let mut w = BitWriter::new();
176        w.ue(0);
177        w.ue(0);
178        w.flag(self.cabac);
179        w.flag(false);
180        w.ue(0);
181        w.ue(0);
182        w.ue(0);
183        w.flag(false);
184        w.u(2, 0);
185        w.se(PIC_INIT_QP as i32 - 26);
186        w.se(0);
187        w.se(0);
188        w.flag(false);
189        w.flag(false);
190        w.flag(false);
191        if self.transform_8x8 {
192            w.flag(true);
193            w.flag(false);
194            w.se(0);
195        }
196        w.trailing_bits();
197        self.pps = nal_unit(&[0x68], w.as_bytes());
198
199        let s = &mut self.sequence;
200        *s = unsafe { std::mem::zeroed() };
201        s.seq_parameter_set_id = 0;
202        s.level_idc = self.level_idc.min(n.h264_driver_level) as u8;
203        s.intra_period = i32::MAX as u32;
204        s.intra_idr_period = i32::MAX as u32;
205        s.ip_period = 1;
206        s.bits_per_second = n.bits_per_second;
207        s.max_num_ref_frames = self.dpb;
208        s.picture_width_in_mbs = self.mb_width as u16;
209        s.picture_height_in_mbs = self.mb_height as u16;
210        unsafe {
211            let f = &mut s.seq_fields.bits;
212            f.set_chroma_format_idc(1);
213            f.set_frame_mbs_only_flag(1);
214            f.set_direct_8x8_inference_flag(1);
215            f.set_log2_max_frame_num_minus4(LOG2_MAX_FRAME_NUM_MINUS4);
216            f.set_pic_order_cnt_type(2);
217            let v = &mut s.vui_fields.bits;
218            v.set_timing_info_present_flag(1);
219            v.set_bitstream_restriction_flag(1);
220            v.set_log2_max_mv_length_horizontal(15);
221            v.set_log2_max_mv_length_vertical(15);
222        }
223        s.frame_cropping_flag = u8::from(crop);
224        if crop {
225            s.frame_crop_right_offset = (self.mb_width * 16 - n.width) / 2;
226            s.frame_crop_bottom_offset = (self.mb_height * 16 - n.height) / 2;
227        }
228        s.vui_parameters_present_flag = 1;
229        s.num_units_in_tick = n.fps.den;
230        s.time_scale = 2 * n.fps.num;
231    }
232
233    pub(super) fn sequence(&mut self, n: &Negotiated, out: &mut Buffers) {
234        out.push(VAEncSequenceParameterBufferType, &self.sequence);
235        if n.packed & VA_ENC_PACKED_HEADER_SEQUENCE != 0 {
236            let mut bytes = self.sps.clone();
237            bytes.extend_from_slice(&self.pps);
238            let bits = 8 * bytes.len() as u32;
239            out.push_packed(VAEncPackedHeaderSequence, &bytes, bits);
240        }
241    }
242
243    /// The picture: its frame number and order count since the key frame, the decoded picture
244    /// buffer as `ReferenceFrames`, and every slice with its packed header, whose reference
245    /// list modification moves any reference but the newest short-term frame, which the
246    /// default list puts first, to the front: a short-term one by its distance in frame numbers,
247    /// an anchor by its long-term index. In a session keeping anchors the key frame is marked
248    /// long-term under index 0 and each anchor after it under its index (memory management
249    /// operation 6), the one it is marked over leaving the buffer.
250    pub(super) fn picture(
251        &mut self,
252        n: &Negotiated,
253        frame: &Frame,
254        out: &mut Buffers,
255    ) -> Result<(), String> {
256        let frame_num = ((frame.pts - frame.key_pts) % u64::from(self.frame_num_range())) as u32;
257        let poc = 2 * (frame.pts - frame.key_pts) as i32;
258        if frame.key {
259            self.idr_pic_id = self
260                .idr_pic_id
261                .wrapping_add(if frame.pts == 0 { 0 } else { 1 });
262        }
263        let long_term = |pts: u64| frame.long_term.iter().find(|l| l.0 == pts).map(|l| l.1);
264        let picture = |pts: u64, surface: VASurfaceID| VAPictureH264 {
265            picture_id: surface,
266            frame_idx: long_term(pts).map_or(
267                ((pts - frame.key_pts) % u64::from(self.frame_num_range())) as u32,
268                u32::from,
269            ),
270            flags: if long_term(pts).is_some() {
271                VA_PICTURE_H264_LONG_TERM_REFERENCE
272            } else {
273                VA_PICTURE_H264_SHORT_TERM_REFERENCE
274            },
275            TopFieldOrderCnt: 2 * (pts - frame.key_pts) as i32,
276            BottomFieldOrderCnt: 2 * (pts - frame.key_pts) as i32,
277            va_reserved: [0; 4],
278        };
279        let invalid = VAPictureH264 {
280            picture_id: VA_INVALID_ID,
281            frame_idx: 0,
282            flags: VA_PICTURE_H264_INVALID,
283            TopFieldOrderCnt: 0,
284            BottomFieldOrderCnt: 0,
285            va_reserved: [0; 4],
286        };
287
288        let marked = if frame.key && frame.anchored {
289            Some(0)
290        } else {
291            frame.anchor
292        };
293        let mut pic: VAEncPictureParameterBufferH264 = unsafe { std::mem::zeroed() };
294        pic.CurrPic = VAPictureH264 {
295            picture_id: frame.recon,
296            frame_idx: marked.map_or(frame_num, u32::from),
297            flags: if marked.is_some() {
298                VA_PICTURE_H264_LONG_TERM_REFERENCE
299            } else {
300                0
301            },
302            TopFieldOrderCnt: poc,
303            BottomFieldOrderCnt: poc,
304            va_reserved: [0; 4],
305        };
306        pic.ReferenceFrames = [invalid; 16];
307        let mut held: Vec<(u64, VASurfaceID)> = frame
308            .held
309            .iter()
310            .filter(|h| !frame.key && h.0 >= frame.key_pts)
311            .map(|h| (h.0, h.1))
312            .collect();
313        held.sort_by_key(|h| std::cmp::Reverse(h.0));
314        for (i, &(pts, surface)) in held.iter().take(16).enumerate() {
315            pic.ReferenceFrames[i] = picture(pts, surface);
316        }
317        pic.coded_buf = frame.coded;
318        pic.frame_num = frame_num as u16;
319        pic.pic_init_qp = PIC_INIT_QP as u8;
320        unsafe {
321            let f = &mut pic.pic_fields.bits;
322            f.set_idr_pic_flag(u32::from(frame.key));
323            f.set_reference_pic_flag(1);
324            f.set_entropy_coding_mode_flag(u32::from(self.cabac));
325            f.set_transform_8x8_mode_flag(u32::from(self.transform_8x8));
326        }
327        out.push(VAEncPictureParameterBufferType, &pic);
328
329        let reference = frame.reference.map(|(pts, surface)| picture(pts, surface));
330        if !frame.key && reference.is_none() {
331            return Err("a predicted H.264 picture without a reference".into());
332        }
333        let slice_qp = if n.rc_mode == VA_RC_CQP {
334            frame.qp.clamp(MIN_QP, 51)
335        } else {
336            PIC_INIT_QP
337        };
338        let slice_qp_delta = slice_qp as i32 - PIC_INIT_QP as i32;
339        let (slice_count, slice_rows) = self.slices;
340        for i in 0..slice_count {
341            let first_row = i * slice_rows;
342            let rows = if i + 1 == slice_count {
343                self.mb_height - first_row
344            } else {
345                slice_rows
346            };
347            let first_mb = first_row * self.mb_width;
348            let mbs = rows * self.mb_width;
349            let slice_type = if frame.key { 7 } else { 5 };
350
351            if n.packed & VA_ENC_PACKED_HEADER_SLICE != 0 {
352                let mut w = BitWriter::new();
353                w.ue(first_mb);
354                w.ue(slice_type);
355                w.ue(0);
356                w.u(4 + LOG2_MAX_FRAME_NUM_MINUS4, u64::from(frame_num));
357                if frame.key {
358                    w.ue(u32::from(self.idr_pic_id));
359                }
360                if !frame.key {
361                    w.flag(false);
362                    let (ref_pts, _) = frame.reference.unwrap();
363                    if let Some(index) = long_term(ref_pts) {
364                        w.flag(true);
365                        w.ue(2);
366                        w.ue(u32::from(index));
367                        w.ue(3);
368                    } else {
369                        let ref_frame_num =
370                            ((ref_pts - frame.key_pts) % u64::from(self.frame_num_range())) as u32;
371                        let newest = held.iter().find(|h| long_term(h.0).is_none()).map(|h| h.0)
372                            == Some(ref_pts);
373                        w.flag(!newest);
374                        if !newest {
375                            let diff = (frame_num + self.frame_num_range() - ref_frame_num)
376                                % self.frame_num_range();
377                            w.ue(0);
378                            w.ue(diff - 1);
379                            w.ue(3);
380                        }
381                    }
382                }
383                if frame.key {
384                    w.flag(false);
385                    w.flag(frame.anchored);
386                } else if let Some(index) = frame.anchor {
387                    w.flag(true);
388                    w.ue(6);
389                    w.ue(u32::from(index));
390                    w.ue(0);
391                } else {
392                    w.flag(false);
393                }
394                if self.cabac && !frame.key {
395                    w.ue(0);
396                }
397                w.se(slice_qp_delta);
398                let header = nal_unit(if frame.key { &[0x65] } else { &[0x41] }, w.as_bytes());
399                let bits = 8 * (header.len() as u32 - w.as_bytes().len() as u32) + w.len() as u32;
400                out.push_packed(VAEncPackedHeaderSlice, &header, bits);
401            }
402
403            let mut slice: VAEncSliceParameterBufferH264 = unsafe { std::mem::zeroed() };
404            slice.macroblock_address = first_mb;
405            slice.num_macroblocks = mbs;
406            slice.macroblock_info = VA_INVALID_ID;
407            slice.slice_type = (slice_type % 5) as u8;
408            slice.pic_parameter_set_id = 0;
409            slice.idr_pic_id = self.idr_pic_id;
410            slice.pic_order_cnt_lsb = 0;
411            slice.direct_spatial_mv_pred_flag = 1;
412            slice.RefPicList0 = [invalid; 32];
413            slice.RefPicList1 = [invalid; 32];
414            if let Some(r) = reference {
415                slice.RefPicList0[0] = r;
416            }
417            slice.slice_qp_delta = slice_qp_delta as i8;
418            out.push(VAEncSliceParameterBufferType, &slice);
419        }
420        Ok(())
421    }
422}
pixelflux/encoders/vaapi/h264.rs

pixelflux/encoders/vaapi/
h264.rs
