Merge pull request #18 from DatCaptainHorse/feat/cleaner-pnext

feat: Cleaner p_next handling

Author: hgaiser

README.md

@@ -24,14 +24,9 @@ A Vulkan-based video encoding library for Rust, supporting H.264, H.265, and AV1
 
 | Codec | Encode |
 |-------|--------|
-| H.264/AVC | ✓ |
+| H.264/AVC  | ✓ |
 | H.265/HEVC | ✓ |
-| AV1 | ✓ (experimental) |
-
-> ⚠️ **AV1 Warning**: AV1 encoding is experimental. On NVIDIA GPUs, P-frames cannot
-> reference other P-frames, causing all P-frames to reference the I-frame instead. This
-> leads to progressively larger frame sizes over time. Consider using H.264 or HEVC
-> until this is resolved.
+| AV1        | ✓ |
 
 ## Requirements
 

examples/query_capabilities.rs

@@ -4,6 +4,7 @@
 //! from the Vulkan video extensions.
 
 use ash::vk;
+use ash::vk::TaggedStructure;
 use pixelforge::{Codec, VideoContextBuilder};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer};
 
@@ -37,7 +38,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("Codec Support:");
     println!("--------------");
 
-    let codecs = [Codec::H264, Codec::H265];
+    let codecs = [Codec::H264, Codec::H265, Codec::AV1];
 
     for codec in codecs {
         println!("\n{:?}:", codec);
@@ -95,7 +96,7 @@ fn query_detailed_capabilities(
         println!("    Checking {}: ", desc);
 
         // Construct profile info
-        let (mut profile_info, mut h264_profile, mut h265_profile) = match codec {
+        let (mut profile_info, mut h264_profile, mut h265_profile, mut av1_profile) = match codec {
             Codec::H264 => {
                 let profile_idc = if subsampling == vk::VideoChromaSubsamplingFlagsKHR::TYPE_444 {
                     ash::vk::native::StdVideoH264ProfileIdc_STD_VIDEO_H264_PROFILE_IDC_HIGH_444_PREDICTIVE
@@ -110,7 +111,7 @@ fn query_detailed_capabilities(
                     .chroma_subsampling(subsampling)
                     .luma_bit_depth(bit_depth)
                     .chroma_bit_depth(bit_depth);
-                (info, Some(h264), None)
+                (info, Some(h264), None, None)
             }
             Codec::H265 => {
                 let profile_idc = if subsampling == vk::VideoChromaSubsamplingFlagsKHR::TYPE_444 {
@@ -128,30 +129,46 @@ fn query_detailed_capabilities(
                     .chroma_subsampling(subsampling)
                     .luma_bit_depth(bit_depth)
                     .chroma_bit_depth(bit_depth);
-                (info, None, Some(h265))
+                (info, None, Some(h265), None)
+            }
+            Codec::AV1 => {
+                let profile = if subsampling == vk::VideoChromaSubsamplingFlagsKHR::TYPE_444 {
+                    ash::vk::native::StdVideoAV1Profile_STD_VIDEO_AV1_PROFILE_HIGH
+                } else {
+                    ash::vk::native::StdVideoAV1Profile_STD_VIDEO_AV1_PROFILE_MAIN
+                };
+
+                let av1 = vk::VideoEncodeAV1ProfileInfoKHR::default().std_profile(profile);
+                let info = vk::VideoProfileInfoKHR::default()
+                    .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_AV1)
+                    .chroma_subsampling(subsampling)
+                    .luma_bit_depth(bit_depth)
+                    .chroma_bit_depth(bit_depth);
+                (info, None, None, Some(av1))
             }
-            _ => return Ok(()),
         };
 
         if let Some(h264) = &mut h264_profile {
-            profile_info.p_next = (h264 as *mut vk::VideoEncodeH264ProfileInfoKHR).cast();
+            profile_info = profile_info.push(h264);
         }
         if let Some(h265) = &mut h265_profile {
-            profile_info.p_next = (h265 as *mut vk::VideoEncodeH265ProfileInfoKHR).cast();
+            profile_info = profile_info.push(h265);
+        }
+        if let Some(av1) = &mut av1_profile {
+            profile_info = profile_info.push(av1);
         }
 
         // 1. Query Video Capabilities
-        let mut caps = vk::VideoCapabilitiesKHR::default();
         let mut encode_caps = vk::VideoEncodeCapabilitiesKHR::default();
-        caps.p_next = (&mut encode_caps as *mut vk::VideoEncodeCapabilitiesKHR).cast();
+        let mut caps = vk::VideoCapabilitiesKHR::default().push(&mut encode_caps);
 
         let mut h264_caps = vk::VideoEncodeH264CapabilitiesKHR::default();
         let mut h265_caps = vk::VideoEncodeH265CapabilitiesKHR::default();
-
-        if codec == Codec::H264 {
-            encode_caps.p_next = (&mut h264_caps as *mut vk::VideoEncodeH264CapabilitiesKHR).cast();
-        } else if codec == Codec::H265 {
-            encode_caps.p_next = (&mut h265_caps as *mut vk::VideoEncodeH265CapabilitiesKHR).cast();
+        let mut av1_caps = vk::VideoEncodeAV1CapabilitiesKHR::default();
+        match codec {
+            Codec::H264 => caps = caps.push(&mut h264_caps),
+            Codec::H265 => caps = caps.push(&mut h265_caps),
+            Codec::AV1 => caps = caps.push(&mut av1_caps),
         }
 
         let result = unsafe {
@@ -185,9 +202,9 @@ fn query_detailed_capabilities(
             vk::VideoProfileListInfoKHR::default().profiles(std::slice::from_ref(&profile_info));
 
         // Check for Input Image support (VIDEO_ENCODE_SRC_KHR)
-        let mut format_info = vk::PhysicalDeviceVideoFormatInfoKHR::default()
-            .image_usage(vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR);
-        format_info.p_next = (&mut format_props_list as *mut vk::VideoProfileListInfoKHR).cast();
+        let format_info = vk::PhysicalDeviceVideoFormatInfoKHR::default()
+            .image_usage(vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR)
+            .push(&mut format_props_list);
 
         let result = unsafe {
             (video_queue_fn
@@ -220,9 +237,9 @@ fn query_detailed_capabilities(
         }
 
         // Check for DPB Image support (VIDEO_ENCODE_DPB_KHR)
-        let mut format_info = vk::PhysicalDeviceVideoFormatInfoKHR::default()
-            .image_usage(vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR);
-        format_info.p_next = (&mut format_props_list as *mut vk::VideoProfileListInfoKHR).cast();
+        let format_info = vk::PhysicalDeviceVideoFormatInfoKHR::default()
+            .image_usage(vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR)
+            .push(&mut format_props_list);
 
         let mut format_props_count = 0;
         let result = unsafe {

src/converter/pipeline.rs

@@ -8,6 +8,7 @@ use crate::encoder::resources::find_memory_type;
 use crate::error::{PixelForgeError, Result};
 use crate::vulkan::VideoContext;
 use ash::vk;
+use ash::vk::TaggedStructure;
 
 /// Precompiled SPIR-V bytecode for the color conversion compute shader.
 const COLOR_CONVERT_SPIRV_BYTES: &[u8] = include_bytes!("../../shader/color_convert.spv");
@@ -153,10 +154,7 @@ pub fn create_converter(
     // runtime descriptor population. The correct sizes for in-buffer
     // descriptors are the regular `*_descriptor_size` fields.
     let mut db_props = vk::PhysicalDeviceDescriptorBufferPropertiesEXT::default();
-    let mut props = vk::PhysicalDeviceProperties2 {
-        p_next: &mut db_props as *mut _ as *mut _,
-        ..Default::default()
-    };
+    let mut props = vk::PhysicalDeviceProperties2::default().push(&mut db_props);
     unsafe {
         instance.get_physical_device_properties2(physical_device, &mut props);
     }

src/encoder/av1/api.rs

@@ -69,10 +69,6 @@ impl AV1Encoder {
             self.order_hint = 0;
             // Reset references for key frames.
             self.references.clear();
-            // Reset DPB slot activation tracking on key frame - all slots become inactive.
-            for active in &mut self.dpb_slot_active {
-                *active = false;
-            }
         }
 
         let mut encoded_data = Vec::new();
@@ -107,29 +103,42 @@ impl AV1Encoder {
         self.frame_num += 1;
         self.order_hint = (self.order_hint + 1) & 0xFF; // 8-bit order hint
 
-        // Only KEY frames are stored as references. P frames all reference the KEY frame
-        // and don't update any reference buffer, avoiding P→P which produces corrupt output
-        // on NVIDIA AV1 encoders.
+        // 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 {
-            let ref_info = super::ReferenceInfo {
-                dpb_slot: self.current_dpb_slot,
-                order_hint: encoded_order_hint,
-                frame_type: ash::vk::native::StdVideoAV1FrameType_STD_VIDEO_AV1_FRAME_TYPE_KEY,
-            };
             self.references.clear();
-            self.references.push(ref_info);
+        }
+        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<u8> = self.references.iter().map(|r| r.dpb_slot).collect();
+        let mut next_slot = (self.current_dpb_slot + 1) % self.dpb_slot_count as u8;
 
-            // KEY frame uses the current DPB slot; pick a different one for P frames.
-            let used_slots: Vec<u8> = self.references.iter().map(|r| r.dpb_slot).collect();
+        // 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) {
-                    self.current_dpb_slot = i;
+                    next_slot = i;
                     break;
                 }
             }
         }
-        // P frames reuse the same scratch DPB slot (current_dpb_slot stays unchanged
-        // between P frames since it's always different from the KEY frame's slot).
+
+        self.current_dpb_slot = next_slot;
 
         Ok(EncodedPacket {
             data: encoded_data,