cros_codecs/encoder/stateless/av1/
vaapi.rs

1// Copyright 2024 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::num::TryFromIntError;
6use std::rc::Rc;
7
8use libva::constants::VA_INVALID_ID;
9use libva::AV1EncLoopFilterFlags;
10use libva::AV1EncLoopRestorationFlags;
11use libva::AV1EncModeControlFlags;
12use libva::AV1EncPictureFlags;
13use libva::AV1EncQMatrixFlags;
14use libva::AV1EncSeqFields;
15use libva::AV1EncTileGroupObuHdrInfo;
16use libva::EncCodedBuffer;
17use libva::EncPictureParameterBufferAV1;
18use libva::EncSegParamAV1;
19use libva::EncSegParamFlagsAV1;
20use libva::EncSequenceParameterBufferAV1;
21use libva::EncTileGroupBufferAV1;
22use libva::Picture;
23use libva::RefFrameCtrlAV1;
24use libva::Surface;
25use libva::SurfaceMemoryDescriptor;
26use libva::VAProfile::VAProfileAV1Profile0;
27use libva::VAProfile::VAProfileAV1Profile1;
28use libva::VaError;
29
30use crate::backend::vaapi::encoder::CodedOutputPromise;
31use crate::backend::vaapi::encoder::Reconstructed;
32use crate::backend::vaapi::encoder::VaapiBackend;
33use crate::codec::av1::parser::Profile;
34use crate::codec::av1::parser::ReferenceFrameType;
35use crate::codec::av1::parser::CDEF_MAX;
36use crate::codec::av1::parser::MAX_SEGMENTS;
37use crate::codec::av1::parser::MAX_TILE_COLS;
38use crate::codec::av1::parser::MAX_TILE_ROWS;
39use crate::codec::av1::parser::REFS_PER_FRAME;
40use crate::codec::av1::parser::SEG_LVL_MAX;
41use crate::encoder::av1::EncoderConfig;
42use crate::encoder::av1::AV1;
43use crate::encoder::stateless::av1::predictor::MAX_BASE_QINDEX;
44use crate::encoder::stateless::av1::predictor::MIN_BASE_QINDEX;
45use crate::encoder::stateless::av1::BackendRequest;
46use crate::encoder::stateless::av1::StatelessAV1EncoderBackend;
47use crate::encoder::stateless::ReadyPromise;
48use crate::encoder::stateless::StatelessBackendError;
49use crate::encoder::stateless::StatelessBackendResult;
50use crate::encoder::stateless::StatelessEncoder;
51use crate::encoder::stateless::StatelessVideoEncoderBackend;
52use crate::encoder::EncodeError;
53use crate::encoder::EncodeResult;
54use crate::encoder::RateControl;
55use crate::BlockingMode;
56use crate::Fourcc;
57use crate::Resolution;
58
59type Request<H> = BackendRequest<H, Reconstructed>;
60
61#[derive(thiserror::Error, Debug)]
62pub enum BackendError {
63    #[error(transparent)]
64    ConversionError(#[from] TryFromIntError),
65
66    #[error("vaBeginPicture failed: {0}")]
67    BeginPictureError(VaError),
68    #[error("vaRenderPicture failed: {0}")]
69    RenderPictureError(VaError),
70    #[error("vaRenderPicture failed: {0}")]
71    EndPictureError(VaError),
72}
73
74impl From<BackendError> for StatelessBackendError {
75    fn from(value: BackendError) -> Self {
76        StatelessBackendError::Other(anyhow::anyhow!(value))
77    }
78}
79
80type Result<T> = std::result::Result<T, BackendError>;
81
82impl<M, Handle> StatelessVideoEncoderBackend<AV1> for VaapiBackend<M, Handle>
83where
84    M: SurfaceMemoryDescriptor,
85    Handle: std::borrow::Borrow<Surface<M>>,
86{
87    type Picture = Handle;
88    type Reconstructed = Reconstructed;
89    type CodedPromise = CodedOutputPromise<M, Handle>;
90    type ReconPromise = ReadyPromise<Self::Reconstructed>;
91}
92
93impl<M, H> VaapiBackend<M, H>
94where
95    M: SurfaceMemoryDescriptor,
96    H: std::borrow::Borrow<Surface<M>> + 'static,
97{
98    fn build_seq_param(request: &Request<H>) -> Result<EncSequenceParameterBufferAV1> {
99        assert!(
100            request.sequence.operating_points_cnt_minus_1 == 0,
101            "Only a single operating point is supported for now"
102        );
103        const OPERATING_POINT: usize = 0;
104
105        let seq_profile = request.sequence.seq_profile as u8;
106        let seq_level_idx = request.sequence.operating_points[OPERATING_POINT].seq_level_idx;
107        let seq_tier = request.sequence.operating_points[OPERATING_POINT].seq_tier;
108        let hierarchical_flag = 0;
109
110        // TODO: Enable bitrate control
111        let bits_per_second = 0;
112
113        // AV1 5.5.2
114        let bit_depth_minus8 = if request.sequence.seq_profile == Profile::Profile2
115            && request.sequence.color_config.high_bitdepth
116        {
117            if request.sequence.color_config.twelve_bit {
118                12
119            } else {
120                10
121            }
122        } else if request.sequence.color_config.high_bitdepth {
123            10
124        } else {
125            8
126        };
127
128        let order_hint_bits_minus_1 = u8::try_from(request.sequence.order_hint_bits_minus_1)?;
129
130        Ok(EncSequenceParameterBufferAV1::new(
131            seq_profile,
132            seq_level_idx,
133            seq_tier,
134            hierarchical_flag,
135            request.intra_period,
136            request.ip_period,
137            bits_per_second,
138            &AV1EncSeqFields::new(
139                request.sequence.still_picture,
140                request.sequence.use_128x128_superblock,
141                request.sequence.enable_filter_intra,
142                request.sequence.enable_intra_edge_filter,
143                request.sequence.enable_interintra_compound,
144                request.sequence.enable_masked_compound,
145                request.sequence.enable_warped_motion,
146                request.sequence.enable_dual_filter,
147                request.sequence.enable_order_hint,
148                request.sequence.enable_jnt_comp,
149                request.sequence.enable_ref_frame_mvs,
150                request.sequence.enable_superres,
151                request.sequence.enable_cdef,
152                request.sequence.enable_restoration,
153                bit_depth_minus8,
154                request.sequence.color_config.subsampling_x,
155                request.sequence.color_config.subsampling_y,
156                request.sequence.color_config.mono_chrome,
157            ),
158            order_hint_bits_minus_1,
159        ))
160    }
161
162    fn build_ref_ctrl(
163        refs: &[Option<Rc<Reconstructed>>; REFS_PER_FRAME],
164        ctrl: &[ReferenceFrameType; REFS_PER_FRAME],
165    ) -> RefFrameCtrlAV1 {
166        let ctrl = ctrl.map(|type_| {
167            if type_ == ReferenceFrameType::Intra {
168                return 0;
169            }
170
171            let idx = type_ as u32 - ReferenceFrameType::Last as u32;
172            if refs[idx as usize].is_none() {
173                return 0;
174            }
175
176            type_ as u32
177        });
178
179        RefFrameCtrlAV1::new(
180            ctrl[0], ctrl[1], ctrl[2], ctrl[3], ctrl[4], ctrl[5], ctrl[6],
181        )
182    }
183
184    fn build_pic_param(
185        request: &Request<H>,
186        recon: &Reconstructed,
187        coded: &EncCodedBuffer,
188    ) -> Result<EncPictureParameterBufferAV1> {
189        let coded_buf = coded.id();
190        let reconstructed_frame = recon.surface_id();
191
192        let mut reference_frames = [VA_INVALID_ID; 8];
193
194        for (i, frame) in reference_frames.iter_mut().enumerate().take(REFS_PER_FRAME) {
195            let Some(ref_frame) = &request.references[i] else {
196                continue;
197            };
198
199            *frame = ref_frame.surface_id();
200        }
201
202        let mut ref_frame_idx = [0; 7];
203        for (i, idx) in ref_frame_idx.iter_mut().enumerate() {
204            *idx = request.frame.ref_frame_idx[i];
205        }
206
207        let frame_width_minus_1 = u16::try_from(request.frame.frame_width - 1)?;
208        let frame_height_minus_1 = u16::try_from(request.frame.frame_height - 1)?;
209
210        // Single temporal layer is used.
211        const HIERARCHICAL_LEVEL_PLUS1: u8 = 0;
212
213        let primary_ref_frame = u8::try_from(request.frame.primary_ref_frame)?;
214        let order_hint = u8::try_from(request.frame.order_hint)?;
215        let refresh_frame_flags = u8::try_from(request.frame.refresh_frame_flags)?;
216
217        let ref_frame_ctrl_l0 =
218            Self::build_ref_ctrl(&request.references, &request.ref_frame_ctrl_l0);
219        let ref_frame_ctrl_l1 =
220            Self::build_ref_ctrl(&request.references, &request.ref_frame_ctrl_l1);
221
222        let frame_type = request.frame.frame_type as u32;
223
224        // Export Frame Obu rather then TileGroup Obu
225        const ENABLE_FRAME_OBU: bool = false;
226
227        // We don't use long term reference frames for now.
228        const LONG_TERM_REFERENCE: bool = false;
229
230        // Current we always expect the reconstructed frame.
231        const DISABLE_FRAME_RECON: bool = false;
232
233        // Palette mode is not used.
234        const PALETTE_MODE_ENABLE: bool = false;
235
236        // Use 16x16 block size for now.
237        // TODO: Use maximum available
238        const SEG_ID_BLOCK_SIZE: u8 = 0;
239
240        // Use single tile group;
241        const NUM_TILE_GROUPS_MINUS1: u8 = 0;
242
243        let filter_level = [
244            request.frame.loop_filter_params.loop_filter_level[0],
245            request.frame.loop_filter_params.loop_filter_level[1],
246        ];
247
248        let filter_level_u = request.frame.loop_filter_params.loop_filter_level[2];
249        let filter_level_v = request.frame.loop_filter_params.loop_filter_level[3];
250
251        let superres_scale_denominator = u8::try_from(request.frame.superres_denom)?;
252
253        let interpolation_filter = request.frame.interpolation_filter as u8;
254
255        let mut loop_filter_ref_deltas = [0; 8];
256        for (i, delta) in loop_filter_ref_deltas.iter_mut().enumerate() {
257            *delta = request.frame.loop_filter_params.loop_filter_ref_deltas[i];
258        }
259
260        let base_qindex = u8::try_from(request.frame.quantization_params.base_q_idx)?;
261        let y_dc_delta_q = i8::try_from(request.frame.quantization_params.delta_q_y_dc)?;
262        let u_dc_delta_q = i8::try_from(request.frame.quantization_params.delta_q_u_dc)?;
263        let u_ac_delta_q = i8::try_from(request.frame.quantization_params.delta_q_u_ac)?;
264        let v_dc_delta_q = i8::try_from(request.frame.quantization_params.delta_q_v_dc)?;
265        let v_ac_delta_q = i8::try_from(request.frame.quantization_params.delta_q_v_ac)?;
266
267        // Clamp tunings's quaility range to correct range
268        let min_base_qindex = request.tunings.min_quality.max(MIN_BASE_QINDEX);
269        let min_base_qindex = u8::try_from(min_base_qindex)?;
270        let max_base_qindex = request.tunings.max_quality.min(MAX_BASE_QINDEX);
271        let max_base_qindex = u8::try_from(max_base_qindex)?;
272
273        let qm_y = u16::try_from(request.frame.quantization_params.qm_y)?;
274        let qm_u = u16::try_from(request.frame.quantization_params.qm_u)?;
275        let qm_v = u16::try_from(request.frame.quantization_params.qm_v)?;
276
277        let tx_mode = request.frame.tx_mode as u32;
278
279        // Make driver make decision use single reference or compound reference.
280        const REFERENCE_MODE: u32 = 0 /* REFERENCE_MODE_SELECT */;
281
282        let segmentation_temporal_update = request
283            .frame
284            .segmentation_params
285            .segmentation_temporal_update;
286
287        const SEGMENT_NUMBER: u8 = 0;
288        assert!(
289            !request.frame.segmentation_params.segmentation_enabled,
290            "Unsupported segmentation_enabled=1"
291        );
292
293        // Segementation feature mask
294        let mut feature_mask = [0u8; MAX_SEGMENTS];
295        for (seg, mask) in feature_mask.iter_mut().enumerate() {
296            for lvl in 0..u8::try_from(SEG_LVL_MAX)? {
297                if request.frame.segmentation_params.feature_enabled[seg][lvl as usize] {
298                    *mask |= 1u8 << lvl;
299                }
300            }
301        }
302
303        assert!(
304            request.frame.tile_info.tile_cols == 1
305                && request.frame.tile_info.tile_cols_log2 == 0
306                && request.frame.tile_info.tile_rows == 1
307                && request.frame.tile_info.tile_rows_log2 == 0,
308            "Single tile is only supported for now"
309        );
310        let tile_cols = u8::try_from(request.frame.tile_info.tile_cols)?;
311        let tile_rows = u8::try_from(request.frame.tile_info.tile_rows)?;
312
313        let mut width_in_sbs_minus_1 = [0u16; MAX_TILE_COLS - 1];
314        for (i, width) in width_in_sbs_minus_1.iter_mut().enumerate() {
315            *width = u16::try_from(request.frame.tile_info.width_in_sbs_minus_1[i])?;
316        }
317
318        let mut height_in_sbs_minus_1 = [0u16; MAX_TILE_ROWS - 1];
319        for (i, height) in height_in_sbs_minus_1.iter_mut().enumerate() {
320            *height = u16::try_from(request.frame.tile_info.height_in_sbs_minus_1[i])?;
321        }
322
323        let context_update_tile_id = u16::try_from(request.frame.tile_info.context_update_tile_id)?;
324
325        let cdef_damping_minus_3 = u8::try_from(request.frame.cdef_params.cdef_damping - 3)?;
326
327        let cdef_bits = u8::try_from(request.frame.cdef_params.cdef_bits)?;
328        let mut cdef_y_strengths = [0u8; CDEF_MAX];
329        for (i, strength) in cdef_y_strengths.iter_mut().enumerate() {
330            *strength = u8::try_from(request.frame.cdef_params.cdef_y_pri_strength[i])?;
331        }
332
333        let mut cdef_uv_strengths = [0u8; CDEF_MAX];
334        for (i, strength) in cdef_uv_strengths.iter_mut().enumerate() {
335            *strength = u8::try_from(request.frame.cdef_params.cdef_uv_pri_strength[i])?;
336        }
337
338        let yframe_restoration_type =
339            request.frame.loop_restoration_params.frame_restoration_type[0] as u16;
340        let cbframe_restoration_type =
341            request.frame.loop_restoration_params.frame_restoration_type[1] as u16;
342        let crframe_restoration_type =
343            request.frame.loop_restoration_params.frame_restoration_type[2] as u16;
344
345        let lr_unit_shift = u16::from(request.frame.loop_restoration_params.lr_unit_shift);
346        let lr_uv_shift = request.frame.loop_restoration_params.lr_uv_shift != 0;
347
348        // Warped motion params
349        let wm = [Default::default(); REFS_PER_FRAME];
350
351        // Ignore by driver
352        const BIT_OFFSET_QINDEX: u32 = 0;
353        const BIT_OFFSET_SEGMENTATION: u32 = 0;
354        const BIT_OFFSET_LOOPFILTER_PARAMS: u32 = 0;
355        const BIT_OFFSET_CDEF_PARAMS: u32 = 0;
356        const SIZE_IN_BITS_CDEF_PARAMS: u32 = 0;
357
358        // TODO: use packed header and fill for bitrate control
359        const BYTE_OFFSET_FRAME_HDR_OBU_SIZE: u32 = 0;
360        const SIZE_IN_BITS_FRAME_HDR_OBU: u32 = 0;
361
362        let temporal_id = u8::try_from(request.frame.obu_header.temporal_id)?;
363        let spatial_id = u8::try_from(request.frame.obu_header.spatial_id)?;
364
365        const NUMBER_SKIP_FRAMES: u8 = 0;
366        const SKIP_FRAMES_REDUCED_SIZE: i32 = 0;
367
368        Ok(EncPictureParameterBufferAV1::new(
369            frame_width_minus_1,
370            frame_height_minus_1,
371            reconstructed_frame,
372            coded_buf,
373            reference_frames,
374            ref_frame_idx,
375            HIERARCHICAL_LEVEL_PLUS1,
376            primary_ref_frame,
377            order_hint,
378            refresh_frame_flags,
379            &ref_frame_ctrl_l0,
380            &ref_frame_ctrl_l1,
381            &AV1EncPictureFlags::new(
382                frame_type,
383                request.frame.error_resilient_mode,
384                request.frame.disable_cdf_update,
385                request.frame.use_superres,
386                request.frame.allow_high_precision_mv,
387                request.frame.use_ref_frame_mvs,
388                request.frame.disable_frame_end_update_cdf,
389                request.frame.reduced_tx_set,
390                ENABLE_FRAME_OBU,
391                LONG_TERM_REFERENCE,
392                DISABLE_FRAME_RECON,
393                request.frame.allow_intrabc,
394                PALETTE_MODE_ENABLE,
395            ),
396            SEG_ID_BLOCK_SIZE,
397            NUM_TILE_GROUPS_MINUS1,
398            temporal_id,
399            filter_level,
400            filter_level_u,
401            filter_level_v,
402            &AV1EncLoopFilterFlags::new(
403                request.frame.loop_filter_params.loop_filter_sharpness,
404                request.frame.loop_filter_params.loop_filter_delta_enabled,
405                request.frame.loop_filter_params.loop_filter_delta_update,
406            ),
407            superres_scale_denominator,
408            interpolation_filter,
409            loop_filter_ref_deltas,
410            request.frame.loop_filter_params.loop_filter_mode_deltas,
411            base_qindex,
412            y_dc_delta_q,
413            u_dc_delta_q,
414            u_ac_delta_q,
415            v_dc_delta_q,
416            v_ac_delta_q,
417            min_base_qindex,
418            max_base_qindex,
419            &AV1EncQMatrixFlags::new(
420                request.frame.quantization_params.using_qmatrix,
421                qm_y,
422                qm_u,
423                qm_v,
424            ),
425            &AV1EncModeControlFlags::new(
426                request.frame.quantization_params.delta_q_present,
427                request.frame.quantization_params.delta_q_res,
428                request.frame.loop_filter_params.delta_lf_present,
429                request.frame.loop_filter_params.delta_lf_res as u32,
430                request.frame.loop_filter_params.delta_lf_multi,
431                tx_mode,
432                REFERENCE_MODE,
433                request.frame.skip_mode_present,
434            ),
435            &EncSegParamAV1::new(
436                &EncSegParamFlagsAV1::new(
437                    request.frame.segmentation_params.segmentation_enabled,
438                    request.frame.segmentation_params.segmentation_update_map,
439                    segmentation_temporal_update,
440                ),
441                SEGMENT_NUMBER,
442                request.frame.segmentation_params.feature_data,
443                feature_mask,
444            ),
445            tile_cols,
446            tile_rows,
447            width_in_sbs_minus_1,
448            height_in_sbs_minus_1,
449            context_update_tile_id,
450            cdef_damping_minus_3,
451            cdef_bits,
452            cdef_y_strengths,
453            cdef_uv_strengths,
454            &AV1EncLoopRestorationFlags::new(
455                yframe_restoration_type,
456                cbframe_restoration_type,
457                crframe_restoration_type,
458                lr_unit_shift,
459                lr_uv_shift,
460            ),
461            wm,
462            BIT_OFFSET_QINDEX,
463            BIT_OFFSET_SEGMENTATION,
464            BIT_OFFSET_LOOPFILTER_PARAMS,
465            BIT_OFFSET_CDEF_PARAMS,
466            SIZE_IN_BITS_CDEF_PARAMS,
467            BYTE_OFFSET_FRAME_HDR_OBU_SIZE,
468            SIZE_IN_BITS_FRAME_HDR_OBU,
469            &AV1EncTileGroupObuHdrInfo::new(
470                request.frame.obu_header.extension_flag,
471                request.frame.obu_header.has_size_field,
472                temporal_id,
473                spatial_id,
474            ),
475            NUMBER_SKIP_FRAMES,
476            SKIP_FRAMES_REDUCED_SIZE,
477        ))
478    }
479
480    fn build_tile_group_param() -> EncTileGroupBufferAV1 {
481        // Single tile is only supported for now.
482        EncTileGroupBufferAV1::new(0, 0)
483    }
484}
485
486impl<M, H> StatelessAV1EncoderBackend for VaapiBackend<M, H>
487where
488    M: SurfaceMemoryDescriptor,
489    H: std::borrow::Borrow<Surface<M>> + 'static,
490{
491    fn encode_tile_group(
492        &mut self,
493        request: BackendRequest<Self::Picture, Self::Reconstructed>,
494    ) -> StatelessBackendResult<(Self::ReconPromise, Self::CodedPromise)> {
495        let coded_buf = self.new_coded_buffer(&request.tunings.rate_control)?;
496        let recon = self.new_scratch_picture()?;
497
498        let seq_param = Self::build_seq_param(&request)?;
499        let seq_param =
500            libva::BufferType::EncSequenceParameter(libva::EncSequenceParameter::AV1(seq_param));
501
502        let pic_param = Self::build_pic_param(&request, &recon, &coded_buf)?;
503        let pic_param =
504            libva::BufferType::EncPictureParameter(libva::EncPictureParameter::AV1(pic_param));
505
506        let tg_param = Self::build_tile_group_param();
507        let tg_param =
508            libva::BufferType::EncSliceParameter(libva::EncSliceParameter::AV1(tg_param));
509
510        let mut references = Vec::new();
511
512        for ref_frame in &request.references {
513            let Some(ref_frame) = ref_frame else {
514                continue;
515            };
516
517            references.push(ref_frame.clone() as Rc<dyn std::any::Any>);
518        }
519
520        let mut picture = Picture::new(
521            request.input_meta.timestamp,
522            Rc::clone(self.context()),
523            request.input,
524        );
525
526        picture.add_buffer(self.context().create_buffer(seq_param)?);
527        picture.add_buffer(self.context().create_buffer(pic_param)?);
528        picture.add_buffer(self.context().create_buffer(tg_param)?);
529
530        // Start processing the picture encoding
531        let picture = picture.begin().map_err(BackendError::BeginPictureError)?;
532        let picture = picture.render().map_err(BackendError::RenderPictureError)?;
533        let picture = picture.end().map_err(BackendError::EndPictureError)?;
534
535        // libva will handle the synchronization of reconstructed surface with implicit fences.
536        // Therefore return the reconstructed frame immediately.
537        let reference_promise = ReadyPromise::from(recon);
538
539        let bitstream_promise =
540            CodedOutputPromise::new(picture, references, coded_buf, request.coded_output);
541
542        Ok((reference_promise, bitstream_promise))
543    }
544}
545
546impl<M, H> StatelessEncoder<AV1, H, VaapiBackend<M, H>>
547where
548    M: SurfaceMemoryDescriptor,
549    H: std::borrow::Borrow<libva::Surface<M>> + 'static,
550{
551    pub fn new_vaapi(
552        display: Rc<libva::Display>,
553        config: EncoderConfig,
554        fourcc: Fourcc,
555        coded_size: Resolution,
556        low_power: bool,
557        blocking_mode: BlockingMode,
558    ) -> EncodeResult<Self> {
559        let va_profile = match config.profile {
560            Profile::Profile0 => VAProfileAV1Profile0,
561            Profile::Profile1 => VAProfileAV1Profile1,
562            _ => return Err(StatelessBackendError::UnsupportedProfile.into()),
563        };
564
565        if !matches!(
566            config.initial_tunings.rate_control,
567            RateControl::ConstantQuality(_)
568        ) {
569            return Err(EncodeError::Unsupported);
570        }
571
572        let backend = VaapiBackend::new(
573            display,
574            va_profile,
575            fourcc,
576            coded_size,
577            libva::constants::VA_RC_CQP,
578            low_power,
579        )?;
580
581        Self::new_av1(backend, config, blocking_mode)
582    }
583}
584
585#[cfg(test)]
586mod tests {
587    use libva::constants::VA_RT_FORMAT_YUV420;
588    use libva::constants::VA_RT_FORMAT_YUV420_10;
589    use libva::Display;
590    use libva::UsageHint;
591    use libva::VAEntrypoint::VAEntrypointEncSliceLP;
592    use libva::VAProfile::VAProfileAV1Profile0;
593
594    use super::*;
595    use crate::backend::vaapi::encoder::tests::upload_test_frame_nv12;
596    use crate::backend::vaapi::encoder::tests::TestFrameGenerator;
597    use crate::backend::vaapi::surface_pool::PooledVaSurface;
598    use crate::backend::vaapi::surface_pool::VaSurfacePool;
599    use crate::codec::av1::parser::BitDepth;
600    use crate::codec::av1::parser::CdefParams;
601    use crate::codec::av1::parser::ColorConfig;
602    use crate::codec::av1::parser::FrameHeaderObu;
603    use crate::codec::av1::parser::FrameType;
604    use crate::codec::av1::parser::ObuHeader;
605    use crate::codec::av1::parser::ObuType;
606    use crate::codec::av1::parser::OperatingPoint;
607    use crate::codec::av1::parser::QuantizationParams;
608    use crate::codec::av1::parser::SequenceHeaderObu;
609    use crate::codec::av1::parser::TemporalDelimiterObu;
610    use crate::codec::av1::parser::TileInfo;
611    use crate::codec::av1::parser::TxMode;
612    use crate::codec::av1::parser::MAX_NUM_OPERATING_POINTS;
613    use crate::codec::av1::parser::PRIMARY_REF_NONE;
614    use crate::codec::av1::parser::SELECT_INTEGER_MV;
615    use crate::codec::av1::parser::SUPERRES_NUM;
616    use crate::codec::av1::synthesizer::Synthesizer;
617    use crate::decoder::FramePool;
618    use crate::encoder::simple_encode_loop;
619    use crate::encoder::stateless::BackendPromise;
620    use crate::encoder::stateless::StatelessEncoderBackendImport;
621    use crate::encoder::FrameMetadata;
622    use crate::encoder::RateControl;
623    use crate::encoder::Tunings;
624    use crate::utils::IvfFileHeader;
625    use crate::utils::IvfFrameHeader;
626    use crate::FrameLayout;
627    use crate::PlaneLayout;
628    use crate::Resolution;
629
630    #[test]
631    // Ignore this test by default as it requires libva-compatible hardware.
632    #[ignore]
633    fn test_single_frame() {
634        let _ = env_logger::try_init();
635
636        type Descriptor = ();
637        type Surface = libva::Surface<Descriptor>;
638        const WIDTH: u32 = 512;
639        const HEIGHT: u32 = 512;
640        let fourcc = b"NV12".into();
641
642        let frame_layout = FrameLayout {
643            format: (fourcc, 0),
644            size: Resolution {
645                width: WIDTH,
646                height: HEIGHT,
647            },
648            planes: vec![
649                PlaneLayout {
650                    buffer_index: 0,
651                    offset: 0,
652                    stride: WIDTH as usize,
653                },
654                PlaneLayout {
655                    buffer_index: 0,
656                    offset: (WIDTH * HEIGHT) as usize,
657                    stride: WIDTH as usize,
658                },
659            ],
660        };
661
662        let display = Display::open().unwrap();
663        let entrypoints = display
664            .query_config_entrypoints(VAProfileAV1Profile0)
665            .unwrap();
666        let low_power = entrypoints.contains(&VAEntrypointEncSliceLP);
667
668        let mut backend = VaapiBackend::<Descriptor, Surface>::new(
669            Rc::clone(&display),
670            VAProfileAV1Profile0,
671            fourcc,
672            Resolution {
673                width: WIDTH,
674                height: HEIGHT,
675            },
676            libva::constants::VA_RC_CQP,
677            low_power,
678        )
679        .unwrap();
680
681        let mut surfaces = display
682            .create_surfaces(
683                VA_RT_FORMAT_YUV420,
684                Some(frame_layout.format.0 .0),
685                WIDTH,
686                HEIGHT,
687                Some(UsageHint::USAGE_HINT_ENCODER),
688                vec![()],
689            )
690            .unwrap();
691
692        let surface = surfaces.pop().unwrap();
693
694        upload_test_frame_nv12(&display, &surface, 0.0);
695
696        let input_meta = FrameMetadata {
697            layout: frame_layout,
698            force_keyframe: false,
699            timestamp: 0,
700        };
701
702        let input = backend.import_picture(&input_meta, surface).unwrap();
703
704        let seq = SequenceHeaderObu {
705            obu_header: ObuHeader {
706                obu_type: ObuType::SequenceHeader,
707                extension_flag: false,
708                has_size_field: true,
709                temporal_id: 0,
710                spatial_id: 0,
711            },
712
713            seq_profile: Profile::Profile0,
714            num_planes: 3,
715
716            frame_width_bits_minus_1: 16 - 1,
717            frame_height_bits_minus_1: 16 - 1,
718            max_frame_width_minus_1: (WIDTH - 1) as u16,
719            max_frame_height_minus_1: (HEIGHT - 1) as u16,
720
721            enable_order_hint: true,
722            order_hint_bits: 8,
723            order_hint_bits_minus_1: 7,
724            seq_force_integer_mv: SELECT_INTEGER_MV as u32,
725
726            operating_points: {
727                let mut ops: [OperatingPoint; MAX_NUM_OPERATING_POINTS] = Default::default();
728                ops[0].seq_level_idx = 7;
729                ops
730            },
731
732            color_config: ColorConfig {
733                subsampling_x: true,
734                subsampling_y: true,
735                ..Default::default()
736            },
737
738            ..Default::default()
739        };
740
741        let frame = FrameHeaderObu {
742            obu_header: ObuHeader {
743                obu_type: ObuType::FrameHeader,
744                extension_flag: false,
745                has_size_field: true,
746                temporal_id: 0,
747                spatial_id: 0,
748            },
749
750            frame_type: FrameType::KeyFrame,
751            frame_is_intra: true,
752            primary_ref_frame: PRIMARY_REF_NONE,
753            refresh_frame_flags: 0xff,
754            error_resilient_mode: true,
755
756            reduced_tx_set: true,
757            tx_mode_select: 1,
758            tx_mode: TxMode::Select,
759
760            quantization_params: QuantizationParams {
761                base_q_idx: 128,
762                ..Default::default()
763            },
764            tile_info: TileInfo {
765                uniform_tile_spacing_flag: true,
766                tile_cols: 1,
767                tile_rows: 1,
768                tile_cols_log2: 0,
769                tile_rows_log2: 0,
770                width_in_sbs_minus_1: {
771                    let mut value = [0u32; MAX_TILE_COLS];
772                    value[0] = WIDTH / 64 - 1;
773                    value
774                },
775                height_in_sbs_minus_1: {
776                    let mut value = [0u32; MAX_TILE_ROWS];
777                    value[0] = HEIGHT / 64 - 1;
778                    value
779                },
780                ..Default::default()
781            },
782            cdef_params: CdefParams {
783                cdef_damping: 3,
784                ..Default::default()
785            },
786            superres_denom: SUPERRES_NUM as u32,
787            upscaled_width: WIDTH,
788            frame_width: WIDTH,
789            frame_height: HEIGHT,
790            render_width: WIDTH,
791            render_height: HEIGHT,
792
793            ..Default::default()
794        };
795
796        let request = Request {
797            sequence: seq.clone(),
798            frame: frame.clone(),
799            input,
800            input_meta,
801            references: [None, None, None, None, None, None, None],
802            ref_frame_ctrl_l0: [ReferenceFrameType::Intra; REFS_PER_FRAME],
803            ref_frame_ctrl_l1: [ReferenceFrameType::Intra; REFS_PER_FRAME],
804            intra_period: 4,
805            ip_period: 1,
806            tunings: Default::default(),
807            coded_output: Vec::new(),
808        };
809
810        let (_, output) = backend.encode_tile_group(request).unwrap();
811        let output = output.sync().unwrap();
812
813        let write_to_file = std::option_env!("CROS_CODECS_TEST_WRITE_TO_FILE") == Some("true");
814        if write_to_file {
815            use std::io::Write;
816
817            let mut out = std::fs::File::create("test_single_frame.av1.ivf").unwrap();
818
819            let td = TemporalDelimiterObu {
820                obu_header: ObuHeader {
821                    obu_type: ObuType::TemporalDelimiter,
822                    extension_flag: false,
823                    has_size_field: true,
824                    temporal_id: 0,
825                    spatial_id: 0,
826                },
827            };
828
829            let file_header = IvfFileHeader::new(
830                IvfFileHeader::CODEC_AV1,
831                WIDTH as u16,
832                HEIGHT as u16,
833                30,
834                10,
835            );
836
837            file_header.writo_into(&mut out).unwrap();
838
839            {
840                let mut hdr_buf = Vec::new();
841
842                Synthesizer::<'_, TemporalDelimiterObu, _>::synthesize(&td, &mut hdr_buf).unwrap();
843                Synthesizer::<'_, SequenceHeaderObu, _>::synthesize(&seq, &mut hdr_buf).unwrap();
844                Synthesizer::<'_, FrameHeaderObu, _>::synthesize(&frame, &seq, &mut hdr_buf)
845                    .unwrap();
846
847                let frame_header = IvfFrameHeader {
848                    frame_size: hdr_buf.len() as u32 + output.len() as u32,
849                    timestamp: 0,
850                };
851
852                frame_header.writo_into(&mut out).unwrap();
853                out.write_all(&hdr_buf).unwrap();
854                out.write_all(&output).unwrap();
855            }
856
857            out.flush().unwrap();
858        }
859    }
860
861    #[test]
862    // Ignore this test by default as it requires libva-compatible hardware.
863    #[ignore]
864    fn test_vaapi_encoder() {
865        type VaapiAv1Encoder<'l> =
866            StatelessEncoder<AV1, PooledVaSurface<()>, VaapiBackend<(), PooledVaSurface<()>>>;
867
868        const WIDTH: usize = 512;
869        const HEIGHT: usize = 512;
870        const FRAME_COUNT: u64 = 100;
871
872        let _ = env_logger::try_init();
873
874        let display = libva::Display::open().unwrap();
875        let entrypoints = display
876            .query_config_entrypoints(VAProfileAV1Profile0)
877            .unwrap();
878        let low_power = entrypoints.contains(&VAEntrypointEncSliceLP);
879
880        let config = EncoderConfig {
881            profile: Profile::Profile0,
882            resolution: Resolution {
883                width: WIDTH as u32,
884                height: HEIGHT as u32,
885            },
886            initial_tunings: Tunings {
887                rate_control: RateControl::ConstantQuality(128),
888                framerate: 30,
889                ..Default::default()
890            },
891            ..Default::default()
892        };
893
894        let frame_layout = FrameLayout {
895            format: (b"NV12".into(), 0),
896            size: Resolution {
897                width: WIDTH as u32,
898                height: HEIGHT as u32,
899            },
900            planes: vec![
901                PlaneLayout {
902                    buffer_index: 0,
903                    offset: 0,
904                    stride: WIDTH,
905                },
906                PlaneLayout {
907                    buffer_index: 0,
908                    offset: WIDTH * HEIGHT,
909                    stride: WIDTH,
910                },
911            ],
912        };
913
914        let mut encoder = VaapiAv1Encoder::new_vaapi(
915            Rc::clone(&display),
916            config,
917            frame_layout.format.0,
918            frame_layout.size,
919            low_power,
920            BlockingMode::Blocking,
921        )
922        .unwrap();
923
924        let mut pool = VaSurfacePool::new(
925            Rc::clone(&display),
926            VA_RT_FORMAT_YUV420,
927            Some(UsageHint::USAGE_HINT_ENCODER),
928            Resolution {
929                width: WIDTH as u32,
930                height: HEIGHT as u32,
931            },
932        );
933
934        pool.add_frames(vec![(); 16]).unwrap();
935
936        let mut frame_producer = TestFrameGenerator::new(FRAME_COUNT, display, pool, frame_layout);
937
938        let mut bitstream = Vec::new();
939
940        let file_header = IvfFileHeader::new(
941            IvfFileHeader::CODEC_AV1,
942            WIDTH as u16,
943            HEIGHT as u16,
944            30,
945            FRAME_COUNT as u32,
946        );
947
948        file_header.writo_into(&mut bitstream).unwrap();
949
950        simple_encode_loop(&mut encoder, &mut frame_producer, |coded| {
951            let header = IvfFrameHeader {
952                timestamp: coded.metadata.timestamp,
953                frame_size: coded.bitstream.len() as u32,
954            };
955
956            header.writo_into(&mut bitstream).unwrap();
957            bitstream.extend(coded.bitstream);
958        })
959        .unwrap();
960
961        let write_to_file = std::option_env!("CROS_CODECS_TEST_WRITE_TO_FILE") == Some("true");
962        if write_to_file {
963            use std::io::Write;
964            let mut out = std::fs::File::create("test_vaapi_encoder.av1.ivf").unwrap();
965            out.write_all(&bitstream).unwrap();
966            out.flush().unwrap();
967        }
968    }
969
970    #[ignore]
971    // Ignore this test by default as it requires libva-compatible hardware.
972    #[test]
973    fn test_vaapi_encoder_p010() {
974        type VaapiAv1Encoder<'l> =
975            StatelessEncoder<AV1, PooledVaSurface<()>, VaapiBackend<(), PooledVaSurface<()>>>;
976
977        const WIDTH: usize = 512;
978        const HEIGHT: usize = 512;
979        const FRAME_COUNT: u64 = 100;
980
981        let _ = env_logger::try_init();
982
983        let display = libva::Display::open().unwrap();
984        let entrypoints = display
985            .query_config_entrypoints(VAProfileAV1Profile0)
986            .unwrap();
987        let low_power = entrypoints.contains(&VAEntrypointEncSliceLP);
988
989        let config = EncoderConfig {
990            profile: Profile::Profile0,
991            resolution: Resolution {
992                width: WIDTH as u32,
993                height: HEIGHT as u32,
994            },
995            bit_depth: BitDepth::Depth10,
996            initial_tunings: Tunings {
997                rate_control: RateControl::ConstantQuality(128),
998                framerate: 30,
999                ..Default::default()
1000            },
1001            ..Default::default()
1002        };
1003
1004        let frame_layout = FrameLayout {
1005            format: (b"P010".into(), 0),
1006            size: Resolution {
1007                width: WIDTH as u32,
1008                height: HEIGHT as u32,
1009            },
1010            planes: vec![
1011                PlaneLayout {
1012                    buffer_index: 0,
1013                    offset: 0,
1014                    stride: WIDTH,
1015                },
1016                PlaneLayout {
1017                    buffer_index: 0,
1018                    offset: WIDTH * HEIGHT,
1019                    stride: WIDTH,
1020                },
1021            ],
1022        };
1023
1024        let mut encoder = VaapiAv1Encoder::new_vaapi(
1025            Rc::clone(&display),
1026            config,
1027            frame_layout.format.0,
1028            frame_layout.size,
1029            low_power,
1030            BlockingMode::Blocking,
1031        )
1032        .unwrap();
1033
1034        let mut pool = VaSurfacePool::new(
1035            Rc::clone(&display),
1036            VA_RT_FORMAT_YUV420_10,
1037            Some(UsageHint::USAGE_HINT_ENCODER),
1038            Resolution {
1039                width: WIDTH as u32,
1040                height: HEIGHT as u32,
1041            },
1042        );
1043
1044        pool.add_frames(vec![(); 16]).unwrap();
1045
1046        let mut frame_producer = TestFrameGenerator::new(FRAME_COUNT, display, pool, frame_layout);
1047
1048        let mut bitstream = Vec::new();
1049
1050        let file_header = IvfFileHeader::new(
1051            IvfFileHeader::CODEC_AV1,
1052            WIDTH as u16,
1053            HEIGHT as u16,
1054            30,
1055            FRAME_COUNT as u32,
1056        );
1057
1058        file_header.writo_into(&mut bitstream).unwrap();
1059
1060        simple_encode_loop(&mut encoder, &mut frame_producer, |coded| {
1061            let header = IvfFrameHeader {
1062                timestamp: coded.metadata.timestamp,
1063                frame_size: coded.bitstream.len() as u32,
1064            };
1065
1066            header.writo_into(&mut bitstream).unwrap();
1067            bitstream.extend(coded.bitstream);
1068        })
1069        .unwrap();
1070
1071        let write_to_file = std::option_env!("CROS_CODECS_TEST_WRITE_TO_FILE") == Some("true");
1072        if write_to_file {
1073            use std::io::Write;
1074            let mut out = std::fs::File::create("test_vaapi_encoder_p010.av1.ivf").unwrap();
1075            out.write_all(&bitstream).unwrap();
1076            out.flush().unwrap();
1077        }
1078    }
1079}