Skip to main content

alpha2_native/hwencode/
vulkan.rs

1//! The Vulkan backend: `h264_vulkan` and `hevc_vulkan` on the renderer's shared device.
2//!
3//! The frames pool, its image usage, images per frame and owning queue family are decided by
4//! `FFmpeg` and must be observed, not assumed; the probes observe them on any device.
5
6use std::collections::HashSet;
7use std::ffi::CString;
8
9use alpha2_media::{ColorMetadata, EncodeError};
10use alpha2_time::MediaTime;
11use rsmpeg::avcodec::{AVCodec, AVCodecContext};
12use rsmpeg::avutil::{AVFrame, AVHWDeviceContext, AVHWFramesContext, AVRational, av_rescale_q, ra};
13use rsmpeg::error::RsmpegError;
14use rsmpeg::ffi;
15
16use super::{
17    EncoderInputHandoff, EncoderInputLayout, EncoderPoolStats, EncoderSurfaceDescriptor,
18    HardwareEncodeBackend, HardwareEncodeCapability, HardwareEncodeOpenError,
19    HardwareEncodeRequest, InputFrame,
20};
21use crate::NativeVideoCodec;
22use crate::bitstream::EncodedVideoFilter;
23use crate::convert::{encode_adapter_error, tick_base};
24use crate::encode::{EncodeStageTimings, timed};
25use crate::mux::Muxer;
26
27/// What one Vulkan encode session turned out to require and produce.
28#[derive(Clone, Debug, Eq, PartialEq)]
29pub struct VulkanEncodeProbe {
30    /// `FFmpeg` encoder that opened.
31    pub encoder: String,
32    /// Coded size the pool allocated.
33    pub coded_size: (u32, u32),
34    /// Image usage the pool was created with. `FFmpeg` keeps it as given; an unsupported flag fails
35    /// the pool or session rather than disappearing.
36    pub frames_usage: u32,
37    /// `VkImage` handles backing one frame. One means a single multi-planar
38    /// image and a per-plane copy.
39    pub images_per_frame: usize,
40    /// Initial layout `FFmpeg` left the frame in.
41    pub initial_layout: u32,
42    /// Queue family owning the frame, or `u32::MAX` for `VK_QUEUE_FAMILY_IGNORED`.
43    pub queue_family: u32,
44    /// Whether the frame exposes a timeline semaphore for the handoff.
45    pub has_timeline_semaphore: bool,
46    /// Packets drained from one submitted frame plus a flush.
47    pub packets: u32,
48    /// Bytes in those packets.
49    pub packet_bytes: u32,
50}
51
52/// Opens `h264_vulkan` against `device` and encodes one untouched pool frame.
53///
54/// # Errors
55/// Returns an adapter failure when the encoder is absent, the pool cannot be created, the session
56/// cannot open, or no packet is produced.
57pub fn probe_vulkan_encode(
58    device: &AVHWDeviceContext,
59    width: u32,
60    height: u32,
61) -> Result<VulkanEncodeProbe, EncodeError> {
62    let encoder = AVCodec::find_encoder_by_name(c"h264_vulkan").ok_or_else(|| {
63        EncodeError::Adapter("this FFmpeg runtime has no h264_vulkan encoder".into())
64    })?;
65    let coded_width =
66        i32::try_from(width).map_err(|_| EncodeError::Adapter("width is out of range".into()))?;
67    let coded_height =
68        i32::try_from(height).map_err(|_| EncodeError::Adapter("height is out of range".into()))?;
69
70    let mut frames = vulkan_encode_frames(device, coded_width, coded_height)?;
71    let frames_usage = alpha2_ffmpeg_bridge::vulkan_frames_usage(&frames);
72
73    let mut context = AVCodecContext::new(&encoder);
74    context.set_width(coded_width);
75    context.set_height(coded_height);
76    context.set_pix_fmt(ffi::AV_PIX_FMT_VULKAN);
77    // `ff_encode_preinit` copies the pool's `sw_format` into `sw_pix_fmt` and rejects a different
78    // caller-set value, so it stays unset.
79    context.set_time_base(ra(1, 30));
80    context.set_framerate(ra(30, 1));
81    context.set_bit_rate(12_000_000);
82    context.set_gop_size(12);
83    context.set_max_b_frames(0);
84    context.set_hw_frames_ctx(frames.clone());
85    context
86        .open(None)
87        .map_err(|error| encode_adapter_error_at_open("h264_vulkan", &error))?;
88
89    let mut frame = AVFrame::new();
90    frames
91        .get_buffer(&mut frame)
92        .map_err(encode_adapter_error)?;
93    let descriptor =
94        alpha2_ffmpeg_bridge::vulkan_frame_info(&frame).map_err(encode_adapter_error)?;
95    frame.set_pts(0);
96
97    context
98        .send_frame(Some(&frame))
99        .map_err(encode_adapter_error)?;
100    context.send_frame(None).map_err(encode_adapter_error)?;
101    let mut packets = 0;
102    let mut packet_bytes = 0_u32;
103    loop {
104        match context.receive_packet() {
105            Ok(packet) => {
106                packets += 1;
107                packet_bytes = packet_bytes.saturating_add(packet.size.unsigned_abs());
108            }
109            Err(RsmpegError::EncoderDrainError | RsmpegError::EncoderFlushedError) => break,
110            Err(error) => return Err(encode_adapter_error(error)),
111        }
112    }
113    if packets == 0 {
114        return Err(EncodeError::Adapter(
115            "h264_vulkan opened but produced no packet".into(),
116        ));
117    }
118    Ok(VulkanEncodeProbe {
119        encoder: encoder.name().to_string_lossy().into_owned(),
120        coded_size: (
121            descriptor.coded_width.unsigned_abs(),
122            descriptor.coded_height.unsigned_abs(),
123        ),
124        frames_usage,
125        images_per_frame: descriptor.plane_count as usize,
126        initial_layout: descriptor.layouts[0],
127        queue_family: descriptor.queue_families[0],
128        has_timeline_semaphore: descriptor.semaphores[0] != 0,
129        packets,
130        packet_bytes,
131    })
132}
133
134/// Opens a Vulkan encoder with one rate-control request and encodes one pool frame, on
135/// a device `FFmpeg` creates.
136///
137/// It passes if it opened, consumed every option and produced a packet. `FFmpeg` checks the mode
138/// against `VkVideoEncodeCapabilitiesKHR::rateControlModes` at open.
139pub(crate) fn probe_vulkan_rate_control(
140    encoder: &rsmpeg::avcodec::AVCodecRef<'static>,
141    width: u32,
142    height: u32,
143    request: &crate::rate_control::MechanismRequest,
144) -> Result<(u32, String), String> {
145    // Pass the pinned device by loader index: `FFmpeg` reads a leading digit as an index, so an
146    // `ALPHA2_GPU_ADAPTER` name cannot pass through. Without a pin, walk the indices until one
147    // accepts or refuses the mode, since only some devices have an encode queue.
148    let pinned = alpha2_native_gpu::adapter_pinned();
149    let shared = alpha2_native_gpu::shared_device_index();
150    let candidates = match (pinned, shared) {
151        (true, Ok(index)) => vec![index.to_string()],
152        (true, Err(error)) => return Err(error.to_string()),
153        (false, shared) => {
154            let first = shared.ok();
155            first
156                .into_iter()
157                .chain((0..8).filter(|index| Some(*index) != first))
158                .map(|index| index.to_string())
159                .collect::<Vec<_>>()
160        }
161    };
162    let mut last = "no Vulkan device".to_owned();
163    for device in candidates {
164        match probe_vulkan_rate_control_on(encoder, &device, width, height, request) {
165            Ok(packets) => return Ok((packets, device)),
166            Err(VulkanProbeError::NoDevice(error)) => {
167                last = error;
168                break;
169            }
170            Err(VulkanProbeError::NoEncodeQueue(error)) if !pinned => last = error,
171            Err(VulkanProbeError::NoEncodeQueue(error) | VulkanProbeError::Refused(error)) => {
172                return Err(format!("on Vulkan device {device}: {error}"));
173            }
174        }
175    }
176    Err(last)
177}
178
179enum VulkanProbeError {
180    /// No device at this index; the walk is over.
181    NoDevice(String),
182    /// The device has no video-encode queue; another one might.
183    NoEncodeQueue(String),
184    /// The device got as far as the encoder, which refused.
185    Refused(String),
186}
187
188fn probe_vulkan_rate_control_on(
189    encoder: &rsmpeg::avcodec::AVCodecRef<'static>,
190    device_name: &str,
191    width: u32,
192    height: u32,
193    request: &crate::rate_control::MechanismRequest,
194) -> Result<u32, VulkanProbeError> {
195    let device_name = CString::new(device_name)
196        .map_err(|_| VulkanProbeError::NoDevice("device name contains a NUL byte".to_owned()))?;
197    // `FFmpeg` enables video-encode extensions on its own device only when asked; the shared export
198    // device asks through wgpu.
199    let extensions = rsmpeg::avutil::AVDictionary::new(
200        c"device_extensions",
201        c"VK_KHR_video_queue+VK_KHR_video_encode_queue+VK_KHR_video_encode_h264+VK_KHR_video_encode_h265+VK_KHR_video_encode_av1+VK_KHR_video_maintenance1",
202        0,
203    );
204    let device = AVHWDeviceContext::create(
205        ffi::AV_HWDEVICE_TYPE_VULKAN,
206        Some(&device_name),
207        Some(&extensions),
208        0,
209    )
210    .map_err(|error| VulkanProbeError::NoDevice(format!("no Vulkan device: {error}")))?;
211    let refused = |error: String| VulkanProbeError::Refused(error);
212    let coded_width =
213        i32::try_from(width).map_err(|_| refused("width is out of range".to_owned()))?;
214    let coded_height =
215        i32::try_from(height).map_err(|_| refused("height is out of range".to_owned()))?;
216    let mut frames = vulkan_encode_frames(&device, coded_width, coded_height)
217        .map_err(|error| VulkanProbeError::NoEncodeQueue(error.to_string()))?;
218    let mut context = AVCodecContext::new(encoder);
219    context.set_width(coded_width);
220    context.set_height(coded_height);
221    context.set_pix_fmt(ffi::AV_PIX_FMT_VULKAN);
222    context.set_time_base(ra(1, 30));
223    context.set_framerate(ra(30, 1));
224    context.set_bit_rate(request.bit_rate);
225    context.set_gop_size(12);
226    context.set_max_b_frames(0);
227    context.set_hw_frames_ctx(frames.clone());
228    let unused = alpha2_ffmpeg_bridge::open_codec_context(
229        &mut context,
230        request.dictionary(crate::convert::color_options(ColorMetadata::EXPORT)),
231    )
232    .map_err(|error| {
233        // `ENOSYS` means no encode queue at all, which says nothing about the mode.
234        let message = error.to_string();
235        if matches!(error, RsmpegError::AVError(code) if code == -38) {
236            VulkanProbeError::NoEncodeQueue(message)
237        } else {
238            refused(message)
239        }
240    })?;
241    crate::convert::reject_unused_encode_options(unused)
242        .map_err(|error| refused(error.to_string()))?;
243    let mut frame = AVFrame::new();
244    frames
245        .get_buffer(&mut frame)
246        .map_err(|error| refused(error.to_string()))?;
247    frame.set_pts(0);
248    context
249        .send_frame(Some(&frame))
250        .map_err(|error| refused(error.to_string()))?;
251    context
252        .send_frame(None)
253        .map_err(|error| refused(error.to_string()))?;
254    let mut packets = 0;
255    loop {
256        match context.receive_packet() {
257            Ok(_) => packets += 1,
258            Err(RsmpegError::EncoderDrainError | RsmpegError::EncoderFlushedError) => break,
259            Err(error) => return Err(refused(error.to_string())),
260        }
261    }
262    if packets == 0 {
263        return Err(refused("opened but produced no packet".to_owned()));
264    }
265    Ok(packets)
266}
267
268/// The `FFmpeg` Vulkan encoder of `codec`, and the name the device's codec list gives it.
269const fn vulkan_encoder(codec: NativeVideoCodec) -> Option<(&'static str, &'static str)> {
270    match codec {
271        NativeVideoCodec::H264 => Some(("h264_vulkan", "h264")),
272        NativeVideoCodec::Hevc => Some(("hevc_vulkan", "hevc")),
273        _ => None,
274    }
275}
276
277/// The capability a report leads with: the first of H.264 and HEVC this device encodes.
278pub(super) fn preferred_capability(
279    shared: &alpha2_native_gpu::SharedVulkanContext,
280) -> Result<HardwareEncodeCapability, String> {
281    capability(shared, NativeVideoCodec::H264)
282        .or_else(|h264| capability(shared, NativeVideoCodec::Hevc).map_err(|_| h264))
283}
284
285/// What this device's Vulkan encoder of `codec` accepts, from the device's own answers.
286pub(super) fn capability(
287    shared: &alpha2_native_gpu::SharedVulkanContext,
288    codec: NativeVideoCodec,
289) -> Result<HardwareEncodeCapability, String> {
290    if shared.encode_codecs().is_empty() {
291        return Err("this Vulkan device has no video-encode queue".to_owned());
292    }
293    let Some((encoder, listed)) = vulkan_encoder(codec) else {
294        return Err(format!(
295            "Vulkan Video encodes h264 and hevc, not {}",
296            codec.as_str()
297        ));
298    };
299    if !shared.encode_codecs().contains(&listed) {
300        return Err(format!(
301            "this Vulkan device has no video encoder for {} (it encodes {})",
302            codec.as_str(),
303            shared.encode_codecs().join(", ")
304        ));
305    }
306    let extent = shared.encode_extent(listed).ok_or_else(|| {
307        format!(
308            "this device did not report its {} encode extents",
309            codec.as_str()
310        )
311    })?;
312    Ok(HardwareEncodeCapability {
313        backend: HardwareEncodeBackend::Vulkan,
314        encoder,
315        codec,
316        gpu_input_only: true,
317        input_layout: EncoderInputLayout::Nv12,
318        handoff: EncoderInputHandoff::GpuCopy,
319        shares_renderer_device: true,
320        minimum: extent.minimum,
321        maximum: extent.maximum,
322        granularity: extent.granularity,
323        max_p_picture_l0_references: Some(extent.max_p_picture_l0_references),
324        max_b_picture_l0_references: Some(extent.max_b_picture_l0_references),
325        max_l1_references: Some(extent.max_l1_references),
326        rate_control_modes: extent.rate_control_modes.clone(),
327        plane_ring_slots: alpha2_native_gpu::VulkanEncodeCopier::slots(),
328        // The copier keeps each copied frame until its slot is reused.
329        handoff_retained_surfaces: alpha2_native_gpu::VulkanEncodeCopier::slots(),
330        // Measured: the pool never exceeded the host's leases plus the copier's slots.
331        codec_retained_surfaces: 0,
332        default_channel_depth: DEFAULT_CHANNEL_DEPTH,
333    })
334}
335
336/// Depth 2 keeps a frame queued behind the one encoding, so uneven renders do not stall the
337/// encoder. Measured on a discrete GPU, depth 0 was about 1 % slower on a heavy 4K composition, at
338/// a cost of two extra surfaces (about 25 MiB at 4K).
339const DEFAULT_CHANNEL_DEPTH: usize = 2;
340
341/// Opens the Vulkan encoder of `request.codec` on the shared device, its input half, and the
342/// output.
343pub(super) fn open(
344    shared: &alpha2_native_gpu::SharedVulkanContext,
345    request: &HardwareEncodeRequest<'_>,
346    generation: u64,
347) -> Result<(VulkanEncodeSession, VulkanEncodeInput), HardwareEncodeOpenError> {
348    let refused = |error: EncodeError| HardwareEncodeOpenError::Refused(error.to_string());
349    let (width, height) = (request.width, request.height);
350    let frame_rate = super::session_frame_rate(request.frame_rate).map_err(refused)?;
351    let (encoder, _) = vulkan_encoder(request.codec).ok_or_else(|| {
352        HardwareEncodeOpenError::Refused(format!(
353            "Vulkan Video does not encode {}",
354            request.codec.as_str()
355        ))
356    })?;
357    let mark = alpha2_ffmpeg_bridge::log_mark();
358    let (frames, context, selection) = open_encoder(
359        &shared.ffmpeg_device(),
360        encoder,
361        width,
362        height,
363        request.rate_control,
364        frame_rate,
365        request.color,
366    )
367    .map_err(|error| super::open_refusal(encoder, mark, &error))?;
368    let time_base = ra(frame_rate.1, frame_rate.0);
369    // HEVC uses the bitstream filter chain: `hevc_metadata` writes the colour into every parameter
370    // set and `extract_extradata` keeps them out of the packets, as `hvc1` requires. H.264 keeps
371    // the session's own in-band and header parameter sets.
372    let bitstream = (request.codec == NativeVideoCodec::Hevc)
373        .then(|| EncodedVideoFilter::new(&context.extract_codecpar(), request.color, time_base))
374        .transpose()
375        .map_err(refused)?;
376    let copier = alpha2_native_gpu::VulkanEncodeCopier::new(shared)
377        .map_err(|error| HardwareEncodeOpenError::Refused(format!("{error:?}")))?;
378    let gpu = shared.gpu();
379    // One pair per copier slot: the coordinator runs ahead of the encoder, so one pair would let
380    // the next render overwrite planes a copy is still reading.
381    let planes = (0..alpha2_native_gpu::VulkanEncodeCopier::slots())
382        .map(|_| Nv12PlanePair {
383            luma: plane_texture(
384                gpu.device(),
385                "alpha2 NV12 luma",
386                width,
387                height,
388                wgpu::TextureFormat::R8Unorm,
389            ),
390            chroma: plane_texture(
391                gpu.device(),
392                "alpha2 NV12 chroma",
393                width / 2,
394                height / 2,
395                wgpu::TextureFormat::Rg8Unorm,
396            ),
397        })
398        .collect();
399    let input = VulkanEncodeInput {
400        copier,
401        planes,
402        frames,
403        size: (width, height),
404        time_base,
405        color: request.color,
406        stats: EncoderPoolStats::default(),
407        images: HashSet::new(),
408    };
409    // From here a failure is the export's, since the partial output exists. A filtered stream's
410    // header waits for the parameter sets in its first packet.
411    let muxer = if bitstream.is_some() {
412        Muxer::open_deferred(
413            request.output,
414            request.container,
415            &context,
416            time_base,
417            request.audio,
418        )
419    } else {
420        Muxer::open(
421            request.output,
422            request.container,
423            &context,
424            time_base,
425            request.audio,
426        )
427    }
428    .map_err(HardwareEncodeOpenError::Failed)?;
429    Ok((
430        VulkanEncodeSession {
431            context,
432            time_base,
433            bitstream,
434            muxer,
435            stages: EncodeStageTimings::default(),
436            selection,
437            generation,
438        },
439        input,
440    ))
441}
442
443/// One renderer-owned NV12 plane target for the encode handoff.
444///
445/// `COPY_SRC` gives the image `VK_IMAGE_USAGE_TRANSFER_SRC` for the raw copy.
446fn plane_texture(
447    device: &wgpu::Device,
448    label: &'static str,
449    width: u32,
450    height: u32,
451    format: wgpu::TextureFormat,
452) -> wgpu::Texture {
453    device.create_texture(&wgpu::TextureDescriptor {
454        label: Some(label),
455        size: wgpu::Extent3d {
456            width,
457            height,
458            depth_or_array_layers: 1,
459        },
460        mip_level_count: 1,
461        sample_count: 1,
462        dimension: wgpu::TextureDimension::D2,
463        format,
464        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
465        view_formats: &[],
466    })
467}
468
469struct Nv12PlanePair {
470    luma: wgpu::Texture,
471    chroma: wgpu::Texture,
472}
473
474/// Opens `encoder` (`h264_vulkan` or `hevc_vulkan`) on `device` with an NV12 encoder-input pool.
475fn open_encoder(
476    device: &AVHWDeviceContext,
477    encoder_name: &'static str,
478    width: u32,
479    height: u32,
480    rate_control: alpha2_media::VideoRateControl,
481    frame_rate: (i32, i32),
482    color: ColorMetadata,
483) -> Result<
484    (
485        AVHWFramesContext,
486        AVCodecContext,
487        crate::NativeVideoEncoderSelection,
488    ),
489    EncodeError,
490> {
491    let encoder = CString::new(encoder_name)
492        .ok()
493        .and_then(|name| AVCodec::find_encoder_by_name(&name))
494        .ok_or_else(|| {
495            EncodeError::Adapter(format!("this FFmpeg runtime has no {encoder_name} encoder"))
496        })?;
497    let coded_width =
498        i32::try_from(width).map_err(|_| EncodeError::Adapter("width is out of range".into()))?;
499    let coded_height =
500        i32::try_from(height).map_err(|_| EncodeError::Adapter("height is out of range".into()))?;
501    let frames = vulkan_encode_frames(device, coded_width, coded_height)?;
502    let rate = u32::try_from(frame_rate.0)
503        .ok()
504        .zip(u32::try_from(frame_rate.1).ok())
505        .and_then(|(num, den)| alpha2_time::FrameRate::new(num, den).ok())
506        .ok_or_else(|| EncodeError::Adapter("frame rate is out of range".into()))?;
507    let open = |quantized: Option<&crate::rate_control::MechanismRequest>| {
508        let mut context = AVCodecContext::new(&encoder);
509        context.set_width(coded_width);
510        context.set_height(coded_height);
511        context.set_pix_fmt(ffi::AV_PIX_FMT_VULKAN);
512        // `ff_encode_preinit` copies the pool's `sw_format` into `sw_pix_fmt` and rejects a
513        // different caller-set value.
514        context.set_time_base(ra(frame_rate.1, frame_rate.0));
515        context.set_framerate(ra(frame_rate.0, frame_rate.1));
516        context.set_bit_rate(quantized.map_or_else(
517            || i64::try_from(rate_control.bitrate).unwrap_or(i64::MAX),
518            |quantized| quantized.bit_rate,
519        ));
520        context.set_gop_size(crate::capabilities::keyframe_interval(rate));
521        // No B pictures: at a fixed rate two B frames cost up to about 0.6 dB PSNR-Y against none.
522        context.set_max_b_frames(0);
523        context.set_hw_frames_ctx(frames.clone());
524        // The flag is read at open, so set it first or the parameter sets travel in-band only.
525        context.set_flags(
526            context.flags | i32::try_from(ffi::AV_CODEC_FLAG_GLOBAL_HEADER).unwrap_or(0),
527        );
528        // Without these the stream has no VUI and decoders guess the colour.
529        let color = crate::convert::color_options(color);
530        let options = match quantized {
531            Some(quantized) => quantized.dictionary(color),
532            None => crate::rate_control::rate_control_options(encoder_name, rate_control, color),
533        };
534        let unused = alpha2_ffmpeg_bridge::open_codec_context(&mut context, options)
535            .map_err(|error| encode_adapter_error_at_open(encoder_name, &error))?;
536        crate::convert::reject_unused_encode_options(unused)?;
537        if context.extradata_size <= 0 {
538            return Err(EncodeError::Adapter(format!(
539                "{encoder_name} opened without parameter-set extradata for the MP4 header"
540            )));
541        }
542        Ok::<_, EncodeError>(context)
543    };
544    // A refused quantizer falls back to bitrate on the same encoder, and the selection reports it.
545    let (context, applied, quantizer_refused) =
546        match crate::rate_control::quantizer_request(encoder_name, rate_control) {
547            Some(Ok(quantized)) => match open(Some(&quantized)) {
548                Ok(context) => (context, crate::AppliedRateControl::Quantizer, None),
549                Err(error) => (
550                    open(None)?,
551                    crate::AppliedRateControl::Bitrate,
552                    Some(format!("{} refused: {error}", quantized.describe())),
553                ),
554            },
555            Some(Err(reason)) => (
556                open(None)?,
557                crate::AppliedRateControl::Bitrate,
558                Some(reason),
559            ),
560            None => (open(None)?, crate::AppliedRateControl::Bitrate, None),
561        };
562    if let Some(reason) = &quantizer_refused {
563        alpha2_os::diagnostic!(
564            "alpha2 native export: {encoder_name} runs the bitrate fallback: {reason}"
565        );
566    }
567    // No candidate list: the encoder was chosen from the device.
568    let selection = crate::NativeVideoEncoderSelection {
569        encoder_name: encoder_name.to_owned(),
570        class: crate::NativeEncoderClass::Hardware,
571        hardware_accelerated: true,
572        rejected: Vec::new(),
573        rate_control: applied,
574        quantizer_refused,
575        software_threads: None,
576        pixels: alpha2_media::PixelLayout::YUV420P,
577    };
578    Ok((frames, context, selection))
579}
580
581/// The encode worker's half of a Vulkan session: the codec and the muxer.
582pub(super) struct VulkanEncodeSession {
583    context: AVCodecContext,
584    time_base: AVRational,
585    /// The filter chain an HEVC stream's packets pass through before the muxer.
586    bitstream: Option<EncodedVideoFilter>,
587    muxer: Muxer,
588    stages: EncodeStageTimings,
589    selection: crate::NativeVideoEncoderSelection,
590    generation: u64,
591}
592
593/// The coordinator's half of a Vulkan session.
594///
595/// The copy stays on the thread that rendered the planes, which owns the wgpu textures and render
596/// signal. The encoder-input pool is reference counted, so frames stay valid across threads.
597///
598/// Fields drop in declaration order and the copier comes first: dropping it waits on every
599/// in-flight copy (wgpu cannot track these raw copies), retaining resources if the wait times out.
600/// Dropping the planes first would free them under a running copy.
601pub(super) struct VulkanEncodeInput {
602    copier: alpha2_native_gpu::VulkanEncodeCopier,
603    planes: Vec<Nv12PlanePair>,
604    frames: AVHWFramesContext,
605    size: (u32, u32),
606    time_base: AVRational,
607    color: ColorMetadata,
608    stats: EncoderPoolStats,
609    /// The distinct images the pool has handed out, which equals what it allocated.
610    images: HashSet<u64>,
611}
612
613/// One frame taken from a Vulkan encoder-input pool.
614pub(super) struct VulkanEncoderFrame(AVFrame);
615
616impl VulkanEncoderFrame {
617    pub(super) fn set_pts(&mut self, pts: i64) {
618        self.0.set_pts(pts);
619    }
620}
621
622impl VulkanEncodeInput {
623    pub(super) fn descriptor(&self, generation: u64) -> EncoderSurfaceDescriptor {
624        EncoderSurfaceDescriptor {
625            backend: HardwareEncodeBackend::Vulkan,
626            pool_generation: generation,
627            layout: EncoderInputLayout::Nv12,
628            visible: self.size,
629            color: self.color,
630        }
631    }
632
633    /// Takes one frame from the encoder-input pool.
634    ///
635    /// The pool recycles an image once the host, encoder and copier (which keeps each frame until
636    /// its slot comes round) have all dropped it.
637    pub(super) fn acquire_frame(&mut self) -> Result<VulkanEncoderFrame, EncodeError> {
638        let mut frame = AVFrame::new();
639        self.frames
640            .get_buffer(&mut frame)
641            .map_err(encode_adapter_error)?;
642        let image = alpha2_ffmpeg_bridge::vulkan_frame_info(&frame)
643            .map_err(encode_adapter_error)?
644            .images[0];
645        self.images.insert(image);
646        self.stats.acquired = self.stats.acquired.saturating_add(1);
647        self.stats.allocated = u64::try_from(self.images.len()).unwrap_or(u64::MAX);
648        Ok(VulkanEncoderFrame(frame))
649    }
650
651    /// Waits until the plane pair the next copy reads is free. Submitting the pack render before
652    /// [`Self::submit`]'s copy on the same queue is what orders the two.
653    pub(super) fn begin_frame(&mut self) -> Result<InputFrame, EncodeError> {
654        let slot = self.copier.await_slot().map_err(encode_adapter_error)?;
655        if slot >= self.planes.len() {
656            return Err(EncodeError::Adapter(format!(
657                "encode copier reported slot {slot}"
658            )));
659        }
660        Ok(InputFrame { slot })
661    }
662
663    pub(super) fn planes(&self, frame: &InputFrame) -> (&wgpu::Texture, &wgpu::Texture) {
664        let pair = &self.planes[frame.slot];
665        (&pair.luma, &pair.chroma)
666    }
667
668    /// Copies `frame`'s planes into `backing`, signalling the frame's timeline semaphore
669    /// for the encoder to wait on.
670    // Taking the token by value enforces one submit per begun frame.
671    #[allow(clippy::needless_pass_by_value)]
672    pub(super) fn submit(
673        &mut self,
674        frame: InputFrame,
675        backing: &mut VulkanEncoderFrame,
676    ) -> Result<(), EncodeError> {
677        let InputFrame { slot } = frame;
678        let pair = &self.planes[slot];
679        let mut access = alpha2_ffmpeg_bridge::VulkanEncodeFrameAccess::begin(&mut backing.0)
680            .map_err(encode_adapter_error)?;
681        self.copier
682            .submit(&mut access, &pair.luma, &pair.chroma, self.size)
683            .map_err(encode_adapter_error)?;
684        self.stats.copy_operations = self.stats.copy_operations.saturating_add(2);
685        self.stats.copy_bytes = self
686            .stats
687            .copy_bytes
688            .saturating_add(EncoderInputLayout::Nv12.frame_bytes(self.size.0, self.size.1));
689        Ok(())
690    }
691
692    /// Presentation timestamp for the frame at `time`, in encoder ticks.
693    pub(super) fn timestamp(&self, time: MediaTime) -> i64 {
694        av_rescale_q(time.ticks(), tick_base(), self.time_base)
695    }
696
697    pub(super) const fn stats(&self) -> EncoderPoolStats {
698        self.stats
699    }
700
701    /// Waits on every copy's fence before the plane textures it reads drop.
702    pub(super) fn finish(self) -> Result<(), EncodeError> {
703        let Self { copier, planes, .. } = self;
704        copier.finish().map_err(encode_adapter_error)?;
705        drop(planes);
706        Ok(())
707    }
708}
709
710impl VulkanEncodeSession {
711    pub(super) const fn selection(&self) -> &crate::NativeVideoEncoderSelection {
712        &self.selection
713    }
714
715    pub(super) fn stage_timings(&self) -> EncodeStageTimings {
716        self.muxer.with_audio_timings(self.stages)
717    }
718
719    /// Submits `frame` and muxes whatever packets it produced.
720    pub(super) fn encode(
721        &mut self,
722        generation: u64,
723        frame: &VulkanEncoderFrame,
724    ) -> Result<(), EncodeError> {
725        if generation != self.generation {
726            return Err(EncodeError::Adapter(format!(
727                "encoder surface from pool {generation} reached session {}",
728                self.generation
729            )));
730        }
731        timed(&mut self.stages.video_send_frame, || {
732            self.context.send_frame(Some(&frame.0))
733        })
734        .map_err(encode_adapter_error)?;
735        self.drain()
736    }
737
738    pub(super) fn write_audio(
739        &mut self,
740        block: &alpha2_media::AudioBlock,
741    ) -> Result<(), EncodeError> {
742        self.muxer.write_audio(block, &mut self.stages)
743    }
744
745    /// Flushes audio and encoder, then closes the output, which renames
746    /// the partial file into place.
747    pub(super) fn finish(&mut self, cancel: &alpha2_media::CancelToken) -> Result<(), EncodeError> {
748        // Every audio packet must exist before the muxer decides the final interleave.
749        self.muxer.finish_audio()?;
750        crate::encode::beat_encode_call();
751        self.context
752            .send_frame(None)
753            .map_err(encode_adapter_error)?;
754        self.drain()?;
755        if let Some(bitstream) = &mut self.bitstream {
756            for mut packet in bitstream.flush()? {
757                self.muxer.write_video(&mut packet, self.time_base)?;
758            }
759        }
760        self.muxer.write_ready_audio(true, &mut self.stages)?;
761        self.muxer.close(cancel)
762    }
763
764    pub(super) fn publish_report(&self) -> Option<crate::PublishReport> {
765        self.muxer.publish_report()
766    }
767
768    pub(super) fn abort(&mut self) {
769        self.muxer.abort();
770    }
771
772    fn drain(&mut self) -> Result<(), EncodeError> {
773        loop {
774            let received = timed(&mut self.stages.video_receive_packet, || {
775                self.context.receive_packet()
776            });
777            match received {
778                Ok(mut packet) => match &mut self.bitstream {
779                    Some(bitstream) => {
780                        for mut packet in bitstream.filter(packet)? {
781                            self.muxer.write_video(&mut packet, self.time_base)?;
782                        }
783                    }
784                    None => self.muxer.write_video(&mut packet, self.time_base)?,
785                },
786                Err(RsmpegError::EncoderDrainError | RsmpegError::EncoderFlushedError) => {
787                    return self.muxer.write_ready_audio(false, &mut self.stages);
788                }
789                Err(error) => return Err(encode_adapter_error(error)),
790            }
791        }
792    }
793}
794
795/// Builds the encoder-input pool.
796///
797/// `avcodec_get_hw_frames_parameters` answers `ENOENT` for encoders, so the pool is described here.
798/// The Vulkan encoder requests no image usage, so the pool gets the export's
799/// (`vulkan_frames_add_usage`) before init.
800fn vulkan_encode_frames(
801    device: &AVHWDeviceContext,
802    width: i32,
803    height: i32,
804) -> Result<AVHWFramesContext, EncodeError> {
805    let mut frames = device.hwframe_ctx_alloc();
806    {
807        let data = frames.data();
808        data.format = ffi::AV_PIX_FMT_VULKAN;
809        data.sw_format = ffi::AV_PIX_FMT_NV12;
810        data.width = width;
811        data.height = height;
812        data.initial_pool_size = 0;
813    }
814    alpha2_ffmpeg_bridge::vulkan_frames_add_usage(
815        &mut frames,
816        alpha2_ffmpeg_bridge::VULKAN_ENCODE_INPUT_USAGE,
817    )
818    .map_err(encode_adapter_error)?;
819    frames.init().map_err(encode_adapter_error)?;
820    Ok(frames)
821}
822
823fn encode_adapter_error_at_open(encoder: &str, error: &RsmpegError) -> EncodeError {
824    EncodeError::Adapter(format!("{encoder} session could not open: {error}"))
825}