Skip to main content

vk_video/
vulkan_encoder.rs

1use std::{
2    collections::VecDeque,
3    num::NonZeroU32,
4    sync::{Arc, Mutex},
5};
6
7use ash::vk;
8
9use crate::{
10    EncodedOutputChunk, InputFrame, RawFrameData, VulkanCommonError,
11    device::{ColorRange, ColorSpace, EncodingDevice, Rational},
12    parameters::H264Profile,
13    wrappers::{
14        Buffer, CommandBufferPool, CommandBufferPoolStorage, DecodedPicturesBuffer, Image,
15        ImageLayoutTracker, ImageView, OpenCommandBuffer, ProfileInfo, QueryPool,
16        SemaphoreWaitValue, Tracker, TrackerKind, VideoEncodeQueueExt, VideoQueueExt, VideoSession,
17        VideoSessionParameters, VkPictureParameterSet, VkSequenceParameterSet,
18    },
19};
20
21const MB: u64 = 1024 * 1024;
22
23#[derive(Debug, thiserror::Error)]
24pub enum VulkanEncoderError {
25    #[error("Vulkan error: {0}")]
26    VkError(#[from] ash::vk::Result),
27
28    #[error("Cannot find enough memory of the right type on the device")]
29    NoMemory,
30
31    #[error(transparent)]
32    VulkanCommonError(#[from] VulkanCommonError),
33
34    #[error("The device does not support vulkan h264 encoding")]
35    VulkanEncoderUnsupported,
36
37    #[error(
38        "The byte length of the provided frame ({bytes}) is not the same as the picture size calculated from the dimensions ({size_from_resolution})"
39    )]
40    InconsistentPictureByteSize {
41        bytes: usize,
42        size_from_resolution: usize,
43    },
44
45    #[error("The profile '{0:?}' is not supported by this device")]
46    ProfileUnsupported(H264Profile),
47
48    #[error("This device does not support the required capabilities: {0}")]
49    UnsupportedDeviceCapabilities(&'static str),
50
51    #[error("Encode operation failed with status {0:?}")]
52    EncodeOperationFailed(vk::QueryResultStatusKHR),
53
54    #[error("Invalid encoder parameters, field: {field} - problem: {problem}")]
55    ParametersError {
56        field: &'static str,
57        problem: String,
58    },
59
60    #[error("Framerate numerator * 2 must fit in u32")]
61    FramerateOverflow,
62
63    #[cfg(feature = "wgpu")]
64    #[error(transparent)]
65    WgpuTextureEcoderError(#[from] WgpuTextureEncoderError),
66}
67
68#[cfg(feature = "wgpu")]
69#[derive(Debug, thiserror::Error)]
70pub enum WgpuTextureEncoderError {
71    #[error("The supplied texture's format is {0:?}, when it should be NV12")]
72    NotNV12Texture(wgpu::TextureFormat),
73
74    #[error("The supplied texture does not have COPY_SRC usage. Texture's usages: {0:?}")]
75    NoCopySrcTextureUsage(wgpu::TextureUsages),
76
77    #[error(
78        "The dimensions of the provided frame ({provided_dimensions:?}) are not the same as the expected dimensions ({expected_dimensions:?})"
79    )]
80    InconsistentPictureDimensions {
81        provided_dimensions: wgpu::Extent3d,
82        expected_dimensions: wgpu::Extent3d,
83    },
84
85    #[error("Wgpu device error: {0}")]
86    WgpuDeviceError(#[from] wgpu::hal::DeviceError),
87
88    #[error(transparent)]
89    VulkanCommonError(#[from] VulkanCommonError),
90}
91
92struct VideoSessionResources<'a> {
93    max_dpb_slots: u32,
94    video_session: VideoSession,
95    parameters: VideoSessionParameters,
96    dpb: DecodedPicturesBuffer<'a>,
97    quality_level: u32,
98    rate_control: RateControl,
99    framerate: Rational,
100}
101
102impl VideoSessionResources<'_> {
103    fn new(
104        encoding_device: &EncodingDevice,
105        command_buffer: &mut OpenCommandBuffer,
106        image_tracker: Arc<Mutex<ImageLayoutTracker>>,
107        parameters: FullEncoderParameters,
108        profile_info: &vk::VideoProfileInfoKHR,
109    ) -> Result<Self, VulkanEncoderError> {
110        let encode_capabilities = encoding_device
111            .native_encode_capabilities
112            .profile(parameters.profile)
113            .ok_or(VulkanEncoderError::ProfileUnsupported(parameters.profile))?;
114
115        let extent = vk::Extent2D {
116            width: parameters.width.get(),
117            height: parameters.height.get(),
118        };
119
120        let max_references = parameters.max_references.get();
121        let max_dpb_slots = max_references + 1; // +1 for current picture
122
123        let video_session = VideoSession::new(
124            &encoding_device.vulkan_device,
125            &encoding_device.h264_encode_queues,
126            profile_info,
127            extent,
128            max_dpb_slots,
129            max_references,
130            vk::VideoSessionCreateFlagsKHR::ALLOW_ENCODE_PARAMETER_OPTIMIZATIONS,
131            &encode_capabilities.video_capabilities.std_header_version,
132        )?;
133
134        let use_separate_images = encoding_device
135            .native_encode_capabilities
136            .profile(parameters.profile)
137            .unwrap()
138            .video_capabilities
139            .flags
140            .contains(vk::VideoCapabilityFlagsKHR::SEPARATE_REFERENCE_IMAGES);
141
142        let dpb = DecodedPicturesBuffer::new(
143            encoding_device,
144            command_buffer,
145            image_tracker,
146            use_separate_images,
147            profile_info,
148            vk::ImageUsageFlags::VIDEO_ENCODE_DPB_KHR,
149            &encode_capabilities.encode_dpb_properties[0],
150            extent,
151            max_dpb_slots,
152            None,
153            vk::ImageLayout::VIDEO_ENCODE_DPB_KHR,
154        )?;
155
156        let sps = VkSequenceParameterSet::new_encode(
157            parameters.profile,
158            extent.width,
159            extent.height,
160            max_references,
161            parameters.color_space,
162            parameters.color_range,
163            parameters.framerate,
164        )?;
165        let pps = VkPictureParameterSet::new_encode();
166
167        let session_parameters = VideoSessionParameters::new(
168            encoding_device.vulkan_device.device.clone(),
169            video_session.session,
170            &[sps.sps],
171            &[pps.pps],
172            None,
173            Some(parameters.quality_level),
174        )?;
175
176        Ok(Self {
177            video_session,
178            dpb,
179            parameters: session_parameters,
180            max_dpb_slots,
181            quality_level: parameters.quality_level,
182            rate_control: RateControl::EncoderDefault,
183            framerate: parameters.framerate,
184        })
185    }
186}
187
188pub(crate) type H264EncodeProfileInfo<'a> = ProfileInfo<'a>;
189
190impl H264EncodeProfileInfo<'_> {
191    pub(crate) fn new_encode(parameters: &FullEncoderParameters) -> Self {
192        let h264_profile = vk::VideoEncodeH264ProfileInfoKHR::default()
193            .std_profile_idc(parameters.profile.to_profile_idc());
194
195        let profile = vk::VideoProfileInfoKHR::default()
196            .video_codec_operation(vk::VideoCodecOperationFlagsKHR::ENCODE_H264)
197            .chroma_subsampling(vk::VideoChromaSubsamplingFlagsKHR::TYPE_420)
198            .luma_bit_depth(vk::VideoComponentBitDepthFlagsKHR::TYPE_8)
199            .chroma_bit_depth(vk::VideoComponentBitDepthFlagsKHR::TYPE_8);
200
201        let h264_profile = Box::new(h264_profile);
202
203        let usage_info = vk::VideoEncodeUsageInfoKHR::default()
204            .video_usage_hints(parameters.usage_flags)
205            .tuning_mode(parameters.tuning_mode)
206            .video_content_hints(parameters.content_flags);
207
208        let usage_info = Box::new(usage_info);
209
210        ProfileInfo::new(profile, vec![h264_profile, usage_info])
211    }
212}
213
214struct EncodingQueryPool {
215    pool: QueryPool,
216}
217
218impl std::ops::Deref for EncodingQueryPool {
219    type Target = QueryPool;
220
221    fn deref(&self) -> &Self::Target {
222        &self.pool
223    }
224}
225
226impl EncodingQueryPool {
227    pub(crate) fn new(
228        encoding_device: &EncodingDevice,
229        profile: H264Profile,
230        profile_info: vk::VideoProfileInfoKHR,
231    ) -> Result<Self, VulkanEncoderError> {
232        let encode_capabilities = encoding_device
233            .native_encode_capabilities
234            .profile(profile)
235            .ok_or(VulkanEncoderError::ProfileUnsupported(profile))?;
236
237        if !encode_capabilities
238            .encode_capabilities
239            .supported_encode_feedback_flags
240            .contains(vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN)
241        {
242            return Err(VulkanEncoderError::UnsupportedDeviceCapabilities(
243                "VkVideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN",
244            ));
245        }
246
247        let pool = QueryPool::new(
248            encoding_device.vulkan_device.device.clone(),
249            vk::QueryType::VIDEO_ENCODE_FEEDBACK_KHR,
250            1,
251            Some(profile_info),
252            Some(
253                vk::QueryPoolVideoEncodeFeedbackCreateInfoKHR::default().encode_feedback_flags(
254                    vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BYTES_WRITTEN
255                        | vk::VideoEncodeFeedbackFlagsKHR::BITSTREAM_BUFFER_OFFSET,
256                ),
257            ),
258        )?;
259
260        Ok(Self { pool })
261    }
262
263    pub(crate) fn get_result_blocking(&self) -> Result<EncodeFeedback, VulkanEncoderError> {
264        let mut result = [EncodeFeedback {
265            offset: 0,
266            bytes_written: 0,
267            status: vk::QueryResultStatusKHR::NOT_READY,
268        }];
269
270        unsafe {
271            self.pool.device.get_query_pool_results(
272                self.pool.pool,
273                0,
274                &mut result,
275                vk::QueryResultFlags::WAIT | vk::QueryResultFlags::WITH_STATUS_KHR,
276            )?
277        };
278
279        Ok(result[0])
280    }
281}
282
283#[repr(C)]
284#[derive(Debug, Clone, Copy)]
285struct EncodeFeedback {
286    offset: u32,
287    bytes_written: u32,
288    status: vk::QueryResultStatusKHR,
289}
290
291pub(crate) enum EncoderTrackerWaitState {
292    InitializeEncoder,
293    #[cfg_attr(not(feature = "transcoder"), allow(dead_code))]
294    ResizeInput,
295    CopyBufferToImage,
296    #[cfg_attr(not(feature = "wgpu"), allow(dead_code))]
297    CopyImageToImage,
298    Encode,
299}
300
301pub(crate) struct EncoderCommandBufferPools {
302    transfer: CommandBufferPool,
303    encode: CommandBufferPool,
304}
305
306impl EncoderCommandBufferPools {
307    fn new(device: &EncodingDevice) -> Result<Self, VulkanEncoderError> {
308        let transfer = CommandBufferPool::new(
309            device.vulkan_device.clone(),
310            device.queues.transfer.family_index,
311        )?;
312        let encode = CommandBufferPool::new(
313            device.vulkan_device.clone(),
314            device.h264_encode_queues.family_index,
315        )?;
316
317        Ok(Self { transfer, encode })
318    }
319}
320
321impl CommandBufferPoolStorage for EncoderCommandBufferPools {
322    fn mark_submitted_as_free(&mut self, last_waited_for: SemaphoreWaitValue) {
323        self.transfer.mark_submitted_as_free(last_waited_for);
324        self.encode.mark_submitted_as_free(last_waited_for);
325    }
326}
327
328pub(crate) struct EncoderTrackerKind {}
329
330impl TrackerKind for EncoderTrackerKind {
331    type WaitState = EncoderTrackerWaitState;
332
333    type CommandBufferPools = EncoderCommandBufferPools;
334}
335
336pub(crate) type EncoderTracker = Tracker<EncoderTrackerKind>;
337
338pub(crate) struct EncodeSubmission<'borrow, 'encoder> {
339    pub(crate) is_idr: bool,
340    pub(crate) wait_value: SemaphoreWaitValue,
341    pub(crate) encoder: &'borrow mut VulkanEncoder<'encoder>,
342    pub(crate) pts: Option<u64>,
343    _image: Arc<Image>,
344}
345
346impl<'a, 'b> EncodeSubmission<'a, 'b> {
347    pub(crate) fn download(self) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
348        self.encoder.download_output(self.is_idr, self.pts)
349    }
350
351    #[cfg_attr(not(feature = "transcoder"), allow(dead_code))]
352    pub(crate) fn mark_waited(&mut self) {
353        self.encoder.tracker.mark_waited(self.wait_value);
354    }
355
356    pub(crate) fn wait(&mut self, timeout: u64) -> Result<(), VulkanEncoderError> {
357        self.encoder.tracker.wait_for(self.wait_value, timeout)?;
358        Ok(())
359    }
360}
361
362pub(crate) struct UnwaitedEncodeSubmission<'a, 'b>(pub(crate) EncodeSubmission<'a, 'b>);
363
364impl<'a, 'b> UnwaitedEncodeSubmission<'a, 'b> {
365    #[cfg_attr(not(feature = "transcoder"), allow(dead_code))]
366    pub(crate) fn mark_waited(mut self) -> WaitedEncodeSubmission<'a, 'b> {
367        self.0.mark_waited();
368        WaitedEncodeSubmission(self.0)
369    }
370
371    pub(crate) fn wait(
372        mut self,
373        timeout: u64,
374    ) -> Result<WaitedEncodeSubmission<'a, 'b>, VulkanEncoderError> {
375        self.0.wait(timeout)?;
376        Ok(WaitedEncodeSubmission(self.0))
377    }
378
379    pub(crate) fn wait_and_download(
380        self,
381        timeout: u64,
382    ) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
383        let waited = self.wait(timeout)?;
384        waited.download()
385    }
386}
387
388pub struct WaitedEncodeSubmission<'a, 'b>(pub(crate) EncodeSubmission<'a, 'b>);
389
390impl<'a, 'b> WaitedEncodeSubmission<'a, 'b> {
391    pub(crate) fn download(self) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
392        self.0.download()
393    }
394}
395
396pub struct VulkanEncoder<'a> {
397    pub(crate) tracker: EncoderTracker,
398    query_pool: EncodingQueryPool,
399    profile: H264Profile,
400    pub(crate) profile_info: H264EncodeProfileInfo<'a>,
401    session_resources: VideoSessionResources<'a>,
402    idr_period_counter: u32,
403    idr_period: u32,
404    #[allow(dead_code)]
405    input_image: Arc<Image>,
406    output_buffer: Buffer,
407    idr_pic_id: u16,
408    frame_num: u32,
409    pic_order_cnt: u8,
410    active_reference_slots: VecDeque<(usize, vk::native::StdVideoEncodeH264ReferenceInfo)>,
411    rate_control: RateControl,
412    inline_stream_params: bool,
413    encoding_device: Arc<EncodingDevice>,
414}
415
416#[derive(Debug, Clone, Copy)]
417pub struct FullEncoderParameters {
418    pub(crate) idr_period: NonZeroU32,
419    pub(crate) width: NonZeroU32,
420    pub(crate) height: NonZeroU32,
421    pub(crate) rate_control: RateControl,
422    pub(crate) max_references: NonZeroU32,
423    pub(crate) profile: H264Profile,
424    pub(crate) quality_level: u32,
425    pub(crate) framerate: Rational,
426    pub(crate) usage_flags: vk::VideoEncodeUsageFlagsKHR,
427    pub(crate) tuning_mode: vk::VideoEncodeTuningModeKHR,
428    pub(crate) content_flags: vk::VideoEncodeContentFlagsKHR,
429    pub(crate) inline_stream_params: bool,
430    pub(crate) color_space: ColorSpace,
431    pub(crate) color_range: ColorRange,
432}
433
434impl<'a> VulkanEncoder<'a> {
435    const OUTPUT_BUFFER_LEN: u64 = 4 * MB;
436
437    pub(crate) fn new(
438        encoding_device: Arc<EncodingDevice>,
439        parameters: FullEncoderParameters,
440    ) -> Result<Self, VulkanEncoderError> {
441        let profile_info = H264EncodeProfileInfo::new_encode(&parameters);
442
443        let command_buffer_pools = EncoderCommandBufferPools::new(&encoding_device)?;
444        let mut tracker = EncoderTracker::new(
445            encoding_device.device.clone(),
446            command_buffer_pools,
447            Some("encoder"),
448        )?;
449
450        let query_pool = EncodingQueryPool::new(
451            &encoding_device,
452            parameters.profile,
453            profile_info.profile_info,
454        )?;
455
456        // TODO: this buffer should grow when necessary
457        let output_buffer = Buffer::new_encode(
458            encoding_device.allocator.clone(),
459            Self::OUTPUT_BUFFER_LEN,
460            &profile_info,
461        )?;
462
463        let mut buffer = tracker.command_buffer_pools.encode.begin_buffer()?;
464
465        let session_resources = VideoSessionResources::new(
466            &encoding_device,
467            &mut buffer,
468            tracker.image_layout_tracker.clone(),
469            parameters,
470            &profile_info.profile_info,
471        )?;
472
473        encoding_device.h264_encode_queues.submit_chain_semaphore(
474            buffer.end()?,
475            &mut tracker,
476            vk::PipelineStageFlags2::ALL_COMMANDS,
477            vk::PipelineStageFlags2::ALL_COMMANDS,
478            EncoderTrackerWaitState::InitializeEncoder,
479        )?;
480
481        let encode_image = Image::new_encode(
482            &encoding_device,
483            session_resources.video_session.max_coded_extent.into(),
484            &profile_info,
485            encoding_device.queues.wgpu.family_index as u32,
486            tracker.image_layout_tracker.clone(),
487        )?;
488
489        Ok(Self {
490            idr_period_counter: 0,
491            idr_pic_id: 0,
492            frame_num: 0,
493            pic_order_cnt: 0,
494            active_reference_slots: VecDeque::with_capacity(session_resources.dpb.len as usize),
495            profile: parameters.profile,
496            profile_info,
497            encoding_device,
498            input_image: Arc::new(encode_image),
499            tracker,
500            query_pool,
501            session_resources,
502            idr_period: parameters.idr_period.get(),
503            output_buffer,
504            rate_control: parameters.rate_control,
505            inline_stream_params: parameters.inline_stream_params,
506        })
507    }
508
509    fn begin_video_coding(&self, buffer: vk::CommandBuffer) {
510        let mut h264_layers =
511            self.h264_rate_control_layers_for(self.session_resources.rate_control);
512        let layers = self.rate_control_layers_for(
513            self.session_resources.rate_control,
514            h264_layers.as_mut().map(|o| &mut o[..]),
515        );
516        let mut h264_rate_control = self.h264_rate_control(layers.as_ref().map(|o| &o[..]));
517        let mut encode_rate_control = self.encoder_rate_control_for(
518            self.session_resources.rate_control,
519            layers.as_ref().map(|o| &o[..]),
520        );
521
522        let mut reference_slot_info = self.session_resources.dpb.reference_slot_info();
523        reference_slot_info.sort_by_key(|s| {
524            if s.slot_index == -1 {
525                return usize::MAX;
526            }
527
528            let (i, _) = self
529                .active_reference_slots
530                .iter()
531                .enumerate()
532                .find(|(_, (slot_idx, _))| (*slot_idx) as i32 == s.slot_index)
533                .unwrap();
534
535            i
536        });
537
538        // Absolutely crucial for nvidia GPUs, nothing works without this.
539        reference_slot_info.reverse();
540
541        let mut begin_info = vk::VideoBeginCodingInfoKHR::default()
542            .video_session(self.session_resources.video_session.session)
543            .video_session_parameters(self.session_resources.parameters.parameters)
544            .reference_slots(&reference_slot_info);
545
546        if let (Some(encode_rate_control), Some(h264_rate_control)) =
547            (encode_rate_control.as_mut(), h264_rate_control.as_mut())
548        {
549            begin_info = begin_info
550                .push_next(encode_rate_control)
551                .push_next(h264_rate_control);
552        }
553
554        unsafe {
555            self.encoding_device
556                .vulkan_device
557                .device
558                .video_queue_ext
559                .cmd_begin_video_coding_khr(buffer, &begin_info);
560        }
561    }
562
563    fn issue_coding_control_reset_for(
564        &mut self,
565        buffer: vk::CommandBuffer,
566        rate_control: RateControl,
567    ) {
568        let mut quality_level = vk::VideoEncodeQualityLevelInfoKHR::default()
569            .quality_level(self.session_resources.quality_level);
570
571        let mut h264_layers = self.h264_rate_control_layers_for(rate_control);
572        let layers =
573            self.rate_control_layers_for(rate_control, h264_layers.as_mut().map(|o| &mut o[..]));
574        let mut h264_rate_control = self.h264_rate_control(layers.as_ref().map(|o| &o[..]));
575        let mut encode_rate_control =
576            self.encoder_rate_control_for(rate_control, layers.as_ref().map(|o| &o[..]));
577
578        let mut flags = vk::VideoCodingControlFlagsKHR::RESET
579            | vk::VideoCodingControlFlagsKHR::ENCODE_QUALITY_LEVEL;
580
581        if encode_rate_control.is_some() {
582            flags |= vk::VideoCodingControlFlagsKHR::ENCODE_RATE_CONTROL;
583        }
584
585        let mut control_info = vk::VideoCodingControlInfoKHR::default()
586            .flags(flags)
587            .push_next(&mut quality_level);
588
589        if let (Some(encode_rate_control), Some(h264_rate_control)) =
590            (encode_rate_control.as_mut(), h264_rate_control.as_mut())
591        {
592            control_info = control_info
593                .push_next(h264_rate_control)
594                .push_next(encode_rate_control);
595        }
596
597        unsafe {
598            self.encoding_device
599                .vulkan_device
600                .device
601                .video_queue_ext
602                .cmd_control_video_coding_khr(buffer, &control_info);
603        }
604
605        self.session_resources.rate_control = rate_control;
606    }
607
608    // TODO: Maybe we should reuse `input_image` here, instead of creating a new image
609    fn transfer_buffer_to_image(
610        &mut self,
611        frame: &InputFrame<RawFrameData>,
612    ) -> Result<(Image, Buffer), VulkanEncoderError> {
613        let extent = vk::Extent3D {
614            width: frame.data.width,
615            height: frame.data.height,
616            depth: 1,
617        };
618
619        if frame.data.width as usize * frame.data.height as usize * 3 / 2 != frame.data.frame.len()
620        {
621            return Err(VulkanEncoderError::InconsistentPictureByteSize {
622                bytes: frame.data.frame.len(),
623                size_from_resolution: frame.data.width as usize * frame.data.height as usize * 3
624                    / 2,
625            });
626        }
627
628        let mut profile_list_info = vk::VideoProfileListInfoKHR::default()
629            .profiles(std::slice::from_ref(&self.profile_info.profile_info));
630
631        let queue_family_indices = [
632            self.encoding_device.queues.transfer.family_index as u32,
633            self.encoding_device.h264_encode_queues.family_index as u32,
634        ];
635
636        let image_create_info = vk::ImageCreateInfo::default()
637            .flags(vk::ImageCreateFlags::empty())
638            .image_type(vk::ImageType::TYPE_2D)
639            .format(vk::Format::G8_B8R8_2PLANE_420_UNORM)
640            .extent(extent)
641            .mip_levels(1)
642            .array_layers(1)
643            .samples(vk::SampleCountFlags::TYPE_1)
644            .tiling(vk::ImageTiling::OPTIMAL)
645            .usage(vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR | vk::ImageUsageFlags::TRANSFER_DST)
646            .sharing_mode(vk::SharingMode::CONCURRENT)
647            .initial_layout(vk::ImageLayout::UNDEFINED)
648            .queue_family_indices(&queue_family_indices)
649            .push_next(&mut profile_list_info);
650
651        let image = Image::new(
652            self.encoding_device.allocator.clone(),
653            &image_create_info,
654            self.tracker.image_layout_tracker.clone(),
655        )?;
656
657        let mut cmd_buffer = self.tracker.command_buffer_pools.transfer.begin_buffer()?;
658
659        image.transition_layout_single_layer(
660            &mut cmd_buffer,
661            vk::PipelineStageFlags2::NONE..vk::PipelineStageFlags2::COPY,
662            vk::AccessFlags2::NONE..vk::AccessFlags2::TRANSFER_WRITE,
663            vk::ImageLayout::TRANSFER_DST_OPTIMAL,
664            0,
665        )?;
666
667        let buffer = Buffer::new_transfer_with_data(
668            self.encoding_device.allocator.clone(),
669            &frame.data.frame,
670        )?;
671
672        unsafe {
673            self.encoding_device
674                .vulkan_device
675                .device
676                .cmd_copy_buffer_to_image(
677                    cmd_buffer.buffer(),
678                    *buffer,
679                    *image,
680                    vk::ImageLayout::TRANSFER_DST_OPTIMAL,
681                    &[
682                        vk::BufferImageCopy::default()
683                            .buffer_offset(0)
684                            .buffer_row_length(0)
685                            .buffer_image_height(0)
686                            .image_subresource(vk::ImageSubresourceLayers {
687                                aspect_mask: vk::ImageAspectFlags::PLANE_0,
688                                layer_count: 1,
689                                base_array_layer: 0,
690                                mip_level: 0,
691                            })
692                            .image_offset(vk::Offset3D { x: 0, y: 0, z: 0 })
693                            .image_extent(vk::Extent3D {
694                                width: frame.data.width,
695                                height: frame.data.height,
696                                depth: 1,
697                            }),
698                        vk::BufferImageCopy::default()
699                            .buffer_offset(frame.data.width as u64 * frame.data.height as u64)
700                            .buffer_row_length(0)
701                            .buffer_image_height(0)
702                            .image_subresource(vk::ImageSubresourceLayers {
703                                aspect_mask: vk::ImageAspectFlags::PLANE_1,
704                                layer_count: 1,
705                                base_array_layer: 0,
706                                mip_level: 0,
707                            })
708                            .image_offset(vk::Offset3D { x: 0, y: 0, z: 0 })
709                            .image_extent(vk::Extent3D {
710                                width: frame.data.width / 2,
711                                height: frame.data.height / 2,
712                                depth: 1,
713                            }),
714                    ],
715                );
716        }
717
718        self.encoding_device
719            .queues
720            .transfer
721            .submit_chain_semaphore(
722                cmd_buffer.end()?,
723                &mut self.tracker,
724                vk::PipelineStageFlags2::COPY,
725                vk::PipelineStageFlags2::COPY,
726                EncoderTrackerWaitState::CopyBufferToImage,
727            )?;
728
729        Ok((image, buffer))
730    }
731
732    #[cfg(feature = "wgpu")]
733    fn copy_wgpu_texture_to_image(
734        &mut self,
735        frame: &InputFrame<wgpu::Texture>,
736    ) -> Result<wgpu::hal::vulkan::CommandEncoder, WgpuTextureEncoderError> {
737        use wgpu::hal::{CommandEncoder, Device, Queue, vulkan::Api as VkApi};
738
739        let encode_texture_extent = wgpu::Extent3d {
740            width: self.input_image.extent.width,
741            height: self.input_image.extent.height,
742            depth_or_array_layers: self.input_image.extent.depth,
743        };
744
745        if !frame.data.usage().contains(wgpu::TextureUsages::COPY_SRC) {
746            return Err(WgpuTextureEncoderError::NoCopySrcTextureUsage(
747                frame.data.usage(),
748            ));
749        }
750        if frame.data.format() != wgpu::TextureFormat::NV12 {
751            return Err(WgpuTextureEncoderError::NotNV12Texture(frame.data.format()));
752        }
753        if frame.data.size() != encode_texture_extent {
754            return Err(WgpuTextureEncoderError::InconsistentPictureDimensions {
755                provided_dimensions: frame.data.size(),
756                expected_dimensions: encode_texture_extent,
757            });
758        }
759
760        let hal_device = unsafe {
761            self.encoding_device
762                .wgpu_device()
763                .as_hal::<VkApi>()
764                .unwrap()
765        };
766        let hal_queue = unsafe { self.encoding_device.wgpu_queue().as_hal::<VkApi>().unwrap() };
767
768        let input_image_clone = self.input_image.clone();
769        let hal_texture = unsafe {
770            hal_device.texture_from_raw(
771                self.input_image.image,
772                &wgpu::hal::TextureDescriptor {
773                    label: None,
774                    size: encode_texture_extent,
775                    mip_level_count: 1,
776                    sample_count: 1,
777                    dimension: wgpu::TextureDimension::D2,
778                    format: wgpu::TextureFormat::NV12,
779                    usage: wgpu::TextureUses::COPY_DST,
780                    memory_flags: wgpu::hal::MemoryFlags::empty(),
781                    view_formats: Vec::new(),
782                },
783                Some(Box::new(move || {
784                    drop(input_image_clone);
785                })),
786                wgpu::hal::vulkan::TextureMemory::External,
787            )
788        };
789
790        let texture = unsafe {
791            self.encoding_device
792                .wgpu_device()
793                .create_texture_from_hal::<VkApi>(
794                    hal_texture,
795                    &wgpu::TextureDescriptor {
796                        label: None,
797                        size: encode_texture_extent,
798                        mip_level_count: 1,
799                        sample_count: 1,
800                        dimension: wgpu::TextureDimension::D2,
801                        format: wgpu::TextureFormat::NV12,
802                        usage: wgpu::TextureUsages::COPY_DST,
803                        view_formats: &[],
804                    },
805                )
806        };
807
808        // Copy is on the wgpu core queue because it will handle `frame.data` layout transitions for us
809        let mut encoder = self
810            .encoding_device
811            .wgpu_device()
812            .create_command_encoder(&Default::default());
813        encoder.copy_texture_to_texture(
814            frame.data.as_image_copy(),
815            texture.as_image_copy(),
816            encode_texture_extent,
817        );
818
819        // TODO: dont wait for all here
820        self.tracker.wait_for_all(u64::MAX)?;
821        self.encoding_device.wgpu_queue().submit([encoder.finish()]);
822
823        self.tracker
824            .image_layout_tracker
825            .lock()
826            .unwrap()
827            .map
828            .insert(
829                self.input_image.key(),
830                vec![vk::ImageLayout::TRANSFER_DST_OPTIMAL].into_boxed_slice(),
831            );
832
833        // wgpu core queue makes it impossible to specify signal semaphores
834        // so we have to make an empty submit on the wgpu hal queue just for the synchronization
835        //
836        // TODO: it'd be better to create one encoder and just reuse it
837        //       because it creates a new command pool every time it's created
838        let mut hal_encoder = unsafe {
839            hal_device
840                .create_command_encoder(&wgpu::hal::CommandEncoderDescriptor {
841                    label: Some("vulkan video synchronize with wgpu"),
842                    queue: &hal_queue,
843                })
844                .unwrap()
845        };
846        let command_buffer = unsafe {
847            hal_encoder.begin_encoding(None)?;
848            hal_encoder.end_encoding()?
849        };
850
851        let mut semaphore_submit_info = self
852            .tracker
853            .semaphore_tracker
854            .next_submit_info(EncoderTrackerWaitState::CopyImageToImage);
855        unsafe {
856            hal_queue.submit(
857                &[&command_buffer],
858                &[],
859                semaphore_submit_info.wgpu_wait_info(),
860            )?;
861        }
862
863        semaphore_submit_info.mark_submitted();
864
865        Ok(hal_encoder)
866    }
867
868    pub(crate) fn encode<'b>(
869        &'b mut self,
870        image: Arc<Image>,
871        force_idr: bool,
872        pts: Option<u64>,
873    ) -> Result<UnwaitedEncodeSubmission<'b, 'a>, VulkanEncoderError> {
874        let is_idr = force_idr || self.idr_period_counter == 0;
875        let mut idr_pic_id = 0;
876
877        if is_idr {
878            self.idr_period_counter = 0;
879            idr_pic_id = self.idr_pic_id;
880            self.idr_pic_id = self.idr_pic_id.wrapping_add(1);
881            self.frame_num = 0;
882            self.pic_order_cnt = 0;
883            self.active_reference_slots.clear();
884            self.session_resources.dpb.reset_all_allocations();
885        } else if self.active_reference_slots.len() == self.session_resources.max_dpb_slots as usize
886        {
887            if let Some((oldest_reference, _)) = self.active_reference_slots.pop_front() {
888                self.session_resources
889                    .dpb
890                    .free_reference_picture(oldest_reference);
891            }
892        }
893
894        let mut cmd_buffer = self.tracker.command_buffer_pools.encode.begin_buffer()?;
895
896        image.transition_layout_single_layer(
897            &mut cmd_buffer,
898            vk::PipelineStageFlags2::NONE..vk::PipelineStageFlags2::VIDEO_ENCODE_KHR,
899            vk::AccessFlags2::NONE..vk::AccessFlags2::VIDEO_ENCODE_READ_KHR,
900            vk::ImageLayout::VIDEO_ENCODE_SRC_KHR,
901            0,
902        )?;
903
904        let mut view_usage_create_info = vk::ImageViewUsageCreateInfo::default()
905            .usage(vk::ImageUsageFlags::VIDEO_ENCODE_SRC_KHR);
906
907        let view_create_info = vk::ImageViewCreateInfo::default()
908            .flags(vk::ImageViewCreateFlags::empty())
909            .image(image.image)
910            .view_type(vk::ImageViewType::TYPE_2D)
911            .format(vk::Format::G8_B8R8_2PLANE_420_UNORM)
912            .components(vk::ComponentMapping::default())
913            .subresource_range(vk::ImageSubresourceRange {
914                aspect_mask: vk::ImageAspectFlags::COLOR,
915                level_count: 1,
916                base_mip_level: 0,
917                layer_count: 1,
918                base_array_layer: 0,
919            })
920            .push_next(&mut view_usage_create_info);
921
922        let view = ImageView::new(
923            self.encoding_device.vulkan_device.device.clone(),
924            image.clone(),
925            &view_create_info,
926        )?;
927
928        self.query_pool.reset(cmd_buffer.buffer());
929
930        self.begin_video_coding(cmd_buffer.buffer());
931
932        if is_idr {
933            // TODO: controllable rate control, framerate and all stream parameters
934            self.issue_coding_control_reset_for(cmd_buffer.buffer(), self.rate_control);
935        }
936
937        let frame_num = self.frame_num;
938        self.frame_num = self.frame_num.wrapping_add(1);
939
940        let pic_order_cnt = self.pic_order_cnt;
941        self.pic_order_cnt = self.pic_order_cnt.wrapping_add(2);
942
943        // bugs in nvidia driver I encountered on this journey:
944        //
945        // bug1: if primary pic type is set to I instead of IDR, the encode command will submit
946        // successfully, the fence will trigger, signalling it has been executed, but if you then
947        // query the implementation for the status of the operation, it will behave as though the
948        // operation never happened (which means it will not return an error!). The division
949        // between I and IDR is invented in the vulkan spec, in h264 the values are equivalent.
950        //
951        // bug2: when rate control is disabled, you have to specify the temporal layer count to 0.
952        // You pass a table length and a pointer to a bable with temporal layer descriptions. Even
953        // when the length is set to 0, the pointer will be dereferenced. If you set it to NULL,
954        // the program will (obviously) segfault.
955        //
956        // bug3: each dpb reference picture has to be in a separate VkImage, even though the spec
957        // says these can be different layers of the same image (even though using layers of one
958        // picture works in the decoder)
959        //
960        // bug4: when you pass the information about which decoded pictures buffer slots contain
961        // references, the spec does not specify the order in which they should be arranged. The
962        // internal implementation expects a very specific order though: from the most recent to
963        // the oldest. It was natural for me to keep references in a FIFO queue, where I append
964        // new pictures to the back and pop old pictures from the front when they're no longer
965        // needed. After hours of trying to figure out what the problem was I jokingly said to a
966        // colleague that we should just try reversing the order we have. It ended up working.
967        // I don't know how anyone is supposed to find this.
968
969        let primary_pic_type = if is_idr {
970            vk::native::StdVideoH264PictureType_STD_VIDEO_H264_PICTURE_TYPE_IDR
971        } else {
972            vk::native::StdVideoH264PictureType_STD_VIDEO_H264_PICTURE_TYPE_P
973        };
974
975        let slice_header = vk::native::StdVideoEncodeH264SliceHeader {
976            flags: vk::native::StdVideoEncodeH264SliceHeaderFlags {
977                _bitfield_align_1: [],
978                _bitfield_1: vk::native::StdVideoEncodeH264SliceHeaderFlags::new_bitfield_1(
979                    1, // TODO: b-frames
980                    1, // TODO: don't override always
981                    0,
982                ),
983            },
984            first_mb_in_slice: 0,
985            slice_type: if is_idr {
986                vk::native::StdVideoH264SliceType_STD_VIDEO_H264_SLICE_TYPE_I
987            } else {
988                vk::native::StdVideoH264SliceType_STD_VIDEO_H264_SLICE_TYPE_P
989            }, // TODO: b-frames
990            slice_alpha_c0_offset_div2: 0,
991            slice_beta_offset_div2: 0,
992            slice_qp_delta: 0, // TODO: check whether this will be overwritten in the bitstream
993            reserved1: 0,
994            cabac_init_idc: vk::native::StdVideoH264CabacInitIdc_STD_VIDEO_H264_CABAC_INIT_IDC_0, // TODO: check whether this will be overwritten in the bitstream
995            disable_deblocking_filter_idc: 0, // TODO: enable for fast decoding?
996            pWeightTable: std::ptr::null(),
997        };
998
999        let mut nalu_slice_entries =
1000            [vk::VideoEncodeH264NaluSliceInfoKHR::default().std_slice_header(&slice_header)];
1001
1002        if let RateControl::Disabled = self.rate_control {
1003            if let Some(caps) = self
1004                .encoding_device
1005                .native_encode_capabilities
1006                .profile(self.profile)
1007            {
1008                let quality_properties =
1009                    &caps.quality_level_properties[self.session_resources.quality_level as usize];
1010
1011                if !quality_properties.zeroed() {
1012                    let qp = quality_properties
1013                        .h264_quality_level_properties
1014                        .preferred_constant_qp;
1015
1016                    if is_idr {
1017                        nalu_slice_entries[0].constant_qp = qp.qp_i;
1018                    } else {
1019                        nalu_slice_entries[0].constant_qp = qp.qp_p;
1020                    }
1021                }
1022            }
1023        }
1024
1025        let mut ref_list0 = [0xff; 32];
1026        for (i, (slot, _)) in self.active_reference_slots.iter().rev().enumerate() {
1027            ref_list0[i] = *slot as u8;
1028        }
1029
1030        let ref_lists = vk::native::StdVideoEncodeH264ReferenceListsInfo {
1031            flags: vk::native::StdVideoEncodeH264ReferenceListsInfoFlags {
1032                _bitfield_align_1: [],
1033                _bitfield_1: vk::native::StdVideoEncodeH264ReferenceListsInfoFlags::new_bitfield_1(
1034                    0, 0, 0,
1035                ),
1036            },
1037            num_ref_idx_l0_active_minus1: self.active_reference_slots.len().saturating_sub(1) as u8,
1038            num_ref_idx_l1_active_minus1: 0,
1039            RefPicList0: ref_list0,
1040            RefPicList1: [0xff; 32],
1041            refList0ModOpCount: 0,
1042            refList1ModOpCount: 0,
1043            refPicMarkingOpCount: 0,
1044            reserved1: [0; 7],
1045            pRefList0ModOperations: std::ptr::null(),
1046            pRefList1ModOperations: std::ptr::null(),
1047            pRefPicMarkingOperations: std::ptr::null(),
1048        };
1049
1050        let std_h264_encode_info = vk::native::StdVideoEncodeH264PictureInfo {
1051            flags: vk::native::StdVideoEncodeH264PictureInfoFlags {
1052                _bitfield_align_1: [],
1053                _bitfield_1: vk::native::StdVideoEncodeH264PictureInfoFlags::new_bitfield_1(
1054                    is_idr as u32,
1055                    1, // TODO
1056                    is_idr as u32,
1057                    0, // long term refs
1058                    0, // adaptive reference control
1059                    0,
1060                ),
1061            },
1062            seq_parameter_set_id: 0,
1063            pic_parameter_set_id: 0,
1064            idr_pic_id,
1065            primary_pic_type,
1066            frame_num,
1067            PicOrderCnt: pic_order_cnt as i32,
1068            temporal_id: 0,
1069            reserved1: [0; 3],
1070            pRefLists: &ref_lists,
1071        };
1072
1073        let mut h264_encode_info = vk::VideoEncodeH264PictureInfoKHR::default()
1074            .nalu_slice_entries(&nalu_slice_entries)
1075            .generate_prefix_nalu(false)
1076            .std_picture_info(&std_h264_encode_info);
1077
1078        let setup_reference_slot_idx = self.session_resources.dpb.allocate_reference_picture()?;
1079
1080        let mut reference_slots = self
1081            .session_resources
1082            .dpb
1083            .reference_slot_info()
1084            .into_iter()
1085            .filter(|i| i.slot_index >= 0 && i.slot_index != setup_reference_slot_idx as i32)
1086            .collect::<Vec<_>>();
1087
1088        let mut std_reference_info = self
1089            .active_reference_slots
1090            .iter()
1091            .rev()
1092            .map(|(i, info)| {
1093                (
1094                    *i,
1095                    vk::VideoEncodeH264DpbSlotInfoKHR::default().std_reference_info(info),
1096                )
1097            })
1098            .collect::<Vec<_>>();
1099
1100        std_reference_info.iter_mut().for_each(|(i, std_info)| {
1101            let slot = reference_slots
1102                .iter_mut()
1103                .find(|reference_slot| reference_slot.slot_index == (*i) as i32)
1104                .unwrap();
1105            *slot = slot.push_next(std_info);
1106        });
1107
1108        let std_new_slot_reference_info = vk::native::StdVideoEncodeH264ReferenceInfo {
1109            flags: vk::native::StdVideoEncodeH264ReferenceInfoFlags {
1110                _bitfield_align_1: [],
1111                _bitfield_1: vk::native::StdVideoEncodeH264ReferenceInfoFlags::new_bitfield_1(0, 0),
1112            },
1113            primary_pic_type,
1114            FrameNum: frame_num,
1115            PicOrderCnt: pic_order_cnt as i32,
1116            long_term_pic_num: 0,
1117            long_term_frame_idx: 0,
1118            temporal_id: 0,
1119        };
1120
1121        let mut new_slot_reference_info = vk::VideoEncodeH264DpbSlotInfoKHR::default()
1122            .std_reference_info(&std_new_slot_reference_info);
1123
1124        let setup_reference_slot_video_resource_info = self
1125            .session_resources
1126            .dpb
1127            .video_resource_info(setup_reference_slot_idx)
1128            .unwrap();
1129
1130        let setup_reference_slot = vk::VideoReferenceSlotInfoKHR::default()
1131            .slot_index(setup_reference_slot_idx as i32)
1132            .picture_resource(setup_reference_slot_video_resource_info)
1133            .push_next(&mut new_slot_reference_info);
1134
1135        let extent = image.extent;
1136
1137        let src_picture_resource = vk::VideoPictureResourceInfoKHR::default()
1138            .coded_offset(vk::Offset2D::default())
1139            .coded_extent(vk::Extent2D {
1140                width: extent.width,
1141                height: extent.height,
1142            })
1143            .base_array_layer(0)
1144            .image_view_binding(view.view);
1145
1146        let mut encode_info = vk::VideoEncodeInfoKHR::default()
1147            .dst_buffer(self.output_buffer.buffer)
1148            .dst_buffer_range(Self::OUTPUT_BUFFER_LEN)
1149            .dst_buffer_offset(0)
1150            .src_picture_resource(src_picture_resource)
1151            .setup_reference_slot(&setup_reference_slot)
1152            .push_next(&mut h264_encode_info);
1153
1154        if !reference_slots.is_empty() {
1155            encode_info = encode_info.reference_slots(&reference_slots);
1156        }
1157
1158        self.query_pool.begin_query(cmd_buffer.buffer());
1159
1160        unsafe {
1161            self.encoding_device
1162                .vulkan_device
1163                .device
1164                .video_encode_queue_ext
1165                .cmd_encode_video_khr(cmd_buffer.buffer(), &encode_info);
1166        }
1167
1168        self.query_pool.end_query(cmd_buffer.buffer());
1169
1170        unsafe {
1171            self.encoding_device
1172                .vulkan_device
1173                .device
1174                .video_queue_ext
1175                .cmd_end_video_coding_khr(
1176                    cmd_buffer.buffer(),
1177                    &vk::VideoEndCodingInfoKHR::default(),
1178                );
1179        }
1180
1181        let wait_value = self
1182            .encoding_device
1183            .h264_encode_queues
1184            .submit_chain_semaphore(
1185                cmd_buffer.end()?,
1186                &mut self.tracker,
1187                vk::PipelineStageFlags2::ALL_COMMANDS,
1188                vk::PipelineStageFlags2::ALL_COMMANDS,
1189                EncoderTrackerWaitState::Encode,
1190            )?;
1191
1192        self.active_reference_slots
1193            .push_back((setup_reference_slot_idx, std_new_slot_reference_info));
1194
1195        self.idr_period_counter += 1;
1196        self.idr_period_counter %= self.idr_period;
1197
1198        Ok(UnwaitedEncodeSubmission(EncodeSubmission {
1199            is_idr,
1200            encoder: self,
1201            wait_value,
1202            pts,
1203            _image: image,
1204        }))
1205    }
1206
1207    pub fn download_output(
1208        &mut self,
1209        is_idr: bool,
1210        pts: Option<u64>,
1211    ) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
1212        let feedback = self.query_pool.get_result_blocking()?;
1213
1214        if feedback.status != vk::QueryResultStatusKHR::COMPLETE {
1215            return Err(VulkanEncoderError::EncodeOperationFailed(feedback.status));
1216        }
1217
1218        let mut output = if is_idr && self.inline_stream_params {
1219            self.stream_parameters(true, true)?
1220        } else {
1221            Vec::new()
1222        };
1223
1224        let encoded = unsafe {
1225            self.output_buffer
1226                .download_data_from_buffer(feedback.bytes_written as usize)?
1227        };
1228
1229        output.extend_from_slice(&encoded);
1230
1231        Ok(EncodedOutputChunk {
1232            data: output,
1233            pts,
1234            is_keyframe: is_idr,
1235        })
1236    }
1237
1238    pub fn stream_parameters(
1239        &self,
1240        write_sps: bool,
1241        write_pps: bool,
1242    ) -> Result<Vec<u8>, VulkanEncoderError> {
1243        if !write_sps && !write_pps {
1244            return Ok(Vec::new());
1245        }
1246
1247        let mut h264_get_info = vk::VideoEncodeH264SessionParametersGetInfoKHR::default()
1248            .write_std_sps(write_sps)
1249            .write_std_pps(write_pps)
1250            .std_sps_id(0)
1251            .std_pps_id(0);
1252
1253        let get_info = vk::VideoEncodeSessionParametersGetInfoKHR::default()
1254            .video_session_parameters(self.session_resources.parameters.parameters)
1255            .push_next(&mut h264_get_info);
1256
1257        let data = unsafe {
1258            self.encoding_device
1259                .vulkan_device
1260                .device
1261                .video_encode_queue_ext
1262                .get_encoded_video_session_parameters_khr(&get_info, None)?
1263        };
1264
1265        Ok(data)
1266    }
1267
1268    pub fn encode_bytes(
1269        &mut self,
1270        frame: &InputFrame<RawFrameData>,
1271        force_idr: bool,
1272    ) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
1273        let (image, _buffer) = self.transfer_buffer_to_image(frame)?;
1274        let image = Arc::new(image);
1275
1276        self.encode(image, force_idr, frame.pts)?
1277            .wait_and_download(u64::MAX)
1278    }
1279
1280    #[cfg(feature = "wgpu")]
1281    pub fn encode_texture(
1282        &mut self,
1283        frame: InputFrame<wgpu::Texture>,
1284        force_idr: bool,
1285    ) -> Result<EncodedOutputChunk<Vec<u8>>, VulkanEncoderError> {
1286        let _cmd_encoder = self.copy_wgpu_texture_to_image(&frame)?;
1287
1288        self.encode(self.input_image.clone(), force_idr, frame.pts)?
1289            .wait_and_download(u64::MAX)
1290    }
1291
1292    fn encoder_rate_control_for<'b>(
1293        &self,
1294        rate_control: RateControl,
1295        layers: Option<&'b [vk::VideoEncodeRateControlLayerInfoKHR]>,
1296    ) -> Option<vk::VideoEncodeRateControlInfoKHR<'b>> {
1297        let layers = layers?;
1298
1299        match rate_control {
1300            RateControl::EncoderDefault => None,
1301
1302            RateControl::VariableBitrate {
1303                virtual_buffer_size,
1304                ..
1305            } => Some(
1306                vk::VideoEncodeRateControlInfoKHR::default()
1307                    .rate_control_mode(vk::VideoEncodeRateControlModeFlagsKHR::VBR)
1308                    .layers(layers)
1309                    .virtual_buffer_size_in_ms(virtual_buffer_size.as_millis() as u32)
1310                    .initial_virtual_buffer_size_in_ms(0),
1311            ),
1312
1313            RateControl::ConstantBitrate {
1314                virtual_buffer_size,
1315                ..
1316            } => Some(
1317                vk::VideoEncodeRateControlInfoKHR::default()
1318                    .rate_control_mode(vk::VideoEncodeRateControlModeFlagsKHR::CBR)
1319                    .layers(layers)
1320                    .virtual_buffer_size_in_ms(virtual_buffer_size.as_millis() as u32)
1321                    .initial_virtual_buffer_size_in_ms(0),
1322            ),
1323
1324            RateControl::Disabled => {
1325                let mut rate_control = vk::VideoEncodeRateControlInfoKHR::default()
1326                    .rate_control_mode(vk::VideoEncodeRateControlModeFlagsKHR::DISABLED)
1327                    .layers(layers);
1328
1329                rate_control.layer_count = 0;
1330                Some(rate_control)
1331            }
1332        }
1333    }
1334
1335    fn h264_rate_control_layers_for(
1336        &self,
1337        rate_control: RateControl,
1338    ) -> Option<Vec<vk::VideoEncodeH264RateControlLayerInfoKHR<'_>>> {
1339        let layer_info = vk::VideoEncodeH264RateControlLayerInfoKHR::default()
1340            .use_min_qp(false)
1341            .use_max_qp(false)
1342            .use_max_frame_size(false);
1343
1344        match rate_control {
1345            RateControl::EncoderDefault => return None,
1346            RateControl::VariableBitrate { .. } => {}
1347            RateControl::ConstantBitrate { .. } => {}
1348            RateControl::Disabled => {}
1349        }
1350
1351        Some(vec![layer_info])
1352    }
1353
1354    fn rate_control_layers_for<'b>(
1355        &self,
1356        rate_control: RateControl,
1357        h264_layer_info: Option<&'b mut [vk::VideoEncodeH264RateControlLayerInfoKHR<'b>]>,
1358    ) -> Option<Vec<vk::VideoEncodeRateControlLayerInfoKHR<'b>>> {
1359        let h264_layer_info = h264_layer_info?;
1360        let mut layer_info = vk::VideoEncodeRateControlLayerInfoKHR::default()
1361            .frame_rate_numerator(self.session_resources.framerate.numerator)
1362            .frame_rate_denominator(self.session_resources.framerate.denominator.get());
1363
1364        match rate_control {
1365            RateControl::EncoderDefault => return None,
1366            RateControl::VariableBitrate {
1367                average_bitrate,
1368                max_bitrate,
1369                ..
1370            } => {
1371                layer_info = layer_info
1372                    .average_bitrate(average_bitrate)
1373                    .max_bitrate(max_bitrate)
1374                    .push_next(&mut h264_layer_info[0])
1375            }
1376
1377            RateControl::ConstantBitrate { bitrate, .. } => {
1378                layer_info = layer_info
1379                    .average_bitrate(bitrate)
1380                    .max_bitrate(bitrate)
1381                    .push_next(&mut h264_layer_info[0])
1382            }
1383
1384            RateControl::Disabled => layer_info = layer_info.push_next(&mut h264_layer_info[0]),
1385        }
1386
1387        Some(vec![layer_info])
1388    }
1389
1390    fn h264_rate_control(
1391        &self,
1392        layers: Option<&[vk::VideoEncodeRateControlLayerInfoKHR]>,
1393    ) -> Option<vk::VideoEncodeH264RateControlInfoKHR<'_>> {
1394        let layers = layers?;
1395
1396        Some(
1397            vk::VideoEncodeH264RateControlInfoKHR::default()
1398                .temporal_layer_count(layers.len() as u32)
1399                .flags(
1400                    vk::VideoEncodeH264RateControlFlagsKHR::REGULAR_GOP
1401                        | vk::VideoEncodeH264RateControlFlagsKHR::REFERENCE_PATTERN_FLAT,
1402                )
1403                .consecutive_b_frame_count(0)
1404                .gop_frame_count(self.idr_period)
1405                .idr_period(self.idr_period),
1406        )
1407    }
1408}
1409
1410/// The rate control algorithm to be used by the encoder.
1411///
1412/// Note: `EncoderDefault` is not a good default! For most implementations it is the same as
1413/// specifying `Disabled`.
1414///
1415/// For most use cases, `Vbr` is the correct option
1416#[derive(Debug, Clone, Copy)]
1417pub enum RateControl {
1418    /// Use the default setting of the encoder implementation.
1419    EncoderDefault,
1420
1421    /// Variable bitrate rate control. This setting fits most use cases. The encoder will try to
1422    /// keep the bitrate around the average, but may increase it temporarily up to the max when
1423    /// necessary, in `virtual_buffer_size`-length windows. Bitrate is measured in bits/second.
1424    VariableBitrate {
1425        average_bitrate: u64,
1426        max_bitrate: u64,
1427        virtual_buffer_size: std::time::Duration,
1428    },
1429
1430    /// Constant bitrate rate control. This setting is for environments that are more
1431    /// bandwidth-constrained. The encoder will keep the bitrate at the specified value, in
1432    /// `virtual_buffer_size`-length windows. Bitrate is measured in bits/second.
1433    ConstantBitrate {
1434        bitrate: u64,
1435        virtual_buffer_size: std::time::Duration,
1436    },
1437
1438    /// Rate control is turned off, frames are compressed with a constant rate. A more complicated
1439    /// frame will just be bigger.
1440    Disabled,
1441}
1442
1443impl RateControl {
1444    pub(crate) fn to_vk(self) -> vk::VideoEncodeRateControlModeFlagsKHR {
1445        match self {
1446            RateControl::EncoderDefault => vk::VideoEncodeRateControlModeFlagsKHR::DEFAULT,
1447            RateControl::VariableBitrate { .. } => vk::VideoEncodeRateControlModeFlagsKHR::VBR,
1448            RateControl::ConstantBitrate { .. } => vk::VideoEncodeRateControlModeFlagsKHR::CBR,
1449            RateControl::Disabled => vk::VideoEncodeRateControlModeFlagsKHR::DISABLED,
1450        }
1451    }
1452}