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/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/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" + ); + } +} 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}" + ); + } +}