diff --git a/Cargo.toml b/Cargo.toml index 950508a..9ca429a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "moq-vaapi" version = "0.1.1" -description = "(AI GENERATED) VA-API H.264 hardware encoder and decoder, derived from discord/cros-libva + discord/cros-codecs" +description = "(AI GENERATED) VA-API H.264/H.265 hardware encoder and H.264 decoder, derived from discord/cros-libva + discord/cros-codecs" authors = ["The ChromiumOS Authors", "Discord", "Luke Curley "] repository = "https://github.com/moq-dev/vaapi" license = "BSD-3-Clause" diff --git a/README.md b/README.md index 3699868..70696cb 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ > it. The encoder and the decoder have both been checked against ffmpeg's > software codecs on Intel hardware (see below). -A small, self-contained **VA-API H.264 hardware codec** for Linux (Intel / AMD), +A small, self-contained **VA-API H.264 / H.265 hardware codec** for Linux (Intel / AMD), derived from [discord/cros-libva @ discord-0.0.13](https://github.com/discord/cros-libva/tree/discord-0.0.13) and @@ -20,9 +20,13 @@ pulled in as a git dependency. ## What it does -- **H.264 encode** over VA-API: an Annex-B elementary stream out (packed - SPS/PPS + slice headers, low-latency IPPP, CBR rate control whose bitrate can - change between frames, the color space labelled in the VUI). Three inputs: +- **H.264 and H.265 encode** over VA-API: an Annex-B elementary stream out + (packed parameter sets, low-latency IPPP, CBR rate control whose bitrate can + change between frames, the color space labelled in the VUI). H.265 is Main + profile with its VPS, SPS, PPS, and slice headers written here, its coding + tools following what the driver reports, and its coded size following the + driver's surface alignment, so the crop holds on radeonsi, which rewrites the + SPS at that alignment. Three inputs: tightly-packed NV12 from the CPU, an NV12 DMA-BUF at the encoder's size encoded in place, and any other importable DMA-BUF (packed RGB, YUYV, NV12 at another size or in another color space) converted and scaled into the @@ -30,8 +34,13 @@ pulled in as a git dependency. formats go through the same video processor call untested. Checked on Intel Meteor Lake (iHD 26.1.5) by decoding with ffmpeg: luma above 35 dB PSNR against the source, and RGB converted to BT.601 and BT.709 - to within a few code values. These run as tests wherever ffmpeg and a VA-API - device are both present. + to within a few code values. These run as tests, for both codecs, wherever + ffmpeg and a VA-API device are both present. H.265 passes them on Intel + Arrow Lake (iHD 26.1.2). On AMD Radeon RX 9070 (radeonsi, Mesa 26.0.8) it + passes all but the color conversion and bitrate checks: chroma converted by + the video processor lands about 7 code values off, for either codec, and + radeonsi's rate control runs about 4x over 250 kbps on per-frame noise. H.264 encode on that card does not yet + produce a decodable stream. - **DMA-BUF import and video post-processing.** `dmabuf` wraps a DRM PRIME buffer as a surface without copying it; `vpp` blits one surface into another through `VAEntrypointVideoProc`, scaling, converting, and re-tiling on the @@ -87,7 +96,8 @@ advisories checked, not a guarantee that no undisclosed vulnerabilities exist. ## Layout - `src/` — libva bindings (`bindings`, `display`, `surface`, `buffer`, ...), the - H.264 bitstream layer (`bitstream_utils`, `codec::h264`), the thin encode + H.264 bitstream layer (`bitstream_utils`, `codec::h264`), the H.265 header + writer (`encode/h265.rs`), the thin encode and decode drivers (`encode`, `decode`), DMA-BUF import (`dmabuf`), video post-processing (`vpp`), and the color space both share (`color`). - `libva/` — vendored libva headers (checked in; refreshed by `just vendor`). diff --git a/src/codec.rs b/src/codec.rs index 6d1d83c..896f40b 100644 --- a/src/codec.rs +++ b/src/codec.rs @@ -1,2 +1,11 @@ // Vendored from discord/cros-codecs (BSD-3-Clause), trimmed to H.264. pub mod h264; + +/// A video coding format. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum Codec { + /// H.264 (AVC). + H264, + /// H.265 (HEVC). + H265, +} diff --git a/src/encode.rs b/src/encode.rs index d156644..fbb3be2 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -5,13 +5,15 @@ // Adapted from discord/cros-codecs (encoder/stateless/h264/vaapi.rs), itself // BSD-3-Clause / Copyright The ChromiumOS Authors. See LICENSE.cros-codecs. -//! Thin VA-API H.264 encoder built on the vendored libva binding. +//! Thin VA-API H.264 and H.265 encoder built on the vendored libva binding. //! -//! Reuses the backend-agnostic bitstream layer ([`crate::codec::h264`]) from -//! discord/cros-codecs (BSD-3-Clause) for SPS/PPS/slice synthesis, and drives -//! libva directly for surface upload + slice submission, rather than vendoring -//! cros-codecs's generic multi-backend encoder framework. The per-frame VA -//! buffer population is ported from cros-codecs's `encoder/stateless/h264/vaapi.rs`. +//! For H.264 it reuses the backend-agnostic bitstream layer +//! ([`crate::codec::h264`]) from discord/cros-codecs (BSD-3-Clause) for +//! SPS/PPS/slice synthesis, and drives libva directly for surface upload + +//! slice submission, rather than vendoring cros-codecs's generic multi-backend +//! encoder framework. The per-frame VA buffer population is ported from +//! cros-codecs's `encoder/stateless/h264/vaapi.rs`. H.265 writes its own +//! parameter sets and slice headers (see `h265.rs`). //! //! Low-latency only: IPPP (no B-frames), one reference frame, matching the //! VideoToolbox / Media Foundation / NVENC backends in moq-video. @@ -38,6 +40,7 @@ use std::sync::Arc; use crate::codec::h264::parser::SpsBuilder; use crate::codec::h264::parser::{Level, Pps, PpsBuilder, Profile, SliceHeader, SliceHeaderBuilder, SliceType, Sps}; use crate::codec::h264::synthesizer::Synthesizer; +use crate::codec::Codec; use crate::dmabuf::DmaBuf; use crate::vpp::Processor; use crate::{ @@ -47,11 +50,13 @@ use crate::{ EncSliceParameter, EncSliceParameterBufferH264, H264EncFrameCropOffsets, H264EncPicFields, H264EncSeqFields, H264VuiFields, Image, MappedCodedBuffer, Picture, PictureH264, RcFlags, Surface, SurfaceMemoryDescriptor, UsageHint, VAConfigAttrib, VAConfigAttribType, VAEntrypoint, VAProfile, VA_ATTRIB_NOT_SUPPORTED, - VA_ENC_PACKED_HEADER_PICTURE, VA_ENC_PACKED_HEADER_SEQUENCE, VA_FOURCC_ARGB, VA_FOURCC_BGRA, VA_FOURCC_BGRX, - VA_FOURCC_NV12, VA_FOURCC_RGBA, VA_FOURCC_RGBX, VA_FOURCC_XRGB, VA_INVALID_ID, VA_PICTURE_H264_INVALID, - VA_PICTURE_H264_SHORT_TERM_REFERENCE, VA_RC_CBR, VA_RT_FORMAT_YUV420, + VA_ENC_PACKED_HEADER_PICTURE, VA_ENC_PACKED_HEADER_SEQUENCE, VA_ENC_PACKED_HEADER_SLICE, VA_FOURCC_ARGB, + VA_FOURCC_BGRA, VA_FOURCC_BGRX, VA_FOURCC_NV12, VA_FOURCC_RGBA, VA_FOURCC_RGBX, VA_FOURCC_XRGB, VA_INVALID_ID, + VA_PICTURE_H264_INVALID, VA_PICTURE_H264_SHORT_TERM_REFERENCE, VA_RC_CBR, VA_RT_FORMAT_YUV420, }; +mod h265; + /// Whether a frame is used as a reference, and for how long. /// (Vendored from discord/cros-codecs, BSD-3-Clause; used by the slice synthesizer.) #[derive(Copy, Clone, PartialEq, Eq, Debug)] @@ -77,9 +82,11 @@ pub const RATE_CONTROL_WINDOW_MS: u32 = 1_500; /// `max_frame_num` upper bound for the low-delay GOP; matches cros-codecs. const LIMIT: u32 = 2048; -/// H.264 encoder configuration. Dimensions must be even (4:2:0 chroma). +/// Encoder configuration. Dimensions must be even (4:2:0 chroma). #[derive(Clone, Debug)] pub struct Config { + /// The format to encode, H.264 unless set. Main profile either way. + pub codec: Codec, pub width: u32, pub height: u32, pub framerate: u32, @@ -89,8 +96,9 @@ pub struct Config { pub gop: u32, /// DRM render node to open (e.g. `/dev/dri/renderD128`). /// - /// `None`, the default, opens the first node whose driver can encode H.264, - /// and [`Encoder::config`] then names the node it picked. + /// `None`, the default, opens the first node whose driver can encode + /// [`codec`](Self::codec), and [`Encoder::config`] then names the node it + /// picked. pub device: Option, /// Color space written into the SPS, and the target of a conversion from RGB. /// @@ -100,8 +108,10 @@ pub struct Config { } impl Config { + /// Returns an H.264 configuration; set [`codec`](Self::codec) for another format. pub fn new(width: u32, height: u32, framerate: u32, bitrate: u32, gop: u32) -> Self { Self { + codec: Codec::H264, width, height, framerate, @@ -120,20 +130,65 @@ struct FrameMeta { frame_num: u32, } -/// A VA-API H.264 encoder. Built once, fed NV12 frames, emits Annex-B H.264. +impl FrameMeta { + /// Returns the H.264 order count and frame number of the picture `count` pictures after an IDR. + /// + /// frame_num counts modulo MaxFrameNum, and the slice header carries only + /// the low bits of the picture order count; the full count goes to the + /// driver in the picture parameters. + fn new(count: u32) -> Self { + Self { + poc: count.wrapping_mul(2), + frame_num: count % LIMIT, + } + } +} + +/// The picture a P picture predicts from: the previous one. +#[derive(Copy, Clone)] +struct Reference { + /// Its reconstructed surface. + surface: crate::bindings::VASurfaceID, + /// How many pictures after the IDR it came. + count: u32, +} + +/// What one picture's parameter buffers are built from. +struct PictureInfo<'a> { + idr: bool, + /// How many pictures after the IDR this one comes. + count: u32, + /// The surface its reconstruction is written to. + recon: crate::bindings::VASurfaceID, + reference: Option, + coded: &'a EncCodedBuffer, + bitrate: u32, +} + +/// The bitstream syntax of the stream being encoded. +enum Syntax { + H264 { + sps: Rc, + pps: Rc, + }, + H265 { + stream: h265::Stream, + /// The packed header kinds the driver takes from us; it writes the rest. + packed: u32, + }, +} + +/// A VA-API H.264 or H.265 encoder. Built once, fed NV12 frames, emits an Annex-B stream. pub struct Encoder { config: Config, - sps: Rc, - pps: Rc, - width_mbs: u16, - height_mbs: u16, + syntax: Syntax, // Source surface (reused; NV12 re-uploaded per frame) and a 2-deep // reconstructed surface ring (current + previous reference for IPPP). input: Option>, recon: Vec>, - /// Reference from the previous encoded frame: (recon index, metadata). - reference: Option<(usize, FrameMeta)>, + /// Reference from the previous encoded frame: (recon index, pictures after the IDR). + reference: Option<(usize, u32)>, counter: u32, @@ -153,26 +208,28 @@ pub struct Encoder { } impl Encoder { - /// Opens [`Config::device`], or the first render node that encodes H.264, and sets up an encode session for `config`. + /// Opens [`Config::device`], or the first render node that encodes [`Config::codec`], and sets up an encode session for `config`. /// /// # Errors /// /// Fails when libva cannot be loaded, when the named node cannot be opened - /// or does not encode H.264 (a named node is never swapped for another), when - /// no node encodes H.264, and when the driver refuses the session. + /// or does not encode the codec (a named node is never swapped for + /// another), when no node encodes it, and when the driver refuses the + /// session or, for H.265, requires a coding tool this encoder does not write + /// or takes no packed sequence header. pub fn new(mut config: Config) -> anyhow::Result { + let codec = config.codec; let (device, display) = match config.device.take() { Some(device) => { let display = Display::open_drm_display(&device) .map_err(|e| anyhow::anyhow!("open DRM display {device:?}: {e:?}"))?; (device, display) } - None => { - crate::display::open_first(probe).map_err(|e| e.context("find a render node that encodes H.264"))? - } + None => crate::display::open_first(|display| probe(display, codec)) + .map_err(|e| e.context(format!("find a render node that encodes {codec:?}")))?, }; - let entrypoint = encode_entrypoint(&display)?; + let entrypoint = encode_entrypoint(&display, codec)?; let mut attrs = vec![ VAConfigAttrib { type_: VAConfigAttribType::VAConfigAttribRTFormat, @@ -183,24 +240,54 @@ impl Encoder { value: VA_RC_CBR, }, ]; - // The SPS and PPS go out as packed headers, so the VUI (color, frame - // buffering) and the level are ours. iHD takes them without being told; - // a driver that follows the spec writes its own unless the config asks. - let packed = packed_headers(&display, entrypoint); + // The parameter sets go out as packed headers, so the VUI (color, frame + // buffering) and the level are ours. For H.264 that is the SPS and PPS; + // iHD takes them without being told, and a driver that follows the spec + // writes its own unless the config asks. For H.265 it is the VPS, SPS and + // PPS as one sequence header, plus each slice header, as ffmpeg packs + // them; the driver writes the slice headers if it does not take them. + let wanted = match codec { + Codec::H264 => VA_ENC_PACKED_HEADER_SEQUENCE | VA_ENC_PACKED_HEADER_PICTURE, + Codec::H265 => VA_ENC_PACKED_HEADER_SEQUENCE | VA_ENC_PACKED_HEADER_SLICE, + }; + let packed = packed_headers(&display, codec, entrypoint) & wanted; + // VA's H.265 sequence parameters have no conformance window and carry no + // VUI here, so an SPS the driver writes would lose the crop and the color. + if codec == Codec::H265 && packed & VA_ENC_PACKED_HEADER_SEQUENCE == 0 { + anyhow::bail!("the driver takes no packed H.265 sequence header, which the crop and color need"); + } if packed != 0 { attrs.push(VAConfigAttrib { type_: VAConfigAttribType::VAConfigAttribEncPackedHeaders, value: packed, }); } else { - log::warn!("the driver takes no packed SPS/PPS; the stream carries the driver's own headers"); + log::warn!("the driver takes no packed {codec:?} headers; the stream carries the driver's own"); } - let va_config = display - .create_config(attrs, VAProfile::VAProfileH264Main, entrypoint) + let mut va_config = display + .create_config(attrs, profile(codec), entrypoint) .map_err(|e| anyhow::anyhow!("create VA config: {e:?}"))?; + let syntax = match codec { + Codec::H264 => { + let (sps, pps) = build_sps_pps(&config); + Syntax::H264 { sps, pps } + } + Codec::H265 => Syntax::H265 { + stream: h265::Stream::new(&display, entrypoint, &mut va_config, &config)?, + packed, + }, + }; + // H.264 sizes the context and its references to the frame, which has + // always worked; H.265 to whole minimum coding blocks, as ffmpeg does, + // and to the driver's surface alignment. + let size = match &syntax { + Syntax::H264 { .. } => (config.width, config.height), + Syntax::H265 { stream, .. } => stream.coded(), + }; + let context = display - .create_context::<()>(&va_config, config.width, config.height, None, true) + .create_context::<()>(&va_config, size.0, size.1, None, true) .map_err(|e| anyhow::anyhow!("create VA context: {e:?}"))?; let coded_size = coded_size(config.bitrate); @@ -208,15 +295,11 @@ impl Encoder { .create_enc_coded(coded_size) .map_err(|e| anyhow::anyhow!("create coded buffer: {e:?}"))?; - let input = make_surface(&display, &config)?; - let recon = vec![make_surface(&display, &config)?, make_surface(&display, &config)?]; - - let (sps, pps) = build_sps_pps(&config); - let width_mbs = sps.pic_width_in_mbs_minus1 + 1; - let height_mbs = sps.pic_height_in_map_units_minus1 + 1; + let input = make_surface(&display, (config.width, config.height))?; + let recon = vec![make_surface(&display, size)?, make_surface(&display, size)?]; log::info!( - "opened VA-API H.264 encoder on {:?}: {}x{} @ {}fps, {} bps, low power: {}", + "opened VA-API {codec:?} encoder on {:?}: {}x{} @ {}fps, {} bps, low power: {}", device, config.width, config.height, @@ -227,10 +310,7 @@ impl Encoder { config.device = Some(device); Ok(Self { config, - sps, - pps, - width_mbs, - height_mbs, + syntax, input: Some(input), recon, reference: None, @@ -245,8 +325,8 @@ impl Encoder { }) } - /// Encodes one tightly-packed NV12 frame, returning Annex-B H.264 (with inline - /// SPS/PPS on each IDR). `keyframe` forces an IDR. + /// Encodes one tightly-packed NV12 frame, returning Annex-B (with the + /// parameter sets inline on each IDR). `keyframe` forces an IDR. /// /// # Errors /// @@ -265,13 +345,14 @@ impl Encoder { result } - /// Encodes one frame held in a DMA-BUF, returning Annex-B H.264 like + /// Encodes one frame held in a DMA-BUF, returning Annex-B like /// [`encode_nv12`](Self::encode_nv12). /// - /// An NV12 buffer at the encoder's size, whose size is a whole number of - /// macroblocks and whose [`color`](DmaBuf::color) is the stream's or - /// unknown, is imported and encoded in place. Any other mapped format or - /// size goes through the video processor into the encoder's own surface + /// An NV12 buffer at the encoder's size, which the encoder reads without + /// padding (whole macroblocks for H.264; for H.265, whole minimum coding + /// blocks at the driver's surface alignment), and whose + /// [`color`](DmaBuf::color) is the stream's or unknown, is imported and + /// encoded in place. Any other mapped format or size goes through the video processor into the encoder's own surface /// first: scaled to the encoder's size and converted to [`Config::color`]. /// A YUV buffer with no color of its own is taken to be in that space /// already. Nothing reaches the CPU on either path. @@ -296,10 +377,14 @@ impl Encoder { .fourcc() .ok_or_else(|| anyhow::anyhow!("no VA-API format for DRM format {:#010x}", buffer.drm_format))?; let color = self.config.color; + let size = (self.config.width, self.config.height); + let unpadded = match &self.syntax { + Syntax::H264 { .. } => size.0 % 16 == 0 && size.1 % 16 == 0, + Syntax::H265 { stream, .. } => stream.coded() == size, + }; let in_place = fourcc == VA_FOURCC_NV12 - && (buffer.width, buffer.height) == (self.config.width, self.config.height) - && buffer.width % 16 == 0 - && buffer.height % 16 == 0 + && (buffer.width, buffer.height) == size + && unpadded && buffer.color.is_none_or(|own| own == color); if in_place { let source = buffer.import(&self.display, Some(UsageHint::USAGE_HINT_ENCODER))?; @@ -376,65 +461,48 @@ impl Encoder { ) -> anyhow::Result> { let is_idr = keyframe || self.counter % self.config.gop == 0; if is_idr { - // IDR resets the H.264 frame numbering and the reference list. + // An IDR restarts the picture count and drops the reference. self.counter = 0; self.reference = None; } - // frame_num counts modulo MaxFrameNum, and the slice header carries only - // the low bits of the picture order count; the full count goes to the - // driver in the picture parameters. - let meta = FrameMeta { - poc: self.counter.wrapping_mul(2), - frame_num: self.counter % LIMIT, - }; let recon_idx = (self.counter as usize) % self.recon.len(); - - let slice_type = if is_idr { SliceType::I } else { SliceType::P }; - let header = SliceHeaderBuilder::new(&self.pps) - .slice_type(slice_type) - .first_mb_in_slice(0) - .frame_num(meta.frame_num as u16) - .pic_order_cnt_lsb((meta.poc % (LIMIT * 2)) as u16) - .build(); - - let num_macroblocks = self.width_mbs as u32 * self.height_mbs as u32; - let bits_per_second = self.config.bitrate; - - // Reference (previous reconstructed frame) for P slices: surface id + meta. + // The reference is the previous picture's reconstruction, for a P picture. let reference = match (is_idr, self.reference) { - (false, Some((ref_idx, ref_meta))) => Some((self.recon[ref_idx].id(), ref_meta)), + (false, Some((index, count))) => Some(Reference { + surface: self.recon[index].id(), + count, + }), _ => None, }; - - let seq_param = build_enc_seq_param(&self.sps, bits_per_second, LIMIT, 0); - let pic_param = build_enc_pic_param( - &self.pps, - &self.coded, - self.recon[recon_idx].id(), - meta, - is_idr, + let info = PictureInfo { + idr: is_idr, + count: self.counter, + recon: self.recon[recon_idx].id(), reference, - ); - let slice_param = build_enc_slice_param(&self.pps, &header, reference, num_macroblocks); + coded: &self.coded, + bitrate: self.config.bitrate, + }; + let buffers = match &self.syntax { + Syntax::H264 { sps, pps } => { + let mut buffers = h264_buffers(sps, pps, &info); + if is_idr { + buffers.extend(packed_parameter_sets(sps, pps)?); + } + buffers + } + Syntax::H265 { stream, packed } => stream.buffers(&self.config, &info, *packed), + }; - let mut picture = Picture::new(meta.frame_num as u64, Rc::clone(&self.context), source); + let mut picture = Picture::new(self.counter as u64, Rc::clone(&self.context), source); // VA-API spec buffer order: sequence, picture, slice, then packed headers, // then rate-control misc params. - picture.add_buffer(self.create(seq_param)?); - picture.add_buffer(self.create(pic_param)?); - picture.add_buffer(self.create(slice_param)?); - - if is_idr { - let (sps_param, sps_data) = packed_header(EncPackedHeaderType::Sequence, &self.packed_sps()?); - let (pps_param, pps_data) = packed_header(EncPackedHeaderType::Picture, &self.packed_pps()?); - picture.add_buffer(self.create(sps_param)?); - picture.add_buffer(self.create(sps_data)?); - picture.add_buffer(self.create(pps_param)?); - picture.add_buffer(self.create(pps_data)?); + for buffer in buffers { + picture.add_buffer(self.create(buffer)?); } + let bits_per_second = self.config.bitrate; let rc = EncMiscParameterRateControl::new( bits_per_second, 100, // target_percentage (CBR) @@ -467,7 +535,7 @@ impl Encoder { // The reconstructed surface syncs implicitly; read the coded bitstream. let bitstream = self.read_coded()?; - self.reference = Some((recon_idx, meta)); + self.reference = Some((recon_idx, self.counter)); self.counter += 1; self.reset_rate_control = false; Ok(bitstream) @@ -479,20 +547,6 @@ impl Encoder { .map_err(|e| anyhow::anyhow!("create VA buffer: {e:?}")) } - fn packed_sps(&self) -> anyhow::Result> { - let mut buf = Vec::new(); - Synthesizer::<'_, Sps, _>::synthesize(3, &self.sps, &mut buf, true) - .map_err(|e| anyhow::anyhow!("synthesize SPS: {e:?}"))?; - Ok(buf) - } - - fn packed_pps(&self) -> anyhow::Result> { - let mut buf = Vec::new(); - Synthesizer::<'_, Pps, _>::synthesize(3, &self.pps, &mut buf, true) - .map_err(|e| anyhow::anyhow!("synthesize PPS: {e:?}"))?; - Ok(buf) - } - fn read_coded(&self) -> anyhow::Result> { let mapped = MappedCodedBuffer::new(&self.coded).map_err(|e| anyhow::anyhow!("map coded buffer: {e:?}"))?; let mut out = Vec::new(); @@ -503,24 +557,32 @@ impl Encoder { } } -/// Checks that `display`'s driver can encode H.264, which is what [`Encoder::new`] needs of a render node. +/// Checks that `display`'s driver can encode `codec`, which is what [`Encoder::new`] needs of a render node. /// /// # Errors /// -/// Fails when the driver offers neither the full nor the low-power H.264 Main -/// encode entrypoint, or cannot be asked. -pub fn probe(display: &Display) -> anyhow::Result<()> { - encode_entrypoint(display).map(drop) +/// Fails when the driver offers neither the full nor the low-power Main +/// profile encode entrypoint for `codec`, or cannot be asked. +pub fn probe(display: &Display, codec: Codec) -> anyhow::Result<()> { + encode_entrypoint(display, codec).map(drop) +} + +/// Returns the VA profile `codec` is encoded in. +fn profile(codec: Codec) -> VAProfile::Type { + match codec { + Codec::H264 => VAProfile::VAProfileH264Main, + Codec::H265 => h265::PROFILE, + } } -/// Returns the entrypoint to encode H.264 Main with: the full one, or the low-power one where it is all there is. +/// Returns the entrypoint to encode `codec` with: the full one, or the low-power one where it is all there is. /// -/// Some Intel parts expose H.264 encode only through `VAEntrypointEncSliceLP`. -fn encode_entrypoint(display: &Display) -> anyhow::Result { +/// Some Intel parts expose encode only through `VAEntrypointEncSliceLP`. +fn encode_entrypoint(display: &Display, codec: Codec) -> anyhow::Result { let entrypoints = display - .query_config_entrypoints(VAProfile::VAProfileH264Main) - .map_err(|e| anyhow::anyhow!("query H.264 Main entrypoints: {e:?}"))?; - pick_entrypoint(&entrypoints).ok_or_else(|| anyhow::anyhow!("the device has no H.264 Main encode entrypoint")) + .query_config_entrypoints(profile(codec)) + .map_err(|e| anyhow::anyhow!("query {codec:?} Main entrypoints: {e:?}"))?; + pick_entrypoint(&entrypoints).ok_or_else(|| anyhow::anyhow!("the device has no {codec:?} Main encode entrypoint")) } /// Returns the encode entrypoint to use out of those a device offers: the full one when present. @@ -530,22 +592,22 @@ fn pick_entrypoint(offered: &[VAEntrypoint::Type]) -> Option .find(|entrypoint| offered.contains(entrypoint)) } -/// Returns the packed header kinds to ask for: sequence and picture, where the driver takes them. +/// Returns the packed header kinds the driver takes for `codec` at `entrypoint`. /// -/// Zero when the driver takes neither, or when asking fails. -fn packed_headers(display: &Display, entrypoint: VAEntrypoint::Type) -> u32 { +/// Zero when it takes none, or when asking fails. +fn packed_headers(display: &Display, codec: Codec, entrypoint: VAEntrypoint::Type) -> u32 { let mut attrs = [VAConfigAttrib { type_: VAConfigAttribType::VAConfigAttribEncPackedHeaders, value: 0, }]; if display - .get_config_attributes(VAProfile::VAProfileH264Main, entrypoint, &mut attrs) + .get_config_attributes(profile(codec), entrypoint, &mut attrs) .is_err() || attrs[0].value & VA_ATTRIB_NOT_SUPPORTED != 0 { return 0; } - attrs[0].value & (VA_ENC_PACKED_HEADER_SEQUENCE | VA_ENC_PACKED_HEADER_PICTURE) + attrs[0].value } /// Returns the lowest H.264 level, from 4 up, whose frame size and macroblock rate cover the stream. @@ -591,13 +653,13 @@ fn is_rgb(fourcc: u32) -> bool { } /// Allocate one NV12 encode surface. -fn make_surface(display: &Arc, config: &Config) -> anyhow::Result> { +fn make_surface(display: &Arc, (width, height): (u32, u32)) -> anyhow::Result> { let mut surfaces = display .create_surfaces::<()>( VA_RT_FORMAT_YUV420, Some(VA_FOURCC_NV12), - config.width, - config.height, + width, + height, Some(UsageHint::USAGE_HINT_ENCODER), vec![()], ) @@ -689,6 +751,42 @@ fn build_sps_pps(config: &Config) -> (Rc, Rc) { (sps, pps) } +/// Returns the sequence, picture and slice parameters of one H.264 picture. +fn h264_buffers(sps: &Sps, pps: &Rc, picture: &PictureInfo<'_>) -> Vec { + let meta = FrameMeta::new(picture.count); + let slice_type = if picture.idr { SliceType::I } else { SliceType::P }; + let header = SliceHeaderBuilder::new(pps) + .slice_type(slice_type) + .first_mb_in_slice(0) + .frame_num(meta.frame_num as u16) + .pic_order_cnt_lsb((meta.poc % (LIMIT * 2)) as u16) + .build(); + let macroblocks = (sps.pic_width_in_mbs_minus1 as u32 + 1) * (sps.pic_height_in_map_units_minus1 as u32 + 1); + let reference = picture + .reference + .map(|reference| (reference.surface, FrameMeta::new(reference.count))); + + vec![ + build_enc_seq_param(sps, picture.bitrate, LIMIT, 0), + build_enc_pic_param(pps, picture.coded, picture.recon, meta, picture.idr, reference), + build_enc_slice_param(pps, &header, reference, macroblocks), + ] +} + +/// Returns the SPS and PPS as packed headers, which go out with every IDR. +fn packed_parameter_sets(sps: &Sps, pps: &Pps) -> anyhow::Result<[BufferType; 4]> { + let mut sps_nal = Vec::new(); + Synthesizer::<'_, Sps, _>::synthesize(3, sps, &mut sps_nal, true) + .map_err(|e| anyhow::anyhow!("synthesize SPS: {e:?}"))?; + let mut pps_nal = Vec::new(); + Synthesizer::<'_, Pps, _>::synthesize(3, pps, &mut pps_nal, true) + .map_err(|e| anyhow::anyhow!("synthesize PPS: {e:?}"))?; + + let (sps_param, sps_data) = packed_header(EncPackedHeaderType::Sequence, &sps_nal); + let (pps_param, pps_data) = packed_header(EncPackedHeaderType::Picture, &pps_nal); + Ok([sps_param, sps_data, pps_param, pps_data]) +} + fn build_invalid_pic() -> PictureH264 { PictureH264::new(VA_INVALID_ID, 0, VA_PICTURE_H264_INVALID, 0, 0) } @@ -898,24 +996,38 @@ mod tests { use crate::dmabuf::DmaBuf; use crate::{VA_RT_FORMAT_RGB32, VA_RT_FORMAT_YUV420}; - const WIDTH: u32 = 320; - const HEIGHT: u32 = 240; + // A size every driver encodes: radeonsi's H.265 minimum on VCN 5 is 384 wide. + const WIDTH: u32 = 640; + const HEIGHT: u32 = 480; + const CODECS: [Codec; 2] = [Codec::H264, Codec::H265]; - fn config() -> Config { - Config::new(WIDTH, HEIGHT, 30, 2_000_000, 30) + fn config(codec: Codec) -> Config { + Config { + codec, + ..Config::new(WIDTH, HEIGHT, 30, 2_000_000, 30) + } } - /// Opens an encoder, or `None` to skip on a machine without VA-API H.264 encode. + /// Opens an encoder, or `None` to skip on a machine without VA-API encode of its codec. fn encoder(config: Config) -> Option { + let codec = config.codec; match Encoder::new(config) { Ok(encoder) => Some(encoder), Err(err) => { - eprintln!("skipping: no VA-API H.264 encoder: {err:#}"); + eprintln!("skipping: no VA-API {codec:?} encoder: {err:#}"); None } } } + /// The NAL unit types that open an IDR access unit: the IDR slice, then the parameter sets it carries. + fn idr_types(codec: Codec) -> (u8, &'static [u8]) { + match codec { + Codec::H264 => (5, &[7, 8]), + Codec::H265 => (19, &[32, 33, 34]), + } + } + /// A diagonal gradient with a block moving across it, so a stale or /// misplaced reference shows up as the block in the wrong place. fn nv12_frame(width: u32, height: u32, step: u32) -> Vec { @@ -939,18 +1051,24 @@ mod tests { /// /// The tests run as threads of one process, so the process id alone would /// let two of them share a file. - fn temp_path(label: &str) -> PathBuf { + /// + /// Its extension names `codec`, for ffmpeg to pick the demuxer by. + fn temp_path(codec: Codec, label: &str) -> PathBuf { static NEXT: AtomicU32 = AtomicU32::new(0); let index = NEXT.fetch_add(1, Ordering::Relaxed); - std::env::temp_dir().join(format!("moq-vaapi-{label}-{}-{index}.264", std::process::id())) + let extension = match codec { + Codec::H264 => "264", + Codec::H265 => "265", + }; + std::env::temp_dir().join(format!("moq-vaapi-{label}-{}-{index}.{extension}", std::process::id())) } /// Decodes an Annex-B stream with ffmpeg's software decoder into NV12 /// pictures, or `None` to skip without ffmpeg. - fn ffmpeg_decode(stream: &[u8], name: &str) -> Option> { + fn ffmpeg_decode(codec: Codec, stream: &[u8], name: &str) -> Option> { // ffmpeg reads the elementary stream from a file rather than a pipe, so // the test never has to drain its output while still writing its input. - let path = temp_path(name); + let path = temp_path(codec, name); std::fs::write(&path, stream).expect("write the elementary stream"); let output = Command::new("ffmpeg") .args(["-v", "error", "-i"]) @@ -969,8 +1087,8 @@ mod tests { } /// Asks ffprobe for the stream's `color_range,color_space`. - fn ffprobe_color(stream: &[u8]) -> Option { - let path = temp_path("probe"); + fn ffprobe_color(codec: Codec, stream: &[u8]) -> Option { + let path = temp_path(codec, "probe"); std::fs::write(&path, stream).expect("write the elementary stream"); let output = Command::new("ffprobe") .args([ @@ -1020,12 +1138,15 @@ mod tests { } /// Returns the NAL unit types in an Annex-B access unit. - fn nal_types(unit: &[u8]) -> Vec { + fn nal_types(codec: Codec, unit: &[u8]) -> Vec { let mut types = Vec::new(); let mut i = 0; while i + 3 < unit.len() { if unit[i..i + 3] == [0, 0, 1] { - types.push(unit[i + 3] & 0x1f); + types.push(match codec { + Codec::H264 => unit[i + 3] & 0x1f, + Codec::H265 => (unit[i + 3] >> 1) & 0x3f, + }); i += 3; } i += 1; @@ -1075,103 +1196,126 @@ mod tests { DmaBuf::from_prime(surface.export_prime().expect("export the surface")).expect("one object") } - /// With no node named, the encoder opens one that encodes H.264 and records - /// which. A named node is used or refused, never swapped for another. + /// With no node named, the encoder opens one that encodes its codec and + /// records which. A named node is used or refused, never swapped for another. #[test] fn the_render_node_is_found_or_named() { - let Some(encoder) = encoder(config()) else { return }; - let device = encoder.config().device.clone().expect("the picked node is recorded"); - let display = Display::open_drm_display(&device).expect("reopen the picked node"); - probe(&display).expect("the picked node encodes H.264"); - - let named = Config { - device: Some(PathBuf::from("/dev/null")), - ..config() - }; - assert!( - Encoder::new(named).is_err(), - "a named node that is not a render node was swapped for one that is" - ); + for codec in CODECS { + let Some(encoder) = encoder(config(codec)) else { + continue; + }; + let device = encoder.config().device.clone().expect("the picked node is recorded"); + let display = Display::open_drm_display(&device).expect("reopen the picked node"); + probe(&display, codec).expect("the picked node encodes the codec"); + + let named = Config { + device: Some(PathBuf::from("/dev/null")), + ..config(codec) + }; + assert!( + Encoder::new(named).is_err(), + "a named node that is not a render node was swapped for one that is" + ); + } } - /// CPU NV12 in, H.264 out that an independent decoder reads back as the same picture. + /// CPU NV12 in, a stream out that an independent decoder reads back as the + /// same pictures, at a size that is whole coding blocks and at one that is not. #[test] fn nv12_round_trips_through_a_software_decoder() { - let Some(mut encoder) = encoder(config()) else { return }; - let mut stream = Vec::new(); - let mut frames = Vec::new(); - for step in 0..10 { - let frame = nv12_frame(WIDTH, HEIGHT, step); - stream.extend(encoder.encode_nv12(&frame, false).expect("encode")); - frames.push(frame); - } - let Some(decoded) = ffmpeg_decode(&stream, "nv12") else { - return; - }; - let len = frames[0].len(); - assert_eq!(decoded.len(), len * frames.len(), "ffmpeg decoded a different count"); - for (index, frame) in frames.iter().enumerate() { - let psnr = luma_psnr(frame, &decoded[index * len..], WIDTH, HEIGHT); - assert!(psnr > 35.0, "picture {index} decoded at {psnr:.1} dB"); + for codec in CODECS { + for (width, height) in [(WIDTH, HEIGHT), (852, 478)] { + let Some(mut encoder) = encoder(Config { + width, + height, + ..config(codec) + }) else { + continue; + }; + let mut stream = Vec::new(); + let mut frames = Vec::new(); + for step in 0..10 { + let frame = nv12_frame(width, height, step); + stream.extend(encoder.encode_nv12(&frame, false).expect("encode")); + frames.push(frame); + } + let Some(decoded) = ffmpeg_decode(codec, &stream, "nv12") else { + return; + }; + let len = frames[0].len(); + assert_eq!( + decoded.len(), + len * frames.len(), + "{codec:?} at {width}x{height}: ffmpeg decoded a different count or size" + ); + for (index, frame) in frames.iter().enumerate() { + let psnr = luma_psnr(frame, &decoded[index * len..], width, height); + assert!( + psnr > 35.0, + "{codec:?} at {width}x{height}: picture {index} decoded at {psnr:.1} dB" + ); + } + } } } /// The SPS names the color space the encoder was configured with. #[test] fn the_stream_is_labelled_with_its_color_space() { - for (config, expected) in [ - (config(), "tv,smpte170m"), - ( - Config { - color: Color::BT709, - ..config() - }, - "tv,bt709", - ), - ] { - let Some(mut encoder) = encoder(config) else { return }; - let unit = encoder - .encode_nv12(&nv12_frame(WIDTH, HEIGHT, 0), true) - .expect("encode"); - let Some(color) = ffprobe_color(&unit) else { - eprintln!("skipping: no ffprobe"); - return; - }; - assert_eq!(color, expected); + for codec in CODECS { + for (color, expected) in [(Color::BT601, "tv,smpte170m"), (Color::BT709, "tv,bt709")] { + let Some(mut encoder) = encoder(Config { color, ..config(codec) }) else { + break; + }; + let unit = encoder + .encode_nv12(&nv12_frame(WIDTH, HEIGHT, 0), true) + .expect("encode"); + let Some(found) = ffprobe_color(codec, &unit) else { + eprintln!("skipping: no ffprobe"); + return; + }; + assert_eq!(found, expected, "{codec:?}"); + } } } - /// An IDR opens every group, and only there. + /// An IDR opens every group, and only there, carrying the parameter sets. #[test] fn keyframes_follow_the_group_length() { - let Some(mut encoder) = encoder(Config { gop: 10, ..config() }) else { - return; - }; - for step in 0..25 { - let unit = encoder - .encode_nv12(&nv12_frame(WIDTH, HEIGHT, step), false) - .expect("encode"); - let types = nal_types(&unit); - let idr = types.contains(&5); - assert_eq!(idr, step % 10 == 0, "frame {step} has NAL types {types:?}"); - if idr { - assert!( - types.contains(&7) && types.contains(&8), - "an IDR carries its SPS and PPS" + for codec in CODECS { + let Some(mut encoder) = encoder(Config { + gop: 10, + ..config(codec) + }) else { + continue; + }; + let (idr_type, parameter_sets) = idr_types(codec); + for step in 0..25 { + let unit = encoder + .encode_nv12(&nv12_frame(WIDTH, HEIGHT, step), false) + .expect("encode"); + let types = nal_types(codec, &unit); + let idr = types.contains(&idr_type); + assert_eq!(idr, step % 10 == 0, "{codec:?} frame {step} has NAL types {types:?}"); + if idr { + assert!( + parameter_sets.iter().all(|kind| types.contains(kind)), + "a {codec:?} IDR carries its parameter sets: {types:?}" + ); + } + } + // A forced keyframe restarts the count: the next IDR is a group later. + for step in 0..=10 { + let unit = encoder + .encode_nv12(&nv12_frame(WIDTH, HEIGHT, step), step == 0) + .expect("encode"); + assert_eq!( + nal_types(codec, &unit).contains(&idr_type), + step % 10 == 0, + "{codec:?} frame {step} after the forced keyframe" ); } } - // A forced keyframe restarts the count: the next IDR is a group later. - for step in 0..=10 { - let unit = encoder - .encode_nv12(&nv12_frame(WIDTH, HEIGHT, step), step == 0) - .expect("encode"); - assert_eq!( - nal_types(&unit).contains(&5), - step % 10 == 0, - "frame {step} after the forced keyframe" - ); - } } #[test] @@ -1188,7 +1332,9 @@ mod tests { /// CPU fallback can encode the same frame. #[test] fn a_refused_import_leaves_the_encoder_usable() { - let Some(mut encoder) = encoder(config()) else { return }; + let Some(mut encoder) = encoder(config(Codec::H264)) else { + return; + }; let Some(display) = Display::open() else { eprintln!("skipping: no VA-API display"); return; @@ -1208,42 +1354,56 @@ mod tests { let unit = encoder .encode_nv12(&nv12_frame(WIDTH, HEIGHT, 0), true) .expect("the CPU path still encodes"); - let Some(decoded) = ffmpeg_decode(&unit, "refused") else { + let Some(decoded) = ffmpeg_decode(Codec::H264, &unit, "refused") else { return; }; let psnr = luma_psnr(&nv12_frame(WIDTH, HEIGHT, 0), &decoded, WIDTH, HEIGHT); assert!(psnr > 35.0, "decoded at {psnr:.1} dB"); } - /// A 1080-row buffer is not encoded in place: the encoder reads 1088 rows, - /// so it goes through the video processor into the padded input surface. + /// A buffer short of whole coding blocks is not encoded in place: the + /// encoder reads whole blocks, so it goes through the video processor into + /// the padded input surface. That is 1080 rows for H.264, which reads 1088, + /// and 1076 for H.265, whose minimum coding blocks are 8 or 16 rows. #[test] - fn a_buffer_short_of_whole_macroblocks_is_blitted() { - let (width, height) = (1920, 1080); - let Some(mut encoder) = encoder(Config::new(width, height, 30, 4_000_000, 30)) else { - return; - }; - let Some(display) = Display::open() else { - eprintln!("skipping: no VA-API display"); - return; - }; - let unit = encoder - .encode_dmabuf(nv12_dmabuf(&display, (width, height), 0), true) - .expect("encode the DMA-BUF"); - let Some(decoded) = ffmpeg_decode(&unit, "1080") else { - return; - }; - let psnr = luma_psnr(&nv12_frame(width, height, 0), &decoded, width, height); - assert!(psnr > 35.0, "decoded at {psnr:.1} dB"); + fn a_buffer_short_of_whole_coding_blocks_is_blitted() { + for (codec, (width, height)) in [(Codec::H264, (1920, 1080)), (Codec::H265, (1920, 1076))] { + let Some(mut encoder) = encoder(Config { + codec, + ..Config::new(width, height, 30, 4_000_000, 30) + }) else { + continue; + }; + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + let unit = encoder + .encode_dmabuf(nv12_dmabuf(&display, (width, height), 0), true) + .expect("encode the DMA-BUF"); + let Some(decoded) = ffmpeg_decode(codec, &unit, "short") else { + return; + }; + let frame = nv12_frame(width, height, 0); + assert_eq!(decoded.len(), frame.len(), "{codec:?} decoded at another size"); + let psnr = luma_psnr(&frame, &decoded, width, height); + assert!(psnr > 35.0, "{codec:?} decoded at {psnr:.1} dB"); + } } /// A YUV buffer labelled with another color space is converted into the /// stream's, not relabelled. #[test] fn a_yuv_dmabuf_in_another_space_is_converted() { + for codec in CODECS { + yuv_dmabuf_in_another_space_is_converted(codec); + } + } + + fn yuv_dmabuf_in_another_space_is_converted(codec: Codec) { let Some(mut encoder) = encoder(Config { color: Color::BT601, - ..config() + ..config(codec) }) else { return; }; @@ -1276,7 +1436,7 @@ mod tests { ..DmaBuf::from_prime(surface.export_prime().expect("export")).expect("one object") }; let unit = encoder.encode_dmabuf(buffer, true).expect("encode"); - let Some(decoded) = ffmpeg_decode(&unit, "yuv-color") else { + let Some(decoded) = ffmpeg_decode(codec, &unit, "yuv-color") else { return; }; @@ -1291,7 +1451,7 @@ mod tests { for (channel, (a, e)) in actual.iter().zip(expected).enumerate() { assert!( (a - e).abs() < 6.0, - "channel {channel}: decoded {a:.1}, expected BT.601 {e:.1}" + "{codec:?} channel {channel}: decoded {a:.1}, expected BT.601 {e:.1}" ); } } @@ -1299,36 +1459,48 @@ mod tests { /// An NV12 DMA-BUF at the encoder's size is encoded from the producer's memory. #[test] fn an_nv12_dmabuf_encodes_in_place() { - let Some(mut encoder) = encoder(config()) else { return }; - let Some(display) = Display::open() else { - eprintln!("skipping: no VA-API display"); - return; - }; - let mut stream = Vec::new(); - for step in 0..5 { - let buffer = nv12_dmabuf(&display, (WIDTH, HEIGHT), step); - stream.extend(encoder.encode_dmabuf(buffer, false).expect("encode the DMA-BUF")); - } - let Some(decoded) = ffmpeg_decode(&stream, "dmabuf") else { - return; - }; - let len = nv12_frame(WIDTH, HEIGHT, 0).len(); - assert_eq!(decoded.len(), 5 * len); - for step in 0..5 { - let psnr = luma_psnr( - &nv12_frame(WIDTH, HEIGHT, step), - &decoded[step as usize * len..], - WIDTH, - HEIGHT, - ); - assert!(psnr > 35.0, "picture {step} decoded at {psnr:.1} dB"); + for codec in CODECS { + let Some(mut encoder) = encoder(config(codec)) else { + continue; + }; + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + let mut stream = Vec::new(); + for step in 0..5 { + let buffer = nv12_dmabuf(&display, (WIDTH, HEIGHT), step); + stream.extend(encoder.encode_dmabuf(buffer, false).expect("encode the DMA-BUF")); + } + let Some(decoded) = ffmpeg_decode(codec, &stream, "dmabuf") else { + return; + }; + let len = nv12_frame(WIDTH, HEIGHT, 0).len(); + assert_eq!(decoded.len(), 5 * len, "{codec:?}"); + for step in 0..5 { + let psnr = luma_psnr( + &nv12_frame(WIDTH, HEIGHT, step), + &decoded[step as usize * len..], + WIDTH, + HEIGHT, + ); + assert!(psnr > 35.0, "{codec:?} picture {step} decoded at {psnr:.1} dB"); + } } } /// An NV12 DMA-BUF at another size is scaled to the encoder's on the GPU. #[test] fn a_dmabuf_at_another_size_is_scaled() { - let Some(mut encoder) = encoder(config()) else { return }; + for codec in CODECS { + dmabuf_at_another_size_is_scaled(codec); + } + } + + fn dmabuf_at_another_size_is_scaled(codec: Codec) { + let Some(mut encoder) = encoder(config(codec)) else { + return; + }; let Some(display) = Display::open() else { eprintln!("skipping: no VA-API display"); return; @@ -1338,11 +1510,11 @@ mod tests { let buffer = nv12_dmabuf(&display, (WIDTH * 2, HEIGHT * 2), step); stream.extend(encoder.encode_dmabuf(buffer, false).expect("encode the DMA-BUF")); } - let Some(decoded) = ffmpeg_decode(&stream, "scaled") else { + let Some(decoded) = ffmpeg_decode(codec, &stream, "scaled") else { return; }; let len = nv12_frame(WIDTH, HEIGHT, 0).len(); - assert_eq!(decoded.len(), 3 * len, "the output is at the encoder's size"); + assert_eq!(decoded.len(), 3 * len, "{codec:?}: the output is at the encoder's size"); // Compare against the source scaled on the CPU by point sampling, which // differs from the GPU filter only at edges. 25 dB separates a scale // from a crop, which scores about 16 dB on this gradient. @@ -1355,7 +1527,7 @@ mod tests { } } let psnr = luma_psnr(&expected, &decoded[step as usize * len..], WIDTH, HEIGHT); - assert!(psnr > 25.0, "picture {step} decoded at {psnr:.1} dB"); + assert!(psnr > 25.0, "{codec:?} picture {step} decoded at {psnr:.1} dB"); } } @@ -1364,9 +1536,12 @@ mod tests { fn an_rgb_dmabuf_is_converted_to_the_configured_space() { // Four flat quadrants, so the check reads colors rather than edges. const QUADRANTS: [[u8; 3]; 4] = [[255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 255]]; - for color in [Color::BT601, Color::BT709] { - let Some(mut encoder) = encoder(Config { color, ..config() }) else { - return; + for (codec, color) in CODECS + .into_iter() + .flat_map(|codec| [(codec, Color::BT601), (codec, Color::BT709)]) + { + let Some(mut encoder) = encoder(Config { color, ..config(codec) }) else { + continue; }; let Some(display) = Display::open() else { eprintln!("skipping: no VA-API display"); @@ -1394,7 +1569,7 @@ mod tests { let buffer = DmaBuf::from_prime(exported).expect("one object"); assert_eq!(buffer.drm_format, crate::dmabuf::drm::XRGB8888); let unit = encoder.encode_dmabuf(buffer, false).expect("encode the RGB DMA-BUF"); - let Some(decoded) = ffmpeg_decode(&unit, "rgb") else { + let Some(decoded) = ffmpeg_decode(codec, &unit, "rgb") else { return; }; @@ -1421,7 +1596,7 @@ mod tests { for (channel, (a, e)) in actual.iter().zip(expected).enumerate() { assert!( (a - e).abs() < 6.0, - "{color:?} quadrant {quadrant} channel {channel}: decoded {a:.1}, expected {e:.1}" + "{codec:?} {color:?} quadrant {quadrant} channel {channel}: decoded {a:.1}, expected {e:.1}" ); } } @@ -1456,9 +1631,19 @@ mod tests { /// Lowering the bitrate mid-stream shrinks the pictures without an IDR. #[test] fn a_bitrate_change_takes_effect_without_a_keyframe() { - let Some(mut encoder) = encoder(Config::new(WIDTH, HEIGHT, 30, 2_000_000, 1_000)) else { + for codec in CODECS { + bitrate_change_takes_effect_without_a_keyframe(codec); + } + } + + fn bitrate_change_takes_effect_without_a_keyframe(codec: Codec) { + let Some(mut encoder) = encoder(Config { + codec, + ..Config::new(WIDTH, HEIGHT, 30, 2_000_000, 1_000) + }) else { return; }; + let (idr_type, _) = idr_types(codec); // The moving block over light noise: enough detail that 2 Mbps is // spent, not so much that the quantizer bound decides the size instead // of the rate controller. @@ -1481,7 +1666,10 @@ mod tests { .map(|step| { let unit = encoder.encode_nv12(&frame(step), false).expect("encode"); if step > 0 { - assert!(!nal_types(&unit).contains(&5), "picture {step} is an IDR"); + assert!( + !nal_types(codec, &unit).contains(&idr_type), + "{codec:?} picture {step} is an IDR" + ); } unit.len() }) @@ -1494,10 +1682,10 @@ mod tests { assert_eq!(encoder.config().bitrate, 250_000); assert!(encoder.set_bitrate(0).is_err(), "a zero bitrate is refused"); let low = average(&mut encoder, 60); - eprintln!("mean picture at 2 Mbps: {high} bytes, at 250 kbps: {low} bytes"); + eprintln!("{codec:?} mean picture at 2 Mbps: {high} bytes, at 250 kbps: {low} bytes"); assert!( low * 3 < high, - "the rate did not drop: {high} -> {low} bytes per picture" + "{codec:?}: the rate did not drop: {high} -> {low} bytes per picture" ); } } diff --git a/src/encode/h265.rs b/src/encode/h265.rs new file mode 100644 index 0000000..026048b --- /dev/null +++ b/src/encode/h265.rs @@ -0,0 +1,874 @@ +//! The H.265 half of the encoder: Main profile parameter sets and slice headers +//! written here, and the VA-API buffers that describe each picture. +//! +//! The parameter sets and slice headers are written here and handed to the +//! driver packed, so it produces only the slice data. That is what ffmpeg's +//! `hevc_vaapi` does, and the path the Intel and AMD drivers are exercised on. +//! A driver must take the packed parameter sets, since VA's sequence parameters +//! cannot carry the conformance window or the VUI; one that does not take +//! packed slice headers writes them itself, from parameter buffers that +//! describe the same stream. +//! +//! The coding tools follow what the driver reports through +//! `VAConfigAttribEncHEVCFeatures` and `VAConfigAttribEncHEVCBlockSizes`: a +//! tool it supports or requires is on, and one it requires that this writer +//! does not emit fails the open. A driver that can predict only from +//! non-empty lists in both directions (`VA_PREDICTION_DIRECTION_BI_NOT_EMPTY`, +//! Intel's low-power encoder among them) gets each P picture as a generalized +//! B picture: a B slice whose list 1 repeats list 0. + +use std::rc::Rc; + +use anyhow::{anyhow, bail}; + +use super::{Config, Reference}; +use crate::{ + bindings, BufferType, Config as VaConfig, Display, EncCodedBuffer, EncPackedHeaderType, EncPictureParameter, + EncPictureParameterBufferHEVC, EncSequenceParameter, EncSequenceParameterBufferHEVC, EncSliceParameter, + EncSliceParameterBufferHEVC, GenericValue, HEVCEncPicFields, HEVCEncSeqFields, HevcEncPicSccFields, + HevcEncSeqSccFields, HevcEncSliceFields, PictureHEVC, VAConfigAttrib, VAConfigAttribType, VAEntrypoint, VAProfile, + VASurfaceAttribType, VA_ATTRIB_NOT_SUPPORTED, VA_ENC_PACKED_HEADER_SEQUENCE, VA_ENC_PACKED_HEADER_SLICE, + VA_INVALID_ID, VA_PICTURE_HEVC_INVALID, VA_PICTURE_HEVC_RPS_ST_CURR_BEFORE, VA_PREDICTION_DIRECTION_BI_NOT_EMPTY, +}; + +/// NAL unit types, H.265 Table 7-1. +const TRAIL_R: u8 = 1; +const IDR_W_RADL: u8 = 19; +const VPS: u8 = 32; +const SPS: u8 = 33; +const PPS: u8 = 34; + +/// Slice types, H.265 Table 7-7. +const SLICE_B: u8 = 0; +const SLICE_P: u8 = 1; +const SLICE_I: u8 = 2; + +/// `general_profile_idc` of the Main profile. +const MAIN: u8 = 1; + +/// `log2_max_pic_order_cnt_lsb_minus4`: twelve bits of order count per slice, as ffmpeg writes. +const LOG2_MAX_POC_LSB_MINUS4: u32 = 8; + +/// The QP the PPS opens each picture at; rate control moves off it per picture. +const INIT_QP: i32 = 26; + +/// The VA profile this encoder opens. +pub(super) const PROFILE: VAProfile::Type = VAProfile::VAProfileHEVCMain; + +/// How a feature of `VAConfigAttribValEncHEVCFeatures` is offered. +const UNSUPPORTED: u32 = 0; +const REQUIRED: u32 = 2; + +/// The coding tools and block sizes a stream is written with, settled once from what the driver reports. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct Tools { + log2_min_cb: u32, + log2_ctb: u32, + log2_min_tb: u32, + log2_max_tb: u32, + depth_inter: u32, + depth_intra: u32, + amp: bool, + sao: bool, + temporal_mvp: bool, + transform_skip: bool, + cu_qp_delta: bool, + /// Whether P pictures go out as B slices whose list 1 repeats list 0. + gpb: bool, +} + +impl Tools { + /// What ffmpeg assumes of a driver too old to report its block sizes and features: the first Intel HEVC encoder's. + const FALLBACK: Self = Self { + log2_min_cb: 4, + log2_ctb: 5, + log2_min_tb: 2, + log2_max_tb: 5, + depth_inter: 3, + depth_intra: 3, + amp: true, + sao: false, + temporal_mvp: false, + transform_skip: false, + cu_qp_delta: true, + gpb: false, + }; + + /// Asks `display` what its H.265 Main encoder on `entrypoint` supports. + /// + /// # Errors + /// + /// Fails when the driver cannot be asked, when it predicts from no + /// previous picture, and when it requires a tool this writer does not emit. + fn query(display: &Display, entrypoint: VAEntrypoint::Type) -> anyhow::Result { + let mut attrs = [ + VAConfigAttribType::VAConfigAttribEncHEVCFeatures, + VAConfigAttribType::VAConfigAttribEncHEVCBlockSizes, + VAConfigAttribType::VAConfigAttribEncMaxRefFrames, + VAConfigAttribType::VAConfigAttribPredictionDirection, + ] + .map(|type_| VAConfigAttrib { type_, value: 0 }); + display + .get_config_attributes(PROFILE, entrypoint, &mut attrs) + .map_err(|e| anyhow!("query H.265 encode attributes: {e:?}"))?; + let reported = |index: usize| Some(attrs[index].value).filter(|&value| value != VA_ATTRIB_NOT_SUPPORTED); + let (features, blocks, references, direction) = (reported(0), reported(1), reported(2), reported(3)); + + let mut tools = Self::FALLBACK; + if let Some(blocks) = blocks { + let blocks = bindings::VAConfigAttribValEncHEVCBlockSizes { value: blocks }; + // Safe because both members of the union are plain integers over the same 32 bits. + let blocks = unsafe { blocks.bits }; + tools.log2_ctb = blocks.log2_max_coding_tree_block_size_minus3() + 3; + tools.log2_min_cb = blocks.log2_min_luma_coding_block_size_minus3() + 3; + tools.log2_min_tb = blocks.log2_min_luma_transform_block_size_minus2() + 2; + tools.log2_max_tb = blocks.log2_max_luma_transform_block_size_minus2() + 2; + tools.depth_inter = blocks.max_max_transform_hierarchy_depth_inter(); + tools.depth_intra = blocks.max_max_transform_hierarchy_depth_intra(); + } + if let Some(features) = features { + let features = bindings::VAConfigAttribValEncHEVCFeatures { value: features }; + // Safe because both members of the union are plain integers over the same 32 bits. + let features = unsafe { features.bits }; + tools.amp = features.amp() != UNSUPPORTED; + tools.sao = features.sao() != UNSUPPORTED; + tools.temporal_mvp = features.temporal_mvp() != UNSUPPORTED; + tools.transform_skip = features.transform_skip() != UNSUPPORTED; + tools.cu_qp_delta = features.cu_qp_delta() != UNSUPPORTED; + for (name, offered) in [ + ("separate colour planes", features.separate_colour_planes()), + ("scaling lists", features.scaling_lists()), + ("PCM", features.pcm()), + ("strong intra smoothing", features.strong_intra_smoothing()), + ("dependent slices", features.dependent_slices()), + ("sign data hiding", features.sign_data_hiding()), + ("constrained intra prediction", features.constrained_intra_pred()), + ("weighted prediction", features.weighted_prediction()), + ("transquant bypass", features.transquant_bypass()), + ("a disabled deblocking filter", features.deblocking_filter_disable()), + ] { + if offered == REQUIRED { + bail!("the driver requires {name}, which this encoder does not write"); + } + } + } + if tools.log2_min_cb > tools.log2_ctb + || tools.log2_min_tb > tools.log2_max_tb + || tools.log2_max_tb > tools.log2_ctb.min(5) + || tools.log2_min_tb >= tools.log2_min_cb + { + bail!("the driver reports inconsistent H.265 block sizes: {tools:?}"); + } + + let (list0, list1) = match references { + Some(value) => (value & 0xffff, value >> 16 & 0xffff), + None => (1, 0), + }; + if list0 == 0 { + bail!("the driver's H.265 encoder predicts from no reference picture"); + } + tools.gpb = + list1 > 0 && direction.is_some_and(|direction| direction & VA_PREDICTION_DIRECTION_BI_NOT_EMPTY != 0); + Ok(tools) + } +} + +/// An H.265 stream's settled parameters: everything per picture derives from these. +pub(super) struct Stream { + tools: Tools, + /// The coded size: the frame rounded up to whole minimum coding blocks and to the driver's surface alignment. + coded: (u32, u32), + level_idc: u8, + /// The VPS, SPS, and PPS, each behind a start code, sent ahead of every IDR. + headers: Vec, +} + +impl Stream { + /// Settles the parameters for `config` on `display`'s encoder at `entrypoint`, configured as `va_config`. + /// + /// # Errors + /// + /// Fails as [`Tools::query`] does, and for a frame too large for any level. + pub(super) fn new( + display: &Display, + entrypoint: VAEntrypoint::Type, + va_config: &mut VaConfig, + config: &Config, + ) -> anyhow::Result { + let tools = Tools::query(display, entrypoint)?; + Self::with_tools(tools, alignment(va_config), config) + } + + /// `alignment` is the driver's surface alignment, in samples across and down. + fn with_tools(tools: Tools, alignment: (u32, u32), config: &Config) -> anyhow::Result { + let min_cb = 1 << tools.log2_min_cb; + let coded = ( + config.width.next_multiple_of(min_cb.max(alignment.0)), + config.height.next_multiple_of(min_cb.max(alignment.1)), + ); + let level_idc = level(coded, config.framerate).ok_or_else(|| { + anyhow!( + "{}x{} at {} fps exceeds every H.265 level", + coded.0, + coded.1, + config.framerate + ) + })?; + let mut stream = Self { + tools, + coded, + level_idc, + headers: Vec::new(), + }; + stream.headers = [stream.vps(config), stream.sps(config), stream.pps()].concat(); + Ok(stream) + } + + /// Returns the coded size: what the context, the reconstructed surfaces, and the SPS are sized to. + pub(super) fn coded(&self) -> (u32, u32) { + self.coded + } + + /// Returns the VA buffers for one picture, in the order the driver takes them: sequence, picture, slice, then packed headers. + pub(super) fn buffers(&self, config: &Config, picture: &super::PictureInfo<'_>, packed: u32) -> Vec { + let tools = &self.tools; + let poc = picture.count as i32; + let current = PictureHEVC::new(picture.recon, poc, 0); + let reference = picture.reference.map(|Reference { surface, count }| { + PictureHEVC::new(surface, count as i32, VA_PICTURE_HEVC_RPS_ST_CURR_BEFORE) + }); + let slice_type = match (reference.is_some(), tools.gpb) { + (false, _) => SLICE_I, + (true, false) => SLICE_P, + (true, true) => SLICE_B, + }; + let nal_type = if picture.idr { IDR_W_RADL } else { TRAIL_R }; + + let mut buffers = vec![ + self.sequence_parameters(config, picture.bitrate), + self.picture_parameters(picture.coded, current, reference.as_ref(), nal_type), + self.slice_parameters(slice_type, reference.as_ref()), + ]; + if picture.idr && packed & VA_ENC_PACKED_HEADER_SEQUENCE != 0 { + buffers.extend(packed_header(EncPackedHeaderType::Sequence, &self.headers)); + } + if packed & VA_ENC_PACKED_HEADER_SLICE != 0 { + let header = self.slice_header(nal_type, slice_type, picture.count); + buffers.extend(packed_header(EncPackedHeaderType::Slice, &header)); + } + buffers + } + + fn ctb_count(&self) -> u32 { + let ctb = 1 << self.tools.log2_ctb; + self.coded.0.div_ceil(ctb) * self.coded.1.div_ceil(ctb) + } + + fn sequence_parameters(&self, config: &Config, bitrate: u32) -> BufferType { + let tools = &self.tools; + let fields = HEVCEncSeqFields::new( + 1, // chroma_format_idc: 4:2:0 + 0, + 0, + 0, + 0, + 0, + tools.amp as u32, + tools.sao as u32, + 0, + 0, + tools.temporal_mvp as u32, + 0, + 0, + ); + BufferType::EncSequenceParameter(EncSequenceParameter::HEVC(EncSequenceParameterBufferHEVC::new( + MAIN, + self.level_idc, + 0, // Main tier + config.gop, + config.gop, + 1, // ip_period: no B pictures + bitrate, + self.coded.0 as u16, + self.coded.1 as u16, + &fields, + (tools.log2_min_cb - 3) as u8, + (tools.log2_ctb - tools.log2_min_cb) as u8, + (tools.log2_min_tb - 2) as u8, + (tools.log2_max_tb - tools.log2_min_tb) as u8, + tools.depth_inter as u8, + tools.depth_intra as u8, + 0, + 0, + 0, + 0, + // The VUI is in the packed SPS, which is what goes out. + None, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + &HevcEncSeqSccFields::new(0), + ))) + } + + fn picture_parameters( + &self, + coded: &EncCodedBuffer, + current: PictureHEVC, + reference: Option<&PictureHEVC>, + nal_type: u8, + ) -> BufferType { + let tools = &self.tools; + let idr = nal_type == IDR_W_RADL; + let fields = HEVCEncPicFields::new( + idr as u32, + if reference.is_some() { 2 } else { 1 }, // coding_type: P (or GPB), else I + 1, // reference_pic_flag + 0, + 0, + 0, + tools.transform_skip as u32, + tools.cu_qp_delta as u32, + 0, + 0, + 0, + 0, + 0, + 0, + 1, // pps_loop_filter_across_slices_enabled_flag + 0, + 0, + 0, + 0, + ); + let mut references: [PictureHEVC; 15] = std::array::from_fn(|_| invalid()); + if let Some(reference) = reference { + references[0] = reference.clone(); + } + BufferType::EncPictureParameter(EncPictureParameter::HEVC(EncPictureParameterBufferHEVC::new( + current, + references, + coded.id(), + if tools.temporal_mvp { 0 } else { 0xff }, + 0, + INIT_QP as u8, + self.diff_cu_qp_delta_depth() as u8, + 0, + 0, + 0, + 0, + [0; 19], + [0; 21], + 0, + 0, + 0, + 0, + 0, + nal_type, + &fields, + 0, + 0, + &HevcEncPicSccFields::new(0), + ))) + } + + fn slice_parameters(&self, slice_type: u8, reference: Option<&PictureHEVC>) -> BufferType { + let tools = &self.tools; + let mut list0: [PictureHEVC; 15] = std::array::from_fn(|_| invalid()); + let mut list1: [PictureHEVC; 15] = std::array::from_fn(|_| invalid()); + if let Some(reference) = reference { + list0[0] = reference.clone(); + if slice_type == SLICE_B { + list1[0] = reference.clone(); + } + } + let fields = HevcEncSliceFields::new( + 1, // last_slice_of_pic_flag + 0, + 0, + (tools.temporal_mvp && reference.is_some()) as u32, + tools.sao as u32, + tools.sao as u32, + 0, + 0, + 0, + 0, + 0, + 1, // collocated_from_l0_flag + ); + BufferType::EncSliceParameter(EncSliceParameter::HEVC(EncSliceParameterBufferHEVC::new( + 0, + self.ctb_count(), + slice_type, + 0, + 0, + 0, + list0, + list1, + 0, + 0, + [0; 15], + [0; 15], + [[0; 2]; 15], + [[0; 2]; 15], + [0; 15], + [0; 15], + [[0; 2]; 15], + [[0; 2]; 15], + 5, // max_num_merge_cand + 0, + 0, + 0, + 0, + 0, + &fields, + 0, + 0, + ))) + } + + fn diff_cu_qp_delta_depth(&self) -> u32 { + match self.tools.cu_qp_delta { + true => self.tools.log2_ctb - self.tools.log2_min_cb, + false => 0, + } + } + + /// Writes `profile_tier_level` for Main with no sub-layers, H.265 7.3.3. + fn profile_tier_level(&self, w: &mut Rbsp) { + w.bits(2, 0); // general_profile_space + w.flag(false); // general_tier_flag: Main + w.bits(5, MAIN.into()); + // Main, and Main 10 which every Main stream also conforms to. + w.bits(32, 1 << (31 - 1) | 1 << (31 - 2)); + w.flag(true); // general_progressive_source_flag + w.flag(false); // general_interlaced_source_flag + w.flag(true); // general_non_packed_constraint_flag + w.flag(true); // general_frame_only_constraint_flag + // The 43 reserved bits and general_inbld_flag of a Main profile. + w.bits(32, 0); + w.bits(12, 0); + w.bits(8, self.level_idc.into()); + } + + /// Writes the video parameter set, H.265 7.3.2.1. + fn vps(&self, config: &Config) -> Vec { + let mut w = Rbsp::default(); + w.bits(4, 0); // vps_video_parameter_set_id + w.flag(true); // vps_base_layer_internal_flag + w.flag(true); // vps_base_layer_available_flag + w.bits(6, 0); // vps_max_layers_minus1 + w.bits(3, 0); // vps_max_sub_layers_minus1 + w.flag(true); // vps_temporal_id_nesting_flag + w.bits(16, 0xffff); + self.profile_tier_level(&mut w); + w.flag(false); // vps_sub_layer_ordering_info_present_flag + ordering(&mut w); + w.bits(6, 0); // vps_max_layer_id + w.ue(0); // vps_num_layer_sets_minus1 + w.flag(true); // vps_timing_info_present_flag + timing(&mut w, config.framerate); + w.ue(0); // vps_num_hrd_parameters + w.flag(false); // vps_extension_flag + w.trailing(); + nal(VPS, &w.finish()) + } + + /// Writes the sequence parameter set, H.265 7.3.2.2, with its VUI. + fn sps(&self, config: &Config) -> Vec { + let tools = &self.tools; + let mut w = Rbsp::default(); + w.bits(4, 0); // sps_video_parameter_set_id + w.bits(3, 0); // sps_max_sub_layers_minus1 + w.flag(true); // sps_temporal_id_nesting_flag + self.profile_tier_level(&mut w); + w.ue(0); // sps_seq_parameter_set_id + w.ue(1); // chroma_format_idc: 4:2:0 + w.ue(self.coded.0); + w.ue(self.coded.1); + // The conformance window crops the coded size back to the frame, in + // chroma samples, so in units of two luma samples each way. + let crop = (self.coded.0 - config.width, self.coded.1 - config.height); + w.flag(crop != (0, 0)); + if crop != (0, 0) { + w.ue(0); + w.ue(crop.0 / 2); + w.ue(0); + w.ue(crop.1 / 2); + } + w.ue(0); // bit_depth_luma_minus8 + w.ue(0); // bit_depth_chroma_minus8 + w.ue(LOG2_MAX_POC_LSB_MINUS4); + w.flag(false); // sps_sub_layer_ordering_info_present_flag + ordering(&mut w); + w.ue(tools.log2_min_cb - 3); + w.ue(tools.log2_ctb - tools.log2_min_cb); + w.ue(tools.log2_min_tb - 2); + w.ue(tools.log2_max_tb - tools.log2_min_tb); + w.ue(tools.depth_inter); + w.ue(tools.depth_intra); + w.flag(false); // scaling_list_enabled_flag + w.flag(tools.amp); + w.flag(tools.sao); + w.flag(false); // pcm_enabled_flag + w.ue(0); // num_short_term_ref_pic_sets: each slice carries its own + w.flag(false); // long_term_ref_pics_present_flag + w.flag(tools.temporal_mvp); + w.flag(false); // strong_intra_smoothing_enabled_flag + w.flag(true); // vui_parameters_present_flag + self.vui(&mut w, config); + w.flag(false); // sps_extension_present_flag + w.trailing(); + nal(SPS, &w.finish()) + } + + /// Writes the VUI, H.265 E.2.1: the color space, the frame rate, and the restrictions IPPP coding keeps. + fn vui(&self, w: &mut Rbsp, config: &Config) { + let (primaries, transfer, matrix) = config.color.vui(); + w.flag(false); // aspect_ratio_info_present_flag + w.flag(false); // overscan_info_present_flag + w.flag(true); // video_signal_type_present_flag + w.bits(3, 5); // video_format: unspecified + w.flag(config.color.full_range); + w.flag(true); // colour_description_present_flag + w.bits(8, primaries.into()); + w.bits(8, transfer.into()); + w.bits(8, matrix.into()); + w.flag(false); // chroma_loc_info_present_flag + w.flag(false); // neutral_chroma_indication_flag + w.flag(false); // field_seq_flag + w.flag(false); // frame_field_info_present_flag + w.flag(false); // default_display_window_flag + w.flag(true); // vui_timing_info_present_flag + timing(w, config.framerate); + w.flag(false); // vui_hrd_parameters_present_flag + w.flag(true); // bitstream_restriction_flag + w.flag(false); // tiles_fixed_structure_flag + w.flag(true); // motion_vectors_over_pic_boundaries_flag + w.flag(true); // restricted_ref_pic_lists_flag + w.ue(0); // min_spatial_segmentation_idc + w.ue(0); // max_bytes_per_pic_denom + w.ue(0); // max_bits_per_min_cu_denom + w.ue(15); // log2_max_mv_length_horizontal + w.ue(15); // log2_max_mv_length_vertical + } + + /// Writes the picture parameter set, H.265 7.3.2.3. + fn pps(&self) -> Vec { + let tools = &self.tools; + let mut w = Rbsp::default(); + w.ue(0); // pps_pic_parameter_set_id + w.ue(0); // pps_seq_parameter_set_id + w.flag(false); // dependent_slice_segments_enabled_flag + w.flag(false); // output_flag_present_flag + w.bits(3, 0); // num_extra_slice_header_bits + w.flag(false); // sign_data_hiding_enabled_flag + w.flag(false); // cabac_init_present_flag + w.ue(0); // num_ref_idx_l0_default_active_minus1 + w.ue(0); // num_ref_idx_l1_default_active_minus1 + w.se(INIT_QP - 26); + w.flag(false); // constrained_intra_pred_flag + w.flag(tools.transform_skip); + w.flag(tools.cu_qp_delta); + if tools.cu_qp_delta { + w.ue(self.diff_cu_qp_delta_depth()); + } + w.se(0); // pps_cb_qp_offset + w.se(0); // pps_cr_qp_offset + w.flag(false); // pps_slice_chroma_qp_offsets_present_flag + w.flag(false); // weighted_pred_flag + w.flag(false); // weighted_bipred_flag + w.flag(false); // transquant_bypass_enabled_flag + w.flag(false); // tiles_enabled_flag + w.flag(false); // entropy_coding_sync_enabled_flag + w.flag(true); // pps_loop_filter_across_slices_enabled_flag + w.flag(false); // deblocking_filter_control_present_flag + w.flag(false); // pps_scaling_list_data_present_flag + w.flag(false); // lists_modification_present_flag + w.ue(0); // log2_parallel_merge_level_minus2 + w.flag(false); // slice_segment_header_extension_present_flag + w.flag(false); // pps_extension_present_flag + w.trailing(); + nal(PPS, &w.finish()) + } + + /// Writes the header of a picture's one slice segment, H.265 7.3.6.1. + /// + /// A P or GPB picture refers to the one before it, so its short-term + /// reference set is that single picture, one order count back. + fn slice_header(&self, nal_type: u8, slice_type: u8, count: u32) -> Vec { + let tools = &self.tools; + let mut w = Rbsp::default(); + w.flag(true); // first_slice_segment_in_pic_flag + if nal_type == IDR_W_RADL { + w.flag(false); // no_output_of_prior_pics_flag + } + w.ue(0); // slice_pic_parameter_set_id + w.ue(slice_type.into()); + if nal_type != IDR_W_RADL { + let lsb_bits = LOG2_MAX_POC_LSB_MINUS4 + 4; + w.bits(lsb_bits, count & ((1 << lsb_bits) - 1)); + w.flag(false); // short_term_ref_pic_set_sps_flag + // st_ref_pic_set(0): no inter prediction of the set at index 0. + w.ue(1); // num_negative_pics + w.ue(0); // num_positive_pics + w.ue(0); // delta_poc_s0_minus1: the previous picture + w.flag(true); // used_by_curr_pic_s0_flag + if tools.temporal_mvp { + w.flag(true); // slice_temporal_mvp_enabled_flag + } + } + if tools.sao { + w.flag(true); // slice_sao_luma_flag + w.flag(true); // slice_sao_chroma_flag + } + if slice_type != SLICE_I { + w.flag(false); // num_ref_idx_active_override_flag + if slice_type == SLICE_B { + w.flag(false); // mvd_l1_zero_flag + } + if tools.temporal_mvp && slice_type == SLICE_B { + w.flag(true); // collocated_from_l0_flag + } + w.ue(0); // five_minus_max_num_merge_cand + } + w.se(0); // slice_qp_delta + // Loop filtering is on, so this is present; with one slice it filters nothing. + w.flag(false); // slice_loop_filter_across_slices_enabled_flag + w.byte_alignment(); + nal(nal_type, &w.finish()) + } +} + +/// Returns an empty reference slot. +fn invalid() -> PictureHEVC { + PictureHEVC::new(VA_INVALID_ID, 0, VA_PICTURE_HEVC_INVALID) +} + +/// Writes the one set of sub-layer ordering limits, for IPPP coding: one reference and the current picture, nothing reordered. +fn ordering(w: &mut Rbsp) { + w.ue(1); // max_dec_pic_buffering_minus1 + w.ue(0); // max_num_reorder_pics + w.ue(0); // max_latency_increase_plus1 +} + +/// Returns the surface alignment the driver encodes at, in samples across and down: 1 where it names none. +/// +/// radeonsi encodes a picture rounded up to it (64 by 16) and rewrites the SPS +/// at that size, keeping our conformance window, so a coded size short of it +/// would decode with the wrong crop. +fn alignment(va_config: &mut VaConfig) -> (u32, u32) { + let named = va_config + .query_surface_attributes_by_type(VASurfaceAttribType::VASurfaceAttribAlignmentSize) + .ok() + .and_then(|values| values.into_iter().next()); + match named { + // VASurfaceAttribAlignmentStruct: log2 width in bits 0-3, log2 height in bits 4-7. + Some(GenericValue::Integer(value)) => (1 << (value & 0xf), 1 << ((value >> 4) & 0xf)), + _ => (1, 1), + } +} + +/// Writes the timing a VPS and a VUI share: one tick per frame, and an order count step per picture. +fn timing(w: &mut Rbsp, framerate: u32) { + w.bits(32, 1); // num_units_in_tick + w.bits(32, framerate.max(1)); // time_scale + w.flag(true); // poc_proportional_to_timing_flag + w.ue(0); // num_ticks_poc_diff_one_minus1 +} + +/// Returns the packed-header parameter and data buffers for `data`, which carries its own emulation prevention. +fn packed_header(kind: EncPackedHeaderType, data: &[u8]) -> [BufferType; 2] { + let (param, payload) = super::packed_header(kind, data); + [param, payload] +} + +/// Returns `general_level_idc` for the lowest Main tier level whose picture size and sample rate cover a `coded` picture at `framerate`. +/// +/// Only those limits, as for H.264: the level's maximum bitrate is not +/// checked, since the bitrate can change while the SPS cannot. +fn level((width, height): (u32, u32), framerate: u32) -> Option { + // H.265 Table A.8: (level_idc, MaxLumaPs, MaxLumaSr). + const LEVELS: [(u8, u64, u64); 13] = [ + (30, 36_864, 552_960), + (60, 122_880, 3_686_400), + (63, 245_760, 7_372_800), + (90, 552_960, 16_588_800), + (93, 983_040, 33_177_600), + (120, 2_228_224, 66_846_720), + (123, 2_228_224, 133_693_440), + (150, 8_912_896, 267_386_880), + (153, 8_912_896, 534_773_760), + (156, 8_912_896, 1_069_547_520), + (180, 35_651_584, 1_069_547_520), + (183, 35_651_584, 2_139_095_040), + (186, 35_651_584, 4_278_190_080), + ]; + let size = u64::from(width) * u64::from(height); + let rate = size * u64::from(framerate.max(1)); + // A.4.1: neither side may exceed sqrt(8 * MaxLumaPs). + let side = u64::from(width.max(height)).pow(2); + LEVELS + .into_iter() + .find(|&(_, max_size, max_rate)| size <= max_size && side <= 8 * max_size && rate <= max_rate) + .map(|(level, _, _)| level) +} + +/// Writes NAL unit type `kind` behind a four-byte start code, with emulation prevention. +fn nal(kind: u8, rbsp: &[u8]) -> Vec { + // nal_unit_header: forbidden_zero_bit, nal_unit_type, nuh_layer_id 0, nuh_temporal_id_plus1 1. + let mut out = vec![0, 0, 0, 1, kind << 1, 1]; + let mut zeros = 0; + for &byte in rbsp { + if zeros == 2 && byte <= 3 { + out.push(3); + zeros = 0; + } + out.push(byte); + zeros = if byte == 0 { zeros + 1 } else { 0 }; + } + out +} + +/// A raw byte sequence payload, written a field at a time. +#[derive(Default)] +struct Rbsp { + bytes: Vec, + /// The partial byte, its bits filled from the top. + current: u8, + /// How many bits of `current` are filled. + filled: u32, +} + +impl Rbsp { + /// Writes the low `count` bits of `value`, most significant first. + fn bits(&mut self, count: u32, value: u32) { + debug_assert!( + count <= 32 && (count == 32 || value >> count == 0), + "{value} in {count} bits" + ); + for bit in (0..count).rev() { + self.flag(value >> bit & 1 == 1); + } + } + + fn flag(&mut self, bit: bool) { + self.current |= (bit as u8) << (7 - self.filled); + self.filled += 1; + if self.filled == 8 { + self.bytes.push(self.current); + self.current = 0; + self.filled = 0; + } + } + + /// Writes `value` as an unsigned Exp-Golomb code, H.265 9.2. + fn ue(&mut self, value: u32) { + let coded = u64::from(value) + 1; + let length = 64 - coded.leading_zeros(); + self.bits(length - 1, 0); + for bit in (0..length).rev() { + self.flag(coded >> bit & 1 == 1); + } + } + + /// Writes `value` as a signed Exp-Golomb code, H.265 9.2.2. + fn se(&mut self, value: i32) { + let mapped = match value > 0 { + true => 2 * value.unsigned_abs() - 1, + false => 2 * value.unsigned_abs(), + }; + self.ue(mapped); + } + + /// Writes `rbsp_trailing_bits`: a one, then zeros to the byte boundary. + fn trailing(&mut self) { + self.flag(true); + while self.filled != 0 { + self.flag(false); + } + } + + /// Writes `byte_alignment` at the end of a slice header, which is the same bits. + fn byte_alignment(&mut self) { + self.trailing(); + } + + fn finish(self) -> Vec { + debug_assert_eq!(self.filled, 0, "an RBSP ends on a byte boundary"); + self.bytes + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn written(write: impl FnOnce(&mut Rbsp)) -> Vec { + let mut w = Rbsp::default(); + write(&mut w); + w.trailing(); + w.finish() + } + + #[test] + fn exp_golomb_codes_match_the_spec_table() { + // 9.2: 0 -> 1, 1 -> 010, 2 -> 011, 3 -> 00100, then the stop bit. + assert_eq!(written(|w| w.ue(0)), [0b1100_0000]); + assert_eq!(written(|w| w.ue(1)), [0b0101_0000]); + assert_eq!(written(|w| w.ue(2)), [0b0111_0000]); + assert_eq!(written(|w| w.ue(3)), [0b0010_0100]); + // 9.2.2: k > 0 maps to 2k - 1, k <= 0 to -2k. + assert_eq!(written(|w| w.se(1)), written(|w| w.ue(1))); + assert_eq!(written(|w| w.se(-1)), written(|w| w.ue(2))); + assert_eq!(written(|w| w.se(-26)), written(|w| w.ue(52))); + assert_eq!(written(|w| w.ue(u32::MAX - 1)).len(), 8, "63 bits and the stop bit"); + assert_eq!(written(|w| w.ue(u32::MAX)).len(), 9, "65 bits and the stop bit"); + } + + /// 7.4.2: two zeros then a byte below 4 take an emulation prevention byte, + /// and the zeros count again from the byte after it. + #[test] + fn start_code_patterns_are_escaped() { + assert_eq!( + nal(PPS, &[0, 0, 1, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0]), + [0, 0, 0, 1, 0x44, 1, 0, 0, 3, 1, 0, 0, 3, 0, 0, 4, 0, 0, 3, 0, 0, 3, 0] + ); + } + + #[test] + fn the_level_covers_the_picture_size_and_sample_rate() { + assert_eq!(level((640, 360), 30), Some(63)); + assert_eq!(level((1280, 720), 30), Some(93)); + assert_eq!(level((1920, 1080), 30), Some(120)); + assert_eq!(level((1920, 1080), 60), Some(123)); + assert_eq!(level((3840, 2160), 30), Some(150)); + assert_eq!(level((3840, 2160), 60), Some(153)); + // A sliver within level 3's area but longer than its widest side, and level 3.1's. + assert_eq!(level((4096, 64), 30), Some(120)); + assert_eq!(level((16384, 16384), 30), None); + } + + /// The coded size rounds up to whole minimum coding blocks, and the SPS crops it back. + #[test] + fn the_coded_size_is_whole_coding_blocks() { + let config = Config::new(852, 478, 30, 1_000_000, 30); + let stream = Stream::with_tools(Tools::FALLBACK, (1, 1), &config).unwrap(); + assert_eq!(stream.coded(), (864, 480)); + let tools = Tools { + log2_min_cb: 3, + ..Tools::FALLBACK + }; + assert_eq!(Stream::with_tools(tools, (1, 1), &config).unwrap().coded(), (856, 480)); + // radeonsi's 64 by 16 surface alignment. + assert_eq!( + Stream::with_tools(tools, (64, 16), &config).unwrap().coded(), + (896, 480) + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 5d33ab2..525ecc7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -42,6 +42,7 @@ pub use bindings::_VADRMPRIMESurfaceDescriptor__bindgen_ty_1 as VADRMPRIMESurfac pub use bindings::_VADRMPRIMESurfaceDescriptor__bindgen_ty_2 as VADRMPRIMESurfaceDescriptorLayer; pub use bindings::*; pub use buffer::*; +pub use codec::Codec; pub use color::{Color, Matrix}; pub use config::*; pub use context::*;