diff --git a/.gitignore b/.gitignore index 8c5dcfe..7f7920e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,5 @@ output.* /target/ + +# Generated at runtime by examples/verify_all.rs (ensure_test_data); dimension-keyed. +testdata/test_frames_*x*_*.yuv diff --git a/Cargo.lock b/Cargo.lock index 8bb1429..ff92edf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -104,6 +104,21 @@ dependencies = [ "log", ] +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -204,11 +219,19 @@ version = "0.6.0" dependencies = [ "ash", "env_logger", + "futures-channel", + "pollster", "thiserror", "tracing", "tracing-subscriber", ] +[[package]] +name = "pollster" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" + [[package]] name = "portable-atomic" version = "1.13.1" diff --git a/Cargo.toml b/Cargo.toml index 4864988..209988d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,9 +15,11 @@ dmabuf = [] [dependencies] ash = { git = "https://github.com/ash-rs/ash", rev = "55dd56906bbb5760e9e9e6c56f45be67f67e0649", features = ["loaded"] } +futures-channel = { version = "0.3", default-features = false, features = ["std"] } thiserror = "1.0" tracing = "0.1" [dev-dependencies] env_logger = "0.11" +pollster = "0.4" tracing-subscriber = { version = "0.3.22", features = ["env-filter"] } diff --git a/examples/encode_av1.rs b/examples/encode_av1.rs index 2e08bca..dd7e1e8 100644 --- a/examples/encode_av1.rs +++ b/examples/encode_av1.rs @@ -85,6 +85,26 @@ fn main() -> Result<(), Box> { let mut output = File::create("output.av1")?; let mut total_bytes = 0; + // Each `encode()` returns a future that resolves with that frame's packet. + // Keep a few in flight (so capture/upload overlaps GPU encode) and drain the + // oldest once the pipeline is full, preserving submission order. + let mut pending: std::collections::VecDeque = + std::collections::VecDeque::new(); + + let mut write_packet = |packet: pixelforge::EncodedPacket, total: &mut usize| { + *total += packet.data.len(); + output.write_all(&packet.data)?; + println!( + " pts={:<2} dts={:<2}: {:>5} bytes, {:?}{}", + packet.pts, + packet.dts, + packet.data.len(), + packet.frame_type, + if packet.is_key_frame { " [KEY]" } else { "" } + ); + Ok::<(), Box>(()) + }; + // Encode frames. for i in 0..num_frames { let frame = &yuv_data[i * frame_size..(i + 1) * frame_size]; @@ -92,34 +112,21 @@ fn main() -> Result<(), Box> { // Upload YUV420 data to the input image. input_image.upload_yuv420(frame)?; - // Encode the image (passing InputImage's image, which triggers - // an internal copy to the encoder's input image with proper - // layout transitions). - for packet in encoder.encode(input_image.image())? { - total_bytes += packet.data.len(); - output.write_all(&packet.data)?; - println!( - " pts={:<2} dts={:<2}: {:>5} bytes, {:?}{}", - packet.pts, - packet.dts, - packet.data.len(), - packet.frame_type, - if packet.is_key_frame { " [KEY]" } else { "" } - ); + // Submit the frame (async) and keep the pipeline at most ~2 deep. Passing + // the InputImage's image triggers an internal copy into the encoder's slot + // image with proper layout transitions. + pending.push_back(encoder.encode(input_image.image())?); + while pending.len() > 2 { + let packet = pollster::block_on(pending.pop_front().unwrap())?; + write_packet(packet, &mut total_bytes)?; } } - // Flush remaining frames. - for packet in encoder.flush()? { - total_bytes += packet.data.len(); - output.write_all(&packet.data)?; - println!( - " pts={:<2} dts={:<2}: {:>5} bytes, {:?} (flushed)", - packet.pts, - packet.dts, - packet.data.len(), - packet.frame_type - ); + // Flush remaining frames: barrier, then drain the outstanding futures in order. + encoder.flush()?; + while let Some(future) = pending.pop_front() { + let packet = pollster::block_on(future)?; + write_packet(packet, &mut total_bytes)?; } let ratio = (num_frames * frame_size) as f64 / total_bytes as f64; diff --git a/examples/encode_h264.rs b/examples/encode_h264.rs index 64bfc24..349c1f6 100644 --- a/examples/encode_h264.rs +++ b/examples/encode_h264.rs @@ -85,6 +85,26 @@ fn main() -> Result<(), Box> { let mut output = File::create("output.h264")?; let mut total_bytes = 0; + // Each `encode()` returns a future that resolves with that frame's packet. + // Keep a few in flight (so capture/upload overlaps GPU encode) and drain the + // oldest once the pipeline is full, preserving submission order. + let mut pending: std::collections::VecDeque = + std::collections::VecDeque::new(); + + let mut write_packet = |packet: pixelforge::EncodedPacket, total: &mut usize| { + *total += packet.data.len(); + output.write_all(&packet.data)?; + println!( + " pts={:<2} dts={:<2}: {:>5} bytes, {:?}{}", + packet.pts, + packet.dts, + packet.data.len(), + packet.frame_type, + if packet.is_key_frame { " [KEY]" } else { "" } + ); + Ok::<(), Box>(()) + }; + // Encode frames. for i in 0..num_frames { let frame = &yuv_data[i * frame_size..(i + 1) * frame_size]; @@ -92,32 +112,19 @@ fn main() -> Result<(), Box> { // Upload YUV420 data to the input image. input_image.upload_yuv420(frame)?; - // Encode the image. - for packet in encoder.encode(input_image.image())? { - total_bytes += packet.data.len(); - output.write_all(&packet.data)?; - println!( - " pts={:<2} dts={:<2}: {:>5} bytes, {:?}{}", - packet.pts, - packet.dts, - packet.data.len(), - packet.frame_type, - if packet.is_key_frame { " [KEY]" } else { "" } - ); + // Submit the frame (async) and keep the pipeline at most ~2 deep. + pending.push_back(encoder.encode(input_image.image())?); + while pending.len() > 2 { + let packet = pollster::block_on(pending.pop_front().unwrap())?; + write_packet(packet, &mut total_bytes)?; } } - // Flush remaining frames. - for packet in encoder.flush()? { - total_bytes += packet.data.len(); - output.write_all(&packet.data)?; - println!( - " pts={:<2} dts={:<2}: {:>5} bytes, {:?} (flushed)", - packet.pts, - packet.dts, - packet.data.len(), - packet.frame_type - ); + // Flush remaining frames: barrier, then drain the outstanding futures in order. + encoder.flush()?; + while let Some(future) = pending.pop_front() { + let packet = pollster::block_on(future)?; + write_packet(packet, &mut total_bytes)?; } let ratio = (num_frames * frame_size) as f64 / total_bytes as f64; diff --git a/examples/encode_h265.rs b/examples/encode_h265.rs index 1ac971c..9adab69 100644 --- a/examples/encode_h265.rs +++ b/examples/encode_h265.rs @@ -91,6 +91,26 @@ fn main() -> Result<(), Box> { let mut output = File::create("output.h265")?; let mut total_bytes = 0; + // Each `encode()` returns a future that resolves with that frame's packet. + // Keep a few in flight (so capture/upload overlaps GPU encode) and drain the + // oldest once the pipeline is full, preserving submission order. + let mut pending: std::collections::VecDeque = + std::collections::VecDeque::new(); + + let mut write_packet = |packet: pixelforge::EncodedPacket, total: &mut usize| { + *total += packet.data.len(); + output.write_all(&packet.data)?; + println!( + " pts={:<2} dts={:<2}: {:>5} bytes, {:?}{}", + packet.pts, + packet.dts, + packet.data.len(), + packet.frame_type, + if packet.is_key_frame { " [KEY]" } else { "" } + ); + Ok::<(), Box>(()) + }; + // Encode frames. for i in 0..num_frames { let frame = &yuv_data[i * frame_size..(i + 1) * frame_size]; @@ -98,32 +118,19 @@ fn main() -> Result<(), Box> { // Upload YUV420 data to the input image. input_image.upload_yuv420(frame)?; - // Encode the image. - for packet in encoder.encode(input_image.image())? { - total_bytes += packet.data.len(); - output.write_all(&packet.data)?; - println!( - " pts={:<2} dts={:<2}: {:>5} bytes, {:?}{}", - packet.pts, - packet.dts, - packet.data.len(), - packet.frame_type, - if packet.is_key_frame { " [KEY]" } else { "" } - ); + // Submit the frame (async) and keep the pipeline at most ~2 deep. + pending.push_back(encoder.encode(input_image.image())?); + while pending.len() > 2 { + let packet = pollster::block_on(pending.pop_front().unwrap())?; + write_packet(packet, &mut total_bytes)?; } } - // Flush remaining frames. - for packet in encoder.flush()? { - total_bytes += packet.data.len(); - output.write_all(&packet.data)?; - println!( - " pts={:<2} dts={:<2}: {:>5} bytes, {:?} (flushed)", - packet.pts, - packet.dts, - packet.data.len(), - packet.frame_type - ); + // Flush remaining frames: barrier, then drain the outstanding futures in order. + encoder.flush()?; + while let Some(future) = pending.pop_front() { + let packet = pollster::block_on(future)?; + write_packet(packet, &mut total_bytes)?; } let ratio = (num_frames * frame_size) as f64 / total_bytes as f64; diff --git a/examples/rfi.rs b/examples/rfi.rs new file mode 100644 index 0000000..6d9d36c --- /dev/null +++ b/examples/rfi.rs @@ -0,0 +1,271 @@ +//! Example / verification: Reference Frame Invalidation (RFI) +//! +//! Encodes a stream and, partway through, simulates packet loss by calling +//! [`Encoder::invalidate_reference_frames`] for the two most recent references. +//! Because the encoder keeps a window of references (default 4), older ones +//! survive, so the next frame recovers by predicting from a surviving reference +//! instead of emitting a full keyframe. +//! +//! Correctness is checked two ways for each codec: +//! - the recovery frame must NOT be a keyframe (proving RFI engaged rather than +//! falling back to an IDR), and +//! - ffmpeg must decode the whole stream with high PSNR (proving the encoder's +//! reference re-signaling — H.264 MMCO, H.265 RPS, AV1 ref mapping — keeps the +//! decoder's DPB in sync with the encoder's). + +use pixelforge::{ + Codec, EncodeBitDepth, EncodeConfig, Encoder, InputImage, PixelFormat, RateControlMode, + VideoContextBuilder, +}; +use std::collections::VecDeque; +use std::fs::File; +use std::io::{Read, Write}; +use std::process::Command; + +const WIDTH: u32 = 320; +const HEIGHT: u32 = 240; +const FRAMES: u64 = 30; +/// Frame at which we simulate loss and invalidate references. +const INVALIDATE_AT: u64 = 12; +/// Number of most-recent references to invalidate (must stay below the encoder's +/// reference window so older references survive and recovery stays a P-frame). +const INVALIDATE_COUNT: u64 = 2; +/// PSNR below this means the decoder desynced from the encoder after recovery. +const MIN_PSNR: f64 = 30.0; + +fn main() -> Result<(), Box> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")), + ) + .init(); + + println!("PixelForge Reference Frame Invalidation verification\n"); + + let input_path = format!("testdata/test_frames_{WIDTH}x{HEIGHT}_yuv420p.yuv"); + ensure_test_data("yuv420p", &input_path)?; + + let context = VideoContextBuilder::new() + .app_name("RFI Verification") + .enable_validation(cfg!(debug_assertions)) + .build()?; + + let mut any_failure = false; + for codec in [Codec::H264, Codec::H265, Codec::AV1] { + if !context.supports_encode(codec) { + println!("{codec:?}: skipped (encode not supported)"); + continue; + } + match run_codec(&context, codec, &input_path) { + Ok(()) => {} + Err(e) => { + println!("{codec:?}: FAIL: {e}"); + any_failure = true; + } + } + } + + if any_failure { + std::process::exit(1); + } + Ok(()) +} + +fn run_codec( + context: &pixelforge::VideoContext, + codec: Codec, + input_path: &str, +) -> Result<(), Box> { + let ext = if codec == Codec::AV1 { "obu" } else { "bin" }; + let output_filename = format!("output_rfi_{codec:?}.{ext}"); + let decoded_filename = format!("decoded_rfi_{codec:?}.yuv"); + + let config = match codec { + Codec::H264 => EncodeConfig::h264(WIDTH, HEIGHT), + Codec::H265 => EncodeConfig::h265(WIDTH, HEIGHT), + Codec::AV1 => EncodeConfig::av1(WIDTH, HEIGHT), + } + .with_rate_control(RateControlMode::Cqp) + .with_quality_level(10) + .with_pixel_format(PixelFormat::Yuv420) + .with_bit_depth(EncodeBitDepth::Eight) + // Infinite GOP so the recovery can only come from RFI, never a periodic IDR. + .with_gop_size(0) + .with_b_frames(0); + + let mut encoder = Encoder::new(context.clone(), config)?; + let mut input_image = InputImage::new( + context.clone(), + codec, + WIDTH, + HEIGHT, + EncodeBitDepth::Eight, + PixelFormat::Yuv420, + )?; + + let mut yuv_data = Vec::new(); + File::open(input_path)?.read_to_end(&mut yuv_data)?; + let frame_size = (WIDTH * HEIGHT * 3 / 2) as usize; + + let mut output_file = File::create(&output_filename)?; + let mut pending: VecDeque = VecDeque::new(); + let mut recovery_is_key: Option = None; + + let drain_one = |pending: &mut VecDeque, + output_file: &mut File, + recovery_is_key: &mut Option| + -> Result<(), Box> { + let packet = pollster::block_on(pending.pop_front().unwrap())?; + if packet.pts == INVALIDATE_AT { + *recovery_is_key = Some(packet.is_key_frame); + } + output_file.write_all(&packet.data)?; + Ok(()) + }; + + for i in 0..FRAMES { + let start = (i as usize) * frame_size; + let end = start + frame_size; + if end > yuv_data.len() { + break; + } + + // Simulate the client losing the `INVALIDATE_COUNT` most recent frames + // just before encoding frame `INVALIDATE_AT`. + if i == INVALIDATE_AT { + let first_lost = INVALIDATE_AT - INVALIDATE_COUNT; + encoder.invalidate_reference_frames(first_lost); + } + + let encoder_image = encoder.input_image(); + input_image.upload_yuv420_to(encoder_image, &yuv_data[start..end])?; + pending.push_back(encoder.encode(encoder_image)?); + while pending.len() > 2 { + drain_one(&mut pending, &mut output_file, &mut recovery_is_key)?; + } + } + + encoder.flush()?; + while !pending.is_empty() { + drain_one(&mut pending, &mut output_file, &mut recovery_is_key)?; + } + drop(output_file); + + match recovery_is_key { + Some(true) => { + return Err(format!( + "recovery frame {INVALIDATE_AT} fell back to a keyframe (RFI did not engage)" + ) + .into()) + } + None => return Err(format!("recovery frame {INVALIDATE_AT} was never produced").into()), + Some(false) => {} + } + + let psnr = decode_and_psnr(&output_filename, &decoded_filename, input_path)?; + std::fs::remove_file(&output_filename).ok(); + std::fs::remove_file(&decoded_filename).ok(); + + if psnr < MIN_PSNR { + return Err(format!( + "post-recovery PSNR {psnr:.2} dB below {MIN_PSNR} dB (decoder desynced)" + ) + .into()); + } + + println!("{codec:?}: PASS — recovered with a P-frame, full-stream PSNR {psnr:.2} dB"); + Ok(()) +} + +/// Decode the bitstream to raw YUV and return its PSNR against the source. +fn decode_and_psnr( + bitstream: &str, + decoded: &str, + source: &str, +) -> Result> { + let status = Command::new("ffmpeg") + .args([ + "-hide_banner", + "-loglevel", + "error", + "-y", + "-i", + bitstream, + "-pix_fmt", + "yuv420p", + "-f", + "rawvideo", + decoded, + ]) + .output()?; + if !status.status.success() { + return Err(format!( + "ffmpeg decode failed: {}", + String::from_utf8_lossy(&status.stderr) + ) + .into()); + } + + let size = format!("{WIDTH}x{HEIGHT}"); + let output = Command::new("ffmpeg") + .args([ + "-hide_banner", + "-loglevel", + "info", + "-s", + &size, + "-pix_fmt", + "yuv420p", + "-f", + "rawvideo", + "-i", + source, + "-s", + &size, + "-pix_fmt", + "yuv420p", + "-f", + "rawvideo", + "-i", + decoded, + "-lavfi", + "psnr", + "-f", + "null", + "-", + ]) + .output()?; + let stderr = String::from_utf8_lossy(&output.stderr); + let pos = stderr + .find("average:") + .ok_or_else(|| format!("could not parse PSNR: {stderr}"))?; + let rest = &stderr[pos + 8..]; + let end = rest.find(' ').unwrap_or(rest.len()); + Ok(rest[..end].parse()?) +} + +fn ensure_test_data(pix_fmt: &str, path: &str) -> Result<(), Box> { + if std::path::Path::new(path).exists() { + return Ok(()); + } + println!("Generating {path}..."); + let status = Command::new("ffmpeg") + .args([ + "-f", + "lavfi", + "-i", + &format!("testsrc=duration=1:size={WIDTH}x{HEIGHT}:rate=30"), + "-pix_fmt", + pix_fmt, + "-f", + "rawvideo", + "-y", + path, + ]) + .output()?; + if !status.status.success() { + return Err(format!("failed to generate test data: {status:?}").into()); + } + Ok(()) +} diff --git a/examples/verify_all.rs b/examples/verify_all.rs index dd80a63..ea352b4 100644 --- a/examples/verify_all.rs +++ b/examples/verify_all.rs @@ -153,6 +153,10 @@ fn run_test( let mut output_file = File::create(&output_filename)?; + // Futures for in-flight encodes, drained in submission order. + let mut pending: std::collections::VecDeque = + std::collections::VecDeque::new(); + for i in 0..FRAMES { let start = (i as usize) * frame_size; let end = start + frame_size; @@ -171,10 +175,19 @@ fn run_test( _ => return Err("Unsupported format".into()), } - for packet in encoder.encode(encoder_image)? { + pending.push_back(encoder.encode(encoder_image)?); + while pending.len() > 2 { + let packet = pollster::block_on(pending.pop_front().unwrap())?; output_file.write_all(&packet.data)?; } } + + // Barrier, then drain the outstanding futures in submission order. + encoder.flush()?; + while let Some(future) = pending.pop_front() { + let packet = pollster::block_on(future)?; + output_file.write_all(&packet.data)?; + } } // 2. Decode to raw YUV diff --git a/src/encoder/av1/api.rs b/src/encoder/av1/api.rs deleted file mode 100644 index 6096b82..0000000 --- a/src/encoder/av1/api.rs +++ /dev/null @@ -1,268 +0,0 @@ -use super::AV1Encoder; - -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::{ColorDescription, EncodedPacket}; -use crate::error::{PixelForgeError, Result}; -use ash::vk; -use tracing::debug; - -impl AV1Encoder { - /// Get the internal input image. - /// - /// This image can be used as a target for `ColorConverter::convert` to avoid - /// an intermediate copy. - pub fn input_image(&self) -> vk::Image { - self.input_image - } - - /// Encode a frame from a GPU image. - /// - /// This accepts a source image on the GPU and encodes it directly without - /// any CPU-side data copies. The source image must be in the correct format - /// with the same dimensions as the encoder configuration, and should be in GENERAL layout. - /// - /// # Panics - /// - /// The encoder will panic at creation time if B-frames are enabled (b_frame_count > 0), - /// as B-frame encoding is not yet supported. - pub fn encode(&mut self, src_image: vk::Image) -> Result> { - let gop_position = self.gop.get_next_frame(); - let display_order = self.input_frame_num; - self.input_frame_num += 1; - - debug!( - "AV1 encode: frame {} from GPU image, type={:?}", - display_order, gop_position.frame_type - ); - - // Upload from GPU image. - self.upload_from_image(src_image)?; - - // Encode immediately. - let packet = self.encode_current_frame(&gop_position, display_order)?; - - Ok(vec![packet]) - } - - /// Internal method to encode the current frame already uploaded to input_image. - fn encode_current_frame( - &mut self, - gop_position: &GopPosition, - display_order: u64, - ) -> Result { - let is_key_frame = - gop_position.frame_type.is_idr() || gop_position.frame_type == GopFrameType::I; - let is_reference = gop_position.is_reference; - let frame_type = match gop_position.frame_type { - GopFrameType::Idr | GopFrameType::I => crate::encoder::FrameType::I, - GopFrameType::P => crate::encoder::FrameType::P, - GopFrameType::B => crate::encoder::FrameType::B, - }; - - debug!( - "Encoding frame: display_order={}, type={:?}, key={}, ref={}", - display_order, frame_type, is_key_frame, is_reference - ); - - if is_key_frame { - self.frame_num = 0; - self.order_hint = 0; - // Reset references for key frames. - self.references.clear(); - } - - let mut encoded_data = Vec::new(); - - // AV1 Temporal Delimiter OBU: type=2, has_size=1, size=0. - // Required as the first OBU in each temporal unit for conformant bitstreams. - // This enables ffmpeg's AV1 demuxer to detect frame boundaries in raw OBU streams. - encoded_data.extend_from_slice(&[0x12, 0x00]); - - // For key frames, prepend the AV1 Sequence Header OBU. - // This is required for AV1 decoders to initialize (equivalent to H.265 VPS/SPS/PPS). - if is_key_frame { - if self.header_data.is_none() { - let header = self.get_av1_sequence_header()?; - debug!( - "AV1 sequence header ({} bytes): {:02X?}", - header.len(), - &header[..std::cmp::min(32, header.len())] - ); - self.header_data = Some(header); - } - if let Some(ref header) = self.header_data { - encoded_data.extend_from_slice(header); - } - } - - encoded_data.extend_from_slice(&self.encode_frame_internal(gop_position, is_key_frame)?); - - // Save the order_hint used during encoding BEFORE incrementing. - let encoded_order_hint = self.order_hint; - self.encode_frame_num += 1; - self.frame_num += 1; - self.order_hint = (self.order_hint + 1) & 0xFF; // 8-bit order hint - - // Update reference frames and DPB slot management - // Key frames refresh all reference slots. Inter frames refresh only their own slot, - // becoming the new LAST_FRAME for the next inter frame - let ref_info = super::ReferenceInfo { - dpb_slot: self.current_dpb_slot, - order_hint: encoded_order_hint, - frame_type: if is_key_frame { - ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_KEY - } else { - ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_INTER - }, - }; - - // Store the encoded frame as the most recent reference - if is_key_frame { - self.references.clear(); - } - self.references.insert(0, ref_info); - // Keep only the most recent reference for single-reference prediction - self.references.truncate(1); - - // Cycle to next available DPB slot for the next frame - let used_slots: Vec = self.references.iter().map(|r| r.dpb_slot).collect(); - let mut next_slot = (self.current_dpb_slot + 1) % self.dpb_slot_count as u8; - - // If the next slot is in use, find the first available slot - if used_slots.contains(&next_slot) { - for i in 0..self.dpb_slot_count as u8 { - if !used_slots.contains(&i) { - next_slot = i; - break; - } - } - } - - self.current_dpb_slot = next_slot; - - Ok(EncodedPacket { - data: encoded_data, - frame_type, - is_key_frame, - pts: display_order, - dts: self.encode_frame_num - 1, - }) - } - - /// Flush the encoder and get any remaining packets. - pub fn flush(&mut self) -> Result> { - // No buffered frames in the current implementation. - Ok(Vec::new()) - } - - /// Request that the next frame be an IDR/key frame. - pub fn request_idr(&mut self) { - self.gop.request_idr(); - } - - /// Retrieve encoded AV1 Sequence Header OBU from video session parameters. - /// - /// Uses vkGetEncodedVideoSessionParametersKHR to get the driver-generated OBU. - /// The driver's sequence header must be used because the frame OBUs it produces - /// reference values from its internal sequence header (not ours). - fn get_av1_sequence_header(&self) -> Result> { - let get_info = vk::VideoEncodeSessionParametersGetInfoKHR { - video_session_parameters: self.session_params, - ..Default::default() - }; - - let mut data = vec![0u8; 4096]; - let mut data_size: usize = data.len(); - let mut feedback = vk::VideoEncodeSessionParametersFeedbackInfoKHR::default(); - - let mut attempts = 0; - loop { - attempts += 1; - let result = unsafe { - (self - .video_encode_fn - .fp() - .get_encoded_video_session_parameters_khr)( - self.context.device().handle(), - &get_info, - &mut feedback, - &mut data_size, - data.as_mut_ptr() as *mut std::ffi::c_void, - ) - }; - - match result { - vk::Result::SUCCESS => { - if data_size == 0 { - return Err(PixelForgeError::SessionParametersCreation( - "AV1 sequence header size is 0".to_string(), - )); - } - data.truncate(data_size); - debug!("Retrieved AV1 sequence header: {} bytes", data.len()); - return Ok(data); - } - vk::Result::INCOMPLETE if attempts < 3 => { - let new_size = data_size.max(data.len() * 2).max(1); - data.resize(new_size, 0); - data_size = data.len(); - } - err => { - return Err(PixelForgeError::SessionParametersCreation(format!( - "Failed to get AV1 sequence header: {:?}", - err - ))); - } - } - } - } - - /// Update the color description in the AV1 sequence header. - /// - /// This recreates the video session parameters with a new sequence header - /// containing the updated color configuration. The next encoded frame will - /// be a key frame with the new sequence header prepended. - pub fn set_color_description(&mut self, desc: ColorDescription) -> Result<()> { - // Wait for any in-flight encode to complete before modifying session params. - // Do NOT reset the fence here — submit_encode_and_read_bitstream() resets it - // before queue_submit. Leaving the fence signaled allows consecutive - // set_color_description() calls without deadlock. - unsafe { - self.context - .device() - .wait_for_fences(&[self.encode_fence], true, u64::MAX) - .map_err(|e| { - PixelForgeError::Synchronization(format!( - "Failed to wait for encode fence: {:?}", - e - )) - })?; - } - - // Save old handle so we can destroy it after successful creation. - let old_session_params = self.session_params; - - let new_session_params = self.create_session_params(&desc)?; - - // Destroy old session parameters now that the new ones are created. - unsafe { - self.video_queue_fn - .destroy_video_session_parameters(old_session_params, None); - } - - self.session_params = new_session_params; - self.config.color_description = Some(desc); - self.header_data = None; // Invalidate cached sequence header - self.gop.request_idr(); - - debug!( - "AV1 color description updated: primaries={}, transfer={}, matrix={}, full_range={}", - desc.color_primaries, - desc.transfer_characteristics, - desc.matrix_coefficients, - desc.full_range - ); - - Ok(()) - } -} diff --git a/src/encoder/av1/encode.rs b/src/encoder/av1/encode.rs deleted file mode 100644 index 273a6af..0000000 --- a/src/encoder/av1/encode.rs +++ /dev/null @@ -1,568 +0,0 @@ -use super::AV1Encoder; - -use crate::encoder::gop::GopPosition; -use crate::encoder::resources::{ - prepare_encode_command_buffer, record_dpb_barriers, record_post_encode_dpb_barrier, - submit_encode_and_read_bitstream, -}; -use crate::error::{PixelForgeError, Result}; -use ash::vk; -use ash::vk::TaggedStructure; -use tracing::debug; - -impl AV1Encoder { - pub(super) fn encode_frame_internal( - &mut self, - _gop_position: &GopPosition, - is_key_frame: bool, - ) -> Result> { - // All frames need a setup reference slot (DPB write) per Vulkan spec when maxDpbSlots > 0. - let is_reference = true; - - debug!( - "encode_frame_internal: key={}, ref={}, refs_len={}, dpb_slot={}", - is_key_frame, - is_reference, - self.references.len(), - self.current_dpb_slot - ); - - // Rate control setup (matches H265 pattern: CQP/Disabled uses DISABLED mode). - let (rc_mode, average_bitrate, max_bitrate, qp) = match self.config.rate_control_mode { - crate::encoder::RateControlMode::Cqp | crate::encoder::RateControlMode::Disabled => ( - vk::VideoEncodeRateControlModeFlagsKHR::DISABLED, - 0, - 0, - self.config.quality_level, - ), - crate::encoder::RateControlMode::Cbr => ( - vk::VideoEncodeRateControlModeFlagsKHR::CBR, - self.config.target_bitrate, - self.config.target_bitrate, - 128u32, - ), - crate::encoder::RateControlMode::Vbr => ( - vk::VideoEncodeRateControlModeFlagsKHR::VBR, - self.config.target_bitrate, - self.config.max_bitrate, - 128u32, - ), - }; - - // Prepare command buffer for recording. - unsafe { - prepare_encode_command_buffer( - self.context.device(), - self.encode_command_buffer, - self.query_pool, - )?; - } - - // Transition DPB images for encode. - let ref_dpb_slots: Vec = self.references.iter().map(|r| r.dpb_slot).collect(); - unsafe { - record_dpb_barriers( - self.context.device(), - self.encode_command_buffer, - &self.dpb_images, - false, // AV1 does not use layered DPB - self.current_dpb_slot, - &ref_dpb_slots, - self.dpb_slot_active[self.current_dpb_slot as usize], - ); - } - - // AV1 frame type. - let frame_type = if is_key_frame { - ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_KEY - } else { - ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_INTER - }; - - // Build picture info flags using ash's accessor methods. - // show_frame must be set for all frames; error_resilient_mode for key frames (match FFmpeg). - let mut picture_info_flags = ash::vk::native::StdVideoEncodeAV1PictureInfoFlags { - _bitfield_align_1: [], - _bitfield_1: Default::default(), - }; - picture_info_flags.set_show_frame(1); - if is_key_frame { - picture_info_flags.set_error_resilient_mode(1); - } else { - picture_info_flags.set_showable_frame(1); - } - - // Frame extent uses display dimensions for all picture resources. - // Per Vulkan spec, without MOTION_VECTOR_SCALING support, all picture resource - // codedExtent values must match, and srcPictureResource.codedExtent must equal - // the sequence header's max_frame_width/height. - let frame_extent = vk::Extent2D { - width: self.config.dimensions.width, - height: self.config.dimensions.height, - }; - - // Setup reconstructed picture (DPB slot for output). - let setup_picture_resource = vk::VideoPictureResourceInfoKHR::default() - .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(frame_extent) - .base_array_layer(0) - .image_view_binding(self.dpb_image_views[self.current_dpb_slot as usize]); - - // AV1 reference info for the setup slot. - let reference_info_flags = ash::vk::native::StdVideoEncodeAV1ReferenceInfoFlags { - _bitfield_align_1: [], - _bitfield_1: ash::vk::native::StdVideoEncodeAV1ReferenceInfoFlags::new_bitfield_1( - 0, 0, 0, - ), - }; - - let std_reference_info = ash::vk::native::StdVideoEncodeAV1ReferenceInfo { - flags: reference_info_flags, - frame_type: if is_key_frame { - ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_KEY - } else { - ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_INTER - }, - RefFrameId: self.current_dpb_slot as u32, - OrderHint: self.order_hint as u8, - reserved1: [0; 3], - pExtensionHeader: std::ptr::null(), - }; - - // AV1 DPB slot info for the setup reference slot (the slot being written). - let mut setup_av1_dpb_info = - vk::VideoEncodeAV1DpbSlotInfoKHR::default().std_reference_info(&std_reference_info); - - let mut setup_av1_dpb_info_ref0 = setup_av1_dpb_info; - let setup_reference_slot = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(self.current_dpb_slot as i32) - .picture_resource(&setup_picture_resource) - .push(&mut setup_av1_dpb_info_ref0); - - // Reference frames for inter frames. - let mut reference_slots = Vec::new(); - let mut av1_reference_infos = Vec::new(); - let mut ref_picture_resources = Vec::new(); - let mut ref_std_infos = Vec::new(); // Store std info to keep it alive - - if !is_key_frame && !self.references.is_empty() { - // Use the most recent reference frame. - let ref_info = &self.references[0]; - - // Create StdVideoEncodeAV1ReferenceInfo for the reference slot. - let ref_std_info = ash::vk::native::StdVideoEncodeAV1ReferenceInfo { - flags: reference_info_flags, - frame_type: ref_info.frame_type, - RefFrameId: ref_info.dpb_slot as u32, - OrderHint: ref_info.order_hint as u8, - reserved1: [0; 3], - pExtensionHeader: std::ptr::null(), - }; - ref_std_infos.push(ref_std_info); - - // Create AV1 DPB slot info for the reference (without pointer first). - let av1_ref_info = vk::VideoEncodeAV1DpbSlotInfoKHR::default(); - av1_reference_infos.push(av1_ref_info); - // Now set the pointer after it's in the vector at its final location. - av1_reference_infos[0] = av1_reference_infos[0].std_reference_info(&ref_std_infos[0]); - - let ref_picture_resource = vk::VideoPictureResourceInfoKHR::default() - .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(frame_extent) - .base_array_layer(0) - .image_view_binding(self.dpb_image_views[ref_info.dpb_slot as usize]); - ref_picture_resources.push(ref_picture_resource); - - // Create reference slot (without pNext first). - let ref_slot = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(ref_info.dpb_slot as i32) - .picture_resource(&ref_picture_resources[0]); - reference_slots.push(ref_slot); - // Now set pNext after it's in the vector at its final location. - reference_slots[0] = reference_slots[0].push(&mut av1_reference_infos[0]); - } - - // AV1 quantization parameters - required structure. - // Start with a moderate QP that the rate controller can adjust. - let quantization_flags = ash::vk::native::StdVideoAV1QuantizationFlags { - _bitfield_align_1: [], - _bitfield_1: ash::vk::native::StdVideoAV1QuantizationFlags::new_bitfield_1( - 0, // using_qmatrix - 0, // diff_uv_delta - 0, // reserved - ), - }; - - let quantization = ash::vk::native::StdVideoAV1Quantization { - flags: quantization_flags, - base_q_idx: qp as u8, // Use the same QP as constant_q_index - DeltaQYDc: 0, - DeltaQUDc: 0, - DeltaQUAc: 0, - DeltaQVDc: 0, - DeltaQVAc: 0, - qm_y: 0, - qm_u: 0, - qm_v: 0, - }; - - // CDEF (Constrained Directional Enhancement Filter) - required since we enabled it in sequence header. - // Match FFmpeg's default initialization (all zeros). - let cdef = ash::vk::native::StdVideoAV1CDEF { - cdef_damping_minus_3: 0, // Match FFmpeg: damping = 3 - cdef_bits: 0, // 1 CDEF strength combination (2^0) - cdef_y_pri_strength: [0, 0, 0, 0, 0, 0, 0, 0], // Match FFmpeg: all zeros - cdef_y_sec_strength: [0, 0, 0, 0, 0, 0, 0, 0], - cdef_uv_pri_strength: [0, 0, 0, 0, 0, 0, 0, 0], - cdef_uv_sec_strength: [0, 0, 0, 0, 0, 0, 0, 0], - }; - - // Loop filter - deblocking filter parameters. - // Match FFmpeg's default initialization. - let loop_filter_flags = ash::vk::native::StdVideoAV1LoopFilterFlags { - _bitfield_align_1: [], - _bitfield_1: ash::vk::native::StdVideoAV1LoopFilterFlags::new_bitfield_1( - 0, // loop_filter_delta_enabled - 0, // loop_filter_delta_update - 0, // reserved - ), - }; - - let loop_filter = ash::vk::native::StdVideoAV1LoopFilter { - flags: loop_filter_flags, - loop_filter_level: [0, 0, 0, 0], // Match FFmpeg: disable filter initially - loop_filter_sharpness: 0, - update_ref_delta: 0, - // Match FFmpeg's default_loop_filter_ref_deltas: { 1, 0, 0, 0, -1, 0, -1, -1 } - loop_filter_ref_deltas: [1, 0, 0, 0, -1, 0, -1, -1], - update_mode_delta: 1, // Match FFmpeg: set to 1 - loop_filter_mode_deltas: [0; 2], - }; - - // Tile info - FFmpeg has this commented out with "TODO FIX" at line 340. - // Match FFmpeg: don't provide tile info (set to null). - // Note: If this causes issues, we may need to re-enable it with proper values. - - // Build ref_frame_idx, ref_order_hint, refresh_frame_flags, and primary_ref_frame. - // For key frames: refresh all slots, all refs point to slot 0. - // For inter frames: LAST_FRAME points to most recent past frame, refresh only current slot. - let (ref_frame_idx, ref_order_hint, primary_ref_frame, refresh_frame_flags) = - self.calculate_reference_frame_mapping(is_key_frame); - - // AV1 encode picture info. - let std_picture_info = ash::vk::native::StdVideoEncodeAV1PictureInfo { - flags: picture_info_flags, - frame_type, - frame_presentation_time: self.frame_num, - current_frame_id: self.current_dpb_slot as u32, // Match FFmpeg: slot index - order_hint: self.order_hint as u8, - primary_ref_frame, - refresh_frame_flags, - coded_denom: 0, - render_width_minus_1: (self.config.dimensions.width - 1) as u16, - render_height_minus_1: (self.config.dimensions.height - 1) as u16, - interpolation_filter: ash::vk::native::StdVideoAV1InterpolationFilter_STD_VIDEO_AV1_INTERPOLATION_FILTER_EIGHTTAP, - TxMode: ash::vk::native::StdVideoAV1TxMode_STD_VIDEO_AV1_TX_MODE_SELECT, - delta_q_res: 0, - delta_lf_res: 0, - ref_order_hint, - ref_frame_idx, - reserved1: [0; 3], - delta_frame_id_minus_1: [0; 7], - pTileInfo: std::ptr::null(), - pQuantization: &quantization, - pSegmentation: std::ptr::null(), - pLoopFilter: &loop_filter, - pCDEF: &cdef, - pLoopRestoration: std::ptr::null(), - pGlobalMotion: std::ptr::null(), - pExtensionHeader: std::ptr::null(), - pBufferRemovalTimes: std::ptr::null(), - }; - - // Reference name slot indices - maps AV1 reference names to Vulkan DPB slot indices. - // Only set entries for reference names that appear in pReferenceSlots. - // For SINGLE_REFERENCE mode, only LAST_FRAME (index 0) is used. - let mut reference_name_slot_indices = [-1i32; 7]; - - if !is_key_frame && !self.references.is_empty() { - // Map LAST_FRAME to the reference's DPB slot. - let ref_info = &self.references[0]; - reference_name_slot_indices[0] = ref_info.dpb_slot as i32; - } - - // Set prediction mode and rate control group based on frame type. - let (prediction_mode, rate_control_group) = if is_key_frame { - ( - vk::VideoEncodeAV1PredictionModeKHR::INTRA_ONLY, - vk::VideoEncodeAV1RateControlGroupKHR::INTRA, - ) - } else { - ( - vk::VideoEncodeAV1PredictionModeKHR::SINGLE_REFERENCE, - vk::VideoEncodeAV1RateControlGroupKHR::PREDICTIVE, - ) - }; - - let mut av1_picture_info = vk::VideoEncodeAV1PictureInfoKHR::default() - .std_picture_info(&std_picture_info) - .prediction_mode(prediction_mode) - .rate_control_group(rate_control_group) - .reference_name_slot_indices(reference_name_slot_indices); - - // For DISABLED rate control mode, set constant_q_index on the picture info. - if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { - av1_picture_info = av1_picture_info.constant_q_index(qp); - } - - // AV1-specific rate control layer info. - let mut av1_rc_layer_info = vk::VideoEncodeAV1RateControlLayerInfoKHR::default(); - if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { - let q_index = vk::VideoEncodeAV1QIndexKHR { - intra_q_index: qp, - predictive_q_index: qp, - bipredictive_q_index: qp, - }; - av1_rc_layer_info = av1_rc_layer_info - .use_min_q_index(true) - .min_q_index(q_index) - .use_max_q_index(true) - .max_q_index(q_index); - } else { - // In CBR/VBR, let device handle QP - av1_rc_layer_info = av1_rc_layer_info - .use_min_q_index(false) - .use_max_q_index(false); - } - - let rc_layer_info = vk::VideoEncodeRateControlLayerInfoKHR::default() - .average_bitrate(average_bitrate as u64) - .max_bitrate(max_bitrate as u64) - .frame_rate_numerator(self.config.frame_rate_numerator) - .frame_rate_denominator(self.config.frame_rate_denominator) - .push(&mut av1_rc_layer_info); - - let rc_layers = [rc_layer_info]; - - // Rate control info (matches H265 pattern: only add layers/buffer for non-DISABLED modes). - let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc_mode); - - if rc_mode != vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { - rc_info = rc_info - .layers(&rc_layers) - .virtual_buffer_size_in_ms(self.config.virtual_buffer_size_ms) - .initial_virtual_buffer_size_in_ms(self.config.initial_virtual_buffer_size_ms); - } - - // Video begin coding info. - // Include the setup slot (with slot_index -1 to indicate it's not yet active) - // and any reference slots that will be used for reading during encoding. - let mut all_reference_slots = Vec::new(); - - if is_reference { - // Build a separate setup slot for begin coding with slot_index = -1. - // This tells the implementation the slot is being set up, not yet active. - let setup_slot_for_begin = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(-1) - .picture_resource(&setup_picture_resource) - .push(&mut setup_av1_dpb_info); - all_reference_slots.push(setup_slot_for_begin); - } - - // Add reference slots (already active slots we're reading from) - all_reference_slots.extend_from_slice(&reference_slots); - - debug!( - "Begin coding: {} reference slots (setup={}, refs={})", - all_reference_slots.len(), - if is_reference { 1 } else { 0 }, - reference_slots.len() - ); - - // Begin video coding with rate control info for non-first frames. - let is_first_frame = self.encode_frame_num == 0; - - // AV1-specific rate control info. - let mut av1_rc_info = vk::VideoEncodeAV1RateControlInfoKHR::default() - .gop_frame_count(self.config.gop_size) - .key_frame_period(self.config.gop_size) - .consecutive_bipredictive_frame_count(0) - .temporal_layer_count(1); - - let begin_coding_info = if is_first_frame { - vk::VideoBeginCodingInfoKHR::default() - .video_session(self.session) - .video_session_parameters(self.session_params) - .reference_slots(&all_reference_slots) - .push(&mut av1_rc_info) - } else { - vk::VideoBeginCodingInfoKHR::default() - .video_session(self.session) - .video_session_parameters(self.session_params) - .reference_slots(&all_reference_slots) - .push(&mut rc_info) - .push(&mut av1_rc_info) - }; - - unsafe { - self.video_queue_fn - .cmd_begin_video_coding(self.encode_command_buffer, &begin_coding_info); - } - - // Reset video coding state for the first frame. - // Combine RESET + RATE_CONTROL + QUALITY_LEVEL into a single control command. - // This matches the H265 approach and is required for AMD RADV. - if is_first_frame { - let mut quality_level_info = - vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); - - let control_info = vk::VideoCodingControlInfoKHR::default() - .flags( - vk::VideoCodingControlFlagsKHR::RESET - | vk::VideoCodingControlFlagsKHR::ENCODE_RATE_CONTROL - | vk::VideoCodingControlFlagsKHR::ENCODE_QUALITY_LEVEL, - ) - .push(&mut rc_info) - .push(&mut av1_rc_info) - .push(&mut quality_level_info); - - unsafe { - self.video_queue_fn - .cmd_control_video_coding(self.encode_command_buffer, &control_info); - } - } - - // Encode info. - let src_picture_resource = vk::VideoPictureResourceInfoKHR::default() - .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(frame_extent) - .base_array_layer(0) - .image_view_binding(self.input_image_view); - - let mut encode_info = vk::VideoEncodeInfoKHR::default() - .src_picture_resource(src_picture_resource) - .dst_buffer(self.bitstream_buffer) - .dst_buffer_offset(0) - .dst_buffer_range(self.bitstream_buffer_size as u64); - - if is_reference { - encode_info = encode_info.setup_reference_slot(&setup_reference_slot); - } - - if !reference_slots.is_empty() { - encode_info = encode_info.reference_slots(&reference_slots); - } - - encode_info = encode_info.push(&mut av1_picture_info); - - // Begin query to capture encode feedback (bitstream size, status). - unsafe { - self.context.device().cmd_begin_query( - self.encode_command_buffer, - self.query_pool, - 0, - vk::QueryControlFlags::empty(), - ); - } - - unsafe { - self.video_encode_fn - .cmd_encode_video(self.encode_command_buffer, &encode_info); - } - - // End query. - unsafe { - self.context - .device() - .cmd_end_query(self.encode_command_buffer, self.query_pool, 0); - } - - // Add DPB synchronization barrier after encoding. - unsafe { - record_post_encode_dpb_barrier( - self.context.device(), - self.encode_command_buffer, - &self.dpb_images, - false, // AV1 does not use layered DPB - self.current_dpb_slot, - ); - } - - // End video coding. - let end_coding_info = vk::VideoEndCodingInfoKHR::default(); - unsafe { - self.video_queue_fn - .cmd_end_video_coding(self.encode_command_buffer, &end_coding_info); - } - - // End command buffer. - unsafe { - self.context - .device() - .end_command_buffer(self.encode_command_buffer) - } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - - // Submit, wait, and read bitstream. - let encode_queue = self.context.video_encode_queue().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) - })?; - - debug!( - "Submitting frame {} to GPU: key={}, num_refs={}, cur_slot={}", - self.encode_frame_num, - is_key_frame, - self.references.len(), - self.current_dpb_slot - ); - - let gpu_start = std::time::Instant::now(); - - let encoded_data = unsafe { - submit_encode_and_read_bitstream( - self.context.device(), - self.encode_command_buffer, - self.encode_fence, - encode_queue, - self.query_pool, - self.bitstream_buffer_ptr, - )? - }; - - debug!("GPU encode took {:?}", gpu_start.elapsed()); - - // Mark current DPB slot as active. - self.dpb_slot_active[self.current_dpb_slot as usize] = true; - - Ok(encoded_data) - } - - /// Calculate proper reference frame mapping for AV1 encoding. - fn calculate_reference_frame_mapping(&self, is_key_frame: bool) -> ([i8; 7], [u8; 8], u8, u8) { - if is_key_frame { - // Key frame refreshes all 8 reference slots - // All named references (LAST_FRAME, LAST2_FRAME, etc.) point to slot 0 - ([0i8; 7], [0u8; 8], 7u8, 0xFFu8) - } else if !self.references.is_empty() { - let mut ref_frame_idx = [0i8; 7]; - let mut ref_order_hint = [0u8; 8]; - - // Map LAST_FRAME (index 0) to our most recent reference's DPB slot - let last_ref = &self.references[0]; - ref_frame_idx[0] = last_ref.dpb_slot as i8; - ref_order_hint[0] = last_ref.order_hint as u8; - - // Other reference slots remain 0 (unused/pointing to nothing active) - // They'll be ignored since we're using SINGLE_REFERENCE prediction mode - - // Refresh only the current DPB slot so this frame becomes the new LAST_FRAME - let refresh_flags = 1u8 << self.current_dpb_slot; - - // primary_ref_frame = 0 means LAST_FRAME is our primary reference - (ref_frame_idx, ref_order_hint, 0u8, refresh_flags) - } else { - // No references available (shouldn't happen for inter frames) - ([0i8; 7], [0u8; 8], 7u8, 0x00u8) - } - } -} diff --git a/src/encoder/av1/init.rs b/src/encoder/av1/init.rs index 4ffcd70..d66eb8d 100644 --- a/src/encoder/av1/init.rs +++ b/src/encoder/av1/init.rs @@ -1,393 +1,94 @@ -use super::{AV1Encoder, MIN_BITSTREAM_BUFFER_SIZE, SUPERBLOCK_SIZE}; +use super::{Av1, MIN_BITSTREAM_BUFFER_SIZE, SUPERBLOCK_SIZE}; -use crate::encoder::gop::GopStructure; -use crate::encoder::resources::{ - allocate_session_memory, clear_input_image, create_bitstream_buffer, create_command_resources, - create_dpb_images, create_image, get_video_format, make_codec_name, map_bitstream_buffer, - query_supported_video_formats, ClearImageParams, +use crate::encoder::codec::{ + build_encoder_common, query_video_caps, CodecEncoder, CommonInitRequest, }; -use crate::encoder::{ColorDescription, PixelFormat}; -use crate::error::{PixelForgeError, Result}; +use crate::encoder::{ColorDescription, EncodeConfig, PixelFormat}; +use crate::error::Result; use crate::vulkan::VideoContext; use ash::vk; use ash::vk::TaggedStructure; -use std::ptr; -use tracing::{debug, info}; +use tracing::info; -impl AV1Encoder { +impl Av1 { /// Create a new AV1 encoder. - pub fn new(context: VideoContext, config: crate::encoder::EncodeConfig) -> Result { - // B-frames are not yet supported. - if config.b_frame_count > 0 { - panic!( - "B-frame encoding is not yet supported. Set b_frame_count=0 in encoder config. \ - Got b_frame_count={}", - config.b_frame_count - ); - } - - let width = config.dimensions.width; - let height = config.dimensions.height; + pub fn create(context: VideoContext, config: EncodeConfig) -> Result> { + assert!( + config.b_frame_count == 0, + "B-frame encoding is not yet supported; set b_frame_count=0 (got {})", + config.b_frame_count + ); info!( - "Creating AV1 encoder: requested {}x{}, pixel_format={:?}", - width, height, config.pixel_format + "Creating AV1 encoder: {}x{}, pixel_format={:?}", + config.dimensions.width, config.dimensions.height, config.pixel_format ); - // Load video queue extension functions. - let video_queue_fn = - ash::khr::video_queue::Device::load(context.instance(), context.device()); - let video_encode_fn = - ash::khr::video_encode_queue::Device::load(context.instance(), context.device()); - - // Get chroma subsampling from pixel format. - let chroma_subsampling: vk::VideoChromaSubsamplingFlagsKHR = config.pixel_format.into(); - let luma_bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - let chroma_bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - - // AV1 profile selection based on chroma subsampling (not bit depth). - // Main profile: 8/10-bit, 4:2:0 only. - // High profile: 8/10-bit, 4:2:0 and 4:4:4. + // AV1 profile from chroma subsampling: Main is 4:2:0, High adds 4:4:4. let profile = match config.pixel_format { PixelFormat::Yuv420 => ash::vk::native::StdVideoAV1Profile_STD_VIDEO_AV1_PROFILE_MAIN, _ => ash::vk::native::StdVideoAV1Profile_STD_VIDEO_AV1_PROFILE_HIGH, }; - // Preferred input format based on pixel format and bit depth. - let preferred_src_format = get_video_format(config.pixel_format, config.bit_depth); + let chroma_subsampling: vk::VideoChromaSubsamplingFlagsKHR = config.pixel_format.into(); + let bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - // Create AV1 encode profile. let mut av1_profile_info = vk::VideoEncodeAV1ProfileInfoKHR::default().std_profile(profile); - - let mut profile_info = vk::VideoProfileInfoKHR::default() + let profile_info = vk::VideoProfileInfoKHR::default() .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_AV1) .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(luma_bit_depth) - .chroma_bit_depth(chroma_bit_depth) + .luma_bit_depth(bit_depth) + .chroma_bit_depth(bit_depth) .push(&mut av1_profile_info); - // Query encode capabilities. - let video_queue_instance = - ash::khr::video_queue::Instance::load(context.entry(), context.instance()); - let mut av1_capabilities = vk::VideoEncodeAV1CapabilitiesKHR::default(); - let mut encode_capabilities = vk::VideoEncodeCapabilitiesKHR::default(); - let mut capabilities = vk::VideoCapabilitiesKHR::default() - .push(&mut encode_capabilities) - .push(&mut av1_capabilities); - - let result = unsafe { - (video_queue_instance - .fp() - .get_physical_device_video_capabilities_khr)( - context.physical_device(), - &profile_info, - &mut capabilities, - ) - }; - if result != vk::Result::SUCCESS { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Failed to query Vulkan Video encode capabilities for AV1: {:?}", - result - ))); - } - - // Helper functions for alignment calculations. - let gcd = |mut a: u32, mut b: u32| { - while b != 0 { - let tmp = a % b; - a = b; - b = tmp; - } - a - }; - let lcm = |a: u32, b: u32| { - if a == 0 || b == 0 { - 0 - } else { - a / gcd(a, b) * b - } - }; - let align_up = |value: u32, alignment: u32| { - if alignment <= 1 { - value - } else { - value.div_ceil(alignment) * alignment - } - }; - - let gran_w = capabilities.picture_access_granularity.width.max(1); - let gran_h = capabilities.picture_access_granularity.height.max(1); - let align_w = lcm(SUPERBLOCK_SIZE, gran_w); - let align_h = lcm(SUPERBLOCK_SIZE, gran_h); + let bitstream_buffer_size = MIN_BITSTREAM_BUFFER_SIZE + .max(config.dimensions.width as usize * config.dimensions.height as usize); - let mut aligned_width = align_up(width, align_w); - let mut aligned_height = align_up(height, align_h); - - aligned_width = aligned_width.max(capabilities.min_coded_extent.width); - aligned_height = aligned_height.max(capabilities.min_coded_extent.height); - - if aligned_width > capabilities.max_coded_extent.width - || aligned_height > capabilities.max_coded_extent.height - { - return Err(PixelForgeError::InvalidInput(format!( - "Requested coded extent {}x{} (aligned to {}x{}) exceeds device max {}x{}", - width, - height, - aligned_width, - aligned_height, - capabilities.max_coded_extent.width, - capabilities.max_coded_extent.height - ))); - } - - info!( - "Using coded extent {}x{} (granularity {}x{})", - aligned_width, aligned_height, gran_w, gran_h - ); - - // Query supported formats. - let supported_src_formats = query_supported_video_formats( - &context, - &profile_info, - vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR, - )?; - let supported_dpb_formats = query_supported_video_formats( - &context, - &profile_info, - vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR, - )?; - - if supported_src_formats.is_empty() { - return Err(PixelForgeError::NoSuitableDevice( - "No supported Vulkan Video SRC formats for AV1".to_string(), - )); - } - if supported_dpb_formats.is_empty() { - return Err(PixelForgeError::NoSuitableDevice( - "No supported Vulkan Video DPB formats for AV1".to_string(), - )); - } - - info!("Supported SRC formats: {:?}", supported_src_formats); - info!("Supported DPB formats: {:?}", supported_dpb_formats); - - let picture_format = if supported_src_formats.contains(&preferred_src_format) { - preferred_src_format - } else { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Preferred input format {:?} not supported for AV1", - preferred_src_format - ))); - }; - - let reference_picture_format = supported_dpb_formats - .iter() - .copied() - .find(|f| *f == picture_format) - .unwrap_or(supported_dpb_formats[0]); - - debug!( - "Selected formats: picture={:?}, reference={:?}", - picture_format, reference_picture_format - ); - - // Get encode queue family. - let encode_queue_family = context.video_encode_queue_family().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue family available".to_string()) + // Query device capabilities with the AV1-specific capability struct + // chained in (required by the driver). + let mut av1_caps = vk::VideoEncodeAV1CapabilitiesKHR::default(); + let mut encode_caps = vk::VideoEncodeCapabilitiesKHR::default(); + let mut capabilities = vk::VideoCapabilitiesKHR::default() + .push(&mut encode_caps) + .push(&mut av1_caps); + let caps = query_video_caps(&context, &profile_info, &mut capabilities)?; + + let init = build_encoder_common(&CommonInitRequest { + context: &context, + config: &config, + profile_info: &profile_info, + caps: &caps, + align_unit: SUPERBLOCK_SIZE, + // AV1 has NUM_REF_FRAMES = 8 DPB slots, and `refresh_frame_flags` + // is an 8-bit syntax element. `build_encoder_common` adds one setup + // slot on top of the active count, so cap active refs at 7 to keep + // every DPB slot index in 0..=7 (the range the [u8; 8] order-hint + // array and `1 << slot` refresh mask can represent). + max_active_refs_cap: 7, + bitstream_buffer_size, + // AV1 reference handling here does not use a layered DPB. + allow_layered_dpb: false, })?; + let active_reference_count = init.active_reference_count; + let common = init.common; - // Create video session. - let std_header_version = vk::ExtensionProperties { - extension_name: make_codec_name(b"VK_STD_vulkan_video_codec_av1_encode"), - spec_version: vk::make_api_version(0, 1, 0, 0), - }; - - // Calculate DPB slots and active references. - let max_dpb_slots_supported = capabilities.max_dpb_slots as usize; - let max_active_reference_pictures_supported = - capabilities.max_active_reference_pictures as usize; - - if max_dpb_slots_supported < 2 { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Device reports max_dpb_slots={}, need at least 2", - max_dpb_slots_supported - ))); - } - - let mut target_active_refs = (config.max_reference_frames as usize) - .min(max_active_reference_pictures_supported) - .min(8); // AV1 supports up to 8 reference frames - - if target_active_refs < 1 && max_active_reference_pictures_supported >= 1 { - target_active_refs = 1; - } - - // AV1 typically needs: active refs + 1 for current frame being setup - let requested_dpb_slots = (target_active_refs + 1).min(max_dpb_slots_supported); - - info!( - "DPB configuration: slots={}, active_refs={} (max_supported: slots={}, refs={})", - requested_dpb_slots, - target_active_refs, - max_dpb_slots_supported, - max_active_reference_pictures_supported - ); - - let session_create_info = vk::VideoSessionCreateInfoKHR::default() - .queue_family_index(encode_queue_family) - .video_profile(&profile_info) - .picture_format(picture_format) - .max_coded_extent(vk::Extent2D { - width: aligned_width, - height: aligned_height, - }) - .reference_picture_format(reference_picture_format) - .max_dpb_slots(requested_dpb_slots as u32) - .max_active_reference_pictures(target_active_refs as u32) - .std_header_version(&std_header_version); - - let mut session = vk::VideoSessionKHR::null(); - let result = unsafe { - (video_queue_fn.fp().create_video_session_khr)( - context.device().handle(), - &session_create_info, - ptr::null(), - &mut session, - ) - }; - if result != vk::Result::SUCCESS { - return Err(PixelForgeError::VideoSessionCreation(format!( - "{:?}", - result - ))); - } - - // Allocate session memory. - let session_memory = allocate_session_memory(&context, session, &video_queue_fn)?; - - let color_desc = config - .color_description - .unwrap_or(ColorDescription::bt709()); - - // Create input image. - let (input_image, input_image_memory, input_image_view) = create_image( - &context, - aligned_width, - aligned_height, - picture_format, - false, // is_dpb - &profile_info, - )?; - let input_image_layout = vk::ImageLayout::UNDEFINED; - - // Create DPB images. - let (dpb_images, dpb_image_memories, dpb_image_views) = create_dpb_images( - &context, - aligned_width, - aligned_height, - reference_picture_format, - requested_dpb_slots, - &profile_info, - false, - )?; - // Create bitstream buffer. - let bitstream_buffer_size = MIN_BITSTREAM_BUFFER_SIZE.max(width as usize * height as usize); - let (bitstream_buffer, bitstream_buffer_memory) = - create_bitstream_buffer(&context, bitstream_buffer_size, &profile_info)?; - // Map bitstream buffer persistently. - let bitstream_buffer_ptr = - map_bitstream_buffer(&context, bitstream_buffer_memory, bitstream_buffer_size)?; - // Create command resources. - let upload_queue_family = context.transfer_queue_family(); - let cmd_resources = - create_command_resources(&context, encode_queue_family, upload_queue_family)?; - let command_pool = cmd_resources.command_pool; - let upload_command_buffer = cmd_resources.upload_command_buffer; - let upload_fence = cmd_resources.upload_fence; - let encode_command_buffer = cmd_resources.encode_command_buffer; - let encode_fence = cmd_resources.encode_fence; - // Clear the input image so padding between user dimensions and the - // aligned coded extent is zero-initialized. - clear_input_image( - &context, - &ClearImageParams { - command_buffer: upload_command_buffer, - fence: upload_fence, - queue: context.transfer_queue(), - image: input_image, - width: aligned_width, - height: aligned_height, - pixel_format: config.pixel_format, - bit_depth: config.bit_depth, - }, - )?; - // Create query pool for bitstream size queries. - // Need 1 query to capture bitstream offset and size. - // Need to provide profile info and feedback flags in pNext chain. - let mut query_feedback_info = vk::QueryPoolVideoEncodeFeedbackCreateInfoKHR::default() - .encode_feedback_flags( - vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BUFFER_OFFSET - | vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN, - ); - - let query_pool_create_info = unsafe { - vk::QueryPoolCreateInfo::default() - .query_type(vk::QueryType::VIDEO_ENCODE_FEEDBACK_KHR) - .query_count(1) - .push(&mut query_feedback_info) - .extend(&mut profile_info) - }; - - let query_pool = unsafe { - context - .device() - .create_query_pool(&query_pool_create_info, None) - .map_err(|e| PixelForgeError::QueryPool(e.to_string()))? - }; - - // Initialize GOP structure. - let gop = GopStructure::new(config.gop_size, config.b_frame_count, config.gop_size); - - let mut encoder = Self { - context, - config, - gop, - video_queue_fn, - video_encode_fn, - session, - session_params: vk::VideoSessionParametersKHR::null(), - session_memory, - input_frame_num: 0, - encode_frame_num: 0, + let codec = Av1 { frame_num: 0, order_hint: 0, - input_image, - input_image_memory, - input_image_view, - input_image_layout, - dpb_images, - dpb_image_memories, - dpb_image_views, - dpb_slot_count: requested_dpb_slots, - dpb_slot_active: vec![false; requested_dpb_slots], - bitstream_buffer, - bitstream_buffer_memory, - bitstream_buffer_size, - bitstream_buffer_ptr, - command_pool, - upload_command_pool: cmd_resources.upload_command_pool, - upload_command_buffer, - upload_fence, - encode_command_buffer, - encode_fence, - query_pool, header_data: None, - current_dpb_slot: 0, references: Vec::new(), + active_reference_count, }; - encoder.session_params = encoder.create_session_params(&color_desc)?; + let mut encoder = CodecEncoder { common, codec }; + let color_desc = config + .color_description + .unwrap_or(ColorDescription::bt709()); + encoder.common.session_params = encoder + .codec + .build_session_params(&encoder.common, &color_desc)?; + info!("AV1 encoder created successfully"); Ok(encoder) } } diff --git a/src/encoder/av1/mod.rs b/src/encoder/av1/mod.rs index 2edcde4..5dacdd7 100644 --- a/src/encoder/av1/mod.rs +++ b/src/encoder/av1/mod.rs @@ -1,26 +1,29 @@ -//! AV1 encoder implementation using Vulkan Video. +//! AV1 codec: the differences from the generic encoder. //! -//! This module implements AV1 video encoding using Vulkan Video extensions. +//! Shared machinery lives in `crate::encoder::codec`; this folder holds only +//! AV1's reference tracking (`Av1`), its per-frame StdVideo* graph (`record`), +//! and its sequence-header generation (`session_params`). +//! +//! Inter frames predict from up to `Av1::active_reference_count` recent +//! references, mapped onto the forward AV1 reference names. Keeping a window of +//! references (rather than just the previous frame) is what lets reference frame +//! invalidation recover from loss by re-anchoring to an older surviving frame +//! instead of forcing a key frame. -mod api; -mod encode; mod init; +mod record; mod session_params; -use ash::vk; -use tracing::debug; - -use crate::encoder::resources::{upload_image_to_input, UploadParams}; +use crate::encoder::codec::{EncoderCommon, FramePlan, PictureSetup, VideoCodec}; +use crate::encoder::pipeline::EncodeFuture; +use crate::encoder::ColorDescription; use crate::error::Result; - -use crate::encoder::gop::GopStructure; -use crate::encoder::EncodeConfig; -use crate::vulkan::VideoContext; +use ash::vk; /// Minimum bitstream buffer size. const MIN_BITSTREAM_BUFFER_SIZE: usize = 2 * 1024 * 1024; -/// AV1 superblock size in pixels (64x64, matching use_128x128_superblock=0 in the sequence header). +/// AV1 superblock size in pixels (64x64, matching use_128x128_superblock=0). pub const SUPERBLOCK_SIZE: u32 = 64; #[derive(Clone, Copy, Debug)] @@ -28,162 +31,128 @@ pub(crate) struct ReferenceInfo { pub dpb_slot: u8, pub order_hint: u32, pub frame_type: u32, + /// Display order (`pts`) of this reference, for reference frame invalidation. + pub display_order: u64, } -/// AV1 encoder. -pub struct AV1Encoder { - context: VideoContext, - config: EncodeConfig, - gop: GopStructure, - - // Video session. - video_queue_fn: ash::khr::video_queue::Device, - video_encode_fn: ash::khr::video_encode_queue::Device, - session: vk::VideoSessionKHR, - session_params: vk::VideoSessionParametersKHR, - session_memory: Vec, - - // Frame counters. - input_frame_num: u64, - encode_frame_num: u64, +/// AV1-specific encoder state (multi-reference prediction). +pub(crate) struct Av1 { frame_num: u32, order_hint: u32, - - // Resources - input_image: vk::Image, - input_image_memory: vk::DeviceMemory, - input_image_view: vk::ImageView, - /// Current Vulkan image layout of `input_image` (tracked to avoid UB when transitioning). - input_image_layout: vk::ImageLayout, - /// DPB images for reference frames. - dpb_images: Vec, - dpb_image_memories: Vec, - dpb_image_views: Vec, - /// Number of DPB slots allocated. - dpb_slot_count: usize, - /// Whether each DPB slot has been activated (written to at least once). - dpb_slot_active: Vec, - bitstream_buffer: vk::Buffer, - bitstream_buffer_memory: vk::DeviceMemory, - /// Size of the allocated bitstream buffer in bytes. - bitstream_buffer_size: usize, - /// Persistently mapped pointer to the bitstream buffer (avoids per-frame map/unmap). - bitstream_buffer_ptr: *mut u8, - - // Command resources. - command_pool: vk::CommandPool, - upload_command_pool: vk::CommandPool, - upload_command_buffer: vk::CommandBuffer, - upload_fence: vk::Fence, - encode_command_buffer: vk::CommandBuffer, - encode_fence: vk::Fence, - query_pool: vk::QueryPool, - - // Cached AV1 sequence header OBU (retrieved from session parameters). + /// Cached sequence-header OBU (invalidated when session params change). header_data: Option>, - - // Reference picture tracking. - /// Current DPB slot to use for setup (the reconstructed picture). - current_dpb_slot: u8, - /// Active reference pictures. Ordered from most recent to oldest. + /// Active references, most-recent first. references: Vec, + /// Negotiated number of active references kept for prediction. + active_reference_count: u32, } -impl AV1Encoder { - /// Upload input frame from a GPU image. - /// - /// This copies from a source image directly to the encoder's input image, - /// avoiding any CPU-side data copies. The source image must match the - /// encoder's configured pixel format and dimensions, and should be in - /// GENERAL layout. - fn upload_from_image(&mut self, src_image: vk::Image) -> Result<()> { - if src_image == self.input_image { - debug!("Source image is the encoder's input image, skipping upload copy"); - return Ok(()); +impl VideoCodec for Av1 { + fn begin_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + let is_key = plan.is_idr(); + if is_key { + self.frame_num = 0; + self.order_hint = 0; + self.references.clear(); } - let params = UploadParams { - upload_command_buffer: self.upload_command_buffer, - upload_fence: self.upload_fence, - src_image, - dst_image: self.input_image, - width: self.config.dimensions.width, - height: self.config.dimensions.height, - pixel_format: self.config.pixel_format, - input_image_layout: self.input_image_layout, - upload_queue: self.context.transfer_queue(), - }; - - upload_image_to_input(&self.context, ¶ms)?; + // Every temporal unit starts with a Temporal Delimiter OBU; key frames + // also carry the sequence header (so decoders can initialize). + let mut header = vec![0x12, 0x00]; + if is_key { + if self.header_data.is_none() { + self.header_data = Some(self.build_header(common)?); + } + if let Some(seq) = &self.header_data { + header.extend_from_slice(seq); + } + } - // Update tracked layout. - self.input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + Ok(PictureSetup { + frame_type: plan.frame_type(), + header: Some(header), + }) + } - Ok(()) + fn record_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + self.record(common, plan) } -} -// SAFETY: The raw pointer bitstream_buffer_ptr is only used within the encoder's -// thread and is properly synchronized via Vulkan fences before access -unsafe impl Send for AV1Encoder {} - -impl Drop for AV1Encoder { - fn drop(&mut self) { - unsafe { - let _ = self.context.device().device_wait_idle(); - self.context - .device() - .destroy_query_pool(self.query_pool, None); - self.context.device().destroy_fence(self.upload_fence, None); - self.context.device().destroy_fence(self.encode_fence, None); - self.context - .device() - .destroy_command_pool(self.command_pool, None); - if self.upload_command_pool != self.command_pool { - self.context - .device() - .destroy_command_pool(self.upload_command_pool, None); - } - self.context - .device() - .destroy_buffer(self.bitstream_buffer, None); - // Unmap the persistently mapped bitstream buffer before freeing memory. - self.context - .device() - .unmap_memory(self.bitstream_buffer_memory); - self.context - .device() - .free_memory(self.bitstream_buffer_memory, None); - self.context - .device() - .destroy_image_view(self.input_image_view, None); - self.context.device().destroy_image(self.input_image, None); - self.context - .device() - .free_memory(self.input_image_memory, None); - - for i in 0..self.dpb_images.len() { - self.context - .device() - .destroy_image_view(self.dpb_image_views[i], None); - self.context - .device() - .destroy_image(self.dpb_images[i], None); - self.context - .device() - .free_memory(self.dpb_image_memories[i], None); - } + fn end_picture(&mut self, common: &mut EncoderCommon, plan: &FramePlan) { + let is_key = plan.is_idr(); + // Order hint used while recording this frame (pre-increment). + let encoded_order_hint = self.order_hint; + self.frame_num += 1; + self.order_hint = (self.order_hint + 1) & 0xFF; - if self.session_params != vk::VideoSessionParametersKHR::null() { - self.video_queue_fn - .destroy_video_session_parameters(self.session_params, None); - } - self.video_queue_fn - .destroy_video_session(self.session, None); - for mem in &self.session_memory { - self.context.device().free_memory(*mem, None); + let ref_info = ReferenceInfo { + dpb_slot: common.current_dpb_slot, + order_hint: encoded_order_hint, + display_order: plan.display_order, + frame_type: if is_key { + ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_KEY + } else { + ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_INTER + }, + }; + + if is_key { + self.references.clear(); + } + self.references.insert(0, ref_info); + // Keep a sliding window of the most recent references for prediction and + // loss recovery. + while self.references.len() > self.active_reference_count.max(1) as usize { + self.references.pop(); + } + + // Cycle to the next available DPB slot. + let used_slots: Vec = self.references.iter().map(|r| r.dpb_slot).collect(); + let mut next_slot = (common.current_dpb_slot + 1) % common.dpb_slot_count as u8; + if used_slots.contains(&next_slot) { + for i in 0..common.dpb_slot_count as u8 { + if !used_slots.contains(&i) { + next_slot = i; + break; + } } } + common.current_dpb_slot = next_slot; + } + + fn invalidate_references( + &mut self, + _common: &mut EncoderCommon, + first_lost_display_order: u64, + ) -> bool { + // Drop every reference at or after the first lost frame: those are + // transitively undecodable on the client. Older survivors are kept, so + // the next P-frame re-anchors LAST_FRAME to the most recent of them (see + // `calculate_reference_frame_mapping`). Returns false only when no + // reference survives, in which case the caller forces a full key frame. + self.references + .retain(|r| r.display_order < first_lost_display_order); + !self.references.is_empty() + } + + fn create_session_params( + &self, + common: &EncoderCommon, + desc: &ColorDescription, + ) -> Result { + self.build_session_params(common, desc) + } + + fn invalidate_header_cache(&mut self) { + self.header_data = None; } } @@ -198,14 +167,12 @@ mod tests { #[test] fn test_superblock_alignment() { - // Dimensions should be aligned up to superblock boundaries. let align = |v: u32| (v + SUPERBLOCK_SIZE - 1) & !(SUPERBLOCK_SIZE - 1); - - assert_eq!(align(1920), 1920); // Already aligned - assert_eq!(align(1080), 1088); // 1080 rounds up to 1088 - assert_eq!(align(2560), 2560); // Already aligned - assert_eq!(align(1440), 1472); // 1440 rounds up to 1472 - assert_eq!(align(1), 64); // Minimum is one superblock + assert_eq!(align(1920), 1920); + assert_eq!(align(1080), 1088); + assert_eq!(align(2560), 2560); + assert_eq!(align(1440), 1472); + assert_eq!(align(1), 64); } #[test] @@ -214,12 +181,10 @@ mod tests { dpb_slot: 2, order_hint: 42, frame_type: 0, + display_order: 0, }; - assert_eq!(ref_info.dpb_slot, 2); assert_eq!(ref_info.order_hint, 42); - - // Should be Copy + Clone. let copied = ref_info; assert_eq!(copied.dpb_slot, ref_info.dpb_slot); assert_eq!(copied.order_hint, ref_info.order_hint); @@ -227,46 +192,37 @@ mod tests { #[test] fn test_order_hint_wrapping() { - // AV1 order hints are 8-bit, wrapping at 256. let mut order_hint: u32 = 254; for _ in 0..4 { order_hint = (order_hint + 1) & 0xFF; } - // 254 -> 255 -> 0 -> 1 -> 2 assert_eq!(order_hint, 2); } #[test] fn test_reference_tracking() { - // Simulate building up a reference list like the encoder does. let mut references: Vec = Vec::new(); let max_refs = 4usize; - for i in 0..6u8 { let ref_info = ReferenceInfo { dpb_slot: i % max_refs as u8, order_hint: i as u32, frame_type: 0, + display_order: 0, }; references.insert(0, ref_info); while references.len() > max_refs { references.pop(); } } - - // Should have exactly max_refs entries. assert_eq!(references.len(), max_refs); - // Most recent should be first. assert_eq!(references[0].order_hint, 5); assert_eq!(references[max_refs - 1].order_hint, 2); } #[test] fn test_key_frame_clears_references() { - // Simulate the encoder's key frame behavior: references.clear() on IDR. let mut references: Vec = Vec::new(); - - // Build up some references (like a sequence of P-frames). for i in 0..3u8 { references.insert( 0, @@ -274,22 +230,20 @@ mod tests { dpb_slot: i, order_hint: i as u32, frame_type: 0, + display_order: 0, }, ); } assert_eq!(references.len(), 3); - - // Key frame resets everything. references.clear(); assert!(references.is_empty()); - - // First frame after key should start fresh at slot 0. references.insert( 0, ReferenceInfo { dpb_slot: 0, order_hint: 0, frame_type: 0, + display_order: 0, }, ); assert_eq!(references.len(), 1); @@ -299,14 +253,11 @@ mod tests { #[test] fn test_dpb_slot_reuse() { - // Simulate the encoder's DPB slot allocation: after encoding a reference - // frame, find the first slot not used by any active reference. let max_refs = 2usize; - let dpb_slot_count = 3u8; // active_refs + 1 + let dpb_slot_count = 3u8; let mut references: Vec = Vec::new(); let mut current_dpb_slot: u8 = 0; - // Helper: find next free slot (mirrors api.rs logic). let find_free_slot = |refs: &[ReferenceInfo], slot_count: u8| -> u8 { let used: Vec = refs.iter().map(|r| r.dpb_slot).collect(); for i in 0..slot_count { @@ -314,10 +265,9 @@ mod tests { return i; } } - 0 // fallback (shouldn't happen with correct slot_count) + 0 }; - // Frame 0 (IDR): uses slot 0. assert_eq!(current_dpb_slot, 0); references.insert( 0, @@ -325,145 +275,66 @@ mod tests { dpb_slot: 0, order_hint: 0, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { references.pop(); } current_dpb_slot = find_free_slot(&references, dpb_slot_count); - assert_eq!(current_dpb_slot, 1); // slot 0 is used, next free is 1 + assert_eq!(current_dpb_slot, 1); - // Frame 1 (P): uses slot 1. references.insert( 0, ReferenceInfo { dpb_slot: 1, order_hint: 1, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { references.pop(); } current_dpb_slot = find_free_slot(&references, dpb_slot_count); - assert_eq!(current_dpb_slot, 2); // slots 0,1 used, next free is 2 + assert_eq!(current_dpb_slot, 2); - // Frame 2 (P): uses slot 2. Now all 3 slots have been touched, - // but max_refs=2 means the oldest reference (slot 0) gets evicted. references.insert( 0, ReferenceInfo { dpb_slot: 2, order_hint: 2, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { references.pop(); } - // references = [{slot:2, hint:2}, {slot:1, hint:1}] - slot 0 evicted assert_eq!(references.len(), 2); assert_eq!(references[0].dpb_slot, 2); assert_eq!(references[1].dpb_slot, 1); current_dpb_slot = find_free_slot(&references, dpb_slot_count); - assert_eq!(current_dpb_slot, 0); // slot 0 is now free again (reuse!) + assert_eq!(current_dpb_slot, 0); - // Frame 3 (P): uses recycled slot 0. references.insert( 0, ReferenceInfo { dpb_slot: 0, order_hint: 3, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { references.pop(); } - // references = [{slot:0, hint:3}, {slot:2, hint:2}] - slot 1 evicted - current_dpb_slot = find_free_slot(&references, dpb_slot_count); - assert_eq!(current_dpb_slot, 1); // slot 1 recycled - } - - #[test] - fn test_idr_p_p_idr_cycle() { - // Full GOP cycle: IDR -> P -> P -> IDR, verifying DPB slot allocation - // and reference list state at each step. - let max_refs = 2usize; - let dpb_slot_count = 3u8; - let mut references: Vec = Vec::new(); - let mut current_dpb_slot: u8 = 0; - - let find_free_slot = |refs: &[ReferenceInfo], slot_count: u8| -> u8 { - let used: Vec = refs.iter().map(|r| r.dpb_slot).collect(); - for i in 0..slot_count { - if !used.contains(&i) { - return i; - } - } - 0 - }; - - // IDR frame: clears refs, writes to slot 0. - references.clear(); - references.insert( - 0, - ReferenceInfo { - dpb_slot: current_dpb_slot, - order_hint: 0, - frame_type: 0, - }, - ); current_dpb_slot = find_free_slot(&references, dpb_slot_count); - - assert_eq!(references.len(), 1); - assert_eq!(references[0].order_hint, 0); assert_eq!(current_dpb_slot, 1); - - // P frame 1: writes to slot 1. - references.insert( - 0, - ReferenceInfo { - dpb_slot: current_dpb_slot, - order_hint: 1, - frame_type: 0, - }, - ); - while references.len() > max_refs { - references.pop(); - } - current_dpb_slot = find_free_slot(&references, dpb_slot_count); - - assert_eq!(references.len(), 2); - assert_eq!(references[0].order_hint, 1); - assert_eq!(current_dpb_slot, 2); - - // P frame 2: writes to slot 2, evicts oldest ref (slot 0). - references.insert( - 0, - ReferenceInfo { - dpb_slot: current_dpb_slot, - order_hint: 2, - frame_type: 0, - }, - ); - while references.len() > max_refs { - references.pop(); - } - current_dpb_slot = find_free_slot(&references, dpb_slot_count); - - assert_eq!(references.len(), 2); - assert_eq!(references[0].order_hint, 2); - assert_eq!(current_dpb_slot, 0); // slot 0 recycled - - // Second IDR: everything resets. - references.clear(); - assert!(references.is_empty()); } #[test] fn test_single_reference_slot() { - // Edge case: only 1 active reference with 2 DPB slots (minimum viable). let max_refs = 1usize; let dpb_slot_count = 2u8; let mut references: Vec = Vec::new(); @@ -478,13 +349,13 @@ mod tests { 0 }; - // Frame 0: slot 0. references.insert( 0, ReferenceInfo { dpb_slot: 0, order_hint: 0, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -493,13 +364,13 @@ mod tests { let mut current_dpb_slot = find_free_slot(&references, dpb_slot_count); assert_eq!(current_dpb_slot, 1); - // Frame 1: slot 1. Old ref (slot 0) evicted since max_refs=1. references.insert( 0, ReferenceInfo { dpb_slot: 1, order_hint: 1, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -508,21 +379,21 @@ mod tests { assert_eq!(references.len(), 1); assert_eq!(references[0].dpb_slot, 1); current_dpb_slot = find_free_slot(&references, dpb_slot_count); - assert_eq!(current_dpb_slot, 0); // ping-pong between 0 and 1 + assert_eq!(current_dpb_slot, 0); - // Frame 2: slot 0 again. references.insert( 0, ReferenceInfo { dpb_slot: 0, order_hint: 2, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { references.pop(); } current_dpb_slot = find_free_slot(&references, dpb_slot_count); - assert_eq!(current_dpb_slot, 1); // ping-pong back + assert_eq!(current_dpb_slot, 1); } } diff --git a/src/encoder/av1/record.rs b/src/encoder/av1/record.rs new file mode 100644 index 0000000..4b7b8b6 --- /dev/null +++ b/src/encoder/av1/record.rs @@ -0,0 +1,464 @@ +//! AV1 per-frame encode recording: builds the StdVideo* graph for one frame and +//! submits it. AV1 uses single-reference prediction. + +use super::Av1; + +use crate::encoder::codec::{EncoderCommon, FramePlan, RateControlPlan}; +use crate::encoder::pipeline::EncodeFuture; +use crate::encoder::resources::{ + prepare_encode_command_buffer, record_dpb_barriers, record_post_encode_dpb_barrier, +}; +use crate::error::{PixelForgeError, Result}; +use ash::vk; +use ash::vk::TaggedStructure; +use tracing::debug; + +impl Av1 { + pub(super) fn record( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + let is_key_frame = plan.is_idr(); + // All frames need a setup reference slot (DPB write) per the Vulkan spec + // when maxDpbSlots > 0. + let is_reference = true; + let current_dpb_slot = common.current_dpb_slot; + + let slot = common.pipeline.current(); + let command_buffer = slot.encode_command_buffer; + let query_pool = slot.query_pool; + let bitstream_buffer = slot.bitstream_buffer; + let bitstream_buffer_size = slot.bitstream_buffer_size; + let input_image_view = slot.input_image_view; + + debug!( + "av1 record: key={}, refs_len={}, dpb_slot={}", + is_key_frame, + self.references.len(), + current_dpb_slot + ); + + let rc = RateControlPlan::new(&common.config, 128); + let qp = rc.qp; + + unsafe { + prepare_encode_command_buffer(common.device(), command_buffer, query_pool)?; + } + let ref_dpb_slots: Vec = self.references.iter().map(|r| r.dpb_slot).collect(); + unsafe { + record_dpb_barriers( + common.device(), + command_buffer, + &common.dpb_images, + common.use_layered_dpb, + current_dpb_slot, + &ref_dpb_slots, + common.dpb_slot_active[current_dpb_slot as usize], + ); + } + + let frame_type = if is_key_frame { + ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_KEY + } else { + ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_INTER + }; + + // show_frame for all frames; error_resilient_mode on key frames (FFmpeg). + let mut picture_info_flags = ash::vk::native::StdVideoEncodeAV1PictureInfoFlags { + _bitfield_align_1: [], + _bitfield_1: Default::default(), + }; + picture_info_flags.set_show_frame(1); + if is_key_frame { + picture_info_flags.set_error_resilient_mode(1); + } else { + picture_info_flags.set_showable_frame(1); + } + + // Without MOTION_VECTOR_SCALING all picture resource codedExtents must + // match and equal the sequence header's max_frame dimensions. + let frame_extent = vk::Extent2D { + width: common.config.dimensions.width, + height: common.config.dimensions.height, + }; + + let setup_picture_resource = vk::VideoPictureResourceInfoKHR::default() + .coded_offset(vk::Offset2D { x: 0, y: 0 }) + .coded_extent(frame_extent) + .base_array_layer(0) + .image_view_binding(common.dpb_image_views[current_dpb_slot as usize]); + + let reference_info_flags = ash::vk::native::StdVideoEncodeAV1ReferenceInfoFlags { + _bitfield_align_1: [], + _bitfield_1: ash::vk::native::StdVideoEncodeAV1ReferenceInfoFlags::new_bitfield_1( + 0, 0, 0, + ), + }; + let std_reference_info = ash::vk::native::StdVideoEncodeAV1ReferenceInfo { + flags: reference_info_flags, + frame_type, + RefFrameId: current_dpb_slot as u32, + OrderHint: self.order_hint as u8, + reserved1: [0; 3], + pExtensionHeader: std::ptr::null(), + }; + + let mut setup_av1_dpb_info = + vk::VideoEncodeAV1DpbSlotInfoKHR::default().std_reference_info(&std_reference_info); + let mut setup_av1_dpb_info_ref0 = setup_av1_dpb_info; + let setup_reference_slot = vk::VideoReferenceSlotInfoKHR::default() + .slot_index(current_dpb_slot as i32) + .picture_resource(&setup_picture_resource) + .push(&mut setup_av1_dpb_info_ref0); + + // Reference frames for inter frames: declare each distinct DPB slot we + // still hold, most-recent first. The vectors hold the data so the + // pointers the slot infos take into them stay stable. + // + // `distinct_refs` dedups by DPB slot (a reference can be mapped to more + // than one AV1 reference name below, but each slot is declared once). + let mut distinct_refs: Vec<&super::ReferenceInfo> = Vec::new(); + if !is_key_frame { + for r in &self.references { + if !distinct_refs.iter().any(|d| d.dpb_slot == r.dpb_slot) { + distinct_refs.push(r); + } + } + } + + let mut ref_std_infos = Vec::with_capacity(distinct_refs.len()); + for ref_info in &distinct_refs { + ref_std_infos.push(ash::vk::native::StdVideoEncodeAV1ReferenceInfo { + flags: reference_info_flags, + frame_type: ref_info.frame_type, + RefFrameId: ref_info.dpb_slot as u32, + OrderHint: ref_info.order_hint as u8, + reserved1: [0; 3], + pExtensionHeader: std::ptr::null(), + }); + } + + let mut ref_picture_resources = Vec::with_capacity(distinct_refs.len()); + for ref_info in &distinct_refs { + ref_picture_resources.push( + vk::VideoPictureResourceInfoKHR::default() + .coded_offset(vk::Offset2D { x: 0, y: 0 }) + .coded_extent(frame_extent) + .base_array_layer(0) + .image_view_binding(common.dpb_image_views[ref_info.dpb_slot as usize]), + ); + } + + let mut av1_reference_infos = Vec::with_capacity(distinct_refs.len()); + for std_info in &ref_std_infos { + av1_reference_infos + .push(vk::VideoEncodeAV1DpbSlotInfoKHR::default().std_reference_info(std_info)); + } + + let mut reference_slots = Vec::with_capacity(distinct_refs.len()); + for ((ref_info, resource), dpb_info) in distinct_refs + .iter() + .zip(ref_picture_resources.iter()) + .zip(av1_reference_infos.iter_mut()) + { + reference_slots.push( + vk::VideoReferenceSlotInfoKHR::default() + .slot_index(ref_info.dpb_slot as i32) + .picture_resource(resource) + .push(dpb_info), + ); + } + + // Quantization (FFmpeg-style defaults). + let quantization = ash::vk::native::StdVideoAV1Quantization { + flags: ash::vk::native::StdVideoAV1QuantizationFlags { + _bitfield_align_1: [], + _bitfield_1: ash::vk::native::StdVideoAV1QuantizationFlags::new_bitfield_1(0, 0, 0), + }, + base_q_idx: qp as u8, + DeltaQYDc: 0, + DeltaQUDc: 0, + DeltaQUAc: 0, + DeltaQVDc: 0, + DeltaQVAc: 0, + qm_y: 0, + qm_u: 0, + qm_v: 0, + }; + + let cdef = ash::vk::native::StdVideoAV1CDEF { + cdef_damping_minus_3: 0, + cdef_bits: 0, + cdef_y_pri_strength: [0; 8], + cdef_y_sec_strength: [0; 8], + cdef_uv_pri_strength: [0; 8], + cdef_uv_sec_strength: [0; 8], + }; + + let loop_filter = ash::vk::native::StdVideoAV1LoopFilter { + flags: ash::vk::native::StdVideoAV1LoopFilterFlags { + _bitfield_align_1: [], + _bitfield_1: ash::vk::native::StdVideoAV1LoopFilterFlags::new_bitfield_1(0, 0, 0), + }, + loop_filter_level: [0; 4], + loop_filter_sharpness: 0, + update_ref_delta: 0, + loop_filter_ref_deltas: [1, 0, 0, 0, -1, 0, -1, -1], + update_mode_delta: 1, + loop_filter_mode_deltas: [0; 2], + }; + + let (ref_frame_idx, ref_order_hint, primary_ref_frame, refresh_frame_flags) = + self.calculate_reference_frame_mapping(is_key_frame, current_dpb_slot); + + let std_picture_info = ash::vk::native::StdVideoEncodeAV1PictureInfo { + flags: picture_info_flags, + frame_type, + frame_presentation_time: self.frame_num, + current_frame_id: current_dpb_slot as u32, + order_hint: self.order_hint as u8, + primary_ref_frame, + refresh_frame_flags, + coded_denom: 0, + render_width_minus_1: (common.config.dimensions.width - 1) as u16, + render_height_minus_1: (common.config.dimensions.height - 1) as u16, + interpolation_filter: ash::vk::native::StdVideoAV1InterpolationFilter_STD_VIDEO_AV1_INTERPOLATION_FILTER_EIGHTTAP, + TxMode: ash::vk::native::StdVideoAV1TxMode_STD_VIDEO_AV1_TX_MODE_SELECT, + delta_q_res: 0, + delta_lf_res: 0, + ref_order_hint, + ref_frame_idx, + reserved1: [0; 3], + delta_frame_id_minus_1: [0; 7], + pTileInfo: std::ptr::null(), + pQuantization: &quantization, + pSegmentation: std::ptr::null(), + pLoopFilter: &loop_filter, + pCDEF: &cdef, + pLoopRestoration: std::ptr::null(), + pGlobalMotion: std::ptr::null(), + pExtensionHeader: std::ptr::null(), + pBufferRemovalTimes: std::ptr::null(), + }; + + // Map the 7 AV1 reference names to DPB slots. The N most-recent + // references take the first N forward names (LAST, LAST2, LAST3, + // GOLDEN, ...); any remaining names reuse the oldest surviving + // reference so every name resolves to a valid, decodable DPB slot. This + // must mirror `ref_frame_idx` from `calculate_reference_frame_mapping`. + let mut reference_name_slot_indices = [-1i32; 7]; + if !is_key_frame && !self.references.is_empty() { + let n = self.references.len(); + for (name, idx) in reference_name_slot_indices.iter_mut().enumerate() { + *idx = self.references[name.min(n - 1)].dpb_slot as i32; + } + } + + let (prediction_mode, rate_control_group) = if is_key_frame { + ( + vk::VideoEncodeAV1PredictionModeKHR::INTRA_ONLY, + vk::VideoEncodeAV1RateControlGroupKHR::INTRA, + ) + } else { + ( + vk::VideoEncodeAV1PredictionModeKHR::SINGLE_REFERENCE, + vk::VideoEncodeAV1RateControlGroupKHR::PREDICTIVE, + ) + }; + + let mut av1_picture_info = vk::VideoEncodeAV1PictureInfoKHR::default() + .std_picture_info(&std_picture_info) + .prediction_mode(prediction_mode) + .rate_control_group(rate_control_group) + .reference_name_slot_indices(reference_name_slot_indices); + if rc.is_disabled() { + av1_picture_info = av1_picture_info.constant_q_index(qp); + } + + let mut av1_rc_layer_info = vk::VideoEncodeAV1RateControlLayerInfoKHR::default(); + if rc.is_disabled() { + let q_index = vk::VideoEncodeAV1QIndexKHR { + intra_q_index: qp, + predictive_q_index: qp, + bipredictive_q_index: qp, + }; + av1_rc_layer_info = av1_rc_layer_info + .use_min_q_index(true) + .min_q_index(q_index) + .use_max_q_index(true) + .max_q_index(q_index); + } else { + av1_rc_layer_info = av1_rc_layer_info + .use_min_q_index(false) + .use_max_q_index(false); + } + + let rc_layer_info = vk::VideoEncodeRateControlLayerInfoKHR::default() + .average_bitrate(rc.average_bitrate as u64) + .max_bitrate(rc.max_bitrate as u64) + .frame_rate_numerator(common.config.frame_rate_numerator) + .frame_rate_denominator(common.config.frame_rate_denominator) + .push(&mut av1_rc_layer_info); + let rc_layers = [rc_layer_info]; + + let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc.mode); + if !rc.is_disabled() { + rc_info = rc_info + .layers(&rc_layers) + .virtual_buffer_size_in_ms(common.config.virtual_buffer_size_ms) + .initial_virtual_buffer_size_in_ms(common.config.initial_virtual_buffer_size_ms); + } + + // Begin coding: setup slot (slot_index -1) plus any active reference slots. + let mut all_reference_slots = Vec::new(); + if is_reference { + let setup_slot_for_begin = vk::VideoReferenceSlotInfoKHR::default() + .slot_index(-1) + .picture_resource(&setup_picture_resource) + .push(&mut setup_av1_dpb_info); + all_reference_slots.push(setup_slot_for_begin); + } + all_reference_slots.extend_from_slice(&reference_slots); + + let is_first_frame = plan.is_first_frame(); + let mut av1_rc_info = vk::VideoEncodeAV1RateControlInfoKHR::default() + .gop_frame_count(common.config.gop_size) + .key_frame_period(common.config.gop_size) + .consecutive_bipredictive_frame_count(0) + .temporal_layer_count(1); + + let begin_coding_info = if is_first_frame { + vk::VideoBeginCodingInfoKHR::default() + .video_session(common.session) + .video_session_parameters(common.session_params) + .reference_slots(&all_reference_slots) + .push(&mut av1_rc_info) + } else { + vk::VideoBeginCodingInfoKHR::default() + .video_session(common.session) + .video_session_parameters(common.session_params) + .reference_slots(&all_reference_slots) + .push(&mut rc_info) + .push(&mut av1_rc_info) + }; + + unsafe { + common + .video_queue_fn + .cmd_begin_video_coding(command_buffer, &begin_coding_info); + } + + if is_first_frame { + let mut quality_level_info = + vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); + let control_info = vk::VideoCodingControlInfoKHR::default() + .flags( + vk::VideoCodingControlFlagsKHR::RESET + | vk::VideoCodingControlFlagsKHR::ENCODE_RATE_CONTROL + | vk::VideoCodingControlFlagsKHR::ENCODE_QUALITY_LEVEL, + ) + .push(&mut rc_info) + .push(&mut av1_rc_info) + .push(&mut quality_level_info); + unsafe { + common + .video_queue_fn + .cmd_control_video_coding(command_buffer, &control_info); + } + } + + let src_picture_resource = vk::VideoPictureResourceInfoKHR::default() + .coded_offset(vk::Offset2D { x: 0, y: 0 }) + .coded_extent(frame_extent) + .base_array_layer(0) + .image_view_binding(input_image_view); + + let mut encode_info = vk::VideoEncodeInfoKHR::default() + .src_picture_resource(src_picture_resource) + .dst_buffer(bitstream_buffer) + .dst_buffer_offset(0) + .dst_buffer_range(bitstream_buffer_size as u64); + if is_reference { + encode_info = encode_info.setup_reference_slot(&setup_reference_slot); + } + if !reference_slots.is_empty() { + encode_info = encode_info.reference_slots(&reference_slots); + } + encode_info = encode_info.push(&mut av1_picture_info); + + unsafe { + let device = common.device(); + device.cmd_begin_query( + command_buffer, + query_pool, + 0, + vk::QueryControlFlags::empty(), + ); + common + .video_encode_fn + .cmd_encode_video(command_buffer, &encode_info); + device.cmd_end_query(command_buffer, query_pool, 0); + + record_post_encode_dpb_barrier( + device, + command_buffer, + &common.dpb_images, + common.use_layered_dpb, + current_dpb_slot, + ); + + let end_coding_info = vk::VideoEndCodingInfoKHR::default(); + common + .video_queue_fn + .cmd_end_video_coding(command_buffer, &end_coding_info); + + device + .end_command_buffer(command_buffer) + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + } + + common.submit_frame() + } + + /// Build `ref_frame_idx`, `ref_order_hint`, `primary_ref_frame` and + /// `refresh_frame_flags` for the current frame. + /// + /// `ref_frame_idx[name]` is the DPB frame-buffer slot for AV1 reference + /// `name` (0 = LAST … 6 = ALTREF) and must mirror the + /// `reference_name_slot_indices` handed to the encoder. `ref_order_hint` is + /// indexed by DPB slot (not reference name) and carries each stored + /// reference's order hint, which the decoder needs because the sequence + /// header enables order hints. `STD_VIDEO_AV1_PRIMARY_REF_NONE` is 7. + fn calculate_reference_frame_mapping( + &self, + is_key_frame: bool, + current_dpb_slot: u8, + ) -> ([i8; 7], [u8; 8], u8, u8) { + const PRIMARY_REF_NONE: u8 = 7; + if is_key_frame { + // Key frame refreshes all 8 reference slots; all names point to slot 0. + return ([0i8; 7], [0u8; 8], PRIMARY_REF_NONE, 0xFFu8); + } + if self.references.is_empty() { + return ([0i8; 7], [0u8; 8], PRIMARY_REF_NONE, 0x00u8); + } + + let mut ref_frame_idx = [0i8; 7]; + let mut ref_order_hint = [0u8; 8]; + let n = self.references.len(); + // Most-recent references take the first forward names; surplus names + // reuse the oldest survivor. Matches `reference_name_slot_indices`. + for (name, idx) in ref_frame_idx.iter_mut().enumerate() { + *idx = self.references[name.min(n - 1)].dpb_slot as i8; + } + // Order hints keyed by DPB slot. + for r in &self.references { + ref_order_hint[r.dpb_slot as usize] = r.order_hint as u8; + } + // Load the frame context from the most recent (LAST) reference, and + // refresh only the current slot so this frame becomes the new LAST. + let refresh_flags = 1u8 << current_dpb_slot; + (ref_frame_idx, ref_order_hint, 0u8, refresh_flags) + } +} diff --git a/src/encoder/av1/session_params.rs b/src/encoder/av1/session_params.rs index da16cb3..2e1448b 100644 --- a/src/encoder/av1/session_params.rs +++ b/src/encoder/av1/session_params.rs @@ -1,29 +1,29 @@ -use super::AV1Encoder; +//! AV1 sequence-header generation and OBU retrieval. +use super::Av1; + +use crate::encoder::codec::EncoderCommon; +use crate::encoder::resources::get_encoded_session_params; use crate::encoder::{BitDepth, ColorDescription, PixelFormat}; use crate::error::{PixelForgeError, Result}; use ash::vk; use ash::vk::TaggedStructure; use std::ptr; -impl AV1Encoder { - /// Build AV1 sequence header and create Vulkan video session parameters. - /// - /// This is used both during initial encoder creation and by - /// `set_color_description()` to rebuild session parameters with - /// updated color configuration. Keeping a single implementation - /// ensures the sequence header stays bit-for-bit consistent. - pub(crate) fn create_session_params( +impl Av1 { + /// Build the AV1 sequence header and create Vulkan video session parameters. + pub(super) fn build_session_params( &self, + common: &EncoderCommon, desc: &ColorDescription, ) -> Result { - let width = self.config.dimensions.width; - let height = self.config.dimensions.height; + let config = &common.config; + let width = config.dimensions.width; + let height = config.dimensions.height; let frame_width_bits = 32 - (width - 1).leading_zeros(); let frame_height_bits = 32 - (height - 1).leading_zeros(); - // Map ColorDescription to AV1 enum constants. let av1_color_primaries = match desc.color_primaries { 9 => ash::vk::native::StdVideoAV1ColorPrimaries_STD_VIDEO_AV1_COLOR_PRIMARIES_BT_2020, _ => ash::vk::native::StdVideoAV1ColorPrimaries_STD_VIDEO_AV1_COLOR_PRIMARIES_BT_709, @@ -49,16 +49,15 @@ impl AV1Encoder { ), }; - let bit_depth = match self.config.bit_depth { + let bit_depth = match config.bit_depth { BitDepth::Eight => 8, BitDepth::Ten => 10, }; - // Chroma subsampling based on pixel format. - let (subsampling_x, subsampling_y) = match self.config.pixel_format { - PixelFormat::Yuv420 => (1u8, 1u8), // 4:2:0 - PixelFormat::Yuv444 => (0u8, 0u8), // 4:4:4 - _ => (1u8, 1u8), // Default to 4:2:0 + let (subsampling_x, subsampling_y) = match config.pixel_format { + PixelFormat::Yuv420 => (1u8, 1u8), + PixelFormat::Yuv444 => (0u8, 0u8), + _ => (1u8, 1u8), }; let color_config = ash::vk::native::StdVideoAV1ColorConfig { @@ -73,18 +72,18 @@ impl AV1Encoder { chroma_sample_position: ash::vk::native::StdVideoAV1ChromaSamplePosition_STD_VIDEO_AV1_CHROMA_SAMPLE_POSITION_UNKNOWN, }; - let profile = match self.config.pixel_format { + let profile = match config.pixel_format { PixelFormat::Yuv420 => ash::vk::native::StdVideoAV1Profile_STD_VIDEO_AV1_PROFILE_MAIN, _ => ash::vk::native::StdVideoAV1Profile_STD_VIDEO_AV1_PROFILE_HIGH, }; - // AV1 sequence header flags - use minimal set to avoid driver issues. + // Minimal sequence-header flags (avoid driver issues). let seq_flags = ash::vk::native::StdVideoAV1SequenceHeaderFlags { _bitfield_align_1: [], _bitfield_1: ash::vk::native::StdVideoAV1SequenceHeaderFlags::new_bitfield_1( 0, // still_picture 0, // reduced_still_picture_header - 0, // use_128x128_superblock (use 64x64 superblocks) + 0, // use_128x128_superblock 0, // enable_filter_intra 0, // enable_intra_edge_filter 0, // enable_interintra_compound @@ -122,7 +121,6 @@ impl AV1Encoder { pTimingInfo: ptr::null(), }; - // Create decoder model info (zero-initialized like FFmpeg). let decoder_model_info = ash::vk::native::StdVideoEncodeAV1DecoderModelInfo { buffer_delay_length_minus_1: 0, buffer_removal_time_length_minus_1: 0, @@ -131,7 +129,6 @@ impl AV1Encoder { num_units_in_decoding_tick: 0, }; - // Create operating point info (single operating point like FFmpeg). let operating_point = ash::vk::native::StdVideoEncodeAV1OperatingPointInfo { flags: ash::vk::native::StdVideoEncodeAV1OperatingPointInfoFlags { _bitfield_align_1: [], @@ -156,21 +153,35 @@ impl AV1Encoder { let mut quality_info = vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); let session_params_create_info = vk::VideoSessionParametersCreateInfoKHR::default() - .video_session(self.session) + .video_session(common.session) .push(&mut quality_info) .push(&mut av1_session_params_create_info); - let session_params = unsafe { - self.video_queue_fn + unsafe { + common + .video_queue_fn .create_video_session_parameters(&session_params_create_info, None) .map_err(|e| { PixelForgeError::SessionParametersCreation(format!( "Failed to create AV1 session parameters: {:?}", e )) - })? - }; + }) + } + } - Ok(session_params) + /// Retrieve the driver-generated sequence-header OBU for key frames. + pub(super) fn build_header(&self, common: &EncoderCommon) -> Result> { + let get_info = vk::VideoEncodeSessionParametersGetInfoKHR { + video_session_parameters: common.session_params, + ..Default::default() + }; + let mut feedback = vk::VideoEncodeSessionParametersFeedbackInfoKHR::default(); + get_encoded_session_params( + &common.context, + &common.video_encode_fn, + &get_info, + &mut feedback, + ) } } diff --git a/src/encoder/codec.rs b/src/encoder/codec.rs new file mode 100644 index 0000000..2e651bd --- /dev/null +++ b/src/encoder/codec.rs @@ -0,0 +1,711 @@ +//! The codec-generic encoder. +//! +//! Everything that is identical across H.264, H.265 and AV1 lives here: +//! the shared per-encoder state ([`EncoderCommon`]), the generic driver +//! ([`CodecEncoder`]) that owns the public `encode`/`flush`/`set_color_description` +//! flow, the shared initialization scaffolding ([`build_encoder_common`]) and the +//! shared rate-control decision ([`RateControlPlan`]). +//! +//! A codec is anything that implements [`VideoCodec`]: a small state struct that +//! plugs its *differences* into the generic flow — building the codec-specific +//! StdVideo* graph for a frame, tracking its reference pictures, and emitting its +//! parameter sets. Reading a codec's folder shows only those differences; the +//! scaffolding is here. + +use ash::vk; + +use crate::encoder::dpb::MAX_DPB_SLOTS; +use crate::encoder::gop::{GopFrameType, GopPosition, GopStructure}; +use crate::encoder::pipeline::{EncodeFuture, EncodePipeline, PipelineConfig, SlotPacketMetadata}; +use crate::encoder::resources::{ + align_up, allocate_session_memory, create_command_resources, create_dpb_images, + destroy_encoder_resources, get_video_format, lcm, query_supported_video_formats, + upload_image_to_input, EncoderTeardown, UploadParams, +}; +use crate::encoder::{ColorDescription, EncodeConfig, FrameType, RateControlMode}; +use crate::error::{PixelForgeError, Result}; +use crate::vulkan::VideoContext; + +/// Per-encoder state shared by every codec. +/// +/// Holds the Vulkan video session, the DPB images, the async [`EncodePipeline`], +/// the GOP structure and the upload path — none of which differ between codecs. +/// Codec-specific state (reference lists, syntax counters, parameter caches) +/// lives in the [`VideoCodec`] implementor instead. +pub(crate) struct EncoderCommon { + pub context: VideoContext, + pub config: EncodeConfig, + + pub video_queue_fn: ash::khr::video_queue::Device, + pub video_encode_fn: ash::khr::video_encode_queue::Device, + pub session: vk::VideoSessionKHR, + pub session_params: vk::VideoSessionParametersKHR, + pub session_memory: Vec, + + /// Coded extent (display dimensions aligned to the codec block size and the + /// device's picture-access granularity). + pub aligned_width: u32, + pub aligned_height: u32, + + /// Async slot rotation + bitstream readback. + pub pipeline: EncodePipeline, + /// Frame-type/POC schedule. + pub gop: GopStructure, + /// Monotonic display-order counter (presentation order). + pub input_frame_num: u64, + /// Monotonic encode-order counter (decode order / DTS; `0` => first frame). + pub encode_frame_num: u64, + + pub dpb_images: Vec, + pub dpb_image_memories: Vec, + pub dpb_image_views: Vec, + pub dpb_slot_count: usize, + /// Whether each DPB slot has been written at least once (governs the + /// UNDEFINED-vs-DPB old layout in the pre-encode barrier). + pub dpb_slot_active: Vec, + /// Single layered DPB image (true) vs one image per slot (false). + pub use_layered_dpb: bool, + /// DPB slot the current frame reconstructs into. + pub current_dpb_slot: u8, + + pub command_pool: vk::CommandPool, + pub upload_command_pool: vk::CommandPool, + pub upload_command_buffer: vk::CommandBuffer, + pub upload_fence: vk::Fence, +} + +impl EncoderCommon { + /// Prologue: wait until the slot we are about to record over has been read + /// back, pull the next GOP position, and snapshot the counters this frame + /// will encode with. + pub fn begin_frame(&mut self) -> FramePlan { + self.pipeline.wait_current_free(); + let gop = self.gop.get_next_frame(); + let display_order = self.input_frame_num; + self.input_frame_num += 1; + FramePlan { + gop, + display_order, + encode_index: self.encode_frame_num, + } + } + + /// Copy a source image into the current slot's input image. No-op when the + /// source already *is* the slot image (e.g. converted in place). + pub fn upload(&mut self, src_image: vk::Image) -> Result<()> { + let (dst_image, input_image_layout) = { + let slot = self.pipeline.current(); + (slot.input_image, slot.input_image_layout) + }; + if src_image == dst_image { + return Ok(()); + } + + let params = UploadParams { + upload_command_buffer: self.upload_command_buffer, + upload_fence: self.upload_fence, + src_image, + dst_image, + width: self.config.dimensions.width, + height: self.config.dimensions.height, + pixel_format: self.config.pixel_format, + input_image_layout, + upload_queue: self.context.transfer_queue(), + }; + upload_image_to_input(&self.context, ¶ms)?; + self.pipeline.current_mut().input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + Ok(()) + } + + /// Record the metadata for the packet the current slot will produce. + pub fn set_pending_metadata(&mut self, metadata: SlotPacketMetadata) { + self.pipeline.set_pending_metadata(metadata); + } + + /// Submit the recorded command buffer for the current slot, mark its DPB + /// slot active, and return the future that resolves with its packet. + pub fn submit_frame(&mut self) -> Result { + let encode_queue = self.context.video_encode_queue().ok_or_else(|| { + PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) + })?; + let future = self + .pipeline + .submit_current(self.context.device(), encode_queue)?; + self.dpb_slot_active[self.current_dpb_slot as usize] = true; + Ok(future) + } + + /// Epilogue: advance to the next pipeline slot. + pub fn advance(&mut self) { + self.pipeline.advance(); + } + + /// The `ash` device handle. + pub fn device(&self) -> &ash::Device { + self.context.device() + } +} + +/// Everything a single frame needs to know about its schedule, snapshotted by +/// [`EncoderCommon::begin_frame`] so the codec hooks all see a consistent view. +pub(crate) struct FramePlan { + pub gop: GopPosition, + /// Presentation order (PTS). + pub display_order: u64, + /// Encode order (DTS); `0` marks the very first encoded frame. + pub encode_index: u64, +} + +impl FramePlan { + pub fn is_idr(&self) -> bool { + self.gop.frame_type.is_idr() + } + pub fn is_reference(&self) -> bool { + self.gop.is_reference + } + pub fn is_b_frame(&self) -> bool { + self.gop.frame_type == GopFrameType::B + } + pub fn is_first_frame(&self) -> bool { + self.encode_index == 0 + } + pub fn pic_order_cnt(&self) -> i32 { + self.gop.pic_order_cnt + } + /// The stream-level frame type for packet metadata. + pub fn frame_type(&self) -> FrameType { + match self.gop.frame_type { + GopFrameType::Idr | GopFrameType::I => FrameType::I, + GopFrameType::P => FrameType::P, + GopFrameType::B => FrameType::B, + } + } +} + +/// The codec header (if any) and stream frame type produced by +/// [`VideoCodec::begin_picture`]. +pub(crate) struct PictureSetup { + pub frame_type: FrameType, + /// Codec header to prepend (SPS/PPS, VPS/SPS/PPS, AV1 sequence header). + /// `Some` only for frames that carry one (typically the IDR/key frame). + pub header: Option>, +} + +/// The resolved rate-control decision for a frame. +/// +/// The CQP/CBR/VBR selection logic is identical across codecs; only the default +/// QP the controller starts from differs (H.264/H.265 use 26, AV1 uses 128), so +/// the codec passes that in. Each codec then wires these values into its own +/// `VideoEncode*RateControl*InfoKHR` structs. +pub(crate) struct RateControlPlan { + pub mode: vk::VideoEncodeRateControlModeFlagsKHR, + pub average_bitrate: u32, + pub max_bitrate: u32, + /// QP/q-index: the configured quality level for CQP/Disabled, otherwise the + /// codec's default starting point for the bitrate controller. + pub qp: u32, +} + +impl RateControlPlan { + pub fn new(config: &EncodeConfig, controller_default_qp: u32) -> Self { + match config.rate_control_mode { + RateControlMode::Cqp | RateControlMode::Disabled => Self { + mode: vk::VideoEncodeRateControlModeFlagsKHR::DISABLED, + average_bitrate: 0, + max_bitrate: 0, + qp: config.quality_level, + }, + RateControlMode::Cbr => Self { + mode: vk::VideoEncodeRateControlModeFlagsKHR::CBR, + average_bitrate: config.target_bitrate, + max_bitrate: config.target_bitrate, + qp: controller_default_qp, + }, + RateControlMode::Vbr => Self { + mode: vk::VideoEncodeRateControlModeFlagsKHR::VBR, + average_bitrate: config.target_bitrate, + max_bitrate: config.max_bitrate, + qp: controller_default_qp, + }, + } + } + + pub fn is_disabled(&self) -> bool { + self.mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED + } +} + +/// A video codec plugged into the generic [`CodecEncoder`]. +/// +/// Implementors are small state structs (reference lists, syntax counters, +/// parameter caches). The generic driver calls these hooks in order around each +/// frame; everything else (slot rotation, upload, readback, teardown) is shared. +pub(crate) trait VideoCodec: Sized + Send { + /// Per-frame prologue: codec IDR/key-frame resets and header retrieval. + /// Runs before recording; returns the packet's frame type and header. + fn begin_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result; + + /// Record the codec-specific encode commands and submit the frame. Builds + /// the StdVideo* graph on its own stack (so the FFI pointers stay valid until + /// `cmd_encode_video`) and finishes via [`EncoderCommon::submit_frame`]. + fn record_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result; + + /// Per-frame epilogue: advance reference lists, syntax counters and the next + /// DPB slot. Runs after submission. + fn end_picture(&mut self, common: &mut EncoderCommon, plan: &FramePlan); + + /// Reference frame invalidation: drop every held reference the client can no + /// longer decode and arrange for the next frame to predict from a surviving + /// one. + /// + /// `first_lost_display_order` is the display order (the `pts` reported on + /// encoded packets) of the earliest frame the client lost. In a linear P + /// chain every reference at or after that frame is transitively undecodable, + /// so all of them are dropped. Returns `true` if a usable reference survives + /// (the next P-frame can be predicted from it); `false` if none remain and + /// the caller must fall back to an IDR. + fn invalidate_references( + &mut self, + common: &mut EncoderCommon, + first_lost_display_order: u64, + ) -> bool; + + /// Build the codec parameter sets and create Vulkan session parameters. + /// Used at init and by `set_color_description`. + fn create_session_params( + &self, + common: &EncoderCommon, + desc: &ColorDescription, + ) -> Result; + + /// Drop any cached header after the session parameters change. + fn invalidate_header_cache(&mut self); +} + +/// The codec-generic encoder: shared state plus a codec. +/// +/// This owns the entire public encode flow; the codec only supplies its +/// differences through [`VideoCodec`]. +pub struct CodecEncoder { + pub(crate) common: EncoderCommon, + pub(crate) codec: C, +} + +// SAFETY: the only non-Send state is the persistently-mapped bitstream pointer +// inside the pipeline slots, which is synchronized via Vulkan fences and only +// read on the dedicated readback thread (see `pipeline`). The codec state is +// `Send` by the trait bound. +unsafe impl Send for CodecEncoder {} + +impl CodecEncoder { + /// The internal input image for the current slot (a `ColorConverter::convert` + /// target that avoids an intermediate copy). + pub fn input_image(&self) -> vk::Image { + self.common.pipeline.input_image() + } + + /// Encode one frame, returning a future for its packet. See [`crate::Encoder::encode`]. + pub fn encode(&mut self, src_image: vk::Image) -> Result { + let plan = self.common.begin_frame(); + self.common.upload(src_image)?; + + let setup = self.codec.begin_picture(&mut self.common, &plan)?; + self.common.set_pending_metadata(SlotPacketMetadata { + frame_type: setup.frame_type, + is_key_frame: plan.is_idr(), + pts: plan.display_order, + dts: plan.encode_index, + header: setup.header, + }); + + let future = self.codec.record_picture(&mut self.common, &plan)?; + self.common.encode_frame_num += 1; + self.codec.end_picture(&mut self.common, &plan); + self.common.advance(); + Ok(future) + } + + /// End-of-stream barrier: wait for all in-flight frames to be read back. + pub fn flush(&mut self) -> Result<()> { + self.common.pipeline.flush(); + Ok(()) + } + + /// Force the next frame to be an IDR/key frame. + pub fn request_idr(&mut self) { + self.common.gop.request_idr(); + } + + /// Reference frame invalidation: drop references the client lost and predict + /// the next frame from a surviving reference, falling back to an IDR when no + /// reference survives. See [`crate::Encoder::invalidate_reference_frames`]. + pub fn invalidate_reference_frames(&mut self, first_lost_display_order: u64) { + if !self + .codec + .invalidate_references(&mut self.common, first_lost_display_order) + { + self.common.gop.request_idr(); + } + } + + /// Rebuild session parameters with a new color description; the next frame is + /// an IDR/key frame carrying the updated header. + pub fn set_color_description(&mut self, desc: ColorDescription) -> Result<()> { + // Drain in-flight encodes before mutating shared session parameters. + self.common.pipeline.wait_all_free(); + + let old_session_params = self.common.session_params; + let new_session_params = self.codec.create_session_params(&self.common, &desc)?; + unsafe { + self.common + .video_queue_fn + .destroy_video_session_parameters(old_session_params, None); + } + + self.common.session_params = new_session_params; + self.common.config.color_description = Some(desc); + self.codec.invalidate_header_cache(); + self.common.gop.request_idr(); + Ok(()) + } +} + +impl Drop for CodecEncoder { + fn drop(&mut self) { + unsafe { + let common = &mut self.common; + let device = common.context.device(); + // Wait on just the queues this encoder used, not the whole device. + let _ = device.queue_wait_idle(common.context.transfer_queue()); + if let Some(q) = common.context.video_encode_queue() { + let _ = device.queue_wait_idle(q); + } + + common.pipeline.destroy(device); + + destroy_encoder_resources( + device, + &common.video_queue_fn, + &EncoderTeardown { + command_pool: common.command_pool, + upload_command_pool: common.upload_command_pool, + upload_fence: common.upload_fence, + dpb_images: &common.dpb_images, + dpb_image_views: &common.dpb_image_views, + dpb_image_memories: &common.dpb_image_memories, + session: common.session, + session_params: common.session_params, + session_memory: &common.session_memory, + }, + ); + } + } +} + +/// What a codec passes to [`build_encoder_common`]; the codec owns the parts that +/// genuinely differ (its profile, block size, reference cap), the builder owns +/// the rest. +pub(crate) struct CommonInitRequest<'a> { + pub context: &'a VideoContext, + pub config: &'a EncodeConfig, + /// The codec profile, with its `VideoEncode*ProfileInfoKHR` already chained. + pub profile_info: &'a vk::VideoProfileInfoKHR<'a>, + /// Device capabilities for this profile, resolved by [`query_video_caps`]. + pub caps: &'a DeviceVideoCaps, + /// Codec block size for coded-extent alignment (macroblock / CTB / superblock). + pub align_unit: u32, + /// Upper bound on active reference pictures the codec's syntax allows. + pub max_active_refs_cap: usize, + pub bitstream_buffer_size: usize, + /// Whether the codec can use a layered DPB image when the driver lacks + /// `SEPARATE_REFERENCE_IMAGES` (H.264/H.265 yes; AV1 no). + pub allow_layered_dpb: bool, +} + +/// Result of [`build_encoder_common`]: the assembled common state plus the +/// negotiated active-reference count the codec needs for its parameter sets. +pub(crate) struct CommonInit { + pub common: EncoderCommon, + pub active_reference_count: u32, +} + +/// The device-capability fields the generic init needs, resolved once by the +/// codec. Copied out of the Vulkan query so they outlive its pointer chain. +pub(crate) struct DeviceVideoCaps { + pub picture_access_granularity: vk::Extent2D, + pub min_coded_extent: vk::Extent2D, + pub max_coded_extent: vk::Extent2D, + pub max_dpb_slots: u32, + pub max_active_reference_pictures: u32, + pub flags: vk::VideoCapabilityFlagsKHR, + pub std_header_version: vk::ExtensionProperties, +} + +/// Run `vkGetPhysicalDeviceVideoCapabilitiesKHR` and resolve the fields the +/// generic init needs. +/// +/// The driver *requires* the codec's `VideoEncode*CapabilitiesKHR` chained into +/// `pNext`, so the codec builds the `capabilities` chain (the only codec-specific +/// part) and passes it in already populated with its encode/codec caps structs. +pub(crate) fn query_video_caps( + context: &VideoContext, + profile_info: &vk::VideoProfileInfoKHR, + capabilities: &mut vk::VideoCapabilitiesKHR, +) -> Result { + let video_queue_instance = + ash::khr::video_queue::Instance::load(context.entry(), context.instance()); + let result = unsafe { + (video_queue_instance + .fp() + .get_physical_device_video_capabilities_khr)( + context.physical_device(), + profile_info, + capabilities, + ) + }; + if result != vk::Result::SUCCESS { + return Err(PixelForgeError::NoSuitableDevice(format!( + "Failed to query Vulkan Video encode capabilities: {:?}", + result + ))); + } + Ok(DeviceVideoCaps { + picture_access_granularity: capabilities.picture_access_granularity, + min_coded_extent: capabilities.min_coded_extent, + max_coded_extent: capabilities.max_coded_extent, + max_dpb_slots: capabilities.max_dpb_slots, + max_active_reference_pictures: capabilities.max_active_reference_pictures, + flags: capabilities.flags, + std_header_version: capabilities.std_header_version, + }) +} + +/// Query capabilities, create the video session, DPB images, command resources, +/// the encode pipeline and the GOP structure — the ~85% of encoder +/// initialization that is identical across codecs. +pub(crate) fn build_encoder_common(req: &CommonInitRequest) -> Result { + let context = req.context; + let config = req.config; + let width = config.dimensions.width; + let height = config.dimensions.height; + + let video_queue_fn = ash::khr::video_queue::Device::load(context.instance(), context.device()); + let video_encode_fn = + ash::khr::video_encode_queue::Device::load(context.instance(), context.device()); + + // Capabilities were queried by the codec (which chains its codec-specific + // capability struct, required by the driver) via [`query_video_caps`]. + let capabilities = req.caps; + + // Align the coded extent to lcm(codec block size, device granularity), then + // clamp up to the device minimum and re-align. + let gran_w = capabilities.picture_access_granularity.width.max(1); + let gran_h = capabilities.picture_access_granularity.height.max(1); + let align_w = lcm(req.align_unit, gran_w); + let align_h = lcm(req.align_unit, gran_h); + let aligned_width = align_up( + align_up(width, align_w).max(capabilities.min_coded_extent.width), + align_w, + ); + let aligned_height = align_up( + align_up(height, align_h).max(capabilities.min_coded_extent.height), + align_h, + ); + if aligned_width > capabilities.max_coded_extent.width + || aligned_height > capabilities.max_coded_extent.height + { + return Err(PixelForgeError::InvalidInput(format!( + "Requested coded extent {}x{} (aligned to {}x{} with granularity {}x{}) exceeds device max {}x{} for this profile", + width, height, aligned_width, aligned_height, gran_w, gran_h, + capabilities.max_coded_extent.width, capabilities.max_coded_extent.height + ))); + } + tracing::info!( + "Using coded extent {}x{} (granularity {}x{}, min {}x{}, max {}x{})", + aligned_width, + aligned_height, + gran_w, + gran_h, + capabilities.min_coded_extent.width, + capabilities.min_coded_extent.height, + capabilities.max_coded_extent.width, + capabilities.max_coded_extent.height + ); + + // Pick input (SRC) and reference (DPB) formats. + let preferred_src_format = get_video_format(config.pixel_format, config.bit_depth); + let supported_src_formats = query_supported_video_formats( + context, + req.profile_info, + vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR, + )?; + let supported_dpb_formats = query_supported_video_formats( + context, + req.profile_info, + vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR, + )?; + if supported_src_formats.is_empty() { + return Err(PixelForgeError::NoSuitableDevice( + "No supported Vulkan Video SRC formats for this profile".to_string(), + )); + } + if supported_dpb_formats.is_empty() { + return Err(PixelForgeError::NoSuitableDevice( + "No supported Vulkan Video DPB formats for this profile".to_string(), + )); + } + if !supported_src_formats.contains(&preferred_src_format) { + return Err(PixelForgeError::NoSuitableDevice(format!( + "Preferred input format {:?} is not supported for VIDEO_ENCODE_SRC_KHR. Supported: {:?}", + preferred_src_format, supported_src_formats + ))); + } + let picture_format = preferred_src_format; + let reference_picture_format = supported_dpb_formats + .iter() + .copied() + .find(|f| *f == picture_format) + .unwrap_or(supported_dpb_formats[0]); + + // Negotiate DPB slots and active references. + let max_dpb_slots_supported = capabilities.max_dpb_slots as usize; + let max_active_supported = capabilities.max_active_reference_pictures as usize; + if max_dpb_slots_supported < 2 { + return Err(PixelForgeError::NoSuitableDevice(format!( + "Device reports max_dpb_slots={} for this profile; need at least 2", + max_dpb_slots_supported + ))); + } + let mut target_active_refs = (config.max_reference_frames as usize) + .min(max_active_supported) + .min(req.max_active_refs_cap); + if target_active_refs < 1 && max_active_supported >= 1 { + target_active_refs = 1; + } + // B-frames are not yet supported, so a reconstructed-frame slot plus the + // active references is enough. + let dpb_slot_count = (target_active_refs + 1) + .min(max_dpb_slots_supported) + .min(MAX_DPB_SLOTS); + let max_active_reference_pictures = target_active_refs.min(dpb_slot_count.saturating_sub(1)); + + let encode_queue_family = context.video_encode_queue_family().ok_or_else(|| { + PixelForgeError::NoSuitableDevice("No video encode queue family available".to_string()) + })?; + + // Use the driver-reported std header version for this profile. + let std_header_version = capabilities.std_header_version; + let session_create_info = vk::VideoSessionCreateInfoKHR::default() + .queue_family_index(encode_queue_family) + .flags(vk::VideoSessionCreateFlagsKHR::empty()) + .video_profile(req.profile_info) + .picture_format(picture_format) + .max_coded_extent(vk::Extent2D { + width: aligned_width, + height: aligned_height, + }) + .reference_picture_format(reference_picture_format) + .max_dpb_slots(dpb_slot_count as u32) + .max_active_reference_pictures(max_active_reference_pictures as u32) + .std_header_version(&std_header_version); + + let mut session = vk::VideoSessionKHR::null(); + let result = unsafe { + (video_queue_fn.fp().create_video_session_khr)( + context.device().handle(), + &session_create_info, + std::ptr::null(), + &mut session, + ) + }; + if result != vk::Result::SUCCESS { + return Err(PixelForgeError::VideoSessionCreation(format!( + "{:?}", + result + ))); + } + let session_memory = allocate_session_memory(context, session, &video_queue_fn)?; + + // Use a layered DPB only when allowed and the driver lacks separate + // reference images (AMD RADV). + let supports_separate_dpb = capabilities + .flags + .contains(vk::VideoCapabilityFlagsKHR::SEPARATE_REFERENCE_IMAGES); + let use_layered_dpb = req.allow_layered_dpb && !supports_separate_dpb; + if use_layered_dpb { + tracing::info!("Using layered DPB (driver does not support separate reference images)"); + } + + let (dpb_images, dpb_image_memories, dpb_image_views) = create_dpb_images( + context, + aligned_width, + aligned_height, + reference_picture_format, + dpb_slot_count, + req.profile_info, + use_layered_dpb, + )?; + + let upload_queue_family = context.transfer_queue_family(); + let cmd = create_command_resources(context, encode_queue_family, upload_queue_family)?; + + let pipeline = EncodePipeline::new(&PipelineConfig { + context, + aligned_width, + aligned_height, + picture_format, + pixel_format: config.pixel_format, + bit_depth: config.bit_depth, + bitstream_buffer_size: req.bitstream_buffer_size, + profile_info: req.profile_info, + command_pool: cmd.command_pool, + upload_command_buffer: cmd.upload_command_buffer, + upload_fence: cmd.upload_fence, + })?; + + // B-frames are not yet supported, so an I-P GOP. Set the SPS-matching + // counters (no-ops for AV1, which keys off order hints). + let mut gop = GopStructure::new_ip_only(config.gop_size); + gop.set_max_frame_num(4); + gop.set_max_poc_lsb(4); + + let common = EncoderCommon { + context: context.clone(), + config: config.clone(), + video_queue_fn, + video_encode_fn, + session, + session_params: vk::VideoSessionParametersKHR::null(), + session_memory, + aligned_width, + aligned_height, + pipeline, + gop, + input_frame_num: 0, + encode_frame_num: 0, + dpb_images, + dpb_image_memories, + dpb_image_views, + dpb_slot_count, + dpb_slot_active: vec![false; dpb_slot_count], + use_layered_dpb, + current_dpb_slot: 0, + command_pool: cmd.command_pool, + upload_command_pool: cmd.upload_command_pool, + upload_command_buffer: cmd.upload_command_buffer, + upload_fence: cmd.upload_fence, + }; + + Ok(CommonInit { + common, + active_reference_count: max_active_reference_pictures as u32, + }) +} diff --git a/src/encoder/h264/api.rs b/src/encoder/h264/api.rs deleted file mode 100644 index e79d33a..0000000 --- a/src/encoder/h264/api.rs +++ /dev/null @@ -1,289 +0,0 @@ -use super::H264Encoder; - -use crate::encoder::dpb::{DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType}; -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::{ColorDescription, EncodedPacket}; -use crate::error::Result; -use crate::PixelForgeError; -use ash::vk; -use ash::vk::TaggedStructure; -use tracing::debug; - -impl H264Encoder { - /// Get the internal input image. - /// - /// This image can be used as a target for `ColorConverter::convert` to avoid - /// an intermediate copy. - pub fn input_image(&self) -> vk::Image { - self.input_image - } - - /// Encode a frame from a GPU image. - /// - /// This accepts a source NV12 image on the GPU and encodes it directly without. - /// any CPU-side data copies. The source image must be in NV12 format with the - /// same dimensions as the encoder configuration, and should be in GENERAL layout. - /// - /// # Panics - /// - /// The encoder will panic at creation time if B-frames are enabled (b_frame_count > 0), - /// as B-frame encoding is not yet supported. - pub fn encode(&mut self, src_image: vk::Image) -> Result> { - let gop_position = self.gop.get_next_frame(); - let display_order = self.input_frame_num; - self.input_frame_num += 1; - - debug!( - "Encoding frame {} from GPU image: type={:?}, poc={}", - display_order, gop_position.frame_type, gop_position.pic_order_cnt - ); - - // Upload from GPU image. - self.upload_from_image(src_image)?; - - // Encode immediately. - let packet = self.encode_current_frame(&gop_position, display_order)?; - - Ok(vec![packet]) - } - - /// Internal method to encode the current frame already uploaded to input_image. - fn encode_current_frame( - &mut self, - gop_position: &GopPosition, - display_order: u64, - ) -> Result { - let is_idr = gop_position.frame_type.is_idr(); - let is_reference = gop_position.is_reference; - let is_b_frame = gop_position.frame_type == GopFrameType::B; - let frame_type = match gop_position.frame_type { - GopFrameType::Idr | GopFrameType::I => crate::encoder::FrameType::I, - GopFrameType::P => crate::encoder::FrameType::P, - GopFrameType::B => crate::encoder::FrameType::B, - }; - - debug!( - "Encoding frame: display_order={}, type={:?}, idr={}, ref={}", - display_order, frame_type, is_idr, is_reference - ); - - if is_idr { - self.frame_num_syntax = 0; - self.idr_pic_id = (self.idr_pic_id + 1) & 1; - // Reset DPB by calling sequence_start with new config for IDR. - let dpb_config = DpbConfig { - dpb_size: self.dpb_slot_count as u32, - max_num_ref_frames: self.config.max_reference_frames, - use_multiple_references: self.config.b_frame_count > 0, - log2_max_frame_num_minus4: 4, - log2_max_pic_order_cnt_lsb_minus4: 4, - ..Default::default() - }; - self.dpb.h264.sequence_start(dpb_config); - self.l0_references.clear(); - self.has_backward_reference = false; - } - - let pic_order_cnt = gop_position.pic_order_cnt; - let frame_num = self.frame_num_syntax; - - let mut encoded_data = Vec::new(); - if is_idr { - encoded_data.extend_from_slice(&self.get_h264_header()?); - self.sps_written = true; - } - - encoded_data.extend_from_slice(&self.encode_frame_internal( - gop_position, - frame_num, - pic_order_cnt, - is_idr, - )?); - - self.encode_frame_num += 1; - if is_reference && !is_b_frame { - self.frame_num_syntax = (self.frame_num_syntax + 1) % 256; - } - - if is_reference { - let pic_type = if is_idr { - PictureType::Idr - } else if is_b_frame { - PictureType::B - } else { - PictureType::P - }; - let pic_info = PictureStartInfo { - frame_id: display_order, - pic_order_cnt, - frame_num, - pic_type, - is_reference, - ..Default::default() - }; - self.dpb.h264.picture_start(pic_info); - self.dpb.h264.picture_end(is_reference); - - // Update reference tracking for the next P-frame. - // The current frame becomes the reference for subsequent frames. - let ref_info = super::ReferenceInfo { - dpb_slot: self.current_dpb_slot, - frame_num, - poc: pic_order_cnt, - }; - self.l0_references.insert(0, ref_info); - // Limit to negotiated max active internal references - while self.l0_references.len() > self.active_reference_count as usize { - self.l0_references.pop(); - } - - // Find a free DPB slot for the next frame. - let used_slots: Vec = self.l0_references.iter().map(|r| r.dpb_slot).collect(); - for i in 0..self.dpb_slot_count as u8 { - if !used_slots.contains(&i) { - self.current_dpb_slot = i; - break; - } - } - } - - Ok(EncodedPacket { - data: encoded_data, - frame_type, - is_key_frame: is_idr, - pts: display_order, - dts: self.encode_frame_num - 1, - }) - } - - /// Flush the encoder and get any remaining packets. - pub fn flush(&mut self) -> Result> { - // No buffered frames in the current implementation. - Ok(Vec::new()) - } - - /// Request that the next frame be an IDR frame. - pub fn request_idr(&mut self) { - self.gop.request_idr(); - } - - /// Retrieve encoded SPS/PPS header data from video session parameters. - /// This uses vkGetEncodedVideoSessionParametersKHR to get the NAL units. - fn get_h264_header(&self) -> Result> { - // H.264-specific get info requesting SPS and PPS. - let mut h264_get_info = vk::VideoEncodeH264SessionParametersGetInfoKHR::default() - .write_std_sps(true) - .write_std_pps(true) - .std_sps_id(0) - .std_pps_id(0); - - let get_info = vk::VideoEncodeSessionParametersGetInfoKHR::default() - .video_session_parameters(self.session_params) - .push(&mut h264_get_info); - - // Some implementations misbehave for a size-only query (pData = NULL), especially with - // 4:4:4 profiles. vk_video_samples avoids that by providing a preallocated buffer. - let mut data = vec![0u8; 4096]; - let mut data_size: usize = data.len(); - let mut h264_feedback = vk::VideoEncodeH264SessionParametersFeedbackInfoKHR::default(); - let mut feedback = - vk::VideoEncodeSessionParametersFeedbackInfoKHR::default().push(&mut h264_feedback); - - let mut attempts = 0; - loop { - attempts += 1; - let result = unsafe { - (self - .video_encode_fn - .fp() - .get_encoded_video_session_parameters_khr)( - self.context.device().handle(), - &get_info, - &mut feedback, - &mut data_size, - data.as_mut_ptr() as *mut std::ffi::c_void, - ) - }; - - match result { - vk::Result::SUCCESS => { - if data_size == 0 { - return Err(PixelForgeError::SessionParametersCreation( - "Encoded parameters size is 0".to_string(), - )); - } - data.truncate(data_size); - return Ok(data); - } - vk::Result::INCOMPLETE if attempts < 3 => { - // Driver indicates the buffer was too small; resize to the reported required - // size (or grow conservatively if the size is not provided). - let new_size = data_size.max(data.len() * 2).max(1); - data.resize(new_size, 0); - data_size = data.len(); - } - err => { - return Err(PixelForgeError::SessionParametersCreation(format!( - "Failed to get encoded parameters: {:?}", - err - ))); - } - } - } - } - - /// Update the color description (VUI parameters) in the encoded stream. - /// - /// This recreates the video session parameters with a new SPS containing the - /// updated VUI color primaries, transfer characteristics, and matrix coefficients. - /// The next encoded frame will be an IDR with the new SPS/PPS prepended. - pub fn set_color_description(&mut self, desc: ColorDescription) -> Result<()> { - // Wait for any in-flight encode to complete before modifying session params. - // Do NOT reset the fence here — submit_encode_and_read_bitstream() resets it - // before queue_submit. Leaving the fence signaled allows consecutive - // set_color_description() calls without deadlock. - unsafe { - self.context - .device() - .wait_for_fences(&[self.encode_fence], true, u64::MAX) - .map_err(|e| { - PixelForgeError::Synchronization(format!( - "Failed to wait for encode fence: {:?}", - e - )) - })?; - } - - // Save old handle so we can destroy it after successful creation. - let old_session_params = self.session_params; - - let new_session_params = self.create_session_params(&desc)?; - - // Destroy old session parameters now that the new ones are created. - unsafe { - (self - .video_queue_fn - .fp() - .destroy_video_session_parameters_khr)( - self.context.device().handle(), - old_session_params, - std::ptr::null(), - ); - } - - self.session_params = new_session_params; - self.config.color_description = Some(desc); - self.sps_written = false; - self.gop.request_idr(); - - debug!( - "H.264 color description updated: primaries={}, transfer={}, matrix={}, full_range={}", - desc.color_primaries, - desc.transfer_characteristics, - desc.matrix_coefficients, - desc.full_range - ); - - Ok(()) - } -} diff --git a/src/encoder/h264/init.rs b/src/encoder/h264/init.rs index cdd3b65..2527704 100644 --- a/src/encoder/h264/init.rs +++ b/src/encoder/h264/init.rs @@ -1,57 +1,33 @@ -use super::{H264Encoder, MB_SIZE}; +use super::{H264, MB_SIZE}; -use crate::encoder::dpb::{DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig}; -use crate::encoder::gop::GopStructure; -use crate::encoder::resources::{ - align_up, allocate_session_memory, clear_input_image, create_bitstream_buffer, - create_command_resources, create_dpb_images, create_image, get_video_format, lcm, - map_bitstream_buffer, query_supported_video_formats, ClearImageParams, - MIN_BITSTREAM_BUFFER_SIZE, +use crate::encoder::codec::{ + build_encoder_common, query_video_caps, CodecEncoder, CommonInitRequest, }; -use crate::encoder::ColorDescription; -use crate::encoder::PixelFormat; -use crate::error::{PixelForgeError, Result}; +use crate::encoder::dpb::{DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig}; +use crate::encoder::resources::MIN_BITSTREAM_BUFFER_SIZE; +use crate::encoder::{ColorDescription, EncodeConfig, PixelFormat}; +use crate::error::Result; use crate::vulkan::VideoContext; use ash::vk; use ash::vk::TaggedStructure; -use std::ptr; use tracing::{debug, info}; -impl H264Encoder { +impl H264 { /// Create a new H.264 encoder. - pub fn new(context: VideoContext, config: crate::encoder::EncodeConfig) -> Result { + pub fn create(context: VideoContext, config: EncodeConfig) -> Result> { // B-frames are not yet supported. - if config.b_frame_count > 0 { - panic!( - "B-frame encoding is not yet supported. Set b_frame_count=0 in encoder config. \ - Got b_frame_count={}", - config.b_frame_count - ); - } - - let width = config.dimensions.width; - let height = config.dimensions.height; + assert!( + config.b_frame_count == 0, + "B-frame encoding is not yet supported; set b_frame_count=0 (got {})", + config.b_frame_count + ); info!( - "Creating H.264 encoder: requested {}x{}, pixel_format={:?}", - width, height, config.pixel_format + "Creating H.264 encoder: {}x{}, pixel_format={:?}", + config.dimensions.width, config.dimensions.height, config.pixel_format ); - // Load video queue extension functions. - let video_queue_fn = - ash::khr::video_queue::Device::load(context.instance(), context.device()); - let video_encode_fn = - ash::khr::video_encode_queue::Device::load(context.instance(), context.device()); - - // Get chroma subsampling from pixel format via `From` impl. - let chroma_subsampling: vk::VideoChromaSubsamplingFlagsKHR = config.pixel_format.into(); - - let luma_bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - let chroma_bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - - // Select H.264 profile based on pixel format. - // - High profile for YUV420 - // - High 4:4:4 Predictive profile for YUV444 + // Select the profile based on chroma subsampling. let profile_idc = match config.pixel_format { PixelFormat::Yuv444 => { ash::vk::native::StdVideoH264ProfileIdc_STD_VIDEO_H264_PROFILE_IDC_HIGH_444_PREDICTIVE @@ -59,541 +35,130 @@ impl H264Encoder { _ => ash::vk::native::StdVideoH264ProfileIdc_STD_VIDEO_H264_PROFILE_IDC_HIGH, }; - // Preferred input format based on pixel format and bit depth. - // Note: the DPB format may differ and must be queried separately. - let preferred_src_format = get_video_format(config.pixel_format, config.bit_depth); + let chroma_subsampling: vk::VideoChromaSubsamplingFlagsKHR = config.pixel_format.into(); + let bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - // Create H.264 encode profile. let mut h264_profile_info = vk::VideoEncodeH264ProfileInfoKHR::default().std_profile_idc(profile_idc); - let profile_info = vk::VideoProfileInfoKHR::default() .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H264) .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(luma_bit_depth) - .chroma_bit_depth(chroma_bit_depth) + .luma_bit_depth(bit_depth) + .chroma_bit_depth(bit_depth) .push(&mut h264_profile_info); - // Query encode capabilities for the selected profile and use them to derive a safe - // coded extent and DPB limits. This mirrors vk_video_samples and avoids device loss - // when the implementation requires larger picture access granularity (commonly for 4:4:4). - let video_queue_instance = - ash::khr::video_queue::Instance::load(context.entry(), context.instance()); - let mut h264_capabilities = vk::VideoEncodeH264CapabilitiesKHR::default(); - let mut encode_capabilities = vk::VideoEncodeCapabilitiesKHR::default(); - - let h264_capabilities_dbg = h264_capabilities; - let encode_capabilities_dbg = encode_capabilities; + // The driver's preferred entropy mode is H.264-specific (some require + // CAVLC for High 4:4:4 Predictive). + let preferred_entropy_cabac = query_preferred_entropy_cabac(&context, &profile_info); + // Query device capabilities with the H.264-specific capability struct + // chained in (required by the driver). + let mut h264_caps = vk::VideoEncodeH264CapabilitiesKHR::default(); + let mut encode_caps = vk::VideoEncodeCapabilitiesKHR::default(); let mut capabilities = vk::VideoCapabilitiesKHR::default() - .push(&mut encode_capabilities) - .push(&mut h264_capabilities); - - let result = unsafe { - (video_queue_instance - .fp() - .get_physical_device_video_capabilities_khr)( - context.physical_device(), - &profile_info, - &mut capabilities, - ) - }; - if result != vk::Result::SUCCESS { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Failed to query Vulkan Video encode capabilities for requested H.264 profile: {:?}", - result - ))); - } - - debug!( - "H.264 capabilities: maxLevelIdc={}, maxSliceCount={}, maxPPictureL0ReferenceCount={}, maxBPictureL0ReferenceCount={}, maxL1ReferenceCount={}, maxTemporalLayerCount={}, prefersGopRemainingFrames={}, requiresGopRemainingFrames={}, stdSyntaxFlags={:#010x}", - h264_capabilities_dbg.max_level_idc, - h264_capabilities_dbg.max_slice_count, - h264_capabilities_dbg.max_p_picture_l0_reference_count, - h264_capabilities_dbg.max_b_picture_l0_reference_count, - h264_capabilities_dbg.max_l1_reference_count, - h264_capabilities_dbg.max_temporal_layer_count, - h264_capabilities_dbg.prefers_gop_remaining_frames, - h264_capabilities_dbg.requires_gop_remaining_frames, - h264_capabilities_dbg.std_syntax_flags.as_raw(), - ); - debug!( - "Encode capabilities: encodeInputPictureGranularity={}x{}, supportedEncodeFeedbackFlags={:#010x}, maxQualityLevels={}", - encode_capabilities_dbg.encode_input_picture_granularity.width, - encode_capabilities_dbg.encode_input_picture_granularity.height, - encode_capabilities_dbg.supported_encode_feedback_flags.as_raw(), - encode_capabilities_dbg.max_quality_levels, - ); - debug!( - "Video capabilities: flags={:#010x}, minBitstreamBufferOffsetAlignment={}, minBitstreamBufferSizeAlignment={}, minCodedExtent={}x{}, maxCodedExtent={}x{}, maxDpbSlots={}, maxActiveReferencePictures={}, pictureAccessGranularity={}x{}", - capabilities.flags.as_raw(), - capabilities.min_bitstream_buffer_offset_alignment, - capabilities.min_bitstream_buffer_size_alignment, - capabilities.min_coded_extent.width, - capabilities.min_coded_extent.height, - capabilities.max_coded_extent.width, - capabilities.max_coded_extent.height, - capabilities.max_dpb_slots, - capabilities.max_active_reference_pictures, - capabilities.picture_access_granularity.width, - capabilities.picture_access_granularity.height, - ); - - // Query quality level properties to get the driver's preferred settings. - let video_encode_instance = - ash::khr::video_encode_queue::Instance::load(context.entry(), context.instance()); - let mut h264_quality_level_properties = - vk::VideoEncodeH264QualityLevelPropertiesKHR::default(); - let mut quality_level_properties = vk::VideoEncodeQualityLevelPropertiesKHR::default() - .push(&mut h264_quality_level_properties); - let quality_level_info = vk::PhysicalDeviceVideoEncodeQualityLevelInfoKHR::default() - .video_profile(&profile_info) - .quality_level(0); - let ql_result = unsafe { - (video_encode_instance - .fp() - .get_physical_device_video_encode_quality_level_properties_khr)( - context.physical_device(), - &quality_level_info, - &mut quality_level_properties, - ) - }; - let preferred_entropy_cabac = if ql_result == vk::Result::SUCCESS { - debug!( - "H.264 quality level 0: preferredStdEntropyCodingModeFlag={}, preferredMaxL0ReferenceCount={}, preferredMaxL1ReferenceCount={}", - h264_quality_level_properties.preferred_std_entropy_coding_mode_flag, - h264_quality_level_properties.preferred_max_l0_reference_count, - h264_quality_level_properties.preferred_max_l1_reference_count, - ); - h264_quality_level_properties.preferred_std_entropy_coding_mode_flag != 0 - } else { - debug!( - "Failed to query quality level properties: {:?}, defaulting to CABAC", - ql_result - ); - true - }; - - let gran_w = capabilities.picture_access_granularity.width.max(1); - let gran_h = capabilities.picture_access_granularity.height.max(1); - let align_w = lcm(MB_SIZE, gran_w); - let align_h = lcm(MB_SIZE, gran_h); - - let mut aligned_width = align_up(width, align_w); - let mut aligned_height = align_up(height, align_h); - - aligned_width = align_up( - aligned_width.max(capabilities.min_coded_extent.width), - align_w, - ); - aligned_height = align_up( - aligned_height.max(capabilities.min_coded_extent.height), - align_h, - ); - - if aligned_width > capabilities.max_coded_extent.width - || aligned_height > capabilities.max_coded_extent.height - { - return Err(PixelForgeError::InvalidInput(format!( - "Requested coded extent {}x{} (aligned to {}x{} with granularity {}x{}) exceeds device max {}x{} for this profile", - width, - height, - aligned_width, - aligned_height, - gran_w, - gran_h, - capabilities.max_coded_extent.width, - capabilities.max_coded_extent.height - ))); - } - - info!( - "Using coded extent {}x{} (granularity {}x{}, min {}x{}, max {}x{})", - aligned_width, - aligned_height, - gran_w, - gran_h, - capabilities.min_coded_extent.width, - capabilities.min_coded_extent.height, - capabilities.max_coded_extent.width, - capabilities.max_coded_extent.height - ); - - // Query supported formats separately for SRC and DPB usage (vk_video_samples-style). - // Using an unsupported DPB format is a common cause of device loss, especially for 4:4:4. - let supported_src_formats = query_supported_video_formats( - &context, - &profile_info, - vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR, - )?; - let supported_dpb_formats = query_supported_video_formats( - &context, - &profile_info, - vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR, - )?; - - if supported_src_formats.is_empty() { - return Err(PixelForgeError::NoSuitableDevice( - "No supported Vulkan Video SRC formats returned for this profile".to_string(), - )); - } - info!("Supported SRC formats: {:?}", supported_src_formats); - if supported_dpb_formats.is_empty() { - return Err(PixelForgeError::NoSuitableDevice( - "No supported Vulkan Video DPB formats returned for this profile".to_string(), - )); - } - info!("Supported DPB formats: {:?}", supported_dpb_formats); - - // For input uploads, we currently require the preferred 2-plane formats. - let picture_format = if supported_src_formats.contains(&preferred_src_format) { - preferred_src_format - } else { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Preferred input format {:?} is not supported for VIDEO_ENCODE_SRC_KHR. Supported: {:?}", - preferred_src_format, supported_src_formats - ))); - }; - - // DPB format can differ from the input format; prefer matching when possible. - let reference_picture_format = supported_dpb_formats - .iter() - .copied() - .find(|f| *f == picture_format) - .unwrap_or(supported_dpb_formats[0]); - - debug!( - "Selected Vulkan Video formats: picture_format={:?}, reference_picture_format={:?} (preferred_src={:?})", - picture_format, - reference_picture_format, - preferred_src_format - ); - - // Create video session. - // Use the STD header version reported by the driver capabilities. - let std_header_version = capabilities.std_header_version; - - // Calculate required DPB slots and active references. - let max_dpb_slots_supported = capabilities.max_dpb_slots as usize; - let max_active_reference_pictures_supported = - capabilities.max_active_reference_pictures as usize; - - if max_dpb_slots_supported < 2 { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Device reports max_dpb_slots={} for this profile; need at least 2", - max_dpb_slots_supported - ))); - } - - // Target number of active reference pictures. - // H.264 L0 list can theoretically handle more, but we clamp to config and device limits. - let mut target_active_refs = (config.max_reference_frames as usize) - .min(max_active_reference_pictures_supported) - .min(32); - - // Ensure we have at least 1 active ref if supported. - if target_active_refs < 1 && max_active_reference_pictures_supported >= 1 { - target_active_refs = 1; - } - - // Calculate required DPB slots. - let requested_dpb_slots = if config.b_frame_count > 0 { - // For B-frames: Active Refs + B-frame buffer + Setup slot + Margin - target_active_refs + config.b_frame_count as usize + 2 - } else { - // For P-frames: Active Refs + Setup slot - // We use target_active_refs + 1 (setup), and maybe +1 for safety if parallel operations occur. - target_active_refs + 1 - }; - - let dpb_slot_count = requested_dpb_slots - .min(max_dpb_slots_supported) - .min(crate::encoder::dpb::MAX_DPB_SLOTS); - - // Finalize active reference count based on what we actually allocated. - // We need at least 1 slot for the current setup frame. - let max_active_reference_pictures = - target_active_refs.min(dpb_slot_count.saturating_sub(1)); // Ensure room for setup - - debug!( - "Allocating {} DPB slots (requested {}, device max {}), max_active_reference_pictures={} (target {}, device max {})", - dpb_slot_count, - requested_dpb_slots, - max_dpb_slots_supported, - max_active_reference_pictures, - target_active_refs, - max_active_reference_pictures_supported - ); - - let encode_queue_family = context.video_encode_queue_family().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue family available".to_string()) + .push(&mut encode_caps) + .push(&mut h264_caps); + let caps = query_video_caps(&context, &profile_info, &mut capabilities)?; + + let init = build_encoder_common(&CommonInitRequest { + context: &context, + config: &config, + profile_info: &profile_info, + caps: &caps, + align_unit: MB_SIZE, + max_active_refs_cap: 32, + bitstream_buffer_size: MIN_BITSTREAM_BUFFER_SIZE, + allow_layered_dpb: true, })?; + let active_reference_count = init.active_reference_count; + let common = init.common; - let session_create_info = vk::VideoSessionCreateInfoKHR::default() - .queue_family_index(encode_queue_family) - .flags(vk::VideoSessionCreateFlagsKHR::empty()) - .video_profile(&profile_info) - .picture_format(picture_format) - .max_coded_extent(vk::Extent2D { - width: aligned_width, - height: aligned_height, - }) - .reference_picture_format(reference_picture_format) - .max_dpb_slots(dpb_slot_count as u32) - .max_active_reference_pictures(max_active_reference_pictures as u32) - .std_header_version(&std_header_version); - - let mut session = vk::VideoSessionKHR::null(); - let result = unsafe { - (video_queue_fn.fp().create_video_session_khr)( - context.device().handle(), - &session_create_info, - ptr::null(), - &mut session, - ) - }; - if result != vk::Result::SUCCESS { - return Err(PixelForgeError::VideoSessionCreation(format!( - "{:?}", - result - ))); - } - - // Query and allocate session memory. - let session_memory = allocate_session_memory(&context, session, &video_queue_fn)?; - - // Create SPS and PPS. - let pic_width_in_mbs = aligned_width / 16; - let pic_height_in_map_units = aligned_height / 16; - - // Cropping offsets are expressed in units that depend on chroma subsampling. - // For progressive frames (frame_mbs_only_flag=1): - // - 4:2:0 => crop_unit_x=2, crop_unit_y=2 - // - 4:4:4 => crop_unit_x=1, crop_unit_y=1 - let (crop_unit_x, crop_unit_y) = match config.pixel_format { - PixelFormat::Yuv420 => (2u32, 2u32), - PixelFormat::Yuv444 => (1u32, 1u32), - _ => { - return Err(PixelForgeError::InvalidInput(format!( - "Unsupported pixel format for H.264: {:?}", - config.pixel_format - ))); - } - }; - - let coded_width = pic_width_in_mbs * 16; - let coded_height = pic_height_in_map_units * 16; - let crop_right_pixels = coded_width.saturating_sub(width); - let crop_bottom_pixels = coded_height.saturating_sub(height); - - if !crop_right_pixels.is_multiple_of(crop_unit_x) { - return Err(PixelForgeError::InvalidInput(format!( - "Width {} is not representable for {:?} with coded width {} (crop_unit_x={}): crop delta {} must be divisible by crop unit", - width, config.pixel_format, coded_width, crop_unit_x, crop_right_pixels - ))); - } - if !crop_bottom_pixels.is_multiple_of(crop_unit_y) { - return Err(PixelForgeError::InvalidInput(format!( - "Height {} is not representable for {:?} with coded height {} (crop_unit_y={}): crop delta {} must be divisible by crop unit", - height, config.pixel_format, coded_height, crop_unit_y, crop_bottom_pixels - ))); - } - - let color_desc = config - .color_description - .unwrap_or(ColorDescription::bt709()); - - // Create profile info for images/buffers. - let mut h264_profile_for_resources = - vk::VideoEncodeH264ProfileInfoKHR::default().std_profile_idc(profile_idc); - let profile_for_resources = vk::VideoProfileInfoKHR::default() - .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H264) - .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(luma_bit_depth) - .chroma_bit_depth(chroma_bit_depth) - .push(&mut h264_profile_for_resources); - - // Create input image. - let (input_image, input_image_memory, input_image_view) = create_image( - &context, - aligned_width, - aligned_height, - picture_format, - false, - &profile_for_resources, - )?; - - // Determine DPB mode: use layered DPB when the driver does not advertise - // support for separate reference images (required for AMD RADV). - let supports_separate_dpb = capabilities - .flags - .contains(vk::VideoCapabilityFlagsKHR::SEPARATE_REFERENCE_IMAGES); - let use_layered_dpb = !supports_separate_dpb; - if use_layered_dpb { - info!("Using layered DPB (driver does not support separate reference images)"); - } - - // Create DPB images. - let (dpb_images, dpb_image_memories, dpb_image_views) = create_dpb_images( - &context, - aligned_width, - aligned_height, - reference_picture_format, - dpb_slot_count, - &profile_for_resources, - use_layered_dpb, - )?; - - // Create bitstream buffer. - let (bitstream_buffer, bitstream_buffer_memory) = - create_bitstream_buffer(&context, MIN_BITSTREAM_BUFFER_SIZE, &profile_for_resources)?; - - // Persistently map the bitstream buffer to avoid per-frame map/unmap overhead. - let bitstream_buffer_ptr = - map_bitstream_buffer(&context, bitstream_buffer_memory, MIN_BITSTREAM_BUFFER_SIZE)?; - - // Create command pool, buffers, and fences. - // Use the transfer queue family for upload commands when the encode queue - // doesn't support transfer operations (AMD RADV). - let upload_queue_family = context.transfer_queue_family(); - let cmd_resources = - create_command_resources(&context, encode_queue_family, upload_queue_family)?; - let command_pool = cmd_resources.command_pool; - let upload_command_pool = cmd_resources.upload_command_pool; - let upload_command_buffer = cmd_resources.upload_command_buffer; - let encode_command_buffer = cmd_resources.encode_command_buffer; - let upload_fence = cmd_resources.upload_fence; - let encode_fence = cmd_resources.encode_fence; - - // Clear the input image so padding between user dimensions and the - // aligned coded extent is zero-initialized. - clear_input_image( - &context, - &ClearImageParams { - command_buffer: upload_command_buffer, - fence: upload_fence, - queue: context.transfer_queue(), - image: input_image, - width: aligned_width, - height: aligned_height, - pixel_format: config.pixel_format, - bit_depth: config.bit_depth, - }, - )?; - - // Create query pool. - let mut h264_profile_info_query = - vk::VideoEncodeH264ProfileInfoKHR::default().std_profile_idc(profile_idc); - - let mut profile_info_query = vk::VideoProfileInfoKHR::default() - .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H264) - .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(luma_bit_depth) - .chroma_bit_depth(chroma_bit_depth) - .push(&mut h264_profile_info_query); - - let mut encode_feedback_create = vk::QueryPoolVideoEncodeFeedbackCreateInfoKHR::default() - .encode_feedback_flags( - vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BUFFER_OFFSET - | vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN, - ); - - let query_pool_create_info = unsafe { - vk::QueryPoolCreateInfo::default() - .query_type(vk::QueryType::VIDEO_ENCODE_FEEDBACK_KHR) - .query_count(1) - .extend(&mut profile_info_query) - .push(&mut encode_feedback_create) - }; - - let query_pool = unsafe { - context - .device() - .create_query_pool(&query_pool_create_info, None) - } - .map_err(|e| PixelForgeError::QueryPool(e.to_string()))?; - - // Create DPB and GOP structure. - // The DPB size should match the actual number of allocated DPB slots. + // H.264 reference marking bookkeeping. let mut dpb = DecodedPictureBuffer::new(); - let dpb_config = DpbConfig { - dpb_size: dpb_slot_count as u32, + dpb.h264.sequence_start(DpbConfig { + dpb_size: common.dpb_slot_count as u32, max_num_ref_frames: if config.b_frame_count > 0 { 2 } else { 1 }, use_multiple_references: config.b_frame_count > 0, max_long_term_refs: 0, - log2_max_frame_num_minus4: 4, // max_frame_num = 256 - log2_max_pic_order_cnt_lsb_minus4: 4, // max_poc_lsb = 256 + log2_max_frame_num_minus4: 4, + log2_max_pic_order_cnt_lsb_minus4: 4, num_temporal_layers: 1, - }; - dpb.h264.sequence_start(dpb_config); - - let mut gop = if config.b_frame_count > 0 { - GopStructure::new(config.gop_size, config.b_frame_count, config.gop_size) - } else { - GopStructure::new_ip_only(config.gop_size) - }; + }); - // Set GOP parameters to match SPS values. - // log2_max_frame_num_minus4 = 4, so max_frame_num = 2^8 = 256 - gop.set_max_frame_num(4); - // log2_max_pic_order_cnt_lsb_minus4 = 4, so max_poc_lsb = 2^8 = 256 - gop.set_max_poc_lsb(4); - - info!("H.264 encoder created successfully"); - - let mut encoder = Self { - context, - config: config.clone(), + let codec = H264 { dpb, - gop, - aligned_width, - aligned_height, - video_queue_fn, - video_encode_fn, - session, - session_params: vk::VideoSessionParametersKHR::null(), - session_memory, - input_frame_num: 0, - encode_frame_num: 0, frame_num_syntax: 0, idr_pic_id: 0, - input_image, - input_image_memory, - input_image_view, - input_image_layout: vk::ImageLayout::VIDEO_ENCODE_SRC_KHR, - dpb_images, - dpb_image_memories, - dpb_image_views, - dpb_slot_count, - use_layered_dpb, - dpb_slot_active: vec![false; dpb_slot_count], - current_dpb_slot: 0, - l0_references: Vec::new(), - active_reference_count: max_active_reference_pictures as u32, - bitstream_buffer, - bitstream_buffer_memory, - bitstream_buffer_ptr, - command_pool, - upload_command_pool, - upload_command_buffer, - upload_fence, - encode_command_buffer, - encode_fence, - query_pool, - sps_written: false, - // has_reference: false, // removed - // reference_frame_num: 0, // removed - // reference_poc: 0, // removed has_backward_reference: false, backward_reference_frame_num: 0, backward_reference_poc: 0, backward_reference_dpb_slot: 2, + l0_references: Vec::new(), + active_reference_count, + pending_unmark_frame_nums: Vec::new(), profile_idc, preferred_entropy_cabac, }; - encoder.session_params = encoder.create_session_params(&color_desc)?; + let mut encoder = CodecEncoder { common, codec }; + let color_desc = config + .color_description + .unwrap_or(ColorDescription::bt709()); + encoder.common.session_params = encoder + .codec + .build_session_params(&encoder.common, &color_desc)?; + info!("H.264 encoder created successfully"); Ok(encoder) } + + /// Profile IDC (used when rebuilding parameter sets). + pub(super) fn profile_idc(&self) -> u32 { + self.profile_idc + } + + /// Preferred entropy coding mode (CABAC if true, CAVLC otherwise). + pub(super) fn preferred_entropy_cabac(&self) -> bool { + self.preferred_entropy_cabac + } + + /// Negotiated active reference count. + pub(super) fn active_reference_count(&self) -> u32 { + self.active_reference_count + } +} + +/// Query the driver's preferred H.264 entropy coding mode for quality level 0. +/// Defaults to CABAC if the query is unsupported. +fn query_preferred_entropy_cabac( + context: &VideoContext, + profile_info: &vk::VideoProfileInfoKHR, +) -> bool { + let video_encode_instance = + ash::khr::video_encode_queue::Instance::load(context.entry(), context.instance()); + let mut h264_quality_level_properties = vk::VideoEncodeH264QualityLevelPropertiesKHR::default(); + let mut quality_level_properties = vk::VideoEncodeQualityLevelPropertiesKHR::default() + .push(&mut h264_quality_level_properties); + let quality_level_info = vk::PhysicalDeviceVideoEncodeQualityLevelInfoKHR::default() + .video_profile(profile_info) + .quality_level(0); + let result = unsafe { + (video_encode_instance + .fp() + .get_physical_device_video_encode_quality_level_properties_khr)( + context.physical_device(), + &quality_level_info, + &mut quality_level_properties, + ) + }; + if result == vk::Result::SUCCESS { + debug!( + "H.264 quality level 0: preferredStdEntropyCodingModeFlag={}", + h264_quality_level_properties.preferred_std_entropy_coding_mode_flag + ); + h264_quality_level_properties.preferred_std_entropy_coding_mode_flag != 0 + } else { + debug!("Quality level query failed ({result:?}); defaulting to CABAC"); + true + } } diff --git a/src/encoder/h264/mod.rs b/src/encoder/h264/mod.rs index 3d339d3..528892d 100644 --- a/src/encoder/h264/mod.rs +++ b/src/encoder/h264/mod.rs @@ -1,24 +1,22 @@ -//! H.264 encoder implementation using Vulkan Video. +//! H.264/AVC codec: the differences from the generic encoder. //! -//! This module implements H.264 video encoding using Vulkan Video extensions. +//! The shared session/DPB/pipeline machinery lives in `crate::encoder::codec`; +//! this folder holds only what is specific to H.264: its reference-picture +//! tracking and syntax counters (`H264`), the per-frame StdVideo* graph +//! (`record`), and the SPS/PPS generation (`session_params`). -mod api; -mod encode; mod init; +mod record; mod session_params; -use ash::vk; -use tracing::debug; - -use crate::encoder::resources::{ - destroy_encoder_resources, upload_image_to_input, EncoderResources, UploadParams, +use crate::encoder::codec::{EncoderCommon, FramePlan, PictureSetup, VideoCodec}; +use crate::encoder::dpb::{ + DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType, }; +use crate::encoder::pipeline::EncodeFuture; +use crate::encoder::ColorDescription; use crate::error::Result; - -use crate::encoder::dpb::DecodedPictureBuffer; -use crate::encoder::gop::GopStructure; -use crate::encoder::EncodeConfig; -use crate::vulkan::VideoContext; +use ash::vk; /// H.264 macroblock size in pixels. pub const MB_SIZE: u32 = 16; @@ -28,159 +26,171 @@ pub(crate) struct ReferenceInfo { pub dpb_slot: u8, pub frame_num: u32, pub poc: i32, + /// Display order (`pts`) of this reference, for reference frame invalidation. + pub display_order: u64, } -/// H.264 encoder. -pub struct H264Encoder { - context: VideoContext, - config: EncodeConfig, +/// H.264-specific encoder state (everything the generic encoder doesn't own). +pub(crate) struct H264 { + /// Decoded-picture-buffer bookkeeping (reference marking). dpb: DecodedPictureBuffer, - gop: GopStructure, - - /// Aligned width (macroblock + granularity aligned). - aligned_width: u32, - /// Aligned height (macroblock + granularity aligned). - aligned_height: u32, - - // Video session. - video_queue_fn: ash::khr::video_queue::Device, - video_encode_fn: ash::khr::video_encode_queue::Device, - session: vk::VideoSessionKHR, - session_params: vk::VideoSessionParametersKHR, - session_memory: Vec, - - // Frame counters. - input_frame_num: u64, - encode_frame_num: u64, + /// `frame_num` syntax element (mod `max_frame_num`). frame_num_syntax: u32, + /// IDR picture id, toggled each IDR. idr_pic_id: u32, - - // Resources - input_image: vk::Image, - input_image_memory: vk::DeviceMemory, - input_image_view: vk::ImageView, - /// Current Vulkan image layout of `input_image` (tracked to avoid UB when transitioning). - input_image_layout: vk::ImageLayout, - /// DPB images (up to MAX_DPB_SLOTS for B-frame and long-term reference support). - dpb_images: Vec, - dpb_image_memories: Vec, - dpb_image_views: Vec, - /// Number of DPB slots allocated. - dpb_slot_count: usize, - /// Whether the DPB uses a single layered image (true) or separate images (false). - use_layered_dpb: bool, - /// Tracks which DPB slots have been activated (used at least once). - dpb_slot_active: Vec, - bitstream_buffer: vk::Buffer, - bitstream_buffer_memory: vk::DeviceMemory, - /// Persistently mapped pointer to the bitstream buffer (avoids per-frame map/unmap). - bitstream_buffer_ptr: *mut u8, - - // Command resources. - command_pool: vk::CommandPool, - upload_command_pool: vk::CommandPool, - upload_command_buffer: vk::CommandBuffer, - upload_fence: vk::Fence, - encode_command_buffer: vk::CommandBuffer, - encode_fence: vk::Fence, - query_pool: vk::QueryPool, - - // SPS/PPS written flag. - sps_written: bool, - - // Reference picture tracking. - /// Whether we have a backward reference (for B-frames, L1). + /// Whether an L1 (backward) reference is available, for B-frames. has_backward_reference: bool, - /// Frame number of the L1 (backward) reference picture (for B-frames). backward_reference_frame_num: u32, - /// POC of the L1 (backward) reference picture. backward_reference_poc: i32, - /// DPB slot for L1 (backward) reference. backward_reference_dpb_slot: u8, - /// Current DPB slot to use for setup (the reconstructed picture). - current_dpb_slot: u8, - /// Active L0 reference pictures (for P-frames). Ordered from most recent to oldest. + /// Active L0 references, most-recent first. l0_references: Vec, - /// Number of active reference frames (as configured/negotiated). + /// Negotiated number of active references. active_reference_count: u32, - /// H.264 profile IDC (cached from initialization for session parameter recreation). + /// `frame_num`s of references dropped by reference frame invalidation that + /// still need to be marked unused-for-reference (via MMCO) in the next + /// P-frame's slice header, so the decoder's DPB matches the encoder's. + pending_unmark_frame_nums: Vec, + /// Profile IDC (cached for parameter-set recreation). profile_idc: u32, - /// Whether CABAC entropy coding is preferred (cached from quality level query). + /// Whether CABAC entropy coding is preferred (from quality-level query). preferred_entropy_cabac: bool, } -impl H264Encoder { - /// Upload input frame from a GPU image. - /// - /// This copies from a source NV12 image directly to the encoder's input image, - /// avoiding any CPU-side data copies. The source image must be in NV12 format - /// with the same dimensions as the encoder configuration. The source image - /// should be in GENERAL layout. - fn upload_from_image(&mut self, src_image: vk::Image) -> Result<()> { - if src_image == self.input_image { - debug!("Source image is the encoder's input image, skipping upload copy"); - return Ok(()); +impl VideoCodec for H264 { + fn begin_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + if plan.is_idr() { + self.frame_num_syntax = 0; + self.idr_pic_id = (self.idr_pic_id + 1) & 1; + // Reset the DPB for the new coded video sequence. + let dpb_config = DpbConfig { + dpb_size: common.dpb_slot_count as u32, + max_num_ref_frames: common.config.max_reference_frames, + use_multiple_references: common.config.b_frame_count > 0, + log2_max_frame_num_minus4: 4, + log2_max_pic_order_cnt_lsb_minus4: 4, + ..Default::default() + }; + self.dpb.h264.sequence_start(dpb_config); + self.l0_references.clear(); + self.has_backward_reference = false; + // An IDR resets the decoder's whole DPB, so any deferred + // unused-reference markings are moot. + self.pending_unmark_frame_nums.clear(); } - let params = UploadParams { - upload_command_buffer: self.upload_command_buffer, - upload_fence: self.upload_fence, - src_image, - dst_image: self.input_image, - width: self.config.dimensions.width, - height: self.config.dimensions.height, - pixel_format: self.config.pixel_format, - input_image_layout: self.input_image_layout, - upload_queue: self.context.transfer_queue(), + let header = if plan.is_idr() { + Some(self.build_header(common)?) + } else { + None }; - upload_image_to_input(&self.context, ¶ms)?; - - // Update tracked layout. - self.input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + Ok(PictureSetup { + frame_type: plan.frame_type(), + header, + }) + } - Ok(()) + fn record_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + self.record(common, plan) } -} -// SAFETY: The raw pointer bitstream_buffer_ptr is only used within the encoder's -// thread and is properly synchronized via Vulkan fences before access -unsafe impl Send for H264Encoder {} - -impl Drop for H264Encoder { - fn drop(&mut self) { - unsafe { - // Wait on the queues used by the encoder rather than stalling - // the entire device. - let _ = self - .context - .device() - .queue_wait_idle(self.context.transfer_queue()); - if let Some(q) = self.context.video_encode_queue() { - let _ = self.context.device().queue_wait_idle(q); - } - destroy_encoder_resources( - self.context.device(), - &self.video_queue_fn, - &EncoderResources { - query_pool: self.query_pool, - upload_fence: self.upload_fence, - encode_fence: self.encode_fence, - command_pool: self.command_pool, - upload_command_pool: self.upload_command_pool, - bitstream_buffer: self.bitstream_buffer, - bitstream_buffer_memory: self.bitstream_buffer_memory, - input_image: self.input_image, - input_image_memory: self.input_image_memory, - input_image_view: self.input_image_view, - dpb_images: &self.dpb_images, - dpb_image_memories: &self.dpb_image_memories, - dpb_image_views: &self.dpb_image_views, - session: self.session, - session_params: self.session_params, - session_memory: &self.session_memory, + fn end_picture(&mut self, common: &mut EncoderCommon, plan: &FramePlan) { + // `frame_num` carried by this frame (pre-increment), reused below for the + // reference entry even after the syntax counter advances. + let frame_num = self.frame_num_syntax; + let pic_order_cnt = plan.pic_order_cnt(); + + if plan.is_reference() && !plan.is_b_frame() { + self.frame_num_syntax = (frame_num + 1) % 256; + } + + if plan.is_reference() { + let pic_type = if plan.is_idr() { + PictureType::Idr + } else if plan.is_b_frame() { + PictureType::B + } else { + PictureType::P + }; + let pic_info = PictureStartInfo { + frame_id: plan.display_order, + pic_order_cnt, + frame_num, + pic_type, + is_reference: true, + ..Default::default() + }; + self.dpb.h264.picture_start(pic_info); + self.dpb.h264.picture_end(true); + + // The current frame becomes a reference for subsequent frames. + self.l0_references.insert( + 0, + ReferenceInfo { + dpb_slot: common.current_dpb_slot, + frame_num, + poc: pic_order_cnt, + display_order: plan.display_order, }, ); + while self.l0_references.len() > self.active_reference_count as usize { + self.l0_references.pop(); + } + + // Pick the next free DPB slot for the reconstructed picture. + let used_slots: Vec = self.l0_references.iter().map(|r| r.dpb_slot).collect(); + for i in 0..common.dpb_slot_count as u8 { + if !used_slots.contains(&i) { + common.current_dpb_slot = i; + break; + } + } } } + + fn invalidate_references( + &mut self, + _common: &mut EncoderCommon, + first_lost_display_order: u64, + ) -> bool { + // Keep references older than the loss; every newer one is transitively + // undecodable on the client. H.264 maintains the DPB with an implicit + // sliding window, which would keep those newer pictures and desync the + // decoder's reference list from ours — so each dropped reference is + // queued to be explicitly marked unused-for-reference (MMCO) in the next + // P-frame (see `record`). + let mut survivors = Vec::with_capacity(self.l0_references.len()); + for r in self.l0_references.drain(..) { + if r.display_order < first_lost_display_order { + survivors.push(r); + } else { + self.pending_unmark_frame_nums.push(r.frame_num); + } + } + self.l0_references = survivors; + // A backward (L1) reference may also be tainted; only B-frames use it. + self.has_backward_reference = false; + !self.l0_references.is_empty() + } + + fn create_session_params( + &self, + common: &EncoderCommon, + desc: &ColorDescription, + ) -> Result { + self.build_session_params(common, desc) + } + + fn invalidate_header_cache(&mut self) { + // H.264 regenerates SPS/PPS on every IDR, so there is nothing cached. + } } diff --git a/src/encoder/h264/encode.rs b/src/encoder/h264/record.rs similarity index 58% rename from src/encoder/h264/encode.rs rename to src/encoder/h264/record.rs index e03de2c..4eef7e9 100644 --- a/src/encoder/h264/encode.rs +++ b/src/encoder/h264/record.rs @@ -1,81 +1,99 @@ -use super::H264Encoder; +//! H.264 per-frame encode recording: builds the StdVideo* graph for one frame +//! and submits it. This is the part that genuinely differs from the other codecs. -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::resources::{ - prepare_encode_command_buffer, record_dpb_barriers, submit_encode_and_read_bitstream, - MIN_BITSTREAM_BUFFER_SIZE, -}; +use super::H264; + +use crate::encoder::codec::{EncoderCommon, FramePlan, RateControlPlan}; +use crate::encoder::pipeline::EncodeFuture; +use crate::encoder::resources::{prepare_encode_command_buffer, record_dpb_barriers}; use crate::error::{PixelForgeError, Result}; use ash::vk; use ash::vk::TaggedStructure; use tracing::debug; -impl H264Encoder { - pub(super) fn encode_frame_internal( +impl H264 { + pub(super) fn record( &mut self, - gop_position: &GopPosition, - frame_num: u32, - pic_order_cnt: i32, - is_idr: bool, - ) -> Result> { - let is_b_frame = gop_position.frame_type == GopFrameType::B; - let is_reference = gop_position.is_reference; + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + let is_idr = plan.is_idr(); + let is_reference = plan.is_reference(); + let is_b_frame = plan.is_b_frame(); + let pic_order_cnt = plan.pic_order_cnt(); + let frame_num = self.frame_num_syntax; + + let slot = common.pipeline.current(); + let command_buffer = slot.encode_command_buffer; + let query_pool = slot.query_pool; + let bitstream_buffer = slot.bitstream_buffer; + let bitstream_buffer_size = slot.bitstream_buffer_size; + let input_image_view = slot.input_image_view; debug!( - "encode_frame_internal: frame_num={}, poc={}, is_idr={}, refs_len={}, current_dpb_slot={}", + "h264 record: frame_num={}, poc={}, is_idr={}, refs_len={}, cur_slot={}", frame_num, pic_order_cnt, is_idr, self.l0_references.len(), - self.current_dpb_slot + common.current_dpb_slot ); - // Rate control setup. - let (rc_mode, average_bitrate, max_bitrate, qp) = match self.config.rate_control_mode { - crate::encoder::RateControlMode::Cqp | crate::encoder::RateControlMode::Disabled => ( - vk::VideoEncodeRateControlModeFlagsKHR::DISABLED, - 0, - 0, - self.config.quality_level as i32, - ), - crate::encoder::RateControlMode::Cbr => ( - vk::VideoEncodeRateControlModeFlagsKHR::CBR, - self.config.target_bitrate, - self.config.target_bitrate, - 26, // Default QP for rate control to adjust from - ), - crate::encoder::RateControlMode::Vbr => ( - vk::VideoEncodeRateControlModeFlagsKHR::VBR, - self.config.target_bitrate, - self.config.max_bitrate, - 26, // Default QP for rate control to adjust from - ), - }; + let rc = RateControlPlan::new(&common.config, 26); + + // Reference frame invalidation: emit MMCO "unmark short-term" operations + // for references dropped since the last frame, so the decoder evicts the + // same pictures we did and its default reference list re-anchors to the + // surviving reference. Only valid on a non-IDR reference picture (an IDR + // resets the DPB on its own). Kept alive for the whole `record` call so + // the pointer handed to the encoder stays valid. + let ref_pic_marking_ops: Vec = + if !is_idr && is_reference && !self.pending_unmark_frame_nums.is_empty() { + self.pending_unmark_frame_nums + .iter() + .map(|&tainted_frame_num| { + // CurrPicNum is frame_num for a (non-field) frame. A + // short-term ref's PicNumX is its FrameNumWrap; a tainted + // frame_num greater than the current one has wrapped past + // MaxFrameNum (256 here, log2_max_frame_num_minus4 = 4). + let curr = frame_num as i32; + let mut pic_num_x = tainted_frame_num as i32; + if pic_num_x > curr { + pic_num_x -= 256; + } + let diff = (curr - pic_num_x).max(1); + ash::vk::native::StdVideoEncodeH264RefPicMarkingEntry { + memory_management_control_operation: + ash::vk::native::StdVideoH264MemMgmtControlOp_STD_VIDEO_H264_MEM_MGMT_CONTROL_OP_UNMARK_SHORT_TERM, + difference_of_pic_nums_minus1: (diff - 1) as u16, + long_term_pic_num: 0, + long_term_frame_idx: 0, + max_long_term_frame_idx_plus1: 0, + } + }) + .collect() + } else { + Vec::new() + }; + let use_adaptive_marking = !ref_pic_marking_ops.is_empty(); - // Prepare command buffer for recording. + // Prepare command buffer and transition DPB images for encode. unsafe { - prepare_encode_command_buffer( - self.context.device(), - self.encode_command_buffer, - self.query_pool, - )?; + prepare_encode_command_buffer(common.device(), command_buffer, query_pool)?; } - - // Transition DPB images for encode. let ref_dpb_slots: Vec = self.l0_references.iter().map(|r| r.dpb_slot).collect(); unsafe { record_dpb_barriers( - self.context.device(), - self.encode_command_buffer, - &self.dpb_images, - self.use_layered_dpb, - self.current_dpb_slot, + common.device(), + command_buffer, + &common.dpb_images, + common.use_layered_dpb, + common.current_dpb_slot, &ref_dpb_slots, - self.dpb_slot_active[self.current_dpb_slot as usize], + common.dpb_slot_active[common.current_dpb_slot as usize], ); } - // Set up H.264 specific encode info. let slice_type = if is_idr { ash::vk::native::StdVideoH264SliceType_STD_VIDEO_H264_SLICE_TYPE_I } else if is_b_frame { @@ -107,9 +125,9 @@ impl H264Encoder { ), }; - let slice_qp_delta = match self.config.rate_control_mode { + let slice_qp_delta = match common.config.rate_control_mode { crate::encoder::RateControlMode::Cqp | crate::encoder::RateControlMode::Disabled => { - (self.config.quality_level as i32 - 26) as i8 + (common.config.quality_level as i32 - 26) as i8 } _ => 0, }; @@ -128,7 +146,6 @@ impl H264Encoder { pWeightTable: std::ptr::null(), }; - // Build StdVideoEncodeH264PictureInfo. let picture_info_flags = ash::vk::native::StdVideoEncodeH264PictureInfoFlags { _bitfield_align_1: [], _bitfield_1: ash::vk::native::StdVideoEncodeH264PictureInfoFlags::new_bitfield_1( @@ -136,7 +153,7 @@ impl H264Encoder { if is_reference { 1 } else { 0 }, // is_reference 0, // no_output_of_prior_pics_flag 0, // long_term_reference_flag - 0, // adaptive_ref_pic_marking_mode_flag + use_adaptive_marking as u32, // adaptive_ref_pic_marking_mode_flag 0, // reserved ), }; @@ -147,16 +164,6 @@ impl H264Encoder { let mut ref_list0: [u8; 32] = [NO_REFERENCE_PICTURE; 32]; let mut ref_list1: [u8; 32] = [NO_REFERENCE_PICTURE; 32]; - // Reference list modification operations (not used) - let _ref_pic_list_mod_flags = ash::vk::native::StdVideoEncodeH264RefPicMarkingEntry { - memory_management_control_operation: - ash::vk::native::StdVideoH264MemMgmtControlOp_STD_VIDEO_H264_MEM_MGMT_CONTROL_OP_END, - difference_of_pic_nums_minus1: 0, - long_term_pic_num: 0, - long_term_frame_idx: 0, - max_long_term_frame_idx_plus1: 0, - }; - let ref_lists_info_flags = ash::vk::native::StdVideoEncodeH264ReferenceListsInfoFlags { _bitfield_align_1: [], _bitfield_1: ash::vk::native::StdVideoEncodeH264ReferenceListsInfoFlags::new_bitfield_1( @@ -167,10 +174,7 @@ impl H264Encoder { }; // Set up reference lists for P-frames and B-frames. - // P-frames: L0 only (reference previous frame) - // B-frames: L0 (forward/past) and L1 (backward/future) let (num_ref_l0, num_ref_l1) = if is_b_frame && self.has_backward_reference { - // B-frame: both L0 and L1 references. if let Some(first_ref) = self.l0_references.first() { ref_list0[0] = first_ref.dpb_slot; ref_list1[0] = self.backward_reference_dpb_slot; @@ -179,27 +183,19 @@ impl H264Encoder { (0, 0) } } else if !is_idr && !self.l0_references.is_empty() { - // P-frame: L0 reference list using actual available references, - // clamped to the negotiated active count (which is itself ≤32 at init time). let actual_count = self .l0_references .len() .min(self.active_reference_count as usize) .min(32); - for (i, ref_info) in self.l0_references.iter().take(actual_count).enumerate() { ref_list0[i] = ref_info.dpb_slot; } - (actual_count, 0) } else { - // IDR: no references. (0, 0) }; - // num_ref_idx_l0_active_minus1 tells the driver how many entries in RefPicList0 - // are valid and is encoded into the slice header (because we set - // num_ref_idx_active_override_flag=1 above). let ref_lists_info = ash::vk::native::StdVideoEncodeH264ReferenceListsInfo { flags: ref_lists_info_flags, num_ref_idx_l0_active_minus1: if num_ref_l0 > 0 { @@ -216,11 +212,15 @@ impl H264Encoder { RefPicList1: ref_list1, refList0ModOpCount: 0, refList1ModOpCount: 0, - refPicMarkingOpCount: 0, + refPicMarkingOpCount: ref_pic_marking_ops.len() as u8, reserved1: [0; 7], pRefList0ModOperations: std::ptr::null(), pRefList1ModOperations: std::ptr::null(), - pRefPicMarkingOperations: std::ptr::null(), + pRefPicMarkingOperations: if ref_pic_marking_ops.is_empty() { + std::ptr::null() + } else { + ref_pic_marking_ops.as_ptr() + }, }; let picture_info = ash::vk::native::StdVideoEncodeH264PictureInfo { @@ -233,50 +233,38 @@ impl H264Encoder { PicOrderCnt: pic_order_cnt, temporal_id: 0, reserved1: [0; 3], - pRefLists: if !is_idr && !self.l0_references.is_empty() { + pRefLists: if !is_idr && (!self.l0_references.is_empty() || use_adaptive_marking) { &ref_lists_info } else { std::ptr::null() }, }; - // Create slice NAL unit entry. - // constant_qp should only be set when rate control is DISABLED. - let constant_qp = if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { - qp - } else { - 0 - }; + let constant_qp = if rc.is_disabled() { rc.qp as i32 } else { 0 }; let nalu_slice_entries = [vk::VideoEncodeH264NaluSliceInfoKHR::default() .constant_qp(constant_qp) .std_slice_header(&slice_header)]; - // Create H.264 picture info. let mut h264_picture_info = vk::VideoEncodeH264PictureInfoKHR::default() .nalu_slice_entries(&nalu_slice_entries) .std_picture_info(&picture_info); - // Set up source picture resource. + let coded_extent = vk::Extent2D { + width: common.aligned_width, + height: common.aligned_height, + }; + let src_picture_resource = vk::VideoPictureResourceInfoKHR::default() .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) + .coded_extent(coded_extent) .base_array_layer(0) - .image_view_binding(self.input_image_view); + .image_view_binding(input_image_view); - // Set up DPB slot for reconstructed picture (setup slot) let setup_picture_resource = vk::VideoPictureResourceInfoKHR::default() .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) + .coded_extent(coded_extent) .base_array_layer(0) - .image_view_binding(self.dpb_image_views[self.current_dpb_slot as usize]); - - // Set up reference picture resources and info. + .image_view_binding(common.dpb_image_views[common.current_dpb_slot as usize]); let std_reference_info_flags = ash::vk::native::StdVideoEncodeH264ReferenceInfoFlags { _bitfield_align_1: [], @@ -286,23 +274,18 @@ impl H264Encoder { ), }; - // We use vectors to hold the data to ensure stable memory addresses for pointers. + // Vectors hold the data so pointers into them stay stable. let mut l0_resources = Vec::with_capacity(self.l0_references.len()); let mut l0_std_infos = Vec::with_capacity(self.l0_references.len()); - // 1. Populate data for L0 references (P-frames and B-frames) for ref_info in &self.l0_references { l0_resources.push( vk::VideoPictureResourceInfoKHR::default() .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) + .coded_extent(coded_extent) .base_array_layer(0) - .image_view_binding(self.dpb_image_views[ref_info.dpb_slot as usize]), + .image_view_binding(common.dpb_image_views[ref_info.dpb_slot as usize]), ); - l0_std_infos.push(ash::vk::native::StdVideoEncodeH264ReferenceInfo { flags: std_reference_info_flags, primary_pic_type: @@ -315,15 +298,12 @@ impl H264Encoder { }); } - // 2. Create DPB slot infos for L0. - // MUST be done after l0_std_infos is fully populated so addresses don't change. let mut l0_dpb_slot_infos = Vec::with_capacity(l0_std_infos.len()); for std_info in &l0_std_infos { l0_dpb_slot_infos .push(vk::VideoEncodeH264DpbSlotInfoKHR::default().std_reference_info(std_info)); } - // 3. Create the L0 reference slots. let mut l0_slots = Vec::with_capacity(l0_resources.len()); for (i, (resource, dpb_info)) in l0_resources .iter() @@ -339,18 +319,14 @@ impl H264Encoder { ); } - // 4. Handle Backward Ref (L1) for B-frames + // Backward (L1) reference for B-frames. let (backward_resource, backward_std_info) = if is_b_frame && self.has_backward_reference { - let image_view = self.dpb_image_views[self.backward_reference_dpb_slot as usize]; + let image_view = common.dpb_image_views[self.backward_reference_dpb_slot as usize]; let resource = vk::VideoPictureResourceInfoKHR::default() .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) + .coded_extent(coded_extent) .base_array_layer(0) .image_view_binding(image_view); - let std_info = ash::vk::native::StdVideoEncodeH264ReferenceInfo { flags: std_reference_info_flags, primary_pic_type: @@ -383,7 +359,7 @@ impl H264Encoder { None }; - // Create H.264 reference info for the setup slot (this frame being encoded) + // Reference info for the setup slot (this frame). let std_reference_info = ash::vk::native::StdVideoEncodeH264ReferenceInfo { flags: std_reference_info_flags, primary_pic_type: if is_idr { @@ -398,28 +374,23 @@ impl H264Encoder { temporal_id: 0, }; - // Create H.264 DPB slot info for setup. let mut h264_dpb_slot_info = vk::VideoEncodeH264DpbSlotInfoKHR::default().std_reference_info(&std_reference_info); - let setup_reference_slot = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(self.current_dpb_slot as i32) + .slot_index(common.current_dpb_slot as i32) .picture_resource(&setup_picture_resource) .push(&mut h264_dpb_slot_info); - // Also create a setup slot for begin_info (slotIndex = -1) let mut h264_begin_dpb_slot_info = vk::VideoEncodeH264DpbSlotInfoKHR::default().std_reference_info(&std_reference_info); - let setup_slot_for_begin = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(-1) // Marked as inactive for begin + .slot_index(-1) .picture_resource(&setup_picture_resource) .push(&mut h264_begin_dpb_slot_info); - // Collect final reference slots for VIDEO ENCODE INFO + // Reference slots for the encode command. let mut encode_ref_slots = Vec::new(); if is_b_frame && self.has_backward_reference { - // B-Frames: Use first L0 (forward) + L1 (backward) if let Some(l0) = l0_slots.first() { encode_ref_slots.push(*l0); } @@ -427,24 +398,21 @@ impl H264Encoder { encode_ref_slots.push(l1); } } else if !is_idr { - // P-Frames: Use all L0 encode_ref_slots.extend_from_slice(&l0_slots); } let mut encode_info = vk::VideoEncodeInfoKHR::default() - .dst_buffer(self.bitstream_buffer) + .dst_buffer(bitstream_buffer) .dst_buffer_offset(0) - .dst_buffer_range(MIN_BITSTREAM_BUFFER_SIZE as vk::DeviceSize) + .dst_buffer_range(bitstream_buffer_size as vk::DeviceSize) .src_picture_resource(src_picture_resource) .setup_reference_slot(&setup_reference_slot); - if !encode_ref_slots.is_empty() { encode_info = encode_info.reference_slots(&encode_ref_slots); } - encode_info = encode_info.push(&mut h264_picture_info); - // Collect reference slots for BEGIN VIDEO CODING + // Reference slots for begin coding (setup slot is marked inactive, -1). let mut reference_slots_for_begin = vec![setup_slot_for_begin]; if is_b_frame && self.has_backward_reference { if let Some(l0) = l0_slots.first() { @@ -457,33 +425,18 @@ impl H264Encoder { reference_slots_for_begin.extend_from_slice(&l0_slots); } - let min_qp_val = if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED - || self.config.rate_control_mode == crate::encoder::RateControlMode::Cqp - || self.config.rate_control_mode == crate::encoder::RateControlMode::Disabled - { - qp // Clamp to fixed QP for CQP/Disabled simulation - } else { - 18 // Allow high quality for H.264 - }; - let max_qp_val = if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED - || self.config.rate_control_mode == crate::encoder::RateControlMode::Cqp - || self.config.rate_control_mode == crate::encoder::RateControlMode::Disabled - { - qp // Clamp to fixed QP for CQP/Disabled simulation - } else { - 42 // Allow lower quality when needed - }; - + // Rate control. + let qp_bounds = if rc.is_disabled() { rc.qp as i32 } else { 18 }; + let qp_bounds_max = if rc.is_disabled() { rc.qp as i32 } else { 42 }; let min_qp = vk::VideoEncodeH264QpKHR { - qp_i: min_qp_val, - qp_p: min_qp_val, - qp_b: min_qp_val, + qp_i: qp_bounds, + qp_p: qp_bounds, + qp_b: qp_bounds, }; - let max_qp = vk::VideoEncodeH264QpKHR { - qp_i: max_qp_val, - qp_p: max_qp_val, - qp_b: max_qp_val, + qp_i: qp_bounds_max, + qp_p: qp_bounds_max, + qp_b: qp_bounds_max, }; let mut h264_rc_layer_info = vk::VideoEncodeH264RateControlLayerInfoKHR::default() @@ -491,63 +444,54 @@ impl H264Encoder { .min_qp(min_qp) .use_max_qp(true) .max_qp(max_qp); - let rc_layer_info = vk::VideoEncodeRateControlLayerInfoKHR::default() - .average_bitrate(average_bitrate as u64) - .max_bitrate(max_bitrate as u64) - .frame_rate_numerator(self.config.frame_rate_numerator) - .frame_rate_denominator(self.config.frame_rate_denominator) + .average_bitrate(rc.average_bitrate as u64) + .max_bitrate(rc.max_bitrate as u64) + .frame_rate_numerator(common.config.frame_rate_numerator) + .frame_rate_denominator(common.config.frame_rate_denominator) .push(&mut h264_rc_layer_info); - let rc_layers = [rc_layer_info]; let mut h264_rc_info = vk::VideoEncodeH264RateControlInfoKHR::default() - .gop_frame_count(self.config.gop_size) - .idr_period(self.config.gop_size) - .consecutive_b_frame_count(self.config.b_frame_count); + .gop_frame_count(common.config.gop_size) + .idr_period(common.config.gop_size) + .consecutive_b_frame_count(common.config.b_frame_count); - let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc_mode); - - if rc_mode != vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { + let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc.mode); + if !rc.is_disabled() { rc_info = rc_info .layers(&rc_layers) - .virtual_buffer_size_in_ms(self.config.virtual_buffer_size_ms) - .initial_virtual_buffer_size_in_ms(self.config.initial_virtual_buffer_size_ms); + .virtual_buffer_size_in_ms(common.config.virtual_buffer_size_ms) + .initial_virtual_buffer_size_in_ms(common.config.initial_virtual_buffer_size_ms); } - // Begin video coding. - // For the first frame, don't include rate control in begin_coding - set it via control command after RESET. - let is_first_frame = self.encode_frame_num == 0; - + // For the first frame, configure rate control via the control command + // after RESET rather than in begin_coding. + let is_first_frame = plan.is_first_frame(); let begin_info = if is_first_frame { vk::VideoBeginCodingInfoKHR::default() - .video_session(self.session) - .video_session_parameters(self.session_params) + .video_session(common.session) + .video_session_parameters(common.session_params) .reference_slots(&reference_slots_for_begin) .push(&mut h264_rc_info) } else { vk::VideoBeginCodingInfoKHR::default() - .video_session(self.session) - .video_session_parameters(self.session_params) + .video_session(common.session) + .video_session_parameters(common.session_params) .reference_slots(&reference_slots_for_begin) .push(&mut h264_rc_info) .push(&mut rc_info) }; unsafe { - (self.video_queue_fn.fp().cmd_begin_video_coding_khr)( - self.encode_command_buffer, - &begin_info, - ); + (common.video_queue_fn.fp().cmd_begin_video_coding_khr)(command_buffer, &begin_info); } - // Reset video coding state for the first frame. - // Combine RESET + RATE_CONTROL + QUALITY_LEVEL into a single control command. - // This matches FFmpeg's approach and is required for AMD RADV. + // RESET + RATE_CONTROL + QUALITY_LEVEL in one control command on the first + // frame (matches FFmpeg; required for AMD RADV). if is_first_frame { let mut quality_level_info = vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); - let control_info = vk::VideoCodingControlInfoKHR::default() .flags( vk::VideoCodingControlFlagsKHR::RESET @@ -557,76 +501,41 @@ impl H264Encoder { .push(&mut rc_info) .push(&mut h264_rc_info) .push(&mut quality_level_info); - unsafe { - (self.video_queue_fn.fp().cmd_control_video_coding_khr)( - self.encode_command_buffer, + (common.video_queue_fn.fp().cmd_control_video_coding_khr)( + command_buffer, &control_info, ); } } - // Begin query. unsafe { - self.context.device().cmd_begin_query( - self.encode_command_buffer, - self.query_pool, + let device = common.device(); + device.cmd_begin_query( + command_buffer, + query_pool, 0, vk::QueryControlFlags::empty(), ); - } + (common.video_encode_fn.fp().cmd_encode_video_khr)(command_buffer, &encode_info); + device.cmd_end_query(command_buffer, query_pool, 0); - // Encode - unsafe { - (self.video_encode_fn.fp().cmd_encode_video_khr)( - self.encode_command_buffer, - &encode_info, - ); - } + let end_info = vk::VideoEndCodingInfoKHR::default(); + (common.video_queue_fn.fp().cmd_end_video_coding_khr)(command_buffer, &end_info); - // End query. - unsafe { - self.context - .device() - .cmd_end_query(self.encode_command_buffer, self.query_pool, 0); + device + .end_command_buffer(command_buffer) + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; } - // End video coding. - let end_info = vk::VideoEndCodingInfoKHR::default(); - unsafe { - (self.video_queue_fn.fp().cmd_end_video_coding_khr)( - self.encode_command_buffer, - &end_info, - ); - } - - // End command buffer. - unsafe { - self.context - .device() - .end_command_buffer(self.encode_command_buffer) + let future = common.submit_frame()?; + // Clear the unmark queue only after the encode is committed to the GPU. + // If any fallible step above had failed, the MMCO ops would still live + // in `pending_unmark_frame_nums` so a retry re-emits them — otherwise the + // encoder would silently desync from the decoder's reference state. + if use_adaptive_marking { + self.pending_unmark_frame_nums.clear(); } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - - // Submit, wait, and read bitstream. - let encode_queue = self.context.video_encode_queue().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) - })?; - - let encoded_data = unsafe { - submit_encode_and_read_bitstream( - self.context.device(), - self.encode_command_buffer, - self.encode_fence, - encode_queue, - self.query_pool, - self.bitstream_buffer_ptr, - )? - }; - - // Mark DPB slot as active. - self.dpb_slot_active[self.current_dpb_slot as usize] = true; - - Ok(encoded_data) + Ok(future) } } diff --git a/src/encoder/h264/session_params.rs b/src/encoder/h264/session_params.rs index f3f1db5..d20ccd7 100644 --- a/src/encoder/h264/session_params.rs +++ b/src/encoder/h264/session_params.rs @@ -1,39 +1,40 @@ -use super::H264Encoder; +//! H.264 parameter-set generation: SPS/PPS construction and header retrieval. +use super::H264; + +use crate::encoder::codec::EncoderCommon; +use crate::encoder::resources::get_encoded_session_params; use crate::encoder::{BitDepth, ColorDescription, PixelFormat}; use crate::error::{PixelForgeError, Result}; use ash::vk; use ash::vk::TaggedStructure; use std::ptr; -impl H264Encoder { +impl H264 { /// Build SPS/PPS and create Vulkan video session parameters. /// - /// This is used both during initial encoder creation and by - /// `set_color_description()` to rebuild session parameters with - /// updated VUI color metadata. Keeping a single implementation - /// ensures the parameter sets stay bit-for-bit consistent. - pub(crate) fn create_session_params( + /// Used both at init and by `set_color_description`; a single implementation + /// keeps the parameter sets bit-for-bit consistent. + pub(super) fn build_session_params( &self, + common: &EncoderCommon, desc: &ColorDescription, ) -> Result { - let width = self.config.dimensions.width; - let height = self.config.dimensions.height; + let config = &common.config; + let width = config.dimensions.width; + let height = config.dimensions.height; - let pic_width_in_mbs = self.aligned_width / 16; - let pic_height_in_map_units = self.aligned_height / 16; + let pic_width_in_mbs = common.aligned_width / 16; + let pic_height_in_map_units = common.aligned_height / 16; - // Cropping offsets are expressed in units that depend on chroma subsampling. - // For progressive frames (frame_mbs_only_flag=1): - // - 4:2:0 => crop_unit_x=2, crop_unit_y=2 - // - 4:4:4 => crop_unit_x=1, crop_unit_y=1 - let (crop_unit_x, crop_unit_y) = match self.config.pixel_format { + // Crop units depend on chroma subsampling (frame_mbs_only_flag=1). + let (crop_unit_x, crop_unit_y) = match config.pixel_format { PixelFormat::Yuv420 => (2u32, 2u32), PixelFormat::Yuv444 => (1u32, 1u32), _ => { return Err(PixelForgeError::InvalidInput(format!( "Unsupported pixel format for H.264: {:?}", - self.config.pixel_format + config.pixel_format ))); } }; @@ -52,7 +53,7 @@ impl H264Encoder { } sps_flags.set_vui_parameters_present_flag(1); - let chroma_format_idc = match self.config.pixel_format { + let chroma_format_idc = match config.pixel_format { PixelFormat::Yuv420 => { ash::vk::native::StdVideoH264ChromaFormatIdc_STD_VIDEO_H264_CHROMA_FORMAT_IDC_420 } @@ -62,20 +63,19 @@ impl H264Encoder { _ => unreachable!("Pixel format validated above"), }; - let (bit_depth_luma_minus8, bit_depth_chroma_minus8) = match self.config.bit_depth { + let (bit_depth_luma_minus8, bit_depth_chroma_minus8) = match config.bit_depth { BitDepth::Eight => (0u8, 0u8), BitDepth::Ten => (2u8, 2u8), }; - let max_active = self.active_reference_count as u8; + let max_active = self.active_reference_count() as u8; let mut vui_flags: ash::vk::native::StdVideoH264SpsVuiFlags = unsafe { std::mem::zeroed() }; vui_flags.set_aspect_ratio_info_present_flag(1); vui_flags.set_video_signal_type_present_flag(1); vui_flags.set_video_full_range_flag(if desc.full_range { 1 } else { 0 }); vui_flags.set_color_description_present_flag(1); - // Do not set HRD parameters when rate control is disabled/CQP. - // HRD with zeroed bitrate values causes device loss on some drivers (AMD). + // Avoid HRD with zeroed bitrates (device loss on some drivers). vui_flags.set_nal_hrd_parameters_present_flag(0); vui_flags.set_bitstream_restriction_flag(1); @@ -91,7 +91,7 @@ impl H264Encoder { matrix_coefficients: desc.matrix_coefficients, num_units_in_tick: 0, time_scale: 0, - max_num_reorder_frames: if self.config.b_frame_count > 0 { 1 } else { 0 }, + max_num_reorder_frames: if config.b_frame_count > 0 { 1 } else { 0 }, max_dec_frame_buffering: max_active + 1, chroma_sample_loc_type_top_field: 0, chroma_sample_loc_type_bottom_field: 0, @@ -101,14 +101,14 @@ impl H264Encoder { let sps = ash::vk::native::StdVideoH264SequenceParameterSet { flags: sps_flags, - profile_idc: self.profile_idc, + profile_idc: self.profile_idc(), level_idc: ash::vk::native::StdVideoH264LevelIdc_STD_VIDEO_H264_LEVEL_IDC_4_1, chroma_format_idc, seq_parameter_set_id: 0, bit_depth_luma_minus8, bit_depth_chroma_minus8, log2_max_frame_num_minus4: 4, - pic_order_cnt_type: if self.config.b_frame_count > 0 { + pic_order_cnt_type: if config.b_frame_count > 0 { ash::vk::native::StdVideoH264PocType_STD_VIDEO_H264_POC_TYPE_0 } else { ash::vk::native::StdVideoH264PocType_STD_VIDEO_H264_POC_TYPE_2 @@ -132,20 +132,15 @@ impl H264Encoder { }; let mut pps_flags: ash::vk::native::StdVideoH264PpsFlags = unsafe { std::mem::zeroed() }; - // Enable 8x8 transform for High profile and above (required by some - // drivers for High 4:4:4 Predictive SPS/PPS generation). - let transform_8x8 = self.profile_idc + // Enable 8x8 transform for High profile and above. + let transform_8x8 = self.profile_idc() >= ash::vk::native::StdVideoH264ProfileIdc_STD_VIDEO_H264_PROFILE_IDC_HIGH; pps_flags.set_transform_8x8_mode_flag(transform_8x8 as u32); - // Use the driver's preferred entropy coding mode from quality level properties. - // Some drivers (e.g., NVIDIA for H.264 High 4:4:4 Predictive) require CAVLC. - pps_flags.set_entropy_coding_mode_flag(self.preferred_entropy_cabac as u32); + pps_flags.set_entropy_coding_mode_flag(self.preferred_entropy_cabac() as u32); pps_flags.set_deblocking_filter_control_present_flag(1); - // vk_video_samples sets chroma QP offsets to 6 for 4:4:4 unless lossless. - // This improves driver compatibility for SPS/PPS generation. - let (chroma_qp_index_offset, second_chroma_qp_index_offset) = match self.config.pixel_format - { + // vk_video_samples sets chroma QP offsets to 6 for 4:4:4 (driver compat). + let (chroma_qp_index_offset, second_chroma_qp_index_offset) = match config.pixel_format { PixelFormat::Yuv444 => (6i8, 6i8), _ => (0i8, 0i8), }; @@ -178,19 +173,21 @@ impl H264Encoder { .max_std_pps_count(1) .parameters_add_info(&h264_add_info); - // Chain quality level info into session parameters creation. - // This is required by AMD RADV and matches FFmpeg's approach. + // Chain quality level info (required by AMD RADV; matches FFmpeg). let mut quality_level_info = vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); let params_create_info = vk::VideoSessionParametersCreateInfoKHR::default() - .video_session(self.session) + .video_session(common.session) .push(&mut h264_params_create_info) .push(&mut quality_level_info); let mut session_params = vk::VideoSessionParametersKHR::null(); let result = unsafe { - (self.video_queue_fn.fp().create_video_session_parameters_khr)( - self.context.device().handle(), + (common + .video_queue_fn + .fp() + .create_video_session_parameters_khr)( + common.device().handle(), ¶ms_create_info, ptr::null(), &mut session_params, @@ -202,7 +199,27 @@ impl H264Encoder { result ))); } - Ok(session_params) } + + /// Retrieve the encoded SPS/PPS NAL units to prepend to an IDR frame. + pub(super) fn build_header(&self, common: &EncoderCommon) -> Result> { + let mut h264_get_info = vk::VideoEncodeH264SessionParametersGetInfoKHR::default() + .write_std_sps(true) + .write_std_pps(true) + .std_sps_id(0) + .std_pps_id(0); + let get_info = vk::VideoEncodeSessionParametersGetInfoKHR::default() + .video_session_parameters(common.session_params) + .push(&mut h264_get_info); + let mut h264_feedback = vk::VideoEncodeH264SessionParametersFeedbackInfoKHR::default(); + let mut feedback = + vk::VideoEncodeSessionParametersFeedbackInfoKHR::default().push(&mut h264_feedback); + get_encoded_session_params( + &common.context, + &common.video_encode_fn, + &get_info, + &mut feedback, + ) + } } diff --git a/src/encoder/h265/api.rs b/src/encoder/h265/api.rs deleted file mode 100644 index d09b7ca..0000000 --- a/src/encoder/h265/api.rs +++ /dev/null @@ -1,323 +0,0 @@ -use super::H265Encoder; - -use crate::encoder::dpb::{DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType}; -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::{ColorDescription, EncodedPacket}; -use crate::error::Result; -use crate::PixelForgeError; -use ash::vk; -use ash::vk::TaggedStructure; -use tracing::debug; - -impl H265Encoder { - /// Get the internal input image. - /// - /// This image can be used as a target for `ColorConverter::convert` to avoid - /// an intermediate copy. - pub fn input_image(&self) -> vk::Image { - self.input_image - } - - /// Encode a frame from a GPU image. - /// - /// This accepts a source NV12 image on the GPU and encodes it directly without. - /// any CPU-side data copies. The source image must be in NV12 format with the - /// same dimensions as the encoder configuration, and should be in GENERAL layout. - /// - /// # Panics - /// - /// The encoder will panic at creation time if B-frames are enabled (b_frame_count > 0), - /// as B-frame encoding is not yet supported. - pub fn encode(&mut self, src_image: vk::Image) -> Result> { - let gop_position = self.gop.get_next_frame(); - let display_order = self.input_frame_num; - self.input_frame_num += 1; - - debug!( - "Encoding frame {} from GPU image: type={:?}, poc={}", - display_order, gop_position.frame_type, gop_position.pic_order_cnt - ); - - // Upload from GPU image. - self.upload_from_image(src_image)?; - - // Encode immediately. - let packet = self.encode_current_frame(&gop_position, display_order)?; - - Ok(vec![packet]) - } - - /// Internal method to encode the current frame already uploaded to input_image. - fn encode_current_frame( - &mut self, - gop_position: &GopPosition, - display_order: u64, - ) -> Result { - let is_idr = gop_position.frame_type.is_idr(); - let is_reference = gop_position.is_reference; - let is_b_frame = gop_position.frame_type == GopFrameType::B; - - debug!( - "Encoding frame: display_order={}, type={:?}, idr={}, ref={}, poc={}, l0_refs={:?}", - display_order, - gop_position.frame_type, - is_idr, - is_reference, - gop_position.pic_order_cnt, - self.l0_references - .iter() - .map(|r| (r.dpb_slot, r.poc)) - .collect::>() - ); - - // Determine frame_type. - let frame_type = if is_idr { - crate::encoder::FrameType::I - } else { - match gop_position.frame_type { - GopFrameType::Idr | GopFrameType::I => crate::encoder::FrameType::I, - GopFrameType::P => crate::encoder::FrameType::P, - GopFrameType::B => crate::encoder::FrameType::B, - } - }; - - if is_idr { - // Reset DPB by calling sequence_start with new config for IDR. - let dpb_config = DpbConfig { - dpb_size: self.dpb_slot_count as u32, - max_num_ref_frames: self.config.max_reference_frames, - use_multiple_references: self.config.b_frame_count > 0, - log2_max_frame_num_minus4: 0, - log2_max_pic_order_cnt_lsb_minus4: 4, - ..Default::default() - }; - self.dpb.h265.sequence_start(dpb_config); - self.l0_references.clear(); - self.has_backward_reference = false; - // Reset DPB slot activation tracking on IDR - all slots become inactive. - for active in &mut self.dpb_slot_active { - *active = false; - } - } - - let pic_order_cnt = gop_position.pic_order_cnt; - - let mut encoded_data = Vec::new(); - - // For IDR frames, prepend VPS/SPS/PPS header. - if is_idr { - if self.header_data.is_none() { - let header = self.get_h265_header()?; - // Debug: print first few bytes of header. - debug!( - "H.265 header ({} bytes): {:02X?}", - header.len(), - &header[..std::cmp::min(32, header.len())] - ); - self.header_data = Some(header); - } - if let Some(ref header) = self.header_data { - encoded_data.extend_from_slice(header); - } - } - - let slice_data = self.encode_frame_internal(gop_position, pic_order_cnt, is_idr)?; - // Debug: print first few bytes of slice data. - debug!( - "H.265 slice ({} bytes): {:02X?}", - slice_data.len(), - &slice_data[..std::cmp::min(16, slice_data.len())] - ); - encoded_data.extend_from_slice(&slice_data); - - self.encode_frame_num += 1; - - if is_reference { - let dpb_pic_type = if is_idr { - PictureType::Idr - } else if is_b_frame { - PictureType::B - } else { - PictureType::P - }; - let pic_info = PictureStartInfo { - frame_id: display_order, - pic_order_cnt, - frame_num: 0, - pic_type: dpb_pic_type, - is_reference, - ..Default::default() - }; - self.dpb.h265.picture_start(pic_info); - self.dpb.h265.picture_end(is_reference); - - // Update reference tracking for the next P-frame. - // The current frame becomes the reference for subsequent frames. - let ref_info = super::ReferenceInfo { - dpb_slot: self.current_dpb_slot, - poc: pic_order_cnt, - }; - self.l0_references.insert(0, ref_info); - - // Limit to active_reference_count - while self.l0_references.len() > self.active_reference_count as usize { - self.l0_references.pop(); - } - - if !is_b_frame { - // Find a free DPB slot for the next frame. - // Avoid slots currently in l0_references. - // B-frames don't consume a slot permanently (in this simple implementation they are not reference), - // but if we support B-frame references we need more complex logic. - // Assuming B-frames are not reference for now (is_reference=false for B in this impl usually). - - let used_slots: Vec = self.l0_references.iter().map(|r| r.dpb_slot).collect(); - for i in 0..self.dpb_slot_count as u8 { - if !used_slots.contains(&i) { - self.current_dpb_slot = i; - break; - } - } - } - } - - Ok(EncodedPacket { - data: encoded_data, - frame_type, - is_key_frame: is_idr, - pts: display_order, - dts: self.encode_frame_num - 1, - }) - } - - /// Flush the encoder and get any remaining packets. - pub fn flush(&mut self) -> Result> { - // No buffered frames in the current implementation. - Ok(Vec::new()) - } - - /// Request that the next frame be an IDR frame. - pub fn request_idr(&mut self) { - self.gop.request_idr(); - } - - /// Retrieve encoded VPS/SPS/PPS header data from video session parameters. - /// This uses vkGetEncodedVideoSessionParametersKHR to get the NAL units. - fn get_h265_header(&self) -> Result> { - // H.265-specific get info requesting VPS, SPS and PPS. - let mut h265_get_info = vk::VideoEncodeH265SessionParametersGetInfoKHR::default() - .write_std_vps(true) - .write_std_sps(true) - .write_std_pps(true) - .std_vps_id(0) - .std_sps_id(0) - .std_pps_id(0); - - let get_info = vk::VideoEncodeSessionParametersGetInfoKHR::default() - .video_session_parameters(self.session_params) - .push(&mut h265_get_info); - - // Some implementations misbehave for a size-only query (pData = NULL). Use a - // preallocated buffer and retry on INCOMPLETE (vk_video_samples-style). - let mut data = vec![0u8; 4096]; - let mut data_size: usize = data.len(); - let mut h265_feedback = vk::VideoEncodeH265SessionParametersFeedbackInfoKHR::default(); - let mut feedback = - vk::VideoEncodeSessionParametersFeedbackInfoKHR::default().push(&mut h265_feedback); - - let mut attempts = 0; - loop { - attempts += 1; - let result = unsafe { - (self - .video_encode_fn - .fp() - .get_encoded_video_session_parameters_khr)( - self.context.device().handle(), - &get_info, - &mut feedback, - &mut data_size, - data.as_mut_ptr() as *mut std::ffi::c_void, - ) - }; - - match result { - vk::Result::SUCCESS => { - if data_size == 0 { - return Err(PixelForgeError::SessionParametersCreation( - "Encoded parameters size is 0".to_string(), - )); - } - data.truncate(data_size); - return Ok(data); - } - vk::Result::INCOMPLETE if attempts < 3 => { - let new_size = data_size.max(data.len() * 2).max(1); - data.resize(new_size, 0); - data_size = data.len(); - } - err => { - return Err(PixelForgeError::SessionParametersCreation(format!( - "Failed to get encoded parameters: {:?}", - err - ))); - } - } - } - } - - /// Update the color description (VUI parameters) in the encoded stream. - /// - /// This recreates the video session parameters with a new SPS containing the - /// updated VUI color primaries, transfer characteristics, and matrix coefficients. - /// The next encoded frame will be an IDR with the new VPS/SPS/PPS prepended. - pub fn set_color_description(&mut self, desc: ColorDescription) -> Result<()> { - // Wait for any in-flight encode to complete before modifying session params. - // Do NOT reset the fence here — submit_encode_and_read_bitstream() resets it - // before queue_submit. Leaving the fence signaled allows consecutive - // set_color_description() calls without deadlock. - unsafe { - self.context - .device() - .wait_for_fences(&[self.encode_fence], true, u64::MAX) - .map_err(|e| { - PixelForgeError::Synchronization(format!( - "Failed to wait for encode fence: {:?}", - e - )) - })?; - } - - // Save old handle so we can destroy it after successful creation. - let old_session_params = self.session_params; - - let new_session_params = self.create_session_params(&desc)?; - - // Destroy old session parameters now that the new ones are created. - unsafe { - (self - .video_queue_fn - .fp() - .destroy_video_session_parameters_khr)( - self.context.device().handle(), - old_session_params, - std::ptr::null(), - ); - } - - self.session_params = new_session_params; - self.config.color_description = Some(desc); - self.header_data = None; // Invalidate cached VPS/SPS/PPS header - self.gop.request_idr(); - - debug!( - "H.265 color description updated: primaries={}, transfer={}, matrix={}, full_range={}", - desc.color_primaries, - desc.transfer_characteristics, - desc.matrix_coefficients, - desc.full_range - ); - - Ok(()) - } -} diff --git a/src/encoder/h265/init.rs b/src/encoder/h265/init.rs index deac373..9d2211c 100644 --- a/src/encoder/h265/init.rs +++ b/src/encoder/h265/init.rs @@ -1,59 +1,31 @@ -use super::H265Encoder; +use super::{CTB_SIZE, H265}; -use crate::encoder::dpb::{DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig}; -use crate::encoder::gop::GopStructure; -use crate::encoder::resources::{ - align_up, allocate_session_memory, clear_input_image, create_bitstream_buffer, - create_command_resources, create_dpb_images, create_image, get_video_format, lcm, - make_codec_name, map_bitstream_buffer, query_supported_video_formats, ClearImageParams, - MIN_BITSTREAM_BUFFER_SIZE, +use crate::encoder::codec::{ + build_encoder_common, query_video_caps, CodecEncoder, CommonInitRequest, }; -use crate::encoder::{BitDepth, ColorDescription, PixelFormat}; +use crate::encoder::dpb::{DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig}; +use crate::encoder::resources::MIN_BITSTREAM_BUFFER_SIZE; +use crate::encoder::{BitDepth, ColorDescription, EncodeConfig, PixelFormat}; use crate::error::{PixelForgeError, Result}; use crate::vulkan::VideoContext; use ash::vk; use ash::vk::TaggedStructure; -use std::ptr; -use tracing::{debug, info}; +use tracing::info; -impl H265Encoder { +impl H265 { /// Create a new H.265/HEVC encoder. - pub fn new(context: VideoContext, config: crate::encoder::EncodeConfig) -> Result { - // B-frames are not yet supported. - if config.b_frame_count > 0 { - panic!( - "B-frame encoding is not yet supported. Set b_frame_count=0 in encoder config. \ - Got b_frame_count={}", - config.b_frame_count - ); - } - - let width = config.dimensions.width; - let height = config.dimensions.height; + pub fn create(context: VideoContext, config: EncodeConfig) -> Result> { + assert!( + config.b_frame_count == 0, + "B-frame encoding is not yet supported; set b_frame_count=0 (got {})", + config.b_frame_count + ); info!( "Creating H.265 encoder: {}x{}, pixel_format={:?}", - width, height, config.pixel_format + config.dimensions.width, config.dimensions.height, config.pixel_format ); - // Load video queue extension functions. - let video_queue_fn = - ash::khr::video_queue::Device::load(context.instance(), context.device()); - let video_encode_fn = - ash::khr::video_encode_queue::Device::load(context.instance(), context.device()); - - // Get chroma subsampling from pixel format via `From` impl - let chroma_subsampling: vk::VideoChromaSubsamplingFlagsKHR = config.pixel_format.into(); - - // Get bit depth flags from config - let bit_depth_flags: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); - let video_format = get_video_format(config.pixel_format, config.bit_depth); - - // Select profile based on pixel format and bit depth: - // - Main for YUV420 8-bit - // - Main 10 for YUV420 10-bit - // - Main 4:4:4 for YUV444 8-bit - // - Main 4:4:4 10 for YUV444 10-bit let profile_idc = match (config.pixel_format, config.bit_depth) { (PixelFormat::Yuv420, BitDepth::Eight) => { ash::vk::native::StdVideoH265ProfileIdc_STD_VIDEO_H265_PROFILE_IDC_MAIN @@ -61,427 +33,87 @@ impl H265Encoder { (PixelFormat::Yuv420, BitDepth::Ten) => { ash::vk::native::StdVideoH265ProfileIdc_STD_VIDEO_H265_PROFILE_IDC_MAIN_10 } - (PixelFormat::Yuv444, BitDepth::Eight) => { - ash::vk::native::StdVideoH265ProfileIdc_STD_VIDEO_H265_PROFILE_IDC_FORMAT_RANGE_EXTENSIONS - } - (PixelFormat::Yuv444, BitDepth::Ten) => { + (PixelFormat::Yuv444, _) => { ash::vk::native::StdVideoH265ProfileIdc_STD_VIDEO_H265_PROFILE_IDC_FORMAT_RANGE_EXTENSIONS } _ => { return Err(PixelForgeError::InvalidInput(format!( - "Unsupported pixel format / bit depth combination for H.265: {:?} / {:?}", + "Unsupported pixel format / bit depth for H.265: {:?} / {:?}", config.pixel_format, config.bit_depth ))); } }; - // Create H.265 encode profile + let chroma_subsampling: vk::VideoChromaSubsamplingFlagsKHR = config.pixel_format.into(); + let bit_depth: vk::VideoComponentBitDepthFlagsKHR = config.bit_depth.into(); + let mut h265_profile_info = vk::VideoEncodeH265ProfileInfoKHR::default().std_profile_idc(profile_idc); - let profile_info = vk::VideoProfileInfoKHR::default() .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H265) .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(bit_depth_flags) - .chroma_bit_depth(bit_depth_flags) + .luma_bit_depth(bit_depth) + .chroma_bit_depth(bit_depth) .push(&mut h265_profile_info); - // Query capabilities to determine limits. - let video_queue_instance = - ash::khr::video_queue::Instance::load(context.entry(), context.instance()); - let mut h265_capabilities = vk::VideoEncodeH265CapabilitiesKHR::default(); - let mut encode_capabilities = vk::VideoEncodeCapabilitiesKHR::default(); + // Query device capabilities with the H.265-specific capability struct + // chained in (required by the driver). + let mut h265_caps = vk::VideoEncodeH265CapabilitiesKHR::default(); + let mut encode_caps = vk::VideoEncodeCapabilitiesKHR::default(); let mut capabilities = vk::VideoCapabilitiesKHR::default() - .push(&mut h265_capabilities) - .push(&mut encode_capabilities); - - let result = unsafe { - (video_queue_instance - .fp() - .get_physical_device_video_capabilities_khr)( - context.physical_device(), - &profile_info, - &mut capabilities, - ) - }; - if result != vk::Result::SUCCESS { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Failed to query H.265 capabilities: {:?}", - result - ))); - } - - // Compute aligned coded extent using capabilities (picture_access_granularity, - // min/max_coded_extent) - required for AMD which reports granularity 64x16 and - // min_coded_extent 130x128. - let ctb_size = super::CTB_SIZE; - let gran_w = capabilities.picture_access_granularity.width.max(1); - let gran_h = capabilities.picture_access_granularity.height.max(1); - let align_w = lcm(ctb_size, gran_w); - let align_h = lcm(ctb_size, gran_h); - - let mut aligned_width = align_up(width, align_w); - let mut aligned_height = align_up(height, align_h); - - aligned_width = align_up( - aligned_width.max(capabilities.min_coded_extent.width), - align_w, - ); - aligned_height = align_up( - aligned_height.max(capabilities.min_coded_extent.height), - align_h, - ); - - if aligned_width > capabilities.max_coded_extent.width - || aligned_height > capabilities.max_coded_extent.height - { - return Err(PixelForgeError::InvalidInput(format!( - "Requested coded extent {}x{} (aligned to {}x{} with granularity {}x{}) exceeds device max {}x{} for this profile", - width, - height, - aligned_width, - aligned_height, - gran_w, - gran_h, - capabilities.max_coded_extent.width, - capabilities.max_coded_extent.height - ))); - } - - info!( - "Using coded extent {}x{} (granularity {}x{}, min {}x{}, max {}x{})", - aligned_width, - aligned_height, - gran_w, - gran_h, - capabilities.min_coded_extent.width, - capabilities.min_coded_extent.height, - capabilities.max_coded_extent.width, - capabilities.max_coded_extent.height - ); - - // Query supported formats for SRC and DPB usage. - let supported_src_formats = query_supported_video_formats( - &context, - &profile_info, - vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR, - )?; - let supported_dpb_formats = query_supported_video_formats( - &context, - &profile_info, - vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR, - )?; - - if supported_src_formats.is_empty() { - return Err(PixelForgeError::NoSuitableDevice( - "No supported Vulkan Video SRC formats returned for this H.265 profile".to_string(), - )); - } - info!("Supported SRC formats: {:?}", supported_src_formats); - if supported_dpb_formats.is_empty() { - return Err(PixelForgeError::NoSuitableDevice( - "No supported Vulkan Video DPB formats returned for this H.265 profile".to_string(), - )); - } - info!("Supported DPB formats: {:?}", supported_dpb_formats); - - let picture_format = if supported_src_formats.contains(&video_format) { - video_format - } else { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Preferred input format {:?} is not supported for VIDEO_ENCODE_SRC_KHR. Supported: {:?}", - video_format, supported_src_formats - ))); - }; - - let reference_picture_format = supported_dpb_formats - .iter() - .copied() - .find(|f| *f == picture_format) - .unwrap_or(supported_dpb_formats[0]); - - debug!( - "Selected Vulkan Video formats: picture_format={:?}, reference_picture_format={:?}", - picture_format, reference_picture_format - ); - - let max_dpb_slots_supported = capabilities.max_dpb_slots as usize; - let max_active_reference_pictures_supported = - capabilities.max_active_reference_pictures as usize; - - if max_dpb_slots_supported < 2 { - return Err(PixelForgeError::NoSuitableDevice(format!( - "Device reports max_dpb_slots={} for this profile; need at least 2", - max_dpb_slots_supported - ))); - } - - // Target number of active reference pictures. - let mut target_active_refs = (config.max_reference_frames as usize) - .min(max_active_reference_pictures_supported) - .min(15); // H.265 limit - - if target_active_refs < 1 && max_active_reference_pictures_supported >= 1 { - target_active_refs = 1; - } - - // Calculate needed DPB slots - let needed_dpb_slots = if config.b_frame_count > 0 { - target_active_refs + config.b_frame_count as usize + 2 - } else { - target_active_refs + 1 - }; - - let dpb_slot_count = needed_dpb_slots - .min(max_dpb_slots_supported) - .min(crate::encoder::dpb::MAX_DPB_SLOTS); - - // Final clamp - let max_active_reference_pictures = - target_active_refs.min(dpb_slot_count.saturating_sub(1)); - - debug!( - "Allocating {} DPB slots (req {}, max {}), active refs {} (req {}, max {})", - dpb_slot_count, - needed_dpb_slots, - max_dpb_slots_supported, - max_active_reference_pictures, - target_active_refs, - max_active_reference_pictures_supported - ); - - // Create video session. - let std_header_version = vk::ExtensionProperties { - extension_name: make_codec_name(b"VK_STD_vulkan_video_codec_h265_encode"), - spec_version: vk::make_api_version(0, 1, 0, 0), - }; - - let encode_queue_family = context.video_encode_queue_family().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue family available".to_string()) + .push(&mut encode_caps) + .push(&mut h265_caps); + let caps = query_video_caps(&context, &profile_info, &mut capabilities)?; + + let init = build_encoder_common(&CommonInitRequest { + context: &context, + config: &config, + profile_info: &profile_info, + caps: &caps, + align_unit: CTB_SIZE, + max_active_refs_cap: 15, + bitstream_buffer_size: MIN_BITSTREAM_BUFFER_SIZE, + allow_layered_dpb: true, })?; + let active_reference_count = init.active_reference_count; + let common = init.common; - let session_create_info = vk::VideoSessionCreateInfoKHR::default() - .queue_family_index(encode_queue_family) - .flags(vk::VideoSessionCreateFlagsKHR::empty()) - .video_profile(&profile_info) - .picture_format(picture_format) - .max_coded_extent(vk::Extent2D { - width: aligned_width, - height: aligned_height, - }) - .reference_picture_format(reference_picture_format) - .max_dpb_slots(dpb_slot_count as u32) - .max_active_reference_pictures(max_active_reference_pictures as u32) - .std_header_version(&std_header_version); - - let mut session = vk::VideoSessionKHR::null(); - let result = unsafe { - (video_queue_fn.fp().create_video_session_khr)( - context.device().handle(), - &session_create_info, - ptr::null(), - &mut session, - ) - }; - if result != vk::Result::SUCCESS { - return Err(PixelForgeError::VideoSessionCreation(format!( - "{:?}", - result - ))); - } - - // Query and allocate session memory. - let session_memory = allocate_session_memory(&context, session, &video_queue_fn)?; - - // Build VPS/SPS/PPS and session parameters via shared helper. - let color_desc = config - .color_description - .unwrap_or(ColorDescription::bt709()); - - // Create profile info for images/buffers - let mut h265_profile_for_resources = - vk::VideoEncodeH265ProfileInfoKHR::default().std_profile_idc(profile_idc); - let profile_for_resources = vk::VideoProfileInfoKHR::default() - .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H265) - .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(bit_depth_flags) - .chroma_bit_depth(bit_depth_flags) - .push(&mut h265_profile_for_resources); - - // Create input image - let (input_image, input_image_memory, input_image_view) = create_image( - &context, - aligned_width, - aligned_height, - picture_format, - false, - &profile_for_resources, - )?; - - // Determine DPB mode: use layered DPB when the driver does not advertise - // support for separate reference images (required for AMD RADV). - let supports_separate_dpb = capabilities - .flags - .contains(vk::VideoCapabilityFlagsKHR::SEPARATE_REFERENCE_IMAGES); - let use_layered_dpb = !supports_separate_dpb; - if use_layered_dpb { - info!("Using layered DPB (driver does not support separate reference images)"); - } - - // Create DPB images. - let (dpb_images, dpb_image_memories, dpb_image_views) = create_dpb_images( - &context, - aligned_width, - aligned_height, - reference_picture_format, - dpb_slot_count, - &profile_for_resources, - use_layered_dpb, - )?; - - // Create bitstream buffer. - let (bitstream_buffer, bitstream_buffer_memory) = - create_bitstream_buffer(&context, MIN_BITSTREAM_BUFFER_SIZE, &profile_for_resources)?; - - // Persistently map the bitstream buffer to avoid per-frame map/unmap overhead. - let bitstream_buffer_ptr = - map_bitstream_buffer(&context, bitstream_buffer_memory, MIN_BITSTREAM_BUFFER_SIZE)?; - - // Create command pool, buffers, and fences. - // Use the transfer queue family for upload commands when the encode queue - // doesn't support transfer operations (AMD RADV). - let upload_queue_family = context.transfer_queue_family(); - let cmd_resources = - create_command_resources(&context, encode_queue_family, upload_queue_family)?; - let command_pool = cmd_resources.command_pool; - let upload_command_pool = cmd_resources.upload_command_pool; - let upload_command_buffer = cmd_resources.upload_command_buffer; - let encode_command_buffer = cmd_resources.encode_command_buffer; - let upload_fence = cmd_resources.upload_fence; - let encode_fence = cmd_resources.encode_fence; - - // Clear the input image so padding between user dimensions and the - // aligned coded extent is zero-initialized. - clear_input_image( - &context, - &ClearImageParams { - command_buffer: upload_command_buffer, - fence: upload_fence, - queue: context.transfer_queue(), - image: input_image, - width: aligned_width, - height: aligned_height, - pixel_format: config.pixel_format, - bit_depth: config.bit_depth, - }, - )?; - - // Create query pool - let mut h265_profile_info_query = - vk::VideoEncodeH265ProfileInfoKHR::default().std_profile_idc(profile_idc); - - let mut profile_info_query = vk::VideoProfileInfoKHR::default() - .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H265) - .chroma_subsampling(chroma_subsampling) - .luma_bit_depth(bit_depth_flags) - .chroma_bit_depth(bit_depth_flags) - .push(&mut h265_profile_info_query); - - let mut encode_feedback_create = vk::QueryPoolVideoEncodeFeedbackCreateInfoKHR::default() - .encode_feedback_flags( - vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BUFFER_OFFSET - | vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN, - ); - - let query_pool_create_info = unsafe { - vk::QueryPoolCreateInfo::default() - .query_type(vk::QueryType::VIDEO_ENCODE_FEEDBACK_KHR) - .query_count(1) - .extend(&mut profile_info_query) - .push(&mut encode_feedback_create) - }; - - let query_pool = unsafe { - context - .device() - .create_query_pool(&query_pool_create_info, None) - } - .map_err(|e| PixelForgeError::QueryPool(e.to_string()))?; - - // Create DPB and GOP structure let mut dpb = DecodedPictureBuffer::new(); - let dpb_config = DpbConfig { - dpb_size: dpb_slot_count as u32, + dpb.h265.sequence_start(DpbConfig { + dpb_size: common.dpb_slot_count as u32, max_num_ref_frames: if config.b_frame_count > 0 { 2 } else { 1 }, use_multiple_references: config.b_frame_count > 0, max_long_term_refs: 0, - log2_max_frame_num_minus4: 0, // Not used in H.265 - log2_max_pic_order_cnt_lsb_minus4: 4, // max_poc_lsb = 256 + log2_max_frame_num_minus4: 0, + log2_max_pic_order_cnt_lsb_minus4: 4, num_temporal_layers: 1, - }; - dpb.h265.sequence_start(dpb_config); - - let mut gop = if config.b_frame_count > 0 { - GopStructure::new(config.gop_size, config.b_frame_count, config.gop_size) - } else { - GopStructure::new_ip_only(config.gop_size) - }; - - // Set GOP parameters to match SPS values - // log2_max_pic_order_cnt_lsb_minus4 = 4, so max_poc_lsb = 2^8 = 256 - gop.set_max_frame_num(4); // Not used in H.265 but set for compatibility - gop.set_max_poc_lsb(4); + }); - // Initialize DPB slot activation tracking - let dpb_slot_active = vec![false; dpb_slot_count]; - - info!("H.265 encoder created successfully"); - - let mut encoder = Self { - context, - config: config.clone(), + let codec = H265 { dpb, - gop, - aligned_width, - aligned_height, - video_queue_fn, - video_encode_fn, - session, - session_params: vk::VideoSessionParametersKHR::null(), - session_memory, - input_frame_num: 0, - encode_frame_num: 0, - input_image, - input_image_memory, - input_image_view, - input_image_layout: vk::ImageLayout::VIDEO_ENCODE_SRC_KHR, - dpb_images, - dpb_image_memories, - dpb_image_views, - dpb_slot_count, - use_layered_dpb, - bitstream_buffer, - bitstream_buffer_memory, - bitstream_buffer_ptr, - command_pool, - upload_command_pool, - upload_command_buffer, - upload_fence, - encode_command_buffer, - encode_fence, - query_pool, header_data: None, has_backward_reference: false, backward_reference_poc: 0, backward_reference_dpb_slot: 2, - current_dpb_slot: 0, l0_references: Vec::new(), - active_reference_count: max_active_reference_pictures as u32, + active_reference_count, profile_idc, - dpb_slot_active, }; - encoder.session_params = encoder.create_session_params(&color_desc)?; + let mut encoder = CodecEncoder { common, codec }; + let color_desc = config + .color_description + .unwrap_or(ColorDescription::bt709()); + encoder.common.session_params = encoder + .codec + .build_session_params(&encoder.common, &color_desc)?; + + info!("H.265 encoder created successfully"); Ok(encoder) } + + /// Profile IDC (used when rebuilding parameter sets). + pub(super) fn profile_idc(&self) -> u32 { + self.profile_idc + } } diff --git a/src/encoder/h265/mod.rs b/src/encoder/h265/mod.rs index 522f262..67dd3bb 100644 --- a/src/encoder/h265/mod.rs +++ b/src/encoder/h265/mod.rs @@ -1,23 +1,21 @@ -//! H.265/HEVC encoder implementation using Vulkan Video. +//! H.265/HEVC codec: the differences from the generic encoder. //! -//! This module implements H.265/HEVC video encoding using Vulkan Video extensions. +//! Shared machinery lives in `crate::encoder::codec`; this folder holds only +//! H.265's reference tracking (`H265`), its per-frame StdVideo* graph +//! (`record`), and its VPS/SPS/PPS generation (`session_params`). -mod api; -mod encode; mod init; +mod record; mod session_params; -use ash::vk; -use tracing::debug; - -use crate::encoder::dpb::DecodedPictureBuffer; -use crate::encoder::gop::GopStructure; -use crate::encoder::resources::{ - destroy_encoder_resources, upload_image_to_input, EncoderResources, UploadParams, +use crate::encoder::codec::{EncoderCommon, FramePlan, PictureSetup, VideoCodec}; +use crate::encoder::dpb::{ + DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType, }; -use crate::encoder::EncodeConfig; +use crate::encoder::pipeline::EncodeFuture; +use crate::encoder::ColorDescription; use crate::error::Result; -use crate::vulkan::VideoContext; +use ash::vk; /// H.265 Coding Tree Block (CTB) size in pixels. pub const CTB_SIZE: u32 = 32; @@ -26,155 +24,143 @@ pub const CTB_SIZE: u32 = 32; pub(crate) struct ReferenceInfo { pub dpb_slot: u8, pub poc: i32, + /// Display order (`pts`) of this reference, for reference frame invalidation. + pub display_order: u64, } -/// H.265 encoder. -pub struct H265Encoder { - context: VideoContext, - config: EncodeConfig, +/// H.265-specific encoder state. +pub(crate) struct H265 { dpb: DecodedPictureBuffer, - gop: GopStructure, - - /// Aligned width (to CTB size). - aligned_width: u32, - /// Aligned height (to CTB size). - aligned_height: u32, - - // Video session. - video_queue_fn: ash::khr::video_queue::Device, - video_encode_fn: ash::khr::video_encode_queue::Device, - session: vk::VideoSessionKHR, - session_params: vk::VideoSessionParametersKHR, - session_memory: Vec, - - // Frame counters. - input_frame_num: u64, - encode_frame_num: u64, - - // Resources - input_image: vk::Image, - input_image_memory: vk::DeviceMemory, - input_image_view: vk::ImageView, - /// Current Vulkan image layout of `input_image` (tracked to avoid UB when transitioning). - input_image_layout: vk::ImageLayout, - /// DPB images. - dpb_images: Vec, - dpb_image_memories: Vec, - dpb_image_views: Vec, - /// Number of DPB slots allocated. - dpb_slot_count: usize, - /// Whether the DPB uses a single layered image (true) or separate images (false). - use_layered_dpb: bool, - bitstream_buffer: vk::Buffer, - bitstream_buffer_memory: vk::DeviceMemory, - /// Persistently mapped pointer to the bitstream buffer (avoids per-frame map/unmap). - bitstream_buffer_ptr: *mut u8, - - // Command resources. - command_pool: vk::CommandPool, - upload_command_pool: vk::CommandPool, - upload_command_buffer: vk::CommandBuffer, - upload_fence: vk::Fence, - encode_command_buffer: vk::CommandBuffer, - encode_fence: vk::Fence, - query_pool: vk::QueryPool, - - // Parameter sets - cached header data (VPS/SPS/PPS) + /// Cached VPS/SPS/PPS header (invalidated when session params change). header_data: Option>, - - // Reference picture tracking. - /// Whether we have a backward reference (for B-frames, L1). has_backward_reference: bool, - /// POC of the L1 (backward) reference picture. backward_reference_poc: i32, - /// DPB slot for L1 (backward) reference. backward_reference_dpb_slot: u8, - /// Current DPB slot to use for setup (the reconstructed picture). - current_dpb_slot: u8, - /// Active L0 reference pictures (for P-frames). l0_references: Vec, - /// Number of active reference frames. active_reference_count: u32, - /// H.265 profile IDC (cached from initialization for session parameter recreation). profile_idc: u32, - - // DPB slot activation tracking. - /// Tracks which DPB slots have been activated (used at least once). - dpb_slot_active: Vec, } -impl H265Encoder { - /// Upload input frame from a GPU image. - /// - /// This copies from a source NV12 image directly to the encoder's input image, - /// avoiding any CPU-side data copies. The source image must be in NV12 format - /// with the same dimensions as the encoder configuration. The source image - /// should be in GENERAL layout. - fn upload_from_image(&mut self, src_image: vk::Image) -> Result<()> { - if src_image == self.input_image { - debug!("Source image is the encoder's input image, skipping upload copy"); - return Ok(()); +impl VideoCodec for H265 { + fn begin_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + if plan.is_idr() { + let dpb_config = DpbConfig { + dpb_size: common.dpb_slot_count as u32, + max_num_ref_frames: common.config.max_reference_frames, + use_multiple_references: common.config.b_frame_count > 0, + log2_max_frame_num_minus4: 0, + log2_max_pic_order_cnt_lsb_minus4: 4, + ..Default::default() + }; + self.dpb.h265.sequence_start(dpb_config); + self.l0_references.clear(); + self.has_backward_reference = false; + // All DPB slots become inactive at the start of a coded video sequence. + for active in &mut common.dpb_slot_active { + *active = false; + } } - let params = UploadParams { - upload_command_buffer: self.upload_command_buffer, - upload_fence: self.upload_fence, - src_image, - dst_image: self.input_image, - width: self.config.dimensions.width, - height: self.config.dimensions.height, - pixel_format: self.config.pixel_format, - input_image_layout: self.input_image_layout, - upload_queue: self.context.transfer_queue(), + let header = if plan.is_idr() { + if self.header_data.is_none() { + self.header_data = Some(self.build_header(common)?); + } + self.header_data.clone() + } else { + None }; - upload_image_to_input(&self.context, ¶ms)?; - - // Update tracked layout. - self.input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + Ok(PictureSetup { + frame_type: plan.frame_type(), + header, + }) + } - Ok(()) + fn record_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + self.record(common, plan) } -} -// SAFETY: The raw pointer bitstream_buffer_ptr is only used within the encoder's -// thread and is properly synchronized via Vulkan fences before access -unsafe impl Send for H265Encoder {} + fn end_picture(&mut self, common: &mut EncoderCommon, plan: &FramePlan) { + if !plan.is_reference() { + return; + } + let pic_order_cnt = plan.pic_order_cnt(); + let pic_type = if plan.is_idr() { + PictureType::Idr + } else if plan.is_b_frame() { + PictureType::B + } else { + PictureType::P + }; + let pic_info = PictureStartInfo { + frame_id: plan.display_order, + pic_order_cnt, + frame_num: 0, + pic_type, + is_reference: true, + ..Default::default() + }; + self.dpb.h265.picture_start(pic_info); + self.dpb.h265.picture_end(true); + + self.l0_references.insert( + 0, + ReferenceInfo { + dpb_slot: common.current_dpb_slot, + poc: pic_order_cnt, + display_order: plan.display_order, + }, + ); + while self.l0_references.len() > self.active_reference_count as usize { + self.l0_references.pop(); + } -impl Drop for H265Encoder { - fn drop(&mut self) { - unsafe { - // Wait on the queues used by the encoder rather than stalling - // the entire device. - let _ = self - .context - .device() - .queue_wait_idle(self.context.transfer_queue()); - if let Some(q) = self.context.video_encode_queue() { - let _ = self.context.device().queue_wait_idle(q); + if !plan.is_b_frame() { + let used_slots: Vec = self.l0_references.iter().map(|r| r.dpb_slot).collect(); + for i in 0..common.dpb_slot_count as u8 { + if !used_slots.contains(&i) { + common.current_dpb_slot = i; + break; + } } - destroy_encoder_resources( - self.context.device(), - &self.video_queue_fn, - &EncoderResources { - query_pool: self.query_pool, - upload_fence: self.upload_fence, - encode_fence: self.encode_fence, - command_pool: self.command_pool, - upload_command_pool: self.upload_command_pool, - bitstream_buffer: self.bitstream_buffer, - bitstream_buffer_memory: self.bitstream_buffer_memory, - input_image: self.input_image, - input_image_memory: self.input_image_memory, - input_image_view: self.input_image_view, - dpb_images: &self.dpb_images, - dpb_image_memories: &self.dpb_image_memories, - dpb_image_views: &self.dpb_image_views, - session: self.session, - session_params: self.session_params, - session_memory: &self.session_memory, - }, - ); } } + + fn invalidate_references( + &mut self, + _common: &mut EncoderCommon, + first_lost_display_order: u64, + ) -> bool { + // Every reference at or after the first lost frame is transitively + // undecodable on the client; keep only older survivors. H.265 declares + // its short-term reference picture set explicitly every frame from + // `l0_references` (see `record`), so dropping entries here is enough for + // the decoder to follow — no extra bitstream marking is required. + self.l0_references + .retain(|r| r.display_order < first_lost_display_order); + // A backward (L1) reference may also be tainted; B-frames are the only + // consumer of it, so simply forget it on invalidation. + self.has_backward_reference = false; + !self.l0_references.is_empty() + } + + fn create_session_params( + &self, + common: &EncoderCommon, + desc: &ColorDescription, + ) -> Result { + self.build_session_params(common, desc) + } + + fn invalidate_header_cache(&mut self) { + self.header_data = None; + } } diff --git a/src/encoder/h265/encode.rs b/src/encoder/h265/record.rs similarity index 56% rename from src/encoder/h265/encode.rs rename to src/encoder/h265/record.rs index d49eccc..8fc350c 100644 --- a/src/encoder/h265/encode.rs +++ b/src/encoder/h265/record.rs @@ -1,59 +1,52 @@ -//! H.265/HEVC encoder internal encoding implementation. -//! -//! This module handles the actual frame encoding using Vulkan Video. +//! H.265 per-frame encode recording: builds the StdVideo* graph for one frame +//! and submits it. -use super::H265Encoder; +use super::H265; -use crate::encoder::gop::{GopFrameType, GopPosition}; +use crate::encoder::codec::{EncoderCommon, FramePlan, RateControlPlan}; +use crate::encoder::pipeline::EncodeFuture; use crate::encoder::resources::{ prepare_encode_command_buffer, record_dpb_barriers, record_post_encode_dpb_barrier, - submit_encode_and_read_bitstream, MIN_BITSTREAM_BUFFER_SIZE, }; use crate::error::{PixelForgeError, Result}; use ash::vk; use ash::vk::TaggedStructure; use tracing::debug; -impl H265Encoder { - /// Encode a frame that has already been uploaded to the input image. - /// - /// This function: - /// 1. Records the video encode command buffer - /// 2. Sets up reference picture information - /// 3. Executes the encode operation - /// 4. Returns the encoded bitstream data - pub(super) fn encode_frame_internal( +impl H265 { + pub(super) fn record( &mut self, - gop_position: &GopPosition, - pic_order_cnt: i32, - is_idr: bool, - ) -> Result> { - // Prepare command buffer for recording. + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + let is_idr = plan.is_idr(); + let is_reference = plan.is_reference(); + let is_b_frame = plan.is_b_frame(); + let pic_order_cnt = plan.pic_order_cnt(); + + let slot = common.pipeline.current(); + let command_buffer = slot.encode_command_buffer; + let query_pool = slot.query_pool; + let bitstream_buffer = slot.bitstream_buffer; + let bitstream_buffer_size = slot.bitstream_buffer_size; + let input_image_view = slot.input_image_view; + unsafe { - prepare_encode_command_buffer( - self.context.device(), - self.encode_command_buffer, - self.query_pool, - )?; + prepare_encode_command_buffer(common.device(), command_buffer, query_pool)?; } - - // Transition DPB images for encode. let ref_dpb_slots: Vec = self.l0_references.iter().map(|r| r.dpb_slot).collect(); unsafe { record_dpb_barriers( - self.context.device(), - self.encode_command_buffer, - &self.dpb_images, - self.use_layered_dpb, - self.current_dpb_slot, + common.device(), + command_buffer, + &common.dpb_images, + common.use_layered_dpb, + common.current_dpb_slot, &ref_dpb_slots, - self.dpb_slot_active[self.current_dpb_slot as usize], + common.dpb_slot_active[common.current_dpb_slot as usize], ); } - // Determine picture type. - let is_b_frame = gop_position.frame_type == GopFrameType::B; - let is_reference = gop_position.is_reference; let slice_type = if is_idr { ash::vk::native::StdVideoH265SliceType_STD_VIDEO_H265_SLICE_TYPE_I } else if is_b_frame { @@ -61,7 +54,6 @@ impl H265Encoder { } else { ash::vk::native::StdVideoH265SliceType_STD_VIDEO_H265_SLICE_TYPE_P }; - let picture_type = if is_idr { ash::vk::native::StdVideoH265PictureType_STD_VIDEO_H265_PICTURE_TYPE_IDR } else if is_b_frame { @@ -70,7 +62,6 @@ impl H265Encoder { ash::vk::native::StdVideoH265PictureType_STD_VIDEO_H265_PICTURE_TYPE_P }; - // Build StdVideoEncodeH265SliceSegmentHeader. let slice_header_flags = ash::vk::native::StdVideoEncodeH265SliceSegmentHeaderFlags { _bitfield_align_1: [], _bitfield_1: ash::vk::native::StdVideoEncodeH265SliceSegmentHeaderFlags::new_bitfield_1( @@ -89,7 +80,6 @@ impl H265Encoder { 0, // reserved ), }; - let slice_header = ash::vk::native::StdVideoEncodeH265SliceSegmentHeader { flags: slice_header_flags, slice_type, @@ -108,66 +98,51 @@ impl H265Encoder { pWeightTable: std::ptr::null(), }; - // Build short-term reference picture set. - // For B-frames, we need both negative (past) and positive (future) references. - // For P-frames, we only need negative references. - let mut delta_poc_s0_minus1 = [0u16; 16]; // negative refs (past) - let mut delta_poc_s1_minus1 = [0u16; 16]; // positive refs (future) + // Short-term reference picture set: negative (past) refs for P, plus a + // positive (future) ref for B. + let mut delta_poc_s0_minus1 = [0u16; 16]; + let mut delta_poc_s1_minus1 = [0u16; 16]; let mut num_negative_pics: u8 = 0; let mut num_positive_pics: u8 = 0; let mut used_by_curr_pic_s0_flag: u16 = 0; let mut used_by_curr_pic_s1_flag: u16 = 0; if !is_idr && !self.l0_references.is_empty() { - // L0 references (negative/past) - // Calculate max_poc from config (2^(log2_max_pic_order_cnt_lsb_minus4 + 4) * 2). - // With log2_max_pic_order_cnt_lsb_minus4 = 4, max_poc = 2^8 * 2 = 512. - let max_poc = 1i32 << 9; // 512 - + // max_poc = 2^(log2_max_pic_order_cnt_lsb_minus4 + 4) * 2 = 512. + let max_poc = 1i32 << 9; let mut prev_delta_poc = 0; - for (i, ref_info) in self.l0_references.iter().enumerate() { if i >= 15 { break; - } // limit to 15 neg pics - - // Calculate delta POC with wraparound handling. - // delta_poc should be negative (reference is in the past). + } let mut delta_poc = ref_info.poc - pic_order_cnt; - // If delta_poc is positive and large, it means POC wrapped around. - // Adjust by subtracting max_poc to get the correct negative delta. if delta_poc > max_poc / 2 { delta_poc -= max_poc; } else if delta_poc < -max_poc / 2 { delta_poc += max_poc; } - let diff = prev_delta_poc - delta_poc; delta_poc_s0_minus1[num_negative_pics as usize] = (diff - 1).max(0) as u16; prev_delta_poc = delta_poc; - used_by_curr_pic_s0_flag |= 1 << num_negative_pics; num_negative_pics += 1; } - - // For B-frames, add L1 reference (positive/future) if is_b_frame && self.has_backward_reference { let delta_poc_l1 = self.backward_reference_poc - pic_order_cnt; - // delta_poc_s1 should be positive delta_poc_s1_minus1[0] = (delta_poc_l1 - 1).max(0) as u16; num_positive_pics = 1; - used_by_curr_pic_s1_flag = 1; // First positive reference is used + used_by_curr_pic_s1_flag = 1; } } - let rps_flags = ash::vk::native::StdVideoH265ShortTermRefPicSetFlags { - _bitfield_align_1: [], - _bitfield_1: ash::vk::native::StdVideoH265ShortTermRefPicSetFlags::new_bitfield_1(0, 0), - __bindgen_padding_0: [0; 3], - }; - let frame_rps = ash::vk::native::StdVideoH265ShortTermRefPicSet { - flags: rps_flags, + flags: ash::vk::native::StdVideoH265ShortTermRefPicSetFlags { + _bitfield_align_1: [], + _bitfield_1: ash::vk::native::StdVideoH265ShortTermRefPicSetFlags::new_bitfield_1( + 0, 0, + ), + __bindgen_padding_0: [0; 3], + }, delta_idx_minus1: 0, use_delta_flag: 0, abs_delta_rps_minus1: 0, @@ -183,15 +158,14 @@ impl H265Encoder { delta_poc_s1_minus1, }; - // Empty RPS for IDR frames. - let empty_rps_flags = ash::vk::native::StdVideoH265ShortTermRefPicSetFlags { - _bitfield_align_1: [], - _bitfield_1: ash::vk::native::StdVideoH265ShortTermRefPicSetFlags::new_bitfield_1(0, 0), - __bindgen_padding_0: [0; 3], - }; - let empty_rps = ash::vk::native::StdVideoH265ShortTermRefPicSet { - flags: empty_rps_flags, + flags: ash::vk::native::StdVideoH265ShortTermRefPicSetFlags { + _bitfield_align_1: [], + _bitfield_1: ash::vk::native::StdVideoH265ShortTermRefPicSetFlags::new_bitfield_1( + 0, 0, + ), + __bindgen_padding_0: [0; 3], + }, delta_idx_minus1: 0, use_delta_flag: 0, abs_delta_rps_minus1: 0, @@ -207,7 +181,6 @@ impl H265Encoder { delta_poc_s1_minus1: [0; 16], }; - // Build picture info flags. let picture_info_flags = ash::vk::native::StdVideoEncodeH265PictureInfoFlags { _bitfield_align_1: [], _bitfield_1: ash::vk::native::StdVideoEncodeH265PictureInfoFlags::new_bitfield_1( @@ -224,13 +197,11 @@ impl H265Encoder { ), }; - // Set up reference lists. const NO_REFERENCE_PICTURE: u8 = 0xFF; let mut ref_list0: [u8; 15] = [NO_REFERENCE_PICTURE; 15]; let mut ref_list1: [u8; 15] = [NO_REFERENCE_PICTURE; 15]; let (num_ref_l0, num_ref_l1) = if is_b_frame && self.has_backward_reference { - // B-frame logic if let Some(first_ref) = self.l0_references.first() { ref_list0[0] = first_ref.dpb_slot; ref_list1[0] = self.backward_reference_dpb_slot; @@ -239,29 +210,25 @@ impl H265Encoder { (0, 0) } } else if !is_idr && !self.l0_references.is_empty() { - // P-frame logic let count = self.l0_references.len(); for (i, ref_info) in self.l0_references.iter().enumerate() { if i < 15 { ref_list0[i] = ref_info.dpb_slot; } } - - // We use all available references as active (count.min(15), 0) } else { (0, 0) }; - let ref_lists_info_flags = ash::vk::native::StdVideoEncodeH265ReferenceListsInfoFlags { - _bitfield_align_1: [], - _bitfield_1: ash::vk::native::StdVideoEncodeH265ReferenceListsInfoFlags::new_bitfield_1( - 0, 0, 0, - ), - }; - let ref_lists_info = ash::vk::native::StdVideoEncodeH265ReferenceListsInfo { - flags: ref_lists_info_flags, + flags: ash::vk::native::StdVideoEncodeH265ReferenceListsInfoFlags { + _bitfield_align_1: [], + _bitfield_1: + ash::vk::native::StdVideoEncodeH265ReferenceListsInfoFlags::new_bitfield_1( + 0, 0, 0, + ), + }, num_ref_idx_l0_active_minus1: if num_ref_l0 > 0 { (num_ref_l0 - 1) as u8 } else { @@ -303,50 +270,43 @@ impl H265Encoder { short_term_ref_pic_set_idx: 0, }; - // Create slice NAL unit entry. - let constant_qp = match self.config.rate_control_mode { - crate::encoder::RateControlMode::Cqp | crate::encoder::RateControlMode::Disabled => { - self.config.quality_level as i32 - } - _ => 0, + let rc = RateControlPlan::new(&common.config, 26); + let constant_qp = if rc.is_disabled() { + common.config.quality_level as i32 + } else { + 0 }; let nalu_slice_entries = [vk::VideoEncodeH265NaluSliceSegmentInfoKHR::default() .constant_qp(constant_qp) .std_slice_segment_header(&slice_header)]; - // Create H.265 picture info. let mut h265_picture_info = vk::VideoEncodeH265PictureInfoKHR::default() .nalu_slice_segment_entries(&nalu_slice_entries) .std_picture_info(&picture_info); - // Set up source picture resource. + let coded_extent = vk::Extent2D { + width: common.aligned_width, + height: common.aligned_height, + }; + let src_picture_resource = vk::VideoPictureResourceInfoKHR::default() .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) + .coded_extent(coded_extent) .base_array_layer(0) - .image_view_binding(self.input_image_view); + .image_view_binding(input_image_view); - // Set up setup picture resource (reconstructed picture) let setup_picture_resource = vk::VideoPictureResourceInfoKHR::default() .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) + .coded_extent(coded_extent) .base_array_layer(0) - .image_view_binding(self.dpb_image_views[self.current_dpb_slot as usize]); + .image_view_binding(common.dpb_image_views[common.current_dpb_slot as usize]); - // Create reference info for setup slot. let std_reference_info_flags = ash::vk::native::StdVideoEncodeH265ReferenceInfoFlags { _bitfield_align_1: [], _bitfield_1: ash::vk::native::StdVideoEncodeH265ReferenceInfoFlags::new_bitfield_1( 0, 0, 0, ), }; - let std_reference_info = ash::vk::native::StdVideoEncodeH265ReferenceInfo { flags: std_reference_info_flags, PicOrderCntVal: pic_order_cnt, @@ -356,52 +316,38 @@ impl H265Encoder { let mut h265_setup_dpb_slot_info = vk::VideoEncodeH265DpbSlotInfoKHR::default().std_reference_info(&std_reference_info); - let setup_slot_info = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(self.current_dpb_slot as i32) + .slot_index(common.current_dpb_slot as i32) .picture_resource(&setup_picture_resource) .push(&mut h265_setup_dpb_slot_info); - // Setup slot for begin - always use -1 to indicate it's not yet active. - // (it will be written to during encoding) let mut h265_begin_dpb_slot_info = vk::VideoEncodeH265DpbSlotInfoKHR::default().std_reference_info(&std_reference_info); - let setup_slot_for_begin = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(-1) // Always -1 for setup slot + .slot_index(-1) .picture_resource(&setup_picture_resource) .push(&mut h265_begin_dpb_slot_info); - // Set up reference slots. - - // Storage vectors to keep data alive while pointers are used. - // We need capacity sufficient for all potential references (L0 + L1). + // Storage to keep pointers alive across the encode call. let mut ref_resources = Vec::with_capacity(16); let mut std_ref_infos = Vec::with_capacity(16); let mut h265_slot_infos = Vec::with_capacity(16); - - // Final slot lists let mut reference_slots = Vec::with_capacity(16); - let mut reference_slots_for_begin = Vec::with_capacity(17); // +1 for setup slot + let mut reference_slots_for_begin = Vec::with_capacity(17); reference_slots_for_begin.push(setup_slot_for_begin); let has_l0_ref = !is_idr && !self.l0_references.is_empty(); let has_l1_ref = is_b_frame && self.has_backward_reference; - // Phase 1: Populate data storage (resources and std infos) if has_l0_ref { for ref_info in &self.l0_references { - let ref_resource = vk::VideoPictureResourceInfoKHR::default() - .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) - .base_array_layer(0) - .image_view_binding(self.dpb_image_views[ref_info.dpb_slot as usize]); - - ref_resources.push(ref_resource); - + ref_resources.push( + vk::VideoPictureResourceInfoKHR::default() + .coded_offset(vk::Offset2D { x: 0, y: 0 }) + .coded_extent(coded_extent) + .base_array_layer(0) + .image_view_binding(common.dpb_image_views[ref_info.dpb_slot as usize]), + ); let mut std_info = std_reference_info; std_info.PicOrderCntVal = ref_info.poc; std_info.pic_type = @@ -409,21 +355,16 @@ impl H265Encoder { std_ref_infos.push(std_info); } } - if has_l1_ref { - let ref_resource = vk::VideoPictureResourceInfoKHR::default() - .coded_offset(vk::Offset2D { x: 0, y: 0 }) - .coded_extent(vk::Extent2D { - width: self.aligned_width, - height: self.aligned_height, - }) - .base_array_layer(0) - .image_view_binding( - self.dpb_image_views[self.backward_reference_dpb_slot as usize], - ); - - ref_resources.push(ref_resource); - + ref_resources.push( + vk::VideoPictureResourceInfoKHR::default() + .coded_offset(vk::Offset2D { x: 0, y: 0 }) + .coded_extent(coded_extent) + .base_array_layer(0) + .image_view_binding( + common.dpb_image_views[self.backward_reference_dpb_slot as usize], + ), + ); let mut std_info = std_reference_info; std_info.PicOrderCntVal = self.backward_reference_poc; std_info.pic_type = @@ -431,20 +372,13 @@ impl H265Encoder { std_ref_infos.push(std_info); } - // Phase 2: Populate chain structs (referencing std_ref_infos) for std_info in &std_ref_infos { - let h265_ref_slot_info = - vk::VideoEncodeH265DpbSlotInfoKHR::default().std_reference_info(std_info); - h265_slot_infos.push(h265_ref_slot_info); + h265_slot_infos + .push(vk::VideoEncodeH265DpbSlotInfoKHR::default().std_reference_info(std_info)); } - // Phase 3: Populate reference slots (referencing ref_resources and h265_slot_infos) - // We know the mapping corresponds 1:1 by index because we pushed in the same order. - // Re-construct the list of DPB slots to assign correct slot_index. - let mut stored_indices_count = 0; let mut slot_infos_iter = h265_slot_infos.iter_mut(); - if has_l0_ref { for ref_info in &self.l0_references { if let Some(h265_slot_info) = slot_infos_iter.next() { @@ -452,130 +386,85 @@ impl H265Encoder { .slot_index(ref_info.dpb_slot as i32) .picture_resource(&ref_resources[stored_indices_count]) .push(h265_slot_info); - reference_slots.push(ref_slot); reference_slots_for_begin.push(ref_slot); stored_indices_count += 1; } } } - if has_l1_ref { if let Some(h265_slot_info) = slot_infos_iter.next() { let ref_slot = vk::VideoReferenceSlotInfoKHR::default() .slot_index(self.backward_reference_dpb_slot as i32) .picture_resource(&ref_resources[stored_indices_count]) .push(h265_slot_info); - reference_slots.push(ref_slot); reference_slots_for_begin.push(ref_slot); } } - // Rate control setup. - let (rc_mode, average_bitrate, max_bitrate, qp) = match self.config.rate_control_mode { - crate::encoder::RateControlMode::Cqp | crate::encoder::RateControlMode::Disabled => ( - vk::VideoEncodeRateControlModeFlagsKHR::DISABLED, - 0, - 0, - self.config.quality_level as i32, - ), - crate::encoder::RateControlMode::Cbr => ( - vk::VideoEncodeRateControlModeFlagsKHR::CBR, - self.config.target_bitrate, - self.config.target_bitrate, - 26, - ), - crate::encoder::RateControlMode::Vbr => ( - vk::VideoEncodeRateControlModeFlagsKHR::VBR, - self.config.target_bitrate, - self.config.max_bitrate, - 26, - ), - }; - - let min_qp_val = if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { - qp - } else { - 26 - }; - let max_qp_val = if rc_mode == vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { - qp - } else { - 51 - }; - + // Rate control. + let qp_min = if rc.is_disabled() { rc.qp as i32 } else { 26 }; + let qp_max = if rc.is_disabled() { rc.qp as i32 } else { 51 }; let min_qp = vk::VideoEncodeH265QpKHR { - qp_i: min_qp_val, - qp_p: min_qp_val, - qp_b: min_qp_val, + qp_i: qp_min, + qp_p: qp_min, + qp_b: qp_min, }; - let max_qp = vk::VideoEncodeH265QpKHR { - qp_i: max_qp_val, - qp_p: max_qp_val, - qp_b: max_qp_val, + qp_i: qp_max, + qp_p: qp_max, + qp_b: qp_max, }; - let mut h265_rc_layer_info = vk::VideoEncodeH265RateControlLayerInfoKHR::default() .min_qp(min_qp) .max_qp(max_qp); - let rc_layer_info = vk::VideoEncodeRateControlLayerInfoKHR::default() - .average_bitrate(average_bitrate as u64) - .max_bitrate(max_bitrate as u64) - .frame_rate_numerator(self.config.frame_rate_numerator) - .frame_rate_denominator(self.config.frame_rate_denominator) + .average_bitrate(rc.average_bitrate as u64) + .max_bitrate(rc.max_bitrate as u64) + .frame_rate_numerator(common.config.frame_rate_numerator) + .frame_rate_denominator(common.config.frame_rate_denominator) .push(&mut h265_rc_layer_info); - let rc_layers = [rc_layer_info]; let mut h265_rc_info = vk::VideoEncodeH265RateControlInfoKHR::default() - .gop_frame_count(self.config.gop_size) - .idr_period(self.config.gop_size) - .consecutive_b_frame_count(self.config.b_frame_count); - - let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc_mode); - - if rc_mode != vk::VideoEncodeRateControlModeFlagsKHR::DISABLED { + .gop_frame_count(common.config.gop_size) + .idr_period(common.config.gop_size) + .consecutive_b_frame_count(common.config.b_frame_count); + let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc.mode); + if !rc.is_disabled() { rc_info = rc_info .layers(&rc_layers) - .virtual_buffer_size_in_ms(self.config.virtual_buffer_size_ms) - .initial_virtual_buffer_size_in_ms(self.config.initial_virtual_buffer_size_ms); + .virtual_buffer_size_in_ms(common.config.virtual_buffer_size_ms) + .initial_virtual_buffer_size_in_ms(common.config.initial_virtual_buffer_size_ms); } - // Begin video coding. - let is_first_frame = self.encode_frame_num == 0; - + let is_first_frame = plan.is_first_frame(); let begin_coding_info = if is_first_frame { vk::VideoBeginCodingInfoKHR::default() - .video_session(self.session) - .video_session_parameters(self.session_params) + .video_session(common.session) + .video_session_parameters(common.session_params) .reference_slots(&reference_slots_for_begin) .push(&mut h265_rc_info) } else { vk::VideoBeginCodingInfoKHR::default() - .video_session(self.session) - .video_session_parameters(self.session_params) + .video_session(common.session) + .video_session_parameters(common.session_params) .reference_slots(&reference_slots_for_begin) .push(&mut h265_rc_info) .push(&mut rc_info) }; unsafe { - (self.video_queue_fn.fp().cmd_begin_video_coding_khr)( - self.encode_command_buffer, + (common.video_queue_fn.fp().cmd_begin_video_coding_khr)( + command_buffer, &begin_coding_info, ); } - // Reset video coding state for the first frame. - // Combine RESET + RATE_CONTROL + QUALITY_LEVEL into a single control command. - // This matches FFmpeg's approach and is required for AMD RADV. if is_first_frame { let mut quality_level_info = vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); - let control_info = vk::VideoCodingControlInfoKHR::default() .flags( vk::VideoCodingControlFlagsKHR::RESET @@ -585,103 +474,59 @@ impl H265Encoder { .push(&mut rc_info) .push(&mut h265_rc_info) .push(&mut quality_level_info); - unsafe { - (self.video_queue_fn.fp().cmd_control_video_coding_khr)( - self.encode_command_buffer, + (common.video_queue_fn.fp().cmd_control_video_coding_khr)( + command_buffer, &control_info, ); } } - // Encode command. let encode_info = vk::VideoEncodeInfoKHR::default() .flags(vk::VideoEncodeFlagsKHR::empty()) .src_picture_resource(src_picture_resource) .setup_reference_slot(&setup_slot_info) .reference_slots(&reference_slots) - .dst_buffer(self.bitstream_buffer) + .dst_buffer(bitstream_buffer) .dst_buffer_offset(0) - .dst_buffer_range(MIN_BITSTREAM_BUFFER_SIZE as u64) + .dst_buffer_range(bitstream_buffer_size as u64) .push(&mut h265_picture_info); + debug!( + "h265 submit frame {}: idr={}, num_refs={}, cur_slot={}", + plan.encode_index, + is_idr, + self.l0_references.len(), + common.current_dpb_slot + ); + unsafe { - self.context.device().cmd_begin_query( - self.encode_command_buffer, - self.query_pool, + let device = common.device(); + device.cmd_begin_query( + command_buffer, + query_pool, 0, vk::QueryControlFlags::empty(), ); + (common.video_encode_fn.fp().cmd_encode_video_khr)(command_buffer, &encode_info); + device.cmd_end_query(command_buffer, query_pool, 0); - (self.video_encode_fn.fp().cmd_encode_video_khr)( - self.encode_command_buffer, - &encode_info, - ); - - self.context - .device() - .cmd_end_query(self.encode_command_buffer, self.query_pool, 0); - } - - // Add DPB synchronization barrier after encoding. - unsafe { record_post_encode_dpb_barrier( - self.context.device(), - self.encode_command_buffer, - &self.dpb_images, - self.use_layered_dpb, - self.current_dpb_slot, + device, + command_buffer, + &common.dpb_images, + common.use_layered_dpb, + common.current_dpb_slot, ); - } - // End video coding. - let end_coding_info = vk::VideoEndCodingInfoKHR::default(); - unsafe { - (self.video_queue_fn.fp().cmd_end_video_coding_khr)( - self.encode_command_buffer, - &end_coding_info, - ); - } + let end_coding_info = vk::VideoEndCodingInfoKHR::default(); + (common.video_queue_fn.fp().cmd_end_video_coding_khr)(command_buffer, &end_coding_info); - // End command buffer. - unsafe { - self.context - .device() - .end_command_buffer(self.encode_command_buffer) + device + .end_command_buffer(command_buffer) + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - - // Submit, wait, and read bitstream. - let encode_queue = self.context.video_encode_queue().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) - })?; - - debug!( - "Submitting frame {} to GPU: idr={}, num_refs={}, cur_slot={}", - self.encode_frame_num, - is_idr, - self.l0_references.len(), - self.current_dpb_slot - ); - - let gpu_start = std::time::Instant::now(); - - let encoded_data = unsafe { - submit_encode_and_read_bitstream( - self.context.device(), - self.encode_command_buffer, - self.encode_fence, - encode_queue, - self.query_pool, - self.bitstream_buffer_ptr, - )? - }; - - debug!("GPU encode took {:?}", gpu_start.elapsed()); - - // Mark DPB slot as active. - self.dpb_slot_active[self.current_dpb_slot as usize] = true; - Ok(encoded_data) + common.submit_frame() } } diff --git a/src/encoder/h265/session_params.rs b/src/encoder/h265/session_params.rs index b6082d9..4deca2f 100644 --- a/src/encoder/h265/session_params.rs +++ b/src/encoder/h265/session_params.rs @@ -1,40 +1,40 @@ -use super::H265Encoder; +//! H.265 parameter-set generation: VPS/SPS/PPS construction and header retrieval. +use super::H265; + +use crate::encoder::codec::EncoderCommon; +use crate::encoder::resources::get_encoded_session_params; use crate::encoder::{BitDepth, ColorDescription, PixelFormat}; use crate::error::{PixelForgeError, Result}; use ash::vk; use ash::vk::TaggedStructure; use std::ptr; -impl H265Encoder { +impl H265 { /// Build VPS/SPS/PPS and create Vulkan video session parameters. - /// - /// This is used both during initial encoder creation and by - /// `set_color_description()` to rebuild session parameters with - /// updated VUI color metadata. Keeping a single implementation - /// ensures the parameter sets stay bit-for-bit consistent. - pub(crate) fn create_session_params( + pub(super) fn build_session_params( &self, + common: &EncoderCommon, desc: &ColorDescription, ) -> Result { - // CTB/CB size constants. + let config = &common.config; + let ctb_log2_size_y: u8 = 5; let min_cb_log2_size_y: u8 = 4; let log2_min_transform_block_size: u8 = 2; let log2_max_transform_block_size: u8 = 5; - let pic_width_in_luma_samples = self.aligned_width; - let pic_height_in_luma_samples = self.aligned_height; + let pic_width_in_luma_samples = common.aligned_width; + let pic_height_in_luma_samples = common.aligned_height; - let (sub_width_c, sub_height_c) = match self.config.pixel_format { + let (sub_width_c, sub_height_c) = match config.pixel_format { PixelFormat::Yuv420 => (2u32, 2u32), PixelFormat::Yuv444 => (1u32, 1u32), _ => (2u32, 2u32), }; - let conf_win_right_offset = - (self.aligned_width - self.config.dimensions.width) / sub_width_c; + let conf_win_right_offset = (common.aligned_width - config.dimensions.width) / sub_width_c; let conf_win_bottom_offset = - (self.aligned_height - self.config.dimensions.height) / sub_height_c; + (common.aligned_height - config.dimensions.height) / sub_height_c; let conformance_window_flag = conf_win_right_offset > 0 || conf_win_bottom_offset > 0; let profile_tier_level = ash::vk::native::StdVideoH265ProfileTierLevel { @@ -45,13 +45,13 @@ impl H265Encoder { ), __bindgen_padding_0: [0; 3], }, - general_profile_idc: self.profile_idc, + general_profile_idc: self.profile_idc(), general_level_idc: ash::vk::native::StdVideoH265LevelIdc_STD_VIDEO_H265_LEVEL_IDC_5_1, }; let dec_pic_buf_mgr = ash::vk::native::StdVideoH265DecPicBufMgr { max_latency_increase_plus1: [0; 7], - max_dec_pic_buffering_minus1: [(self.dpb_slot_count - 1) as u8, 0, 0, 0, 0, 0, 0], + max_dec_pic_buffering_minus1: [(common.dpb_slot_count - 1) as u8, 0, 0, 0, 0, 0, 0], max_num_reorder_pics: [0; 7], }; @@ -174,12 +174,12 @@ impl H265Encoder { pHrdParameters: ptr::null(), }; - let bit_depth_minus8: u8 = match self.config.bit_depth { + let bit_depth_minus8: u8 = match config.bit_depth { BitDepth::Eight => 0, BitDepth::Ten => 2, }; - let chroma_format_idc = match self.config.pixel_format { + let chroma_format_idc = match config.pixel_format { PixelFormat::Yuv420 => { ash::vk::native::StdVideoH265ChromaFormatIdc_STD_VIDEO_H265_CHROMA_FORMAT_IDC_420 } @@ -189,7 +189,7 @@ impl H265Encoder { _ => { return Err(PixelForgeError::InvalidInput(format!( "Unsupported pixel format for H.265: {:?}", - self.config.pixel_format + config.pixel_format ))); } }; @@ -336,14 +336,17 @@ impl H265Encoder { let mut quality_level_info = vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); let params_create_info = vk::VideoSessionParametersCreateInfoKHR::default() - .video_session(self.session) + .video_session(common.session) .push(&mut h265_params_create_info) .push(&mut quality_level_info); let mut session_params = vk::VideoSessionParametersKHR::null(); let result = unsafe { - (self.video_queue_fn.fp().create_video_session_parameters_khr)( - self.context.device().handle(), + (common + .video_queue_fn + .fp() + .create_video_session_parameters_khr)( + common.device().handle(), ¶ms_create_info, ptr::null(), &mut session_params, @@ -355,7 +358,29 @@ impl H265Encoder { result ))); } - Ok(session_params) } + + /// Retrieve the encoded VPS/SPS/PPS NAL units to prepend to an IDR frame. + pub(super) fn build_header(&self, common: &EncoderCommon) -> Result> { + let mut h265_get_info = vk::VideoEncodeH265SessionParametersGetInfoKHR::default() + .write_std_vps(true) + .write_std_sps(true) + .write_std_pps(true) + .std_vps_id(0) + .std_sps_id(0) + .std_pps_id(0); + let get_info = vk::VideoEncodeSessionParametersGetInfoKHR::default() + .video_session_parameters(common.session_params) + .push(&mut h265_get_info); + let mut h265_feedback = vk::VideoEncodeH265SessionParametersFeedbackInfoKHR::default(); + let mut feedback = + vk::VideoEncodeSessionParametersFeedbackInfoKHR::default().push(&mut h265_feedback); + get_encoded_session_params( + &common.context, + &common.video_encode_fn, + &get_info, + &mut feedback, + ) + } } diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs index 983ba12..fbb1a02 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -7,10 +7,12 @@ pub mod av1; pub mod bitwriter; +pub(crate) mod codec; pub mod dpb; pub mod gop; pub mod h264; pub mod h265; +pub(crate) mod pipeline; pub mod reorder; pub mod resources; @@ -385,6 +387,8 @@ impl EncodeConfig { } } +pub use pipeline::EncodeFuture; + /// Encoded video packet. #[derive(Debug, Clone)] pub struct EncodedPacket { @@ -400,45 +404,64 @@ pub struct EncodedPacket { pub dts: u64, } -/// Video encoder supporting multiple codecs. +/// The codec-erased operations every [`codec::CodecEncoder`] exposes. /// -/// The encoder is implemented as an enum to dispatch to codec-specific implementations. -// Allow large_enum_variant: H265Encoder is currently a stub. When fully implemented, -// it will be similar in size to H264Encoder, making the size difference negligible. -#[allow(clippy::large_enum_variant)] -pub enum Encoder { - /// H.264/AVC encoder. - H264(self::h264::H264Encoder), - /// H.265/HEVC encoder. - H265(self::h265::H265Encoder), - /// AV1 encoder. - AV1(self::av1::AV1Encoder), +/// One blanket impl covers all codecs, so [`Encoder`] can hold any of them +/// behind a single boxed pointer instead of an enum that re-dispatches by hand. +trait EncoderApi: Send { + fn input_image(&self) -> vk::Image; + fn encode(&mut self, src_image: vk::Image) -> Result; + fn flush(&mut self) -> Result<()>; + fn request_idr(&mut self); + fn invalidate_reference_frames(&mut self, first_lost_display_order: u64); + fn set_color_description(&mut self, desc: ColorDescription) -> Result<()>; +} + +impl EncoderApi for codec::CodecEncoder { + fn input_image(&self) -> vk::Image { + codec::CodecEncoder::input_image(self) + } + fn encode(&mut self, src_image: vk::Image) -> Result { + codec::CodecEncoder::encode(self, src_image) + } + fn flush(&mut self) -> Result<()> { + codec::CodecEncoder::flush(self) + } + fn request_idr(&mut self) { + codec::CodecEncoder::request_idr(self) + } + fn invalidate_reference_frames(&mut self, first_lost_display_order: u64) { + codec::CodecEncoder::invalidate_reference_frames(self, first_lost_display_order) + } + fn set_color_description(&mut self, desc: ColorDescription) -> Result<()> { + codec::CodecEncoder::set_color_description(self, desc) + } } +/// Video encoder supporting multiple codecs. +/// +/// Constructed via [`Encoder::new`], which selects the codec from the config and +/// boxes the corresponding `codec::CodecEncoder`. All codecs share one generic +/// implementation; this type just erases which one is in use. +pub struct Encoder(Box); + impl Encoder { + /// Create a new encoder for the codec named in `config`. + pub fn new(context: VideoContext, config: EncodeConfig) -> Result { + let inner: Box = match config.codec { + Codec::H264 => Box::new(self::h264::H264::create(context, config)?), + Codec::H265 => Box::new(self::h265::H265::create(context, config)?), + Codec::AV1 => Box::new(self::av1::Av1::create(context, config)?), + }; + Ok(Encoder(inner)) + } + /// Get the internal input image. /// /// This image can be used as a target for `ColorConverter::convert` to avoid /// an intermediate copy. pub fn input_image(&self) -> vk::Image { - match self { - Encoder::H264(encoder) => encoder.input_image(), - Encoder::H265(encoder) => encoder.input_image(), - Encoder::AV1(encoder) => encoder.input_image(), - } - } - - /// Create a new encoder. - pub fn new(context: VideoContext, config: EncodeConfig) -> Result { - match config.codec { - Codec::H264 => Ok(Encoder::H264(self::h264::H264Encoder::new( - context, config, - )?)), - Codec::H265 => Ok(Encoder::H265(self::h265::H265Encoder::new( - context, config, - )?)), - Codec::AV1 => Ok(Encoder::AV1(self::av1::AV1Encoder::new(context, config)?)), - } + self.0.input_image() } /// Encode a frame from a GPU image. @@ -446,6 +469,11 @@ impl Encoder { /// This accepts a source NV12 (YUV420) or planar YUV444 image on the GPU and encodes it directly. /// The source image must match the format and dimensions in the encoder configuration. /// + /// Encoding is asynchronous: this submits the frame without blocking and + /// returns an [`EncodeFuture`] that resolves with the encoded packet once the + /// GPU finishes and a background readback thread reads it back. The call only + /// blocks if every pipeline slot is still in flight (backpressure). + /// /// Use `InputImage` to create an image from YUV data: /// ```no_run /// use pixelforge::{InputImage, Encoder, EncodeConfig, EncodeBitDepth, PixelFormat, VideoContext, Codec}; @@ -466,35 +494,52 @@ impl Encoder { /// # let yuv_data = vec![0u8; 1920 * 1080 * 3 / 2]; /// input.upload_yuv420(&yuv_data)?; /// - /// // Encode the image - /// let packets = encoder.encode(input.image())?; + /// // Submit the frame and await its packet. + /// let future = encoder.encode(input.image())?; + /// let packet = pollster::block_on(future)?; + /// // ... use packet.data ... /// # Ok(()) /// # } /// ``` - pub fn encode(&mut self, src_image: vk::Image) -> Result> { - match self { - Encoder::H264(encoder) => encoder.encode(src_image), - Encoder::H265(encoder) => encoder.encode(src_image), - Encoder::AV1(encoder) => encoder.encode(src_image), - } + pub fn encode(&mut self, src_image: vk::Image) -> Result { + self.0.encode(src_image) } - /// Flush the encoder and get remaining packets. - pub fn flush(&mut self) -> Result> { - match self { - Encoder::H264(encoder) => encoder.flush(), - Encoder::H265(encoder) => encoder.flush(), - Encoder::AV1(encoder) => encoder.flush(), - } + /// Wait for all in-flight frames to finish encoding (end-of-stream barrier). + /// + /// Packets are delivered through the [`EncodeFuture`]s returned by + /// [`Encoder::encode`]; await those to obtain them. Once this returns, every + /// outstanding future has been resolved. + pub fn flush(&mut self) -> Result<()> { + self.0.flush() } /// Request that the next frame be an IDR frame. pub fn request_idr(&mut self) { - match self { - Encoder::H264(encoder) => encoder.request_idr(), - Encoder::H265(encoder) => encoder.request_idr(), - Encoder::AV1(encoder) => encoder.request_idr(), - } + self.0.request_idr() + } + + /// Recover from packet loss without a full IDR (reference frame invalidation). + /// + /// When a client reports that it could not decode a range of frames, every + /// reference picture the encoder still holds from the earliest lost frame + /// onward is transitively undecodable on the client. This drops those + /// references so the next P-frame is predicted from the most recent + /// *surviving* reference instead — a much cheaper recovery than re-sending a + /// full keyframe. If no reference survives (the loss covers the encoder's + /// whole reference window), it transparently falls back to forcing an IDR. + /// + /// `first_lost_display_order` is the display order — the `pts` reported on + /// [`EncodedPacket`]s — of the earliest frame the client lost. + /// + /// Effective recovery requires the encoder to keep more than one reference + /// (see [`EncodeConfig::with_max_reference_frames`]); with a single + /// reference this necessarily falls back to an IDR. This applies uniformly + /// across H.264, H.265, and AV1 — each keeps a multi-reference window and + /// re-anchors prediction to the most recent survivor, falling back to an + /// IDR (AV1: key frame) only when the loss covers the entire window. + pub fn invalidate_reference_frames(&mut self, first_lost_display_order: u64) { + self.0.invalidate_reference_frames(first_lost_display_order) } /// Update the color description (VUI parameters) for the encoder. @@ -503,11 +548,7 @@ impl Encoder { /// header containing the new color description. The next frame will be encoded as /// an IDR/key frame with the new parameters. pub fn set_color_description(&mut self, desc: ColorDescription) -> Result<()> { - match self { - Encoder::H264(encoder) => encoder.set_color_description(desc), - Encoder::H265(encoder) => encoder.set_color_description(desc), - Encoder::AV1(encoder) => encoder.set_color_description(desc), - } + self.0.set_color_description(desc) } } diff --git a/src/encoder/pipeline.rs b/src/encoder/pipeline.rs new file mode 100644 index 0000000..3bd5e47 --- /dev/null +++ b/src/encoder/pipeline.rs @@ -0,0 +1,527 @@ +//! Asynchronous (push) encode pipelining, shared by all codecs. +//! +//! Each in-flight frame owns an [`EncodeSlot`] (its own input image, bitstream +//! buffer, encode command buffer, fence and query pool). [`EncodePipeline`] +//! rotates through the slots so that the CPU can record and submit frame N+1 +//! while the GPU is still encoding frame N, instead of blocking on a fence after +//! every frame. +//! +//! Bitstream readback is performed off the calling thread: a single background +//! *completion thread* waits on each submission's fence, copies the bitstream +//! out, and resolves that frame's [`EncodeFuture`] the moment the GPU signals — +//! rather than deferring readback until the slot is reused. This delivers each +//! packet at roughly the GPU encode time instead of one or two `encode()` calls +//! later. Each `encode()` returns its own future (backed by a oneshot channel), +//! so packet delivery is paired with its submission by construction rather than +//! by a shared channel and a separate ordering convention. +//! +//! The DPB images and video session are shared across slots, so encode +//! submissions must still run in DPB order on the GPU. That ordering is enforced +//! with a single timeline semaphore (each submit waits on the previous submit's +//! value and signals its own). Only the calling thread ever touches the queue or +//! the timeline; the completion thread only waits on fences and reads bitstream +//! buffers, so the two never race on the same Vulkan object. +//! +//! Slot reuse is coordinated with a per-slot "busy" flag ([`SlotSync`]): a slot +//! is busy from submit until the completion thread has finished reading its +//! bitstream. The calling thread waits for the current slot to be free before +//! converting into / recording over it, which also covers the write-after-read +//! hazard on the shared input image. + +use std::future::Future; +use std::pin::Pin; +use std::sync::mpsc::{Receiver, Sender}; +use std::sync::{Arc, Condvar, Mutex}; +use std::task::{Context, Poll}; +use std::thread::JoinHandle; + +use ash::vk; +use futures_channel::oneshot; + +use crate::encoder::resources::{ + clear_input_image, create_bitstream_buffer, create_encode_feedback_query_pool, create_image, + create_timeline_semaphore, map_bitstream_buffer, submit_encode_only, wait_and_read_bitstream, + ClearImageParams, +}; +use crate::encoder::{BitDepth, EncodedPacket, FrameType, PixelFormat}; +use crate::error::{PixelForgeError, Result}; +use crate::vulkan::VideoContext; + +/// A handle to the packet a single `EncodePipeline::submit_current` will +/// eventually produce. +/// +/// Returned by `Encoder::encode`, one per submitted frame. Awaiting it yields +/// that frame's [`EncodedPacket`] once the GPU finishes encoding and the +/// completion thread reads the bitstream back. Futures resolve in submission +/// order (one submission → one readback → one packet) as long as you await them +/// in the order `encode` returned them. +/// +/// Dropping the future before it resolves is harmless: the completion thread +/// still reads the slot back (its send simply finds no receiver) and frees it. +pub struct EncodeFuture { + rx: oneshot::Receiver>, +} + +impl Future for EncodeFuture { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match Pin::new(&mut self.rx).poll(cx) { + Poll::Ready(Ok(result)) => Poll::Ready(result), + // The sender was dropped without sending. The completion thread only + // drops a sender after sending, so this indicates encoder teardown. + Poll::Ready(Err(oneshot::Canceled)) => { + Poll::Ready(Err(PixelForgeError::CommandBuffer( + "encode cancelled: encoder shut down before the frame was read back" + .to_string(), + ))) + } + Poll::Pending => Poll::Pending, + } + } +} + +/// Number of encode submissions allowed to be in flight at once. +/// +/// See the module docs and `docs` discussion for why 2 is the sweet spot: it +/// fully overlaps capture/convert/upload of the next frame with the GPU encode +/// of the current one, while keeping the GPU-serialized DPB chain from growing. +pub(crate) const ENCODE_PIPELINE_DEPTH: usize = 2; + +/// Per-frame packet info captured at submit time and attached to the bitstream +/// when the completion thread reads the slot back. +pub(crate) struct SlotPacketMetadata { + pub frame_type: FrameType, + pub is_key_frame: bool, + pub pts: u64, + pub dts: u64, + /// Codec header (SPS/PPS, VPS/SPS/PPS, or AV1 sequence header) to prepend. + /// `Some` only for frames that carry one (e.g. the IDR/key frame). + pub header: Option>, +} + +/// A raw pointer wrapper asserting `Send` so the persistently-mapped bitstream +/// pointer can be handed to the completion thread. The memory is only read by +/// that thread, and only after the encode fence has signalled. +struct SendPtr(*const u8); +// SAFETY: the pointed-to bitstream buffer is host-coherent, persistently mapped +// for the lifetime of the slot, and read exclusively by the completion thread +// after the encode fence signals. The calling thread does not touch the buffer +// until the slot is marked free again (after this read completes). +unsafe impl Send for SendPtr {} + +/// A submission handed from the calling thread to the completion thread. +struct WorkItem { + slot_index: usize, + fence: vk::Fence, + query_pool: vk::QueryPool, + bitstream_ptr: SendPtr, + metadata: SlotPacketMetadata, + /// Resolves this frame's [`EncodeFuture`] once the bitstream is read back. + result_tx: oneshot::Sender>, +} + +/// Cross-thread per-slot readiness. A slot is "busy" from the moment its encode +/// is submitted until the completion thread has finished reading its bitstream. +struct SlotSync { + busy: Mutex>, + cv: Condvar, +} + +impl SlotSync { + fn new(slot_count: usize) -> Self { + Self { + busy: Mutex::new(vec![false; slot_count]), + cv: Condvar::new(), + } + } + + /// Block until slot `index` is free (its previous encode has been read back). + fn wait_free(&self, index: usize) { + let mut busy = self.busy.lock().unwrap(); + while busy[index] { + busy = self.cv.wait(busy).unwrap(); + } + } + + /// Block until every slot is free (no submissions in flight). + fn wait_all_free(&self) { + let mut busy = self.busy.lock().unwrap(); + while busy.iter().any(|b| *b) { + busy = self.cv.wait(busy).unwrap(); + } + } + + /// Mark a slot busy at submit time. No notify: nobody waits to *enter* busy. + fn set_busy(&self, index: usize) { + self.busy.lock().unwrap()[index] = true; + } + + /// Mark a slot free once its bitstream has been read; wake any waiters. + fn set_free(&self, index: usize) { + self.busy.lock().unwrap()[index] = false; + self.cv.notify_all(); + } +} + +/// All per-frame resources that must be private to a single in-flight encode. +pub(crate) struct EncodeSlot { + pub input_image: vk::Image, + pub input_image_memory: vk::DeviceMemory, + pub input_image_view: vk::ImageView, + /// Tracked layout of `input_image` (to avoid UB when transitioning). + pub input_image_layout: vk::ImageLayout, + + pub bitstream_buffer: vk::Buffer, + pub bitstream_buffer_memory: vk::DeviceMemory, + pub bitstream_buffer_size: usize, + /// Persistently mapped pointer to the bitstream buffer. + pub bitstream_buffer_ptr: *mut u8, + + pub encode_command_buffer: vk::CommandBuffer, + pub encode_fence: vk::Fence, + pub query_pool: vk::QueryPool, + + /// Packet metadata recorded before submit, moved to the completion thread + /// with the work item. + pub pending_metadata: Option, +} + +/// Configuration for building an [`EncodePipeline`]. +pub(crate) struct PipelineConfig<'a> { + pub context: &'a VideoContext, + pub aligned_width: u32, + pub aligned_height: u32, + pub picture_format: vk::Format, + pub pixel_format: PixelFormat, + pub bit_depth: BitDepth, + pub bitstream_buffer_size: usize, + /// Codec profile (with the codec-specific profile chained in) used for the + /// input images, bitstream buffers and feedback query pools. + pub profile_info: &'a vk::VideoProfileInfoKHR<'a>, + pub command_pool: vk::CommandPool, + /// Transfer command buffer/fence reused to zero-initialize each input image. + pub upload_command_buffer: vk::CommandBuffer, + pub upload_fence: vk::Fence, +} + +/// Rotating set of [`EncodeSlot`]s plus the timeline semaphore that orders their +/// encode submissions and the completion thread that reads bitstreams back. +pub(crate) struct EncodePipeline { + slots: Vec, + current_slot: usize, + /// Orders encode submissions that share DPB state. + timeline: vk::Semaphore, + /// Value the next submit will signal. + next_value: u64, + /// Value the most recent submit signaled (0 = none yet). + last_value: u64, + + /// Per-slot busy flags shared with the completion thread. + slot_sync: Arc, + /// Sends submitted work to the completion thread. Dropped on shutdown to end + /// the thread. + work_tx: Option>, + /// The completion thread handle, joined on shutdown. + completion_thread: Option>, +} + +impl EncodePipeline { + /// Allocate the timeline semaphore, `ENCODE_PIPELINE_DEPTH` slots and spawn + /// the bitstream-readback completion thread. + pub(crate) fn new(config: &PipelineConfig) -> Result { + let context = config.context; + let device = context.device(); + + let timeline = create_timeline_semaphore(context)?; + + let alloc_info = vk::CommandBufferAllocateInfo::default() + .command_pool(config.command_pool) + .level(vk::CommandBufferLevel::PRIMARY) + .command_buffer_count(ENCODE_PIPELINE_DEPTH as u32); + let command_buffers = unsafe { device.allocate_command_buffers(&alloc_info) } + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + + let mut slots = Vec::with_capacity(ENCODE_PIPELINE_DEPTH); + for &encode_command_buffer in &command_buffers { + let (input_image, input_image_memory, input_image_view) = create_image( + context, + config.aligned_width, + config.aligned_height, + config.picture_format, + false, + config.profile_info, + )?; + + let (bitstream_buffer, bitstream_buffer_memory) = create_bitstream_buffer( + context, + config.bitstream_buffer_size, + config.profile_info, + )?; + let bitstream_buffer_ptr = map_bitstream_buffer( + context, + bitstream_buffer_memory, + config.bitstream_buffer_size, + )?; + + // Zero the padding between the user dimensions and the aligned coded + // extent so the first frame has no undefined samples. + clear_input_image( + context, + &ClearImageParams { + command_buffer: config.upload_command_buffer, + fence: config.upload_fence, + queue: context.transfer_queue(), + image: input_image, + width: config.aligned_width, + height: config.aligned_height, + pixel_format: config.pixel_format, + bit_depth: config.bit_depth, + }, + )?; + + // Created signaled so it is safe to wait on before the first encode; + // `submit_encode_only` resets it before each submit. + let fence_create_info = + vk::FenceCreateInfo::default().flags(vk::FenceCreateFlags::SIGNALED); + let encode_fence = unsafe { device.create_fence(&fence_create_info, None) } + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + + let mut profile = *config.profile_info; + let query_pool = create_encode_feedback_query_pool(context, &mut profile)?; + + slots.push(EncodeSlot { + input_image, + input_image_memory, + input_image_view, + input_image_layout: vk::ImageLayout::VIDEO_ENCODE_SRC_KHR, + bitstream_buffer, + bitstream_buffer_memory, + bitstream_buffer_size: config.bitstream_buffer_size, + bitstream_buffer_ptr, + encode_command_buffer, + encode_fence, + query_pool, + pending_metadata: None, + }); + } + + let slot_sync = Arc::new(SlotSync::new(slots.len())); + let (work_tx, work_rx) = std::sync::mpsc::channel::(); + + // The completion thread only needs a handle to the device; the Vulkan + // device handle is internally shared and safe to use from this thread + // for fence waits, query reads and host-coherent buffer reads. + let thread_device = device.clone(); + let thread_sync = slot_sync.clone(); + let completion_thread = std::thread::Builder::new() + .name("pixelforge-encode-readback".to_string()) + .spawn(move || { + run_completion_thread(thread_device, work_rx, thread_sync); + }) + .map_err(|e| PixelForgeError::CommandBuffer(format!("spawn readback thread: {e}")))?; + + Ok(Self { + slots, + current_slot: 0, + timeline, + next_value: 1, + last_value: 0, + slot_sync, + work_tx: Some(work_tx), + completion_thread: Some(completion_thread), + }) + } + + /// The slot the next frame will be encoded into. + pub(crate) fn current(&self) -> &EncodeSlot { + &self.slots[self.current_slot] + } + + pub(crate) fn current_mut(&mut self) -> &mut EncodeSlot { + &mut self.slots[self.current_slot] + } + + /// Return the current slot's input image, first waiting until the slot is + /// free so it is safe to use as a convert/upload target (write-after-read on + /// the shared input image). + pub(crate) fn input_image(&self) -> vk::Image { + self.slot_sync.wait_free(self.current_slot); + self.slots[self.current_slot].input_image + } + + /// Wait until the current slot is free to record over and submit. + pub(crate) fn wait_current_free(&self) { + self.slot_sync.wait_free(self.current_slot); + } + + /// Wait until every in-flight submission has been read back. Used before + /// mutating shared session state and at teardown. + pub(crate) fn wait_all_free(&self) { + self.slot_sync.wait_all_free(); + } + + /// Record the metadata for the packet that the current slot will produce. + /// Must be called before [`EncodePipeline::submit_current`]. + pub(crate) fn set_pending_metadata(&mut self, metadata: SlotPacketMetadata) { + self.slots[self.current_slot].pending_metadata = Some(metadata); + } + + /// Submit the current slot's recorded command buffer without waiting, and + /// hand the slot to the completion thread for bitstream readback. + /// + /// Chains onto the timeline semaphore so the GPU keeps encodes in DPB order, + /// and marks the slot busy until the completion thread reads it back. Returns + /// the [`EncodeFuture`] that resolves with this frame's packet. + pub(crate) fn submit_current( + &mut self, + device: &ash::Device, + encode_queue: vk::Queue, + ) -> Result { + let wait = (self.last_value > 0).then_some((self.timeline, self.last_value)); + let signal_value = self.next_value; + let slot_index = self.current_slot; + + // Capture the Copy handles + metadata, releasing the slot borrow before + // touching the cross-thread channels. + let (command_buffer, fence, query_pool, bitstream_ptr, metadata) = { + let slot = &mut self.slots[slot_index]; + let metadata = slot.pending_metadata.take().ok_or_else(|| { + PixelForgeError::CommandBuffer( + "submit_current called without pending packet metadata".to_string(), + ) + })?; + ( + slot.encode_command_buffer, + slot.encode_fence, + slot.query_pool, + slot.bitstream_buffer_ptr as *const u8, + metadata, + ) + }; + + unsafe { + submit_encode_only( + device, + command_buffer, + fence, + encode_queue, + wait, + Some((self.timeline, signal_value)), + )?; + } + + self.last_value = signal_value; + self.next_value = signal_value + 1; + + // Mark busy *before* handing the work off, so the completion thread can + // never clear the flag before it is set. + self.slot_sync.set_busy(slot_index); + + let (result_tx, result_rx) = oneshot::channel::>(); + let work = WorkItem { + slot_index, + fence, + query_pool, + bitstream_ptr: SendPtr(bitstream_ptr), + metadata, + result_tx, + }; + if let Some(tx) = &self.work_tx { + // The receiver only disconnects during shutdown, after the queue is + // idle; a failed send there is benign. + let _ = tx.send(work); + } + + Ok(EncodeFuture { rx: result_rx }) + } + + /// Advance to the next slot after a frame has been submitted. + pub(crate) fn advance(&mut self) { + self.current_slot = (self.current_slot + 1) % self.slots.len(); + } + + /// Barrier: wait for every in-flight frame to be read back. + /// + /// The completion thread resolves each frame's [`EncodeFuture`] before it + /// marks the slot free, so once every slot is free every outstanding future + /// has already been resolved. Callers await the futures `encode` returned to + /// obtain the packets themselves. + pub(crate) fn flush(&mut self) { + self.slot_sync.wait_all_free(); + } + + /// Stop the completion thread and wait for it to finish any in-flight + /// readback. Safe to call more than once. + fn shutdown(&mut self) { + // Dropping the sender ends the thread's `for work in rx` loop once it + // has drained outstanding items. + self.work_tx.take(); + if let Some(handle) = self.completion_thread.take() { + let _ = handle.join(); + } + } + + /// Destroy all slot resources and the timeline semaphore. + /// + /// # Safety + /// + /// All queues that may reference these resources must be idle. + pub(crate) unsafe fn destroy(&mut self, device: &ash::Device) { + // Join the readback thread before freeing the fences/buffers it reads. + self.shutdown(); + + for slot in &mut self.slots { + if !slot.bitstream_buffer_ptr.is_null() { + device.unmap_memory(slot.bitstream_buffer_memory); + slot.bitstream_buffer_ptr = std::ptr::null_mut(); + } + device.destroy_query_pool(slot.query_pool, None); + device.destroy_fence(slot.encode_fence, None); + device.destroy_buffer(slot.bitstream_buffer, None); + device.free_memory(slot.bitstream_buffer_memory, None); + device.destroy_image_view(slot.input_image_view, None); + device.destroy_image(slot.input_image, None); + device.free_memory(slot.input_image_memory, None); + } + device.destroy_semaphore(self.timeline, None); + } +} + +/// Completion-thread body: wait on each submission's fence, copy its bitstream +/// out, resolve the frame's future, then mark the slot free. +fn run_completion_thread( + device: ash::Device, + work_rx: Receiver, + slot_sync: Arc, +) { + for work in work_rx { + // Start the packet from any codec header, then read the encoded bitstream + // straight onto it — a single copy out of the mapped buffer. + let mut data = work.metadata.header.unwrap_or_default(); + let result = unsafe { + wait_and_read_bitstream( + &device, + work.fence, + work.query_pool, + work.bitstream_ptr.0, + &mut data, + ) + }; + + let packet = result.map(|()| EncodedPacket { + data, + frame_type: work.metadata.frame_type, + is_key_frame: work.metadata.is_key_frame, + pts: work.metadata.pts, + dts: work.metadata.dts, + }); + + // Resolve the future *before* freeing the slot. This ordering means that + // once all slots are observed free, every packet has already been + // delivered, which `flush` relies on for completeness. A dropped + // receiver (future cancelled) makes the send a no-op, which is fine. + let _ = work.result_tx.send(packet); + slot_sync.set_free(work.slot_index); + } +} diff --git a/src/encoder/resources.rs b/src/encoder/resources.rs index d738129..5f465e1 100644 --- a/src/encoder/resources.rs +++ b/src/encoder/resources.rs @@ -123,20 +123,6 @@ pub(crate) fn get_video_format(pixel_format: PixelFormat, bit_depth: BitDepth) - } } -/// Create a codec name array for Vulkan from a string. -/// -/// This creates a null-terminated c_char array of 256 bytes for use with Vulkan -/// video extensions. -pub(crate) fn make_codec_name(codec_name: &[u8]) -> [std::ffi::c_char; 256] { - let mut name = [0 as std::ffi::c_char; 256]; - for (i, &byte) in codec_name.iter().enumerate() { - if i < 255 { - name[i] = byte as std::ffi::c_char; - } - } - name -} - /// Create a buffer that requires device addresses (SHADER_DEVICE_ADDRESS usage). /// /// This allocates memory with `VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT` so that @@ -250,6 +236,43 @@ pub(crate) fn create_bitstream_buffer( Ok((buffer, memory)) } + +pub(crate) fn create_timeline_semaphore(context: &VideoContext) -> Result { + let mut type_info = vk::SemaphoreTypeCreateInfo::default() + .semaphore_type(vk::SemaphoreType::TIMELINE) + .initial_value(0); + let create_info = vk::SemaphoreCreateInfo::default().push(&mut type_info); + + unsafe { context.device().create_semaphore(&create_info, None) } + .map_err(|e| PixelForgeError::Synchronization(e.to_string())) +} + +pub(crate) fn create_encode_feedback_query_pool( + context: &VideoContext, + profile_info: &mut vk::VideoProfileInfoKHR, +) -> Result { + let mut encode_feedback_create = vk::QueryPoolVideoEncodeFeedbackCreateInfoKHR::default() + .encode_feedback_flags( + vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BUFFER_OFFSET + | vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN, + ); + + let query_pool_create_info = unsafe { + vk::QueryPoolCreateInfo::default() + .query_type(vk::QueryType::VIDEO_ENCODE_FEEDBACK_KHR) + .query_count(1) + .extend(profile_info) + .push(&mut encode_feedback_create) + }; + + unsafe { + context + .device() + .create_query_pool(&query_pool_create_info, None) + } + .map_err(|e| PixelForgeError::QueryPool(e.to_string())) +} + /// Create an image for video encoding (input or DPB). /// /// This creates a VkImage suitable for use with a video encoder. @@ -483,10 +506,6 @@ pub(crate) struct CommandResources { pub upload_command_buffer: vk::CommandBuffer, /// Fence for upload synchronization. pub upload_fence: vk::Fence, - /// Command buffer for encode operations. - pub encode_command_buffer: vk::CommandBuffer, - /// Fence for encode synchronization. - pub encode_fence: vk::Fence, } /// Create command resources for encoding. @@ -511,16 +530,6 @@ pub(crate) fn create_command_resources( } .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - // Allocate encode command buffer. - let alloc_info = vk::CommandBufferAllocateInfo::default() - .command_pool(command_pool) - .level(vk::CommandBufferLevel::PRIMARY) - .command_buffer_count(1); - - let encode_command_buffers = unsafe { context.device().allocate_command_buffers(&alloc_info) } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - let encode_command_buffer = encode_command_buffers[0]; - // Create command pool for upload commands (may be the same family). let upload_command_pool = if upload_queue_family == encode_queue_family { command_pool @@ -550,22 +559,17 @@ pub(crate) fn create_command_resources( .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; let upload_command_buffer = upload_command_buffers[0]; - // Create fences. The encode fence is created signaled so that - // set_color_description() can safely wait on it before the first encode. + // Upload fence is created unsignaled. Per-slot encode fences are created by + // the encode pipeline (see `encoder::pipeline`). let fence_create_info = vk::FenceCreateInfo::default(); let upload_fence = unsafe { context.device().create_fence(&fence_create_info, None) } .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - let signaled_fence_info = vk::FenceCreateInfo::default().flags(vk::FenceCreateFlags::SIGNALED); - let encode_fence = unsafe { context.device().create_fence(&signaled_fence_info, None) } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; Ok(CommandResources { command_pool, upload_command_pool, upload_command_buffer, upload_fence, - encode_command_buffer, - encode_fence, }) } @@ -1161,77 +1165,6 @@ pub(crate) fn upload_image_to_input( Ok(()) } -/// Parameters for cleaning up shared encoder resources. -pub(crate) struct EncoderResources<'a> { - pub query_pool: vk::QueryPool, - pub upload_fence: vk::Fence, - pub encode_fence: vk::Fence, - pub command_pool: vk::CommandPool, - pub upload_command_pool: vk::CommandPool, - pub bitstream_buffer: vk::Buffer, - pub bitstream_buffer_memory: vk::DeviceMemory, - pub input_image: vk::Image, - pub input_image_memory: vk::DeviceMemory, - pub input_image_view: vk::ImageView, - pub dpb_images: &'a [vk::Image], - pub dpb_image_memories: &'a [vk::DeviceMemory], - pub dpb_image_views: &'a [vk::ImageView], - pub session: vk::VideoSessionKHR, - pub session_params: vk::VideoSessionParametersKHR, - pub session_memory: &'a [vk::DeviceMemory], -} - -/// Destroy all shared encoder resources. -/// -/// # Safety -/// -/// All queues that may reference these resources (transfer and video encode) -/// must be idle before calling this function. -pub(crate) unsafe fn destroy_encoder_resources( - device: &ash::Device, - video_queue_fn: &ash::khr::video_queue::Device, - res: &EncoderResources, -) { - device.destroy_query_pool(res.query_pool, None); - device.destroy_fence(res.upload_fence, None); - device.destroy_fence(res.encode_fence, None); - device.destroy_command_pool(res.command_pool, None); - if res.upload_command_pool != res.command_pool { - device.destroy_command_pool(res.upload_command_pool, None); - } - - device.unmap_memory(res.bitstream_buffer_memory); - device.destroy_buffer(res.bitstream_buffer, None); - device.free_memory(res.bitstream_buffer_memory, None); - - device.destroy_image_view(res.input_image_view, None); - device.destroy_image(res.input_image, None); - device.free_memory(res.input_image_memory, None); - - for view in res.dpb_image_views { - device.destroy_image_view(*view, None); - } - for image in res.dpb_images { - device.destroy_image(*image, None); - } - for memory in res.dpb_image_memories { - device.free_memory(*memory, None); - } - - if res.session_params != vk::VideoSessionParametersKHR::null() { - (video_queue_fn.fp().destroy_video_session_parameters_khr)( - device.handle(), - res.session_params, - std::ptr::null(), - ); - } - (video_queue_fn.fp().destroy_video_session_khr)(device.handle(), res.session, std::ptr::null()); - - for memory in res.session_memory { - device.free_memory(*memory, None); - } -} - /// Record DPB image barriers for encode. /// /// Transitions the setup DPB slot from UNDEFINED to VIDEO_ENCODE_DPB and @@ -1399,39 +1332,79 @@ pub(crate) unsafe fn record_post_encode_dpb_barrier( ); } -/// Submit an encode command buffer and wait for completion. +/// Submit an encode command buffer without waiting for completion. /// -/// Submits the command buffer to the encode queue, waits for the fence, -/// then reads query results and copies the encoded bitstream data. -/// The fence is reset before submission so it may be in any state on entry. +/// Timeline semaphore waits/signals are used to keep shared DPB access ordered +/// across pipelined slots while leaving bitstream readback delayed. Uses +/// Synchronization2 (`queue_submit2`) so the timeline wait/signal are scoped to +/// the `VIDEO_ENCODE` stage rather than `ALL_COMMANDS`, avoiding unnecessary +/// over-synchronization. /// /// # Safety /// /// The command buffer must have been ended. -/// The bitstream buffer pointer must be valid and the buffer must be persistently mapped. -pub(crate) unsafe fn submit_encode_and_read_bitstream( +pub(crate) unsafe fn submit_encode_only( device: &ash::Device, command_buffer: vk::CommandBuffer, fence: vk::Fence, encode_queue: vk::Queue, - query_pool: vk::QueryPool, - bitstream_buffer_ptr: *const u8, -) -> Result> { - let submit_info = - vk::SubmitInfo::default().command_buffers(std::slice::from_ref(&command_buffer)); + wait_timeline: Option<(vk::Semaphore, u64)>, + signal_timeline: Option<(vk::Semaphore, u64)>, +) -> Result<()> { + let command_buffer_info = vk::CommandBufferSubmitInfo::default().command_buffer(command_buffer); + let command_buffer_infos = [command_buffer_info]; + + let wait_infos: Vec = wait_timeline + .into_iter() + .map(|(semaphore, value)| { + vk::SemaphoreSubmitInfo::default() + .semaphore(semaphore) + .value(value) + .stage_mask(vk::PipelineStageFlags2::VIDEO_ENCODE_KHR) + }) + .collect(); + let signal_infos: Vec = signal_timeline + .into_iter() + .map(|(semaphore, value)| { + vk::SemaphoreSubmitInfo::default() + .semaphore(semaphore) + .value(value) + .stage_mask(vk::PipelineStageFlags2::VIDEO_ENCODE_KHR) + }) + .collect(); + + let submit_info = vk::SubmitInfo2::default() + .wait_semaphore_infos(&wait_infos) + .command_buffer_infos(&command_buffer_infos) + .signal_semaphore_infos(&signal_infos); - // Reset the fence before submit (it may be signaled from a previous encode - // or from initial creation with SIGNALED_BIT). This ensures the fence is - // unsignaled for queue_submit, and after wait_for_fences it stays signaled — - // which lets set_color_description() safely wait on it between encodes. device .reset_fences(&[fence]) .map_err(|e| PixelForgeError::Synchronization(e.to_string()))?; device - .queue_submit(encode_queue, &[submit_info], fence) + .queue_submit2(encode_queue, &[submit_info], fence) .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + Ok(()) +} + +/// Wait for a submitted encode and copy its bitstream data. +/// +/// # Safety +/// +/// The bitstream buffer pointer must be valid and persistently mapped. +/// Wait for the encode to finish, then append its bitstream directly onto `dst` +/// (which already holds any codec header). Appending in place means the encoded +/// bytes are copied out of the mapped buffer exactly once — there is no +/// intermediate owned `Vec`. +pub(crate) unsafe fn wait_and_read_bitstream( + device: &ash::Device, + fence: vk::Fence, + query_pool: vk::QueryPool, + bitstream_buffer_ptr: *const u8, + dst: &mut Vec, +) -> Result<()> { device .wait_for_fences(&[fence], true, u64::MAX) .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; @@ -1468,11 +1441,122 @@ pub(crate) unsafe fn submit_encode_and_read_bitstream( tracing::debug!("Encoded frame: offset={}, size={}", offset, size); - let mut encoded_data = vec![0u8; size]; let src = std::slice::from_raw_parts(bitstream_buffer_ptr.add(offset), size); - encoded_data.copy_from_slice(src); + dst.extend_from_slice(src); + + Ok(()) +} + +/// The per-codec Vulkan resources torn down on `Drop`, other than the encode +/// pipeline (which the engine owns and destroys). Bundled so the identical +/// teardown lives in one place; see [`destroy_encoder_resources`]. +pub(crate) struct EncoderTeardown<'a> { + pub command_pool: vk::CommandPool, + pub upload_command_pool: vk::CommandPool, + pub upload_fence: vk::Fence, + pub dpb_images: &'a [vk::Image], + pub dpb_image_views: &'a [vk::ImageView], + pub dpb_image_memories: &'a [vk::DeviceMemory], + pub session: vk::VideoSessionKHR, + pub session_params: vk::VideoSessionParametersKHR, + pub session_memory: &'a [vk::DeviceMemory], +} + +/// Destroy the common encoder resources (command pools, upload fence, DPB +/// images, and the video session). Identical across codecs. +/// +/// # Safety +/// +/// All queues that may reference these resources must be idle. +pub(crate) unsafe fn destroy_encoder_resources( + device: &ash::Device, + video_queue_fn: &ash::khr::video_queue::Device, + res: &EncoderTeardown, +) { + device.destroy_fence(res.upload_fence, None); + device.destroy_command_pool(res.command_pool, None); + if res.upload_command_pool != res.command_pool { + device.destroy_command_pool(res.upload_command_pool, None); + } - Ok(encoded_data) + for &view in res.dpb_image_views { + device.destroy_image_view(view, None); + } + for &image in res.dpb_images { + device.destroy_image(image, None); + } + for &memory in res.dpb_image_memories { + device.free_memory(memory, None); + } + + if res.session_params != vk::VideoSessionParametersKHR::null() { + (video_queue_fn.fp().destroy_video_session_parameters_khr)( + device.handle(), + res.session_params, + std::ptr::null(), + ); + } + (video_queue_fn.fp().destroy_video_session_khr)(device.handle(), res.session, std::ptr::null()); + for &memory in res.session_memory { + device.free_memory(memory, None); + } +} + +/// Retrieve driver-encoded video session parameters (H.264 SPS/PPS, H.265 +/// VPS/SPS/PPS, or AV1 sequence header) into a byte buffer, retrying on +/// `INCOMPLETE`. +/// +/// The codec-specific `*GetInfoKHR` and feedback structs are built by the caller +/// and chained into `get_info`/`feedback`; this owns only the identical query +/// loop. A preallocated buffer is always provided because some drivers misbehave +/// on a size-only query (`pData == NULL`), notably for 4:4:4 profiles. +pub(crate) fn get_encoded_session_params( + context: &VideoContext, + video_encode_fn: &ash::khr::video_encode_queue::Device, + get_info: &vk::VideoEncodeSessionParametersGetInfoKHR, + feedback: &mut vk::VideoEncodeSessionParametersFeedbackInfoKHR, +) -> Result> { + let mut data = vec![0u8; 4096]; + let mut data_size: usize = data.len(); + let mut attempts = 0; + loop { + attempts += 1; + let result = unsafe { + (video_encode_fn + .fp() + .get_encoded_video_session_parameters_khr)( + context.device().handle(), + get_info, + feedback, + &mut data_size, + data.as_mut_ptr() as *mut std::ffi::c_void, + ) + }; + + match result { + vk::Result::SUCCESS => { + if data_size == 0 { + return Err(PixelForgeError::SessionParametersCreation( + "Encoded session parameters size is 0".to_string(), + )); + } + data.truncate(data_size); + return Ok(data); + } + // Driver indicates the buffer was too small; resize to the reported + // required size (or grow conservatively if it is not provided). + vk::Result::INCOMPLETE if attempts < 3 => { + let new_size = data_size.max(data.len() * 2).max(1); + data.resize(new_size, 0); + data_size = data.len(); + } + err => { + return Err(PixelForgeError::SessionParametersCreation(format!( + "Failed to get encoded session parameters: {err:?}" + ))); + } + } + } } #[cfg(test)] diff --git a/src/lib.rs b/src/lib.rs index 499273f..77c6432 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -197,9 +197,10 @@ pub(crate) const fn align4(size: usize) -> usize { pub use converter::{ColorConverter, ColorConverterConfig, ColorSpace, InputFormat, OutputFormat}; pub use encoder::{ - BitDepth as EncodeBitDepth, Codec, ColorDescription, EncodeConfig, EncodedPacket, Encoder, - FrameType, PixelFormat, RateControlMode, DEFAULT_FRAME_RATE, DEFAULT_GOP_SIZE, DEFAULT_H264_QP, - DEFAULT_H265_QP, DEFAULT_MAX_BITRATE, DEFAULT_MAX_REFERENCE_FRAMES, DEFAULT_TARGET_BITRATE, + BitDepth as EncodeBitDepth, Codec, ColorDescription, EncodeConfig, EncodeFuture, EncodedPacket, + Encoder, FrameType, PixelFormat, RateControlMode, DEFAULT_FRAME_RATE, DEFAULT_GOP_SIZE, + DEFAULT_H264_QP, DEFAULT_H265_QP, DEFAULT_MAX_BITRATE, DEFAULT_MAX_REFERENCE_FRAMES, + DEFAULT_TARGET_BITRATE, }; pub use error::PixelForgeError; pub use image::InputImage; diff --git a/src/vulkan.rs b/src/vulkan.rs index 339de1f..9257c72 100644 --- a/src/vulkan.rs +++ b/src/vulkan.rs @@ -454,6 +454,22 @@ impl VideoContext { let mut sync2_features = vk::PhysicalDeviceSynchronization2Features::default().synchronization2(true); + let mut supported_timeline_features = + vk::PhysicalDeviceTimelineSemaphoreFeatures::default(); + let mut timeline_feature_query = + vk::PhysicalDeviceFeatures2::default().push(&mut supported_timeline_features); + unsafe { + instance.get_physical_device_features2(physical_device, &mut timeline_feature_query); + } + if supported_timeline_features.timeline_semaphore == 0 { + return Err(PixelForgeError::NoSuitableDevice( + "Timeline semaphores are required for pipelined video encode synchronization" + .to_string(), + )); + } + let mut timeline_features = + vk::PhysicalDeviceTimelineSemaphoreFeatures::default().timeline_semaphore(true); + // Enable sampler YCbCr conversion feature (required for YUV image views with SAMPLED flag). let mut ycbcr_features = vk::PhysicalDeviceSamplerYcbcrConversionFeatures::default() .sampler_ycbcr_conversion(true); @@ -515,7 +531,8 @@ impl VideoContext { let mut device_create_info = vk::DeviceCreateInfo::default() .queue_create_infos(&queue_create_infos) .enabled_extension_names(&extension_names) - .push(&mut sync2_features); + .push(&mut sync2_features) + .push(&mut timeline_features); if supported_encode_codecs.contains(&Codec::AV1) { device_create_info = device_create_info.push(&mut av1_encode_features);