Fix a couple of bugs in the hardware decoder.

vk-video/Cargo.toml

@@ -18,6 +18,7 @@ ash = "0.38.0"
 bytes = "1"
 derivative = "2.2.0"
 h264-reader = { git = "https://github.com/membraneframework-labs/h264-reader.git", branch = "live-compositor" }
+memchr = "2.7.4"
 thiserror = "1.0.59"
 tracing = "0.1.40"
 vk-mem = "0.4.0"

vk-video/examples/player/player/shader.wgsl

@@ -29,9 +29,10 @@ fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> {
     var u = uv.x;
     var v = uv.y;
 
-    let r = y + 1.40200 * (v - 128.0 / 255.0);
-    let g = y - 0.34414 * (u - 128.0 / 255.0) - 0.71414 * (v - 128.0 / 255.0);
-    let b = y + 1.77200 * (u - 128.0 / 255.0);
+    // https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion
+    let r = 1.1643828125 * y + 1.59602734375 * v - 0.87078515625;
+    let g = 1.1643828125 * y - 0.39176171875 * u - 0.81296875000 * v + 0.52959375;
+    let b = 1.1643828125 * y + 2.01723437500 * u - 1.08139062500;
 
     return vec4<f32>(clamp(r, 0.0, 1.0), clamp(g, 0.0, 1.0), clamp(b, 0.0, 1.0), 1.0);
 }

vk-video/src/parser/nalu_splitter.rs

@@ -1,22 +1,33 @@
+use std::sync::LazyLock;
+
 use bytes::{BufMut, BytesMut};
+use memchr::memmem::Finder;
 
 #[derive(Debug, Default)]
 pub(crate) struct NALUSplitter {
     buffer: BytesMut,
     pts: Option<u64>,
 }
 
-fn find_nalu_start_code(buf: &[u8]) -> Option<usize> {
-    if buf.is_empty() {
+fn find_start_of_next_nalu(buf: &[u8]) -> Option<usize> {
+    static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(&[0, 0, 1]));
+
+    if buf.len() < 4 {
         return None;
     };
 
-    buf.windows(3)
-        .enumerate()
-        .filter(|(_, window)| **window == [0, 0, 1])
-        .filter(|(i, window)| !(*i == 0 || (*i == 1 && window[0] == 0)))
-        .map(|(i, _)| i + 3)
-        .next()
+    // If the start code is at the beginning of nalu, we need to skip it, because this means we
+    // would give the parser only the start code, without a nal unit before it.
+    //
+    // The code can be either 0, 0, 0, 1 or 0, 0, 1.
+    // We're looking for the sequence 0, 0, 1. If we find it starting at the second byte of input,
+    // this could mean either that it's the longer 0, 0, 0, 1 code, or that there is one byte of
+    // the previous nalu and the shorter code. This is what is checked here.
+    if buf[0] != 0 && buf[1..4] == [0, 0, 1] {
+        return Some(5);
+    }
+
+    FINDER.find(&buf[2..]).map(|i| i + 5)
 }
 
 impl NALUSplitter {
@@ -34,7 +45,7 @@ impl NALUSplitter {
         self.buffer.put(bytestream);
         let mut result = Vec::new();
 
-        while let Some(i) = find_nalu_start_code(&self.buffer) {
+        while let Some(i) = find_start_of_next_nalu(&self.buffer) {
             let nalu = self.buffer.split_to(i);
             result.push((nalu.to_vec(), output_pts));
             output_pts = pts;

vk-video/src/vulkan_decoder.rs

@@ -94,6 +94,9 @@ pub enum VulkanDecoderError {
     #[error("Level changed during the stream")]
     LevelChangeUnsupported,
 
+    #[error("Monochrome video is not supported")]
+    MonochromeChromaFormatUnsupported,
+
     #[error(transparent)]
     VulkanCtxError(#[from] VulkanCtxError),
 }
@@ -522,10 +525,7 @@ impl VulkanDecoder<'_> {
             .get(&decode_information.sps_id)
             .ok_or(VulkanDecoderError::NoSession)?;
 
-        let dimensions = vk::Extent2D {
-            width: sps.width()?,
-            height: sps.height()?,
-        };
+        let dimensions = sps.size()?;
 
         Ok(DecodeSubmission {
             image: target_image,

vk-video/src/vulkan_decoder/session_resources.rs

@@ -71,9 +71,7 @@ impl VideoSessionResources<'_> {
             ));
         }
 
-        let width = sps.width()?;
-        let height = sps.height()?;
-        let max_coded_extent = vk::Extent2D { width, height };
+        let max_coded_extent = sps.size()?;
         // +1 for current frame
         let max_dpb_slots = sps.max_num_ref_frames + 1;
         let max_active_references = sps.max_num_ref_frames;
@@ -146,16 +144,13 @@ impl VideoSessionResources<'_> {
             return Err(VulkanDecoderError::LevelChangeUnsupported);
         }
 
-        let width = sps.width()?;
-        let height = sps.height()?;
-
-        let max_coded_extent = vk::Extent2D { width, height };
+        let max_coded_extent = sps.size()?;
         // +1 for current frame
         let max_dpb_slots = sps.max_num_ref_frames + 1;
         let max_active_references = sps.max_num_ref_frames;
 
-        if self.video_session.max_coded_extent.width >= width
-            && self.video_session.max_coded_extent.height >= height
+        if self.video_session.max_coded_extent.width >= max_coded_extent.width
+            && self.video_session.max_coded_extent.height >= max_coded_extent.height
             && self.video_session.max_dpb_slots >= max_dpb_slots
         {
             // no need to change the session

vk-video/src/vulkan_decoder/vulkan_ctx.rs

@@ -48,10 +48,10 @@ pub enum VulkanCtxError {
 }
 
 pub struct VulkanInstance {
+    pub wgpu_instance: Arc<wgpu::Instance>,
     _entry: Arc<Entry>,
     instance: Arc<Instance>,
     _debug_messenger: Option<DebugMessenger>,
-    pub wgpu_instance: Arc<wgpu::Instance>,
 }
 
 impl VulkanInstance {
@@ -323,6 +323,9 @@ impl std::fmt::Debug for VulkanInstance {
 }
 
 pub struct VulkanDevice {
+    pub wgpu_device: Arc<wgpu::Device>,
+    pub wgpu_queue: Arc<wgpu::Queue>,
+    pub wgpu_adapter: Arc<wgpu::Adapter>,
     _physical_device: vk::PhysicalDevice,
     pub(crate) device: Arc<Device>,
     pub(crate) allocator: Arc<Allocator>,
@@ -331,9 +334,6 @@ pub struct VulkanDevice {
     pub(crate) h264_dpb_format_properties: vk::VideoFormatPropertiesKHR<'static>,
     pub(crate) h264_dst_format_properties: Option<vk::VideoFormatPropertiesKHR<'static>>,
     pub(crate) h264_caps: vk::VideoDecodeH264CapabilitiesKHR<'static>,
-    pub wgpu_device: Arc<wgpu::Device>,
-    pub wgpu_queue: Arc<wgpu::Queue>,
-    pub wgpu_adapter: Arc<wgpu::Adapter>,
 }
 
 impl VulkanDevice {

vk-video/src/vulkan_decoder/wrappers/command.rs

@@ -91,3 +91,13 @@ impl std::ops::Deref for CommandBuffer {
         &self.buffer
     }
 }
+
+impl Drop for CommandBuffer {
+    fn drop(&mut self) {
+        unsafe {
+            self.pool
+                .device
+                .free_command_buffers(**self.pool, &[self.buffer])
+        };
+    }
+}

vk-video/src/vulkan_decoder/wrappers/parameter_sets.rs

@@ -1,32 +1,64 @@
 use ash::vk;
-use h264_reader::nal::sps::SeqParameterSet;
+use h264_reader::nal::sps::{FrameMbsFlags, SeqParameterSet};
 
 use crate::VulkanDecoderError;
 
 const MACROBLOCK_SIZE: u32 = 16;
 
 pub(crate) trait SeqParameterSetExt {
-    fn width(&self) -> Result<u32, VulkanDecoderError>;
-    fn height(&self) -> Result<u32, VulkanDecoderError>;
+    fn size(&self) -> Result<vk::Extent2D, VulkanDecoderError>;
 }
 
 impl SeqParameterSetExt for SeqParameterSet {
-    fn width(&self) -> Result<u32, VulkanDecoderError> {
-        match self.frame_cropping {
-            None => Ok((self.pic_width_in_mbs_minus1 + 1) * MACROBLOCK_SIZE),
-            Some(_) => Err(VulkanDecoderError::FrameCroppingNotSupported),
-        }
-    }
+    #[allow(non_snake_case)]
+    fn size(&self) -> Result<vk::Extent2D, VulkanDecoderError> {
+        let chroma_array_type = if self.chroma_info.separate_colour_plane_flag {
+            0
+        } else {
+            self.chroma_info.chroma_format.to_chroma_format_idc()
+        };
 
-    fn height(&self) -> Result<u32, VulkanDecoderError> {
-        match self.frame_mbs_flags {
-            h264_reader::nal::sps::FrameMbsFlags::Frames => {
-                Ok((self.pic_height_in_map_units_minus1 + 1) * MACROBLOCK_SIZE)
+        let (SubWidthC, SubHeightC) = match self.chroma_info.chroma_format {
+            h264_reader::nal::sps::ChromaFormat::Monochrome => {
+                return Err(VulkanDecoderError::MonochromeChromaFormatUnsupported)
             }
-            h264_reader::nal::sps::FrameMbsFlags::Fields { .. } => {
-                Err(VulkanDecoderError::FieldsNotSupported)
+            h264_reader::nal::sps::ChromaFormat::YUV420 => (2, 2),
+            h264_reader::nal::sps::ChromaFormat::YUV422 => (2, 1),
+            h264_reader::nal::sps::ChromaFormat::YUV444 => (1, 1),
+            h264_reader::nal::sps::ChromaFormat::Invalid(x) => {
+                return Err(VulkanDecoderError::InvalidInputData(format!(
+                    "Invalid chroma_format_idc: {x}"
+                )))
             }
-        }
+        };
+
+        let (CropUnitX, CropUnitY) = match chroma_array_type {
+            0 => (
+                1,
+                2 - (self.frame_mbs_flags == FrameMbsFlags::Frames) as u32,
+            ),
+
+            _ => (
+                SubWidthC,
+                SubHeightC * (2 - (self.frame_mbs_flags == FrameMbsFlags::Frames) as u32),
+            ),
+        };
+
+        let (width_offset, height_offset) = match &self.frame_cropping {
+            None => (0, 0),
+            Some(frame_cropping) => (
+                (frame_cropping.left_offset + frame_cropping.right_offset) * CropUnitX,
+                (frame_cropping.top_offset + frame_cropping.bottom_offset) * CropUnitY,
+            ),
+        };
+
+        let width = (self.pic_width_in_mbs_minus1 + 1) * MACROBLOCK_SIZE - width_offset;
+        let height = (self.pic_height_in_map_units_minus1 + 1)
+            * (2 - (self.frame_mbs_flags == FrameMbsFlags::Frames) as u32)
+            * MACROBLOCK_SIZE
+            - height_offset;
+
+        Ok(vk::Extent2D { width, height })
     }
 }