1use 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#[derive(Clone, Debug, Eq, PartialEq)]
20pub struct VulkanDecodeProbe {
21 pub decoder: String,
23 pub hardware_format: String,
25 pub software_format: String,
27 pub display_size: (i32, i32),
29 pub coded_size: (i32, i32),
31 pub hardware_frames: u64,
33 pub unique_images: usize,
35 pub reused_images: u64,
37 pub all_images_sampled: bool,
39 pub all_images_synchronized: bool,
41 pub seek_frames: u64,
43 pub device: u64,
45 pub physical_device: u64,
47 pub max_plane_images: u32,
49 pub interop_submissions: u64,
51 pub source_cpu_plane_copy_bytes: u64,
53 pub source_upload_bytes: u64,
55 pub source_download_bytes: u64,
57}
58
59#[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 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 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(¶meters).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 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}