Skip to main content

alpha2_native/
hwdecode.rs

1//! Standalone Linux Vulkan hardware-decode probe for the CLI.
2
3use std::collections::HashSet;
4use std::ffi::CString;
5use std::path::Path;
6
7use alpha2_ffmpeg_bridge::{VulkanFrameInfo, VulkanFrameLease, pixel_format_name};
8use alpha2_media::DecodeError;
9use alpha2_native_gpu::probe_vulkan_submission;
10use rsmpeg::avcodec::{AVCodec, AVCodecContext};
11use rsmpeg::avutil::{AVFrame, AVHWDeviceContext};
12use rsmpeg::error::RsmpegError;
13use rsmpeg::ffi;
14
15use crate::convert::adapter_error;
16use crate::decode::{open_input, probe_context};
17
18/// Result of a bounded Vulkan hardware decode and surface probe.
19#[derive(Clone, Debug, Eq, PartialEq)]
20pub struct VulkanDecodeProbe {
21    /// `FFmpeg` decoder implementation.
22    pub decoder: String,
23    /// Selected opaque hardware pixel format.
24    pub hardware_format: String,
25    /// Actual decoded surface layout.
26    pub software_format: String,
27    /// Display size reported by the decoded frames.
28    pub display_size: (i32, i32),
29    /// Allocated surface extent, padded to whole macroblocks by the codec.
30    pub coded_size: (i32, i32),
31    /// Number of hardware frames received before the bound.
32    pub hardware_frames: u64,
33    /// Number of distinct `VkImage` identities observed.
34    pub unique_images: usize,
35    /// Frames whose primary image had already appeared (pool reuse).
36    pub reused_images: u64,
37    /// Whether every decoded image included sampled-image usage.
38    pub all_images_sampled: bool,
39    /// Whether every image exposed a ready timeline semaphore.
40    pub all_images_synchronized: bool,
41    /// Number of hardware frames received after seeking back to the start.
42    pub seek_frames: u64,
43    /// Device identity shared by all observed surfaces.
44    pub device: u64,
45    /// Physical-device identity shared by all observed surfaces.
46    pub physical_device: u64,
47    /// Maximum number of Vulkan images backing one decoded frame.
48    pub max_plane_images: u32,
49    /// GPU-only synchronization/layout submissions completed by the probe.
50    pub interop_submissions: u64,
51    /// CPU source-plane copies performed by this probe.
52    pub source_cpu_plane_copy_bytes: u64,
53    /// CPU-to-GPU uploads performed by this probe.
54    pub source_upload_bytes: u64,
55    /// GPU-to-CPU downloads performed by this probe.
56    pub source_download_bytes: u64,
57}
58
59/// Decodes Vulkan frames without mapping, downloading, or uploading their planes.
60///
61/// Checks image, layout and timeline state, pool reuse, and decode after a seek. It opens
62/// its own device and decoder instead of [`crate::NativeExportDecoder`] to observe
63/// `FFmpeg`'s surfaces directly.
64///
65/// # Errors
66/// Returns an error if Vulkan decode is unavailable or a frame breaks
67/// the resident-surface contract.
68#[allow(clippy::too_many_lines)]
69pub fn probe_vulkan_decode(
70    path: impl AsRef<Path>,
71    frame_limit: usize,
72) -> Result<VulkanDecodeProbe, DecodeError> {
73    if frame_limit < 2 {
74        return Err(DecodeError::Unsupported(
75            "Vulkan decode probe requires at least two frames".into(),
76        ));
77    }
78    let mut input = open_input(path.as_ref())?;
79    let selected = probe_context(&input, path.as_ref())?
80        .selected_video
81        .ok_or_else(|| DecodeError::Unsupported("source has no video stream".into()))?;
82    input.discard_all_but(selected);
83    let parameters = input
84        .streams()
85        .get(selected)
86        .ok_or(DecodeError::OutOfRange)?
87        .codecpar();
88    let preferred = AVCodec::find_decoder(parameters.codec_id)
89        .ok_or_else(|| DecodeError::Unsupported("selected video codec has no decoder".into()))?;
90    // The preferred decoder may be software (libdav1d for AV1); the Vulkan configuration lives on
91    // `FFmpeg`'s own decoder for the codec.
92    let (codec, config) = crate::decode::registered_video_decoders(parameters.codec_id)
93        .find_map(|codec| {
94            let config = (0..)
95                .map_while(|index| codec.hw_config(index))
96                .find(|config| {
97                    config.device_type == ffi::AV_HWDEVICE_TYPE_VULKAN
98                        && config.pix_fmt == ffi::AV_PIX_FMT_VULKAN
99                        && crate::decode::has_hw_config_method(
100                            config.methods,
101                            ffi::AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX,
102                        )
103                })?;
104            Some((codec, config))
105        })
106        .ok_or_else(|| {
107            DecodeError::Unsupported(format!(
108                "decoder {} has no Vulkan HW_DEVICE_CTX configuration",
109                preferred.name().to_string_lossy()
110            ))
111        })?;
112    debug_assert_eq!(config.pix_fmt, ffi::AV_PIX_FMT_VULKAN);
113
114    // Pass the device by loader index: `FFmpeg` reads a leading digit as an index, so an
115    // `ALPHA2_GPU_ADAPTER` name passed unchanged could select another device.
116    let requested_device = CString::new(
117        alpha2_native_gpu::shared_device_index()
118            .map_err(|error| DecodeError::Unsupported(error.to_string()))?
119            .to_string(),
120    )
121    .map_err(|_| DecodeError::Unsupported("a device index held a NUL byte".into()))?;
122    let device = AVHWDeviceContext::create(
123        ffi::AV_HWDEVICE_TYPE_VULKAN,
124        Some(&requested_device),
125        None,
126        0,
127    )
128    .map_err(adapter_error)?;
129    let mut decoder = AVCodecContext::new(&codec);
130    decoder.apply_codecpar(&parameters).map_err(adapter_error)?;
131    drop(parameters);
132    decoder.set_pkt_timebase(input.streams()[selected].time_base);
133    decoder.set_get_format(Some(alpha2_ffmpeg_bridge::vulkan_get_format_callback()));
134    decoder.set_hw_device_ctx(device);
135    decoder.open(None).map_err(adapter_error)?;
136
137    let mut demux_eof = false;
138    let mut flushed = false;
139    let mut images = HashSet::new();
140    let mut reuse = 0_u64;
141    let mut first_info = None;
142    let mut all_sampled = true;
143    let mut all_synchronized = true;
144    let mut max_plane_images = 0;
145    let mut frames = 0_u64;
146    while frames < frame_limit as u64 {
147        let Some(frame) = decode_next_raw(
148            &mut input,
149            &mut decoder,
150            selected,
151            &mut demux_eof,
152            &mut flushed,
153        )?
154        else {
155            break;
156        };
157        let lease = VulkanFrameLease::new(frame);
158        let info = validate_frame(&lease, first_info.as_ref())?;
159        probe_vulkan_submission(&lease).map_err(adapter_error)?;
160        let primary = info
161            .images
162            .iter()
163            .copied()
164            .find(|image| *image != 0)
165            .ok_or_else(|| DecodeError::Adapter("Vulkan frame has no image".into()))?;
166        if !images.insert(primary) {
167            reuse += 1;
168        }
169        all_sampled &= info.sampled;
170        all_synchronized &= synchronized(&info);
171        max_plane_images = max_plane_images.max(info.plane_count);
172        first_info.get_or_insert(info);
173        frames += 1;
174    }
175    let first_info = first_info.ok_or_else(|| {
176        DecodeError::Unsupported("decoder returned no Vulkan hardware frame".into())
177    })?;
178    if !all_sampled {
179        return Err(DecodeError::Unsupported(
180            "decoded Vulkan image is not sampleable".into(),
181        ));
182    }
183    if !all_synchronized {
184        return Err(DecodeError::Adapter(
185            "decoded Vulkan image has no ready timeline semaphore".into(),
186        ));
187    }
188
189    input
190        .seek(
191            i32::try_from(selected).map_err(|_| DecodeError::OutOfRange)?,
192            0,
193            ffi::AVSEEK_FLAG_BACKWARD.cast_signed(),
194        )
195        .map_err(adapter_error)?;
196    decoder.flush_buffers();
197    demux_eof = false;
198    flushed = false;
199    let seek_frame = decode_next_raw(
200        &mut input,
201        &mut decoder,
202        selected,
203        &mut demux_eof,
204        &mut flushed,
205    )?
206    .ok_or_else(|| DecodeError::Adapter("no Vulkan frame after seek".into()))?;
207    let seek_lease = VulkanFrameLease::new(seek_frame);
208    validate_frame(&seek_lease, Some(&first_info))?;
209    probe_vulkan_submission(&seek_lease).map_err(adapter_error)?;
210
211    Ok(VulkanDecodeProbe {
212        decoder: codec.name().to_string_lossy().into_owned(),
213        hardware_format: format_name(first_info.hardware_format),
214        software_format: format_name(first_info.software_format),
215        display_size: (first_info.width, first_info.height),
216        coded_size: (first_info.coded_width, first_info.coded_height),
217        hardware_frames: frames,
218        unique_images: images.len(),
219        reused_images: reuse,
220        all_images_sampled: all_sampled,
221        all_images_synchronized: all_synchronized,
222        seek_frames: 1,
223        device: first_info.device,
224        physical_device: first_info.physical_device,
225        max_plane_images,
226        interop_submissions: frames + 1,
227        source_cpu_plane_copy_bytes: 0,
228        source_upload_bytes: 0,
229        source_download_bytes: 0,
230    })
231}
232
233fn decode_next_raw(
234    input: &mut alpha2_ffmpeg_bridge::DemuxInput,
235    decoder: &mut AVCodecContext,
236    stream_index: usize,
237    demux_eof: &mut bool,
238    flushed: &mut bool,
239) -> Result<Option<AVFrame>, DecodeError> {
240    loop {
241        match decoder.receive_frame() {
242            Ok(frame) => return Ok(Some(frame)),
243            Err(RsmpegError::DecoderDrainError) => {}
244            Err(RsmpegError::DecoderFlushedError) => return Ok(None),
245            Err(error) => return Err(adapter_error(error)),
246        }
247        if *demux_eof {
248            if !*flushed {
249                decoder.send_packet(None).map_err(adapter_error)?;
250                *flushed = true;
251                continue;
252            }
253            return Ok(None);
254        }
255        let Some(packet) = input.read_packet().map_err(adapter_error)? else {
256            *demux_eof = true;
257            continue;
258        };
259        if usize::try_from(packet.stream_index).ok() == Some(stream_index) {
260            decoder.send_packet(Some(&packet)).map_err(adapter_error)?;
261        }
262    }
263}
264
265fn validate_frame(
266    frame: &VulkanFrameLease,
267    expected: Option<&VulkanFrameInfo>,
268) -> Result<VulkanFrameInfo, DecodeError> {
269    let info = frame.info().map_err(adapter_error)?;
270    if info.hardware_format != ffi::AV_PIX_FMT_VULKAN {
271        return Err(DecodeError::Adapter(format!(
272            "decoder returned {} instead of vulkan",
273            format_name(info.hardware_format)
274        )));
275    }
276    if let Some(expected) = expected {
277        if info.device != expected.device || info.physical_device != expected.physical_device {
278            return Err(DecodeError::Adapter(
279                "Vulkan decoder changed device while probing".into(),
280            ));
281        }
282        if info.software_format != expected.software_format {
283            return Err(DecodeError::Adapter(
284                "Vulkan decoder changed surface format while probing".into(),
285            ));
286        }
287        if info.coded_width != expected.coded_width || info.coded_height != expected.coded_height {
288            return Err(DecodeError::Adapter(
289                "Vulkan decoder changed surface extent while probing".into(),
290            ));
291        }
292    }
293    // Presentation normalizes against the allocated extent, so a smaller surface
294    // would rescale every sample.
295    if info.coded_width < info.width || info.coded_height < info.height {
296        return Err(DecodeError::Adapter(
297            "Vulkan surface is smaller than the frame it carries".into(),
298        ));
299    }
300    Ok(info)
301}
302
303fn synchronized(info: &VulkanFrameInfo) -> bool {
304    info.images
305        .iter()
306        .enumerate()
307        .filter(|(_, image)| **image != 0)
308        .all(|(index, _)| info.semaphores[index] != 0)
309}
310
311fn format_name(format: ffi::AVPixelFormat) -> String {
312    pixel_format_name(format).unwrap_or_else(|| format!("unknown({format})"))
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318
319    #[test]
320    fn synchronization_ignores_unused_plane_slots() {
321        let info = VulkanFrameInfo {
322            hardware_format: ffi::AV_PIX_FMT_VULKAN,
323            software_format: ffi::AV_PIX_FMT_NV12,
324            width: 16,
325            height: 16,
326            coded_width: 16,
327            coded_height: 16,
328            plane_count: 1,
329            sampled: true,
330            instance: 1,
331            physical_device: 2,
332            device: 3,
333            images: [4, 0, 0, 0, 0, 0, 0, 0],
334            formats: [1; 8],
335            layouts: [0; 8],
336            access: [0; 8],
337            semaphores: [5, 0, 0, 0, 0, 0, 0, 0],
338            semaphore_values: [0; 8],
339            queue_families: [0; 8],
340        };
341        assert!(synchronized(&info));
342    }
343}