feat: encode from DMA-BUFs, with GPU scale and convert, and change the bitrate

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.
 

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;

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

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<Self> {
+		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