diff --git a/lightningbeam-ui/gpu-video-encoder/src/decoder.rs b/lightningbeam-ui/gpu-video-encoder/src/decoder.rs new file mode 100644 index 0000000..f53c251 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/src/decoder.rs @@ -0,0 +1,216 @@ +//! VAAPI hardware video decode → wgpu textures. The mirror of [`crate::encoder`]: the codec +//! decodes into a VAAPI NV12 surface, which is mapped to a DRM-PRIME DMA-BUF and imported as two +//! wgpu plane textures via [`crate::dmabuf::import_raw`] — the exact same path the encoder uses, +//! in the read direction. Stays GPU-resident: no CPU frame copy. + +use crate::dmabuf::{self, ImportedNv12, Nv12DmaBuf}; +use crate::vk_device::{self, DrmDevice}; +use ffmpeg_sys_next as ff; +use std::ffi::CString; +use std::path::Path; +use std::ptr; + +#[inline] +fn averror(e: i32) -> i32 { + -e +} + +/// `get_format` callback: pick VAAPI surfaces so the decoder outputs hardware frames. With +/// `hw_device_ctx` set, FFmpeg auto-allocates the matching frames context. +unsafe extern "C" fn get_vaapi_format( + _ctx: *mut ff::AVCodecContext, + mut fmts: *const ff::AVPixelFormat, +) -> ff::AVPixelFormat { + while *fmts != ff::AVPixelFormat::AV_PIX_FMT_NONE { + if *fmts == ff::AVPixelFormat::AV_PIX_FMT_VAAPI { + return ff::AVPixelFormat::AV_PIX_FMT_VAAPI; + } + fmts = fmts.add(1); + } + ff::AVPixelFormat::AV_PIX_FMT_NONE +} + +/// Hardware decoder for a single video file/stream. Frames come back as importable NV12 textures +/// on [`Self::device`]. +pub struct VaapiDecoder { + drm: DrmDevice, + hw_device: *mut ff::AVBufferRef, + fmt: *mut ff::AVFormatContext, + dec: *mut ff::AVCodecContext, + pkt: *mut ff::AVPacket, + frame: *mut ff::AVFrame, + stream_index: i32, + flushing: bool, +} + +// Owns its FFmpeg/Vulkan handles exclusively; only moved, never shared (same as the encoder). +unsafe impl Send for VaapiDecoder {} + +impl VaapiDecoder { + /// Open `input_path` and set up VAAPI hardware decoding of its best video stream. + pub fn new(input_path: &Path) -> Result { + let drm = vk_device::create()?; + unsafe { + let mut hw_device = crate::vaapi::create_device()?; + let cleanup_hw = |hw: *mut ff::AVBufferRef| { + let mut h = hw; + ff::av_buffer_unref(&mut h); + }; + + let path_c = CString::new(input_path.to_string_lossy().as_ref()).unwrap(); + let mut fmt: *mut ff::AVFormatContext = ptr::null_mut(); + if ff::avformat_open_input(&mut fmt, path_c.as_ptr(), ptr::null_mut(), ptr::null_mut()) < 0 { + cleanup_hw(hw_device); + return Err(format!("avformat_open_input {input_path:?} failed")); + } + if ff::avformat_find_stream_info(fmt, ptr::null_mut()) < 0 { + ff::avformat_close_input(&mut fmt); + cleanup_hw(hw_device); + return Err("avformat_find_stream_info failed".into()); + } + + let mut decoder: *const ff::AVCodec = ptr::null(); + let stream_index = ff::av_find_best_stream( + fmt, + ff::AVMediaType::AVMEDIA_TYPE_VIDEO, + -1, + -1, + &mut decoder, + 0, + ); + if stream_index < 0 || decoder.is_null() { + ff::avformat_close_input(&mut fmt); + cleanup_hw(hw_device); + return Err("no decodable video stream".into()); + } + + let dec = ff::avcodec_alloc_context3(decoder); + let stream = *(*fmt).streams.add(stream_index as usize); + if ff::avcodec_parameters_to_context(dec, (*stream).codecpar) < 0 { + ff::avcodec_free_context(&mut (dec as *mut _)); + ff::avformat_close_input(&mut fmt); + cleanup_hw(hw_device); + return Err("avcodec_parameters_to_context failed".into()); + } + (*dec).hw_device_ctx = ff::av_buffer_ref(hw_device); + (*dec).get_format = Some(get_vaapi_format); + + if ff::avcodec_open2(dec, decoder, ptr::null_mut()) < 0 { + ff::avcodec_free_context(&mut (dec as *mut _)); + ff::avformat_close_input(&mut fmt); + cleanup_hw(hw_device); + return Err("avcodec_open2 (vaapi decode) failed".into()); + } + + // `mut` only to satisfy the move into the struct; the binding above is consumed. + let _ = &mut hw_device; + Ok(Self { + drm, + hw_device, + fmt, + dec, + pkt: ff::av_packet_alloc(), + frame: ff::av_frame_alloc(), + stream_index, + flushing: false, + }) + } + } + + /// The wgpu device the decoded textures live on (the DMA-BUF-import device). + pub fn device(&self) -> &wgpu::Device { + &self.drm.device + } + pub fn queue(&self) -> &wgpu::Queue { + &self.drm.queue + } + + /// Decode the next frame and import it as NV12 plane textures, or `Ok(None)` at end of stream. + pub fn next_frame(&mut self) -> Result, String> { + unsafe { + loop { + let r = ff::avcodec_receive_frame(self.dec, self.frame); + if r == 0 { + let imported = self.map_current(); + ff::av_frame_unref(self.frame); + return imported.map(Some); + } + if r == ff::AVERROR_EOF { + return Ok(None); + } + if r != averror(libc::EAGAIN) { + return Err(format!("avcodec_receive_frame failed: {r}")); + } + if self.flushing { + return Ok(None); // already drained the flush + } + + // Decoder wants more input: pump one packet (or signal EOF to flush). + let rp = ff::av_read_frame(self.fmt, self.pkt); + if rp < 0 { + self.flushing = true; + ff::avcodec_send_packet(self.dec, ptr::null()); + continue; + } + if (*self.pkt).stream_index == self.stream_index { + let rs = ff::avcodec_send_packet(self.dec, self.pkt); + ff::av_packet_unref(self.pkt); + if rs < 0 && rs != averror(libc::EAGAIN) { + return Err(format!("avcodec_send_packet failed: {rs}")); + } + } else { + ff::av_packet_unref(self.pkt); + } + } + } + } + + /// Map the just-decoded VAAPI surface (`self.frame`) to a DRM-PRIME DMA-BUF and import it. + unsafe fn map_current(&self) -> Result { + let drm_f = ff::av_frame_alloc(); + (*drm_f).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; + if ff::av_hwframe_map(drm_f, self.frame, flags) < 0 { + ff::av_frame_free(&mut (drm_f as *mut _)); + return Err("av_hwframe_map failed".into()); + } + let desc = (*drm_f).data[0] as *const ff::AVDRMFrameDescriptor; + let obj = &(*desc).objects[0]; + let width = (*self.frame).width as u32; + let height = (*self.frame).height as u32; + // NV12: Y then UV — either as two layers (one plane each) or one layer with two planes. + let (y_pl, uv_pl) = if (*desc).nb_layers >= 2 { + (&(*desc).layers[0].planes[0], &(*desc).layers[1].planes[0]) + } else { + (&(*desc).layers[0].planes[0], &(*desc).layers[0].planes[1]) + }; + let buf = Nv12DmaBuf { + fd: obj.fd, + size: obj.size as u64, + modifier: obj.format_modifier, + width, + height, + y_offset: y_pl.offset as u64, + y_pitch: y_pl.pitch as u64, + uv_offset: uv_pl.offset as u64, + uv_pitch: uv_pl.pitch as u64, + }; + let imported = dmabuf::import_raw(&self.drm, &buf); + ff::av_frame_free(&mut (drm_f as *mut _)); // the fd was dup'd into Vulkan + imported + } +} + +impl Drop for VaapiDecoder { + fn drop(&mut self) { + unsafe { + ff::av_frame_free(&mut (self.frame as *mut _)); + ff::av_packet_free(&mut (self.pkt as *mut _)); + ff::avcodec_free_context(&mut (self.dec as *mut _)); + if !self.fmt.is_null() { + ff::avformat_close_input(&mut self.fmt); + } + ff::av_buffer_unref(&mut self.hw_device); + } + } +} diff --git a/lightningbeam-ui/gpu-video-encoder/src/lib.rs b/lightningbeam-ui/gpu-video-encoder/src/lib.rs index 0d76fe4..b1c5c3b 100644 --- a/lightningbeam-ui/gpu-video-encoder/src/lib.rs +++ b/lightningbeam-ui/gpu-video-encoder/src/lib.rs @@ -30,6 +30,10 @@ pub mod dmabuf; #[cfg(target_os = "linux")] pub mod encoder; +/// VAAPI hardware decode → wgpu textures (Linux). +#[cfg(target_os = "linux")] +pub mod decoder; + #[cfg(test)] mod probe_tests { /// Confirm a headless GPU adapter is reachable (Vulkan on Linux/Intel). This gates diff --git a/lightningbeam-ui/gpu-video-encoder/tests/decode_roundtrip.rs b/lightningbeam-ui/gpu-video-encoder/tests/decode_roundtrip.rs new file mode 100644 index 0000000..81f9dc8 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/tests/decode_roundtrip.rs @@ -0,0 +1,91 @@ +//! Round-trip: encode solid frames with the zero-copy encoder, then hardware-decode them back +//! into a wgpu texture and read the Y plane. Verifies the VAAPI decode → DMA-BUF → wgpu import +//! path produces real pixels on the GPU. Skips when VAAPI is unavailable. + +#![cfg(target_os = "linux")] + +use gpu_video_encoder::decoder::VaapiDecoder; +use gpu_video_encoder::encoder::ZeroCopyEncoder; + +#[test] +fn vaapi_decode_roundtrip() { + // 256-wide so the R8 Y readback row (256 B) is already 256-aligned. + let (w, h) = (256u32, 256u32); + let out = std::env::temp_dir().join("gpu_video_encoder_decode_rt.mp4"); + let _ = std::fs::remove_file(&out); + + // --- Encode 10 frames of solid mid-gray. Full range → Y == luma ≈ 128. --- + { + let mut enc = match ZeroCopyEncoder::new(w, h, 30, 4000, &out, true) { + Ok(e) => e, + Err(e) => { + eprintln!("[decode-rt] encode unavailable, skipping: {e}"); + return; + } + }; + let device = enc.device(); + let src = device.create_texture(&wgpu::TextureDescriptor { + label: Some("gray"), + size: wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: wgpu::TextureFormat::Rgba8Unorm, + usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST, + view_formats: &[], + }); + let gray = vec![128u8; (w * h * 4) as usize]; + enc.queue().write_texture( + wgpu::TexelCopyTextureInfo { texture: &src, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All }, + &gray, + wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(w * 4), rows_per_image: Some(h) }, + wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, + ); + for _ in 0..10 { + enc.encode_rgba(&src).expect("encode_rgba"); + } + enc.finish().expect("finish"); + } + + // --- Decode it back on the GPU. --- + let mut dec = match VaapiDecoder::new(&out) { + Ok(d) => d, + Err(e) => { + eprintln!("[decode-rt] decode unavailable, skipping: {e}"); + return; + } + }; + let frame = dec.next_frame().expect("next_frame").expect("expected at least one frame"); + assert_eq!(frame.y().width(), w, "decoded Y width"); + assert_eq!(frame.y().height(), h, "decoded Y height"); + + // Read back the Y plane (R8) and check it's ≈ the gray we encoded. + let device = dec.device(); + let buf = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("y_readback"), + size: (w * h) as u64, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut cmd = device.create_command_encoder(&Default::default()); + cmd.copy_texture_to_buffer( + wgpu::TexelCopyTextureInfo { texture: frame.y(), mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All }, + wgpu::TexelCopyBufferInfo { + buffer: &buf, + layout: wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(w), rows_per_image: Some(h) }, + }, + wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, + ); + dec.queue().submit(Some(cmd.finish())); + buf.slice(..).map_async(wgpu::MapMode::Read, |_| {}); + let _ = device.poll(wgpu::PollType::wait_indefinitely()); + + let data = buf.slice(..).get_mapped_range(); + let mean = data.iter().map(|&b| b as f64).sum::() / data.len() as f64; + eprintln!("[decode-rt] decoded {w}x{h}, mean Y = {mean:.1}"); + assert!( + (mean - 128.0).abs() < 12.0, + "mean Y {mean} not ≈ 128 — decode produced wrong pixels" + ); + eprintln!("[decode-rt] ✅ VAAPI decode → wgpu texture verified"); +}