Files
punktfunk/crates/punktfunk-host/src/encode/linux/mod.rs
T
enricobuehler 0dacb37088 feat(linux): zero-copy 4:4:4 — the EGL worker converts to planar YUV444 on the GPU
A 4:4:4 session no longer falls to the CPU path (SHM capture + swscale
RGB→YUV444P + re-upload — the fps-ceiling triple tax). The zero-copy worker
grows a Yuv444Blit: three full-res R8 GL passes (the proven BT.709
coefficients; studio or full range per PUNKTFUNK_444_FULLRANGE, read by both
processes so pixels and VUI flip together) into ONE stacked 3-plane pitched
CUDA allocation — which keeps the worker↔host wire and IPC single-plane. The
encoder copies the planes into ffmpeg's yuv444p CUDA surface and hevc_nvenc
emits Range-Extensions 4:4:4 natively.

ImportKind::Tiled444 is APPENDED to the worker protocol (a worker outliving a
replaced host binary must keep the old tags stable; an old worker just errors
the import, which the fail machinery already handles). A 4:4:4 session on a
LINEAR/gamescope capture — no convert wired there — fails with a clear message
instead of letting hevc_nvenc silently subsample. caps().chroma_444 now keys
off the session (it missed the GPU path when keyed off the swscale's
existence).

Live-verified on the CachyOS VM (RTX 5070 Ti): per-frame "imported to CUDA
yuv444=true", stream Rext/yuv444p/bt709 in both tv and pc range, no CPU-path
warning.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-10 18:09:28 +02:00

843 lines
42 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! NVENC encoder via `ffmpeg-next` (binds the system FFmpeg — `ffmpeg-sys-next` auto-detects the
//! installed version, so this builds against FFmpeg 7.x/libavcodec 61 *or* 8.x/libavcodec 62;
//! validated live on Ubuntu 26.04 (FFmpeg 8) and Bazzite F43 (FFmpeg 7.1)).
//!
//! Input is a packed RGB/BGR CPU frame; `*_nvenc` accepts `rgb0`/`bgr0`/`rgba`/`bgra`
//! directly and does the RGB→YUV conversion on the GPU, so the host stays off the
//! colour-conversion path. The portal commonly negotiates packed 24-bit `RGB`, which NVENC
//! does *not* accept — we expand it to `rgb0` (one padding byte/pixel, no colour math).
//! The encoder is opened *without* a global header so VPS/SPS/PPS are emitted in-band on
//! every IDR — the output is both a playable raw Annex-B stream and self-contained AUs.
// Every `unsafe` block in this file carries a `// SAFETY:` proof; enforce it (unsafe-proof program).
#![deny(clippy::undocumented_unsafe_blocks)]
use super::{ChromaFormat, Codec, EncodedFrame, Encoder};
use crate::capture::{CapturedFrame, FramePayload, PixelFormat};
use anyhow::{anyhow, bail, Context, Result};
use ffmpeg::format::Pixel;
use ffmpeg::util::frame::Video as VideoFrame;
use ffmpeg::{codec, encoder, Dictionary, Packet, Rational};
use ffmpeg_next as ffmpeg;
use std::os::raw::c_int;
use std::ptr;
use ffmpeg::ffi; // = ffmpeg_sys_next
/// swscale: nearest-neighbour scaler flag (`SWS_POINT`). We never rescale (src dims == dst dims), so
/// the resampler choice only governs the colour-conversion path; POINT is the cheapest.
const SWS_POINT: c_int = 0x10;
/// swscale colorspace id for ITU-R BT.709 (`SWS_CS_ITU709`) — the CSC coefficients for our RGB→YUV.
const SWS_CS_ITU709: c_int = 1;
/// The swscale *source* pixel format for a captured packed RGB/BGR layout (the real byte order, not
/// the NVENC-padded `*0` form). Used by the 4:4:4 RGB→YUV444P conversion path. Mirrors the VAAPI
/// CPU-input mapping; YUV/10-bit inputs can't feed this path (the 4:4:4 session forces packed RGB).
fn sws_src_pixel(format: PixelFormat) -> Result<Pixel> {
Ok(match format {
PixelFormat::Bgrx => Pixel::BGRZ, // bgr0
PixelFormat::Rgbx => Pixel::RGBZ, // rgb0
PixelFormat::Bgra => Pixel::BGRA,
PixelFormat::Rgba => Pixel::RGBA,
PixelFormat::Rgb => Pixel::RGB24,
PixelFormat::Bgr => Pixel::BGR24,
PixelFormat::Nv12 | PixelFormat::P010 | PixelFormat::Rgb10a2 | PixelFormat::Yuv444 => {
bail!("NVENC 4:4:4 CPU-input path supports packed RGB/BGR only; got {format:?}")
}
})
}
/// `AVCUDADeviceContext` (libavutil/hwcontext_cuda.h) — not in the ffmpeg-sys bindings (the
/// crate doesn't allowlist that header), so mirror its stable 3-pointer layout. We set the
/// first field to *our* `CUcontext` so NVENC shares the context the EGL importer maps into.
#[repr(C)]
struct AVCUDADeviceContext {
cuda_ctx: *mut std::ffi::c_void, // CUcontext
stream: *mut std::ffi::c_void, // CUstream (null = default)
internal: *mut std::ffi::c_void, // filled by ctx_init
}
/// CUDA hardware-frame contexts that wrap our shared `CUcontext`, so `hevc_nvenc` reads the
/// imported device buffer directly. Owns two `AVBufferRef`s, unref'd on drop.
struct CudaHw {
device_ref: *mut ffi::AVBufferRef,
frames_ref: *mut ffi::AVBufferRef,
}
impl CudaHw {
/// Build a CUDA hwdevice wrapping `cu_ctx` and a frames pool (`sw_format` = `pixel`).
unsafe fn new(cu_ctx: *mut std::ffi::c_void, sw_format: Pixel, w: u32, h: u32) -> Result<Self> {
let mut device_ref = ffi::av_hwdevice_ctx_alloc(ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA);
if device_ref.is_null() {
bail!("av_hwdevice_ctx_alloc(CUDA) failed");
}
let dev_ctx = (*device_ref).data as *mut ffi::AVHWDeviceContext;
let cu = (*dev_ctx).hwctx as *mut AVCUDADeviceContext;
(*cu).cuda_ctx = cu_ctx; // share the importer's context
let r = ffi::av_hwdevice_ctx_init(device_ref);
if r < 0 {
ffi::av_buffer_unref(&mut device_ref);
bail!("av_hwdevice_ctx_init failed ({r})");
}
let mut frames_ref = ffi::av_hwframe_ctx_alloc(device_ref);
if frames_ref.is_null() {
ffi::av_buffer_unref(&mut device_ref);
bail!("av_hwframe_ctx_alloc failed");
}
let fc = (*frames_ref).data as *mut ffi::AVHWFramesContext;
(*fc).format = ffi::AVPixelFormat::AV_PIX_FMT_CUDA;
(*fc).sw_format = pixel_to_av(sw_format);
(*fc).width = w as c_int;
(*fc).height = h as c_int;
(*fc).initial_pool_size = 0; // we supply the device pointers
let r = ffi::av_hwframe_ctx_init(frames_ref);
if r < 0 {
ffi::av_buffer_unref(&mut frames_ref);
ffi::av_buffer_unref(&mut device_ref);
bail!("av_hwframe_ctx_init failed ({r})");
}
Ok(CudaHw {
device_ref,
frames_ref,
})
}
}
impl Drop for CudaHw {
fn drop(&mut self) {
// SAFETY: `frames_ref`/`device_ref` are the two non-null `AVBufferRef`s `CudaHw::new` created
// (it bails before returning `Self` if either alloc fails, so a live `CudaHw` always holds
// both). `av_buffer_unref` drops one reference and nulls the pointer through the `&mut`. This
// `Drop` runs exactly once and `CudaHw` owns these refs exclusively → no double-free /
// use-after-free. Frames are unref'd before the device (the frames ctx internally refs the
// device; refcounted, so the order is sound regardless).
unsafe {
ffi::av_buffer_unref(&mut self.frames_ref);
ffi::av_buffer_unref(&mut self.device_ref);
}
}
}
/// `ffmpeg::format::Pixel` → raw `AVPixelFormat`.
fn pixel_to_av(p: Pixel) -> ffi::AVPixelFormat {
// `Pixel` is `#[repr(i32)]`-compatible with `AVPixelFormat` (the bindgen enum) via this
// documented conversion in ffmpeg-next.
ffi::AVPixelFormat::from(p)
}
/// Map a captured layout to the NVENC input pixel format, and whether a 3→4 byte expand is
/// needed (packed RGB/BGR have no padding byte; the NVENC `*0` formats do).
fn nvenc_input(format: PixelFormat) -> (Pixel, bool) {
match format {
PixelFormat::Bgrx => (Pixel::BGRZ, false), // bgr0
PixelFormat::Rgbx => (Pixel::RGBZ, false), // rgb0
PixelFormat::Bgra => (Pixel::BGRA, false),
PixelFormat::Rgba => (Pixel::RGBA, false),
PixelFormat::Rgb => (Pixel::RGBZ, true), // RGB -> rgb0
PixelFormat::Bgr => (Pixel::BGRZ, true), // BGR -> bgr0
// NV12 is native YUV: NVENC encodes it with NO internal RGB→YUV CSC (the Tier 2A win). On
// Linux it's produced by the GPU convert on the zero-copy tiled path (`PUNKTFUNK_NV12`); on
// Windows by the D3D11 video processor.
PixelFormat::Nv12 => (Pixel::NV12, false),
// Planar YUV444 from the zero-copy worker's GPU convert (a 4:4:4 session) — native
// full-chroma YUV in, `hevc_nvenc` emits Range-Extensions 4:4:4.
PixelFormat::Yuv444 => (Pixel::YUV444P, false),
// Rgb10a2 (HDR) and P010 (the Windows 10-bit video-processor output) are produced only by
// the Windows paths; the Linux capturer never emits them. Map to BGRA so the match is
// exhaustive — unreachable here.
PixelFormat::Rgb10a2 | PixelFormat::P010 => (Pixel::BGRA, false),
}
}
pub struct NvencEncoder {
enc: encoder::video::Encoder,
/// Reusable 4-bpp CPU input frame (CPU path only; `None` for the zero-copy/CUDA path).
/// Mutating it in place across frames is sound only because the encoder is opened with
/// `delay=0`/`bf=0`/`max_b_frames=0` and the caller drains `poll()` after each `submit`,
/// so libavcodec holds no reference to the previous frame's buffer when we overwrite it.
frame: Option<VideoFrame>,
/// Zero-copy path: CUDA hwdevice/hwframes contexts (the encoder takes `AV_PIX_FMT_CUDA`).
cuda: Option<CudaHw>,
/// 4:4:4 CPU path only: swscale context converting the captured packed RGB/BGR → planar
/// YUV444P into [`Self::frame`], because `hevc_nvenc` only emits 4:4:4 from a YUV444 *input*
/// (RGB-in is always 4:2:0). `None` on the 4:2:0 paths AND on the zero-copy 4:4:4 path (the
/// worker's GPU convert delivers YUV444 CUDA frames). Freed in `Drop`.
sws_444: Option<*mut ffi::SwsContext>,
/// This session opened as full-chroma 4:4:4 (FREXT) — via either input path.
want_444: bool,
src_format: PixelFormat,
expand: bool,
width: u32,
height: u32,
fps: u32,
/// Monotonic presentation index, in `1/fps` time-base units.
frame_idx: i64,
/// Force the next submitted frame to be an IDR (set by [`request_keyframe`]).
force_kf: bool,
/// Opened in intra-refresh mode (surfaced via [`caps`](Encoder::caps) so the session glue
/// rate-limits forced IDRs — the wave heals loss without them).
intra_refresh: bool,
}
// `CudaHw` holds raw `AVBufferRef`s and `sws_444` a raw `SwsContext`; the encoder lives on a single
// thread. The CPU encoder is already `Send` via ffmpeg-next; assert it for the raw fields too.
// SAFETY: `NvencEncoder` owns an ffmpeg-next `Encoder`/`VideoFrame` (already `Send`) plus a `CudaHw`
// holding raw `AVBufferRef`s and an optional raw `SwsContext`, none of which are `Send` by default.
// The `SwsContext` is a self-contained swscale state object with no thread affinity, touched only
// through `&mut self` on the one encode thread. The encoder is owned and driven by
// exactly ONE thread — the per-session encode thread it is moved to — and is only touched through
// `&mut self` methods, so it is never aliased or accessed concurrently. The wrapped libav contexts
// (and the shared `CUcontext` the `CudaHw` references) have no thread affinity, so transferring
// ownership across threads is sound. This asserts `Send` (transfer) only, extending ffmpeg-next's
// existing `Send` to the raw CUDA fields; `Sync` (shared `&`) is deliberately NOT implemented.
unsafe impl Send for NvencEncoder {}
/// Latched true once an intra-refresh open failed with the device-capability error (ENOSYS from
/// `NV_ENC_CAPS_SUPPORT_INTRA_REFRESH`), so later sessions skip the doomed attempt. Never set by
/// other open failures (a bitrate EINVAL must not permanently disable the feature).
static IR_UNSUPPORTED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
/// Whether this open should run the NVENC **intra-refresh** loss-recovery mode
/// (`PUNKTFUNK_INTRA_REFRESH` truthy, opt-in until on-glass validated): a moving intra band +
/// recovery-point SEI refreshes the whole picture every [`intra_refresh_period`] frames, so
/// FEC-unrecoverable loss heals without the 20-40× full-IDR spike (which under loss causes more
/// loss — the cascade). The session glue then rate-limits client keyframe requests
/// ([`EncoderCaps::intra_refresh`](super::EncoderCaps)).
fn intra_refresh_requested() -> bool {
std::env::var("PUNKTFUNK_INTRA_REFRESH")
.map(|v| matches!(v.trim(), "1" | "true" | "yes" | "on"))
.unwrap_or(false)
&& !IR_UNSUPPORTED.load(std::sync::atomic::Ordering::Relaxed)
}
/// The intra-refresh wave length in frames — ffmpeg derives `intraRefreshPeriod`/`Cnt` from
/// `gop_size` before forcing the real GOP infinite, so this is what `gop_size` is set to in IR
/// mode. Default = half a second of frames (heals fast, spreads the intra cost to ~2-3% per
/// frame); `PUNKTFUNK_IR_PERIOD_FRAMES` overrides.
fn intra_refresh_period(fps: u32) -> i32 {
std::env::var("PUNKTFUNK_IR_PERIOD_FRAMES")
.ok()
.and_then(|s| s.parse::<i32>().ok())
.filter(|v| *v >= 2)
.unwrap_or_else(|| (fps.max(16) / 2) as i32)
}
impl NvencEncoder {
#[allow(clippy::too_many_arguments)]
pub fn open(
codec: Codec,
format: PixelFormat,
width: u32,
height: u32,
fps: u32,
bitrate_bps: u64,
cuda: bool,
bit_depth: u8,
chroma: ChromaFormat,
) -> Result<Self> {
// TODO(hdr): Linux 10-bit parity. Unlike the Windows raw-SDK path (which upconverts 8-bit
// ARGB → Main10 via pixelBitDepthMinus8), libavcodec hevc_nvenc needs a 10-bit input pixel
// format (p010) for Main10, so it's a bigger change; deferred until a Linux GPU box is
// available to validate. The Linux host stays 8-bit for now.
if bit_depth != 8 {
tracing::warn!(
bit_depth,
"Linux NVENC 10-bit not yet wired — encoding 8-bit"
);
}
// Full-chroma 4:4:4 (HEVC Range Extensions). `hevc_nvenc` only emits 4:4:4 from a YUV444
// *input* frame — feeding RGB always subsamples to 4:2:0 regardless of profile (verified on
// the RTX 5070 Ti). Two ways to produce that input: the zero-copy worker's GPU convert
// (planar-YUV444 CUDA frames — `cuda` true), or the CPU path's swscale RGB→YUV444P. Both
// feed `profile=rext`; the range follows `PUNKTFUNK_444_FULLRANGE` in both.
let want_444 = chroma.is_444() && codec == Codec::H265;
ffmpeg::init().context("ffmpeg init")?;
if std::env::var_os("PUNKTFUNK_FFMPEG_DEBUG").is_some() {
// SAFETY: `av_log_set_level` sets libav's global integer log level; `48` (= AV_LOG_DEBUG)
// is a valid level with no pointer args, and libav was just initialized by `ffmpeg::init()`
// above — always sound.
unsafe { ffi::av_log_set_level(48) }; // AV_LOG_DEBUG — surface NVENC hw-frame rejects
}
let name = codec.nvenc_name();
let av_codec = encoder::find_by_name(name)
.ok_or_else(|| anyhow!("{name} not built into libavcodec"))?;
let (rgb_pixel, rgb_expand) = nvenc_input(format);
// 4:4:4 feeds NVENC a planar YUV444P frame we produce by swscale; the ordinary path feeds the
// captured RGB straight in and lets NVENC's internal CSC subsample to 4:2:0.
let (nvenc_pixel, expand) = if want_444 {
(Pixel::YUV444P, false)
} else {
(rgb_pixel, rgb_expand)
};
let mut video = codec::context::Context::new_with_codec(av_codec)
.encoder()
.video()
.context("alloc video encoder")?;
video.set_width(width);
video.set_height(height);
video.set_format(nvenc_pixel); // NVENC converts RGB→YUV internally
video.set_time_base(Rational(1, fps as i32));
video.set_frame_rate(Some(Rational(fps as i32, 1)));
video.set_bit_rate(bitrate_bps as usize);
video.set_max_bit_rate(bitrate_bps as usize);
// VBV/HRD buffer — bound the SIZE of any single frame. Under CBR with no buffer set, NVENC
// uses a loose default VBV, so a high-motion P-frame is allowed to balloon to many times the
// average; those extra packets overflow the bounded send queue + kernel socket buffer and
// get dropped, which the client sees as framedrops/jitter (and, on the infinite-GOP path, as
// old/stale frames flashing until the next RFI). A tight ~1-frame buffer makes the encoder
// hold frame size roughly constant and absorb motion as a momentary QP (quality) dip instead
// — the trade we want. Default = 1 frame of bits (bitrate/fps); PUNKTFUNK_VBV_FRAMES tunes it
// (larger = better motion quality but bigger per-frame bursts).
let vbv_frames = std::env::var("PUNKTFUNK_VBV_FRAMES")
.ok()
.and_then(|s| s.parse::<f32>().ok())
.filter(|v| v.is_finite() && *v > 0.0)
.unwrap_or(1.0);
let vbv_bits = ((bitrate_bps as f64 / fps.max(1) as f64) * vbv_frames as f64)
.clamp(1.0, i32::MAX as f64);
// SAFETY: `video` is the ffmpeg-next encoder builder wrapping a freshly-allocated
// `AVCodecContext` that we hold by value and have not opened yet; `video.as_mut_ptr()` returns
// that non-null, properly-aligned, exclusively-owned context. Writing the plain `rc_buffer_size`
// int field before `open_with` is the supported way to set a field ffmpeg-next exposes no
// setter for. Sole owner → no aliasing; synchronous in-bounds scalar write.
unsafe {
(*video.as_mut_ptr()).rc_buffer_size = vbv_bits as i32;
}
video.set_max_b_frames(0);
// Infinite GOP — NO periodic IDR. A keyframe at 5120x1440 is ~20-40x a P-frame, so a
// periodic IDR is a recurring multi-millisecond encode+packetize+send spike — the ~2s
// "freeze". NVENC emits one IDR at stream start, then P-frames only; `forced-idr` (below)
// turns a client recovery request (RFI, via `request_keyframe`) into an IDR on demand.
// This is the Moonlight/Sunshine low-latency model.
// In intra-refresh mode the GOP is still infinite — ffmpeg reads `gop_size` as the refresh
// WAVE length (`intraRefreshPeriod`/`Cnt`) and then forces `gopLength` infinite itself, so
// a positive `gop_size` here does NOT reintroduce periodic IDRs.
let intra_refresh = intra_refresh_requested();
// SAFETY: same `video` builder as above — a non-null, properly-aligned, sole-owned, not-yet-
// opened `AVCodecContext`. We write the plain `gop_size` int field (-1 = infinite GOP, or the
// intra-refresh wave length) before `open_with`, which ffmpeg-next has no setter for. No
// aliasing; synchronous scalar write.
unsafe {
(*video.as_mut_ptr()).gop_size = if intra_refresh {
intra_refresh_period(fps)
} else {
-1
};
}
// NV12 / 4:4:4 paths: we do the RGB→YUV conversion ourselves as BT.709 (swscale), so
// signal that in the bitstream VUI (colorspace/range/primaries/transfer) — otherwise the
// client decoder assumes a default and the picture comes out washed-out / wrong-contrast.
// The RGB-input 4:2:0 path leaves these unset (NVENC's internal CSC writes its own VUI).
// Matches the Windows NV12 path's BT.709 limited-range signalling.
//
// PUNKTFUNK_444_FULLRANGE=1 (experimental, 4:4:4-only): convert AND signal FULL range —
// recovers the ~12% of code space limited-range quantization gives up, for the exact
// text/UI chroma 4:4:4 exists for. Every punktfunk client honors the signaled range
// (csc_rows / the Apple rows port); ship as default only if the on-glass A/B shows a
// visible win. Linux-only: the Windows path's NVENC-internal CSC range is unmeasured.
let full_range_444 =
want_444 && std::env::var("PUNKTFUNK_444_FULLRANGE").is_ok_and(|v| v.trim() == "1");
if matches!(format, PixelFormat::Nv12) || want_444 {
// SAFETY: same `video` builder — `raw = video.as_mut_ptr()` is the non-null, properly-
// aligned, sole-owned, not-yet-opened `AVCodecContext`. We set its four VUI colour enum
// fields to valid `AVColorSpace`/`AVColorRange`/`AVColorPrimaries`/`AVColorTransfer-
// Characteristic` variants before `open_with`. Sole owner → no aliasing; synchronous writes.
unsafe {
let raw = video.as_mut_ptr();
(*raw).colorspace = ffi::AVColorSpace::AVCOL_SPC_BT709;
(*raw).color_range = if full_range_444 {
ffi::AVColorRange::AVCOL_RANGE_JPEG // full
} else {
ffi::AVColorRange::AVCOL_RANGE_MPEG // limited/studio
};
(*raw).color_primaries = ffi::AVColorPrimaries::AVCOL_PRI_BT709;
(*raw).color_trc = ffi::AVColorTransferCharacteristic::AVCOL_TRC_BT709;
}
}
// For the zero-copy path, take CUDA surfaces: wrap the shared CUcontext in CUDA
// hwdevice/hwframes contexts and set `pix_fmt = CUDA` on the raw encoder context
// *before* open (NVENC derives the device from `hw_frames_ctx`).
let cuda_hw = if cuda {
let cu_ctx = crate::zerocopy::cuda::context().context("shared CUDA context")?;
// SAFETY: `CudaHw::new` (an `unsafe fn`) requires libav initialized (the `ffmpeg::init()`
// above ran) and a valid `CUcontext`; `cu_ctx` is the shared importer context from
// `zerocopy::cuda::context()?`, non-null on the `Ok` path. `nvenc_pixel` is a valid `Pixel`
// and `width`/`height` are the validated positive dims. It returns a RAII `CudaHw` wrapping
// (not owning) `cu_ctx` and owning two `AVBufferRef`s freed on drop.
let hw = unsafe { CudaHw::new(cu_ctx, nvenc_pixel, width, height)? };
// SAFETY: `raw = video.as_mut_ptr()` is the non-null, sole-owned, not-yet-opened
// `AVCodecContext`. We set `pix_fmt = CUDA` and attach NEW refs (`av_buffer_ref`) of
// `hw.device_ref`/`hw.frames_ref` — both non-null (`CudaHw::new` guarantees) and from the
// live `hw`, which is moved into `NvencEncoder.cuda` next to `enc` and so outlives the
// encoder. The context owns its own refs (freed when the context closes). No aliasing.
unsafe {
let raw = video.as_mut_ptr();
(*raw).pix_fmt = ffi::AVPixelFormat::AV_PIX_FMT_CUDA;
(*raw).hw_device_ctx = ffi::av_buffer_ref(hw.device_ref);
(*raw).hw_frames_ctx = ffi::av_buffer_ref(hw.frames_ref);
}
Some(hw)
} else {
None
};
// 4:4:4 CPU path: build the RGB→YUV444P swscale (BT.709, range per the flag; no rescale).
// Mirrors the VAAPI CPU path's RGB→NV12 scaler, but the dst is full-chroma planar 4:4:4.
// Skipped on the zero-copy path (`cuda`): the worker's GPU convert already delivers
// planar YUV444 CUDA frames — no CPU pixels exist to scale.
let sws_444 = if want_444 && !cuda {
let src_av = pixel_to_av(sws_src_pixel(format)?);
// SAFETY: `sws_getContext` allocates a swscale context for the given src/dst dims + pixel
// formats. Both dims are the encoder's positive `width`/`height` as `c_int`; `src_av` is a
// valid `AVPixelFormat` (from the `sws_src_pixel`-validated, packed-RGB-only source), the
// dst is YUV444P. The trailing filter/param pointers are null = "use defaults" (documented
// as accepted). No Rust memory is borrowed; the returned pointer is null-checked below.
let sws = unsafe {
ffi::sws_getContext(
width as c_int,
height as c_int,
src_av,
width as c_int,
height as c_int,
ffi::AVPixelFormat::AV_PIX_FMT_YUV444P,
SWS_POINT,
ptr::null_mut(),
ptr::null_mut(),
ptr::null(),
)
};
if sws.is_null() {
bail!("sws_getContext(RGB→YUV444P) failed");
}
// SAFETY: `sws` is the non-null context from the call above (null-checked). The ITU-709
// coefficient table from `sws_getCoefficients` is a process-lifetime libswscale static,
// reused for src+dst matrices; `sws_setColorspaceDetails` only reads it and writes scalar
// CSC settings into `sws` (dstRange matches the VUI: 0 = limited, 1 = the
// PUNKTFUNK_444_FULLRANGE experiment). No Rust memory is passed.
unsafe {
let cs709 = ffi::sws_getCoefficients(SWS_CS_ITU709);
let dst_range = i32::from(full_range_444);
ffi::sws_setColorspaceDetails(sws, cs709, 1, cs709, dst_range, 0, 1 << 16, 1 << 16);
}
Some(sws)
} else {
None
};
// Low-latency NVENC tuning (plan §7 / linux-setup doc).
let mut opts = Dictionary::new();
opts.set("preset", "p1"); // fastest
opts.set("tune", "ull"); // ultra-low-latency
opts.set("rc", "cbr");
opts.set("bf", "0");
opts.set("delay", "0");
opts.set("forced-idr", "1"); // RFI/request_keyframe → real IDR under the infinite GOP
if intra_refresh {
// Moving intra band + recovery-point SEI (period set via gop_size above). Loss now
// self-heals within the wave; forced IDRs remain available (rate-limited by the glue).
opts.set("intra-refresh", "1");
}
if want_444 {
// HEVC Range Extensions — the profile that carries chroma_format_idc=3. With a YUV444P
// input `hevc_nvenc` auto-selects it, but pin it explicitly so the chroma is never silently
// dropped on a future libavcodec.
opts.set("profile", "rext");
}
// Split-frame encode across both NVENC engines (GB203 has 2) when the pixel rate exceeds
// a single engine's HEVC capacity (~1 Gpix/s); e.g. 5120x1440@240 = 1.77 Gpix/s needs it,
// @120 = 0.88 Gpix/s does not. HEVC/AV1 only (not H.264). AUTO won't engage below ~2112px
// height, so we force `2`; below the threshold we leave it AUTO (split costs ~2% BD-rate).
// Output is standard HEVC — transparent to the client. Override with PUNKTFUNK_SPLIT_ENCODE.
let pix_rate = width as u64 * height as u64 * fps as u64;
let split = std::env::var("PUNKTFUNK_SPLIT_ENCODE").ok();
match split.as_deref() {
Some(mode) => opts.set("split_encode_mode", mode),
None if matches!(codec, Codec::H265 | Codec::Av1) && pix_rate > 1_000_000_000 => {
opts.set("split_encode_mode", "2");
tracing::info!(
pix_rate,
"NVENC: forcing 2-way split encode (high pixel rate)"
);
}
None => {}
}
let enc = match video.open_with(opts) {
Ok(enc) => enc,
// The GPU lacks NV_ENC_CAPS_SUPPORT_INTRA_REFRESH — ffmpeg fails the open with
// ENOSYS ("Function not implemented"). Latch it (skip the doomed attempt on later
// sessions) and reopen this session without intra-refresh; any other failure — and
// any failure when IR wasn't requested — propagates untouched (the bitrate probe
// keys on EINVAL, which must not trip the latch).
Err(e) if intra_refresh && format!("{e:#}").contains("Function not implemented") => {
tracing::warn!(
encoder = name,
"NVENC intra-refresh not supported by this GPU — falling back to IDR-only \
recovery"
);
IR_UNSUPPORTED.store(true, std::sync::atomic::Ordering::Relaxed);
return Self::open(
codec,
format,
width,
height,
fps,
bitrate_bps,
cuda,
bit_depth,
chroma,
);
}
Err(e) => {
return Err(e).with_context(|| {
format!("open {name} ({width}x{height}@{fps}, {bitrate_bps} bps)")
})
}
};
if intra_refresh {
tracing::info!(
encoder = name,
period_frames = intra_refresh_period(fps),
"NVENC intra-refresh recovery active (no periodic IDR; wave heals loss)"
);
}
let frame = if cuda {
None
} else {
Some(VideoFrame::new(nvenc_pixel, width, height))
};
Ok(NvencEncoder {
enc,
frame,
cuda: cuda_hw,
sws_444,
want_444,
src_format: format,
expand,
width,
height,
fps,
frame_idx: 0,
force_kf: false,
intra_refresh,
})
}
}
impl Encoder for NvencEncoder {
fn caps(&self) -> super::EncoderCaps {
super::EncoderCaps {
// 4:4:4 iff this session opened FREXT — the CPU swscale path or the zero-copy GPU
// convert. RFI/HDR-SEI stay unsupported on libavcodec NVENC (the trait defaults).
chroma_444: self.want_444,
intra_refresh: self.intra_refresh,
..super::EncoderCaps::default()
}
}
fn submit(&mut self, captured: &CapturedFrame) -> Result<()> {
anyhow::ensure!(
captured.width == self.width && captured.height == self.height,
"captured frame {}x{} != encoder {}x{}",
captured.width,
captured.height,
self.width,
self.height
);
let pts = self.frame_idx;
self.frame_idx += 1;
// Force an IDR when requested (client RFI); otherwise let NVENC pick (GOP/P-frame).
let idr = self.force_kf;
self.force_kf = false;
match &captured.payload {
FramePayload::Cuda(buf) => self.submit_cuda(buf, pts, idr),
FramePayload::Cpu(bytes) => self.submit_cpu(bytes, captured.format, pts, idr),
FramePayload::Dmabuf(_) => {
bail!("NVENC got a VAAPI dmabuf frame — capture/encoder backend mismatch")
}
}
}
fn request_keyframe(&mut self) {
self.force_kf = true;
}
fn poll(&mut self) -> Result<Option<EncodedFrame>> {
let mut pkt = Packet::empty();
match self.enc.receive_packet(&mut pkt) {
Ok(()) => {
let data = pkt.data().map(|d| d.to_vec()).unwrap_or_default();
let pts = pkt.pts().unwrap_or(0).max(0) as u64;
let pts_ns = pts * 1_000_000_000 / self.fps as u64;
Ok(Some(EncodedFrame {
data,
pts_ns,
keyframe: pkt.is_key(),
}))
}
// No packet ready yet (need another input frame).
Err(ffmpeg::Error::Other { errno })
if errno == ffmpeg::util::error::EAGAIN
|| errno == ffmpeg::util::error::EWOULDBLOCK =>
{
Ok(None)
}
// Fully drained after flush().
Err(ffmpeg::Error::Eof) => Ok(None),
Err(e) => Err(e).context("receive_packet"),
}
}
fn flush(&mut self) -> Result<()> {
self.enc.send_eof().context("send_eof")?;
Ok(())
}
}
impl NvencEncoder {
/// CPU path: expand/copy the packed RGB/BGR bytes into the reusable 4-bpp frame, then send.
fn submit_cpu(&mut self, bytes: &[u8], format: PixelFormat, pts: i64, idr: bool) -> Result<()> {
anyhow::ensure!(
format == self.src_format,
"captured format {:?} != encoder source {:?}",
format,
self.src_format
);
let w = self.width as usize;
let h = self.height as usize;
let src_bpp = self.src_format.bytes_per_pixel();
let src_row = w * src_bpp;
anyhow::ensure!(
bytes.len() >= src_row * h,
"captured buffer {} bytes < required {}",
bytes.len(),
src_row * h
);
// 4:4:4: swscale the packed RGB straight into the planar YUV444P input frame (BT.709 limited),
// then send it — no byte-expand. The 4:2:0 RGB path (below) feeds NVENC packed RGB directly.
if let Some(sws) = self.sws_444 {
let frame = self
.frame
.as_mut()
.context("CPU frame missing (encoder opened in CUDA mode)")?;
// SAFETY: `format == self.src_format` and `bytes.len() >= src_row * h` (the `ensure!`s
// above), so `sws_scale` reads `h` rows of `src_row` bytes from `src_data[0] = bytes`
// (packed RGB is single-plane; the other src planes are null/0) — all in bounds. `sws` is
// the non-null context built in `open`. The dst is `frame`'s underlying `AVFrame`: its
// `data`/`linesize` in-struct arrays were sized for YUV444P by `VideoFrame::new`, and the
// 3 planes are each `width`×`height`. All pointers are live locals for this synchronous
// call; the encoder runs only on this thread (`unsafe impl Send`), so no aliasing/race.
unsafe {
let dst_av = frame.as_mut_ptr();
let src_data: [*const u8; 4] =
[bytes.as_ptr(), ptr::null(), ptr::null(), ptr::null()];
let src_stride: [c_int; 4] = [src_row as c_int, 0, 0, 0];
let r = ffi::sws_scale(
sws,
src_data.as_ptr(),
src_stride.as_ptr(),
0,
h as c_int,
(*dst_av).data.as_ptr(),
(*dst_av).linesize.as_ptr(),
);
if r < 0 {
bail!("sws_scale(RGB→YUV444P) failed ({r})");
}
}
frame.set_pts(Some(pts));
frame.set_kind(if idr {
ffmpeg::picture::Type::I
} else {
ffmpeg::picture::Type::None
});
self.enc.send_frame(frame).context("send_frame(444)")?;
return Ok(());
}
let frame = self
.frame
.as_mut()
.context("CPU frame missing (encoder opened in CUDA mode)")?;
let stride = frame.stride(0); // dst is 4-bpp, aligned
let dst = frame.data_mut(0);
if self.expand {
// packed 3-bpp RGB/BGR → 4-bpp *0 (copy 3 bytes, zero the pad byte)
for y in 0..h {
let s = &bytes[y * src_row..y * src_row + src_row];
let drow = &mut dst[y * stride..y * stride + w * 4];
for x in 0..w {
drow[x * 4..x * 4 + 3].copy_from_slice(&s[x * 3..x * 3 + 3]);
drow[x * 4 + 3] = 0;
}
}
} else {
// 4-bpp → 4-bpp, honoring the (possibly larger) dst stride
for y in 0..h {
dst[y * stride..y * stride + src_row]
.copy_from_slice(&bytes[y * src_row..y * src_row + src_row]);
}
}
frame.set_pts(Some(pts));
frame.set_kind(if idr {
ffmpeg::picture::Type::I
} else {
ffmpeg::picture::Type::None
});
self.enc.send_frame(frame).context("send_frame")?;
Ok(())
}
/// Zero-copy path: hand the imported CUDA device buffer to NVENC with no CPU touch.
///
/// We take a *pooled* surface from the CUDA hwframes context (`av_hwframe_get_buffer`) and
/// device→device-copy our imported buffer into it, rather than wrapping our own pointer in a
/// bare frame. Two reasons: (1) NVENC's `nvenc_send_frame` ignores frames whose `buf[0]` is
/// null and the generic encode path's `av_frame_ref` needs a refcounted buffer — a bare
/// frame is rejected with `EINVAL`; (2) NVENC caches CUDA-resource *registrations* keyed by
/// device pointer with a bounded table, so a fresh pointer every frame would thrash/overflow
/// it — the pool recycles a small set of pointers. The extra copy is device-local (~8 MB at
/// 1080p, sub-millisecond on the GPU) and keeps the host fully off the pixel path.
fn submit_cuda(
&mut self,
buf: &crate::zerocopy::DeviceBuffer,
pts: i64,
idr: bool,
) -> Result<()> {
let frames_ref = self
.cuda
.as_ref()
.context("CUDA hw context missing (encoder opened in CPU mode)")?
.frames_ref;
// The device→device copy below uses our shared context directly; make it current on the
// encode thread (ffmpeg pushes its own around the pool alloc, so order is fine).
crate::zerocopy::cuda::make_current().context("CUDA context current (encode thread)")?;
// SAFETY: `frames_ref` is the non-null CUDA frames ctx from `self.cuda` (unwrapped via
// `.context(..)?` above), and the shared CUDA context was just made current on THIS thread
// (`make_current()?`), the precondition for the device-pointer copies below.
// * `av_frame_alloc` → `f` (null-checked). `av_hwframe_get_buffer(frames_ref, f, 0)` fills `f`
// with a pooled CUDA surface (sets `data[]`/`linesize[]`/`buf[0]`/`hw_frames_ctx`); on
// failure we free `f` and bail.
// * For NV12 we read `(*f).data[0..2]` / `linesize[0..2]` (Y + interleaved UV), else
// `data[0]`/`linesize[0]` — in-struct fields of the non-null `f`, valid for the surface dims
// ffmpeg allocated — and pass them to the cuda copy helpers, which device→device copy `buf`
// (the imported `DeviceBuffer`, owned by the caller and live for this call) into the surface.
// * On copy error we free `f` and return. Otherwise we write `pts`/`pict_type` through `f` and
// `avcodec_send_frame` it into the live owned `self.enc` context (which takes its own ref of
// the pooled surface), then free our `f` ref exactly once. Single-threaded encoder → no race.
unsafe {
let mut f = ffi::av_frame_alloc();
if f.is_null() {
bail!("av_frame_alloc failed");
}
// Pooled CUDA surface: sets format, width/height, data[0]/linesize[0], buf[0] and
// hw_frames_ctx. Reused across frames (the pool recycles), keeping NVENC's
// registration cache warm.
let r = ffi::av_hwframe_get_buffer(frames_ref, f, 0);
if r < 0 {
ffi::av_frame_free(&mut f);
bail!("av_hwframe_get_buffer(CUDA) failed ({r})");
}
// NV12 surfaces are two-plane (Y in data[0], interleaved UV in data[1]); YUV444
// surfaces are three-plane (`yuv444p` frames ctx — data[0..3]); the RGB surfaces are
// single-plane. Copy the matching layout into NVENC's pooled surface. A 4:4:4 session
// whose buffer ISN'T YUV444 (a LINEAR/gamescope capture the worker can't convert)
// fails loudly here rather than letting `hevc_nvenc` silently subsample RGB to 4:2:0.
let copy_res = if buf.yuv444 {
let dsts = core::array::from_fn(|i| {
(
(*f).data[i] as crate::zerocopy::cuda::CUdeviceptr,
(*f).linesize[i] as usize,
)
});
crate::zerocopy::cuda::copy_yuv444_to_device(buf, dsts)
} else if self.want_444 {
ffi::av_frame_free(&mut f);
bail!(
"4:4:4 session but the zero-copy frame is not YUV444 (LINEAR/gamescope \
capture has no GPU 4:4:4 convert) — unset PUNKTFUNK_ZEROCOPY to use the \
CPU 4:4:4 path on this compositor"
);
} else if buf.is_nv12() {
let y_ptr = (*f).data[0] as crate::zerocopy::cuda::CUdeviceptr;
let y_pitch = (*f).linesize[0] as usize;
let uv_ptr = (*f).data[1] as crate::zerocopy::cuda::CUdeviceptr;
let uv_pitch = (*f).linesize[1] as usize;
crate::zerocopy::cuda::copy_nv12_to_device(buf, y_ptr, y_pitch, uv_ptr, uv_pitch)
} else {
let dst_ptr = (*f).data[0] as crate::zerocopy::cuda::CUdeviceptr;
let dst_pitch = (*f).linesize[0] as usize;
crate::zerocopy::cuda::copy_device_to_device(buf, dst_ptr, dst_pitch)
};
if let Err(e) = copy_res {
ffi::av_frame_free(&mut f);
return Err(e).context("copy imported buffer into NVENC surface");
}
(*f).pts = pts;
(*f).pict_type = if idr {
ffi::AVPictureType::AV_PICTURE_TYPE_I
} else {
ffi::AVPictureType::AV_PICTURE_TYPE_NONE
};
let r = ffi::avcodec_send_frame(self.enc.as_mut_ptr(), f);
ffi::av_frame_free(&mut f);
if r < 0 {
bail!("avcodec_send_frame(CUDA) failed ({r})");
}
}
Ok(())
}
}
impl Drop for NvencEncoder {
fn drop(&mut self) {
if let Some(sws) = self.sws_444.take() {
// SAFETY: `sws` is the non-null `SwsContext` allocated by `sws_getContext` in `open` and
// owned exclusively by this encoder (taken out of the field so it can't be freed twice).
// `sws_freeContext` frees it; nothing else references it after this single-threaded drop.
unsafe { ffi::sws_freeContext(sws) };
}
}
}
/// Probe whether this NVIDIA GPU + driver + libavcodec can actually encode HEVC **4:4:4** (Range
/// Extensions). Opens a tiny real `hevc_nvenc` 4:4:4 session — the exact path [`NvencEncoder::open`]
/// takes for a live 4:4:4 stream — and reports whether it succeeded. HEVC-only; the result is cached
/// by the caller ([`crate::encode::can_encode_444`]). A GPU/driver/ffmpeg without RExt 4:4:4 fails
/// the open here, so the host resolves the session to 4:2:0 before the Welcome (honest downgrade).
pub fn probe_can_encode_444(codec: Codec) -> bool {
if codec != Codec::H265 {
return false;
}
if ffmpeg::init().is_err() {
return false;
}
// Quiet ffmpeg's open error on a GPU that lacks 4:4:4 — the probe failing is an expected outcome.
// SAFETY: libav initialized above; `av_log_{get,set}_level` only read/write the global int level
// (no pointer args) and are always sound post-init.
let prev = unsafe {
let p = ffi::av_log_get_level();
ffi::av_log_set_level(ffi::AV_LOG_FATAL);
p
};
let ok = NvencEncoder::open(
codec,
PixelFormat::Bgra,
640,
480,
30,
2_000_000,
false, // CPU input (the 4:4:4 path never uses CUDA)
8,
ChromaFormat::Yuv444,
)
.is_ok();
// SAFETY: restore the saved global log level (scalar arg, no pointers).
unsafe { ffi::av_log_set_level(prev) };
ok
}