Lightningbeam/lightningbeam-ui/gpu-video-encoder/src/vaapi.rs

514 lines
20 KiB
Rust

//! VAAPI hardware H.264 encoding (Linux/Intel/AMD).
//!
//! Level 1 (this module first): a CPU-fed encoder — upload NV12 frames to VAAPI
//! surfaces (`av_hwframe_transfer_data`) and encode with `h264_vaapi`. This proves the
//! encoder works and establishes the FFI scaffolding. Level 2 (zero-copy: GPU writes
//! NV12 straight into the VAAPI surface via DMA-BUF) builds on this.
//!
//! All `unsafe` FFmpeg FFI is contained here.
use ffmpeg_sys_next as ff;
use std::ffi::CString;
use std::ptr;
#[inline]
fn averror(e: i32) -> i32 {
-e
}
/// Create a VAAPI hwdevice on `/dev/dri/renderD128`, trying driver names in turn.
///
/// libva's auto-selection can pick a driver that doesn't support the GPU — notably it
/// chooses the legacy `i965` driver on newer Intel parts (Gen 11+) where the modern `iHD`
/// driver is required. Each `av_hwdevice_ctx_create` opens a fresh VADisplay, so
/// `LIBVA_DRIVER_NAME` is re-read per attempt. We try `iHD` first (modern Intel), then the
/// caller's original setting, then `i965` (older Intel) and `radeonsi` (AMD). On success the
/// working driver name is left in the env; on total failure the original value is restored.
pub fn create_device() -> Result<*mut ff::AVBufferRef, String> {
unsafe {
let node = CString::new("/dev/dri/renderD128").unwrap();
let original = std::env::var_os("LIBVA_DRIVER_NAME");
let attempts: [Option<&str>; 4] = [Some("iHD"), None, Some("i965"), Some("radeonsi")];
for drv in attempts {
match drv {
Some(d) => std::env::set_var("LIBVA_DRIVER_NAME", d),
// `None` = the caller's original setting (or libva auto if unset).
None => match &original {
Some(v) => std::env::set_var("LIBVA_DRIVER_NAME", v),
None => std::env::remove_var("LIBVA_DRIVER_NAME"),
},
}
let mut hw: *mut ff::AVBufferRef = ptr::null_mut();
if ff::av_hwdevice_ctx_create(
&mut hw,
ff::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
node.as_ptr(),
ptr::null_mut(),
0,
) >= 0
{
return Ok(hw);
}
}
match &original {
Some(v) => std::env::set_var("LIBVA_DRIVER_NAME", v),
None => std::env::remove_var("LIBVA_DRIVER_NAME"),
}
Err("av_hwdevice_ctx_create(VAAPI) failed for all drivers (iHD/i965/radeonsi)".into())
}
}
/// Copy tight NV12 (`Y` then interleaved `UV`) into an AVFrame's planes, respecting
/// each plane's linesize (which FFmpeg may pad).
unsafe fn fill_nv12(frame: *mut ff::AVFrame, nv12: &[u8], width: u32, height: u32) {
let w = width as usize;
let h = height as usize;
// Y plane: h rows of w bytes.
let dst_y = (*frame).data[0];
let ls_y = (*frame).linesize[0] as usize;
for row in 0..h {
let src = &nv12[row * w..row * w + w];
ptr::copy_nonoverlapping(src.as_ptr(), dst_y.add(row * ls_y), w);
}
// UV plane: h/2 rows of w bytes (interleaved U,V), source offset starts at w*h.
let dst_uv = (*frame).data[1];
let ls_uv = (*frame).linesize[1] as usize;
let uv_off = w * h;
for row in 0..h / 2 {
let src = &nv12[uv_off + row * w..uv_off + row * w + w];
ptr::copy_nonoverlapping(src.as_ptr(), dst_uv.add(row * ls_uv), w);
}
}
/// A VAAPI NV12 surface mapped to a DMA-BUF, with its layout extracted for Vulkan import.
/// Keeps the FFmpeg handles alive; the `fd` stays valid until drop (dup it for Vulkan).
pub struct MappedSurface {
hw_device: *mut ff::AVBufferRef,
frames_ref: *mut ff::AVBufferRef,
surf: *mut ff::AVFrame,
drm: *mut ff::AVFrame,
pub width: u32,
pub height: u32,
pub fd: i32,
pub size: u64,
pub modifier: u64,
pub y_offset: u64,
pub y_pitch: u64,
pub uv_offset: u64,
pub uv_pitch: u64,
}
impl MappedSurface {
/// Allocate a VAAPI NV12 surface and map it to DRM-PRIME.
pub fn alloc(width: u32, height: u32) -> Result<Self, String> {
unsafe {
let mut hw_device: *mut ff::AVBufferRef = ptr::null_mut();
let node = CString::new("/dev/dri/renderD128").unwrap();
if ff::av_hwdevice_ctx_create(
&mut hw_device,
ff::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
node.as_ptr(),
ptr::null_mut(),
0,
) < 0
{
return Err("av_hwdevice_ctx_create failed".into());
}
let frames_ref = ff::av_hwframe_ctx_alloc(hw_device);
if frames_ref.is_null() {
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_ctx_alloc failed".into());
}
{
let fctx = (*frames_ref).data as *mut ff::AVHWFramesContext;
(*fctx).format = ff::AVPixelFormat::AV_PIX_FMT_VAAPI;
(*fctx).sw_format = ff::AVPixelFormat::AV_PIX_FMT_NV12;
(*fctx).width = width as i32;
(*fctx).height = height as i32;
(*fctx).initial_pool_size = 4;
}
if ff::av_hwframe_ctx_init(frames_ref) < 0 {
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_ctx_init failed".into());
}
let surf = ff::av_frame_alloc();
if ff::av_hwframe_get_buffer(frames_ref, surf, 0) < 0 {
ff::av_frame_free(&mut (surf as *mut _));
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_get_buffer failed".into());
}
let drm = ff::av_frame_alloc();
(*drm).format = ff::AVPixelFormat::AV_PIX_FMT_DRM_PRIME as i32;
let flags = ff::AV_HWFRAME_MAP_DIRECT as i32
| ff::AV_HWFRAME_MAP_READ as i32
| ff::AV_HWFRAME_MAP_WRITE as i32;
if ff::av_hwframe_map(drm, surf, flags) < 0 {
ff::av_frame_free(&mut (drm as *mut _));
ff::av_frame_free(&mut (surf as *mut _));
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_map failed".into());
}
let desc = (*drm).data[0] as *const ff::AVDRMFrameDescriptor;
// Expect 1 object, 2 layers (Y=R8, UV=GR88).
if (*desc).nb_objects != 1 || (*desc).nb_layers != 2 {
return Err(format!(
"unexpected DRM layout: {} objects, {} layers",
(*desc).nb_objects, (*desc).nb_layers
));
}
let obj = &(*desc).objects[0];
let y = &(*desc).layers[0].planes[0];
let uv = &(*desc).layers[1].planes[0];
Ok(MappedSurface {
hw_device,
frames_ref,
surf,
drm,
width,
height,
fd: obj.fd,
size: obj.size as u64,
modifier: obj.format_modifier,
y_offset: y.offset as u64,
y_pitch: y.pitch as u64,
uv_offset: uv.offset as u64,
uv_pitch: uv.pitch as u64,
})
}
}
/// The underlying VASurface AVFrame (to hand to the encoder).
pub fn av_frame(&self) -> *mut ff::AVFrame {
self.surf
}
/// Read the surface back to tight CPU NV12 (for verifying what the GPU wrote).
pub fn readback_nv12(&self) -> Result<Vec<u8>, String> {
unsafe {
let sw = ff::av_frame_alloc();
(*sw).format = ff::AVPixelFormat::AV_PIX_FMT_NV12 as i32;
(*sw).width = self.width as i32;
(*sw).height = self.height as i32;
if ff::av_frame_get_buffer(sw, 0) < 0 {
ff::av_frame_free(&mut (sw as *mut _));
return Err("av_frame_get_buffer failed".into());
}
if ff::av_hwframe_transfer_data(sw, self.surf, 0) < 0 {
ff::av_frame_free(&mut (sw as *mut _));
return Err("av_hwframe_transfer_data (download) failed".into());
}
let w = self.width as usize;
let h = self.height as usize;
let mut out = vec![0u8; w * h + w * (h / 2)];
let ls_y = (*sw).linesize[0] as usize;
for row in 0..h {
let src = (*sw).data[0].add(row * ls_y);
ptr::copy_nonoverlapping(src, out.as_mut_ptr().add(row * w), w);
}
let ls_uv = (*sw).linesize[1] as usize;
let uv_off = w * h;
for row in 0..h / 2 {
let src = (*sw).data[1].add(row * ls_uv);
ptr::copy_nonoverlapping(src, out.as_mut_ptr().add(uv_off + row * w), w);
}
ff::av_frame_free(&mut (sw as *mut _));
Ok(out)
}
}
}
impl Drop for MappedSurface {
fn drop(&mut self) {
unsafe {
ff::av_frame_free(&mut (self.drm as *mut _));
ff::av_frame_free(&mut (self.surf as *mut _));
let mut fr = self.frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut self.hw_device);
}
}
}
/// Allocate one VAAPI NV12 surface, map it to a DRM-PRIME descriptor, and return a
/// human-readable dump of its DMA-BUF layout (object fds/size/modifier; layer fourcc;
/// per-plane object/offset/pitch). The format **modifier** decides the zero-copy path:
/// `0` = LINEAR (compute can write a linear NV12 buffer/image), anything else = tiled
/// (needs a GPU copy into the tiled surface, or a linear import VAAPI accepts).
pub fn probe_surface_drm(width: u32, height: u32) -> Result<String, String> {
unsafe {
let mut hw_device: *mut ff::AVBufferRef = ptr::null_mut();
let node = CString::new("/dev/dri/renderD128").unwrap();
if ff::av_hwdevice_ctx_create(
&mut hw_device,
ff::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
node.as_ptr(),
ptr::null_mut(),
0,
) < 0
{
return Err("av_hwdevice_ctx_create(VAAPI) failed".into());
}
let frames_ref = ff::av_hwframe_ctx_alloc(hw_device);
if frames_ref.is_null() {
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_ctx_alloc failed".into());
}
{
let fctx = (*frames_ref).data as *mut ff::AVHWFramesContext;
(*fctx).format = ff::AVPixelFormat::AV_PIX_FMT_VAAPI;
(*fctx).sw_format = ff::AVPixelFormat::AV_PIX_FMT_NV12;
(*fctx).width = width as i32;
(*fctx).height = height as i32;
(*fctx).initial_pool_size = 2;
}
if ff::av_hwframe_ctx_init(frames_ref) < 0 {
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_ctx_init failed".into());
}
let surf = ff::av_frame_alloc();
if ff::av_hwframe_get_buffer(frames_ref, surf, 0) < 0 {
ff::av_frame_free(&mut (surf as *mut _));
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
return Err("av_hwframe_get_buffer failed".into());
}
let drm = ff::av_frame_alloc();
(*drm).format = ff::AVPixelFormat::AV_PIX_FMT_DRM_PRIME as i32;
let flags = ff::AV_HWFRAME_MAP_DIRECT as i32
| ff::AV_HWFRAME_MAP_READ as i32
| ff::AV_HWFRAME_MAP_WRITE as i32;
let r = ff::av_hwframe_map(drm, surf, flags);
if r < 0 {
ff::av_frame_free(&mut (drm as *mut _));
ff::av_frame_free(&mut (surf as *mut _));
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
return Err(format!("av_hwframe_map(DRM_PRIME) failed: {r}"));
}
let desc = (*drm).data[0] as *const ff::AVDRMFrameDescriptor;
let mut s = format!("VAAPI NV12 {width}x{height} surface as DRM-PRIME:\n");
s += &format!(" nb_objects = {}\n", (*desc).nb_objects);
for o in 0..(*desc).nb_objects as usize {
let obj = &(*desc).objects[o];
s += &format!(
" object[{o}]: fd={} size={} format_modifier=0x{:016x}{}\n",
obj.fd,
obj.size,
obj.format_modifier,
if obj.format_modifier == 0 { " (LINEAR)" } else { " (tiled)" },
);
}
s += &format!(" nb_layers = {}\n", (*desc).nb_layers);
for l in 0..(*desc).nb_layers as usize {
let lay = &(*desc).layers[l];
let f = lay.format;
let fourcc = [(f & 0xff) as u8, ((f >> 8) & 0xff) as u8, ((f >> 16) & 0xff) as u8, ((f >> 24) & 0xff) as u8];
s += &format!(
" layer[{l}]: format='{}' (0x{:08x}) nb_planes={}\n",
String::from_utf8_lossy(&fourcc),
f,
lay.nb_planes,
);
for p in 0..lay.nb_planes as usize {
let pl = &lay.planes[p];
s += &format!(
" plane[{p}]: object_index={} offset={} pitch={}\n",
pl.object_index, pl.offset, pl.pitch,
);
}
}
ff::av_frame_free(&mut (drm as *mut _));
ff::av_frame_free(&mut (surf as *mut _));
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::av_buffer_unref(&mut hw_device);
Ok(s)
}
}
/// Encode NV12 frames with `h264_vaapi` and write the raw Annex-B H.264 to `out_path`.
/// Returns the number of encoded packets. `Err` (rather than panic) when VAAPI/the
/// encoder is unavailable, so callers can fall back.
pub fn encode_nv12_to_file(
width: u32,
height: u32,
frames: &[Vec<u8>],
framerate: i32,
out_path: &str,
) -> Result<usize, String> {
unsafe {
// 1. VAAPI device.
let mut hw_device: *mut ff::AVBufferRef = ptr::null_mut();
let node = CString::new("/dev/dri/renderD128").unwrap();
let r = ff::av_hwdevice_ctx_create(
&mut hw_device,
ff::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
node.as_ptr(),
ptr::null_mut(),
0,
);
if r < 0 {
return Err(format!("av_hwdevice_ctx_create(VAAPI) failed: {r}"));
}
let cleanup_dev = |dev: *mut ff::AVBufferRef| {
let mut d = dev;
ff::av_buffer_unref(&mut d);
};
// 2. Encoder.
let name = CString::new("h264_vaapi").unwrap();
let codec = ff::avcodec_find_encoder_by_name(name.as_ptr());
if codec.is_null() {
cleanup_dev(hw_device);
return Err("encoder h264_vaapi not found in this FFmpeg build".into());
}
let enc = ff::avcodec_alloc_context3(codec);
if enc.is_null() {
cleanup_dev(hw_device);
return Err("avcodec_alloc_context3 failed".into());
}
(*enc).width = width as i32;
(*enc).height = height as i32;
(*enc).time_base = ff::AVRational { num: 1, den: framerate };
(*enc).framerate = ff::AVRational { num: framerate, den: 1 };
(*enc).pix_fmt = ff::AVPixelFormat::AV_PIX_FMT_VAAPI;
// 3. HW frames context (VAAPI surfaces with NV12 sw layout).
let frames_ref = ff::av_hwframe_ctx_alloc(hw_device);
if frames_ref.is_null() {
ff::avcodec_free_context(&mut (enc as *mut _));
cleanup_dev(hw_device);
return Err("av_hwframe_ctx_alloc failed".into());
}
{
let fctx = (*frames_ref).data as *mut ff::AVHWFramesContext;
(*fctx).format = ff::AVPixelFormat::AV_PIX_FMT_VAAPI;
(*fctx).sw_format = ff::AVPixelFormat::AV_PIX_FMT_NV12;
(*fctx).width = width as i32;
(*fctx).height = height as i32;
(*fctx).initial_pool_size = 8;
}
let r = ff::av_hwframe_ctx_init(frames_ref);
if r < 0 {
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::avcodec_free_context(&mut (enc as *mut _));
cleanup_dev(hw_device);
return Err(format!("av_hwframe_ctx_init failed: {r}"));
}
(*enc).hw_frames_ctx = ff::av_buffer_ref(frames_ref);
// 4. Open.
let r = ff::avcodec_open2(enc, codec, ptr::null_mut());
if r < 0 {
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::avcodec_free_context(&mut (enc as *mut _));
cleanup_dev(hw_device);
return Err(format!("avcodec_open2(h264_vaapi) failed: {r}"));
}
let mut out: Vec<u8> = Vec::new();
let pkt = ff::av_packet_alloc();
let mut count = 0usize;
// Drain helper: pull packets and append to `out`.
let drain = |enc: *mut ff::AVCodecContext, out: &mut Vec<u8>, count: &mut usize| -> Result<(), String> {
loop {
let r = ff::avcodec_receive_packet(enc, pkt);
if r == averror(libc::EAGAIN) || r == ff::AVERROR_EOF {
break;
}
if r < 0 {
return Err(format!("avcodec_receive_packet failed: {r}"));
}
let data = std::slice::from_raw_parts((*pkt).data, (*pkt).size as usize);
out.extend_from_slice(data);
*count += 1;
ff::av_packet_unref(pkt);
}
Ok(())
};
let mut err: Option<String> = None;
for (i, nv12) in frames.iter().enumerate() {
// Software NV12 frame.
let sw = ff::av_frame_alloc();
(*sw).format = ff::AVPixelFormat::AV_PIX_FMT_NV12 as i32;
(*sw).width = width as i32;
(*sw).height = height as i32;
if ff::av_frame_get_buffer(sw, 0) < 0 {
err = Some("av_frame_get_buffer(sw) failed".into());
ff::av_frame_free(&mut (sw as *mut _));
break;
}
fill_nv12(sw, nv12, width, height);
// VAAPI surface frame + upload.
let hw = ff::av_frame_alloc();
if ff::av_hwframe_get_buffer(frames_ref, hw, 0) < 0 {
err = Some("av_hwframe_get_buffer failed".into());
ff::av_frame_free(&mut (sw as *mut _));
ff::av_frame_free(&mut (hw as *mut _));
break;
}
if ff::av_hwframe_transfer_data(hw, sw, 0) < 0 {
err = Some("av_hwframe_transfer_data failed".into());
ff::av_frame_free(&mut (sw as *mut _));
ff::av_frame_free(&mut (hw as *mut _));
break;
}
(*hw).pts = i as i64;
let r = ff::avcodec_send_frame(enc, hw);
ff::av_frame_free(&mut (sw as *mut _));
ff::av_frame_free(&mut (hw as *mut _));
if r < 0 {
err = Some(format!("avcodec_send_frame failed: {r}"));
break;
}
if let Err(e) = drain(enc, &mut out, &mut count) {
err = Some(e);
break;
}
}
// Flush.
if err.is_none() {
ff::avcodec_send_frame(enc, ptr::null_mut());
if let Err(e) = drain(enc, &mut out, &mut count) {
err = Some(e);
}
}
// Cleanup.
ff::av_packet_free(&mut (pkt as *mut _));
let mut fr = frames_ref;
ff::av_buffer_unref(&mut fr);
ff::avcodec_free_context(&mut (enc as *mut _));
cleanup_dev(hw_device);
if let Some(e) = err {
return Err(e);
}
std::fs::write(out_path, &out).map_err(|e| format!("write {out_path}: {e}"))?;
Ok(count)
}
}