gpu-video-encoder: complete zero-copy H.264 encode pipeline
Build the full end-to-end zero-copy encoder, validated on Intel/VAAPI: - render_nv12: fragment-shader RGBA->NV12 that renders luma/chroma into the imported R8/RG8 plane render targets (compute storage can't write the DMA-BUF-backed planes; render attachments can). - dmabuf: import_raw imports an NV12 DMA-BUF by explicit layout; the two plane images + shared memory are now destroyed by wgpu via texture_from_raw drop callbacks (Arc MemoryGuard frees the memory once both images are gone, in wgpu's wait-idle'd deferred pass) -- fixes the teardown segfault. - encoder::ZeroCopyEncoder: renders an RGBA texture straight into a pooled VAAPI surface (imports cached by VASurface id) and encodes with h264_vaapi. encode_rgba + finish; the caller renders on device(). Tests: real-frame render into the surface matches the CPU NV12 reference, and a 30-frame encode produces valid H.264 (ffprobe-verified) with clean teardown. Not yet wired into the editor.
This commit is contained in:
Skyler Lehmkuhl
parent
5917ce7921
commit
ba897eaea2
|
|
@ -1,38 +1,90 @@
|
|||
//! 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,
|
||||
/// Plane layout for a single-object NV12 DMA-BUF (the common VAAPI case).
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct Nv12DmaBuf {
|
||||
pub fd: i32,
|
||||
pub size: u64,
|
||||
pub modifier: u64,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub y_offset: u64,
|
||||
pub y_pitch: u64,
|
||||
pub uv_offset: u64,
|
||||
pub uv_pitch: u64,
|
||||
}
|
||||
|
||||
/// Frees the shared imported `VkDeviceMemory` once both plane images are gone. Held by
|
||||
/// both textures' drop callbacks (via `Arc`); the last one to run frees the memory —
|
||||
/// after wgpu has destroyed the images, in its wait-idle'd deferred-destruction pass.
|
||||
struct MemoryGuard {
|
||||
device: ash::Device,
|
||||
memory: vk::DeviceMemory,
|
||||
}
|
||||
impl Drop for MemoryGuard {
|
||||
fn drop(&mut self) {
|
||||
unsafe { self.device.free_memory(self.memory, None) };
|
||||
}
|
||||
}
|
||||
|
||||
/// A VAAPI surface imported as two wgpu plane textures. The underlying Vulkan image/
|
||||
/// memory are destroyed by wgpu (via drop callbacks) when these textures drop.
|
||||
pub struct ImportedNv12 {
|
||||
y: wgpu::Texture,
|
||||
uv: wgpu::Texture,
|
||||
}
|
||||
|
||||
impl ImportedNv12 {
|
||||
pub fn y(&self) -> &wgpu::Texture {
|
||||
&self.y
|
||||
}
|
||||
pub fn uv(&self) -> &wgpu::Texture {
|
||||
&self.uv
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience: map a freshly-allocated `MappedSurface` and import it.
|
||||
pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result<ImportedNv12, String> {
|
||||
import_raw(
|
||||
drm,
|
||||
&Nv12DmaBuf {
|
||||
fd: surf.fd,
|
||||
size: surf.size,
|
||||
modifier: surf.modifier,
|
||||
width: surf.width,
|
||||
height: surf.height,
|
||||
y_offset: surf.y_offset,
|
||||
y_pitch: surf.y_pitch,
|
||||
uv_offset: surf.uv_offset,
|
||||
uv_pitch: surf.uv_pitch,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Import an NV12 DMA-BUF (described by `buf`) as two wgpu plane textures. The fd is
|
||||
/// duplicated, so the caller keeps ownership of theirs.
|
||||
pub fn import_raw(drm: &DrmDevice, buf: &Nv12DmaBuf) -> Result<ImportedNv12, String> {
|
||||
unsafe {
|
||||
let device = &drm.raw_device;
|
||||
let device = drm.raw_device.clone();
|
||||
let instance = &drm.raw_instance;
|
||||
|
||||
let dup_fd = libc::dup(surf.fd);
|
||||
let dup_fd = libc::dup(buf.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<vk::Image, String> {
|
||||
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)
|
||||
.drm_format_modifier(buf.modifier)
|
||||
.plane_layouts(&plane_layouts);
|
||||
let info = vk::ImageCreateInfo::default()
|
||||
.image_type(vk::ImageType::TYPE_2D)
|
||||
|
|
@ -56,16 +108,14 @@ pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result<ImportedNv12, Str
|
|||
.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)?;
|
||||
let img_y = make_image(vk::Format::R8_UNORM, buf.width, buf.height, buf.y_pitch)?;
|
||||
let img_uv = make_image(vk::Format::R8G8_UNORM, buf.width / 2, buf.height / 2, buf.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 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;
|
||||
|
|
@ -78,7 +128,7 @@ pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result<ImportedNv12, Str
|
|||
.handle_type(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT)
|
||||
.fd(dup_fd);
|
||||
let alloc = vk::MemoryAllocateInfo::default()
|
||||
.allocation_size(surf.size)
|
||||
.allocation_size(buf.size)
|
||||
.memory_type_index(mem_type)
|
||||
.push_next(&mut import_info);
|
||||
let memory = device
|
||||
|
|
@ -86,19 +136,28 @@ pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result<ImportedNv12, Str
|
|||
.map_err(|e| format!("vkAllocateMemory(import dma-buf) failed: {e:?}"))?;
|
||||
|
||||
device
|
||||
.bind_image_memory(img_y, memory, surf.y_offset)
|
||||
.bind_image_memory(img_y, memory, buf.y_offset)
|
||||
.map_err(|e| format!("bind Y plane: {e:?}"))?;
|
||||
device
|
||||
.bind_image_memory(img_uv, memory, surf.uv_offset)
|
||||
.bind_image_memory(img_uv, memory, buf.uv_offset)
|
||||
.map_err(|e| format!("bind UV plane: {e:?}"))?;
|
||||
|
||||
// --- wrap each VkImage as a wgpu texture ---
|
||||
// Shared guard: frees `memory` once both images' drop callbacks have run.
|
||||
let mem_guard = std::sync::Arc::new(MemoryGuard { device: device.clone(), memory });
|
||||
|
||||
let hal_device = drm
|
||||
.device
|
||||
.as_hal::<wgpu_hal::vulkan::Api>()
|
||||
.ok_or("device is not Vulkan")?;
|
||||
|
||||
let wrap = |img: vk::Image, format: wgpu::TextureFormat, w: u32, h: u32| -> wgpu::Texture {
|
||||
let mut wrap = |img: vk::Image, format: wgpu::TextureFormat, w: u32, h: u32| -> wgpu::Texture {
|
||||
// wgpu destroys the image (after wait-idle) when the texture drops; the
|
||||
// captured Arc<MemoryGuard> frees the shared memory once both have run.
|
||||
let dev = device.clone();
|
||||
let guard = mem_guard.clone();
|
||||
let cb: wgpu_hal::DropCallback = Box::new(move || {
|
||||
unsafe { dev.destroy_image(img, None) };
|
||||
drop(guard);
|
||||
});
|
||||
let hal_desc = wgpu_hal::TextureDescriptor {
|
||||
label: Some("vaapi-plane"),
|
||||
size: wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 },
|
||||
|
|
@ -110,7 +169,7 @@ pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result<ImportedNv12, Str
|
|||
memory_flags: wgpu_hal::MemoryFlags::empty(),
|
||||
view_formats: vec![],
|
||||
};
|
||||
let hal_tex = hal_device.texture_from_raw(img, &hal_desc, None);
|
||||
let hal_tex = hal_device.texture_from_raw(img, &hal_desc, Some(cb));
|
||||
drm.device.create_texture_from_hal::<wgpu_hal::vulkan::Api>(
|
||||
hal_tex,
|
||||
&wgpu::TextureDescriptor {
|
||||
|
|
@ -125,11 +184,10 @@ pub fn import(drm: &DrmDevice, surf: &MappedSurface) -> Result<ImportedNv12, Str
|
|||
},
|
||||
)
|
||||
};
|
||||
let y = wrap(img_y, wgpu::TextureFormat::R8Unorm, buf.width, buf.height);
|
||||
let uv = wrap(img_uv, wgpu::TextureFormat::Rg8Unorm, buf.width / 2, buf.height / 2);
|
||||
drop(hal_device);
|
||||
|
||||
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 })
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,224 @@
|
|||
//! End-to-end zero-copy H.264 encoder: render an RGBA wgpu texture straight into a VAAPI
|
||||
//! NV12 surface (no CPU copy) and encode it with `h264_vaapi`. The caller renders frames
|
||||
//! on [`ZeroCopyEncoder::device`] (the custom Vulkan device with DMA-BUF import enabled).
|
||||
//!
|
||||
//! Imports are cached by VASurface id, so the pooled surfaces are imported once each.
|
||||
|
||||
use crate::dmabuf::{self, ImportedNv12, Nv12DmaBuf};
|
||||
use crate::render_nv12::Rgba2Nv12;
|
||||
use crate::vk_device::{self, DrmDevice};
|
||||
use ffmpeg_sys_next as ff;
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::CString;
|
||||
use std::ptr;
|
||||
|
||||
#[inline]
|
||||
fn averror(e: i32) -> i32 {
|
||||
-e
|
||||
}
|
||||
|
||||
pub struct ZeroCopyEncoder {
|
||||
drm: DrmDevice,
|
||||
renderer: Rgba2Nv12,
|
||||
hw_device: *mut ff::AVBufferRef,
|
||||
frames_ref: *mut ff::AVBufferRef,
|
||||
enc: *mut ff::AVCodecContext,
|
||||
pkt: *mut ff::AVPacket,
|
||||
width: u32,
|
||||
height: u32,
|
||||
pts: i64,
|
||||
cache: HashMap<usize, ImportedNv12>,
|
||||
out: Vec<u8>,
|
||||
}
|
||||
|
||||
impl ZeroCopyEncoder {
|
||||
/// Build a zero-copy `h264_vaapi` encoder, or `Err` if VAAPI/the device is unavailable.
|
||||
pub fn new(width: u32, height: u32, framerate: i32, bitrate_kbps: u32) -> Result<Self, String> {
|
||||
let drm = vk_device::create()?;
|
||||
let renderer = Rgba2Nv12::new(&drm.device);
|
||||
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 name = CString::new("h264_vaapi").unwrap();
|
||||
let codec = ff::avcodec_find_encoder_by_name(name.as_ptr());
|
||||
if codec.is_null() {
|
||||
ff::av_buffer_unref(&mut hw_device);
|
||||
return Err("h264_vaapi not found".into());
|
||||
}
|
||||
let enc = ff::avcodec_alloc_context3(codec);
|
||||
(*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;
|
||||
(*enc).bit_rate = (bitrate_kbps as i64) * 1000;
|
||||
|
||||
let frames_ref = ff::av_hwframe_ctx_alloc(hw_device);
|
||||
{
|
||||
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 = 16;
|
||||
}
|
||||
if ff::av_hwframe_ctx_init(frames_ref) < 0 {
|
||||
let mut fr = frames_ref;
|
||||
ff::av_buffer_unref(&mut fr);
|
||||
ff::avcodec_free_context(&mut (enc as *mut _));
|
||||
ff::av_buffer_unref(&mut hw_device);
|
||||
return Err("av_hwframe_ctx_init failed".into());
|
||||
}
|
||||
(*enc).hw_frames_ctx = ff::av_buffer_ref(frames_ref);
|
||||
|
||||
if ff::avcodec_open2(enc, codec, ptr::null_mut()) < 0 {
|
||||
let mut fr = frames_ref;
|
||||
ff::av_buffer_unref(&mut fr);
|
||||
ff::avcodec_free_context(&mut (enc as *mut _));
|
||||
ff::av_buffer_unref(&mut hw_device);
|
||||
return Err("avcodec_open2(h264_vaapi) failed".into());
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
drm,
|
||||
renderer,
|
||||
hw_device,
|
||||
frames_ref,
|
||||
enc,
|
||||
pkt: ff::av_packet_alloc(),
|
||||
width,
|
||||
height,
|
||||
pts: 0,
|
||||
cache: HashMap::new(),
|
||||
out: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The wgpu device frames must be rendered on (so the RGBA texture is importable).
|
||||
pub fn device(&self) -> &wgpu::Device {
|
||||
&self.drm.device
|
||||
}
|
||||
pub fn queue(&self) -> &wgpu::Queue {
|
||||
&self.drm.queue
|
||||
}
|
||||
|
||||
/// Render `rgba` (an `Rgba8Unorm` texture on [`Self::device`], `TEXTURE_BINDING`)
|
||||
/// into a VAAPI surface and encode it. Appends any produced packets internally.
|
||||
pub fn encode_rgba(&mut self, rgba: &wgpu::Texture) -> Result<(), String> {
|
||||
unsafe {
|
||||
let surf = ff::av_frame_alloc();
|
||||
if ff::av_hwframe_get_buffer(self.frames_ref, surf, 0) < 0 {
|
||||
ff::av_frame_free(&mut (surf as *mut _));
|
||||
return Err("av_hwframe_get_buffer failed".into());
|
||||
}
|
||||
let id = (*surf).data[3] as usize; // VASurfaceID
|
||||
|
||||
if !self.cache.contains_key(&id) {
|
||||
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
|
||||
| ff::AV_HWFRAME_MAP_WRITE as i32;
|
||||
if ff::av_hwframe_map(drm_f, surf, flags) < 0 {
|
||||
ff::av_frame_free(&mut (drm_f as *mut _));
|
||||
ff::av_frame_free(&mut (surf 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 y = &(*desc).layers[0].planes[0];
|
||||
let uv = &(*desc).layers[1].planes[0];
|
||||
let buf = Nv12DmaBuf {
|
||||
fd: obj.fd,
|
||||
size: obj.size as u64,
|
||||
modifier: obj.format_modifier,
|
||||
width: self.width,
|
||||
height: self.height,
|
||||
y_offset: y.offset as u64,
|
||||
y_pitch: y.pitch as u64,
|
||||
uv_offset: uv.offset as u64,
|
||||
uv_pitch: uv.pitch as u64,
|
||||
};
|
||||
let imported = match dmabuf::import_raw(&self.drm, &buf) {
|
||||
Ok(i) => i,
|
||||
Err(e) => {
|
||||
ff::av_frame_free(&mut (drm_f as *mut _));
|
||||
ff::av_frame_free(&mut (surf as *mut _));
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
ff::av_frame_free(&mut (drm_f as *mut _)); // fd was dup'd into Vulkan
|
||||
self.cache.insert(id, imported);
|
||||
}
|
||||
|
||||
// Render RGBA -> NV12 directly into the surface planes.
|
||||
let imp = self.cache.get(&id).unwrap();
|
||||
let rgba_view = rgba.create_view(&Default::default());
|
||||
let y_view = imp.y().create_view(&Default::default());
|
||||
let uv_view = imp.uv().create_view(&Default::default());
|
||||
let mut cmd = self.drm.device.create_command_encoder(&Default::default());
|
||||
self.renderer.convert(&self.drm.device, &mut cmd, &rgba_view, &y_view, &uv_view);
|
||||
self.drm.queue.submit(Some(cmd.finish()));
|
||||
let _ = self.drm.device.poll(wgpu::PollType::wait_indefinitely());
|
||||
|
||||
// Encode the surface.
|
||||
(*surf).pts = self.pts;
|
||||
self.pts += 1;
|
||||
let r = ff::avcodec_send_frame(self.enc, surf);
|
||||
ff::av_frame_free(&mut (surf as *mut _));
|
||||
if r < 0 {
|
||||
return Err(format!("avcodec_send_frame failed: {r}"));
|
||||
}
|
||||
self.drain()
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn drain(&mut self) -> Result<(), String> {
|
||||
loop {
|
||||
let r = ff::avcodec_receive_packet(self.enc, self.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((*self.pkt).data, (*self.pkt).size as usize);
|
||||
self.out.extend_from_slice(data);
|
||||
ff::av_packet_unref(self.pkt);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Flush the encoder and return the accumulated Annex-B H.264 bitstream.
|
||||
pub fn finish(mut self) -> Result<Vec<u8>, String> {
|
||||
unsafe {
|
||||
ff::avcodec_send_frame(self.enc, ptr::null_mut());
|
||||
self.drain()?;
|
||||
}
|
||||
Ok(std::mem::take(&mut self.out))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ZeroCopyEncoder {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
self.cache.clear(); // frees imported Vulkan resources first
|
||||
ff::av_packet_free(&mut (self.pkt as *mut _));
|
||||
ff::avcodec_free_context(&mut (self.enc as *mut _));
|
||||
let mut fr = self.frames_ref;
|
||||
ff::av_buffer_unref(&mut fr);
|
||||
ff::av_buffer_unref(&mut self.hw_device);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -11,6 +11,9 @@
|
|||
|
||||
pub mod nv12;
|
||||
|
||||
/// Fragment-shader RGBA→NV12 conversion that renders into plane textures.
|
||||
pub mod render_nv12;
|
||||
|
||||
/// VAAPI hardware encode (Linux-only; libva).
|
||||
#[cfg(target_os = "linux")]
|
||||
pub mod vaapi;
|
||||
|
|
@ -23,6 +26,10 @@ pub mod vk_device;
|
|||
#[cfg(target_os = "linux")]
|
||||
pub mod dmabuf;
|
||||
|
||||
/// End-to-end zero-copy `h264_vaapi` encoder (Linux).
|
||||
#[cfg(target_os = "linux")]
|
||||
pub mod encoder;
|
||||
|
||||
#[cfg(test)]
|
||||
mod probe_tests {
|
||||
/// Confirm a headless GPU adapter is reachable (Vulkan on Linux/Intel). This gates
|
||||
|
|
|
|||
|
|
@ -0,0 +1,145 @@
|
|||
//! Fragment-shader RGBA→NV12 conversion that **renders** luma/chroma into the encoder
|
||||
//! surface's plane textures (R8 Y, RG8 UV). Render targets (not compute storage) so it
|
||||
//! works with the DMA-BUF-imported plane images, which aren't storage-writable.
|
||||
//!
|
||||
//! BT.709 full-range, matching `nv12::cpu_reference` and the encoder's color tags.
|
||||
|
||||
/// Converts a bound RGBA texture into a Y plane (R8) and a UV plane (RG8) via two passes.
|
||||
pub struct Rgba2Nv12 {
|
||||
y_pipeline: wgpu::RenderPipeline,
|
||||
uv_pipeline: wgpu::RenderPipeline,
|
||||
bgl: wgpu::BindGroupLayout,
|
||||
}
|
||||
|
||||
impl Rgba2Nv12 {
|
||||
pub fn new(device: &wgpu::Device) -> Self {
|
||||
let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||||
label: Some("rgba2nv12_bgl"),
|
||||
entries: &[wgpu::BindGroupLayoutEntry {
|
||||
binding: 0,
|
||||
visibility: wgpu::ShaderStages::FRAGMENT,
|
||||
ty: wgpu::BindingType::Texture {
|
||||
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||||
view_dimension: wgpu::TextureViewDimension::D2,
|
||||
multisampled: false,
|
||||
},
|
||||
count: None,
|
||||
}],
|
||||
});
|
||||
let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
|
||||
label: Some("rgba2nv12_pl"),
|
||||
bind_group_layouts: &[&bgl],
|
||||
push_constant_ranges: &[],
|
||||
});
|
||||
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||||
label: Some("rgba2nv12_shader"),
|
||||
source: wgpu::ShaderSource::Wgsl(SHADER.into()),
|
||||
});
|
||||
let mk = |fs: &str, fmt: wgpu::TextureFormat| {
|
||||
device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
|
||||
label: Some("rgba2nv12_pipeline"),
|
||||
layout: Some(&layout),
|
||||
vertex: wgpu::VertexState {
|
||||
module: &shader,
|
||||
entry_point: Some("vs_main"),
|
||||
buffers: &[],
|
||||
compilation_options: Default::default(),
|
||||
},
|
||||
fragment: Some(wgpu::FragmentState {
|
||||
module: &shader,
|
||||
entry_point: Some(fs),
|
||||
targets: &[Some(fmt.into())],
|
||||
compilation_options: Default::default(),
|
||||
}),
|
||||
primitive: wgpu::PrimitiveState {
|
||||
topology: wgpu::PrimitiveTopology::TriangleList,
|
||||
..Default::default()
|
||||
},
|
||||
depth_stencil: None,
|
||||
multisample: Default::default(),
|
||||
multiview: None,
|
||||
cache: None,
|
||||
})
|
||||
};
|
||||
Self {
|
||||
y_pipeline: mk("y_fs", wgpu::TextureFormat::R8Unorm),
|
||||
uv_pipeline: mk("uv_fs", wgpu::TextureFormat::Rg8Unorm),
|
||||
bgl,
|
||||
}
|
||||
}
|
||||
|
||||
/// Record both plane passes. `y_view`/`uv_view` are the R8/RG8 plane render targets.
|
||||
pub fn convert(
|
||||
&self,
|
||||
device: &wgpu::Device,
|
||||
encoder: &mut wgpu::CommandEncoder,
|
||||
rgba_view: &wgpu::TextureView,
|
||||
y_view: &wgpu::TextureView,
|
||||
uv_view: &wgpu::TextureView,
|
||||
) {
|
||||
let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: Some("rgba2nv12_bg"),
|
||||
layout: &self.bgl,
|
||||
entries: &[wgpu::BindGroupEntry {
|
||||
binding: 0,
|
||||
resource: wgpu::BindingResource::TextureView(rgba_view),
|
||||
}],
|
||||
});
|
||||
for (pipeline, view) in [(&self.y_pipeline, y_view), (&self.uv_pipeline, uv_view)] {
|
||||
let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
|
||||
label: Some("rgba2nv12_pass"),
|
||||
color_attachments: &[Some(wgpu::RenderPassColorAttachment {
|
||||
view,
|
||||
resolve_target: None,
|
||||
depth_slice: None,
|
||||
ops: wgpu::Operations {
|
||||
load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
|
||||
store: wgpu::StoreOp::Store,
|
||||
},
|
||||
})],
|
||||
depth_stencil_attachment: None,
|
||||
timestamp_writes: None,
|
||||
occlusion_query_set: None,
|
||||
});
|
||||
pass.set_pipeline(pipeline);
|
||||
pass.set_bind_group(0, &bg, &[]);
|
||||
pass.draw(0..3, 0..1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const SHADER: &str = r#"
|
||||
@group(0) @binding(0) var input_rgba: texture_2d<f32>;
|
||||
|
||||
// Fullscreen triangle.
|
||||
@vertex
|
||||
fn vs_main(@builtin(vertex_index) vi: u32) -> @builtin(position) vec4<f32> {
|
||||
let x = f32((vi << 1u) & 2u);
|
||||
let y = f32(vi & 2u);
|
||||
return vec4<f32>(x * 2.0 - 1.0, 1.0 - y * 2.0, 0.0, 1.0);
|
||||
}
|
||||
|
||||
fn load(p: vec2<i32>) -> vec3<f32> {
|
||||
return textureLoad(input_rgba, p, 0).rgb;
|
||||
}
|
||||
|
||||
// Y plane (full res): one luma byte per pixel.
|
||||
@fragment
|
||||
fn y_fs(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
|
||||
let c = load(vec2<i32>(i32(pos.x), i32(pos.y)));
|
||||
let y = 0.2126 * c.r + 0.7152 * c.g + 0.0722 * c.b;
|
||||
return vec4<f32>(y, 0.0, 0.0, 1.0);
|
||||
}
|
||||
|
||||
// UV plane (half res): 2x2 box-averaged chroma, interleaved into RG.
|
||||
@fragment
|
||||
fn uv_fs(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
|
||||
let sx = 2 * i32(pos.x);
|
||||
let sy = 2 * i32(pos.y);
|
||||
let a = (load(vec2<i32>(sx, sy)) + load(vec2<i32>(sx + 1, sy))
|
||||
+ load(vec2<i32>(sx, sy + 1)) + load(vec2<i32>(sx + 1, sy + 1))) * 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;
|
||||
return vec4<f32>(u, v, 0.0, 1.0);
|
||||
}
|
||||
"#;
|
||||
|
|
@ -4,7 +4,81 @@
|
|||
|
||||
#![cfg(target_os = "linux")]
|
||||
|
||||
use gpu_video_encoder::{dmabuf, vaapi, vk_device};
|
||||
use gpu_video_encoder::{dmabuf, nv12, render_nv12, vaapi, vk_device};
|
||||
|
||||
/// Render a real RGBA frame into the VAAPI surface (zero-copy) and verify the surface's
|
||||
/// NV12 matches the CPU reference for that frame.
|
||||
#[test]
|
||||
fn zerocopy_real_frame_render() {
|
||||
let drm = match vk_device::create() {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
eprintln!("[zerocopy-real] no Vulkan, skipping: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let (w, h) = (640u32, 480u32);
|
||||
let surf = match vaapi::MappedSurface::alloc(w, h) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
eprintln!("[zerocopy-real] no VAAPI, skipping: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let imported = dmabuf::import(&drm, &surf).expect("import");
|
||||
|
||||
// A varied RGBA pattern.
|
||||
let mut rgba = Vec::with_capacity((w * h * 4) as usize);
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
rgba.push(((x * 3 + y) % 256) as u8);
|
||||
rgba.push(((x + y * 2) % 256) as u8);
|
||||
rgba.push(((x * 2 + y * 3) % 256) as u8);
|
||||
rgba.push(255);
|
||||
}
|
||||
}
|
||||
let src = drm.device.create_texture(&wgpu::TextureDescriptor {
|
||||
label: Some("rgba_src"),
|
||||
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: &[],
|
||||
});
|
||||
drm.queue.write_texture(
|
||||
wgpu::TexelCopyTextureInfo { texture: &src, 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 conv = render_nv12::Rgba2Nv12::new(&drm.device);
|
||||
let src_view = src.create_view(&Default::default());
|
||||
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());
|
||||
conv.convert(&drm.device, &mut enc, &src_view, &y_view, &uv_view);
|
||||
drm.queue.submit(Some(enc.finish()));
|
||||
let _ = drm.device.poll(wgpu::PollType::wait_indefinitely());
|
||||
|
||||
let got = surf.readback_nv12().expect("readback");
|
||||
let want = nv12::cpu_reference(&rgba, w, h);
|
||||
assert_eq!(got.len(), want.len());
|
||||
let mut max_diff = 0i32;
|
||||
let mut nbad = 0;
|
||||
for (g, c) in got.iter().zip(want.iter()) {
|
||||
let d = (*g as i32 - *c as i32).abs();
|
||||
max_diff = max_diff.max(d);
|
||||
if d > 2 {
|
||||
nbad += 1;
|
||||
}
|
||||
}
|
||||
eprintln!("[zerocopy-real] {}x{} real-frame render, max diff={max_diff}, bad={nbad}/{}", w, h, got.len());
|
||||
assert!(nbad * 100 < got.len(), "too many bytes differ from CPU NV12 reference");
|
||||
eprintln!("[zerocopy-real] ✅ real RGBA frame rendered into VAAPI surface, NV12 matches reference");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zerocopy_render_into_vaapi_surface() {
|
||||
|
|
@ -34,8 +108,8 @@ fn zerocopy_render_into_vaapi_surface() {
|
|||
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 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 {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,75 @@
|
|||
//! Capstone: encode RGBA frames fully zero-copy (GPU render → VAAPI surface → h264_vaapi)
|
||||
//! and verify the output is real H.264. Skips when VAAPI is unavailable.
|
||||
|
||||
#![cfg(target_os = "linux")]
|
||||
|
||||
use gpu_video_encoder::encoder::ZeroCopyEncoder;
|
||||
|
||||
#[test]
|
||||
fn zerocopy_encode_h264() {
|
||||
let (w, h) = (640u32, 480u32);
|
||||
let mut enc = match ZeroCopyEncoder::new(w, h, 30, 4000) {
|
||||
Ok(e) => e,
|
||||
Err(e) => {
|
||||
eprintln!("[zc-encode] unavailable, skipping: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Build one reusable RGBA source texture; update it per frame with a moving pattern.
|
||||
let device = enc.device();
|
||||
let src = device.create_texture(&wgpu::TextureDescriptor {
|
||||
label: Some("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: &[],
|
||||
});
|
||||
|
||||
let n = 30;
|
||||
for f in 0..n {
|
||||
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 * 8) % 256) as u8);
|
||||
rgba.push(((y + f * 4) % 256) as u8);
|
||||
rgba.push(((x + y) % 256) as u8);
|
||||
rgba.push(255);
|
||||
}
|
||||
}
|
||||
enc.queue().write_texture(
|
||||
wgpu::TexelCopyTextureInfo { texture: &src, 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 },
|
||||
);
|
||||
enc.encode_rgba(&src).expect("encode_rgba");
|
||||
}
|
||||
|
||||
let h264 = enc.finish().expect("finish");
|
||||
eprintln!("[zc-encode] {} frames -> {} bytes H.264", n, h264.len());
|
||||
assert!(h264.len() > 1000, "implausibly small output");
|
||||
assert!(
|
||||
h264.starts_with(&[0, 0, 0, 1]) || h264.starts_with(&[0, 0, 1]),
|
||||
"not Annex-B H.264"
|
||||
);
|
||||
|
||||
// Write it out and ffprobe-verify if ffprobe is present.
|
||||
let out = std::env::temp_dir().join("gpu_video_encoder_zerocopy.h264");
|
||||
std::fs::write(&out, &h264).unwrap();
|
||||
eprintln!("[zc-encode] wrote {}", out.display());
|
||||
if let Ok(o) = std::process::Command::new("ffprobe")
|
||||
.args(["-hide_banner", "-v", "error", "-show_entries", "stream=codec_name,width,height", "-of", "default=noprint_wrappers=1"])
|
||||
.arg(&out)
|
||||
.output()
|
||||
{
|
||||
let s = String::from_utf8_lossy(&o.stdout);
|
||||
eprintln!("[zc-encode] ffprobe:\n{s}");
|
||||
assert!(s.contains("codec_name=h264"), "ffprobe didn't see H.264");
|
||||
assert!(s.contains(&format!("width={w}")), "wrong width");
|
||||
}
|
||||
eprintln!("[zc-encode] ✅ zero-copy H.264 encode verified");
|
||||
}
|
||||
Loading…
Reference in New Issue