pito_codecs/gpu/
mod.rs

1#[cfg(any(target_os = "linux", target_os = "windows"))]
2#[cfg_attr(target_os = "windows", allow(dead_code))]
3mod biplanar;
4mod nal;
5
6#[cfg(test)]
7mod tests;
8
9#[cfg(target_os = "linux")]
10mod vaapi;
11
12#[cfg(any(
13    target_os = "windows",
14    all(target_os = "linux", target_pointer_width = "64")
15))]
16mod nvenc;
17
18#[cfg(target_os = "macos")]
19mod videotoolbox;
20
21#[cfg(target_os = "windows")]
22mod media_foundation;
23
24use crate::media::{self, Codec, Error, Frame, Packet, Result, Rgba, TimeBase};
25use std::sync::Arc;
26
27#[cfg(target_os = "linux")]
28pub use vaapi::{DmaBuf, DmaLayer, DmaObject, DmaPlane};
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
31pub enum Direction {
32    Decode,
33    Encode,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
37pub enum Backend {
38    VaApi,
39    Nvenc,
40    VideoToolbox,
41    MediaFoundation,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
45pub enum Profile {
46    Baseline,
47    Main,
48    High,
49    Main10,
50}
51
52#[deprecated(note = "use pito_codecs::media::Rate")]
53pub type Rate = crate::media::Rate;
54
55#[deprecated(note = "use pito_codecs::media::Bitstream")]
56pub type Bitstream = crate::media::Bitstream;
57
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub struct EncoderConfig {
60    pub width: u32,
61    pub height: u32,
62    pub fps: (u32, u32),
63    pub profile: Profile,
64    pub rate: media::Rate,
65    pub gop: u32,
66    pub b_frames: u8,
67    pub bitstream: media::Bitstream,
68}
69
70impl Default for EncoderConfig {
71    fn default() -> Self {
72        Self {
73            width: 1920,
74            height: 1080,
75            fps: (30, 1),
76            profile: Profile::High,
77            rate: media::Rate::Crf(20),
78            gop: 0,
79            b_frames: 0,
80            bitstream: media::Bitstream::LengthPrefixed,
81        }
82    }
83}
84
85pub fn available() -> Vec<(Codec, Direction, Backend)> {
86    let mut out = Vec::new();
87    #[cfg(target_os = "linux")]
88    out.extend(
89        vaapi::available()
90            .into_iter()
91            .map(|(c, d)| (c, d, Backend::VaApi)),
92    );
93    #[cfg(any(
94        target_os = "windows",
95        all(target_os = "linux", target_pointer_width = "64")
96    ))]
97    out.extend(
98        nvenc::available()
99            .into_iter()
100            .map(|(c, d)| (c, d, Backend::Nvenc)),
101    );
102    #[cfg(target_os = "macos")]
103    out.extend(
104        videotoolbox::available()
105            .into_iter()
106            .map(|(c, d)| (c, d, Backend::VideoToolbox)),
107    );
108    #[cfg(target_os = "windows")]
109    out.extend(
110        media_foundation::available()
111            .into_iter()
112            .map(|(c, d)| (c, d, Backend::MediaFoundation)),
113    );
114    out
115}
116
117pub fn driver(backend: Backend) -> Result<String> {
118    match backend {
119        #[cfg(target_os = "linux")]
120        Backend::VaApi => vaapi::driver(),
121        b => Err(Error::Unsupported(format!("{b:?} on this system"))),
122    }
123}
124
125#[cfg_attr(
126    not(all(
127        feature = "h264",
128        feature = "hevc",
129        feature = "aac",
130        feature = "mp4",
131        feature = "filter",
132        feature = "mkv",
133        feature = "vp9"
134    )),
135    allow(dead_code)
136)]
137pub(crate) fn color(codec: Codec, extradata: &[u8]) -> Option<(u8, bool)> {
138    let sets = nal::parse_extradata(codec, extradata).ok()?.parameter_sets;
139    match codec {
140        Codec::H264 => {
141            let nal = sets
142                .iter()
143                .find(|n| n.first().is_some_and(|h| h & 0x1f == 7))?;
144            let rbsp = crate::syntax::bits::unescape_rbsp(nal.get(1..)?);
145            let vui = crate::syntax::h264::params::Sps::parse(&rbsp).ok()?.vui?;
146            Some((vui.matrix_coefficients, vui.full_range))
147        }
148        Codec::Hevc => {
149            let nal = sets
150                .iter()
151                .find(|n| n.first().is_some_and(|h| (h >> 1) & 0x3f == 33))?;
152            let rbsp = crate::syntax::bits::unescape_rbsp(nal.get(2..)?);
153            let vui = crate::syntax::hevc::params::Sps::parse(&rbsp).ok()?.vui?;
154            Some((vui.matrix_coefficients, vui.full_range))
155        }
156        _ => None,
157    }
158}
159
160pub(crate) trait DecoderBackend: Send {
161    fn backend(&self) -> Backend;
162    fn name(&self) -> &str;
163    fn decode(&mut self, packet: &Packet) -> Result<Vec<Frame>>;
164    fn flush(&mut self) -> Result<Vec<Frame>>;
165    fn reset(&mut self);
166    fn stills(&mut self) {}
167    fn decode_gpu(&mut self, packet: &Packet) -> Result<Vec<GpuFrame>> {
168        Ok(self
169            .decode(packet)?
170            .into_iter()
171            .map(GpuFrame::cpu)
172            .collect())
173    }
174    fn flush_gpu(&mut self) -> Result<Vec<GpuFrame>> {
175        Ok(self.flush()?.into_iter().map(GpuFrame::cpu).collect())
176    }
177}
178
179pub(crate) trait EncoderBackend: Send {
180    fn backend(&self) -> Backend;
181    fn name(&self) -> &str;
182    fn b_frames(&self) -> u8;
183    fn encode(&mut self, frame: &Frame) -> Result<Vec<Packet>>;
184    fn flush(&mut self) -> Result<Vec<Packet>>;
185    fn extradata(&self) -> Option<&[u8]>;
186    fn encode_gpu(&mut self, frame: &GpuFrame) -> Result<Vec<Packet>> {
187        self.encode(&frame.download()?)
188    }
189    fn composes(&self) -> bool {
190        false
191    }
192    fn encode_gpu_over(&mut self, _frame: &GpuFrame, _layer: &Arc<Rgba>) -> Result<Vec<Packet>> {
193        Err(Error::Unsupported(
194            "this GPU encoder can't draw an overlay on the GPU".into(),
195        ))
196    }
197}
198
199type DecoderMaker = fn(Codec, &[u8]) -> Result<Box<dyn DecoderBackend>>;
200type EncoderMaker = fn(Codec, &EncoderConfig) -> Result<Box<dyn EncoderBackend>>;
201
202#[allow(clippy::vec_init_then_push)]
203fn decoders() -> Vec<(&'static str, DecoderMaker)> {
204    let mut out: Vec<(&'static str, DecoderMaker)> = Vec::new();
205    #[cfg(target_os = "linux")]
206    out.push(("VA-API", vaapi::decoder));
207    #[cfg(target_os = "macos")]
208    out.push(("VideoToolbox", videotoolbox::decoder));
209    #[cfg(target_os = "windows")]
210    out.push(("Media Foundation", media_foundation::decoder));
211    out
212}
213
214#[allow(clippy::vec_init_then_push)]
215fn encoders() -> Vec<(&'static str, EncoderMaker)> {
216    let mut out: Vec<(&'static str, EncoderMaker)> = Vec::new();
217    #[cfg(target_os = "linux")]
218    out.push(("VA-API", vaapi::encoder));
219    #[cfg(any(
220        target_os = "windows",
221        all(target_os = "linux", target_pointer_width = "64")
222    ))]
223    out.push(("NVENC", nvenc::encoder));
224    #[cfg(target_os = "macos")]
225    out.push(("VideoToolbox", videotoolbox::encoder));
226    out
227}
228
229fn first<T>(makers: Vec<(&'static str, impl Fn() -> Result<T>)>, none: &str) -> Result<T> {
230    let mut why = Vec::new();
231    for (name, make) in makers {
232        match make() {
233            Ok(v) => return Ok(v),
234            Err(e) => why.push(format!("{name}: {e}")),
235        }
236    }
237    if why.is_empty() {
238        why.push(none.to_string());
239    }
240    Err(Error::Unsupported(why.join("; ")))
241}
242
243pub struct HwDecoder {
244    inner: Box<dyn DecoderBackend>,
245    time_base: TimeBase,
246}
247
248impl HwDecoder {
249    pub fn new(codec: Codec, extradata: &[u8]) -> Result<Self> {
250        let makers = decoders()
251            .into_iter()
252            .map(|(name, make)| (name, move || make(codec, extradata)))
253            .collect();
254        let inner = first(makers, "no GPU decoder on this system")?;
255        Ok(Self {
256            inner,
257            time_base: TimeBase::new(1, 90000),
258        })
259    }
260
261    pub fn with_time_base(mut self, time_base: TimeBase) -> Self {
262        self.time_base = time_base;
263        self
264    }
265
266    pub fn backend(&self) -> Backend {
267        self.inner.backend()
268    }
269
270    pub fn name(&self) -> &str {
271        self.inner.name()
272    }
273
274    pub fn decode(&mut self, packet: &Packet) -> Result<Vec<Frame>> {
275        let frames = self.inner.decode(packet)?;
276        Ok(self.stamp(frames))
277    }
278
279    pub fn flush(&mut self) -> Result<Vec<Frame>> {
280        let frames = self.inner.flush()?;
281        Ok(self.stamp(frames))
282    }
283
284    pub fn decode_gpu(&mut self, packet: &Packet) -> Result<Vec<GpuFrame>> {
285        let frames = self.inner.decode_gpu(packet)?;
286        Ok(self.stamp_gpu(frames))
287    }
288
289    pub fn flush_gpu(&mut self) -> Result<Vec<GpuFrame>> {
290        let frames = self.inner.flush_gpu()?;
291        Ok(self.stamp_gpu(frames))
292    }
293
294    pub fn reset(&mut self) {
295        self.inner.reset();
296    }
297
298    pub(crate) fn stills(&mut self) {
299        self.inner.stills();
300    }
301
302    fn stamp(&self, mut frames: Vec<Frame>) -> Vec<Frame> {
303        for f in &mut frames {
304            f.time_base = self.time_base;
305        }
306        frames
307    }
308
309    fn stamp_gpu(&self, mut frames: Vec<GpuFrame>) -> Vec<GpuFrame> {
310        for f in &mut frames {
311            f.time_base = self.time_base;
312        }
313        frames
314    }
315}
316
317pub(crate) enum GpuInner {
318    Cpu(Frame),
319    #[cfg(target_os = "linux")]
320    VaApi(vaapi::GpuSurface),
321}
322
323pub struct GpuFrame {
324    pub(crate) inner: GpuInner,
325    time_base: TimeBase,
326}
327
328impl GpuFrame {
329    pub(crate) fn cpu(frame: Frame) -> Self {
330        let time_base = frame.time_base;
331        Self {
332            inner: GpuInner::Cpu(frame),
333            time_base,
334        }
335    }
336
337    #[cfg(target_os = "linux")]
338    pub(crate) fn vaapi(surface: vaapi::GpuSurface) -> Self {
339        Self {
340            inner: GpuInner::VaApi(surface),
341            time_base: TimeBase::new(1, 90000),
342        }
343    }
344
345    pub fn pts(&self) -> i64 {
346        match &self.inner {
347            GpuInner::Cpu(f) => f.pts,
348            #[cfg(target_os = "linux")]
349            GpuInner::VaApi(g) => g.pts,
350        }
351    }
352
353    pub fn time_base(&self) -> TimeBase {
354        self.time_base
355    }
356
357    pub fn size(&self) -> (u32, u32) {
358        match &self.inner {
359            GpuInner::Cpu(f) => (f.width, f.height),
360            #[cfg(target_os = "linux")]
361            GpuInner::VaApi(g) => (g.width(), g.height()),
362        }
363    }
364
365    pub fn on_gpu(&self) -> bool {
366        !matches!(self.inner, GpuInner::Cpu(_))
367    }
368
369    pub fn surface(&self) -> Option<(u64, u32)> {
370        match &self.inner {
371            GpuInner::Cpu(_) => None,
372            #[cfg(target_os = "linux")]
373            GpuInner::VaApi(g) => Some(g.key()),
374        }
375    }
376
377    pub fn sync(&self) -> Result<()> {
378        match &self.inner {
379            GpuInner::Cpu(_) => Ok(()),
380            #[cfg(target_os = "linux")]
381            GpuInner::VaApi(g) => g.sync().map_err(Error::Backend),
382        }
383    }
384
385    #[cfg(target_os = "linux")]
386    pub fn dmabuf(&self, separate_layers: bool) -> Result<DmaBuf> {
387        match &self.inner {
388            GpuInner::VaApi(g) => g.dmabuf(separate_layers).map_err(Error::Backend),
389            GpuInner::Cpu(_) => Err(Error::Unsupported(
390                "this frame is in system memory, not on the GPU".into(),
391            )),
392        }
393    }
394
395    pub fn download(&self) -> Result<Frame> {
396        let frame = match &self.inner {
397            GpuInner::Cpu(f) => f.clone(),
398            #[cfg(target_os = "linux")]
399            GpuInner::VaApi(g) => g.download().map_err(Error::Backend)?,
400        };
401        Ok(Frame {
402            time_base: self.time_base,
403            ..frame
404        })
405    }
406}
407
408#[cfg_attr(
409    not(all(
410        feature = "h264",
411        feature = "hevc",
412        feature = "aac",
413        feature = "mp4",
414        feature = "filter",
415        feature = "mkv",
416        feature = "vp9"
417    )),
418    allow(dead_code)
419)]
420impl GpuFrame {
421    pub(crate) fn retime(&mut self, pts: i64, time_base: TimeBase) {
422        self.time_base = time_base;
423        match &mut self.inner {
424            GpuInner::Cpu(f) => {
425                f.pts = pts;
426                f.time_base = time_base;
427            }
428            #[cfg(target_os = "linux")]
429            GpuInner::VaApi(g) => g.pts = pts,
430        }
431    }
432
433    pub(crate) fn into_cpu(self) -> Option<Frame> {
434        match self.inner {
435            GpuInner::Cpu(f) => Some(f),
436            #[cfg(target_os = "linux")]
437            GpuInner::VaApi(_) => None,
438        }
439    }
440}
441
442pub struct GpuScaler {
443    #[cfg(target_os = "linux")]
444    inner: vaapi::Scaler,
445    size: (u32, u32),
446}
447
448impl GpuScaler {
449    pub fn new(width: u32, height: u32, ten_bit: bool) -> Result<Self> {
450        Self::with_frames(width, height, ten_bit, 8)
451    }
452
453    pub fn with_frames(width: u32, height: u32, ten_bit: bool, frames: usize) -> Result<Self> {
454        if width == 0 || height == 0 {
455            return Err(Error::Invalid("scaler size must be set".into()));
456        }
457        #[cfg(target_os = "linux")]
458        {
459            let bits = if ten_bit { 10 } else { 8 };
460            let inner =
461                vaapi::Scaler::new((width, height), bits, frames).map_err(Error::Backend)?;
462            Ok(Self {
463                inner,
464                size: (width, height),
465            })
466        }
467        #[cfg(not(target_os = "linux"))]
468        {
469            let _ = (ten_bit, frames);
470            Err(Error::Unsupported("no GPU scaler on this system".into()))
471        }
472    }
473
474    pub fn size(&self) -> (u32, u32) {
475        self.size
476    }
477
478    pub fn scale(&self, frame: &GpuFrame) -> Result<GpuFrame> {
479        match &frame.inner {
480            #[cfg(target_os = "linux")]
481            GpuInner::VaApi(g) => {
482                let out = self.inner.scale(g).map_err(Error::Backend)?;
483                Ok(GpuFrame {
484                    inner: GpuInner::VaApi(out),
485                    time_base: frame.time_base,
486                })
487            }
488            GpuInner::Cpu(_) => Err(Error::Unsupported(
489                "the GPU scaler takes frames on the GPU".into(),
490            )),
491        }
492    }
493}
494
495pub struct HwEncoder {
496    inner: Box<dyn EncoderBackend>,
497}
498
499impl HwEncoder {
500    pub fn new(codec: Codec, config: &EncoderConfig) -> Result<Self> {
501        let makers = encoders()
502            .into_iter()
503            .map(|(name, make)| (name, move || make(codec, config)))
504            .collect();
505        let inner = first(makers, "no GPU encoder on this system")?;
506        Ok(Self { inner })
507    }
508
509    pub fn backend(&self) -> Backend {
510        self.inner.backend()
511    }
512
513    pub fn name(&self) -> &str {
514        self.inner.name()
515    }
516
517    pub fn b_frames(&self) -> u8 {
518        self.inner.b_frames()
519    }
520
521    pub fn encode(&mut self, frame: &Frame) -> Result<Vec<Packet>> {
522        self.inner.encode(frame)
523    }
524
525    pub fn encode_gpu(&mut self, frame: &GpuFrame) -> Result<Vec<Packet>> {
526        self.inner.encode_gpu(frame)
527    }
528
529    pub fn composes(&self) -> bool {
530        self.inner.composes()
531    }
532
533    pub fn encode_gpu_over(&mut self, frame: &GpuFrame, layer: &Arc<Rgba>) -> Result<Vec<Packet>> {
534        self.inner.encode_gpu_over(frame, layer)
535    }
536
537    pub fn flush(&mut self) -> Result<Vec<Packet>> {
538        self.inner.flush()
539    }
540
541    pub fn extradata(&self) -> Option<&[u8]> {
542        self.inner.extradata()
543    }
544}
pito_codecs/gpu/mod.rs

pito_codecs/gpu/
mod.rs
