diff --git a/Cargo.lock b/Cargo.lock index 005a6e24..5afef537 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2303,6 +2303,7 @@ dependencies = [ "libc", "log", "nesprotocol", + "nespyro", "once_cell", "pixelforge", "pollster", @@ -2393,6 +2394,18 @@ dependencies = [ "serde_json", ] +[[package]] +name = "nespyro" +version = "0.1.0" +dependencies = [ + "ash", + "futures-channel", + "nesprotocol", + "pollster", + "thiserror 2.0.21", + "tracing", +] + [[package]] name = "neswire" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index e661c596..b0b06432 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ members = [ "apps/nesinit", "apps/neswire", "crates/nesprotocol", + "crates/nespyro", ] # One pin per dependency, here. A member never writes a version — that is what diff --git a/apps/nescapture/Cargo.toml b/apps/nescapture/Cargo.toml index 69217304..58675961 100644 --- a/apps/nescapture/Cargo.toml +++ b/apps/nescapture/Cargo.toml @@ -39,3 +39,6 @@ libc = "0.2" # Shared IPC protocol nesprotocol = { path = "../../crates/nesprotocol" } + +# PyroWave, the LAN codec, on the same device as pixelforge. +nespyro = { path = "../../crates/nespyro" } diff --git a/apps/nescapture/src/encode.rs b/apps/nescapture/src/encode.rs index ed9da0f0..6a54ffbd 100644 --- a/apps/nescapture/src/encode.rs +++ b/apps/nescapture/src/encode.rs @@ -14,6 +14,14 @@ // Where the game's device cannot host the encoder, it gets a device of its // own, and the slot is read back on the CPU and uploaded there instead. // +// PyroWave takes a shorter path, on the game's device only: +// +// nespyro::Encoder::encode_after(slot, ...) waits on the blit, compute +// shaders, RGB in, packets out +// IPC send to neshub, as a run of messages +// +// It is never chosen by negotiation, only by a client asking for it. +// // Environment variables // ────────────────────── // NESCAPTURE_CODEC "h264" | "h265" | "av1" (default: best available) @@ -46,9 +54,9 @@ use std::thread; use std::time::Instant; use nesprotocol::{ - CODEC_AV1, CODEC_H264, CODEC_H265, CODEC_KEEP, ClientCaps, FLAG_KEYFRAME, FLAG_RECONFIG, - MSG_CLIENT_CAPS, MSG_ENCODE_SETTINGS, MSG_IDR_REQUEST, STREAM_VIDEO, decode_client_caps, - decode_encode_settings, encode_ipc_frame, + CODEC_AV1, CODEC_H264, CODEC_H265, CODEC_KEEP, CODEC_PYROWAVE, ClientCaps, FLAG_KEYFRAME, + FLAG_RECONFIG, MSG_CLIENT_CAPS, MSG_ENCODE_SETTINGS, MSG_IDR_REQUEST, STREAM_VIDEO, + decode_client_caps, decode_encode_settings_ext, encode_ipc_frame, }; use pixelforge::{ Codec, ColorConverter, ColorConverterConfig, ColorRange, ColorSpec, EncodeBitDepth, @@ -397,8 +405,18 @@ fn spawn_stall_watchdog( } struct EncodedFrame { - future: EncodeFuture, + future: Pending, present_time: Instant, + /// The frame's own size, for the IPC header. It used to be the size the + /// pipeline was created at, which a game changing resolution left stale. + width: u32, + height: u32, +} + +/// An encode in flight, from whichever encoder took the frame. +enum Pending { + Hw(EncodeFuture), + Pyro(nespyro::EncodeFuture), } pub enum FrameSource { @@ -437,6 +455,80 @@ impl HwCodec { } } +/// What the stream is being encoded as. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamCodec { + /// A Vulkan Video codec, through pixelforge. + Hw(HwCodec), + /// PyroWave, through nespyro: the LAN codec, used only when a client asks. + PyroWave, +} + +impl StreamCodec { + fn to_protocol_codec(self) -> u8 { + match self { + Self::Hw(c) => c.to_protocol_codec(), + Self::PyroWave => CODEC_PYROWAVE, + } + } +} + +/// The protocol id of what is being encoded, or `CODEC_KEEP` for nothing. +fn protocol_codec(codec: Option) -> u8 { + codec.map_or(CODEC_KEEP, StreamCodec::to_protocol_codec) +} + +/// What a PyroWave request named, kept across rebuilds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PyroSettings { + pub kbps: u32, + /// One of `nesprotocol::pyrowave::CHROMA_*`. + pub chroma: u8, + /// Largest packet the hub can carry in one datagram. + pub packet_size: u16, +} + +/// A PyroWave request, as the command socket delivered it, checked. +/// +/// The hub fills in the rate and the packet size before forwarding, so a +/// request missing either did not come through one; it is refused rather +/// than encoded at a guess. `pyro_available` is whether this device has a +/// nespyro context at all. +fn pyro_request( + settings: &nesprotocol::EncodeSettings, + pyro_available: bool, +) -> Result { + if !pyro_available { + return Err("PyroWave runs only on the game's own device, and this is not one".into()); + } + if settings.rate_control != nesprotocol::RC_CBR || settings.value == 0 { + return Err("a PyroWave request needs a rate".into()); + } + let packet_size = settings + .packet_size + .ok_or("a PyroWave request needs a packet size")?; + if usize::from(packet_size) < MIN_PYRO_PACKET { + return Err(format!( + "a {packet_size}-byte packet is below the {MIN_PYRO_PACKET} the packetizer takes" + )); + } + let chroma = settings.chroma.unwrap_or(nesprotocol::pyrowave::CHROMA_420); + if !matches!( + chroma, + nesprotocol::pyrowave::CHROMA_420 | nesprotocol::pyrowave::CHROMA_444 + ) { + return Err(format!("unknown chroma format {chroma}")); + } + Ok(PyroSettings { + kbps: settings.value, + chroma, + packet_size, + }) +} + +/// nespyro's floor on a packet: a block header and something after it. +const MIN_PYRO_PACKET: usize = 64; + fn probe_any() -> Option<(HwCodec, pixelforge::VideoContext)> { probe_specific(HwCodec::AV1) .or_else(|| probe_specific(HwCodec::H265)) @@ -486,8 +578,19 @@ fn resolve_codec(requested: Option<&str>) -> Option<(HwCodec, pixelforge::VideoC /// Ten bits is offered for everything but H.264, which matches /// [`depth_for_codec`]: the encoder refuses ten-bit H.264, so advertising it /// would let a negotiation settle on something that is then quietly downgraded. -fn host_caps(ctx: &pixelforge::VideoContext) -> ClientCaps { +/// +/// PyroWave is in it when nespyro is: [`ClientCaps::best`] never picks it, so +/// it is there to describe the host, not to steer the negotiation. +fn host_caps(ctx: Option<&pixelforge::VideoContext>, pyro: bool) -> ClientCaps { let mut caps = ClientCaps::empty(); + if pyro { + caps = caps + .with(CODEC_PYROWAVE, nesprotocol::DEPTH_8) + .with(CODEC_PYROWAVE, nesprotocol::DEPTH_10); + } + let Some(ctx) = ctx else { + return caps; + }; for (codec, id) in [ (HwCodec::AV1, nesprotocol::CODEC_AV1), (HwCodec::H265, nesprotocol::CODEC_H265), @@ -828,7 +931,6 @@ pub struct PipelineHandle { progress_epoch: Instant, idr_requested: Arc, shutdown: Arc, - pub codec: HwCodec, pub capture_fps: Arc, pub encode_avg_ms: Arc, pub capture_ms: Arc, @@ -839,21 +941,41 @@ pub struct PipelineHandle { pub timing: Arc, } +/// The encoders the game's own device was set up for. +pub struct SharedEncoders { + pub video: Option, + pub pyro: Option, +} + impl PipelineHandle { - /// Build the pipeline, on `shared` when given: a context on the game's own - /// device. Without one the encoder gets a device of its own. - pub fn new( - config: PipelineConfig, - shared: Option, - ) -> Result { - let (codec, ctx) = match shared { - Some(ctx) => { - resolve_codec_on(&ctx, config.codec_request.as_deref()).map(|codec| (codec, ctx)) + /// Build the pipeline, on `shared` when given: contexts on the game's own + /// device. Without them the encoder gets a device of its own, which hosts + /// Vulkan Video only. + pub fn new(config: PipelineConfig, shared: Option) -> Result { + let (codec, ctx, pyro_ctx) = match shared { + Some(SharedEncoders { video, pyro }) => { + let codec = video + .as_ref() + .and_then(|ctx| resolve_codec_on(ctx, config.codec_request.as_deref())); + if codec.is_none() && pyro.is_none() { + return Err("no encoder on the game's device".into()); + } + (codec, video, pyro) } - None => resolve_codec(config.codec_request.as_deref()), + None => { + let (codec, ctx) = resolve_codec(config.codec_request.as_deref()) + .ok_or_else(|| "no hardware video encoder found on this GPU".to_string())?; + (Some(codec), Some(ctx), None) + } + }; + if codec.is_none() { + log::warn!( + "no hardware video encoder on this GPU; nothing is encoded until a client asks \ + for PyroWave" + ); } - .ok_or_else(|| "no hardware video encoder found on this GPU".to_string())?; - let device_caps = host_caps(&ctx); + let device_caps = host_caps(ctx.as_ref(), pyro_ctx.is_some()); + let pyro_available = pyro_ctx.is_some(); // The surface colour as the compositor last stated it, for the clients. // They need it for two things the host cannot do for them: picking a // swapchain that reads the values the right way, and telling their own @@ -906,7 +1028,7 @@ impl PipelineHandle { let present_attempts = Arc::new(AtomicU32::new(0)); let capture_attempts = Arc::new(AtomicU32::new(0)); let timing = Arc::new(crate::timing::PresentTiming::default()); - let current_codec = Arc::new(AtomicU8::new(codec.to_protocol_codec())); + let current_codec = Arc::new(AtomicU8::new(protocol_codec(codec.map(StreamCodec::Hw)))); let needs_reconfig_flag = Arc::new(AtomicBool::new(false)); let enc_shutdown = shutdown.clone(); @@ -920,8 +1042,10 @@ impl PipelineHandle { idr_interval: config.idr_interval, encoder_tuning_mode: config.encoder_tuning_mode, pixel_format: config.pixel_format, - codec, + codec: codec.map(StreamCodec::Hw), ctx, + pyro_ctx, + pyro: None, idr_requested: idr_requested.clone(), reconfig_rx, current_codec: current_codec.clone(), @@ -953,9 +1077,8 @@ impl PipelineHandle { ipc_path: config.ipc_path, current_codec: current_codec.clone(), needs_reconfig_flag: needs_reconfig_flag, - width: config.width as u16, - height: config.height as u16, encode_ms: encode_avg_ms.clone(), + dropped: dropped_frames.clone(), idr_requested: idr_requested.clone(), epoch: Instant::now(), reconfig_tx: reconfig_tx.clone(), @@ -1079,10 +1202,11 @@ impl PipelineHandle { client.bits() ); let change = EncodeSettingsChange { - codec: Some(codec), + codec: Some(StreamCodec::Hw(codec)), rate_control_mode: None, value: 0, bit_depth: Some(depth), + pyro: None, }; if reconfig_tx.send(change).is_err() { log::warn!( @@ -1099,14 +1223,32 @@ impl PipelineHandle { } } Ok(n) if n >= 2 && buf[0] == MSG_ENCODE_SETTINGS => { - if let Some((codec_id, rc, value, depth)) = - decode_encode_settings(&buf[1..n]) - { + if let Some(settings) = decode_encode_settings_ext(&buf[1..n]) { + let (codec_id, rc, value, depth) = ( + settings.codec, + settings.rate_control, + settings.value, + settings.bit_depth, + ); + let mut pyro = None; let codec = match codec_id { CODEC_KEEP => None, - CODEC_H264 => Some(HwCodec::H264), - CODEC_H265 => Some(HwCodec::H265), - CODEC_AV1 => Some(HwCodec::AV1), + CODEC_H264 => Some(StreamCodec::Hw(HwCodec::H264)), + CODEC_H265 => Some(StreamCodec::Hw(HwCodec::H265)), + CODEC_AV1 => Some(StreamCodec::Hw(HwCodec::AV1)), + CODEC_PYROWAVE => match pyro_request(&settings, pyro_available) + { + Ok(p) => { + pyro = Some(p); + Some(StreamCodec::PyroWave) + } + Err(why) => { + // Once per request, which is rare: + // a person asked for this. + log::warn!("refusing PyroWave: {why}"); + continue; + } + }, _ => { log::warn!( "unknown codec id {codec_id} in encode settings" @@ -1140,6 +1282,7 @@ impl PipelineHandle { rate_control_mode: rate_control, value, bit_depth, + pyro, }; if reconfig_tx.send(change).is_err() { log::warn!("reconfig channel closed, stopping cmd listener"); @@ -1175,7 +1318,6 @@ impl PipelineHandle { progress_epoch, idr_requested, shutdown, - codec, capture_fps, encode_avg_ms, capture_ms, @@ -1271,8 +1413,13 @@ struct EncoderConfig { idr_interval: u32, encoder_tuning_mode: EncoderTuningMode, pixel_format: PixelFormat, - codec: HwCodec, - ctx: pixelforge::VideoContext, + /// `None` until there is something to encode: a GPU with no hardware + /// encoder waits here for a client to ask for PyroWave. + codec: Option, + ctx: Option, + pyro_ctx: Option, + /// What the last PyroWave request named. + pyro: Option, idr_requested: Arc, reconfig_rx: mpsc::Receiver, current_codec: Arc, @@ -1288,20 +1435,16 @@ struct EncoderConfig { capture_ms: Arc, } -struct EncodedPacket { - data: Vec, - is_key_frame: bool, - frame_number: u32, -} - #[derive(Debug, Clone)] pub struct EncodeSettingsChange { - pub codec: Option, + pub codec: Option, /// `None` leaves rate control exactly as it is, which is what a message /// that only means to change something else says. pub rate_control_mode: Option, pub value: u32, pub bit_depth: Option, + /// Set exactly when `codec` is PyroWave. + pub pyro: Option, } fn encoder_thread( @@ -1312,18 +1455,25 @@ fn encoder_thread( epoch: Instant, progress: Arc, ) { - let ctx = cfg.ctx; + let ctx = cfg.ctx.take(); + let pyro_ctx = cfg.pyro_ctx.take(); + // At most one of these exists at a time: a codec switch drops whichever + // was running before the other is built. let mut encoder_state: Option = None; + let mut pyro_state: Option = None; // Last VkFormat we refused, so the error is logged on change rather than // once per frame. let mut unsupported_format: Option = None; + // So "nothing to encode yet" is said once rather than per frame. + let mut said_idle = false; let mut frame_number = 0u32; - // The slot of the last frame converted in place, until that conversion - // is known to be done. See `HeldSlot`. - let mut held: Option = None; + // Slots whose image an encoder may still be reading, oldest first, until + // that is known to be done. See `HeldSlot`. The hardware path holds one + // at a time; PyroWave up to its two frames in flight. + let mut held: std::collections::VecDeque = std::collections::VecDeque::new(); let wanted_depth = std::env::var("NESCAPTURE_DEPTH"); // So the ten-bit refusal is said once per codec rather than per frame. @@ -1340,11 +1490,7 @@ fn encoder_thread( // Also on idle ticks, so the last frame before a pause gives its slot // back: a ring rebuild waits for every slot. - if held.as_ref().is_some_and(HeldSlot::done) - && let Some(done) = held.take() - { - done.release(); - } + release_done(&mut held); // Check for dynamic encode settings changes if let Ok(change) = cfg.reconfig_rx.try_recv() { @@ -1360,16 +1506,29 @@ fn encoder_thread( // a rebuild and an IDR, and the negotiation sends one every time a // client connects -- including the common case where the client // wants exactly what this encoder is already producing. - if changes_nothing(&change, cfg.codec, cfg.wanted_depth_override) { + if changes_nothing(&change, cfg.codec, cfg.wanted_depth_override, cfg.pyro) { log::debug!("reconfig asks for the current settings; nothing to do"); continue; } - if let Some(kbps) = bitrate_only_change( - &change, - cfg.rate_control, - cfg.codec, - cfg.wanted_depth_override, - ) && let Some(state) = encoder_state.as_mut() + if cfg.codec == Some(StreamCodec::PyroWave) + && let Some(kbps) = pyro_bitrate_only(&change, cfg.pyro, cfg.wanted_depth_override) + && let (Some(state), Some(settings)) = (pyro_state.as_mut(), cfg.pyro.as_mut()) + { + state + .encoder + .set_target_bitrate(u64::from(kbps).saturating_mul(1_000)); + settings.kbps = kbps; + log::trace!("PyroWave rate → {kbps} kbps (no rebuild)"); + continue; + } + if matches!(cfg.codec, Some(StreamCodec::Hw(_))) + && let Some(kbps) = bitrate_only_change( + &change, + cfg.rate_control, + cfg.codec, + cfg.wanted_depth_override, + ) + && let Some(state) = encoder_state.as_mut() { match state.encoder.set_target_bitrate(bps(kbps)) { Ok(()) => { @@ -1384,48 +1543,72 @@ fn encoder_thread( Err(e) => log::info!("live retune refused ({e}), rebuilding"), } } + // A codec this device has no encoder for is refused here, before + // anything else in the change is applied: half a change would + // leave the running encode configured for a codec it is not. + let unavailable = match change.codec { + Some(StreamCodec::PyroWave) => pyro_ctx.is_none(), + Some(StreamCodec::Hw(_)) => ctx.is_none(), + None => false, + }; + if unavailable { + log::warn!( + "reconfig to {:?} refused: this device has no encoder for it", + change.codec + ); + continue; + } log::info!( "reconfig: codec={:?}, rc={:?}, value={}", change.codec, change.rate_control_mode, change.value, ); - match change.rate_control_mode { - // Said nothing about rate control, so nothing changes. The - // rebuild below still happens: a depth change needs one. - None => {} - // `retargeted` rather than an outright CBR, so a VBR encode - // keeps its ceiling across a rebuild it is having for some - // other reason -- a codec change, say. The settings message has - // no way to name VBR, so every bitrate arrives labelled CBR; - // reading that label as a mode would mean a ceiling asked for - // at launch survived only until the first codec switch. - Some(RateControlMode::Cbr) => { - cfg.rate_control = cfg.rate_control.retargeted(change.value); - } - Some(RateControlMode::Cqp) => { - cfg.rate_control = RateControl::Cqp { qp: change.value }; - } - _ => { - log::warn!( - "unsupported rate control mode {:?}, keeping current", - change.rate_control_mode - ); + // A PyroWave request's rate is PyroWave's alone. Folding it into + // the hardware rate control would leave the next hardware encode + // aiming at hundreds of megabits. + if change.codec == Some(StreamCodec::PyroWave) { + cfg.pyro = change.pyro; + } else { + match change.rate_control_mode { + // Said nothing about rate control, so nothing changes. The + // rebuild below still happens: a depth change needs one. + None => {} + // `retargeted` rather than an outright CBR, so a VBR encode + // keeps its ceiling across a rebuild it is having for some + // other reason -- a codec change, say. The settings message + // has no way to name VBR, so every bitrate arrives labelled + // CBR; reading that label as a mode would mean a ceiling + // asked for at launch survived only until the first codec + // switch. + Some(RateControlMode::Cbr) => { + cfg.rate_control = cfg.rate_control.retargeted(change.value); + } + Some(RateControlMode::Cqp) => { + cfg.rate_control = RateControl::Cqp { qp: change.value }; + } + _ => { + log::warn!( + "unsupported rate control mode {:?}, keeping current", + change.rate_control_mode + ); + } } } if let Some(codec) = change.codec { - cfg.codec = codec; + cfg.codec = Some(codec); } if let Some(depth) = change.bit_depth { cfg.wanted_depth_override = Some(depth); } // Drop old encoder state to force re-creation with new settings - drop_encoder(&mut encoder_state, &mut held); + drop_encoder(&mut encoder_state, &mut pyro_state, &mut held); + init_failed_at = None; // Signal IPC thread to set FLAG_RECONFIG on next frame cfg.needs_reconfig_flag.store(true, Ordering::Relaxed); // Update IPC thread with new codec cfg.current_codec - .store(cfg.codec.to_protocol_codec(), Ordering::Relaxed); + .store(protocol_codec(cfg.codec), Ordering::Relaxed); // Request IDR so first frame after reconfig has new SPS/PPS cfg.idr_requested.store(true, Ordering::Relaxed); } @@ -1437,6 +1620,16 @@ fn encoder_thread( Err(mpsc::RecvTimeoutError::Disconnected) => break, }; + let Some(codec) = cfg.codec else { + // Dropping the frame gives its slot straight back. + if !said_idle { + log::info!("no encoder chosen yet; frames are dropped until a client asks for one"); + said_idle = true; + } + continue; + }; + said_idle = false; + // Wait for the blit and export the buffer. This used to be a thread of // its own between the present hook and here; it is cheaper on this one, // because the blit it waits for was submitted a frame earlier and has @@ -1487,14 +1680,7 @@ fn encoder_thread( _ => EncodeBitDepth::Eight, } }; - let bit_depth = depth_for_codec(cfg.codec, requested_depth, &mut warned_depth); - let out_fmt = output_format(cfg.pixel_format, bit_depth); - // The geometry joins the guard. It used to be absent, and `cfg.width` / - // `cfg.height` were whatever the *first* frame happened to be, so a game - // that changed resolution kept being encoded at the old one: shrinking - // left a band of the previous picture down the right edge and along the - // bottom, and growing had no surface large enough to hold the frame. // Resolved per frame, because the compositor can say what a surface // is after the swapchain that carries it was created -- and on the // path that needs this, it always does. @@ -1506,98 +1692,265 @@ fn encoder_thread( cfg.stream_colorspace .store(frame_colorspace, Ordering::Relaxed); - let state = match encoder_state.as_mut() { - Some(s) - if encoder_still_serves( - ( - s.width, - s.height, - s.bit_depth, - s.pixel_format, - s.colorspace, - s.input_fmt, - ), - ( - raw.width, - raw.height, - bit_depth, - cfg.pixel_format, - frame_colorspace, - input_fmt, - ), - ) => - { - s - } - _ => { - // Building an encoder is expensive and a configuration the - // device cannot do will not start working on the next frame. - // Without this, a refused profile is retried sixty times a - // second forever -- which is how an unsupported ten-bit H.264 - // request turned into a log with nothing else in it and a - // session that never recovered. - if let Some(failed_at) = init_failed_at - && failed_at.elapsed() < INIT_RETRY_INTERVAL - { - frame_number += 1; + // Which ring slot the frame is in, for the slot hold below. + let buffer_index = raw.slot.as_ref().map(|s| s.index()).unwrap_or(0); + + let future = match codec { + StreamCodec::PyroWave => { + let (Some(pyro_ctx), Some(settings)) = (pyro_ctx.as_ref(), cfg.pyro) else { + // Refused before it got here; unreachable unless that + // check is lost. + log::error!("PyroWave chosen with no nespyro context or no settings"); + cfg.codec = None; continue; - } - if let Some(old) = encoder_state.as_ref() { - if old.width != raw.width || old.height != raw.height { - log::info!( - "resolution changed {}x{} -> {}x{}, rebuilding the encoder", - old.width, - old.height, - raw.width, - raw.height, - ); - } - // Said out loud because a game can change this without - // changing anything else about the surface -- Cyberpunk's - // HDR10 and scRGB modes differ in the format and in - // nothing the old check looked at. - if old.input_fmt != input_fmt { - log::info!( - "surface format changed {:?} -> {:?}, rebuilding the encoder", - old.input_fmt, - input_fmt, - ); - } - } - // The old converter may still be reading a held slot, and the - // held slot's point is on the old converter's timeline. - drop_encoder(&mut encoder_state, &mut held); - match PerFrameEncoder::new( - &ctx, - cfg.codec.to_pixelforge(), - raw.width, - raw.height, - cfg.fps, - cfg.rate_control, - cfg.idr_interval, - cfg.encoder_tuning_mode, - cfg.pixel_format, - bit_depth, + }; + let FrameSource::Shared { image, blit } = source else { + log::error!("PyroWave chosen on a device of the encoder's own"); + cfg.codec = None; + continue; + }; + let key = PyroKey { + width: raw.width, + height: raw.height, input_fmt, - out_fmt, - frame_colorspace, - matches!(source, FrameSource::Shared { .. }), + colorspace: frame_colorspace, + ten_bit: requested_depth == EncodeBitDepth::Ten, + chroma: settings.chroma, + packet_size: settings.packet_size, + }; + if pyro_state.as_ref().is_none_or(|s| s.key != key) { + if let Some(failed_at) = init_failed_at + && failed_at.elapsed() < INIT_RETRY_INTERVAL + { + frame_number += 1; + continue; + } + drop_encoder(&mut encoder_state, &mut pyro_state, &mut held); + match PyroEncoder::new(pyro_ctx, key, settings.kbps, cfg.fps) { + Ok(s) => { + init_failed_at = None; + last_init_error = None; + pyro_state = Some(s); + } + Err(e) => { + if last_init_error.as_deref() != Some(e.as_str()) { + log::error!("PyroWave encoder (re)init: {e}"); + last_init_error = Some(e); + } + init_failed_at = Some(std::time::Instant::now()); + frame_number += 1; + continue; + } + } + } + let state = pyro_state.as_mut().expect("built above"); + // Every PyroWave frame is a key frame, so a request for one + // has already been met. + cfg.idr_requested.store(false, Ordering::Relaxed); + + progress.note(epoch, 2); + let wait = nespyro::TimelinePoint::new(blit.semaphore, blit.value); + match state.encoder.encode_after( + image, + input_fmt.vk_format(), + ash::vk::ImageLayout::GENERAL, + &[wait], ) { - Ok(s) => { - init_failed_at = None; - last_init_error = None; - encoder_state = Some(s); - encoder_state.as_mut().unwrap() + Ok(future) => { + // The slot stays the encoder's until its timeline says + // the frame is done with: nespyro reads the image in + // place, as late as its own queue gets to it. + let done = state.encoder.last_submission(); + if let Some(guard) = raw.slot.take() { + held.push_back(HeldSlot { + guard, + slot: buffer_index, + converted: pixelforge::TimelinePoint::new( + done.semaphore, + done.value, + ), + ds: ds.clone(), + }); + } + Pending::Pyro(future) } Err(e) => { - // Said once per distinct failure. The same refusal every - // second says nothing the first one did not, and buries - // the one that is different. - if last_init_error.as_deref() != Some(e.as_str()) { - log::error!("encoder (re)init: {e}"); - last_init_error = Some(e); + log::warn!("PyroWave encode frame {frame_number}: {e}"); + frame_number += 1; + continue; + } + } + } + StreamCodec::Hw(hw) => { + let Some(ctx) = ctx.as_ref() else { + log::error!("{hw:?} chosen with no Vulkan Video context"); + cfg.codec = None; + continue; + }; + let bit_depth = depth_for_codec(hw, requested_depth, &mut warned_depth); + let out_fmt = output_format(cfg.pixel_format, bit_depth); + + // The geometry joins the guard. It used to be absent, and + // `cfg.width` / `cfg.height` were whatever the *first* frame + // happened to be, so a game that changed resolution kept being + // encoded at the old one: shrinking left a band of the previous + // picture down the right edge and along the bottom, and growing + // had no surface large enough to hold the frame. + let state = match encoder_state.as_mut() { + Some(s) + if encoder_still_serves( + ( + s.width, + s.height, + s.bit_depth, + s.pixel_format, + s.colorspace, + s.input_fmt, + ), + ( + raw.width, + raw.height, + bit_depth, + cfg.pixel_format, + frame_colorspace, + input_fmt, + ), + ) => + { + s + } + _ => { + // Building an encoder is expensive and a configuration + // the device cannot do will not start working on the + // next frame. Without this, a refused profile is + // retried sixty times a second forever -- which is how + // an unsupported ten-bit H.264 request turned into a + // log with nothing else in it and a session that never + // recovered. + if let Some(failed_at) = init_failed_at + && failed_at.elapsed() < INIT_RETRY_INTERVAL + { + frame_number += 1; + continue; } - init_failed_at = Some(std::time::Instant::now()); + if let Some(old) = encoder_state.as_ref() { + if old.width != raw.width || old.height != raw.height { + log::info!( + "resolution changed {}x{} -> {}x{}, rebuilding the encoder", + old.width, + old.height, + raw.width, + raw.height, + ); + } + // Said out loud because a game can change this + // without changing anything else about the surface + // -- Cyberpunk's HDR10 and scRGB modes differ in + // the format and in nothing the old check looked at. + if old.input_fmt != input_fmt { + log::info!( + "surface format changed {:?} -> {:?}, rebuilding the encoder", + old.input_fmt, + input_fmt, + ); + } + } + // The old converter may still be reading a held slot, + // and the held slot's point is on the old converter's + // timeline. + drop_encoder(&mut encoder_state, &mut pyro_state, &mut held); + match PerFrameEncoder::new( + ctx, + hw.to_pixelforge(), + raw.width, + raw.height, + cfg.fps, + cfg.rate_control, + cfg.idr_interval, + cfg.encoder_tuning_mode, + cfg.pixel_format, + bit_depth, + input_fmt, + out_fmt, + frame_colorspace, + matches!(source, FrameSource::Shared { .. }), + ) { + Ok(s) => { + init_failed_at = None; + last_init_error = None; + encoder_state = Some(s); + encoder_state.as_mut().unwrap() + } + Err(e) => { + // Said once per distinct failure. The same + // refusal every second says nothing the first + // one did not, and buries the one that is + // different. + if last_init_error.as_deref() != Some(e.as_str()) { + log::error!("encoder (re)init: {e}"); + last_init_error = Some(e); + } + init_failed_at = Some(std::time::Instant::now()); + frame_number += 1; + continue; + } + } + } + }; + + // Only when something asked. Periodic key frames are the + // encoder's own schedule, set by `with_gop_size` -- counting + // frames here as well meant two mechanisms driving one thing, + // and the encoder's schedule being the one that could be turned + // off. Turning it off changed nothing, because this kept asking + // every four seconds regardless. + if cfg.idr_requested.swap(false, Ordering::Relaxed) { + state.encoder.request_idr(); + } + + progress.note(epoch, 2); + let result = match &mut source { + // The encoder converts RGB itself: it copies the slot and + // waits for the copy before returning, so the slot needs no + // holding. + FrameSource::Shared { image, blit } if state.converter.is_none() => state + .encoder + .encode_after(*image, &[*blit]) + .map_err(|e| anyhow::anyhow!("Encoder::encode_after: {e}")), + FrameSource::Shared { image, blit } => { + let converter = state.converter.as_mut().expect("guarded above"); + let (result, converted) = + shared_encode_frame(converter, &mut state.encoder, *image, *blit); + // Replacing the held slot is what releases it: a + // successful conversion started only once the previous + // one had finished. + if let (Some(converted), Some(guard)) = (converted, raw.slot.take()) { + while let Some(prev) = held.pop_front() { + prev.release(); + } + held.push_back(HeldSlot { + guard, + slot: buffer_index, + converted, + ds: ds.clone(), + }); + } + result + } + FrameSource::Pixels(pixels) => cpu_encode_frame( + ctx, + hw.to_pixelforge(), + &mut state.encoder, + pixels.as_slice(), + raw.width, + raw.height, + raw.vk_format, + ), + }; + match result { + Ok(future) => Pending::Hw(future), + Err(e) => { + log::warn!("encode frame {frame_number}: {e}"); frame_number += 1; continue; } @@ -1605,76 +1958,24 @@ fn encoder_thread( } }; - // Only when something asked. Periodic key frames are the encoder's own - // schedule, set by `with_gop_size` -- counting frames here as well - // meant two mechanisms driving one thing, and the encoder's schedule - // being the one that could be turned off. Turning it off changed - // nothing, because this kept asking every four seconds regardless. - if cfg.idr_requested.swap(false, Ordering::Relaxed) { - state.encoder.request_idr(); - } - - // Which ring slot the frame is in, for the slot hold below. - let buffer_index = raw.slot.as_ref().map(|s| s.index()).unwrap_or(0); - - progress.note(epoch, 2); - let result = match &mut source { - // The encoder converts RGB itself: it copies the slot and waits for - // the copy before returning, so the slot needs no holding. - FrameSource::Shared { image, blit } if state.converter.is_none() => state - .encoder - .encode_after(*image, &[*blit]) - .map_err(|e| anyhow::anyhow!("Encoder::encode_after: {e}")), - FrameSource::Shared { image, blit } => { - let converter = state.converter.as_mut().expect("guarded above"); - let (result, converted) = - shared_encode_frame(converter, &mut state.encoder, *image, *blit); - // Replacing the held slot is what releases it: a successful - // conversion started only once the previous one had finished. - if let (Some(converted), Some(guard)) = (converted, raw.slot.take()) - && let Some(prev) = held.replace(HeldSlot { - guard, - slot: buffer_index, - converted, - ds: ds.clone(), - }) - { - prev.release(); - } - result - } - FrameSource::Pixels(pixels) => cpu_encode_frame( - &ctx, - &mut state.encoder, - pixels.as_slice(), - raw.width, - raw.height, - raw.vk_format, - ), + // Blocking, deliberately. Dropping an encoded frame does not just + // waste the encode — it breaks the reference chain. The encoder's DPB + // believes the frame exists and codes later frames against it, so a + // decoder that never receives it shows corruption until the next IDR. + // Blocking here pushes back through `frame_rx` to `push_frame`, where a + // drop is free: that frame never entered the encoder and no later + // frame refers to it. + let pending = EncodedFrame { + future, + present_time: raw.present_time, + width: raw.width, + height: raw.height, }; - - match result { - Err(e) => log::warn!("encode frame {frame_number}: {e}"), - Ok(future) => { - // Blocking, deliberately. Dropping an encoded frame does not - // just waste the encode — it breaks the reference chain. The - // encoder's DPB believes the frame exists and codes later - // frames against it, so a decoder that never receives it shows - // corruption until the next IDR. Blocking here pushes back - // through `frame_rx` to `push_frame`, where a drop is free: - // that frame never entered the encoder and no later frame - // refers to it. - let pending = EncodedFrame { - future, - present_time: raw.present_time, - }; - // The blocking send. If the thread draining this stops, every - // frame after the second one waits here forever. - progress.note(epoch, 3); - if encoded_tx.send(pending).is_err() { - break; - } - } + // The blocking send. If the thread draining this stops, every frame + // after the second one waits here forever. + progress.note(epoch, 3); + if encoded_tx.send(pending).is_err() { + break; } frame_number += 1; @@ -1684,11 +1985,154 @@ fn encoder_thread( if let Some(state) = encoder_state.as_mut() { let _ = state.encoder.flush(); } - drop_encoder(&mut encoder_state, &mut held); + drop_encoder(&mut encoder_state, &mut pyro_state, &mut held); log::info!("encoder thread exited"); } +/// The configuration a PyroWave encoder was built for. A frame wanting any +/// other needs a new encoder. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct PyroKey { + width: u32, + height: u32, + input_fmt: InputFormat, + colorspace: u32, + /// Whether ten bits were asked for. An HDR stream gets sixteen-bit planes + /// either way. + ten_bit: bool, + chroma: u8, + packet_size: u16, +} + +struct PyroEncoder { + encoder: nespyro::Encoder, + key: PyroKey, +} + +impl PyroEncoder { + fn new(ctx: &nespyro::Context, key: PyroKey, kbps: u32, fps: u32) -> Result { + let config = pyro_config(key, kbps, fps)?; + log::info!( + "(re)init PyroWave encoder: {}x{} {:?} {:?} {:?} → {:?}, {kbps} kbps, packets of {} bytes", + key.width, + key.height, + config.chroma, + config.depth, + config.source, + if config.colour.is_hdr() { + "BT.2020 PQ" + } else { + "BT.709" + }, + key.packet_size, + ); + let encoder = nespyro::Encoder::new(ctx.clone(), config) + .map_err(|e| format!("nespyro::Encoder::new: {e}"))?; + Ok(Self { encoder, key }) + } +} + +/// nespyro's configuration for a key, decided by the same `source_spec` and +/// `stream_spec` the hardware path's converter uses, so the two codecs cannot +/// disagree about what a surface is. +fn pyro_config(key: PyroKey, kbps: u32, fps: u32) -> Result { + let source = source_spec(key.colorspace, key.input_fmt); + let stream = stream_spec(source); + let (colour, depth) = match stream { + ColorSpec::Bt2020Pq => ( + nespyro::ColourDescription::bt2020_pq(), + nespyro::Depth::Sixteen, + ), + ColorSpec::Srgb => ( + nespyro::ColourDescription::bt709(), + if key.ten_bit { + nespyro::Depth::Sixteen + } else { + nespyro::Depth::Eight + }, + ), + other => return Err(format!("{other:?} is not a stream PyroWave can carry")), + }; + let source = match source { + ColorSpec::Srgb => nespyro::Source::Srgb, + ColorSpec::Bt709Linear => nespyro::Source::Bt709Linear, + ColorSpec::Bt2020Linear => nespyro::Source::Bt2020Linear, + ColorSpec::Bt2020Pq => nespyro::Source::Bt2020Pq, + }; + let chroma = match key.chroma { + nesprotocol::pyrowave::CHROMA_444 => nespyro::Chroma::Yuv444, + _ => nespyro::Chroma::Yuv420, + }; + // pixelforge's own figure for the same source, so a surface is exactly as + // bright whichever encoder carries it: 80 nits for scRGB, 203 for the rest. + let white = source_spec(key.colorspace, key.input_fmt).reference_white_nits(); + let config = nespyro::EncodeConfig::new(key.width, key.height); + let config = match white { + Some(nits) => config.with_reference_white(nits), + None => config, + }; + Ok(config + .with_chroma(chroma) + .with_depth(depth) + .with_colour(colour, source) + .with_frame_rate(fps.max(1), 1) + .with_target_bitrate(u64::from(kbps).saturating_mul(1_000)) + .with_packet_size(usize::from(key.packet_size))) +} + +/// The new PyroWave rate when a change is nothing but one. +fn pyro_bitrate_only( + change: &EncodeSettingsChange, + current: Option, + current_depth: Option, +) -> Option { + let current = current?; + if change.rate_control_mode != Some(RateControlMode::Cbr) || change.value == 0 { + return None; + } + if change.codec.is_some_and(|c| c != StreamCodec::PyroWave) { + return None; + } + if change.bit_depth.is_some_and(|d| Some(d) != current_depth) { + return None; + } + if change + .pyro + .is_some_and(|p| p.chroma != current.chroma || p.packet_size != current.packet_size) + { + return None; + } + Some(change.value) +} + +/// Give back every held slot whose encoder is done with it, oldest first. +/// +/// In order, and stopping at the first that is not done: points on one +/// timeline are reached in order, so a later one being reached says nothing +/// that the earlier one has not. +fn release_done(held: &mut std::collections::VecDeque) { + while held.front().is_some_and(H::done) { + held.pop_front().expect("checked").release(); + } +} + +/// Something held until a GPU point is reached: a [`HeldSlot`], or a stand-in +/// in tests, where there is no device to hold a slot of. +trait Hold { + fn done(&self) -> bool; + fn release(self); +} + +impl Hold for HeldSlot { + fn done(&self) -> bool { + HeldSlot::done(self) + } + fn release(self) { + HeldSlot::release(self) + } +} + struct PerFrameEncoder { encoder: Encoder, /// The colour space this encoder and its converter were built for. @@ -1975,13 +2419,17 @@ fn intra_refresh_enabled() -> bool { /// which is never a no-op. fn changes_nothing( change: &EncodeSettingsChange, - current_codec: HwCodec, + current_codec: Option, current_depth: Option, + current_pyro: Option, ) -> bool { if change.rate_control_mode.is_some() { return false; } - let codec_same = change.codec.is_none_or(|c| c == current_codec); + if change.pyro.is_some() && change.pyro != current_pyro { + return false; + } + let codec_same = change.codec.is_none_or(|c| Some(c) == current_codec); // `None` for the current depth means nothing has overridden the default, // which is eight bits -- so a request for eight bits is a no-op and a // request for ten is not. @@ -1994,7 +2442,7 @@ fn changes_nothing( fn bitrate_only_change( change: &EncodeSettingsChange, current: RateControl, - current_codec: HwCodec, + current_codec: Option, current_depth_override: Option, ) -> Option { if change.rate_control_mode != Some(RateControlMode::Cbr) { @@ -2010,7 +2458,7 @@ fn bitrate_only_change( // ceiling asked for at launch would last exactly until the path moved. // What the message carries is a target; the mode is what was asked for. current.target_kbps()?; - if change.codec.is_some_and(|c| c != current_codec) { + if change.codec.is_some_and(|c| Some(c) != current_codec) { return None; } if change @@ -2109,14 +2557,20 @@ impl HeldSlot { } } -/// Drop the encoder and its converter, then the held slot. +/// Drop whichever encoder there is, then the held slots. /// -/// In that order: dropping the converter waits for its last conversion, after -/// which the held slot is free and its point, on the converter's timeline, -/// would be a dangling handle anyway. -fn drop_encoder(state: &mut Option, held: &mut Option) { +/// In that order: dropping the converter waits for its last conversion, and +/// dropping a PyroWave encoder for its last frame, after which the held slots +/// are free and their points, on the dropped encoder's timeline, would be +/// dangling handles anyway. +fn drop_encoder( + state: &mut Option, + pyro: &mut Option, + held: &mut std::collections::VecDeque, +) { *state = None; - if let Some(slot) = held.take() { + *pyro = None; + while let Some(slot) = held.pop_front() { slot.release(); } } @@ -2157,6 +2611,7 @@ fn map_vk_format_raw(vk_format: u32) -> ash::vk::Format { fn cpu_encode_frame( ctx: &pixelforge::VideoContext, + codec: Codec, encoder: &mut Encoder, pixels: &[u8], width: u32, @@ -2183,9 +2638,11 @@ fn cpu_encode_frame( let yuv = bgra_to_yuv420(pixels, width, height, vk_format); + // The encoder's own codec. This was `H264` whatever the encoder was, which + // sized and aligned the upload for the wrong codec's blocks. let mut input_image = InputImage::new( ctx.clone(), - Codec::H264, + codec, width, height, EncodeBitDepth::Eight, @@ -2246,13 +2703,42 @@ fn bgra_to_yuv420(pixels: &[u8], width: u32, height: u32, vk_format: u32) -> Vec // IPC send thread // ───────────────────────────────────────────────────────────────────────────── +/// A PyroWave frame as the IPC messages that carry it to the hub. +/// +/// The frame index is the encoder's own count, cut to the 32 bits the header +/// has: it only has to tell one frame from the next. +fn pyro_messages( + frame: &nespyro::EncodedFrame, + ts_ms: u32, + width: u16, + height: u16, +) -> Result>, nesprotocol::pyrowave::ChunkError> { + nesprotocol::pyrowave::chunk_frame( + frame.index as u32, + nesprotocol::pyrowave::FrameMeta { + ts_ms, + width, + height, + }, + &frame.data, + &frame.packets, + frame.critical_packets, + ) +} + +/// A finished encode, from whichever encoder made it. +enum Encoded { + Hw(pixelforge::EncodedPacket), + Pyro(nespyro::EncodedFrame), +} + struct IpcConfig { ipc_path: std::path::PathBuf, current_codec: Arc, needs_reconfig_flag: Arc, - width: u16, - height: u16, encode_ms: Arc, + /// Frames that never reached the hub, for the stats. + dropped: Arc, /// Shared with the encoder thread, which honours it on the next frame. idr_requested: Arc, /// Zero point for wire timestamps. @@ -2317,6 +2803,8 @@ fn ipc_send_thread( // even though frames are stamped from their own present. let start_time = cfg.epoch; let mut frame_count: u64 = 0; + // So a PyroWave encoder short of scratch space is said once per episode. + let mut scratch_overflowing = false; 'outer: loop { if shutdown.load(Ordering::Relaxed) { @@ -2410,23 +2898,24 @@ fn ipc_send_thread( // thread -- a whole pipeline stalled behind one frame, with // nothing said anywhere. progress.note(epoch, 4); - let result = pollster::block_on(pending.future); + let encoded = match pending.future { + Pending::Hw(future) => pollster::block_on(future) + .map(Encoded::Hw) + .map_err(|e| e.to_string()), + Pending::Pyro(future) => pollster::block_on(future) + .map(Encoded::Pyro) + .map_err(|e| e.to_string()), + }; let awaited = encode_start.elapsed(); let waited = queued + awaited; - let pkt = match result { - Ok(p) => p, + let encoded = match encoded { + Ok(e) => e, Err(e) => { log::warn!("encoder future receive: {e}"); continue; } }; - if let Some(stats) = &pkt.stats { - let enc_ms: f32 = (std::time::Duration::from_nanos(stats.gpu_time_ns).as_secs_f64() - * 1000.0) as f32; - cfg.encode_ms.store(enc_ms.to_bits(), Ordering::Relaxed); - } - // From the present, not from here. Stamping at send time folded // however long the frame spent queued for the encoder into the // timestamp, so the receiver could not tell capture time from @@ -2434,27 +2923,75 @@ fn ipc_send_thread( let timestamp_ms = present_time .saturating_duration_since(start_time) .as_millis() as u32; - let mut flags = if pkt.is_key_frame { FLAG_KEYFRAME } else { 0 }; // Set FLAG_RECONFIG on the first frame after an encoder reconfig. - // Clear it after setting so only the first frame is marked. - if cfg.needs_reconfig_flag.swap(false, Ordering::Relaxed) { - flags |= FLAG_RECONFIG; - } + // Clear it after setting so only the first frame is marked. Taken + // for a PyroWave frame too, which carries no flags, so that the + // mark cannot outlive the switch and land on a later one. + let reconfig = cfg.needs_reconfig_flag.swap(false, Ordering::Relaxed); let protocol_codec = cfg.current_codec.load(Ordering::Relaxed); + // The frame's own size: the IPC header has sixteen bits for it. + let (width, height) = (pending.width as u16, pending.height as u16); - let ipc_frame = encode_ipc_frame( - STREAM_VIDEO, - protocol_codec, - flags, - timestamp_ms, - cfg.width, - cfg.height, - &pkt.data, - ); + let (messages, is_key_frame, bytes) = match &encoded { + Encoded::Hw(pkt) => { + if let Some(stats) = &pkt.stats { + let enc_ms: f32 = (std::time::Duration::from_nanos(stats.gpu_time_ns) + .as_secs_f64() + * 1000.0) as f32; + cfg.encode_ms.store(enc_ms.to_bits(), Ordering::Relaxed); + } + let mut flags = if pkt.is_key_frame { FLAG_KEYFRAME } else { 0 }; + if reconfig { + flags |= FLAG_RECONFIG; + } + let ipc_frame = encode_ipc_frame( + STREAM_VIDEO, + protocol_codec, + flags, + timestamp_ms, + width, + height, + &pkt.data, + ); + (vec![ipc_frame], pkt.is_key_frame, pkt.data.len()) + } + Encoded::Pyro(frame) => { + let enc_ms = (frame.stats.gpu_ns() / 1e6) as f32; + cfg.encode_ms.store(enc_ms.to_bits(), Ordering::Relaxed); + // Said on the edges. It means rate control asked more of + // the quantizer's scratch space than it has, so blocks were + // lost before the frame ever left this box -- and nothing + // downstream can count what was never sent. + let overflowed = frame.stats.scratch_overflowed(); + if overflowed && !scratch_overflowing { + log::warn!( + "PyroWave encoder ran out of scratch space ({} of {} bytes); \ + frames are losing blocks", + frame.stats.scratch_bytes, + frame.stats.scratch_capacity + ); + } else if !overflowed && scratch_overflowing { + log::info!("PyroWave encoder has scratch space again"); + } + scratch_overflowing = overflowed; + match pyro_messages(frame, timestamp_ms, width, height) { + // Every PyroWave frame stands alone, so none is a key + // frame in the sense the counters below mean. + Ok(messages) => (messages, false, frame.data.len()), + Err(e) => { + log::warn!("PyroWave frame {} cannot be sent: {e}", frame.index); + cfg.dropped.fetch_add(1, Ordering::Relaxed); + continue; + } + } + } + }; - if let Some(probe) = rate_probe.as_mut() { + if let Some(probe) = rate_probe.as_mut() + && matches!(encoded, Encoded::Hw(_)) + { let now = Instant::now(); - probe.observe(now, pkt.data.len() as u32, pkt.is_key_frame); + probe.observe(now, bytes as u32, is_key_frame); if let Some(kbps) = probe.due_step(now) { log::info!("rate probe: stepping to {kbps} kbps"); let change = EncodeSettingsChange { @@ -2462,6 +2999,7 @@ fn ipc_send_thread( rate_control_mode: Some(RateControlMode::Cbr), value: kbps, bit_depth: None, + pyro: None, }; let _ = cfg.reconfig_tx.send(change); } @@ -2505,7 +3043,7 @@ fn ipc_send_thread( } } - if pkt.is_key_frame { + if is_key_frame { keyframes += 1; worst_key_wait = worst_key_wait.max(waited); } @@ -2514,7 +3052,14 @@ fn ipc_send_thread( let send_start = Instant::now(); progress.note(epoch, 5); - if let Err(e) = socket.send(&ipc_frame) { + let mut failure = None; + for message in &messages { + if let Err(e) = socket.send(message) { + failure = Some(e); + break; + } + } + if let Some(e) = failure { // Said on the edges only. A stalled consumer fails every frame, // and sixty identical lines a second bury whatever else is // being said about why it stalled. @@ -2525,6 +3070,9 @@ fn ipc_send_thread( log::warn!("IPC consumer is not reading; dropping frames until it does"); blocked_consumer = true; } + // The rest of a PyroWave frame goes with it: the hub + // abandons a frame with a message missing anyway. + cfg.dropped.fetch_add(1, Ordering::Relaxed); frame_count += 1; continue; } @@ -3399,7 +3947,9 @@ mod depth_tests { #[cfg(test)] mod bitrate_only_tests { - use super::{HwCodec, RateControl, bitrate_only_change, changes_nothing, negotiated}; + use super::{ + HwCodec, RateControl, StreamCodec, bitrate_only_change, changes_nothing, negotiated, + }; use crate::encode::EncodeSettingsChange; use pixelforge::{EncodeBitDepth, RateControlMode}; @@ -3411,16 +3961,17 @@ mod bitrate_only_tests { bit_depth: Option, ) -> EncodeSettingsChange { EncodeSettingsChange { - codec, + codec: codec.map(StreamCodec::Hw), rate_control_mode: Some(mode), value, bit_depth, + pyro: None, } } /// The running encode for these: CBR at 8 Mbps, H.264, no depth override. fn running(c: &EncodeSettingsChange) -> Option { - bitrate_only_change(c, CBR_8M, HwCodec::H264, None) + bitrate_only_change(c, CBR_8M, Some(StreamCodec::Hw(HwCodec::H264)), None) } fn caps(pairs: &[(u8, u8)]) -> nesprotocol::ClientCaps { @@ -3492,25 +4043,29 @@ mod bitrate_only_tests { #[test] fn asking_for_the_current_settings_is_a_no_op() { let same = EncodeSettingsChange { - codec: Some(HwCodec::AV1), + codec: Some(StreamCodec::Hw(HwCodec::AV1)), rate_control_mode: None, value: 0, bit_depth: Some(EncodeBitDepth::Ten), + pyro: None, }; assert!(changes_nothing( &same, - HwCodec::AV1, - Some(EncodeBitDepth::Ten) + Some(StreamCodec::Hw(HwCodec::AV1)), + Some(EncodeBitDepth::Ten), + None )); assert!(!changes_nothing( &same, - HwCodec::H265, - Some(EncodeBitDepth::Ten) + Some(StreamCodec::Hw(HwCodec::H265)), + Some(EncodeBitDepth::Ten), + None )); assert!(!changes_nothing( &same, - HwCodec::AV1, - Some(EncodeBitDepth::Eight) + Some(StreamCodec::Hw(HwCodec::AV1)), + Some(EncodeBitDepth::Eight), + None )); } @@ -3523,28 +4078,41 @@ mod bitrate_only_tests { rate_control_mode: None, value: 0, bit_depth: Some(EncodeBitDepth::Eight), + pyro: None, }; - assert!(changes_nothing(&eight, HwCodec::AV1, None)); + assert!(changes_nothing( + &eight, + Some(StreamCodec::Hw(HwCodec::AV1)), + None, + None + )); let ten = EncodeSettingsChange { bit_depth: Some(EncodeBitDepth::Ten), ..eight.clone() }; - assert!(!changes_nothing(&ten, HwCodec::AV1, None)); + assert!(!changes_nothing( + &ten, + Some(StreamCodec::Hw(HwCodec::AV1)), + None, + None + )); } /// A rate-control change is never a no-op, whatever else it carries. #[test] fn a_rate_control_change_always_counts() { let rc = EncodeSettingsChange { - codec: Some(HwCodec::AV1), + codec: Some(StreamCodec::Hw(HwCodec::AV1)), rate_control_mode: Some(RateControlMode::Cbr), value: 5_000, bit_depth: Some(EncodeBitDepth::Ten), + pyro: None, }; assert!(!changes_nothing( &rc, - HwCodec::AV1, - Some(EncodeBitDepth::Ten) + Some(StreamCodec::Hw(HwCodec::AV1)), + Some(EncodeBitDepth::Ten), + None )); } @@ -3558,6 +4126,7 @@ mod bitrate_only_tests { rate_control_mode: None, value: 0, bit_depth: Some(EncodeBitDepth::Ten), + pyro: None, }; assert_eq!( running(&depth_only), @@ -3608,7 +4177,12 @@ mod bitrate_only_tests { Some(EncodeBitDepth::Eight), ); assert_eq!( - bitrate_only_change(&c, CBR_8M, HwCodec::H264, Some(EncodeBitDepth::Eight)), + bitrate_only_change( + &c, + CBR_8M, + Some(StreamCodec::Hw(HwCodec::H264)), + Some(EncodeBitDepth::Eight) + ), Some(2_000), ); } @@ -3627,7 +4201,10 @@ mod bitrate_only_tests { None, Some(EncodeBitDepth::Eight), ); - assert_eq!(bitrate_only_change(&c, CBR_8M, HwCodec::H264, None), None); + assert_eq!( + bitrate_only_change(&c, CBR_8M, Some(StreamCodec::Hw(HwCodec::H264)), None), + None + ); } #[test] @@ -3657,7 +4234,12 @@ mod bitrate_only_tests { // retarget -- it has to become an encode that has one. let c = change(RateControlMode::Cbr, 2_000, None, None); assert_eq!( - bitrate_only_change(&c, RateControl::Cqp { qp: 26 }, HwCodec::H264, None), + bitrate_only_change( + &c, + RateControl::Cqp { qp: 26 }, + Some(StreamCodec::Hw(HwCodec::H264)), + None + ), None, ); } @@ -3670,7 +4252,7 @@ mod bitrate_only_tests { fn a_bitrate_under_vbr_is_a_retune_not_a_mode_change() { let c = change(RateControlMode::Cbr, 2_000, None, None); assert_eq!( - bitrate_only_change(&c, VBR_8M, HwCodec::H264, None), + bitrate_only_change(&c, VBR_8M, Some(StreamCodec::Hw(HwCodec::H264)), None), Some(2_000), ); } @@ -3680,10 +4262,16 @@ mod bitrate_only_tests { #[test] fn a_vbr_encode_rebuilds_for_the_same_reasons_any_other_does() { let c = change(RateControlMode::Cbr, 2_000, Some(HwCodec::AV1), None); - assert_eq!(bitrate_only_change(&c, VBR_8M, HwCodec::H264, None), None); + assert_eq!( + bitrate_only_change(&c, VBR_8M, Some(StreamCodec::Hw(HwCodec::H264)), None), + None + ); let c = change(RateControlMode::Cqp, 28, None, None); - assert_eq!(bitrate_only_change(&c, VBR_8M, HwCodec::H264, None), None); + assert_eq!( + bitrate_only_change(&c, VBR_8M, Some(StreamCodec::Hw(HwCodec::H264)), None), + None + ); } } @@ -3957,3 +4545,347 @@ mod stream_colour_tests { assert!(stream_colour(&declared, &resolved).is_none()); } } + +#[cfg(test)] +mod pyrowave_tests { + use super::*; + use nesprotocol::pyrowave::{CHROMA_420, CHROMA_444}; + + fn request(value: u32, packet_size: Option) -> nesprotocol::EncodeSettings { + nesprotocol::EncodeSettings { + codec: CODEC_PYROWAVE, + rate_control: nesprotocol::RC_CBR, + value, + bit_depth: Some(nesprotocol::DEPTH_8), + chroma: Some(CHROMA_444), + packet_size, + } + } + + #[test] + fn a_request_from_the_hub_is_taken() { + assert_eq!( + pyro_request(&request(300_000, Some(1440)), true), + Ok(PyroSettings { + kbps: 300_000, + chroma: CHROMA_444, + packet_size: 1440, + }) + ); + } + + /// The hub fills in the rate and the packet size. A request without them + /// did not come through one, and encoding it at a guessed size would + /// produce packets the path may not carry. + #[test] + fn a_request_the_hub_did_not_complete_is_refused() { + assert!(pyro_request(&request(0, Some(1440)), true).is_err()); + assert!(pyro_request(&request(300_000, None), true).is_err()); + assert!(pyro_request(&request(300_000, Some(63)), true).is_err()); + let mut cqp = request(300_000, Some(1440)); + cqp.rate_control = nesprotocol::RC_CQP; + assert!(pyro_request(&cqp, true).is_err()); + let mut odd = request(300_000, Some(1440)); + odd.chroma = Some(7); + assert!(pyro_request(&odd, true).is_err()); + } + + /// The fallback device has no nespyro context: PyroWave runs on the + /// game's device or not at all. + #[test] + fn no_pyrowave_on_a_device_of_its_own() { + assert!(pyro_request(&request(300_000, Some(1440)), false).is_err()); + } + + #[test] + fn a_missing_chroma_means_420() { + let mut r = request(300_000, Some(1440)); + r.chroma = None; + assert_eq!(pyro_request(&r, true).unwrap().chroma, CHROMA_420); + } + + #[test] + fn caps_name_pyrowave_only_where_it_runs() { + let caps = host_caps(None, true); + assert!(caps.supports(CODEC_PYROWAVE, nesprotocol::DEPTH_8)); + assert!(caps.supports(CODEC_PYROWAVE, nesprotocol::DEPTH_10)); + assert!(!host_caps(None, false).supports_codec(CODEC_PYROWAVE)); + // And negotiation still never lands on it. + assert_eq!(caps.best(caps), None); + } + + fn key(colorspace: u32, input_fmt: InputFormat) -> PyroKey { + PyroKey { + width: 1920, + height: 1080, + input_fmt, + colorspace, + ten_bit: false, + chroma: CHROMA_420, + packet_size: 1188, + } + } + + #[test] + fn an_sdr_surface_is_bt709_at_eight_bits() { + let c = pyro_config( + key(VK_COLOR_SPACE_SRGB_NONLINEAR_KHR, InputFormat::BGRA), + 200_000, + 60, + ) + .unwrap(); + assert_eq!(c.colour, nespyro::ColourDescription::bt709()); + assert_eq!(c.source, nespyro::Source::Srgb); + assert_eq!(c.depth, nespyro::Depth::Eight); + assert_eq!(c.chroma, nespyro::Chroma::Yuv420); + assert_eq!(c.target_bitrate, 200_000_000); + assert_eq!(c.packet_size, 1188); + } + + #[test] + fn ten_bits_asked_for_sdr_is_sixteen_bit_planes() { + let mut k = key(VK_COLOR_SPACE_SRGB_NONLINEAR_KHR, InputFormat::BGRA); + k.ten_bit = true; + assert_eq!( + pyro_config(k, 1, 60).unwrap().depth, + nespyro::Depth::Sixteen + ); + } + + #[test] + fn hdr10_is_carried_as_it_is() { + let c = pyro_config( + key(VK_COLOR_SPACE_HDR10_ST2084_EXT, InputFormat::ABGR2101010), + 1, + 60, + ) + .unwrap(); + assert_eq!(c.colour, nespyro::ColourDescription::bt2020_pq()); + assert_eq!(c.source, nespyro::Source::Bt2020Pq); + assert_eq!(c.depth, nespyro::Depth::Sixteen, "HDR at eight bits bands"); + } + + /// The same rule the hardware path follows, through the same function: a + /// float surface declared PQ is scRGB, not PQ. + #[test] + fn a_float_surface_is_scrgb_whatever_it_declares() { + let c = pyro_config( + key(VK_COLOR_SPACE_HDR10_ST2084_EXT, InputFormat::RGBA16F), + 1, + 60, + ) + .unwrap(); + assert_eq!(c.source, nespyro::Source::Bt709Linear); + assert_eq!(c.colour, nespyro::ColourDescription::bt2020_pq()); + // scRGB's white is 80 nits. At the 203 other linear sources use, + // Cyberpunk's scRGB mode came out far too bright on PyroWave alone. + assert_eq!(c.reference_white_nits, Some(80.0)); + } + + #[test] + fn a_444_request_is_444() { + let mut k = key(VK_COLOR_SPACE_SRGB_NONLINEAR_KHR, InputFormat::BGRA); + k.chroma = CHROMA_444; + assert_eq!( + pyro_config(k, 1, 60).unwrap().chroma, + nespyro::Chroma::Yuv444 + ); + } + + fn settings() -> PyroSettings { + PyroSettings { + kbps: 300_000, + chroma: CHROMA_420, + packet_size: 1188, + } + } + + fn pyro_change(kbps: u32, pyro: Option) -> EncodeSettingsChange { + EncodeSettingsChange { + codec: Some(StreamCodec::PyroWave), + rate_control_mode: Some(RateControlMode::Cbr), + value: kbps, + bit_depth: None, + pyro, + } + } + + #[test] + fn a_new_pyrowave_rate_is_not_a_rebuild() { + let to = PyroSettings { + kbps: 150_000, + ..settings() + }; + assert_eq!( + pyro_bitrate_only(&pyro_change(150_000, Some(to)), Some(settings()), None), + Some(150_000) + ); + } + + /// A packet size or a chroma format is the shape of the stream, and only + /// a new encoder makes that. + #[test] + fn a_new_packet_size_or_chroma_is_a_rebuild() { + let bigger = PyroSettings { + packet_size: 1400, + ..settings() + }; + assert_eq!( + pyro_bitrate_only(&pyro_change(300_000, Some(bigger)), Some(settings()), None), + None + ); + let full = PyroSettings { + chroma: CHROMA_444, + ..settings() + }; + assert_eq!( + pyro_bitrate_only(&pyro_change(300_000, Some(full)), Some(settings()), None), + None + ); + } + + #[test] + fn a_switch_away_is_not_a_rate_change() { + let mut c = pyro_change(8_000, None); + c.codec = Some(StreamCodec::Hw(HwCodec::AV1)); + assert_eq!(pyro_bitrate_only(&c, Some(settings()), None), None); + // And there is nothing to retune before PyroWave has run at all. + assert_eq!(pyro_bitrate_only(&pyro_change(1, None), None, None), None); + } + + #[test] + fn a_different_pyrowave_request_is_never_nothing() { + let other = PyroSettings { + packet_size: 1400, + ..settings() + }; + let mut c = pyro_change(300_000, Some(other)); + c.rate_control_mode = None; + assert!(!changes_nothing( + &c, + Some(StreamCodec::PyroWave), + None, + Some(settings()) + )); + c.pyro = Some(settings()); + assert!(changes_nothing( + &c, + Some(StreamCodec::PyroWave), + None, + Some(settings()) + )); + } + + /// What the encoder hands over is what the hub puts back together: + /// every packet, in order, the critical count and the frame's own size. + #[test] + fn a_pyrowave_frame_reaches_the_hub_whole() { + let mut data = Vec::new(); + let mut packets = Vec::new(); + for i in 0..3000usize { + let start = data.len(); + data.extend((0..1100).map(|k| (i + k) as u8)); + packets.push(start..data.len()); + } + let frame = nespyro::EncodedFrame { + data, + packets, + critical_packets: 4, + index: 41, + sequence: 1, + stats: Default::default(), + }; + let messages = pyro_messages(&frame, 1234, 2560, 1440).expect("chunks"); + assert!(messages.len() > 1, "a 3 MB frame in one message"); + + let mut asm = nesprotocol::pyrowave::FrameAssembler::new(); + let mut got = None; + for m in &messages { + let ipc = nesprotocol::decode_ipc_frame(m).expect("an IPC frame"); + got = asm.push(&ipc, 0).expect("assembles").or(got); + } + let got = got.expect("the hub completed the frame"); + assert_eq!(got.frame, 41); + assert_eq!(got.critical, 4); + assert_eq!( + (got.meta.ts_ms, got.meta.width, got.meta.height), + (1234, 2560, 1440) + ); + assert_eq!(got.packets.len(), frame.packets.len()); + for (i, r) in frame.packets.iter().enumerate() { + assert_eq!(got.packet(i), &frame.data[r.clone()], "packet {i}"); + } + } + + /// A slot held by a frame the encoder has not finished with. + struct Fake { + done: std::rc::Rc>, + released: std::rc::Rc>>, + id: u32, + } + + impl Hold for Fake { + fn done(&self) -> bool { + self.done.get() + } + fn release(self) { + self.released.borrow_mut().push(self.id); + } + } + + /// Two PyroWave frames in flight hold two slots, and each goes back when + /// its own frame is done -- never sooner, which is the blit racing a read. + #[test] + fn held_slots_go_back_in_order_once_done() { + let released = std::rc::Rc::new(std::cell::RefCell::new(Vec::new())); + let flags: Vec<_> = (0..2) + .map(|_| std::rc::Rc::new(std::cell::Cell::new(false))) + .collect(); + let mut held: std::collections::VecDeque = (0..2) + .map(|i| Fake { + done: flags[i].clone(), + released: released.clone(), + id: i as u32, + }) + .collect(); + + release_done(&mut held); + assert!( + released.borrow().is_empty(), + "a slot went back while in use" + ); + + flags[0].set(true); + release_done(&mut held); + assert_eq!(*released.borrow(), [0]); + assert_eq!(held.len(), 1); + + flags[1].set(true); + release_done(&mut held); + assert_eq!(*released.borrow(), [0, 1]); + assert!(held.is_empty()); + } + + /// A later point on one timeline is never reached before an earlier one, + /// so a done slot behind one still in use waits with it. + #[test] + fn a_done_slot_behind_a_busy_one_waits() { + let released = std::rc::Rc::new(std::cell::RefCell::new(Vec::new())); + let busy = std::rc::Rc::new(std::cell::Cell::new(false)); + let done = std::rc::Rc::new(std::cell::Cell::new(true)); + let mut held = std::collections::VecDeque::from([ + Fake { + done: busy, + released: released.clone(), + id: 0, + }, + Fake { + done, + released: released.clone(), + id: 1, + }, + ]); + release_done(&mut held); + assert!(released.borrow().is_empty()); + } +} diff --git a/apps/nescapture/src/instance.rs b/apps/nescapture/src/instance.rs index b04b1b31..ae67273f 100644 --- a/apps/nescapture/src/instance.rs +++ b/apps/nescapture/src/instance.rs @@ -41,10 +41,9 @@ pub unsafe extern "system" fn vkCreateInstance( let next_create: PFN_vkCreateInstance = unsafe { load_instance_fn(next_gipa, vk::Instance::null(), b"vkCreateInstance\0") }; - // The encoder needs Vulkan 1.1 on the game's instance to run on the game's - // device. An application asking for 1.0 is raised to 1.1 where the loader - // has it: 1.1 only adds to 1.0, so nothing the application can do behaves - // any differently, and without it the encoder needs a device of its own. + // The encoders run on the game's device, so the game's instance has to be + // new enough for them: 1.1 for pixelforge, 1.3 for nespyro. See + // `raised_version`. let asked = unsafe { let app = (*p_create_info).p_application_info; if app.is_null() { @@ -57,7 +56,7 @@ pub unsafe extern "system" fn vkCreateInstance( let mut raised_ci; let mut create_info = p_create_info; let mut api_version = asked; - if asked < vk::API_VERSION_1_1 && unsafe { loader_version(next_gipa) } >= vk::API_VERSION_1_1 { + if let Some(raised) = raised_version(asked, unsafe { loader_version(next_gipa) }) { raised_app = unsafe { let app = (*p_create_info).p_application_info; if app.is_null() { @@ -66,13 +65,17 @@ pub unsafe extern "system" fn vkCreateInstance( *app } }; - raised_app.api_version = vk::API_VERSION_1_1; + raised_app.api_version = raised; raised_ci = unsafe { *p_create_info }; raised_ci.p_application_info = &raised_app; create_info = &raised_ci; - api_version = vk::API_VERSION_1_1; + api_version = raised; log::info!( - "instance asked for Vulkan 1.0; created as 1.1 so the encoder can share its devices" + "instance asked for Vulkan {}.{}; created as {}.{} so the encoders can share its devices", + vk::api_version_major(asked), + vk::api_version_minor(asked), + vk::api_version_major(raised), + vk::api_version_minor(raised), ); } @@ -115,6 +118,27 @@ pub unsafe extern "system" fn vkCreateInstance( vk::Result::SUCCESS } +/// The version to create the game's instance at instead of the one it asked +/// for, or `None` to leave it alone. +/// +/// As high as 1.3, which nespyro needs, and as the loader allows. A newer +/// version only adds to an older one: a game cannot depend on behaviour it +/// never asked for, and the only side of a session that has to care about old +/// Vulkan is the client, which this layer is not. A loader below 1.1 is left +/// alone, since then neither encoder can share the device anyway. The patch is +/// dropped, as the version names a major.minor and its patch level is the +/// loader's business. +fn raised_version(asked: u32, loader: u32) -> Option { + let target = vk::API_VERSION_1_3.min(loader); + let target = vk::make_api_version( + 0, + vk::api_version_major(target), + vk::api_version_minor(target), + 0, + ); + (target >= vk::API_VERSION_1_1 && asked < target).then_some(target) +} + /// The highest instance version the loader below supports; 1.0 when it cannot /// say, since vkEnumerateInstanceVersion is itself a 1.1 addition. unsafe fn loader_version(next_gipa: crate::dispatch::PFN_vkGetInstanceProcAddr) -> u32 { @@ -145,3 +169,44 @@ pub unsafe extern "system" fn vkDestroyInstance( unsafe { (state.destroy_instance)(instance, p_allocator) }; } } + +#[cfg(test)] +mod tests { + use super::raised_version; + use ash::vk; + + const LOADER_1_4: u32 = vk::make_api_version(0, 1, 4, 321); + + #[test] + fn an_old_instance_is_raised_to_1_3() { + for asked in [ + 0, + vk::API_VERSION_1_0, + vk::API_VERSION_1_1, + vk::API_VERSION_1_2, + ] { + assert_eq!(raised_version(asked, LOADER_1_4), Some(vk::API_VERSION_1_3)); + } + } + + /// Raising is only ever upward: a game asking for more keeps it. + #[test] + fn a_new_enough_instance_is_left_alone() { + assert_eq!(raised_version(vk::API_VERSION_1_3, LOADER_1_4), None); + assert_eq!(raised_version(LOADER_1_4, LOADER_1_4), None); + } + + #[test] + fn an_older_loader_gets_what_it_has() { + let loader = vk::make_api_version(0, 1, 2, 198); + assert_eq!( + raised_version(vk::API_VERSION_1_0, loader), + Some(vk::API_VERSION_1_2) + ); + assert_eq!(raised_version(vk::API_VERSION_1_2, loader), None); + assert_eq!( + raised_version(vk::API_VERSION_1_0, vk::API_VERSION_1_0), + None + ); + } +} diff --git a/apps/nescapture/src/present.rs b/apps/nescapture/src/present.rs index c615d4e6..c5650223 100644 --- a/apps/nescapture/src/present.rs +++ b/apps/nescapture/src/present.rs @@ -326,14 +326,24 @@ fn encoder_ready(ds: &crate::state::DeviceState, ds_key: usize, width: u32, heig // On the game's own device where it was created for that, and // on a device of the encoder's own where it was not, or where // adopting it fails. - let shared = ds.shared.as_ref().and_then(|s| match s.video_context() { - Ok(ctx) => Some(ctx), - Err(e) => { - log::warn!( - "could not encode on the game's device ({e}); using a device of its own" - ); - None + // Each encoder the device was set up for is adopted on its own; a + // failure in one costs only that one. Only when neither is left + // does the encoder move to a device of its own, which can then + // host Vulkan Video and nothing else. + let shared = ds.shared.as_ref().and_then(|s| { + let video = s.video_context().and_then(|r| { + r.map_err(|e| log::warn!("Vulkan Video on the game's device failed ({e})")) + .ok() + }); + let pyro = s.pyro_context().and_then(|r| { + r.map_err(|e| log::warn!("PyroWave on the game's device failed ({e})")) + .ok() + }); + if video.is_none() && pyro.is_none() { + log::warn!("no encoder on the game's device; using a device of its own"); + return None; } + Some(crate::encode::SharedEncoders { video, pyro }) }); let on_shared = shared.is_some(); match PipelineHandle::new(cfg, shared) { diff --git a/apps/nescapture/src/shared.rs b/apps/nescapture/src/shared.rs index bd80719f..29ba3eee 100644 --- a/apps/nescapture/src/shared.rs +++ b/apps/nescapture/src/shared.rs @@ -6,6 +6,10 @@ // because two devices share no timeline. Here the game's device is created // with what the encoder needs, and the encoder is handed that device instead. // +// Two encoders can live there: pixelforge's Vulkan Video encoder and +// nespyro's PyroWave. Each is checked on its own, and the device gets the +// additions of whichever it can host — one, both, or, if neither, none. +// // Three things are added to the game's vkCreateDevice, and none of them // changes what the game gets: // @@ -33,9 +37,13 @@ use std::collections::BTreeMap; /// the device's queue create infos. #[derive(Debug, Clone, PartialEq, Eq)] pub struct QueuePlan { - pub encode: DeviceQueue, + /// Vulkan Video encode, when pixelforge is hosted. + pub encode: Option, + /// pixelforge's converter and all of nespyro. Both encoders submit from the + /// one encoder thread, so they share it freely. pub compute: DeviceQueue, - pub transfer: DeviceQueue, + /// pixelforge's copies, when pixelforge is hosted. + pub transfer: Option, /// How many queues to create in each family, where that is more than the /// game asked for. pub counts: BTreeMap, @@ -71,7 +79,32 @@ fn rank(flags: vk::QueueFlags) -> u32 { r } -/// Pick a queue for every role. +/// Which roles a plan has to fill. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Wants { + /// An encode queue, and the family the encoder would like it from. + pub encode: Option>, + /// A transfer queue. + pub transfer: bool, +} + +impl Wants { + /// What pixelforge needs: encode, compute and transfer. + pub fn video(preferred_encode: Option) -> Self { + Self { + encode: Some(preferred_encode), + transfer: true, + } + } + + /// What nespyro alone needs: compute, which every plan has. + pub const PYRO_ONLY: Self = Self { + encode: None, + transfer: false, + }; +} + +/// Pick a queue for every role `wants` names, and a compute queue always. /// /// `families` is the device's queue families, `game` how many queues the game /// asked for in each (and with which flags). The encoder takes at most one @@ -83,7 +116,7 @@ fn rank(flags: vk::QueueFlags) -> u32 { pub fn plan_queues( families: &[vk::QueueFamilyProperties], game: &BTreeMap, - preferred_encode: Option, + wants: Wants, internally_synchronized_queues: bool, ) -> Option { let mut counts: BTreeMap = BTreeMap::new(); @@ -146,13 +179,16 @@ pub fn plan_queues( None }; - let encode = pick( - vk::QueueFlags::VIDEO_ENCODE_KHR, - preferred_encode, - &mut counts, - &mut ours, - &mut shared, - )?; + let encode = match wants.encode { + Some(preferred) => Some(pick( + vk::QueueFlags::VIDEO_ENCODE_KHR, + preferred, + &mut counts, + &mut ours, + &mut shared, + )?), + None => None, + }; let compute = pick( vk::QueueFlags::COMPUTE, None, @@ -160,13 +196,17 @@ pub fn plan_queues( &mut ours, &mut shared, )?; - let transfer = pick( - vk::QueueFlags::TRANSFER, - Some(compute.family), - &mut counts, - &mut ours, - &mut shared, - )?; + let transfer = if wants.transfer { + Some(pick( + vk::QueueFlags::TRANSFER, + Some(compute.family), + &mut counts, + &mut ours, + &mut shared, + )?) + } else { + None + }; Some(QueuePlan { encode, @@ -191,6 +231,52 @@ pub enum Feature { VideoEncodeIntraRefresh, VideoEncodeQuantizationMap, InternallySynchronizedQueues, + // nespyro's. The first two are Vulkan 1.0 core features, which live in + // `VkPhysicalDeviceFeatures` rather than in a struct of their own. + ShaderInt16, + ShaderStorageImageWriteWithoutFormat, + StorageBuffer16BitAccess, + StorageBuffer8BitAccess, + ShaderInt8, + ShaderFloat16, + BufferDeviceAddress, + SubgroupSizeControl, + ComputeFullSubgroups, +} + +impl Feature { + /// Whether this is a Vulkan 1.0 feature, set in `VkPhysicalDeviceFeatures` + /// rather than in a pNext struct of its own. + fn is_core(self) -> bool { + matches!( + self, + Self::ShaderInt16 | Self::ShaderStorageImageWriteWithoutFormat + ) + } +} + +/// The features nespyro asks for. +pub fn pyro_features(f: &nespyro::DeviceFeatures) -> Vec { + use Feature as F; + [ + (f.shader_int16, F::ShaderInt16), + ( + f.shader_storage_image_write_without_format, + F::ShaderStorageImageWriteWithoutFormat, + ), + (f.storage_buffer_16bit_access, F::StorageBuffer16BitAccess), + (f.storage_buffer_8bit_access, F::StorageBuffer8BitAccess), + (f.shader_int8, F::ShaderInt8), + (f.shader_float16, F::ShaderFloat16), + (f.timeline_semaphore, F::TimelineSemaphore), + (f.buffer_device_address, F::BufferDeviceAddress), + (f.synchronization2, F::Synchronization2), + (f.subgroup_size_control, F::SubgroupSizeControl), + (f.compute_full_subgroups, F::ComputeFullSubgroups), + ] + .into_iter() + .filter_map(|(on, feature)| on.then_some(feature)) + .collect() } /// The features `f` asks for, plus the one queue sharing needs. @@ -307,6 +393,79 @@ unsafe fn field( vk::PhysicalDeviceInternallySynchronizedQueuesFeaturesKHR, internally_synchronized_queues ), + (S::PHYSICAL_DEVICE_FEATURES_2, F::ShaderInt16) => Some(unsafe { + &raw mut (*(p as *mut vk::PhysicalDeviceFeatures2)) + .features + .shader_int16 + }), + (S::PHYSICAL_DEVICE_FEATURES_2, F::ShaderStorageImageWriteWithoutFormat) => Some(unsafe { + &raw mut (*(p as *mut vk::PhysicalDeviceFeatures2)) + .features + .shader_storage_image_write_without_format + }), + (S::PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, F::StorageBuffer16BitAccess) => { + at!( + vk::PhysicalDeviceVulkan11Features, + storage_buffer16_bit_access + ) + } + (S::PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES, F::StorageBuffer16BitAccess) => { + at!( + vk::PhysicalDevice16BitStorageFeatures, + storage_buffer16_bit_access + ) + } + (S::PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, F::StorageBuffer8BitAccess) => { + at!( + vk::PhysicalDeviceVulkan12Features, + storage_buffer8_bit_access + ) + } + (S::PHYSICAL_DEVICE_8BIT_STORAGE_FEATURES, F::StorageBuffer8BitAccess) => { + at!( + vk::PhysicalDevice8BitStorageFeatures, + storage_buffer8_bit_access + ) + } + (S::PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, F::ShaderInt8) => { + at!(vk::PhysicalDeviceVulkan12Features, shader_int8) + } + (S::PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES, F::ShaderInt8) => { + at!(vk::PhysicalDeviceShaderFloat16Int8Features, shader_int8) + } + (S::PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, F::ShaderFloat16) => { + at!(vk::PhysicalDeviceVulkan12Features, shader_float16) + } + (S::PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES, F::ShaderFloat16) => { + at!(vk::PhysicalDeviceShaderFloat16Int8Features, shader_float16) + } + (S::PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, F::BufferDeviceAddress) => { + at!(vk::PhysicalDeviceVulkan12Features, buffer_device_address) + } + (S::PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES, F::BufferDeviceAddress) => { + at!( + vk::PhysicalDeviceBufferDeviceAddressFeatures, + buffer_device_address + ) + } + (S::PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, F::SubgroupSizeControl) => { + at!(vk::PhysicalDeviceVulkan13Features, subgroup_size_control) + } + (S::PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES, F::SubgroupSizeControl) => { + at!( + vk::PhysicalDeviceSubgroupSizeControlFeatures, + subgroup_size_control + ) + } + (S::PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, F::ComputeFullSubgroups) => { + at!(vk::PhysicalDeviceVulkan13Features, compute_full_subgroups) + } + (S::PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES, F::ComputeFullSubgroups) => { + at!( + vk::PhysicalDeviceSubgroupSizeControlFeatures, + compute_full_subgroups + ) + } _ => None, } } @@ -338,6 +497,11 @@ enum OwnedFeature { IntraRefresh(vk::PhysicalDeviceVideoEncodeIntraRefreshFeaturesKHR<'static>), QpMap(vk::PhysicalDeviceVideoEncodeQuantizationMapFeaturesKHR<'static>), SharedQueues(vk::PhysicalDeviceInternallySynchronizedQueuesFeaturesKHR<'static>), + Storage16(vk::PhysicalDevice16BitStorageFeatures<'static>), + Storage8(vk::PhysicalDevice8BitStorageFeatures<'static>), + Float16Int8(vk::PhysicalDeviceShaderFloat16Int8Features<'static>), + Bda(vk::PhysicalDeviceBufferDeviceAddressFeatures<'static>), + SubgroupSize(vk::PhysicalDeviceSubgroupSizeControlFeatures<'static>), } impl OwnedFeature { @@ -377,6 +541,33 @@ impl OwnedFeature { vk::PhysicalDeviceInternallySynchronizedQueuesFeaturesKHR::default() .internally_synchronized_queues(true), ), + F::StorageBuffer16BitAccess => Self::Storage16( + vk::PhysicalDevice16BitStorageFeatures::default().storage_buffer16_bit_access(true), + ), + F::StorageBuffer8BitAccess => Self::Storage8( + vk::PhysicalDevice8BitStorageFeatures::default().storage_buffer8_bit_access(true), + ), + F::ShaderInt8 => Self::Float16Int8( + vk::PhysicalDeviceShaderFloat16Int8Features::default().shader_int8(true), + ), + F::ShaderFloat16 => Self::Float16Int8( + vk::PhysicalDeviceShaderFloat16Int8Features::default().shader_float16(true), + ), + F::BufferDeviceAddress => Self::Bda( + vk::PhysicalDeviceBufferDeviceAddressFeatures::default() + .buffer_device_address(true), + ), + F::SubgroupSizeControl => Self::SubgroupSize( + vk::PhysicalDeviceSubgroupSizeControlFeatures::default() + .subgroup_size_control(true), + ), + F::ComputeFullSubgroups => Self::SubgroupSize( + vk::PhysicalDeviceSubgroupSizeControlFeatures::default() + .compute_full_subgroups(true), + ), + F::ShaderInt16 | F::ShaderStorageImageWriteWithoutFormat => { + unreachable!("core features have no struct of their own") + } } } @@ -391,6 +582,11 @@ impl OwnedFeature { Self::IntraRefresh(s) => std::ptr::from_mut(s).cast(), Self::QpMap(s) => std::ptr::from_mut(s).cast(), Self::SharedQueues(s) => std::ptr::from_mut(s).cast(), + Self::Storage16(s) => std::ptr::from_mut(s).cast(), + Self::Storage8(s) => std::ptr::from_mut(s).cast(), + Self::Float16Int8(s) => std::ptr::from_mut(s).cast(), + Self::Bda(s) => std::ptr::from_mut(s).cast(), + Self::SubgroupSize(s) => std::ptr::from_mut(s).cast(), }; p.cast() } @@ -398,8 +594,13 @@ impl OwnedFeature { impl FeaturePatch { /// Turn on every feature in `wanted` for a device whose create info has - /// `chain` as its pNext. Returns the patch and the pNext the create info - /// should carry instead. + /// `chain` as its pNext. Returns the patch, the pNext the create info + /// should carry instead, and the Vulkan 1.0 features the chain had no + /// `VkPhysicalDeviceFeatures2` to hold — those belong in + /// `pEnabledFeatures`, see [`core_features`]. + /// + /// Structs of ours are searched like the game's, so two features sharing + /// one struct type land in one struct: a chain may hold each type once. /// /// # Safety /// @@ -408,16 +609,18 @@ impl FeaturePatch { pub unsafe fn apply( chain: *const std::ffi::c_void, wanted: &[Feature], - ) -> (Self, *const std::ffi::c_void) { + ) -> (Self, *const std::ffi::c_void, Vec) { let mut patch = Self { restores: Vec::new(), owned: Vec::new(), }; - let mut missing: Vec = Vec::new(); + let mut core_missing: Vec = Vec::new(); + // Ours sit in front of the game's chain, in the order they were added. + let mut head = chain; for &feature in wanted { let mut found = false; - let mut p = chain as *mut vk::BaseOutStructure<'_>; + let mut p = head as *mut vk::BaseOutStructure<'_>; while !p.is_null() { let s_type = unsafe { (*p).s_type }; if let Some(bit) = unsafe { field(s_type, p, feature) } { @@ -427,20 +630,23 @@ impl FeaturePatch { } p = unsafe { (*p).p_next }; } - if !found { - missing.push(feature); + if found { + continue; + } + if feature.is_core() { + core_missing.push(feature); + continue; } - } - - let mut head = chain; - for feature in missing.into_iter().rev() { let mut owned = Box::new(OwnedFeature::new(feature)); let base = owned.base(); - unsafe { (*base).p_next = head as *mut _ }; - head = base as *const _; + unsafe { (*base).p_next = chain as *mut _ }; + match patch.owned.last_mut() { + Some(last) => unsafe { (*last.base()).p_next = base }, + None => head = base as *const _, + } patch.owned.push(owned); } - (patch, head) + (patch, head, core_missing) } /// Put the application's chain back as it was. @@ -455,6 +661,28 @@ impl FeaturePatch { } } +/// `pEnabledFeatures` with `missing` turned on: the game's own, copied, or a +/// fresh set when it gave none. +/// +/// Only used when the chain holds no `VkPhysicalDeviceFeatures2`, since the +/// two may not be given together, and with one the bits went there instead. +pub fn core_features( + game: Option<&vk::PhysicalDeviceFeatures>, + missing: &[Feature], +) -> vk::PhysicalDeviceFeatures { + let mut out = game.copied().unwrap_or_default(); + for feature in missing { + match feature { + Feature::ShaderInt16 => out.shader_int16 = vk::TRUE, + Feature::ShaderStorageImageWriteWithoutFormat => { + out.shader_storage_image_write_without_format = vk::TRUE + } + other => debug_assert!(!other.is_core(), "{other:?} has no place here"), + } + } + out +} + // ── Extensions ──────────────────────────────────────────────────────────────── /// `game` with every name in `add` it does not already have, as pointers @@ -475,28 +703,51 @@ pub fn merged_extensions( out } -/// Everything the device needs beyond what the game asked for: the pixelforge -/// requirements, and the queue plan that satisfies them. +/// Everything the device needs beyond what the game asked for: each hosted +/// encoder's requirements, and the queue plan that satisfies them. pub struct Additions { - pub requirements: DeviceRequirements, + /// pixelforge's, when the device can encode Vulkan Video. + pub video: Option, + /// nespyro's, when the device can run PyroWave. + pub pyro: Option, pub queues: QueuePlan, } impl Additions { - /// The extensions to add, including the one queue sharing needs. + /// The extensions to add, including the one queue sharing needs. Each + /// named once, however many encoders want it. pub fn extensions(&self) -> Vec<&'static std::ffi::CStr> { - let mut names = self.requirements.extensions.clone(); + let mut names: Vec<&'static std::ffi::CStr> = Vec::new(); + let mut add = |name: &'static std::ffi::CStr| { + if !names.contains(&name) { + names.push(name); + } + }; + for name in self.video.iter().flat_map(|r| r.extensions.iter()) { + add(name); + } + for name in self.pyro.iter().flat_map(|r| r.extensions.iter()) { + add(name); + } if !self.queues.internally_synchronized.is_empty() { - names.push(ash::khr::internally_synchronized_queues::NAME); + add(ash::khr::internally_synchronized_queues::NAME); } names } + /// Each named once, however many encoders want it. pub fn features(&self) -> Vec { - wanted_features( - &self.requirements.features, - !self.queues.internally_synchronized.is_empty(), - ) + let share = !self.queues.internally_synchronized.is_empty(); + let mut out = match &self.video { + Some(r) => wanted_features(&r.features, share), + None => wanted_features(&DeviceFeatures::default(), share), + }; + for f in self.pyro.iter().flat_map(|r| pyro_features(&r.features)) { + if !out.contains(&f) { + out.push(f); + } + } + out } } @@ -630,6 +881,8 @@ pub struct PreparedDevice { extensions: Vec<*const std::ffi::c_char>, patch: FeaturePatch, p_next: *const std::ffi::c_void, + /// The 1.0 features, when they had to go in `pEnabledFeatures`. + enabled_features: Option>, } /// Work out what to add to the game's device so the encoder can run on it. @@ -664,17 +917,44 @@ pub unsafe fn prepare( return None; } let (entry, instance) = instance_view(istate, next_gdpa); - let requirements = match pixelforge::VideoContextBuilder::new().encode_device_requirements( + let video = match pixelforge::VideoContextBuilder::new().encode_device_requirements( &entry, &instance, physical_device, ) { - Ok(r) => r, + Ok(r) => Some(r), Err(e) => { - log::info!("the game's device cannot host the encoder ({e})"); - return None; + log::info!("the game's device cannot host Vulkan Video encode ({e})"); + None } }; + // nespyro uses Vulkan 1.3 commands, which a device only has when its + // instance asked for 1.3 as well as the GPU having it. + let pyro = if istate.api_version < vk::API_VERSION_1_3 { + log::info!("instance is below Vulkan 1.3; no PyroWave on the game's device"); + None + } else { + match nespyro::DeviceRequirements::query(&instance, physical_device, nespyro::Roles::ENCODE) + { + Ok(r) => Some(r), + Err(e) => { + log::info!("the game's device cannot host PyroWave ({e})"); + None + } + } + }; + if video.is_none() && pyro.is_none() { + log::info!("neither encoder fits the game's device; encoding on a device of its own"); + return None; + } + let internally_synchronized_queues = match &video { + Some(r) => r.internally_synchronized_queues, + None => supports_internally_synchronized_queues(&instance, physical_device), + }; + let wants = match &video { + Some(r) => Wants::video(r.queues.encode), + None => Wants::PYRO_ONLY, + }; let families = unsafe { instance.get_physical_device_queue_family_properties(physical_device) }; let game_infos: &[vk::DeviceQueueCreateInfo<'_>] = if ci.queue_create_info_count == 0 @@ -691,12 +971,7 @@ pub unsafe fn prepare( .map(|q| (q.queue_family_index, (q.queue_count, q.flags))) .collect(); - let Some(queues) = plan_queues( - &families, - &game, - requirements.queues.encode, - requirements.internally_synchronized_queues, - ) else { + let Some(queues) = plan_queues(&families, &game, wants, internally_synchronized_queues) else { log::info!( "no queue for the encoder that the game does not submit to; \ encoding on a device of its own" @@ -704,7 +979,12 @@ pub unsafe fn prepare( return None; }; log::info!( - "encoder on the game's device: encode queue {:?}, compute queue {:?}, transfer queue {:?}{}", + "encoders on the game's device ({}): encode queue {:?}, compute queue {:?}, transfer queue {:?}{}", + match (video.is_some(), pyro.is_some()) { + (true, true) => "Vulkan Video and PyroWave", + (true, false) => "Vulkan Video", + _ => "PyroWave", + }, queues.encode, queues.compute, queues.transfer, @@ -759,11 +1039,17 @@ pub unsafe fn prepare( } let additions = Additions { - requirements, + video, + pyro, queues, }; let extensions = merged_extensions(extensions, &additions.extensions()); - let (patch, p_next) = unsafe { FeaturePatch::apply(ci.p_next, &additions.features()) }; + let (patch, p_next, core_missing) = + unsafe { FeaturePatch::apply(ci.p_next, &additions.features()) }; + let enabled_features = (!core_missing.is_empty()).then(|| { + let game = unsafe { ci.p_enabled_features.as_ref() }; + Box::new(core_features(game, &core_missing)) + }); Some(PreparedDevice { additions, @@ -774,9 +1060,33 @@ pub unsafe fn prepare( extensions, patch, p_next, + enabled_features, }) } +/// Whether `physical_device` supports `VK_KHR_internally_synchronized_queues`, +/// extension and feature both. pixelforge answers this when it is hosted; this +/// is the same question for a device hosting only nespyro. +fn supports_internally_synchronized_queues( + instance: &ash::Instance, + physical_device: vk::PhysicalDevice, +) -> bool { + let Ok(exts) = (unsafe { instance.enumerate_device_extension_properties(physical_device) }) + else { + return false; + }; + if !exts.iter().any(|e| { + e.extension_name_as_c_str() + .is_ok_and(|n| n == ash::khr::internally_synchronized_queues::NAME) + }) { + return false; + } + let mut features = vk::PhysicalDeviceInternallySynchronizedQueuesFeaturesKHR::default(); + let mut query = vk::PhysicalDeviceFeatures2::default().push(&mut features); + unsafe { instance.get_physical_device_features2(physical_device, &mut query) }; + features.internally_synchronized_queues != 0 +} + impl PreparedDevice { /// `ci` with the additions in. Valid while `self` is. pub fn create_info<'a>(&'a self, ci: &vk::DeviceCreateInfo<'a>) -> vk::DeviceCreateInfo<'a> { @@ -786,6 +1096,9 @@ impl PreparedDevice { out.p_queue_create_infos = self.queue_infos.as_ptr().cast(); out.enabled_extension_count = self.extensions.len() as u32; out.pp_enabled_extension_names = self.extensions.as_ptr(); + if let Some(features) = &self.enabled_features { + out.p_enabled_features = &raw const **features; + } out } @@ -811,6 +1124,9 @@ pub struct SharedDevice { /// created below Vulkan 1.2 still has: the encoder enabled the extension. timeline: ash::khr::timeline_semaphore::Device, pub queues: QueuePlan, + /// Which encoders the device was created for. + video: bool, + pyro: bool, } impl SharedDevice { @@ -860,6 +1176,8 @@ impl SharedDevice { device, timeline, queues: additions.queues, + video: additions.video.is_some(), + pyro: additions.pyro.is_some(), } } @@ -902,26 +1220,56 @@ impl SharedDevice { } /// A pixelforge context on the game's device, submitting to the queues - /// the plan set aside. - pub fn video_context(&self) -> Result { - pixelforge::VideoContextBuilder::new() - .app_name("nescapture") - .with_encode_queue(self.queues.encode) - .with_compute_queue(self.queues.compute) - .with_transfer_queue(self.queues.transfer) - .build_from_existing_encode( - self.entry.clone(), + /// the plan set aside. `None` when the device was not created for it. + pub fn video_context( + &self, + ) -> Option> { + if !self.video { + return None; + } + let (Some(encode), Some(transfer)) = (self.queues.encode, self.queues.transfer) else { + return None; + }; + Some( + pixelforge::VideoContextBuilder::new() + .app_name("nescapture") + .with_encode_queue(encode) + .with_compute_queue(self.queues.compute) + .with_transfer_queue(transfer) + .build_from_existing_encode( + self.entry.clone(), + self.instance.clone(), + self.physical_device, + self.device.clone(), + ), + ) + } + + /// A nespyro context on the game's device, on the plan's compute queue. + /// `None` when the device was not created for it. + /// + /// No queue lock: the encoder thread is the only submitter to that queue + /// apart from a game sharing it, and a queue shared with the game was + /// created internally synchronized. + pub fn pyro_context(&self) -> Option> { + self.pyro.then(|| { + nespyro::Context::from_existing( self.instance.clone(), self.physical_device, self.device.clone(), + nespyro::DeviceQueue::new(self.queues.compute.family, self.queues.compute.index), + nespyro::QueueLock::none(), + nespyro::Roles::ENCODE, ) + }) } /// The queue families that touch a capture image: the one the game /// presents on, where the blit runs, and the encoder's. pub fn image_families(&self, present_family: u32) -> Vec { let mut out = vec![present_family]; - for f in [self.queues.compute.family, self.queues.transfer.family] { + let transfer = self.queues.transfer.map(|q| q.family); + for f in std::iter::once(self.queues.compute.family).chain(transfer) { if !out.contains(&f) { out.push(f); } @@ -966,19 +1314,25 @@ mod tests { #[test] fn spare_queues_keep_the_encoder_off_the_games_queue() { - let plan = plan_queues(&amd(), &game(&[(0, 1)]), Some(2), false).unwrap(); - assert_eq!(plan.encode, DeviceQueue::new(2, 0)); + let plan = plan_queues(&amd(), &game(&[(0, 1)]), Wants::video(Some(2)), false).unwrap(); + assert_eq!(plan.encode, Some(DeviceQueue::new(2, 0))); // Compute on the compute family, not on the game's graphics queue. assert_eq!(plan.compute, DeviceQueue::new(1, 0)); // And copies share it: one thread submits both. - assert_eq!(plan.transfer, plan.compute); + assert_eq!(plan.transfer, Some(plan.compute)); assert!(plan.internally_synchronized.is_empty()); } #[test] fn a_family_the_game_already_uses_gets_one_more_queue() { // The game took two of the four compute queues for itself. - let plan = plan_queues(&amd(), &game(&[(0, 1), (1, 2)]), Some(2), false).unwrap(); + let plan = plan_queues( + &amd(), + &game(&[(0, 1), (1, 2)]), + Wants::video(Some(2)), + false, + ) + .unwrap(); assert_eq!(plan.compute, DeviceQueue::new(1, 2)); assert_eq!(plan.counts.get(&1), Some(&3)); } @@ -987,11 +1341,14 @@ mod tests { fn with_every_capable_family_full_nothing_is_planned_unless_sharing_is_safe() { // The game took all four compute queues and its one graphics queue. let full = game(&[(0, 1), (1, 4)]); - assert_eq!(plan_queues(&amd(), &full, Some(2), false), None); + assert_eq!( + plan_queues(&amd(), &full, Wants::video(Some(2)), false), + None + ); // Sharing picks the compute family over the graphics one, since the // game renders on the latter. - let plan = plan_queues(&amd(), &full, Some(2), true).unwrap(); + let plan = plan_queues(&amd(), &full, Wants::video(Some(2)), true).unwrap(); assert_eq!(plan.compute, DeviceQueue::new(1, 0)); assert_eq!(plan.internally_synchronized, vec![1]); } @@ -1009,17 +1366,17 @@ mod tests { #[test] fn without_a_spare_the_games_queue_is_shared_where_that_is_safe() { - let plan = plan_queues(&intel(), &game(&[(0, 1)]), Some(1), true).unwrap(); - assert_eq!(plan.encode, DeviceQueue::new(1, 0)); + let plan = plan_queues(&intel(), &game(&[(0, 1)]), Wants::video(Some(1)), true).unwrap(); + assert_eq!(plan.encode, Some(DeviceQueue::new(1, 0))); assert_eq!(plan.compute, DeviceQueue::new(0, 0)); - assert_eq!(plan.transfer, DeviceQueue::new(0, 0)); + assert_eq!(plan.transfer, Some(DeviceQueue::new(0, 0))); assert_eq!(plan.internally_synchronized, vec![0]); } #[test] fn without_a_spare_or_a_safe_share_there_is_no_plan() { assert_eq!( - plan_queues(&intel(), &game(&[(0, 1)]), Some(1), false), + plan_queues(&intel(), &game(&[(0, 1)]), Wants::video(Some(1)), false), None ); } @@ -1028,15 +1385,21 @@ mod tests { fn a_protected_queue_is_never_shared() { let mut g = game(&[(0, 1)]); g.insert(0, (1, vk::DeviceQueueCreateFlags::PROTECTED)); - assert_eq!(plan_queues(&intel(), &g, Some(1), true), None); + assert_eq!(plan_queues(&intel(), &g, Wants::video(Some(1)), true), None); } #[test] fn a_game_that_already_encodes_keeps_its_encode_queue() { let mut families = amd(); families[2].queue_count = 2; - let plan = plan_queues(&families, &game(&[(0, 1), (2, 1)]), Some(2), false).unwrap(); - assert_eq!(plan.encode, DeviceQueue::new(2, 1)); + let plan = plan_queues( + &families, + &game(&[(0, 1), (2, 1)]), + Wants::video(Some(2)), + false, + ) + .unwrap(); + assert_eq!(plan.encode, Some(DeviceQueue::new(2, 1))); } #[test] @@ -1046,7 +1409,7 @@ mod tests { v13.p_next = (&raw mut v12).cast(); let chain = (&raw const v13).cast(); - let (patch, head) = unsafe { + let (patch, head, _) = unsafe { FeaturePatch::apply( chain, &[Feature::Synchronization2, Feature::TimelineSemaphore], @@ -1067,7 +1430,7 @@ mod tests { fn a_feature_with_no_home_gets_a_struct_in_front_of_the_chain() { let v12 = vk::PhysicalDeviceVulkan12Features::default(); let chain = (&raw const v12).cast(); - let (patch, head) = unsafe { FeaturePatch::apply(chain, &[Feature::Synchronization2]) }; + let (patch, head, _) = unsafe { FeaturePatch::apply(chain, &[Feature::Synchronization2]) }; assert_ne!(head, chain); let first = head as *const vk::BaseInStructure<'_>; unsafe { @@ -1087,7 +1450,7 @@ mod tests { Feature::TimelineSemaphore, Feature::VideoEncodeAv1, ]; - let (patch, head) = unsafe { FeaturePatch::apply(std::ptr::null(), &wanted) }; + let (patch, head, _) = unsafe { FeaturePatch::apply(std::ptr::null(), &wanted) }; let mut seen = Vec::new(); let mut p = head as *const vk::BaseInStructure<'_>; while !p.is_null() { @@ -1118,6 +1481,183 @@ mod tests { assert_eq!(merged.len(), 2); } + // ── PyroWave beside, or instead of, Vulkan Video ───────────────────────── + + /// A GPU with no video encode family still gets a compute queue of its + /// own for nespyro, and no encode or transfer queue it would not use. + #[test] + fn pyrowave_alone_plans_only_compute() { + let families = vec![family(GFX, 1), family(COMPUTE, 2)]; + let plan = plan_queues(&families, &game(&[(0, 1)]), Wants::PYRO_ONLY, false).unwrap(); + assert_eq!(plan.compute, DeviceQueue::new(1, 0)); + assert_eq!(plan.encode, None); + assert_eq!(plan.transfer, None); + assert!(plan.internally_synchronized.is_empty()); + } + + /// Without the video family there is nothing to plan an encode queue on, + /// and asking for one must fail rather than quietly skip it. + #[test] + fn vulkan_video_without_an_encode_family_has_no_plan() { + let families = vec![family(GFX, 1), family(COMPUTE, 2)]; + assert_eq!( + plan_queues(&families, &game(&[(0, 1)]), Wants::video(None), true), + None + ); + } + + /// Two features of one struct type must land in one struct: a pNext chain + /// may hold each type once, and int8 and float16 share + /// `VkPhysicalDeviceShaderFloat16Int8Features`. + #[test] + fn features_sharing_a_struct_share_one_of_ours() { + let wanted = [ + Feature::ShaderInt8, + Feature::ShaderFloat16, + Feature::SubgroupSizeControl, + Feature::ComputeFullSubgroups, + ]; + let (patch, head, core) = unsafe { FeaturePatch::apply(std::ptr::null(), &wanted) }; + assert!(core.is_empty()); + let mut seen = Vec::new(); + let mut p = head as *const vk::BaseInStructure<'_>; + while !p.is_null() { + seen.push(unsafe { (*p).s_type }); + p = unsafe { (*p).p_next }; + } + assert_eq!( + seen, + vec![ + vk::StructureType::PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES, + vk::StructureType::PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES, + ] + ); + let f = unsafe { &*(head as *const vk::PhysicalDeviceShaderFloat16Int8Features<'_>) }; + assert_eq!((f.shader_int8, f.shader_float16), (vk::TRUE, vk::TRUE)); + unsafe { patch.restore() }; + } + + /// Ours go in front of the game's chain, which is then still reachable + /// from the head and left as it was. + #[test] + fn our_structs_lead_into_the_games_chain() { + let v11 = vk::PhysicalDeviceVulkan11Features::default(); + let chain = (&raw const v11).cast(); + let (patch, head, _) = unsafe { + FeaturePatch::apply( + chain, + &[ + Feature::BufferDeviceAddress, + Feature::StorageBuffer16BitAccess, + ], + ) + }; + let first = head as *const vk::BaseInStructure<'_>; + unsafe { + assert_eq!( + (*first).s_type, + vk::StructureType::PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES + ); + assert_eq!((*first).p_next.cast(), chain, "the game's chain was lost"); + } + // 16-bit storage had a home in the game's Vulkan 1.1 struct. + assert_eq!(v11.storage_buffer16_bit_access, vk::TRUE); + unsafe { patch.restore() }; + assert_eq!(v11.storage_buffer16_bit_access, vk::FALSE); + } + + /// The 1.0 features live in `VkPhysicalDeviceFeatures`. Inside a chained + /// `VkPhysicalDeviceFeatures2` they are set there and put back after. + #[test] + fn core_features_go_in_a_chained_features2() { + let mut f2 = vk::PhysicalDeviceFeatures2::default(); + f2.features.shader_int16 = vk::FALSE; + let chain = (&raw const f2).cast(); + let (patch, head, core) = unsafe { + FeaturePatch::apply( + chain, + &[ + Feature::ShaderInt16, + Feature::ShaderStorageImageWriteWithoutFormat, + ], + ) + }; + assert!(core.is_empty(), "Features2 was there to hold them"); + assert_eq!(head, chain); + assert_eq!(f2.features.shader_int16, vk::TRUE); + assert_eq!( + f2.features.shader_storage_image_write_without_format, + vk::TRUE + ); + unsafe { patch.restore() }; + assert_eq!(f2.features.shader_int16, vk::FALSE); + } + + /// Without one they are reported, never chained as a struct of their own: + /// a Features2 of ours beside the game's `pEnabledFeatures` is invalid. + #[test] + fn core_features_without_features2_are_handed_back() { + let (patch, head, core) = unsafe { + FeaturePatch::apply( + std::ptr::null(), + &[Feature::ShaderInt16, Feature::ShaderInt8], + ) + }; + assert_eq!(core, vec![Feature::ShaderInt16]); + let first = head as *const vk::BaseInStructure<'_>; + unsafe { + assert_eq!( + (*first).s_type, + vk::StructureType::PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES + ); + assert!((*first).p_next.is_null()); + } + unsafe { patch.restore() }; + } + + /// The game's own `pEnabledFeatures` is copied, never written: it is the + /// game's memory, and the copy keeps every bit it asked for. + #[test] + fn enabled_features_are_a_copy_with_ours_added() { + let game = vk::PhysicalDeviceFeatures { + geometry_shader: vk::TRUE, + ..Default::default() + }; + let ours = core_features( + Some(&game), + &[ + Feature::ShaderInt16, + Feature::ShaderStorageImageWriteWithoutFormat, + ], + ); + assert_eq!(ours.geometry_shader, vk::TRUE); + assert_eq!(ours.shader_int16, vk::TRUE); + assert_eq!(ours.shader_storage_image_write_without_format, vk::TRUE); + assert_eq!(game.shader_int16, vk::FALSE); + + let fresh = core_features(None, &[Feature::ShaderInt16]); + assert_eq!(fresh.shader_int16, vk::TRUE); + assert_eq!(fresh.geometry_shader, vk::FALSE); + } + + #[test] + fn nespyro_features_map_to_ours() { + let f = nespyro::DeviceFeatures { + shader_int16: true, + shader_float16: true, + subgroup_size_control: true, + ..Default::default() + }; + assert_eq!( + pyro_features(&f), + vec![ + Feature::ShaderInt16, + Feature::ShaderFloat16, + Feature::SubgroupSizeControl + ] + ); + } + // ── The null instance a layer below need not survive ───────────────────── /// On construction `ash::Entry` asks for the instance-global commands with diff --git a/apps/neshub/examples/pyro_bench.rs b/apps/neshub/examples/pyro_bench.rs new file mode 100644 index 00000000..0907e60f --- /dev/null +++ b/apps/neshub/examples/pyro_bench.rs @@ -0,0 +1,256 @@ +//! How much PyroWave an iroh connection carries. +//! +//! Sends PyroWave-shaped frames at a set rate through the same packing and the +//! same skip rule the hub uses, and collects them with the same collector the +//! client uses, then says what got through. Run it before trusting a rate: +//! +//! ```text +//! # CPU cost of the transport alone, both ends in this process: +//! cargo run --release -p neshub --example pyro_bench -- loopback --mbps 300 +//! +//! # The real thing, across a LAN. On the sending machine: +//! cargo run --release -p neshub --example pyro_bench -- serve --mbps 300 +//! # and on the receiving one, with what `serve` printed: +//! cargo run --release -p neshub --example pyro_bench -- connect +//! ``` +//! +//! Options: `--mbps`, `--fps` (60), `--secs` (10). + +use std::time::{Duration, Instant}; + +use anyhow::{Context, Result, bail}; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use bytes::Bytes; +use iroh::endpoint::{Connection, QuicTransportConfig, presets}; + +use nesprotocol::ALPN_VIDEO; +use nesprotocol::datagram::DGRAM_BUFFER_BYTES; +use nesprotocol::pyrowave::{ + FrameMeta, PYRO_DGRAM_HDR_LEN, PyroCollector, PyroFrame, decode_pyro_datagram, pack_datagrams, + should_skip, +}; + +struct Opts { + mbps: u64, + fps: u64, + secs: u64, +} + +fn transport() -> QuicTransportConfig { + QuicTransportConfig::builder() + .datagram_send_buffer_size(DGRAM_BUFFER_BYTES) + .datagram_receive_buffer_size(Some(DGRAM_BUFFER_BYTES)) + .build() +} + +async fn endpoint(alpn: bool) -> Result { + let mut b = iroh::Endpoint::builder(presets::N0) + .relay_mode(iroh::endpoint::RelayMode::Disabled) + .transport_config(transport()); + if alpn { + b = b.alpns(vec![ALPN_VIDEO.to_vec()]); + } + Ok(b.bind().await?) +} + +/// A frame of whole packets at `packet` bytes, behind an 8-byte stand-in for +/// the sequence header. +fn frame(bytes: usize, packet: usize) -> PyroFrame { + let mut data = Vec::with_capacity(bytes + packet); + let mut packets = Vec::with_capacity(bytes / packet + 2); + packets.push(0..8); + data.extend_from_slice(&[0; 8]); + while data.len() < bytes { + let start = data.len(); + data.extend((0..packet).map(|k| k as u8)); + packets.push(start..data.len()); + } + PyroFrame { + frame: 0, + meta: FrameMeta::default(), + critical: 3.min(packets.len()), + data, + packets, + } +} + +async fn send(conn: Connection, o: &Opts) -> Result<()> { + let max = conn + .max_datagram_size() + .context("peer does not take datagrams")?; + let packet = max - PYRO_DGRAM_HDR_LEN; + let frame_bytes = (o.mbps * 1_000_000 / 8 / o.fps) as usize; + let mut f = frame(frame_bytes, packet); + println!( + "send: {} Mbit/s at {} fps, {} KB frames of {} packets, {packet} B each", + o.mbps, + o.fps, + frame_bytes / 1000, + f.packets.len() + ); + + let mut tick = tokio::time::interval(Duration::from_micros(1_000_000 / o.fps)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let end = Instant::now() + Duration::from_secs(o.secs); + let (mut seq, mut sent, mut skipped) = (0u16, 0u64, 0u64); + let mut busy = Duration::ZERO; + let mut second = Instant::now(); + while Instant::now() < end { + tick.tick().await; + let t0 = Instant::now(); + f.meta.ts_ms = seq as u32; + // Read per frame, as the hub does: with several paths up, the limit + // is the smallest of their MTUs and moves as paths come and go. + let now_max = conn.max_datagram_size().context("datagrams went away")?; + let d = pack_datagrams(seq, &f, now_max)?; + let backlog = DGRAM_BUFFER_BYTES.saturating_sub(conn.datagram_send_buffer_space()); + if should_skip(backlog, d.buf.len(), DGRAM_BUFFER_BYTES).is_some() { + skipped += 1; + } else { + let buf = Bytes::from(d.buf); + let mut at = 0; + for len in d.lens { + conn.send_datagram(buf.slice(at..at + len))?; + at += len; + } + seq = seq.wrapping_add(1); + sent += 1; + } + busy += t0.elapsed(); + + if second.elapsed() >= Duration::from_secs(1) { + let path = conn.paths().iter().find(|p| p.is_selected()).map(|p| { + let s = p.stats(); + (p.rtt(), s.cwnd) + }); + println!( + "send: {sent} frames, {skipped} skipped, {:.1} ms of work per frame, path {path:?}", + busy.as_secs_f64() * 1000.0 / (sent + skipped).max(1) as f64 + ); + (sent, skipped, busy, second) = (0, 0, Duration::ZERO, Instant::now()); + } + } + // Let the queue drain before closing, or its tail reads as loss. + while conn.datagram_send_buffer_space() < DGRAM_BUFFER_BYTES { + tokio::time::sleep(Duration::from_millis(10)).await; + } + tokio::time::sleep(Duration::from_millis(200)).await; + conn.close(0u32.into(), b"done"); + Ok(()) +} + +async fn receive(conn: Connection) -> Result<()> { + let start = Instant::now(); + let mut c = PyroCollector::::new(); + let (mut bytes, mut datagrams) = (0u64, 0u64); + let (mut all_bytes, mut whole, mut partial, mut lost) = (0u64, 0u64, 0u64, 0u64); + let mut second = Instant::now(); + let mut busy = Duration::ZERO; + loop { + let d = match conn.read_datagram().await { + Ok(d) => d, + Err(_) => break, + }; + let t0 = Instant::now(); + datagrams += 1; + bytes += d.len() as u64; + if let Some((h, p)) = decode_pyro_datagram(&d) { + let at = d.len() - p.len(); + c.push(h, d.slice(at..), start.elapsed().as_micros() as u64); + } + busy += t0.elapsed(); + if second.elapsed() >= Duration::from_secs(1) { + c.poll(start.elapsed().as_micros() as u64); + let s = c.take_stats(); + println!( + "recv: {:.0} Mbit/s, {datagrams} datagrams, {} whole, {} partial, {} lost, \ + {} recovered, {:.2} µs of work per datagram", + bytes as f64 * 8.0 / second.elapsed().as_secs_f64() / 1e6, + s.whole, + s.partial, + s.lost, + s.recovered, + busy.as_secs_f64() * 1e6 / datagrams.max(1) as f64 + ); + all_bytes += bytes; + (whole, partial, lost) = (whole + s.whole, partial + s.partial, lost + s.lost); + (bytes, datagrams, busy, second) = (0, 0, Duration::ZERO, Instant::now()); + } + } + c.poll(u64::MAX); + let s = c.take_stats(); + (whole, partial, lost) = (whole + s.whole, partial + s.partial, lost + s.lost); + all_bytes += bytes; + println!( + "recv total: {:.0} Mbit/s average, {whole} whole, {partial} partial, {lost} datagrams lost", + all_bytes as f64 * 8.0 / start.elapsed().as_secs_f64() / 1e6 + ); + Ok(()) +} + +fn parse() -> Result<(String, Option, Opts)> { + let mut args = std::env::args().skip(1); + let mode = args.next().context("mode: loopback, serve or connect")?; + let mut target = None; + let mut o = Opts { + mbps: 300, + fps: 60, + secs: 10, + }; + while let Some(a) = args.next() { + let mut num = |name: &str| -> Result { + args.next() + .with_context(|| format!("{name} needs a value"))? + .parse() + .with_context(|| format!("{name} is a number")) + }; + match a.as_str() { + "--mbps" => o.mbps = num("--mbps")?, + "--fps" => o.fps = num("--fps")?, + "--secs" => o.secs = num("--secs")?, + other if target.is_none() && !other.starts_with("--") => target = Some(a), + other => bail!("unknown argument {other}"), + } + } + Ok((mode, target, o)) +} + +#[tokio::main] +async fn main() -> Result<()> { + let (mode, target, o) = parse()?; + match mode.as_str() { + "loopback" => { + let (hub, client) = (endpoint(true).await?, endpoint(false).await?); + let addr = hub.addr(); + let accept = tokio::spawn(async move { + let conn = hub.accept().await.context("accept")?.await?; + Ok::<_, anyhow::Error>((hub, conn)) + }); + let rx = client.connect(addr, ALPN_VIDEO).await?; + let (_hub, tx) = accept.await??; + let r = tokio::spawn(receive(rx)); + send(tx, &o).await?; + r.await??; + } + "serve" => { + let hub = endpoint(true).await?; + // Direct addresses are learned over the first moments. + tokio::time::sleep(Duration::from_secs(1)).await; + let addr = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&hub.addr())?); + println!("connect with:\n{addr}"); + let conn = hub.accept().await.context("accept")?.await?; + send(conn, &o).await?; + } + "connect" => { + let addr = target.context("connect needs the address serve printed")?; + let addr: iroh::EndpointAddr = + serde_json::from_slice(&URL_SAFE_NO_PAD.decode(addr.as_bytes())?)?; + let client = endpoint(false).await?; + let conn = client.connect(addr, ALPN_VIDEO).await?; + receive(conn).await?; + } + other => bail!("unknown mode {other}"), + } + Ok(()) +} diff --git a/apps/neshub/src/dgram.rs b/apps/neshub/src/dgram.rs index 3781858f..bd8964fb 100644 --- a/apps/neshub/src/dgram.rs +++ b/apps/neshub/src/dgram.rs @@ -7,19 +7,35 @@ use std::fmt; use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; -use bytes::BytesMut; +use bytes::{Bytes, BytesMut}; use iroh::endpoint::{Connection, QuicTransportConfig, SendDatagramError}; use tracing::{debug, warn}; use crate::keyframe::KeyframeSender; +use crate::pyro::PyroSender; use nesprotocol::MSG_DATA; use nesprotocol::datagram::{ DGRAM_BUFFER_BYTES, DGRAM_HDR_LEN, MAX_FRAGMENTS_PER_FRAME, MIN_DGRAM_PAYLOAD, fragment_count, write_datagram_header, }; +use nesprotocol::pyrowave::PyroFrame; use nesprotocol::reliable::video_wants_reliable; +/// What a session's media writer is handed. +/// +/// Both are refcounted, so a frame broadcast to several clients is shared +/// rather than copied per client — which at PyroWave's frame sizes stops being +/// a rounding error. +#[derive(Debug, Clone)] +pub enum MediaItem { + /// A hardware-codec frame or an audio packet, laid out as the hub + /// broadcasts it. + Coded(Bytes), + /// A whole PyroWave frame, which travels by rules of its own. + PyroWave(Arc), +} + /// Transport settings for an endpoint carrying media datagrams. /// /// See [`DGRAM_BUFFER_BYTES`] for why the defaults are not enough. @@ -225,7 +241,7 @@ pub async fn run_datagram_writer( conn: Connection, kind: u8, label: &'static str, - mut rx: tokio::sync::mpsc::UnboundedReceiver>, + mut rx: tokio::sync::mpsc::UnboundedReceiver, relay_ms: Option>, keyframes_reliable: bool, // Set while this client has asked for a keyframe and not yet been sent @@ -234,7 +250,10 @@ pub async fn run_datagram_writer( // Counts frames withheld during a resynchronisation, so the bitrate // controller can tell a second it starved from a second the path did. withheld: Option>, + // Counts PyroWave frames skipped for want of room. Video only. + pyro_skipped: Option>, ) { + let mut pyro = pyro_skipped.map(|skipped| PyroSender::new(conn.clone(), skipped)); let sender = DatagramSender::new(conn.clone(), kind); let keyframes = keyframes_reliable.then(|| KeyframeSender::new(conn, sender.clone())); let mut seq: u16 = 0; @@ -245,8 +264,25 @@ pub async fn run_datagram_writer( let mut warned_unsupported = false; let mut resync = ResyncGate::default(); - while let Some(payload) = rx.recv().await { + while let Some(item) = rx.recv().await { let t0 = std::time::Instant::now(); + let payload = match item { + MediaItem::Coded(payload) => payload, + // No keyframe stream and no resync gate: every PyroWave frame + // stands alone, so there is nothing to resynchronise and no frame + // worth a reliable stream more than another. + MediaItem::PyroWave(frame) => { + match pyro.as_mut() { + Some(pyro) => pyro.send(&frame, label), + None => debug!("{label}: a PyroWave frame on a carrier that has none"), + } + if let Some(ref relay) = relay_ms { + let elapsed = t0.elapsed().as_secs_f32() * 1000.0; + relay.store(elapsed.to_bits(), Ordering::Relaxed); + } + continue; + } + }; body.clear(); nesprotocol::encode_frame_body(&mut body, MSG_DATA, seq, &payload); diff --git a/apps/neshub/src/e2e_tests.rs b/apps/neshub/src/e2e_tests.rs new file mode 100644 index 00000000..d2294406 --- /dev/null +++ b/apps/neshub/src/e2e_tests.rs @@ -0,0 +1,311 @@ +//! PyroWave through the hub for real: IPC messages into the video listener, +//! out of the datagram writer, over an iroh connection, into a collector. + +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use iroh::endpoint::{Connection, presets}; +use tokio::net::UnixDatagram; + +use nesprotocol::pyrowave::{ + CollectedFrame, FLAG_DUPLICATE, FrameMeta, PyroCollector, PyroDatagramHeader, PyroFrame, + chunk_frame, decode_pyro_datagram, +}; +use nesprotocol::{ALPN_VIDEO, Carrier}; + +use crate::control::{Controller, Limits}; +use crate::session::SessionManager; + +/// A frame of synthetic packets, each filled with bytes particular to it, so +/// a packet delivered into the wrong place or the wrong frame shows. Packet 0 +/// stands in for the sequence header; one in every fifty is an oversized +/// block that has to travel in pieces. +pub(crate) fn synthetic_frame(id: u32, bytes: usize, packet: usize) -> PyroFrame { + let mut data = Vec::with_capacity(bytes); + let mut packets = Vec::new(); + let mut i = 0usize; + while data.len() < bytes { + let len = match i { + 0 => 8, + _ if i.is_multiple_of(50) => nesprotocol::pyrowave::LARGEST_BLOCK_BYTES, + _ => packet - (i * 13) % 200, + }; + let start = data.len(); + let seed = (id as usize).wrapping_mul(7919) ^ i; + data.extend((0..len).map(|k| (seed.wrapping_mul(31).wrapping_add(k) % 251) as u8)); + packets.push(start..data.len()); + i += 1; + } + PyroFrame { + frame: id, + meta: FrameMeta { + ts_ms: id * 16, + width: 1920, + height: 1080, + }, + critical: 3, + data, + packets, + } +} + +struct Rig { + socket_path: PathBuf, + sender: UnixDatagram, + client: Connection, + mgr: Arc, + _hub: iroh::Endpoint, + _client_ep: iroh::Endpoint, +} + +impl Drop for Rig { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.socket_path); + } +} + +/// `client_buffer` sizes the client's datagram receive buffer; `None` is the +/// client's real one. +async fn rig_with(client_buffer: Option) -> Rig { + static NEXT: AtomicU32 = AtomicU32::new(0); + let socket_path = std::env::temp_dir().join(format!( + "neshub-e2e-{}-{}.sock", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + let mgr = Arc::new(SessionManager::new(300_000)); + tokio::spawn({ + let mgr = mgr.clone(); + let path = socket_path.clone(); + async move { crate::ipc_listener::run_video_listener(path, mgr).await } + }); + + let hub = iroh::Endpoint::builder(presets::N0) + .alpns(vec![ALPN_VIDEO.to_vec()]) + .relay_mode(iroh::endpoint::RelayMode::Disabled) + .transport_config(crate::dgram::media_transport_config()) + .bind() + .await + .expect("hub endpoint"); + let client_transport = match client_buffer { + Some(bytes) => iroh::endpoint::QuicTransportConfig::builder() + .datagram_send_buffer_size(nesprotocol::datagram::DGRAM_BUFFER_BYTES) + .datagram_receive_buffer_size(Some(bytes)) + .build(), + None => crate::dgram::media_transport_config(), + }; + let client_ep = iroh::Endpoint::builder(presets::N0) + .relay_mode(iroh::endpoint::RelayMode::Disabled) + .transport_config(client_transport) + .bind() + .await + .expect("client endpoint"); + + let accept = tokio::spawn({ + let hub = hub.clone(); + let mgr = mgr.clone(); + async move { + let conn = hub.accept().await.expect("incoming").await.expect("conn"); + let (input, _) = tokio::sync::broadcast::channel(4); + let (cmd, _cmd_rx) = tokio::sync::mpsc::unbounded_channel(); + let controller = Arc::new(tokio::sync::Mutex::new(Controller::new(Limits::new(4_000)))); + mgr.attach( + conn.remote_id(), + Carrier::Video, + conn, + input, + cmd, + controller, + ) + .await; + } + }); + let client = client_ep + .connect(hub.addr(), ALPN_VIDEO) + .await + .expect("dial the hub"); + accept.await.expect("accept"); + + // The listener binds on its own task; wait for the socket to exist. + let sender = UnixDatagram::unbound().expect("unbound socket"); + for _ in 0..200 { + if socket_path.exists() { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + sender + .connect(&socket_path) + .expect("connect to the listener"); + + Rig { + socket_path, + sender, + client, + mgr, + _hub: hub, + _client_ep: client_ep, + } +} + +async fn rig() -> Rig { + rig_with(None).await +} + +async fn send_frame(sender: &UnixDatagram, f: &PyroFrame) { + for m in chunk_frame(f.frame, f.meta, &f.data, &f.packets, f.critical).expect("chunks") { + sender.send(&m).await.expect("ipc send"); + } +} + +/// Read datagrams until `want` frames have come out of the collector or the +/// line goes quiet. `drop_datagram` decides, by arrival order and header, +/// which ones the "network" loses. +async fn collect( + conn: &Connection, + want: usize, + mut drop_datagram: impl FnMut(usize, &PyroDatagramHeader) -> bool, +) -> Vec> { + let start = Instant::now(); + let now = || start.elapsed().as_micros() as u64; + let mut collector = PyroCollector::new(); + let mut frames = Vec::new(); + let mut arrivals = 0usize; + while frames.len() < want { + let datagram = + match tokio::time::timeout(Duration::from_millis(500), conn.read_datagram()).await { + Ok(Ok(d)) => d, + Ok(Err(e)) => panic!("connection failed: {e}"), + Err(_) => break, + }; + arrivals += 1; + let (h, payload) = decode_pyro_datagram(&datagram).expect("a PyroWave datagram"); + if drop_datagram(arrivals, &h) { + continue; + } + let at = datagram.len() - payload.len(); + frames.extend(collector.push(h, datagram.slice(at..), now())); + } + // Whatever is still in flight leaves at its deadline. + frames.extend(collector.poll(u64::MAX)); + frames +} + +fn packets(f: &CollectedFrame) -> Vec> { + f.packets.iter().map(|p| p.as_ref().to_vec()).collect() +} + +fn originals(f: &PyroFrame) -> Vec> { + (0..f.packets.len()).map(|i| f.packet(i).to_vec()).collect() +} + +#[tokio::test(flavor = "multi_thread")] +async fn frames_cross_the_hub_byte_for_byte() { + let rig = rig().await; + let sent: Vec<_> = (0..20).map(|i| synthetic_frame(i, 120_000, 1188)).collect(); + let reader = tokio::spawn({ + let conn = rig.client.clone(); + async move { collect(&conn, 20, |_, _| false).await } + }); + for f in &sent { + send_frame(&rig.sender, f).await; + // About 60 Mbit/s: this is about the bytes, not about how fast they + // go, and an unoptimised build's QUIC cannot carry much more. The + // bench is where rates are measured. + tokio::time::sleep(Duration::from_millis(16)).await; + } + let got = reader.await.expect("reader"); + + assert_eq!( + rig.mgr.take_pyro_skipped().await, + 0, + "a quiet path skipped frames" + ); + assert_eq!(got.len(), sent.len()); + for (i, (g, s)) in got.iter().zip(&sent).enumerate() { + assert!(g.is_whole(), "frame {i}: {}/{}", g.received, g.total); + assert_eq!(g.seq, i as u16); + assert_eq!(g.ts_ms, s.meta.ts_ms); + assert_eq!(packets(g), originals(s), "frame {i} changed in transit"); + } +} + +/// With datagrams lost, what arrives is still exactly what was sent, minus +/// the lost packets — never a packet stitched from two, never one from +/// another frame — and the critical copies stand in for lost originals. +#[tokio::test(flavor = "multi_thread")] +async fn a_lossy_path_loses_packets_and_nothing_else() { + let rig = rig().await; + let sent: Vec<_> = (0..20).map(|i| synthetic_frame(i, 120_000, 1188)).collect(); + let reader = tokio::spawn({ + let conn = rig.client.clone(); + // One in 37: about 4% loss, landing on every position in turn, the + // critical datagrams at the head of a frame included. Copies are + // spared, so that a lost critical original is always covered and the + // assertion below can be exact. + async move { + collect(&conn, 20, |n, h| { + h.flags & FLAG_DUPLICATE == 0 && n.is_multiple_of(37) + }) + .await + } + }); + for f in &sent { + send_frame(&rig.sender, f).await; + tokio::time::sleep(Duration::from_millis(16)).await; + } + let got = reader.await.expect("reader"); + assert_eq!(got.len(), sent.len()); + + let mut partial = 0; + for (g, s) in got.iter().zip(&sent) { + let all = originals(s); + let got = packets(g); + // Every critical packet arrives: the copies cover a single loss. + assert_eq!(&got[..s.critical], &all[..s.critical], "frame {}", s.frame); + // And what arrived is an ordered subset of what was sent. + let mut rest = all.iter(); + for p in &got { + assert!( + rest.any(|o| o == p), + "frame {}: a packet arrived that was never sent", + s.frame + ); + } + partial += usize::from(!g.is_whole()); + } + assert!(partial > 0, "the loss never reached the collector"); +} + +/// Handed frames faster than the path drains them, the hub skips whole +/// frames rather than letting QUIC evict pieces of two. +/// +/// The client's receive buffer is made larger than the whole burst, so a +/// torn frame here can only be the sender's doing: a client too slow to read +/// tears frames too, by its own buffer's eviction, and that is not what this +/// is about. +#[tokio::test(flavor = "multi_thread")] +async fn a_burst_is_skipped_whole_not_torn() { + let rig = rig_with(Some(64 << 20)).await; + let n = 30u32; + let reader = tokio::spawn({ + let conn = rig.client.clone(); + async move { collect(&conn, n as usize, |_, _| false).await } + }); + for i in 0..n { + send_frame(&rig.sender, &synthetic_frame(i, 1_500_000, 1188)).await; + } + let got = reader.await.expect("reader"); + let skipped = rig.mgr.take_pyro_skipped().await as usize; + + let torn = got.iter().filter(|f| !f.is_whole()).count(); + assert_eq!(torn, 0, "{torn} frames arrived damaged ({skipped} skipped)"); + assert_eq!( + got.len() + skipped, + n as usize, + "frames went missing unaccounted" + ); +} diff --git a/apps/neshub/src/ipc_listener.rs b/apps/neshub/src/ipc_listener.rs index 6b5d738a..624e346f 100644 --- a/apps/neshub/src/ipc_listener.rs +++ b/apps/neshub/src/ipc_listener.rs @@ -4,9 +4,11 @@ use std::sync::Arc; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{UnixDatagram, UnixListener}; use tokio::sync::broadcast::error::RecvError; +use tracing::debug; use crate::session::SessionManager; -use nesprotocol::{CODEC_OPUS, STREAM_AUDIO, STREAM_VIDEO, decode_ipc_frame}; +use nesprotocol::pyrowave::FrameAssembler; +use nesprotocol::{CODEC_OPUS, STREAM_AUDIO, STREAM_PYROWAVE, STREAM_VIDEO, decode_ipc_frame}; fn set_recv_buffer(socket: &UnixDatagram, size: libc::c_int) { unsafe { @@ -46,10 +48,28 @@ pub async fn run_video_listener(socket_path: PathBuf, session_manager: Arc { if let Some(decoded) = decode_ipc_frame(&buf[..n]) { + if decoded.stream_type == STREAM_PYROWAVE { + let now_ms = started.elapsed().as_millis() as u64; + match pyrowave.push(&decoded, now_ms) { + Ok(Some(frame)) => session_manager.broadcast_pyrowave(frame).await, + Ok(None) => {} + Err(e) => debug!("dropping a PyroWave IPC message: {e}"), + } + // Per frame, so debug: an encoder whose writes keep + // timing out would say this sixty times a second. + let abandoned = pyrowave.take_abandoned(); + if abandoned > 0 { + debug!("abandoned {abandoned} PyroWave frame(s) missing a message"); + } + continue; + } if decoded.stream_type != STREAM_VIDEO { tracing::warn!( "unexpected stream type on video socket: {}", diff --git a/apps/neshub/src/main.rs b/apps/neshub/src/main.rs index c8825076..1e678c17 100644 --- a/apps/neshub/src/main.rs +++ b/apps/neshub/src/main.rs @@ -1,7 +1,10 @@ mod control; mod dgram; +#[cfg(test)] +mod e2e_tests; mod ipc_listener; mod keyframe; +mod pyro; mod screenshot; mod session; mod ticket; @@ -84,6 +87,15 @@ struct Args { #[arg(long, env = "NESTRI_MAX_BITRATE")] max_bitrate_kbps: Option, + /// The PyroWave rate, in kbps, for a client that asks for PyroWave without + /// naming one. + /// + /// PyroWave is the LAN codec and runs at a set rate, outside the bitrate + /// controller and above the ceiling: it is only ever used because a client + /// asked for it. + #[arg(long, env = "NESTRI_PYROWAVE_KBPS", default_value_t = 300_000)] + pyrowave_kbps: u32, + /// Path for the gamepad IPC socket (neshub ↔ nesgamepad). neshub /// listens; the gamepad service dials in. #[arg( @@ -110,6 +122,41 @@ struct Args { /// chosen and the encoder's own default stood in for one. const DEFAULT_MAX_BITRATE_KBPS: u32 = 4_000; +/// What a PyroWave packet is sized to with no client to measure: a datagram +/// at QUIC's minimum MTU always carries this much. +const FALLBACK_PYRO_PACKET_BYTES: usize = 1100; + +/// Name the PyroWave packet size in an encode command bound for nescapture. +/// +/// Filled in here, as the command leaves, because what fits a datagram is a +/// fact about the path *now*. Any other command passes untouched. +fn fill_pyrowave_packet_size(cmd: &mut Vec, limit: Option) { + if cmd.first() != Some(&nesprotocol::MSG_ENCODE_SETTINGS) { + return; + } + let Some(mut settings) = nesprotocol::decode_encode_settings_ext(&cmd[1..]) else { + return; + }; + if settings.codec != nesprotocol::CODEC_PYROWAVE { + return; + } + let size = limit + .unwrap_or(FALLBACK_PYRO_PACKET_BYTES) + .min(usize::from(u16::MAX)); + settings.packet_size = Some(size as u16); + // Every field before it has to be present for it to be read. + settings.bit_depth.get_or_insert(nesprotocol::DEPTH_8); + settings + .chroma + .get_or_insert(nesprotocol::pyrowave::CHROMA_420); + cmd.truncate(1); + nesprotocol::encode_settings(cmd, &settings); + tracing::info!( + "PyroWave at {} kbps, packets of {size} bytes", + settings.value + ); +} + #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt() @@ -169,7 +216,7 @@ async fn main() -> Result<()> { let (nescope_stats_tx, mut nescope_stats_rx) = tokio::sync::mpsc::unbounded_channel::>(); - let session_manager = Arc::new(SessionManager::new()); + let session_manager = Arc::new(SessionManager::new(args.pyrowave_kbps)); // One controller for the box, not one per client: there is one encoder, so // there is one bitrate, and the client having the worst time is the one it @@ -182,9 +229,11 @@ async fn main() -> Result<()> { // IDR / encode settings command channel: input reader → nescapture let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::>(); { + let mgr = session_manager.clone(); tokio::spawn(async move { let cmd_path = std::path::PathBuf::from("/tmp/nescapture-cmd.sock"); - while let Some(bytes) = cmd_rx.recv().await { + while let Some(mut bytes) = cmd_rx.recv().await { + fill_pyrowave_packet_size(&mut bytes, mgr.pyro_packet_size().await); if let Ok(sock) = std::os::unix::net::UnixDatagram::unbound() { if sock.send_to(&bytes, &cmd_path).is_err() { tracing::warn!("nescapture cmd send failed at {}", cmd_path.display()); @@ -262,6 +311,7 @@ async fn main() -> Result<()> { } let mut interval = tokio::time::interval(std::time::Duration::from_secs(1)); + let mut pyro_skipping = false; loop { interval.tick().await; @@ -291,6 +341,22 @@ async fn main() -> Result<()> { } } + // Said when skipping starts and when it stops, and otherwise + // only at debug: a path that cannot carry the rate skips every + // second until somebody lowers it. + let skipped = mgr.take_pyro_skipped().await; + if skipped > 0 && !pyro_skipping { + tracing::warn!( + "skipping PyroWave frames: the path is not carrying the rate \ + ({skipped} this second)" + ); + } else if skipped == 0 && pyro_skipping { + tracing::info!("no longer skipping PyroWave frames"); + } else if skipped > 0 { + tracing::debug!("{skipped} PyroWave frames skipped this second"); + } + pyro_skipping = skipped > 0; + let clients = mgr.client_count().await as u8; let (key_bps, delta_bps, keyframes) = mgr.video_breakdown(); let (pipeline_p50_ms, pipeline_p95_ms, pipeline_max_ms) = mgr.pipeline_delays(); @@ -469,3 +535,54 @@ async fn main() -> Result<()> { let _ = std::fs::remove_file(&args.ticket_ipc); Ok(()) } + +#[cfg(test)] +mod tests { + use super::{FALLBACK_PYRO_PACKET_BYTES, fill_pyrowave_packet_size}; + use nesprotocol::{ + CODEC_AV1, CODEC_PYROWAVE, EncodeSettings, MSG_ENCODE_SETTINGS, MSG_IDR_REQUEST, RC_CBR, + }; + + fn command(codec: u8) -> Vec { + let mut cmd = vec![MSG_ENCODE_SETTINGS]; + nesprotocol::encode_settings( + &mut cmd, + &EncodeSettings { + codec, + rate_control: RC_CBR, + value: 300_000, + bit_depth: Some(0), + chroma: Some(1), + packet_size: None, + }, + ); + cmd + } + + #[test] + fn a_pyrowave_command_leaves_with_the_paths_packet_size() { + let mut cmd = command(CODEC_PYROWAVE); + fill_pyrowave_packet_size(&mut cmd, Some(1440)); + let s = nesprotocol::decode_encode_settings_ext(&cmd[1..]).unwrap(); + assert_eq!(s.packet_size, Some(1440)); + assert_eq!( + (s.value, s.chroma), + (300_000, Some(1)), + "nothing else moved" + ); + + let mut cmd = command(CODEC_PYROWAVE); + fill_pyrowave_packet_size(&mut cmd, None); + let s = nesprotocol::decode_encode_settings_ext(&cmd[1..]).unwrap(); + assert_eq!(s.packet_size, Some(FALLBACK_PYRO_PACKET_BYTES as u16)); + } + + #[test] + fn every_other_command_passes_untouched() { + for cmd in [command(CODEC_AV1), vec![MSG_IDR_REQUEST], vec![]] { + let mut sent = cmd.clone(); + fill_pyrowave_packet_size(&mut sent, Some(1440)); + assert_eq!(sent, cmd); + } + } +} diff --git a/apps/neshub/src/pyro.rs b/apps/neshub/src/pyro.rs new file mode 100644 index 00000000..5763e289 --- /dev/null +++ b/apps/neshub/src/pyro.rs @@ -0,0 +1,87 @@ +//! Sending PyroWave frames. +//! +//! See `nesprotocol::pyrowave` for the wire and why PyroWave has a path of its +//! own. This is the hub's half: pack each frame as one packet per datagram and +//! decide whether to send it at all. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use bytes::Bytes; +use iroh::endpoint::Connection; +use tracing::{debug, warn}; + +use nesprotocol::datagram::DGRAM_BUFFER_BYTES; +use nesprotocol::pyrowave::{PYRO_DGRAM_HDR_LEN, PyroFrame, Skip, pack_datagrams, should_skip}; + +/// One connection's PyroWave sender. +pub struct PyroSender { + conn: Connection, + /// Its own counter, apart from the hardware frames': a receiver collects + /// each kind by its own sequence. + seq: u16, + skipped: Arc, + warned_too_large: bool, +} + +impl PyroSender { + pub fn new(conn: Connection, skipped: Arc) -> Self { + Self { + conn, + seq: 0, + skipped, + warned_too_large: false, + } + } + + /// Send one frame, or skip it. + /// + /// A skipped frame does not use up a sequence number: nothing went out + /// under it, and a hole would read at the far end as a frame the path + /// lost. + pub fn send(&mut self, frame: &PyroFrame, label: &str) { + let Some(max) = self.conn.max_datagram_size() else { + debug!("{label}: peer does not support QUIC datagrams"); + return; + }; + let datagrams = match pack_datagrams(self.seq, frame, max) { + Ok(d) => d, + Err(e) => { + debug!("{label}: dropping PyroWave frame {}: {e}", frame.frame); + return; + } + }; + + let backlog = DGRAM_BUFFER_BYTES.saturating_sub(self.conn.datagram_send_buffer_space()); + if let Some(skip) = should_skip(backlog, datagrams.buf.len(), DGRAM_BUFFER_BYTES) { + self.skipped.fetch_add(1, Ordering::Relaxed); + if skip == Skip::TooLarge && !self.warned_too_large { + warn!( + "{label}: a {} byte PyroWave frame cannot fit the {DGRAM_BUFFER_BYTES} byte \ + send buffer; frames this size are never sent. Lower the rate.", + datagrams.buf.len() + ); + self.warned_too_large = true; + } else { + debug!("{label}: skipped PyroWave frame {}: {skip:?}", frame.frame); + } + return; + } + + // Taking the `Vec` is free, and each datagram is a slice of it. + let buf = Bytes::from(datagrams.buf); + let mut at = 0; + for len in datagrams.lens { + debug_assert!(len >= PYRO_DGRAM_HDR_LEN); + let datagram = buf.slice(at..at + len); + at += len; + // Not `send_datagram_wait`, for the reason `dgram::send_frame` + // gives; the guard above is what keeps this from evicting. + if let Err(e) = self.conn.send_datagram(datagram) { + debug!("{label}: PyroWave frame {} cut short: {e}", frame.frame); + break; + } + } + self.seq = self.seq.wrapping_add(1); + } +} diff --git a/apps/neshub/src/session.rs b/apps/neshub/src/session.rs index 532ffd27..1cb17296 100644 --- a/apps/neshub/src/session.rs +++ b/apps/neshub/src/session.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use iroh::endpoint::Connection; use tokio::sync::Mutex; -use tracing::{debug, info}; +use tracing::{debug, info, warn}; use nesprotocol::datagram::{DGRAM_AUDIO, DGRAM_BUFFER_BYTES, DGRAM_VIDEO}; use nesprotocol::input::{INPUT_KEY, INPUT_MOUSE_BUTTON, INPUT_MOUSE_MOVE, INPUT_MOUSE_WHEEL}; @@ -17,7 +17,9 @@ use nesprotocol::{ use crate::control::{Controller, PathView}; -use crate::dgram::run_datagram_writer; +use crate::dgram::{MediaItem, run_datagram_writer}; +use bytes::Bytes; +use nesprotocol::pyrowave::PyroFrame; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; @@ -58,6 +60,13 @@ pub struct ClientSession { /// Overwritten rather than queued. A report describes the second that just /// passed, and an older one is not evidence about now. latest_report: Arc>>, + /// PyroWave frames this client was not sent because its queue was full. + pyro_skipped: Arc, + /// Whether this client said it decodes the PyroWave format this hub + /// carries. Only then is a request for PyroWave passed on. + pyrowave_ok: Arc, + /// The rate a PyroWave request that names none runs at. + pyrowave_kbps: u32, relay_ms: Arc, input_broadcast: tokio::sync::broadcast::Sender>, /// This client's number on the gamepad socket, so the box's side can tell @@ -70,13 +79,13 @@ pub struct ClientSession { pending_gamepad_feedback: Option>>, idr_cmd_tx: tokio::sync::mpsc::UnboundedSender>, controller: Arc>, - send_video: tokio::sync::mpsc::UnboundedSender>, - send_audio: tokio::sync::mpsc::UnboundedSender>, + send_video: tokio::sync::mpsc::UnboundedSender, + send_audio: tokio::sync::mpsc::UnboundedSender, send_cursor: tokio::sync::mpsc::UnboundedSender>, send_stats: tokio::sync::mpsc::UnboundedSender>, /// Receiving ends, held until the carrier that drains them connects. - pending_video: Option>>, - pending_audio: Option>>, + pending_video: Option>, + pending_audio: Option>, pending_cursor: Option>>, pending_stats: Option>>, tasks: Vec>, @@ -105,10 +114,11 @@ impl ClientSession { relay_ms: Arc, idr_cmd_tx: tokio::sync::mpsc::UnboundedSender>, controller: Arc>, + pyrowave_kbps: u32, ) -> Self { let (gamepad_feedback_tx, gamepad_feedback_rx) = tokio::sync::mpsc::unbounded_channel(); - let (video_tx, video_rx) = tokio::sync::mpsc::unbounded_channel::>(); - let (audio_tx, audio_rx) = tokio::sync::mpsc::unbounded_channel::>(); + let (video_tx, video_rx) = tokio::sync::mpsc::unbounded_channel::(); + let (audio_tx, audio_rx) = tokio::sync::mpsc::unbounded_channel::(); let (cursor_tx, cursor_rx) = tokio::sync::mpsc::unbounded_channel::>(); let (stats_tx, stats_rx) = tokio::sync::mpsc::unbounded_channel::>(); Self { @@ -117,6 +127,9 @@ impl ClientSession { awaiting_keyframe: Arc::new(AtomicBool::new(false)), withheld: Arc::new(AtomicU64::new(0)), latest_report: Arc::new(std::sync::Mutex::new(None)), + pyro_skipped: Arc::new(AtomicU64::new(0)), + pyrowave_ok: Arc::new(AtomicBool::new(false)), + pyrowave_kbps, relay_ms, input_broadcast, gamepad_session, @@ -161,6 +174,7 @@ impl ClientSession { let relay_ms = self.relay_ms.clone(); let awaiting = self.awaiting_keyframe.clone(); let withheld = self.withheld.clone(); + let pyro_skipped = self.pyro_skipped.clone(); // Keyframes go on reliable streams of their own, because a lost // keyframe freezes the picture until the next one instead of // costing a single frame. See `nesprotocol::reliable`. @@ -174,6 +188,7 @@ impl ClientSession { true, Some(awaiting), Some(withheld), + Some(pyro_skipped), ) .await })); @@ -187,8 +202,18 @@ impl ClientSession { // No reliable path offered: audio has no keyframes to promote, // and on its own connection it is no longer queued behind any. self.tasks.push(tokio::spawn(async move { - run_datagram_writer(conn, DGRAM_AUDIO, "audio", rx, None, false, None, None) - .await + run_datagram_writer( + conn, + DGRAM_AUDIO, + "audio", + rx, + None, + false, + None, + None, + None, + ) + .await })); } Carrier::Input => { @@ -226,8 +251,12 @@ impl ClientSession { let awaiting = self.awaiting_keyframe.clone(); let idr = self.idr_cmd_tx.clone(); let controller = self.controller.clone(); + let pyrowave = PyroWaveGate { + ok: self.pyrowave_ok.clone(), + default_kbps: self.pyrowave_kbps, + }; self.tasks.push(tokio::spawn(async move { - run_control_reader(conn, idr, reports, controller, awaiting).await + run_control_reader(conn, idr, reports, controller, awaiting, pyrowave).await })); } } @@ -247,6 +276,17 @@ impl ClientSession { self.latest_report.lock().ok()?.take() } + /// PyroWave frames skipped since the last call, cleared as it is read. + pub fn take_pyro_skipped(&self) -> u64 { + self.pyro_skipped.swap(0, Ordering::Relaxed) + } + + /// The largest PyroWave packet one of this client's datagrams carries now. + pub fn pyro_packet_limit(&self) -> Option { + let max = self.video_conn.as_ref()?.max_datagram_size()?; + Some(max.saturating_sub(nesprotocol::pyrowave::PYRO_DGRAM_HDR_LEN)) + } + /// What this end can see of the path, from the route actually in use. /// /// A connection can hold several paths at once -- typically one through a @@ -284,14 +324,14 @@ impl ClientSession { /// only way it fails is a closed channel, which means the writer is already /// gone, and sixty warnings a second on the way down bury whatever actually /// ended the session. - pub fn send_video_frame(&self, data: Vec) { - if let Err(e) = self.send_video.send(data) { + pub fn send_video_frame(&self, item: MediaItem) { + if let Err(e) = self.send_video.send(item) { debug!("failed to send video data: {e}"); } } - pub fn send_audio_packet(&self, data: Vec) { - if let Err(e) = self.send_audio.send(data) { + pub fn send_audio_packet(&self, data: Bytes) { + if let Err(e) = self.send_audio.send(MediaItem::Coded(data)) { debug!("failed to send audio data: {e}"); } } @@ -499,6 +539,44 @@ async fn run_input_reader( .await; } +/// What the control reader needs to decide about a PyroWave request. +#[derive(Clone)] +struct PyroWaveGate { + /// Set once the client's capabilities named PyroWave in this hub's format. + ok: Arc, + default_kbps: u32, +} + +/// Turn a client's PyroWave request into the command nescapture gets, or say +/// why it cannot have one. +/// +/// A rate of zero means "the hub's default". The packet size is left for the +/// command forwarder, which fills it in as it sends: only the hub sees the +/// path, and what fits a datagram is a fact about the moment the encoder is +/// told, not the moment the client asked. +fn pyrowave_request( + request: nesprotocol::EncodeSettings, + gate: &PyroWaveGate, +) -> Result { + if !gate.ok.load(Ordering::Relaxed) { + return Err("the client never said it decodes this hub's PyroWave format"); + } + if request.rate_control != nesprotocol::RC_CBR { + return Err("PyroWave runs at a set rate, and the request named none"); + } + Ok(nesprotocol::EncodeSettings { + value: if request.value == 0 { + gate.default_kbps + } else { + request.value + }, + bit_depth: Some(request.bit_depth.unwrap_or(nesprotocol::DEPTH_8)), + chroma: Some(request.chroma.unwrap_or(nesprotocol::pyrowave::CHROMA_420)), + packet_size: None, + ..request + }) +} + /// Everything the client says that is not an input event. async fn run_control_reader( conn: Connection, @@ -506,12 +584,14 @@ async fn run_control_reader( latest_report: Arc>>, controller: Arc>, awaiting_keyframe: Arc, + pyrowave: PyroWaveGate, ) { run_framed_reader(conn, BIDI_CONTROL, "control", None, |msg_type, payload| { let idr_cmd_tx = idr_cmd_tx.clone(); let latest_report = latest_report.clone(); let controller = controller.clone(); let awaiting_keyframe = awaiting_keyframe.clone(); + let pyrowave = pyrowave.clone(); async move { match msg_type { MSG_IDR_REQUEST => { @@ -535,6 +615,21 @@ async fn run_control_reader( // decode. match nesprotocol::decode_client_caps(&payload) { Some(caps) => { + // Recorded whatever the mode: it is a fact about the + // client, consulted only if it asks for PyroWave. + let format = nesprotocol::decode_pyrowave_format(&payload); + let pyro = caps.supports_codec(nesprotocol::CODEC_PYROWAVE); + pyrowave.ok.store( + pyro && format == Some(nesprotocol::pyrowave::PYROWAVE_FORMAT), + Ordering::Relaxed, + ); + if pyro && format != Some(nesprotocol::pyrowave::PYROWAVE_FORMAT) { + info!( + "client decodes PyroWave format {format:?}, this hub carries \ + {}; it will not be sent PyroWave", + nesprotocol::pyrowave::PYROWAVE_FORMAT + ); + } // Manual means a person is choosing, and the panel // they chose in sets the codec and the depth // together. A client joining afterwards must not @@ -563,6 +658,22 @@ async fn run_control_reader( "received encode settings from client ({} bytes)", payload.len() ); + let mut payload = payload; + if let Some(request) = nesprotocol::decode_encode_settings_ext(&payload) + && request.codec == nesprotocol::CODEC_PYROWAVE + { + match pyrowave_request(request, &pyrowave) { + Ok(settings) => { + info!("client asked for PyroWave at {} kbps", settings.value); + payload.clear(); + nesprotocol::encode_settings(&mut payload, &settings); + } + Err(why) => { + warn!("refusing a PyroWave request: {why}"); + return; + } + } + } // A person set this by hand, so the controller stops // deciding until it is told otherwise. Overriding a // person's setting a second later would take away the only @@ -778,10 +889,12 @@ pub struct SessionManager { audio_bytes: AtomicU64, last_audio_bytes: AtomicU64, relay_ms: Arc, // latest relay latency (f32 bits) + /// The rate a PyroWave request that names none runs at. + pyrowave_kbps: u32, } impl SessionManager { - pub fn new() -> Self { + pub fn new(pyrowave_kbps: u32) -> Self { let (gamepad_tx, _) = tokio::sync::broadcast::channel(256); Self { sessions: Arc::new(Mutex::new(HashMap::new())), @@ -797,6 +910,7 @@ impl SessionManager { audio_bytes: AtomicU64::new(0), last_audio_bytes: AtomicU64::new(0), relay_ms: Arc::new(AtomicU32::new(0)), + pyrowave_kbps, } } @@ -825,6 +939,7 @@ impl SessionManager { self.relay_ms.clone(), idr_cmd_tx, controller, + self.pyrowave_kbps, ) }); session.attach(carrier, conn); @@ -885,7 +1000,10 @@ impl SessionManager { let Some(ts_bytes) = payload.get(2..6) else { return; }; - let ts_ms = u32::from_le_bytes(ts_bytes.try_into().unwrap()); + self.note_pipeline_ts(u32::from_le_bytes(ts_bytes.try_into().unwrap())); + } + + fn note_pipeline_ts(&self, ts_ms: u32) { let now_ms = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_millis() as u64) @@ -924,25 +1042,57 @@ impl SessionManager { self.video_delta_bytes .fetch_add(data.len() as u64, Ordering::Relaxed); } + let data = MediaItem::Coded(Bytes::from(data)); let sessions = self.sessions.lock().await; - if sessions.is_empty() { - return; - } for session in sessions.values() { session.send_video_frame(data.clone()); } } + /// Hand a PyroWave frame to every client. + /// + /// Counted with the delta bytes: every PyroWave frame is intra, and + /// counting them as keyframes would report sixty keyframes a second, which + /// describes nothing. + pub async fn broadcast_pyrowave(&self, frame: PyroFrame) { + self.note_pipeline_ts(frame.meta.ts_ms); + self.video_delta_bytes + .fetch_add(frame.data.len() as u64, Ordering::Relaxed); + let item = MediaItem::PyroWave(Arc::new(frame)); + let sessions = self.sessions.lock().await; + for session in sessions.values() { + session.send_video_frame(item.clone()); + } + } + + /// PyroWave frames skipped across every client since the last call. + pub async fn take_pyro_skipped(&self) -> u64 { + let sessions = self.sessions.lock().await; + sessions + .values() + .map(ClientSession::take_pyro_skipped) + .sum() + } + + /// The PyroWave packet size every streaming client can carry: the smallest + /// of theirs, since one encoder serves them all. `None` with no client + /// streaming. + pub async fn pyro_packet_size(&self) -> Option { + let sessions = self.sessions.lock().await; + sessions + .values() + .filter_map(ClientSession::pyro_packet_limit) + .min() + } + pub async fn broadcast_audio(&self, data: Vec) { // Counted before the early return, so the figure measures what neswire // delivered rather than what a client happened to be around for. A hub // with no client still knows whether audio is being produced. self.audio_bytes .fetch_add(data.len() as u64, Ordering::Relaxed); + let data = Bytes::from(data); let sessions = self.sessions.lock().await; - if sessions.is_empty() { - return; - } for session in sessions.values() { session.send_audio_packet(data.clone()); } @@ -1098,7 +1248,7 @@ mod pipeline_delay_tests { #[test] fn a_second_with_no_frames_reports_nothing() { - let mgr = SessionManager::new(); + let mgr = SessionManager::new(0); assert_eq!(mgr.pipeline_delays(), (0, 0, 0)); } @@ -1109,7 +1259,7 @@ mod pipeline_delay_tests { // same time on every frame must read as zero, not as that constant -- // which is what a plain subtraction would have reported, pinned at the // maximum the wire can carry. - let mgr = SessionManager::new(); + let mgr = SessionManager::new(0); for i in 0..60u32 { mgr.note_pipeline_delay(&payload(i * 16)); } @@ -1123,7 +1273,7 @@ mod pipeline_delay_tests { #[test] fn draining_means_each_answer_describes_one_second() { - let mgr = SessionManager::new(); + let mgr = SessionManager::new(0); mgr.note_pipeline_delay(&payload(0)); assert!(mgr.pipeline_delays().0 == 0); assert_eq!(mgr.pipeline_delays(), (0, 0, 0)); @@ -1131,7 +1281,7 @@ mod pipeline_delay_tests { #[test] fn a_truncated_payload_is_ignored_rather_than_misread() { - let mgr = SessionManager::new(); + let mgr = SessionManager::new(0); mgr.note_pipeline_delay(&[0u8, 0u8, 1u8]); assert_eq!(mgr.pipeline_delays(), (0, 0, 0)); } @@ -1213,3 +1363,67 @@ mod input_batch_tests { assert!(events.is_empty()); } } + +#[cfg(test)] +mod pyrowave_request_tests { + use super::{PyroWaveGate, pyrowave_request}; + use nesprotocol::pyrowave::{CHROMA_420, CHROMA_444}; + use nesprotocol::{CODEC_PYROWAVE, DEPTH_8, DEPTH_10, EncodeSettings, RC_CBR, RC_CQP}; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + + fn gate(ok: bool) -> PyroWaveGate { + PyroWaveGate { + ok: Arc::new(AtomicBool::new(ok)), + default_kbps: 300_000, + } + } + + fn request(value: u32) -> EncodeSettings { + EncodeSettings { + codec: CODEC_PYROWAVE, + rate_control: RC_CBR, + value, + bit_depth: None, + chroma: None, + packet_size: Some(9999), + } + } + + /// The opt-in is only as safe as this: a client that never said it + /// decodes the format must not get it by asking. + #[test] + fn a_client_without_the_format_is_refused() { + assert!(pyrowave_request(request(100_000), &gate(false)).is_err()); + } + + #[test] + fn no_rate_means_the_hubs_default() { + let s = pyrowave_request(request(0), &gate(true)).unwrap(); + assert_eq!(s.value, 300_000); + let s = pyrowave_request(request(150_000), &gate(true)).unwrap(); + assert_eq!(s.value, 150_000, "a named rate is not overruled"); + } + + /// The client cannot see the path, so whatever packet size it sent is not + /// passed on; the forwarder names one. + #[test] + fn the_clients_packet_size_is_not_trusted() { + let s = pyrowave_request(request(0), &gate(true)).unwrap(); + assert_eq!(s.packet_size, None); + assert_eq!((s.bit_depth, s.chroma), (Some(DEPTH_8), Some(CHROMA_420))); + + let mut r = request(0); + r.bit_depth = Some(DEPTH_10); + r.chroma = Some(CHROMA_444); + let s = pyrowave_request(r, &gate(true)).unwrap(); + assert_eq!((s.bit_depth, s.chroma), (Some(DEPTH_10), Some(CHROMA_444))); + } + + #[test] + fn constant_quality_is_not_a_pyrowave_mode() { + let mut r = request(0); + r.rate_control = RC_CQP; + assert!(pyrowave_request(r, &gate(true)).is_err()); + } +} diff --git a/apps/nesinit/src/services.rs b/apps/nesinit/src/services.rs index 59ae5a87..c2c61392 100644 --- a/apps/nesinit/src/services.rs +++ b/apps/nesinit/src/services.rs @@ -148,6 +148,7 @@ const READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15); /// them listens on and answer in its place — and the answer that matters is the /// address a client is told to connect to. See `ticket::Untrusted`. pub const SERVICE_UID: u32 = 1000; +pub const SERVICE_GID: u32 = 1000; /// Where audio's socket lives, for both the services and the workload. /// @@ -168,18 +169,6 @@ pub const SERVICE_UID: u32 = 1000; /// workload, which is the property [`crate::ticket::Untrusted`] depends on. pub const AUDIO_DIR: &str = "/run/pipewire"; -/// Where the PulseAudio protocol is served, in [`AUDIO_DIR`] for the reason -/// audio's own socket is. -/// -/// pipewire-pulse is told this path by its configuration file in the image, -/// which is not something this program can pass it, so the two are compared by -/// a test rather than trusted to agree. -pub const PULSE_SOCKET: &str = "/run/pipewire/pulse-native"; - -/// [`PULSE_SOCKET`] as a PulseAudio client is told it. -pub const PULSE_SERVER: &str = "unix:/run/pipewire/pulse-native"; -pub const SERVICE_GID: u32 = 1000; - /// Where a service's runtime sockets live. pub const RUNTIME_DIR: &str = "/run/user/1000"; @@ -292,26 +281,6 @@ pub const STACK: &[Service] = &[ umask: None, ready: None, }, - Service { - name: "pipewire-pulse", - argv: &["/usr/bin/pipewire-pulse"], - env: &[ - ("XDG_RUNTIME_DIR", RUNTIME_DIR), - ("PIPEWIRE_RUNTIME_DIR", AUDIO_DIR), - ("DBUS_SESSION_BUS_ADDRESS", "unix:path=/run/user/1000/bus"), - ], - user: Some((SERVICE_UID, SERVICE_GID)), - // Optional for the reason the sender is: everything that speaks - // PipeWire itself is unaffected, and a session with sound missing is - // degraded rather than unusable. - cost: "anything that only speaks PulseAudio plays silently, and Wine is one", - required: false, - // The workload is its client, and is not this user. - umask: Some(0), - // Nothing in this table connects to it, but the workload does, and the - // workload is started after the table. - ready: Some(PULSE_SOCKET), - }, Service { name: "neswire", argv: &["/usr/bin/neswire"], @@ -932,7 +901,7 @@ mod tests { /// owner -- and the workload is not the owner. #[test] fn the_audio_socket_is_reachable_by_a_user_who_does_not_own_it() { - for name in ["pipewire", "pipewire-pulse"] { + for name in ["pipewire"] { let service = STACK .iter() .find(|s| s.name == name) @@ -945,32 +914,6 @@ mod tests { } } - /// The PulseAudio socket's path is written three times: here, in the - /// client address the workload is given, and in pipewire-pulse's own - /// configuration in the image. If any of them moves on its own, the game - /// finds no server and plays silently, and nothing fails. - #[test] - fn pulse_is_served_where_the_workload_is_told_to_look() { - assert!( - PULSE_SOCKET.starts_with(AUDIO_DIR), - "the workload can only reach sockets in {AUDIO_DIR}" - ); - assert_eq!(PULSE_SERVER, format!("unix:{PULSE_SOCKET}")); - - let pulse = STACK - .iter() - .find(|s| s.name == "pipewire-pulse") - .expect("pulse is in the table"); - assert_eq!(pulse.ready, Some(PULSE_SOCKET)); - - let config = - include_str!("../../../build/etc/pipewire/pipewire-pulse.conf.d/50-nestri.conf"); - assert!( - config.contains(&format!("\"{PULSE_SERVER}\"")), - "pipewire-pulse is configured to listen somewhere other than {PULSE_SOCKET}" - ); - } - /// The two that failed on a real boot, and the reason each matters. #[test] fn every_service_has_somewhere_to_write() { diff --git a/apps/nesinit/src/workload.rs b/apps/nesinit/src/workload.rs index d260def9..a32845b4 100644 --- a/apps/nesinit/src/workload.rs +++ b/apps/nesinit/src/workload.rs @@ -303,9 +303,6 @@ const GRAPHICS: &[(&str, &str)] = &[ // reach and both are told where. Without this a game renders and plays // silently, having looked under its own uid and found nothing. ("PIPEWIRE_RUNTIME_DIR", crate::services::AUDIO_DIR), - // The same, for a client that speaks PulseAudio instead. It does not read - // the variable above, and its default is under its own runtime directory. - ("PULSE_SERVER", crate::services::PULSE_SERVER), ]; /// Mount one share where the descriptor says to put it. diff --git a/build/Containerfile b/build/Containerfile index ae939881..9f316440 100644 --- a/build/Containerfile +++ b/build/Containerfile @@ -31,7 +31,7 @@ # scoped to one stage instead, `FROM ${PROTON_IMAGE}` expands to nothing, and # the build fails with "no FROM statement found" — which says nothing about # the actual mistake. See the Proton stage below for what this is. -ARG PROTON_IMAGE=ghcr.io/nestrilabs/proton-ge:11-7 +ARG PROTON_IMAGE=ghcr.io/nestrilabs/proton-cachyos-native-wow64:11.0-20260703 # ─────────────────────────────────────────────────────────── # initial / builder @@ -40,7 +40,7 @@ ARG PROTON_IMAGE=ghcr.io/nestrilabs/proton-ge:11-7 # guest was Artix, chosen for an init system this image no longer contains. # Only build artifacts leave these stages. # ─────────────────────────────────────────────────────────── -FROM docker.io/archlinux:base-devel AS initial +FROM docker.io/cachyos/cachyos-v3:latest AS initial RUN pacman -Syu --noconfirm FROM initial AS builder @@ -52,8 +52,8 @@ RUN pacman -S --noconfirm --needed \ libx11 libxext libxrandr libxshmfence libxfixes libxxf86vm libxcb \ xcb-util-keysyms xorgproto \ wayland wayland-protocols \ - expat zlib zstd libxml2 lm_sensors \ - llvm clang libclc spirv-tools spirv-llvm-translator glslang \ + expat zlib-ng-compat zstd libxml2 lm_sensors \ + llvm clang libclc spirv-tools spirv-llvm-translator glslang shader-slang \ elfutils libva libdrm directx-headers \ rust rust-bindgen cbindgen \ curl openssl \ @@ -131,6 +131,7 @@ FROM builder AS nestri-src WORKDIR /build/nestri COPY Cargo.toml Cargo.lock ./ COPY crates/nesprotocol crates/nesprotocol +COPY crates/nespyro crates/nespyro # Not a guest component and not installed below — it runs on a stranger's own # machine. It is here because `cargo` loads every workspace member's manifest # before it builds anything, so a member missing from the context fails the @@ -209,7 +210,7 @@ FROM ${PROTON_IMAGE} AS proton # obscurely, and a previous attempt at this image lost two services to exactly # that. # ─────────────────────────────────────────────────────────── -FROM docker.io/archlinux:base AS os-base +FROM docker.io/cachyos/cachyos-v3:latest AS os-base # The base image ships an unsigned local keyring, so upgrading # `archlinux-keyring` runs a hook that reports `There is no secret key @@ -264,12 +265,12 @@ RUN pacman -Syu --noconfirm --needed \ dbus \ iptables iproute2 \ libglvnd libdrm libepoxy libxxf86vm libinput wayland \ - expat zlib llvm-libs lm_sensors elfutils libva shaderc vulkan-icd-loader \ + expat zlib-ng-compat llvm-libs lm_sensors elfutils libva shaderc vulkan-icd-loader \ pixman libxkbcommon xcb-util-keysyms xorg-xwayland \ - pipewire pipewire-audio pipewire-pulse libpulse wireplumber opus \ + pipewire pipewire-audio wireplumber opus \ python libunwind sdl2-compat \ && rm -f /usr/share/libalpm/hooks/dbus-reload.hook \ - && pacman -Rdd --noconfirm systemd systemd-sysvcompat \ + && pacman -Rdd --noconfirm systemd systemd-sysvcompat openssh \ && pacman -Scc --noconfirm # `libunwind` is Wine's, not ours. `ntdll.so` links it, so without it every @@ -278,15 +279,6 @@ RUN pacman -Syu --noconfirm --needed \ # after the prefix had already been created -- so the session got past every # check that Proton was present and usable. # -# `pipewire-pulse` and `libpulse` are Wine's audio. Its driver speaks only the -# PulseAudio protocol, so without the translation layer a game starts, renders -# and plays silently, with nothing failing to say why. `libpulse` is what that -# driver, `winepulse.so`, links, and it would arrive as a dependency of -# pipewire-pulse anyway. It is named because the dependency is Wine's and not -# pipewire-pulse's, and a library that only arrives as somebody else's -# dependency goes away when they do. The check below covers `winepulse.so`, so a -# missing library fails the build rather than the session. - # `python` is not a build dependency here -- the builder stage has its own for # Mesa -- it is a *runtime* one. The compatibility tool's entry point is a # Python script, so a box without an interpreter starts a game and the launch @@ -508,7 +500,7 @@ RUN failed=0; \ : > /tmp/missing-libs; \ for f in /usr/bin/nesinit /usr/bin/nescope /usr/bin/neshub /usr/bin/neswire \ /usr/bin/nesgamepad /usr/lib/libnescapture_layer.so /usr/lib/libSDL2-2.0.so.0 \ - /usr/bin/dbus-daemon /usr/bin/pipewire /usr/bin/pipewire-pulse \ + /usr/bin/dbus-daemon /usr/bin/pipewire \ /usr/bin/wireplumber /usr/bin/ip \ /usr/lib/libgallium-*.so /usr/lib/libEGL_mesa.so.0 \ /usr/lib/libvulkan_*.so /usr/lib/dri/*.so /usr/lib/gbm/*.so \ @@ -531,7 +523,7 @@ RUN failed=0; \ RUN for required in /usr/bin/nesinit /usr/bin/nescope /usr/bin/neshub /usr/bin/neswire \ /usr/bin/nesgamepad /usr/lib/libSDL2-2.0.so.0 \ - /usr/bin/dbus-daemon /usr/bin/pipewire /usr/bin/pipewire-pulse \ + /usr/bin/dbus-daemon /usr/bin/pipewire \ /usr/bin/wireplumber /usr/bin/ip /usr/bin/python3 \ /usr/share/steam/compatibilitytools.d/proton-cachyos/proton; do \ test -x "$required" || { echo "the image is missing $required" >&2; exit 1; }; \ diff --git a/build/etc/pipewire/pipewire-pulse.conf.d/50-nestri.conf b/build/etc/pipewire/pipewire-pulse.conf.d/50-nestri.conf deleted file mode 100644 index 98917a0c..00000000 --- a/build/etc/pipewire/pipewire-pulse.conf.d/50-nestri.conf +++ /dev/null @@ -1,66 +0,0 @@ -# pipewire-pulse, the PulseAudio protocol on top of PipeWire, for clients that -# speak nothing else. Wine's audio driver is one of them. -# -# The socket moves out of the service user's runtime directory. That directory -# is 0700 and the workload runs as a different user, so a client there finds -# nothing and plays silently, which is the failure PipeWire's own socket was -# moved for. The workload is pointed here with PULSE_SERVER, and nesinit starts -# this under a umask that lets a user who does not own the socket open it. This -# path is written in nesinit too, and a test there compares the two. -# -# A drop-in replaces a key of pulse.properties rather than appending to it, so -# the default "unix:native" socket is not created as well. -pulse.properties = { - server.address = [ "unix:/run/pipewire/pulse-native" ] - - # module-always-sink creates a "Dummy Output" sink whenever no other sink - # exists, and while neswire's sink is not up yet that is always. A second - # sink is a second default-node candidate, and a game that picks it plays - # into nothing. neswire has to be the only sink; see the WirePlumber - # drop-in. - pulse.cmd.always-sink = false - - # The root is read-only and every box starts from the same image, so there - # is nothing to restore and nowhere to save, which is the reasoning behind - # turning off WirePlumber's state hooks. Left on, these are one more place - # that could pin a default. - pulse.cmd.device-manager = false - pulse.cmd.device-restore = false - pulse.cmd.stream-restore = false - - # Wine sizes its whole audio path from this number. At startup its driver - # opens a probe stream asking for a 1-frame request size, reads back what - # the server grants, which is this floor, and makes the device period ten - # times that. Streams then queue three periods. The default floor of 256 - # frames therefore becomes a 53 ms period and about 160 ms of audio in - # flight, which is desktop playback and not a game being streamed. - # - # 48 frames (1 ms) makes that period 10 ms, the default period Windows - # itself uses, so a game gets the buffering it was written for. A stream - # then asks for 480-frame requests, which is what is granted, so the floor - # affects the probe and nothing else Wine does. - pulse.min.req = 48/48000 - - # What a client gets when it asks for no particular buffering. The shipped - # default is two seconds, which is only harmless where nobody is listening - # live. 10 ms requests and 40 ms of target buffer are enough for a client - # that does not say. - pulse.default.req = 480/48000 - pulse.default.tlength = 1920/48000 -} - -# pipewire-pulse.conf raises the minimum node latency to 1024 frames (21.3 ms) -# inside a virtual machine, which a box is. That puts every Pulse client, and -# therefore Wine, at least 21 ms behind, whatever it asked for. Floor it where -# the graph itself is floored instead; see the pipewire.conf drop-in for why -# 256 frames. A drop-in appends to this array, and every matching rule applies -# in order, so this runs after the shipped rule and wins. -pulse.properties.rules = [ - { matches = [ { cpu.vm.name = !null } ] - actions = { - update-props = { - pulse.min.quantum = 256/48000 - } - } - } -] diff --git a/build/etc/wireplumber/wireplumber.conf.d/50-nestri.conf b/build/etc/wireplumber/wireplumber.conf.d/50-nestri.conf index c5fcf4d5..8f0068d0 100644 --- a/build/etc/wireplumber/wireplumber.conf.d/50-nestri.conf +++ b/build/etc/wireplumber/wireplumber.conf.d/50-nestri.conf @@ -3,14 +3,6 @@ # The goal is that `neswire` is the *only* Audio/Sink in the graph, so # default-node selection has exactly one candidate and cannot pick wrong. # -# The "Dummy Output" (`auto_null`) sink is not disabled here, because it does -# not come from this side. pipewire-pulse creates it with `module-always-sink`, -# and the pipewire-pulse drop-in turns that off with -# `pulse.cmd.always-sink = false`. WirePlumber ships scripts/fallback-sink.lua, -# which creates a node by the same name, but no component in wireplumber.conf -# references it, so it never loads. If auto_null appears, look at the -# pipewire-pulse drop-in first. -# # pipewire.conf's Dummy-Driver / Freewheel-Driver are a third thing again: # support.node.driver objects with no ports -- drivers, not sinks, so nothing # can play into them. Dummy-Driver is load-bearing here. With no ALSA hardware diff --git a/crates/nesprotocol/src/datagram.rs b/crates/nesprotocol/src/datagram.rs index 165d7283..03409aee 100644 --- a/crates/nesprotocol/src/datagram.rs +++ b/crates/nesprotocol/src/datagram.rs @@ -40,6 +40,9 @@ pub const DGRAM_HDR_LEN: usize = 7; /// so a frame's identity does not change with the transport under it. pub const DGRAM_VIDEO: u8 = 0; pub const DGRAM_AUDIO: u8 = 1; +/// PyroWave, whose datagrams have a header of their own: see +/// [`crate::pyrowave`]. +pub const DGRAM_PYROWAVE: u8 = crate::pyrowave::DGRAM_PYROWAVE; /// Bounds reassembly work. A 4K keyframe runs to a few hundred fragments at a /// typical MTU; anything past this did not come from a sender of ours. diff --git a/crates/nesprotocol/src/lib.rs b/crates/nesprotocol/src/lib.rs index ed3a3597..a0176bf1 100644 --- a/crates/nesprotocol/src/lib.rs +++ b/crates/nesprotocol/src/lib.rs @@ -8,6 +8,7 @@ pub mod gamepad; pub mod input; #[cfg(feature = "lifecycle")] pub mod lifecycle; +pub mod pyrowave; pub mod reliable; pub mod stats; @@ -161,6 +162,9 @@ pub const STREAM_STATS: u8 = 4; /// One keyframe, on a stream of its own, closed after it. The exception to /// "media travels as datagrams" — see the [`reliable`] module for why. pub const STREAM_KEYFRAME: u8 = 5; +/// One message of a PyroWave frame, on the nescapture→hub video socket only. +/// A frame is several of these: see the [`pyrowave`] module. +pub const STREAM_PYROWAVE: u8 = 7; // ── Bidi stream types (desktop ↔ hub over QUIC bidi) ─────────── @@ -178,6 +182,9 @@ pub const CODEC_H264: u8 = 0; pub const CODEC_H265: u8 = 1; pub const CODEC_AV1: u8 = 2; pub const CODEC_OPUS: u8 = 3; +/// The LAN codec. Advertised like the others but never chosen by +/// negotiation: see [`pyrowave`]. +pub const CODEC_PYROWAVE: u8 = 4; pub const CODEC_KEEP: u8 = 0xFF; // "keep current" sentinel for dynamic encoder settings // Rate control modes (encode settings) @@ -277,6 +284,7 @@ pub fn codec_name(codec: u8) -> &'static str { CODEC_H265 => "h265", CODEC_AV1 => "av1", CODEC_OPUS => "opus", + CODEC_PYROWAVE => "pyrowave", _ => "unknown", } } @@ -418,7 +426,7 @@ impl ClientCaps { // `CODEC_KEEP` and the audio codec have no place in a video capability // set, and shifting by them would be nonsense rather than a small // error. - if !matches!(codec, CODEC_H264 | CODEC_H265 | CODEC_AV1) { + if !matches!(codec, CODEC_H264 | CODEC_H265 | CODEC_AV1 | CODEC_PYROWAVE) { return None; } if !matches!(depth, DEPTH_8 | DEPTH_10) { @@ -456,6 +464,10 @@ impl ClientCaps { /// `None` when nothing overlaps, which is a real possibility rather than a /// theoretical one: an old client that sends no capabilities at all reads /// as empty. The caller keeps whatever it was already doing. + /// + /// Never PyroWave, whatever both ends support: it is not in + /// [`CODEC_PREFERENCE`]. It saturates anything but a LAN, and nothing here + /// can tell a LAN apart, so only a client asking for it gets it. pub fn best(self, host: Self) -> Option<(u8, u8)> { if self.is_empty() || host.is_empty() { return None; @@ -471,9 +483,22 @@ impl ClientCaps { } } -/// Encode a capability set: `[2B bits LE]`. +/// Encode a capability set: `[2B bits LE]`, then `[1B PYROWAVE_FORMAT]` when +/// the set includes PyroWave. +/// +/// The format byte is appended rather than folded into the bits because an +/// older hub reads exactly two bytes and ignores the rest, so a newer client +/// still says what it decodes to it. pub fn encode_client_caps(buf: &mut Vec, caps: ClientCaps) { buf.extend_from_slice(&caps.bits().to_le_bytes()); + if caps.supports_codec(CODEC_PYROWAVE) { + buf.push(pyrowave::PYROWAVE_FORMAT); + } +} + +/// The PyroWave bitstream format a capability payload declares, if any. +pub fn decode_pyrowave_format(payload: &[u8]) -> Option { + payload.get(2).copied() } /// Decode a capability set. `None` when the payload is too short to be one. @@ -498,6 +523,52 @@ pub fn encode_depth_only(buf: &mut Vec, bit_depth: u8) { encode_encode_settings(buf, CODEC_KEEP, RC_KEEP, 0, bit_depth); } +/// An encode-settings payload with every field it can carry. +/// +/// The fields past `bit_depth` were appended for PyroWave and are `None` when +/// the payload stops before them, which every payload from before them does. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct EncodeSettings { + pub codec: u8, + pub rate_control: u8, + pub value: u32, + pub bit_depth: Option, + /// One of the `pyrowave::CHROMA_*` constants. + pub chroma: Option, + /// The largest PyroWave packet the hub can carry in one datagram. A client + /// sends 0; the hub fills it in, because only the hub can see the path. + pub packet_size: Option, +} + +/// `[codec][rc][value u32 LE][depth][chroma][packet_size u16 LE]`. +pub fn encode_settings(buf: &mut Vec, s: &EncodeSettings) { + buf.push(s.codec); + buf.push(s.rate_control); + buf.extend_from_slice(&s.value.to_le_bytes()); + // Each optional field is written only if everything before it is, since a + // reader finds a field by its offset. + let Some(depth) = s.bit_depth else { return }; + buf.push(depth); + let Some(chroma) = s.chroma else { return }; + buf.push(chroma); + let Some(size) = s.packet_size else { return }; + buf.extend_from_slice(&size.to_le_bytes()); +} + +/// Decode every field of an encode-settings payload. `None` when it is too +/// short to be one. +pub fn decode_encode_settings_ext(payload: &[u8]) -> Option { + let (codec, rate_control, value, bit_depth) = decode_encode_settings(payload)?; + Some(EncodeSettings { + codec, + rate_control, + value, + bit_depth, + chroma: payload.get(7).copied(), + packet_size: payload.get(8..10).map(|b| u16::from_le_bytes([b[0], b[1]])), + }) +} + /// Decode an encode-settings payload. Returns `(codec_id, rate_control_mode, value, bit_depth)`. /// bit_depth is None for 6-byte (old client) payloads, Some(n) for 7+ byte payloads. pub fn decode_encode_settings(payload: &[u8]) -> Option<(u8, u8, u32, Option)> { @@ -949,4 +1020,136 @@ mod media_control_tests { } } } + + /// The negative the whole opt-in design rests on: however much both ends + /// support it, negotiation never lands a session on PyroWave. + #[test] + fn negotiation_never_picks_pyrowave() { + let both = host_all() + .with(CODEC_PYROWAVE, DEPTH_8) + .with(CODEC_PYROWAVE, DEPTH_10); + assert_eq!(both.best(both), Some((CODEC_AV1, DEPTH_10))); + + let only = ClientCaps::empty() + .with(CODEC_PYROWAVE, DEPTH_8) + .with(CODEC_PYROWAVE, DEPTH_10); + assert_eq!(only.best(only), None, "PyroWave alone is no answer either"); + assert!(!CODEC_PREFERENCE.contains(&CODEC_PYROWAVE)); + } + + #[test] + fn pyrowave_takes_bits_of_its_own() { + let caps = ClientCaps::empty().with(CODEC_PYROWAVE, DEPTH_10); + assert!(caps.supports(CODEC_PYROWAVE, DEPTH_10)); + assert!(!caps.supports(CODEC_PYROWAVE, DEPTH_8)); + assert_eq!(caps.bits(), 1 << 9); + for codec in [CODEC_H264, CODEC_H265, CODEC_AV1] { + assert!(!caps.supports_codec(codec)); + } + } + + /// The format byte rides behind the two an older hub reads, so it changes + /// nothing for one. + #[test] + fn caps_carry_the_pyrowave_format_only_when_pyrowave_is_there() { + let mut buf = Vec::new(); + encode_client_caps(&mut buf, host_all()); + assert_eq!(buf.len(), 2); + assert_eq!(decode_pyrowave_format(&buf), None); + + let caps = host_all().with(CODEC_PYROWAVE, DEPTH_8); + let mut buf = Vec::new(); + encode_client_caps(&mut buf, caps); + assert_eq!(decode_client_caps(&buf), Some(caps)); + assert_eq!( + decode_pyrowave_format(&buf), + Some(pyrowave::PYROWAVE_FORMAT) + ); + } + + #[test] + fn extended_settings_survive_the_wire() { + let full = EncodeSettings { + codec: CODEC_PYROWAVE, + rate_control: RC_CBR, + value: 300_000, + bit_depth: Some(DEPTH_10), + chroma: Some(pyrowave::CHROMA_444), + packet_size: Some(1440), + }; + let mut buf = Vec::new(); + encode_settings(&mut buf, &full); + assert_eq!(buf.len(), 10); + assert_eq!(decode_encode_settings_ext(&buf), Some(full)); + // And an older reader still sees the first four fields. + assert_eq!( + decode_encode_settings(&buf), + Some((CODEC_PYROWAVE, RC_CBR, 300_000, Some(DEPTH_10))) + ); + } + + /// A field is found by its offset, so one can only be written if all + /// those before it are. + #[test] + fn a_missing_field_ends_the_payload() { + let mut s = EncodeSettings { + codec: CODEC_PYROWAVE, + rate_control: RC_CBR, + value: 1, + bit_depth: Some(DEPTH_8), + chroma: None, + packet_size: Some(1200), + }; + let mut buf = Vec::new(); + encode_settings(&mut buf, &s); + assert_eq!(buf.len(), 7, "a packet size with no chroma before it"); + s.packet_size = None; + assert_eq!(decode_encode_settings_ext(&buf), Some(s)); + + let mut buf = Vec::new(); + encode_bitrate_only(&mut buf, 5); + let s = decode_encode_settings_ext(&buf).unwrap(); + assert_eq!((s.bit_depth, s.chroma, s.packet_size), (None, None, None)); + } + + #[test] + fn every_stream_type_is_its_own_number() { + // Stream and bidi types share the byte that opens a stream. + let types = [ + ("STREAM_VIDEO", STREAM_VIDEO), + ("STREAM_AUDIO", STREAM_AUDIO), + ("BIDI_INPUT", BIDI_INPUT), + ("STREAM_CURSOR", STREAM_CURSOR), + ("STREAM_STATS", STREAM_STATS), + ("STREAM_KEYFRAME", STREAM_KEYFRAME), + ("BIDI_CONTROL", BIDI_CONTROL), + ("STREAM_PYROWAVE", STREAM_PYROWAVE), + ]; + for (i, (name, value)) in types.iter().enumerate() { + for (other_name, other_value) in &types[i + 1..] { + assert_ne!( + value, other_value, + "{name} and {other_name} are both {value}" + ); + } + } + } + + #[test] + fn every_codec_is_its_own_number() { + let codecs = [ + CODEC_H264, + CODEC_H265, + CODEC_AV1, + CODEC_OPUS, + CODEC_PYROWAVE, + CODEC_KEEP, + ]; + for (i, a) in codecs.iter().enumerate() { + for b in &codecs[i + 1..] { + assert_ne!(a, b); + } + } + assert_eq!(codec_name(CODEC_PYROWAVE), "pyrowave"); + } } diff --git a/crates/nesprotocol/src/pyrowave.rs b/crates/nesprotocol/src/pyrowave.rs new file mode 100644 index 00000000..36bc97a5 --- /dev/null +++ b/crates/nesprotocol/src/pyrowave.rs @@ -0,0 +1,1499 @@ +//! Carrying PyroWave, the LAN codec. +//! +//! PyroWave frames are intra-only and already cut into packets that each +//! decode on their own, and its decoder shows a frame that is missing some of +//! them. Everything the hardware-codec path does to survive the internet works +//! against that: byte-slicing a frame across datagrams makes every packet +//! depend on its neighbours' fragments, a reassembler that needs every +//! fragment throws partial frames away, and keyframe streams and resync gates +//! exist for prediction chains PyroWave does not have. +//! +//! So PyroWave gets a path of its own, and this module is all of its logic: +//! +//! - **nescapture → hub:** a frame is a run of IPC messages of whole packets +//! ([`chunk_frame`]), because one unix datagram cannot hold a frame — the +//! kernel caps it near 416 KB by default and a frame at 1 Gbit/s is 2 MB. +//! [`FrameAssembler`] puts them back together at the hub. +//! - **hub → client:** one packet per datagram ([`pack_datagrams`]), never a +//! slice across two unless a single block is larger than a datagram. The +//! critical packets go out first and again last, so a burst that eats one +//! copy usually spares the other. +//! - **client:** [`PyroCollector`] gathers a frame's datagrams and lets it go +//! whole, or partial once it is plainly not going to complete. Whether a +//! partial frame is worth decoding is the decoder's question, not this one's. +//! +//! None of it does I/O or reads a clock: time is passed in, so every rule here +//! is tested as a rule. +//! +//! # Selection is opt-in +//! +//! PyroWave runs at hundreds of Mbit/s. Nothing on this side can tell a LAN +//! from a fast WAN, so negotiation never picks it ([`crate::ClientCaps::best`] +//! leaves it out); a session uses it only when the client asks. + +use std::collections::VecDeque; +use std::ops::Range; + +use crate::datagram::{MIN_DGRAM_PAYLOAD, seq_older_than}; +use crate::{CODEC_PYROWAVE, DecodedIpcFrame, STREAM_PYROWAVE, encode_ipc_frame}; + +/// The PyroWave bitstream this build speaks. +/// +/// The bitstream carries no version of its own and upstream marks it draft, so +/// both ends state one here. 1 is upstream `89f7e47`, as frozen by nespyro. +/// Bump it whenever nespyro's bitstream changes; a client with another value is +/// never sent PyroWave. +pub const PYROWAVE_FORMAT: u8 = 1; + +/// Chroma format, in encode settings. +pub const CHROMA_420: u8 = 0; +pub const CHROMA_444: u8 = 1; + +/// The largest single PyroWave block, in bytes: nespyro's +/// `LARGEST_BLOCK_WORDS * 4`. A packet is at most this or the packet size, +/// whichever is larger, since the packetizer never splits a block. +pub const LARGEST_BLOCK_BYTES: usize = 2488; + +/// Most data one IPC message carries, header included. +/// +/// Far below the kernel's unix datagram ceiling, so nothing has to be tuned on +/// the box for it, and large enough that a 2 MB frame is a few dozen sends. +pub const IPC_PYROWAVE_MAX: usize = 64 * 1024; + +/// `[frame u32][msg index u16][msg count u16][critical u16][packet count u16]`. +pub const IPC_PYRO_HDR_LEN: usize = 12; + +/// A frame still incomplete this long after its first message is abandoned. +/// The messages of one frame are written back to back, so a gap this long +/// means the sender stopped halfway. +pub const ABANDON_MS: u64 = 100; + +/// Datagram kind, beside `DGRAM_VIDEO` and `DGRAM_AUDIO`. +pub const DGRAM_PYROWAVE: u8 = 2; + +/// `[kind][seq u16][index u16][total u16][critical u16][ts_ms u32][flags]`. +pub const PYRO_DGRAM_HDR_LEN: usize = 14; + +/// Bound on datagrams per frame, about 19 MB at a typical MTU. Past it a +/// frame did not come from a sender of ours, and the bound keeps a corrupt +/// `total` from sizing an allocation. +pub const PYRO_MAX_DATAGRAMS: usize = 16384; + +/// A second copy of a critical datagram, under the same index. +pub const FLAG_DUPLICATE: u8 = 1 << 0; +/// A piece of one packet too large for a datagram. Pieces take consecutive +/// indexes, from one marked [`FLAG_FIRST_PART`] to one marked +/// [`FLAG_LAST_PART`]. +pub const FLAG_PART: u8 = 1 << 1; +pub const FLAG_FIRST_PART: u8 = 1 << 2; +pub const FLAG_LAST_PART: u8 = 1 << 3; +const KNOWN_FLAGS: u8 = FLAG_DUPLICATE | FLAG_PART | FLAG_FIRST_PART | FLAG_LAST_PART; + +/// How long a frame waits for stragglers once a newer frame has begun. +/// +/// On a LAN, datagrams of one frame do not arrive after the next frame's +/// unless they are lost. A little slack covers reordering inside a NIC queue. +pub const GRACE_US: u64 = 2_000; + +/// How long a frame waits once its datagrams stop arriving. +/// +/// Measured from the *last* datagram, not the first. A frame on a saturated +/// path takes about a frame interval to arrive by definition — the rate is a +/// frame per interval — so a deadline from its first datagram tears every +/// frame the moment the path is slower than the deadline assumes. Silence +/// says a frame is done arriving whatever the rate; on a LAN its datagrams are +/// microseconds apart while it is still coming. +pub const IDLE_US: u64 = 20_000; + +/// Frames in collection at once. More means the receiver is far behind, and +/// the oldest is let go to make room. +pub const MAX_IN_FLIGHT: usize = 8; + +// ── IPC ───────────────────────────────────────────────────────────────── + +/// The per-frame fields of the IPC header. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct FrameMeta { + pub ts_ms: u32, + pub width: u16, + pub height: u16, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ChunkError { + /// A frame always has at least its sequence header. + Empty, + /// More critical packets than packets. + Critical { critical: usize, packets: usize }, + /// A packet that cannot fit one message, which no packetizer of ours makes. + PacketTooLarge { index: usize, len: usize }, + /// More messages than the header can count. + TooLarge { messages: usize }, +} + +impl std::fmt::Display for ChunkError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Empty => write!(f, "frame has no packets"), + Self::Critical { critical, packets } => { + write!(f, "{critical} critical packets of {packets}") + } + Self::PacketTooLarge { index, len } => write!(f, "packet {index} is {len} bytes"), + Self::TooLarge { messages } => write!(f, "frame needs {messages} messages"), + } + } +} + +impl std::error::Error for ChunkError {} + +/// Room for packets in one message, after both headers. +const IPC_PACKET_ROOM: usize = IPC_PYROWAVE_MAX - crate::IPC_HEADER_LEN - IPC_PYRO_HDR_LEN; + +/// Cut a frame into IPC messages, each a complete datagram for the video +/// socket. +/// +/// `packets` are ranges of `data` in frame order; the first `critical` of them +/// are the ones whose loss the decoder cannot mask. +pub fn chunk_frame( + frame: u32, + meta: FrameMeta, + data: &[u8], + packets: &[Range], + critical: usize, +) -> Result>, ChunkError> { + if packets.is_empty() { + return Err(ChunkError::Empty); + } + if critical > packets.len() || critical > usize::from(u16::MAX) { + return Err(ChunkError::Critical { + critical, + packets: packets.len(), + }); + } + + // Decide the cuts first: every message names how many there are. + let mut cuts = Vec::new(); + let mut used = 0usize; + for (index, p) in packets.iter().enumerate() { + let need = 2 + p.len(); + if p.len() > usize::from(u16::MAX) || need > IPC_PACKET_ROOM { + return Err(ChunkError::PacketTooLarge { + index, + len: p.len(), + }); + } + if used + need > IPC_PACKET_ROOM { + cuts.push(index); + used = 0; + } + used += need; + } + let count = cuts.len() + 1; + if count > usize::from(u16::MAX) { + return Err(ChunkError::TooLarge { messages: count }); + } + + let bounds = std::iter::once(0) + .chain(cuts.iter().copied()) + .zip(cuts.iter().copied().chain(std::iter::once(packets.len()))); + let mut messages = Vec::with_capacity(count); + let mut body = Vec::with_capacity(IPC_PYROWAVE_MAX); + for (index, (start, end)) in bounds.enumerate() { + body.clear(); + body.extend_from_slice(&frame.to_le_bytes()); + body.extend_from_slice(&(index as u16).to_le_bytes()); + body.extend_from_slice(&(count as u16).to_le_bytes()); + body.extend_from_slice(&(critical as u16).to_le_bytes()); + body.extend_from_slice(&((end - start) as u16).to_le_bytes()); + for p in &packets[start..end] { + body.extend_from_slice(&(p.len() as u16).to_le_bytes()); + body.extend_from_slice(&data[p.clone()]); + } + messages.push(encode_ipc_frame( + STREAM_PYROWAVE, + CODEC_PYROWAVE, + 0, + meta.ts_ms, + meta.width, + meta.height, + &body, + )); + } + Ok(messages) +} + +/// A whole frame, as the hub has it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PyroFrame { + /// The sender's frame counter. + pub frame: u32, + pub meta: FrameMeta, + /// Leading packets whose loss cannot be masked. + pub critical: usize, + /// Every packet, back to back. + pub data: Vec, + pub packets: Vec>, +} + +impl PyroFrame { + pub fn packet(&self, index: usize) -> &[u8] { + &self.data[self.packets[index].clone()] + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AssembleError { + /// Not a PyroWave message at all. + NotPyroWave, + /// A PyroWave message that does not parse. + Malformed(&'static str), +} + +impl std::fmt::Display for AssembleError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::NotPyroWave => write!(f, "not a PyroWave message"), + Self::Malformed(what) => write!(f, "malformed PyroWave message: {what}"), + } + } +} + +impl std::error::Error for AssembleError {} + +struct Partial { + frame: u32, + meta: FrameMeta, + count: u16, + next: u16, + critical: u16, + started_ms: u64, + data: Vec, + packets: Vec>, +} + +/// Puts IPC messages back into frames. +/// +/// Messages must arrive in order, which a unix datagram socket with one +/// sender guarantees. A gap therefore means one was dropped — the sender's +/// write timed out — and the frame is abandoned rather than patched: the hub +/// has no business inventing a partial frame the encoder never finished +/// sending. Packets land straight in the frame's buffer, so a frame costs one +/// copy here. +#[derive(Default)] +pub struct FrameAssembler { + partial: Option, + abandoned: u64, + /// The frame most recently counted as abandoned, so a frame dropped + /// mid-way counts once and not once per remaining message. + last_abandoned: Option, +} + +impl FrameAssembler { + pub fn new() -> Self { + Self::default() + } + + /// Take one message. Returns the frame it completed, if it did. + pub fn push( + &mut self, + ipc: &DecodedIpcFrame<'_>, + now_ms: u64, + ) -> Result, AssembleError> { + if ipc.stream_type != STREAM_PYROWAVE || ipc.codec != CODEC_PYROWAVE { + return Err(AssembleError::NotPyroWave); + } + let d = ipc.data; + if d.len() < IPC_PYRO_HDR_LEN { + return Err(AssembleError::Malformed("short header")); + } + let u16_at = |o: usize| u16::from_le_bytes([d[o], d[o + 1]]); + let frame = u32::from_le_bytes([d[0], d[1], d[2], d[3]]); + let index = u16_at(4); + let count = u16_at(6); + let critical = u16_at(8); + let n = usize::from(u16_at(10)); + if count == 0 || index >= count { + return Err(AssembleError::Malformed("message index outside its count")); + } + + // Every packet is located before anything is kept, so a message that + // fails halfway leaves no half of itself behind. + let mut ranges = Vec::with_capacity(n); + let mut at = IPC_PYRO_HDR_LEN; + for _ in 0..n { + let len_bytes = d + .get(at..at + 2) + .ok_or(AssembleError::Malformed("truncated packet length"))?; + let len = usize::from(u16::from_le_bytes([len_bytes[0], len_bytes[1]])); + at += 2; + if at + len > d.len() { + return Err(AssembleError::Malformed("truncated packet")); + } + ranges.push(at..at + len); + at += len; + } + if at != d.len() { + return Err(AssembleError::Malformed("bytes after the last packet")); + } + + let meta = FrameMeta { + ts_ms: ipc.timestamp_ms, + width: ipc.width, + height: ipc.height, + }; + let continues = self.partial.as_ref().is_some_and(|p| { + p.frame == frame + && p.next == index + && p.count == count + && p.critical == critical + && p.meta == meta + && now_ms.saturating_sub(p.started_ms) <= ABANDON_MS + }); + if !continues { + if let Some(p) = self.partial.take() { + self.abandon(p.frame); + } + if index != 0 { + // The middle of a frame whose start is gone. + self.abandon(frame); + return Ok(None); + } + self.partial = Some(Partial { + frame, + meta, + count, + next: 0, + critical, + started_ms: now_ms, + data: Vec::new(), + packets: Vec::new(), + }); + } + + let p = self.partial.as_mut().expect("set above"); + for r in ranges { + let start = p.data.len(); + p.data.extend_from_slice(&d[r]); + p.packets.push(start..p.data.len()); + } + p.next += 1; + if p.next < p.count { + return Ok(None); + } + + let p = self.partial.take().expect("set above"); + if usize::from(p.critical) > p.packets.len() || p.packets.is_empty() { + self.abandon(p.frame); + return Err(AssembleError::Malformed( + "more critical packets than packets", + )); + } + Ok(Some(PyroFrame { + frame: p.frame, + meta: p.meta, + critical: usize::from(p.critical), + data: p.data, + packets: p.packets, + })) + } + + fn abandon(&mut self, frame: u32) { + if self.last_abandoned != Some(frame) { + self.abandoned += 1; + self.last_abandoned = Some(frame); + } + } + + /// Frames abandoned since the last call. + pub fn take_abandoned(&mut self) -> u64 { + std::mem::take(&mut self.abandoned) + } +} + +// ── Datagrams ─────────────────────────────────────────────────────────── + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PyroDatagramHeader { + /// The frame, counted by the hub per connection. + pub seq: u16, + /// This datagram's place in the frame, `0..total`. + pub index: u16, + /// Datagrams in the frame, not counting duplicates. + pub total: u16, + /// How many leading indexes carry the critical packets. + /// + /// On the wire because nothing else can say it. A receiver cannot tell + /// whether the coarse bands arrived whole from the blocks it has: a block + /// the encoder quantized to nothing is never sent, and its absence looks + /// exactly like a loss. Coarse high-pass bands of noisy content do that + /// routinely, so a depacketizer counting blocks declares the coarse bands + /// broken on frames that arrived complete. The sender knows which + /// datagrams those bands went in. + pub critical: u16, + pub ts_ms: u32, + pub flags: u8, +} + +/// Write a header into the first [`PYRO_DGRAM_HDR_LEN`] bytes of `buf`. +pub fn write_pyro_header(buf: &mut [u8], h: &PyroDatagramHeader) { + buf[0] = DGRAM_PYROWAVE; + buf[1..3].copy_from_slice(&h.seq.to_le_bytes()); + buf[3..5].copy_from_slice(&h.index.to_le_bytes()); + buf[5..7].copy_from_slice(&h.total.to_le_bytes()); + buf[7..9].copy_from_slice(&h.critical.to_le_bytes()); + buf[9..13].copy_from_slice(&h.ts_ms.to_le_bytes()); + buf[13] = h.flags; +} + +/// Split a PyroWave datagram into header and payload. `None` for anything that +/// is not one or is inconsistent with itself. +pub fn decode_pyro_datagram(buf: &[u8]) -> Option<(PyroDatagramHeader, &[u8])> { + if buf.len() < PYRO_DGRAM_HDR_LEN || buf[0] != DGRAM_PYROWAVE { + return None; + } + let u16_at = |o: usize| u16::from_le_bytes([buf[o], buf[o + 1]]); + let h = PyroDatagramHeader { + seq: u16_at(1), + index: u16_at(3), + total: u16_at(5), + critical: u16_at(7), + ts_ms: u32::from_le_bytes([buf[9], buf[10], buf[11], buf[12]]), + flags: buf[13], + }; + if h.total == 0 || usize::from(h.total) > PYRO_MAX_DATAGRAMS || h.index >= h.total { + return None; + } + if h.critical > h.total { + return None; + } + if h.flags & !KNOWN_FLAGS != 0 { + return None; + } + if h.flags & FLAG_PART == 0 && h.flags & (FLAG_FIRST_PART | FLAG_LAST_PART) != 0 { + return None; + } + Some((h, &buf[PYRO_DGRAM_HDR_LEN..])) +} + +/// A frame laid out as datagrams, in send order, in one buffer. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Datagrams { + pub buf: Vec, + /// Each datagram's length, in order; they lie back to back in `buf`. + pub lens: Vec, + /// Datagrams in the frame, as the headers say. `lens.len()` is this plus + /// the duplicates. + pub total: u16, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PackError { + /// The path admits less than a useful payload per datagram. + PathTooSmall { payload: usize }, + /// More datagrams than a receiver will collect. + TooMany { datagrams: usize }, +} + +impl std::fmt::Display for PackError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::PathTooSmall { payload } => write!(f, "path admits {payload} bytes a datagram"), + Self::TooMany { datagrams } => { + write!( + f, + "frame needs {datagrams} datagrams, limit {PYRO_MAX_DATAGRAMS}" + ) + } + } + } +} + +impl std::error::Error for PackError {} + +/// Lay a frame out as datagrams no larger than `max_datagram`. +/// +/// One packet per datagram. A packet larger than the payload — only a single +/// oversized block can be, and the hub sizes packets to fit — goes in pieces +/// on consecutive indexes. The critical packets lead the frame, and their +/// datagrams are repeated at the end under the same indexes. +pub fn pack_datagrams( + seq: u16, + frame: &PyroFrame, + max_datagram: usize, +) -> Result { + let payload = max_datagram.saturating_sub(PYRO_DGRAM_HDR_LEN); + if payload < MIN_DGRAM_PAYLOAD { + return Err(PackError::PathTooSmall { payload }); + } + + // (packet, byte range within it, flags), one per index. + let mut entries: Vec<(usize, Range, u8)> = Vec::with_capacity(frame.packets.len()); + let mut critical_entries = 0; + for (i, p) in frame.packets.iter().enumerate() { + let len = p.len(); + if len <= payload { + entries.push((i, 0..len, 0)); + } else { + let pieces = len.div_ceil(payload); + for k in 0..pieces { + let mut flags = FLAG_PART; + if k == 0 { + flags |= FLAG_FIRST_PART; + } + if k == pieces - 1 { + flags |= FLAG_LAST_PART; + } + entries.push((i, k * payload..((k + 1) * payload).min(len), flags)); + } + } + if i < frame.critical { + critical_entries = entries.len(); + } + } + if entries.len() > PYRO_MAX_DATAGRAMS { + return Err(PackError::TooMany { + datagrams: entries.len(), + }); + } + let total = entries.len() as u16; + + let order = (0..entries.len()).chain(0..critical_entries); + let sends = entries.len() + critical_entries; + let bytes: usize = entries.iter().map(|e| e.1.len()).sum::() + + entries[..critical_entries] + .iter() + .map(|e| e.1.len()) + .sum::() + + sends * PYRO_DGRAM_HDR_LEN; + let mut buf = Vec::with_capacity(bytes); + let mut lens = Vec::with_capacity(sends); + let mut header = [0u8; PYRO_DGRAM_HDR_LEN]; + for (n, index) in order.enumerate() { + let (packet, ref range, flags) = entries[index]; + let duplicate = if n >= entries.len() { + FLAG_DUPLICATE + } else { + 0 + }; + write_pyro_header( + &mut header, + &PyroDatagramHeader { + seq, + index: index as u16, + total, + critical: critical_entries as u16, + ts_ms: frame.meta.ts_ms, + flags: flags | duplicate, + }, + ); + buf.extend_from_slice(&header); + buf.extend_from_slice(&frame.packet(packet)[range.clone()]); + lens.push(PYRO_DGRAM_HDR_LEN + range.len()); + } + Ok(Datagrams { buf, lens, total }) +} + +// ── Whether to send at all ────────────────────────────────────────────── + +/// Why a frame was not sent. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Skip { + /// More than a frame is still queued from before. + Behind, + /// The frame would not fit beside what is queued. + NoRoom, + /// The frame would not fit the send buffer even empty. + TooLarge, +} + +/// Whether a frame of `frame_bytes` should go out on a queue holding +/// `backlog` of a `buffer`-byte send buffer. +/// +/// QUIC makes room in a full datagram buffer by evicting the *oldest* queued +/// datagrams and reporting success (see `nestri/docs/media-transport.md`). For +/// PyroWave that is the worst outcome available: the evicted datagrams are the +/// tail of the frame before, and the head of this one may follow them, so two +/// frames are damaged where skipping would have cost one. Every PyroWave frame +/// stands alone, so skipping one is clean, and the queue never holds much more +/// than a frame waiting behind the one going out. +pub fn should_skip(backlog: usize, frame_bytes: usize, buffer: usize) -> Option { + if frame_bytes > buffer { + Some(Skip::TooLarge) + } else if backlog > frame_bytes { + Some(Skip::Behind) + } else if backlog + frame_bytes > buffer { + Some(Skip::NoRoom) + } else { + None + } +} + +// ── Collection, at the client ─────────────────────────────────────────── + +/// One packet of a collected frame. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CollectedPacket

{ + /// It arrived in one datagram, and this is that datagram's payload. + Whole(P), + /// It arrived in pieces, joined here. + Joined(Vec), +} + +impl> AsRef<[u8]> for CollectedPacket

{ + fn as_ref(&self) -> &[u8] { + match self { + Self::Whole(p) => p.as_ref(), + Self::Joined(v) => v, + } + } +} + +/// A frame let go by the collector, whole or not. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CollectedFrame

{ + pub seq: u16, + pub ts_ms: u32, + /// Datagrams that arrived, out of `total`. + pub received: u16, + pub total: u16, + /// Whether every critical datagram arrived, by the sender's own count. + /// + /// The answer to "are the coarse bands whole" that a depacketizer cannot + /// give; see [`PyroDatagramHeader::critical`]. A decoder should use this + /// rather than count blocks. + pub critical_complete: bool, + /// What arrived, in frame order. A packet whose pieces did not all arrive + /// is not here. + pub packets: Vec>, +} + +impl

CollectedFrame

{ + pub fn is_whole(&self) -> bool { + self.received == self.total + } +} + +/// What collection looked like, since the last [`PyroCollector::take_stats`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct CollectorStats { + /// Frames let go with every datagram. + pub whole: u64, + /// Frames let go with some missing. Not lost — the decoder may well show + /// them — and not fine either. + pub partial: u64, + /// Datagrams that never arrived, across the partial frames. + pub lost: u64, + /// Datagrams that arrived when their index already had. Expected: every + /// critical datagram is sent twice. + pub duplicates: u64, + /// Critical copies that filled a hole the first copy left. + pub recovered: u64, + /// Datagrams for a frame already let go. + pub stale: u64, + /// Datagrams disagreeing with their frame about its size. + pub inconsistent: u64, + /// Oversized packets dropped because a piece was missing. + pub broken_parts: u64, +} + +struct Flight

{ + seq: u16, + ts_ms: u32, + total: u16, + critical: u16, + slots: Vec>, + received: u16, + /// When its most recent datagram arrived. + last_us: u64, + /// When a newer frame began, if one has. + superseded_us: Option, +} + +impl

Flight

{ + fn due_at(&self) -> u64 { + let deadline = self.last_us + IDLE_US; + match self.superseded_us { + Some(t) => deadline.min(t + GRACE_US), + None => deadline, + } + } +} + +/// Gathers PyroWave datagrams into frames. +/// +/// Generic over the payload so the client can keep slices of the datagrams it +/// received rather than copying each into a frame buffer. +pub struct PyroCollector

{ + /// In frame order, oldest first. + flights: VecDeque>, + newest_released: Option, + stats: CollectorStats, +} + +impl> Default for PyroCollector

{ + fn default() -> Self { + Self::new() + } +} + +impl> PyroCollector

{ + pub fn new() -> Self { + Self { + flights: VecDeque::new(), + newest_released: None, + stats: CollectorStats::default(), + } + } + + /// Take one datagram. Returns every frame now let go, oldest first. + pub fn push( + &mut self, + h: PyroDatagramHeader, + payload: P, + now_us: u64, + ) -> Vec> { + let mut out = self.poll(now_us); + + if let Some(newest) = self.newest_released + && !seq_older_than(newest, h.seq) + { + self.stats.stale += 1; + return out; + } + + let pos = match self.flights.iter().position(|f| f.seq == h.seq) { + Some(pos) => pos, + None => { + // Sorted insert. Whatever is older than the new frame has now + // been overtaken, and the new frame itself is overtaken if it + // arrived behind a newer one. + let pos = self + .flights + .iter() + .position(|f| seq_older_than(h.seq, f.seq)) + .unwrap_or(self.flights.len()); + for f in self.flights.iter_mut().take(pos) { + f.superseded_us.get_or_insert(now_us); + } + let superseded = (pos < self.flights.len()).then_some(now_us); + self.flights.insert( + pos, + Flight { + seq: h.seq, + ts_ms: h.ts_ms, + total: h.total, + critical: h.critical, + slots: (0..h.total).map(|_| None).collect(), + received: 0, + last_us: now_us, + superseded_us: superseded, + }, + ); + if self.flights.len() > MAX_IN_FLIGHT { + self.release_through(0, &mut out); + pos - 1 + } else { + pos + } + } + }; + + let flight = &mut self.flights[pos]; + flight.last_us = now_us; + if flight.total != h.total || flight.critical != h.critical { + self.stats.inconsistent += 1; + return out; + } + let slot = &mut flight.slots[usize::from(h.index)]; + if slot.is_some() { + self.stats.duplicates += 1; + return out; + } + if h.flags & FLAG_DUPLICATE != 0 { + self.stats.recovered += 1; + } + *slot = Some((h.flags & !FLAG_DUPLICATE, payload)); + flight.received += 1; + + if flight.received == flight.total { + self.release_through(pos, &mut out); + } + out + } + + /// Let go of every frame whose time is up. + pub fn poll(&mut self, now_us: u64) -> Vec> { + let mut out = Vec::new(); + // A due frame takes every older one with it: frames leave in order. + if let Some(last) = self.flights.iter().rposition(|f| f.due_at() <= now_us) { + self.release_through(last, &mut out); + } + out + } + + /// When [`Self::poll`] next has something to do, if ever. + pub fn next_deadline(&self) -> Option { + self.flights.iter().map(Flight::due_at).min() + } + + pub fn take_stats(&mut self) -> CollectorStats { + std::mem::take(&mut self.stats) + } + + fn release_through(&mut self, pos: usize, out: &mut Vec>) { + for _ in 0..=pos { + let flight = self.flights.pop_front().expect("pos is in range"); + self.newest_released = Some(flight.seq); + let frame = self.finish(flight); + out.push(frame); + } + } + + fn finish(&mut self, flight: Flight

) -> CollectedFrame

{ + if flight.received == flight.total { + self.stats.whole += 1; + } else { + self.stats.partial += 1; + self.stats.lost += u64::from(flight.total - flight.received); + } + + let critical_complete = flight.slots[..usize::from(flight.critical)] + .iter() + .all(Option::is_some); + let mut packets = Vec::new(); + let mut slots = flight.slots.into_iter().peekable(); + while let Some(slot) = slots.next() { + let Some((flags, payload)) = slot else { + continue; + }; + if flags & FLAG_PART == 0 { + packets.push(CollectedPacket::Whole(payload)); + continue; + } + if flags & FLAG_FIRST_PART == 0 { + // The run's start is gone; skip what is left of it. + self.stats.broken_parts += 1; + while let Some(Some((f, _))) = slots.peek() + && f & FLAG_PART != 0 + && f & FLAG_FIRST_PART == 0 + { + slots.next(); + } + continue; + } + let mut joined = payload.as_ref().to_vec(); + let mut done = flags & FLAG_LAST_PART != 0; + while !done { + match slots.peek() { + Some(Some((f, _))) if f & FLAG_PART != 0 && f & FLAG_FIRST_PART == 0 => { + let (f, p) = slots.next().flatten().expect("peeked"); + joined.extend_from_slice(p.as_ref()); + done = f & FLAG_LAST_PART != 0; + } + _ => break, + } + } + if done { + packets.push(CollectedPacket::Joined(joined)); + } else { + self.stats.broken_parts += 1; + // A missing piece leaves the rest of the run orphaned. Skip it + // here, so the same packet is not counted broken twice. + loop { + match slots.peek() { + Some(None) => {} + Some(Some((f, _))) if f & FLAG_PART != 0 && f & FLAG_FIRST_PART == 0 => {} + _ => break, + } + slots.next(); + } + } + } + + CollectedFrame { + seq: flight.seq, + ts_ms: flight.ts_ms, + received: flight.received, + total: flight.total, + critical_complete, + packets, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::decode_ipc_frame; + + /// A frame of `n` packets of the given sizes, each filled with its own + /// index so a misplaced byte shows. + fn frame_of(sizes: &[usize], critical: usize) -> PyroFrame { + let mut data = Vec::new(); + let mut packets = Vec::new(); + for (i, &len) in sizes.iter().enumerate() { + let start = data.len(); + data.extend((0..len).map(|k| (i * 31 + k) as u8)); + packets.push(start..data.len()); + } + PyroFrame { + frame: 7, + meta: FrameMeta { + ts_ms: 1234, + width: 1920, + height: 1080, + }, + critical, + data, + packets, + } + } + + fn roundtrip_ipc(frame: &PyroFrame) -> (usize, PyroFrame) { + let messages = chunk_frame( + frame.frame, + frame.meta, + &frame.data, + &frame.packets, + frame.critical, + ) + .expect("chunks"); + let mut asm = FrameAssembler::new(); + let mut done = None; + for (i, m) in messages.iter().enumerate() { + assert!( + m.len() <= IPC_PYROWAVE_MAX, + "message {i} is {} bytes", + m.len() + ); + let ipc = decode_ipc_frame(m).expect("valid ipc"); + let got = asm.push(&ipc, 0).expect("assembles"); + assert_eq!(got.is_some(), i == messages.len() - 1); + done = got.or(done); + } + (messages.len(), done.expect("completed")) + } + + #[test] + fn a_frame_survives_the_socket() { + let frame = frame_of(&[8, 1200, 1200, 37, 2488, 1200], 2); + let (messages, back) = roundtrip_ipc(&frame); + assert_eq!(messages, 1); + assert_eq!(back, frame); + } + + /// The reason the socket carries more than one message per frame. + #[test] + fn a_frame_far_past_one_datagram_survives_the_socket() { + let sizes: Vec = (0..4000).map(|i| 900 + (i * 7) % 600).collect(); + let frame = frame_of(&sizes, 3); + assert!(frame.data.len() > 4 * 1024 * 1024); + let (messages, back) = roundtrip_ipc(&frame); + assert!(messages > 60); + assert_eq!(back, frame); + } + + #[test] + fn chunking_refuses_what_it_cannot_carry() { + let f = frame_of(&[10], 0); + assert_eq!( + chunk_frame(0, f.meta, &f.data, &[], 0), + Err(ChunkError::Empty) + ); + assert!(matches!( + chunk_frame(0, f.meta, &f.data, &f.packets, 2), + Err(ChunkError::Critical { .. }) + )); + let big = vec![0u8; IPC_PYROWAVE_MAX]; + assert!(matches!( + chunk_frame(0, f.meta, &big, std::slice::from_ref(&(0..big.len())), 0), + Err(ChunkError::PacketTooLarge { index: 0, .. }) + )); + } + + fn messages_of(frame: &PyroFrame) -> Vec> { + chunk_frame( + frame.frame, + frame.meta, + &frame.data, + &frame.packets, + frame.critical, + ) + .unwrap() + } + + fn big_frame(id: u32) -> PyroFrame { + let mut f = frame_of(&vec![1000; 200], 1); + f.frame = id; + f + } + + /// The encoder never finished sending this frame; the hub must not + /// invent one from what it got. + #[test] + fn a_frame_missing_a_message_is_abandoned_not_patched() { + let frame = big_frame(1); + let messages = messages_of(&frame); + assert!(messages.len() >= 3); + let mut asm = FrameAssembler::new(); + for (i, m) in messages.iter().enumerate() { + if i == 1 { + continue; + } + let got = asm.push(&decode_ipc_frame(m).unwrap(), 0).unwrap(); + assert!(got.is_none(), "a frame with a hole was completed"); + } + assert_eq!(asm.take_abandoned(), 1, "counted once, not per message"); + + // And the next frame is unaffected. + let next = big_frame(2); + let mut got = None; + for m in messages_of(&next) { + got = asm.push(&decode_ipc_frame(&m).unwrap(), 0).unwrap().or(got); + } + assert_eq!(got, Some(next)); + assert_eq!(asm.take_abandoned(), 0); + } + + #[test] + fn a_new_frame_abandons_an_unfinished_one() { + let (a, b) = (big_frame(1), big_frame(2)); + let mut asm = FrameAssembler::new(); + asm.push(&decode_ipc_frame(&messages_of(&a)[0]).unwrap(), 0) + .unwrap(); + let mut got = None; + for m in messages_of(&b) { + got = asm.push(&decode_ipc_frame(&m).unwrap(), 0).unwrap().or(got); + } + assert_eq!(got, Some(b)); + assert_eq!(asm.take_abandoned(), 1); + } + + #[test] + fn a_stalled_frame_is_abandoned() { + let a = big_frame(1); + let ms = messages_of(&a); + let mut asm = FrameAssembler::new(); + asm.push(&decode_ipc_frame(&ms[0]).unwrap(), 0).unwrap(); + for m in &ms[1..] { + assert_eq!( + asm.push(&decode_ipc_frame(m).unwrap(), ABANDON_MS + 1) + .unwrap(), + None + ); + } + assert_eq!(asm.take_abandoned(), 1); + } + + #[test] + fn a_malformed_message_is_refused_whole() { + let f = frame_of(&[100, 100], 0); + let mut m = messages_of(&f).remove(0); + // Claim one packet more than there is. + let at = crate::IPC_HEADER_LEN + 10; + m[at] += 1; + let mut asm = FrameAssembler::new(); + assert!(matches!( + asm.push(&decode_ipc_frame(&m).unwrap(), 0), + Err(AssembleError::Malformed(_)) + )); + + let other = crate::encode_ipc_frame(crate::STREAM_VIDEO, crate::CODEC_AV1, 0, 0, 1, 1, &[]); + assert_eq!( + asm.push(&decode_ipc_frame(&other).unwrap(), 0), + Err(AssembleError::NotPyroWave) + ); + } + + // ── packing ── + + fn unpack(d: &Datagrams) -> Vec<(PyroDatagramHeader, Vec)> { + let mut out = Vec::new(); + let mut at = 0; + for &len in &d.lens { + let (h, p) = decode_pyro_datagram(&d.buf[at..at + len]).expect("parses"); + out.push((h, p.to_vec())); + at += len; + } + assert_eq!(at, d.buf.len()); + out + } + + #[test] + fn a_packet_that_fits_takes_one_datagram_whole() { + let frame = frame_of(&[8, 500, 1200 - PYRO_DGRAM_HDR_LEN, 300], 0); + let d = pack_datagrams(3, &frame, 1200).unwrap(); + let got = unpack(&d); + assert_eq!(d.total, 4); + assert_eq!(got.len(), 4); + for (i, (h, p)) in got.iter().enumerate() { + assert_eq!(h.index as usize, i); + assert_eq!(h.flags, 0); + assert_eq!(h.seq, 3); + assert_eq!(h.ts_ms, 1234); + assert_eq!( + p.as_slice(), + frame.packet(i), + "packet {i} not carried whole" + ); + } + } + + #[test] + fn critical_datagrams_go_twice_under_the_same_index() { + let frame = frame_of(&[8, 900, 900, 900, 900], 2); + let got = unpack(&pack_datagrams(0, &frame, 1200).unwrap()); + assert_eq!(got.len(), 5 + 2); + let tail: Vec<_> = got[5..].iter().map(|(h, _)| (h.index, h.flags)).collect(); + assert_eq!(tail, [(0, FLAG_DUPLICATE), (1, FLAG_DUPLICATE)]); + assert_eq!(got[5].1, got[0].1); + assert_eq!(got[6].1, got[1].1); + } + + #[test] + fn an_oversized_block_goes_in_marked_pieces() { + let frame = frame_of(&[8, 2488, 100], 2); + let got = unpack(&pack_datagrams(0, &frame, 1200).unwrap()); + let flags: Vec = got.iter().map(|(h, _)| h.flags).collect(); + let part = FLAG_PART; + assert_eq!( + flags, + [ + 0, + part | FLAG_FIRST_PART, + part, + part | FLAG_LAST_PART, + 0, + // The critical copies: the header packet and all three pieces. + FLAG_DUPLICATE, + FLAG_DUPLICATE | part | FLAG_FIRST_PART, + FLAG_DUPLICATE | part, + FLAG_DUPLICATE | part | FLAG_LAST_PART, + ] + ); + let joined: Vec = got[1..4].iter().flat_map(|(_, p)| p.clone()).collect(); + assert_eq!(joined, frame.packet(1)); + } + + #[test] + fn packing_refuses_a_path_too_small_or_a_frame_too_large() { + let frame = frame_of(&[8], 0); + assert!(matches!( + pack_datagrams(0, &frame, PYRO_DGRAM_HDR_LEN + MIN_DGRAM_PAYLOAD - 1), + Err(PackError::PathTooSmall { .. }) + )); + let frame = frame_of(&vec![100; PYRO_MAX_DATAGRAMS + 1], 0); + assert!(matches!( + pack_datagrams(0, &frame, 1200), + Err(PackError::TooMany { .. }) + )); + } + + #[test] + fn a_datagram_inconsistent_with_itself_is_refused() { + let mut buf = [0u8; PYRO_DGRAM_HDR_LEN]; + let ok = PyroDatagramHeader { + seq: 1, + index: 0, + total: 2, + critical: 1, + ts_ms: 0, + flags: 0, + }; + write_pyro_header(&mut buf, &ok); + assert!(decode_pyro_datagram(&buf).is_some()); + assert!(decode_pyro_datagram(&buf[..PYRO_DGRAM_HDR_LEN - 1]).is_none()); + + for bad in [ + PyroDatagramHeader { total: 0, ..ok }, + PyroDatagramHeader { index: 2, ..ok }, + PyroDatagramHeader { + total: (PYRO_MAX_DATAGRAMS + 1) as u16, + ..ok + }, + PyroDatagramHeader { flags: 0x80, ..ok }, + PyroDatagramHeader { critical: 3, ..ok }, + PyroDatagramHeader { + flags: FLAG_FIRST_PART, + ..ok + }, + ] { + write_pyro_header(&mut buf, &bad); + assert!(decode_pyro_datagram(&buf).is_none(), "{bad:?} accepted"); + } + // Another kind entirely. + write_pyro_header(&mut buf, &ok); + buf[0] = crate::datagram::DGRAM_VIDEO; + assert!(decode_pyro_datagram(&buf).is_none()); + } + + // ── collection ── + + type Wire = Vec<(PyroDatagramHeader, Vec)>; + + fn wire(seq: u16, frame: &PyroFrame) -> Wire { + unpack(&pack_datagrams(seq, frame, 1200).unwrap()) + } + + fn packets_of(f: &CollectedFrame>) -> Vec> { + f.packets.iter().map(|p| p.as_ref().to_vec()).collect() + } + + fn all_packets(frame: &PyroFrame) -> Vec> { + (0..frame.packets.len()) + .map(|i| frame.packet(i).to_vec()) + .collect() + } + + #[test] + fn a_whole_frame_leaves_the_moment_it_completes() { + let frame = frame_of(&[8, 900, 2488, 900], 1); + let mut c = PyroCollector::new(); + let mut out = Vec::new(); + // The duplicates come last; the frame must not wait for them. + let w = wire(0, &frame); + let originals = w.len() - 1; + for (h, p) in w.into_iter().take(originals) { + out.extend(c.push(h, p, 0)); + } + assert_eq!(out.len(), 1); + assert!(out[0].is_whole()); + assert_eq!(packets_of(&out[0]), all_packets(&frame)); + let s = c.take_stats(); + assert_eq!((s.whole, s.partial, s.lost), (1, 0, 0)); + } + + #[test] + fn order_within_a_frame_does_not_matter() { + let frame = frame_of(&[8, 900, 2488, 900, 50], 2); + let mut w = wire(0, &frame); + w.reverse(); + let mut c = PyroCollector::new(); + let out: Vec<_> = w.into_iter().flat_map(|(h, p)| c.push(h, p, 0)).collect(); + assert_eq!(out.len(), 1); + assert_eq!(packets_of(&out[0]), all_packets(&frame)); + } + + #[test] + fn a_late_copy_is_counted_and_not_kept() { + let frame = frame_of(&[8, 900], 1); + let mut c = PyroCollector::new(); + for (h, p) in wire(0, &frame) { + c.push(h, p, 0); + } + let s = c.take_stats(); + // The copy of datagram 0 arrives after the frame left: stale, since + // the frame is gone, and not a duplicate. + assert_eq!((s.whole, s.stale, s.duplicates), (1, 1, 0)); + } + + #[test] + fn a_copy_fills_the_hole_its_original_left() { + let frame = frame_of(&[8, 900, 900], 1); + let w = wire(0, &frame); + let mut c = PyroCollector::new(); + let mut out = Vec::new(); + for (h, p) in w.into_iter().skip(1) { + out.extend(c.push(h, p, 0)); + } + assert_eq!(out.len(), 1); + assert!(out[0].is_whole(), "the critical copy did not stand in"); + assert_eq!(c.take_stats().recovered, 1); + } + + #[test] + fn a_partial_frame_leaves_once_the_next_begins_and_grace_passes() { + let (a, b) = (frame_of(&[8, 900, 900, 900], 1), frame_of(&[8, 900], 1)); + let mut c = PyroCollector::new(); + let mut wa = wire(0, &a); + wa.retain(|(h, _)| h.index != 2); + for (h, p) in wa { + assert!(c.push(h, p, 0).is_empty()); + } + let wb = wire(1, &b); + let (h, p) = wb[0].clone(); + assert!(c.push(h, p, 100).is_empty(), "left before grace"); + assert_eq!(c.next_deadline(), Some(100 + GRACE_US)); + assert!(c.poll(100 + GRACE_US - 1).is_empty()); + let out = c.poll(100 + GRACE_US); + assert_eq!(out.len(), 1); + assert_eq!(out[0].seq, 0); + assert!(!out[0].is_whole()); + assert_eq!((out[0].received, out[0].total), (3, 4)); + let mut expect = all_packets(&a); + expect.remove(2); + assert_eq!(packets_of(&out[0]), expect); + let s = c.take_stats(); + assert_eq!((s.partial, s.lost), (1, 1)); + } + + #[test] + fn a_frame_nobody_follows_leaves_once_it_goes_quiet() { + let a = frame_of(&[8, 900, 900], 0); + let w = wire(0, &a); + let mut c = PyroCollector::new(); + c.push(w[0].0, w[0].1.clone(), 50); + c.push(w[1].0, w[1].1.clone(), 900); + // Timed from the last datagram, not the first. + assert!(c.poll(900 + IDLE_US - 1).is_empty()); + let out = c.poll(900 + IDLE_US); + assert_eq!(out.len(), 1); + assert_eq!(out[0].received, 2); + } + + /// The failure a deadline from the first datagram had: a frame that takes + /// longer than the deadline to arrive, as every frame does on a path + /// slower than the deadline assumes, was cut off while still arriving. + #[test] + fn a_slow_frame_is_not_cut_off_while_it_is_still_arriving() { + let frame = frame_of(&vec![900; 400], 0); + let mut c = PyroCollector::new(); + let mut out = Vec::new(); + let gap = IDLE_US / 2; + for (i, (h, p)) in wire(0, &frame).into_iter().enumerate() { + out.extend(c.push(h, p, i as u64 * gap)); + } + assert_eq!(out.len(), 1); + assert!(out[0].is_whole(), "{}/{}", out[0].received, out[0].total); + } + + #[test] + fn a_complete_frame_takes_older_unfinished_ones_with_it() { + let (a, b) = (frame_of(&[8, 900, 900], 0), frame_of(&[8, 900], 0)); + let mut c = PyroCollector::new(); + let (h, p) = wire(0, &a)[0].clone(); + c.push(h, p, 0); + let out: Vec<_> = wire(1, &b) + .into_iter() + .flat_map(|(h, p)| c.push(h, p, 10)) + .collect(); + assert_eq!(out.iter().map(|f| f.seq).collect::>(), [0, 1]); + // And the rest of the older frame is stale now. + let (h, p) = wire(0, &a)[1].clone(); + assert!(c.push(h, p, 20).is_empty()); + assert_eq!(c.take_stats().stale, 1); + } + + #[test] + fn a_broken_run_drops_its_packet_and_keeps_the_rest() { + let frame = frame_of(&[8, 2488, 2488, 100], 0); + let w = wire(0, &frame); + // Index 2 is the middle piece of packet 1; 4 is the first of packet 2. + for missing in [2u16, 4] { + let mut c = PyroCollector::new(); + for (h, p) in w.iter().filter(|(h, _)| h.index != missing).cloned() { + c.push(h, p, 0); + } + let out = c.poll(IDLE_US); + assert_eq!(out.len(), 1); + let mut expect = all_packets(&frame); + expect.remove(if missing == 2 { 1 } else { 2 }); + assert_eq!(packets_of(&out[0]), expect, "missing {missing}"); + assert_eq!(c.take_stats().broken_parts, 1); + } + } + + #[test] + fn critical_completeness_is_the_senders_count() { + let frame = frame_of(&[8, 2488, 900, 900, 900], 2); + let w = wire(0, &frame); + assert!( + w.iter().all(|(h, _)| h.critical == 4), + "header packet and three pieces" + ); + + // Everything but one non-critical datagram: critical, not whole. + let mut c = PyroCollector::new(); + for (h, p) in w.iter().filter(|(h, _)| h.index != 5).cloned() { + c.push(h, p, 0); + } + let f = c.poll(IDLE_US).pop().unwrap(); + assert!(!f.is_whole()); + assert!(f.critical_complete); + + // A critical piece lost along with its copy: not critical. + let mut c = PyroCollector::new(); + for (h, p) in w.iter().filter(|(h, _)| h.index != 2).cloned() { + c.push(h, p, 0); + } + let f = c.poll(IDLE_US).pop().unwrap(); + assert!(!f.critical_complete, "a lost critical piece went unnoticed"); + } + + #[test] + fn a_disagreeing_total_is_counted_and_ignored() { + let frame = frame_of(&[8, 900], 0); + let w = wire(0, &frame); + let mut c = PyroCollector::new(); + c.push(w[0].0, w[0].1.clone(), 0); + let mut h = w[1].0; + h.total = 3; + c.push(h, w[1].1.clone(), 0); + assert_eq!(c.take_stats().inconsistent, 1); + assert_eq!(c.push(w[1].0, w[1].1.clone(), 0).len(), 1); + } + + #[test] + fn too_many_frames_in_flight_lets_the_oldest_go() { + let f = frame_of(&[8, 900], 0); + let mut c = PyroCollector::new(); + let mut out = Vec::new(); + for seq in 0..=MAX_IN_FLIGHT as u16 { + let (h, p) = wire(seq, &f)[0].clone(); + out.extend(c.push(h, p, 0)); + } + assert_eq!(out.len(), 1); + assert_eq!(out[0].seq, 0); + } + + #[test] + fn collection_survives_sequence_wraparound() { + let f = frame_of(&[8, 900], 0); + let mut c = PyroCollector::new(); + let mut seqs = Vec::new(); + for seq in [65534u16, 65535, 0, 1] { + for (h, p) in wire(seq, &f).into_iter().take(2) { + seqs.extend(c.push(h, p, 0).into_iter().map(|f| f.seq)); + } + } + assert_eq!(seqs, [65534, 65535, 0, 1]); + assert_eq!(c.take_stats().stale, 0); + } +} + +#[cfg(test)] +mod skip_tests { + use super::*; + + const BUF: usize = 4 << 20; + + #[test] + fn an_empty_queue_sends() { + assert_eq!(should_skip(0, 400_000, BUF), None); + assert_eq!(should_skip(0, BUF, BUF), None); + } + + /// One frame waiting behind the one going out is the steady state at a + /// rate the path carries; it must not trip the guard. + #[test] + fn a_frame_in_flight_is_not_a_backlog() { + assert_eq!(should_skip(400_000, 400_000, BUF), None); + assert_eq!(should_skip(399_999, 400_000, BUF), None); + } + + #[test] + fn more_than_a_frame_behind_skips() { + assert_eq!(should_skip(400_001, 400_000, BUF), Some(Skip::Behind)); + } + + #[test] + fn a_frame_that_would_evict_skips() { + let frame = 3 << 20; + assert_eq!(should_skip(2 << 20, frame, BUF), Some(Skip::NoRoom)); + } + + #[test] + fn a_frame_larger_than_the_buffer_never_goes() { + assert_eq!(should_skip(0, BUF + 1, BUF), Some(Skip::TooLarge)); + } +} diff --git a/crates/nespyro/Cargo.toml b/crates/nespyro/Cargo.toml new file mode 100644 index 00000000..47e766d8 --- /dev/null +++ b/crates/nespyro/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "nespyro" +version = "0.1.0" +description = "PyroWave, a compute-shader wavelet video codec, on ash" +edition.workspace = true +# PyroWave's own licence, not the workspace's: this crate is a port of it. +license = "MIT" +repository.workspace = true + +[dependencies] +# The same rev nescapture and nescore pin, so a `vk::Device` handed in by +# either is the type this crate expects. +ash = { git = "https://github.com/ash-rs/ash", rev = "f4c2ca3e4f6b998d5254ad101a32f024d87cdec2", default-features = false, features = ["std", "debug"] } +futures-channel = "0.3" +thiserror.workspace = true +tracing.workspace = true + +[dev-dependencies] +ash = { git = "https://github.com/ash-rs/ash", rev = "f4c2ca3e4f6b998d5254ad101a32f024d87cdec2", features = ["loaded"] } +pollster = "1.0" +# The transport test carries real frames through the wire format. +nesprotocol = { path = "../nesprotocol" } diff --git a/crates/nespyro/LICENSE b/crates/nespyro/LICENSE new file mode 100644 index 00000000..3392479d --- /dev/null +++ b/crates/nespyro/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2025 Hans-Kristian Arntzen + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/crates/nespyro/build.rs b/crates/nespyro/build.rs new file mode 100644 index 00000000..fe51d01c --- /dev/null +++ b/crates/nespyro/build.rs @@ -0,0 +1,149 @@ +//! Compiles every shader to SPIR-V at build time. +//! +//! `slangc` is a hard build dependency, as in nesrecon: a build that cannot +//! compile the shaders must not produce a binary that discovers it on a +//! customer's machine. Nothing compiles at runtime. +//! +//! Every module is also run through `spirv-val` (or `$SPIRV_VAL`), and that +//! is required too: slangc has been seen to emit SPIR-V that does not +//! validate (see `shuffle_xor` in common.slang), and a check that quietly +//! skips when its tool is missing looks exactly like one that passed. The +//! guest image builder has it from spirv-tools. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +/// Every module the crate creates, as `(source, entry point, defines, name)`. +/// Listed rather than discovered, so a new variant is a deliberate line here. +const MODULES: &[(&str, &str, &[&str], &str)] = &[ + ("rgb_to_ycbcr.slang", "rgb_to_ycbcr", &[], "rgb_to_ycbcr"), + // The encoder requires shaderFloat16, so only the FP16 tile exists. + ("dwt.slang", "dwt", &["FP16=1"], "dwt"), + ("wavelet_quant.slang", "wavelet_quant", &[], "wavelet_quant"), + ( + "analyze_rate_control.slang", + "analyze_rate_control", + &[], + "analyze_rate_control", + ), + ( + "analyze_rate_control.slang", + "analyze_rate_control_finalize", + &[], + "analyze_rate_control_finalize", + ), + // One per subgroup size the resolve pass can run at: the workgroup is + // exactly one subgroup. + ( + "resolve_rate_control.slang", + "resolve_rate_control", + &["WG_SIZE=16"], + "resolve_rate_control_16", + ), + ( + "resolve_rate_control.slang", + "resolve_rate_control", + &["WG_SIZE=32"], + "resolve_rate_control_32", + ), + ( + "resolve_rate_control.slang", + "resolve_rate_control", + &["WG_SIZE=64"], + "resolve_rate_control_64", + ), + ("block_packing.slang", "block_packing", &[], "block_packing"), + ( + "wavelet_dequant.slang", + "wavelet_dequant", + &[], + "wavelet_dequant", + ), + // The decoder does not require shaderFloat16 and picks at runtime. + ("idwt.slang", "idwt", &["FP16=1"], "idwt_fp16"), + ("idwt.slang", "idwt", &["FP16=0"], "idwt_fp32"), +]; + +fn main() { + let manifest = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let shader_dir = manifest.join("shaders"); + let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); + + println!("cargo:rerun-if-changed={}", shader_dir.display()); + println!("cargo:rerun-if-env-changed=SLANGC"); + println!("cargo:rerun-if-env-changed=SPIRV_VAL"); + + let slangc = std::env::var("SLANGC").unwrap_or_else(|_| "slangc".into()); + let spirv_val = std::env::var("SPIRV_VAL").unwrap_or_else(|_| "spirv-val".into()); + + let mut table = String::new(); + for (file, entry, defines, name) in MODULES { + let src = shader_dir.join(file); + let out = out_dir.join(format!("{name}.spv")); + + let mut cmd = Command::new(&slangc); + cmd.arg(&src) + .args(["-target", "spirv"]) + .args(["-profile", "spirv_1_5"]) + .arg("-emit-spirv-directly") + // Pipelines look the entry point up by this name. + .arg("-fvk-use-entrypoint-name") + // Push constants are written by Rust `#[repr(C)]` structs. + .arg("-fvk-use-c-layout") + // slangc raises the profile for subgroup operations and says so on + // every module; the capabilities it adds are what the shader uses, + // and the device check is what decides whether they are there. + .args(["-warnings-disable", "41012"]) + .args(["-I", shader_dir.to_str().unwrap()]) + .args(["-entry", entry]) + .args(["-stage", "compute"]) + .args(["-o", out.to_str().unwrap()]); + for d in *defines { + cmd.arg(format!("-D{d}")); + } + + let output = cmd.output().unwrap_or_else(|e| { + panic!( + "could not run `{slangc}`: {e}\n\ + nespyro compiles its shaders at build time, so Slang is \ + required to build it. Install shader-slang, or point $SLANGC \ + at the binary." + ) + }); + let diagnostics = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "slangc failed on {file}:{entry} {defines:?}\n{diagnostics}" + ); + // Warnings are real here: an uninitialised read or an implicit + // narrowing in a codec shader is a bug, not noise. + assert!( + !diagnostics.contains("warning"), + "slangc warned on {file}:{entry} {defines:?}\n{diagnostics}" + ); + + let status = Command::new(&spirv_val) + .args(["--target-env", "vulkan1.3"]) + .arg(&out) + .status() + .unwrap_or_else(|e| { + panic!( + "could not run `{spirv_val}`: {e}\n\ + nespyro validates its shaders at build time. Install \ + spirv-tools, or point $SPIRV_VAL at the binary." + ) + }); + assert!( + status.success(), + "{name}.spv from {file}:{entry} {defines:?} failed spirv-val" + ); + + table.push_str(&format!( + "pub(crate) static {}: &[u8] = include_bytes!({:?});\n", + name.to_uppercase(), + out.display().to_string() + )); + } + + std::fs::write(Path::new(&out_dir).join("shaders.rs"), table).unwrap(); +} diff --git a/crates/nespyro/reference/build.sh b/crates/nespyro/reference/build.sh new file mode 100755 index 00000000..ed803922 --- /dev/null +++ b/crates/nespyro/reference/build.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Builds the C++ reference harness the cross-check tests run against. +# +# crates/nespyro/reference/build.sh +# +# Clones PyroWave and Granite at the commits nespyro is frozen at into the +# build directory, builds PyroWave's C API, and links the harness against +# it. Prints the harness's path; point NESPYRO_REFERENCE at it. +# +# It builds with two jobs by default, so it can share a machine; set JOBS to +# change that. +set -euo pipefail + +PYROWAVE_COMMIT=89f7e47d4abbf650c91fae766728af866c5e32a0 +PYROWAVE_REPO=${PYROWAVE_REPO:-https://github.com/Themaister/pyrowave} +JOBS=${JOBS:-2} + +here=$(cd "$(dirname "$0")" && pwd) +out=${1:?usage: build.sh } +mkdir -p "$out" +out=$(cd "$out" && pwd) + +if [ ! -d "$out/pyrowave/.git" ]; then + git clone "$PYROWAVE_REPO" "$out/pyrowave" +fi +git -C "$out/pyrowave" fetch -q origin +git -C "$out/pyrowave" checkout -q "$PYROWAVE_COMMIT" +# The script pins Granite's commit itself. +(cd "$out/pyrowave" && ./checkout_granite.sh) + +cmake -S "$out/pyrowave" -B "$out/build" -G Ninja -DCMAKE_BUILD_TYPE=Release +cmake --build "$out/build" --target pyrowave-shared -j "$JOBS" + +lib=$(find "$out/build" -maxdepth 2 -name 'libpyrowave-shared.so*' -type f | head -n1) +c++ -std=c++17 -O2 "$here/harness.cpp" -o "$out/harness" \ + -I "$out/pyrowave" \ + -I "$out/pyrowave/Granite/third_party/khronos/vulkan-headers/include" \ + -L "$(dirname "$lib")" -lpyrowave-shared -Wl,-rpath,"$(dirname "$lib")" + +echo "$out/harness" diff --git a/crates/nespyro/reference/harness.cpp b/crates/nespyro/reference/harness.cpp new file mode 100644 index 00000000..43e035b8 --- /dev/null +++ b/crates/nespyro/reference/harness.cpp @@ -0,0 +1,215 @@ +// The C++ reference, for nespyro's cross-check tests. +// +// A thin command-line wrapper over PyroWave's C API at the commit nespyro's +// bitstream is frozen at. It moves 8-bit planar YCbCr and packet files, so the +// Rust tests can hand it their input and compare what comes out. +// +// A packet file is a sequence of packets, each a little-endian u32 length and +// that many bytes. +// +// harness encode W H 420|444 MAX_BYTES PACKET_SIZE in.yuv out.pkts +// harness decode W H 420|444 in.pkts out.yuv +// harness bench W H 420|444 MAX_BYTES FRAMES in.yuv +// +// NESPYRO_REFERENCE_VID picks the device by PCI vendor ID (hex); otherwise the +// first Vulkan device. + +#include +#include "pyrowave.h" + +#include +#include +#include +#include +#include +#include + +static void die(const char *what, int code = 1) +{ + fprintf(stderr, "harness: %s\n", what); + exit(code); +} + +static std::vector read_file(const char *path) +{ + FILE *f = fopen(path, "rb"); + if (!f) + die("cannot open input"); + std::vector data; + uint8_t buf[65536]; + size_t n; + while ((n = fread(buf, 1, sizeof(buf), f)) > 0) + data.insert(data.end(), buf, buf + n); + fclose(f); + return data; +} + +static void write_file(const char *path, const std::vector &data) +{ + FILE *f = fopen(path, "wb"); + if (!f || fwrite(data.data(), 1, data.size(), f) != data.size()) + die("cannot write output"); + fclose(f); +} + +struct Planes +{ + int width, height; + bool yuv444; + std::vector data; + + int chroma_width() const { return yuv444 ? width : width / 2; } + int chroma_height() const { return yuv444 ? height : height / 2; } + size_t luma_size() const { return size_t(width) * height; } + size_t chroma_size() const { return size_t(chroma_width()) * chroma_height(); } + size_t total() const { return luma_size() + 2 * chroma_size(); } + + pyrowave_cpu_buffer buffer() + { + pyrowave_cpu_buffer b = {}; + b.format = yuv444 ? PYROWAVE_CPU_BUFFER_FORMAT_YUV444P : PYROWAVE_CPU_BUFFER_FORMAT_YUV420P; + b.width = width; + b.height = height; + b.data[0] = data.data(); + b.data[1] = data.data() + luma_size(); + b.data[2] = data.data() + luma_size() + chroma_size(); + b.row_stride_in_bytes[0] = width; + b.row_stride_in_bytes[1] = chroma_width(); + b.row_stride_in_bytes[2] = chroma_width(); + b.plane_size_in_bytes[0] = luma_size(); + b.plane_size_in_bytes[1] = chroma_size(); + b.plane_size_in_bytes[2] = chroma_size(); + return b; + } +}; + +static pyrowave_device make_device() +{ + uint32_t vid = 0; + if (const char *v = getenv("NESPYRO_REFERENCE_VID")) + vid = uint32_t(strtoul(v, nullptr, 16)); + pyrowave_device device = nullptr; + if (pyrowave_create_device_by_compat(vid, 0, nullptr, nullptr, nullptr, &device) != PYROWAVE_SUCCESS) + die("no device", 2); + return device; +} + +static pyrowave_encoder make_encoder(pyrowave_device device, int w, int h, bool yuv444) +{ + pyrowave_encoder_create_info info = {}; + info.device = device; + info.width = w; + info.height = h; + info.chroma = yuv444 ? PYROWAVE_CHROMA_SUBSAMPLING_444 : PYROWAVE_CHROMA_SUBSAMPLING_420; + pyrowave_encoder encoder = nullptr; + if (pyrowave_encoder_create(&info, &encoder) != PYROWAVE_SUCCESS) + die("encoder create failed", 3); + return encoder; +} + +static void print_stats(void *, const char *msg) +{ + fputs(msg, stdout); +} + +int main(int argc, char **argv) +{ + if (argc < 5) + die("usage: harness encode|decode|bench W H 420|444 ..."); + std::string cmd = argv[1]; + Planes p; + p.width = atoi(argv[2]); + p.height = atoi(argv[3]); + p.yuv444 = std::string(argv[4]) == "444"; + + pyrowave_device device = make_device(); + + if (cmd == "encode" || cmd == "bench") + { + if (argc < 8) + die("usage: harness encode W H C MAX_BYTES PACKET_SIZE|FRAMES in out?"); + size_t max_bytes = strtoull(argv[5], nullptr, 10); + p.data = read_file(argv[7]); + if (p.data.size() != p.total()) + die("input is not W x H planar YCbCr"); + pyrowave_encoder encoder = make_encoder(device, p.width, p.height, p.yuv444); + pyrowave_cpu_buffer buffer = p.buffer(); + pyrowave_rate_control rate = { max_bytes }; + + if (cmd == "bench") + { + int frames = atoi(argv[6]); + for (int i = 0; i < frames; i++) + if (pyrowave_encoder_encode_cpu_synchronous(encoder, &buffer, &rate) != PYROWAVE_SUCCESS) + die("encode failed", 4); + // Upstream's own GPU timestamps, per pass. + pyrowave_device_report_performance_stats(device, print_stats, nullptr, true); + pyrowave_encoder_destroy(encoder); + return 0; + } + + if (argc < 9) + die("usage: harness encode W H C MAX_BYTES PACKET_SIZE in.yuv out.pkts"); + size_t packet_size = strtoull(argv[6], nullptr, 10); + if (pyrowave_encoder_encode_cpu_synchronous(encoder, &buffer, &rate) != PYROWAVE_SUCCESS) + die("encode failed", 4); + size_t count = 0; + if (pyrowave_encoder_compute_num_packets(encoder, packet_size, &count) != PYROWAVE_SUCCESS) + die("packet count failed", 4); + std::vector packets(count); + std::vector bitstream(max_bytes + (1 << 20)); + if (pyrowave_encoder_packetize(encoder, packets.data(), packet_size, &count, + bitstream.data(), bitstream.size()) != PYROWAVE_SUCCESS) + die("packetize failed", 4); + std::vector out; + for (size_t i = 0; i < count; i++) + { + uint32_t len = uint32_t(packets[i].size); + out.insert(out.end(), reinterpret_cast(&len), reinterpret_cast(&len) + 4); + out.insert(out.end(), bitstream.begin() + packets[i].offset, + bitstream.begin() + packets[i].offset + packets[i].size); + } + write_file(argv[8], out); + pyrowave_encoder_destroy(encoder); + return 0; + } + + if (cmd == "decode") + { + if (argc < 7) + die("usage: harness decode W H C in.pkts out.yuv"); + std::vector pkts = read_file(argv[5]); + pyrowave_decoder_create_info info = {}; + info.device = device; + info.width = p.width; + info.height = p.height; + info.chroma = p.yuv444 ? PYROWAVE_CHROMA_SUBSAMPLING_444 : PYROWAVE_CHROMA_SUBSAMPLING_420; + info.fragment_path = false; + pyrowave_decoder decoder = nullptr; + if (pyrowave_decoder_create(&info, &decoder) != PYROWAVE_SUCCESS) + die("decoder create failed", 3); + size_t at = 0; + while (at + 4 <= pkts.size()) + { + uint32_t len; + memcpy(&len, pkts.data() + at, 4); + at += 4; + if (at + len > pkts.size()) + die("truncated packet file"); + if (pyrowave_decoder_push_packet(decoder, pkts.data() + at, len) != PYROWAVE_SUCCESS) + die("the reference decoder refused a packet", 5); + at += len; + } + if (!pyrowave_decoder_decode_is_ready(decoder, false)) + die("the reference decoder says the frame is incomplete", 5); + p.data.resize(p.total()); + pyrowave_cpu_buffer buffer = p.buffer(); + if (pyrowave_decoder_decode_cpu_buffer_synchronous(decoder, &buffer) != PYROWAVE_SUCCESS) + die("decode failed", 4); + write_file(argv[6], p.data); + pyrowave_decoder_destroy(decoder); + return 0; + } + + die("unknown command"); +} diff --git a/crates/nespyro/shaders/analyze_rate_control.slang b/crates/nespyro/shaders/analyze_rate_control.slang new file mode 100644 index 00000000..3769690d --- /dev/null +++ b/crates/nespyro/shaders/analyze_rate_control.slang @@ -0,0 +1,234 @@ +// Rate control, first half: sort every block's possible quantization steps +// into buckets by how much distortion they add per byte they save. +// +// Ported from PyroWave's analyze_rate_control.comp and +// analyze_rate_control_finalize.comp (Copyright (c) 2025 Hans-Kristian +// Arntzen, MIT). One workgroup per 32×32 block sums its sixteen 8×8 blocks' +// statistics, turns each extra quantization step into a (saving, bucket) +// pair, and adds the saving to its bucket. The finalize pass prefix-sums the +// buckets, so resolve can find how many buckets it must spend. + +import glsl; +#include "common.slang" + +struct Registers +{ + // The bucket buffer: two counters, then the per-bucket savings from byte + // 64, then the list of operations. + uint *buckets; + uint *total_savings_per_bucket; + RDOperation *rdo_operations; + BlockStats *block_stats; + int2 resolution; + int2 resolution_8x8_blocks; + int block_offset_8x8; + int block_stride_8x8; + int block_offset_32x32; + int block_stride_32x32; + uint total_wg_count; + uint num_blocks_aligned; + uint block_index_shamt; +}; +[[vk::push_constant]] ConstantBuffer registers; + +groupshared uint shared_rate_cost[16]; +groupshared float shared_distortion[16]; + +// Operations that add less distortion per byte go first. +uint distortion_to_bucket_index(float d, float cost, float d_base, float cost_base) +{ + if (cost == cost_base) + return 0; + + // A large range compressed into buckets of about 1.5 dB each: the greedy + // pass chases the least weighted distortion added per byte removed. + float index = 60.0 + 2.0 * log2(max(d - d_base, 0.0) / (cost_base - cost)); + return uint(max(index + 0.5, 0.0)); +} + +uint inclusive_max_clustered16(uint v) +{ + // Never above 127. + v = min(v, 128 - 16 + gl_SubgroupInvocationID); + + for (uint i = 1; i < 16; i *= 2) + { + // Buckets progress monotonically, and each quant level sits at least + // one bucket above the one before. + uint up = subgroupShuffleUp(v, i) + i; + v = max(v, gl_SubgroupInvocationID >= i ? up : 0); + } + + return v; +} + +RDOperation rd_operation(int quant, uint block_offset_saving) +{ + RDOperation op; + op.quant = quant; + op.block_offset_saving = block_offset_saving; + return op; +} + +void emit_rdo_operations() +{ + // Lanes past 16 exist on wider subgroups and only carry dummy values. + // Upstream reads shared_rate_cost at their lane index, past the end of the + // array; the value is never used, so the index is clamped instead. + uint lane = min(gl_SubgroupInvocationID, 15u); + float cost = float(shared_rate_cost[lane]); + float distortion = gl_SubgroupInvocationID < 16 ? shared_distortion[lane] : 1e30; + + uint bucket_index = distortion_to_bucket_index(distortion, cost, shared_distortion[0], float(shared_rate_cost[0])); + if (gl_SubgroupInvocationID == 0) + bucket_index = 0; + + // Bucket index for Q1 is below Q2's if Q1 < Q2. When a high quant sees + // very favourable rate-distortion, lower quants get lower buckets too. + uint inclusive_bucket_index = inclusive_max_clustered16(bucket_index); + + if (gl_SubgroupInvocationID == 0) + { + uint unquantized_cost = shared_rate_cost[0]; + uint ignored; + InterlockedAdd(registers.buckets[1], unquantized_cost, ignored); + } + else if (gl_SubgroupInvocationID < 16) + { + uint saving = shared_rate_cost[gl_SubgroupInvocationID - 1] - shared_rate_cost[gl_SubgroupInvocationID]; + + if (saving != 0) + { + int2 block32x32_index = int2(gl_WorkGroupID.xy); + int block_index = registers.block_offset_32x32 + + block32x32_index.y * registers.block_stride_32x32 + block32x32_index.x; + uint subdivision = uint(block_index) >> registers.block_index_shamt; + uint ignored; + InterlockedAdd( + registers.total_savings_per_bucket[inclusive_bucket_index * BLOCK_SPACE_SUBDIVISION + subdivision], + saving, ignored); + registers.rdo_operations[uint(block_index) + inclusive_bucket_index * registers.num_blocks_aligned] = + rd_operation(int(gl_SubgroupInvocationID), uint(block_index) | (saving << 16)); + } + } +} + +[shader("compute")] +[numthreads(64, 1, 1)] +void analyze_rate_control() +{ + // Each workgroup sums one 32×32 block's cost and distortion at every + // possible quant step. + uint index = gl_SubgroupInvocationID + gl_SubgroupSize * gl_SubgroupID; + int2 block32x32_index = int2(gl_WorkGroupID.xy); + int2 local_block_index = int2(int(bitfieldExtract(index, 0, 2)), int(bitfieldExtract(index, 2, 2))); + int2 block8x8_index = 4 * block32x32_index + local_block_index; + + uint num_active_planes = 0; + + bool block8x8_in_range = all(block8x8_index < registers.resolution_8x8_blocks); + int block_index_8x8 = registers.block_offset_8x8 + + registers.block_stride_8x8 * block8x8_index.y + + block8x8_index.x; + + if (block8x8_in_range) + num_active_planes = registers.block_stats[block_index_8x8].num_planes; + + uint bit_index = index >> 4; + + for (uint i = bit_index; i < 16; i += 4) + { + float dist = 0.0; + uint cost = 0; + + if (block8x8_in_range) + { + QuantStats stats = registers.block_stats[block_index_8x8].errors[min(i, num_active_planes)]; + dist = float(stats.square_error); + cost = uint(stats.payload_cost); + } + + // 16 bits of control codes and 8 of Q bits and quant scale. Cost is in + // bits; 8×8 blocks decode in isolation. + if (cost != 0) + cost += 24; + + if (gl_SubgroupSize == 16) + { + cost = subgroupAdd(cost); + dist = subgroupAdd(dist); + } + else + { + cost += shuffle_xor(cost, 1u); + cost += shuffle_xor(cost, 2u); + cost += shuffle_xor(cost, 4u); + cost += shuffle_xor(cost, 8u); + + dist += shuffle_xor(dist, 1u); + dist += shuffle_xor(dist, 2u); + dist += shuffle_xor(dist, 4u); + dist += shuffle_xor(dist, 8u); + } + + if ((index & 15u) == 0u) + { + // A header must be coded, unless everything in the 32×32 block + // decodes to zero and it can be dropped. + if (cost != 0) + cost += 64; + + // Each packet is aligned to 4 bytes. + shared_rate_cost[i] = (cost + 31) >> 5; + shared_distortion[i] = dist; + } + } + + barrier(); + + if (gl_SubgroupID == 0) + emit_rdo_operations(); +} + +groupshared uint shared_scan[512]; + +struct FinalizeRegisters +{ + uint4 *total_savings_per_bucket; +}; +[[vk::push_constant]] ConstantBuffer finalize_registers; + +// Prefix sum over all 128 × 16 buckets, four per lane. +[shader("compute")] +[numthreads(512, 1, 1)] +void analyze_rate_control_finalize() +{ + uint4 v = finalize_registers.total_savings_per_bucket[gl_LocalInvocationIndex]; + v.y += v.x; + v.z += v.y; + v.w += v.z; + shared_scan[gl_LocalInvocationIndex] = v.w; + + barrier(); + + // As upstream: the steps stop short of 256, so the running totals of the + // upper half of the lanes miss the lower half's first contributions and + // come out low. That makes resolve spend more buckets than it needs, so + // frames land under their target rather than over. Kept as is so the + // output stays comparable with the reference; worth raising upstream. + for (uint step = 1u; step < 512u / 2u; step *= 2u) + { + barrier(); + + uint shuffled_up = 0; + if (gl_LocalInvocationIndex >= step) + shuffled_up = shared_scan[gl_LocalInvocationIndex - step]; + + barrier(); + + v += shuffled_up; + shared_scan[gl_LocalInvocationIndex] = v.w; + } + + finalize_registers.total_savings_per_bucket[gl_LocalInvocationIndex] = v; +} diff --git a/crates/nespyro/shaders/block_packing.slang b/crates/nespyro/shaders/block_packing.slang new file mode 100644 index 00000000..80537903 --- /dev/null +++ b/crates/nespyro/shaders/block_packing.slang @@ -0,0 +1,367 @@ +// Write the final bitstream: apply rate control's quant to every block and +// pack it into its finished form. +// +// Ported from PyroWave's block_packing.comp (Copyright (c) 2025 Hans-Kristian +// Arntzen, MIT). A workgroup of 64 lanes handles 2×2 blocks of 32×32, 16 +// lanes per block, one lane per 8×8 block. Each block reserves its space with +// an atomic, so blocks land in the output in whatever order the GPU runs +// them, and a table records where each went; the host packetizes in block +// order from that table. +// +// Two additions. First, bytes the format ignores are written as zero rather +// than left as they were: the unused bits of a block's last sign byte, which +// upstream takes from uninitialised shared memory, and the padding to a word +// boundary, which upstream leaves holding whatever an earlier frame put +// there. Decoders ignore both, so the stream is the same format; it is also +// deterministic, and carries no stale memory onto the network. +// +// Second, the output has a capacity. Upstream relies on rate control +// never overshooting its buffer. Here the buffer is reached through a raw +// device address with no bounds checking, so a block that would not fit is +// recorded as empty instead of being written, and the host sees the overflow +// from the counter, which still counts every word asked for. + +import glsl; +#include "common.slang" + +struct Registers +{ + // One output buffer, addressed as words, halves and bytes. + uint *bitstream_data; + uint16_t *bitstream_data_16b; + uint8_t *bitstream_data_8b; + BitstreamPacket *bitstream_meta; + BlockMeta *block_meta; + // Two counters: the quantizer's scratch allocation, then this pass's. + uint *payload_counter; + uint8_t *payload_data; + BlockStats *block_stats; + int *quant_data; + // What changes per frame, written by the frame's own commands so the + // rest can be recorded once: [0] is the 3-bit sequence number. + uint *frame_params; + int2 resolution; + int2 resolution_32x32_blocks; + int2 resolution_8x8_blocks; + uint quant_resolution_code; + int block_offset_32x32; + int block_stride_32x32; + int block_offset_8x8; + int block_stride_8x8; + uint bitstream_capacity_words; +}; +[[vk::push_constant]] ConstantBuffer registers; + +uint compute_required_8x8_size(uint control_word) +{ + int q_bits = int(bitfieldExtract(control_word, Q_PLANES_OFFSET, Q_PLANES_BITS)); + uint lsbs = control_word & 0x5555u; + uint msbs = control_word & 0xaaaau; + uint msbs_shift = msbs >> 1; + msbs |= msbs_shift; + return bitCount(lsbs) + bitCount(msbs) + uint(q_bits) * 8; +} + +uint quantize_code_word(uint control_word, int quant) +{ + if (quant != 0 && control_word != 0) + { + int q_bits = int(bitfieldExtract(control_word, Q_PLANES_OFFSET, Q_PLANES_BITS)); + int sub_quant = min(q_bits, quant); + q_bits -= sub_quant; + quant -= sub_quant; + + if (quant != 0) + { + quant = min(quant, 3); + + uint plane0 = control_word & 0x5555u; + uint plane1 = (control_word & 0xaaaau) >> 1; + uint plane2 = plane0 & plane1; + + do + { + plane0 = plane1; + plane1 = plane2; + plane2 = 0; + quant--; + } while (quant != 0); + + plane0 &= ~plane1; + + uint new_control_word = plane0 | (plane1 << 1); + control_word = bitfieldInsert(control_word, new_control_word, 0, 16); + } + + control_word = bitfieldInsert(control_word, uint(q_bits), Q_PLANES_OFFSET, Q_PLANES_BITS); + } + + return control_word; +} + +uint copy_bytes(inout uint output_offset, uint input_offset, uint count) +{ + uint significant_mask = 0; + + do + { + uint in_data = uint(registers.payload_data[input_offset]); + // Any 1 in the non-sign planes means it was not deadzone quantized. + significant_mask |= in_data; + registers.bitstream_data_8b[output_offset++] = uint8_t(in_data); + count--; + input_offset++; + } while (count > 0); + + return significant_mask; +} + +uint modify_quant_code(uint code, int quant) +{ + int e = int(bitfieldExtract(code, 3, 5)); + e = max(e - quant, 0); + code = bitfieldInsert(code, uint(e), 3, 5); + return code; +} + +uint inclusive_add_clustered16(uint v) +{ + for (uint i = 1; i < 16; i *= 2) + { + uint up = subgroupShuffleUp(v, i); + v += (gl_SubgroupInvocationID & 15) >= i ? up : 0; + } + return v; +} + +groupshared uint shared_sign_bank[4][1024 / 32]; +static uint pending_sign_write = 0; +static uint pending_sign_mask = 0; + +void append_sign_plane(uint bank, inout uint local_sign_offset, uint sign_mask, uint significant_mask) +{ + // One bit at a time. + while (significant_mask != 0) + { + int bit = findLSB(significant_mask); + significant_mask &= significant_mask - 1; + int out_bit = int(local_sign_offset & 31u); + pending_sign_write = bitfieldInsert(pending_sign_write, bitfieldExtract(sign_mask, bit, 1), out_bit, 1); + pending_sign_mask = bitfieldInsert(pending_sign_mask, 1u, out_bit, 1); + + if (out_bit == 31) + { + if (pending_sign_mask == ~0u) + { + shared_sign_bank[bank][local_sign_offset / 32] = pending_sign_write; + } + else + { + atomicAnd(shared_sign_bank[bank][local_sign_offset / 32], ~pending_sign_mask); + atomicOr(shared_sign_bank[bank][local_sign_offset / 32], pending_sign_write & pending_sign_mask); + } + + pending_sign_mask = 0; + } + + local_sign_offset++; + } +} + +void flush_sign_plane(uint bank, uint local_sign_offset) +{ + if (pending_sign_mask != 0) + { + atomicAnd(shared_sign_bank[bank][local_sign_offset / 32], ~pending_sign_mask); + atomicOr(shared_sign_bank[bank][local_sign_offset / 32], pending_sign_write & pending_sign_mask); + pending_sign_mask = 0; + } +} + +BitstreamPacket bitstream_packet(uint offset, uint num_words) +{ + BitstreamPacket p; + p.offset = offset; + p.num_words = num_words; + return p; +} + +[shader("compute")] +[numthreads(64, 1, 1)] +void block_packing() +{ + uint index = gl_SubgroupInvocationID + gl_SubgroupSize * gl_SubgroupID; + uint linear_block_32x32_index = index >> 4; + int2 block32x32_index = 2 * int2(gl_WorkGroupID.xy); + block32x32_index.x += int(bitfieldExtract(index, 4, 1)); + block32x32_index.y += int(bitfieldExtract(index, 5, 1)); + int2 local_block_index = int2(int(bitfieldExtract(index, 0, 2)), int(bitfieldExtract(index, 2, 2))); + int2 block8x8_index = 4 * block32x32_index + local_block_index; + + // Each 32×32 block's 16 lanes share one bank and sit in one subgroup, so a + // subgroup barrier orders this before the bank is used. + shared_sign_bank[linear_block_32x32_index][2 * (index & 15u) + 0] = 0; + shared_sign_bank[linear_block_32x32_index][2 * (index & 15u) + 1] = 0; + subgroupBarrier(); + + BlockMeta meta; + int quant; + + bool in_range_8x8 = all(block8x8_index < registers.resolution_8x8_blocks); + bool in_range_32x32 = all(block32x32_index < registers.resolution_32x32_blocks); + uint num_bits_for_q = 0; + + if (in_range_32x32) + { + int block_index = registers.block_offset_32x32 + + registers.block_stride_32x32 * block32x32_index.y + + block32x32_index.x; + quant = registers.quant_data[block_index]; + } + else + { + quant = 0; + } + + if (in_range_8x8) + { + int block_index = registers.block_offset_8x8 + + registers.block_stride_8x8 * block8x8_index.y + + block8x8_index.x; + meta = registers.block_meta[block_index]; + uint num_planes = registers.block_stats[block_index].num_planes; + num_bits_for_q = uint(registers.block_stats[block_index].errors[min(num_planes, uint(quant))].payload_cost); + } + else + { + meta.code_word = 0; + meta.offset = 0; + } + + uint code_word = quantize_code_word(meta.code_word, quant); + bool active_code_word = (code_word & 0xffffu) != 0; + + uint4 code_word_ballot = subgroupBallot(active_code_word); + uint local_ballot = gl_SubgroupSize >= 64 && linear_block_32x32_index >= 2 ? code_word_ballot.y : code_word_ballot.x; + local_ballot = gl_SubgroupSize >= 32 ? bitfieldExtract(local_ballot, int(16u * (linear_block_32x32_index & 1u)), 16) : local_ballot; + + uint required_plane_bytes = compute_required_8x8_size(code_word); + uint required_sign_bits = num_bits_for_q - required_plane_bytes * 8; + + uint required_bits_with_meta = num_bits_for_q; + if (required_bits_with_meta != 0) + required_bits_with_meta += 24; + + static const uint HeaderSize = 2; + bool writes_header = + all(block32x32_index < registers.resolution_32x32_blocks) && (index & 15u) == 15u; + + uint payload_total_bits = subgroupClusteredAdd(required_bits_with_meta, 16u); + uint payload_total_words = (payload_total_bits + 31) / 32; + if (payload_total_words != 0) + payload_total_words += HeaderSize; + + uint global_payload_offset = 0; + if (writes_header && payload_total_words != 0) + InterlockedAdd(registers.payload_counter[1], payload_total_words, global_payload_offset); + global_payload_offset = subgroupShuffle(global_payload_offset, gl_SubgroupInvocationID | 15u); + + // Every lane of the block agrees on this: both terms are the same across + // its 16 lanes. + bool fits = global_payload_offset + payload_total_words <= registers.bitstream_capacity_words; + + if (writes_header) + { + uint block_index = uint(registers.block_offset_32x32 + + block32x32_index.y * registers.block_stride_32x32 + block32x32_index.x); + + if (payload_total_words != 0 && fits) + { + registers.bitstream_data[global_payload_offset + 0] = + local_ballot | (payload_total_words << 16) | (registers.frame_params[0] << 28); + registers.bitstream_data[global_payload_offset + 1] = + modify_quant_code(registers.quant_resolution_code, quant) | (block_index << 8); + } + + registers.bitstream_meta[block_index] = + bitstream_packet(global_payload_offset, fits ? payload_total_words : 0); + } + + uint total_subblocks = bitCount(local_ballot); + + uint total_sign_bits = inclusive_add_clustered16(required_sign_bits); + uint local_planes_offset = inclusive_add_clustered16(required_plane_bytes) - required_plane_bytes; + uint local_sign_offset = total_sign_bits - required_sign_bits; + uint global_planes_offset = 4 * global_payload_offset + 3 * total_subblocks + 4 * HeaderSize; + uint global_sign_offset = global_planes_offset + subgroupClusteredAdd(required_plane_bytes, 16u); + global_planes_offset += local_planes_offset; + + uint total_sign_bytes = (subgroupShuffle(total_sign_bits, gl_SubgroupInvocationID | 15u) + 7) / 8; + + // Then N code words, one per active 8×8 block in ballot order. + if (active_code_word && fits) + { + uint block_header_offset = bitCount(bitfieldExtract( + local_ballot, 0, local_block_index.y * 4 + local_block_index.x)); + + uint in_q_bits = bitfieldExtract(meta.code_word, Q_PLANES_OFFSET, Q_PLANES_BITS); + uint out_q_bits = bitfieldExtract(code_word, Q_PLANES_OFFSET, Q_PLANES_BITS); + uint input_offset = meta.offset; + uint output_offset = global_planes_offset; + + for (int bit_offset = 0; bit_offset < 16; bit_offset += 2) + { + uint out_planes = bitfieldExtract(code_word, bit_offset, 2) + out_q_bits; + uint in_planes = bitfieldExtract(meta.code_word, bit_offset, 2) + in_q_bits; + if (in_planes != 0) + in_planes++; + + uint sign_plane = uint(registers.payload_data[input_offset]); + + if (out_planes != 0) + { + uint significant_mask = copy_bytes(output_offset, input_offset + 1, out_planes); + append_sign_plane(linear_block_32x32_index, local_sign_offset, sign_plane, significant_mask); + } + + input_offset += in_planes; + } + + flush_sign_plane(linear_block_32x32_index, local_sign_offset); + + registers.bitstream_data_16b[2 * global_payload_offset + block_header_offset + 2 * HeaderSize] = + uint16_t(code_word); + registers.bitstream_data_8b[4 * global_payload_offset + 2 * total_subblocks + block_header_offset + 4 * HeaderSize] = + uint8_t(code_word >> 16); + } + + subgroupBarrier(); + + if (!fits) + return; + + // Copy out every sign plane of the block. + for (uint i = index & 15u; i < total_sign_bytes / 4; i += 16) + { + uint sign_word = shared_sign_bank[linear_block_32x32_index][i]; + uint offset_8b = global_sign_offset + 4 * i; + registers.bitstream_data_8b[offset_8b + 0] = uint8_t(sign_word >> 0); + registers.bitstream_data_8b[offset_8b + 1] = uint8_t(sign_word >> 8); + registers.bitstream_data_8b[offset_8b + 2] = uint8_t(sign_word >> 16); + registers.bitstream_data_8b[offset_8b + 3] = uint8_t(sign_word >> 24); + } + + // And any stragglers. + for (uint i = (total_sign_bytes & ~3u) + (index & 15u); i < total_sign_bytes; i += 16) + { + uint sign_word = shared_sign_bank[linear_block_32x32_index][i / 4]; + uint offset_8b = global_sign_offset + i; + registers.bitstream_data_8b[offset_8b] = uint8_t(sign_word >> (8 * (i & 3u))); + } + + // Zero the padding up to the block's last word. + uint block_end = 4 * (global_payload_offset + payload_total_words); + if (payload_total_words != 0 && (index & 15u) == 0u) + for (uint b = global_sign_offset + total_sign_bytes; b < block_end; b++) + registers.bitstream_data_8b[b] = uint8_t(0); +} diff --git a/crates/nespyro/shaders/common.slang b/crates/nespyro/shaders/common.slang new file mode 100644 index 00000000..284b9f60 --- /dev/null +++ b/crates/nespyro/shaders/common.slang @@ -0,0 +1,110 @@ +// Shared by every PyroWave pass. +// +// Ported from PyroWave's constants.h, dwt_swizzle.h and dwt_quant_scale.h +// (Copyright (c) 2025 Hans-Kristian Arntzen, MIT). Upstream's GLSL builtins +// come from Slang's glsl module, so the ports read line for line against the +// originals; only declarations and types change. + +#pragma once + +// Where a quantizer code word keeps its two extra fields above the 16 bits of +// per-subblock plane deltas. +static const int Q_PLANES_OFFSET = 16; +static const int Q_PLANES_BITS = 4; +static const int QUANT_SCALE_OFFSET = 20; +static const int QUANT_SCALE_BITS = 4; + +// subgroupShuffleXor, declaring the capability it needs. +// +// Slang 2026.17's glsl module emits OpGroupNonUniformShuffleXor without +// declaring GroupNonUniformShuffle, so a shader whose only shuffle is an XOR +// one fails spirv-val. This emits the same instruction with the capability. +T shuffle_xor(T value, uint mask) +{ + return spirv_asm { + OpCapability GroupNonUniformShuffle; + result:$$T = OpGroupNonUniformShuffleXor Subgroup $value $mask + }; +} + +// Coefficient order inside a 4×8 and an 8×8 block. +int2 unswizzle4x8(uint index) +{ + uint y = bitfieldExtract(index, 0, 1); + uint x = bitfieldExtract(index, 1, 2); + y |= bitfieldExtract(index, 3, 2) << 1; + return int2(int(x), int(y)); +} + +int2 unswizzle8x8(uint index) +{ + uint y = bitfieldExtract(index, 0, 1); + uint x = bitfieldExtract(index, 1, 2); + y |= bitfieldExtract(index, 3, 2) << 1; + x |= bitfieldExtract(index, 5, 1) << 2; + return int2(int(x), int(y)); +} + +// The per-8×8 adaptive scale: a 4-bit code, 0.25 to about 2.2. +float decode_quant_scale(uint code) +{ + return float(code) / 8.0 + 0.25; +} + +static const uint ENCODE_QUANT_IDENTITY = 6; + +uint encode_quant_scale(float scale) +{ + // Round the scale up so that the quantizer scale effectively rounds down. + return uint(ceil((scale - 0.25) * 8.0)); +} + +// The per-32×32 quant code: a small float with a 5-bit exponent and 3-bit +// mantissa, for numbers in (0, 16). +static const int MaxScaleExp = 4; + +float decode_quant(uint quant_code) +{ + int e = MaxScaleExp - int(quant_code >> 3); + int m = int(quant_code) & 0x7; + float inv_quant = (1.0 / (8.0 * 1024.0 * 1024.0)) * float((8 + m) * (1 << (20 + e))); + return inv_quant; +} + +// Shared layouts. Every one matches a `#[repr(C)]` type on the Rust side. + +struct BlockMeta +{ + uint code_word; + uint offset; +}; + +struct QuantStats +{ + half square_error; + uint16_t payload_cost; +}; + +// 64 bytes: the distortion and cost of dropping 0 to 15 bit planes. +struct BlockStats +{ + uint num_planes; + QuantStats errors[15]; +}; + +struct RDOperation +{ + int quant; + uint block_offset_saving; +}; + +struct BitstreamPacket +{ + uint offset; + uint num_words; +}; + +// Rate control sorts savings into this many buckets per subdivision of the +// block space. +static const int BLOCK_SPACE_SUBDIVISION = 16; +static const int NUM_RDO_BUCKETS = 128; diff --git a/crates/nespyro/shaders/dwt.slang b/crates/nespyro/shaders/dwt.slang new file mode 100644 index 00000000..b603e714 --- /dev/null +++ b/crates/nespyro/shaders/dwt.slang @@ -0,0 +1,199 @@ +// Forward CDF 9/7 wavelet transform of one 32×32 tile into the four bands of +// the next level. +// +// Ported from PyroWave's dwt.comp (Copyright (c) 2025 Hans-Kristian Arntzen, +// MIT). One workgroup of 64 lanes reads its tile plus a 4-texel apron through +// a mirrored-repeat sampler, lifts rows, then columns, and writes a 16×16 +// quarter to each band layer. Edges mirror through the sampler rather than +// through branches; `generate_mirror_uv` bends the coordinates so the +// padding past the real image mirrors the way the format expects. + +import glsl; +#include "dwt_common.slang" + +[[vk::binding(0, 0)]] Sampler2D uTexture; +[[vk::binding(1, 0)]] [[vk::image_format("unknown")]] WTexture2DArray uOutput; + +// The first transform of a plane moves it from [0, 1] to [-0.5, 0.5]. +[[vk::constant_id(0)]] const bool DCShift = false; + +struct Registers +{ + int2 resolution; + float2 inv_resolution; + int2 aligned_resolution; +}; +[[vk::push_constant]] ConstantBuffer registers; + +static uint local_index; + +float2 generate_mirror_uv(int2 coord) +{ + coord -= int2(coord < int2(0)); + coord += 1; + int2 end_mirrored_clamp = (2 * registers.aligned_resolution) - registers.resolution; + int2 past_wrapped_coord = coord + 2 * (registers.resolution - registers.aligned_resolution) + 1; + coord = select(coord >= end_mirrored_clamp, past_wrapped_coord, min(coord, registers.resolution)); + return float2(coord) * registers.inv_resolution; +} + +float4 gather(int2 coord) +{ + float4 texels = uTexture.Gather(generate_mirror_uv(coord)).wzxy; + if (DCShift) + texels -= 0.5; + return texels; +} + +void load_image_with_apron() +{ + int2 base_coord = int2(gl_WorkGroupID.xy) * int2(BLOCK_SIZE, BLOCK_SIZE) - APRON; + int2 local_coord0 = 2 * unswizzle8x8(local_index); + int2 coord0 = base_coord + local_coord0; + + float4 texels0 = gather(coord0); + float4 texels1 = gather(coord0 + int2(16, 0)); + float4 texels2 = gather(coord0 + int2(0, 16)); + float4 texels3 = gather(coord0 + int2(16, 16)); + + int local_coord0_y_half = local_coord0.y >> 1; + + // Pack two lines together in one vec2, so packed maths processes two lines + // in parallel. + store_shared(local_coord0_y_half + 0, local_coord0.x + 0, texels0.xz); + store_shared(local_coord0_y_half + 0, local_coord0.x + 1, texels0.yw); + store_shared(local_coord0_y_half + 0, local_coord0.x + 16, texels1.xz); + store_shared(local_coord0_y_half + 0, local_coord0.x + 17, texels1.yw); + store_shared(local_coord0_y_half + 8, local_coord0.x + 0, texels2.xz); + store_shared(local_coord0_y_half + 8, local_coord0.x + 1, texels2.yw); + store_shared(local_coord0_y_half + 8, local_coord0.x + 16, texels3.xz); + store_shared(local_coord0_y_half + 8, local_coord0.x + 17, texels3.yw); + + // The top-right apron. + { + int2 local_coord = int2(BLOCK_SIZE + 2 * int(local_index % 4u), 2 * int(local_index / 4u)); + float4 texels = gather(base_coord + local_coord); + store_shared(local_coord.y >> 1, local_coord.x + 0, texels.xz); + store_shared(local_coord.y >> 1, local_coord.x + 1, texels.yw); + } + + // The bottom-left apron. + { + int2 local_coord = int2(2 * int(local_index % 16u), BLOCK_SIZE + 2 * int(local_index / 16u)); + float4 texels = gather(base_coord + local_coord); + store_shared(local_coord.y >> 1, local_coord.x + 0, texels.xz); + store_shared(local_coord.y >> 1, local_coord.x + 1, texels.yw); + } + + if (local_index < 16) + { + // The bottom-right apron. + int2 local_coord = int2(BLOCK_SIZE + 2 * int(local_index % 4u), BLOCK_SIZE + 2 * int(local_index / 4u)); + float4 texels = gather(base_coord + local_coord); + store_shared(local_coord.y >> 1, local_coord.x + 0, texels.xz); + store_shared(local_coord.y >> 1, local_coord.x + 1, texels.yw); + } +} + +// The four CDF 9/7 lifting steps over one padded line, then the rescale and a +// transposing write back, so the next pass lifts the other direction. +void lift_and_transpose(bool active_lane, int2 local_coord) +{ + const int PADDED_SIZE = SIZE + 2 * APRON; + const int PADDED_SIZE_HALF = PADDED_SIZE / 2; + // Zeroed only so the compiler can see it is never read uninitialised; + // inactive lanes write nothing back. + float2 values[SIZE + 2 * APRON] = {}; + + if (active_lane) + { + for (int i = 0; i < PADDED_SIZE; i++) + values[i] = load_shared(local_coord.y, local_coord.x + i); + + for (int i = 1; i < PADDED_SIZE - 1; i += 2) + values[i] += ALPHA * (values[i - 1] + values[i + 1]); + for (int i = 2; i < PADDED_SIZE - 2; i += 2) + values[i] += BETA * (values[i - 1] + values[i + 1]); + for (int i = 3; i < PADDED_SIZE - 3; i += 2) + values[i] += GAMMA * (values[i - 1] + values[i + 1]); + for (int i = 4; i < PADDED_SIZE - 4; i += 2) + values[i] += DELTA * (values[i - 1] + values[i + 1]); + } + + // Everyone has read before anyone writes. + barrier(); + + if (active_lane) + { + for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++) + { + float2 a = values[2 * i + 0]; + float2 b = values[2 * i + 1]; + + // Filter kernel rescale. + a *= inv_K; + b *= K; + + // Transpose the 2×2 block. + float2 t0 = float2(a.x, b.x); + float2 t1 = float2(a.y, b.y); + + int y_coord = (local_coord.x >> 1) + (i - APRON_HALF); + store_shared(y_coord, 2 * local_coord.y + 0, t0); + store_shared(y_coord, 2 * local_coord.y + 1, t1); + } + } +} + +void forward_transform8x2() +{ + lift_and_transpose<8>(true, int2(8 * int(local_index % 4u), int(local_index / 4u))); +} + +void forward_transform4x2(bool active_lane, int y_offset) +{ + lift_and_transpose<4>(active_lane, int2(4 * int(local_index % 8u), int(local_index / 8u) + y_offset)); +} + +[shader("compute")] +[numthreads(64, 1, 1)] +void dwt() +{ + local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID; + + load_image_with_apron(); + + barrier(); + + // Horizontal transform. + forward_transform8x2(); + + // The apron needs transforming too. + forward_transform4x2(local_index < 32, BLOCK_SIZE_HALF); + + barrier(); + + // Vertical transform. + forward_transform8x2(); + + barrier(); + + int2 local_coord = unswizzle8x8(local_index); + for (int y = local_coord.y; y < BLOCK_SIZE_HALF; y += 8) + { + for (int x = local_coord.x * 2; x < BLOCK_SIZE; x += 16) + { + float2 v0 = load_shared(y, x + 0); + float2 v1 = load_shared(y, x + 1); + + int img_x = x >> 1; + int img_y = y; + + int2 base_image_coord = int2(gl_WorkGroupID.xy) * (BLOCK_SIZE / 2) + int2(img_x, img_y); + uOutput.Store(int3(base_image_coord, 0), v0.xxxx); + uOutput.Store(int3(base_image_coord, 2), v0.yyyy); + uOutput.Store(int3(base_image_coord, 1), v1.xxxx); + uOutput.Store(int3(base_image_coord, 3), v1.yyyy); + } + } +} diff --git a/crates/nespyro/shaders/dwt_common.slang b/crates/nespyro/shaders/dwt_common.slang new file mode 100644 index 00000000..72fdf43e --- /dev/null +++ b/crates/nespyro/shaders/dwt_common.slang @@ -0,0 +1,41 @@ +// The CDF 9/7 lifting constants and the shared tile the forward and inverse +// transforms work in. +// +// Ported from PyroWave's dwt_common.h (Copyright (c) 2025 Hans-Kristian +// Arntzen, MIT) at precision level 1, upstream's default: the lifting maths is +// FP32 and the shared tile is FP16. `FP16` says how the tile is stored: as +// native half2 where the device has shaderFloat16, packed into a uint through +// packHalf2x16 where it does not. The values held are the same either way. + +#pragma once + +#include "common.slang" + +static const int APRON = 4; +static const int APRON_HALF = APRON / 2; +static const int BLOCK_SIZE = 32; +static const int BLOCK_SIZE_HALF = BLOCK_SIZE >> 1; + +static const float ALPHA = -1.586134342059924; +static const float BETA = -0.052980118572961; +static const float GAMMA = 0.882911075530934; +static const float DELTA = 0.443506852043971; +static const float K = 1.230174104914001; +static const float inv_K = 1.0 / 1.230174104914001; + +// Two image rows packed side by side in each element, so one lane lifts two +// lines at once. +#if FP16 +groupshared half2 shared_block[(BLOCK_SIZE + 2 * APRON) / 2][(BLOCK_SIZE + 2 * APRON) + 1]; +float2 load_shared(uint y, uint x) { return float2(shared_block[y][x]); } +void store_shared(uint y, uint x, float2 v) { shared_block[y][x] = half2(v); } +#else +groupshared uint shared_block[(BLOCK_SIZE + 2 * APRON) / 2][(BLOCK_SIZE + 2 * APRON) + 1]; +float2 load_shared(uint y, uint x) { return unpackHalf2x16(shared_block[y][x]); } +void store_shared(uint y, uint x, float2 v) { shared_block[y][x] = packHalf2x16(v); } +#endif + +bool2 band(bool2 a, bool2 b) +{ + return bool2(a.x && b.x, a.y && b.y); +} diff --git a/crates/nespyro/shaders/idwt.slang b/crates/nespyro/shaders/idwt.slang new file mode 100644 index 00000000..abdece57 --- /dev/null +++ b/crates/nespyro/shaders/idwt.slang @@ -0,0 +1,176 @@ +// Inverse CDF 9/7 wavelet transform: four bands of one level back into a +// 32×32 tile of the level above. +// +// Ported from PyroWave's idwt.comp (Copyright (c) 2025 Hans-Kristian Arntzen, +// MIT), the compute path. Upstream also has a fragment-shader path for +// mobile GPUs; this crate targets desktop GPUs and does not carry it. +// +// Everything here is transposed: the input is sampled with x and y swapped, +// dispatched over the transposed resolution, and written back swapped, which +// lets the same lifting code serve both directions. + +import glsl; +#include "dwt_common.slang" + +[[vk::binding(0, 0)]] Sampler2DArray uTexture; +[[vk::binding(1, 0)]] [[vk::image_format("unknown")]] WTexture2D uOutput; + +// The final transform of a plane moves it back from [-0.5, 0.5] to [0, 1]. +[[vk::constant_id(0)]] const bool DCShift = false; + +struct Registers +{ + int2 resolution; + float2 inv_resolution; +}; +[[vk::push_constant]] ConstantBuffer registers; + +static uint local_index; + +float2 generate_mirror_uv(int2 coord, bool even_x, bool even_y) +{ + coord -= int2(band(bool2(even_x, even_y), coord < int2(0))); + coord += 1; + coord += int2(band(bool2(!even_x, !even_y), coord >= registers.resolution)); + float2 uv = float2(coord) * registers.inv_resolution; + return uv.yx; // Transpose on load. +} + +void write_shared_4x4(int2 coord, float4 texels0, float4 texels1, float4 texels2, float4 texels3) +{ + store_shared(coord.y + 0, 2 * coord.x + 0, float2(texels0.x, texels2.x)); + store_shared(coord.y + 0, 2 * coord.x + 1, float2(texels1.x, texels3.x)); + store_shared(coord.y + 0, 2 * coord.x + 2, float2(texels0.y, texels2.y)); + store_shared(coord.y + 0, 2 * coord.x + 3, float2(texels1.y, texels3.y)); + store_shared(coord.y + 1, 2 * coord.x + 0, float2(texels0.z, texels2.z)); + store_shared(coord.y + 1, 2 * coord.x + 1, float2(texels1.z, texels3.z)); + store_shared(coord.y + 1, 2 * coord.x + 2, float2(texels0.w, texels2.w)); + store_shared(coord.y + 1, 2 * coord.x + 3, float2(texels1.w, texels3.w)); +} + +// The four bands of the 2×2 quad at `coord`: LL, HL, LH and HH are layers 0, +// 1, 2 and 3, and the parity of each band decides which way its edge mirrors. +void load_quad(int2 coord) +{ + float4 texels0 = uTexture.Gather(float3(generate_mirror_uv(coord, true, true), 0.0)).wxzy; + float4 texels1 = uTexture.Gather(float3(generate_mirror_uv(coord, false, true), 2.0)).wxzy; + float4 texels2 = uTexture.Gather(float3(generate_mirror_uv(coord, true, false), 1.0)).wxzy; + float4 texels3 = uTexture.Gather(float3(generate_mirror_uv(coord, false, false), 3.0)).wxzy; + write_shared_4x4(coord - (int2(gl_WorkGroupID.xy) * int2(BLOCK_SIZE_HALF) - APRON_HALF), + texels0, texels1, texels2, texels3); +} + +void load_image_with_apron() +{ + int2 base_coord = int2(gl_WorkGroupID.xy) * int2(BLOCK_SIZE_HALF) - APRON_HALF; + int2 local_coord0 = 2 * unswizzle8x8(local_index); + load_quad(base_coord + local_coord0); + + int2 local_coord_horiz = int2(BLOCK_SIZE_HALF + 2 * int(local_index % 2u), 2 * int(local_index / 2u)); + if (local_coord_horiz.y < BLOCK_SIZE_HALF + 2 * APRON_HALF) + load_quad(base_coord + local_coord_horiz); + + int2 local_coord_vert = local_coord_horiz.yx; + if (local_coord_vert.x < BLOCK_SIZE_HALF) + load_quad(base_coord + local_coord_vert); + + barrier(); +} + +// The inverse lifting steps over one padded line, then a transposing write +// back, so the next pass lifts the other direction. +void lift_and_transpose(bool active_lane, int2 local_coord) +{ + const int PADDED_SIZE = SIZE + 2 * APRON; + const int PADDED_SIZE_HALF = PADDED_SIZE / 2; + // Zeroed only so the compiler can see it is never read uninitialised; + // inactive lanes write nothing back. + float2 values[SIZE + 2 * APRON] = {}; + + if (active_lane) + { + for (int i = 0; i < PADDED_SIZE; i += 2) + { + float2 v0 = load_shared(local_coord.y, local_coord.x + i + 0); + float2 v1 = load_shared(local_coord.y, local_coord.x + i + 1); + values[i + 0] = v0 * K; + values[i + 1] = v1 * inv_K; + } + + for (int i = 2; i < PADDED_SIZE - 1; i += 2) + values[i] -= DELTA * (values[i - 1] + values[i + 1]); + for (int i = 3; i < PADDED_SIZE - 2; i += 2) + values[i] -= GAMMA * (values[i - 1] + values[i + 1]); + for (int i = 4; i < PADDED_SIZE - 3; i += 2) + values[i] -= BETA * (values[i - 1] + values[i + 1]); + for (int i = 5; i < PADDED_SIZE - 4; i += 2) + values[i] -= ALPHA * (values[i - 1] + values[i + 1]); + } + + // Everyone has read before anyone writes. + barrier(); + + if (active_lane) + { + for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++) + { + float2 a = values[2 * i + 0]; + float2 b = values[2 * i + 1]; + + // Transpose the 2×2 block. + float2 t0 = float2(a.x, b.x); + float2 t1 = float2(a.y, b.y); + + int y_coord = (local_coord.x >> 1) + (i - APRON_HALF); + store_shared(y_coord, 2 * local_coord.y + 0, t0); + store_shared(y_coord, 2 * local_coord.y + 1, t1); + } + } +} + +void inverse_transform8x2() +{ + lift_and_transpose<8>(true, int2(8 * int(local_index % 4u), int(local_index / 4u))); +} + +void inverse_transform4x2(bool active_lane, int y_offset) +{ + lift_and_transpose<4>(active_lane, int2(4 * int(local_index % 8u), int(local_index / 8u) + y_offset)); +} + +[shader("compute")] +[numthreads(64, 1, 1)] +void idwt() +{ + local_index = gl_LocalInvocationIndex; + + load_image_with_apron(); + + // Horizontal transform. + inverse_transform8x2(); + + // The apron needs transforming too. + inverse_transform4x2(local_index < 32, BLOCK_SIZE_HALF); + + barrier(); + + // Vertical transform. + inverse_transform8x2(); + + barrier(); + + int2 local_coord = unswizzle8x8(local_index); + + for (int y = local_coord.y; y < BLOCK_SIZE_HALF; y += 8) + { + for (int x = local_coord.x; x < BLOCK_SIZE; x += 8) + { + float2 v = load_shared(y, x); + if (DCShift) + v += 0.5; + int2 base = BLOCK_SIZE * int2(gl_WorkGroupID.yx); + uOutput.Store(int2(2 * y + 0, x) + base, v.xxxx); + uOutput.Store(int2(2 * y + 1, x) + base, v.yyyy); + } + } +} diff --git a/crates/nespyro/shaders/resolve_rate_control.slang b/crates/nespyro/shaders/resolve_rate_control.slang new file mode 100644 index 00000000..b3ec1519 --- /dev/null +++ b/crates/nespyro/shaders/resolve_rate_control.slang @@ -0,0 +1,79 @@ +// Rate control, second half: spend the cheapest buckets until the frame fits. +// +// Ported from PyroWave's resolve_rate_control.comp (Copyright (c) 2025 +// Hans-Kristian Arntzen, MIT). One workgroup per bucket and block-space +// subdivision, in order of rising distortion per byte. A workgroup whose +// bucket the earlier ones already cover does nothing; the one where the +// running total crosses the target applies just enough of its operations. +// Applying an operation raises its block's quant, so a block ends at the +// hardest quant any bucket asked of it. +// +// The workgroup is exactly one subgroup. Upstream sets its size with a +// specialisation constant; here it is compiled once per size as WG_SIZE, and +// the pipeline requires that subgroup size. + +import glsl; +#include "common.slang" + +struct Registers +{ + uint *buckets; + int *total_savings_per_bucket; + RDOperation *rdo_operations; + int *quant_data; + // What changes per frame, written by the frame's own commands so the + // rest can be recorded once: [1] is the target size in words. + uint *frame_params; + uint num_blocks_per_subdivision; +}; +[[vk::push_constant]] ConstantBuffer registers; + +[shader("compute")] +[numthreads(WG_SIZE, 1, 1)] +void resolve_rate_control() +{ + int required_savings_per_bucket = int(registers.buckets[1]) - int(registers.frame_params[1]); + if (gl_WorkGroupID.x != 0) + { + int prev_bucket_total = registers.total_savings_per_bucket[gl_WorkGroupID.x - 1]; + // This bucket is empty. + if (registers.total_savings_per_bucket[gl_WorkGroupID.x] == prev_bucket_total) + return; + + required_savings_per_bucket -= prev_bucket_total; + } + else + { + // This bucket is empty. + if (registers.total_savings_per_bucket[gl_WorkGroupID.x] == 0) + return; + } + + // Earlier buckets already do the job. + if (required_savings_per_bucket <= 0) + return; + + uint total_saved = 0; + + for (uint i = 0; i < registers.num_blocks_per_subdivision && total_saved < uint(required_savings_per_bucket); + i += gl_SubgroupSize) + { + RDOperation op = { 0, 0 }; + if (i + gl_SubgroupInvocationID < registers.num_blocks_per_subdivision) + op = registers.rdo_operations[gl_WorkGroupID.x * registers.num_blocks_per_subdivision + i + gl_SubgroupInvocationID]; + + uint saving = bitfieldExtract(op.block_offset_saving, 16, 16); + uint block_offset = bitfieldExtract(op.block_offset_saving, 0, 16); + + uint scan_saving = subgroupInclusiveAdd(saving); + + bool should_apply_quant = total_saved + scan_saving - saving < uint(required_savings_per_bucket); + if (should_apply_quant && saving != 0) + { + int ignored; + InterlockedMax(registers.quant_data[block_offset], op.quant, ignored); + } + + total_saved += subgroupShuffle(scan_saving, gl_SubgroupSize - 1); + } +} diff --git a/crates/nespyro/shaders/rgb_to_ycbcr.slang b/crates/nespyro/shaders/rgb_to_ycbcr.slang new file mode 100644 index 00000000..848ce804 --- /dev/null +++ b/crates/nespyro/shaders/rgb_to_ycbcr.slang @@ -0,0 +1,153 @@ +// The caller's RGB image to the three full-range YCbCr planes PyroWave codes. +// +// Not a port: upstream's C API does this in Granite's video scaler, which is +// not part of PyroWave. The colour maths follows pixelforge's +// color_convert.comp, so a stream looks the same whichever of the two +// encoded it, with one difference: output is always full range. Limited +// range is never produced. +// +// Each lane converts a 2×2 quad. For 4:2:0 the quad's four chroma samples +// are averaged into one, which sites chroma at the centre of the quad; the +// sequence header says so. + +import glsl; + +[[vk::binding(0, 0)]] Texture2D uInput; +[[vk::binding(1, 0)]] [[vk::image_format("unknown")]] WTexture2D uY; +[[vk::binding(2, 0)]] [[vk::image_format("unknown")]] WTexture2D uCb; +[[vk::binding(3, 0)]] [[vk::image_format("unknown")]] WTexture2D uCr; + +[[vk::constant_id(0)]] const bool Chroma420 = true; + +// What the input already is. The discriminants match nespyro's `Source`. +static const uint SOURCE_SRGB = 0; +static const uint SOURCE_BT709_LINEAR = 1; +static const uint SOURCE_BT2020_LINEAR = 2; +static const uint SOURCE_BT2020_PQ = 3; + +struct Registers +{ + uint width; + uint height; + uint source; + // Zero for SDR BT.709, one for HDR10: BT.2020 NCL with PQ. + uint hdr; + // Luminance of source value 1.0, in nits. Unused when the source is PQ. + float reference_white_nits; +}; +[[vk::push_constant]] ConstantBuffer registers; + +static const float PQ_M1 = 0.1593017578125; +static const float PQ_M2 = 78.84375; +static const float PQ_C1 = 0.8359375; +static const float PQ_C2 = 18.8515625; +static const float PQ_C3 = 18.6875; + +// Linear light, 1.0 = 10 000 nits, to a PQ signal. +float3 linear_to_pq(float3 l) +{ + l = max(l, float3(0.0)); + float3 lm1 = pow(l, float3(PQ_M1)); + return pow((PQ_C1 + PQ_C2 * lm1) / (1.0 + PQ_C3 * lm1), float3(PQ_M2)); +} + +float3 srgb_to_linear(float3 srgb) +{ + // Clamped first: FP16 sources can carry values pow() turns into NaN. + srgb = clamp(srgb, 0.0, 1.0); + return select(srgb <= float3(0.04045), srgb / 12.92, pow((srgb + 0.055) / 1.055, float3(2.4))); +} + +// BT.709 to BT.2020 primaries, in linear light. +float3 bt709_to_bt2020(float3 c) +{ + return float3( + 0.6274 * c.r + 0.3293 * c.g + 0.0433 * c.b, + 0.0691 * c.r + 0.9195 * c.g + 0.0114 * c.b, + 0.0164 * c.r + 0.0880 * c.g + 0.8956 * c.b); +} + +// The signal to put through the matrix. +float3 read_rgb(int2 coord) +{ + float3 rgb = uInput.Load(int3(coord, 0)).rgb; + + if (registers.hdr != 0) + { + if (registers.source == SOURCE_BT2020_PQ) + return rgb; + + float3 l = rgb; + if (registers.source == SOURCE_SRGB) + l = srgb_to_linear(l); + if (registers.source != SOURCE_BT2020_LINEAR) + l = bt709_to_bt2020(l); + return linear_to_pq(l * (registers.reference_white_nits / 10000.0)); + } + + // An SDR target only ever sees an sRGB source; the host refuses the rest. + return rgb; +} + +// Full range: Y in [0, 1], Cb and Cr in [0, 1] centred on 0.5. +float3 matrix(float3 rgb) +{ + float3 y_row, cb_row, cr_row; + if (registers.hdr != 0) + { + // BT.2020 NCL. + y_row = float3(0.2627, 0.6780, 0.0593); + cb_row = float3(-0.1396, -0.3604, 0.5000); + cr_row = float3(0.5000, -0.4598, -0.0402); + } + else + { + // BT.709. + y_row = float3(0.2126, 0.7152, 0.0722); + cb_row = float3(-0.1146, -0.3854, 0.5000); + cr_row = float3(0.5000, -0.4542, -0.0458); + } + float3 ycbcr = float3(dot(y_row, rgb), 0.5 + dot(cb_row, rgb), 0.5 + dot(cr_row, rgb)); + return clamp(ycbcr, 0.0, 1.0); +} + +[shader("compute")] +[numthreads(8, 8, 1)] +void rgb_to_ycbcr() +{ + int2 quad = int2(gl_GlobalInvocationID.xy); + int2 base = 2 * quad; + if (base.x >= int(registers.width) || base.y >= int(registers.height)) + return; + + int2 last = int2(int(registers.width) - 1, int(registers.height) - 1); + float2 chroma_sum = float2(0.0); + + for (int dy = 0; dy < 2; dy++) + { + for (int dx = 0; dx < 2; dx++) + { + int2 p = base + int2(dx, dy); + // A quad past an odd edge repeats the edge pixel, which only a + // 4:4:4 stream can have. + float3 ycbcr = matrix(read_rgb(min(p, last))); + if (all(p <= last)) + { + uY.Store(p, float4(ycbcr.x)); + if (!Chroma420) + { + uCb.Store(p, float4(ycbcr.y)); + uCr.Store(p, float4(ycbcr.z)); + } + } + chroma_sum += ycbcr.yz; + } + } + + if (Chroma420) + { + float2 c = chroma_sum * 0.25; + uCb.Store(quad, float4(c.x)); + uCr.Store(quad, float4(c.y)); + } +} diff --git a/crates/nespyro/shaders/wavelet_dequant.slang b/crates/nespyro/shaders/wavelet_dequant.slang new file mode 100644 index 00000000..f0290eb0 --- /dev/null +++ b/crates/nespyro/shaders/wavelet_dequant.slang @@ -0,0 +1,289 @@ +// Turn one band's received blocks back into wavelet coefficients. +// +// Ported from PyroWave's wavelet_dequant.comp (Copyright (c) 2025 +// Hans-Kristian Arntzen, MIT), storage mode 0: the payload is read straight +// from a buffer with 8- and 16-bit loads. Upstream's texel-buffer and linear +// image modes exist for mobile GPUs without 8-bit storage; this crate targets +// desktop GPUs and does not carry them. +// +// A workgroup of 128 lanes decodes one 32×32 block, 8 coefficients a lane. A +// block whose offset is ~0 never arrived and decodes as zeros. + +import glsl; +#include "common.slang" + +[[vk::binding(0, 0)]] [[vk::image_format("unknown")]] WTexture2DArray uDequantImg; + +struct Registers +{ + uint *payload_offsets; + // One payload buffer, addressed as words, halves and bytes. + uint *payload_data_u32; + uint16_t *payload_data_u16; + uint8_t *payload_data_u8; + int2 resolution; + int output_layer; + int block_offset_32x32; + int block_stride_32x32; +}; +[[vk::push_constant]] ConstantBuffer registers; + +float4 decode_payload(uint code_word, uint q_bits, uint offset, uint block_index, out float4 m1) +{ + bool empty_block = code_word == 0; + if (empty_block) + { + m1 = float4(0.0); + return float4(0.0); + } + + int bit_offset = 2 * int(block_index); + + // Where this 4×2 subblock's planes start. + uint lsbs = code_word & 0x5555u; + uint msbs = code_word & 0xaaaau; + uint msbs_shift = msbs >> 1; + msbs |= msbs_shift; + + uint byte_offset = + bitCount(bitfieldExtract(lsbs, 0, bit_offset)) + + bitCount(bitfieldExtract(msbs, 0, bit_offset)) + + q_bits * block_index + offset; + + // Load early to keep latency down. + uint payload = uint(registers.payload_data_u8[byte_offset]); + + uint local_control_word = bitfieldExtract(code_word, bit_offset, 2); + int decoded_abs[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + int plane_iterations = int(q_bits + local_control_word); + + for (int q = plane_iterations - 1; q >= 0; q--) + { + for (int b = 0; b < 8; b++) + { + int decoded = int(bitfieldExtract(payload, b, 1)); + decoded_abs[b] = bitfieldInsert(decoded_abs[b], decoded, q, 1); + } + byte_offset++; + payload = uint(registers.payload_data_u8[byte_offset]); + } + + // Columns 0 and 1 of upstream's mat2x4: coefficient i*2+j goes to column + // j, row i. + float4 m[2]; + for (int i = 0; i < 4; i++) + { + for (int j = 0; j < 2; j++) + { + float v = float(decoded_abs[i * 2 + j]); + if (v != 0.0) + v += 0.5; + m[j][i] = v; + } + } + + m1 = m[1]; + return m[0]; +} + +groupshared uint shared_sign_offset; +groupshared uint shared_plane_byte_offsets[16]; +groupshared uint shared_sign_scan[128 / 4]; + +uint scan_subgroups(uint v) +{ + for (uint i = 1; i < gl_NumSubgroups; i *= 2) + { + uint up = subgroupShuffleUp(v, i); + v += gl_SubgroupInvocationID >= i ? up : 0; + } + return v; +} + +void scan_subgroups_fallback(uint local_index) +{ + barrier(); + + // Slow shared-memory fallback for subgroups narrower than 16. + bool active_lane = local_index < gl_NumSubgroups; + + uint v = 0; + if (active_lane) + v = shared_sign_scan[local_index]; + + for (uint i = 1; i < gl_NumSubgroups; i *= 2) + { + uint up = 0; + bool do_work = local_index >= i && active_lane; + if (do_work) + up = shared_sign_scan[local_index - i]; + + // Resolve the write-after-read hazard. + barrier(); + + if (do_work) + { + v += up; + shared_sign_scan[local_index] = v; + } + + barrier(); + } +} + +[shader("compute")] +[numthreads(128, 1, 1)] +void wavelet_dequant() +{ + uint local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID; + + int block_index_32x32 = int(uint(registers.block_offset_32x32) + + gl_WorkGroupID.y * uint(registers.block_stride_32x32) + + gl_WorkGroupID.x); + + uint block_local_index = bitfieldExtract(local_index, 0, 3); + uint block_x = bitfieldExtract(local_index, 3, 2); + uint block_y = bitfieldExtract(local_index, 5, 2); + uint linear_block = block_y * 4 + block_x; + + // Each lane decodes 8 values. + int2 local_coord = unswizzle8x8(block_local_index << 3); + + int2 coord = int2(gl_WorkGroupID.xy) * 32; + coord += 8 * int2(int(block_x), int(block_y)); + coord += local_coord; + + uint offset_u32 = registers.payload_offsets[block_index_32x32]; + + if (offset_u32 == ~0u) + { + for (int j = 0; j < 2; j++) + for (int i = 0; i < 4; i++) + uDequantImg.Store(int3(coord + int2(i, j), registers.output_layer), float4(0.0)); + return; + } + + uint ballot = registers.payload_data_u32[offset_u32] & 0xffff; + uint q_code = registers.payload_data_u32[offset_u32 + 1] & 0xff; + + if (local_index < 16) + { + uint control_word = 0; + uint q_bits = 0; + + if (bitfieldExtract(ballot, int(local_index), 1) != 0) + { + uint local_code_offset = bitCount(bitfieldExtract(ballot, 0, int(local_index))); + control_word = uint(registers.payload_data_u16[offset_u32 * 2 + 4 + local_code_offset]); + q_bits = uint(registers.payload_data_u8[offset_u32 * 4 + 8 + bitCount(ballot) * 2 + local_code_offset]) & 0xfu; + } + + uint lsbs = control_word & 0x5555u; + uint msbs = control_word & 0xaaaau; + uint msbs_shift = msbs >> 1; + msbs |= msbs_shift; + uint byte_cost = bitCount(lsbs) + bitCount(msbs) + q_bits * 8; + + uint byte_scan = offset_u32 * 4 + 8 + 3 * bitCount(ballot) + subgroupInclusiveAdd(byte_cost); + if (local_index == 15) + shared_sign_offset = 8 * byte_scan; + shared_plane_byte_offsets[local_index] = byte_scan - byte_cost; + } + + barrier(); + + float4 v0, v1; + int significant_count; + + if (bitfieldExtract(ballot, int(linear_block), 1) != 0) + { + uint local_code_offset = bitCount(bitfieldExtract(ballot, 0, int(linear_block))); + + uint control_word = uint(registers.payload_data_u16[offset_u32 * 2 + 4 + local_code_offset]); + uint control_word2 = uint(registers.payload_data_u8[offset_u32 * 4 + 8 + bitCount(ballot) * 2 + local_code_offset]); + + v0 = decode_payload(control_word, control_word2 & 0xfu, + shared_plane_byte_offsets[linear_block], block_local_index, v1); + + significant_count = 0; + for (int i = 0; i < 4; i++) + significant_count += int(v0[i] != 0.0) + int(v1[i] != 0.0); + + float q = decode_quant(q_code); + float inv_scale = q * decode_quant_scale(bitfieldExtract(control_word2, QUANT_SCALE_OFFSET - 16, QUANT_SCALE_BITS)); + + v0 *= inv_scale; + v1 *= inv_scale; + } + else + { + v0 = float4(0.0); + v1 = float4(0.0); + significant_count = 0; + } + + // How many significant coefficients come before this lane. + int significant_scan = subgroupInclusiveAdd(significant_count); + if (gl_SubgroupInvocationID == gl_SubgroupSize - 1) + shared_sign_scan[gl_SubgroupID] = uint(significant_scan); + + if (gl_NumSubgroups <= 8) + { + barrier(); + if (gl_SubgroupSize <= 32) + { + // More robust: not every compiler understands the shuffle-up pattern. + if (local_index < gl_NumSubgroups) + shared_sign_scan[local_index] = subgroupInclusiveAdd(shared_sign_scan[local_index]); + } + else + { + if (local_index < gl_NumSubgroups) + shared_sign_scan[local_index] = scan_subgroups(shared_sign_scan[local_index]); + } + barrier(); + } + else + { + scan_subgroups_fallback(local_index); + } + + // Where this lane's sign bits start. + uint sign_offset = shared_sign_offset + uint(significant_scan - significant_count); + if (gl_SubgroupID != 0) + sign_offset += shared_sign_scan[gl_SubgroupID - 1]; + + // Every sign bit this lane could need. + uint sign_word = registers.payload_data_u32[sign_offset / 32 + 0]; + uint sign_word_upper = registers.payload_data_u32[sign_offset / 32 + 1]; + + uint masked_sign_offset = sign_offset & 31u; + if (masked_sign_offset != 0) + { + sign_word >>= masked_sign_offset; + sign_word |= sign_word_upper << (32 - masked_sign_offset); + } + + int sign_counter = 0; + + // Clock out the sign bits, in upstream's order: row i, then column j. + for (int i = 0; i < 4; i++) + { + if (v0[i] != 0.0) + { + v0[i] *= 1.0 - 2.0 * float(bitfieldExtract(sign_word, sign_counter, 1)); + sign_counter++; + } + if (v1[i] != 0.0) + { + v1[i] *= 1.0 - 2.0 * float(bitfieldExtract(sign_word, sign_counter, 1)); + sign_counter++; + } + } + + for (int i = 0; i < 4; i++) + { + uDequantImg.Store(int3(coord + int2(i, 0), registers.output_layer), float4(v0[i])); + uDequantImg.Store(int3(coord + int2(i, 1), registers.output_layer), float4(v1[i])); + } +} diff --git a/crates/nespyro/shaders/wavelet_quant.slang b/crates/nespyro/shaders/wavelet_quant.slang new file mode 100644 index 00000000..818d44b7 --- /dev/null +++ b/crates/nespyro/shaders/wavelet_quant.slang @@ -0,0 +1,317 @@ +// Quantize one band's coefficients and record, for every 8×8 block, what each +// possible further quantization would cost and save. +// +// Ported from PyroWave's wavelet_quant.comp (Copyright (c) 2025 Hans-Kristian +// Arntzen, MIT). A workgroup of 128 lanes covers a 32×32 region; each lane +// holds 8 coefficients and 8 lanes form an 8×8 block, so an 8×8 block's +// reductions are clustered subgroup operations of width 8. +// +// Bit planes go to a scratch payload in their final order, but sign bits are +// kept per 4×2 subblock for now: which coefficients stay significant is not +// known until rate control has chosen how hard to quantize. +// +// One addition: the scratch payload has a capacity. Upstream sizes it by +// estimate and trusts content never to exceed it. Here it is reached through +// a raw device address with no bounds checking, so an 8×8 block whose +// allocation would not fit is recorded as all zero instead of being written; +// the counter still counts every byte asked for, which is how the host sees +// it happened. + +import glsl; +#include "common.slang" + +[[vk::binding(0, 0)]] Sampler2DArray uTexture; + +[[vk::constant_id(1)]] const bool SkipQuantScale = false; + +struct Registers +{ + BlockMeta *block_meta; + BlockStats *block_stats; + // Two counters, then the scratch payload bytes from offset 8. + uint *payload_counter; + uint8_t *payload_data; + int2 resolution; + int2 resolution_8x8_blocks; + float2 inv_resolution; + float input_layer; + float quant_resolution; + int block_offset; + int block_stride; + float rdo_distortion_scale; + uint payload_capacity; +}; +[[vk::push_constant]] ConstantBuffer registers; + +float max4(float4 v) +{ + float2 v2 = max(v.xy, v.zw); + return max(v2.x, v2.y); +} + +int scan_clustered8(int v) +{ + for (uint i = 1; i < 8; i *= 2) + { + int up = subgroupShuffleUp(v, i); + v += (gl_SubgroupInvocationID & 7u) >= i ? up : 0; + } + return v; +} + +void compute_quant_scale(float max_wave_texels, out uint quant_code, out float quant_scale) +{ + if (SkipQuantScale || max_wave_texels < 1.0) + { + quant_code = ENCODE_QUANT_IDENTITY; + quant_scale = 1.0; + } + else + { + int e; + frexp(max_wave_texels - 0.25, e); + float target_max = float(1 << e) - 0.25; + float inv_scale = max_wave_texels / target_max; + quant_code = encode_quant_scale(inv_scale); + quant_scale = 1.0 / decode_quant_scale(quant_code); + } +} + +float compute_square_error(float4 v[2], int q, out uint num_significant_values) +{ + v[0] = abs(v[0]); + v[1] = abs(v[1]); + float4 iv[2] = { floor(ldexp(v[0], int4(-q))), trunc(ldexp(v[1], int4(-q))) }; + num_significant_values = 0; + for (int j = 0; j < 2; j++) + for (int i = 0; i < 4; i++) + if (iv[j][i] != 0.0) + num_significant_values++; + iv[0] += select(iv[0] != float4(0.0), float4(0.5), float4(0.0)); + iv[1] += select(iv[1] != float4(0.0), float4(0.5), float4(0.0)); + iv[0] = trunc(ldexp(iv[0], int4(q))); + iv[1] = trunc(ldexp(iv[1], int4(q))); + float4 err0 = v[0] - iv[0]; + float4 err1 = v[1] - iv[1]; + num_significant_values = subgroupClusteredAdd(num_significant_values, 8u); + return (dot(err0, err0) + dot(err1, err1)) * registers.rdo_distortion_scale; +} + +struct QuantResult +{ + float square_error; + int encode_cost_early; + int block4x2_shifted; + int encode_cost_late_bits; + int quality_planes; +}; + +QuantResult compute_quant_stats(float4 v[2], int q, int msb, int block4x2_max, float inv_quant_squared) +{ + block4x2_max >>= q; + + uint wave8_num_significants; + QuantResult result; + + result.square_error = compute_square_error(v, q, wave8_num_significants) * inv_quant_squared; + result.block4x2_shifted = block4x2_max; + + result.encode_cost_early = block4x2_max > 0 ? 1 : 0; + msb -= q; + + result.quality_planes = 0; + + if (msb >= 3) + { + result.quality_planes = msb - 2; + // The sign plane must be coded if there are quality planes. + result.encode_cost_early = result.quality_planes + 1; + result.block4x2_shifted >>= result.quality_planes; + } + + result.encode_cost_early += findMSB(result.block4x2_shifted) + 1; + result.encode_cost_late_bits = + 8 * subgroupClusteredAdd(max(result.encode_cost_early - 1, 0), 8u) + int(wave8_num_significants); + return result; +} + +QuantStats quant_stats(float square_error, uint cost) +{ + QuantStats s; + s.square_error = half(square_error); + s.payload_cost = uint16_t(cost); + return s; +} + +BlockMeta block_meta(uint code_word, uint offset) +{ + BlockMeta m; + m.code_word = code_word; + m.offset = offset; + return m; +} + +void encode_payload(int2 block_index_8x8, float4 texels[2]) +{ + precise float max_subblock_texel = max(max4(abs(texels[0])), max4(abs(texels[1]))); + precise float max_wave_texels = subgroupClusteredMax(max_subblock_texel, 8u); + float quant_scale; + uint quant_code; + compute_quant_scale(max_wave_texels, quant_code, quant_scale); + texels[0] *= quant_scale; + texels[1] *= quant_scale; + max_wave_texels *= quant_scale; + max_subblock_texel *= quant_scale; + + float overall_quant_scale = registers.quant_resolution * quant_scale; + float inv_quant = 1.0 / overall_quant_scale; + float inv_quant_squared = inv_quant * inv_quant; + int4 abs_quant_texels0 = abs(int4(texels[0])); + int4 abs_quant_texels1 = abs(int4(texels[1])); + int max_absolute_value = int(max_wave_texels); + int block4x2_max = int(max_subblock_texel); + + uint block_index = uint(registers.block_offset + block_index_8x8.y * registers.block_stride + block_index_8x8.x); + + // The entire block quantizes to zero. + if (max_absolute_value == 0) + { + if ((gl_SubgroupInvocationID & 7) == 0) + { + registers.block_meta[block_index] = block_meta(0, 0); + registers.block_stats[block_index].num_planes = 0; + registers.block_stats[block_index].errors[0] = quant_stats(0.0, 0); + } + return; + } + + int msb = findMSB(max_absolute_value); + + QuantResult result = compute_quant_stats(texels, 0, msb, block4x2_max, inv_quant_squared); + int scan = scan_clustered8(result.encode_cost_early); + + uint global_offset = 0; + + // Allocate the block's payload, and report where it went. + if ((gl_SubgroupInvocationID & 7u) == 7u) + InterlockedAdd(registers.payload_counter[0], uint(scan), global_offset); + global_offset = subgroupShuffle(global_offset, gl_SubgroupInvocationID | 7u); + uint block_bytes = uint(subgroupShuffle(scan, gl_SubgroupInvocationID | 7u)); + + // Past the scratch capacity: the block is dropped whole, as if it had + // quantized to zero. All 8 lanes of the block agree. + if (global_offset + block_bytes > registers.payload_capacity) + { + if ((gl_SubgroupInvocationID & 7) == 0) + { + registers.block_meta[block_index] = block_meta(0, 0); + registers.block_stats[block_index].num_planes = 0; + registers.block_stats[block_index].errors[0] = quant_stats(0.0, 0); + } + return; + } + + scan -= result.encode_cost_early; + + // First, the code word. + int quality_planes = result.quality_planes; + uint code_word = uint(quality_planes) << Q_PLANES_OFFSET; + code_word = bitfieldInsert(code_word, quant_code, QUANT_SCALE_OFFSET, QUANT_SCALE_BITS); + uint plane_code = uint(findMSB(result.block4x2_shifted) + 1); + + uint merged_plane_code = plane_code << ((gl_SubgroupInvocationID & 7u) * 2u); + merged_plane_code |= shuffle_xor(merged_plane_code, 1u); + merged_plane_code |= shuffle_xor(merged_plane_code, 2u); + merged_plane_code |= shuffle_xor(merged_plane_code, 4u); + code_word |= merged_plane_code; + + if ((gl_SubgroupInvocationID & 7u) == 0u) + { + registers.block_meta[block_index] = block_meta(code_word, global_offset); + registers.block_stats[block_index].num_planes = uint(msb + 1); + // No distortion for quant 0: that decision has been made already. + registers.block_stats[block_index].errors[0] = quant_stats(0.0, uint(result.encode_cost_late_bits)); + } + + for (int q = 1; q <= msb; q++) + { + QuantResult quant_result = compute_quant_stats(texels, q, msb, block4x2_max, inv_quant_squared); + float square_error = subgroupClusteredAdd(quant_result.square_error, 8u); + + if ((gl_SubgroupInvocationID & 7u) == 0) + { + registers.block_stats[block_index].errors[q] = + quant_stats(min(square_error, 60000.0), uint(quant_result.encode_cost_late_bits)); + } + } + + // The distortion of throwing everything away. + float square_error = subgroupClusteredAdd( + (dot(texels[0], texels[0]) + dot(texels[1], texels[1])) * inv_quant_squared, 8u); + if ((gl_SubgroupInvocationID & 7u) == 0) + registers.block_stats[block_index].errors[msb + 1] = quant_stats(min(60000.0, square_error), 0); + + uint byte_offset = uint(scan) + global_offset; + bool need_sign = result.block4x2_shifted != 0 || quality_planes != 0; + + // Signs are not packed until the final pass, when the quantization is known. + if (need_sign) + { + uint4 s0 = uint4(texels[0] < float4(0.0)) << uint4(0, 1, 2, 3); + uint4 s1 = uint4(texels[1] < float4(0.0)) << uint4(4, 5, 6, 7); + uint s = s0.x | s0.y | s0.z | s0.w | s1.x | s1.y | s1.z | s1.w; + registers.payload_data[byte_offset++] = uint8_t(s); + + int plane_iterations = quality_planes + int(plane_code); + int q = plane_iterations - 1; + do + { + s0 = uint4( + bitfieldExtract(uint(abs_quant_texels0.x), q, 1), + bitfieldExtract(uint(abs_quant_texels0.y), q, 1), + bitfieldExtract(uint(abs_quant_texels0.z), q, 1), + bitfieldExtract(uint(abs_quant_texels0.w), q, 1)); + s1 = uint4( + bitfieldExtract(uint(abs_quant_texels1.x), q, 1), + bitfieldExtract(uint(abs_quant_texels1.y), q, 1), + bitfieldExtract(uint(abs_quant_texels1.z), q, 1), + bitfieldExtract(uint(abs_quant_texels1.w), q, 1)); + s0 <<= uint4(0, 1, 2, 3); + s1 <<= uint4(4, 5, 6, 7); + s = s0.x | s0.y | s0.z | s0.w | s1.x | s1.y | s1.z | s1.w; + registers.payload_data[byte_offset++] = uint8_t(s); + q--; + } while (q >= 0); + } +} + +[shader("compute")] +[numthreads(128, 1, 1)] +void wavelet_quant() +{ + uint local_index = gl_SubgroupID * gl_SubgroupSize + gl_SubgroupInvocationID; + uint block_local_index = bitfieldExtract(local_index, 0, 3); + uint block_x = bitfieldExtract(local_index, 3, 2); + uint block_y = bitfieldExtract(local_index, 5, 2); + + // Each lane encodes 8 values. + int2 local_coord = unswizzle8x8(block_local_index << 3); + + int2 coord = int2(gl_WorkGroupID.xy) * 32; + coord += 8 * int2(int(block_x), int(block_y)); + coord += local_coord; + + int2 block_index = 4 * int2(gl_WorkGroupID.xy) + int2(int(block_x), int(block_y)); + + float3 uv = float3(float2(coord) * registers.inv_resolution, registers.input_layer); + float4 texels0 = uTexture.Gather(uv, int2(1, 1)).wxzy; + float4 texels1 = uTexture.Gather(uv, int2(3, 1)).wxzy; + precise float4 scaled_texels0 = texels0 * registers.quant_resolution; + precise float4 scaled_texels1 = texels1 * registers.quant_resolution; + bool in_bounds = all(block_index < registers.resolution_8x8_blocks); + if (in_bounds) + { + float4 scaled[2] = { scaled_texels0, scaled_texels1 }; + encode_payload(block_index, scaled); + } +} diff --git a/crates/nespyro/src/bitstream/depacketize.rs b/crates/nespyro/src/bitstream/depacketize.rs new file mode 100644 index 00000000..f9d79a01 --- /dev/null +++ b/crates/nespyro/src/bitstream/depacketize.rs @@ -0,0 +1,338 @@ +//! Collecting a frame's blocks as its packets arrive. +//! +//! Packets can come in any order, some never, and a newer frame's can arrive +//! before an older frame is done. The 3-bit sequence number decides which frame +//! a packet belongs to, exactly as upstream's decoder does: a packet from an +//! older frame is dropped, one from a newer frame abandons the current one. +//! +//! Every block is checked here, on the CPU, before it can reach the GPU. The +//! dequantizer reads plane bytes and sign bits at offsets computed from each +//! block's own code words, so a block whose contents disagree with its +//! declared size would make it read past its end. This is network input; the +//! check is what keeps a corrupt or hostile packet from becoming an +//! out-of-bounds read on the device. + +use super::{ + BitstreamError, BlockHeader, Chroma, ColourDescription, HEADER_SIZE, Header, Layout, + SEQUENCE_MASK, SequenceHeader, packetize::CRITICAL_BANDS, +}; + +/// The parts of a sequence header that fix the decoder's resources. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StreamInfo { + pub width: u32, + pub height: u32, + pub chroma: Chroma, +} + +/// What became of a packet. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Push { + /// Its blocks were taken. + Accepted, + /// It belongs to a frame older than the one being collected, and was + /// dropped. + Stale, + /// Its sequence header describes a different size or chroma format than + /// this decoder was built for. Nothing was taken; the caller rebuilds the + /// decoder for `StreamInfo` and pushes the packet again. + Reconfigure(StreamInfo), +} + +/// How much of the current frame has arrived. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Readiness { + /// Whether a start-of-frame header has been seen for this frame. + pub has_header: bool, + /// Distinct blocks received. + pub received: u32, + /// Non-empty blocks the frame has, from its header, or every possible + /// block if the header has not arrived. + pub expected: u32, + /// Whether every block of the two coarsest bands is here. + /// + /// A guess, and a pessimistic one: a block the encoder quantized to + /// nothing is never sent, and this cannot tell it from a lost one. The + /// coarse high-pass bands of noisy content are often empty, so this reads + /// false on frames that arrived complete. A transport that knows which + /// packets were critical should say so instead, through + /// [`Readiness::is_complete_enough_with`]. + pub critical_complete: bool, + /// Whether this frame was already decoded. + pub decoded: bool, +} + +impl Readiness { + /// Every block has arrived. + pub fn is_complete(&self) -> bool { + !self.decoded && self.has_header && self.received >= self.expected + } + + /// Upstream's rule for decoding a partial frame: everything, or more than + /// 90% of it with the two coarsest bands intact. Losing blocks from finer + /// bands only blurs their region; losing the coarse ones does not mask. + pub fn is_complete_enough(&self) -> bool { + self.is_complete_enough_with(self.critical_complete) + } + + /// The same rule, with whether the coarse bands are whole decided by the + /// caller: by a transport that knows which packets carried them, which + /// this cannot know. See [`Readiness::critical_complete`]. + pub fn is_complete_enough_with(&self, critical_complete: bool) -> bool { + if self.decoded || !self.has_header { + return false; + } + self.received >= self.expected + || (critical_complete && self.received as f32 > self.expected as f32 * 0.9) + } +} + +/// Collects blocks for one frame at a time. +pub struct Depacketizer { + layout: Layout, + sequence: Option, + header: Option, + /// Word offset of each block within `payload`, or `u32::MAX` if missing. + /// Uploaded as is: the dequantizer reads `u32::MAX` as "decode zeros". + offsets: Vec, + payload: Vec, + received: u32, + decoded: bool, +} + +/// The dequantizer marks a missing block this way. +pub(crate) const MISSING: u32 = u32::MAX; + +impl Depacketizer { + pub fn new(layout: Layout) -> Self { + let blocks = layout.blocks_32x32() as usize; + let mut this = Self { + layout, + sequence: None, + header: None, + offsets: vec![MISSING; blocks], + payload: Vec::with_capacity(256 * 1024), + received: 0, + decoded: false, + }; + this.clear(); + this + } + + pub fn layout(&self) -> &Layout { + &self.layout + } + + /// Takes one packet: any mix of start-of-frame headers and blocks. + /// + /// A packet that fails to parse is rejected whole, but blocks taken from + /// it before the failure stay: each was checked on its own and is sound. + pub fn push(&mut self, packet: &[u8]) -> Result { + let mut rest = packet; + while !rest.is_empty() { + match Header::parse(rest)? { + Header::Sequence(seq) => { + if self.is_stale(seq.sequence) { + return Ok(Push::Stale); + } + let info = StreamInfo { + width: seq.width, + height: seq.height, + chroma: seq.chroma, + }; + if info.width != self.layout.width + || info.height != self.layout.height + || info.chroma != self.layout.chroma + { + return Ok(Push::Reconfigure(info)); + } + self.enter(seq.sequence); + self.header = Some(seq); + rest = &rest[HEADER_SIZE..]; + } + Header::Block(block) => { + let size = usize::from(block.payload_words) * 4; + if size < HEADER_SIZE { + return Err(BitstreamError::BlockTooSmall(block.payload_words.into())); + } + if size > rest.len() { + return Err(BitstreamError::Truncated { + needed: size, + left: rest.len(), + }); + } + if self.is_stale(block.sequence) { + return Ok(Push::Stale); + } + self.enter(block.sequence); + self.take(&block, &rest[..size])?; + rest = &rest[size..]; + } + } + } + Ok(Push::Accepted) + } + + /// Sequence arithmetic modulo 8, as upstream: anything more than half the + /// ring behind is from the past. + fn is_stale(&self, sequence: u8) -> bool { + match self.sequence { + None => false, + Some(current) => { + let diff = sequence.wrapping_sub(current) & SEQUENCE_MASK; + diff > SEQUENCE_MASK / 2 + } + } + } + + /// Makes `sequence` the frame being collected, abandoning any other. + fn enter(&mut self, sequence: u8) { + if self.sequence != Some(sequence) { + self.clear(); + self.sequence = Some(sequence); + } + } + + fn take(&mut self, block: &BlockHeader, bytes: &[u8]) -> Result<(), BitstreamError> { + let count = self.layout.blocks_32x32(); + if block.block_index >= count { + return Err(BitstreamError::BlockIndexOutOfRange { + index: block.block_index, + count, + }); + } + validate_block(&self.layout, block, bytes)?; + + let slot = &mut self.offsets[block.block_index as usize]; + if *slot != MISSING { + // The format allows duplicates, for crude redundancy. + return Ok(()); + } + *slot = self.payload.len() as u32; + let (words, _) = bytes.as_chunks::<4>(); + self.payload + .extend(words.iter().map(|w| u32::from_le_bytes(*w))); + self.received += 1; + Ok(()) + } + + /// Forgets the current frame. + pub fn clear(&mut self) { + self.offsets.fill(MISSING); + self.payload.clear(); + self.header = None; + self.sequence = None; + self.received = 0; + self.decoded = false; + } + + pub fn readiness(&self) -> Readiness { + Readiness { + has_header: self.header.is_some(), + received: self.received, + expected: self + .header + .map_or(self.layout.blocks_32x32(), |h| h.total_blocks), + critical_complete: self.critical_complete(), + decoded: self.decoded, + } + } + + /// Whether every block of the coarsest bands arrived. + /// + /// A block missing because the encoder sent nothing for it is + /// indistinguishable from one lost in transit, so this can reject a frame + /// whose coarse bands were legitimately empty. Upstream accepts the same + /// false rejection; in the coarsest bands an all-zero block is vanishingly + /// rare. + fn critical_complete(&self) -> bool { + let limit = self.layout.leading_blocks(CRITICAL_BANDS) as usize; + self.offsets[..limit].iter().all(|&o| o != MISSING) + } + + /// The colour the current frame's header describes. + pub fn colour(&self) -> Option { + self.header.map(|h| h.colour) + } + + /// Distinct blocks the current frame's header promised but that have not + /// arrived. + pub fn missing_blocks(&self) -> u32 { + self.readiness().expected.saturating_sub(self.received) + } + + /// The block offsets and payload words to upload, and marks the frame + /// decoded so it is not decoded twice. + pub(crate) fn take_for_decode(&mut self) -> (&[u32], &[u32]) { + self.decoded = true; + (&self.offsets, &self.payload) + } +} + +/// Checks that a block's contents are exactly as long as its header says. +/// +/// The walk is upstream's `validate_bitstream`: for every 8×8 block the ballot +/// says is present, its code word gives the bit planes of each 4×2 subblock; +/// the planes' union gives how many coefficients are significant, and so how +/// many sign bits follow. The total, padded to a word, must equal +/// `payload_words`. +pub(crate) fn validate_block( + layout: &Layout, + block: &BlockHeader, + bytes: &[u8], +) -> Result<(), BitstreamError> { + let index = block.block_index; + let mapping = layout + .mapping(index) + .expect("block index was range checked"); + let n = block.ballot.count_ones() as usize; + let declared = u32::from(block.payload_words); + + let mut offset = HEADER_SIZE + 3 * n; + if offset > bytes.len() { + return Err(BitstreamError::BlockSizeMismatch { + index, + computed: offset.div_ceil(4) as u32, + declared, + }); + } + let code_words = &bytes[HEADER_SIZE..HEADER_SIZE + 2 * n]; + let q_scales = &bytes[HEADER_SIZE + 2 * n..HEADER_SIZE + 3 * n]; + + let mut significant = 0u32; + for (i, bit) in (0..16u32) + .filter(|b| block.ballot >> b & 1 != 0) + .enumerate() + { + let (x, y) = (bit & 3, bit >> 2); + if x >= mapping.width_8x8 || y >= mapping.height_8x8 { + return Err(BitstreamError::BallotOutOfBand { index, x, y }); + } + let code_word = u16::from_le_bytes([code_words[2 * i], code_words[2 * i + 1]]); + let q_bits = usize::from(q_scales[i] & 0xf); + for subblock in 0..8 { + let planes = q_bits + usize::from(code_word >> (2 * subblock) & 3); + let Some(plane_bytes) = bytes.get(offset..offset + planes) else { + return Err(BitstreamError::BlockSizeMismatch { + index, + computed: (offset + planes).div_ceil(4) as u32, + declared, + }); + }; + let union = plane_bytes.iter().fold(0u8, |a, &b| a | b); + significant += union.count_ones(); + offset += planes; + } + } + offset += (significant as usize).div_ceil(8); + + let computed = offset.div_ceil(4) as u32; + if computed != declared { + return Err(BitstreamError::BlockSizeMismatch { + index, + computed, + declared, + }); + } + Ok(()) +} diff --git a/crates/nespyro/src/bitstream/layout.rs b/crates/nespyro/src/bitstream/layout.rs new file mode 100644 index 00000000..7475a4d6 --- /dev/null +++ b/crates/nespyro/src/bitstream/layout.rs @@ -0,0 +1,207 @@ +//! Where every block of a frame lives. +//! +//! The block index of a 32×32 block is its position in a fixed walk: levels +//! from coarsest to finest, then Y, Cb, Cr, then band, then row-major within +//! the band (`bitstream.md`, "Block index ordering"). Low frequencies come +//! first, which is what lets a packetizer put the blocks that matter most in +//! the first packets. +//! +//! 8×8 blocks are numbered the same way and index the encoder's per-block +//! statistics. They never appear on the wire. + +use super::{ + ALIGNMENT, BitstreamError, Chroma, DECOMPOSITION_LEVELS, MAXIMUM_IMAGE_SIZE, + MINIMUM_IMAGE_SIZE, NUM_BANDS, NUM_COMPONENTS, +}; + +/// One sub-band of one component at one level. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct BandLayout { + /// Size of the band in coefficients. + pub width: u32, + pub height: u32, + /// Index of the band's first 8×8 block, and 8×8 blocks per row. + pub offset_8x8: u32, + pub stride_8x8: u32, + /// Index of the band's first 32×32 block, 32×32 blocks per row, and how + /// many it has. + pub offset_32x32: u32, + pub stride_32x32: u32, + pub count_32x32: u32, +} + +impl BandLayout { + /// 8×8 blocks across and down. + pub fn blocks_8x8(&self) -> (u32, u32) { + (self.width.div_ceil(8), self.height.div_ceil(8)) + } + + /// 32×32 blocks across and down. + pub fn blocks_32x32(&self) -> (u32, u32) { + (self.width.div_ceil(32), self.height.div_ceil(32)) + } +} + +/// The 8×8 blocks a 32×32 block covers. Edge blocks cover fewer than 4×4. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlockMapping { + pub offset_8x8: u32, + pub stride_8x8: u32, + pub width_8x8: u32, + pub height_8x8: u32, +} + +/// Every band's position in the block numbering, for one frame size and +/// chroma format. +#[derive(Debug, Clone)] +pub struct Layout { + pub width: u32, + pub height: u32, + pub chroma: Chroma, + /// `width` and `height` padded to [`ALIGNMENT`] and at least + /// [`MINIMUM_IMAGE_SIZE`]. + pub aligned_width: u32, + pub aligned_height: u32, + bands: [[[Option; NUM_BANDS]; DECOMPOSITION_LEVELS]; NUM_COMPONENTS], + mapping: Vec, + blocks_8x8: u32, +} + +impl Layout { + pub fn new(width: u32, height: u32, chroma: Chroma) -> Result { + if !(1..=MAXIMUM_IMAGE_SIZE).contains(&width) || !(1..=MAXIMUM_IMAGE_SIZE).contains(&height) + { + return Err(BitstreamError::BadDimensions { width, height }); + } + if chroma == Chroma::Yuv420 && (!width.is_multiple_of(2) || !height.is_multiple_of(2)) { + return Err(BitstreamError::OddYuv420 { width, height }); + } + + let aligned_width = width.next_multiple_of(ALIGNMENT).max(MINIMUM_IMAGE_SIZE); + let aligned_height = height.next_multiple_of(ALIGNMENT).max(MINIMUM_IMAGE_SIZE); + + let mut layout = Self { + width, + height, + chroma, + aligned_width, + aligned_height, + bands: Default::default(), + mapping: Vec::new(), + blocks_8x8: 0, + }; + + for level in (0..DECOMPOSITION_LEVELS).rev() { + for component in 0..NUM_COMPONENTS { + if !chroma.has_level(component, level) { + continue; + } + for band in first_band(level)..NUM_BANDS { + let (w, h) = layout.level_size(level); + let blocks_x_8x8 = w.div_ceil(8); + let blocks_y_8x8 = h.div_ceil(8); + let blocks_x_32x32 = w.div_ceil(32); + let blocks_y_32x32 = h.div_ceil(32); + + layout.bands[component][level][band] = Some(BandLayout { + width: w, + height: h, + offset_8x8: layout.blocks_8x8, + stride_8x8: blocks_x_8x8, + offset_32x32: layout.mapping.len() as u32, + stride_32x32: blocks_x_32x32, + count_32x32: blocks_x_32x32 * blocks_y_32x32, + }); + + for y in 0..blocks_y_32x32 { + for x in 0..blocks_x_32x32 { + layout.mapping.push(BlockMapping { + offset_8x8: layout.blocks_8x8 + 4 * y * blocks_x_8x8 + 4 * x, + stride_8x8: blocks_x_8x8, + width_8x8: 4.min(blocks_x_8x8 - 4 * x), + height_8x8: 4.min(blocks_y_8x8 - 4 * y), + }); + } + } + layout.blocks_8x8 += blocks_x_8x8 * blocks_y_8x8; + } + } + } + + Ok(layout) + } + + /// Size of every band at `level`, in coefficients. Level 0 is half the + /// aligned image; each level halves again. + pub fn level_size(&self, level: usize) -> (u32, u32) { + ( + self.aligned_width >> (level + 1), + self.aligned_height >> (level + 1), + ) + } + + /// The band, or `None` where the format has none: the LL band exists only + /// at the coarsest level, and 4:2:0 chroma has no level 0. + pub fn band(&self, component: usize, level: usize, band: usize) -> Option<&BandLayout> { + self.bands[component][level][band].as_ref() + } + + /// Every band that exists, as `(component, level, band, layout)`, in the + /// order the GPU passes dispatch them: finest level first. + pub fn bands(&self) -> impl Iterator { + (0..DECOMPOSITION_LEVELS).flat_map(move |level| { + (0..NUM_COMPONENTS).flat_map(move |component| { + (0..NUM_BANDS).filter_map(move |band| { + self.band(component, level, band) + .map(|b| (component, level, band, b)) + }) + }) + }) + } + + /// Total 32×32 blocks, the range of valid block indices. + pub fn blocks_32x32(&self) -> u32 { + self.mapping.len() as u32 + } + + /// Total 8×8 blocks. + pub fn blocks_8x8(&self) -> u32 { + self.blocks_8x8 + } + + /// The 8×8 blocks a 32×32 block covers. + pub fn mapping(&self, block_index: u32) -> Option<&BlockMapping> { + self.mapping.get(block_index as usize) + } + + /// How many leading block indices make up the `bands` coarsest bands, by + /// upstream's definition (`get_num_active_blocks`). + /// + /// The range ends at Cr's last band of that depth, and blocks run + /// component by component within a level, so `bands = 1` ends at Cr's LL + /// band and includes Y's and Cb's high-pass bands at the coarsest level + /// with it. From `bands = 2` on it covers whole levels: 2 is all of the + /// coarsest level, 3 adds the next level's high-pass bands, and so on. + pub fn leading_blocks(&self, bands: usize) -> u32 { + assert!(bands < DECOMPOSITION_LEVELS); + if bands == 0 { + return 0; + } + let last_level = DECOMPOSITION_LEVELS - (bands - 1).max(1); + let last_band = if bands == 1 { 0 } else { NUM_BANDS - 1 }; + let b = self + .band(NUM_COMPONENTS - 1, last_level, last_band) + .expect("coarse bands exist for every component"); + b.offset_32x32 + b.count_32x32 + } +} + +/// The LL band is only coded at the coarsest level; every finer level's LL is +/// reconstructed from the level below it. +pub(crate) fn first_band(level: usize) -> usize { + if level == DECOMPOSITION_LEVELS - 1 { + 0 + } else { + 1 + } +} diff --git a/crates/nespyro/src/bitstream/mod.rs b/crates/nespyro/src/bitstream/mod.rs new file mode 100644 index 00000000..e806bd04 --- /dev/null +++ b/crates/nespyro/src/bitstream/mod.rs @@ -0,0 +1,341 @@ +//! PyroWave's bitstream, with no Vulkan in it. +//! +//! Everything here is the format as upstream defines it in +//! `bitstream/bitstream.md` at PyroWave `89f7e47`, byte for byte: a stream this +//! crate writes decodes in upstream's decoder, and one upstream writes decodes +//! here. Nothing is added to it. The format has no version field of its own, so +//! the version travels in nesprotocol instead. +//! +//! A frame is a start-of-frame header followed by independently decodable +//! 32×32 coefficient blocks. A block needs no context from any other block, +//! which is what makes the format loss tolerant: a missing block decodes as +//! zeros and only blurs its own region. + +mod depacketize; +mod layout; +mod packetize; + +pub use depacketize::{Depacketizer, Push, Readiness, StreamInfo}; +pub use layout::{BandLayout, BlockMapping, Layout}; +pub use packetize::{GpuPacket, Packetized, packetize}; + +use thiserror::Error; + +/// Wavelet decomposition levels. Fixed by the format. +pub const DECOMPOSITION_LEVELS: usize = 5; +/// Y, Cb and Cr. +pub const NUM_COMPONENTS: usize = 3; +/// LL, HL, LH and HH. +pub const NUM_BANDS: usize = 4; +/// Image dimensions are padded to a multiple of this, one per decomposition level. +pub const ALIGNMENT: u32 = 1 << DECOMPOSITION_LEVELS; +/// Below this the coarsest band is so small that the mirrored edge extension +/// starts mirroring twice, so smaller images are padded up to it. +pub const MINIMUM_IMAGE_SIZE: u32 = 4 << DECOMPOSITION_LEVELS; +/// Width and height are 14-bit fields holding the value minus one. +pub const MAXIMUM_IMAGE_SIZE: u32 = 1 << 14; +/// The frame counter is three bits wide. +pub const SEQUENCE_MASK: u8 = 0x7; +/// Both header kinds are eight bytes. +pub const HEADER_SIZE: usize = 8; +/// `payload_words` is a 12-bit field. +pub const MAX_BLOCK_WORDS: u32 = (1 << 12) - 1; + +/// The largest a 32×32 block can possibly be, in 32-bit words. +/// +/// Sixteen 8×8 blocks, each with a 2-byte code word, a 1-byte scale, eight +/// 4×2 subblocks of at most 15 + 3 bit planes one byte each, and one sign bit +/// per coefficient, after the 8-byte header. It is far below +/// [`MAX_BLOCK_WORDS`], which is why an encoder never has to cap a block to +/// fit the field. +pub const LARGEST_BLOCK_WORDS: u32 = { + let bytes = HEADER_SIZE as u32 + 16 * (2 + 1 + 8 * 18) + (32 * 32) / 8; + bytes.div_ceil(4) +}; +const _: () = assert!(LARGEST_BLOCK_WORDS <= MAX_BLOCK_WORDS); + +/// Chroma resolution. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Chroma { + /// Chroma at half resolution in both directions. Needs even dimensions. + Yuv420, + /// Chroma at full resolution. + Yuv444, +} + +impl Chroma { + /// Whether this component is present at this decomposition level. + /// + /// 4:2:0 chroma planes are half size, so they enter the transform at level + /// 1 and have no level 0 bands at all. + pub fn has_level(self, component: usize, level: usize) -> bool { + !(self == Chroma::Yuv420 && component != 0 && level == 0) + } +} + +/// Colour primaries, as the sequence header signals them. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Primaries { + Bt709, + Bt2020, +} + +/// Transfer function, as the sequence header signals it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Transfer { + Bt709, + Pq, +} + +/// YCbCr matrix. BT.2020 is the non-constant-luminance variant. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Matrix { + Bt709, + Bt2020Ncl, +} + +/// YCbCr range. +/// +/// This crate only ever writes [`Range::Full`]. `Limited` exists because a +/// stream from elsewhere can say so, and a decoder must report what the stream +/// says rather than what it expected. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Range { + Full, + Limited, +} + +/// Where 4:2:0 chroma samples sit relative to luma. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Siting { + Center, + Left, +} + +/// How to interpret decoded Y, Cb and Cr. Signalled in every frame's sequence +/// header; it does not affect decoding itself. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct ColourDescription { + pub primaries: Primaries, + pub transfer: Transfer, + pub matrix: Matrix, + pub range: Range, + pub siting: Siting, +} + +impl ColourDescription { + /// SDR: BT.709 primaries, transfer and matrix, full range. + pub const fn bt709() -> Self { + Self { + primaries: Primaries::Bt709, + transfer: Transfer::Bt709, + matrix: Matrix::Bt709, + range: Range::Full, + siting: Siting::Center, + } + } + + /// HDR10: BT.2020 primaries and matrix, PQ transfer, full range. + pub const fn bt2020_pq() -> Self { + Self { + primaries: Primaries::Bt2020, + transfer: Transfer::Pq, + matrix: Matrix::Bt2020Ncl, + range: Range::Full, + siting: Siting::Center, + } + } + + /// Whether this is an HDR description. + pub fn is_hdr(&self) -> bool { + self.transfer == Transfer::Pq + } +} + +impl Default for ColourDescription { + fn default() -> Self { + Self::bt709() + } +} + +/// The 8-byte header every 32×32 block starts with. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BlockHeader { + /// One bit per 8×8 block in the 4×4 grid, row major. A zero bit means that + /// 8×8 block is all zero and is not transmitted. + pub ballot: u16, + /// Size of the whole block including this header, in 32-bit words. + pub payload_words: u16, + /// Frame counter, modulo 8. + pub sequence: u8, + /// How to scale this block's coefficients back to floating point. + pub quant_code: u8, + /// The block's position in [`Layout`]'s ordering. 24 bits. + pub block_index: u32, +} + +/// The start-of-frame header, a reinterpretation of the block header when its +/// `extended` bit is set. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SequenceHeader { + pub width: u32, + pub height: u32, + /// Frame counter, modulo 8. + pub sequence: u8, + /// Number of non-empty blocks in this frame. + pub total_blocks: u32, + pub chroma: Chroma, + pub colour: ColourDescription, +} + +/// Either kind of header, as read from the wire. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Header { + Block(BlockHeader), + Sequence(SequenceHeader), +} + +/// Start-of-frame is the only extended code the format defines. +const EXTENDED_CODE_START_OF_FRAME: u32 = 0; + +fn word(bytes: &[u8], index: usize) -> u32 { + u32::from_le_bytes(bytes[4 * index..4 * index + 4].try_into().unwrap()) +} + +impl BlockHeader { + /// The header as the eight bytes that go on the wire. + pub fn to_bytes(&self) -> [u8; HEADER_SIZE] { + let w0 = u32::from(self.ballot) + | (u32::from(self.payload_words) & 0xfff) << 16 + | (u32::from(self.sequence & SEQUENCE_MASK)) << 28; + let w1 = u32::from(self.quant_code) | (self.block_index & 0xff_ffff) << 8; + let mut out = [0; HEADER_SIZE]; + out[..4].copy_from_slice(&w0.to_le_bytes()); + out[4..].copy_from_slice(&w1.to_le_bytes()); + out + } + + /// Reads a header from two words already known not to be extended. + fn from_words(w0: u32, w1: u32) -> Self { + Self { + ballot: w0 as u16, + payload_words: ((w0 >> 16) & 0xfff) as u16, + sequence: ((w0 >> 28) & 0x7) as u8, + quant_code: w1 as u8, + block_index: w1 >> 8, + } + } +} + +impl SequenceHeader { + /// The header as the eight bytes that go on the wire. + pub fn to_bytes(&self) -> [u8; HEADER_SIZE] { + debug_assert!((1..=MAXIMUM_IMAGE_SIZE).contains(&self.width)); + debug_assert!((1..=MAXIMUM_IMAGE_SIZE).contains(&self.height)); + let c = &self.colour; + let w0 = (self.width - 1) & 0x3fff + | ((self.height - 1) & 0x3fff) << 14 + | u32::from(self.sequence & SEQUENCE_MASK) << 28 + | 1 << 31; + let w1 = self.total_blocks & 0xff_ffff + | EXTENDED_CODE_START_OF_FRAME << 24 + | u32::from(self.chroma == Chroma::Yuv444) << 26 + | u32::from(c.primaries == Primaries::Bt2020) << 27 + | u32::from(c.transfer == Transfer::Pq) << 28 + | u32::from(c.matrix == Matrix::Bt2020Ncl) << 29 + | u32::from(c.range == Range::Limited) << 30 + | u32::from(c.siting == Siting::Left) << 31; + let mut out = [0; HEADER_SIZE]; + out[..4].copy_from_slice(&w0.to_le_bytes()); + out[4..].copy_from_slice(&w1.to_le_bytes()); + out + } +} + +impl Header { + /// Reads whichever header the first eight bytes hold. + pub fn parse(bytes: &[u8]) -> Result { + if bytes.len() < HEADER_SIZE { + return Err(BitstreamError::Truncated { + needed: HEADER_SIZE, + left: bytes.len(), + }); + } + let w0 = word(bytes, 0); + let w1 = word(bytes, 1); + + if w0 >> 31 == 0 { + return Ok(Header::Block(BlockHeader::from_words(w0, w1))); + } + + let code = (w1 >> 24) & 0x3; + if code != EXTENDED_CODE_START_OF_FRAME { + return Err(BitstreamError::UnknownExtendedCode(code)); + } + let bit = |n: u32| (w1 >> n) & 1 != 0; + Ok(Header::Sequence(SequenceHeader { + width: (w0 & 0x3fff) + 1, + height: ((w0 >> 14) & 0x3fff) + 1, + sequence: ((w0 >> 28) & 0x7) as u8, + total_blocks: w1 & 0xff_ffff, + chroma: if bit(26) { + Chroma::Yuv444 + } else { + Chroma::Yuv420 + }, + colour: ColourDescription { + primaries: if bit(27) { + Primaries::Bt2020 + } else { + Primaries::Bt709 + }, + transfer: if bit(28) { + Transfer::Pq + } else { + Transfer::Bt709 + }, + matrix: if bit(29) { + Matrix::Bt2020Ncl + } else { + Matrix::Bt709 + }, + range: if bit(30) { Range::Limited } else { Range::Full }, + siting: if bit(31) { + Siting::Left + } else { + Siting::Center + }, + }, + })) + } +} + +/// Something in a stream that the format does not allow. +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum BitstreamError { + #[error("truncated: needed {needed} bytes, {left} left")] + Truncated { needed: usize, left: usize }, + #[error("extended header code {0} is reserved")] + UnknownExtendedCode(u32), + #[error("block header says {0} words, less than the header itself")] + BlockTooSmall(u32), + #[error("block index {index} is out of range (the layout has {count})")] + BlockIndexOutOfRange { index: u32, count: u32 }, + #[error("block {index}: 8×8 block ({x}, {y}) is outside the band")] + BallotOutOfBand { index: u32, x: u32, y: u32 }, + #[error("block {index}: its contents need {computed} words, its header says {declared}")] + BlockSizeMismatch { + index: u32, + computed: u32, + declared: u32, + }, + #[error("4:2:0 needs even dimensions, got {width}x{height}")] + OddYuv420 { width: u32, height: u32 }, + #[error("dimensions {width}x{height} are outside 1..={MAXIMUM_IMAGE_SIZE}")] + BadDimensions { width: u32, height: u32 }, + #[error("the encoder's output for block {index} is inconsistent: {what}")] + EncoderOutput { index: u32, what: &'static str }, +} + +#[cfg(test)] +mod tests; diff --git a/crates/nespyro/src/bitstream/packetize.rs b/crates/nespyro/src/bitstream/packetize.rs new file mode 100644 index 00000000..4f47bb28 --- /dev/null +++ b/crates/nespyro/src/bitstream/packetize.rs @@ -0,0 +1,175 @@ +//! Turning the encoder's GPU output into a frame. +//! +//! The GPU writes finished 32×32 blocks wherever an atomic counter puts them, +//! so their order in its output buffer changes from run to run. Beside them it +//! writes a table giving each block's position. Packetizing walks that table +//! in block-index order, which is what makes the result deterministic: the same +//! input gives the same bytes however the GPU scheduled its work. + +use std::ops::Range; + +use super::{BitstreamError, BlockHeader, HEADER_SIZE, Header, Layout, SequenceHeader}; + +/// One entry of the GPU's block table: where block `i` landed, in 32-bit +/// words, and how many words it is. Zero words means the block quantized to +/// nothing and is not sent. Matches `BitstreamPacket` in `block_packing`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +#[repr(C)] +pub struct GpuPacket { + pub offset_words: u32, + pub num_words: u32, +} + +/// A frame, ready for the network. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Packetized { + /// The start-of-frame header, then every non-empty block in index order. + pub data: Vec, + /// `data` cut into packets. Every one parses on its own; none splits a + /// block, and none exceeds the packet size unless a single block does. + pub packets: Vec>, + /// How many leading packets carry the two coarsest bands, whose loss the + /// decoder cannot mask. These deserve the most reliable delivery. + pub critical_packets: usize, + /// Non-empty blocks in the frame. + pub blocks: u32, +} + +/// The bands whose loss is not maskable: the coarsest LL and the high-pass +/// bands of the same level. Upstream's decoder refuses a partial frame without +/// them, and so does [`super::Readiness::is_complete_enough`]. +pub(crate) const CRITICAL_BANDS: usize = 2; + +/// Builds a frame from the GPU's block table and bitstream. +/// +/// `bitstream` is the part of the GPU's output that was read back. A block +/// lying outside it is an error rather than a truncated frame: it means the +/// read-back was too small for what rate control produced. +pub fn packetize( + layout: &Layout, + header: SequenceHeader, + table: &[GpuPacket], + bitstream: &[u32], + packet_size: usize, +) -> Result { + assert_eq!(table.len(), layout.blocks_32x32() as usize); + assert!(packet_size >= HEADER_SIZE); + + let mut blocks = 0u32; + for (index, entry) in table.iter().enumerate() { + let index = index as u32; + if entry.num_words == 0 { + continue; + } + let words = block_words(index, entry, bitstream)?; + let Header::Block(block) = Header::parse(bytes_of(&words[..2]))? else { + return Err(BitstreamError::EncoderOutput { + index, + what: "block has its extended bit set", + }); + }; + check_block(index, entry, &block, header.sequence)?; + blocks += 1; + } + + let total_bytes: usize = HEADER_SIZE + + table + .iter() + .map(|e| e.num_words as usize * 4) + .sum::(); + let mut data = Vec::with_capacity(total_bytes); + data.extend_from_slice( + &SequenceHeader { + total_blocks: blocks, + ..header + } + .to_bytes(), + ); + + let critical_limit = layout.leading_blocks(CRITICAL_BANDS); + let mut packets = Vec::new(); + let mut critical_packets = 1; + let mut packet_start = 0; + + for (index, entry) in table.iter().enumerate() { + if entry.num_words == 0 { + continue; + } + let size = entry.num_words as usize * 4; + if data.len() - packet_start + size > packet_size && data.len() > packet_start { + packets.push(packet_start..data.len()); + packet_start = data.len(); + if (index as u32) < critical_limit { + critical_packets = packets.len() + 1; + } + } + let words = block_words(index as u32, entry, bitstream)?; + data.extend_from_slice(bytes_of(words)); + } + if data.len() > packet_start { + packets.push(packet_start..data.len()); + } + + Ok(Packetized { + data, + packets, + critical_packets, + blocks, + }) +} + +fn block_words<'a>( + index: u32, + entry: &GpuPacket, + bitstream: &'a [u32], +) -> Result<&'a [u32], BitstreamError> { + let start = entry.offset_words as usize; + let end = start + entry.num_words as usize; + if entry.num_words < 2 { + return Err(BitstreamError::BlockTooSmall(entry.num_words)); + } + bitstream + .get(start..end) + .ok_or(BitstreamError::EncoderOutput { + index, + what: "block lies past the end of the read-back bitstream", + }) +} + +/// The checks upstream's packetizer makes in debug builds, made always: a +/// table and a bitstream that disagree mean a GPU-side bug, and sending the +/// result would hand the client a stream it cannot parse. +fn check_block( + index: u32, + entry: &GpuPacket, + block: &BlockHeader, + sequence: u8, +) -> Result<(), BitstreamError> { + if block.block_index != index { + return Err(BitstreamError::EncoderOutput { + index, + what: "header names a different block index", + }); + } + if u32::from(block.payload_words) != entry.num_words { + return Err(BitstreamError::EncoderOutput { + index, + what: "header and table disagree on the block's size", + }); + } + if block.sequence != sequence { + return Err(BitstreamError::EncoderOutput { + index, + what: "header carries another frame's sequence number", + }); + } + Ok(()) +} + +fn bytes_of(words: &[u32]) -> &[u8] { + // SAFETY: u32 has no padding and any bit pattern is a valid u8; the length + // is exact. The wire format is little-endian, as every target this crate + // builds for is. + const _: () = assert!(cfg!(target_endian = "little")); + unsafe { std::slice::from_raw_parts(words.as_ptr().cast(), std::mem::size_of_val(words)) } +} diff --git a/crates/nespyro/src/bitstream/tests.rs b/crates/nespyro/src/bitstream/tests.rs new file mode 100644 index 00000000..b7462bff --- /dev/null +++ b/crates/nespyro/src/bitstream/tests.rs @@ -0,0 +1,683 @@ +use super::*; + +/// xorshift, so the tests need no dependency and every run is the same. +struct Rng(u64); +impl Rng { + fn next(&mut self) -> u32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 >> 16) as u32 + } +} + +/// The block-index walk exactly as `bitstream.md` writes it, kept separate +/// from `Layout` so the two can be compared. +fn spec_walk(w: u32, h: u32, chroma: Chroma) -> Vec<(usize, usize, usize, u32)> { + let aw = w.next_multiple_of(32).max(128); + let ah = h.next_multiple_of(32).max(128); + let mut out = Vec::new(); + let mut index = 0; + for level in (0..5).rev() { + for component in 0..3 { + if level == 0 && component != 0 && chroma == Chroma::Yuv420 { + continue; + } + for band in (if level == 4 { 0 } else { 1 })..4 { + let bw = (aw >> (level + 1)).div_ceil(32); + let bh = (ah >> (level + 1)).div_ceil(32); + out.push((component, level, band, index)); + index += bw * bh; + } + } + } + out.push((9, 9, 9, index)); + out +} + +#[test] +fn layout_follows_the_spec_walk() { + for &(w, h, chroma) in &[ + (128, 128, Chroma::Yuv420), + (128, 128, Chroma::Yuv444), + (1920, 1080, Chroma::Yuv420), + (1366, 768, Chroma::Yuv444), + (3840, 2160, Chroma::Yuv420), + (64, 48, Chroma::Yuv420), + ] { + let layout = Layout::new(w, h, chroma).unwrap(); + let walk = spec_walk(w, h, chroma); + for &(c, l, b, index) in &walk[..walk.len() - 1] { + let band = layout.band(c, l, b).unwrap(); + assert_eq!( + band.offset_32x32, index, + "{w}x{h} {chroma:?} c{c} l{l} b{b}" + ); + } + assert_eq!(layout.blocks_32x32(), walk.last().unwrap().3); + } +} + +#[test] +fn smallest_frame_block_counts() { + // 128×128 is the minimum: bands are 64, 32, 16, 8 and 4 wide, so only + // level 0 has more than one 32×32 block per band (2×2 of them). + let l420 = Layout::new(128, 128, Chroma::Yuv420).unwrap(); + assert_eq!(l420.blocks_32x32(), 12 + 9 + 9 + 9 + 3 * 4); + let l444 = Layout::new(128, 128, Chroma::Yuv444).unwrap(); + assert_eq!(l444.blocks_32x32(), 12 + 9 + 9 + 9 + 9 * 4); + + // Level 4 runs Y's four bands, Cb's four, Cr's four; one band deep ends + // at Cr's LL, two deep at the end of the level. + assert_eq!(l420.leading_blocks(1), 9); + assert_eq!(l420.leading_blocks(2), 12); + assert_eq!(l420.leading_blocks(3), 21); +} + +#[test] +fn small_frames_pad_to_the_minimum() { + let l = Layout::new(40, 20, Chroma::Yuv444).unwrap(); + assert_eq!((l.aligned_width, l.aligned_height), (128, 128)); + let l = Layout::new(1366, 768, Chroma::Yuv420).unwrap(); + assert_eq!((l.aligned_width, l.aligned_height), (1376, 768)); +} + +#[test] +fn layout_refuses_what_the_format_cannot_carry() { + assert!(matches!( + Layout::new(1921, 1080, Chroma::Yuv420), + Err(BitstreamError::OddYuv420 { .. }) + )); + assert!(Layout::new(1921, 1081, Chroma::Yuv444).is_ok()); + assert!(matches!( + Layout::new(0, 10, Chroma::Yuv444), + Err(BitstreamError::BadDimensions { .. }) + )); + assert!(matches!( + Layout::new(16385, 10, Chroma::Yuv444), + Err(BitstreamError::BadDimensions { .. }) + )); +} + +#[test] +fn edge_blocks_cover_fewer_8x8_blocks() { + // 1366 pads to 1376: level 0 is 688 wide, 86 8×8 blocks, so the last + // 32×32 block in a row covers 86 - 21·4 = 2 of them. + let l = Layout::new(1366, 768, Chroma::Yuv420).unwrap(); + let band = l.band(0, 0, 1).unwrap(); + let last_in_row = band.offset_32x32 + band.stride_32x32 - 1; + let m = l.mapping(last_in_row).unwrap(); + assert_eq!((m.width_8x8, m.height_8x8), (2, 4)); +} + +#[test] +fn largest_block_fits_the_size_field() { + assert_eq!(LARGEST_BLOCK_WORDS, 622); +} + +#[test] +fn block_header_bytes() { + let h = BlockHeader { + ballot: 0x8001, + payload_words: 0xabc, + sequence: 5, + quant_code: 0x3d, + block_index: 0x12_3456, + }; + let b = h.to_bytes(); + // ballot, then payload_words:12 | sequence:3 | extended:1. + assert_eq!(b[..4], (0x8001u32 | 0xabc << 16 | 5 << 28).to_le_bytes()); + assert_eq!(b[4..], (0x3du32 | 0x12_3456 << 8).to_le_bytes()); + assert_eq!(Header::parse(&b).unwrap(), Header::Block(h)); +} + +#[test] +fn sequence_header_bytes() { + let h = SequenceHeader { + width: 1920, + height: 1080, + sequence: 3, + total_blocks: 4242, + chroma: Chroma::Yuv444, + colour: ColourDescription::bt2020_pq(), + }; + let b = h.to_bytes(); + assert_eq!( + b[..4], + (1919u32 | 1079 << 14 | 3 << 28 | 1 << 31).to_le_bytes() + ); + // chroma 444, primaries 2020, transfer PQ, matrix 2020, full, centre. + assert_eq!( + b[4..], + (4242u32 | 1 << 26 | 1 << 27 | 1 << 28 | 1 << 29).to_le_bytes() + ); + assert_eq!(Header::parse(&b).unwrap(), Header::Sequence(h)); + + // Every colour field must round-trip on its own, not just as a bundle. + let odd = ColourDescription { + primaries: Primaries::Bt709, + transfer: Transfer::Pq, + matrix: Matrix::Bt709, + range: Range::Limited, + siting: Siting::Left, + }; + let h = SequenceHeader { colour: odd, ..h }; + assert_eq!(Header::parse(&h.to_bytes()).unwrap(), Header::Sequence(h)); +} + +#[test] +fn reserved_extended_codes_are_refused() { + let mut b = SequenceHeader { + width: 128, + height: 128, + sequence: 0, + total_blocks: 0, + chroma: Chroma::Yuv420, + colour: ColourDescription::bt709(), + } + .to_bytes(); + b[7] |= 1; // code = 1 + assert_eq!( + Header::parse(&b), + Err(BitstreamError::UnknownExtendedCode(1)) + ); +} + +/// One structurally valid block, built independently of `validate_block`. +/// +/// For each present 8×8 block: a code word of random 2-bit plane deltas, a +/// random q_bits, the plane bytes, and one sign bit per coefficient that any +/// plane marks significant. +fn make_block(layout: &Layout, index: u32, sequence: u8, rng: &mut Rng) -> Vec { + let m = layout.mapping(index).unwrap(); + let mut ballot = 0u16; + for y in 0..m.height_8x8 { + for x in 0..m.width_8x8 { + if !rng.next().is_multiple_of(3) { + ballot |= 1 << (y * 4 + x); + } + } + } + if ballot == 0 { + ballot = 1; + } + let n = ballot.count_ones() as usize; + let mut code_words = Vec::new(); + let mut q_scales = Vec::new(); + let mut planes = Vec::new(); + let mut significant = 0; + for _ in 0..n { + let code_word = rng.next() as u16; + let q_bits = rng.next() % 4; + code_words.extend_from_slice(&code_word.to_le_bytes()); + q_scales.push(q_bits as u8 | ((rng.next() % 16) as u8) << 4); + for sub in 0..8 { + let count = q_bits as u16 + (code_word >> (2 * sub) & 3); + let mut union = 0u8; + for _ in 0..count { + let b = rng.next() as u8; + union |= b; + planes.push(b); + } + significant += union.count_ones(); + } + } + let mut body = Vec::new(); + body.extend(code_words); + body.extend(q_scales); + body.extend(planes); + body.extend((0..significant.div_ceil(8)).map(|_| rng.next() as u8)); + let words = (HEADER_SIZE + body.len()).div_ceil(4); + body.resize(words * 4 - HEADER_SIZE, 0); + + let header = BlockHeader { + ballot, + payload_words: words as u16, + sequence, + quant_code: (rng.next() & 0xff) as u8, + block_index: index, + }; + let mut out = header.to_bytes().to_vec(); + out.extend(body); + out +} + +/// A GPU table and word buffer holding `blocks`, scattered in the buffer in +/// an order unrelated to their index, as the GPU's atomics would leave them. +fn gpu_output( + layout: &Layout, + blocks: &[(u32, Vec)], + rng: &mut Rng, +) -> (Vec, Vec) { + let mut table = vec![GpuPacket::default(); layout.blocks_32x32() as usize]; + let mut order: Vec = (0..blocks.len()).collect(); + for i in (1..order.len()).rev() { + order.swap(i, rng.next() as usize % (i + 1)); + } + let mut words = Vec::new(); + for i in order { + let (index, bytes) = &blocks[i]; + table[*index as usize] = GpuPacket { + offset_words: words.len() as u32, + num_words: (bytes.len() / 4) as u32, + }; + words.extend( + bytes + .as_chunks::<4>() + .0 + .iter() + .map(|w| u32::from_le_bytes(*w)), + ); + } + (table, words) +} + +fn seq(layout: &Layout, sequence: u8) -> SequenceHeader { + SequenceHeader { + width: layout.width, + height: layout.height, + sequence, + total_blocks: 0, + chroma: layout.chroma, + colour: ColourDescription::bt709(), + } +} + +fn sample_frame( + layout: &Layout, + sequence: u8, + keep: impl Fn(u32) -> bool, + rng: &mut Rng, +) -> Vec<(u32, Vec)> { + (0..layout.blocks_32x32()) + .filter(|&i| keep(i)) + .map(|i| (i, make_block(layout, i, sequence, rng))) + .collect() +} + +#[test] +fn a_hand_built_block_validates() { + // One 8×8 block present. Its code word gives subblock 0 one plane and the + // rest none, q_bits 0. The plane byte has two bits set, so two sign bits, + // one byte. 8 header + 2 code word + 1 scale + 1 plane + 1 sign = 13 + // bytes, 4 words. + let layout = Layout::new(128, 128, Chroma::Yuv420).unwrap(); + let header = BlockHeader { + ballot: 1, + payload_words: 4, + sequence: 0, + quant_code: 0, + block_index: 40, + }; + let mut b = header.to_bytes().to_vec(); + b.extend([0b01, 0, 0x60, 0b1010_0000, 0b10, 0, 0, 0]); + depacketize::validate_block(&layout, &header, &b).unwrap(); + + let wrong = BlockHeader { + payload_words: 5, + ..header + }; + assert_eq!( + depacketize::validate_block(&layout, &wrong, &[b.clone(), vec![0; 4]].concat()), + Err(BitstreamError::BlockSizeMismatch { + index: 40, + computed: 4, + declared: 5 + }) + ); +} + +#[test] +fn packetize_then_parse_reproduces_every_block() { + let mut rng = Rng(0x9e37_79b9_7f4a_7c15); + for &(w, h, chroma) in &[ + (1920, 1080, Chroma::Yuv420), + (1366, 768, Chroma::Yuv444), + (128, 128, Chroma::Yuv420), + ] { + let layout = Layout::new(w, h, chroma).unwrap(); + let blocks = sample_frame(&layout, 2, |i| i % 7 != 3, &mut rng); + let (table, words) = gpu_output(&layout, &blocks, &mut rng); + let frame = packetize(&layout, seq(&layout, 2), &table, &words, 1200).unwrap(); + + assert_eq!(frame.blocks as usize, blocks.len()); + // The frame is the header and the blocks in index order, nothing else. + let mut expected = seq(&layout, 2); + expected.total_blocks = blocks.len() as u32; + let mut concat = expected.to_bytes().to_vec(); + for (_, b) in &blocks { + concat.extend(b); + } + assert_eq!(frame.data, concat); + + // Packets tile the frame exactly. + assert_eq!(frame.packets.first().unwrap().start, 0); + assert_eq!(frame.packets.last().unwrap().end, frame.data.len()); + for pair in frame.packets.windows(2) { + assert_eq!(pair[0].end, pair[1].start); + } + + let mut d = Depacketizer::new(layout.clone()); + for p in &frame.packets { + assert!(p.len() <= 1200 || p.len() <= LARGEST_BLOCK_WORDS as usize * 4 + 8); + assert_eq!(d.push(&frame.data[p.clone()]).unwrap(), Push::Accepted); + } + let r = d.readiness(); + assert!(r.is_complete(), "{r:?}"); + assert_eq!(d.missing_blocks(), 0); + + // Every block comes back byte for byte at the offset the table says. + let (offsets, payload) = d.take_for_decode(); + for (index, bytes) in &blocks { + let at = offsets[*index as usize] as usize; + let got: Vec = payload[at..at + bytes.len() / 4] + .iter() + .flat_map(|w| w.to_le_bytes()) + .collect(); + assert_eq!(&got, bytes); + } + let absent = offsets + .iter() + .filter(|&&o| o == depacketize::MISSING) + .count(); + assert_eq!(absent, layout.blocks_32x32() as usize - blocks.len()); + assert!(!d.readiness().is_complete_enough(), "decoded twice"); + } +} + +#[test] +fn every_packet_parses_on_its_own_in_any_order() { + let mut rng = Rng(7); + let layout = Layout::new(1920, 1080, Chroma::Yuv420).unwrap(); + let blocks = sample_frame(&layout, 6, |_| true, &mut rng); + let (table, words) = gpu_output(&layout, &blocks, &mut rng); + let frame = packetize(&layout, seq(&layout, 6), &table, &words, 1400).unwrap(); + assert!(frame.packets.len() > 10); + + let mut d = Depacketizer::new(layout); + for p in frame.packets.iter().rev() { + assert_eq!(d.push(&frame.data[p.clone()]).unwrap(), Push::Accepted); + } + assert!(d.readiness().is_complete()); +} + +#[test] +fn packetizing_does_not_depend_on_the_gpu_order() { + let layout = Layout::new(1366, 768, Chroma::Yuv444).unwrap(); + let blocks = sample_frame(&layout, 1, |i| i % 3 != 0, &mut Rng(99)); + let (t1, w1) = gpu_output(&layout, &blocks, &mut Rng(1)); + let (t2, w2) = gpu_output(&layout, &blocks, &mut Rng(2)); + assert_ne!( + w1, w2, + "the two scatters must differ for this to mean anything" + ); + let a = packetize(&layout, seq(&layout, 1), &t1, &w1, 1200).unwrap(); + let b = packetize(&layout, seq(&layout, 1), &t2, &w2, 1200).unwrap(); + assert_eq!(a, b); +} + +#[test] +fn packets_never_split_a_block_and_respect_the_size() { + let mut rng = Rng(3); + let layout = Layout::new(1920, 1080, Chroma::Yuv420).unwrap(); + let blocks = sample_frame(&layout, 0, |_| true, &mut rng); + let (table, words) = gpu_output(&layout, &blocks, &mut rng); + for size in [64, 300, 1200, 9000] { + let frame = packetize(&layout, seq(&layout, 0), &table, &words, size).unwrap(); + for p in &frame.packets { + let bytes = &frame.data[p.clone()]; + // Walking the packet's headers must land exactly on its end. + let mut at = 0; + let mut count = 0; + while at < bytes.len() { + match Header::parse(&bytes[at..]).unwrap() { + Header::Sequence(_) => at += HEADER_SIZE, + Header::Block(b) => at += usize::from(b.payload_words) * 4, + } + count += 1; + } + assert_eq!(at, bytes.len()); + assert!( + bytes.len() <= size || count == 1 || (count == 2 && p.start == 0), + "packet of {} bytes over {size} holds {count} headers", + bytes.len() + ); + } + } +} + +#[test] +fn critical_packets_cover_the_coarse_bands() { + let mut rng = Rng(11); + let layout = Layout::new(3840, 2160, Chroma::Yuv420).unwrap(); + let blocks = sample_frame(&layout, 0, |_| true, &mut rng); + let (table, words) = gpu_output(&layout, &blocks, &mut rng); + let frame = packetize(&layout, seq(&layout, 0), &table, &words, 1200).unwrap(); + let limit = layout.leading_blocks(2); + assert!(frame.critical_packets > 1 && frame.critical_packets < frame.packets.len()); + + // Delivering only the critical packets gives every coarse block; + // delivering all but the last of them does not. + let mut d = Depacketizer::new(layout.clone()); + for p in &frame.packets[..frame.critical_packets] { + d.push(&frame.data[p.clone()]).unwrap(); + } + assert!(d.readiness().critical_complete); + assert!(d.readiness().received >= limit); + + let mut d = Depacketizer::new(layout); + for p in &frame.packets[..frame.critical_packets - 1] { + d.push(&frame.data[p.clone()]).unwrap(); + } + assert!(!d.readiness().critical_complete); +} + +#[test] +fn readiness_follows_upstreams_partial_rule() { + let mut rng = Rng(5); + let layout = Layout::new(1920, 1080, Chroma::Yuv420).unwrap(); + let blocks = sample_frame(&layout, 0, |_| true, &mut rng); + let (table, words) = gpu_output(&layout, &blocks, &mut rng); + let frame = packetize(&layout, seq(&layout, 0), &table, &words, 1200).unwrap(); + let n = frame.packets.len(); + + // Drop 5% of the non-critical packets: decodable, not complete. + let mut d = Depacketizer::new(layout.clone()); + for (i, p) in frame.packets.iter().enumerate() { + if i < frame.critical_packets || i % 20 != 7 { + d.push(&frame.data[p.clone()]).unwrap(); + } + } + let r = d.readiness(); + assert!(!r.is_complete() && r.is_complete_enough(), "{r:?}"); + assert!(d.missing_blocks() > 0); + + // Drop a quarter: not decodable yet. + let mut d = Depacketizer::new(layout.clone()); + for (i, p) in frame.packets.iter().enumerate() { + if i < frame.critical_packets || i % 4 != 1 { + d.push(&frame.data[p.clone()]).unwrap(); + } + } + assert!(!d.readiness().is_complete_enough()); + + // Everything but the start-of-frame header: nothing is decodable, because + // nothing says how many blocks the frame has. + let mut d = Depacketizer::new(layout); + let first = &frame.data[frame.packets[0].clone()]; + d.push(&first[HEADER_SIZE..]).unwrap(); + for p in &frame.packets[1..n] { + d.push(&frame.data[p.clone()]).unwrap(); + } + assert!(!d.readiness().has_header); + assert!(!d.readiness().is_complete_enough()); +} + +#[test] +fn sequence_numbers_wrap_and_stale_frames_drop() { + let mut rng = Rng(21); + let layout = Layout::new(128, 128, Chroma::Yuv420).unwrap(); + let packet = |s: u8, rng: &mut Rng| { + let mut p = seq(&layout, s).to_bytes().to_vec(); + p.extend(make_block(&layout, 0, s, rng)); + p + }; + let mut d = Depacketizer::new(layout.clone()); + assert_eq!(d.push(&packet(7, &mut rng)).unwrap(), Push::Accepted); + assert_eq!(d.readiness().received, 1); + // 7 → 0 is one frame forward across the wrap: the old frame is abandoned. + assert_eq!(d.push(&packet(0, &mut rng)).unwrap(), Push::Accepted); + assert_eq!(d.readiness().received, 1); + // 6 is behind 0: stale, and it must not disturb the current frame. + assert_eq!(d.push(&packet(6, &mut rng)).unwrap(), Push::Stale); + assert_eq!(d.readiness().received, 1); + // A duplicate block is allowed and counted once. + let dup = make_block(&layout, 0, 0, &mut rng); + assert_eq!(d.push(&dup).unwrap(), Push::Accepted); + assert_eq!(d.readiness().received, 1); +} + +#[test] +fn a_different_stream_asks_for_a_new_decoder() { + let layout = Layout::new(1920, 1080, Chroma::Yuv420).unwrap(); + let mut d = Depacketizer::new(layout); + let other = SequenceHeader { + width: 1280, + height: 720, + sequence: 0, + total_blocks: 1, + chroma: Chroma::Yuv444, + colour: ColourDescription::bt709(), + }; + assert_eq!( + d.push(&other.to_bytes()).unwrap(), + Push::Reconfigure(StreamInfo { + width: 1280, + height: 720, + chroma: Chroma::Yuv444 + }) + ); + assert!(!d.readiness().has_header); +} + +#[test] +fn malformed_packets_are_refused_not_skipped() { + let mut rng = Rng(31); + let layout = Layout::new(128, 128, Chroma::Yuv420).unwrap(); + let good = make_block(&layout, 5, 0, &mut rng); + + let mut d = Depacketizer::new(layout.clone()); + // A block cut short. + assert!(matches!( + d.push(&good[..good.len() - 4]), + Err(BitstreamError::Truncated { .. }) + )); + // A stray partial header at the end. The block before it is sound and is + // kept; only what follows it is refused. + let mut kept = Depacketizer::new(layout.clone()); + assert!(matches!( + kept.push(&[good.clone(), vec![1, 2, 3]].concat()), + Err(BitstreamError::Truncated { .. }) + )); + assert_eq!(kept.readiness().received, 1); + // A header claiming fewer words than itself. + let mut tiny = good.clone(); + tiny[2] = 1; + tiny[3] &= 0xf0; + assert_eq!(d.push(&tiny), Err(BitstreamError::BlockTooSmall(1))); + // A block index past the layout. + let mut far = good.clone(); + let far_index = layout.blocks_32x32(); + far[5..8].copy_from_slice(&far_index.to_le_bytes()[..3]); + assert!(matches!( + d.push(&far), + Err(BitstreamError::BlockIndexOutOfRange { .. }) + )); + // An 8×8 block outside the band: level 4 is 4 coefficients wide, one 8×8 + // block, so ballot bit 1 names a block that cannot exist. + let edge = BlockHeader { + ballot: 0b11, + payload_words: 4, + sequence: 0, + quant_code: 0, + block_index: 0, + }; + let mut b = edge.to_bytes().to_vec(); + b.extend([0u8; 8]); + assert!(matches!( + d.push(&b), + Err(BitstreamError::BallotOutOfBand { + index: 0, + x: 1, + y: 0 + }) + )); + // Contents longer than the header admits. + let mut long = good.clone(); + let words = u16::from_le_bytes([long[2], long[3]]) & 0xfff; + let shrunk = (words - 1) | (u16::from_le_bytes([long[2], long[3]]) & 0xf000); + long[2..4].copy_from_slice(&shrunk.to_le_bytes()); + assert!(matches!( + d.push(&long[..long.len() - 4]), + Err(BitstreamError::BlockSizeMismatch { .. }) + )); + + // None of that left anything behind. + assert_eq!(d.readiness().received, 0); + // And the good block still goes in. + assert_eq!(d.push(&good).unwrap(), Push::Accepted); + assert_eq!(d.readiness().received, 1); +} + +#[test] +fn random_bytes_never_panic_and_never_pass_as_blocks() { + let mut rng = Rng(0xdead_beef); + let layout = Layout::new(640, 480, Chroma::Yuv420).unwrap(); + let mut d = Depacketizer::new(layout); + let mut accepted = 0; + for _ in 0..20_000 { + let len = (rng.next() % 64) as usize; + let bytes: Vec = (0..len).map(|_| rng.next() as u8).collect(); + if let Ok(Push::Accepted) = d.push(&bytes) { + accepted += d.readiness().received; + } + } + // Random bytes that happen to form a consistent block are possible but + // rare; what matters is that nothing panicked and almost nothing passed. + assert!(accepted < 20, "{accepted} random packets passed validation"); +} + +#[test] +fn packetize_refuses_inconsistent_gpu_output() { + let mut rng = Rng(41); + let layout = Layout::new(128, 128, Chroma::Yuv420).unwrap(); + let blocks = sample_frame(&layout, 4, |i| i < 3, &mut rng); + let (table, words) = gpu_output(&layout, &blocks, &mut rng); + + // Another frame's sequence number in a block. + assert!(matches!( + packetize(&layout, seq(&layout, 5), &table, &words, 1200), + Err(BitstreamError::EncoderOutput { .. }) + )); + // A table entry pointing past what was read back. + let mut t = table.clone(); + t[1].offset_words = words.len() as u32; + assert!(matches!( + packetize(&layout, seq(&layout, 4), &t, &words, 1200), + Err(BitstreamError::EncoderOutput { .. }) + )); + // A table entry naming a size its header does not. + let mut t = table.clone(); + t[2].num_words -= 1; + assert!(matches!( + packetize(&layout, seq(&layout, 4), &t, &words, 1200), + Err(BitstreamError::EncoderOutput { .. }) + )); + // Two table entries pointing at the same block. + let mut t = table.clone(); + t[0] = t[1]; + assert!(matches!( + packetize(&layout, seq(&layout, 4), &t, &words, 1200), + Err(BitstreamError::EncoderOutput { .. }) + )); +} diff --git a/crates/nespyro/src/decoder.rs b/crates/nespyro/src/decoder.rs new file mode 100644 index 00000000..f4703b4f --- /dev/null +++ b/crates/nespyro/src/decoder.rs @@ -0,0 +1,682 @@ +//! The decoder: packets in as they arrive, three planes out. +//! +//! Packets go to [`Decoder::push_packet`] in any order. When +//! [`Decoder::readiness`] says enough of the frame is there, and when that is +//! enough is the caller's decision, [`Decoder::decode_after`] submits the +//! frame and returns its planes with the timeline point they are ready at. +//! Nothing waits on the CPU: the consumer waits on that point in its own +//! submission. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use ash::vk; + +use crate::bitstream::{ + Chroma, ColourDescription, DECOMPOSITION_LEVELS, Depacketizer, Layout, NUM_COMPONENTS, Push, + Readiness, +}; +use crate::device::Context; +use crate::encoder::Depth; +use crate::error::{Error, Result}; +use crate::gpu::{ + Buffer, Commands, Image, ImageDesc, Location, Timeline, Timestamps, compute_barrier, + memory_barrier, to_general, +}; +use crate::pipeline::{Bind, Pipeline, PipelineDesc, SubgroupSize, entry}; +use crate::shaders; +use crate::sync::TimelinePoint; +use crate::wavelet::Wavelet; + +#[derive(Debug, Clone)] +pub struct DecodeConfig { + pub width: u32, + pub height: u32, + pub chroma: Chroma, + /// Depth of the output planes. The stream does not carry one: PyroWave + /// decodes in floating point, so this is the consumer's choice. + pub depth: Depth, + /// The queue family that reads the planes, when it is not nespyro's. + pub consumer_family: Option, +} + +impl DecodeConfig { + pub fn new(width: u32, height: u32, chroma: Chroma, depth: Depth) -> Self { + Self { + width, + height, + chroma, + depth, + consumer_family: None, + } + } + + /// Output planes are shared with `family`, so a consumer there samples + /// them without an ownership transfer. + pub fn with_consumer_queue_family(mut self, family: u32) -> Self { + self.consumer_family = Some(family); + self + } +} + +/// One output plane. +/// +/// The image is padded to the codec's alignment, so the last inverse +/// transform never writes out of bounds; only `width` × `height` from the +/// top-left corner is the picture. Sample by texel coordinate, or scale +/// normalised coordinates by `width / image_width`. +#[derive(Debug, Clone, Copy)] +pub struct PlaneView { + pub image: vk::Image, + pub view: vk::ImageView, + pub format: vk::Format, + pub width: u32, + pub height: u32, + pub image_width: u32, + pub image_height: u32, +} + +/// A decoded frame: its planes stay the caller's until this is dropped. +/// +/// The planes are in `GENERAL`, full-range YCbCr as `colour` describes, and +/// valid once `ready` is reached. A consumer waits on `ready` in its own +/// submission and drops the frame only after that submission is done with +/// the planes, since dropping hands them back for the next decode. +pub struct DecodedFrame { + pub y: PlaneView, + pub cb: PlaneView, + pub cr: PlaneView, + pub width: u32, + pub height: u32, + pub chroma: Chroma, + pub depth: Depth, + pub colour: ColourDescription, + /// Blocks the frame's header promised that never arrived, decoded as + /// zeros. Zero for a whole frame. + pub missing_blocks: u32, + pub ready: TimelinePoint, + _hold: SlotHold, +} + +impl std::fmt::Debug for DecodedFrame { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DecodedFrame") + .field("width", &self.width) + .field("height", &self.height) + .field("chroma", &self.chroma) + .field("colour", &self.colour) + .field("missing_blocks", &self.missing_blocks) + .field("ready", &self.ready) + .finish_non_exhaustive() + } +} + +/// Marks an output slot busy until dropped. +struct SlotHold(Arc); + +impl Drop for SlotHold { + fn drop(&mut self) { + self.0.store(false, Ordering::Release); + } +} + +/// Output slots. One being decoded, one being shown, one spare. +const OUTPUTS: usize = 3; +/// Uploads in flight. +const UPLOADS: usize = 2; + +struct Output { + planes: [Image; 3], + views: [vk::ImageView; 3], + held: Arc, +} + +/// GPU time of one decode, in nanoseconds. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct DecodeStats { + /// Copying the payload to device-local memory, when it is not there + /// already. + pub upload_ns: f64, + pub dequant_ns: f64, + pub idwt_ns: f64, +} + +impl DecodeStats { + pub fn gpu_ns(&self) -> f64 { + self.upload_ns + self.dequant_ns + self.idwt_ns + } +} + +/// Timestamps around the upload, the dequantizer and the inverse transform. +const TIMESTAMPS: u32 = 4; + +/// One frame's upload: the block offset table, then the payload words, then +/// padding the dequantizer may read one word into. +struct Upload { + timestamps: Option, + host: Buffer, + /// Where the shaders read from: `host` itself when it is device local, + /// otherwise a device-local copy. + device: Option, + /// The timeline value of the last decode that read this upload. + last_use: u64, +} + +/// The dequantizer reads a word past the end of the last block's signs. +const PAYLOAD_PADDING: u64 = 16; + +#[repr(C)] +#[derive(Clone, Copy)] +struct DequantPush { + payload_offsets: u64, + payload_data_u32: u64, + payload_data_u16: u64, + payload_data_u8: u64, + resolution: [i32; 2], + output_layer: i32, + block_offset_32x32: i32, + block_stride_32x32: i32, + _pad: i32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct IdwtPush { + resolution: [i32; 2], + inv_resolution: [f32; 2], +} + +pub struct Decoder { + ctx: Context, + config: DecodeConfig, + layout: Layout, + depacketizer: Depacketizer, + dequant: Pipeline, + /// Without and with the DC shift of a plane's final transform. + idwt: [Pipeline; 2], + wavelet: Wavelet, + outputs: [Output; OUTPUTS], + uploads: [Upload; UPLOADS], + next_upload: usize, + /// A prologue per upload, then a body per upload and output. + commands: Commands, + /// Which bodies are recorded against the current upload buffers. + recorded: [[bool; OUTPUTS]; UPLOADS], + timeline: Timeline, + submitted: u64, +} + +impl Decoder { + pub fn new(ctx: Context, config: DecodeConfig) -> Result { + if !ctx.supports_decode() { + return Err(Error::Config( + "the context's device was not created for decoding (Roles::DECODE)".into(), + )); + } + let layout = Layout::new(config.width, config.height, config.chroma)?; + let s = ctx.inner().subgroups; + + // Upstream prefers 16 and wider, and falls back to anything from 4. + let dequant_size = s.pick(16, 128).or_else(|| s.pick(4, 128)).ok_or_else(|| { + Error::Unsupported(vec!["a subgroup size for the dequantizer".into()]) + })?; + let dequant = Pipeline::new( + &ctx, + &PipelineDesc { + spirv: shaders::WAVELET_DEQUANT, + entry: entry::WAVELET_DEQUANT, + bindings: &[vk::DescriptorType::STORAGE_IMAGE], + push_size: std::mem::size_of::() as u32, + specialization: &[], + subgroup: dequant_size, + }, + )?; + let idwt_spirv = if ctx.inner().enabled.shader_float16 { + shaders::IDWT_FP16 + } else { + shaders::IDWT_FP32 + }; + let [idwt0, idwt1] = [0, 1].map(|shift| { + Pipeline::new( + &ctx, + &PipelineDesc { + spirv: idwt_spirv, + entry: entry::IDWT, + bindings: &[ + vk::DescriptorType::COMBINED_IMAGE_SAMPLER, + vk::DescriptorType::STORAGE_IMAGE, + ], + push_size: std::mem::size_of::() as u32, + specialization: &[(0, shift)], + subgroup: SubgroupSize::Any, + }, + ) + }); + let wavelet = Wavelet::new(&ctx, &layout)?; + + let mut families = vec![ctx.queue_family()]; + families.extend(config.consumer_family); + // Planes the size the transform writes, not the picture: see + // `PlaneView`. + let (aw, ah) = (layout.aligned_width, layout.aligned_height); + let (acw, ach) = match config.chroma { + Chroma::Yuv420 => (aw / 2, ah / 2), + Chroma::Yuv444 => (aw, ah), + }; + let make_output = || -> Result { + let plane = |w, h| { + Image::new( + &ctx, + &ImageDesc { + format: config.depth.format(), + width: w, + height: h, + layers: 1, + mips: 1, + usage: vk::ImageUsageFlags::SAMPLED + | vk::ImageUsageFlags::STORAGE + | vk::ImageUsageFlags::TRANSFER_SRC, + families: &families, + }, + ) + }; + let mut planes = [plane(aw, ah)?, plane(acw, ach)?, plane(acw, ach)?]; + let mut views = [vk::ImageView::null(); 3]; + for (v, p) in views.iter_mut().zip(planes.iter_mut()) { + *v = p.view(vk::ImageViewType::TYPE_2D, 0, 0, 1)?; + } + Ok(Output { + planes, + views, + held: Arc::new(AtomicBool::new(false)), + }) + }; + let outputs = [make_output()?, make_output()?, make_output()?]; + + let make_upload = || Self::upload_buffers(&ctx, &layout, 256 * 1024); + let uploads = [make_upload()?, make_upload()?]; + + Ok(Self { + depacketizer: Depacketizer::new(layout.clone()), + commands: Commands::new(&ctx, (UPLOADS + UPLOADS * OUTPUTS) as u32)?, + recorded: [[false; OUTPUTS]; UPLOADS], + timeline: Timeline::new(&ctx)?, + ctx, + config, + layout, + dequant, + idwt: [idwt0?, idwt1?], + wavelet, + outputs, + uploads, + next_upload: 0, + submitted: 0, + }) + } + + fn upload_buffers(ctx: &Context, layout: &Layout, payload_bytes: u64) -> Result { + let size = u64::from(layout.blocks_32x32()) * 4 + payload_bytes + PAYLOAD_PADDING; + let usage = vk::BufferUsageFlags::STORAGE_BUFFER | vk::BufferUsageFlags::TRANSFER_SRC; + let host = Buffer::new(ctx, size, usage, Location::Upload)?; + let device = if host.device_local { + None + } else { + Some(Buffer::new( + ctx, + size, + vk::BufferUsageFlags::STORAGE_BUFFER | vk::BufferUsageFlags::TRANSFER_DST, + Location::Device, + )?) + }; + Ok(Upload { + timestamps: if ctx.inner().timestamps { + Some(Timestamps::new(ctx, TIMESTAMPS)?) + } else { + None + }, + host, + device, + last_use: 0, + }) + } + + pub fn config(&self) -> &DecodeConfig { + &self.config + } + + /// Takes one packet. See [`Depacketizer::push`]. + pub fn push_packet(&mut self, packet: &[u8]) -> Result { + Ok(self.depacketizer.push(packet)?) + } + + pub fn readiness(&self) -> Readiness { + self.depacketizer.readiness() + } + + /// GPU time of the most recent decode, once it is done; `None` before it + /// is, or on a queue without timestamps. + pub fn stats(&self) -> Result> { + if self.submitted == 0 || self.timeline_value()? < self.submitted { + return Ok(None); + } + let last = (self.next_upload + UPLOADS - 1) % UPLOADS; + let Some(t) = &self.uploads[last].timestamps else { + return Ok(None); + }; + let passes = t.read()?; + Ok(Some(DecodeStats { + upload_ns: passes[0], + dequant_ns: passes[1], + idwt_ns: passes[2], + })) + } + + fn timeline_value(&self) -> Result { + Ok(unsafe { + self.ctx + .device() + .get_semaphore_counter_value(self.timeline.semaphore)? + }) + } + + /// Forgets the frame being collected. + pub fn clear(&mut self) { + self.depacketizer.clear(); + } + + /// Decodes what has arrived of the current frame once every point in + /// `wait` is reached. Missing blocks decode as zeros. + pub fn decode_after(&mut self, wait: &[TimelinePoint]) -> Result { + let r = self.depacketizer.readiness(); + if !r.has_header { + return Err(Error::NotReady("no start-of-frame header has arrived")); + } + if r.decoded { + return Err(Error::NotReady("this frame was already decoded")); + } + let output = self + .outputs + .iter() + .position(|o| !o.held.load(Ordering::Acquire)) + .ok_or(Error::Busy)?; + + let colour = self.depacketizer.colour().unwrap_or_default(); + let missing_blocks = self.depacketizer.missing_blocks(); + + // The upload slot's last decode must be done before it is rewritten. + let u = self.next_upload; + self.next_upload = (u + 1) % UPLOADS; + self.timeline.wait(self.uploads[u].last_use)?; + + let (offsets, payload) = self.depacketizer.take_for_decode(); + let offsets_bytes = offsets.len() as u64 * 4; + let payload_bytes = payload.len() as u64 * 4; + if offsets_bytes + payload_bytes + PAYLOAD_PADDING > self.uploads[u].host.size { + let grown = (payload_bytes * 2).max(256 * 1024); + self.uploads[u] = Self::upload_buffers(&self.ctx, &self.layout, grown)?; + // The bodies recorded for this upload read the old buffer. + self.recorded[u] = [false; OUTPUTS]; + } + { + let host = &mut self.uploads[u].host; + let bytes = host.bytes_mut(); + let (table, rest) = bytes.split_at_mut(offsets_bytes as usize); + for (dst, src) in table.as_chunks_mut::<4>().0.iter_mut().zip(offsets) { + *dst = src.to_le_bytes(); + } + for (dst, src) in rest.as_chunks_mut::<4>().0.iter_mut().zip(payload) { + *dst = src.to_le_bytes(); + } + let pad = payload_bytes as usize; + rest[pad..pad + PAYLOAD_PADDING as usize].fill(0); + host.flush()?; + } + let upload_size = offsets_bytes + payload_bytes + PAYLOAD_PADDING; + + // Only the prologue is recorded per frame: the copy to device memory, + // when the upload is not there already, is sized to this frame. The + // body, tens of dispatches, is recorded once per upload and output. + let prologue = self.commands.buffers[u]; + let main = self.commands.buffers[UPLOADS + u * OUTPUTS + output]; + self.record_prologue(prologue, u, upload_size)?; + if !self.recorded[u][output] { + self.record_main(main, u, output, offsets_bytes)?; + self.recorded[u][output] = true; + } + + let value = self.submitted + 1; + let mut waits: Vec<_> = wait.iter().map(TimelinePoint::wait_info).collect(); + // Decodes share the wavelet images: each starts after the last. + waits.push(TimelinePoint::new(self.timeline.semaphore, self.submitted).wait_info()); + let signals = [vk::SemaphoreSubmitInfo::default() + .semaphore(self.timeline.semaphore) + .value(value) + .stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS)]; + let cmds = + [prologue, main].map(|c| vk::CommandBufferSubmitInfo::default().command_buffer(c)); + let submit = vk::SubmitInfo2::default() + .wait_semaphore_infos(&waits) + .command_buffer_infos(&cmds) + .signal_semaphore_infos(&signals); + { + let inner = self.ctx.inner(); + let _guard = inner.queue_lock.lock(); + unsafe { + inner + .device + .queue_submit2(inner.queue, &[submit], vk::Fence::null())? + }; + } + self.submitted = value; + self.uploads[u].last_use = value; + + let o = &self.outputs[output]; + o.held.store(true, Ordering::Release); + let (cw, ch) = match self.config.chroma { + Chroma::Yuv420 => (self.config.width / 2, self.config.height / 2), + Chroma::Yuv444 => (self.config.width, self.config.height), + }; + let plane = |i: usize| PlaneView { + image: o.planes[i].image, + view: o.views[i], + format: o.planes[i].format, + width: if i == 0 { self.config.width } else { cw }, + height: if i == 0 { self.config.height } else { ch }, + image_width: o.planes[i].width, + image_height: o.planes[i].height, + }; + Ok(DecodedFrame { + y: plane(0), + cb: plane(1), + cr: plane(2), + width: self.config.width, + height: self.config.height, + chroma: self.config.chroma, + depth: self.config.depth, + colour, + missing_blocks, + ready: TimelinePoint::new(self.timeline.semaphore, value), + _hold: SlotHold(o.held.clone()), + }) + } + + fn begin(&self, cmd: vk::CommandBuffer, once: bool) -> Result<()> { + let device = self.ctx.device(); + unsafe { + device.reset_command_buffer(cmd, vk::CommandBufferResetFlags::empty())?; + device.begin_command_buffer( + cmd, + &vk::CommandBufferBeginInfo::default().flags(if once { + vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT + } else { + vk::CommandBufferUsageFlags::empty() + }), + )?; + } + Ok(()) + } + + fn stamp(&self, cmd: vk::CommandBuffer, upload: usize, i: u32) { + if let Some(t) = &self.uploads[upload].timestamps { + unsafe { + self.ctx.device().cmd_write_timestamp2( + cmd, + vk::PipelineStageFlags2::ALL_COMMANDS, + t.pool, + i, + ) + }; + } + } + + /// Per frame: the copy of this frame's upload into device memory, when + /// the upload is not device local. + fn record_prologue( + &self, + cmd: vk::CommandBuffer, + upload: usize, + upload_size: u64, + ) -> Result<()> { + let ctx = &self.ctx; + let device = ctx.device(); + let up = &self.uploads[upload]; + self.begin(cmd, true)?; + if let Some(t) = &up.timestamps { + unsafe { device.cmd_reset_query_pool(cmd, t.pool, 0, t.count) }; + } + self.stamp(cmd, upload, 0); + if let Some(dst) = &up.device { + unsafe { + device.cmd_copy_buffer( + cmd, + up.host.buffer, + dst.buffer, + &[vk::BufferCopy { + src_offset: 0, + dst_offset: 0, + size: upload_size, + }], + ) + }; + memory_barrier( + ctx, + cmd, + vk::PipelineStageFlags2::COPY, + vk::AccessFlags2::TRANSFER_WRITE, + vk::PipelineStageFlags2::COMPUTE_SHADER, + vk::AccessFlags2::SHADER_READ, + ); + } + self.stamp(cmd, upload, 1); + unsafe { device.end_command_buffer(cmd)? }; + Ok(()) + } + + /// Once per upload and output: dequantize every band, then the inverse + /// transforms into the output planes. Nothing in it changes between + /// frames that use the same upload buffer and output. + fn record_main( + &self, + cmd: vk::CommandBuffer, + upload: usize, + output: usize, + offsets_bytes: u64, + ) -> Result<()> { + let ctx = &self.ctx; + let device = ctx.device(); + let layout = &self.layout; + let chroma = self.config.chroma; + let out = &self.outputs[output]; + let up = &self.uploads[upload]; + self.begin(cmd, false)?; + + let source = up.device.as_ref().map_or(up.host.address, |d| d.address); + let offsets = source; + let payload = source + offsets_bytes; + + let mut fresh: Vec<_> = self + .wavelet + .images() + .into_iter() + .map(|i| (i, vk::ImageLayout::UNDEFINED)) + .collect(); + fresh.extend( + out.planes + .iter() + .map(|p| (p.image, vk::ImageLayout::UNDEFINED)), + ); + to_general( + ctx, + cmd, + &fresh, + vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::SHADER_WRITE, + ); + + // Every band's coefficients. + for (component, level, band, b) in layout.bands() { + self.dequant.dispatch( + cmd, + &[Bind::storage(self.wavelet.bands[component][level])], + &DequantPush { + payload_offsets: offsets, + payload_data_u32: payload, + payload_data_u16: payload, + payload_data_u8: payload, + resolution: [b.width as i32, b.height as i32], + output_layer: band as i32, + block_offset_32x32: b.offset_32x32 as i32, + block_stride_32x32: b.stride_32x32 as i32, + _pad: 0, + }, + (b.width.div_ceil(32), b.height.div_ceil(32), 1), + ); + } + compute_barrier(ctx, cmd); + self.stamp(cmd, upload, 2); + + // Inverse transforms, coarsest level first. + let mirror = self.wavelet.mirror.sampler; + for level in (0..DECOMPOSITION_LEVELS).rev() { + let (w, h) = layout.level_size(level); + // Transposed, as the shader expects. + let push = IdwtPush { + resolution: [h as i32, w as i32], + inv_resolution: [1.0 / h as f32, 1.0 / w as f32], + }; + for component in 0..NUM_COMPONENTS { + if !chroma.has_level(component, level) { + continue; + } + // Where this level's reconstruction goes: the next level's + // LL band, or at the end, the output plane with the DC shift + // undone. + let (target, shift) = if level == 0 || !chroma.has_level(component, level - 1) { + (out.views[component], true) + } else { + (self.wavelet.ll[component][level - 1], false) + }; + self.idwt[usize::from(shift)].dispatch( + cmd, + &[ + Bind::sampled(self.wavelet.bands[component][level], mirror), + Bind::storage(target), + ], + &push, + (h.div_ceil(16), w.div_ceil(16), 1), + ); + } + compute_barrier(ctx, cmd); + } + self.stamp(cmd, upload, 3); + + unsafe { device.end_command_buffer(cmd)? }; + Ok(()) + } +} + +impl Drop for Decoder { + fn drop(&mut self) { + // Nothing may be destroyed while the GPU still uses it. + let _ = self.timeline.wait(self.submitted); + } +} diff --git a/crates/nespyro/src/device.rs b/crates/nespyro/src/device.rs new file mode 100644 index 00000000..bbc52163 --- /dev/null +++ b/crates/nespyro/src/device.rs @@ -0,0 +1,499 @@ +//! What nespyro needs from a device, and the handle to one it runs on. +//! +//! nespyro never creates a device. Both of its callers already have one: the +//! capture layer runs on the game's, the client on its renderer's. So the +//! flow is pixelforge's adopted-device one: ask [`DeviceRequirements::query`] +//! what to enable, merge that into the device being created, then wrap the +//! finished device with [`Context::from_existing`]. + +use std::ffi::CStr; +use std::sync::Arc; + +use ash::vk; +use ash::vk::TaggedStructure as _; + +use crate::error::{Error, Result}; +use crate::pipeline::SubgroupSize; +use crate::sync::QueueLock; + +/// Which halves of the codec a device is being set up for. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Roles { + pub encode: bool, + pub decode: bool, +} + +impl Roles { + pub const ENCODE: Self = Self { + encode: true, + decode: false, + }; + pub const DECODE: Self = Self { + encode: false, + decode: true, + }; + pub const BOTH: Self = Self { + encode: true, + decode: true, + }; +} + +/// Device features nespyro needs, as booleans rather than a `pNext` chain. +/// +/// A caller merging these into its own device creation may already chain +/// `VkPhysicalDeviceVulkan12Features` and friends, and Vulkan forbids chaining +/// those alongside the structs they contain, so nespyro names the bits and the +/// caller sets them wherever it keeps them. Field names match Vulkan's. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct DeviceFeatures { + /// Vulkan 1.0. + pub shader_int16: bool, + pub shader_storage_image_write_without_format: bool, + /// Vulkan 1.1. + pub storage_buffer_16bit_access: bool, + /// Vulkan 1.2. + pub storage_buffer_8bit_access: bool, + pub shader_int8: bool, + /// Required by the encoder, whose quantizer and rate control compute in + /// FP16. Asked of a decode-only device only when it has it: the inverse + /// transform then keeps its tile in native half2 rather than packing it. + pub shader_float16: bool, + pub timeline_semaphore: bool, + pub buffer_device_address: bool, + /// Vulkan 1.3. + pub synchronization2: bool, + pub subgroup_size_control: bool, + pub compute_full_subgroups: bool, +} + +impl DeviceFeatures { + fn for_roles(roles: Roles, available: &Self) -> Self { + Self { + shader_int16: true, + shader_storage_image_write_without_format: true, + storage_buffer_16bit_access: true, + storage_buffer_8bit_access: true, + shader_int8: true, + shader_float16: roles.encode || available.shader_float16, + timeline_semaphore: true, + buffer_device_address: true, + synchronization2: true, + subgroup_size_control: true, + compute_full_subgroups: true, + } + } + + fn query(instance: &ash::Instance, physical_device: vk::PhysicalDevice) -> Self { + let mut v11 = vk::PhysicalDeviceVulkan11Features::default(); + let mut v12 = vk::PhysicalDeviceVulkan12Features::default(); + let mut v13 = vk::PhysicalDeviceVulkan13Features::default(); + let mut f2 = vk::PhysicalDeviceFeatures2::default() + .push(&mut v11) + .push(&mut v12) + .push(&mut v13); + unsafe { instance.get_physical_device_features2(physical_device, &mut f2) }; + let f = f2.features; + Self { + shader_int16: f.shader_int16 != 0, + shader_storage_image_write_without_format: f.shader_storage_image_write_without_format + != 0, + storage_buffer_16bit_access: v11.storage_buffer16_bit_access != 0, + storage_buffer_8bit_access: v12.storage_buffer8_bit_access != 0, + shader_int8: v12.shader_int8 != 0, + shader_float16: v12.shader_float16 != 0, + timeline_semaphore: v12.timeline_semaphore != 0, + buffer_device_address: v12.buffer_device_address != 0, + synchronization2: v13.synchronization2 != 0, + subgroup_size_control: v13.subgroup_size_control != 0, + compute_full_subgroups: v13.compute_full_subgroups != 0, + } + } + + /// The names of the features `self` needs that `available` lacks. + fn missing_from(&self, available: &Self) -> Vec<&'static str> { + let pairs = [ + (self.shader_int16, available.shader_int16, "shaderInt16"), + ( + self.shader_storage_image_write_without_format, + available.shader_storage_image_write_without_format, + "shaderStorageImageWriteWithoutFormat", + ), + ( + self.storage_buffer_16bit_access, + available.storage_buffer_16bit_access, + "storageBuffer16BitAccess", + ), + ( + self.storage_buffer_8bit_access, + available.storage_buffer_8bit_access, + "storageBuffer8BitAccess", + ), + (self.shader_int8, available.shader_int8, "shaderInt8"), + ( + self.shader_float16, + available.shader_float16, + "shaderFloat16", + ), + ( + self.timeline_semaphore, + available.timeline_semaphore, + "timelineSemaphore", + ), + ( + self.buffer_device_address, + available.buffer_device_address, + "bufferDeviceAddress", + ), + ( + self.synchronization2, + available.synchronization2, + "synchronization2", + ), + ( + self.subgroup_size_control, + available.subgroup_size_control, + "subgroupSizeControl", + ), + ( + self.compute_full_subgroups, + available.compute_full_subgroups, + "computeFullSubgroups", + ), + ]; + pairs + .into_iter() + .filter(|(need, have, _)| *need && !*have) + .map(|(_, _, name)| name) + .collect() + } +} + +/// Everything a device must be created with for nespyro to run on it. +#[derive(Debug, Clone)] +pub struct DeviceRequirements { + /// A compute-capable family nespyro would like a queue from: one without + /// graphics where the device has it, so the codec runs beside rendering + /// rather than behind it. Any compute-capable family works. + pub queue_family: u32, + /// Device extensions to enable. Merge with your own. + pub extensions: Vec<&'static CStr>, + /// Device features to enable. Merge with your own. + pub features: DeviceFeatures, +} + +/// Subgroup operations each half needs, from upstream's own checks. +fn required_subgroup_ops(roles: Roles) -> vk::SubgroupFeatureFlags { + let mut ops = vk::SubgroupFeatureFlags::BASIC + | vk::SubgroupFeatureFlags::VOTE + | vk::SubgroupFeatureFlags::ARITHMETIC + | vk::SubgroupFeatureFlags::BALLOT + | vk::SubgroupFeatureFlags::SHUFFLE + | vk::SubgroupFeatureFlags::SHUFFLE_RELATIVE; + if roles.encode { + ops |= vk::SubgroupFeatureFlags::CLUSTERED; + } + ops +} + +/// A device's subgroup properties, as far as nespyro cares. +#[derive(Debug, Clone, Copy)] +pub(crate) struct Subgroups { + pub min: u32, + pub max: u32, + pub ops: vk::SubgroupFeatureFlags, + pub stages: vk::ShaderStageFlags, + pub required_size_stages: vk::ShaderStageFlags, +} + +impl Subgroups { + fn query(instance: &ash::Instance, physical_device: vk::PhysicalDevice) -> Self { + let mut v11 = vk::PhysicalDeviceVulkan11Properties::default(); + let mut v13 = vk::PhysicalDeviceVulkan13Properties::default(); + let mut p2 = vk::PhysicalDeviceProperties2::default() + .push(&mut v11) + .push(&mut v13); + unsafe { instance.get_physical_device_properties2(physical_device, &mut p2) }; + Self { + min: v13.min_subgroup_size, + max: v13.max_subgroup_size, + ops: v11.subgroup_supported_operations, + stages: v11.subgroup_supported_stages, + required_size_stages: v13.required_subgroup_size_stages, + } + } + + /// How a pipeline asking for full subgroups of `lo..=hi` lanes is sized, + /// by Granite's rule, which upstream's shaders were tuned under: when the + /// range covers everything the device offers, the driver picks and may + /// vary it; otherwise the smallest size in range the device has. + pub fn pick(&self, lo: u32, hi: u32) -> Option { + if lo <= self.min && hi >= self.max { + return Some(SubgroupSize::Varying); + } + if lo > self.max || hi < self.min { + return None; + } + Some(SubgroupSize::Full(lo.max(self.min))) + } + + /// The first of `sizes` the device can run at exactly. + pub fn pick_exact(&self, sizes: &[u32]) -> Option { + sizes + .iter() + .copied() + .find(|&s| (self.min..=self.max).contains(&s)) + } +} + +/// Every subgroup-size range the passes of `roles` ask for, from upstream. +fn check_subgroup_sizes(roles: Roles, s: &Subgroups, missing: &mut Vec) { + let mut need = |what: &str, ok: bool| { + if !ok { + missing.push(format!( + "a subgroup size for {what} (device offers {}..={})", + s.min, s.max + )); + } + }; + if roles.encode { + need("the forward transform (4..=128)", s.pick(4, 128).is_some()); + need("the quantizer (8..=128)", s.pick(8, 128).is_some()); + need("rate control analysis (16..=64)", s.pick(16, 64).is_some()); + need( + "rate control resolve (exactly 64, 16 or 32)", + s.pick_exact(&[64, 16, 32]).is_some(), + ); + need("block packing (16..=64)", s.pick(16, 64).is_some()); + } + if roles.decode { + need("the dequantizer (4..=128)", s.pick(4, 128).is_some()); + } +} + +impl DeviceRequirements { + /// What a device on `physical_device` must be created with to run the + /// halves of the codec in `roles`, or why it cannot. + /// + /// Checks the device's real features and properties, never its name. + pub fn query( + instance: &ash::Instance, + physical_device: vk::PhysicalDevice, + roles: Roles, + ) -> Result { + let props = unsafe { instance.get_physical_device_properties(physical_device) }; + let mut missing: Vec = Vec::new(); + + if props.api_version < vk::API_VERSION_1_3 { + missing.push(format!( + "Vulkan 1.3 (device is {}.{})", + vk::api_version_major(props.api_version), + vk::api_version_minor(props.api_version) + )); + } + + let available = DeviceFeatures::query(instance, physical_device); + let features = DeviceFeatures::for_roles(roles, &available); + missing.extend( + features + .missing_from(&available) + .into_iter() + .map(String::from), + ); + + let extensions = [ash::khr::push_descriptor::NAME]; + let present = unsafe { instance.enumerate_device_extension_properties(physical_device) } + .map_err(Error::Vulkan)?; + for ext in extensions { + if !present + .iter() + .any(|p| p.extension_name_as_c_str().is_ok_and(|n| n == ext)) + { + missing.push(ext.to_string_lossy().into_owned()); + } + } + + let subgroups = Subgroups::query(instance, physical_device); + let ops = required_subgroup_ops(roles); + if !subgroups.ops.contains(ops) { + missing.push(format!( + "subgroup operations {:?} (device has {:?})", + ops & !subgroups.ops, + subgroups.ops + )); + } + if !subgroups.stages.contains(vk::ShaderStageFlags::COMPUTE) { + missing.push("subgroup operations in compute shaders".into()); + } + if !subgroups + .required_size_stages + .contains(vk::ShaderStageFlags::COMPUTE) + { + missing.push("required subgroup sizes in compute shaders".into()); + } + check_subgroup_sizes(roles, &subgroups, &mut missing); + + // Formats the codec writes as storage images without a declared format, + // and samples. Only R32_SFLOAT storage is guaranteed by core Vulkan. + use vk::FormatFeatureFlags2 as F; + let needed = F::STORAGE_IMAGE | F::STORAGE_WRITE_WITHOUT_FORMAT | F::SAMPLED_IMAGE; + for format in [ + vk::Format::R8_UNORM, + vk::Format::R16_UNORM, + vk::Format::R16_SFLOAT, + vk::Format::R32_SFLOAT, + ] { + let mut props3 = vk::FormatProperties3::default(); + let mut props = vk::FormatProperties2::default().push(&mut props3); + unsafe { + instance.get_physical_device_format_properties2(physical_device, format, &mut props) + }; + if !props3.optimal_tiling_features.contains(needed) { + missing.push(format!( + "{format:?} as a sampled and format-less storage image (has {:?})", + props3.optimal_tiling_features & needed + )); + } + } + + let families = + unsafe { instance.get_physical_device_queue_family_properties(physical_device) }; + let compute = |f: &vk::QueueFamilyProperties| { + f.queue_flags.contains(vk::QueueFlags::COMPUTE) && f.queue_count > 0 + }; + let queue_family = families + .iter() + .position(|f| compute(f) && !f.queue_flags.contains(vk::QueueFlags::GRAPHICS)) + .or_else(|| families.iter().position(compute)); + let Some(queue_family) = queue_family else { + missing.push("a compute queue".into()); + return Err(Error::Unsupported(missing)); + }; + + if !missing.is_empty() { + return Err(Error::Unsupported(missing)); + } + Ok(Self { + queue_family: queue_family as u32, + extensions: extensions.to_vec(), + features, + }) + } +} + +/// One queue on a device: a family and an index within it, as passed to +/// `vkGetDeviceQueue`. The same as pixelforge's. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct DeviceQueue { + pub family: u32, + pub index: u32, +} + +impl DeviceQueue { + pub fn new(family: u32, index: u32) -> Self { + Self { family, index } + } +} + +/// A caller's device, set up for nespyro. Cheap to clone. +#[derive(Clone)] +pub struct Context(pub(crate) Arc); + +pub(crate) struct ContextInner { + pub device: ash::Device, + pub queue: vk::Queue, + pub queue_family: u32, + pub queue_lock: QueueLock, + pub push_descriptor: ash::khr::push_descriptor::Device, + pub memory: vk::PhysicalDeviceMemoryProperties, + pub subgroups: Subgroups, + /// What the device was created with, as the caller said. + pub roles: Roles, + pub enabled: DeviceFeatures, + pub timestamp_period: f32, + pub timestamps: bool, +} + +impl Context { + /// Wraps a device the caller created and still owns. + /// + /// The device must have been created with everything + /// [`DeviceRequirements::query`] returned for `roles`, and a queue at + /// `queue`. Vulkan offers no way to ask a device which features it was + /// created with, so this trusts the caller for those; it does check that + /// the physical device supports them, and an encoder on a context not + /// created for encoding is refused. + /// + /// The context never destroys the device. Everything nespyro creates on + /// it must be dropped before the caller destroys it. + pub fn from_existing( + instance: ash::Instance, + physical_device: vk::PhysicalDevice, + device: ash::Device, + queue: DeviceQueue, + queue_lock: QueueLock, + roles: Roles, + ) -> Result { + let requirements = DeviceRequirements::query(&instance, physical_device, roles)?; + + let families = + unsafe { instance.get_physical_device_queue_family_properties(physical_device) }; + let family = families.get(queue.family as usize).ok_or_else(|| { + Error::Config(format!("queue family {} does not exist", queue.family)) + })?; + if !family.queue_flags.contains(vk::QueueFlags::COMPUTE) { + return Err(Error::Config(format!( + "queue family {} cannot run compute", + queue.family + ))); + } + if queue.index >= family.queue_count { + return Err(Error::Config(format!( + "queue family {} has {} queues, index {} asked for", + queue.family, family.queue_count, queue.index + ))); + } + let timestamps = family.timestamp_valid_bits > 0; + + let vk_queue = unsafe { device.get_device_queue(queue.family, queue.index) }; + let push_descriptor = ash::khr::push_descriptor::Device::load(&instance, &device); + let memory = unsafe { instance.get_physical_device_memory_properties(physical_device) }; + let props = unsafe { instance.get_physical_device_properties(physical_device) }; + + Ok(Self(Arc::new(ContextInner { + subgroups: Subgroups::query(&instance, physical_device), + roles, + enabled: requirements.features, + device, + queue: vk_queue, + queue_family: queue.family, + queue_lock, + push_descriptor, + memory, + timestamp_period: props.limits.timestamp_period, + timestamps, + }))) + } + + /// Whether an encoder can be made on this context. + pub fn supports_encode(&self) -> bool { + self.0.roles.encode + } + + /// Whether a decoder can be made on this context. + pub fn supports_decode(&self) -> bool { + self.0.roles.decode + } + + /// The queue family nespyro submits on. + pub fn queue_family(&self) -> u32 { + self.0.queue_family + } + + pub fn device(&self) -> &ash::Device { + &self.0.device + } + + pub(crate) fn inner(&self) -> &ContextInner { + &self.0 + } +} diff --git a/crates/nespyro/src/encoder.rs b/crates/nespyro/src/encoder.rs new file mode 100644 index 00000000..96222fa4 --- /dev/null +++ b/crates/nespyro/src/encoder.rs @@ -0,0 +1,1528 @@ +//! The encoder: an RGB image in, a packetized frame out. +//! +//! Shaped like pixelforge's encoder so the capture layer drives both the same +//! way: [`Encoder::encode_after`] waits on the caller's timeline points, +//! submits the whole frame as one submission, and returns a future a +//! completion thread resolves once the frame is read back and packetized. +//! +//! PyroWave is intra only, so there is no GOP, no IDR and no reference to +//! invalidate: every frame decodes alone. + +use std::future::Future; +use std::pin::Pin; +use std::sync::mpsc; +use std::sync::{Arc, Mutex}; +use std::task::{Context as TaskContext, Poll}; +use std::thread::JoinHandle; + +use ash::vk; +use futures_channel::oneshot; + +use crate::bitstream::{ + Chroma, ColourDescription, DECOMPOSITION_LEVELS, GpuPacket, Layout, NUM_COMPONENTS, + SEQUENCE_MASK, SequenceHeader, packetize, +}; +use crate::device::Context; +use crate::error::{Error, Result}; +use crate::gpu::{ + Buffer, Commands, Image, ImageDesc, Location, OwnedView, Timeline, Timestamps, compute_barrier, + memory_barrier, to_general, +}; +use crate::pipeline::{Bind, Pipeline, PipelineDesc, SubgroupSize, entry}; +use crate::rate; +use crate::shaders; +use crate::sync::TimelinePoint; +use crate::wavelet::Wavelet; + +/// What the RGB input already is. Mirrors pixelforge's `ColorSpec`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Source { + /// Ordinary SDR content, sRGB encoded. + Srgb = 0, + /// scRGB: sRGB primaries in linear light, values above 1.0 brighter than + /// SDR white. + Bt709Linear = 1, + /// Like `Bt709Linear` with BT.2020 primaries. + Bt2020Linear = 2, + /// Already HDR10: BT.2020 primaries, PQ encoded. + Bt2020Pq = 3, +} + +impl Source { + /// How bright this source's 1.0 is, in nits, unless the caller says + /// otherwise. pixelforge's `ColorSpec::reference_white_nits`, so the two + /// encoders agree on one surface. + /// + /// scRGB defines its white as 80 nits; the BT.2100 reference white of 203 + /// is right for everything else that needs one. Using 203 for scRGB was + /// the bug: a Windows title's HDR, which DXVK hands over as scRGB, came out + /// two and a half times too bright. `None` for PQ, whose values are + /// already absolute. + pub fn reference_white_nits(self) -> Option { + match self { + Self::Srgb | Self::Bt2020Linear => Some(203.0), + Self::Bt709Linear => Some(80.0), + Self::Bt2020Pq => None, + } + } +} + +/// Bits per sample of the YCbCr planes the codec reads. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Depth { + /// `R8_UNORM` planes, for SDR. + Eight, + /// `R16_UNORM` planes, for HDR. + Sixteen, +} + +impl Depth { + pub(crate) fn format(self) -> vk::Format { + match self { + Depth::Eight => vk::Format::R8_UNORM, + Depth::Sixteen => vk::Format::R16_UNORM, + } + } +} + +/// RGB formats the encoder samples directly. +const INPUT_FORMATS: &[vk::Format] = &[ + vk::Format::B8G8R8A8_UNORM, + vk::Format::R8G8B8A8_UNORM, + vk::Format::A2B10G10R10_UNORM_PACK32, + vk::Format::A2R10G10B10_UNORM_PACK32, + vk::Format::R16G16B16A16_SFLOAT, +]; + +#[derive(Debug, Clone)] +pub struct EncodeConfig { + pub width: u32, + pub height: u32, + pub chroma: Chroma, + pub depth: Depth, + /// What the stream is, and what its header says it is. + pub colour: ColourDescription, + /// What the input is. + pub source: Source, + /// Luminance of a linear source's 1.0, in nits, overriding the source's + /// own ([`Source::reference_white_nits`]). Unused for PQ sources and SDR + /// targets. + pub reference_white_nits: Option, + pub frame_rate: (u32, u32), + pub target_bitrate: u64, + /// Largest packet the packetizer cuts, in bytes. + pub packet_size: usize, +} + +impl EncodeConfig { + /// 4:2:0 SDR at 60 fps and 200 Mbit/s, cut for 1200-byte packets. + pub fn new(width: u32, height: u32) -> Self { + Self { + width, + height, + chroma: Chroma::Yuv420, + depth: Depth::Eight, + colour: ColourDescription::bt709(), + source: Source::Srgb, + reference_white_nits: None, + frame_rate: (60, 1), + target_bitrate: 200_000_000, + packet_size: 1200, + } + } + + pub fn with_chroma(mut self, chroma: Chroma) -> Self { + self.chroma = chroma; + self + } + + pub fn with_depth(mut self, depth: Depth) -> Self { + self.depth = depth; + self + } + + pub fn with_colour(mut self, colour: ColourDescription, source: Source) -> Self { + self.colour = colour; + self.source = source; + self + } + + pub fn with_reference_white(mut self, nits: f32) -> Self { + self.reference_white_nits = Some(nits); + self + } + + pub fn with_frame_rate(mut self, numerator: u32, denominator: u32) -> Self { + self.frame_rate = (numerator, denominator); + self + } + + pub fn with_target_bitrate(mut self, bits_per_second: u64) -> Self { + self.target_bitrate = bits_per_second; + self + } + + pub fn with_packet_size(mut self, bytes: usize) -> Self { + self.packet_size = bytes; + self + } + + fn validate(&self) -> Result<()> { + if self.colour != ColourDescription::bt709() + && self.colour != ColourDescription::bt2020_pq() + { + return Err(Error::Config(format!( + "the encoder produces BT.709 or BT.2020 PQ, full range, centre sited; not {:?}", + self.colour + ))); + } + if !self.colour.is_hdr() && self.source != Source::Srgb { + return Err(Error::Config(format!( + "an SDR stream needs an sRGB source; {:?} needs tone mapping this encoder does not do", + self.source + ))); + } + if self.frame_rate.0 == 0 || self.frame_rate.1 == 0 { + return Err(Error::Config("frame rate must be positive".into())); + } + if self.packet_size < 64 { + return Err(Error::Config(format!( + "packet size {} is below 64 bytes", + self.packet_size + ))); + } + Ok(()) + } + + /// The reference white in force: the caller's, else the source's, else a + /// value the shader never reads (a PQ source has no use for one). + fn effective_reference_white_nits(&self) -> f32 { + self.reference_white_nits + .or(self.source.reference_white_nits()) + .unwrap_or(203.0) + } + + /// Bytes rate control aims each frame at, a whole number of words. + fn target_bytes(&self) -> u64 { + let bits = + self.target_bitrate * u64::from(self.frame_rate.1) / u64::from(self.frame_rate.0); + (bits / 8) & !3 + } +} + +/// GPU time spent in each pass of one frame, in nanoseconds. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct EncodeStats { + pub convert_ns: f64, + pub dwt_ns: f64, + /// Includes clearing the buffers rate control accumulates into. + pub quant_ns: f64, + pub analyze_ns: f64, + pub resolve_ns: f64, + pub packing_ns: f64, + /// Bytes rate control aimed at. + pub target_bytes: u64, + /// Words of block data the GPU produced, including any it could not keep. + pub produced_words: u32, + /// Blocks dropped because rate control overshot what was read back. + /// Zero unless something is wrong. + pub dropped_blocks: u32, + /// Bytes the quantizer asked of its scratch payload, and what it has. + /// Asking for more means 8×8 blocks were dropped as all zero before rate + /// control ever saw them. + pub scratch_bytes: u32, + pub scratch_capacity: u32, +} + +impl EncodeStats { + /// Whether the quantizer ran out of scratch payload. + pub fn scratch_overflowed(&self) -> bool { + self.scratch_bytes > self.scratch_capacity + } +} + +impl EncodeStats { + pub fn gpu_ns(&self) -> f64 { + self.convert_ns + + self.dwt_ns + + self.quant_ns + + self.analyze_ns + + self.resolve_ns + + self.packing_ns + } +} + +/// One frame, ready for the network. +#[derive(Debug, Clone)] +pub struct EncodedFrame { + /// The start-of-frame header, then every non-empty block in index order. + pub data: Vec, + /// `data` cut into packets that each parse on their own. + pub packets: Vec>, + /// Leading packets that carry the coarsest bands, whose loss cannot be + /// masked. Worth the most reliable delivery. + pub critical_packets: usize, + /// Frames encoded before this one. + pub index: u64, + /// The 3-bit frame counter in its headers. + pub sequence: u8, + pub stats: EncodeStats, +} + +/// Resolves to the frame once it is read back. Futures resolve in submission +/// order; dropping one is harmless. +pub struct EncodeFuture { + rx: oneshot::Receiver>, +} + +impl Future for EncodeFuture { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll { + match Pin::new(&mut self.rx).poll(cx) { + Poll::Ready(Ok(result)) => Poll::Ready(result), + Poll::Ready(Err(oneshot::Canceled)) => Poll::Ready(Err(Error::Cancelled)), + Poll::Pending => Poll::Pending, + } + } +} + +/// Submissions allowed in flight: one on the GPU while the previous one is +/// packetized, as in pixelforge. +const DEPTH: usize = 2; + +/// Timestamps around the six passes. +const TIMESTAMPS: u32 = 7; + +#[repr(C)] +#[derive(Clone, Copy)] +struct RgbPush { + width: u32, + height: u32, + source: u32, + hdr: u32, + reference_white_nits: f32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct DwtPush { + resolution: [i32; 2], + inv_resolution: [f32; 2], + aligned_resolution: [i32; 2], +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct QuantPush { + block_meta: u64, + block_stats: u64, + payload_counter: u64, + payload_data: u64, + resolution: [i32; 2], + resolution_8x8_blocks: [i32; 2], + inv_resolution: [f32; 2], + input_layer: f32, + quant_resolution: f32, + block_offset: i32, + block_stride: i32, + rdo_distortion_scale: f32, + payload_capacity: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct AnalyzePush { + buckets: u64, + total_savings_per_bucket: u64, + rdo_operations: u64, + block_stats: u64, + resolution: [i32; 2], + resolution_8x8_blocks: [i32; 2], + block_offset_8x8: i32, + block_stride_8x8: i32, + block_offset_32x32: i32, + block_stride_32x32: i32, + total_wg_count: u32, + num_blocks_aligned: u32, + block_index_shamt: u32, + _pad: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct FinalizePush { + total_savings_per_bucket: u64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct ResolvePush { + buckets: u64, + total_savings_per_bucket: u64, + rdo_operations: u64, + quant_data: u64, + frame_params: u64, + num_blocks_per_subdivision: u32, + _pad: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct PackingPush { + bitstream_data: u64, + bitstream_data_16b: u64, + bitstream_data_8b: u64, + bitstream_meta: u64, + block_meta: u64, + payload_counter: u64, + payload_data: u64, + block_stats: u64, + quant_data: u64, + frame_params: u64, + resolution: [i32; 2], + resolution_32x32_blocks: [i32; 2], + resolution_8x8_blocks: [i32; 2], + quant_resolution_code: u32, + block_offset_32x32: i32, + block_stride_32x32: i32, + block_offset_8x8: i32, + block_stride_8x8: i32, + bitstream_capacity_words: u32, +} + +// Vulkan guarantees 128 bytes of push constants and no more. +const _: () = assert!(std::mem::size_of::() <= 128); + +struct Pipelines { + rgb: Pipeline, + /// Without and with the DC shift of a plane's first transform. + dwt: [Pipeline; 2], + quant: Pipeline, + analyze: Pipeline, + finalize: Pipeline, + resolve: Pipeline, + packing: Pipeline, +} + +impl Pipelines { + fn new(ctx: &Context, chroma: Chroma) -> Result { + let s = ctx.inner().subgroups; + let no = || Error::Unsupported(vec!["a usable subgroup size".into()]); + let sampled = vk::DescriptorType::COMBINED_IMAGE_SAMPLER; + let storage = vk::DescriptorType::STORAGE_IMAGE; + let size = |n: usize| n as u32; + + let rgb = Pipeline::new( + ctx, + &PipelineDesc { + spirv: shaders::RGB_TO_YCBCR, + entry: entry::RGB_TO_YCBCR, + bindings: &[vk::DescriptorType::SAMPLED_IMAGE, storage, storage, storage], + push_size: size(std::mem::size_of::()), + specialization: &[(0, u32::from(chroma == Chroma::Yuv420))], + subgroup: SubgroupSize::Any, + }, + )?; + let dwt_size = s.pick(4, 128).ok_or_else(no)?; + let dwt = [0, 1].map(|shift| { + Pipeline::new( + ctx, + &PipelineDesc { + spirv: shaders::DWT, + entry: entry::DWT, + bindings: &[sampled, storage], + push_size: size(std::mem::size_of::()), + specialization: &[(0, shift)], + subgroup: dwt_size, + }, + ) + }); + let [dwt0, dwt1] = dwt; + let quant = Pipeline::new( + ctx, + &PipelineDesc { + spirv: shaders::WAVELET_QUANT, + entry: entry::WAVELET_QUANT, + bindings: &[sampled], + push_size: size(std::mem::size_of::()), + specialization: &[(1, 0)], + subgroup: s.pick(8, 128).ok_or_else(no)?, + }, + )?; + let analyze = Pipeline::new( + ctx, + &PipelineDesc { + spirv: shaders::ANALYZE_RATE_CONTROL, + entry: entry::ANALYZE_RATE_CONTROL, + bindings: &[], + push_size: size(std::mem::size_of::()), + specialization: &[], + subgroup: s.pick(16, 64).ok_or_else(no)?, + }, + )?; + let finalize = Pipeline::new( + ctx, + &PipelineDesc { + spirv: shaders::ANALYZE_RATE_CONTROL_FINALIZE, + entry: entry::ANALYZE_RATE_CONTROL_FINALIZE, + bindings: &[], + push_size: size(std::mem::size_of::()), + specialization: &[], + subgroup: SubgroupSize::Any, + }, + )?; + // Upstream's preference order. + let resolve_size = s.pick_exact(&[64, 16, 32]).ok_or_else(no)?; + let resolve_spirv = match resolve_size { + 16 => shaders::RESOLVE_RATE_CONTROL_16, + 32 => shaders::RESOLVE_RATE_CONTROL_32, + _ => shaders::RESOLVE_RATE_CONTROL_64, + }; + let resolve = Pipeline::new( + ctx, + &PipelineDesc { + spirv: resolve_spirv, + entry: entry::RESOLVE_RATE_CONTROL, + bindings: &[], + push_size: size(std::mem::size_of::()), + specialization: &[], + subgroup: SubgroupSize::Full(resolve_size), + }, + )?; + let packing = Pipeline::new( + ctx, + &PipelineDesc { + spirv: shaders::BLOCK_PACKING, + entry: entry::BLOCK_PACKING, + bindings: &[], + push_size: size(std::mem::size_of::()), + specialization: &[], + subgroup: s.pick(16, 64).ok_or_else(no)?, + }, + )?; + Ok(Self { + rgb, + dwt: [dwt0?, dwt1?], + quant, + analyze, + finalize, + resolve, + packing, + }) + } +} + +/// What a slot's read-back holds, in order. +struct ReadbackLayout { + /// The GPU's block table. + table: u64, + /// The payload buffer's two counters. + counters: u64, + /// The bitstream words. + bitstream: u64, +} + +struct Slot { + /// Sized for the current target plus headroom, and replaced with a + /// larger one when the target grows. Only touched by the thread that + /// holds the slot. + readback: Mutex, + timestamps: Option, +} + +fn readback_buffer(ctx: &Context, layout: &ReadbackLayout, bitstream_bytes: u64) -> Result { + Buffer::new( + ctx, + layout.bitstream + bitstream_bytes, + vk::BufferUsageFlags::TRANSFER_DST, + Location::Readback, + ) +} + +/// What a frame is encoded from. The views are kept alive until the GPU is +/// done with the frame. +enum Input { + /// The caller's RGB image, converted by the encoder. + Rgb { + view: OwnedView, + layout: vk::ImageLayout, + }, + /// The caller's own Y, Cb and Cr planes. + Planes { + views: [OwnedView; 3], + layout: vk::ImageLayout, + }, +} + +/// Handed to the completion thread per frame. +struct Work { + slot: usize, + value: u64, + _input: Input, + header: SequenceHeader, + index: u64, + copied_words: u32, + target_bytes: u64, + scratch_capacity: u32, + tx: oneshot::Sender>, +} + +pub struct Encoder { + ctx: Context, + layout: Arc, + config: EncodeConfig, + pipelines: Pipelines, + wavelet: Wavelet, + planes: [Image; 3], + plane_views: [vk::ImageView; 3], + block_meta: Buffer, + block_stats: Buffer, + payload: Buffer, + quant: Buffer, + buckets: Buffer, + bitstream: Buffer, + table: Buffer, + readback_layout: Arc, + /// What changes per frame: the sequence number and the target size. + frame_params: Buffer, + commands: Commands, + /// Which bodies each slot has recorded. + recorded: [[bool; 2]; DEPTH], + timeline: Arc, + submitted: u64, + sequence: u8, + frames: u64, + slots: Arc<[Slot; DEPTH]>, + free_slots: mpsc::Receiver, + work: Option>, + worker: Option>, +} + +impl Encoder { + pub fn new(ctx: Context, config: EncodeConfig) -> Result { + config.validate()?; + if !ctx.supports_encode() { + return Err(Error::Config( + "the context's device was not created for encoding (Roles::ENCODE)".into(), + )); + } + let layout = Arc::new(Layout::new(config.width, config.height, config.chroma)?); + if layout.blocks_32x32() > u32::from(u16::MAX) { + // Rate control packs a block index into 16 bits. + return Err(Error::Config(format!( + "{}x{} has {} blocks; rate control addresses at most 65535", + config.width, + config.height, + layout.blocks_32x32() + ))); + } + + let pipelines = Pipelines::new(&ctx, config.chroma)?; + let wavelet = Wavelet::new(&ctx, &layout)?; + + let (cw, ch) = match config.chroma { + Chroma::Yuv420 => (config.width / 2, config.height / 2), + Chroma::Yuv444 => (config.width, config.height), + }; + let plane = |w, h| { + Image::new( + &ctx, + &ImageDesc { + format: config.depth.format(), + width: w, + height: h, + layers: 1, + mips: 1, + usage: vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::STORAGE, + families: &[], + }, + ) + }; + let mut planes = [ + plane(config.width, config.height)?, + plane(cw, ch)?, + plane(cw, ch)?, + ]; + let mut plane_views = [vk::ImageView::null(); 3]; + for (view, image) in plane_views.iter_mut().zip(planes.iter_mut()) { + *view = image.view(vk::ImageViewType::TYPE_2D, 0, 0, 1)?; + } + + let storage = vk::BufferUsageFlags::STORAGE_BUFFER; + let transfer = vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::TRANSFER_SRC; + let blocks_8x8 = u64::from(layout.blocks_8x8()); + let blocks_32x32 = u64::from(layout.blocks_32x32()); + let per_sub = u64::from(rate::blocks_per_subdivision(layout.blocks_32x32())); + // The quantizer's scratch holds, per 8×8 block, eight 4×2 subblocks of + // a sign byte and at most 16 bit planes; the block statistics' 15 + // entries already assume fewer. Upstream sizes this by estimate, + // aligned width × height × 2, which high-entropy 4:4:4 content + // exceeds: its encoder then packs blocks from bytes it never wrote. + // The quantizer still refuses to write past this, as a backstop. + let scratch_bytes = blocks_8x8 * 8 * (1 + 16); + // No frame's blocks can exceed this, whatever rate control does. + let bitstream_bytes = blocks_32x32 * u64::from(crate::bitstream::LARGEST_BLOCK_WORDS) * 4; + + let block_meta = Buffer::new(&ctx, blocks_8x8 * 8, storage, Location::Device)?; + let block_stats = Buffer::new(&ctx, blocks_8x8 * 64, storage, Location::Device)?; + let payload = Buffer::new( + &ctx, + 8 + scratch_bytes, + storage | transfer, + Location::Device, + )?; + let quant = Buffer::new(&ctx, blocks_32x32 * 4, storage | transfer, Location::Device)?; + let buckets = Buffer::new( + &ctx, + rate::RDO_BUCKET_OFFSET + + u64::from(rate::NUM_RDO_BUCKETS * rate::BLOCK_SPACE_SUBDIVISION) * 4 + + u64::from(rate::NUM_RDO_BUCKETS) + * per_sub + * u64::from(rate::BLOCK_SPACE_SUBDIVISION) + * 8, + storage | transfer, + Location::Device, + )?; + // Block packing may never write past this; the read-back copies only + // the part rate control aimed at, plus headroom. + let bitstream = Buffer::new(&ctx, bitstream_bytes, storage | transfer, Location::Device)?; + let table = Buffer::new(&ctx, blocks_32x32 * 8, storage | transfer, Location::Device)?; + + let readback_layout = Arc::new(ReadbackLayout { + table: 0, + counters: blocks_32x32 * 8, + bitstream: blocks_32x32 * 8 + 8, + }); + let make_slot = || -> Result { + Ok(Slot { + readback: Mutex::new(readback_buffer( + &ctx, + &readback_layout, + config.target_bytes() + blocks_32x32 * 8, + )?), + timestamps: if ctx.inner().timestamps { + Some(Timestamps::new(&ctx, TIMESTAMPS)?) + } else { + None + }, + }) + }; + let slots = Arc::new([make_slot()?, make_slot()?]); + + let commands = Commands::new(&ctx, (DEPTH * COMMANDS_PER_SLOT) as u32)?; + let frame_params = Buffer::new( + &ctx, + 16, + storage | vk::BufferUsageFlags::TRANSFER_DST, + Location::Device, + )?; + let timeline = Arc::new(Timeline::new(&ctx)?); + + let (free_tx, free_slots) = mpsc::channel(); + for i in 0..DEPTH { + free_tx.send(i).unwrap(); + } + let (work, work_rx) = mpsc::channel::(); + let worker = { + let layout = layout.clone(); + let slots = slots.clone(); + let timeline = timeline.clone(); + let readback_layout = readback_layout.clone(); + let packet_size = config.packet_size; + std::thread::Builder::new() + .name("nespyro-encode".into()) + .spawn(move || { + for w in work_rx { + let slot = w.slot; + let result = complete( + &layout, + &slots[slot], + &timeline, + &readback_layout, + packet_size, + &w, + ); + let _ = w.tx.send(result); + drop(w._input); + // The receiver only goes away with the encoder. + let _ = free_tx.send(slot); + } + }) + .map_err(|e| Error::Config(format!("could not start the completion thread: {e}")))? + }; + + Ok(Self { + ctx, + layout, + config, + pipelines, + wavelet, + planes, + plane_views, + block_meta, + block_stats, + payload, + quant, + buckets, + bitstream, + table, + readback_layout, + frame_params, + commands, + recorded: [[false; 2]; DEPTH], + timeline, + submitted: 0, + sequence: 0, + frames: 0, + slots, + free_slots, + work: Some(work), + worker: Some(worker), + }) + } + + pub fn config(&self) -> &EncodeConfig { + &self.config + } + + /// Takes effect on the next frame. Nothing is rebuilt. + pub fn set_target_bitrate(&mut self, bits_per_second: u64) { + self.config.target_bitrate = bits_per_second; + } + + /// Takes effect on the next frame: the conversion and the header's colour + /// fields change, the planes' depth does not. + pub fn set_colour_description( + &mut self, + colour: ColourDescription, + source: Source, + ) -> Result<()> { + let next = EncodeConfig { + colour, + source, + ..self.config.clone() + }; + next.validate()?; + self.config = next; + Ok(()) + } + + /// Encodes `image` once every point in `wait` is reached. + /// + /// `image` is sampled in `layout` through a view of `format`, which must + /// be one of the UNORM or SFLOAT formats the encoder reads; the caller + /// keeps it in that layout until the returned future resolves, and must + /// not write it until then. + pub fn encode_after( + &mut self, + image: vk::Image, + format: vk::Format, + layout: vk::ImageLayout, + wait: &[TimelinePoint], + ) -> Result { + if !INPUT_FORMATS.contains(&format) { + return Err(Error::Config(format!( + "cannot sample {format:?}; use one of {INPUT_FORMATS:?} (an sRGB format would \ + be linearized by the sampler)" + ))); + } + + let input = Input::Rgb { + view: OwnedView::new(&self.ctx, image, format)?, + layout, + }; + self.submit(input, wait) + } + + /// Encodes the caller's own full-range YCbCr planes once every point in + /// `wait` is reached, skipping the encoder's colour conversion. + /// + /// The planes are sampled in `layout`, and must be single-level images of + /// the configured depth's format (`R8_UNORM` or `R16_UNORM`) and sizes: + /// the configured width and height for Y, and for Cb and Cr the same, or + /// half of each for 4:2:0. The header still carries the configured + /// colour description; the planes must be what it says. + pub fn encode_planes_after( + &mut self, + planes: [vk::Image; 3], + layout: vk::ImageLayout, + wait: &[TimelinePoint], + ) -> Result { + let format = self.config.depth.format(); + let [y, cb, cr] = planes; + let input = Input::Planes { + views: [ + OwnedView::new(&self.ctx, y, format)?, + OwnedView::new(&self.ctx, cb, format)?, + OwnedView::new(&self.ctx, cr, format)?, + ], + layout, + }; + self.submit(input, wait) + } + + fn submit(&mut self, input: Input, wait: &[TimelinePoint]) -> Result { + // A free slot means its previous frame has been read back, so its + // command buffers and read-back buffer are no longer in use. + let slot = self.free_slots.recv().map_err(|_| Error::Cancelled)?; + + self.sequence = (self.sequence + 1) & SEQUENCE_MASK; + let target_bytes = self.config.target_bytes(); + // Rate control aims at the target less the start-of-frame header, and + // the read-back takes the same headroom upstream gives its bitstream + // buffer: one table entry's worth per block. + let headroom = u64::from(self.layout.blocks_32x32()) * 8; + let copy_bytes = (target_bytes + headroom).min(self.bitstream.size); + let copied_words = (copy_bytes / 4) as u32; + { + // The slot is ours until it is handed to the completion thread. + let mut readback = self.slots[slot].readback.lock().unwrap(); + if readback.size < self.readback_layout.bitstream + copy_bytes { + *readback = readback_buffer(&self.ctx, &self.readback_layout, copy_bytes)?; + } + } + + // Only what changes per frame is recorded per frame: the prologue + // (the frame's parameters, and whatever reads the caller's images) + // and the epilogue (the read-back, sized to the target). The body, + // tens of dispatches, is recorded once per slot. + let kind = match input { + Input::Rgb { .. } => MainKind::Rgb, + Input::Planes { .. } => MainKind::Planes, + }; + let prologue = self.commands.buffers[slot * COMMANDS_PER_SLOT]; + let epilogue = self.commands.buffers[slot * COMMANDS_PER_SLOT + 1]; + let main = self.commands.buffers[slot * COMMANDS_PER_SLOT + 2 + kind as usize]; + self.record_prologue(prologue, slot, &input, target_bytes)?; + if !self.recorded[slot][kind as usize] { + self.record_main(main, slot, kind)?; + self.recorded[slot][kind as usize] = true; + } + self.record_epilogue(epilogue, slot, copy_bytes)?; + + let value = self.submitted + 1; + // The previous frame shares every scratch resource with this one, so + // this frame starts after it ends. + let mut waits: Vec<_> = wait.iter().map(TimelinePoint::wait_info).collect(); + waits.push(TimelinePoint::new(self.timeline.semaphore, self.submitted).wait_info()); + let signals = [vk::SemaphoreSubmitInfo::default() + .semaphore(self.timeline.semaphore) + .value(value) + .stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS)]; + let cmds = [prologue, main, epilogue] + .map(|c| vk::CommandBufferSubmitInfo::default().command_buffer(c)); + let submit = vk::SubmitInfo2::default() + .wait_semaphore_infos(&waits) + .command_buffer_infos(&cmds) + .signal_semaphore_infos(&signals); + { + let inner = self.ctx.inner(); + let _guard = inner.queue_lock.lock(); + unsafe { + inner + .device + .queue_submit2(inner.queue, &[submit], vk::Fence::null())? + }; + } + self.submitted = value; + + let (tx, rx) = oneshot::channel(); + let header = SequenceHeader { + width: self.config.width, + height: self.config.height, + sequence: self.sequence, + total_blocks: 0, + chroma: self.config.chroma, + colour: self.config.colour, + }; + let work = Work { + slot, + value, + _input: input, + header, + index: self.frames, + copied_words, + target_bytes, + scratch_capacity: (self.payload.size - 8) as u32, + tx, + }; + self.frames += 1; + self.work + .as_ref() + .expect("the work channel lives as long as the encoder") + .send(work) + .map_err(|_| Error::Cancelled)?; + Ok(EncodeFuture { rx }) + } + + /// The point the most recent submission signals. Waiting on it on the GPU + /// orders work after the encoder is done reading the source image. + pub fn last_submission(&self) -> TimelinePoint { + TimelinePoint::new(self.timeline.semaphore, self.submitted) + } + + fn begin(&self, cmd: vk::CommandBuffer, once: bool) -> Result<()> { + let device = self.ctx.device(); + unsafe { + device.reset_command_buffer(cmd, vk::CommandBufferResetFlags::empty())?; + device.begin_command_buffer( + cmd, + &vk::CommandBufferBeginInfo::default().flags(if once { + vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT + } else { + vk::CommandBufferUsageFlags::empty() + }), + )?; + } + Ok(()) + } + + fn stamp(&self, cmd: vk::CommandBuffer, slot: usize, i: u32) { + if let Some(t) = &self.slots[slot].timestamps { + unsafe { + self.ctx.device().cmd_write_timestamp2( + cmd, + vk::PipelineStageFlags2::ALL_COMMANDS, + t.pool, + i, + ) + }; + } + } + + /// Per frame: fresh scratch, the frame's parameters, and every pass that + /// reads the caller's images, since those change from frame to frame. + fn record_prologue( + &self, + cmd: vk::CommandBuffer, + slot: usize, + input: &Input, + target_bytes: u64, + ) -> Result<()> { + let ctx = &self.ctx; + let device = ctx.device(); + let config = &self.config; + self.begin(cmd, true)?; + if let Some(t) = &self.slots[slot].timestamps { + unsafe { device.cmd_reset_query_pool(cmd, t.pool, 0, t.count) }; + } + + // Every scratch image starts undefined: a frame never reads what the + // one before it left. + let mut fresh: Vec<_> = self + .wavelet + .images() + .into_iter() + .map(|i| (i, vk::ImageLayout::UNDEFINED)) + .collect(); + fresh.extend( + self.planes + .iter() + .map(|p| (p.image, vk::ImageLayout::UNDEFINED)), + ); + to_general( + ctx, + cmd, + &fresh, + vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::SHADER_WRITE, + ); + + let params = [ + u32::from(self.sequence), + (target_bytes.saturating_sub(8) / 4) as u32, + 0, + 0, + ]; + let params: Vec = params.iter().flat_map(|v| v.to_le_bytes()).collect(); + unsafe { device.cmd_update_buffer(cmd, self.frame_params.buffer, 0, ¶ms) }; + memory_barrier( + ctx, + cmd, + vk::PipelineStageFlags2::ALL_TRANSFER, + vk::AccessFlags2::TRANSFER_WRITE, + vk::PipelineStageFlags2::COMPUTE_SHADER, + vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::SHADER_WRITE, + ); + + self.stamp(cmd, slot, 0); + match input { + Input::Rgb { view, layout } => { + self.pipelines.rgb.dispatch( + cmd, + &[ + Bind::texture(view.view, *layout), + Bind::storage(self.plane_views[0]), + Bind::storage(self.plane_views[1]), + Bind::storage(self.plane_views[2]), + ], + &RgbPush { + width: config.width, + height: config.height, + source: config.source as u32, + hdr: u32::from(config.colour.is_hdr()), + reference_white_nits: config.effective_reference_white_nits(), + }, + ( + config.width.div_ceil(2).div_ceil(8), + config.height.div_ceil(2).div_ceil(8), + 1, + ), + ); + compute_barrier(ctx, cmd); + self.stamp(cmd, slot, 1); + } + Input::Planes { views, layout } => { + self.stamp(cmd, slot, 1); + let views = views.each_ref().map(|v| v.view); + self.record_dwt(cmd, &views, *layout, |level, component| { + reads_plane(config.chroma, level, component) + }); + } + } + + unsafe { device.end_command_buffer(cmd)? }; + Ok(()) + } + + /// Once per slot and input kind: every pass after the first transforms + /// of the planes. Nothing in it changes between frames; what does is + /// read from `frame_params`. + fn record_main(&self, cmd: vk::CommandBuffer, slot: usize, kind: MainKind) -> Result<()> { + let ctx = &self.ctx; + let device = ctx.device(); + let layout = &*self.layout; + let config = &self.config; + let p = &self.pipelines; + self.begin(cmd, false)?; + + // Forward transforms. From the encoder's own planes they all run + // here; from the caller's, the ones reading its planes ran in the + // prologue. + match kind { + MainKind::Rgb => { + self.record_dwt(cmd, &self.plane_views, vk::ImageLayout::GENERAL, |_, _| { + true + }) + } + MainKind::Planes => self.record_dwt( + cmd, + &self.plane_views, + vk::ImageLayout::GENERAL, + |level, component| !reads_plane(config.chroma, level, component), + ), + } + self.stamp(cmd, slot, 2); + + // Clear what the quantizer and rate control accumulate into. Here, + // recorded once, rather than per frame: Mesa's ANV spends 12 ms of CPU + // recording a fill of the 8 MB bucket buffer 1080p 4:4:4 has, though + // not the 16 MB one at 4K. + unsafe { + device.cmd_fill_buffer(cmd, self.payload.buffer, 0, 8, 0); + device.cmd_fill_buffer(cmd, self.buckets.buffer, 0, vk::WHOLE_SIZE, 0); + device.cmd_fill_buffer(cmd, self.quant.buffer, 0, vk::WHOLE_SIZE, 0); + } + memory_barrier( + ctx, + cmd, + vk::PipelineStageFlags2::CLEAR, + vk::AccessFlags2::TRANSFER_WRITE, + vk::PipelineStageFlags2::COMPUTE_SHADER, + vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::SHADER_WRITE, + ); + + // Quantize every band. + let border = self.wavelet.border.sampler; + let payload_capacity = (self.payload.size - 8) as u32; + for (component, level, band, b) in layout.bands() { + let res = rate::quant_resolution(level, component, band); + p.quant.dispatch( + cmd, + &[Bind::sampled(self.wavelet.bands[component][level], border)], + &QuantPush { + block_meta: self.block_meta.address, + block_stats: self.block_stats.address, + payload_counter: self.payload.address, + payload_data: self.payload.address + 8, + resolution: [b.width as i32, b.height as i32], + resolution_8x8_blocks: [ + b.width.div_ceil(8) as i32, + b.height.div_ceil(8) as i32, + ], + inv_resolution: [1.0 / b.width as f32, 1.0 / b.height as f32], + input_layer: band as f32, + quant_resolution: 1.0 / rate::decode_quant(rate::encode_quant(1.0 / res)), + block_offset: b.offset_8x8 as i32, + block_stride: b.stride_8x8 as i32, + rdo_distortion_scale: rate::rdo_distortion_scale( + level, + component, + band, + config.chroma, + ) / 256.0, + payload_capacity, + }, + (b.width.div_ceil(32), b.height.div_ceil(32), 1), + ); + } + compute_barrier(ctx, cmd); + self.stamp(cmd, slot, 3); + + // Rate control: sort every possible saving into buckets. + let per_sub = rate::blocks_per_subdivision(layout.blocks_32x32()); + let savings = self.buckets.address + rate::RDO_BUCKET_OFFSET; + let operations = + savings + u64::from(rate::NUM_RDO_BUCKETS * rate::BLOCK_SPACE_SUBDIVISION) * 4; + for (_, _, _, b) in layout.bands() { + p.analyze.dispatch( + cmd, + &[], + &AnalyzePush { + buckets: self.buckets.address, + total_savings_per_bucket: savings, + rdo_operations: operations, + block_stats: self.block_stats.address, + resolution: [b.width as i32, b.height as i32], + resolution_8x8_blocks: [ + b.width.div_ceil(8) as i32, + b.height.div_ceil(8) as i32, + ], + block_offset_8x8: b.offset_8x8 as i32, + block_stride_8x8: b.stride_8x8 as i32, + block_offset_32x32: b.offset_32x32 as i32, + block_stride_32x32: b.stride_32x32 as i32, + total_wg_count: layout.blocks_32x32(), + num_blocks_aligned: per_sub * rate::BLOCK_SPACE_SUBDIVISION, + block_index_shamt: per_sub.trailing_zeros(), + _pad: 0, + }, + (b.width.div_ceil(32), b.height.div_ceil(32), 1), + ); + } + compute_barrier(ctx, cmd); + p.finalize.dispatch( + cmd, + &[], + &FinalizePush { + total_savings_per_bucket: savings, + }, + (1, 1, 1), + ); + compute_barrier(ctx, cmd); + self.stamp(cmd, slot, 4); + + // Rate control: spend buckets until the frame fits. + p.resolve.dispatch( + cmd, + &[], + &ResolvePush { + buckets: self.buckets.address, + total_savings_per_bucket: savings, + rdo_operations: operations, + quant_data: self.quant.address, + frame_params: self.frame_params.address, + num_blocks_per_subdivision: per_sub, + _pad: 0, + }, + (rate::NUM_RDO_BUCKETS * rate::BLOCK_SPACE_SUBDIVISION, 1, 1), + ); + compute_barrier(ctx, cmd); + self.stamp(cmd, slot, 5); + + // Pack the final blocks. + let capacity_words = (self.bitstream.size / 4) as u32; + for (component, level, band, b) in layout.bands() { + let res = rate::quant_resolution(level, component, band); + let (bx, by) = b.blocks_32x32(); + p.packing.dispatch( + cmd, + &[], + &PackingPush { + bitstream_data: self.bitstream.address, + bitstream_data_16b: self.bitstream.address, + bitstream_data_8b: self.bitstream.address, + bitstream_meta: self.table.address, + block_meta: self.block_meta.address, + payload_counter: self.payload.address, + payload_data: self.payload.address + 8, + block_stats: self.block_stats.address, + quant_data: self.quant.address, + frame_params: self.frame_params.address, + resolution: [b.width as i32, b.height as i32], + resolution_32x32_blocks: [bx as i32, by as i32], + resolution_8x8_blocks: [ + b.width.div_ceil(8) as i32, + b.height.div_ceil(8) as i32, + ], + quant_resolution_code: u32::from(rate::encode_quant(1.0 / res)), + block_offset_32x32: b.offset_32x32 as i32, + block_stride_32x32: b.stride_32x32 as i32, + block_offset_8x8: b.offset_8x8 as i32, + block_stride_8x8: b.stride_8x8 as i32, + bitstream_capacity_words: capacity_words, + }, + (bx.div_ceil(2), by.div_ceil(2), 1), + ); + } + memory_barrier( + ctx, + cmd, + vk::PipelineStageFlags2::COMPUTE_SHADER, + vk::AccessFlags2::SHADER_WRITE, + vk::PipelineStageFlags2::COPY, + vk::AccessFlags2::TRANSFER_READ, + ); + self.stamp(cmd, slot, 6); + + unsafe { device.end_command_buffer(cmd)? }; + Ok(()) + } + + /// Per frame: the table, the counters, and the bitstream up to the + /// target plus headroom, into the slot's read-back buffer. + fn record_epilogue(&self, cmd: vk::CommandBuffer, slot: usize, copy_bytes: u64) -> Result<()> { + let ctx = &self.ctx; + let device = ctx.device(); + let rb = &*self.readback_layout; + let readback = self.slots[slot].readback.lock().unwrap().buffer; + self.begin(cmd, true)?; + unsafe { + device.cmd_copy_buffer( + cmd, + self.table.buffer, + readback, + &[vk::BufferCopy { + src_offset: 0, + dst_offset: rb.table, + size: self.table.size, + }], + ); + device.cmd_copy_buffer( + cmd, + self.payload.buffer, + readback, + &[vk::BufferCopy { + src_offset: 0, + dst_offset: rb.counters, + size: 8, + }], + ); + device.cmd_copy_buffer( + cmd, + self.bitstream.buffer, + readback, + &[vk::BufferCopy { + src_offset: 0, + dst_offset: rb.bitstream, + size: copy_bytes, + }], + ); + } + memory_barrier( + ctx, + cmd, + vk::PipelineStageFlags2::COPY, + vk::AccessFlags2::TRANSFER_WRITE, + vk::PipelineStageFlags2::HOST, + vk::AccessFlags2::HOST_READ, + ); + unsafe { device.end_command_buffer(cmd)? }; + Ok(()) + } + + /// The forward transforms `include` selects, level by level, reading the + /// planes through `planes` in `planes_layout`. + fn record_dwt( + &self, + cmd: vk::CommandBuffer, + planes: &[vk::ImageView; 3], + planes_layout: vk::ImageLayout, + include: impl Fn(usize, usize) -> bool, + ) { + let ctx = &self.ctx; + let layout = &*self.layout; + let chroma = self.config.chroma; + let mirror = self.wavelet.mirror.sampler; + for level in 0..DECOMPOSITION_LEVELS { + let mut any = false; + // A component indexes the planes, their views and the wavelet + // bands alike. + #[allow(clippy::needless_range_loop)] + for component in 0..NUM_COMPONENTS { + if !chroma.has_level(component, level) || !include(level, component) { + continue; + } + // What this component's transform at this level reads: its + // plane where it enters the transform, otherwise the LL band + // the level above left. + let (view, view_layout, resolution, aligned, shift) = + if reads_plane(chroma, level, component) { + let aligned = if level == 0 { + [layout.aligned_width, layout.aligned_height] + } else { + [layout.aligned_width >> 1, layout.aligned_height >> 1] + }; + ( + planes[component], + planes_layout, + [self.planes[component].width, self.planes[component].height], + aligned, + true, + ) + } else { + let (w, h) = layout.level_size(level - 1); + ( + self.wavelet.ll[component][level - 1], + vk::ImageLayout::GENERAL, + [w, h], + [w, h], + false, + ) + }; + self.pipelines.dwt[usize::from(shift)].dispatch( + cmd, + &[ + Bind::sampled_in(view, mirror, view_layout), + Bind::storage(self.wavelet.bands[component][level]), + ], + &DwtPush { + resolution: [resolution[0] as i32, resolution[1] as i32], + inv_resolution: [1.0 / resolution[0] as f32, 1.0 / resolution[1] as f32], + aligned_resolution: [aligned[0] as i32, aligned[1] as i32], + }, + (aligned[0].div_ceil(32), aligned[1].div_ceil(32), 1), + ); + any = true; + } + if any { + compute_barrier(ctx, cmd); + } + } + } +} + +/// Whether a component's transform at `level` reads its plane: at level 0, +/// or at level 1 for 4:2:0 chroma, whose half-size plane enters there. +fn reads_plane(chroma: Chroma, level: usize, component: usize) -> bool { + level == 0 || (level == 1 && !chroma.has_level(component, 0)) +} + +/// Which body a frame runs: after the encoder's own conversion, or after the +/// first transforms of the caller's planes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MainKind { + Rgb = 0, + Planes = 1, +} + +/// Prologue, epilogue, and one body for each kind of input. +const COMMANDS_PER_SLOT: usize = 4; + +/// Reads one frame back and packetizes it, on the completion thread. +fn complete( + layout: &Layout, + slot: &Slot, + timeline: &Timeline, + rb: &ReadbackLayout, + packet_size: usize, + w: &Work, +) -> Result { + timeline.wait(w.value)?; + let readback = slot.readback.lock().unwrap(); + readback.invalidate()?; + let bytes = readback.bytes(); + + let blocks = layout.blocks_32x32() as usize; + let words = |offset: u64, count: usize| -> Vec { + let start = offset as usize; + bytes[start..start + 4 * count] + .as_chunks::<4>() + .0 + .iter() + .map(|b| u32::from_le_bytes(*b)) + .collect() + }; + let raw_table = words(rb.table, 2 * blocks); + let counters = words(rb.counters, 2); + let bitstream = words(rb.bitstream, w.copied_words as usize); + + let mut table: Vec = raw_table + .as_chunks::<2>() + .0 + .iter() + .map(|p| GpuPacket { + offset_words: p[0], + num_words: p[1], + }) + .collect(); + + // Blocks rate control produced beyond what was read back are dropped, + // whole; so, earlier, were blocks past the scratch payload's capacity. + // Either is a frame that overshot badly, and says so in the stats. + let mut dropped_blocks = 0; + for entry in &mut table { + if entry.num_words != 0 && entry.offset_words + entry.num_words > w.copied_words { + entry.num_words = 0; + dropped_blocks += 1; + } + } + if dropped_blocks > 0 { + tracing::warn!( + frame = w.index, + dropped_blocks, + produced_words = counters[1], + read_back_words = w.copied_words, + "rate control overshot the read-back; blocks dropped" + ); + } + + let frame = packetize(layout, w.header, &table, &bitstream, packet_size)?; + + let passes = match &slot.timestamps { + Some(t) => t.read()?, + None => vec![0.0; TIMESTAMPS as usize - 1], + }; + let stats = EncodeStats { + convert_ns: passes[0], + dwt_ns: passes[1], + quant_ns: passes[2], + analyze_ns: passes[3], + resolve_ns: passes[4], + packing_ns: passes[5], + target_bytes: w.target_bytes, + produced_words: counters[1], + dropped_blocks, + scratch_bytes: counters[0], + scratch_capacity: w.scratch_capacity, + }; + if stats.scratch_overflowed() { + tracing::warn!( + frame = w.index, + scratch_bytes = stats.scratch_bytes, + scratch_capacity = stats.scratch_capacity, + "the quantizer's scratch payload overflowed; blocks dropped" + ); + } + + Ok(EncodedFrame { + data: frame.data, + packets: frame.packets, + critical_packets: frame.critical_packets, + index: w.index, + sequence: w.header.sequence, + stats, + }) +} + +impl Drop for Encoder { + fn drop(&mut self) { + // Let the completion thread drain what is in flight, then stop it. + drop(self.work.take()); + if let Some(worker) = self.worker.take() { + let _ = worker.join(); + } + // Nothing may be destroyed while the GPU still uses it. + let _ = self.timeline.wait(self.submitted); + } +} diff --git a/crates/nespyro/src/error.rs b/crates/nespyro/src/error.rs new file mode 100644 index 00000000..6cce983d --- /dev/null +++ b/crates/nespyro/src/error.rs @@ -0,0 +1,34 @@ +use ash::vk; +use thiserror::Error; + +use crate::bitstream::BitstreamError; + +pub type Result = std::result::Result; + +#[derive(Debug, Error)] +pub enum Error { + #[error("Vulkan: {0}")] + Vulkan(vk::Result), + /// The device lacks what the codec needs; every missing piece is named. + #[error("device cannot run nespyro, missing: {}", .0.join(", "))] + Unsupported(Vec), + #[error("invalid configuration: {0}")] + Config(String), + #[error(transparent)] + Bitstream(#[from] BitstreamError), + /// Every output slot is held by the caller; drop a frame and try again. + #[error("every output slot is still held")] + Busy, + /// The encoder was shut down before the frame was read back. + #[error("encoder shut down before the frame was read back")] + Cancelled, + /// Nothing to decode: no start-of-frame header yet, or already decoded. + #[error("nothing to decode: {0}")] + NotReady(&'static str), +} + +impl From for Error { + fn from(r: vk::Result) -> Self { + Error::Vulkan(r) + } +} diff --git a/crates/nespyro/src/gpu.rs b/crates/nespyro/src/gpu.rs new file mode 100644 index 00000000..d73fba6b --- /dev/null +++ b/crates/nespyro/src/gpu.rs @@ -0,0 +1,617 @@ +//! The little of Granite nespyro needs: buffers, images, samplers, command +//! buffers and semaphores, each owning what it creates. +//! +//! Every allocation is dedicated. A codec instance makes a few dozen objects +//! once and then reuses them for every frame, so a suballocator would buy +//! nothing here. + +use ash::vk; +use ash::vk::TaggedStructure as _; + +use crate::device::Context; +use crate::error::{Error, Result}; + +/// Where a buffer's memory lives. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Location { + /// Only the GPU touches it. + Device, + /// Written by the CPU, read by the GPU. Device-local when the device has + /// such memory mappable (ReBAR, or an integrated GPU), so shaders read it + /// without a copy. + Upload, + /// Written by the GPU, read by the CPU. Cached, so reading it back is not + /// an uncached crawl. + Readback, +} + +fn find_memory_type( + ctx: &Context, + bits: u32, + preferred: &[vk::MemoryPropertyFlags], +) -> Option<(u32, vk::MemoryPropertyFlags)> { + let memory = &ctx.inner().memory; + preferred.iter().find_map(|&want| { + (0..memory.memory_type_count).find_map(|i| { + let flags = memory.memory_types[i as usize].property_flags; + (bits & (1 << i) != 0 && flags.contains(want)).then_some((i, flags)) + }) + }) +} + +fn allocate( + ctx: &Context, + requirements: vk::MemoryRequirements, + preferred: &[vk::MemoryPropertyFlags], + device_address: bool, + dedicated: vk::MemoryDedicatedAllocateInfo, +) -> Result<(vk::DeviceMemory, vk::MemoryPropertyFlags)> { + let (index, flags) = find_memory_type(ctx, requirements.memory_type_bits, preferred) + .ok_or_else(|| { + Error::Unsupported(vec![format!( + "a memory type with any of {preferred:?} for bits {:#x}", + requirements.memory_type_bits + )]) + })?; + let mut flags_info = + vk::MemoryAllocateFlagsInfo::default().flags(vk::MemoryAllocateFlags::DEVICE_ADDRESS); + let mut dedicated = dedicated; + let mut info = vk::MemoryAllocateInfo::default() + .allocation_size(requirements.size) + .memory_type_index(index) + .push(&mut dedicated); + if device_address { + info = info.push(&mut flags_info); + } + let memory = unsafe { ctx.device().allocate_memory(&info, None)? }; + Ok((memory, flags)) +} + +pub(crate) struct Buffer { + ctx: Context, + pub buffer: vk::Buffer, + memory: vk::DeviceMemory, + pub size: u64, + pub address: u64, + mapped: *mut u8, + coherent: bool, + /// Whether the memory is device local, so shaders read it at full speed. + pub device_local: bool, +} + +// SAFETY: the mapping is plain memory owned by this buffer; synchronising its +// use with the GPU is the owner's business, as with any Vulkan memory. +unsafe impl Send for Buffer {} +unsafe impl Sync for Buffer {} + +impl Buffer { + pub fn new( + ctx: &Context, + size: u64, + usage: vk::BufferUsageFlags, + location: Location, + ) -> Result { + let device = ctx.device(); + let usage = usage | vk::BufferUsageFlags::SHADER_DEVICE_ADDRESS; + let buffer = unsafe { + device.create_buffer( + &vk::BufferCreateInfo::default() + .size(size.max(4)) + .usage(usage) + .sharing_mode(vk::SharingMode::EXCLUSIVE), + None, + )? + }; + let requirements = unsafe { device.get_buffer_memory_requirements(buffer) }; + use vk::MemoryPropertyFlags as M; + let preferred: &[M] = match location { + Location::Device => &[M::DEVICE_LOCAL], + Location::Upload => &[ + M::DEVICE_LOCAL | M::HOST_VISIBLE | M::HOST_COHERENT, + M::HOST_VISIBLE | M::HOST_COHERENT, + ], + Location::Readback => &[ + M::HOST_VISIBLE | M::HOST_CACHED | M::HOST_COHERENT, + M::HOST_VISIBLE | M::HOST_CACHED, + M::HOST_VISIBLE | M::HOST_COHERENT, + ], + }; + let allocated = allocate( + ctx, + requirements, + preferred, + true, + vk::MemoryDedicatedAllocateInfo::default().buffer(buffer), + ); + let (memory, flags) = match allocated { + Ok(m) => m, + Err(e) => { + unsafe { device.destroy_buffer(buffer, None) }; + return Err(e); + } + }; + unsafe { device.bind_buffer_memory(buffer, memory, 0)? }; + let address = unsafe { + device.get_buffer_device_address(&vk::BufferDeviceAddressInfo::default().buffer(buffer)) + }; + let mapped = if flags.contains(M::HOST_VISIBLE) { + unsafe { device.map_memory(memory, 0, vk::WHOLE_SIZE, vk::MemoryMapFlags::empty())? } + .cast() + } else { + std::ptr::null_mut() + }; + Ok(Self { + ctx: ctx.clone(), + buffer, + memory, + size, + address, + mapped, + coherent: flags.contains(M::HOST_COHERENT), + device_local: flags.contains(M::DEVICE_LOCAL), + }) + } + + /// Whether the CPU can reach this buffer's memory directly. + pub fn is_mapped(&self) -> bool { + !self.mapped.is_null() + } + + /// The mapping. Only valid for host-visible buffers. + pub fn bytes(&self) -> &[u8] { + assert!(self.is_mapped()); + unsafe { std::slice::from_raw_parts(self.mapped, self.size as usize) } + } + + pub fn bytes_mut(&mut self) -> &mut [u8] { + assert!(self.is_mapped()); + unsafe { std::slice::from_raw_parts_mut(self.mapped, self.size as usize) } + } + + /// Makes the GPU's writes visible to the CPU, after the fence or + /// semaphore that ordered them has been waited on. + pub fn invalidate(&self) -> Result<()> { + if !self.coherent { + unsafe { + self.ctx.device().invalidate_mapped_memory_ranges(&[ + vk::MappedMemoryRange::default() + .memory(self.memory) + .offset(0) + .size(vk::WHOLE_SIZE), + ])? + }; + } + Ok(()) + } + + /// Makes the CPU's writes visible to the GPU, before the submission that + /// reads them. + pub fn flush(&self) -> Result<()> { + if !self.coherent { + unsafe { + self.ctx + .device() + .flush_mapped_memory_ranges(&[vk::MappedMemoryRange::default() + .memory(self.memory) + .offset(0) + .size(vk::WHOLE_SIZE)])? + }; + } + Ok(()) + } +} + +impl Drop for Buffer { + fn drop(&mut self) { + let device = self.ctx.device(); + unsafe { + device.destroy_buffer(self.buffer, None); + device.free_memory(self.memory, None); + } + } +} + +/// A 2D image, possibly an array with mips, with a view per subresource +/// range asked for. +pub(crate) struct Image { + ctx: Context, + pub image: vk::Image, + memory: vk::DeviceMemory, + pub format: vk::Format, + pub width: u32, + pub height: u32, + views: Vec, +} + +pub(crate) struct ImageDesc<'a> { + pub format: vk::Format, + pub width: u32, + pub height: u32, + pub layers: u32, + pub mips: u32, + pub usage: vk::ImageUsageFlags, + /// Families that share the image. One or none means exclusive. + pub families: &'a [u32], +} + +impl Image { + pub fn new(ctx: &Context, desc: &ImageDesc) -> Result { + let device = ctx.device(); + let mut families: Vec = desc.families.to_vec(); + families.sort_unstable(); + families.dedup(); + let concurrent = families.len() > 1; + let info = vk::ImageCreateInfo::default() + .image_type(vk::ImageType::TYPE_2D) + .format(desc.format) + .extent(vk::Extent3D { + width: desc.width, + height: desc.height, + depth: 1, + }) + .mip_levels(desc.mips) + .array_layers(desc.layers) + .samples(vk::SampleCountFlags::TYPE_1) + .tiling(vk::ImageTiling::OPTIMAL) + .usage(desc.usage) + .initial_layout(vk::ImageLayout::UNDEFINED); + let info = if concurrent { + info.sharing_mode(vk::SharingMode::CONCURRENT) + .queue_family_indices(&families) + } else { + info.sharing_mode(vk::SharingMode::EXCLUSIVE) + }; + let image = unsafe { device.create_image(&info, None)? }; + let requirements = unsafe { device.get_image_memory_requirements(image) }; + let allocated = allocate( + ctx, + requirements, + &[vk::MemoryPropertyFlags::DEVICE_LOCAL], + false, + vk::MemoryDedicatedAllocateInfo::default().image(image), + ); + let (memory, _) = match allocated { + Ok(m) => m, + Err(e) => { + unsafe { device.destroy_image(image, None) }; + return Err(e); + } + }; + unsafe { device.bind_image_memory(image, memory, 0)? }; + Ok(Self { + ctx: ctx.clone(), + image, + memory, + format: desc.format, + width: desc.width, + height: desc.height, + views: Vec::new(), + }) + } + + /// A view of `layers` layers from `base_layer` at mip `mip`. Owned by the + /// image. + pub fn view( + &mut self, + view_type: vk::ImageViewType, + mip: u32, + base_layer: u32, + layers: u32, + ) -> Result { + let view = unsafe { + self.ctx.device().create_image_view( + &vk::ImageViewCreateInfo::default() + .image(self.image) + .view_type(view_type) + .format(self.format) + .subresource_range(vk::ImageSubresourceRange { + aspect_mask: vk::ImageAspectFlags::COLOR, + base_mip_level: mip, + level_count: 1, + base_array_layer: base_layer, + layer_count: layers, + }), + None, + )? + }; + self.views.push(view); + Ok(view) + } +} + +impl Drop for Image { + fn drop(&mut self) { + let device = self.ctx.device(); + unsafe { + for &v in &self.views { + device.destroy_image_view(v, None); + } + device.destroy_image(self.image, None); + device.free_memory(self.memory, None); + } + } +} + +/// A view the caller's image is sampled through, made per frame and dropped +/// once that frame's work is known done. +pub(crate) struct OwnedView { + ctx: Context, + pub view: vk::ImageView, +} + +impl OwnedView { + pub fn new(ctx: &Context, image: vk::Image, format: vk::Format) -> Result { + let view = unsafe { + ctx.device().create_image_view( + &vk::ImageViewCreateInfo::default() + .image(image) + .view_type(vk::ImageViewType::TYPE_2D) + .format(format) + .subresource_range(vk::ImageSubresourceRange { + aspect_mask: vk::ImageAspectFlags::COLOR, + base_mip_level: 0, + level_count: 1, + base_array_layer: 0, + layer_count: 1, + }), + None, + )? + }; + Ok(Self { + ctx: ctx.clone(), + view, + }) + } +} + +impl Drop for OwnedView { + fn drop(&mut self) { + unsafe { self.ctx.device().destroy_image_view(self.view, None) }; + } +} + +/// Nearest-filtered samplers: gathers read exact texels, and the address mode +/// is the only thing that differs. +pub(crate) struct Sampler { + ctx: Context, + pub sampler: vk::Sampler, +} + +impl Sampler { + pub fn new(ctx: &Context, address_mode: vk::SamplerAddressMode) -> Result { + let sampler = unsafe { + ctx.device().create_sampler( + &vk::SamplerCreateInfo::default() + .mag_filter(vk::Filter::NEAREST) + .min_filter(vk::Filter::NEAREST) + .mipmap_mode(vk::SamplerMipmapMode::NEAREST) + .address_mode_u(address_mode) + .address_mode_v(address_mode) + .address_mode_w(address_mode) + .border_color(vk::BorderColor::FLOAT_TRANSPARENT_BLACK) + .max_lod(vk::LOD_CLAMP_NONE), + None, + )? + }; + Ok(Self { + ctx: ctx.clone(), + sampler, + }) + } +} + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { self.ctx.device().destroy_sampler(self.sampler, None) }; + } +} + +/// A command pool with one primary buffer per slot. +pub(crate) struct Commands { + ctx: Context, + pool: vk::CommandPool, + pub buffers: Vec, +} + +impl Commands { + pub fn new(ctx: &Context, count: u32) -> Result { + let device = ctx.device(); + let pool = unsafe { + device.create_command_pool( + &vk::CommandPoolCreateInfo::default() + .queue_family_index(ctx.queue_family()) + .flags(vk::CommandPoolCreateFlags::RESET_COMMAND_BUFFER), + None, + )? + }; + let buffers = unsafe { + device.allocate_command_buffers( + &vk::CommandBufferAllocateInfo::default() + .command_pool(pool) + .level(vk::CommandBufferLevel::PRIMARY) + .command_buffer_count(count), + ) + }; + let buffers = match buffers { + Ok(b) => b, + Err(e) => { + unsafe { device.destroy_command_pool(pool, None) }; + return Err(e.into()); + } + }; + Ok(Self { + ctx: ctx.clone(), + pool, + buffers, + }) + } +} + +impl Drop for Commands { + fn drop(&mut self) { + unsafe { self.ctx.device().destroy_command_pool(self.pool, None) }; + } +} + +pub(crate) struct Timeline { + ctx: Context, + pub semaphore: vk::Semaphore, +} + +impl Timeline { + pub fn new(ctx: &Context) -> Result { + let mut kind = vk::SemaphoreTypeCreateInfo::default() + .semaphore_type(vk::SemaphoreType::TIMELINE) + .initial_value(0); + let semaphore = unsafe { + ctx.device() + .create_semaphore(&vk::SemaphoreCreateInfo::default().push(&mut kind), None)? + }; + Ok(Self { + ctx: ctx.clone(), + semaphore, + }) + } + + /// Blocks until the semaphore reaches `value`. + pub fn wait(&self, value: u64) -> Result<()> { + let semaphores = [self.semaphore]; + let values = [value]; + unsafe { + self.ctx.device().wait_semaphores( + &vk::SemaphoreWaitInfo::default() + .semaphores(&semaphores) + .values(&values), + u64::MAX, + )? + }; + Ok(()) + } +} + +impl Drop for Timeline { + fn drop(&mut self) { + unsafe { self.ctx.device().destroy_semaphore(self.semaphore, None) }; + } +} + +/// A query pool for GPU timestamps around each pass. +pub(crate) struct Timestamps { + ctx: Context, + pub pool: vk::QueryPool, + pub count: u32, +} + +impl Timestamps { + pub fn new(ctx: &Context, count: u32) -> Result { + let pool = unsafe { + ctx.device().create_query_pool( + &vk::QueryPoolCreateInfo::default() + .query_type(vk::QueryType::TIMESTAMP) + .query_count(count), + None, + )? + }; + Ok(Self { + ctx: ctx.clone(), + pool, + count, + }) + } + + /// Nanoseconds between consecutive timestamps, once the work that wrote + /// them is known done. + pub fn read(&self) -> Result> { + let mut raw = vec![0u64; self.count as usize]; + unsafe { + self.ctx.device().get_query_pool_results( + self.pool, + 0, + &mut raw, + vk::QueryResultFlags::TYPE_64, + )? + }; + let period = f64::from(self.ctx.inner().timestamp_period); + Ok(raw + .windows(2) + .map(|w| w[1].wrapping_sub(w[0]) as f64 * period) + .collect()) + } +} + +impl Drop for Timestamps { + fn drop(&mut self) { + unsafe { self.ctx.device().destroy_query_pool(self.pool, None) }; + } +} + +/// A barrier between two compute passes: every shader write before it is +/// visible to every shader read and write after it. +pub(crate) fn compute_barrier(ctx: &Context, cmd: vk::CommandBuffer) { + memory_barrier( + ctx, + cmd, + vk::PipelineStageFlags2::COMPUTE_SHADER, + vk::AccessFlags2::SHADER_WRITE, + vk::PipelineStageFlags2::COMPUTE_SHADER, + vk::AccessFlags2::SHADER_READ | vk::AccessFlags2::SHADER_WRITE, + ); +} + +pub(crate) fn memory_barrier( + ctx: &Context, + cmd: vk::CommandBuffer, + src_stage: vk::PipelineStageFlags2, + src_access: vk::AccessFlags2, + dst_stage: vk::PipelineStageFlags2, + dst_access: vk::AccessFlags2, +) { + let barrier = [vk::MemoryBarrier2::default() + .src_stage_mask(src_stage) + .src_access_mask(src_access) + .dst_stage_mask(dst_stage) + .dst_access_mask(dst_access)]; + unsafe { + ctx.device().cmd_pipeline_barrier2( + cmd, + &vk::DependencyInfo::default().memory_barriers(&barrier), + ) + }; +} + +/// Moves `image` from `old` to `GENERAL`, ordered after `src_stage`. +pub(crate) fn to_general( + ctx: &Context, + cmd: vk::CommandBuffer, + images: &[(vk::Image, vk::ImageLayout)], + dst_access: vk::AccessFlags2, +) { + let barriers: Vec<_> = images + .iter() + .map(|&(image, old)| { + vk::ImageMemoryBarrier2::default() + .src_stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS) + .src_access_mask(vk::AccessFlags2::MEMORY_WRITE) + .dst_stage_mask(vk::PipelineStageFlags2::COMPUTE_SHADER) + .dst_access_mask(dst_access) + .old_layout(old) + .new_layout(vk::ImageLayout::GENERAL) + .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED) + .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED) + .image(image) + .subresource_range(vk::ImageSubresourceRange { + aspect_mask: vk::ImageAspectFlags::COLOR, + base_mip_level: 0, + level_count: vk::REMAINING_MIP_LEVELS, + base_array_layer: 0, + layer_count: vk::REMAINING_ARRAY_LAYERS, + }) + }) + .collect(); + unsafe { + ctx.device().cmd_pipeline_barrier2( + cmd, + &vk::DependencyInfo::default().image_memory_barriers(&barriers), + ) + }; +} diff --git a/crates/nespyro/src/lib.rs b/crates/nespyro/src/lib.rs new file mode 100644 index 00000000..4e99ae9e --- /dev/null +++ b/crates/nespyro/src/lib.rs @@ -0,0 +1,68 @@ +//! PyroWave, a compute-shader wavelet video codec, on `ash`. +//! +//! A port of Hans-Kristian Arntzen's PyroWave (MIT) from C++ and Granite to +//! Rust, with the shaders ported from GLSL to Slang. The bitstream is +//! upstream's exactly, frozen at PyroWave `89f7e47`. +//! +//! PyroWave is intra-only: every frame stands alone, with no reference frames, +//! no entropy coding and no motion search. It spends bandwidth — around +//! 170 Mbit/s at 1080p60 — to buy latency that no hardware encoder matches, +//! which makes it the codec for a direct, wired path. +//! +//! # Using it +//! +//! nespyro runs on a device the caller creates. Ask +//! [`DeviceRequirements::query`] what to enable for the [`Roles`] the device +//! will have, merge that into the device's creation, and wrap the result with +//! [`Context::from_existing`]. +//! +//! - [`Encoder::encode_after`] takes an RGB image and timeline points to wait +//! on, and returns a future of an [`EncodedFrame`]: the frame's bytes, cut +//! into packets that each parse on their own. +//! - [`Decoder::push_packet`] takes those packets in any order; +//! [`Decoder::readiness`] says how much of the frame is there; and +//! [`Decoder::decode_after`] returns its planes with the timeline point +//! they are ready at. +//! - [`bitstream`] is the format itself, with no Vulkan in it. +//! +//! Only desktop GPUs are supported: upstream's fragment-shader and +//! texel-buffer paths for mobile GPUs are not carried. A device without what +//! the compute paths need is refused with everything it lacks named. +//! +//! # Testing +//! +//! `cargo test -p nespyro` runs the bitstream tests, which need no GPU. The +//! rest are `#[ignore]`d and say what they need; they fail rather than skip +//! when it is missing, and every GPU test runs under Khronos validation and +//! fails on any validation error. +//! +//! - `--test encode --test roundtrip -- --ignored`: a Vulkan 1.3 GPU. Pick +//! one with `NESPYRO_TEST_DEVICE=`. +//! - `--test reference --test bench -- --ignored`: also the C++ reference. +//! Build it with `crates/nespyro/reference/build.sh

` and point +//! `NESPYRO_REFERENCE` at the harness it prints; on a machine with several +//! GPUs, `NESPYRO_REFERENCE_VID=` puts it on the same +//! one. Run `bench` with `--release`. + +pub mod bitstream; +mod decoder; +mod device; +mod encoder; +mod error; +mod gpu; +mod pipeline; +mod rate; +mod sync; +mod wavelet; + +pub use bitstream::{Chroma, ColourDescription, Matrix, Primaries, Range, Siting, Transfer}; +pub use decoder::{DecodeConfig, DecodeStats, DecodedFrame, Decoder, PlaneView}; +pub use device::{Context, DeviceFeatures, DeviceQueue, DeviceRequirements, Roles}; +pub use encoder::{Depth, EncodeConfig, EncodeFuture, EncodeStats, EncodedFrame, Encoder, Source}; +pub use error::{Error, Result}; +pub use sync::{QueueLock, TimelinePoint}; + +/// The SPIR-V `build.rs` compiled, one module per pass and variant. +mod shaders { + include!(concat!(env!("OUT_DIR"), "/shaders.rs")); +} diff --git a/crates/nespyro/src/pipeline.rs b/crates/nespyro/src/pipeline.rs new file mode 100644 index 00000000..d7b89f55 --- /dev/null +++ b/crates/nespyro/src/pipeline.rs @@ -0,0 +1,287 @@ +//! Compute pipelines, with the layouts Granite derived by reflection written +//! out by hand. +//! +//! Every pass binds its images and samplers with push descriptors, so there +//! is no descriptor pool, and reaches its buffers through device addresses in +//! push constants, so buffers need no descriptors at all. + +use std::ffi::CStr; + +use ash::vk; +use ash::vk::TaggedStructure as _; + +use crate::device::Context; +use crate::error::{Error, Result}; + +/// How a pipeline's subgroups are sized. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SubgroupSize { + /// Whatever the driver likes; the shader does not care. + Any, + /// Whatever the driver likes, possibly varying, but every subgroup full. + Varying, + /// Exactly this size, every subgroup full. + Full(u32), +} + +pub(crate) struct Pipeline { + ctx: Context, + set_layout: vk::DescriptorSetLayout, + layout: vk::PipelineLayout, + pipeline: vk::Pipeline, + push_size: u32, +} + +pub(crate) struct PipelineDesc<'a> { + pub spirv: &'a [u8], + pub entry: &'a CStr, + pub bindings: &'a [vk::DescriptorType], + pub push_size: u32, + /// `(constant_id, value)` pairs, every value a 32-bit word. + pub specialization: &'a [(u32, u32)], + pub subgroup: SubgroupSize, +} + +impl Pipeline { + pub fn new(ctx: &Context, desc: &PipelineDesc) -> Result { + let device = ctx.device(); + + let words = ash::util::read_spv(&mut std::io::Cursor::new(desc.spirv)) + .map_err(|e| Error::Config(format!("embedded SPIR-V is malformed: {e}")))?; + let module = unsafe { + device + .create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&words), None)? + }; + + let bindings: Vec<_> = desc + .bindings + .iter() + .enumerate() + .map(|(i, &ty)| { + vk::DescriptorSetLayoutBinding::default() + .binding(i as u32) + .descriptor_type(ty) + .descriptor_count(1) + .stage_flags(vk::ShaderStageFlags::COMPUTE) + }) + .collect(); + let set_layout = unsafe { + device.create_descriptor_set_layout( + &vk::DescriptorSetLayoutCreateInfo::default() + .flags(vk::DescriptorSetLayoutCreateFlags::PUSH_DESCRIPTOR_KHR) + .bindings(&bindings), + None, + )? + }; + let ranges = [vk::PushConstantRange::default() + .stage_flags(vk::ShaderStageFlags::COMPUTE) + .offset(0) + .size(desc.push_size)]; + let set_layouts = [set_layout]; + let layout = unsafe { + device.create_pipeline_layout( + &vk::PipelineLayoutCreateInfo::default() + .set_layouts(&set_layouts) + .push_constant_ranges(if desc.push_size > 0 { &ranges } else { &[] }), + None, + )? + }; + + let entries: Vec<_> = desc + .specialization + .iter() + .enumerate() + .map(|(i, &(id, _))| vk::SpecializationMapEntry { + constant_id: id, + offset: 4 * i as u32, + size: 4, + }) + .collect(); + let data: Vec = desc + .specialization + .iter() + .flat_map(|&(_, v)| v.to_ne_bytes()) + .collect(); + let specialization = vk::SpecializationInfo::default() + .map_entries(&entries) + .data(&data); + + let mut required = vk::PipelineShaderStageRequiredSubgroupSizeCreateInfo::default(); + let mut stage = vk::PipelineShaderStageCreateInfo::default() + .stage(vk::ShaderStageFlags::COMPUTE) + .module(module) + .name(desc.entry) + .specialization_info(&specialization); + match desc.subgroup { + SubgroupSize::Any => {} + SubgroupSize::Varying => { + stage = stage.flags( + vk::PipelineShaderStageCreateFlags::REQUIRE_FULL_SUBGROUPS + | vk::PipelineShaderStageCreateFlags::ALLOW_VARYING_SUBGROUP_SIZE, + ); + } + SubgroupSize::Full(size) => { + required = required.required_subgroup_size(size); + stage = stage + .flags(vk::PipelineShaderStageCreateFlags::REQUIRE_FULL_SUBGROUPS) + .push(&mut required); + } + } + + let created = unsafe { + device.create_compute_pipelines( + vk::PipelineCache::null(), + &[vk::ComputePipelineCreateInfo::default() + .stage(stage) + .layout(layout)], + None, + ) + }; + unsafe { device.destroy_shader_module(module, None) }; + let pipeline = match created { + Ok(p) => p[0], + Err((_, e)) => { + unsafe { + device.destroy_pipeline_layout(layout, None); + device.destroy_descriptor_set_layout(set_layout, None); + } + return Err(e.into()); + } + }; + + Ok(Self { + ctx: ctx.clone(), + set_layout, + layout, + pipeline, + push_size: desc.push_size, + }) + } + + /// Binds the pipeline, pushes its descriptors and constants, and + /// dispatches. + pub fn dispatch( + &self, + cmd: vk::CommandBuffer, + images: &[Bind], + push: &P, + groups: (u32, u32, u32), + ) { + assert_eq!(std::mem::size_of::

() as u32, self.push_size); + let ctx = self.ctx.inner(); + let infos: Vec<_> = images + .iter() + .map(|b| { + [vk::DescriptorImageInfo::default() + .image_view(b.view) + .image_layout(b.layout) + .sampler(b.sampler.unwrap_or_default())] + }) + .collect(); + let writes: Vec<_> = images + .iter() + .zip(&infos) + .enumerate() + .map(|(i, (b, info))| { + vk::WriteDescriptorSet::default() + .dst_binding(i as u32) + .descriptor_type(b.ty) + .image_info(info) + }) + .collect(); + unsafe { + ctx.device + .cmd_bind_pipeline(cmd, vk::PipelineBindPoint::COMPUTE, self.pipeline); + if !writes.is_empty() { + ctx.push_descriptor.cmd_push_descriptor_set( + cmd, + vk::PipelineBindPoint::COMPUTE, + self.layout, + 0, + &writes, + ); + } + if self.push_size > 0 { + let bytes = std::slice::from_raw_parts( + (push as *const P).cast::(), + std::mem::size_of::

(), + ); + ctx.device.cmd_push_constants( + cmd, + self.layout, + vk::ShaderStageFlags::COMPUTE, + 0, + bytes, + ); + } + ctx.device.cmd_dispatch(cmd, groups.0, groups.1, groups.2); + } + } +} + +impl Drop for Pipeline { + fn drop(&mut self) { + let device = self.ctx.device(); + unsafe { + device.destroy_pipeline(self.pipeline, None); + device.destroy_pipeline_layout(self.layout, None); + device.destroy_descriptor_set_layout(self.set_layout, None); + } + } +} + +/// One image binding, in binding order. nespyro's own images are always in +/// `GENERAL`; only the caller's source image comes in some other layout. +pub(crate) struct Bind { + pub ty: vk::DescriptorType, + pub view: vk::ImageView, + pub sampler: Option, + pub layout: vk::ImageLayout, +} + +impl Bind { + pub fn sampled(view: vk::ImageView, sampler: vk::Sampler) -> Self { + Self::sampled_in(view, sampler, vk::ImageLayout::GENERAL) + } + + pub fn sampled_in(view: vk::ImageView, sampler: vk::Sampler, layout: vk::ImageLayout) -> Self { + Self { + ty: vk::DescriptorType::COMBINED_IMAGE_SAMPLER, + view, + sampler: Some(sampler), + layout, + } + } + + pub fn texture(view: vk::ImageView, layout: vk::ImageLayout) -> Self { + Self { + ty: vk::DescriptorType::SAMPLED_IMAGE, + view, + sampler: None, + layout, + } + } + + pub fn storage(view: vk::ImageView) -> Self { + Self { + ty: vk::DescriptorType::STORAGE_IMAGE, + view, + sampler: None, + layout: vk::ImageLayout::GENERAL, + } + } +} + +/// Entry point names, as `-fvk-use-entrypoint-name` keeps them. +pub(crate) mod entry { + use std::ffi::CStr; + pub const RGB_TO_YCBCR: &CStr = c"rgb_to_ycbcr"; + pub const DWT: &CStr = c"dwt"; + pub const WAVELET_QUANT: &CStr = c"wavelet_quant"; + pub const ANALYZE_RATE_CONTROL: &CStr = c"analyze_rate_control"; + pub const ANALYZE_RATE_CONTROL_FINALIZE: &CStr = c"analyze_rate_control_finalize"; + pub const RESOLVE_RATE_CONTROL: &CStr = c"resolve_rate_control"; + pub const BLOCK_PACKING: &CStr = c"block_packing"; + pub const WAVELET_DEQUANT: &CStr = c"wavelet_dequant"; + pub const IDWT: &CStr = c"idwt"; +} diff --git a/crates/nespyro/src/rate.rs b/crates/nespyro/src/rate.rs new file mode 100644 index 00000000..402a4336 --- /dev/null +++ b/crates/nespyro/src/rate.rs @@ -0,0 +1,138 @@ +//! The per-band constants rate control starts from, as upstream computes them +//! in `pyrowave_encoder.cpp` at precision level 1. + +use crate::bitstream::{Chroma, DECOMPOSITION_LEVELS}; + +/// Buckets the rate-distortion slopes are sorted into, and how many +/// subdivisions of the block space each has. +pub(crate) const NUM_RDO_BUCKETS: u32 = 128; +pub(crate) const BLOCK_SPACE_SUBDIVISION: u32 = 16; +/// The bucket buffer's two counters are padded to this before the savings. +pub(crate) const RDO_BUCKET_OFFSET: u64 = 64; + +/// 32×32 blocks per subdivision, a power of two so a block's subdivision is a +/// shift of its index. +pub(crate) fn blocks_per_subdivision(blocks_32x32: u32) -> u32 { + blocks_32x32 + .next_multiple_of(BLOCK_SPACE_SUBDIVISION) + .div_euclid(BLOCK_SPACE_SUBDIVISION) + .next_power_of_two() +} + +/// A flat noise spectrum after the transform: CDF 9/7's low pass gains 6 dB, +/// one bit, per level, so each level and each low-pass direction gets a bit +/// more resolution. Chroma enters a level later and gets one less. +pub(crate) fn noise_power_normalized_quant_resolution( + level: usize, + component: usize, + band: usize, +) -> f32 { + let mut bits = 8; + if band == 0 { + bits += 2; + } else if band < 3 { + bits += 1; + } + bits += level as i32; + if component != 0 { + bits -= 1; + } + (1u32 << bits) as f32 +} + +/// FP16 has limited range, and this is a good enough first estimate. +pub(crate) fn quant_resolution(level: usize, component: usize, band: usize) -> f32 { + noise_power_normalized_quant_resolution(level, component, band).min(4096.0) +} + +/// How much a unit of distortion in this band costs, weighted by the +/// contrast sensitivity function at the band's spatial frequency. +pub(crate) fn rdo_distortion_scale( + level: usize, + component: usize, + band: usize, + chroma: Chroma, +) -> f32 { + // Upstream's model: normal PC monitors, viewed between couch and desk. + let horiz_midpoint = if band & 1 != 0 { 0.75f32 } else { 0.25 }; + let vert_midpoint = if band & 2 != 0 { 0.75f32 } else { 0.25 }; + const DPI: f32 = 96.0; + const VIEWING_DISTANCE: f32 = 1.0; + const CPD_NYQUIST: f32 = 0.34 * VIEWING_DISTANCE * DPI; + + let mut cpd = (horiz_midpoint * horiz_midpoint + vert_midpoint * vert_midpoint).sqrt() + * CPD_NYQUIST + * (-(level as f32)).exp2(); + // Never quantize the LL band hard. + cpd = cpd.max(8.0); + + let mut csf = 2.6 * (0.0192 + 0.114 * cpd) * (-(0.114 * cpd).powf(1.1)).exp(); + + // Chroma is heavily discounted, less so when not subsampled. + if component != 0 && level != DECOMPOSITION_LEVELS - 1 && chroma == Chroma::Yuv420 { + csf *= 0.6; + } + + // Distortion in lower bands becomes more noise power after filtering; + // scaling by the resolution evens that out. Distortion is power, so the + // weight is squared. + let weighted = csf * noise_power_normalized_quant_resolution(level, component, band); + weighted * weighted +} + +/// The block header's quant code for a decoder-side scale: a small float with +/// a 5-bit negated exponent and a 3-bit mantissa. +pub(crate) fn encode_quant(decoder_q_scale: f32) -> u8 { + const MAX_SCALE_EXP: i32 = 4; + let v = decoder_q_scale.to_bits(); + let e = ((v >> 23) & 0xff) as i32 - 127 - MAX_SCALE_EXP; + let m = ((v >> 20) & 0x7) as i32; + let e = -e; + assert!( + (0..=20).contains(&e), + "quant scale {decoder_q_scale} outside the code's range" + ); + ((e << 3) | m) as u8 +} + +/// The scale a quant code stands for; the shaders' `decode_quant`. +pub(crate) fn decode_quant(code: u8) -> f32 { + const MAX_SCALE_EXP: i32 = 4; + let e = MAX_SCALE_EXP - i32::from(code >> 3); + let m = i32::from(code & 0x7); + (1.0 / (8.0 * 1024.0 * 1024.0)) * ((8 + m) * (1 << (20 + e))) as f32 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quant_codes_round_trip_the_scales_rate_control_uses() { + for level in 0..DECOMPOSITION_LEVELS { + for component in 0..3 { + for band in 0..4 { + let res = quant_resolution(level, component, band); + let code = encode_quant(1.0 / res); + // Every starting resolution is a power of two, which the + // code holds exactly. + assert_eq!(1.0 / decode_quant(code), res); + } + } + } + } + + #[test] + fn resolution_is_capped_for_fp16() { + assert_eq!(quant_resolution(4, 0, 0), 4096.0); + assert_eq!(quant_resolution(0, 0, 3), 256.0); + assert_eq!(quant_resolution(1, 1, 1), 512.0); + } + + #[test] + fn subdivisions_are_powers_of_two() { + assert_eq!(blocks_per_subdivision(1), 1); + assert_eq!(blocks_per_subdivision(51), 4); + assert_eq!(blocks_per_subdivision(4000), 256); + } +} diff --git a/crates/nespyro/src/sync.rs b/crates/nespyro/src/sync.rs new file mode 100644 index 00000000..c68baf6a --- /dev/null +++ b/crates/nespyro/src/sync.rs @@ -0,0 +1,67 @@ +//! Ordering nespyro's GPU work against the caller's, without the CPU. + +use std::sync::{Arc, Mutex, MutexGuard}; + +use ash::vk; + +/// A point on a timeline semaphore: reached once the semaphore's value is at +/// least `value`. The same shape as pixelforge's, so a caller moving a frame +/// from its own rendering into nespyro and out again chains points the way it +/// does for pixelforge. +/// +/// The semaphore must be a timeline semaphore on the same device as the +/// [`Context`](crate::Context). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TimelinePoint { + pub semaphore: vk::Semaphore, + pub value: u64, +} + +impl TimelinePoint { + pub fn new(semaphore: vk::Semaphore, value: u64) -> Self { + Self { semaphore, value } + } + + /// A wait that holds back every stage of the submission. + /// + /// nespyro's submissions open with layout transitions whose stages the + /// caller cannot know, so a narrower mask could let one of them run + /// before the caller's work is done. + pub(crate) fn wait_info(&self) -> vk::SemaphoreSubmitInfo<'static> { + vk::SemaphoreSubmitInfo::default() + .semaphore(self.semaphore) + .value(self.value) + .stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS) + } +} + +/// Serialises submissions to a queue shared with other threads. +/// +/// A `VkQueue` must not be submitted to from two threads at once. When +/// nespyro has a queue to itself, or the queue was created internally +/// synchronized, pass [`QueueLock::none`]. When it shares a queue with the +/// caller's own submissions, both sides lock the same `QueueLock` around every +/// `vkQueueSubmit2` and `vkQueuePresentKHR`. +#[derive(Debug, Clone)] +pub struct QueueLock(Option>>); + +impl QueueLock { + /// A lock for a queue shared with the caller; clone it to the other side. + pub fn shared() -> Self { + Self(Some(Arc::new(Mutex::new(())))) + } + + /// No locking: the queue is nespyro's alone, or the driver synchronises it. + pub fn none() -> Self { + Self(None) + } + + /// Holds the queue until the guard drops. + pub fn lock(&self) -> Option> { + // A panic while holding the lock poisons it, but the queue itself is + // no worse for it, so the lock is still usable. + self.0 + .as_ref() + .map(|m| m.lock().unwrap_or_else(|p| p.into_inner())) + } +} diff --git a/crates/nespyro/src/wavelet.rs b/crates/nespyro/src/wavelet.rs new file mode 100644 index 00000000..e702aba2 --- /dev/null +++ b/crates/nespyro/src/wavelet.rs @@ -0,0 +1,92 @@ +//! The wavelet coefficient images both halves of the codec work in. +//! +//! Upstream's precision level 1: levels 0 and 1, the large ones, are stored +//! as FP16, and levels 2 to 4, where the energy of the image concentrates and +//! the bandwidth is trivial, as FP32. Each image is an array of 12 layers, +//! four bands for each of three components, with one mip per level. + +use ash::vk; + +use crate::bitstream::{DECOMPOSITION_LEVELS, Layout, NUM_COMPONENTS}; +use crate::device::Context; +use crate::error::Result; +use crate::gpu::{Image, ImageDesc, Sampler}; + +/// Levels stored as FP16; the rest are FP32. +const FP16_LEVELS: usize = 2; + +pub(crate) struct Wavelet { + high: Image, + low: Image, + /// Per component and level, the four bands as a 2D array. + pub bands: [[vk::ImageView; DECOMPOSITION_LEVELS]; NUM_COMPONENTS], + /// Per component and level, the LL band alone as a 2D image. + pub ll: [[vk::ImageView; DECOMPOSITION_LEVELS]; NUM_COMPONENTS], + /// Edges of every transform mirror through this. + pub mirror: Sampler, + /// The quantizer reads past a band's edge as zero through this. + pub border: Sampler, +} + +impl Wavelet { + pub fn new(ctx: &Context, layout: &Layout) -> Result { + let (w, h) = layout.level_size(0); + let usage = vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::STORAGE; + let layers = 4 * NUM_COMPONENTS as u32; + let mut high = Image::new( + ctx, + &ImageDesc { + format: vk::Format::R16_SFLOAT, + width: w, + height: h, + layers, + mips: FP16_LEVELS as u32, + usage, + families: &[], + }, + )?; + let mut low = Image::new( + ctx, + &ImageDesc { + format: vk::Format::R32_SFLOAT, + width: w >> FP16_LEVELS, + height: h >> FP16_LEVELS, + layers, + mips: (DECOMPOSITION_LEVELS - FP16_LEVELS) as u32, + usage, + families: &[], + }, + )?; + + let mut bands = [[vk::ImageView::null(); DECOMPOSITION_LEVELS]; NUM_COMPONENTS]; + let mut ll = bands; + for level in 0..DECOMPOSITION_LEVELS { + let (image, mip) = if level < FP16_LEVELS { + (&mut high, level as u32) + } else { + (&mut low, (level - FP16_LEVELS) as u32) + }; + for component in 0..NUM_COMPONENTS { + let base = 4 * component as u32; + bands[component][level] = + image.view(vk::ImageViewType::TYPE_2D_ARRAY, mip, base, 4)?; + ll[component][level] = image.view(vk::ImageViewType::TYPE_2D, mip, base, 1)?; + } + } + + Ok(Self { + high, + low, + bands, + ll, + mirror: Sampler::new(ctx, vk::SamplerAddressMode::MIRRORED_REPEAT)?, + border: Sampler::new(ctx, vk::SamplerAddressMode::CLAMP_TO_BORDER)?, + }) + } + + /// Both images, for a transition out of `UNDEFINED`: a frame never reads + /// what the previous one left. + pub fn images(&self) -> [vk::Image; 2] { + [self.high.image, self.low.image] + } +} diff --git a/crates/nespyro/tests/bench.rs b/crates/nespyro/tests/bench.rs new file mode 100644 index 00000000..255efcab --- /dev/null +++ b/crates/nespyro/tests/bench.rs @@ -0,0 +1,148 @@ +//! Encode and decode latency on the GPU, ours beside the C++ reference's. +//! +//! GPU time from timestamps, not command recording: an earlier measurement in +//! this project counted CPU time spent recording and called it the cost of a +//! pass. Reports only; the numbers depend on the GPU, and the reference's are +//! printed by the reference itself. +//! +//! NESPYRO_REFERENCE= cargo test -p nespyro --release --test bench -- --ignored --nocapture + +mod common; + +use std::process::Command; +use std::time::Instant; + +use ash::vk; +use common::{Content, Gpu, reference_ycbcr, rgba8}; +use nespyro::{Chroma, DecodeConfig, Decoder, Depth, EncodeConfig, Encoder}; + +const FRAMES: usize = 120; + +fn median(mut v: Vec) -> f64 { + v.sort_by(|a, b| a.partial_cmp(b).unwrap()); + v[v.len() / 2] +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU and the reference harness (NESPYRO_REFERENCE)"] +fn latency_beside_the_reference() { + let harness = std::env::var("NESPYRO_REFERENCE") + .expect("NESPYRO_REFERENCE must name the harness reference/build.sh built"); + let gpu = Gpu::new(); + for &(w, h, chroma) in &[ + (1920, 1080, Chroma::Yuv420), + (3840, 2160, Chroma::Yuv420), + (1920, 1080, Chroma::Yuv444), + ] { + let rgba = rgba8(Content::Edges, w, h); + let config = EncodeConfig::new(w, h).with_chroma(chroma); + let image = gpu.upload(vk::Format::R8G8B8A8_UNORM, w, h, &rgba); + let mut encoder = Encoder::new(gpu.ctx(), config.clone()).unwrap(); + let mut decoder = + Decoder::new(gpu.ctx(), DecodeConfig::new(w, h, chroma, Depth::Eight)).unwrap(); + + let mut passes: Vec<[f64; 6]> = Vec::new(); + let mut encode_ms = Vec::new(); + let mut encode_cpu_ms = Vec::new(); + let mut decode_ms = Vec::new(); + let mut decode_cpu_ms = Vec::new(); + let mut decode_gpu: Vec<[f64; 3]> = Vec::new(); + for _ in 0..FRAMES { + let start = Instant::now(); + let f = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + encode_cpu_ms.push(start.elapsed().as_secs_f64() * 1e3); + let frame = pollster::block_on(f).unwrap(); + encode_ms.push(start.elapsed().as_secs_f64() * 1e3); + let s = frame.stats; + passes.push([ + s.convert_ns, + s.dwt_ns, + s.quant_ns, + s.analyze_ns, + s.resolve_ns, + s.packing_ns, + ]); + + for p in &frame.packets { + decoder.push_packet(&frame.data[p.clone()]).unwrap(); + } + let start = Instant::now(); + let decoded = decoder.decode_after(&[]).unwrap(); + decode_cpu_ms.push(start.elapsed().as_secs_f64() * 1e3); + gpu.wait(decoded.ready); + decode_ms.push(start.elapsed().as_secs_f64() * 1e3); + let d = decoder + .stats() + .unwrap() + .expect("the decode is done, so its stats are"); + decode_gpu.push([d.upload_ns, d.dequant_ns, d.idwt_ns]); + } + let names = ["convert", "dwt", "quant", "analyze", "resolve", "packing"]; + let medians: Vec = (0..6) + .map(|i| median(passes.iter().map(|p| p[i]).collect()) / 1000.0) + .collect(); + eprintln!("ours {w}x{h} {chroma:?}, median of {FRAMES} frames, GPU µs:"); + for (n, m) in names.iter().zip(&medians) { + eprintln!(" {n:8} {m:8.1}"); + } + eprintln!( + " encode {:8.1} (without convert {:.1}); {:.1} submit to packetized, wall clock, of which {:.1} in encode_after", + medians.iter().sum::(), + medians[1..].iter().sum::(), + median(encode_ms) * 1000.0, + median(encode_cpu_ms) * 1000.0 + ); + let d: Vec = (0..3) + .map(|i| median(decode_gpu.iter().map(|p| p[i]).collect()) / 1000.0) + .collect(); + eprintln!(" upload {:8.1}", d[0]); + eprintln!(" dequant {:8.1}", d[1]); + eprintln!(" idwt {:8.1}", d[2]); + eprintln!( + " decode {:8.1} (GPU); {:.1} submit to ready, wall clock, of which {:.1} in decode_after", + d.iter().sum::(), + median(decode_ms) * 1000.0, + median(decode_cpu_ms) * 1000.0 + ); + drop((encoder, decoder)); + + // The reference encodes the same content, converted on the CPU. + let yuv = std::env::temp_dir().join(format!("nespyro-bench-{}.yuv", std::process::id())); + let planes = reference_ycbcr(&rgba, w, h, chroma == Chroma::Yuv420, false); + let bytes: Vec = planes + .iter() + .flatten() + .map(|v| (v * 255.0).round() as u8) + .collect(); + std::fs::write(&yuv, bytes).unwrap(); + let target = (config.target_bitrate / 60 / 8) & !3; + let out = Command::new(&harness) + .args([ + "bench", + &w.to_string(), + &h.to_string(), + if chroma == Chroma::Yuv420 { + "420" + } else { + "444" + }, + &target.to_string(), + &FRAMES.to_string(), + yuv.to_str().unwrap(), + ]) + .output() + .unwrap(); + assert!( + out.status.success(), + "{}", + String::from_utf8_lossy(&out.stderr) + ); + eprintln!( + "reference {w}x{h} {chroma:?}:\n{}", + String::from_utf8_lossy(&out.stdout) + ); + } + gpu.assert_clean(); +} diff --git a/crates/nespyro/tests/common/mod.rs b/crates/nespyro/tests/common/mod.rs new file mode 100644 index 00000000..7b9c6758 --- /dev/null +++ b/crates/nespyro/tests/common/mod.rs @@ -0,0 +1,686 @@ +//! A device for the GPU tests, created the way a real caller would: from +//! `DeviceRequirements`, with Khronos validation on and every validation +//! error recorded, so a test that ran "successfully" with errors fails. +//! +//! `NESPYRO_TEST_DEVICE` picks a device by name substring; otherwise the +//! first discrete GPU, then anything that is not a CPU implementation. A +//! machine without one fails the test rather than skipping it. + +#![allow(dead_code)] + +use std::ffi::c_void; +use std::sync::{Arc, Mutex}; + +use ash::vk; +use ash::vk::TaggedStructure as _; +use nespyro::{Context, DeviceQueue, DeviceRequirements, QueueLock, Roles}; + +pub struct Gpu { + pub entry: ash::Entry, + pub instance: ash::Instance, + pub physical_device: vk::PhysicalDevice, + pub device: ash::Device, + pub queue: vk::Queue, + pub family: u32, + pub lock: QueueLock, + pub ctx: Option, + pub name: String, + errors: Arc>>, + debug: Option<(ash::ext::debug_utils::Instance, vk::DebugUtilsMessengerEXT)>, + pool: vk::CommandPool, +} + +unsafe extern "system" fn on_message( + severity: vk::DebugUtilsMessageSeverityFlagsEXT, + _kind: vk::DebugUtilsMessageTypeFlagsEXT, + data: *const vk::DebugUtilsMessengerCallbackDataEXT<'_>, + user: *mut c_void, +) -> vk::Bool32 { + if severity.contains(vk::DebugUtilsMessageSeverityFlagsEXT::ERROR) { + let errors = unsafe { &*(user as *const Mutex>) }; + let message = unsafe { (*data).message_as_c_str() } + .map(|m| m.to_string_lossy().into_owned()) + .unwrap_or_default(); + eprintln!("validation: {message}"); + errors.lock().unwrap().push(message); + } + vk::FALSE +} + +impl Gpu { + pub fn new() -> Self { + let entry = unsafe { ash::Entry::load() }.expect("no Vulkan loader"); + let errors = Arc::new(Mutex::new(Vec::new())); + + let validation = c"VK_LAYER_KHRONOS_validation"; + let have_validation = unsafe { entry.enumerate_instance_layer_properties() } + .unwrap() + .iter() + .any(|l| l.layer_name_as_c_str().is_ok_and(|n| n == validation)); + assert!( + have_validation || std::env::var_os("NESPYRO_NO_VALIDATION").is_some(), + "the GPU tests run under Khronos validation; install it, or set \ + NESPYRO_NO_VALIDATION=1 to run without" + ); + let layers: Vec<*const i8> = if have_validation { + vec![validation.as_ptr()] + } else { + vec![] + }; + let extensions = [ash::ext::debug_utils::NAME.as_ptr()]; + let app = vk::ApplicationInfo::default().api_version(vk::API_VERSION_1_3); + let instance = unsafe { + entry + .create_instance( + &vk::InstanceCreateInfo::default() + .application_info(&app) + .enabled_layer_names(&layers) + .enabled_extension_names(&extensions), + None, + ) + .unwrap() + }; + + let debug = { + let loader = ash::ext::debug_utils::Instance::load(&entry, &instance); + let info = vk::DebugUtilsMessengerCreateInfoEXT::default() + .message_severity(vk::DebugUtilsMessageSeverityFlagsEXT::ERROR) + .message_type( + vk::DebugUtilsMessageTypeFlagsEXT::VALIDATION + | vk::DebugUtilsMessageTypeFlagsEXT::GENERAL, + ) + .pfn_user_callback(Some(on_message)) + .user_data(Arc::as_ptr(&errors) as *mut c_void); + let messenger = unsafe { loader.create_debug_utils_messenger(&info, None).unwrap() }; + Some((loader, messenger)) + }; + + let wanted = std::env::var("NESPYRO_TEST_DEVICE").ok(); + let devices = unsafe { instance.enumerate_physical_devices().unwrap() }; + let describe = |d: &vk::PhysicalDevice| { + let p = unsafe { instance.get_physical_device_properties(*d) }; + let name = p + .device_name_as_c_str() + .unwrap() + .to_string_lossy() + .into_owned(); + (name, p.device_type) + }; + let pick = devices + .iter() + .find(|d| { + let (name, kind) = describe(d); + match &wanted { + Some(w) => name.contains(w.as_str()), + None => kind == vk::PhysicalDeviceType::DISCRETE_GPU, + } + }) + .or_else(|| { + wanted.is_none().then(|| { + devices + .iter() + .find(|d| describe(d).1 != vk::PhysicalDeviceType::CPU) + })? + }) + .copied() + .unwrap_or_else(|| { + let names: Vec<_> = devices.iter().map(|d| describe(d).0).collect(); + panic!("no GPU to test on (wanted {wanted:?}, have {names:?})") + }); + let physical_device = pick; + let name = describe(&pick).0; + + let req = DeviceRequirements::query(&instance, physical_device, Roles::BOTH) + .unwrap_or_else(|e| panic!("{name} cannot run nespyro: {e}")); + let family = req.queue_family; + + let f = req.features; + let mut v11 = vk::PhysicalDeviceVulkan11Features::default() + .storage_buffer16_bit_access(f.storage_buffer_16bit_access); + let mut v12 = vk::PhysicalDeviceVulkan12Features::default() + .storage_buffer8_bit_access(f.storage_buffer_8bit_access) + .shader_int8(f.shader_int8) + .shader_float16(f.shader_float16) + .timeline_semaphore(f.timeline_semaphore) + .buffer_device_address(f.buffer_device_address); + let mut v13 = vk::PhysicalDeviceVulkan13Features::default() + .synchronization2(f.synchronization2) + .subgroup_size_control(f.subgroup_size_control) + .compute_full_subgroups(f.compute_full_subgroups); + let base = vk::PhysicalDeviceFeatures::default() + .shader_int16(f.shader_int16) + .shader_storage_image_write_without_format(f.shader_storage_image_write_without_format); + let mut features = vk::PhysicalDeviceFeatures2::default().features(base); + let priorities = [1.0]; + let queues = [vk::DeviceQueueCreateInfo::default() + .queue_family_index(family) + .queue_priorities(&priorities)]; + let extensions: Vec<_> = req.extensions.iter().map(|e| e.as_ptr()).collect(); + let device = unsafe { + instance + .create_device( + physical_device, + &vk::DeviceCreateInfo::default() + .queue_create_infos(&queues) + .enabled_extension_names(&extensions) + .push(&mut features) + .push(&mut v11) + .push(&mut v12) + .push(&mut v13), + None, + ) + .unwrap() + }; + let queue = unsafe { device.get_device_queue(family, 0) }; + let lock = QueueLock::shared(); + let ctx = Context::from_existing( + instance.clone(), + physical_device, + device.clone(), + DeviceQueue::new(family, 0), + lock.clone(), + Roles::BOTH, + ) + .unwrap(); + let pool = unsafe { + device + .create_command_pool( + &vk::CommandPoolCreateInfo::default().queue_family_index(family), + None, + ) + .unwrap() + }; + + eprintln!("testing on {name}"); + Self { + entry, + instance, + physical_device, + device, + queue, + family, + lock, + ctx: Some(ctx), + name, + errors, + debug, + pool, + } + } + + pub fn ctx(&self) -> Context { + self.ctx.clone().unwrap() + } + + /// Every validation error so far, taken. + pub fn take_errors(&self) -> Vec { + std::mem::take(&mut *self.errors.lock().unwrap()) + } + + pub fn assert_clean(&self) { + let errors = self.take_errors(); + assert!( + errors.is_empty(), + "{} validation errors:\n{}", + errors.len(), + errors.join("\n") + ); + } + + fn memory_type(&self, bits: u32, want: vk::MemoryPropertyFlags) -> u32 { + let props = unsafe { + self.instance + .get_physical_device_memory_properties(self.physical_device) + }; + (0..props.memory_type_count) + .find(|&i| { + bits & (1 << i) != 0 && props.memory_types[i as usize].property_flags.contains(want) + }) + .expect("no suitable memory type") + } + + /// Blocks until a timeline point is reached. + pub fn wait(&self, point: nespyro::TimelinePoint) { + let semaphores = [point.semaphore]; + let values = [point.value]; + unsafe { + self.device + .wait_semaphores( + &vk::SemaphoreWaitInfo::default() + .semaphores(&semaphores) + .values(&values), + u64::MAX, + ) + .unwrap() + }; + } + + /// Records into a one-off command buffer, submits it and waits. + pub fn run(&self, record: impl FnOnce(vk::CommandBuffer)) { + let device = &self.device; + unsafe { + let cmd = device + .allocate_command_buffers( + &vk::CommandBufferAllocateInfo::default() + .command_pool(self.pool) + .command_buffer_count(1), + ) + .unwrap()[0]; + device + .begin_command_buffer(cmd, &vk::CommandBufferBeginInfo::default()) + .unwrap(); + record(cmd); + device.end_command_buffer(cmd).unwrap(); + let fence = device + .create_fence(&vk::FenceCreateInfo::default(), None) + .unwrap(); + let cmds = [cmd]; + { + let _g = self.lock.lock(); + device + .queue_submit( + self.queue, + &[vk::SubmitInfo::default().command_buffers(&cmds)], + fence, + ) + .unwrap(); + } + device.wait_for_fences(&[fence], true, u64::MAX).unwrap(); + device.destroy_fence(fence, None); + device.free_command_buffers(self.pool, &cmds); + } + } + + pub fn buffer(&self, size: u64, usage: vk::BufferUsageFlags) -> HostBuffer { + let device = &self.device; + unsafe { + let buffer = device + .create_buffer( + &vk::BufferCreateInfo::default().size(size).usage(usage), + None, + ) + .unwrap(); + let req = device.get_buffer_memory_requirements(buffer); + let memory = device + .allocate_memory( + &vk::MemoryAllocateInfo::default() + .allocation_size(req.size) + .memory_type_index(self.memory_type( + req.memory_type_bits, + vk::MemoryPropertyFlags::HOST_VISIBLE + | vk::MemoryPropertyFlags::HOST_COHERENT, + )), + None, + ) + .unwrap(); + device.bind_buffer_memory(buffer, memory, 0).unwrap(); + let ptr = device + .map_memory(memory, 0, vk::WHOLE_SIZE, vk::MemoryMapFlags::empty()) + .unwrap() + .cast(); + HostBuffer { + device: device.clone(), + buffer, + memory, + ptr, + size: size as usize, + } + } + } + + /// An image in `GENERAL`, filled from `bytes` tightly packed. + pub fn upload(&self, format: vk::Format, width: u32, height: u32, bytes: &[u8]) -> GpuImage { + let device = &self.device; + let image = unsafe { + device + .create_image( + &vk::ImageCreateInfo::default() + .image_type(vk::ImageType::TYPE_2D) + .format(format) + .extent(vk::Extent3D { + width, + height, + depth: 1, + }) + .mip_levels(1) + .array_layers(1) + .samples(vk::SampleCountFlags::TYPE_1) + .tiling(vk::ImageTiling::OPTIMAL) + .usage( + vk::ImageUsageFlags::SAMPLED + | vk::ImageUsageFlags::TRANSFER_DST + | vk::ImageUsageFlags::TRANSFER_SRC, + ) + .initial_layout(vk::ImageLayout::UNDEFINED), + None, + ) + .unwrap() + }; + let req = unsafe { device.get_image_memory_requirements(image) }; + let memory = unsafe { + device + .allocate_memory( + &vk::MemoryAllocateInfo::default() + .allocation_size(req.size) + .memory_type_index(self.memory_type( + req.memory_type_bits, + vk::MemoryPropertyFlags::DEVICE_LOCAL, + )), + None, + ) + .unwrap() + }; + unsafe { device.bind_image_memory(image, memory, 0).unwrap() }; + + let staging = self.buffer(bytes.len() as u64, vk::BufferUsageFlags::TRANSFER_SRC); + staging.write(bytes); + self.run(|cmd| unsafe { + barrier( + device, + cmd, + image, + vk::ImageLayout::UNDEFINED, + vk::ImageLayout::TRANSFER_DST_OPTIMAL, + ); + device.cmd_copy_buffer_to_image( + cmd, + staging.buffer, + image, + vk::ImageLayout::TRANSFER_DST_OPTIMAL, + &[vk::BufferImageCopy::default() + .image_subresource(color_layers()) + .image_extent(vk::Extent3D { + width, + height, + depth: 1, + })], + ); + barrier( + device, + cmd, + image, + vk::ImageLayout::TRANSFER_DST_OPTIMAL, + vk::ImageLayout::GENERAL, + ); + }); + GpuImage { + device: device.clone(), + image, + memory, + format, + width, + height, + } + } + + /// Copies a single-layer image in `GENERAL` back to the CPU. + pub fn download(&self, image: vk::Image, width: u32, height: u32, texel_bytes: u32) -> Vec { + let size = u64::from(width * height * texel_bytes); + let staging = self.buffer(size, vk::BufferUsageFlags::TRANSFER_DST); + let device = &self.device; + self.run(|cmd| unsafe { + let b = [vk::MemoryBarrier2::default() + .src_stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS) + .src_access_mask(vk::AccessFlags2::MEMORY_WRITE) + .dst_stage_mask(vk::PipelineStageFlags2::COPY) + .dst_access_mask(vk::AccessFlags2::TRANSFER_READ)]; + device.cmd_pipeline_barrier2(cmd, &vk::DependencyInfo::default().memory_barriers(&b)); + device.cmd_copy_image_to_buffer( + cmd, + image, + vk::ImageLayout::GENERAL, + staging.buffer, + &[vk::BufferImageCopy::default() + .image_subresource(color_layers()) + .image_extent(vk::Extent3D { + width, + height, + depth: 1, + })], + ); + }); + staging.read() + } +} + +impl Drop for Gpu { + fn drop(&mut self) { + self.ctx.take(); + unsafe { + self.device.device_wait_idle().unwrap(); + self.device.destroy_command_pool(self.pool, None); + self.device.destroy_device(None); + if let Some((loader, messenger)) = self.debug.take() { + loader.destroy_debug_utils_messenger(messenger, None); + } + self.instance.destroy_instance(None); + } + } +} + +fn color_layers() -> vk::ImageSubresourceLayers { + vk::ImageSubresourceLayers { + aspect_mask: vk::ImageAspectFlags::COLOR, + mip_level: 0, + base_array_layer: 0, + layer_count: 1, + } +} + +unsafe fn barrier( + device: &ash::Device, + cmd: vk::CommandBuffer, + image: vk::Image, + old: vk::ImageLayout, + new: vk::ImageLayout, +) { + let b = [vk::ImageMemoryBarrier2::default() + .src_stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS) + .src_access_mask(vk::AccessFlags2::MEMORY_WRITE) + .dst_stage_mask(vk::PipelineStageFlags2::ALL_COMMANDS) + .dst_access_mask(vk::AccessFlags2::MEMORY_READ | vk::AccessFlags2::MEMORY_WRITE) + .old_layout(old) + .new_layout(new) + .image(image) + .subresource_range(vk::ImageSubresourceRange { + aspect_mask: vk::ImageAspectFlags::COLOR, + base_mip_level: 0, + level_count: 1, + base_array_layer: 0, + layer_count: 1, + })]; + unsafe { + device.cmd_pipeline_barrier2( + cmd, + &vk::DependencyInfo::default().image_memory_barriers(&b), + ) + }; +} + +pub struct HostBuffer { + device: ash::Device, + pub buffer: vk::Buffer, + memory: vk::DeviceMemory, + ptr: *mut u8, + size: usize, +} + +impl HostBuffer { + pub fn write(&self, bytes: &[u8]) { + assert!(bytes.len() <= self.size); + unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), self.ptr, bytes.len()) }; + } + + pub fn read(&self) -> Vec { + unsafe { std::slice::from_raw_parts(self.ptr, self.size).to_vec() } + } +} + +impl Drop for HostBuffer { + fn drop(&mut self) { + unsafe { + self.device.destroy_buffer(self.buffer, None); + self.device.free_memory(self.memory, None); + } + } +} + +pub struct GpuImage { + device: ash::Device, + pub image: vk::Image, + memory: vk::DeviceMemory, + pub format: vk::Format, + pub width: u32, + pub height: u32, +} + +impl Drop for GpuImage { + fn drop(&mut self) { + unsafe { + self.device.destroy_image(self.image, None); + self.device.free_memory(self.memory, None); + } + } +} + +/// xorshift, so every run sees the same content. +pub struct Rng(pub u64); +impl Rng { + pub fn next(&mut self) -> u32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 >> 16) as u32 + } +} + +/// Synthetic content, RGBA8, chosen to exercise the codec differently. +#[derive(Debug, Clone, Copy)] +pub enum Content { + /// Smooth gradients: almost everything in the coarse bands. + Gradient, + /// Uniform noise: the worst case for a transform codec. + Noise, + /// Hard black and white edges, like text and UI. + Edges, + /// One flat colour: nearly nothing but DC. + Flat, +} + +pub fn rgba8(content: Content, width: u32, height: u32) -> Vec { + let mut rng = Rng(0x1234_5678_9abc_def1); + let mut out = Vec::with_capacity((width * height * 4) as usize); + for y in 0..height { + for x in 0..width { + let px = match content { + Content::Gradient => [ + (x * 255 / width.max(2).saturating_sub(1)) as u8, + (y * 255 / height.max(2).saturating_sub(1)) as u8, + ((x + y) * 255 / (width + height)) as u8, + ], + Content::Noise => [rng.next() as u8, rng.next() as u8, rng.next() as u8], + Content::Edges => { + let on = ((x / 3) + (y / 7)) % 5 == 0 || (x % 16 < 2) || (y % 11 == 0); + if on { [250, 250, 250] } else { [12, 12, 12] } + } + Content::Flat => [40, 120, 200], + }; + out.extend_from_slice(&[px[0], px[1], px[2], 255]); + } + } + out +} + +/// Full-range YCbCr from 8-bit sRGB, as the encoder's conversion computes it, +/// written independently of it. 4:2:0 chroma is the mean of each 2×2 quad. +pub fn reference_ycbcr(rgba: &[u8], w: u32, h: u32, yuv420: bool, bt2020: bool) -> [Vec; 3] { + let (yr, cbr, crr) = if bt2020 { + ( + [0.2627, 0.6780, 0.0593], + [-0.1396, -0.3604, 0.5], + [0.5, -0.4598, -0.0402], + ) + } else { + ( + [0.2126, 0.7152, 0.0722], + [-0.1146, -0.3854, 0.5], + [0.5, -0.4542, -0.0458], + ) + }; + let dot = |a: [f32; 3], b: [f32; 3]| a[0] * b[0] + a[1] * b[1] + a[2] * b[2]; + let at = |x: u32, y: u32| { + let i = ((y * w + x) * 4) as usize; + let c = [rgba[i], rgba[i + 1], rgba[i + 2]].map(|v| f32::from(v) / 255.0); + [ + dot(yr, c).clamp(0.0, 1.0), + (0.5 + dot(cbr, c)).clamp(0.0, 1.0), + (0.5 + dot(crr, c)).clamp(0.0, 1.0), + ] + }; + let mut y = Vec::with_capacity((w * h) as usize); + for j in 0..h { + for i in 0..w { + y.push(at(i, j)[0]); + } + } + let (cw, ch) = if yuv420 { (w / 2, h / 2) } else { (w, h) }; + let mut cb = Vec::with_capacity((cw * ch) as usize); + let mut cr = Vec::with_capacity((cw * ch) as usize); + for j in 0..ch { + for i in 0..cw { + let v = if yuv420 { + let q = [ + at(2 * i, 2 * j), + at(2 * i + 1, 2 * j), + at(2 * i, 2 * j + 1), + at(2 * i + 1, 2 * j + 1), + ]; + [1, 2].map(|k| q.iter().map(|p| p[k]).sum::() / 4.0) + } else { + let p = at(i, j); + [p[1], p[2]] + }; + cb.push(v[0]); + cr.push(v[1]); + } + } + [y, cb, cr] +} + +/// A decoded plane as floats in [0, 1], cropped to the picture. +pub fn plane_values(gpu: &Gpu, plane: &nespyro::PlaneView) -> Vec { + let texel = if plane.format == vk::Format::R16_UNORM { + 2 + } else { + 1 + }; + let raw = gpu.download(plane.image, plane.image_width, plane.image_height, texel); + let mut out = Vec::with_capacity((plane.width * plane.height) as usize); + for y in 0..plane.height { + for x in 0..plane.width { + let i = ((y * plane.image_width + x) * texel) as usize; + out.push(if texel == 2 { + f32::from(u16::from_le_bytes([raw[i], raw[i + 1]])) / 65535.0 + } else { + f32::from(raw[i]) / 255.0 + }); + } + } + out +} + +/// PSNR in dB against a peak of 1.0. +pub fn psnr(a: &[f32], b: &[f32]) -> f64 { + assert_eq!(a.len(), b.len()); + let mse = a + .iter() + .zip(b) + .map(|(x, y)| (f64::from(*x) - f64::from(*y)).powi(2)) + .sum::() + / a.len() as f64; + if mse == 0.0 { + f64::INFINITY + } else { + -10.0 * mse.log10() + } +} diff --git a/crates/nespyro/tests/encode.rs b/crates/nespyro/tests/encode.rs new file mode 100644 index 00000000..c8249b3f --- /dev/null +++ b/crates/nespyro/tests/encode.rs @@ -0,0 +1,66 @@ +//! The encoder on a real GPU, checked without the decoder: every block it +//! writes must pass the depacketizer's structural check, which recomputes +//! each block's size from its own contents. + +mod common; + +use ash::vk; +use common::{Content, Gpu, rgba8}; +use nespyro::bitstream::{Depacketizer, Layout, Push}; +use nespyro::{Chroma, EncodeConfig, Encoder}; + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn encoder_output_is_a_valid_stream() { + let gpu = Gpu::new(); + for &(w, h, chroma) in &[ + (1920, 1080, Chroma::Yuv420), + (1366, 768, Chroma::Yuv444), + (64, 48, Chroma::Yuv420), + ] { + for content in [ + Content::Gradient, + Content::Noise, + Content::Edges, + Content::Flat, + ] { + let image = gpu.upload(vk::Format::R8G8B8A8_UNORM, w, h, &rgba8(content, w, h)); + let config = EncodeConfig::new(w, h).with_chroma(chroma); + let mut encoder = Encoder::new(gpu.ctx(), config.clone()).unwrap(); + let future = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + let frame = pollster::block_on(future).unwrap(); + + let mut d = Depacketizer::new(Layout::new(w, h, chroma).unwrap()); + for p in &frame.packets { + assert_eq!( + d.push(&frame.data[p.clone()]).unwrap(), + Push::Accepted, + "{w}x{h} {chroma:?} {content:?}" + ); + } + let r = d.readiness(); + assert!(r.is_complete(), "{w}x{h} {chroma:?} {content:?}: {r:?}"); + assert_eq!(frame.stats.dropped_blocks, 0); + eprintln!( + "scratch {} of {}", + frame.stats.scratch_bytes, frame.stats.scratch_capacity + ); + assert!( + !frame.stats.scratch_overflowed(), + "{w}x{h} {chroma:?} {content:?}: {:?}", + frame.stats + ); + eprintln!( + "{w}x{h} {chroma:?} {content:?}: {} bytes of {} target, {} packets, {:.1} µs GPU", + frame.data.len(), + frame.stats.target_bytes, + frame.packets.len(), + frame.stats.gpu_ns() / 1000.0 + ); + drop(encoder); + gpu.assert_clean(); + } + } +} diff --git a/crates/nespyro/tests/reference.rs b/crates/nespyro/tests/reference.rs new file mode 100644 index 00000000..3211bdc5 --- /dev/null +++ b/crates/nespyro/tests/reference.rs @@ -0,0 +1,332 @@ +//! nespyro against the C++ reference it is a port of. +//! +//! Needs the harness `reference/build.sh` builds, named by +//! `NESPYRO_REFERENCE`. Unset or missing, the test fails: a cross-check that +//! quietly does not run is indistinguishable from one that passed. +//! +//! Decoding is floating point, so two decoders of one stream may differ by +//! rounding. The tolerance is upstream's own, from +//! `pyrowave_device_validation.cpp`: one step of luma, two of chroma. + +mod common; + +use std::path::PathBuf; +use std::process::Command; + +use ash::vk; +use common::{Content, Gpu, plane_values, psnr, reference_ycbcr, rgba8}; +use nespyro::bitstream::Push; +use nespyro::{Chroma, DecodeConfig, Decoder, Depth, EncodeConfig, EncodedFrame, Encoder}; + +fn harness() -> PathBuf { + let path = std::env::var_os("NESPYRO_REFERENCE").unwrap_or_else(|| { + panic!( + "NESPYRO_REFERENCE is not set. Build the reference with \ + crates/nespyro/reference/build.sh

and point it at the harness it prints." + ) + }); + let path = PathBuf::from(path); + assert!( + path.is_file(), + "NESPYRO_REFERENCE names {path:?}, which does not exist" + ); + path +} + +fn scratch(name: &str) -> PathBuf { + let dir = std::env::temp_dir().join(format!("nespyro-reference-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + dir.join(name) +} + +fn chroma_arg(c: Chroma) -> &'static str { + match c { + Chroma::Yuv420 => "420", + Chroma::Yuv444 => "444", + } +} + +fn run(args: &[&str]) { + let out = Command::new(harness()).args(args).output().unwrap(); + assert!( + out.status.success(), + "harness {args:?} failed: {}", + String::from_utf8_lossy(&out.stderr) + ); +} + +fn write_packets(path: &PathBuf, packets: &[&[u8]]) { + let mut out = Vec::new(); + for p in packets { + out.extend((p.len() as u32).to_le_bytes()); + out.extend(*p); + } + std::fs::write(path, out).unwrap(); +} + +fn read_packets(path: &PathBuf) -> Vec> { + let data = std::fs::read(path).unwrap(); + let mut at = 0; + let mut packets = Vec::new(); + while at < data.len() { + let len = u32::from_le_bytes(data[at..at + 4].try_into().unwrap()) as usize; + packets.push(data[at + 4..at + 4 + len].to_vec()); + at += 4 + len; + } + packets +} + +/// Planes as 8-bit samples, Y then Cb then Cr. +fn to_u8(planes: &[Vec; 3]) -> Vec { + planes + .iter() + .flatten() + .map(|v| (v * 255.0).round() as u8) + .collect() +} + +fn our_decode(gpu: &Gpu, w: u32, h: u32, chroma: Chroma, packets: &[&[u8]]) -> Vec { + try_our_decode(gpu, w, h, chroma, packets).unwrap() +} + +fn try_our_decode( + gpu: &Gpu, + w: u32, + h: u32, + chroma: Chroma, + packets: &[&[u8]], +) -> Result, nespyro::Error> { + let mut decoder = + Decoder::new(gpu.ctx(), DecodeConfig::new(w, h, chroma, Depth::Eight)).unwrap(); + for p in packets { + assert_eq!(decoder.push_packet(p)?, Push::Accepted); + } + assert!(decoder.readiness().is_complete()); + let frame = decoder.decode_after(&[]).unwrap(); + gpu.wait(frame.ready); + Ok(to_u8(&[ + plane_values(gpu, &frame.y), + plane_values(gpu, &frame.cb), + plane_values(gpu, &frame.cr), + ])) +} + +fn their_decode(w: u32, h: u32, chroma: Chroma, packets: &[&[u8]]) -> Vec { + let pkts = scratch("in.pkts"); + let yuv = scratch("out.yuv"); + write_packets(&pkts, packets); + run(&[ + "decode", + &w.to_string(), + &h.to_string(), + chroma_arg(chroma), + pkts.to_str().unwrap(), + yuv.to_str().unwrap(), + ]); + std::fs::read(yuv).unwrap() +} + +/// The largest difference per plane between two decodes of one stream. +fn max_difference(a: &[u8], b: &[u8], w: u32, h: u32, chroma: Chroma) -> [u8; 3] { + assert_eq!(a.len(), b.len()); + let luma = (w * h) as usize; + let c = match chroma { + Chroma::Yuv420 => luma / 4, + Chroma::Yuv444 => luma, + }; + let ranges = [0..luma, luma..luma + c, luma + c..luma + 2 * c]; + ranges.map(|r| { + a[r.clone()] + .iter() + .zip(&b[r]) + .map(|(x, y)| x.abs_diff(*y)) + .max() + .unwrap() + }) +} + +fn within_tolerance(diff: [u8; 3]) -> bool { + diff[0] <= 1 && diff[1] <= 2 && diff[2] <= 2 +} + +const CASES: &[(u32, u32, Chroma)] = &[(1920, 1080, Chroma::Yuv420), (1366, 768, Chroma::Yuv444)]; +const CONTENT: [Content; 3] = [Content::Gradient, Content::Edges, Content::Noise]; + +fn our_encode(gpu: &Gpu, w: u32, h: u32, chroma: Chroma, rgba: &[u8]) -> EncodedFrame { + let image = gpu.upload(vk::Format::R8G8B8A8_UNORM, w, h, rgba); + let mut encoder = Encoder::new(gpu.ctx(), EncodeConfig::new(w, h).with_chroma(chroma)).unwrap(); + let f = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + pollster::block_on(f).unwrap() +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU and the reference harness (NESPYRO_REFERENCE)"] +fn their_decoder_reads_our_stream() { + let gpu = Gpu::new(); + for &(w, h, chroma) in CASES { + for content in CONTENT { + let frame = our_encode(&gpu, w, h, chroma, &rgba8(content, w, h)); + let packets: Vec<&[u8]> = frame + .packets + .iter() + .map(|p| &frame.data[p.clone()]) + .collect(); + // Their decoder refuses a packet it cannot parse and a frame it + // considers incomplete, so getting planes back at all is half of + // the check. + let theirs = their_decode(w, h, chroma, &packets); + let ours = our_decode(&gpu, w, h, chroma, &packets); + let diff = max_difference(&ours, &theirs, w, h, chroma); + eprintln!("ours → theirs {w}x{h} {chroma:?} {content:?}: max difference {diff:?}"); + assert!( + within_tolerance(diff), + "{w}x{h} {chroma:?} {content:?}: {diff:?}" + ); + } + } + gpu.assert_clean(); +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU and the reference harness (NESPYRO_REFERENCE)"] +fn our_decoder_reads_their_stream_and_our_encoder_keeps_up() { + let gpu = Gpu::new(); + for &(w, h, chroma) in CASES { + for content in CONTENT { + let rgba = rgba8(content, w, h); + // The same conversion our encoder makes, done on the CPU and + // rounded to 8 bits, is their input. + let reference = reference_ycbcr(&rgba, w, h, chroma == Chroma::Yuv420, false); + let yuv = scratch("in.yuv"); + let pkts = scratch("out.pkts"); + std::fs::write(&yuv, to_u8(&reference)).unwrap(); + + // Both encoders get exactly these planes. + let input = to_u8(&reference); + let luma = (w * h) as usize; + let (cw, ch) = match chroma { + Chroma::Yuv420 => (w / 2, h / 2), + Chroma::Yuv444 => (w, h), + }; + let c = (cw * ch) as usize; + let y_img = gpu.upload(vk::Format::R8_UNORM, w, h, &input[..luma]); + let cb_img = gpu.upload(vk::Format::R8_UNORM, cw, ch, &input[luma..luma + c]); + let cr_img = gpu.upload(vk::Format::R8_UNORM, cw, ch, &input[luma + c..]); + let mut encoder = + Encoder::new(gpu.ctx(), EncodeConfig::new(w, h).with_chroma(chroma)).unwrap(); + let f = encoder + .encode_planes_after( + [y_img.image, cb_img.image, cr_img.image], + vk::ImageLayout::GENERAL, + &[], + ) + .unwrap(); + let ours_frame = pollster::block_on(f).unwrap(); + drop(encoder); + let target = ours_frame.stats.target_bytes; + run(&[ + "encode", + &w.to_string(), + &h.to_string(), + chroma_arg(chroma), + &target.to_string(), + "1200", + yuv.to_str().unwrap(), + pkts.to_str().unwrap(), + ]); + let their_packets = read_packets(&pkts); + let their_refs: Vec<&[u8]> = their_packets.iter().map(|p| p.as_slice()).collect(); + let their_bytes: usize = their_packets.iter().map(Vec::len).sum(); + + // Upstream sizes its quantizer scratch at aligned width × height + // × 2 bytes, and high-entropy content needs more. Ours asks for + // nearly the same amount on the same planes, so it says when + // theirs overflowed; its encoder then packs blocks from bytes it + // never wrote, and a block's contents stop matching its header. + let layout = nespyro::bitstream::Layout::new(w, h, chroma).unwrap(); + let upstream_scratch = layout.aligned_width * layout.aligned_height * 2; + let upstream_overflowed = ours_frame.stats.scratch_bytes > upstream_scratch; + let by_them = their_decode(w, h, chroma, &their_refs); + let by_us = match try_our_decode(&gpu, w, h, chroma, &their_refs) { + Ok(planes) => planes, + Err(nespyro::Error::Bitstream(e)) if upstream_overflowed => { + // Refused, not decoded: the stream is malformed. + eprintln!( + "theirs → ours {w}x{h} {chroma:?} {content:?}: upstream's scratch overflowed \ + ({} of {upstream_scratch} bytes) and our decoder refused its stream: {e}", + ours_frame.stats.scratch_bytes + ); + continue; + } + Err(e) => { + panic!("{w}x{h} {chroma:?} {content:?}: our decoder refused their stream: {e}") + } + }; + let diff = max_difference(&by_us, &by_them, w, h, chroma); + + // Our encoder beside theirs, each judged against the 8-bit input. + let ours_packets: Vec<&[u8]> = ours_frame + .packets + .iter() + .map(|p| &ours_frame.data[p.clone()]) + .collect(); + let our_decoded = our_decode(&gpu, w, h, chroma, &ours_packets); + let as_f32 = |v: &[u8]| v.iter().map(|&x| f32::from(x) / 255.0).collect::>(); + let our_db = psnr(&as_f32(&our_decoded[..luma]), &as_f32(&input[..luma])); + let their_db = psnr(&as_f32(&by_them[..luma]), &as_f32(&input[..luma])); + let same_blocks = identical_blocks(&ours_frame.data, &their_packets); + + eprintln!( + "theirs → ours {w}x{h} {chroma:?} {content:?}: max difference {diff:?}; \ + encoders: ours {} bytes {our_db:.2} dB, theirs {their_bytes} bytes {their_db:.2} dB, \ + {:.1}% of blocks identical", + ours_frame.data.len(), + same_blocks * 100.0 + ); + assert!( + within_tolerance(diff), + "{w}x{h} {chroma:?} {content:?}: {diff:?}" + ); + // Slang and glslang may round floating point differently, so the + // encoders need not agree byte for byte; ours must not be + // measurably worse. + assert!( + our_db >= their_db - 0.5, + "{w}x{h} {chroma:?} {content:?}: ours {our_db:.2} dB, theirs {their_db:.2} dB" + ); + } + } + gpu.assert_clean(); +} + +/// The fraction of their blocks that also appear, byte for byte, in ours. +/// Sequence numbers are masked, since the two encoders count separately. +fn identical_blocks(ours: &[u8], theirs: &[Vec]) -> f64 { + let blocks = |bytes: &[u8]| { + let mut out = std::collections::HashMap::new(); + let mut at = 0; + while at + 8 <= bytes.len() { + let w0 = u32::from_le_bytes(bytes[at..at + 4].try_into().unwrap()); + if w0 >> 31 != 0 { + at += 8; + continue; + } + let size = ((w0 >> 16) & 0xfff) as usize * 4; + let w1 = u32::from_le_bytes(bytes[at + 4..at + 8].try_into().unwrap()); + let mut b = bytes[at..at + size].to_vec(); + b[3] &= 0x8f; + out.insert(w1 >> 8, b); + at += size; + } + out + }; + let ours = blocks(ours); + let theirs: std::collections::HashMap<_, _> = theirs.iter().flat_map(|p| blocks(p)).collect(); + let same = theirs + .iter() + .filter(|(k, v)| ours.get(*k) == Some(*v)) + .count(); + same as f64 / theirs.len().max(1) as f64 +} diff --git a/crates/nespyro/tests/roundtrip.rs b/crates/nespyro/tests/roundtrip.rs new file mode 100644 index 00000000..1f75b40c --- /dev/null +++ b/crates/nespyro/tests/roundtrip.rs @@ -0,0 +1,502 @@ +//! Encode, packetize, depacketize and decode on a real GPU, and compare what +//! comes out with what went in. + +mod common; + +use ash::vk; +use common::{Content, Gpu, plane_values, psnr, reference_ycbcr, rgba8}; +use nespyro::bitstream::Push; +use nespyro::{Chroma, DecodeConfig, Decoder, Depth, EncodeConfig, EncodedFrame, Encoder}; + +fn encode(gpu: &Gpu, config: EncodeConfig, rgba: &[u8]) -> EncodedFrame { + let image = gpu.upload( + vk::Format::R8G8B8A8_UNORM, + config.width, + config.height, + rgba, + ); + let mut encoder = Encoder::new(gpu.ctx(), config).unwrap(); + let future = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + pollster::block_on(future).unwrap() +} + +/// Decodes `frame`, delivering only the packets `keep` allows, and returns +/// the three planes as floats with the frame's missing-block count. +fn decode( + gpu: &Gpu, + config: &EncodeConfig, + depth: Depth, + frame: &EncodedFrame, + keep: impl Fn(usize) -> bool, +) -> ([Vec; 3], u32) { + let mut decoder = Decoder::new( + gpu.ctx(), + DecodeConfig::new(config.width, config.height, config.chroma, depth), + ) + .unwrap(); + for (i, p) in frame.packets.iter().enumerate() { + if keep(i) { + assert_eq!( + decoder.push_packet(&frame.data[p.clone()]).unwrap(), + Push::Accepted + ); + } + } + let decoded = decoder.decode_after(&[]).unwrap(); + gpu.wait(decoded.ready); + let planes = [ + plane_values(gpu, &decoded.y), + plane_values(gpu, &decoded.cb), + plane_values(gpu, &decoded.cr), + ]; + (planes, decoded.missing_blocks) +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn round_trip_quality() { + let gpu = Gpu::new(); + for &(w, h, chroma) in &[ + (1920, 1080, Chroma::Yuv420), + (1366, 768, Chroma::Yuv444), + (64, 48, Chroma::Yuv420), + ] { + for depth in [Depth::Eight, Depth::Sixteen] { + for content in [ + Content::Gradient, + Content::Noise, + Content::Edges, + Content::Flat, + ] { + let rgba = rgba8(content, w, h); + let config = EncodeConfig::new(w, h) + .with_chroma(chroma) + .with_depth(depth); + let frame = encode(&gpu, config.clone(), &rgba); + let (planes, missing) = decode(&gpu, &config, depth, &frame, |_| true); + assert_eq!(missing, 0); + let reference = reference_ycbcr(&rgba, w, h, chroma == Chroma::Yuv420, false); + let db: Vec = (0..3).map(|i| psnr(&planes[i], &reference[i])).collect(); + eprintln!( + "{w}x{h} {chroma:?} {depth:?} {content:?}: {} bytes, Y {:.1} Cb {:.1} Cr {:.1} dB", + frame.data.len(), + db[0], + db[1], + db[2] + ); + // A frame that fit its budget comes back as well as its depth + // can hold. For 8-bit output that floor is half a code off + // everywhere, 20·log10(255 / 0.5) = 54.2 dB: grey chroma sits + // at exactly 0.5, between two codes. Frames the budget + // constrained are rate_control's business. + let target = frame.stats.target_bytes as usize; + if frame.data.len() * 100 < target * 95 { + let floor = match depth { + Depth::Eight => 54.0, + Depth::Sixteen => 60.0, + }; + assert!( + db.iter().all(|&d| d >= floor), + "{w}x{h} {chroma:?} {depth:?} {content:?}: {db:?} below {floor}" + ); + } + gpu.assert_clean(); + } + } + } +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn rate_control_limits_and_spends() { + let gpu = Gpu::new(); + let (w, h) = (1920, 1080); + for content in [Content::Noise, Content::Edges] { + let rgba = rgba8(content, w, h); + let reference = reference_ycbcr(&rgba, w, h, true, false); + let mut curve: Vec = Vec::new(); + for mbps in [25, 50, 100, 200, 400, 800, 1600, 3200] { + let config = EncodeConfig::new(w, h).with_target_bitrate(mbps * 1_000_000); + let frame = encode(&gpu, config.clone(), &rgba); + let target = frame.stats.target_bytes as usize; + let (planes, _) = decode(&gpu, &config, Depth::Eight, &frame, |_| true); + let db = psnr(&planes[0], &reference[0]); + eprintln!( + "{content:?} {mbps} Mbit/s: {} bytes, Y {db:.1} dB", + frame.data.len() + ); + + // Never over the target. + assert!( + frame.data.len() <= target, + "{content:?} {mbps}: {} > {target}", + frame.data.len() + ); + // And spent, until the picture is as good as 8 bits can carry: a + // frame that came out nearly empty and decoded to grey fails + // here, not only one that overshot. + let saturated = db >= 50.0; + assert!( + saturated || frame.data.len() * 100 >= target * 95, + "{content:?} {mbps}: {} bytes of {target} at {db:.1} dB", + frame.data.len() + ); + // Each doubling buys quality until it saturates. + if let Some(&last) = curve.last() { + assert!( + db > last + 0.2 || last >= 50.0, + "{content:?}: {db:.1} dB at {mbps} after {last:.1}" + ); + } + curve.push(db); + } + assert!( + curve.last().unwrap() - curve[0] > 10.0, + "{content:?}: 128 times the rate bought {:.1} dB", + curve.last().unwrap() - curve[0] + ); + gpu.assert_clean(); + } +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn the_same_input_encodes_to_the_same_bytes() { + let gpu = Gpu::new(); + for &(w, h, chroma) in &[(1920, 1080, Chroma::Yuv420), (1366, 768, Chroma::Yuv444)] { + let rgba = rgba8(Content::Edges, w, h); + let image = gpu.upload(vk::Format::R8G8B8A8_UNORM, w, h, &rgba); + let mut encoder = + Encoder::new(gpu.ctx(), EncodeConfig::new(w, h).with_chroma(chroma)).unwrap(); + let mut frames = Vec::new(); + for _ in 0..3 { + let f = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + frames.push(pollster::block_on(f).unwrap()); + } + // Consecutive frames differ only in the 3-bit sequence number, which + // sits in the top of each header's first word; masked out, the + // bytes must match exactly, however the GPU ordered its atomics. + let masked = |f: &EncodedFrame| { + let mut d = f.data.clone(); + let mut at = 0; + while at < d.len() { + let w0 = u32::from_le_bytes(d[at..at + 4].try_into().unwrap()); + let size = if w0 >> 31 != 0 { + 8 + } else { + ((w0 >> 16) & 0xfff) as usize * 4 + }; + d[at + 3] &= 0x8f; + at += size; + } + d + }; + assert_ne!(frames[0].sequence, frames[1].sequence); + assert_eq!(masked(&frames[0]), masked(&frames[1]), "{w}x{h} {chroma:?}"); + assert_eq!(masked(&frames[1]), masked(&frames[2]), "{w}x{h} {chroma:?}"); + drop(encoder); + gpu.assert_clean(); + } +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn losing_packets_blurs_but_decodes() { + let gpu = Gpu::new(); + let (w, h) = (1920, 1080); + let rgba = rgba8(Content::Gradient, w, h); + let reference = reference_ycbcr(&rgba, w, h, true, false); + let config = EncodeConfig::new(w, h); + let frame = encode(&gpu, config.clone(), &rgba); + let critical = frame.critical_packets; + assert!(critical >= 1 && critical < frame.packets.len()); + + // Drop one in twenty of the packets after the critical ones. + let dropped = |i: usize| i >= critical && i % 20 == 3; + let lost: usize = (0..frame.packets.len()).filter(|&i| dropped(i)).count(); + assert!(lost > 0); + let (whole, missing_whole) = decode(&gpu, &config, Depth::Eight, &frame, |_| true); + let (partial, missing) = decode(&gpu, &config, Depth::Eight, &frame, |i| !dropped(i)); + assert_eq!(missing_whole, 0); + + // Exactly the blocks in the dropped packets are reported missing. + let blocks_in = |i: usize| { + let bytes = &frame.data[frame.packets[i].clone()]; + let mut at = 0; + let mut n = 0; + while at < bytes.len() { + let w0 = u32::from_le_bytes(bytes[at..at + 4].try_into().unwrap()); + if w0 >> 31 != 0 { + at += 8; + } else { + at += ((w0 >> 16) & 0xfff) as usize * 4; + n += 1; + } + } + n + }; + let expected: u32 = (0..frame.packets.len()) + .filter(|&i| dropped(i)) + .map(blocks_in) + .sum(); + assert_eq!(missing, expected); + + let whole_db = psnr(&whole[0], &reference[0]); + let partial_db = psnr(&partial[0], &reference[0]); + eprintln!( + "{lost} packets, {missing} blocks lost: Y {whole_db:.1} dB whole, {partial_db:.1} dB partial" + ); + // Worse, but a picture: the coarse bands all arrived. + assert!(partial_db < whole_db); + assert!(partial_db > 25.0, "{partial_db:.1} dB"); + gpu.assert_clean(); +} + +/// Linear light from an sRGB signal. +fn srgb_to_linear(c: f32) -> f32 { + if c <= 0.04045 { + c / 12.92 + } else { + ((c + 0.055) / 1.055).powf(2.4) + } +} + +fn linear_to_pq(l: f32) -> f32 { + // ST 2084's constants, as the standard defines them. + let (m1, m2) = (2610.0 / 16384.0, 2523.0 / 4096.0 * 128.0); + let (c1, c2, c3) = ( + 3424.0 / 4096.0, + 2413.0 / 4096.0 * 32.0, + 2392.0 / 4096.0 * 32.0, + ); + let lm1 = l.max(0.0).powf(m1); + ((c1 + c2 * lm1) / (1.0 + c3 * lm1)).powf(m2) +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn colour_reaches_the_planes_and_the_header() { + use nespyro::{ColourDescription, Source}; + let gpu = Gpu::new(); + let (w, h) = (256, 128); + let rgba = rgba8(Content::Gradient, w, h); + + // SDR: BT.709 in the header, BT.709 in the planes. + let config = EncodeConfig::new(w, h).with_chroma(Chroma::Yuv444); + let frame = encode(&gpu, config.clone(), &rgba); + let mut decoder = Decoder::new( + gpu.ctx(), + DecodeConfig::new(w, h, Chroma::Yuv444, Depth::Eight), + ) + .unwrap(); + for p in &frame.packets { + decoder.push_packet(&frame.data[p.clone()]).unwrap(); + } + let decoded = decoder.decode_after(&[]).unwrap(); + assert_eq!(decoded.colour, ColourDescription::bt709()); + gpu.wait(decoded.ready); + let reference = reference_ycbcr(&rgba, w, h, false, false); + let y = plane_values(&gpu, &decoded.y); + let cb = plane_values(&gpu, &decoded.cb); + assert!(psnr(&y, &reference[0]) > 50.0 && psnr(&cb, &reference[1]) > 50.0); + drop(decoded); + + // HDR from an sRGB source: into BT.2020 linear light at the reference + // white, PQ encoded, through the BT.2020 matrix; the header says so. + let nits = 203.0; + let config = EncodeConfig::new(w, h) + .with_chroma(Chroma::Yuv444) + .with_depth(Depth::Sixteen) + .with_colour(ColourDescription::bt2020_pq(), Source::Srgb) + .with_reference_white(nits); + let frame = encode(&gpu, config, &rgba); + let mut decoder = Decoder::new( + gpu.ctx(), + DecodeConfig::new(w, h, Chroma::Yuv444, Depth::Sixteen), + ) + .unwrap(); + for p in &frame.packets { + decoder.push_packet(&frame.data[p.clone()]).unwrap(); + } + let decoded = decoder.decode_after(&[]).unwrap(); + assert_eq!(decoded.colour, ColourDescription::bt2020_pq()); + gpu.wait(decoded.ready); + let pq: Vec = rgba + .as_chunks::<4>() + .0 + .iter() + .flat_map(|p| { + let l = [p[0], p[1], p[2]].map(|v| srgb_to_linear(f32::from(v) / 255.0)); + let m = [ + 0.6274 * l[0] + 0.3293 * l[1] + 0.0433 * l[2], + 0.0691 * l[0] + 0.9195 * l[1] + 0.0114 * l[2], + 0.0164 * l[0] + 0.0880 * l[1] + 0.8956 * l[2], + ]; + let e = m.map(|c| (linear_to_pq(c * nits / 10000.0) * 255.0).round() as u8); + [e[0], e[1], e[2], 255] + }) + .collect(); + // The reference is rounded through 8 bits on the way, so the comparison + // is looser than the codec itself. + let reference = reference_ycbcr(&pq, w, h, false, true); + let y = plane_values(&gpu, &decoded.y); + let cr = plane_values(&gpu, &decoded.cr); + let (dy, dcr) = (psnr(&y, &reference[0]), psnr(&cr, &reference[2])); + eprintln!("PQ: Y {dy:.1} dB, Cr {dcr:.1} dB against an 8-bit reference"); + assert!(dy > 45.0 && dcr > 45.0, "{dy:.1} {dcr:.1}"); + // And it really is PQ, not SDR relabelled: SDR white at 203 nits sits + // near 0.58 in PQ, far from 1.0. + let brightest = y.iter().cloned().fold(0.0f32, f32::max); + assert!( + (0.5..0.65).contains(&brightest), + "brightest PQ luma {brightest}" + ); + drop(decoded); + gpu.assert_clean(); +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn configurations_the_encoder_cannot_honour_are_refused() { + use nespyro::{ColourDescription, Error, Source}; + let gpu = Gpu::new(); + let sdr_from_linear = + EncodeConfig::new(128, 128).with_colour(ColourDescription::bt709(), Source::Bt709Linear); + assert!(matches!( + Encoder::new(gpu.ctx(), sdr_from_linear), + Err(Error::Config(_)) + )); + let limited = ColourDescription { + range: nespyro::Range::Limited, + ..ColourDescription::bt709() + }; + assert!(matches!( + Encoder::new( + gpu.ctx(), + EncodeConfig::new(128, 128).with_colour(limited, Source::Srgb) + ), + Err(Error::Config(_)) + )); + assert!(Encoder::new(gpu.ctx(), EncodeConfig::new(1921, 1080)).is_err()); + + // An sRGB-format view would be linearised by the sampler. + let image = gpu.upload( + vk::Format::R8G8B8A8_SRGB, + 128, + 128, + &rgba8(Content::Flat, 128, 128), + ); + let mut encoder = Encoder::new(gpu.ctx(), EncodeConfig::new(128, 128)).unwrap(); + assert!(matches!( + encoder.encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]), + Err(Error::Config(_)) + )); + drop(encoder); + gpu.assert_clean(); +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn a_live_bitrate_change_takes_effect_on_the_next_frame() { + let gpu = Gpu::new(); + let (w, h) = (1920, 1080); + let image = gpu.upload( + vk::Format::R8G8B8A8_UNORM, + w, + h, + &rgba8(Content::Noise, w, h), + ); + let mut encoder = Encoder::new(gpu.ctx(), EncodeConfig::new(w, h)).unwrap(); + // The encoder records its passes once and reuses them; what changes + // between frames reaches the GPU through the frame's own parameters. + for mbps in [200, 50, 400, 50, 200] { + encoder.set_target_bitrate(mbps * 1_000_000); + let f = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + let frame = pollster::block_on(f).unwrap(); + let target = frame.stats.target_bytes as usize; + assert_eq!(target as u64, (mbps * 1_000_000 / 60 / 8) & !3); + assert!( + frame.data.len() <= target && frame.data.len() * 100 >= target * 95, + "{mbps} Mbit/s: {} bytes of {target}", + frame.data.len() + ); + } + drop(encoder); + gpu.assert_clean(); +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn one_decoder_across_many_frames() { + use nespyro::Error; + let gpu = Gpu::new(); + let (w, h) = (1920, 1080); + let config = EncodeConfig::new(w, h); + // Small frames first, then large ones, so the upload buffers have to + // grow and the recorded passes that read them have to follow. + let contents = [ + Content::Flat, + Content::Gradient, + Content::Edges, + Content::Noise, + Content::Gradient, + Content::Flat, + ]; + // One encoder, so the frames carry consecutive sequence numbers as a real + // stream does. + let mut encoder = Encoder::new(gpu.ctx(), config.clone()).unwrap(); + let frames: Vec<_> = contents + .iter() + .map(|&c| { + let rgba = rgba8(c, w, h); + let image = gpu.upload(vk::Format::R8G8B8A8_UNORM, w, h, &rgba); + let f = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + let frame = pollster::block_on(f).unwrap(); + (c, frame, reference_ycbcr(&rgba, w, h, true, false)) + }) + .collect(); + drop(encoder); + + let mut decoder = Decoder::new( + gpu.ctx(), + DecodeConfig::new(w, h, Chroma::Yuv420, Depth::Eight), + ) + .unwrap(); + let mut held = Vec::new(); + for (i, (content, frame, reference)) in frames.iter().enumerate() { + for p in &frame.packets { + assert_eq!( + decoder.push_packet(&frame.data[p.clone()]).unwrap(), + Push::Accepted + ); + } + // Three outputs: a consumer holding all three gets Busy, not a stall + // and not one of its own frames overwritten. + if held.len() == 3 { + assert!(matches!(decoder.decode_after(&[]), Err(Error::Busy))); + held.remove(0); + } + let decoded = decoder.decode_after(&[]).unwrap(); + gpu.wait(decoded.ready); + let y = plane_values(&gpu, &decoded.y); + let db = psnr(&y, &reference[0]); + let floor = if frame.data.len() * 100 < frame.stats.target_bytes as usize * 95 { + 54.0 + } else { + 14.0 + }; + assert!(db >= floor, "frame {i} ({content:?}): {db:.1} dB"); + held.push(decoded); + } + drop(held); + drop(decoder); + gpu.assert_clean(); +} diff --git a/crates/nespyro/tests/transport.rs b/crates/nespyro/tests/transport.rs new file mode 100644 index 00000000..753bf48c --- /dev/null +++ b/crates/nespyro/tests/transport.rs @@ -0,0 +1,156 @@ +//! Real frames through the PyroWave wire, end to end: the IPC messages +//! nescapture writes, the hub's reassembly and datagram packing, the client's +//! collector, and a depacketizer at the end. What arrives must be what the +//! encoder made. +//! +//! The network itself is not here; neshub's own end-to-end tests carry +//! synthetic packets over a real iroh connection. This is about the bytes +//! being real PyroWave, which those cannot be. + +mod common; + +use ash::vk; +use common::{Content, Gpu, rgba8}; +use nesprotocol::pyrowave::{ + CollectedFrame, FLAG_DUPLICATE, FrameAssembler, FrameMeta, PyroCollector, chunk_frame, + decode_pyro_datagram, pack_datagrams, +}; +use nespyro::bitstream::{Depacketizer, Layout, Push, Readiness}; +use nespyro::{Chroma, EncodeConfig, EncodedFrame, Encoder}; + +/// The datagram size of a typical path, and the packet size the hub would +/// then hand the encoder. +const DATAGRAM: usize = 1200; +const PACKET: usize = DATAGRAM - nesprotocol::pyrowave::PYRO_DGRAM_HDR_LEN; + +fn encode(gpu: &Gpu, config: EncodeConfig, content: Content) -> EncodedFrame { + let rgba = rgba8(content, config.width, config.height); + let image = gpu.upload( + vk::Format::R8G8B8A8_UNORM, + config.width, + config.height, + &rgba, + ); + let mut encoder = Encoder::new(gpu.ctx(), config).unwrap(); + let future = encoder + .encode_after(image.image, image.format, vk::ImageLayout::GENERAL, &[]) + .unwrap(); + pollster::block_on(future).unwrap() +} + +fn readiness<'a>(config: &EncodeConfig, packets: impl Iterator) -> Readiness { + let mut d = Depacketizer::new(Layout::new(config.width, config.height, config.chroma).unwrap()); + for p in packets { + assert_eq!(d.push(p).unwrap(), Push::Accepted, "a packet was refused"); + } + d.readiness() +} + +/// The whole wire, with `lose` deciding, by arrival order, which datagrams +/// the network drops. +fn across_the_wire(frame: &EncodedFrame, lose: impl Fn(usize) -> bool) -> CollectedFrame> { + let meta = FrameMeta { + ts_ms: 16, + width: 1920, + height: 1080, + }; + let messages = chunk_frame( + frame.index as u32, + meta, + &frame.data, + &frame.packets, + frame.critical_packets, + ) + .unwrap(); + + let mut hub = FrameAssembler::new(); + let mut whole = None; + for m in &messages { + let ipc = nesprotocol::decode_ipc_frame(m).unwrap(); + whole = hub.push(&ipc, 0).unwrap().or(whole); + } + let whole = whole.expect("the hub never completed the frame"); + + let datagrams = pack_datagrams(0, &whole, DATAGRAM).unwrap(); + let mut client = PyroCollector::>::new(); + let mut out = Vec::new(); + let mut at = 0; + for (n, &len) in datagrams.lens.iter().enumerate() { + let d = &datagrams.buf[at..at + len]; + at += len; + assert!( + len <= DATAGRAM, + "a {len}-byte datagram on a {DATAGRAM}-byte path" + ); + let (h, payload) = decode_pyro_datagram(d).unwrap(); + if lose(n) && h.flags & FLAG_DUPLICATE == 0 { + continue; + } + out.extend(client.push(h, payload.to_vec(), 0)); + } + out.extend(client.poll(u64::MAX)); + assert_eq!(out.len(), 1); + out.pop().unwrap() +} + +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn a_frame_crosses_the_wire_as_the_encoder_made_it() { + let gpu = Gpu::new(); + for chroma in [Chroma::Yuv420, Chroma::Yuv444] { + for content in [Content::Gradient, Content::Noise, Content::Edges] { + let config = EncodeConfig::new(1920, 1080) + .with_chroma(chroma) + .with_target_bitrate(400_000_000) + .with_packet_size(PACKET); + let frame = encode(&gpu, config.clone(), content); + let direct = readiness( + &config, + frame.packets.iter().map(|p| &frame.data[p.clone()]), + ); + assert!(direct.is_complete(), "{chroma:?} {content:?}: {direct:?}"); + + let got = across_the_wire(&frame, |_| false); + assert!(got.is_whole()); + assert!(got.critical_complete); + assert_eq!(got.packets.len(), frame.packets.len()); + let wired = readiness(&config, got.packets.iter().map(|p| p.as_ref())); + assert_eq!(wired, direct, "{chroma:?} {content:?}"); + } + } + gpu.assert_clean(); +} + +/// With datagrams lost, the frame that arrives is short exactly the blocks +/// that were lost and nothing is refused; the critical bands survive on their +/// copies, so the decoder can still show it. +/// +/// Whether they survived is the collector's answer, not the depacketizer's: +/// noise leaves coarse high-pass blocks empty, so counting blocks reads a +/// frame with nothing lost as broken. This test is how that was found. +#[test] +#[ignore = "needs a Vulkan 1.3 GPU that can run nespyro"] +fn a_lossy_wire_costs_blocks_and_never_the_coarse_bands() { + let gpu = Gpu::new(); + let config = EncodeConfig::new(1920, 1080) + .with_target_bitrate(400_000_000) + .with_packet_size(PACKET); + let frame = encode(&gpu, config.clone(), Content::Noise); + let direct = readiness( + &config, + frame.packets.iter().map(|p| &frame.data[p.clone()]), + ); + + // Every critical original lost, and one in fifty of the rest. + let critical = frame.critical_packets; + let got = across_the_wire(&frame, |n| n < critical || n % 50 == 49); + assert!(!got.is_whole()); + assert!(got.critical_complete, "the copies did not stand in"); + let wired = readiness(&config, got.packets.iter().map(|p| p.as_ref())); + assert!(wired.received < direct.received); + assert!( + wired.is_complete_enough_with(got.critical_complete), + "a 2% loss should leave the frame decodable: {wired:?}" + ); + gpu.assert_clean(); +} diff --git a/docs/superpowers/specs/2026-09-27-nespyro-codec-crate-design.md b/docs/superpowers/specs/2026-09-27-nespyro-codec-crate-design.md new file mode 100644 index 00000000..998cd466 --- /dev/null +++ b/docs/superpowers/specs/2026-09-27-nespyro-codec-crate-design.md @@ -0,0 +1,448 @@ +# nespyro: a PyroWave codec crate + +Status: approved design, 2026-09-27. Part 1 of 4. + +## Why + +Nestri streams with hardware H.264, H.265 and AV1 through Vulkan Video. PyroWave +(Hans-Kristian Arntzen, MIT) is an intra-only wavelet codec that runs entirely in +compute shaders. It spends far more bandwidth (about 170 Mbit/s at 1080p60 4:2:0) +and far less time (upstream claims under 0.1 ms encode and decode at 1080p) than +any hardware encoder, which makes it the LAN codec: a direct path, a wired +network, and latency as the only thing that matters. + +## The four parts + +Each part gets its own spec and is built complete; nothing half-wired lands. + +1. **nespyro** (this spec): the codec as a Rust crate, both directions, proven + on one device and against the C++ reference. +2. **Protocol and transport**: a codec id, a self-contained-frame flag so intra + frames do not ride the reliable keyframe path, datagrams cut on PyroWave + packet boundaries, partial frames delivered to the decoder, and a bitrate + ceiling that admits hundreds of Mbit/s. +3. **nescapture**: `PerFrameEncoder` behind an abstraction over pixelforge or + nespyro, and nespyro's device requirements merged into the game's device. +4. **nesrecon / nesvideo**: a nespyro decoder beside the others, a shared device + and compute queue no longer gated on H.264 Vulkan Video decode, a + three-plane convert, and the convert shader's format bug fixed while it is + open (see part 4 notes below). + +Parts 3 and 4 depend on 1. Part 2 is independent of both, but nothing crosses +the network without it. + +## Decisions + +- **Separate crate, not pixelforge.** pixelforge is shaped by Vulkan Video at + every layer: a closed `Codec` enum matched throughout, a `VideoCodec` trait of + sessions and DPBs, and device selection that refuses a GPU without a video + queue. nescore also pins upstream pixelforge rather than the fork, and a + wavelet codec would never go upstream there. +- **Lives at `nestri/crates/nespyro`**, MIT like PyroWave itself. It is open; + there is nothing to gate. +- **Bitstream is exactly upstream's**, frozen at PyroWave `89f7e47` + (2026-09-25). Codec bugs stay reportable upstream, and the C++ reference stays + a usable oracle. nesprotocol (part 2) carries a PyroWave format version, since + the bitstream itself has none and is marked draft. +- **Shaders are Slang**, ported from upstream's GLSL, compiled by `build.rs` + with `slangc` at build time only. Nothing compiles at runtime. +- **Desktop GPUs only** (path A). See "Future work: mobile path". +- **Vulkan 1.3 minimum.** Its core features are used freely. +- **The API mirrors pixelforge** so nescapture plugs it in the same way it + plugs in the pixelforge encoder. + +## Scope + +- Chroma 4:2:0 and 4:4:4. +- 8-bit (R8) and 16-bit (R16) planes: SDR, and HDR10 PQ. +- Precision mode 1 only, which is upstream's default: FP16 wavelet storage for + levels 0 and 1, FP32 for levels 2 to 4. Modes 0 and 2 are tuning knobs that do + not affect the bitstream. +- The decoder carries the non-FP16 variants of the inverse transform and + chooses at runtime. The encoder requires `shaderFloat16`, because upstream's + quantizer and rate-control shaders use it unconditionally. +- Colour: full range always, BT.709 or BT.2020 PQ. Limited range and BT.601 + are not produced. + +## Crate layout + +``` +crates/nespyro/ + build.rs slangc over an explicit entry-point list + shaders/ Slang + src/ + lib.rs + bitstream/ pure Rust, no Vulkan + device.rs requirements, Context + pipeline.rs layouts, spec constants, subgroup sizes, cache + encoder/ EncodeConfig, Encoder, EncodeFuture, EncodedFrame + decoder/ Decoder, DecodedFrame + sync.rs TimelinePoint, QueueLock + tests/ +``` + +### build.rs and shaders + +`build.rs` runs `slangc` (or `$SLANGC`) for each entry point, with +`-target spirv -profile spirv_1_5 -emit-spirv-directly -fvk-use-entrypoint-name`. +It does not pass `-fvk-invert-y`, because everything here is compute. Output +goes to `OUT_DIR`, is embedded with `include_bytes!`, and a missing or failing +`slangc` fails the build with the command it ran. + +The guest image builder (`build/Containerfile`, builder stage) gains +`shader-slang` in its `pacman -S` list. The runtime stages do not. + +Shaders, ported one to one from upstream: + +| Slang | Upstream | +|---|---| +| `common.slang` | `dwt_common.h`, `dwt_swizzle.h`, `dwt_quant_scale.h`, `constants.h` | +| `rgb_to_ycbcr.slang` | none, ours | +| `dwt.slang` | `dwt.comp` | +| `wavelet_quant.slang` | `wavelet_quant.comp` | +| `analyze_rate_control.slang` | `analyze_rate_control.comp`, `analyze_rate_control_finalize.comp` | +| `resolve_rate_control.slang` | `resolve_rate_control.comp` | +| `block_packing.slang` | `block_packing.comp` | +| `wavelet_dequant.slang` | `wavelet_dequant.comp`, SSBO mode only | +| `idwt.slang` | `idwt.comp`, FP16 and FP32 variants | + +Buffers are reached through buffer device addresses carried in push constants. +Images and samplers are bound with push descriptors (`VK_KHR_push_descriptor`), +so no descriptor pool exists. Every pipeline declares the subgroup size +upstream requires, and `computeFullSubgroups` where upstream asks for full +groups. + +### device + +`DeviceRequirements` lists what nespyro needs, in a form a caller merges into a +device it is about to create, as nescapture does for the game's device: + +- Vulkan 1.3 core: `synchronization2`, `subgroupSizeControl`, + `computeFullSubgroups`, `maintenance4`. +- Vulkan 1.2 core: `timelineSemaphore`, `bufferDeviceAddress`, + `storageBuffer8BitAccess`, `shaderFloat16` (encoder only). +- Vulkan 1.1 core: `storageBuffer16BitAccess`. +- Vulkan 1.0: `shaderInt16`, `shaderStorageImageWriteWithoutFormat`. +- `VK_KHR_push_descriptor`. +- Subgroup operations checked against the device's properties: the encoder + needs basic, vote, arithmetic, ballot, shuffle, shuffle relative and + clustered. The decoder needs the same without clustered. Each pipeline's + subgroup size range, taken from upstream (for example 16..64 for block + packing, and exactly 64, 16 or 32 for the resolve pass), must intersect the + device's `minSubgroupSize..maxSubgroupSize`, with compute in + `requiredSubgroupSizeStages`. + +`Context::from_existing(instance, physical_device, device, queue, family, +QueueLock)` wraps a device the caller owns. `Context::supports_encode()` and +`supports_decode()` answer from the device's real features and properties, +never from its name. There is no context that creates its own device; both +callers already have one. + +### sync + +`TimelinePoint { semaphore, value }` has the same shape as pixelforge's. Waits +use `ALL_COMMANDS`, for the reason pixelforge gives: nespyro's submissions begin +with transitions whose stages the caller cannot know. + +`QueueLock` serialises nespyro's submits against anything else on the same +queue. A caller whose queue is not shared passes an uncontended lock. + +## Encode + +```rust +let mut enc = Encoder::new(ctx, EncodeConfig::new(w, h) + .with_chroma(Chroma::Yuv444) + .with_depth(Depth::Sixteen) + .with_colour(ColourDescription::bt2020_pq()) + .with_frame_rate(120, 1) + .with_target_bitrate(400_000_000) + .with_packet_size(1200))?; +let fut = enc.encode_after(src_image, src_format, src_layout, &[blit_done])?; +let frame: EncodedFrame = fut.await?; +``` + +**Input.** The caller's RGB image, sampled in place: `B8G8R8A8`, `R8G8B8A8`, +`A2B10G10R10` or `R16G16B16A16_SFLOAT`. No copy is made first. + +**Per slot** (two in flight): the Y, Cb and Cr plane images, the block +metadata and bitstream buffers, and a persistently mapped host-cached readback +buffer. Upstream recreates its buffers on every call; nespyro allocates them +once. + +**Shared across slots**: the wavelet images, the block statistics, the quant +and bucket buffers, and the pipelines. The GPU runs the encode work in order, +so these need no copies. + +**One submission per frame:** + +1. Wait on the caller's points. +2. `rgb_to_ycbcr` writes the three planes. The matrix and the full range are + applied here, and so is the 4:2:0 subsampling, with centre siting. +3. `dwt`, 13 dispatches for 4:2:0 and 15 for 4:4:4. +4. `wavelet_quant` then `analyze_rate_control`, per band (42 bands for 4:2:0, + 48 for 4:4:4). +5. `analyze_rate_control_finalize`, then `resolve_rate_control`. +6. `block_packing` per band. Finished 32×32 block packets land in GPU order. +7. Copy the metadata table and the first `target_bytes + headroom` bytes of the + bitstream into the readback buffer. Rate control caps the frame at its + target, so the copy is sized to the target, not to the whole buffer. +8. Signal the encoder's timeline. + +**Completion thread.** It waits for that value, then packetizes on the CPU: +the 8-byte sequence header first, carrying the configured colour fields and +the frame counter, then every non-empty block in block-index order. A new +packet starts when the next block would pass `packet_size`, and blocks are +never split. It resolves the future with: + +```rust +struct EncodedFrame { + data: Vec, + packets: Vec>, // each independently parseable + critical_packets: usize, // leading packets covering the lowest bands + pts: u64, + stats: EncodeStats, // GPU timestamps per pass +} +``` + +**Rate control.** The target bits per frame is bitrate ÷ fps. +`set_target_bitrate` takes effect on the next frame with no rebuild. There is +no IDR and no GOP, because every frame stands alone. `set_colour_description` +changes the matrix and the header bits. Changing the size, chroma or depth +means building a new `Encoder`, as in pixelforge. + +**Errors.** A block's header size must equal its size in the metadata table, +or packetizing fails with an error, as upstream checks it. A block cannot +exceed the 12-bit `payload_words` limit: the largest possible 32×32 block is +about 620 words (8-byte header, 16 × 3 bytes of codes, 16 × 8 subblocks × 18 +planes, 128 bytes of signs), well under 4095. A test pins that bound. A lost +device or failed submit is an error from the future. + +## Decode + +```rust +let mut dec = Decoder::new(ctx, DecodeConfig::new(w, h, chroma, depth) + .with_consumer_queue_family(graphics_family))?; +match dec.push_packet(datagram)? { Push::Accepted | Push::Stale => {}, + Push::Reconfigure(info) => { /* rebuild the decoder */ } } +if dec.readiness().is_complete_enough() { + let frame = dec.decode_after(&[])?; +} +``` + +**Packets.** `push_packet(&[u8]) -> Result` parses one packet on the CPU, +copying only block payloads. It handles the 3-bit frame counter as upstream +does: an older frame's packet is `Stale` and dropped, a newer frame's clears +the state. A header that disagrees with the decoder's size, chroma or depth +returns `Reconfigure`, and nothing is rebuilt behind the caller's back. A +truncated packet, or one whose block sizes run past its end, is `Err(Malformed)`, +never skipped. + +**Readiness.** `readiness()` reports blocks received against the header's +total, and whether the two lowest bands are complete. When to decode is the +caller's decision (part 2's policy). `is_complete_enough()` is upstream's rule: +every block, or at least 90% with the two lowest bands complete. + +**One submission per frame:** + +1. Write block payloads and the offset table into the slot's mapped buffer. On + memory that is both device-local and host-visible (ReBAR, or an integrated + GPU) the shader reads it there. Otherwise a staging copy moves it to + device-local memory first. +2. `wavelet_dequant` per band. A missing block has offset `u32::MAX` and + decodes as zeros. +3. `idwt` per level per component, FP16 or FP32 as chosen at creation. The last + level adds the DC offset back and writes straight into the output planes. +4. Signal the decoder's timeline. + +**Output.** + +```rust +struct DecodedFrame { + y: PlaneView, cb: PlaneView, cr: PlaneView, + width: u32, height: u32, chroma: Chroma, depth: Depth, + colour: ColourDescription, // from the sequence header + missing_blocks: u32, + ready: TimelinePoint, +} +``` + +The planes come from a ring of three slots, R8 or R16, in `GENERAL`, with +sampled and storage usage. When the consumer's queue family differs, they are +created `CONCURRENT` across the decode and consumer families, as pixelforge +does for `with_consumer_queue_family`, so no ownership transfer is needed. The +consumer waits on `ready` in its own submit; nothing waits on the CPU. +Dropping the frame returns its slot. If all three slots are held, +`decode_after` returns `Busy` instead of stalling, which makes a consumer +holding frames too long visible. + +A frame decoded with blocks missing succeeds and says so in `missing_blocks`, +so the loss is visible to the caller rather than looking like a whole frame. + +## Testing + +GPU and reference tests carry `#[ignore = "needs …"]` naming exactly what they +need, and they fail, never skip, when that is absent. + +**Bitstream (pure Rust, always run):** +- Header pack and unpack against `bitstream.md`'s field layouts. +- Block-index ordering and the per-level/band tables against values taken from + `metal/pyrowave_bitstream.cpp`. +- Packetize then parse reproduces every block exactly. No packet passes + `packet_size` unless one block alone is larger, and no block is split. +- Frame counter wrap: stale packets drop, a new frame clears state. +- A truncated packet or an impossible `payload_words` returns `Malformed`, and + the parser never reads past the end. +- The largest possible block fits in 12 bits of words. + +**GPU (one device):** +- Round trip over synthetic sources (gradients, noise, hard text-like edges, + flat colour) at 1920×1080, 1366×768 (not a multiple of 32) and one under the + 128 minimum. Each runs in 4:2:0 and 4:4:4, and in R8 and R16. Decoded PSNR + must clear a threshold for each bitrate. +- Rate control really limits: encoded size is at or under the target at every + tested rate, and above a floor, so a nearly empty frame decoding to grey + cannot pass. A low target must give measurably lower PSNR than a high one. +- Determinism: the same input encodes to byte-identical output twice, despite + the GPU's unordered block writes. +- Colour: known RGB values map to the expected full-range YCbCr for BT.709 and + BT.2020 PQ. The header's colour bits match the configuration, and the decoder + reads the same back. +- Loss: dropping non-critical packets still decodes, `missing_blocks` equals + what was dropped, and the lowest bands are intact. Dropping a critical packet + shows in `readiness()`. + +**Reference cross-check.** A small C++ harness is built from PyroWave +`89f7e47` and Granite `1b2d1801` (`checkout_granite.sh`), outside cargo. Its +path is given in `NESPYRO_REFERENCE`, and the test fails if that is set but the +binary is missing. +- Our encoder into the reference decoder, and the reference encoder into ours. + Pass is upstream's own tolerance from `pyrowave_device_validation.cpp`: at + most 1 LSB luma error and at most 2 chroma. +- Our encoder beside theirs on identical YCbCr input. Slang and glslang may + round floating point differently, so byte identity is not required. The test + reports the fraction of blocks that match exactly, and the size and PSNR + difference, as a diagnostic. + +**Latency.** A benchmark takes GPU timestamps around every pass, ours beside +the reference, at 1080p and 4K. It measures GPU time, not command recording. + +## Future work: mobile path (path B) + +Left out on purpose; the device check refuses such devices with a clear error +rather than degrading. Adding them means: + +- The fragment inverse transform: `idwt.vert` and `idwt.frag` with its three + `CHROMA_CONFIG` variants, separate vertical and horizontal passes into + R16F/R32F (RG for chroma) intermediates, and three scissored draws per pass + with the `EDGE_CONDITION` specialization constant. +- Its output views need `COLOR_ATTACHMENT` usage, and the pass uses dynamic + rendering. +- `wavelet_dequant` storage modes 1 (R8/R16/R32 uniform texel buffers, for + devices without 8-bit storage) and 2 (linear-tiled images aliased on one + buffer). +- Upstream's vendor preference table (`pyrowave_decoder.cpp:913-930`) choosing + fragment over compute for Mali, Qualcomm and PanVK. +- Hardware to test on, which this project does not yet have. + +## Part 4 notes, carried forward + +Found while reading nesrecon, to fix when part 4 opens the convert shader: + +- The Vulkan Video path creates its RGBA target as `R8G8B8A8_UNORM` + (`renderer.rs:616,624`), but `yuv_convert.slang:3` writes it as `rgba16f`, + and `probe_signal` reads it as `PASS_FORMAT` (`renderer.rs:1493-1499`). The + DMA-BUF path correctly uses `PASS_FORMAT`. +- One `vk::Queue` per family is shared between the render thread and + pixelforge's decode thread with no lock (`vulkan/mod.rs:555-578`). nespyro's + `QueueLock` is where that gets fixed. + +## Found while building + +Recorded after part 1 was built, for parts 2 to 4 to start from. + +### Where the crate differs from this design + +- `Context::from_existing` takes the `Roles` the device was created for. + Vulkan cannot say which features a device was created with, and the decoder + needs to know whether `shaderFloat16` is on before choosing its FP16 tile. + Decode-only requirements ask for FP16 only where the device has it. +- `Encoder::encode_planes_after` encodes the caller's own YCbCr planes. It + exists so the encoder can be compared with the reference on identical input, + and it lets a caller with its own conversion skip ours. +- Decoder output planes are padded to the codec's alignment. The final inverse + transform writes whole 32×32 tiles; upstream relies on out-of-bounds stores + being discarded, which core Vulkan only promises with `robustImageAccess`. + `PlaneView` gives both the picture size and the image size. +- `Decoder::stats` reports the GPU time of the last finished decode. +- Both halves record their many dispatches once and reuse them. On Mesa's ANV + recording costs 6 to 14 µs per dispatch; per frame it was 416 µs of a + decode. Per-frame values reach the shaders through a small parameters buffer. +- Device requirements also check each plane and wavelet format for sampled + and format-less storage use, which core Vulkan guarantees only for + `R32_SFLOAT`. + +### Measured, Arc A310, Mesa ANV + +| | nespyro | reference | +|---|---|---| +| encode 1080p 4:2:0, GPU, excluding colour conversion | 0.94 ms | 1.05 ms | +| encode 4K 4:2:0 | 3.07 ms | 3.22 ms | +| encode 1080p 4:4:4 | 1.36 ms | 1.50 ms | +| decode 1080p 4:2:0, GPU | 0.56 ms | | +| `encode_after` / `decode_after` CPU, 1080p 4:2:0 | 140 / 134 µs | | + +Both decoders give bit-identical planes on every stream either encoder wrote. +On identical input the encoders tie on smooth content and nespyro is level or +ahead on edges and noise. + +Subgroup sizes follow Granite's rule, which upstream's shaders were tuned +under: varying when a pass's range covers the device's, otherwise the smallest +in range. Taking the largest made rate-control analysis 73% slower on Intel. + +### Upstream bugs, worth reporting to PyroWave + +1. **Scratch payload undersized.** The quantizer's scratch is sized at aligned + width × height × 2 bytes. High-entropy 4:4:4 content needs more (3.4 MB + against 2.1 MB at 1366×768), and the encoder then packs blocks from bytes it + never wrote: a block whose header claims 115 words holds 95. nespyro sizes + it by the real bound, 136 bytes per 8×8 block. +2. **Rate-control prefix sum stops short.** `analyze_rate_control_finalize` + steps while `step < 512 / 2`, so the upper half of its lanes miss part of + the running total. Frames land under target, not over. Ported unchanged. +3. **Out-of-bounds shared read.** `analyze_rate_control` reads + `shared_rate_cost[gl_SubgroupInvocationID]` past its 16 entries on wider + subgroups. The value is unused; nespyro clamps the index. +4. **Stale bytes in the stream.** A block's unused sign bits come from + uninitialised shared memory and its padding from whatever an earlier frame + left. Harmless to decoders, but non-deterministic and leaks old memory; + nespyro writes zeros. + +### Toolchain and driver bugs + +- **Slang 2026.17** emits `OpGroupNonUniformShuffleXor` without declaring + `GroupNonUniformShuffle` when nothing else in the shader does. `spirv-val` + catches it; `shuffle_xor` in `common.slang` works around it. `build.rs` + requires `spirv-val` for that reason. +- **Mesa ANV** spends 12.6 ms of CPU recording `vkCmdFillBuffer` of an + 8,396,864-byte buffer, and 3 µs for a 16,785,472-byte one. nespyro records + its fills once, so it no longer matters here, but it is worth a Mesa report. + +### Found by part 3 + +- **`Readiness::critical_complete` cannot be trusted.** A block the encoder + quantized to nothing is never sent, and its absence is indistinguishable from + a loss. Noise leaves the coarsest high-pass bands empty often enough that a + frame with every packet delivered read as critically incomplete, so no + partial frame would ever have been decoded. Upstream has the same flaw and + answers it with an out-of-band mask of active blocks. Here the transport + knows which datagrams carried the critical packets and says so + (`CollectedFrame::critical_complete`); `Readiness::is_complete_enough_with` + takes that answer. Found by the real-frame transport test. + +### For the parts that follow + +- Part 2: a frame at 200 Mbit/s and 60 fps is about 416 KB in 1200-byte + packets, a few hundred per frame. `EncodedFrame::critical_packets` is small, + typically one to a few packets. +- Part 4: nesrecon's convert shader takes two planes; nespyro gives three, in + `GENERAL`, padded, with the picture size in `PlaneView`. diff --git a/docs/superpowers/specs/2026-10-01-pyrowave-client-decode-design.md b/docs/superpowers/specs/2026-10-01-pyrowave-client-decode-design.md new file mode 100644 index 00000000..6c270b8b --- /dev/null +++ b/docs/superpowers/specs/2026-10-01-pyrowave-client-decode-design.md @@ -0,0 +1,122 @@ +# PyroWave decode in nesrecon / nesvideo + +Status: built, 2026-10-01 (nescore `e20a4d2`, `0e729ef`). Part 4 of 4; parts 1-3 are +`2026-09-27-nespyro-codec-crate-design.md`, +`2026-10-01-pyrowave-transport-design.md`, +`2026-10-01-pyrowave-nescapture-design.md`. Code lives in nescore +(`crates/nesrecon`, `apps/nesvideo`); nespyro is a git dep on nestri at the +same rev nesprotocol is pinned to. + +## Decisions + +- **nespyro decodes on nesrecon's device**, from the decode thread, as one + more backend beside Vulkan Video / VA-API / CPU. Its planes stay on the GPU + and are converted by the render thread like a Vulkan Video picture. +- **The device is no longer gated on H.264 Vulkan Video decode.** + `shared_decode_device` exists when either pixelforge or nespyro fits. +- **Planes are always `Depth::Sixteen`.** PyroWave decodes in float and the + stream carries no depth; R16 costs little and never bands an HDR stream. +- **Converter: a second entry point**, three planes, chroma shift in the + uniform's spare word. The two-plane path is untouched, so its output and + tests do not move. +- **Fix while open:** the Vulkan-path RGBA target is `R8G8B8A8_UNORM` while the + shader writes `rgba16f` and the probe reads `PASS_FORMAT`; it becomes + `PASS_FORMAT`. That changes the Vulkan Video path's output (it was undefined + behaviour), so the PSNR baseline is re-taken for that path. +- **Deblock is off for PyroWave**: it has no block grid. `block_grid_for` + returns none, as for anything without one. +- **Busy is not an error.** nespyro has three output slots and the renderer + can hold all three (hand-off slot, current picture, retire queue); a frame + that finds none free is dropped and counted, and the next decodes. + +## Device (`vulkan/mod.rs`) + +- `nespyro::DeviceRequirements::query(.., Roles::DECODE)` beside pixelforge's, + independently. Device preference also accepts a nespyro-capable device. +- Features: the 1.0 ones in `PhysicalDeviceFeatures`, 16-bit storage in the + existing `Vulkan11Features`, subgroup size control and full subgroups in the + existing `Vulkan13Features`, and a new `Vulkan12Features` for 8-bit + storage, int8, float16 (when available), buffer device address and timeline + semaphores — replacing the standalone timeline struct, which may not be + chained beside it. +- Queue: nespyro gets the compute-only family's queue 0, which nesrecon + never submits to; pixelforge may use it too, but only from the same decode + thread, so the two are serialised by that thread. Without a compute-only + family, a second graphics queue if there is one. Without either -- every + Intel GPU under ANV, whose one family has one queue -- the render queue + itself, under a `QueueLock` the render thread holds around every submit, + present and `vkDeviceWaitIdle` (`VkCtx::queue_lock`, `VkCtx::wait_idle`). + The first build refused PyroWave in that case, which on this machine meant + the client never advertised it and the hub refused every request. +- `SharedVulkanDevice` carries `pyro: Option<(family, index)>`. +- `nesrecon::decodes(PyroWave, _)` answers from a flag set when the device was + created with nespyro's additions; nesvideo advertises PyroWave only then. + +Still open, out of scope: pixelforge's transfer queue can alias the graphics +queue on a device without a transfer or compute family, and pixelforge has no +lock to share. Unchanged by this part. + +## Decode thread (`decode/`) + +- The channel carries `enum VideoUnit { Coded(Vec, VideoCodec), + PyroWave(PyroUnit) }`, `PyroUnit { packets: Vec, critical_complete, + ts_ms }`. `VideoCodec` gains `PyroWave`; every closed match says what it + means there. +- `PyroDecoder` (new backend): built lazily for the stream's size and chroma, + rebuilt on `Push::Reconfigure`. Per unit: push every packet; decode when + `readiness().is_complete_enough_with(critical_complete)`; otherwise drop + and count. Output `DecodedFrame::PyroWave(PyroPicture)` into the same sink. +- `decode_ns` is CPU time as for the others; nespyro's GPU time goes to the + log at trace once a second. + +## Render (`vulkan/renderer.rs`, `yuv_convert.*`) + +- `import_frame` takes `PyroPicture` into the same slot the Vulkan picture + uses; retired the same way, after the frame that sampled it. +- The render submit waits on the picture's `ready` point when it has one. +- `yuv_convert_planar` entry: three `Texture2D` loads, chroma at + `(x >> shift, y >> shift)`, same matrix/range/peak code. R16 full range. +- Matrix and range from the picture's `ColourDescription`, per frame. + +## nesvideo + +- `PyroCollector` frames go to the decode channel as `VideoUnit::PyroWave`. + A full channel drops and counts; no IDR, nothing to resync. +- The sequence header (first packet) gives size and colour: `stream_dims`, and + `bitstream_colour` as a `StreamColour`, so the existing colour logic picks + the swapchain (HDR10 for PQ) exactly as for the other codecs. +- `client_capabilities` asks `nesrecon::decodes` for PyroWave like the rest. +- The overlay's "not decoded yet" note goes; the PyroWave rate line stays. + +## Tests (last step) + +- Unit: channel/unit plumbing, colour mapping, `block_grid_for`, the planar + shader's text checks beside the existing ones. +- Existing nesrecon/nesvideo tests stay green. +- Full stack by eye: PyroWave on screen, switching to and from a hardware + codec, 4:2:0 and 4:4:4, SDR and HDR. + +## Found while building + +- **The channel type stayed.** A PyroWave unit is serialised into the existing + `Vec` (`[flags][count]` then length-prefixed packets, + `nesrecon::encode_pyro_unit`), so PyroWave is one more `VideoDecoder` + variant and nothing else on the channel changed. One copy a frame, about + 37 MB/s at 300 Mbit/s. +- **A picture keeps its decoder alive.** Dropping a `nespyro::Decoder` + destroys its output images even while a `DecodedFrame` of them is held, and + a codec switch or resize drops one while the renderer still shows its last + picture. `PyroPicture` holds an `Arc` of the decoder, released after the + slot. Worth fixing in nespyro itself when it is next opened: the output + images should outlive the decoder or the frame should own them. +- **Capabilities are sent twice.** The first message goes out at connect, + before the renderer has created the device that decides whether PyroWave + decodes. A second follows once it has; the host renegotiates to the codec + it was already on. A `--codec pyrowave` request waits for the second, or the + hub would refuse it. +- **Deblock is suppressed, not switched off.** `PostProcessChain` gained a + per-type suppression, so a person's own pass toggles are untouched and the + pass comes back on a switch to a block codec. +- **A device shared for PyroWave alone** is not offered to the Vulkan Video + decoder (`SharedVulkanDevice::video_decode`), which would otherwise fail and + log an error on every switch to H.264. diff --git a/docs/superpowers/specs/2026-10-01-pyrowave-nescapture-design.md b/docs/superpowers/specs/2026-10-01-pyrowave-nescapture-design.md new file mode 100644 index 00000000..de2b48a1 --- /dev/null +++ b/docs/superpowers/specs/2026-10-01-pyrowave-nescapture-design.md @@ -0,0 +1,89 @@ +# PyroWave in nescapture + +Status: approved design, 2026-10-01. Part 3 of 4. Parts 1 and 2: +`2026-09-27-nespyro-codec-crate-design.md`, `2026-10-01-pyrowave-transport-design.md`. + +## Decisions + +- **Vulkan 1.3 instance.** `vkCreateInstance` raises anything below 1.3 to 1.3 + where the loader has it (today: 1.0 → 1.1). Only the client app has to care + about old Vulkan; nescapture runs on the game's machine, which we control. +- **Shared device only.** nespyro runs on the game's device or not at all. On + the own-device fallback a PyroWave request is refused and logged; a CPU round + trip of a frame defeats the codec's reason to exist. +- **Decoupled from Vulkan Video.** The game's device gets pixelforge's + additions, nespyro's, both, or neither, each checked on its own. A GPU with + no video encode can still stream PyroWave once a client asks; until then + nothing is encoded, and the log says why. + +## Device (`instance.rs`, `shared.rs`, `device.rs`) + +- `prepare` asks pixelforge `encode_device_requirements` and nespyro + `DeviceRequirements::query(Roles::ENCODE)` independently. + `Additions { video: Option, pyro: + Option, queues }`. The effective version + (min of instance and physical device) must be 1.3 for `pyro`. +- `Feature` gains nespyro's bits: Storage16Bit, Storage8Bit, ShaderInt8, + ShaderFloat16, BufferDeviceAddress, SubgroupSizeControl, + ComputeFullSubgroups, and the 1.0 core ShaderInt16 and + ShaderStorageImageWriteWithoutFormat. 1.0 bits live in + `VkPhysicalDeviceFeatures`: patched in place inside a chained + `VkPhysicalDeviceFeatures2`, else in a copy of `*pEnabledFeatures` the create + info points at instead, else in a copy of our own. Restore covers all three. +- Extension `VK_KHR_push_descriptor` added with nespyro. +- `plan_queues` plans encode only with pixelforge, compute with either, + transfer with pixelforge. nespyro uses the compute queue, `QueueLock::none()`: + one encoder thread submits for both, at most one queue of ours per family, + internally synchronized where shared with the game. +- `SharedDevice::pyro_context() -> Option`. + `SharedDevice::video_context()` returns an error when pixelforge's additions + were not made. + +## Encoder thread (`encode.rs`) + +- `enum Codec { Hw(HwCodec), PyroWave }` in `EncoderConfig`, + `EncodeSettingsChange`, `current_codec` (protocol id). `HwCodec` unchanged. + Every match says what PyroWave does; depth clamping, GOP, intra refresh and + IDR apply only to `Hw`. IDR on PyroWave is a no-op. +- `enum Backend { Hw(PerFrameEncoder), Pyro(PyroEncoder) }`. `PyroEncoder` + owns `nespyro::Encoder` and its rebuild key (w, h, depth, chroma, colour, + packet size). Bitrate-only change → `set_target_bitrate`, no rebuild. +- Input: the slot image in `GENERAL`, `encode_after(view, Source, blit)`. + `Source` from the same `source_spec`/`stream_spec`; only BT.709 SDR and + BT.2020 PQ, anything else refused. pixelforge `TimelinePoint` converted at + the boundary. +- Slot lifetime: a PyroWave frame's slot is held until its future resolves; up + to two held (nespyro's in-flight limit). Hardware path unchanged. +- `EncodedFrame.future` becomes an enum of the two futures. +- Codec switches use the existing `drop_encoder` + `needs_reconfig_flag`. + +## Commands, IPC, stats + +- `MSG_ENCODE_SETTINGS` via `decode_encode_settings_ext`. Codec 4 → + `PyroWave`, with chroma and packet size. Refused (warned) with no nespyro + context, no or too small packet size, or rate 0. +- `host_caps` includes PyroWave bits only with a nespyro context. +- IPC: PyroWave frames via `chunk_frame`, each message under the existing 250 + ms timeout. A timed-out message drops the rest of the frame and counts it in + `dropped`; other errors reconnect as today. Width/height from the frame, not + frozen at pipeline creation (fixes a stale-size bug for every codec). +- `enc_ms` from nespyro `gpu_ns`. Scratch overflow warned once per episode. +- Fix while open: the CPU fallback's hardcoded `Codec::H264` in + `InputImage::new`. + +## Tests + +- Unit: planning with one or both sets; 1.0 feature bits in `pEnabledFeatures`, + in `Features2`, in neither, and restoration; the instance raise decision; + command decode and refusals; the two-slot hold rule. +- GPU `#[ignore]`, in `crates/nespyro/tests/transport.rs` beside nespyro's GPU + harness: real frames → `chunk_frame` → `FrameAssembler` → `pack_datagrams` → + `PyroCollector` → `Depacketizer`; readiness equals the encoder's own packets, + clean and with loss. (Moved from part 2.) nescapture's own mapping of an + `EncodedFrame` onto `chunk_frame` (`pyro_messages`) is unit-tested on the + CPU, since `EncodedFrame` can be built by hand. +- Found by it: `Readiness::critical_complete` false-rejects frames with empty + coarse blocks; the datagram header gained a critical count so the collector + can answer exactly (see parts 1 and 2). +- Guest image: rebuild, run a title, send a PyroWave `MSG_ENCODE_SETTINGS` to + `/tmp/nescapture-cmd.sock`, and see PyroWave frames leave neshub. diff --git a/docs/superpowers/specs/2026-10-01-pyrowave-transport-design.md b/docs/superpowers/specs/2026-10-01-pyrowave-transport-design.md new file mode 100644 index 00000000..53f4e368 --- /dev/null +++ b/docs/superpowers/specs/2026-10-01-pyrowave-transport-design.md @@ -0,0 +1,176 @@ +# PyroWave transport + +Status: approved design, 2026-10-01. Part 2 of 4; part 1 is +`2026-09-27-nespyro-codec-crate-design.md`. + +## Why + +The media path was built for hardware codecs over the internet: a frame is one +IPC message, the hub cuts it into byte slices, the client needs every slice, +keyframes ride reliable streams, and a controller ratchets the bitrate. None of +that fits PyroWave: + +- A unix datagram tops out near 416 KB under the default `net.core.wmem_max`, + and the hub reads into 2 MiB. A PyroWave frame is 400 KB at 200 Mbit/s and + 2 MB at 1 Gbit/s. +- Byte slicing throws away what PyroWave is built around: packets that each + decode on their own, and a decoder that shows a frame missing some of them. +- Every PyroWave frame is intra. Flagged as a keyframe it would take the + two-slot reliable path every frame; unflagged, `ResyncGate` withholds it. + +## Decisions + +- **Opt-in only.** `ClientCaps` advertises PyroWave, but `best()` and + `CODEC_PREFERENCE` never return it. A session reaches it only by the client + sending `MSG_ENCODE_SETTINGS` naming it. +- **No controller.** PyroWave runs at a static rate: what the request names, or + the hub's `--pyrowave-kbps` default (300 000) when it names 0. The request is + CBR, which already sets the controller to Manual (`note_manual_target`), so + it stands down with no change to `control.rs`. The existing controller is + slated for replacement and gets nothing new. +- **Critical packets sent twice**, first and again last, as datagrams. No + reliable stream anywhere on the PyroWave path. +- **Shared logic in nesprotocol, pure.** Chunking, reassembly of IPC messages, + datagram packing and the client-side collector are functions and state + machines with an explicit clock. neshub wires them to sockets. nesprotocol + keeps zero dependencies and does not depend on nespyro: packets are byte + ranges to it. +- **The 4 MiB datagram buffer stays** for every codec. A hub-side skip guard + keeps PyroWave's queue to about two frames instead. + +## Scope + +In part 2: nesprotocol, neshub, tests, a throughput bench. Not in part 2: + +- nescapture (part 3) calls `chunk_frame`, reads the packet size from encode + settings, and checks the format version it was built with. +- nesvideo (part 4) dispatches datagram kind 2 to `PyroCollector` and feeds a + decoder. Moved from part 2: without a decoder the client never advertises + PyroWave, so the glue would be unreachable. The collector itself is built and + tested here, end to end through a real iroh connection. + +## Wire + +### Negotiation + +- `CODEC_PYROWAVE = 4`. `ClientCaps` bit `4*2 + DEPTH_8` = decodes PyroWave SDR, + `+ DEPTH_10` = decodes PyroWave HDR. Chroma is not a capability: anything + running nespyro decodes 4:2:0 and 4:4:4. +- `PYROWAVE_FORMAT: u8 = 1`, meaning upstream `89f7e47`'s bitstream. The client + appends it to `MSG_CLIENT_CAPS` as byte 3; `decode_client_caps` keeps reading + the first two. `decode_pyrowave_format(payload) -> Option` reads it. +- Encode settings grow two appended fields, read only when present: + `[codec][rc][value u32][depth][chroma][packet_size u16]`. `CHROMA_420 = 0`, + `CHROMA_444 = 1`. The client sends packet size 0; the hub fills it in. + `decode_encode_settings` keeps its signature; a new + `decode_encode_settings_ext` returns a struct with the two options. + +The hub refuses (logs, does not forward) a PyroWave request from a client +whose caps lacked PyroWave or carried a different `PYROWAVE_FORMAT`. + +### IPC (nescapture → neshub), stream type `STREAM_PYROWAVE = 7` + +The same 20-byte IPC header (codec `CODEC_PYROWAVE`, ts, width, height), and +its data is: + +```text +[frame u32][msg index u16][msg count u16][critical u16][packet count u16] +then per packet: [len u16][bytes] +``` + +`frame` is a counter, `critical` the frame's total of leading critical packets +(the same in every message). Each message holds at most `IPC_PYROWAVE_MAX` +(64 KiB) of data and whole packets only, in frame order. A packet is at most +`max(packet_size, 2488)` bytes, the largest block, so `u16` lengths suffice. + +### Datagram kind `DGRAM_PYROWAVE = 2` + +```text +[1B kind][2B seq][2B index][2B total][2B critical][4B ts_ms][1B flags][payload] +``` + +14 bytes. `seq` is the frame, `index`/`total` count datagrams in it. +`critical` is how many leading indexes carry the critical packets: a receiver +cannot work that out from the blocks it got, because an empty block is never +sent and looks like a lost one (found in part 3; the header was 12 bytes +before). The collector reports `critical_complete` from it. A payload +is one whole packet, or one piece of a packet too large for a datagram: + +- `FLAG_DUPLICATE` (bit 0): a second copy of a critical datagram. Same index. +- `FLAG_PART` (bit 1): a piece of one packet. Pieces occupy consecutive + indexes from one marked `FLAG_FIRST_PART` (bit 2) to one marked + `FLAG_LAST_PART` (bit 3). Both ends are marked so a receiver missing a piece + can tell an orphaned middle from the start of the next run. Only a block + larger than the payload limit does this, at most two or three pieces. + +`total` is capped at `PYRO_MAX_DATAGRAMS = 16384`, about 19 MB a frame. + +## neshub + +- **Packet size.** On forwarding a PyroWave request, the hub fills + `packet_size` with the smallest `max_datagram_size() - 12` across streaming + clients (1100 with none). The command forwarder in `main.rs` does it, since + it already serialises every command to nescapture. If the path later shrinks + below that, `pack_datagrams` sends the now-oversized packets in pieces like + any oversized block; nothing is dropped. +- **Listener.** The video socket accepts `STREAM_PYROWAVE` beside + `STREAM_VIDEO`. `FrameAssembler` collects a frame's messages; an incomplete + frame is abandoned when the next frame starts or after 100 ms, and counted. + The receive buffer stays 2 MiB, which is 32 messages' worth. +- **Broadcast.** The per-session video channel carries an enum: a hardware + frame as `Bytes`, or a PyroWave frame as `Arc`. Both are + refcounted, which also removes today's per-client copy of every frame. +- **Writer.** PyroWave frames skip `ResyncGate` and `KeyframeSender` and keep a + `seq` of their own. `pack_datagrams` lays them out (critical, rest, critical + again) into one `BytesMut`, split without copying. +- **Skip guard** (`pyrowave::should_skip`, shared with the bench). Before sending, with `backlog = DGRAM_BUFFER_BYTES - + datagram_send_buffer_space()`: skip the whole frame if the backlog exceeds + this frame's size, or if backlog plus frame exceeds the buffer. The skipped + frame's `seq` is reused (the transport doc's rule). A frame larger than the + buffer itself can never go out; that is warned once. Skips are counted and + logged at debug. + +## The collector (client side, in nesprotocol) + +`PyroCollector::push(datagram, now) -> Vec`, plus `poll(now)`. +Frames in flight by `seq`, each a slot per index (duplicates dropped by index). +A frame is released when: + +1. every index has arrived, or +2. a newer frame has started and `GRACE` (2 ms) has passed since, or +3. `IDLE` (20 ms) has passed since its *last* datagram. + +Not a deadline from the first datagram: on a saturated path a frame takes a +frame interval to arrive, so such a deadline tears every frame once the path is +slower than it assumes. The end-to-end burst test found this. + +Released frames carry `seq`, `ts_ms`, packets in index order (piece runs joined, +broken runs dropped), received/total, and `critical_complete`. A datagram for a `seq` at or behind +the newest released one is stale and dropped. No reorder window, no IDR: every +frame stands alone. Whether a frame is decodable is the decoder's call. + +Counts per frame: whole, partial, and datagrams lost, duplicated, stale. + +## Testing + +- **nesprotocol units:** chunk/assemble round trips including over 4 MB; + packing never splits a packet that fits; oversized packets come back whole + from pieces and a broken run is dropped; critical datagrams repeat with the + same index; every release trigger; stale and duplicate datagrams; the + negatives (`best()` never picks PyroWave, a wrong format is refused). +- **neshub end to end (CPU):** synthetic packets in as IPC messages through the + real listener and writer, over an iroh loopback connection, into a collector. + Byte-exact, clean and with dropped datagrams; the skip guard engages under a + blocked receiver. +- **End to end with real nespyro frames:** moved to part 3, where nescapture + produces them. Every released frame must give a `Depacketizer` the same + readiness as the encoder's own packets. +- **Throughput:** `examples/pyro_bench.rs` sends PyroWave-shaped frames at a + given rate over iroh. Loopback measures CPU cost; serve/connect across the + LAN measures the real thing. If iroh cannot carry the rates, that comes back + as a design question before part 3. + +## Order + +nestri `feat/pyrowave`: nesprotocol, then neshub, then the bench. nescore is +untouched until part 4 pins the rev.