From fe4188444e49adc38c9c1cf97fcd96f5b0adc6b9 Mon Sep 17 00:00:00 2001 From: Frando Date: Thu, 24 Sep 2026 02:06:34 +0200 Subject: [PATCH 1/2] feat: import DMA-BUFs and post-process surfaces on the GPU This adds the two modules an encoder needs to take frames it did not allocate. They are also the "VPP wrapper" that moq-video's VAAPI quest lists as a prerequisite. `dmabuf::DmaBuf` describes a single-object DRM PRIME buffer by its DRM format, modifier and, where the producer knows it, its YUV color space. It maps the DRM format to the VA-API fourcc, which differs for packed RGB and for planar 4:2:0, and imports the buffer as a surface without copying. The object size comes from `lseek(SEEK_END)` on the descriptor, because iHD refuses a packed RGB import whose size is zero (it works the size out itself for YUV). `DmaBuf::from_prime` turns an export from `Surface::export_prime` back into one. `vpp::Processor` drives `VAEntrypointVideoProc`. A blit stretches the input over the output and converts to the output's format, so one call scales, converts RGB to NV12, converts between YUV color spaces, or re-tiles into the driver's preferred layout. The processor keeps one context per destination size, drops a context after a failed blit, and refuses surfaces from another display, whose ids it would otherwise misread as its own. `color::Color` names the matrix and range on each side of a conversion. `ExportedFrame::from_surface` exports a processed NV12 surface the way the decoder exports its pictures, so a consumer handles both the same way. Tests, which skip without a VA-API device: a retile round trip, and a half-dark, half-bright surface scaled to half size and read back to check it was scaled rather than cropped. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/color.rs | 93 +++++++++++ src/decode.rs | 29 ++++ src/dmabuf.rs | 301 +++++++++++++++++++++++++++++++++ src/lib.rs | 8 + src/vpp.rs | 454 ++++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 885 insertions(+) create mode 100644 src/color.rs create mode 100644 src/dmabuf.rs create mode 100644 src/vpp.rs diff --git a/src/color.rs b/src/color.rs new file mode 100644 index 0000000..b8323e9 --- /dev/null +++ b/src/color.rs @@ -0,0 +1,93 @@ +//! The YUV color space an encoded stream is labelled with. +//! +//! One value serves two jobs that have to agree: the encoder writes it into the +//! SPS video usability information, and the post-processor uses it as the +//! target of every conversion it runs for the encoder. Pixels the encoder +//! converts therefore match their label. Pixels it takes as they are (a CPU +//! upload, a YUV buffer with no color of its own) are labelled, not converted. + +use crate::bindings; + +/// The matrix that maps RGB to Y'CbCr. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Matrix { + /// ITU-R BT.601, the standard-definition matrix. + Bt601, + /// ITU-R BT.709, the high-definition matrix. + Bt709, +} + +/// A YUV color space: the matrix and whether samples use the full 8-bit range. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Color { + /// The RGB to Y'CbCr matrix. + pub matrix: Matrix, + /// Whether luma spans 0..=255 rather than 16..=235. + pub full_range: bool, +} + +impl Color { + /// BT.601 in limited range. + pub const BT601: Self = Self { + matrix: Matrix::Bt601, + full_range: false, + }; + /// BT.709 in limited range. + pub const BT709: Self = Self { + matrix: Matrix::Bt709, + full_range: false, + }; + + /// Returns the conventional space for a stream of `height` lines. + /// + /// Standard definition (576 lines or fewer) is BT.601, anything taller + /// BT.709, both in limited range. That is the guess most players make for + /// an unlabelled stream, though not all of them (ffmpeg's scaler takes + /// BT.601 at every size), which is why the stream is labelled. + pub fn infer(height: u32) -> Self { + match height <= 576 { + true => Self::BT601, + false => Self::BT709, + } + } + + /// Returns the H.264 VUI `(colour_primaries, transfer_characteristics, matrix_coefficients)`. + /// + /// BT.601 goes out as SMPTE 170M primaries and matrix (code point 6) with the + /// BT.709 transfer curve (1), since the two standards differ in primaries + /// and matrix but not in gamma. + pub(crate) fn vui(self) -> (u8, u8, u8) { + match self.matrix { + Matrix::Bt601 => (6, 1, 6), + Matrix::Bt709 => (1, 1, 1), + } + } + + /// Returns the `VAProcColorStandardType` of the matrix. + pub(crate) fn va_standard(self) -> u8 { + (match self.matrix { + Matrix::Bt601 => bindings::_VAProcColorStandardType_VAProcColorStandardBT601, + Matrix::Bt709 => bindings::_VAProcColorStandardType_VAProcColorStandardBT709, + }) as u8 + } + + /// Returns the `VA_SOURCE_RANGE_*` of the range. + pub(crate) fn va_range(self) -> u8 { + (match self.full_range { + true => bindings::VA_SOURCE_RANGE_FULL, + false => bindings::VA_SOURCE_RANGE_REDUCED, + }) as u8 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn inference_splits_at_standard_definition() { + assert_eq!(Color::infer(480), Color::BT601); + assert_eq!(Color::infer(576), Color::BT601); + assert_eq!(Color::infer(720), Color::BT709); + } +} diff --git a/src/decode.rs b/src/decode.rs index c24b811..be4d9f5 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -127,6 +127,35 @@ pub struct ExportedFrame { } impl ExportedFrame { + /// Exports `surface`, an NV12 picture from somewhere other than the decoder, the way the decoder exports its own. + /// + /// A picture the video processor produced, such as a scaled copy of a + /// captured frame, then reaches its consumer with the same descriptor and + /// the same [`download`](Self::download) fallback as a decoded one, so one + /// consumer handles both. `timestamp` is carried through untouched. + /// + /// # Errors + /// + /// Fails when the surface is not NV12, cannot be synced, or the driver will + /// not export it. + pub fn from_surface(surface: Surface<()>, timestamp: u64) -> anyhow::Result { + surface.sync().map_err(|e| anyhow!("surface sync: {e:?}"))?; + let (width, height) = surface.size(); + let descriptor = surface + .export_prime() + .map_err(|e| anyhow!("export the surface: {e:?}"))?; + if descriptor.fourcc != VA_FOURCC_NV12 { + anyhow::bail!("exported fourcc {:#010x}, expected NV12", descriptor.fourcc); + } + Ok(Self { + timestamp, + width, + height, + descriptor, + surface: Arc::new(surface), + }) + } + /// Reads this picture back to client memory as tightly packed NV12. /// /// The way back for a caller that took a descriptor and then found itself diff --git a/src/dmabuf.rs b/src/dmabuf.rs new file mode 100644 index 0000000..1fbf4b0 --- /dev/null +++ b/src/dmabuf.rs @@ -0,0 +1,301 @@ +//! Importing a Linux DMA-BUF as a VA surface, without copying it. +//! +//! A [`DmaBuf`] describes an allocation some other API made (a PipeWire +//! producer, a camera, a V4L2 device, another VA-API context that exported +//! its surface) well enough for the driver to wrap it as a [`Surface`]. The +//! surface reads the producer's memory directly: an encoder given one encodes +//! the producer's pixels, and a post-processor given one converts them, without +//! the frame passing through the CPU. +//! +//! The format is named the DRM way, which is what every producer of a DMA-BUF +//! speaks. [`va_fourcc`] maps it to the VA-API name, which differs for packed +//! RGB (DRM names the channels in little-endian word order, VA-API in byte +//! order) and for planar 4:2:0 (`YU12` against `I420`). + +use std::io::{Seek, SeekFrom}; +use std::os::fd::{AsRawFd, OwnedFd}; +use std::sync::Arc; + +use crate::{ + Color, Display, DrmPrimeSurfaceDescriptor, ExternalBufferDescriptor, MemoryType, Surface, UsageHint, + VADRMPRIMESurfaceDescriptor, VA_FOURCC_ARGB, VA_FOURCC_BGRA, VA_FOURCC_BGRX, VA_FOURCC_I420, VA_FOURCC_NV12, + VA_FOURCC_P010, VA_FOURCC_RGBA, VA_FOURCC_RGBX, VA_FOURCC_XRGB, VA_FOURCC_YUY2, VA_FOURCC_YV12, VA_RT_FORMAT_RGB32, + VA_RT_FORMAT_YUV420, VA_RT_FORMAT_YUV420_10, VA_RT_FORMAT_YUV422, +}; + +/// DRM fourcc codes for the formats [`va_fourcc`] maps. +/// +/// These are the `DRM_FORMAT_*` values from `drm_fourcc.h`, which the VA-API +/// headers do not carry. +pub mod drm { + const fn fourcc(code: &[u8; 4]) -> u32 { + u32::from_le_bytes(*code) + } + + /// Semi-planar 4:2:0, a luma plane and an interleaved Cb/Cr plane. + pub const NV12: u32 = fourcc(b"NV12"); + /// Semi-planar 4:2:0 with 10 bits per sample in 16-bit words. + pub const P010: u32 = fourcc(b"P010"); + /// Planar 4:2:0: luma, then Cb, then Cr. + pub const YUV420: u32 = fourcc(b"YU12"); + /// Planar 4:2:0: luma, then Cr, then Cb. + pub const YVU420: u32 = fourcc(b"YV12"); + /// Packed 4:2:2, `Y0 Cb Y1 Cr` in memory. + pub const YUYV: u32 = fourcc(b"YUYV"); + /// Packed 32-bit, `B G R X` in memory. + pub const XRGB8888: u32 = fourcc(b"XR24"); + /// Packed 32-bit, `B G R A` in memory. + pub const ARGB8888: u32 = fourcc(b"AR24"); + /// Packed 32-bit, `R G B X` in memory. + pub const XBGR8888: u32 = fourcc(b"XB24"); + /// Packed 32-bit, `R G B A` in memory. + pub const ABGR8888: u32 = fourcc(b"AB24"); + /// Packed 32-bit, `X R G B` in memory. + pub const BGRX8888: u32 = fourcc(b"BX24"); + /// Packed 32-bit, `A R G B` in memory. + pub const BGRA8888: u32 = fourcc(b"BA24"); +} + +/// Returns the `VA_FOURCC_*` naming the same layout as the DRM format `drm_format`. +/// +/// Returns `None` for a format this crate has no mapping for. +pub fn va_fourcc(drm_format: u32) -> Option { + Some(match drm_format { + drm::NV12 => VA_FOURCC_NV12, + drm::P010 => VA_FOURCC_P010, + drm::YUV420 => VA_FOURCC_I420, + drm::YVU420 => VA_FOURCC_YV12, + drm::YUYV => VA_FOURCC_YUY2, + drm::XRGB8888 => VA_FOURCC_BGRX, + drm::ARGB8888 => VA_FOURCC_BGRA, + drm::XBGR8888 => VA_FOURCC_RGBX, + drm::ABGR8888 => VA_FOURCC_RGBA, + drm::BGRX8888 => VA_FOURCC_XRGB, + drm::BGRA8888 => VA_FOURCC_ARGB, + _ => return None, + }) +} + +/// Returns the `VA_RT_FORMAT_*` a surface of the VA fourcc `fourcc` is allocated as. +/// +/// Returns `None` for a format this crate has no mapping for. Guessing the +/// render-target format wrong makes `vaCreateSurfaces` fail with an error that +/// names neither, so an unknown format is refused up front instead. +pub fn rt_format(fourcc: u32) -> Option { + match fourcc { + VA_FOURCC_NV12 | VA_FOURCC_I420 | VA_FOURCC_YV12 => Some(VA_RT_FORMAT_YUV420), + VA_FOURCC_P010 => Some(VA_RT_FORMAT_YUV420_10), + VA_FOURCC_YUY2 => Some(VA_RT_FORMAT_YUV422), + VA_FOURCC_BGRA | VA_FOURCC_BGRX | VA_FOURCC_RGBA | VA_FOURCC_RGBX | VA_FOURCC_ARGB | VA_FOURCC_XRGB => { + Some(VA_RT_FORMAT_RGB32) + } + _ => None, + } +} + +/// One plane of a DMA-BUF, within the single object holding them all. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Plane { + /// Byte offset of the plane from the start of the object. + pub offset: u32, + /// Bytes between adjacent rows of the plane. + pub pitch: u32, +} + +/// An externally allocated DMA-BUF, described well enough to import as a surface. +/// +/// One object holding every plane, which is what a VA-API decoder exports and +/// what a PipeWire producer negotiating a single modifier hands out. A buffer +/// split across several objects is not describable here. +/// +/// The descriptor owns the file descriptor and keeps it open for the life of the +/// surface built from it. libva does not take ownership (the driver resolves the +/// buffer during `vaCreateSurfaces` and leaves the fd alone), so holding it +/// costs one descriptor and removes the question of whether closing it early is +/// safe. +#[derive(Debug)] +pub struct DmaBuf { + /// Pixel format, a `DRM_FORMAT_*` code; see [`drm`]. + pub drm_format: u32, + /// Format modifier describing the tiling of the object. + pub modifier: u64, + /// Width of the picture in pixels. + /// + /// The visible width, not the allocated one: the padding lives in each + /// plane's pitch and offset. + pub width: u32, + /// Height of the picture in pixels, likewise the visible height. + pub height: u32, + /// Plane offsets and pitches, in format order. One to four. + pub planes: Vec, + /// The exported descriptor, held open until the surface is destroyed. + pub fd: OwnedFd, + /// The YUV color space of the pixels, where the producer knows it. + /// + /// Ignored for packed RGB. The encoder converts a YUV buffer whose space + /// differs from the stream's, and takes one with `None` to be in the + /// stream's space already. + pub color: Option, +} + +impl DmaBuf { + /// Describes a surface exported by [`Surface::export_prime`] for import elsewhere. + /// + /// The color is left unknown; set [`color`](Self::color) when the + /// producer knows it. + /// + /// # Errors + /// + /// Fails unless the export is one object holding one composed layer, which + /// is what `export_prime` asks the driver for. + pub fn from_prime(mut exported: DrmPrimeSurfaceDescriptor) -> anyhow::Result { + if exported.objects.len() != 1 || exported.layers.len() != 1 { + anyhow::bail!( + "expected one object and one layer, the export has {} and {}", + exported.objects.len(), + exported.layers.len() + ); + } + let object = exported.objects.remove(0); + let layer = &exported.layers[0]; + Ok(Self { + drm_format: layer.drm_format, + modifier: object.drm_format_modifier, + width: exported.width, + height: exported.height, + planes: (0..(layer.num_planes as usize).min(4)) + .map(|i| Plane { + offset: layer.offset[i], + pitch: layer.pitch[i], + }) + .collect(), + fd: object.fd, + color: None, + }) + } + + /// Returns the `VA_FOURCC_*` of this buffer's format, if it has one. + pub fn fourcc(&self) -> Option { + va_fourcc(self.drm_format) + } + + /// Imports this buffer as a surface on `display`, without copying it. + /// + /// The returned surface owns the buffer, so the descriptor stays open + /// exactly as long as the surface referring to it. `usage_hint` tells the + /// driver what the surface is for; some drivers only accept an import as an + /// encoder source when it says so. + /// + /// # Errors + /// + /// Fails when the format has no VA-API mapping, when the plane count is not + /// one to four, and when the driver refuses the import, which it does for a + /// modifier or pitch it cannot address. + pub fn import(self, display: &Arc, usage_hint: Option) -> anyhow::Result> { + let fourcc = self + .fourcc() + .ok_or_else(|| anyhow::anyhow!("no VA-API format for DRM format {:#010x}", self.drm_format))?; + let rt_format = + rt_format(fourcc).ok_or_else(|| anyhow::anyhow!("no VA-API render target format for {fourcc:#010x}"))?; + if self.planes.is_empty() || self.planes.len() > 4 { + anyhow::bail!("a DMA-BUF import takes one to four planes, got {}", self.planes.len()); + } + let (width, height) = (self.width, self.height); + + display + .create_surfaces(rt_format, Some(fourcc), width, height, usage_hint, vec![self]) + .map_err(|e| anyhow::anyhow!("import a {width}x{height} DMA-BUF as a surface: {e:?}"))? + .pop() + .ok_or_else(|| anyhow::anyhow!("vaCreateSurfaces returned no surface")) + } +} + +impl ExternalBufferDescriptor for DmaBuf { + const MEMORY_TYPE: MemoryType = MemoryType::DrmPrime2; + type DescriptorAttribute = VADRMPRIMESurfaceDescriptor; + + fn va_surface_attribute(&mut self) -> Self::DescriptorAttribute { + // `import` has checked the format and the plane count; an unmapped format + // reaching here anyway goes out as zero, which the driver refuses. + let mut descriptor = VADRMPRIMESurfaceDescriptor { + fourcc: self.fourcc().unwrap_or(0), + width: self.width, + height: self.height, + num_objects: 1, + num_layers: 1, + ..Default::default() + }; + + descriptor.objects[0].fd = self.fd.as_raw_fd(); + descriptor.objects[0].size = buffer_size(&self.fd); + descriptor.objects[0].drm_format_modifier = self.modifier; + + let layer = &mut descriptor.layers[0]; + layer.drm_format = self.drm_format; + layer.num_planes = self.planes.len().min(4) as u32; + for (index, plane) in self.planes.iter().take(4).enumerate() { + layer.object_index[index] = 0; + layer.offset[index] = plane.offset; + layer.pitch[index] = plane.pitch; + } + + descriptor + } +} + +/// Returns the size of the DMA-BUF behind `fd`, or zero when the kernel will not say. +/// +/// A DMA-BUF reports its size through `lseek(SEEK_END)`. Intel's iHD driver +/// refuses to import packed RGB with a size of zero, where for YUV formats it +/// works the size out from the layout; zero stays the fallback because it is +/// what the other drivers expect when the size is unknown. The seek moves the +/// file offset shared with the producer's descriptor. The kernel accepts only +/// the start and the end as a DMA-BUF offset (not even `SEEK_CUR`), so the +/// offset was at the start before, and seeking back there restores it. +fn buffer_size(fd: &OwnedFd) -> u32 { + let Ok(clone) = fd.try_clone() else { + return 0; + }; + let mut file = std::fs::File::from(clone); + let size = file + .seek(SeekFrom::End(0)) + .ok() + .and_then(|size| u32::try_from(size).ok()); + let _ = file.seek(SeekFrom::Start(0)); + size.unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn every_mapped_format_has_a_render_target() { + for format in [ + drm::NV12, + drm::P010, + drm::YUV420, + drm::YVU420, + drm::YUYV, + drm::XRGB8888, + drm::ARGB8888, + drm::XBGR8888, + drm::ABGR8888, + drm::BGRX8888, + drm::BGRA8888, + ] { + let fourcc = va_fourcc(format).expect("mapped"); + assert!(rt_format(fourcc).is_some(), "{format:#010x} has no render target"); + } + } + + #[test] + fn packed_rgb_maps_to_the_byte_order_name() { + // DRM XRGB8888 is a little-endian word, so memory holds B G R X, which + // VA-API calls BGRX. + assert_eq!(va_fourcc(drm::XRGB8888), Some(VA_FOURCC_BGRX)); + assert_eq!(va_fourcc(drm::ABGR8888), Some(VA_FOURCC_RGBA)); + assert_eq!(va_fourcc(drm::YUV420), Some(VA_FOURCC_I420)); + assert_eq!(va_fourcc(u32::from_le_bytes(*b"ZZZZ")), None); + } +} diff --git a/src/lib.rs b/src/lib.rs index 36b8439..5d33ab2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -31,10 +31,18 @@ pub mod codec; pub mod decode; pub mod encode; +// Written for this crate: DMA-BUF import, the color space an encoder labels its +// stream with, and the VA-API video post-processing entrypoint, which scales, +// converts, and re-tiles surfaces on the GPU. +pub mod color; +pub mod dmabuf; +pub mod vpp; + pub use bindings::_VADRMPRIMESurfaceDescriptor__bindgen_ty_1 as VADRMPRIMESurfaceDescriptorObject; pub use bindings::_VADRMPRIMESurfaceDescriptor__bindgen_ty_2 as VADRMPRIMESurfaceDescriptorLayer; pub use bindings::*; pub use buffer::*; +pub use color::{Color, Matrix}; pub use config::*; pub use context::*; pub use display::*; diff --git a/src/vpp.rs b/src/vpp.rs new file mode 100644 index 0000000..57f5cb0 --- /dev/null +++ b/src/vpp.rs @@ -0,0 +1,454 @@ +//! Video post-processing: the `VAEntrypointVideoProc` pipeline. +//! +//! The one operation is a blit from one surface into another, which the video +//! processor stretches to fit the destination and converts to the destination's +//! format on the way. That covers three jobs without the pixels leaving the GPU: +//! +//! - Scaling: blit into a smaller or larger surface. +//! - Color conversion: blit a packed RGB or YUYV surface into an NV12 one, the +//! format an encoder takes. [`Processor::blit`] takes the YUV space to write. +//! - Re-tiling: blit into a surface of the same size and format. The destination +//! is one the driver allocated for itself, so the blit rewrites the pixels +//! into whatever tiling the driver prefers for a fresh allocation. That is how +//! a DMA-BUF carrying a format modifier some other API refuses to import +//! becomes one it accepts. Measured on Intel Meteor Lake with iHD 26.1.5 (the +//! `retile_round_trip` test prints it), a decode target exports at +//! `0x100000000000009` (`I915_FORMAT_MOD_4_TILED`), which Mesa's Vulkan driver +//! lists, so that combination does not need it. Older iHD releases wrote +//! `Y_TILED` (`0x100000000000002`), which Mesa does not list. +//! +//! ```no_run +//! # fn example(buffer: moq_vaapi::dmabuf::DmaBuf) -> anyhow::Result<()> { +//! use moq_vaapi::vpp::Processor; +//! +//! let processor = Processor::new("/dev/dri/renderD128")?; +//! let retiled = processor.retile(buffer)?; +//! let exported = retiled.export_prime().map_err(|e| anyhow::anyhow!("export: {e:?}"))?; +//! # let _ = exported; +//! # Ok(()) +//! # } +//! ``` + +use std::cell::RefCell; +use std::collections::HashMap; +use std::path::Path; +use std::rc::Rc; +use std::sync::Arc; + +use crate::dmabuf::{rt_format, DmaBuf}; +use crate::{ + bindings, BufferType, Color, Config as VaConfig, Context, Display, Picture, ProcColorProperties, + ProcPipelineParameterBuffer, Surface, SurfaceMemoryDescriptor, UsageHint, VAConfigAttrib, VAConfigAttribType, + VAEntrypoint, VAProfile, +}; + +/// A video post-processor on one device. +/// +/// Holds the display, the VPP configuration, and one processing context per +/// destination size, which are what cost something to build. A caller +/// converting every frame of a stream, or of a few renditions, pays for each +/// context once. +pub struct Processor { + display: Arc, + // Drop order: the contexts before the config they were created from. + contexts: RefCell>>, + config: VaConfig, +} + +impl Processor { + /// Opens `device` and configures it for video post-processing. + /// + /// `device` is a DRM render node, e.g. `/dev/dri/renderD128`. Pick the one + /// matching the device that will consume the result: the blit only stays on + /// the GPU when both ends are the same device. + pub fn new>(device: P) -> anyhow::Result { + let device = device.as_ref(); + let display = + Display::open_drm_display(device).map_err(|e| anyhow::anyhow!("open DRM display {device:?}: {e}"))?; + Self::with_display(display) + } + + /// Opens the first device whose driver can post-process, if any. + /// + /// Mirrors [`Display::open`], and is what a caller with no opinion about + /// which GPU should use. A machine with more than one wants + /// [`Processor::new`] or [`Processor::with_display`] instead. + pub fn open() -> anyhow::Result { + for device in crate::DrmDeviceIterator::default() { + if let Ok(processor) = Self::new(&device) { + return Ok(processor); + } + } + Err(anyhow::anyhow!("no DRM device supports video post-processing")) + } + + /// Configures an already open display for video post-processing. + /// + /// Sharing the display with a decoder or an encoder is what lets a surface + /// move between them: a surface belongs to the display that created it. + /// + /// # Errors + /// + /// Fails when the device has no `VAEntrypointVideoProc`, which is how a + /// caller finds out before the first frame rather than during it. + pub fn with_display(display: Arc) -> anyhow::Result { + // Asking for the attribute is the probe: a device without the VPP + // entrypoint fails here instead of at vaCreateConfig, with an error that + // says which entrypoint was missing. + let mut attrs = [VAConfigAttrib { + type_: VAConfigAttribType::VAConfigAttribRTFormat, + value: 0, + }]; + display + .get_config_attributes( + VAProfile::VAProfileNone, + VAEntrypoint::VAEntrypointVideoProc, + &mut attrs, + ) + .map_err(|e| anyhow::anyhow!("query VPP config attributes: {e:?}"))?; + + let config = display + .create_config( + attrs.to_vec(), + VAProfile::VAProfileNone, + VAEntrypoint::VAEntrypointVideoProc, + ) + .map_err(|e| anyhow::anyhow!("create VPP config: {e:?}"))?; + + Ok(Self { + display, + contexts: RefCell::new(HashMap::new()), + config, + }) + } + + /// Returns the display this processor runs on, to share with a decoder or an encoder. + pub fn display(&self) -> &Arc { + &self.display + } + + /// Imports an exported DMA-BUF as a surface on this processor's display, without copying it. + /// + /// The returned surface owns `buffer`, so the descriptor stays open exactly + /// as long as the surface referring to it. + /// + /// # Errors + /// + /// See [`DmaBuf::import`]. + pub fn import(&self, buffer: DmaBuf) -> anyhow::Result> { + buffer.import(&self.display, Some(UsageHint::USAGE_HINT_VPP_READ)) + } + + /// Blits `input` into `output`, scaling to the size of `output` and converting to its format. + /// + /// `input_color` and `output_color` name the YUV space on each side. Leave + /// the input side `None` for an RGB input, which has no matrix; set the + /// output side to the space a YUV destination should be written in. A side + /// left `None` is up to the driver, which is right for a blit that only + /// re-tiles, and for a scale that names neither side. Returns once the + /// destination holds the result. + /// + /// # Errors + /// + /// Fails when either surface belongs to another display, and when the + /// driver refuses the blit (a format pair or scale it cannot do). + pub fn blit( + &self, + input: &Surface, + output: &Surface, + input_color: Option, + output_color: Option, + ) -> anyhow::Result<()> + where + D: SurfaceMemoryDescriptor, + E: SurfaceMemoryDescriptor, + { + // A surface id means something only on the display that made it, and + // the driver would read another display's id as one of its own. + if !Arc::ptr_eq(input.display(), &self.display) || !Arc::ptr_eq(output.display(), &self.display) { + anyhow::bail!("a VPP blit between surfaces of another display"); + } + let size = output.size(); + let context = self.context(size)?; + let result = run_blit(&context, input, output, input_color, output_color); + if result.is_err() { + // A picture that failed half way can leave the context mid-picture, + // so the next blit of this size starts on a fresh one. + self.contexts.borrow_mut().remove(&size); + } + result + } + + /// Blits `input` into a new driver-allocated surface of `fourcc`, `width` and `height`. + /// + /// The destination is hinted as post-processing output that will be + /// exported, so [`Surface::export_prime`] works on it. See + /// [`blit`](Self::blit) for the colors. + /// + /// # Errors + /// + /// Fails for a format with no render target mapping, when the destination + /// cannot be allocated, and as [`blit`](Self::blit) does. + pub fn process( + &self, + input: &Surface, + fourcc: u32, + (width, height): (u32, u32), + input_color: Option, + output_color: Option, + ) -> anyhow::Result> { + let rt_format = rt_format(fourcc).ok_or_else(|| anyhow::anyhow!("no RT format for fourcc {fourcc:#x}"))?; + let output = self + .display + .create_surfaces( + rt_format, + Some(fourcc), + width, + height, + Some(UsageHint::USAGE_HINT_VPP_WRITE | UsageHint::USAGE_HINT_EXPORT), + vec![()], + ) + .map_err(|e| anyhow::anyhow!("allocate a {width}x{height} VPP output surface: {e:?}"))? + .pop() + .ok_or_else(|| anyhow::anyhow!("vaCreateSurfaces returned no surface"))?; + + self.blit(input, &output, input_color, output_color)?; + Ok(output) + } + + /// Imports `buffer` and blits it into a driver-allocated surface of the same format and size. + /// + /// This is the whole re-tile. Export the result with + /// [`Surface::export_prime`] and read the modifier off the object to see what + /// the driver chose. + /// Destroying the returned surface does not invalidate an export taken from + /// it, since the exported descriptor holds its own reference on the + /// underlying buffer. + pub fn retile(&self, buffer: DmaBuf) -> anyhow::Result> { + let fourcc = buffer + .fourcc() + .ok_or_else(|| anyhow::anyhow!("no VA-API format for DRM format {:#010x}", buffer.drm_format))?; + let input = self.import(buffer)?; + self.process(&input, fourcc, input.size(), None, None) + } + + /// Returns the processing context for a destination of `size`, building it on first use. + fn context(&self, (width, height): (u32, u32)) -> anyhow::Result> { + let mut contexts = self.contexts.borrow_mut(); + if let Some(context) = contexts.get(&(width, height)) { + return Ok(Rc::clone(context)); + } + // No render targets: a VPP context takes its destination per picture, and + // binding it to one surface would make every other destination a rebuild. + let context = self + .display + .create_context::<()>(&self.config, width, height, None, true) + .map_err(|e| anyhow::anyhow!("create a {width}x{height} VPP context: {e:?}"))?; + contexts.insert((width, height), Rc::clone(&context)); + Ok(context) + } +} + +/// Runs one blit on `context` and waits for it to finish. +fn run_blit( + context: &Rc, + input: &Surface, + output: &Surface, + input_color: Option, + output_color: Option, +) -> anyhow::Result<()> +where + D: SurfaceMemoryDescriptor, + E: SurfaceMemoryDescriptor, +{ + let parameters = pipeline(input.id(), input_color, output_color); + let buffer = context + .create_buffer(BufferType::ProcPipelineParameter(parameters)) + .map_err(|e| anyhow::anyhow!("create VPP pipeline buffer: {e:?}"))?; + + let mut picture = Picture::new(0, Rc::clone(context), output); + picture.add_buffer(buffer); + picture + .begin() + .map_err(|e| anyhow::anyhow!("vaBeginPicture: {e:?}"))? + .render() + .map_err(|e| anyhow::anyhow!("vaRenderPicture: {e:?}"))? + .end() + .map_err(|e| anyhow::anyhow!("vaEndPicture: {e:?}"))? + .sync() + .map_err(|(e, _)| anyhow::anyhow!("vaSyncSurface: {e:?}"))?; + Ok(()) +} + +/// Returns a pipeline parameter buffer describing a plain blit of `surface`. +/// +/// No regions (so the whole input stretched over the whole output), no +/// filters, no references, no rotation or mirroring, and each side's color +/// standard and range as given, or the driver's default where not. +fn pipeline( + surface: bindings::VASurfaceID, + input_color: Option, + output_color: Option, +) -> ProcPipelineParameterBuffer { + let side = |color: Option| match color { + Some(color) => ( + color.va_standard(), + ProcColorProperties::new(0, color.va_range(), 0, 0, 0), + ), + None => (0, ProcColorProperties::default()), + }; + let (input_standard, input_properties) = side(input_color); + let (output_standard, output_properties) = side(output_color); + ProcPipelineParameterBuffer::new( + surface, + None, + input_standard, + None, + 0, + output_standard, + 0, + 0, + None, + None, + None, + 0, + None, + 0, + None, + 0, + 0, + input_properties, + output_properties, + 0, + None, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{VA_FOURCC_NV12, VA_RT_FORMAT_YUV420}; + + /// A driver-allocated NV12 surface survives a round trip through the + /// post-processor, and the modifier the driver picked is printed. + /// + /// Skips without a VA-API device, so it is a no-op on a builder and real + /// coverage on a machine with a GPU. What it asserts is that the whole + /// sequence runs (import an exported DMA-BUF, blit, export the result); the + /// modifier the driver picks is hardware policy, so it is printed rather + /// than asserted. + #[test] + fn retile_round_trip() { + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + let Ok(processor) = Processor::with_display(display) else { + eprintln!("skipping: device has no video post-processing entrypoint"); + return; + }; + + // A decode-hinted surface, which is the allocation whose tiling is the + // problem a re-tile exists for. + let mut surfaces = processor + .display() + .create_surfaces( + VA_RT_FORMAT_YUV420, + Some(VA_FOURCC_NV12), + 640, + 480, + Some(UsageHint::USAGE_HINT_DECODER | UsageHint::USAGE_HINT_EXPORT), + vec![()], + ) + .expect("allocate a decode surface"); + let decoded = surfaces.pop().expect("one surface"); + let exported = decoded.export_prime().expect("export the decode surface"); + let source_modifier = exported.objects[0].drm_format_modifier; + let (fourcc, size, planes) = ( + exported.fourcc, + (exported.width, exported.height), + exported.layers[0].num_planes, + ); + let buffer = DmaBuf::from_prime(exported).expect("a composed export is one object"); + + let retiled = processor.retile(buffer).expect("re-tile the decode surface"); + let output = retiled.export_prime().expect("export the re-tiled surface"); + let target_modifier = output.objects[0].drm_format_modifier; + + eprintln!("modifier {source_modifier:#x} -> {target_modifier:#x}"); + assert_eq!(output.fourcc, fourcc, "the blit preserves the pixel format"); + assert_eq!((output.width, output.height), size); + assert_eq!(output.objects.len(), 1, "a composed export is one object"); + assert_eq!( + output.layers[0].num_planes, planes, + "the blit preserves the plane count" + ); + } + /// A surface scaled by the processor exports and reads back at its new + /// size, with the picture scaled rather than cropped. + #[test] + fn a_scaled_surface_reads_back_at_its_new_size() { + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + let Ok(processor) = Processor::with_display(display) else { + eprintln!("skipping: device has no video post-processing entrypoint"); + return; + }; + + // Left half dark, right half bright, so a crop and a scale differ. + let (width, height) = (640u32, 480u32); + let source = processor + .display() + .create_surfaces( + VA_RT_FORMAT_YUV420, + Some(VA_FOURCC_NV12), + width, + height, + Some(UsageHint::USAGE_HINT_EXPORT), + vec![()], + ) + .expect("allocate a surface") + .pop() + .expect("one surface"); + let format = processor + .display() + .query_image_formats() + .expect("query image formats") + .into_iter() + .find(|f| f.fourcc == VA_FOURCC_NV12) + .expect("an NV12 image format"); + { + let mut image = + crate::Image::create_from(&source, format, (width, height), (width, height)).expect("create an image"); + let va_image = *image.image(); + let data = image.as_mut(); + for y in 0..height as usize { + let row = va_image.offsets[0] as usize + y * va_image.pitches[0] as usize; + for x in 0..width as usize { + data[row + x] = if x < width as usize / 2 { 40 } else { 200 }; + } + } + for y in 0..height as usize / 2 { + let row = va_image.offsets[1] as usize + y * va_image.pitches[1] as usize; + data[row..row + width as usize].fill(128); + } + } + + let scaled = processor + .process(&source, VA_FOURCC_NV12, (width / 2, height / 2), None, None) + .expect("scale the surface"); + let exported = crate::decode::ExportedFrame::from_surface(scaled, 7).expect("export the scaled surface"); + assert_eq!((exported.width, exported.height), (width / 2, height / 2)); + assert_eq!(exported.timestamp, 7); + let frame = exported.download().expect("read the scaled surface back"); + let (w, h) = ((width / 2) as usize, (height / 2) as usize); + assert_eq!(frame.data.len(), w * h * 3 / 2); + let (left, right) = (frame.data[h / 2 * w + w / 8], frame.data[h / 2 * w + w * 7 / 8]); + assert!( + left.abs_diff(40) < 4 && right.abs_diff(200) < 4, + "read back {left} and {right}" + ); + } +} From 8ae97afefb7fcbd22ac643a6bfee4bc4cec17308 Mon Sep 17 00:00:00 2001 From: Frando Date: Thu, 24 Sep 2026 02:06:34 +0200 Subject: [PATCH 2/2] feat: encode from a DMA-BUF, and change the bitrate between frames `Encoder::encode_dmabuf` takes a frame that a capture or a decoder already holds on the GPU. An NV12 buffer at the encoder's size, a whole number of macroblocks, and in the stream's color space is imported and encoded in place. Anything else importable is blitted through the video processor into the encoder's own surface, scaled to its size and converted to its color space. Examples are packed RGB from a screen capture, YUYV, NV12 at another size or in another space, and a 1080-row buffer. The last case matters because the encoder reads 1088 rows, and a producer that puts the chroma plane right after the last visible row has nothing there to read. The encode body now takes any surface, so both paths and the CPU upload share it. `Encoder::set_bitrate` retunes the rate controller from the next frame. The rate control parameters already went out with every frame, so no IDR is needed. The next frame sets the rate control reset flag, and the coded buffer grows when the new rate could overflow it. Also: - `Config::color` is written into the SPS as the VUI colour description. It defaults to BT.601 up to 576 lines and BT.709 above. - Packed SPS and PPS headers are negotiated through `VAConfigAttribEncPackedHeaders` where the driver offers them. iHD took them without being asked, but a driver that follows the spec would write its own SPS, without our VUI and level. - The level follows the frame size and macroblock rate instead of being fixed at 4. Level 4 was too low for 1440p and 4K streams. - frame_num wraps at MaxFrameNum, and the slice header carries only the low bits of the picture order count. Before, a group longer than 2048 frames depended on the driver masking them. - Where the device has no `VAEntrypointEncSlice` for H.264 Main, the encoder opens the low-power one, which is the only one some Intel parts expose. - `encode_nv12` rejects a short buffer instead of panicking on it. Validated on Intel Meteor Lake (iHD 26.1.5). The tests decode with ffmpeg and skip without it or without a device: - CPU and DMA-BUF input decode at 35 dB luma PSNR or better. - A scaled DMA-BUF decodes at 25 dB or better against a point-sampled reference. - A 1080p buffer is encoded through the blit. - BGRX and a BT.709 NV12 buffer convert into BT.601 and BT.709 streams, with every channel within 6 code values. - The VUI label is checked with ffprobe. - IDRs follow the group length and restart after a forced keyframe. - A refused import leaves the encoder usable. - The entrypoint choice is covered. - A bitrate drop from 2 Mbps to 250 kbps shows up in the picture sizes, with no IDR. The low-power entrypoint itself is untested, since this device exposes the full one. Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 30 +- src/codec/h264/parser.rs | 23 + src/encode.rs | 913 +++++++++++++++++++++++++++++++++++++-- 3 files changed, 921 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index 0c3b481..7a7aaf8 100644 --- a/README.md +++ b/README.md @@ -2,10 +2,9 @@ > **⚠️ AI GENERATED.** This crate was written by an AI agent (Claude), using the > repositories below as references. How closely individual files track upstream -> varies and the encoder's emitted bitstream has not been validated at playback — -> treat the whole thing as derived, unverified work and review before relying on -> it. The decoder is the exception: it has been checked against a software -> reference on Intel hardware (see below). +> varies, so treat the whole thing as derived work and review before relying on +> 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), derived from @@ -21,8 +20,22 @@ pulled in as a git dependency. ## What it does -- **H.264 encode** over VA-API: tightly-packed NV12 in, an Annex-B elementary - stream out (packed SPS/PPS + slice headers, low-latency IPPP, rate control). +- **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: + 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 + encoder's surface on the GPU. NV12 and BGRX input are tested; the other + 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. +- **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 + way. - **H.264 decode** over VA-API: one Annex-B access unit in, tightly-packed NV12 out, with in-stream parameter sets, a conformant DPB (reference marking, reordering, frame_num gaps), and mid-stream resolution changes. Progressive @@ -45,8 +58,9 @@ pulled in as a git dependency. ## Layout - `src/` — libva bindings (`bindings`, `display`, `surface`, `buffer`, ...), the - H.264 bitstream layer (`bitstream_utils`, `codec::h264`), and the thin encode - and decode drivers (`encode`, `decode`). + H.264 bitstream layer (`bitstream_utils`, `codec::h264`), 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`). - `build.rs` + `bindgen_gen.rs` + `libva-wrapper.h` — bindgen setup. diff --git a/src/codec/h264/parser.rs b/src/codec/h264/parser.rs index 8927243..74ff755 100644 --- a/src/codec/h264/parser.rs +++ b/src/codec/h264/parser.rs @@ -1161,6 +1161,29 @@ impl SpsBuilder { self } + /// Labels the stream's color space in the VUI. + /// + /// `video_format` goes out as 5 (unspecified), the value a decoder infers + /// when the field is absent, so this only adds the range and the color + /// description. + pub fn video_signal_type( + mut self, + full_range: bool, + colour_primaries: u8, + transfer_characteristics: u8, + matrix_coefficients: u8, + ) -> Self { + self = self.vui_parameters_present(); + self.0.vui_parameters.video_signal_type_present_flag = true; + self.0.vui_parameters.video_format = 5; + self.0.vui_parameters.video_full_range_flag = full_range; + self.0.vui_parameters.colour_description_present_flag = true; + self.0.vui_parameters.colour_primaries = colour_primaries; + self.0.vui_parameters.transfer_characteristics = transfer_characteristics; + self.0.vui_parameters.matrix_coefficients = matrix_coefficients; + self + } + pub fn bitstream_restrictions_present(mut self) -> Self { if self.0.vui_parameters.bitstream_restriction_flag { return self; diff --git a/src/encode.rs b/src/encode.rs index 13f923d..fcbe40d 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -16,10 +16,20 @@ //! Low-latency only: IPPP (no B-frames), one reference frame, matching the //! VideoToolbox / Media Foundation / NVENC backends in moq-video. //! -//! NOTE: compile-verified against libva 1.23, but the emitted bitstream has not -//! been validated at playback. The per-frame param population, the reconstructed -//! surface / reference rotation, and NV12 stride handling are the spots most -//! likely to need hardware tuning. +//! Three ways in, cheapest first: +//! - [`Encoder::encode_dmabuf`] with an NV12 DMA-BUF at the encoder's size +//! encodes the producer's memory in place. +//! - [`Encoder::encode_dmabuf`] with any other importable DMA-BUF (packed RGB +//! from a screen capture, YUYV from a camera, NV12 at another size) converts +//! and scales it into the encoder's own surface with the video processor, on +//! the GPU. +//! - [`Encoder::encode_nv12`] uploads CPU pixels. +//! +//! Checked on Intel Meteor Lake with the iHD driver by decoding with ffmpeg's +//! software decoder: CPU and DMA-BUF input come back above 35 dB luma PSNR, +//! and BGRX converts to BT.601 and BT.709 with every channel within 6 code +//! values of the reference. The low-power entrypoint is untested, since that +//! device exposes the full one. use std::path::PathBuf; use std::rc::Rc; @@ -28,13 +38,17 @@ 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::dmabuf::DmaBuf; +use crate::vpp::Processor; use crate::{ - BufferType, Config as VaConfig, Context, Display, EncCodedBuffer, EncMiscParameter, EncMiscParameterFrameRate, - EncMiscParameterRateControl, EncPackedHeaderParameter, EncPackedHeaderType, EncPictureParameter, - EncPictureParameterBufferH264, EncSequenceParameter, EncSequenceParameterBufferH264, EncSliceParameter, - EncSliceParameterBufferH264, H264EncFrameCropOffsets, H264EncPicFields, H264EncSeqFields, H264VuiFields, Image, - MappedCodedBuffer, Picture, PictureH264, RcFlags, Surface, UsageHint, VAConfigAttrib, VAConfigAttribType, - VAEntrypoint, VAProfile, VA_FOURCC_NV12, VA_INVALID_ID, VA_PICTURE_H264_INVALID, + BufferType, Color, Config as VaConfig, Context, Display, EncCodedBuffer, EncMiscParameter, + EncMiscParameterFrameRate, EncMiscParameterRateControl, EncPackedHeaderParameter, EncPackedHeaderType, + EncPictureParameter, EncPictureParameterBufferH264, EncSequenceParameter, EncSequenceParameterBufferH264, + 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, }; @@ -59,7 +73,7 @@ const MAX_QP: u8 = 36; /// /// This matches cros-codecs. The old `framerate * 1000` value confused frames /// per second with seconds and gave a 30 fps encoder a 30-second window. -const RATE_CONTROL_WINDOW_MS: u32 = 1_500; +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; @@ -75,6 +89,11 @@ pub struct Config { pub gop: u32, /// DRM render node to open (e.g. `/dev/dri/renderD128`). pub device: PathBuf, + /// Color space written into the SPS, and the target of a conversion from RGB. + /// + /// Defaults to [`Color::infer`] for the height, which is what a decoder + /// assumes for an unlabelled stream. + pub color: Color, } impl Config { @@ -86,6 +105,7 @@ impl Config { bitrate, gop, device: PathBuf::from("/dev/dri/renderD128"), + color: Color::infer(height), } } } @@ -93,7 +113,7 @@ impl Config { /// Per-frame reference metadata. (Mirrors cros-codecs `DpbEntryMeta`.) #[derive(Copy, Clone)] struct FrameMeta { - poc: u16, + poc: u32, frame_num: u32, } @@ -114,8 +134,16 @@ pub struct Encoder { counter: u32, + /// The video processor converting and scaling DMA-BUF input, opened on the + /// first frame that needs it. + processor: Option, + /// Set by a bitrate change, and cleared by the next frame, which asks the + /// driver to restart its rate controller at the new target. + reset_rate_control: bool, + // Drop order: encode resources before the display they belong to. coded: EncCodedBuffer, + coded_size: usize, context: Rc, _va_config: VaConfig, display: Arc, @@ -126,7 +154,8 @@ impl Encoder { let display = Display::open_drm_display(&config.device) .map_err(|e| anyhow::anyhow!("open DRM display {:?}: {e:?}", config.device))?; - let attrs = vec![ + let entrypoint = encode_entrypoint(&display)?; + let mut attrs = vec![ VAConfigAttrib { type_: VAConfigAttribType::VAConfigAttribRTFormat, value: VA_RT_FORMAT_YUV420, @@ -136,15 +165,27 @@ 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); + 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"); + } let va_config = display - .create_config(attrs, VAProfile::VAProfileH264Main, VAEntrypoint::VAEntrypointEncSlice) + .create_config(attrs, VAProfile::VAProfileH264Main, entrypoint) .map_err(|e| anyhow::anyhow!("create VA config: {e:?}"))?; let context = display .create_context::<()>(&va_config, config.width, config.height, None, true) .map_err(|e| anyhow::anyhow!("create VA context: {e:?}"))?; - let coded_size = (config.bitrate as usize / 4).max(1_500_000); + let coded_size = coded_size(config.bitrate); let coded = context .create_enc_coded(coded_size) .map_err(|e| anyhow::anyhow!("create coded buffer: {e:?}"))?; @@ -157,11 +198,12 @@ impl Encoder { let height_mbs = sps.pic_height_in_map_units_minus1 + 1; log::info!( - "opened VA-API H.264 encoder: {}x{} @ {}fps, {} bps", + "opened VA-API H.264 encoder: {}x{} @ {}fps, {} bps, low power: {}", config.width, config.height, config.framerate, - config.bitrate + config.bitrate, + entrypoint == VAEntrypoint::VAEntrypointEncSliceLP, ); Ok(Self { config, @@ -173,16 +215,145 @@ impl Encoder { recon, reference: None, counter: 0, + processor: None, + reset_rate_control: false, coded, + coded_size, context, _va_config: va_config, display, }) } - /// Encode one tightly-packed NV12 frame, returning Annex-B H.264 (with inline + /// Encodes one tightly-packed NV12 frame, returning Annex-B H.264 (with inline /// SPS/PPS on each IDR). `keyframe` forces an IDR. + /// + /// # Errors + /// + /// Fails when `nv12` is shorter than a frame at the encoder's size, and when + /// the upload or the encode fails. pub fn encode_nv12(&mut self, nv12: &[u8], keyframe: bool) -> anyhow::Result> { + let (width, height) = (self.config.width as usize, self.config.height as usize); + let expected = width * height + width * height.div_ceil(2); + if nv12.len() < expected { + anyhow::bail!("NV12 frame is {} bytes, {width}x{height} needs {expected}", nv12.len()); + } + let input = self.input.take().expect("input surface present"); + let result = upload_nv12(&self.display, &input, self.config.width, self.config.height, nv12) + .and_then(|()| self.encode_surface(&input, keyframe)); + self.input = Some(input); + result + } + + /// Encodes one frame held in a DMA-BUF, returning Annex-B H.264 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 + /// 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. + /// + /// The macroblock condition is not a formality: the encoder reads whole + /// macroblocks, so it reads 1088 rows of a 1080-row picture, and a producer + /// that places the chroma plane right after the last visible luma row + /// (PipeWire does) has nothing there to read. + /// + /// The producer must have finished writing the buffer; this does not wait + /// on its fence. + /// + /// # Errors + /// + /// Fails when the driver refuses to import the buffer (an unsupported + /// modifier or layout), when the device has no video processor and the + /// buffer needs one, and when the encode itself fails. The encoder stays + /// usable after a refused import, so a caller can fall back to + /// [`encode_nv12`](Self::encode_nv12) for the same frame. + pub fn encode_dmabuf(&mut self, buffer: DmaBuf, keyframe: bool) -> anyhow::Result> { + let fourcc = buffer + .fourcc() + .ok_or_else(|| anyhow::anyhow!("no VA-API format for DRM format {:#010x}", buffer.drm_format))?; + let color = self.config.color; + 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.color.is_none_or(|own| own == color); + if in_place { + let source = buffer.import(&self.display, Some(UsageHint::USAGE_HINT_ENCODER))?; + return self.encode_surface(&source, keyframe); + } + + let input_color = match is_rgb(fourcc) { + true => None, + false => Some(buffer.color.unwrap_or(color)), + }; + let source = buffer.import(&self.display, Some(UsageHint::USAGE_HINT_VPP_READ))?; + let input = self.input.take().expect("input surface present"); + let result = self + .processor() + .and_then(|processor| processor.blit(&source, &input, input_color, Some(color))) + .and_then(|()| self.encode_surface(&input, keyframe)); + self.input = Some(input); + result + } + + /// Changes the target bitrate, in bits per second, from the next frame on. + /// + /// The rate control parameters go out with every frame, so this neither + /// forces an IDR nor rebuilds the session; the next frame also asks the + /// driver to reset its rate controller. It then converges over its window, + /// [`RATE_CONTROL_WINDOW_MS`]. + /// + /// The level in the SPS is not revisited, so a rate far above the one the + /// stream opened with can exceed what that level allows. + /// + /// # Errors + /// + /// Fails for a bitrate of zero, and when a larger coded buffer for the new + /// rate cannot be allocated, in which case the old rate stays. + pub fn set_bitrate(&mut self, bitrate: u32) -> anyhow::Result<()> { + if bitrate == 0 { + anyhow::bail!("a bitrate of zero"); + } + let size = coded_size(bitrate); + if size > self.coded_size { + // A higher rate can produce a picture larger than the buffer sized for + // the old one, and a truncated picture is corrupt rather than short. + self.coded = self + .context + .create_enc_coded(size) + .map_err(|e| anyhow::anyhow!("create coded buffer: {e:?}"))?; + self.coded_size = size; + } + if bitrate != self.config.bitrate { + self.reset_rate_control = true; + } + self.config.bitrate = bitrate; + Ok(()) + } + + /// Returns the configuration in effect, including a bitrate changed since opening. + pub fn config(&self) -> &Config { + &self.config + } + + /// Returns the video processor, opening it on first use. + fn processor(&mut self) -> anyhow::Result<&Processor> { + if self.processor.is_none() { + self.processor = Some(Processor::with_display(Arc::clone(&self.display))?); + } + Ok(self.processor.as_ref().expect("opened above")) + } + + /// Encodes the picture in `source`, which must be NV12 at the encoder's size. + fn encode_surface( + &mut self, + source: &Surface, + keyframe: bool, + ) -> 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. @@ -190,12 +361,12 @@ impl Encoder { self.reference = None; } - let input = self.input.take().expect("input surface present"); - upload_nv12(&self.display, &input, self.config.width, self.config.height, nv12)?; - + // 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 * 2) & 0xffff) as u16, - frame_num: self.counter, + poc: self.counter.wrapping_mul(2), + frame_num: self.counter % LIMIT, }; let recon_idx = (self.counter as usize) % self.recon.len(); @@ -204,7 +375,7 @@ impl Encoder { .slice_type(slice_type) .first_mb_in_slice(0) .frame_num(meta.frame_num as u16) - .pic_order_cnt_lsb(meta.poc) + .pic_order_cnt_lsb((meta.poc % (LIMIT * 2)) as u16) .build(); let num_macroblocks = self.width_mbs as u32 * self.height_mbs as u32; @@ -227,7 +398,7 @@ impl Encoder { ); let slice_param = build_enc_slice_param(&self.pps, &header, reference, num_macroblocks); - let mut picture = Picture::new(meta.frame_num as u64, Rc::clone(&self.context), input); + let mut picture = Picture::new(meta.frame_num as u64, Rc::clone(&self.context), source); // VA-API spec buffer order: sequence, picture, slice, then packed headers, // then rate-control misc params. @@ -253,7 +424,7 @@ impl Encoder { 0, // basic_unit_size // Do not let rate control drop a frame to meet the bitrate. A live // encoder has already admitted this frame into the media timeline. - RcFlags::new(0, 1, 0, 0, 0, 0, 0, 0, 0), + RcFlags::new(u32::from(self.reset_rate_control), 1, 0, 0, 0, 0, 0, 0, 0), 0, // icq_quality_factor MAX_QP as u32, 0, // quality_factor @@ -267,24 +438,18 @@ impl Encoder { let picture = picture.begin().map_err(|e| anyhow::anyhow!("picture begin: {e:?}"))?; let picture = picture.render().map_err(|e| anyhow::anyhow!("picture render: {e:?}"))?; let picture = picture.end().map_err(|e| anyhow::anyhow!("picture end: {e:?}"))?; - // Sync (PictureEnd -> PictureSync) so the surface is reclaimable and the - // coded buffer is ready to read. - let picture = picture + // Sync (PictureEnd -> PictureSync) so the source is no longer read and + // the coded buffer is ready. + picture .sync() .map_err(|(e, _)| anyhow::anyhow!("picture sync: {e:?}"))?; - // Reclaim the source surface for the next frame. - self.input = Some( - picture - .take_surface() - .map_err(|_| anyhow::anyhow!("reclaim input surface (still referenced)"))?, - ); - // The reconstructed surface syncs implicitly; read the coded bitstream. let bitstream = self.read_coded()?; self.reference = Some((recon_idx, meta)); self.counter += 1; + self.reset_rate_control = false; Ok(bitstream) } @@ -318,6 +483,83 @@ impl Encoder { } } +/// Returns the entrypoint to encode H.264 Main 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 { + 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")) +} + +/// Returns the encode entrypoint to use out of those a device offers: the full one when present. +fn pick_entrypoint(offered: &[VAEntrypoint::Type]) -> Option { + [VAEntrypoint::VAEntrypointEncSlice, VAEntrypoint::VAEntrypointEncSliceLP] + .into_iter() + .find(|entrypoint| offered.contains(entrypoint)) +} + +/// Returns the packed header kinds to ask for: sequence and picture, where the driver takes them. +/// +/// Zero when the driver takes neither, or when asking fails. +fn packed_headers(display: &Display, entrypoint: VAEntrypoint::Type) -> u32 { + let mut attrs = [VAConfigAttrib { + type_: VAConfigAttribType::VAConfigAttribEncPackedHeaders, + value: 0, + }]; + if display + .get_config_attributes(VAProfile::VAProfileH264Main, 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) +} + +/// Returns the lowest H.264 level, from 4 up, whose frame size and macroblock rate cover the stream. +/// +/// Only those two limits: the level's maximum bitrate is not checked. +/// +/// Level 4 is the floor because it was the fixed level before, and decoders +/// take the DPB size from `max_dec_frame_buffering` rather than from the level, +/// so a lower one would buy nothing. Above 1080p (a 1440p or 4K screen share) a +/// fixed level 4 would declare a stream its decoder may refuse. +fn level(width: u32, height: u32, framerate: u32) -> Level { + // H.264 Table A-1: (level, MaxMBPS, MaxFS). + const LEVELS: [(Level, u64, u64); 8] = [ + (Level::L4, 245_760, 8_192), + (Level::L4_2, 522_240, 8_704), + (Level::L5, 589_824, 22_080), + (Level::L5_1, 983_040, 36_864), + (Level::L5_2, 2_073_600, 36_864), + (Level::L6, 4_177_920, 139_264), + (Level::L6_1, 8_355_840, 139_264), + (Level::L6_2, 16_711_680, 139_264), + ]; + let frame = u64::from(width.div_ceil(16)) * u64::from(height.div_ceil(16)); + let rate = frame * u64::from(framerate.max(1)); + LEVELS + .into_iter() + .find(|&(_, max_rate, max_frame)| frame <= max_frame && rate <= max_rate) + .map(|(level, _, _)| level) + .unwrap_or(Level::L6_2) +} + +/// Returns the coded buffer size in bytes for `bitrate`: two seconds of it, and at least 1.5 MB. +fn coded_size(bitrate: u32) -> usize { + (bitrate as usize / 4).max(1_500_000) +} + +/// Returns whether the VA fourcc `fourcc` is packed RGB, which has no YUV matrix of its own. +fn is_rgb(fourcc: u32) -> bool { + matches!( + fourcc, + VA_FOURCC_BGRX | VA_FOURCC_BGRA | VA_FOURCC_RGBX | VA_FOURCC_RGBA | VA_FOURCC_XRGB | VA_FOURCC_ARGB + ) +} + /// Allocate one NV12 encode surface. fn make_surface(display: &Arc, config: &Config) -> anyhow::Result> { let mut surfaces = display @@ -381,7 +623,8 @@ fn upload_nv12( /// Build the SPS/PPS for a low-latency IPPP stream. Ported from /// discord/cros-codecs `LowDelayH264Delegate::new_sequence` (BSD-3-Clause). fn build_sps_pps(config: &Config) -> (Rc, Rc) { - let level = Level::L4; + let level = level(config.width, config.height, config.framerate); + let (primaries, transfer, matrix) = config.color.vui(); let sps = SpsBuilder::new() .seq_parameter_set_id(0) .profile_idc(Profile::Main) @@ -398,6 +641,7 @@ fn build_sps_pps(config: &Config) -> (Rc, Rc) { .bit_depth_chroma(8) .aspect_ratio(1, 1) .timing_info(1, config.framerate * 2, false) + .video_signal_type(config.color.full_range, primaries, transfer, matrix) .max_num_reorder_frames(0) .max_dec_frame_buffering(1) .build(); @@ -613,3 +857,598 @@ fn packed_header(kind: EncPackedHeaderType, data: &[u8]) -> (BufferType, BufferT let payload = BufferType::EncPackedHeaderData(data.to_vec()); (param, payload) } + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::process::Command; + use std::sync::atomic::{AtomicU32, Ordering}; + + use super::*; + use crate::dmabuf::DmaBuf; + use crate::{VA_RT_FORMAT_RGB32, VA_RT_FORMAT_YUV420}; + + const WIDTH: u32 = 320; + const HEIGHT: u32 = 240; + + fn config() -> Config { + 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. + fn encoder(config: Config) -> Option { + match Encoder::new(config) { + Ok(encoder) => Some(encoder), + Err(err) => { + eprintln!("skipping: no VA-API H.264 encoder: {err:#}"); + None + } + } + } + + /// 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 { + let (w, h) = (width as usize, height as usize); + let mut data = vec![128u8; w * h + w * h.div_ceil(2)]; + for y in 0..h { + for x in 0..w { + data[y * w + x] = (32 + (x + y) * 160 / (w + h)) as u8; + } + } + let x0 = (step as usize * 8) % (w / 2); + for y in h / 4..h / 2 { + for x in x0..x0 + w / 4 { + data[y * w + x] = 220; + } + } + data + } + + /// Returns a temporary file path unique to this call. + /// + /// 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 { + 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())) + } + + /// 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> { + // 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); + std::fs::write(&path, stream).expect("write the elementary stream"); + let output = Command::new("ffmpeg") + .args(["-v", "error", "-i"]) + .arg(&path) + .args(["-pix_fmt", "nv12", "-f", "rawvideo", "-"]) + .output(); + let _ = std::fs::remove_file(&path); + match output { + Ok(output) if output.status.success() => Some(output.stdout), + Ok(output) => panic!("ffmpeg refused the stream: {}", String::from_utf8_lossy(&output.stderr)), + Err(_) => { + eprintln!("skipping: no ffmpeg"); + None + } + } + } + + /// Asks ffprobe for the stream's `color_range,color_space`. + fn ffprobe_color(stream: &[u8]) -> Option { + let path = temp_path("probe"); + std::fs::write(&path, stream).expect("write the elementary stream"); + let output = Command::new("ffprobe") + .args([ + "-v", + "error", + "-show_entries", + "stream=color_range,color_space", + "-of", + "csv=p=0", + ]) + .arg(&path) + .output() + .ok(); + let _ = std::fs::remove_file(&path); + Some(String::from_utf8(output?.stdout).ok()?.trim().to_string()) + } + + /// Returns the peak signal-to-noise ratio of the luma planes, in dB. + fn luma_psnr(a: &[u8], b: &[u8], width: u32, height: u32) -> f64 { + let len = (width * height) as usize; + let mse = a[..len] + .iter() + .zip(&b[..len]) + .map(|(&x, &y)| (x as f64 - y as f64).powi(2)) + .sum::() + / len as f64; + 10.0 * (255.0f64.powi(2) / mse.max(1e-9)).log10() + } + + /// Returns the mean of the samples in a rectangle of one plane. + fn mean( + plane: &[u8], + stride: usize, + (x0, y0, x1, y1): (usize, usize, usize, usize), + step: usize, + at: usize, + ) -> f64 { + let mut sum = 0.0; + let mut count = 0.0; + for y in y0..y1 { + for x in x0..x1 { + sum += plane[y * stride + x * step + at] as f64; + count += 1.0; + } + } + sum / count + } + + /// Returns the NAL unit types in an Annex-B access unit. + fn nal_types(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); + i += 3; + } + i += 1; + } + types + } + + /// Returns a driver surface of `fourcc` on `display`, filled through an image. + fn surface_with( + display: &Arc, + fourcc: u32, + rt_format: u32, + (width, height): (u32, u32), + fill: impl FnOnce(&mut [u8], &crate::bindings::VAImage), + ) -> Surface<()> { + let surface = display + .create_surfaces( + rt_format, + Some(fourcc), + width, + height, + Some(UsageHint::USAGE_HINT_EXPORT), + vec![()], + ) + .expect("allocate a surface") + .pop() + .expect("one surface"); + let format = display + .query_image_formats() + .expect("query image formats") + .into_iter() + .find(|f| f.fourcc == fourcc) + .expect("the driver has an image format for the surface format"); + let mut image = + Image::create_from(&surface, format, (width, height), (width, height)).expect("create an image"); + let va_image = *image.image(); + fill(image.as_mut(), &va_image); + drop(image); + surface.sync().expect("sync the filled surface"); + surface + } + + /// Returns an NV12 surface holding `nv12_frame`, exported as a DMA-BUF. + fn nv12_dmabuf(display: &Arc, (width, height): (u32, u32), step: u32) -> DmaBuf { + let surface = surface_with(display, VA_FOURCC_NV12, VA_RT_FORMAT_YUV420, (width, height), |_, _| {}); + upload_nv12(display, &surface, width, height, &nv12_frame(width, height, step)).expect("upload"); + DmaBuf::from_prime(surface.export_prime().expect("export the surface")).expect("one object") + } + + /// CPU NV12 in, H.264 out that an independent decoder reads back as the same picture. + #[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"); + } + } + + /// 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); + } + } + + /// An IDR opens every group, and only there. + #[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" + ); + } + } + // 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] + fn the_full_entrypoint_wins_and_low_power_is_the_fallback() { + use VAEntrypoint::{ + VAEntrypointEncSlice as Full, VAEntrypointEncSliceLP as LowPower, VAEntrypointVLD as Decode, + }; + assert_eq!(pick_entrypoint(&[Decode, LowPower, Full]), Some(Full)); + assert_eq!(pick_entrypoint(&[Decode, LowPower]), Some(LowPower)); + assert_eq!(pick_entrypoint(&[Decode]), None); + } + + /// A buffer the driver refuses leaves the encoder usable, so the caller's + /// 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(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + // A format with no mapping, and a real buffer under a modifier no driver has. + let unmapped = DmaBuf { + drm_format: u32::from_le_bytes(*b"ZZZZ"), + ..nv12_dmabuf(&display, (WIDTH, HEIGHT), 0) + }; + assert!(encoder.encode_dmabuf(unmapped, true).is_err()); + let unknown_modifier = DmaBuf { + modifier: 0x00ff_ffff_ffff_fffe, + ..nv12_dmabuf(&display, (WIDTH, HEIGHT), 0) + }; + assert!(encoder.encode_dmabuf(unknown_modifier, true).is_err()); + + 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 { + 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. + #[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"); + } + + /// 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() { + let Some(mut encoder) = encoder(Config { + color: Color::BT601, + ..config() + }) else { + return; + }; + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + // Flat red, written in BT.709. + let [y, u, v] = limited([255, 0, 0], Color::BT709).map(|value| value.round() as u8); + let surface = surface_with( + &display, + VA_FOURCC_NV12, + VA_RT_FORMAT_YUV420, + (WIDTH, HEIGHT), + |data, image| { + for row in 0..HEIGHT as usize { + let at = image.offsets[0] as usize + row * image.pitches[0] as usize; + data[at..at + WIDTH as usize].fill(y); + } + for row in 0..HEIGHT as usize / 2 { + let at = image.offsets[1] as usize + row * image.pitches[1] as usize; + for pair in data[at..at + WIDTH as usize].chunks_exact_mut(2) { + pair.copy_from_slice(&[u, v]); + } + } + }, + ); + let buffer = DmaBuf { + color: Some(Color::BT709), + ..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 { + return; + }; + + let (w, h) = (WIDTH as usize, HEIGHT as usize); + let (luma, chroma) = decoded.split_at(w * h); + let actual = [ + mean(luma, w, (w / 4, h / 4, w * 3 / 4, h * 3 / 4), 1, 0), + mean(chroma, w, (w / 8, h / 8, w * 3 / 8, h * 3 / 8), 2, 0), + mean(chroma, w, (w / 8, h / 8, w * 3 / 8, h * 3 / 8), 2, 1), + ]; + let expected = limited([255, 0, 0], Color::BT601); + 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}" + ); + } + } + + /// 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"); + } + } + + /// 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 }; + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + let mut stream = Vec::new(); + for step in 0..3 { + 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 { + return; + }; + let len = nv12_frame(WIDTH, HEIGHT, 0).len(); + assert_eq!(decoded.len(), 3 * len, "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. + for step in 0..3u32 { + let source = nv12_frame(WIDTH * 2, HEIGHT * 2, step); + let mut expected = vec![0u8; (WIDTH * HEIGHT) as usize]; + for y in 0..HEIGHT as usize { + for x in 0..WIDTH as usize { + expected[y * WIDTH as usize + x] = source[(y * 2) * (WIDTH as usize * 2) + x * 2]; + } + } + let psnr = luma_psnr(&expected, &decoded[step as usize * len..], WIDTH, HEIGHT); + assert!(psnr > 25.0, "picture {step} decoded at {psnr:.1} dB"); + } + } + + /// A packed RGB DMA-BUF is converted on the GPU into the configured color space. + #[test] + 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; + }; + let Some(display) = Display::open() else { + eprintln!("skipping: no VA-API display"); + return; + }; + let surface = surface_with( + &display, + VA_FOURCC_BGRX, + VA_RT_FORMAT_RGB32, + (WIDTH, HEIGHT), + |data, image| { + let pitch = image.pitches[0] as usize; + let offset = image.offsets[0] as usize; + for y in 0..HEIGHT as usize { + for x in 0..WIDTH as usize { + let quadrant = (y >= HEIGHT as usize / 2) as usize * 2 + (x >= WIDTH as usize / 2) as usize; + let [r, g, b] = QUADRANTS[quadrant]; + let at = offset + y * pitch + x * 4; + data[at..at + 4].copy_from_slice(&[b, g, r, 255]); + } + } + }, + ); + let exported = surface.export_prime().expect("export"); + 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 { + return; + }; + + let (w, h) = (WIDTH as usize, HEIGHT as usize); + let (luma, chroma) = decoded.split_at(w * h); + for (quadrant, rgb) in QUADRANTS.iter().enumerate() { + let (qx, qy) = (quadrant % 2, quadrant / 2); + // The middle half of each quadrant, clear of the chroma bleed at its edges. + let area = |scale: usize| { + let (qw, qh) = (w / 2 / scale, h / 2 / scale); + ( + qx * qw + qw / 4, + qy * qh + qh / 4, + qx * qw + qw * 3 / 4, + qy * qh + qh * 3 / 4, + ) + }; + let actual = [ + mean(luma, w, area(1), 1, 0), + mean(chroma, w, area(2), 2, 0), + mean(chroma, w, area(2), 2, 1), + ]; + let expected = limited(*rgb, color); + 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}" + ); + } + } + } + } + + #[test] + fn the_level_covers_the_frame_size_and_rate() { + assert_eq!(level(320, 240, 30), Level::L4); + assert_eq!(level(1920, 1080, 30), Level::L4); + assert_eq!(level(1920, 1080, 60), Level::L4_2); + assert_eq!(level(2560, 1440, 30), Level::L5); + assert_eq!(level(3840, 2160, 30), Level::L5_1); + assert_eq!(level(3840, 2160, 60), Level::L5_2); + } + + /// Returns the limited-range Y'CbCr of an RGB color, per BT.601 or BT.709. + fn limited([r, g, b]: [u8; 3], color: Color) -> [f64; 3] { + let (kr, kb) = match color.matrix { + crate::Matrix::Bt601 => (0.299, 0.114), + crate::Matrix::Bt709 => (0.2126, 0.0722), + }; + let (r, g, b) = (r as f64, g as f64, b as f64); + let y = kr * r + (1.0 - kr - kb) * g + kb * b; + [ + 16.0 + y * 219.0 / 255.0, + 128.0 + (b - y) / (2.0 * (1.0 - kb)) * 224.0 / 255.0, + 128.0 + (r - y) / (2.0 * (1.0 - kr)) * 224.0 / 255.0, + ] + } + + /// 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 { + return; + }; + // 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. + let mut state = 0x2545_f491_u32; + let mut frame = |step: u32| { + let mut frame = nv12_frame(WIDTH, HEIGHT, step); + for sample in &mut frame[..(WIDTH * HEIGHT) as usize] { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + *sample = sample.saturating_add((state >> 28) as u8); + } + frame + }; + // The mean picture size over the second half of 60 frames, once the + // rate controller has settled. + // Also asserts that no picture after the first is an IDR. + let mut average = |encoder: &mut Encoder, from: u32| { + let sizes: Vec = (from..from + 60) + .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"); + } + unit.len() + }) + .collect(); + sizes[30..].iter().sum::() / 30 + }; + + let high = average(&mut encoder, 0); + encoder.set_bitrate(250_000).expect("set the bitrate"); + 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"); + assert!( + low * 3 < high, + "the rate did not drop: {high} -> {low} bytes per picture" + ); + } +}