From 7bb0577fd2ee2c3992d6c3e93afe599f16aef447 Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Sun, 14 Jun 2026 18:14:04 +0800 Subject: [PATCH 1/9] feat: encode pipelining with timeline semaphores --- examples/verify_all.rs | 4 + src/encoder/av1/api.rs | 52 +++--- src/encoder/av1/encode.rs | 68 +++----- src/encoder/av1/init.rs | 96 +++-------- src/encoder/av1/mod.rs | 82 +++------ src/encoder/h264/api.rs | 58 +++---- src/encoder/h264/encode.rs | 76 +++------ src/encoder/h264/init.rs | 100 +++-------- src/encoder/h264/mod.rs | 95 +++++------ src/encoder/h265/api.rs | 63 ++++--- src/encoder/h265/encode.rs | 74 +++------ src/encoder/h265/init.rs | 99 +++-------- src/encoder/h265/mod.rs | 91 +++++----- src/encoder/mod.rs | 4 + src/encoder/pipeline.rs | 331 +++++++++++++++++++++++++++++++++++++ src/encoder/resources.rs | 196 ++++++++++------------ src/vulkan.rs | 19 ++- 17 files changed, 791 insertions(+), 717 deletions(-) create mode 100644 src/encoder/pipeline.rs diff --git a/examples/verify_all.rs b/examples/verify_all.rs index dd80a63..4f6a1ee 100644 --- a/examples/verify_all.rs +++ b/examples/verify_all.rs @@ -175,6 +175,10 @@ fn run_test( output_file.write_all(&packet.data)?; } } + + for packet in encoder.flush()? { + 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 index 6096b82..6b3e536 100644 --- a/src/encoder/av1/api.rs +++ b/src/encoder/av1/api.rs @@ -1,6 +1,7 @@ use super::AV1Encoder; use crate::encoder::gop::{GopFrameType, GopPosition}; +use crate::encoder::pipeline::SlotPacketMetadata; use crate::encoder::{ColorDescription, EncodedPacket}; use crate::error::{PixelForgeError, Result}; use ash::vk; @@ -12,7 +13,7 @@ impl AV1Encoder { /// 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 + self.pipeline.input_image(self.context.device()) } /// Encode a frame from a GPU image. @@ -26,6 +27,7 @@ impl AV1Encoder { /// 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 drained_packet = self.pipeline.drain_current(self.context.device())?; let gop_position = self.gop.get_next_frame(); let display_order = self.input_frame_num; self.input_frame_num += 1; @@ -38,10 +40,11 @@ impl AV1Encoder { // Upload from GPU image. self.upload_from_image(src_image)?; - // Encode immediately. - let packet = self.encode_current_frame(&gop_position, display_order)?; + self.encode_current_frame(&gop_position, display_order)?; - Ok(vec![packet]) + self.pipeline.advance(); + + Ok(drained_packet.into_iter().collect()) } /// Internal method to encode the current frame already uploaded to input_image. @@ -49,7 +52,7 @@ impl AV1Encoder { &mut self, gop_position: &GopPosition, display_order: u64, - ) -> Result { + ) -> Result<()> { let is_key_frame = gop_position.frame_type.is_idr() || gop_position.frame_type == GopFrameType::I; let is_reference = gop_position.is_reference; @@ -71,12 +74,12 @@ impl AV1Encoder { self.references.clear(); } - let mut encoded_data = Vec::new(); + let mut header = 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]); + header.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). @@ -90,19 +93,28 @@ impl AV1Encoder { ); self.header_data = Some(header); } - if let Some(ref header) = self.header_data { - encoded_data.extend_from_slice(header); + if let Some(ref sequence_header) = self.header_data { + header.extend_from_slice(sequence_header); } } - encoded_data.extend_from_slice(&self.encode_frame_internal(gop_position, is_key_frame)?); + 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; + let dts = self.encode_frame_num; self.encode_frame_num += 1; self.frame_num += 1; self.order_hint = (self.order_hint + 1) & 0xFF; // 8-bit order hint + self.pipeline.set_pending_metadata(SlotPacketMetadata { + frame_type, + is_key_frame, + pts: display_order, + dts, + header: Some(header), + }); + // 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 @@ -140,19 +152,12 @@ impl AV1Encoder { 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, - }) + Ok(()) } /// Flush the encoder and get any remaining packets. pub fn flush(&mut self) -> Result> { - // No buffered frames in the current implementation. - Ok(Vec::new()) + self.pipeline.flush(self.context.device()) } /// Request that the next frame be an IDR/key frame. @@ -224,16 +229,15 @@ impl AV1Encoder { /// 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. + // Do not reset the fences here; submit resets each slot fence before reuse. + let fences = self.pipeline.encode_fences(); unsafe { self.context .device() - .wait_for_fences(&[self.encode_fence], true, u64::MAX) + .wait_for_fences(&fences, true, u64::MAX) .map_err(|e| { PixelForgeError::Synchronization(format!( - "Failed to wait for encode fence: {:?}", + "Failed to wait for encode fences: {:?}", e )) })?; diff --git a/src/encoder/av1/encode.rs b/src/encoder/av1/encode.rs index 273a6af..1374d1e 100644 --- a/src/encoder/av1/encode.rs +++ b/src/encoder/av1/encode.rs @@ -3,7 +3,6 @@ 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; @@ -15,7 +14,14 @@ impl AV1Encoder { &mut self, _gop_position: &GopPosition, is_key_frame: bool, - ) -> Result> { + ) -> Result<()> { + let slot = self.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; + // All frames need a setup reference slot (DPB write) per Vulkan spec when maxDpbSlots > 0. let is_reference = true; @@ -51,11 +57,7 @@ impl AV1Encoder { // Prepare command buffer for recording. unsafe { - prepare_encode_command_buffer( - self.context.device(), - self.encode_command_buffer, - self.query_pool, - )?; + prepare_encode_command_buffer(self.context.device(), command_buffer, query_pool)?; } // Transition DPB images for encode. @@ -63,7 +65,7 @@ impl AV1Encoder { unsafe { record_dpb_barriers( self.context.device(), - self.encode_command_buffer, + command_buffer, &self.dpb_images, false, // AV1 does not use layered DPB self.current_dpb_slot, @@ -406,7 +408,7 @@ impl AV1Encoder { unsafe { self.video_queue_fn - .cmd_begin_video_coding(self.encode_command_buffer, &begin_coding_info); + .cmd_begin_video_coding(command_buffer, &begin_coding_info); } // Reset video coding state for the first frame. @@ -428,7 +430,7 @@ impl AV1Encoder { unsafe { self.video_queue_fn - .cmd_control_video_coding(self.encode_command_buffer, &control_info); + .cmd_control_video_coding(command_buffer, &control_info); } } @@ -437,13 +439,13 @@ impl AV1Encoder { .coded_offset(vk::Offset2D { x: 0, y: 0 }) .coded_extent(frame_extent) .base_array_layer(0) - .image_view_binding(self.input_image_view); + .image_view_binding(input_image_view); let mut encode_info = vk::VideoEncodeInfoKHR::default() .src_picture_resource(src_picture_resource) - .dst_buffer(self.bitstream_buffer) + .dst_buffer(bitstream_buffer) .dst_buffer_offset(0) - .dst_buffer_range(self.bitstream_buffer_size as u64); + .dst_buffer_range(bitstream_buffer_size as u64); if is_reference { encode_info = encode_info.setup_reference_slot(&setup_reference_slot); @@ -458,8 +460,8 @@ impl AV1Encoder { // Begin query to capture encode feedback (bitstream size, status). unsafe { self.context.device().cmd_begin_query( - self.encode_command_buffer, - self.query_pool, + command_buffer, + query_pool, 0, vk::QueryControlFlags::empty(), ); @@ -467,21 +469,21 @@ impl AV1Encoder { unsafe { self.video_encode_fn - .cmd_encode_video(self.encode_command_buffer, &encode_info); + .cmd_encode_video(command_buffer, &encode_info); } // End query. unsafe { self.context .device() - .cmd_end_query(self.encode_command_buffer, self.query_pool, 0); + .cmd_end_query(command_buffer, query_pool, 0); } // Add DPB synchronization barrier after encoding. unsafe { record_post_encode_dpb_barrier( self.context.device(), - self.encode_command_buffer, + command_buffer, &self.dpb_images, false, // AV1 does not use layered DPB self.current_dpb_slot, @@ -492,18 +494,14 @@ impl AV1Encoder { let end_coding_info = vk::VideoEndCodingInfoKHR::default(); unsafe { self.video_queue_fn - .cmd_end_video_coding(self.encode_command_buffer, &end_coding_info); + .cmd_end_video_coding(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()))?; + unsafe { self.context.device().end_command_buffer(command_buffer) } + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - // Submit, wait, and read bitstream. + // Submit without waiting; the slot is drained on a later encode call. let encode_queue = self.context.video_encode_queue().ok_or_else(|| { PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) })?; @@ -516,25 +514,13 @@ impl AV1Encoder { 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()); + self.pipeline + .submit_current(self.context.device(), encode_queue)?; // Mark current DPB slot as active. self.dpb_slot_active[self.current_dpb_slot as usize] = true; - Ok(encoded_data) + Ok(()) } /// Calculate proper reference frame mapping for AV1 encoding. diff --git a/src/encoder/av1/init.rs b/src/encoder/av1/init.rs index 4ffcd70..46cc558 100644 --- a/src/encoder/av1/init.rs +++ b/src/encoder/av1/init.rs @@ -1,10 +1,10 @@ use super::{AV1Encoder, MIN_BITSTREAM_BUFFER_SIZE, SUPERBLOCK_SIZE}; use crate::encoder::gop::GopStructure; +use crate::encoder::pipeline::{EncodePipeline, PipelineConfig}; 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, + allocate_session_memory, create_command_resources, create_dpb_images, get_video_format, + make_codec_name, query_supported_video_formats, }; use crate::encoder::{ColorDescription, PixelFormat}; use crate::error::{PixelForgeError, Result}; @@ -59,7 +59,7 @@ impl AV1Encoder { // 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) @@ -269,17 +269,6 @@ impl AV1Encoder { .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, @@ -290,60 +279,33 @@ impl AV1Encoder { &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_pool = cmd_resources.upload_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 bitstream_buffer_size = MIN_BITSTREAM_BUFFER_SIZE.max(width as usize * height as usize); - let query_pool = unsafe { - context - .device() - .create_query_pool(&query_pool_create_info, None) - .map_err(|e| PixelForgeError::QueryPool(e.to_string()))? - }; + // Create the depth-N encode pipeline (per-frame input images, bitstream + // buffers, encode command buffers, fences and query pools). + let pipeline = EncodePipeline::new(&PipelineConfig { + context: &context, + aligned_width, + aligned_height, + picture_format, + pixel_format: config.pixel_format, + bit_depth: config.bit_depth, + bitstream_buffer_size, + profile_info: &profile_info, + command_pool, + upload_command_buffer, + upload_fence, + })?; // Initialize GOP structure. let gop = GopStructure::new(config.gop_size, config.b_frame_count, config.gop_size); @@ -361,26 +323,16 @@ impl AV1Encoder { encode_frame_num: 0, frame_num: 0, order_hint: 0, - input_image, - input_image_memory, - input_image_view, - input_image_layout, + pipeline, 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_pool, upload_command_buffer, upload_fence, - encode_command_buffer, - encode_fence, - query_pool, header_data: None, current_dpb_slot: 0, references: Vec::new(), diff --git a/src/encoder/av1/mod.rs b/src/encoder/av1/mod.rs index 2edcde4..0131135 100644 --- a/src/encoder/av1/mod.rs +++ b/src/encoder/av1/mod.rs @@ -10,6 +10,7 @@ mod session_params; use ash::vk; use tracing::debug; +use crate::encoder::pipeline::EncodePipeline; use crate::encoder::resources::{upload_image_to_input, UploadParams}; use crate::error::Result; @@ -49,12 +50,8 @@ pub struct AV1Encoder { 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, + /// Depth-N encode pipeline (per-frame slots + submission ordering). + pipeline: EncodePipeline, /// DPB images for reference frames. dpb_images: Vec, dpb_image_memories: Vec, @@ -63,21 +60,11 @@ pub struct AV1Encoder { 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). header_data: Option>, @@ -97,7 +84,8 @@ impl AV1Encoder { /// 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 { + let slot = self.pipeline.current(); + if src_image == slot.input_image { debug!("Source image is the encoder's input image, skipping upload copy"); return Ok(()); } @@ -106,18 +94,18 @@ impl AV1Encoder { upload_command_buffer: self.upload_command_buffer, upload_fence: self.upload_fence, src_image, - dst_image: self.input_image, + dst_image: slot.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, + input_image_layout: slot.input_image_layout, upload_queue: self.context.transfer_queue(), }; upload_image_to_input(&self.context, ¶ms)?; // Update tracked layout. - self.input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + self.pipeline.current_mut().input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; Ok(()) } @@ -130,48 +118,24 @@ 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); + let device = self.context.device(); + let _ = device.queue_wait_idle(self.context.transfer_queue()); + if let Some(q) = self.context.video_encode_queue() { + let _ = device.queue_wait_idle(q); + } + + self.pipeline.destroy(device); + + device.destroy_fence(self.upload_fence, None); + 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); + 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); + device.destroy_image_view(self.dpb_image_views[i], None); + device.destroy_image(self.dpb_images[i], None); + device.free_memory(self.dpb_image_memories[i], None); } if self.session_params != vk::VideoSessionParametersKHR::null() { @@ -181,7 +145,7 @@ impl Drop for AV1Encoder { self.video_queue_fn .destroy_video_session(self.session, None); for mem in &self.session_memory { - self.context.device().free_memory(*mem, None); + device.free_memory(*mem, None); } } } diff --git a/src/encoder/h264/api.rs b/src/encoder/h264/api.rs index e79d33a..b2bfa85 100644 --- a/src/encoder/h264/api.rs +++ b/src/encoder/h264/api.rs @@ -2,6 +2,7 @@ use super::H264Encoder; use crate::encoder::dpb::{DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType}; use crate::encoder::gop::{GopFrameType, GopPosition}; +use crate::encoder::pipeline::SlotPacketMetadata; use crate::encoder::{ColorDescription, EncodedPacket}; use crate::error::Result; use crate::PixelForgeError; @@ -15,7 +16,7 @@ impl H264Encoder { /// 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 + self.pipeline.input_image(self.context.device()) } /// Encode a frame from a GPU image. @@ -29,6 +30,7 @@ impl H264Encoder { /// 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 drained_packet = self.pipeline.drain_current(self.context.device())?; let gop_position = self.gop.get_next_frame(); let display_order = self.input_frame_num; self.input_frame_num += 1; @@ -41,10 +43,11 @@ impl H264Encoder { // Upload from GPU image. self.upload_from_image(src_image)?; - // Encode immediately. - let packet = self.encode_current_frame(&gop_position, display_order)?; + self.encode_current_frame(&gop_position, display_order)?; - Ok(vec![packet]) + self.pipeline.advance(); + + Ok(drained_packet.into_iter().collect()) } /// Internal method to encode the current frame already uploaded to input_image. @@ -52,7 +55,7 @@ impl H264Encoder { &mut self, gop_position: &GopPosition, display_order: u64, - ) -> Result { + ) -> 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; @@ -87,24 +90,29 @@ impl H264Encoder { 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()?); + let header = if is_idr { self.sps_written = true; - } + Some(self.get_h264_header()?) + } else { + None + }; - encoded_data.extend_from_slice(&self.encode_frame_internal( - gop_position, - frame_num, - pic_order_cnt, - is_idr, - )?); + self.encode_frame_internal(gop_position, frame_num, pic_order_cnt, is_idr)?; + let dts = self.encode_frame_num; self.encode_frame_num += 1; if is_reference && !is_b_frame { self.frame_num_syntax = (self.frame_num_syntax + 1) % 256; } + self.pipeline.set_pending_metadata(SlotPacketMetadata { + frame_type, + is_key_frame: is_idr, + pts: display_order, + dts, + header, + }); + if is_reference { let pic_type = if is_idr { PictureType::Idr @@ -147,19 +155,12 @@ impl H264Encoder { } } - Ok(EncodedPacket { - data: encoded_data, - frame_type, - is_key_frame: is_idr, - pts: display_order, - dts: self.encode_frame_num - 1, - }) + Ok(()) } /// Flush the encoder and get any remaining packets. pub fn flush(&mut self) -> Result> { - // No buffered frames in the current implementation. - Ok(Vec::new()) + self.pipeline.flush(self.context.device()) } /// Request that the next frame be an IDR frame. @@ -239,16 +240,15 @@ impl H264Encoder { /// 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. + // Do not reset the fences here; submit resets each slot fence before reuse. + let fences = self.pipeline.encode_fences(); unsafe { self.context .device() - .wait_for_fences(&[self.encode_fence], true, u64::MAX) + .wait_for_fences(&fences, true, u64::MAX) .map_err(|e| { PixelForgeError::Synchronization(format!( - "Failed to wait for encode fence: {:?}", + "Failed to wait for encode fences: {:?}", e )) })?; diff --git a/src/encoder/h264/encode.rs b/src/encoder/h264/encode.rs index e03de2c..0e910bc 100644 --- a/src/encoder/h264/encode.rs +++ b/src/encoder/h264/encode.rs @@ -1,10 +1,7 @@ use super::H264Encoder; 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 crate::encoder::resources::{prepare_encode_command_buffer, record_dpb_barriers}; use crate::error::{PixelForgeError, Result}; use ash::vk; use ash::vk::TaggedStructure; @@ -17,7 +14,14 @@ impl H264Encoder { frame_num: u32, pic_order_cnt: i32, is_idr: bool, - ) -> Result> { + ) -> Result<()> { + let slot = self.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; + let is_b_frame = gop_position.frame_type == GopFrameType::B; let is_reference = gop_position.is_reference; @@ -54,11 +58,7 @@ impl H264Encoder { // Prepare command buffer for recording. unsafe { - prepare_encode_command_buffer( - self.context.device(), - self.encode_command_buffer, - self.query_pool, - )?; + prepare_encode_command_buffer(self.context.device(), command_buffer, query_pool)?; } // Transition DPB images for encode. @@ -66,7 +66,7 @@ impl H264Encoder { unsafe { record_dpb_barriers( self.context.device(), - self.encode_command_buffer, + command_buffer, &self.dpb_images, self.use_layered_dpb, self.current_dpb_slot, @@ -264,7 +264,7 @@ impl H264Encoder { height: self.aligned_height, }) .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() @@ -432,9 +432,9 @@ impl H264Encoder { } 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); @@ -535,10 +535,7 @@ impl H264Encoder { }; unsafe { - (self.video_queue_fn.fp().cmd_begin_video_coding_khr)( - self.encode_command_buffer, - &begin_info, - ); + (self.video_queue_fn.fp().cmd_begin_video_coding_khr)(command_buffer, &begin_info); } // Reset video coding state for the first frame. @@ -560,7 +557,7 @@ impl H264Encoder { unsafe { (self.video_queue_fn.fp().cmd_control_video_coding_khr)( - self.encode_command_buffer, + command_buffer, &control_info, ); } @@ -569,8 +566,8 @@ impl H264Encoder { // Begin query. unsafe { self.context.device().cmd_begin_query( - self.encode_command_buffer, - self.query_pool, + command_buffer, + query_pool, 0, vk::QueryControlFlags::empty(), ); @@ -578,55 +575,36 @@ impl H264Encoder { // Encode unsafe { - (self.video_encode_fn.fp().cmd_encode_video_khr)( - self.encode_command_buffer, - &encode_info, - ); + (self.video_encode_fn.fp().cmd_encode_video_khr)(command_buffer, &encode_info); } // End query. unsafe { self.context .device() - .cmd_end_query(self.encode_command_buffer, self.query_pool, 0); + .cmd_end_query(command_buffer, query_pool, 0); } // 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, - ); + (self.video_queue_fn.fp().cmd_end_video_coding_khr)(command_buffer, &end_info); } // End command buffer. - unsafe { - self.context - .device() - .end_command_buffer(self.encode_command_buffer) - } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + unsafe { self.context.device().end_command_buffer(command_buffer) } + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - // Submit, wait, and read bitstream. + // Submit without waiting; the slot is drained on a later encode call. 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, - )? - }; + self.pipeline + .submit_current(self.context.device(), encode_queue)?; // Mark DPB slot as active. self.dpb_slot_active[self.current_dpb_slot as usize] = true; - Ok(encoded_data) + Ok(()) } } diff --git a/src/encoder/h264/init.rs b/src/encoder/h264/init.rs index cdd3b65..a815609 100644 --- a/src/encoder/h264/init.rs +++ b/src/encoder/h264/init.rs @@ -2,11 +2,10 @@ use super::{H264Encoder, MB_SIZE}; use crate::encoder::dpb::{DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig}; use crate::encoder::gop::GopStructure; +use crate::encoder::pipeline::{EncodePipeline, PipelineConfig}; 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, + align_up, allocate_session_memory, create_command_resources, create_dpb_images, + get_video_format, lcm, query_supported_video_formats, MIN_BITSTREAM_BUFFER_SIZE, }; use crate::encoder::ColorDescription; use crate::encoder::PixelFormat; @@ -413,16 +412,6 @@ impl H264Encoder { .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 @@ -444,14 +433,6 @@ impl H264Encoder { 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). @@ -461,57 +442,23 @@ impl H264Encoder { 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 the depth-N encode pipeline (per-frame input images, bitstream + // buffers, encode command buffers, fences and query pools). + let pipeline = EncodePipeline::new(&PipelineConfig { + context: &context, + aligned_width, + aligned_height, + picture_format, + pixel_format: config.pixel_format, + bit_depth: config.bit_depth, + bitstream_buffer_size: MIN_BITSTREAM_BUFFER_SIZE, + profile_info: &profile_for_resources, + command_pool, + upload_command_buffer, + upload_fence, + })?; // Create DPB and GOP structure. // The DPB size should match the actual number of allocated DPB slots. @@ -557,10 +504,7 @@ impl H264Encoder { 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, + pipeline, dpb_images, dpb_image_memories, dpb_image_views, @@ -570,16 +514,10 @@ impl H264Encoder { 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 diff --git a/src/encoder/h264/mod.rs b/src/encoder/h264/mod.rs index 3d339d3..38c3c0e 100644 --- a/src/encoder/h264/mod.rs +++ b/src/encoder/h264/mod.rs @@ -10,9 +10,8 @@ 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::pipeline::EncodePipeline; +use crate::encoder::resources::{upload_image_to_input, UploadParams}; use crate::error::Result; use crate::encoder::dpb::DecodedPictureBuffer; @@ -55,12 +54,8 @@ pub struct H264Encoder { frame_num_syntax: u32, 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, + /// Depth-N encode pipeline (per-frame slots + submission ordering). + pipeline: EncodePipeline, /// DPB images (up to MAX_DPB_SLOTS for B-frame and long-term reference support). dpb_images: Vec, dpb_image_memories: Vec, @@ -71,19 +66,11 @@ pub struct H264Encoder { 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, @@ -117,7 +104,8 @@ impl H264Encoder { /// 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 { + let slot = self.pipeline.current(); + if src_image == slot.input_image { debug!("Source image is the encoder's input image, skipping upload copy"); return Ok(()); } @@ -126,18 +114,18 @@ impl H264Encoder { upload_command_buffer: self.upload_command_buffer, upload_fence: self.upload_fence, src_image, - dst_image: self.input_image, + dst_image: slot.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, + input_image_layout: slot.input_image_layout, upload_queue: self.context.transfer_queue(), }; upload_image_to_input(&self.context, ¶ms)?; // Update tracked layout. - self.input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + self.pipeline.current_mut().input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; Ok(()) } @@ -150,37 +138,50 @@ unsafe impl Send for H264Encoder {} impl Drop for H264Encoder { fn drop(&mut self) { unsafe { + let device = self.context.device(); // 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()); + let _ = 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); + let _ = device.queue_wait_idle(q); + } + + self.pipeline.destroy(device); + + device.destroy_fence(self.upload_fence, None); + device.destroy_command_pool(self.command_pool, None); + if self.upload_command_pool != self.command_pool { + device.destroy_command_pool(self.upload_command_pool, None); + } + + for view in &self.dpb_image_views { + device.destroy_image_view(*view, None); } - 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, - }, + for image in &self.dpb_images { + device.destroy_image(*image, None); + } + for memory in &self.dpb_image_memories { + device.free_memory(*memory, None); + } + + if self.session_params != vk::VideoSessionParametersKHR::null() { + (self + .video_queue_fn + .fp() + .destroy_video_session_parameters_khr)( + device.handle(), + self.session_params, + std::ptr::null(), + ); + } + (self.video_queue_fn.fp().destroy_video_session_khr)( + device.handle(), + self.session, + std::ptr::null(), ); + for memory in &self.session_memory { + device.free_memory(*memory, None); + } } } } diff --git a/src/encoder/h265/api.rs b/src/encoder/h265/api.rs index d09b7ca..bbbb207 100644 --- a/src/encoder/h265/api.rs +++ b/src/encoder/h265/api.rs @@ -2,6 +2,7 @@ use super::H265Encoder; use crate::encoder::dpb::{DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType}; use crate::encoder::gop::{GopFrameType, GopPosition}; +use crate::encoder::pipeline::SlotPacketMetadata; use crate::encoder::{ColorDescription, EncodedPacket}; use crate::error::Result; use crate::PixelForgeError; @@ -15,7 +16,7 @@ impl H265Encoder { /// 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 + self.pipeline.input_image(self.context.device()) } /// Encode a frame from a GPU image. @@ -29,6 +30,7 @@ impl H265Encoder { /// 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 drained_packet = self.pipeline.drain_current(self.context.device())?; let gop_position = self.gop.get_next_frame(); let display_order = self.input_frame_num; self.input_frame_num += 1; @@ -41,10 +43,11 @@ impl H265Encoder { // Upload from GPU image. self.upload_from_image(src_image)?; - // Encode immediately. - let packet = self.encode_current_frame(&gop_position, display_order)?; + self.encode_current_frame(&gop_position, display_order)?; - Ok(vec![packet]) + self.pipeline.advance(); + + Ok(drained_packet.into_iter().collect()) } /// Internal method to encode the current frame already uploaded to input_image. @@ -52,7 +55,7 @@ impl H265Encoder { &mut self, gop_position: &GopPosition, display_order: u64, - ) -> Result { + ) -> 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; @@ -102,10 +105,8 @@ impl H265Encoder { 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 { + let header = if is_idr { if self.header_data.is_none() { let header = self.get_h265_header()?; // Debug: print first few bytes of header. @@ -116,22 +117,24 @@ impl H265Encoder { ); self.header_data = Some(header); } - if let Some(ref header) = self.header_data { - encoded_data.extend_from_slice(header); - } - } + self.header_data.clone() + } else { + None + }; - 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_internal(gop_position, pic_order_cnt, is_idr)?; + let dts = self.encode_frame_num; self.encode_frame_num += 1; + self.pipeline.set_pending_metadata(SlotPacketMetadata { + frame_type, + is_key_frame: is_idr, + pts: display_order, + dts, + header, + }); + if is_reference { let dpb_pic_type = if is_idr { PictureType::Idr @@ -181,19 +184,12 @@ impl H265Encoder { } } - Ok(EncodedPacket { - data: encoded_data, - frame_type, - is_key_frame: is_idr, - pts: display_order, - dts: self.encode_frame_num - 1, - }) + Ok(()) } /// Flush the encoder and get any remaining packets. pub fn flush(&mut self) -> Result> { - // No buffered frames in the current implementation. - Ok(Vec::new()) + self.pipeline.flush(self.context.device()) } /// Request that the next frame be an IDR frame. @@ -273,16 +269,15 @@ impl H265Encoder { /// 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. + // Do not reset the fences here; submit resets each slot fence before reuse. + let fences = self.pipeline.encode_fences(); unsafe { self.context .device() - .wait_for_fences(&[self.encode_fence], true, u64::MAX) + .wait_for_fences(&fences, true, u64::MAX) .map_err(|e| { PixelForgeError::Synchronization(format!( - "Failed to wait for encode fence: {:?}", + "Failed to wait for encode fences: {:?}", e )) })?; diff --git a/src/encoder/h265/encode.rs b/src/encoder/h265/encode.rs index d49eccc..d0cb67c 100644 --- a/src/encoder/h265/encode.rs +++ b/src/encoder/h265/encode.rs @@ -7,7 +7,6 @@ use super::H265Encoder; use crate::encoder::gop::{GopFrameType, GopPosition}; 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; @@ -27,14 +26,17 @@ impl H265Encoder { gop_position: &GopPosition, pic_order_cnt: i32, is_idr: bool, - ) -> Result> { + ) -> Result<()> { + let slot = self.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; + // Prepare command buffer for recording. unsafe { - prepare_encode_command_buffer( - self.context.device(), - self.encode_command_buffer, - self.query_pool, - )?; + prepare_encode_command_buffer(self.context.device(), command_buffer, query_pool)?; } // Transition DPB images for encode. @@ -42,7 +44,7 @@ impl H265Encoder { unsafe { record_dpb_barriers( self.context.device(), - self.encode_command_buffer, + command_buffer, &self.dpb_images, self.use_layered_dpb, self.current_dpb_slot, @@ -327,7 +329,7 @@ impl H265Encoder { height: self.aligned_height, }) .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() @@ -564,7 +566,7 @@ impl H265Encoder { unsafe { (self.video_queue_fn.fp().cmd_begin_video_coding_khr)( - self.encode_command_buffer, + command_buffer, &begin_coding_info, ); } @@ -588,7 +590,7 @@ impl H265Encoder { unsafe { (self.video_queue_fn.fp().cmd_control_video_coding_khr)( - self.encode_command_buffer, + command_buffer, &control_info, ); } @@ -600,34 +602,31 @@ impl H265Encoder { .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); unsafe { self.context.device().cmd_begin_query( - self.encode_command_buffer, - self.query_pool, + command_buffer, + query_pool, 0, vk::QueryControlFlags::empty(), ); - (self.video_encode_fn.fp().cmd_encode_video_khr)( - self.encode_command_buffer, - &encode_info, - ); + (self.video_encode_fn.fp().cmd_encode_video_khr)(command_buffer, &encode_info); self.context .device() - .cmd_end_query(self.encode_command_buffer, self.query_pool, 0); + .cmd_end_query(command_buffer, query_pool, 0); } // Add DPB synchronization barrier after encoding. unsafe { record_post_encode_dpb_barrier( self.context.device(), - self.encode_command_buffer, + command_buffer, &self.dpb_images, self.use_layered_dpb, self.current_dpb_slot, @@ -637,21 +636,14 @@ impl H265Encoder { // 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, - ); + (self.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) - } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + unsafe { self.context.device().end_command_buffer(command_buffer) } + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - // Submit, wait, and read bitstream. + // Submit without waiting; the slot is drained on a later encode call. let encode_queue = self.context.video_encode_queue().ok_or_else(|| { PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) })?; @@ -664,24 +656,12 @@ impl H265Encoder { 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()); + self.pipeline + .submit_current(self.context.device(), encode_queue)?; // Mark DPB slot as active. self.dpb_slot_active[self.current_dpb_slot as usize] = true; - Ok(encoded_data) + Ok(()) } } diff --git a/src/encoder/h265/init.rs b/src/encoder/h265/init.rs index deac373..0b599a7 100644 --- a/src/encoder/h265/init.rs +++ b/src/encoder/h265/init.rs @@ -2,10 +2,10 @@ use super::H265Encoder; use crate::encoder::dpb::{DecodedPictureBuffer, DecodedPictureBufferTrait, DpbConfig}; use crate::encoder::gop::GopStructure; +use crate::encoder::pipeline::{EncodePipeline, PipelineConfig}; 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, + align_up, allocate_session_memory, create_command_resources, create_dpb_images, + get_video_format, lcm, make_codec_name, query_supported_video_formats, MIN_BITSTREAM_BUFFER_SIZE, }; use crate::encoder::{BitDepth, ColorDescription, PixelFormat}; @@ -308,16 +308,6 @@ impl H265Encoder { .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 @@ -339,14 +329,6 @@ impl H265Encoder { 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). @@ -356,57 +338,23 @@ impl H265Encoder { 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 the depth-N encode pipeline (per-frame input images, bitstream + // buffers, encode command buffers, fences and query pools). + let pipeline = EncodePipeline::new(&PipelineConfig { + context: &context, + aligned_width, + aligned_height, + picture_format, + pixel_format: config.pixel_format, + bit_depth: config.bit_depth, + bitstream_buffer_size: MIN_BITSTREAM_BUFFER_SIZE, + profile_info: &profile_for_resources, + command_pool, + upload_command_buffer, + upload_fence, + })?; // Create DPB and GOP structure let mut dpb = DecodedPictureBuffer::new(); @@ -451,25 +399,16 @@ impl H265Encoder { 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, + pipeline, 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, diff --git a/src/encoder/h265/mod.rs b/src/encoder/h265/mod.rs index 522f262..b133cb6 100644 --- a/src/encoder/h265/mod.rs +++ b/src/encoder/h265/mod.rs @@ -12,9 +12,8 @@ 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::pipeline::EncodePipeline; +use crate::encoder::resources::{upload_image_to_input, UploadParams}; use crate::encoder::EncodeConfig; use crate::error::Result; use crate::vulkan::VideoContext; @@ -51,12 +50,8 @@ pub struct H265Encoder { 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, + /// Depth-N encode pipeline (per-frame slots + submission ordering). + pipeline: EncodePipeline, /// DPB images. dpb_images: Vec, dpb_image_memories: Vec, @@ -65,19 +60,11 @@ pub struct H265Encoder { 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) header_data: Option>, @@ -111,7 +98,8 @@ impl H265Encoder { /// 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 { + let slot = self.pipeline.current(); + if src_image == slot.input_image { debug!("Source image is the encoder's input image, skipping upload copy"); return Ok(()); } @@ -120,18 +108,18 @@ impl H265Encoder { upload_command_buffer: self.upload_command_buffer, upload_fence: self.upload_fence, src_image, - dst_image: self.input_image, + dst_image: slot.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, + input_image_layout: slot.input_image_layout, upload_queue: self.context.transfer_queue(), }; upload_image_to_input(&self.context, ¶ms)?; // Update tracked layout. - self.input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; + self.pipeline.current_mut().input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; Ok(()) } @@ -144,37 +132,46 @@ unsafe impl Send for H265Encoder {} impl Drop for H265Encoder { fn drop(&mut self) { unsafe { + let device = self.context.device(); // 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()); + let _ = 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); + let _ = device.queue_wait_idle(q); + } + self.pipeline.destroy(device); + device.destroy_fence(self.upload_fence, None); + device.destroy_command_pool(self.command_pool, None); + if self.upload_command_pool != self.command_pool { + device.destroy_command_pool(self.upload_command_pool, None); + } + for view in &self.dpb_image_views { + device.destroy_image_view(*view, None); + } + for image in &self.dpb_images { + device.destroy_image(*image, None); } - 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, - }, + for memory in &self.dpb_image_memories { + device.free_memory(*memory, None); + } + if self.session_params != vk::VideoSessionParametersKHR::null() { + (self + .video_queue_fn + .fp() + .destroy_video_session_parameters_khr)( + device.handle(), + self.session_params, + std::ptr::null(), + ); + } + (self.video_queue_fn.fp().destroy_video_session_khr)( + device.handle(), + self.session, + std::ptr::null(), ); + for memory in &self.session_memory { + device.free_memory(*memory, None); + } } } } diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs index 983ba12..dfd6998 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -11,6 +11,7 @@ pub mod dpb; pub mod gop; pub mod h264; pub mod h265; +pub(crate) mod pipeline; pub mod reorder; pub mod resources; @@ -445,6 +446,9 @@ 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. + /// Encoders may return packets from earlier submitted frames because encode + /// readback is pipelined. Call [`Encoder::flush`] at end-of-stream to drain + /// any remaining packets. /// /// Use `InputImage` to create an image from YUV data: /// ```no_run diff --git a/src/encoder/pipeline.rs b/src/encoder/pipeline.rs new file mode 100644 index 0000000..f4f085f --- /dev/null +++ b/src/encoder/pipeline.rs @@ -0,0 +1,331 @@ +//! Depth-N 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. +//! +//! 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), while bitstream readback is deferred until the +//! slot is reused — keeping the encode hardware continuously fed. + +use ash::vk; +use tracing::warn; + +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; + +/// 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 (and end-to-end +/// latency) 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 slot is later drained. +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>, +} + +/// 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, + + /// Whether an encode has been submitted to this slot and not yet drained. + pub in_flight: bool, + 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. +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, +} + +impl EncodePipeline { + /// Allocate the timeline semaphore and `ENCODE_PIPELINE_DEPTH` slots. + 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 `wait_for_encode_fences` is safe 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, + in_flight: false, + pending_metadata: None, + }); + } + + Ok(Self { + slots, + current_slot: 0, + timeline, + next_value: 1, + last_value: 0, + }) + } + + /// 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 for any encode still + /// reading it to finish (so it is safe to use as a convert/upload target). + pub(crate) fn input_image(&self, device: &ash::Device) -> vk::Image { + if let Err(e) = self.wait_current_slot(device) { + warn!("Failed to wait for encode input slot: {e}"); + } + self.slots[self.current_slot].input_image + } + + /// Wait for the current slot's in-flight encode (if any) to finish. + pub(crate) fn wait_current_slot(&self, device: &ash::Device) -> Result<()> { + let slot = &self.slots[self.current_slot]; + if slot.in_flight { + unsafe { device.wait_for_fences(&[slot.encode_fence], true, u64::MAX) } + .map_err(|e| PixelForgeError::Synchronization(e.to_string()))?; + } + Ok(()) + } + + /// Record the metadata for the packet that the current slot will produce. + 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. + /// + /// Chains onto the timeline semaphore so the GPU keeps encodes in DPB order, + /// and marks the slot in-flight. The bitstream is drained later, when the + /// slot is reused. + 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 = &self.slots[self.current_slot]; + + unsafe { + submit_encode_only( + device, + slot.encode_command_buffer, + slot.encode_fence, + encode_queue, + wait, + Some((self.timeline, signal_value)), + )?; + } + + self.last_value = signal_value; + self.next_value = signal_value + 1; + self.slots[self.current_slot].in_flight = true; + Ok(()) + } + + /// 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(); + } + + /// Drain the current slot, returning its packet if an encode was in flight. + pub(crate) fn drain_current(&mut self, device: &ash::Device) -> Result> { + Self::drain_slot(device, &mut self.slots[self.current_slot]) + } + + /// Drain every in-flight slot, in submission order starting from the current + /// one. Used to flush remaining packets at end of stream. + pub(crate) fn flush(&mut self, device: &ash::Device) -> Result> { + let mut packets = Vec::new(); + let len = self.slots.len(); + for offset in 0..len { + let index = (self.current_slot + offset) % len; + if let Some(packet) = Self::drain_slot(device, &mut self.slots[index])? { + packets.push(packet); + } + } + Ok(packets) + } + + /// Fences of every slot, for waiting before mutating shared session state. + pub(crate) fn encode_fences(&self) -> Vec { + self.slots.iter().map(|slot| slot.encode_fence).collect() + } + + fn drain_slot(device: &ash::Device, slot: &mut EncodeSlot) -> Result> { + if !slot.in_flight { + return Ok(None); + } + + let bitstream = unsafe { + wait_and_read_bitstream( + device, + slot.encode_fence, + slot.query_pool, + slot.bitstream_buffer_ptr, + )? + }; + slot.in_flight = false; + + let meta = slot.pending_metadata.take().ok_or_else(|| { + PixelForgeError::CommandBuffer( + "Drained encode slot has bitstream data but no packet metadata".to_string(), + ) + })?; + + let mut data = meta.header.unwrap_or_default(); + data.extend_from_slice(&bitstream); + + Ok(Some(EncodedPacket { + data, + frame_type: meta.frame_type, + is_key_frame: meta.is_key_frame, + pts: meta.pts, + dts: meta.dts, + })) + } + + /// 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) { + 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); + } +} diff --git a/src/encoder/resources.rs b/src/encoder/resources.rs index d738129..d1d2305 100644 --- a/src/encoder/resources.rs +++ b/src/encoder/resources.rs @@ -250,6 +250,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 +520,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 +544,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 +573,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 +1179,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,31 +1346,54 @@ 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. /// /// # 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 = + wait_timeline: Option<(vk::Semaphore, u64)>, + signal_timeline: Option<(vk::Semaphore, u64)>, +) -> Result<()> { + let mut wait_semaphores = Vec::new(); + let mut wait_dst_stage_mask = Vec::new(); + let mut wait_values = Vec::new(); + let mut signal_semaphores = Vec::new(); + let mut signal_values = Vec::new(); + + let mut submit_info = vk::SubmitInfo::default().command_buffers(std::slice::from_ref(&command_buffer)); - // 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. + if let Some((semaphore, value)) = wait_timeline { + wait_semaphores.push(semaphore); + wait_dst_stage_mask.push(vk::PipelineStageFlags::ALL_COMMANDS); + wait_values.push(value); + submit_info = submit_info + .wait_semaphores(&wait_semaphores) + .wait_dst_stage_mask(&wait_dst_stage_mask); + } + + if let Some((semaphore, value)) = signal_timeline { + signal_semaphores.push(semaphore); + signal_values.push(value); + submit_info = submit_info.signal_semaphores(&signal_semaphores); + } + + let mut timeline_info = vk::TimelineSemaphoreSubmitInfo::default() + .wait_semaphore_values(&wait_values) + .signal_semaphore_values(&signal_values); + + if !wait_values.is_empty() || !signal_values.is_empty() { + submit_info = submit_info.push(&mut timeline_info); + } + device .reset_fences(&[fence]) .map_err(|e| PixelForgeError::Synchronization(e.to_string()))?; @@ -1432,6 +1402,20 @@ pub(crate) unsafe fn submit_encode_and_read_bitstream( .queue_submit(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. +pub(crate) unsafe fn wait_and_read_bitstream( + device: &ash::Device, + fence: vk::Fence, + query_pool: vk::QueryPool, + bitstream_buffer_ptr: *const u8, +) -> Result> { device .wait_for_fences(&[fence], true, u64::MAX) .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; 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); From b8665d5348469d03217ec9288c76a1480fddc27f Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Sun, 14 Jun 2026 18:34:59 +0800 Subject: [PATCH 2/9] nit: scope encode timeline sync to VIDEO_ENCODE stage; document pipelined encode() --- src/encoder/av1/api.rs | 6 ++++ src/encoder/h264/api.rs | 6 ++++ src/encoder/h265/api.rs | 6 ++++ src/encoder/resources.rs | 62 +++++++++++++++++++--------------------- 4 files changed, 48 insertions(+), 32 deletions(-) diff --git a/src/encoder/av1/api.rs b/src/encoder/av1/api.rs index 6b3e536..b9808f2 100644 --- a/src/encoder/av1/api.rs +++ b/src/encoder/av1/api.rs @@ -22,6 +22,12 @@ impl AV1Encoder { /// 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. /// + /// Encoding is pipelined (depth 2): this submits the given frame to the encode + /// queue without blocking, then returns the packet for a *previously* submitted + /// frame. The first call returns an empty `Vec` while the pipeline fills, and + /// thereafter one packet per call. Call [`Self::flush`] at end-of-stream to drain + /// the remaining in-flight frames. + /// /// # Panics /// /// The encoder will panic at creation time if B-frames are enabled (b_frame_count > 0), diff --git a/src/encoder/h264/api.rs b/src/encoder/h264/api.rs index b2bfa85..92aa82f 100644 --- a/src/encoder/h264/api.rs +++ b/src/encoder/h264/api.rs @@ -25,6 +25,12 @@ impl H264Encoder { /// 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. /// + /// Encoding is pipelined (depth 2): this submits the given frame to the encode + /// queue without blocking, then returns the packet for a *previously* submitted + /// frame. The first call returns an empty `Vec` while the pipeline fills, and + /// thereafter one packet per call. Call [`Self::flush`] at end-of-stream to drain + /// the remaining in-flight frames. + /// /// # Panics /// /// The encoder will panic at creation time if B-frames are enabled (b_frame_count > 0), diff --git a/src/encoder/h265/api.rs b/src/encoder/h265/api.rs index bbbb207..fa949e4 100644 --- a/src/encoder/h265/api.rs +++ b/src/encoder/h265/api.rs @@ -25,6 +25,12 @@ impl H265Encoder { /// 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. /// + /// Encoding is pipelined (depth 2): this submits the given frame to the encode + /// queue without blocking, then returns the packet for a *previously* submitted + /// frame. The first call returns an empty `Vec` while the pipeline fills, and + /// thereafter one packet per call. Call [`Self::flush`] at end-of-stream to drain + /// the remaining in-flight frames. + /// /// # Panics /// /// The encoder will panic at creation time if B-frames are enabled (b_frame_count > 0), diff --git a/src/encoder/resources.rs b/src/encoder/resources.rs index d1d2305..3c9a05d 100644 --- a/src/encoder/resources.rs +++ b/src/encoder/resources.rs @@ -1349,7 +1349,10 @@ pub(crate) unsafe fn record_post_encode_dpb_barrier( /// Submit an encode command buffer without waiting for completion. /// /// Timeline semaphore waits/signals are used to keep shared DPB access ordered -/// across pipelined slots while leaving bitstream readback delayed. +/// 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 /// @@ -1362,44 +1365,39 @@ pub(crate) unsafe fn submit_encode_only( wait_timeline: Option<(vk::Semaphore, u64)>, signal_timeline: Option<(vk::Semaphore, u64)>, ) -> Result<()> { - let mut wait_semaphores = Vec::new(); - let mut wait_dst_stage_mask = Vec::new(); - let mut wait_values = Vec::new(); - let mut signal_semaphores = Vec::new(); - let mut signal_values = Vec::new(); - - let mut submit_info = - vk::SubmitInfo::default().command_buffers(std::slice::from_ref(&command_buffer)); - - if let Some((semaphore, value)) = wait_timeline { - wait_semaphores.push(semaphore); - wait_dst_stage_mask.push(vk::PipelineStageFlags::ALL_COMMANDS); - wait_values.push(value); - submit_info = submit_info - .wait_semaphores(&wait_semaphores) - .wait_dst_stage_mask(&wait_dst_stage_mask); - } - - if let Some((semaphore, value)) = signal_timeline { - signal_semaphores.push(semaphore); - signal_values.push(value); - submit_info = submit_info.signal_semaphores(&signal_semaphores); - } - - let mut timeline_info = vk::TimelineSemaphoreSubmitInfo::default() - .wait_semaphore_values(&wait_values) - .signal_semaphore_values(&signal_values); + 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(); - if !wait_values.is_empty() || !signal_values.is_empty() { - submit_info = submit_info.push(&mut timeline_info); - } + let submit_info = vk::SubmitInfo2::default() + .wait_semaphore_infos(&wait_infos) + .command_buffer_infos(&command_buffer_infos) + .signal_semaphore_infos(&signal_infos); 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(()) From b2a0e160fefebfaad6b428e3482d859033403393 Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Sun, 14 Jun 2026 23:33:45 +0800 Subject: [PATCH 3/9] feat: async push-based encode readback on a completion thread --- examples/encode_av1.rs | 3 +- examples/encode_h264.rs | 5 +- examples/encode_h265.rs | 5 +- examples/verify_all.rs | 3 +- src/encoder/av1/api.rs | 73 +++++---- src/encoder/h264/api.rs | 71 +++++---- src/encoder/h265/api.rs | 65 ++++---- src/encoder/mod.rs | 50 +++++- src/encoder/pipeline.rs | 334 ++++++++++++++++++++++++++++++---------- src/lib.rs | 7 +- 10 files changed, 428 insertions(+), 188 deletions(-) diff --git a/examples/encode_av1.rs b/examples/encode_av1.rs index 2e08bca..8d9809d 100644 --- a/examples/encode_av1.rs +++ b/examples/encode_av1.rs @@ -95,7 +95,8 @@ fn main() -> Result<(), Box> { // 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())? { + encoder.encode(input_image.image())?; + while let Some(packet) = encoder.poll_packet()? { total_bytes += packet.data.len(); output.write_all(&packet.data)?; println!( diff --git a/examples/encode_h264.rs b/examples/encode_h264.rs index 64bfc24..b6ac224 100644 --- a/examples/encode_h264.rs +++ b/examples/encode_h264.rs @@ -92,8 +92,9 @@ 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())? { + // Encode the image (async); drain any packets that are ready. + encoder.encode(input_image.image())?; + while let Some(packet) = encoder.poll_packet()? { total_bytes += packet.data.len(); output.write_all(&packet.data)?; println!( diff --git a/examples/encode_h265.rs b/examples/encode_h265.rs index 1ac971c..59b04b3 100644 --- a/examples/encode_h265.rs +++ b/examples/encode_h265.rs @@ -98,8 +98,9 @@ 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())? { + // Encode the image (async); drain any packets that are ready. + encoder.encode(input_image.image())?; + while let Some(packet) = encoder.poll_packet()? { total_bytes += packet.data.len(); output.write_all(&packet.data)?; println!( diff --git a/examples/verify_all.rs b/examples/verify_all.rs index 4f6a1ee..f4dc491 100644 --- a/examples/verify_all.rs +++ b/examples/verify_all.rs @@ -171,7 +171,8 @@ fn run_test( _ => return Err("Unsupported format".into()), } - for packet in encoder.encode(encoder_image)? { + encoder.encode(encoder_image)?; + while let Some(packet) = encoder.poll_packet()? { output_file.write_all(&packet.data)?; } } diff --git a/src/encoder/av1/api.rs b/src/encoder/av1/api.rs index b9808f2..f4e21de 100644 --- a/src/encoder/av1/api.rs +++ b/src/encoder/av1/api.rs @@ -13,7 +13,21 @@ impl AV1Encoder { /// This image can be used as a target for `ColorConverter::convert` to avoid /// an intermediate copy. pub fn input_image(&self) -> vk::Image { - self.pipeline.input_image(self.context.device()) + self.pipeline.input_image() + } + + /// Try to receive a finished encoded packet without blocking. + /// + /// Encoding is asynchronous (see [`Self::encode`]); packets are produced by + /// a background readback thread and surfaced here. Returns `None` when no + /// packet is currently ready. + pub fn poll_packet(&self) -> Result> { + self.pipeline.poll_packet().transpose() + } + + /// Take ownership of the encoded-packet receiver for out-of-band consumption. + pub fn take_packet_receiver(&mut self) -> Option { + self.pipeline.take_packet_receiver() } /// Encode a frame from a GPU image. @@ -22,18 +36,21 @@ impl AV1Encoder { /// 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. /// - /// Encoding is pipelined (depth 2): this submits the given frame to the encode - /// queue without blocking, then returns the packet for a *previously* submitted - /// frame. The first call returns an empty `Vec` while the pipeline fills, and - /// thereafter one packet per call. Call [`Self::flush`] at end-of-stream to drain - /// the remaining in-flight frames. + /// Encoding is asynchronous and pipelined (depth 2): this submits the given + /// frame to the encode queue without blocking and returns immediately. The + /// encoded packet is produced by a background readback thread and retrieved + /// via [`Self::poll_packet`] (or [`Self::flush`] at end-of-stream). This call + /// only blocks if all pipeline slots are still in flight (backpressure). /// /// # 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 drained_packet = self.pipeline.drain_current(self.context.device())?; + pub fn encode(&mut self, src_image: vk::Image) -> Result<()> { + // Ensure the slot we are about to record over and reuse has been read + // back by the completion thread. + self.pipeline.wait_current_free(); + let gop_position = self.gop.get_next_frame(); let display_order = self.input_frame_num; self.input_frame_num += 1; @@ -50,7 +67,7 @@ impl AV1Encoder { self.pipeline.advance(); - Ok(drained_packet.into_iter().collect()) + Ok(()) } /// Internal method to encode the current frame already uploaded to input_image. @@ -104,15 +121,10 @@ impl AV1Encoder { } } - 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; + // Record packet metadata before submitting: `submit_current` hands it to + // the readback thread together with the bitstream. `dts` reads the + // pre-increment encode counter (also used by `encode_frame_internal`). let dts = self.encode_frame_num; - self.encode_frame_num += 1; - self.frame_num += 1; - self.order_hint = (self.order_hint + 1) & 0xFF; // 8-bit order hint - self.pipeline.set_pending_metadata(SlotPacketMetadata { frame_type, is_key_frame, @@ -121,6 +133,14 @@ impl AV1Encoder { header: Some(header), }); + 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 @@ -163,7 +183,7 @@ impl AV1Encoder { /// Flush the encoder and get any remaining packets. pub fn flush(&mut self) -> Result> { - self.pipeline.flush(self.context.device()) + self.pipeline.flush() } /// Request that the next frame be an IDR/key frame. @@ -234,20 +254,9 @@ impl AV1Encoder { /// 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 fences here; submit resets each slot fence before reuse. - let fences = self.pipeline.encode_fences(); - unsafe { - self.context - .device() - .wait_for_fences(&fences, true, u64::MAX) - .map_err(|e| { - PixelForgeError::Synchronization(format!( - "Failed to wait for encode fences: {:?}", - e - )) - })?; - } + // Wait for all in-flight encodes to be read back before modifying the + // shared session parameters. + self.pipeline.wait_all_free(); // Save old handle so we can destroy it after successful creation. let old_session_params = self.session_params; diff --git a/src/encoder/h264/api.rs b/src/encoder/h264/api.rs index 92aa82f..6bafe49 100644 --- a/src/encoder/h264/api.rs +++ b/src/encoder/h264/api.rs @@ -16,7 +16,21 @@ impl H264Encoder { /// This image can be used as a target for `ColorConverter::convert` to avoid /// an intermediate copy. pub fn input_image(&self) -> vk::Image { - self.pipeline.input_image(self.context.device()) + self.pipeline.input_image() + } + + /// Try to receive a finished encoded packet without blocking. + /// + /// Encoding is asynchronous (see [`Self::encode`]); packets are produced by + /// a background readback thread and surfaced here. Returns `None` when no + /// packet is currently ready. + pub fn poll_packet(&self) -> Result> { + self.pipeline.poll_packet().transpose() + } + + /// Take ownership of the encoded-packet receiver for out-of-band consumption. + pub fn take_packet_receiver(&mut self) -> Option { + self.pipeline.take_packet_receiver() } /// Encode a frame from a GPU image. @@ -25,18 +39,21 @@ impl H264Encoder { /// 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. /// - /// Encoding is pipelined (depth 2): this submits the given frame to the encode - /// queue without blocking, then returns the packet for a *previously* submitted - /// frame. The first call returns an empty `Vec` while the pipeline fills, and - /// thereafter one packet per call. Call [`Self::flush`] at end-of-stream to drain - /// the remaining in-flight frames. + /// Encoding is asynchronous and pipelined (depth 2): this submits the given + /// frame to the encode queue without blocking and returns immediately. The + /// encoded packet is produced by a background readback thread and retrieved + /// via [`Self::poll_packet`] (or [`Self::flush`] at end-of-stream). This call + /// only blocks if all pipeline slots are still in flight (backpressure). /// /// # 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 drained_packet = self.pipeline.drain_current(self.context.device())?; + pub fn encode(&mut self, src_image: vk::Image) -> Result<()> { + // Ensure the slot we are about to record over and reuse has been read + // back by the completion thread. + self.pipeline.wait_current_free(); + let gop_position = self.gop.get_next_frame(); let display_order = self.input_frame_num; self.input_frame_num += 1; @@ -53,7 +70,7 @@ impl H264Encoder { self.pipeline.advance(); - Ok(drained_packet.into_iter().collect()) + Ok(()) } /// Internal method to encode the current frame already uploaded to input_image. @@ -103,14 +120,10 @@ impl H264Encoder { None }; - self.encode_frame_internal(gop_position, frame_num, pic_order_cnt, is_idr)?; - + // Record packet metadata before submitting: `submit_current` hands it to + // the readback thread together with the bitstream. `dts` reads the + // pre-increment encode counter (also used by `encode_frame_internal`). let dts = self.encode_frame_num; - self.encode_frame_num += 1; - if is_reference && !is_b_frame { - self.frame_num_syntax = (self.frame_num_syntax + 1) % 256; - } - self.pipeline.set_pending_metadata(SlotPacketMetadata { frame_type, is_key_frame: is_idr, @@ -119,6 +132,13 @@ impl H264Encoder { header, }); + 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 @@ -166,7 +186,7 @@ impl H264Encoder { /// Flush the encoder and get any remaining packets. pub fn flush(&mut self) -> Result> { - self.pipeline.flush(self.context.device()) + self.pipeline.flush() } /// Request that the next frame be an IDR frame. @@ -245,20 +265,9 @@ impl H264Encoder { /// 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 fences here; submit resets each slot fence before reuse. - let fences = self.pipeline.encode_fences(); - unsafe { - self.context - .device() - .wait_for_fences(&fences, true, u64::MAX) - .map_err(|e| { - PixelForgeError::Synchronization(format!( - "Failed to wait for encode fences: {:?}", - e - )) - })?; - } + // Wait for all in-flight encodes to be read back before modifying the + // shared session parameters. + self.pipeline.wait_all_free(); // Save old handle so we can destroy it after successful creation. let old_session_params = self.session_params; diff --git a/src/encoder/h265/api.rs b/src/encoder/h265/api.rs index fa949e4..01be55d 100644 --- a/src/encoder/h265/api.rs +++ b/src/encoder/h265/api.rs @@ -16,7 +16,21 @@ impl H265Encoder { /// This image can be used as a target for `ColorConverter::convert` to avoid /// an intermediate copy. pub fn input_image(&self) -> vk::Image { - self.pipeline.input_image(self.context.device()) + self.pipeline.input_image() + } + + /// Try to receive a finished encoded packet without blocking. + /// + /// Encoding is asynchronous (see [`Self::encode`]); packets are produced by + /// a background readback thread and surfaced here. Returns `None` when no + /// packet is currently ready. + pub fn poll_packet(&self) -> Result> { + self.pipeline.poll_packet().transpose() + } + + /// Take ownership of the encoded-packet receiver for out-of-band consumption. + pub fn take_packet_receiver(&mut self) -> Option { + self.pipeline.take_packet_receiver() } /// Encode a frame from a GPU image. @@ -25,18 +39,21 @@ impl H265Encoder { /// 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. /// - /// Encoding is pipelined (depth 2): this submits the given frame to the encode - /// queue without blocking, then returns the packet for a *previously* submitted - /// frame. The first call returns an empty `Vec` while the pipeline fills, and - /// thereafter one packet per call. Call [`Self::flush`] at end-of-stream to drain - /// the remaining in-flight frames. + /// Encoding is asynchronous and pipelined (depth 2): this submits the given + /// frame to the encode queue without blocking and returns immediately. The + /// encoded packet is produced by a background readback thread and retrieved + /// via [`Self::poll_packet`] (or [`Self::flush`] at end-of-stream). This call + /// only blocks if all pipeline slots are still in flight (backpressure). /// /// # 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 drained_packet = self.pipeline.drain_current(self.context.device())?; + pub fn encode(&mut self, src_image: vk::Image) -> Result<()> { + // Ensure the slot we are about to record over and reuse has been read + // back by the completion thread. + self.pipeline.wait_current_free(); + let gop_position = self.gop.get_next_frame(); let display_order = self.input_frame_num; self.input_frame_num += 1; @@ -53,7 +70,7 @@ impl H265Encoder { self.pipeline.advance(); - Ok(drained_packet.into_iter().collect()) + Ok(()) } /// Internal method to encode the current frame already uploaded to input_image. @@ -128,11 +145,10 @@ impl H265Encoder { None }; - self.encode_frame_internal(gop_position, pic_order_cnt, is_idr)?; - + // Record packet metadata before submitting: `submit_current` hands it to + // the readback thread together with the bitstream. `dts` reads the + // pre-increment encode counter (also used by `encode_frame_internal`). let dts = self.encode_frame_num; - self.encode_frame_num += 1; - self.pipeline.set_pending_metadata(SlotPacketMetadata { frame_type, is_key_frame: is_idr, @@ -141,6 +157,10 @@ impl H265Encoder { header, }); + self.encode_frame_internal(gop_position, pic_order_cnt, is_idr)?; + + self.encode_frame_num += 1; + if is_reference { let dpb_pic_type = if is_idr { PictureType::Idr @@ -195,7 +215,7 @@ impl H265Encoder { /// Flush the encoder and get any remaining packets. pub fn flush(&mut self) -> Result> { - self.pipeline.flush(self.context.device()) + self.pipeline.flush() } /// Request that the next frame be an IDR frame. @@ -274,20 +294,9 @@ impl H265Encoder { /// 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 fences here; submit resets each slot fence before reuse. - let fences = self.pipeline.encode_fences(); - unsafe { - self.context - .device() - .wait_for_fences(&fences, true, u64::MAX) - .map_err(|e| { - PixelForgeError::Synchronization(format!( - "Failed to wait for encode fences: {:?}", - e - )) - })?; - } + // Wait for all in-flight encodes to be read back before modifying the + // shared session parameters. + self.pipeline.wait_all_free(); // Save old handle so we can destroy it after successful creation. let old_session_params = self.session_params; diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs index dfd6998..d00b342 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -386,6 +386,11 @@ impl EncodeConfig { } } +/// Channel receiver for encoded packets, produced by the encoder's background +/// readback thread. Obtained via [`Encoder::take_packet_receiver`] for +/// out-of-band (push) consumption from a dedicated thread. +pub type EncodedPacketReceiver = std::sync::mpsc::Receiver>; + /// Encoded video packet. #[derive(Debug, Clone)] pub struct EncodedPacket { @@ -446,9 +451,12 @@ 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. - /// Encoders may return packets from earlier submitted frames because encode - /// readback is pipelined. Call [`Encoder::flush`] at end-of-stream to drain - /// any remaining packets. + /// + /// Encoding is asynchronous: this submits the frame without blocking and + /// returns immediately. Encoded packets are produced by a background readback + /// thread and retrieved via [`Encoder::poll_packet`] (or [`Encoder::flush`] + /// at end-of-stream). The call only blocks if every pipeline slot is still in + /// flight (backpressure). /// /// Use `InputImage` to create an image from YUV data: /// ```no_run @@ -470,12 +478,15 @@ 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, then drain any ready packets. + /// encoder.encode(input.image())?; + /// while let Some(packet) = encoder.poll_packet()? { + /// // ... use packet.data ... + /// } /// # Ok(()) /// # } /// ``` - pub fn encode(&mut self, src_image: vk::Image) -> Result> { + 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), @@ -483,6 +494,33 @@ impl Encoder { } } + /// Try to receive a finished encoded packet without blocking. + /// + /// Returns `Ok(None)` when no packet is currently ready. Encoded packets are + /// produced asynchronously after [`Encoder::encode`]; poll until `None` to + /// drain everything currently available. + pub fn poll_packet(&self) -> Result> { + match self { + Encoder::H264(encoder) => encoder.poll_packet(), + Encoder::H265(encoder) => encoder.poll_packet(), + Encoder::AV1(encoder) => encoder.poll_packet(), + } + } + + /// Take ownership of the encoded-packet receiver for out-of-band (push) + /// consumption from a dedicated thread. + /// + /// After this is called, [`Encoder::poll_packet`] and [`Encoder::flush`] no + /// longer yield packets — they are delivered to the returned receiver + /// instead. Returns `None` if the receiver has already been taken. + pub fn take_packet_receiver(&mut self) -> Option { + match self { + Encoder::H264(encoder) => encoder.take_packet_receiver(), + Encoder::H265(encoder) => encoder.take_packet_receiver(), + Encoder::AV1(encoder) => encoder.take_packet_receiver(), + } + } + /// Flush the encoder and get remaining packets. pub fn flush(&mut self) -> Result> { match self { diff --git a/src/encoder/pipeline.rs b/src/encoder/pipeline.rs index f4f085f..9f3feae 100644 --- a/src/encoder/pipeline.rs +++ b/src/encoder/pipeline.rs @@ -1,4 +1,4 @@ -//! Depth-N encode pipelining, shared by all codecs. +//! 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`] @@ -6,21 +6,38 @@ //! 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 pushes the finished [`EncodedPacket`] onto a channel 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. +//! //! 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), while bitstream readback is deferred until the -//! slot is reused — keeping the encode hardware continuously fed. +//! 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::sync::mpsc::{Receiver, Sender}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread::JoinHandle; use ash::vk; -use tracing::warn; 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::encoder::{BitDepth, EncodedPacket, EncodedPacketReceiver, FrameType, PixelFormat}; use crate::error::{PixelForgeError, Result}; use crate::vulkan::VideoContext; @@ -28,12 +45,11 @@ use crate::vulkan::VideoContext; /// /// 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 (and end-to-end -/// latency) from growing. +/// 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 slot is later drained. +/// when the completion thread reads the slot back. pub(crate) struct SlotPacketMetadata { pub frame_type: FrameType, pub is_key_frame: bool, @@ -44,6 +60,68 @@ pub(crate) struct SlotPacketMetadata { 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, +} + +/// 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, @@ -62,8 +140,8 @@ pub(crate) struct EncodeSlot { pub encode_fence: vk::Fence, pub query_pool: vk::QueryPool, - /// Whether an encode has been submitted to this slot and not yet drained. - pub in_flight: bool, + /// Packet metadata recorded before submit, moved to the completion thread + /// with the work item. pub pending_metadata: Option, } @@ -86,7 +164,7 @@ pub(crate) struct PipelineConfig<'a> { } /// Rotating set of [`EncodeSlot`]s plus the timeline semaphore that orders their -/// encode submissions. +/// encode submissions and the completion thread that reads bitstreams back. pub(crate) struct EncodePipeline { slots: Vec, current_slot: usize, @@ -96,10 +174,22 @@ pub(crate) struct EncodePipeline { 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>, + /// Receives finished packets from the completion thread. `None` once the + /// caller has taken ownership via [`EncodePipeline::take_packet_receiver`]. + packet_rx: Option, + /// The completion thread handle, joined on shutdown. + completion_thread: Option>, } impl EncodePipeline { - /// Allocate the timeline semaphore and `ENCODE_PIPELINE_DEPTH` slots. + /// 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(); @@ -145,8 +235,8 @@ impl EncodePipeline { }, )?; - // Created signaled so `wait_for_encode_fences` is safe before the - // first encode; `submit_encode_only` resets it before each submit. + // 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) } @@ -167,17 +257,36 @@ impl EncodePipeline { encode_command_buffer, encode_fence, query_pool, - in_flight: false, pending_metadata: None, }); } + let slot_sync = Arc::new(SlotSync::new(slots.len())); + let (work_tx, work_rx) = std::sync::mpsc::channel::(); + let (packet_tx, packet_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, packet_tx, 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), + packet_rx: Some(packet_rx), + completion_thread: Some(completion_thread), }) } @@ -190,35 +299,36 @@ impl EncodePipeline { &mut self.slots[self.current_slot] } - /// Return the current slot's input image, first waiting for any encode still - /// reading it to finish (so it is safe to use as a convert/upload target). - pub(crate) fn input_image(&self, device: &ash::Device) -> vk::Image { - if let Err(e) = self.wait_current_slot(device) { - warn!("Failed to wait for encode input slot: {e}"); - } + /// 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 for the current slot's in-flight encode (if any) to finish. - pub(crate) fn wait_current_slot(&self, device: &ash::Device) -> Result<()> { - let slot = &self.slots[self.current_slot]; - if slot.in_flight { - unsafe { device.wait_for_fences(&[slot.encode_fence], true, u64::MAX) } - .map_err(|e| PixelForgeError::Synchronization(e.to_string()))?; - } - Ok(()) + /// 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. + /// 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 in-flight. The bitstream is drained later, when the - /// slot is reused. + /// and marks the slot busy until the completion thread reads it back. pub(crate) fn submit_current( &mut self, device: &ash::Device, @@ -226,13 +336,31 @@ impl EncodePipeline { ) -> Result<()> { let wait = (self.last_value > 0).then_some((self.timeline, self.last_value)); let signal_value = self.next_value; - let slot = &self.slots[self.current_slot]; + 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, - slot.encode_command_buffer, - slot.encode_fence, + command_buffer, + fence, encode_queue, wait, Some((self.timeline, signal_value)), @@ -241,7 +369,24 @@ impl EncodePipeline { self.last_value = signal_value; self.next_value = signal_value + 1; - self.slots[self.current_slot].in_flight = true; + + // 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 work = WorkItem { + slot_index, + fence, + query_pool, + bitstream_ptr: SendPtr(bitstream_ptr), + metadata, + }; + 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(()) } @@ -250,61 +395,50 @@ impl EncodePipeline { self.current_slot = (self.current_slot + 1) % self.slots.len(); } - /// Drain the current slot, returning its packet if an encode was in flight. - pub(crate) fn drain_current(&mut self, device: &ash::Device) -> Result> { - Self::drain_slot(device, &mut self.slots[self.current_slot]) + /// Try to receive a finished packet without blocking. + /// + /// Returns `None` if no packet is ready, or if the caller has taken + /// ownership of the receiver via [`EncodePipeline::take_packet_receiver`]. + pub(crate) fn poll_packet(&self) -> Option> { + self.packet_rx.as_ref().and_then(|rx| rx.try_recv().ok()) + } + + /// Take ownership of the packet receiver for out-of-band consumption (e.g. + /// a dedicated sender thread). After this, [`EncodePipeline::poll_packet`] + /// and [`EncodePipeline::flush`] no longer return packets. + pub(crate) fn take_packet_receiver(&mut self) -> Option { + self.packet_rx.take() } - /// Drain every in-flight slot, in submission order starting from the current - /// one. Used to flush remaining packets at end of stream. - pub(crate) fn flush(&mut self, device: &ash::Device) -> Result> { + /// Wait for all in-flight frames to be read back, returning any packets not + /// yet polled. Used to drain at end of stream. + /// + /// If the receiver has been taken, this only acts as a completion barrier + /// and returns an empty `Vec` (packets are delivered to the taken receiver). + pub(crate) fn flush(&mut self) -> Result> { + // Once every slot is free, the completion thread has already sent every + // packet (it sends before marking a slot free), so draining now is + // complete. + self.slot_sync.wait_all_free(); + let mut packets = Vec::new(); - let len = self.slots.len(); - for offset in 0..len { - let index = (self.current_slot + offset) % len; - if let Some(packet) = Self::drain_slot(device, &mut self.slots[index])? { - packets.push(packet); + if let Some(rx) = &self.packet_rx { + while let Ok(result) = rx.try_recv() { + packets.push(result?); } } Ok(packets) } - /// Fences of every slot, for waiting before mutating shared session state. - pub(crate) fn encode_fences(&self) -> Vec { - self.slots.iter().map(|slot| slot.encode_fence).collect() - } - - fn drain_slot(device: &ash::Device, slot: &mut EncodeSlot) -> Result> { - if !slot.in_flight { - return Ok(None); + /// 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(); } - - let bitstream = unsafe { - wait_and_read_bitstream( - device, - slot.encode_fence, - slot.query_pool, - slot.bitstream_buffer_ptr, - )? - }; - slot.in_flight = false; - - let meta = slot.pending_metadata.take().ok_or_else(|| { - PixelForgeError::CommandBuffer( - "Drained encode slot has bitstream data but no packet metadata".to_string(), - ) - })?; - - let mut data = meta.header.unwrap_or_default(); - data.extend_from_slice(&bitstream); - - Ok(Some(EncodedPacket { - data, - frame_type: meta.frame_type, - is_key_frame: meta.is_key_frame, - pts: meta.pts, - dts: meta.dts, - })) } /// Destroy all slot resources and the timeline semaphore. @@ -313,6 +447,9 @@ impl EncodePipeline { /// /// 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); @@ -329,3 +466,36 @@ impl EncodePipeline { device.destroy_semaphore(self.timeline, None); } } + +/// Completion-thread body: wait on each submission's fence, copy its bitstream +/// out, deliver the packet, then mark the slot free. +fn run_completion_thread( + device: ash::Device, + work_rx: Receiver, + packet_tx: Sender>, + slot_sync: Arc, +) { + for work in work_rx { + let result = unsafe { + wait_and_read_bitstream(&device, work.fence, work.query_pool, work.bitstream_ptr.0) + }; + + let packet = result.map(|bitstream| { + let mut data = work.metadata.header.unwrap_or_default(); + data.extend_from_slice(&bitstream); + EncodedPacket { + data, + frame_type: work.metadata.frame_type, + is_key_frame: work.metadata.is_key_frame, + pts: work.metadata.pts, + dts: work.metadata.dts, + } + }); + + // Deliver the packet *before* freeing the slot. This ordering means that + // once all slots are observed free, every packet has already been sent, + // which `flush` relies on for completeness. + let _ = packet_tx.send(packet); + slot_sync.set_free(work.slot_index); + } +} diff --git a/src/lib.rs b/src/lib.rs index 499273f..6cf5c0b 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, EncodedPacket, + EncodedPacketReceiver, 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; From 4d784f5d5d5888293ec3d3b6c118bd25cac68223 Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Tue, 16 Jun 2026 22:38:43 +0800 Subject: [PATCH 4/9] chore: format --- src/encoder/pipeline.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/encoder/pipeline.rs b/src/encoder/pipeline.rs index 9f3feae..dad3c19 100644 --- a/src/encoder/pipeline.rs +++ b/src/encoder/pipeline.rs @@ -214,10 +214,16 @@ impl EncodePipeline { 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)?; + 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. From 6c82a5b0cf2d261e59660f2b97434466368b295c Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Fri, 19 Jun 2026 15:40:49 +0800 Subject: [PATCH 5/9] refactor: unify codecs into generic CodecEncoder with shared EncoderCommon; eliminate redundant hotpath bitstream copy --- .gitignore | 3 + Cargo.lock | 23 + Cargo.toml | 2 + examples/encode_av1.rs | 58 +- examples/encode_h264.rs | 54 +- examples/encode_h265.rs | 54 +- examples/verify_all.rs | 14 +- src/encoder/av1/api.rs | 287 --------- src/encoder/av1/encode.rs | 554 ------------------ src/encoder/av1/init.rs | 354 ++--------- src/encoder/av1/mod.rs | 337 ++++------- src/encoder/av1/record.rs | 415 +++++++++++++ src/encoder/av1/session_params.rs | 69 ++- src/encoder/codec.rs | 683 ++++++++++++++++++++++ src/encoder/h264/api.rs | 304 ---------- src/encoder/h264/init.rs | 588 ++++--------------- src/encoder/h264/mod.rs | 258 ++++---- src/encoder/h264/{encode.rs => record.rs} | 299 +++------- src/encoder/h264/session_params.rs | 99 ++-- src/encoder/h265/api.rs | 333 ----------- src/encoder/h265/init.rs | 425 ++------------ src/encoder/h265/mod.rs | 246 ++++---- src/encoder/h265/{encode.rs => record.rs} | 407 +++++-------- src/encoder/h265/session_params.rs | 73 ++- src/encoder/mod.rs | 156 ++--- src/encoder/pipeline.rs | 144 +++-- src/encoder/resources.rs | 138 ++++- src/lib.rs | 8 +- 28 files changed, 2412 insertions(+), 3973 deletions(-) delete mode 100644 src/encoder/av1/api.rs delete mode 100644 src/encoder/av1/encode.rs create mode 100644 src/encoder/av1/record.rs create mode 100644 src/encoder/codec.rs delete mode 100644 src/encoder/h264/api.rs rename src/encoder/h264/{encode.rs => record.rs} (63%) delete mode 100644 src/encoder/h265/api.rs rename src/encoder/h265/{encode.rs => record.rs} (59%) 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 8d9809d..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,35 +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). - encoder.encode(input_image.image())?; - while let Some(packet) = encoder.poll_packet()? { - 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 b6ac224..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,33 +112,19 @@ fn main() -> Result<(), Box> { // Upload YUV420 data to the input image. input_image.upload_yuv420(frame)?; - // Encode the image (async); drain any packets that are ready. - encoder.encode(input_image.image())?; - while let Some(packet) = encoder.poll_packet()? { - 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 59b04b3..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,33 +118,19 @@ fn main() -> Result<(), Box> { // Upload YUV420 data to the input image. input_image.upload_yuv420(frame)?; - // Encode the image (async); drain any packets that are ready. - encoder.encode(input_image.image())?; - while let Some(packet) = encoder.poll_packet()? { - 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/verify_all.rs b/examples/verify_all.rs index f4dc491..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,13 +175,17 @@ fn run_test( _ => return Err("Unsupported format".into()), } - encoder.encode(encoder_image)?; - while let Some(packet) = encoder.poll_packet()? { + 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)?; } } - for packet in encoder.flush()? { + // 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)?; } } diff --git a/src/encoder/av1/api.rs b/src/encoder/av1/api.rs deleted file mode 100644 index f4e21de..0000000 --- a/src/encoder/av1/api.rs +++ /dev/null @@ -1,287 +0,0 @@ -use super::AV1Encoder; - -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::pipeline::SlotPacketMetadata; -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.pipeline.input_image() - } - - /// Try to receive a finished encoded packet without blocking. - /// - /// Encoding is asynchronous (see [`Self::encode`]); packets are produced by - /// a background readback thread and surfaced here. Returns `None` when no - /// packet is currently ready. - pub fn poll_packet(&self) -> Result> { - self.pipeline.poll_packet().transpose() - } - - /// Take ownership of the encoded-packet receiver for out-of-band consumption. - pub fn take_packet_receiver(&mut self) -> Option { - self.pipeline.take_packet_receiver() - } - - /// 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. - /// - /// Encoding is asynchronous and pipelined (depth 2): this submits the given - /// frame to the encode queue without blocking and returns immediately. The - /// encoded packet is produced by a background readback thread and retrieved - /// via [`Self::poll_packet`] (or [`Self::flush`] at end-of-stream). This call - /// only blocks if all pipeline slots are still in flight (backpressure). - /// - /// # 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<()> { - // Ensure the slot we are about to record over and reuse has been read - // back by the completion thread. - self.pipeline.wait_current_free(); - - 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)?; - - self.encode_current_frame(&gop_position, display_order)?; - - self.pipeline.advance(); - - Ok(()) - } - - /// 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 header = 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. - header.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 sequence_header) = self.header_data { - header.extend_from_slice(sequence_header); - } - } - - // Record packet metadata before submitting: `submit_current` hands it to - // the readback thread together with the bitstream. `dts` reads the - // pre-increment encode counter (also used by `encode_frame_internal`). - let dts = self.encode_frame_num; - self.pipeline.set_pending_metadata(SlotPacketMetadata { - frame_type, - is_key_frame, - pts: display_order, - dts, - header: Some(header), - }); - - 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(()) - } - - /// Flush the encoder and get any remaining packets. - pub fn flush(&mut self) -> Result> { - self.pipeline.flush() - } - - /// 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 all in-flight encodes to be read back before modifying the - // shared session parameters. - self.pipeline.wait_all_free(); - - // 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 1374d1e..0000000 --- a/src/encoder/av1/encode.rs +++ /dev/null @@ -1,554 +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, -}; -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<()> { - let slot = self.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; - - // 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(), command_buffer, 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(), - 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(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(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(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); - - // Begin query to capture encode feedback (bitstream size, status). - unsafe { - self.context.device().cmd_begin_query( - command_buffer, - query_pool, - 0, - vk::QueryControlFlags::empty(), - ); - } - - unsafe { - self.video_encode_fn - .cmd_encode_video(command_buffer, &encode_info); - } - - // End query. - unsafe { - self.context - .device() - .cmd_end_query(command_buffer, query_pool, 0); - } - - // Add DPB synchronization barrier after encoding. - unsafe { - record_post_encode_dpb_barrier( - self.context.device(), - 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(command_buffer, &end_coding_info); - } - - // End command buffer. - unsafe { self.context.device().end_command_buffer(command_buffer) } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - - // Submit without waiting; the slot is drained on a later encode call. - 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 - ); - - self.pipeline - .submit_current(self.context.device(), encode_queue)?; - - // Mark current DPB slot as active. - self.dpb_slot_active[self.current_dpb_slot as usize] = true; - - Ok(()) - } - - /// 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 46cc558..0e1bea7 100644 --- a/src/encoder/av1/init.rs +++ b/src/encoder/av1/init.rs @@ -1,345 +1,87 @@ -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::pipeline::{EncodePipeline, PipelineConfig}; -use crate::encoder::resources::{ - allocate_session_memory, create_command_resources, create_dpb_images, get_video_format, - make_codec_name, query_supported_video_formats, +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 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 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()) - })?; - - // 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 - ))); - } + let bitstream_buffer_size = MIN_BITSTREAM_BUFFER_SIZE + .max(config.dimensions.width as usize * config.dimensions.height as usize); - // 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 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 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_pool = cmd_resources.upload_command_pool; - let upload_command_buffer = cmd_resources.upload_command_buffer; - let upload_fence = cmd_resources.upload_fence; - - let bitstream_buffer_size = MIN_BITSTREAM_BUFFER_SIZE.max(width as usize * height as usize); + // 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)?; - // Create the depth-N encode pipeline (per-frame input images, bitstream - // buffers, encode command buffers, fences and query pools). - let pipeline = EncodePipeline::new(&PipelineConfig { + let init = build_encoder_common(&CommonInitRequest { context: &context, - aligned_width, - aligned_height, - picture_format, - pixel_format: config.pixel_format, - bit_depth: config.bit_depth, - bitstream_buffer_size, + config: &config, profile_info: &profile_info, - command_pool, - upload_command_buffer, - upload_fence, + caps: &caps, + align_unit: SUPERBLOCK_SIZE, + max_active_refs_cap: 8, + bitstream_buffer_size, + // AV1 reference handling here does not use a layered DPB. + allow_layered_dpb: false, })?; + let common = init.common; - // 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, - pipeline, - dpb_images, - dpb_image_memories, - dpb_image_views, - dpb_slot_count: requested_dpb_slots, - dpb_slot_active: vec![false; requested_dpb_slots], - command_pool, - upload_command_pool, - upload_command_buffer, - upload_fence, header_data: None, - current_dpb_slot: 0, references: Vec::new(), }; - 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 0131135..756a43e 100644 --- a/src/encoder/av1/mod.rs +++ b/src/encoder/av1/mod.rs @@ -1,27 +1,23 @@ -//! 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 single-reference tracking ([`Av1`]), its per-frame StdVideo* graph +//! (`record`), and its sequence-header generation (`session_params`). -mod api; -mod encode; mod init; +mod record; mod session_params; -use ash::vk; -use tracing::debug; - -use crate::encoder::pipeline::EncodePipeline; -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)] @@ -31,123 +27,103 @@ pub(crate) struct ReferenceInfo { pub frame_type: u32, } -/// 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 (single-reference prediction). +pub(crate) struct Av1 { frame_num: u32, order_hint: u32, - - /// Depth-N encode pipeline (per-frame slots + submission ordering). - pipeline: EncodePipeline, - /// 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, - // Command resources. - command_pool: vk::CommandPool, - upload_command_pool: vk::CommandPool, - upload_command_buffer: vk::CommandBuffer, - upload_fence: vk::Fence, - - // 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 (only the first is used). references: Vec, } -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<()> { - let slot = self.pipeline.current(); - if src_image == slot.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: slot.input_image, - width: self.config.dimensions.width, - height: self.config.dimensions.height, - pixel_format: self.config.pixel_format, - input_image_layout: slot.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.pipeline.current_mut().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 device = self.context.device(); - let _ = device.queue_wait_idle(self.context.transfer_queue()); - if let Some(q) = self.context.video_encode_queue() { - let _ = device.queue_wait_idle(q); - } + 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; - self.pipeline.destroy(device); + let ref_info = ReferenceInfo { + dpb_slot: common.current_dpb_slot, + order_hint: encoded_order_hint, + 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 + }, + }; - device.destroy_fence(self.upload_fence, None); - device.destroy_command_pool(self.command_pool, None); - if self.upload_command_pool != self.command_pool { - device.destroy_command_pool(self.upload_command_pool, None); + if is_key { + self.references.clear(); + } + self.references.insert(0, ref_info); + // Single-reference prediction keeps only the most recent reference. + self.references.truncate(1); + + // 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; + } - for i in 0..self.dpb_images.len() { - device.destroy_image_view(self.dpb_image_views[i], None); - device.destroy_image(self.dpb_images[i], None); - device.free_memory(self.dpb_image_memories[i], None); - } + fn create_session_params( + &self, + common: &EncoderCommon, + desc: &ColorDescription, + ) -> Result { + self.build_session_params(common, desc) + } - 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 { - device.free_memory(*mem, None); - } - } + fn invalidate_header_cache(&mut self) { + self.header_data = None; } } @@ -162,14 +138,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] @@ -179,11 +153,8 @@ mod tests { order_hint: 42, frame_type: 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); @@ -191,21 +162,17 @@ 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, @@ -217,20 +184,14 @@ mod tests { 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, @@ -242,12 +203,8 @@ mod tests { ); } 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 { @@ -263,14 +220,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 { @@ -278,10 +232,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, @@ -295,9 +248,8 @@ mod tests { 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 { @@ -310,10 +262,8 @@ mod tests { 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 { @@ -325,14 +275,12 @@ mod tests { 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 { @@ -344,90 +292,12 @@ mod tests { 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(); @@ -442,7 +312,6 @@ mod tests { 0 }; - // Frame 0: slot 0. references.insert( 0, ReferenceInfo { @@ -457,7 +326,6 @@ 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 { @@ -472,9 +340,8 @@ 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 { @@ -487,6 +354,6 @@ mod tests { 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..e4eff3d --- /dev/null +++ b/src/encoder/av1/record.rs @@ -0,0 +1,415 @@ +//! 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. + 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(); + + if !is_key_frame && !self.references.is_empty() { + let ref_info = &self.references[0]; + 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(), + }); + av1_reference_infos.push(vk::VideoEncodeAV1DpbSlotInfoKHR::default()); + av1_reference_infos[0] = av1_reference_infos[0].std_reference_info(&ref_std_infos[0]); + + 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 ref_slot = vk::VideoReferenceSlotInfoKHR::default() + .slot_index(ref_info.dpb_slot as i32) + .picture_resource(&ref_picture_resources[0]); + reference_slots.push(ref_slot); + reference_slots[0] = reference_slots[0].push(&mut av1_reference_infos[0]); + } + + // 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 AV1 reference names to DPB slots (SINGLE_REFERENCE uses LAST_FRAME). + let mut reference_name_slot_indices = [-1i32; 7]; + if !is_key_frame && !self.references.is_empty() { + reference_name_slot_indices[0] = self.references[0].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. + fn calculate_reference_frame_mapping( + &self, + is_key_frame: bool, + current_dpb_slot: u8, + ) -> ([i8; 7], [u8; 8], u8, u8) { + if is_key_frame { + // Key frame refreshes all 8 reference slots; all names point to slot 0. + ([0i8; 7], [0u8; 8], 7u8, 0xFFu8) + } else if let Some(last_ref) = self.references.first() { + let mut ref_frame_idx = [0i8; 7]; + let mut ref_order_hint = [0u8; 8]; + ref_frame_idx[0] = last_ref.dpb_slot as i8; + ref_order_hint[0] = last_ref.order_hint as u8; + // Refresh only the current slot so this frame becomes the new LAST_FRAME. + let refresh_flags = 1u8 << current_dpb_slot; + (ref_frame_idx, ref_order_hint, 0u8, refresh_flags) + } else { + ([0i8; 7], [0u8; 8], 7u8, 0x00u8) + } + } +} 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..262addf --- /dev/null +++ b/src/encoder/codec.rs @@ -0,0 +1,683 @@ +//! 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); + + /// 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(); + } + + /// 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 6bafe49..0000000 --- a/src/encoder/h264/api.rs +++ /dev/null @@ -1,304 +0,0 @@ -use super::H264Encoder; - -use crate::encoder::dpb::{DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType}; -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::pipeline::SlotPacketMetadata; -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.pipeline.input_image() - } - - /// Try to receive a finished encoded packet without blocking. - /// - /// Encoding is asynchronous (see [`Self::encode`]); packets are produced by - /// a background readback thread and surfaced here. Returns `None` when no - /// packet is currently ready. - pub fn poll_packet(&self) -> Result> { - self.pipeline.poll_packet().transpose() - } - - /// Take ownership of the encoded-packet receiver for out-of-band consumption. - pub fn take_packet_receiver(&mut self) -> Option { - self.pipeline.take_packet_receiver() - } - - /// 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. - /// - /// Encoding is asynchronous and pipelined (depth 2): this submits the given - /// frame to the encode queue without blocking and returns immediately. The - /// encoded packet is produced by a background readback thread and retrieved - /// via [`Self::poll_packet`] (or [`Self::flush`] at end-of-stream). This call - /// only blocks if all pipeline slots are still in flight (backpressure). - /// - /// # 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<()> { - // Ensure the slot we are about to record over and reuse has been read - // back by the completion thread. - self.pipeline.wait_current_free(); - - 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)?; - - self.encode_current_frame(&gop_position, display_order)?; - - self.pipeline.advance(); - - Ok(()) - } - - /// 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 header = if is_idr { - self.sps_written = true; - Some(self.get_h264_header()?) - } else { - None - }; - - // Record packet metadata before submitting: `submit_current` hands it to - // the readback thread together with the bitstream. `dts` reads the - // pre-increment encode counter (also used by `encode_frame_internal`). - let dts = self.encode_frame_num; - self.pipeline.set_pending_metadata(SlotPacketMetadata { - frame_type, - is_key_frame: is_idr, - pts: display_order, - dts, - header, - }); - - 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(()) - } - - /// Flush the encoder and get any remaining packets. - pub fn flush(&mut self) -> Result> { - self.pipeline.flush() - } - - /// 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 all in-flight encodes to be read back before modifying the - // shared session parameters. - self.pipeline.wait_all_free(); - - // 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 a815609..e1c00e1 100644 --- a/src/encoder/h264/init.rs +++ b/src/encoder/h264/init.rs @@ -1,56 +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::pipeline::{EncodePipeline, PipelineConfig}; -use crate::encoder::resources::{ - align_up, allocate_session_memory, create_command_resources, create_dpb_images, - get_video_format, lcm, query_supported_video_formats, 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 @@ -58,480 +35,129 @@ 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 + .push(&mut encode_caps) + .push(&mut h264_caps); + let caps = query_video_caps(&context, &profile_info, &mut capabilities)?; - 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()) - })?; - - 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); - - // 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 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 upload_fence = cmd_resources.upload_fence; - - // Create the depth-N encode pipeline (per-frame input images, bitstream - // buffers, encode command buffers, fences and query pools). - let pipeline = EncodePipeline::new(&PipelineConfig { + let init = build_encoder_common(&CommonInitRequest { context: &context, - aligned_width, - aligned_height, - picture_format, - pixel_format: config.pixel_format, - bit_depth: config.bit_depth, + config: &config, + profile_info: &profile_info, + caps: &caps, + align_unit: MB_SIZE, + max_active_refs_cap: 32, bitstream_buffer_size: MIN_BITSTREAM_BUFFER_SIZE, - profile_info: &profile_for_resources, - command_pool, - upload_command_buffer, - upload_fence, + allow_layered_dpb: true, })?; + let active_reference_count = init.active_reference_count; + let common = init.common; - // 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, - pipeline, - 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, - command_pool, - upload_command_pool, - upload_command_buffer, - upload_fence, - 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, 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 38c3c0e..7e7a94d 100644 --- a/src/encoder/h264/mod.rs +++ b/src/encoder/h264/mod.rs @@ -1,23 +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::pipeline::EncodePipeline; -use crate::encoder::resources::{upload_image_to_input, 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; @@ -29,159 +28,134 @@ pub(crate) struct ReferenceInfo { pub poc: i32, } -/// 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, - - /// Depth-N encode pipeline (per-frame slots + submission ordering). - pipeline: EncodePipeline, - /// 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, - // Command resources. - command_pool: vk::CommandPool, - upload_command_pool: vk::CommandPool, - upload_command_buffer: vk::CommandBuffer, - upload_fence: vk::Fence, - - // 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). + /// 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<()> { - let slot = self.pipeline.current(); - if src_image == slot.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; } - let params = UploadParams { - upload_command_buffer: self.upload_command_buffer, - upload_fence: self.upload_fence, - src_image, - dst_image: slot.input_image, - width: self.config.dimensions.width, - height: self.config.dimensions.height, - pixel_format: self.config.pixel_format, - input_image_layout: slot.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.pipeline.current_mut().input_image_layout = vk::ImageLayout::VIDEO_ENCODE_SRC_KHR; - - Ok(()) + Ok(PictureSetup { + frame_type: plan.frame_type(), + header, + }) } -} -// 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 { - let device = self.context.device(); - // Wait on the queues used by the encoder rather than stalling - // the entire device. - let _ = device.queue_wait_idle(self.context.transfer_queue()); - if let Some(q) = self.context.video_encode_queue() { - let _ = device.queue_wait_idle(q); - } + fn record_picture( + &mut self, + common: &mut EncoderCommon, + plan: &FramePlan, + ) -> Result { + self.record(common, plan) + } - self.pipeline.destroy(device); + 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(); - device.destroy_fence(self.upload_fence, None); - device.destroy_command_pool(self.command_pool, None); - if self.upload_command_pool != self.command_pool { - device.destroy_command_pool(self.upload_command_pool, None); - } + if plan.is_reference() && !plan.is_b_frame() { + self.frame_num_syntax = (frame_num + 1) % 256; + } - for view in &self.dpb_image_views { - device.destroy_image_view(*view, None); - } - for image in &self.dpb_images { - device.destroy_image(*image, None); - } - for memory in &self.dpb_image_memories { - device.free_memory(*memory, None); + 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, + }, + ); + while self.l0_references.len() > self.active_reference_count as usize { + self.l0_references.pop(); } - if self.session_params != vk::VideoSessionParametersKHR::null() { - (self - .video_queue_fn - .fp() - .destroy_video_session_parameters_khr)( - device.handle(), - self.session_params, - std::ptr::null(), - ); - } - (self.video_queue_fn.fp().destroy_video_session_khr)( - device.handle(), - self.session, - std::ptr::null(), - ); - for memory in &self.session_memory { - device.free_memory(*memory, None); + // 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 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 63% rename from src/encoder/h264/encode.rs rename to src/encoder/h264/record.rs index 0e910bc..f65f6cb 100644 --- a/src/encoder/h264/encode.rs +++ b/src/encoder/h264/record.rs @@ -1,81 +1,63 @@ -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 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 slot = self.pipeline.current(); + 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; - let is_b_frame = gop_position.frame_type == GopFrameType::B; - let is_reference = gop_position.is_reference; - 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); - // Prepare command buffer for recording. + // Prepare command buffer and transition DPB images for encode. unsafe { - prepare_encode_command_buffer(self.context.device(), command_buffer, 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(), + common.device(), command_buffer, - &self.dpb_images, - self.use_layered_dpb, - self.current_dpb_slot, + &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 +89,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 +110,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( @@ -147,16 +128,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 +138,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 +147,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 { @@ -240,43 +200,31 @@ impl H264Encoder { }, }; - // 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(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 +234,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 +258,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 +279,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 +319,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 +334,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,7 +358,6 @@ impl H264Encoder { encode_ref_slots.push(l1); } } else if !is_idr { - // P-Frames: Use all L0 encode_ref_slots.extend_from_slice(&l0_slots); } @@ -437,14 +367,12 @@ impl H264Encoder { .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 +385,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,60 +404,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); - - let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default().rate_control_mode(rc_mode); + .gop_frame_count(common.config.gop_size) + .idr_period(common.config.gop_size) + .consecutive_b_frame_count(common.config.b_frame_count); - 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)(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 @@ -554,57 +461,33 @@ 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)( + (common.video_queue_fn.fp().cmd_control_video_coding_khr)( command_buffer, &control_info, ); } } - // Begin query. unsafe { - self.context.device().cmd_begin_query( + let device = common.device(); + device.cmd_begin_query( command_buffer, query_pool, 0, vk::QueryControlFlags::empty(), ); - } - - // Encode - unsafe { - (self.video_encode_fn.fp().cmd_encode_video_khr)(command_buffer, &encode_info); - } + (common.video_encode_fn.fp().cmd_encode_video_khr)(command_buffer, &encode_info); + device.cmd_end_query(command_buffer, query_pool, 0); - // End query. - unsafe { - self.context - .device() - .cmd_end_query(command_buffer, query_pool, 0); - } + let end_info = vk::VideoEndCodingInfoKHR::default(); + (common.video_queue_fn.fp().cmd_end_video_coding_khr)(command_buffer, &end_info); - // End video coding. - let end_info = vk::VideoEndCodingInfoKHR::default(); - unsafe { - (self.video_queue_fn.fp().cmd_end_video_coding_khr)(command_buffer, &end_info); + device + .end_command_buffer(command_buffer) + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; } - // End command buffer. - unsafe { self.context.device().end_command_buffer(command_buffer) } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; - - // Submit without waiting; the slot is drained on a later encode call. - let encode_queue = self.context.video_encode_queue().ok_or_else(|| { - PixelForgeError::NoSuitableDevice("No video encode queue available".to_string()) - })?; - self.pipeline - .submit_current(self.context.device(), encode_queue)?; - - // Mark DPB slot as active. - self.dpb_slot_active[self.current_dpb_slot as usize] = true; - - Ok(()) + common.submit_frame() } } 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 01be55d..0000000 --- a/src/encoder/h265/api.rs +++ /dev/null @@ -1,333 +0,0 @@ -use super::H265Encoder; - -use crate::encoder::dpb::{DecodedPictureBufferTrait, DpbConfig, PictureStartInfo, PictureType}; -use crate::encoder::gop::{GopFrameType, GopPosition}; -use crate::encoder::pipeline::SlotPacketMetadata; -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.pipeline.input_image() - } - - /// Try to receive a finished encoded packet without blocking. - /// - /// Encoding is asynchronous (see [`Self::encode`]); packets are produced by - /// a background readback thread and surfaced here. Returns `None` when no - /// packet is currently ready. - pub fn poll_packet(&self) -> Result> { - self.pipeline.poll_packet().transpose() - } - - /// Take ownership of the encoded-packet receiver for out-of-band consumption. - pub fn take_packet_receiver(&mut self) -> Option { - self.pipeline.take_packet_receiver() - } - - /// 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. - /// - /// Encoding is asynchronous and pipelined (depth 2): this submits the given - /// frame to the encode queue without blocking and returns immediately. The - /// encoded packet is produced by a background readback thread and retrieved - /// via [`Self::poll_packet`] (or [`Self::flush`] at end-of-stream). This call - /// only blocks if all pipeline slots are still in flight (backpressure). - /// - /// # 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<()> { - // Ensure the slot we are about to record over and reuse has been read - // back by the completion thread. - self.pipeline.wait_current_free(); - - 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)?; - - self.encode_current_frame(&gop_position, display_order)?; - - self.pipeline.advance(); - - Ok(()) - } - - /// 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; - - // For IDR frames, prepend VPS/SPS/PPS header. - let 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); - } - self.header_data.clone() - } else { - None - }; - - // Record packet metadata before submitting: `submit_current` hands it to - // the readback thread together with the bitstream. `dts` reads the - // pre-increment encode counter (also used by `encode_frame_internal`). - let dts = self.encode_frame_num; - self.pipeline.set_pending_metadata(SlotPacketMetadata { - frame_type, - is_key_frame: is_idr, - pts: display_order, - dts, - header, - }); - - self.encode_frame_internal(gop_position, pic_order_cnt, is_idr)?; - - 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(()) - } - - /// Flush the encoder and get any remaining packets. - pub fn flush(&mut self) -> Result> { - self.pipeline.flush() - } - - /// 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 all in-flight encodes to be read back before modifying the - // shared session parameters. - self.pipeline.wait_all_free(); - - // 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 0b599a7..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::pipeline::{EncodePipeline, PipelineConfig}; -use crate::encoder::resources::{ - align_up, allocate_session_memory, create_command_resources, create_dpb_images, - get_video_format, lcm, make_codec_name, query_supported_video_formats, - 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,366 +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 - }; + .push(&mut encode_caps) + .push(&mut h265_caps); + let caps = query_video_caps(&context, &profile_info, &mut capabilities)?; - 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()) - })?; - - 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); - - // 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 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 upload_fence = cmd_resources.upload_fence; - - // Create the depth-N encode pipeline (per-frame input images, bitstream - // buffers, encode command buffers, fences and query pools). - let pipeline = EncodePipeline::new(&PipelineConfig { + let init = build_encoder_common(&CommonInitRequest { context: &context, - aligned_width, - aligned_height, - picture_format, - pixel_format: config.pixel_format, - bit_depth: config.bit_depth, + config: &config, + profile_info: &profile_info, + caps: &caps, + align_unit: CTB_SIZE, + max_active_refs_cap: 15, bitstream_buffer_size: MIN_BITSTREAM_BUFFER_SIZE, - profile_info: &profile_for_resources, - command_pool, - upload_command_buffer, - upload_fence, + allow_layered_dpb: true, })?; + let active_reference_count = init.active_reference_count; + let common = init.common; - // 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, - pipeline, - dpb_images, - dpb_image_memories, - dpb_image_views, - dpb_slot_count, - use_layered_dpb, - command_pool, - upload_command_pool, - upload_command_buffer, - upload_fence, 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 b133cb6..d8cf103 100644 --- a/src/encoder/h265/mod.rs +++ b/src/encoder/h265/mod.rs @@ -1,22 +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::pipeline::EncodePipeline; -use crate::encoder::resources::{upload_image_to_input, UploadParams}; -use crate::encoder::EncodeConfig; +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::vulkan::VideoContext; +use ash::vk; /// H.265 Coding Tree Block (CTB) size in pixels. pub const CTB_SIZE: u32 = 32; @@ -27,151 +26,120 @@ pub(crate) struct ReferenceInfo { pub poc: i32, } -/// 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, - - /// Depth-N encode pipeline (per-frame slots + submission ordering). - pipeline: EncodePipeline, - /// 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, - // Command resources. - command_pool: vk::CommandPool, - upload_command_pool: vk::CommandPool, - upload_command_buffer: vk::CommandBuffer, - upload_fence: vk::Fence, - - // 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<()> { - let slot = self.pipeline.current(); - if src_image == slot.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: slot.input_image, - width: self.config.dimensions.width, - height: self.config.dimensions.height, - pixel_format: self.config.pixel_format, - input_image_layout: slot.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.pipeline.current_mut().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, + }, + ); + while self.l0_references.len() > self.active_reference_count as usize { + self.l0_references.pop(); + } -impl Drop for H265Encoder { - fn drop(&mut self) { - unsafe { - let device = self.context.device(); - // Wait on the queues used by the encoder rather than stalling - // the entire device. - let _ = device.queue_wait_idle(self.context.transfer_queue()); - if let Some(q) = self.context.video_encode_queue() { - let _ = device.queue_wait_idle(q); - } - self.pipeline.destroy(device); - device.destroy_fence(self.upload_fence, None); - device.destroy_command_pool(self.command_pool, None); - if self.upload_command_pool != self.command_pool { - device.destroy_command_pool(self.upload_command_pool, None); - } - for view in &self.dpb_image_views { - device.destroy_image_view(*view, None); - } - for image in &self.dpb_images { - device.destroy_image(*image, None); - } - for memory in &self.dpb_image_memories { - device.free_memory(*memory, None); - } - if self.session_params != vk::VideoSessionParametersKHR::null() { - (self - .video_queue_fn - .fp() - .destroy_video_session_parameters_khr)( - device.handle(), - self.session_params, - std::ptr::null(), - ); - } - (self.video_queue_fn.fp().destroy_video_session_khr)( - device.handle(), - self.session, - std::ptr::null(), - ); - for memory in &self.session_memory { - device.free_memory(*memory, None); + 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; + } } } } + + 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 59% rename from src/encoder/h265/encode.rs rename to src/encoder/h265/record.rs index d0cb67c..8fc350c 100644 --- a/src/encoder/h265/encode.rs +++ b/src/encoder/h265/record.rs @@ -1,10 +1,10 @@ -//! 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, }; @@ -13,49 +13,40 @@ 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<()> { - let slot = self.pipeline.current(); + 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; - // Prepare command buffer for recording. unsafe { - prepare_encode_command_buffer(self.context.device(), command_buffer, 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(), + common.device(), command_buffer, - &self.dpb_images, - self.use_layered_dpb, - self.current_dpb_slot, + &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 { @@ -63,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 { @@ -72,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( @@ -91,7 +80,6 @@ impl H265Encoder { 0, // reserved ), }; - let slice_header = ash::vk::native::StdVideoEncodeH265SliceSegmentHeader { flags: slice_header_flags, slice_type, @@ -110,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, @@ -185,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, @@ -209,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( @@ -226,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; @@ -241,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 { @@ -305,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(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, @@ -358,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 = @@ -411,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 = @@ -433,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() { @@ -454,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)( + (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 @@ -587,16 +474,14 @@ 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)( + (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) @@ -607,61 +492,41 @@ impl H265Encoder { .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( + 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)(command_buffer, &encode_info); - - self.context - .device() - .cmd_end_query(command_buffer, query_pool, 0); - } - - // Add DPB synchronization barrier after encoding. - unsafe { record_post_encode_dpb_barrier( - self.context.device(), + device, command_buffer, - &self.dpb_images, - self.use_layered_dpb, - self.current_dpb_slot, + &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)(command_buffer, &end_coding_info); - } - - // End command buffer. - unsafe { self.context.device().end_command_buffer(command_buffer) } - .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + let end_coding_info = vk::VideoEndCodingInfoKHR::default(); + (common.video_queue_fn.fp().cmd_end_video_coding_khr)(command_buffer, &end_coding_info); - // Submit without waiting; the slot is drained on a later encode call. - 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 - ); - - self.pipeline - .submit_current(self.context.device(), encode_queue)?; - - // Mark DPB slot as active. - self.dpb_slot_active[self.current_dpb_slot as usize] = true; + device + .end_command_buffer(command_buffer) + .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; + } - Ok(()) + 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 d00b342..3d6891d 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -7,6 +7,7 @@ pub mod av1; pub mod bitwriter; +pub(crate) mod codec; pub mod dpb; pub mod gop; pub mod h264; @@ -386,10 +387,7 @@ impl EncodeConfig { } } -/// Channel receiver for encoded packets, produced by the encoder's background -/// readback thread. Obtained via [`Encoder::take_packet_receiver`] for -/// out-of-band (push) consumption from a dedicated thread. -pub type EncodedPacketReceiver = std::sync::mpsc::Receiver>; +pub use pipeline::EncodeFuture; /// Encoded video packet. #[derive(Debug, Clone)] @@ -406,45 +404,60 @@ 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 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 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. @@ -453,10 +466,9 @@ impl Encoder { /// The source image must match the format and dimensions in the encoder configuration. /// /// Encoding is asynchronous: this submits the frame without blocking and - /// returns immediately. Encoded packets are produced by a background readback - /// thread and retrieved via [`Encoder::poll_packet`] (or [`Encoder::flush`] - /// at end-of-stream). The call only blocks if every pipeline slot is still in - /// flight (backpressure). + /// 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 @@ -478,65 +490,29 @@ impl Encoder { /// # let yuv_data = vec![0u8; 1920 * 1080 * 3 / 2]; /// input.upload_yuv420(&yuv_data)?; /// - /// // Submit the frame, then drain any ready packets. - /// encoder.encode(input.image())?; - /// while let Some(packet) = encoder.poll_packet()? { - /// // ... use packet.data ... - /// } + /// // 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) } - /// Try to receive a finished encoded packet without blocking. + /// Wait for all in-flight frames to finish encoding (end-of-stream barrier). /// - /// Returns `Ok(None)` when no packet is currently ready. Encoded packets are - /// produced asynchronously after [`Encoder::encode`]; poll until `None` to - /// drain everything currently available. - pub fn poll_packet(&self) -> Result> { - match self { - Encoder::H264(encoder) => encoder.poll_packet(), - Encoder::H265(encoder) => encoder.poll_packet(), - Encoder::AV1(encoder) => encoder.poll_packet(), - } - } - - /// Take ownership of the encoded-packet receiver for out-of-band (push) - /// consumption from a dedicated thread. - /// - /// After this is called, [`Encoder::poll_packet`] and [`Encoder::flush`] no - /// longer yield packets — they are delivered to the returned receiver - /// instead. Returns `None` if the receiver has already been taken. - pub fn take_packet_receiver(&mut self) -> Option { - match self { - Encoder::H264(encoder) => encoder.take_packet_receiver(), - Encoder::H265(encoder) => encoder.take_packet_receiver(), - Encoder::AV1(encoder) => encoder.take_packet_receiver(), - } - } - - /// 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(), - } + /// 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() } /// Update the color description (VUI parameters) for the encoder. @@ -545,11 +521,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 index dad3c19..80e590d 100644 --- a/src/encoder/pipeline.rs +++ b/src/encoder/pipeline.rs @@ -8,10 +8,12 @@ //! //! Bitstream readback is performed off the calling thread: a single background //! *completion thread* waits on each submission's fence, copies the bitstream -//! out, and pushes the finished [`EncodedPacket`] onto a channel 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. +//! 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 @@ -26,21 +28,59 @@ //! 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, EncodedPacketReceiver, FrameType, PixelFormat}; +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 @@ -77,6 +117,8 @@ struct WorkItem { 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 @@ -180,9 +222,6 @@ pub(crate) struct EncodePipeline { /// Sends submitted work to the completion thread. Dropped on shutdown to end /// the thread. work_tx: Option>, - /// Receives finished packets from the completion thread. `None` once the - /// caller has taken ownership via [`EncodePipeline::take_packet_receiver`]. - packet_rx: Option, /// The completion thread handle, joined on shutdown. completion_thread: Option>, } @@ -269,7 +308,6 @@ impl EncodePipeline { let slot_sync = Arc::new(SlotSync::new(slots.len())); let (work_tx, work_rx) = std::sync::mpsc::channel::(); - let (packet_tx, packet_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 @@ -279,7 +317,7 @@ impl EncodePipeline { let completion_thread = std::thread::Builder::new() .name("pixelforge-encode-readback".to_string()) .spawn(move || { - run_completion_thread(thread_device, work_rx, packet_tx, thread_sync); + run_completion_thread(thread_device, work_rx, thread_sync); }) .map_err(|e| PixelForgeError::CommandBuffer(format!("spawn readback thread: {e}")))?; @@ -291,7 +329,6 @@ impl EncodePipeline { last_value: 0, slot_sync, work_tx: Some(work_tx), - packet_rx: Some(packet_rx), completion_thread: Some(completion_thread), }) } @@ -334,12 +371,13 @@ impl EncodePipeline { /// 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. + /// 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<()> { + ) -> 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; @@ -380,12 +418,14 @@ impl EncodePipeline { // 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 @@ -393,7 +433,7 @@ impl EncodePipeline { let _ = tx.send(work); } - Ok(()) + Ok(EncodeFuture { rx: result_rx }) } /// Advance to the next slot after a frame has been submitted. @@ -401,39 +441,14 @@ impl EncodePipeline { self.current_slot = (self.current_slot + 1) % self.slots.len(); } - /// Try to receive a finished packet without blocking. + /// Barrier: wait for every in-flight frame to be read back. /// - /// Returns `None` if no packet is ready, or if the caller has taken - /// ownership of the receiver via [`EncodePipeline::take_packet_receiver`]. - pub(crate) fn poll_packet(&self) -> Option> { - self.packet_rx.as_ref().and_then(|rx| rx.try_recv().ok()) - } - - /// Take ownership of the packet receiver for out-of-band consumption (e.g. - /// a dedicated sender thread). After this, [`EncodePipeline::poll_packet`] - /// and [`EncodePipeline::flush`] no longer return packets. - pub(crate) fn take_packet_receiver(&mut self) -> Option { - self.packet_rx.take() - } - - /// Wait for all in-flight frames to be read back, returning any packets not - /// yet polled. Used to drain at end of stream. - /// - /// If the receiver has been taken, this only acts as a completion barrier - /// and returns an empty `Vec` (packets are delivered to the taken receiver). - pub(crate) fn flush(&mut self) -> Result> { - // Once every slot is free, the completion thread has already sent every - // packet (it sends before marking a slot free), so draining now is - // complete. + /// 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(); - - let mut packets = Vec::new(); - if let Some(rx) = &self.packet_rx { - while let Ok(result) = rx.try_recv() { - packets.push(result?); - } - } - Ok(packets) } /// Stop the completion thread and wait for it to finish any in-flight @@ -474,34 +489,39 @@ impl EncodePipeline { } /// Completion-thread body: wait on each submission's fence, copy its bitstream -/// out, deliver the packet, then mark the slot free. +/// out, resolve the frame's future, then mark the slot free. fn run_completion_thread( device: ash::Device, work_rx: Receiver, - packet_tx: Sender>, 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) + wait_and_read_bitstream( + &device, + work.fence, + work.query_pool, + work.bitstream_ptr.0, + &mut data, + ) }; - let packet = result.map(|bitstream| { - let mut data = work.metadata.header.unwrap_or_default(); - data.extend_from_slice(&bitstream); - EncodedPacket { - data, - frame_type: work.metadata.frame_type, - is_key_frame: work.metadata.is_key_frame, - pts: work.metadata.pts, - dts: work.metadata.dts, - } + 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, }); - // Deliver the packet *before* freeing the slot. This ordering means that - // once all slots are observed free, every packet has already been sent, - // which `flush` relies on for completeness. - let _ = packet_tx.send(packet); + // 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 3c9a05d..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 @@ -1408,12 +1394,17 @@ pub(crate) unsafe fn submit_encode_only( /// # 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, -) -> Result> { + dst: &mut Vec, +) -> Result<()> { device .wait_for_fences(&[fence], true, u64::MAX) .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; @@ -1450,11 +1441,122 @@ pub(crate) unsafe fn wait_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], +} - Ok(encoded_data) +/// 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); + } + + 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 6cf5c0b..77c6432 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -197,10 +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, - EncodedPacketReceiver, 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; From f6b94233cb176f968d4e9f3b3fe602cbbf4ad713 Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Fri, 19 Jun 2026 16:41:17 +0800 Subject: [PATCH 6/9] feat: reference frame invalidation (RFI) for H.264/H.265/AV1; AV1 multi-reference prediction --- examples/rfi.rs | 223 +++++++++++++++++++++++++++++++++++++ src/encoder/av1/init.rs | 2 + src/encoder/av1/mod.rs | 50 ++++++++- src/encoder/av1/record.rs | 103 ++++++++++++----- src/encoder/codec.rs | 28 +++++ src/encoder/h264/init.rs | 1 + src/encoder/h264/mod.rs | 35 ++++++ src/encoder/h264/record.rs | 51 ++++++++- src/encoder/h265/mod.rs | 21 ++++ src/encoder/mod.rs | 25 +++++ 10 files changed, 502 insertions(+), 37 deletions(-) create mode 100644 examples/rfi.rs diff --git a/examples/rfi.rs b/examples/rfi.rs new file mode 100644 index 0000000..7bf301d --- /dev/null +++ b/examples/rfi.rs @@ -0,0 +1,223 @@ +//! 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/src/encoder/av1/init.rs b/src/encoder/av1/init.rs index 0e1bea7..0061761 100644 --- a/src/encoder/av1/init.rs +++ b/src/encoder/av1/init.rs @@ -64,6 +64,7 @@ impl Av1 { // 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; let codec = Av1 { @@ -71,6 +72,7 @@ impl Av1 { order_hint: 0, header_data: None, references: Vec::new(), + active_reference_count, }; let mut encoder = CodecEncoder { common, codec }; diff --git a/src/encoder/av1/mod.rs b/src/encoder/av1/mod.rs index 756a43e..b6ba640 100644 --- a/src/encoder/av1/mod.rs +++ b/src/encoder/av1/mod.rs @@ -1,8 +1,14 @@ //! AV1 codec: the differences from the generic encoder. //! //! Shared machinery lives in [`crate::encoder::codec`]; this folder holds only -//! AV1's single-reference tracking ([`Av1`]), its per-frame StdVideo* graph -//! (`record`), and its sequence-header generation (`session_params`). +//! 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 init; mod record; @@ -25,16 +31,20 @@ 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-specific encoder state (single-reference prediction). +/// AV1-specific encoder state (multi-reference prediction). pub(crate) struct Av1 { frame_num: u32, order_hint: u32, /// Cached sequence-header OBU (invalidated when session params change). header_data: Option>, - /// Active references, most-recent first (only the first is used). + /// Active references, most-recent first. references: Vec, + /// Negotiated number of active references kept for prediction. + active_reference_count: u32, } impl VideoCodec for Av1 { @@ -86,6 +96,7 @@ impl VideoCodec for Av1 { 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 { @@ -97,8 +108,11 @@ impl VideoCodec for Av1 { self.references.clear(); } self.references.insert(0, ref_info); - // Single-reference prediction keeps only the most recent reference. - self.references.truncate(1); + // 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(); @@ -114,6 +128,19 @@ impl VideoCodec for Av1 { common.current_dpb_slot = next_slot; } + fn invalidate_references( + &mut self, + _common: &mut EncoderCommon, + first_lost_display_order: u64, + ) -> bool { + // This AV1 path predicts from a single reference. If that reference is at + // or after the first lost frame it is undecodable on the client and no + // older survivor is kept, so recovery can only be 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, @@ -152,6 +179,7 @@ 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); @@ -178,6 +206,7 @@ mod tests { 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 { @@ -199,6 +228,7 @@ mod tests { dpb_slot: i, order_hint: i as u32, frame_type: 0, + display_order: 0, }, ); } @@ -211,6 +241,7 @@ mod tests { dpb_slot: 0, order_hint: 0, frame_type: 0, + display_order: 0, }, ); assert_eq!(references.len(), 1); @@ -242,6 +273,7 @@ mod tests { dpb_slot: 0, order_hint: 0, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -256,6 +288,7 @@ mod tests { dpb_slot: 1, order_hint: 1, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -270,6 +303,7 @@ mod tests { dpb_slot: 2, order_hint: 2, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -287,6 +321,7 @@ mod tests { dpb_slot: 0, order_hint: 3, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -318,6 +353,7 @@ mod tests { dpb_slot: 0, order_hint: 0, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -332,6 +368,7 @@ mod tests { dpb_slot: 1, order_hint: 1, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { @@ -348,6 +385,7 @@ mod tests { dpb_slot: 0, order_hint: 2, frame_type: 0, + display_order: 0, }, ); while references.len() > max_refs { diff --git a/src/encoder/av1/record.rs b/src/encoder/av1/record.rs index e4eff3d..4b7b8b6 100644 --- a/src/encoder/av1/record.rs +++ b/src/encoder/av1/record.rs @@ -112,14 +112,23 @@ impl Av1 { .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(); + // 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); + } + } + } - if !is_key_frame && !self.references.is_empty() { - let ref_info = &self.references[0]; + 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, @@ -128,9 +137,10 @@ impl Av1 { reserved1: [0; 3], pExtensionHeader: std::ptr::null(), }); - av1_reference_infos.push(vk::VideoEncodeAV1DpbSlotInfoKHR::default()); - av1_reference_infos[0] = av1_reference_infos[0].std_reference_info(&ref_std_infos[0]); + } + 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 }) @@ -138,12 +148,26 @@ impl Av1 { .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 ref_slot = vk::VideoReferenceSlotInfoKHR::default() - .slot_index(ref_info.dpb_slot as i32) - .picture_resource(&ref_picture_resources[0]); - reference_slots.push(ref_slot); - reference_slots[0] = reference_slots[0].push(&mut av1_reference_infos[0]); + 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). @@ -218,10 +242,17 @@ impl Av1 { pBufferRemovalTimes: std::ptr::null(), }; - // Map AV1 reference names to DPB slots (SINGLE_REFERENCE uses LAST_FRAME). + // 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() { - reference_name_slot_indices[0] = self.references[0].dpb_slot as i32; + 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 { @@ -392,24 +423,42 @@ impl Av1 { /// 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. - ([0i8; 7], [0u8; 8], 7u8, 0xFFu8) - } else if let Some(last_ref) = self.references.first() { - let mut ref_frame_idx = [0i8; 7]; - let mut ref_order_hint = [0u8; 8]; - ref_frame_idx[0] = last_ref.dpb_slot as i8; - ref_order_hint[0] = last_ref.order_hint as u8; - // Refresh only the current slot so this frame becomes the new LAST_FRAME. - let refresh_flags = 1u8 << current_dpb_slot; - (ref_frame_idx, ref_order_hint, 0u8, refresh_flags) - } else { - ([0i8; 7], [0u8; 8], 7u8, 0x00u8) + 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/codec.rs b/src/encoder/codec.rs index 262addf..2e651bd 100644 --- a/src/encoder/codec.rs +++ b/src/encoder/codec.rs @@ -262,6 +262,22 @@ pub(crate) trait VideoCodec: Sized + Send { /// 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( @@ -328,6 +344,18 @@ impl CodecEncoder { 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<()> { diff --git a/src/encoder/h264/init.rs b/src/encoder/h264/init.rs index e1c00e1..2527704 100644 --- a/src/encoder/h264/init.rs +++ b/src/encoder/h264/init.rs @@ -95,6 +95,7 @@ impl H264 { backward_reference_dpb_slot: 2, l0_references: Vec::new(), active_reference_count, + pending_unmark_frame_nums: Vec::new(), profile_idc, preferred_entropy_cabac, }; diff --git a/src/encoder/h264/mod.rs b/src/encoder/h264/mod.rs index 7e7a94d..c7ea722 100644 --- a/src/encoder/h264/mod.rs +++ b/src/encoder/h264/mod.rs @@ -26,6 +26,8 @@ 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-specific encoder state (everything the generic encoder doesn't own). @@ -45,6 +47,10 @@ pub(crate) struct H264 { l0_references: Vec, /// Negotiated number of active references. active_reference_count: u32, + /// `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 (from quality-level query). @@ -72,6 +78,9 @@ impl VideoCodec for H264 { 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 header = if plan.is_idr() { @@ -130,6 +139,7 @@ impl VideoCodec for H264 { 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 { @@ -147,6 +157,31 @@ impl VideoCodec for H264 { } } + 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, diff --git a/src/encoder/h264/record.rs b/src/encoder/h264/record.rs index f65f6cb..544688b 100644 --- a/src/encoder/h264/record.rs +++ b/src/encoder/h264/record.rs @@ -41,6 +41,45 @@ impl H264 { 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(); + if use_adaptive_marking { + self.pending_unmark_frame_nums.clear(); + } + // Prepare command buffer and transition DPB images for encode. unsafe { prepare_encode_command_buffer(common.device(), command_buffer, query_pool)?; @@ -117,7 +156,7 @@ impl H264 { 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 ), }; @@ -176,11 +215,15 @@ impl H264 { 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 { @@ -193,7 +236,7 @@ impl H264 { 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() diff --git a/src/encoder/h265/mod.rs b/src/encoder/h265/mod.rs index d8cf103..19127b2 100644 --- a/src/encoder/h265/mod.rs +++ b/src/encoder/h265/mod.rs @@ -24,6 +24,8 @@ 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-specific encoder state. @@ -114,6 +116,7 @@ impl VideoCodec for H265 { 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 { @@ -131,6 +134,24 @@ impl VideoCodec for H265 { } } + 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, diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs index 3d6891d..f776682 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -413,6 +413,7 @@ trait EncoderApi: Send { 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<()>; } @@ -429,6 +430,9 @@ impl EncoderApi for codec::CodecEncoder { 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) } @@ -515,6 +519,27 @@ impl Encoder { 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. Codecs that predict from + /// a single reference (AV1 here) likewise fall back to an IDR. + 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. /// /// This recreates the video session parameters with an updated SPS/VPS/sequence From 9261a3af12c78ef9b14536bc5fe9631adacdd85b Mon Sep 17 00:00:00 2001 From: urwrstkn8mare <42769125+urwrstkn8mare@users.noreply.github.com> Date: Fri, 19 Jun 2026 23:45:44 +0800 Subject: [PATCH 7/9] chore: format --- examples/rfi.rs | 72 ++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 60 insertions(+), 12 deletions(-) diff --git a/examples/rfi.rs b/examples/rfi.rs index 7bf301d..6d9d36c 100644 --- a/examples/rfi.rs +++ b/examples/rfi.rs @@ -95,8 +95,14 @@ fn run_codec( .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 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)?; @@ -107,8 +113,8 @@ fn run_codec( let mut recovery_is_key: Option = None; let drain_one = |pending: &mut VecDeque, - output_file: &mut File, - recovery_is_key: &mut Option| + 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 { @@ -180,20 +186,54 @@ fn decode_and_psnr( ) -> Result> { let status = Command::new("ffmpeg") .args([ - "-hide_banner", "-loglevel", "error", "-y", "-i", bitstream, "-pix_fmt", "yuv420p", - "-f", "rawvideo", decoded, + "-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()); + 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", "-", + "-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); @@ -212,8 +252,16 @@ fn ensure_test_data(pix_fmt: &str, path: &str) -> Result<(), Box Date: Sat, 20 Jun 2026 23:51:25 +0800 Subject: [PATCH 8/9] nit: minor fixes/tweaks --- src/encoder/av1/init.rs | 7 ++++++- src/encoder/av1/mod.rs | 8 +++++--- src/encoder/h264/record.rs | 13 +++++++++---- src/encoder/mod.rs | 6 ++++-- 4 files changed, 24 insertions(+), 10 deletions(-) diff --git a/src/encoder/av1/init.rs b/src/encoder/av1/init.rs index 0061761..d66eb8d 100644 --- a/src/encoder/av1/init.rs +++ b/src/encoder/av1/init.rs @@ -59,7 +59,12 @@ impl Av1 { profile_info: &profile_info, caps: &caps, align_unit: SUPERBLOCK_SIZE, - max_active_refs_cap: 8, + // 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, diff --git a/src/encoder/av1/mod.rs b/src/encoder/av1/mod.rs index b6ba640..28ad1ce 100644 --- a/src/encoder/av1/mod.rs +++ b/src/encoder/av1/mod.rs @@ -133,9 +133,11 @@ impl VideoCodec for Av1 { _common: &mut EncoderCommon, first_lost_display_order: u64, ) -> bool { - // This AV1 path predicts from a single reference. If that reference is at - // or after the first lost frame it is undecodable on the client and no - // older survivor is kept, so recovery can only be a full key frame. + // 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() diff --git a/src/encoder/h264/record.rs b/src/encoder/h264/record.rs index 544688b..4eef7e9 100644 --- a/src/encoder/h264/record.rs +++ b/src/encoder/h264/record.rs @@ -76,9 +76,6 @@ impl H264 { Vec::new() }; let use_adaptive_marking = !ref_pic_marking_ops.is_empty(); - if use_adaptive_marking { - self.pending_unmark_frame_nums.clear(); - } // Prepare command buffer and transition DPB images for encode. unsafe { @@ -531,6 +528,14 @@ impl H264 { .map_err(|e| PixelForgeError::CommandBuffer(e.to_string()))?; } - common.submit_frame() + 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(); + } + Ok(future) } } diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs index f776682..84cc1d7 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -534,8 +534,10 @@ impl Encoder { /// /// 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. Codecs that predict from - /// a single reference (AV1 here) likewise fall back to an IDR. + /// 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) } From e8b1e66721cb4de9458c62c317d432af374f292e Mon Sep 17 00:00:00 2001 From: Hans Gaiser Date: Fri, 26 Jun 2026 10:29:00 +0200 Subject: [PATCH 9/9] docs: fix broken links --- src/encoder/av1/mod.rs | 6 +++--- src/encoder/h264/mod.rs | 4 ++-- src/encoder/h265/mod.rs | 4 ++-- src/encoder/mod.rs | 2 +- src/encoder/pipeline.rs | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/encoder/av1/mod.rs b/src/encoder/av1/mod.rs index 28ad1ce..5dacdd7 100644 --- a/src/encoder/av1/mod.rs +++ b/src/encoder/av1/mod.rs @@ -1,10 +1,10 @@ //! AV1 codec: the differences from the generic encoder. //! -//! Shared machinery lives in [`crate::encoder::codec`]; this folder holds only -//! AV1's reference tracking ([`Av1`]), its per-frame StdVideo* graph (`record`), +//! 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 +//! 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 diff --git a/src/encoder/h264/mod.rs b/src/encoder/h264/mod.rs index c7ea722..528892d 100644 --- a/src/encoder/h264/mod.rs +++ b/src/encoder/h264/mod.rs @@ -1,8 +1,8 @@ //! H.264/AVC codec: the differences from the generic encoder. //! -//! The shared session/DPB/pipeline machinery lives in [`crate::encoder::codec`]; +//! 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 +//! tracking and syntax counters (`H264`), the per-frame StdVideo* graph //! (`record`), and the SPS/PPS generation (`session_params`). mod init; diff --git a/src/encoder/h265/mod.rs b/src/encoder/h265/mod.rs index 19127b2..67dd3bb 100644 --- a/src/encoder/h265/mod.rs +++ b/src/encoder/h265/mod.rs @@ -1,7 +1,7 @@ //! H.265/HEVC codec: the differences from the generic encoder. //! -//! Shared machinery lives in [`crate::encoder::codec`]; this folder holds only -//! H.265's reference tracking ([`H265`]), its per-frame StdVideo* graph +//! 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 init; diff --git a/src/encoder/mod.rs b/src/encoder/mod.rs index 84cc1d7..fbb1a02 100644 --- a/src/encoder/mod.rs +++ b/src/encoder/mod.rs @@ -441,7 +441,7 @@ impl EncoderApi for codec::CodecEncoder { /// 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 +/// boxes the corresponding `codec::CodecEncoder`. All codecs share one generic /// implementation; this type just erases which one is in use. pub struct Encoder(Box); diff --git a/src/encoder/pipeline.rs b/src/encoder/pipeline.rs index 80e590d..3bd5e47 100644 --- a/src/encoder/pipeline.rs +++ b/src/encoder/pipeline.rs @@ -47,7 +47,7 @@ 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 +/// A handle to the packet a single `EncodePipeline::submit_current` will /// eventually produce. /// /// Returned by `Encoder::encode`, one per submitted frame. Awaiting it yields