1use std::sync::Arc;
2
3use ash::vk;
4
5use h264_reader::nal::{pps::PicParameterSet, sps::SeqParameterSet};
6use rustc_hash::FxHashMap;
7use session_resources::VideoSessionResources;
8
9use crate::{
10 RawFrameData,
11 device::{ColorRange, ColorSpace, DecodingDevice},
12 parser::{
13 decoder_instructions::DecoderInstruction,
14 reference_manager::{DecodeInformation, ReferenceId},
15 },
16};
17use crate::{VulkanCommonError, wrappers::*};
18
19mod frame_sorter;
20mod session_resources;
21
22pub(crate) use frame_sorter::FrameSorter;
23
24pub struct VulkanDecoder<'a> {
25 video_session_resources: Option<VideoSessionResources<'a>>,
26 pub(crate) tracker: DecoderTracker,
27 reference_id_to_dpb_slot_index: FxHashMap<ReferenceId, usize>,
28 decoding_device: Arc<DecodingDevice>,
29 usage_info: vk::VideoDecodeUsageInfoKHR<'a>,
30 image_modifiers: ImageModifiers,
31}
32
33#[derive(Debug, Clone, Copy)]
34pub(crate) struct ImageModifiers {
35 pub(crate) create_flags: vk::ImageCreateFlags,
36 pub(crate) usage_flags: vk::ImageUsageFlags,
37 pub(crate) additional_queue_index: usize,
38}
39
40pub(crate) enum DecoderTrackerWaitState {
41 NewDecodingImagesLayoutTransition,
42 Decode,
43 DownloadImageToBuffer,
44 #[cfg_attr(not(feature = "transcoder"), allow(dead_code))]
45 ExternalProcessing,
46}
47
48pub(crate) struct DecoderTrackerKind {}
49
50impl TrackerKind for DecoderTrackerKind {
51 type WaitState = DecoderTrackerWaitState;
52
53 type CommandBufferPools = DecoderCommandBufferPools;
54}
55
56pub(crate) struct DecoderCommandBufferPools {
57 decode: CommandBufferPool,
58 transfer: CommandBufferPool,
59}
60
61impl CommandBufferPoolStorage for DecoderCommandBufferPools {
62 fn mark_submitted_as_free(&mut self, last_waited_for: SemaphoreWaitValue) {
63 self.decode.mark_submitted_as_free(last_waited_for);
64 self.transfer.mark_submitted_as_free(last_waited_for);
65 }
66}
67
68pub(crate) type DecoderTracker = Tracker<DecoderTrackerKind>;
69
70pub(crate) struct DecodeSubmissionImageInfo {
71 pub(crate) image: Arc<Image>,
72 pub(crate) layer: u32,
73 pub(crate) cropped_extent: vk::Extent2D,
74}
75
76pub(crate) struct DecodeResultMetadata {
77 pub(crate) pts: Option<u64>,
78 pub(crate) pic_order_cnt: i32,
79 pub(crate) max_num_reorder_frames: u64,
80 pub(crate) is_idr: bool,
81 pub(crate) color_space: ColorSpace,
82 pub(crate) color_range: ColorRange,
83}
84
85pub(crate) struct DecodeResult<T> {
86 pub(crate) frame: T,
87 pub(crate) metadata: DecodeResultMetadata,
88}
89
90pub(crate) struct InFlightDecodeResources {
92 _video_session: Arc<VideoSession>,
93 _video_session_params: Arc<VideoSessionParameters>,
94 _dpb_image_with_view: Arc<ImageWithView>,
95 _dst_image_with_view: Option<Arc<ImageWithView>>,
96}
97
98pub(crate) struct DecodeSubmission<'borrow, 'decoder> {
99 pub(crate) decode_result: DecodeResult<DecodeSubmissionImageInfo>,
100 pub(crate) decoder: &'borrow mut VulkanDecoder<'decoder>,
101 pub(crate) input_buffer: DecodeInputBuffer,
102 pub(crate) decode_query_pool: Option<Arc<DecodingQueryPool>>,
103 #[cfg_attr(not(feature = "transcoder"), allow(dead_code))]
104 pub(crate) semaphore_wait_value: SemaphoreWaitValue,
105 #[cfg_attr(not(feature = "transcoder"), allow(dead_code))]
106 pub(crate) in_flight_resources: InFlightDecodeResources,
107}
108
109impl<'a, 'b> DecodeSubmission<'a, 'b> {
110 fn download_output(self) -> Result<DecodeResult<RawFrameData>, VulkanDecoderError> {
111 let raw_frame_data = self.decoder.download_output(&self.decode_result.frame)?;
112 let frame = RawFrameData {
113 frame: raw_frame_data,
114 width: self.decode_result.frame.cropped_extent.width,
115 height: self.decode_result.frame.cropped_extent.height,
116 };
117
118 self.finish(frame)
119 }
120
121 #[cfg(feature = "wgpu")]
122 fn output_to_wgpu_texture(self) -> Result<DecodeResult<wgpu::Texture>, VulkanDecoderError> {
123 let wgpu_texture = self
124 .decoder
125 .output_to_wgpu_texture(&self.decode_result.frame)?;
126
127 self.finish(wgpu_texture)
128 }
129
130 fn finish<T>(self, output: T) -> Result<DecodeResult<T>, VulkanDecoderError> {
131 self.input_buffer.release_to_pool();
132
133 if let Some(query_pool) = self.decode_query_pool {
134 query_pool.check_results_blocking()?;
135 }
136
137 Ok(DecodeResult {
138 frame: output,
139 metadata: self.decode_result.metadata,
140 })
141 }
142}
143
144#[derive(Debug, thiserror::Error)]
145pub enum VulkanDecoderError {
146 #[error("Vulkan error: {0}")]
147 VkError(#[from] vk::Result),
148
149 #[error("The device does not support vulkan h264 decoding")]
150 VulkanDecoderUnsupported,
151
152 #[error(
153 "A NALU requiring a session received before a session was created (probably before receiving first SPS)"
154 )]
155 NoSession,
156
157 #[error(
158 "A picture which is not in the decoded pictures buffer was requested as a reference picture"
159 )]
160 NonExistentReferenceRequested,
161
162 #[error("A vulkan decode operation failed with code {0:?}")]
163 DecodeOperationFailed(vk::QueryResultStatusKHR),
164
165 #[error("Invalid input data for the decoder: {0}.")]
166 InvalidInputData(String),
167
168 #[error("Monochrome video is not supported")]
169 MonochromeChromaFormatUnsupported,
170
171 #[error(transparent)]
172 VulkanCommonError(#[from] VulkanCommonError),
173}
174
175impl VulkanDecoder<'_> {
176 pub fn new(
177 decoding_device: Arc<DecodingDevice>,
178 usage_flags: crate::parameters::DecoderUsageFlags,
179 image_modifiers: ImageModifiers,
180 ) -> Result<Self, VulkanDecoderError> {
181 let command_buffer_pools = DecoderCommandBufferPools {
182 transfer: CommandBufferPool::new(
183 decoding_device.vulkan_device.clone(),
184 decoding_device.vulkan_device.queues.transfer.family_index,
185 )?,
186 decode: CommandBufferPool::new(
187 decoding_device.vulkan_device.clone(),
188 decoding_device.h264_decode_queues.family_index,
189 )?,
190 };
191
192 let tracker = Tracker::new(
193 decoding_device.vulkan_device.device.clone(),
194 command_buffer_pools,
195 Some("decoder"),
196 )?;
197
198 let usage_info = vk::VideoDecodeUsageInfoKHR::default().video_usage_hints(usage_flags);
199
200 Ok(Self {
201 decoding_device,
202 video_session_resources: None,
203 tracker,
204 reference_id_to_dpb_slot_index: Default::default(),
205 usage_info,
206 image_modifiers,
207 })
208 }
209}
210
211impl<'a> VulkanDecoder<'a> {
212 pub fn decode_to_bytes(
213 &mut self,
214 decoder_instructions: &[DecoderInstruction],
215 ) -> Result<Vec<DecodeResult<RawFrameData>>, VulkanDecoderError> {
216 let mut result = Vec::new();
217 for instruction in decoder_instructions {
218 if let Some(output) = self.decode(instruction)? {
219 result.push(output.download_output()?);
220 }
221 }
222
223 Ok(result)
224 }
225
226 #[cfg(feature = "wgpu")]
227 pub fn decode_to_wgpu_textures(
228 &mut self,
229 decoder_instructions: &[DecoderInstruction],
230 ) -> Result<Vec<DecodeResult<wgpu::Texture>>, VulkanDecoderError> {
231 let mut result = Vec::new();
232 for instruction in decoder_instructions {
233 if let Some(output) = self.decode(instruction)? {
234 result.push(output.output_to_wgpu_texture()?);
235 }
236 }
237
238 Ok(result)
239 }
240
241 pub(crate) fn decode<'b>(
242 &'b mut self,
243 instruction: &DecoderInstruction,
244 ) -> Result<Option<DecodeSubmission<'b, 'a>>, VulkanDecoderError> {
245 match instruction {
246 DecoderInstruction::Decode {
247 decode_info,
248 reference_id,
249 } => {
250 return self
251 .process_reference_frame(decode_info, *reference_id)
252 .map(Option::Some);
253 }
254
255 DecoderInstruction::Idr {
256 decode_info,
257 reference_id,
258 } => {
259 return self
260 .process_idr(decode_info, *reference_id)
261 .map(Option::Some);
262 }
263
264 DecoderInstruction::Drop { reference_ids } => {
265 for reference_id in reference_ids {
266 match self.reference_id_to_dpb_slot_index.remove(reference_id) {
267 Some(dpb_idx) => self
268 .video_session_resources
269 .as_mut()
270 .map(|s| s.free_reference_picture(dpb_idx)),
271 None => return Err(VulkanDecoderError::NonExistentReferenceRequested),
272 };
273 }
274 }
275
276 DecoderInstruction::Sps(sps) => self.process_sps(sps)?,
277
278 DecoderInstruction::Pps(pps) => self.process_pps(pps)?,
279 }
280
281 Ok(None)
282 }
283
284 fn process_sps(&mut self, sps: &SeqParameterSet) -> Result<(), VulkanDecoderError> {
285 match self.video_session_resources.as_mut() {
286 Some(session) => session.process_sps(sps.clone(), self.usage_info)?,
287 None => {
288 self.video_session_resources = Some(VideoSessionResources::new_from_sps(
289 &self.decoding_device,
290 self.tracker.command_buffer_pools.decode.begin_buffer()?,
291 sps.clone(),
292 self.usage_info,
293 &mut self.tracker,
294 self.image_modifiers,
295 )?)
296 }
297 }
298
299 Ok(())
300 }
301
302 fn process_pps(&mut self, pps: &PicParameterSet) -> Result<(), VulkanDecoderError> {
303 self.video_session_resources
304 .as_mut()
305 .ok_or(VulkanDecoderError::NoSession)?
306 .process_pps(pps.clone())?;
307
308 Ok(())
309 }
310
311 fn process_idr<'b>(
312 &'b mut self,
313 decode_information: &DecodeInformation,
314 reference_id: ReferenceId,
315 ) -> Result<DecodeSubmission<'b, 'a>, VulkanDecoderError> {
316 self.do_decode(decode_information, reference_id, true, true)
317 }
318
319 fn process_reference_frame<'b>(
320 &'b mut self,
321 decode_information: &DecodeInformation,
322 reference_id: ReferenceId,
323 ) -> Result<DecodeSubmission<'b, 'a>, VulkanDecoderError> {
324 self.do_decode(decode_information, reference_id, false, true)
325 }
326
327 fn do_decode<'b>(
328 &'b mut self,
329 decode_information: &'_ DecodeInformation,
330 reference_id: ReferenceId,
331 is_idr: bool,
332 is_reference: bool,
333 ) -> Result<DecodeSubmission<'b, 'a>, VulkanDecoderError> {
334 let video_session_resources = self
335 .video_session_resources
336 .as_mut()
337 .ok_or(VulkanDecoderError::NoSession)?;
338
339 let sps = video_session_resources
340 .sps
341 .get(&decode_information.sps_id)
342 .ok_or(VulkanDecoderError::InvalidInputData(format!(
343 "Unknown SPS id {}",
344 decode_information.sps_id
345 )))?;
346
347 let cropped_extent = sps.size()?;
348 let color_space = ColorSpace::from(sps);
349 let color_range = ColorRange::from(sps);
350
351 if is_idr {
352 video_session_resources.ensure_session(
353 &self.decoding_device,
354 self.tracker.command_buffer_pools.decode.begin_buffer()?,
355 &mut self.tracker,
356 )?;
357 }
358
359 let size = (decode_information.rbsp_bytes.len() as u64).next_multiple_of(
361 self.decoding_device
362 .profile_capabilities
363 .video_capabilities
364 .min_bitstream_buffer_size_alignment,
365 );
366
367 let mut buffer = video_session_resources.decode_buffer_pool.buffer()?;
368 buffer.upload_data(
369 &decode_information.rbsp_bytes,
370 size,
371 &video_session_resources.parameters.profile_info,
372 )?;
373
374 if is_idr {
377 video_session_resources
378 .decoding_images
379 .reset_all_allocations();
380
381 self.reference_id_to_dpb_slot_index = Default::default();
382 }
383
384 let mut cmd_buffer = self.tracker.command_buffer_pools.decode.begin_buffer()?;
386
387 video_session_resources
388 .decoding_images
389 .dpb
390 .image
391 .image_with_view
392 .transition_layout(
393 &mut cmd_buffer,
394 vk::PipelineStageFlags2::VIDEO_DECODE_KHR
395 ..vk::PipelineStageFlags2::VIDEO_DECODE_KHR,
396 vk::AccessFlags2::VIDEO_DECODE_WRITE_KHR
397 ..vk::AccessFlags2::VIDEO_DECODE_WRITE_KHR
398 | vk::AccessFlags2::VIDEO_DECODE_READ_KHR,
399 vk::ImageLayout::VIDEO_DECODE_DPB_KHR,
400 vk::ImageSubresourceRange {
401 base_array_layer: 0,
402 layer_count: vk::REMAINING_ARRAY_LAYERS,
403 aspect_mask: vk::ImageAspectFlags::COLOR,
404 base_mip_level: 0,
405 level_count: 1,
406 },
407 )?;
408
409 if let Some(dst) = &video_session_resources.decoding_images.dst_image {
410 dst.image_with_view.transition_layout(
411 &mut cmd_buffer,
412 vk::PipelineStageFlags2::VIDEO_DECODE_KHR
413 ..vk::PipelineStageFlags2::VIDEO_DECODE_KHR,
414 vk::AccessFlags2::VIDEO_DECODE_WRITE_KHR
415 ..vk::AccessFlags2::VIDEO_DECODE_WRITE_KHR
416 | vk::AccessFlags2::VIDEO_DECODE_READ_KHR,
417 vk::ImageLayout::VIDEO_DECODE_DST_KHR,
418 vk::ImageSubresourceRange {
419 base_array_layer: 0,
420 layer_count: vk::REMAINING_ARRAY_LAYERS,
421 aspect_mask: vk::ImageAspectFlags::COLOR,
422 base_mip_level: 0,
423 level_count: 1,
424 },
425 )?;
426 }
427
428 let memory_barrier = vk::MemoryBarrier2::default()
429 .src_stage_mask(vk::PipelineStageFlags2::VIDEO_DECODE_KHR)
430 .src_access_mask(vk::AccessFlags2::VIDEO_DECODE_WRITE_KHR)
431 .dst_stage_mask(vk::PipelineStageFlags2::VIDEO_DECODE_KHR)
432 .dst_access_mask(
433 vk::AccessFlags2::VIDEO_DECODE_READ_KHR | vk::AccessFlags2::VIDEO_DECODE_WRITE_KHR,
434 );
435
436 unsafe {
437 self.decoding_device
438 .vulkan_device
439 .device
440 .cmd_pipeline_barrier2(
441 cmd_buffer.buffer(),
442 &vk::DependencyInfo::default().memory_barriers(&[memory_barrier]),
443 )
444 };
445
446 if let Some(pool) = video_session_resources.decode_query_pool.as_ref() {
447 pool.reset(cmd_buffer.buffer());
448 }
449
450 let reference_slots = video_session_resources
451 .decoding_images
452 .reference_slot_info();
453
454 let begin_info = vk::VideoBeginCodingInfoKHR::default()
455 .video_session(video_session_resources.video_session.session)
456 .video_session_parameters(video_session_resources.parameters_manager.parameters())
457 .reference_slots(&reference_slots);
458
459 unsafe {
460 self.decoding_device
461 .vulkan_device
462 .device
463 .video_queue_ext
464 .cmd_begin_video_coding_khr(cmd_buffer.buffer(), &begin_info)
465 };
466
467 if is_idr {
469 let control_info = vk::VideoCodingControlInfoKHR::default()
470 .flags(vk::VideoCodingControlFlagsKHR::RESET);
471
472 unsafe {
473 self.decoding_device
474 .vulkan_device
475 .device
476 .video_queue_ext
477 .cmd_control_video_coding_khr(cmd_buffer.buffer(), &control_info)
478 };
479 }
480
481 let new_reference_slot_index = video_session_resources
483 .decoding_images
484 .allocate_reference_picture()?;
485
486 let new_reference_slot_std_reference_info = decode_information.picture_info.into();
487 let mut new_reference_slot_dpb_slot_info = vk::VideoDecodeH264DpbSlotInfoKHR::default()
488 .std_reference_info(&new_reference_slot_std_reference_info);
489
490 let new_reference_slot_video_picture_resource_info = video_session_resources
491 .decoding_images
492 .video_resource_info(new_reference_slot_index)
493 .unwrap();
494
495 let setup_reference_slot = vk::VideoReferenceSlotInfoKHR::default()
496 .picture_resource(new_reference_slot_video_picture_resource_info)
497 .slot_index(new_reference_slot_index as i32)
498 .push_next(&mut new_reference_slot_dpb_slot_info);
499
500 let reference_slots = video_session_resources
502 .decoding_images
503 .reference_slot_info();
504
505 let references_std_ref_info = Self::prepare_references_std_ref_info(decode_information);
506
507 let mut references_dpb_slot_info =
508 Self::prepare_references_dpb_slot_info(&references_std_ref_info);
509
510 let pic_reference_slots = Self::prepare_reference_list_slot_info(
511 &self.reference_id_to_dpb_slot_index,
512 &reference_slots,
513 &mut references_dpb_slot_info,
514 decode_information,
515 )?;
516
517 let std_picture_info = vk::native::StdVideoDecodeH264PictureInfo {
519 flags: vk::native::StdVideoDecodeH264PictureInfoFlags {
520 _bitfield_align_1: [],
521 __bindgen_padding_0: [0; 3],
522 _bitfield_1: vk::native::StdVideoDecodeH264PictureInfoFlags::new_bitfield_1(
523 matches!(
524 decode_information.header.field_pic,
525 h264_reader::nal::slice::FieldPic::Field(..)
526 )
527 .into(),
528 is_idr.into(),
529 is_idr.into(),
530 0,
531 is_reference.into(),
532 0,
533 ),
534 },
535 PicOrderCnt: decode_information.picture_info.PicOrderCnt_for_decoding,
536 seq_parameter_set_id: decode_information.sps_id,
537 pic_parameter_set_id: decode_information.pps_id,
538 frame_num: decode_information.header.frame_num,
539 idr_pic_id: decode_information
540 .header
541 .idr_pic_id
542 .map(|a| a as u16)
543 .unwrap_or(0),
544 reserved1: 0,
545 reserved2: 0,
546 };
547
548 let slice_offsets = decode_information
549 .slice_indices
550 .iter()
551 .map(|&x| x as u32)
552 .collect::<Vec<_>>();
553
554 let mut decode_h264_picture_info = vk::VideoDecodeH264PictureInfoKHR::default()
555 .std_picture_info(&std_picture_info)
556 .slice_offsets(&slice_offsets);
557
558 let dst_picture_resource_info = &video_session_resources
559 .decoding_images
560 .target_picture_resource_info(new_reference_slot_index)
561 .unwrap();
562
563 let (target_image, target_layer) = video_session_resources
565 .decoding_images
566 .target_info(new_reference_slot_index);
567
568 let decode_info = vk::VideoDecodeInfoKHR::default()
570 .src_buffer(*buffer.buffer)
571 .src_buffer_offset(0)
572 .src_buffer_range(size)
573 .dst_picture_resource(*dst_picture_resource_info)
574 .setup_reference_slot(&setup_reference_slot)
575 .reference_slots(&pic_reference_slots)
576 .push_next(&mut decode_h264_picture_info);
577
578 if let Some(pool) = video_session_resources.decode_query_pool.as_ref() {
579 pool.begin_query(cmd_buffer.buffer());
580 }
581
582 unsafe {
583 self.decoding_device
584 .vulkan_device
585 .device
586 .video_decode_queue_ext
587 .cmd_decode_video_khr(cmd_buffer.buffer(), &decode_info)
588 };
589
590 if let Some(pool) = video_session_resources.decode_query_pool.as_ref() {
591 pool.end_query(cmd_buffer.buffer());
592 }
593
594 unsafe {
595 self.decoding_device
596 .vulkan_device
597 .device
598 .video_queue_ext
599 .cmd_end_video_coding_khr(
600 cmd_buffer.buffer(),
601 &vk::VideoEndCodingInfoKHR::default(),
602 )
603 };
604
605 let semaphore_wait_value = self
606 .decoding_device
607 .h264_decode_queues
608 .submit_chain_semaphore(
609 cmd_buffer.end()?,
610 &mut self.tracker,
611 vk::PipelineStageFlags2::ALL_COMMANDS,
612 vk::PipelineStageFlags2::ALL_COMMANDS,
613 DecoderTrackerWaitState::Decode,
614 )?;
615
616 self.reference_id_to_dpb_slot_index
618 .insert(reference_id, new_reference_slot_index);
619
620 let in_flight_resources = InFlightDecodeResources {
621 _video_session: video_session_resources.video_session.clone(),
622 _video_session_params: video_session_resources
623 .parameters_manager
624 .parameters
625 .clone(),
626 _dpb_image_with_view: video_session_resources
627 .decoding_images
628 .dpb_image_with_view(),
629 _dst_image_with_view: video_session_resources
630 .decoding_images
631 .dst_image_with_view(),
632 };
633
634 Ok(DecodeSubmission {
635 decode_result: DecodeResult {
636 frame: DecodeSubmissionImageInfo {
637 image: target_image,
638 layer: target_layer as u32,
639 cropped_extent,
640 },
641 metadata: DecodeResultMetadata {
642 pic_order_cnt: decode_information.picture_info.PicOrderCnt_for_decoding[0],
643 max_num_reorder_frames: video_session_resources
644 .parameters
645 .max_num_reorder_frames,
646 is_idr,
647 pts: decode_information.pts,
648 color_space,
649 color_range,
650 },
651 },
652 semaphore_wait_value,
653 decode_query_pool: video_session_resources.decode_query_pool.clone(),
654 in_flight_resources,
655 input_buffer: buffer,
656 decoder: self,
657 })
658 }
659
660 #[cfg(feature = "wgpu")]
661 fn output_to_wgpu_texture(
662 &mut self,
663 decode_output: &DecodeSubmissionImageInfo,
664 ) -> Result<wgpu::Texture, VulkanDecoderError> {
665 let wgpu_device = unsafe {
666 self.decoding_device
667 .wgpu_device()
668 .as_hal::<wgpu::hal::vulkan::Api>()
669 .unwrap()
670 };
671 let copy_extent = vk::Extent3D {
672 width: decode_output.cropped_extent.width,
673 height: decode_output.cropped_extent.height,
674 depth: 1,
675 };
676
677 let queue_indices = [
678 self.decoding_device.queues.transfer.family_index as u32,
679 self.decoding_device.queues.wgpu.family_index as u32,
680 ];
681
682 let create_info = vk::ImageCreateInfo::default()
683 .flags(vk::ImageCreateFlags::MUTABLE_FORMAT)
684 .image_type(vk::ImageType::TYPE_2D)
685 .format(vk::Format::G8_B8R8_2PLANE_420_UNORM)
686 .extent(copy_extent)
687 .mip_levels(1)
688 .array_layers(1)
689 .samples(vk::SampleCountFlags::TYPE_1)
690 .tiling(vk::ImageTiling::OPTIMAL)
691 .usage(
692 vk::ImageUsageFlags::SAMPLED
693 | vk::ImageUsageFlags::TRANSFER_DST
694 | vk::ImageUsageFlags::TRANSFER_SRC,
695 )
696 .sharing_mode(vk::SharingMode::CONCURRENT)
697 .queue_family_indices(&queue_indices)
698 .initial_layout(vk::ImageLayout::UNDEFINED);
699
700 let image = Image::new(
701 self.decoding_device.allocator.clone(),
702 &create_info,
703 self.tracker.image_layout_tracker.clone(),
704 )?;
705
706 let mut cmd_buffer = self.tracker.command_buffer_pools.transfer.begin_buffer()?;
707
708 decode_output.image.transition_layout_single_layer(
709 &mut cmd_buffer,
710 vk::PipelineStageFlags2::NONE..vk::PipelineStageFlags2::COPY,
711 vk::AccessFlags2::NONE..vk::AccessFlags2::TRANSFER_READ,
712 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
713 decode_output.layer,
714 )?;
715
716 image.transition_layout_single_layer(
717 &mut cmd_buffer,
718 vk::PipelineStageFlags2::NONE..vk::PipelineStageFlags2::COPY,
719 vk::AccessFlags2::NONE..vk::AccessFlags2::TRANSFER_WRITE,
720 vk::ImageLayout::TRANSFER_DST_OPTIMAL,
721 0,
722 )?;
723
724 let copy_info = [
725 vk::ImageCopy::default()
726 .src_subresource(vk::ImageSubresourceLayers {
727 base_array_layer: decode_output.layer,
728 mip_level: 0,
729 layer_count: 1,
730 aspect_mask: vk::ImageAspectFlags::PLANE_0,
731 })
732 .src_offset(vk::Offset3D::default())
733 .dst_subresource(vk::ImageSubresourceLayers {
734 base_array_layer: 0,
735 mip_level: 0,
736 layer_count: 1,
737 aspect_mask: vk::ImageAspectFlags::PLANE_0,
738 })
739 .dst_offset(vk::Offset3D::default())
740 .extent(copy_extent),
741 vk::ImageCopy::default()
742 .src_subresource(vk::ImageSubresourceLayers {
743 base_array_layer: decode_output.layer,
744 mip_level: 0,
745 layer_count: 1,
746 aspect_mask: vk::ImageAspectFlags::PLANE_1,
747 })
748 .src_offset(vk::Offset3D::default())
749 .dst_subresource(vk::ImageSubresourceLayers {
750 base_array_layer: 0,
751 mip_level: 0,
752 layer_count: 1,
753 aspect_mask: vk::ImageAspectFlags::PLANE_1,
754 })
755 .dst_offset(vk::Offset3D::default())
756 .extent(vk::Extent3D {
757 width: copy_extent.width / 2,
758 height: copy_extent.height / 2,
759 ..copy_extent
760 }),
761 ];
762
763 unsafe {
764 self.decoding_device.vulkan_device.device.cmd_copy_image(
765 cmd_buffer.buffer(),
766 decode_output.image.image,
767 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
768 *image,
769 vk::ImageLayout::TRANSFER_DST_OPTIMAL,
770 ©_info,
771 );
772 }
773
774 image.transition_layout_single_layer(
775 &mut cmd_buffer,
776 vk::PipelineStageFlags2::COPY..vk::PipelineStageFlags2::NONE,
777 vk::AccessFlags2::TRANSFER_WRITE..vk::AccessFlags2::NONE,
778 vk::ImageLayout::GENERAL,
779 0,
780 )?;
781
782 let semaphore_wait_value = self
783 .decoding_device
784 .queues
785 .transfer
786 .submit_chain_semaphore(
787 cmd_buffer.end()?,
788 &mut self.tracker,
789 vk::PipelineStageFlags2::ALL_COMMANDS,
790 vk::PipelineStageFlags2::ALL_COMMANDS,
791 DecoderTrackerWaitState::DownloadImageToBuffer,
792 )?;
793
794 self.tracker.wait_for(semaphore_wait_value, u64::MAX)?;
795
796 let image = Arc::new(image);
797 let image_clone = image.clone();
798
799 let hal_texture = unsafe {
800 wgpu_device.texture_from_raw(
801 **image,
802 &wgpu::hal::TextureDescriptor {
803 label: Some("vulkan video output texture"),
804 usage: wgpu::TextureUses::RESOURCE
805 | wgpu::TextureUses::COPY_DST
806 | wgpu::TextureUses::COPY_SRC,
807 memory_flags: wgpu::hal::MemoryFlags::empty(),
808 size: wgpu::Extent3d {
809 width: copy_extent.width,
810 height: copy_extent.height,
811 depth_or_array_layers: copy_extent.depth,
812 },
813 dimension: wgpu::TextureDimension::D2,
814 sample_count: 1,
815 view_formats: Vec::new(),
816 format: wgpu::TextureFormat::NV12,
817 mip_level_count: 1,
818 },
819 Some(Box::new(move || {
820 drop(image_clone);
821 })),
822 wgpu::hal::vulkan::TextureMemory::External,
823 )
824 };
825
826 let wgpu_texture = unsafe {
827 self.decoding_device
828 .wgpu_device()
829 .create_texture_from_hal::<wgpu::hal::vulkan::Api>(
830 hal_texture,
831 &wgpu::TextureDescriptor {
832 label: Some("vulkan video output texture"),
833 usage: wgpu::TextureUsages::COPY_DST
834 | wgpu::TextureUsages::TEXTURE_BINDING
835 | wgpu::TextureUsages::COPY_SRC,
836 size: wgpu::Extent3d {
837 width: copy_extent.width,
838 height: copy_extent.height,
839 depth_or_array_layers: copy_extent.depth,
840 },
841 dimension: wgpu::TextureDimension::D2,
842 sample_count: 1,
843 view_formats: &[],
844 format: wgpu::TextureFormat::NV12,
845 mip_level_count: 1,
846 },
847 )
848 };
849
850 Ok(wgpu_texture)
851 }
852
853 fn download_output(
854 &mut self,
855 decode_output: &DecodeSubmissionImageInfo,
856 ) -> Result<Vec<u8>, VulkanDecoderError> {
857 let extent = vk::Extent3D {
858 width: decode_output.cropped_extent.width,
859 height: decode_output.cropped_extent.height,
860 depth: 1,
861 };
862 let (mut dst_buffer, wait_value) =
863 self.copy_image_to_buffer(&decode_output.image, extent, decode_output.layer)?;
864
865 self.tracker.wait_for(wait_value, u64::MAX)?;
866
867 let output = unsafe {
868 dst_buffer
869 .download_data_from_buffer(extent.width as usize * extent.height as usize * 3 / 2)?
870 };
871
872 Ok(output)
873 }
874
875 fn prepare_references_std_ref_info(
876 decode_information: &DecodeInformation,
877 ) -> Vec<vk::native::StdVideoDecodeH264ReferenceInfo> {
878 decode_information
879 .reference_list_l0
880 .iter()
881 .flatten()
882 .chain(decode_information.reference_list_l1.iter().flatten())
883 .map(|&ref_info| ref_info.into())
884 .collect::<Vec<_>>()
885 }
886
887 fn prepare_references_dpb_slot_info(
888 references_std_ref_info: &[vk::native::StdVideoDecodeH264ReferenceInfo],
889 ) -> Vec<vk::VideoDecodeH264DpbSlotInfoKHR<'_>> {
890 references_std_ref_info
891 .iter()
892 .map(|info| vk::VideoDecodeH264DpbSlotInfoKHR::default().std_reference_info(info))
893 .collect::<Vec<_>>()
894 }
895
896 fn prepare_reference_list_slot_info<'b>(
897 reference_id_to_dpb_slot_index: &FxHashMap<ReferenceId, usize>,
898 reference_slots: &'b [vk::VideoReferenceSlotInfoKHR<'b>],
899 references_dpb_slot_info: &'b mut [vk::VideoDecodeH264DpbSlotInfoKHR<'b>],
900 decode_information: &'b DecodeInformation,
901 ) -> Result<Vec<vk::VideoReferenceSlotInfoKHR<'b>>, VulkanDecoderError> {
902 let mut pic_reference_slots: Vec<vk::VideoReferenceSlotInfoKHR<'b>> = Vec::new();
903 for (ref_info, dpb_slot_info) in decode_information
904 .reference_list_l0
905 .iter()
906 .flatten()
907 .chain(decode_information.reference_list_l1.iter().flatten())
908 .zip(references_dpb_slot_info.iter_mut())
909 {
910 let i = *reference_id_to_dpb_slot_index
911 .get(&ref_info.id)
912 .ok_or(VulkanDecoderError::NonExistentReferenceRequested)?;
913
914 let reference = *reference_slots
915 .get(i)
916 .ok_or(VulkanDecoderError::NonExistentReferenceRequested)?;
917
918 if reference.slot_index < 0 || reference.p_picture_resource.is_null() {
919 return Err(VulkanDecoderError::NonExistentReferenceRequested);
920 }
921
922 let reference = reference.push_next(dpb_slot_info);
923
924 if pic_reference_slots
925 .iter()
926 .all(|r| r.slot_index != reference.slot_index)
927 {
928 pic_reference_slots.push(reference);
929 }
930 }
931
932 Ok(pic_reference_slots)
933 }
934
935 fn copy_image_to_buffer(
936 &mut self,
937 image: &Image,
938 dimensions: vk::Extent3D,
939 layer: u32,
940 ) -> Result<(Buffer, SemaphoreWaitValue), VulkanDecoderError> {
941 let mut cmd_buffer = self.tracker.command_buffer_pools.transfer.begin_buffer()?;
942
943 image.transition_layout_single_layer(
944 &mut cmd_buffer,
945 vk::PipelineStageFlags2::NONE..vk::PipelineStageFlags2::COPY,
946 vk::AccessFlags2::NONE..vk::AccessFlags2::TRANSFER_READ,
947 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
948 layer,
949 )?;
950
951 let y_plane_size = dimensions.width as u64 * dimensions.height as u64;
952
953 let dst_buffer = Buffer::new_transfer(
954 self.decoding_device.allocator.clone(),
955 y_plane_size * 3 / 2,
956 TransferDirection::GpuToMem,
957 )?;
958
959 let copy_info = [
960 vk::BufferImageCopy::default()
961 .image_subresource(vk::ImageSubresourceLayers {
962 mip_level: 0,
963 layer_count: 1,
964 base_array_layer: layer,
965 aspect_mask: vk::ImageAspectFlags::PLANE_0,
966 })
967 .image_offset(vk::Offset3D { x: 0, y: 0, z: 0 })
968 .image_extent(vk::Extent3D {
969 width: dimensions.width,
970 height: dimensions.height,
971 depth: 1,
972 })
973 .buffer_offset(0)
974 .buffer_row_length(0)
975 .buffer_image_height(0),
976 vk::BufferImageCopy::default()
977 .image_subresource(vk::ImageSubresourceLayers {
978 mip_level: 0,
979 layer_count: 1,
980 base_array_layer: layer,
981 aspect_mask: vk::ImageAspectFlags::PLANE_1,
982 })
983 .image_offset(vk::Offset3D { x: 0, y: 0, z: 0 })
984 .image_extent(vk::Extent3D {
985 width: dimensions.width / 2,
986 height: dimensions.height / 2,
987 depth: 1,
988 })
989 .buffer_offset(y_plane_size)
990 .buffer_row_length(0)
991 .buffer_image_height(0),
992 ];
993
994 unsafe {
995 self.decoding_device
996 .vulkan_device
997 .device
998 .cmd_copy_image_to_buffer(
999 cmd_buffer.buffer(),
1000 **image,
1001 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
1002 *dst_buffer,
1003 ©_info,
1004 )
1005 };
1006
1007 let wait_value = self
1008 .decoding_device
1009 .queues
1010 .transfer
1011 .submit_chain_semaphore(
1012 cmd_buffer.end()?,
1013 &mut self.tracker,
1014 vk::PipelineStageFlags2::ALL_COMMANDS,
1015 vk::PipelineStageFlags2::ALL_COMMANDS,
1016 DecoderTrackerWaitState::DownloadImageToBuffer,
1017 )?;
1018
1019 Ok((dst_buffer, wait_value))
1020 }
1021}