diff --git a/lightningbeam-ui/Cargo.lock b/lightningbeam-ui/Cargo.lock index 7fe5fcf..add318c 100644 --- a/lightningbeam-ui/Cargo.lock +++ b/lightningbeam-ui/Cargo.lock @@ -2849,6 +2849,19 @@ dependencies = [ "bitflags 2.10.0", ] +[[package]] +name = "gpu-video-encoder" +version = "0.1.0" +dependencies = [ + "ash", + "ffmpeg-sys-next", + "libc", + "pollster 0.4.0", + "wgpu", + "wgpu-hal", + "wgpu-types", +] + [[package]] name = "gtk" version = "0.18.2" diff --git a/lightningbeam-ui/Cargo.toml b/lightningbeam-ui/Cargo.toml index 7a66854..1c2c1b1 100644 --- a/lightningbeam-ui/Cargo.toml +++ b/lightningbeam-ui/Cargo.toml @@ -4,6 +4,7 @@ members = [ "lightningbeam-editor", "lightningbeam-core", "beamdsp", + "gpu-video-encoder", ] [workspace.dependencies] diff --git a/lightningbeam-ui/gpu-video-encoder/Cargo.toml b/lightningbeam-ui/gpu-video-encoder/Cargo.toml new file mode 100644 index 0000000..3722e12 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "gpu-video-encoder" +version = "0.1.0" +edition = "2021" +description = "Zero-copy GPU video encoding (RGBA->NV12 compute + hardware encoder interop). Unsafe FFI isolated here." + +[dependencies] +wgpu = { workspace = true } +# Raw Vulkan access for the DMA-BUF import. Versions MUST match what wgpu links +# (wgpu-hal 27.0.4 / ash 0.38) so the hal/ash types unify across the boundary. +wgpu-hal = { version = "27", features = ["vulkan"] } +wgpu-types = "27" +ash = "0.38" +# Raw FFmpeg FFI for VAAPI hwcontext + hardware encode. Matches the editor's +# ffmpeg-next 8.0 / static link so cargo unifies to one libav* across the build. +ffmpeg-sys-next = { version = "8.0", features = ["static"] } +libc = "0.2" + +[dev-dependencies] +pollster = "0.4" diff --git a/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs b/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs new file mode 100644 index 0000000..abcedbf --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/src/dmabuf.rs @@ -0,0 +1,135 @@ +//! Import a tiled VAAPI NV12 DMA-BUF as two wgpu textures (Y = R8, UV = RG8), aliasing +//! the one imported `VkDeviceMemory` at the plane offsets. Two single-format images are +//! used instead of one multi-planar image so each is an ordinary wgpu render target. +//! +//! Spike-grade: leaks the VkImages/memory on drop (process-scoped test). Cleanup +//! ordering (textures before memory) is a follow-up. + +use crate::vaapi::MappedSurface; +use crate::vk_device::DrmDevice; +use ash::vk; + +pub struct ImportedNv12 { + /// Luma plane, `R8Unorm`, full resolution. + pub y: wgpu::Texture, + /// Chroma plane, `Rg8Unorm`, half resolution (interleaved U,V). + pub uv: wgpu::Texture, +} + +pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result { + unsafe { + let device = &drm.raw_device; + let instance = &drm.raw_instance; + + let dup_fd = libc::dup(surf.fd); + if dup_fd < 0 { + return Err("dup(dma-buf fd) failed".into()); + } + + // --- create a single-plane DRM-modifier image --- + let make_image = |format: vk::Format, w: u32, h: u32, pitch: u64| -> Result { + let mut ext = vk::ExternalMemoryImageCreateInfo::default() + .handle_types(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT); + let plane_layouts = [vk::SubresourceLayout::default().offset(0).row_pitch(pitch)]; + let mut drm_info = vk::ImageDrmFormatModifierExplicitCreateInfoEXT::default() + .drm_format_modifier(surf.modifier) + .plane_layouts(&plane_layouts); + let info = vk::ImageCreateInfo::default() + .image_type(vk::ImageType::TYPE_2D) + .format(format) + .extent(vk::Extent3D { width: w, height: h, depth: 1 }) + .mip_levels(1) + .array_layers(1) + .samples(vk::SampleCountFlags::TYPE_1) + .tiling(vk::ImageTiling::DRM_FORMAT_MODIFIER_EXT) + .usage( + vk::ImageUsageFlags::COLOR_ATTACHMENT + | vk::ImageUsageFlags::TRANSFER_SRC + | vk::ImageUsageFlags::TRANSFER_DST, + ) + .sharing_mode(vk::SharingMode::EXCLUSIVE) + .initial_layout(vk::ImageLayout::UNDEFINED) + .push_next(&mut ext) + .push_next(&mut drm_info); + device + .create_image(&info, None) + .map_err(|e| format!("vkCreateImage(modifier) failed: {e:?}")) + }; + + let img_y = make_image(vk::Format::R8_UNORM, surf.width, surf.height, surf.y_pitch)?; + let img_uv = make_image(vk::Format::R8G8_UNORM, surf.width / 2, surf.height / 2, surf.uv_pitch)?; + + // --- import the dma-buf as one VkDeviceMemory, bind both planes --- + let fd_dev = ash::khr::external_memory_fd::Device::new(instance, device); + let mut fd_props = vk::MemoryFdPropertiesKHR::default(); + fd_dev + .get_memory_fd_properties(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT, dup_fd, &mut fd_props) + .map_err(|e| format!("vkGetMemoryFdPropertiesKHR failed: {e:?}"))?; + + let req_y = device.get_image_memory_requirements(img_y); + let req_uv = device.get_image_memory_requirements(img_uv); + let type_bits = fd_props.memory_type_bits & req_y.memory_type_bits & req_uv.memory_type_bits; + if type_bits == 0 { + return Err("no memory type compatible with dma-buf + both plane images".into()); + } + let mem_type = type_bits.trailing_zeros(); + + let mut import_info = vk::ImportMemoryFdInfoKHR::default() + .handle_type(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT) + .fd(dup_fd); + let alloc = vk::MemoryAllocateInfo::default() + .allocation_size(surf.size) + .memory_type_index(mem_type) + .push_next(&mut import_info); + let memory = device + .allocate_memory(&alloc, None) + .map_err(|e| format!("vkAllocateMemory(import dma-buf) failed: {e:?}"))?; + + device + .bind_image_memory(img_y, memory, surf.y_offset) + .map_err(|e| format!("bind Y plane: {e:?}"))?; + device + .bind_image_memory(img_uv, memory, surf.uv_offset) + .map_err(|e| format!("bind UV plane: {e:?}"))?; + + // --- wrap each VkImage as a wgpu texture --- + let hal_device = drm + .device + .as_hal::() + .ok_or("device is not Vulkan")?; + + let wrap = |img: vk::Image, format: wgpu::TextureFormat, w: u32, h: u32| -> wgpu::Texture { + let hal_desc = wgpu_hal::TextureDescriptor { + label: Some("vaapi-plane"), + size: wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format, + usage: wgpu_types::TextureUses::COLOR_TARGET | wgpu_types::TextureUses::COPY_SRC, + memory_flags: wgpu_hal::MemoryFlags::empty(), + view_formats: vec![], + }; + let hal_tex = hal_device.texture_from_raw(img, &hal_desc, None); + drm.device.create_texture_from_hal::( + hal_tex, + &wgpu::TextureDescriptor { + label: Some("vaapi-plane"), + size: wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format, + usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC, + view_formats: &[], + }, + ) + }; + + let y = wrap(img_y, wgpu::TextureFormat::R8Unorm, surf.width, surf.height); + let uv = wrap(img_uv, wgpu::TextureFormat::Rg8Unorm, surf.width / 2, surf.height / 2); + + // NOTE: img_y/img_uv/memory intentionally leaked for the spike (process-scoped). + Ok(ImportedNv12 { y, uv }) + } +} diff --git a/lightningbeam-ui/gpu-video-encoder/src/lib.rs b/lightningbeam-ui/gpu-video-encoder/src/lib.rs new file mode 100644 index 0000000..05aa9ad --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/src/lib.rs @@ -0,0 +1,54 @@ +//! Zero-copy GPU video encoding. +//! +//! Converts a rendered RGBA texture to the encoder's pixel format (NV12) on the GPU +//! and feeds it to a hardware video encoder without a CPU round-trip. All the unsafe +//! GPU↔encoder interop (Vulkan external memory / DMA-BUF → VAAPI on Linux, etc.) is +//! isolated in this crate. +//! +//! Status: scaffolding. Headless GPU probe + (next) NV12 compute live here first so +//! the GPU-side conversion can be validated against a CPU reference before any unsafe +//! interop is written. See `lightningbeam-ui/ZEROCOPY_GPU_ENCODE_PLAN.md`. + +pub mod nv12; + +/// VAAPI hardware encode (Linux-only; libva). +#[cfg(target_os = "linux")] +pub mod vaapi; + +/// Custom Vulkan device with DMA-BUF import extensions (Linux). +#[cfg(target_os = "linux")] +pub mod vk_device; + +/// Import a VAAPI NV12 DMA-BUF as wgpu textures (Linux). +#[cfg(target_os = "linux")] +pub mod dmabuf; + +#[cfg(test)] +mod probe_tests { + /// Confirm a headless GPU adapter is reachable (Vulkan on Linux/Intel). This gates + /// whether the GPU-side conversion can be tested on real hardware in this env. + #[test] + fn headless_adapter_available() { + let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor { + backends: wgpu::Backends::VULKAN | wgpu::Backends::GL, + ..Default::default() + }); + let adapter = pollster::block_on(instance.request_adapter( + &wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::HighPerformance, + force_fallback_adapter: false, + compatible_surface: None, + }, + )); + match adapter { + Ok(a) => { + let info = a.get_info(); + eprintln!( + "[gpu-probe] adapter: {} | backend={:?} | type={:?} | driver={}", + info.name, info.backend, info.device_type, info.driver + ); + } + Err(e) => panic!("no GPU adapter available headless: {e:?}"), + } + } +} diff --git a/lightningbeam-ui/gpu-video-encoder/src/nv12.rs b/lightningbeam-ui/gpu-video-encoder/src/nv12.rs new file mode 100644 index 0000000..44d3e2e --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/src/nv12.rs @@ -0,0 +1,207 @@ +//! GPU RGBA→NV12 conversion (BT.709 full-range), the pixel format hardware video +//! encoders (VAAPI/QSV/NVENC/VideoToolbox) consume. +//! +//! NV12 layout (what this writes, tight-packed into a storage buffer): +//! - `[0, W*H)` Y plane, one byte/pixel, row stride `W` +//! - `[W*H, W*H + W*H/2)` UV plane, interleaved `U,V` at 4:2:0, row stride `W` +//! (`W/2` chroma columns × 2 bytes), `H/2` rows +//! +//! Same BT.709 full-range matrix as the editor's planar YUV420p path, so colors match. +//! Requires `W % 8 == 0 && H % 2 == 0` (the shader packs 4 bytes per `u32`). + +/// `true` when [`Nv12Converter`] can handle these dimensions (else caller pads/falls back). +pub fn supports(width: u32, height: u32) -> bool { + width % 8 == 0 && height % 2 == 0 && width > 0 && height > 0 +} + +/// Tight NV12 byte length for `width`×`height`. +pub fn nv12_len(width: u32, height: u32) -> usize { + (width * height + width * (height / 2)) as usize +} + +/// Compute pipeline: `Rgba8Unorm` texture → tight NV12 storage buffer. +pub struct Nv12Converter { + y_pipeline: wgpu::ComputePipeline, + uv_pipeline: wgpu::ComputePipeline, + bind_group_layout: wgpu::BindGroupLayout, +} + +impl Nv12Converter { + pub fn new(device: &wgpu::Device) -> Self { + let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("nv12_bgl"), + entries: &[ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Texture { + sample_type: wgpu::TextureSampleType::Float { filterable: false }, + view_dimension: wgpu::TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }, + wgpu::BindGroupLayoutEntry { + binding: 1, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: false }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + ], + }); + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("nv12_pl"), + bind_group_layouts: &[&bind_group_layout], + push_constant_ranges: &[], + }); + let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("nv12_shader"), + source: wgpu::ShaderSource::Wgsl(SHADER.into()), + }); + let mk = |entry: &str| { + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("nv12_pipeline"), + layout: Some(&pipeline_layout), + module: &shader, + entry_point: Some(entry), + compilation_options: wgpu::PipelineCompilationOptions::default(), + cache: None, + }) + }; + Self { + y_pipeline: mk("y_main"), + uv_pipeline: mk("uv_main"), + bind_group_layout, + } + } + + /// Record RGBA→NV12 into `encoder`. `out_buffer` must be `STORAGE | COPY_SRC` of at + /// least [`nv12_len`] bytes. Caller must ensure [`supports`]`(width, height)`. + pub fn convert( + &self, + device: &wgpu::Device, + encoder: &mut wgpu::CommandEncoder, + rgba_view: &wgpu::TextureView, + out_buffer: &wgpu::Buffer, + width: u32, + height: u32, + ) { + debug_assert!(supports(width, height)); + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("nv12_bg"), + layout: &self.bind_group_layout, + entries: &[ + wgpu::BindGroupEntry { binding: 0, resource: wgpu::BindingResource::TextureView(rgba_view) }, + wgpu::BindGroupEntry { binding: 1, resource: out_buffer.as_entire_binding() }, + ], + }); + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("nv12_pass"), + timestamp_writes: None, + }); + pass.set_bind_group(0, &bind_group, &[]); + let wg = 8u32; + // Y: one thread per 4 horizontal luma samples. + pass.set_pipeline(&self.y_pipeline); + pass.dispatch_workgroups(((width / 4) + wg - 1) / wg, (height + wg - 1) / wg, 1); + // UV: one thread per 4 interleaved UV bytes = 2 chroma columns; (W/4)×(H/2) threads. + pass.set_pipeline(&self.uv_pipeline); + pass.dispatch_workgroups(((width / 4) + wg - 1) / wg, ((height / 2) + wg - 1) / wg, 1); + } +} + +/// CPU reference producing the exact bytes the shader should — used by tests to verify +/// the GPU output on real hardware. +pub fn cpu_reference(rgba: &[u8], width: u32, height: u32) -> Vec { + let w = width as usize; + let h = height as usize; + let mut out = vec![0u8; nv12_len(width, height)]; + let to_byte = |v: f32| (v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8; + let px = |x: usize, y: usize| { + let i = (y * w + x) * 4; + [rgba[i] as f32 / 255.0, rgba[i + 1] as f32 / 255.0, rgba[i + 2] as f32 / 255.0] + }; + // Y + for y in 0..h { + for x in 0..w { + let p = px(x, y); + out[y * w + x] = to_byte(0.2126 * p[0] + 0.7152 * p[1] + 0.0722 * p[2]); + } + } + // Interleaved UV (2x2 box average) + let y_size = w * h; + for cy in 0..h / 2 { + for cx in 0..w / 2 { + let mut acc = [0.0f32; 3]; + for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] { + let p = px(2 * cx + dx, 2 * cy + dy); + acc[0] += p[0]; acc[1] += p[1]; acc[2] += p[2]; + } + let a = [acc[0] / 4.0, acc[1] / 4.0, acc[2] / 4.0]; + let u = -0.1146 * a[0] - 0.3854 * a[1] + 0.5000 * a[2] + 0.5; + let v = 0.5000 * a[0] - 0.4542 * a[1] - 0.0458 * a[2] + 0.5; + out[y_size + cy * w + 2 * cx] = to_byte(u); + out[y_size + cy * w + 2 * cx + 1] = to_byte(v); + } + } + out +} + +const SHADER: &str = r#" +@group(0) @binding(0) var input_rgba: texture_2d; +@group(0) @binding(1) var out_buf: array; + +fn to_byte(v: f32) -> u32 { return u32(clamp(v, 0.0, 1.0) * 255.0 + 0.5); } + +// Y plane: pack 4 horizontal luma bytes. +@compute @workgroup_size(8, 8, 1) +fn y_main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(input_rgba); + let w = dims.x; + let h = dims.y; + let x4 = gid.x * 4u; + let y = gid.y; + if (x4 >= w || y >= h) { return; } + var packed: u32 = 0u; + for (var i = 0u; i < 4u; i = i + 1u) { + let c = textureLoad(input_rgba, vec2(x4 + i, y), 0).rgb; + let yy = 0.2126 * c.r + 0.7152 * c.g + 0.0722 * c.b; + packed = packed | (to_byte(yy) << (8u * i)); + } + out_buf[(y * w + x4) / 4u] = packed; +} + +// UV plane: each thread writes 4 interleaved bytes = U0 V0 U1 V1 for 2 chroma columns. +@compute @workgroup_size(8, 8, 1) +fn uv_main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(input_rgba); + let w = dims.x; + let h = dims.y; + let k = gid.x; // chroma-column pair index: covers columns 2k, 2k+1 + let cy = gid.y; + if (k * 2u >= w / 2u || cy >= h / 2u) { return; } + let y_size = w * h; + + var packed: u32 = 0u; + for (var j = 0u; j < 2u; j = j + 1u) { + let cx = 2u * k + j; // chroma column + let sx = 2u * cx; + let sy = 2u * cy; + let p00 = textureLoad(input_rgba, vec2(sx, sy), 0).rgb; + let p10 = textureLoad(input_rgba, vec2(sx + 1u, sy), 0).rgb; + let p01 = textureLoad(input_rgba, vec2(sx, sy + 1u), 0).rgb; + let p11 = textureLoad(input_rgba, vec2(sx + 1u, sy + 1u), 0).rgb; + let a = (p00 + p10 + p01 + p11) * 0.25; + let u = -0.1146 * a.r - 0.3854 * a.g + 0.5000 * a.b + 0.5; + let v = 0.5000 * a.r - 0.4542 * a.g - 0.0458 * a.b + 0.5; + packed = packed | (to_byte(u) << (16u * j)); // byte 0 or 2 + packed = packed | (to_byte(v) << (16u * j + 8u)); // byte 1 or 3 + } + // UV row stride is w bytes; this thread writes 4 bytes at column 4k. + out_buf[(y_size + cy * w + 4u * k) / 4u] = packed; +} +"#; diff --git a/lightningbeam-ui/gpu-video-encoder/src/vaapi.rs b/lightningbeam-ui/gpu-video-encoder/src/vaapi.rs new file mode 100644 index 0000000..d2eee0c --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/src/vaapi.rs @@ -0,0 +1,471 @@ +//! 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 +} + +/// 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 { + 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, 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 { + 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], + framerate: i32, + out_path: &str, +) -> Result { + 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 = 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, 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 = 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) + } +} diff --git a/lightningbeam-ui/gpu-video-encoder/src/vk_device.rs b/lightningbeam-ui/gpu-video-encoder/src/vk_device.rs new file mode 100644 index 0000000..946d364 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/src/vk_device.rs @@ -0,0 +1,149 @@ +//! Custom wgpu Vulkan device that additionally enables `VK_EXT_image_drm_format_modifier` +//! (plus the external-memory extensions wgpu-hal already turns on), so we can import a +//! tiled VAAPI NV12 DMA-BUF as a Vulkan image. wgpu's safe API can't add arbitrary device +//! extensions, so we build the `VkDevice` ourselves and wrap it via `device_from_raw`. +//! +//! All `unsafe` is contained here. Returns owned handles the caller must keep alive +//! together (instance → adapter → device/queue). + +use ash::vk; +use std::ffi::CStr; + +/// A wgpu device/queue backed by a hand-built Vulkan device with DMA-BUF import enabled. +pub struct DrmDevice { + // Order matters for drop; wgpu handles refcount internally but we keep these owned. + pub device: wgpu::Device, + pub queue: wgpu::Queue, + pub adapter: wgpu::Adapter, + pub instance: wgpu::Instance, + /// The raw VkDevice (for the ash image-import calls in `dmabuf.rs`). + pub raw_device: ash::Device, + pub raw_physical_device: vk::PhysicalDevice, + pub raw_instance: ash::Instance, +} + +/// Create the device, or `Err` if Vulkan/the extension isn't available (caller falls back). +pub fn create() -> Result { + unsafe { create_inner() } +} + +unsafe fn create_inner() -> Result { + use wgpu_hal::vulkan::Api as Vk; + // Bring the HAL Instance trait into scope for `init` / `enumerate_adapters`. + use wgpu_hal::Instance as _; + + // 1. HAL instance. + let hal_instance = wgpu_hal::vulkan::Instance::init(&wgpu_hal::InstanceDescriptor { + name: "gpu-video-encoder", + flags: wgpu::InstanceFlags::empty(), + memory_budget_thresholds: Default::default(), + backend_options: Default::default(), + }) + .map_err(|e| format!("vulkan instance init failed: {e:?}"))?; + + let ash_instance = hal_instance.shared_instance().raw_instance().clone(); + + // 2. Pick an adapter (prefer the integrated/discrete GPU). + let mut exposed_adapters = hal_instance.enumerate_adapters(None); + if exposed_adapters.is_empty() { + return Err("no Vulkan adapters".into()); + } + // Prefer a real GPU over CPU/llvmpipe. + exposed_adapters.sort_by_key(|a| match a.info.device_type { + wgpu::DeviceType::DiscreteGpu => 0, + wgpu::DeviceType::IntegratedGpu => 1, + _ => 2, + }); + let exposed = exposed_adapters.into_iter().next().unwrap(); + let phys = exposed.adapter.raw_physical_device(); + + // 3. Queue family with graphics + compute. + let qf_props = ash_instance.get_physical_device_queue_family_properties(phys); + let family_index = qf_props + .iter() + .position(|p| { + p.queue_flags + .contains(vk::QueueFlags::GRAPHICS | vk::QueueFlags::COMPUTE) + }) + .ok_or("no graphics+compute queue family")? as u32; + + // 4. Extensions: what wgpu-hal wants + DRM modifier import set. + let mut ext_names: Vec<&'static CStr> = + exposed.adapter.required_device_extensions(exposed.features); + // Only the genuine extensions; external_memory / bind_memory2 / ycbcr / format_list + // are core in Vulkan 1.1+ (this device is 1.3) so they need no enabling. + let extra: &[&'static CStr] = &[ + ash::ext::image_drm_format_modifier::NAME, + ash::khr::external_memory_fd::NAME, + ash::ext::external_memory_dma_buf::NAME, + ash::ext::queue_family_foreign::NAME, + ]; + for e in extra { + if !ext_names.contains(e) { + ext_names.push(e); + } + } + let ext_ptrs: Vec<*const i8> = ext_names.iter().map(|c| c.as_ptr()).collect(); + + // 5. Enable all supported physical-device features (so wgpu has what it needs) plus + // sampler YCbCr conversion (required for the NV12 multi-planar image). + let supported = ash_instance.get_physical_device_features(phys); + let mut ycbcr = + vk::PhysicalDeviceSamplerYcbcrConversionFeatures::default().sampler_ycbcr_conversion(true); + + let priorities = [1.0f32]; + let queue_info = vk::DeviceQueueCreateInfo::default() + .queue_family_index(family_index) + .queue_priorities(&priorities); + let queue_infos = [queue_info]; + + let create_info = vk::DeviceCreateInfo::default() + .queue_create_infos(&queue_infos) + .enabled_extension_names(&ext_ptrs) + .enabled_features(&supported) + .push_next(&mut ycbcr); + + let ash_device = ash_instance + .create_device(phys, &create_info, None) + .map_err(|e| format!("vkCreateDevice failed: {e:?}"))?; + + // 6. Wrap the raw device into a hal OpenDevice, then a wgpu device. + let open_device = exposed + .adapter + .device_from_raw( + ash_device.clone(), + None, + &ext_names, + exposed.features, + &wgpu::MemoryHints::default(), + family_index, + 0, + ) + .map_err(|e| format!("device_from_raw failed: {e:?}"))?; + + let raw_physical_device = phys; + + let wgpu_instance = wgpu::Instance::from_hal::(hal_instance); + let wgpu_adapter = wgpu_instance.create_adapter_from_hal::(exposed); + let (device, queue) = wgpu_adapter + .create_device_from_hal::( + open_device, + &wgpu::DeviceDescriptor { + label: Some("drm-import-device"), + required_features: wgpu::Features::empty(), + required_limits: wgpu::Limits::downlevel_defaults(), + ..Default::default() + }, + ) + .map_err(|e| format!("create_device_from_hal failed: {e:?}"))?; + + Ok(DrmDevice { + device, + queue, + adapter: wgpu_adapter, + instance: wgpu_instance, + raw_device: ash_device, + raw_physical_device, + raw_instance: ash_instance, + }) +} diff --git a/lightningbeam-ui/gpu-video-encoder/tests/drm_device.rs b/lightningbeam-ui/gpu-video-encoder/tests/drm_device.rs new file mode 100644 index 0000000..6441c13 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/tests/drm_device.rs @@ -0,0 +1,42 @@ +//! Step 1 of zero-copy: the custom Vulkan device with DMA-BUF import extensions builds +//! and can do a trivial GPU op. Skips (passes) when Vulkan is unavailable. + +#![cfg(target_os = "linux")] + +#[test] +fn drm_device_creates_and_works() { + let dev = match gpu_video_encoder::vk_device::create() { + Ok(d) => d, + Err(e) => { + eprintln!("[drm-device] unavailable, skipping: {e}"); + return; + } + }; + eprintln!("[drm-device] created custom Vulkan device OK"); + + // Trivial sanity op: write+read a small buffer, proving the wrapped device is usable. + let data: Vec = (0..256u32).map(|i| i as u8).collect(); + let src = wgpu::util::DeviceExt::create_buffer_init( + &dev.device, + &wgpu::util::BufferInitDescriptor { + label: Some("src"), + contents: &data, + usage: wgpu::BufferUsages::COPY_SRC, + }, + ); + let dst = dev.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("dst"), + size: 256, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut enc = dev.device.create_command_encoder(&Default::default()); + enc.copy_buffer_to_buffer(&src, 0, &dst, 0, 256); + dev.queue.submit(Some(enc.finish())); + let slice = dst.slice(..); + slice.map_async(wgpu::MapMode::Read, |_| {}); + let _ = dev.device.poll(wgpu::PollType::wait_indefinitely()); + let got = slice.get_mapped_range().to_vec(); + assert_eq!(got, data, "round-trip through custom device failed"); + eprintln!("[drm-device] buffer round-trip OK on custom device"); +} diff --git a/lightningbeam-ui/gpu-video-encoder/tests/gpu_nv12.rs b/lightningbeam-ui/gpu-video-encoder/tests/gpu_nv12.rs new file mode 100644 index 0000000..9544308 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/tests/gpu_nv12.rs @@ -0,0 +1,117 @@ +//! Real-hardware test: run the RGBA→NV12 compute on the GPU and check it byte-matches +//! the CPU reference. Skips (passes) if no GPU adapter is available. + +use gpu_video_encoder::nv12::{cpu_reference, nv12_len, Nv12Converter}; + +fn device_queue() -> Option<(wgpu::Device, wgpu::Queue)> { + let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor { + backends: wgpu::Backends::VULKAN | wgpu::Backends::GL, + ..Default::default() + }); + let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::HighPerformance, + force_fallback_adapter: false, + compatible_surface: None, + })) + .ok()?; + pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("nv12-test"), + required_features: wgpu::Features::empty(), + required_limits: wgpu::Limits::downlevel_defaults(), + ..Default::default() + })) + .ok() +} + +/// A deterministic, varied RGBA pattern so luma and 2x2 chroma subsampling are exercised. +fn pattern(w: u32, h: u32) -> Vec { + let mut v = Vec::with_capacity((w * h * 4) as usize); + for y in 0..h { + for x in 0..w { + v.push(((x * 37 + y * 11) % 256) as u8); // R + v.push(((x * 5 + y * 53) % 256) as u8); // G + v.push(((x * 97 + y * 17) % 256) as u8); // B + v.push(255); + } + } + v +} + +#[test] +fn gpu_nv12_matches_cpu_reference() { + let Some((device, queue)) = device_queue() else { + eprintln!("[gpu_nv12] no GPU adapter; skipping"); + return; + }; + + let (w, h) = (64u32, 16u32); + let rgba = pattern(w, h); + + // Source RGBA texture. + let tex = device.create_texture(&wgpu::TextureDescriptor { + label: Some("src_rgba"), + 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: &[], + }); + queue.write_texture( + wgpu::TexelCopyTextureInfo { + texture: &tex, + mip_level: 0, + origin: wgpu::Origin3d::ZERO, + aspect: wgpu::TextureAspect::All, + }, + &rgba, + 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 }, + ); + let view = tex.create_view(&Default::default()); + + let len = nv12_len(w, h) as u64; + let out = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("nv12_out"), + size: len, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let staging = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("nv12_staging"), + size: len, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + + let conv = Nv12Converter::new(&device); + let mut enc = device.create_command_encoder(&Default::default()); + conv.convert(&device, &mut enc, &view, &out, w, h); + enc.copy_buffer_to_buffer(&out, 0, &staging, 0, len); + queue.submit(Some(enc.finish())); + + let slice = staging.slice(..); + slice.map_async(wgpu::MapMode::Read, |_| {}); + let _ = device.poll(wgpu::PollType::wait_indefinitely()); + let gpu = slice.get_mapped_range().to_vec(); + + let cpu = cpu_reference(&rgba, w, h); + assert_eq!(gpu.len(), cpu.len(), "length mismatch"); + + // Allow ±1 for rounding differences between GPU and CPU float paths. + let mut max_diff = 0i32; + let mut nbad = 0; + for (i, (g, c)) in gpu.iter().zip(cpu.iter()).enumerate() { + let d = (*g as i32 - *c as i32).abs(); + max_diff = max_diff.max(d); + if d > 1 { + nbad += 1; + if nbad <= 8 { + eprintln!("[gpu_nv12] byte {i}: gpu={g} cpu={c} (diff {d})"); + } + } + } + eprintln!("[gpu_nv12] {}x{} NV12, max byte diff = {max_diff}", w, h); + assert_eq!(nbad, 0, "{nbad} bytes differ from CPU reference by >1"); +} diff --git a/lightningbeam-ui/gpu-video-encoder/tests/vaapi_encode.rs b/lightningbeam-ui/gpu-video-encoder/tests/vaapi_encode.rs new file mode 100644 index 0000000..2249e8a --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/tests/vaapi_encode.rs @@ -0,0 +1,66 @@ +//! Level-1 spike: prove `h264_vaapi` encodes NV12 in this environment. Skips (passes) +//! when VAAPI isn't available so it's a no-op on CI/macOS/Windows. + +#![cfg(target_os = "linux")] + +use gpu_video_encoder::nv12::{cpu_reference, nv12_len}; +use gpu_video_encoder::vaapi::encode_nv12_to_file; + +/// A moving-gradient RGBA pattern → NV12 via the CPU reference, so we feed valid frames. +fn nv12_frames(w: u32, h: u32, n: usize) -> Vec> { + (0..n) + .map(|f| { + let mut rgba = Vec::with_capacity((w * h * 4) as usize); + for y in 0..h { + for x in 0..w { + rgba.push(((x + f as u32 * 4) % 256) as u8); + rgba.push(((y + f as u32 * 2) % 256) as u8); + rgba.push(((x + y) % 256) as u8); + rgba.push(255); + } + } + let v = cpu_reference(&rgba, w, h); + assert_eq!(v.len(), nv12_len(w, h)); + v + }) + .collect() +} + +#[test] +fn vaapi_surface_drm_layout() { + match gpu_video_encoder::vaapi::probe_surface_drm(1920, 1088) { + Ok(s) => eprintln!("[vaapi-drm]\n{s}"), + Err(e) => eprintln!("[vaapi-drm] unavailable, skipping: {e}"), + } +} + +#[test] +fn vaapi_h264_encode_smoke() { + let (w, h) = (320u32, 240u32); + let frames = nv12_frames(w, h, 30); + let out = std::env::temp_dir().join("gpu_video_encoder_vaapi_smoke.h264"); + let out_str = out.to_str().unwrap(); + + match encode_nv12_to_file(w, h, &frames, 30, out_str) { + Ok(packets) => { + let meta = std::fs::metadata(&out).expect("output file missing"); + eprintln!( + "[vaapi] encoded {} packets, {} bytes -> {}", + packets, + meta.len(), + out_str + ); + assert!(packets > 0, "no packets produced"); + assert!(meta.len() > 0, "empty output file"); + // First frame should be an IDR; Annex-B starts with a start code. + let head = std::fs::read(&out).unwrap(); + assert!( + head.starts_with(&[0, 0, 0, 1]) || head.starts_with(&[0, 0, 1]), + "output is not Annex-B H.264 (no start code)" + ); + } + Err(e) => { + eprintln!("[vaapi] unavailable, skipping: {e}"); + } + } +} diff --git a/lightningbeam-ui/gpu-video-encoder/tests/zerocopy.rs b/lightningbeam-ui/gpu-video-encoder/tests/zerocopy.rs new file mode 100644 index 0000000..5e6d126 --- /dev/null +++ b/lightningbeam-ui/gpu-video-encoder/tests/zerocopy.rs @@ -0,0 +1,95 @@ +//! End-to-end zero-copy proof: import a VAAPI NV12 surface as wgpu textures, render +//! known values into them via Vulkan, read the surface back, and verify the bytes — +//! proving the GPU wrote straight into the encoder's surface with no CPU upload. + +#![cfg(target_os = "linux")] + +use gpu_video_encoder::{dmabuf, vaapi, vk_device}; + +#[test] +fn zerocopy_render_into_vaapi_surface() { + let drm = match vk_device::create() { + Ok(d) => d, + Err(e) => { + eprintln!("[zerocopy] no Vulkan device, skipping: {e}"); + return; + } + }; + let surf = match vaapi::MappedSurface::alloc(640, 480) { + Ok(s) => s, + Err(e) => { + eprintln!("[zerocopy] no VAAPI surface, skipping: {e}"); + return; + } + }; + eprintln!( + "[zerocopy] surface: modifier=0x{:016x} y(off={},pitch={}) uv(off={},pitch={}) size={}", + surf.modifier, surf.y_offset, surf.y_pitch, surf.uv_offset, surf.uv_pitch, surf.size + ); + + let imported = match dmabuf::import(&drm, &surf) { + Ok(i) => i, + Err(e) => panic!("dma-buf import failed: {e}"), + }; + eprintln!("[zerocopy] imported surface as wgpu Y(R8) + UV(RG8) textures"); + + // Render known constants via clear: Y=0.5(->128), U=0.25(->64), V=0.75(->191). + let y_view = imported.y.create_view(&Default::default()); + let uv_view = imported.uv.create_view(&Default::default()); + let mut enc = drm.device.create_command_encoder(&Default::default()); + { + enc.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("clear-y"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: &y_view, + resolve_target: None, + depth_slice: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Clear(wgpu::Color { r: 0.5, g: 0.0, b: 0.0, a: 0.0 }), + store: wgpu::StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, + }); + enc.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("clear-uv"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: &uv_view, + resolve_target: None, + depth_slice: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Clear(wgpu::Color { r: 0.25, g: 0.75, b: 0.0, a: 0.0 }), + store: wgpu::StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, + }); + } + drm.queue.submit(Some(enc.finish())); + let _ = drm.device.poll(wgpu::PollType::wait_indefinitely()); + + // Read the VAAPI surface back and check what the GPU wrote. + let nv12 = surf.readback_nv12().expect("readback"); + let (w, h) = (640usize, 480usize); + let y_plane = &nv12[..w * h]; + let uv_plane = &nv12[w * h..]; + + let near = |v: u8, t: i32| (v as i32 - t).abs() <= 3; + let y_ok = y_plane.iter().filter(|&&v| near(v, 128)).count(); + let u_ok = uv_plane.iter().step_by(2).filter(|&&v| near(v, 64)).count(); + let v_ok = uv_plane.iter().skip(1).step_by(2).filter(|&&v| near(v, 191)).count(); + eprintln!( + "[zerocopy] Y~128: {}/{}, U~64: {}/{}, V~191: {}/{}", + y_ok, w * h, u_ok, uv_plane.len() / 2, v_ok, uv_plane.len() / 2 + ); + + let frac = |ok: usize, n: usize| ok as f64 / n as f64; + assert!(frac(y_ok, w * h) > 0.98, "Y plane not the rendered value (sample {:?})", &y_plane[..8]); + assert!(frac(u_ok, uv_plane.len() / 2) > 0.98, "U not rendered value"); + assert!(frac(v_ok, uv_plane.len() / 2) > 0.98, "V not rendered value"); + eprintln!("[zerocopy] ✅ GPU rendered straight into the VAAPI surface (verified via readback)"); +}