dynlink_nvidia_encode/
api.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use crate::ffi::*;
5use crate::load::SharedLibrary;
6use crate::{NvInt, NvencError};
7use std::ptr::addr_of_mut;
8use std::{
9    fmt::{Debug, Formatter},
10    mem::MaybeUninit,
11    pin::Pin,
12    rc::Rc,
13};
14
15macro_rules! api_call {
16    ($expr:expr_2021) => {{
17        let status = $expr;
18        if status != _NVENCSTATUS::NV_ENC_SUCCESS {
19            return Err(NvencError::ErrCode {
20                status,
21                fname: file!(),
22                line_num: line!(),
23                message: crate::error::code_to_string(status),
24            });
25        }
26    }};
27}
28
29macro_rules! load_func {
30    ($inner:expr_2021, $ident:ident) => {{
31        let func = if let Some(func) = $inner.$ident {
32            func
33        } else {
34            return Err(NvencError::NameFFIError {
35                name: stringify!($ident).to_string(),
36            });
37        };
38
39        Ok(func)
40    }};
41}
42
43macro_rules! get_func {
44    ($lib:expr_2021, $name:expr_2021) => {{
45        unsafe { $lib.library.get($name) }.map_err(|source| NvencError::NameFFIError2 {
46            name: String::from_utf8_lossy($name).to_string(),
47            source,
48        })
49        // format!(
50        //     "the name {} could not be opened: {}", String::from_utf8_lossy($name), e))?
51    }};
52}
53
54pub fn init(library: &SharedLibrary) -> Result<Rc<LibNvEncode<'_>>, NvencError> {
55    let lib_nv_encode = LibNvEncode {
56        NvEncodeAPICreateInstance: get_func!(library, b"NvEncodeAPICreateInstance\0")?,
57        NvEncodeAPIGetMaxSupportedVersion: get_func!(
58            library,
59            b"NvEncodeAPIGetMaxSupportedVersion\0"
60        )?,
61    };
62
63    Ok(Rc::new(lib_nv_encode))
64}
65
66#[expect(non_snake_case)]
67pub struct LibNvEncode<'lib> {
68    NvEncodeAPICreateInstance:
69        libloading::Symbol<'lib, extern "C" fn(*mut NV_ENCODE_API_FUNCTION_LIST) -> NVENCSTATUS>,
70    NvEncodeAPIGetMaxSupportedVersion:
71        libloading::Symbol<'lib, extern "C" fn(*mut u32) -> NVENCSTATUS>,
72}
73
74impl<'lib> LibNvEncode<'lib> {
75    pub fn api_create_instance(
76        self_: Rc<Self>,
77    ) -> Result<NvEncodeApiFunctionList<'lib>, NvencError> {
78        let function_list = MaybeUninit::zeroed();
79        let mut function_list: NV_ENCODE_API_FUNCTION_LIST = unsafe { function_list.assume_init() };
80
81        function_list.version = NV_ENCODE_API_FUNCTION_LIST_VER;
82
83        api_call!((*self_.NvEncodeAPICreateInstance)(&mut function_list));
84        Ok(NvEncodeApiFunctionList {
85            inner: function_list,
86            _libnvencode: self_.clone(),
87        })
88    }
89    pub fn api_get_max_supported_version(&self) -> Result<ApiVersion, NvencError> {
90        let mut value = 0;
91        api_call!((*self.NvEncodeAPIGetMaxSupportedVersion)(&mut value));
92        Ok(ApiVersion {
93            major: value >> 4,
94            minor: value & 0xf,
95        })
96    }
97}
98
99#[derive(Debug)]
100pub struct ApiVersion {
101    pub major: u32,
102    pub minor: u32,
103}
104
105/// The lifetime 'lib refers to the shared library.
106#[derive(Clone)]
107pub struct NvEncodeApiFunctionList<'lib> {
108    inner: NV_ENCODE_API_FUNCTION_LIST,
109    _libnvencode: Rc<LibNvEncode<'lib>>,
110}
111
112struct EncoderPtr(*mut std::ffi::c_void);
113
114impl<'lib> NvEncodeApiFunctionList<'lib> {
115    pub fn new_encoder(
116        &self,
117        mut ctx: dynlink_cuda::CudaContext,
118    ) -> Result<Rc<Encoder<'lib>>, NvencError> {
119        let func = load_func!(self.inner, nvEncOpenEncodeSessionEx)?;
120        let params = MaybeUninit::zeroed();
121        let mut params: NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS = unsafe { params.assume_init() };
122
123        params.version = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;
124        params.apiVersion = NVENCAPI_VERSION;
125        params.deviceType = _NV_ENC_DEVICE_TYPE::NV_ENC_DEVICE_TYPE_CUDA;
126        params.device = ctx.as_mut_void_ptr();
127
128        let mut encoder: *mut std::ffi::c_void = std::ptr::null_mut();
129        api_call!(unsafe { func(&mut params, &mut encoder) });
130        Ok(Rc::new(Encoder {
131            parent: self.clone(),
132            inner: EncoderPtr(encoder),
133            destroyed: false,
134        }))
135    }
136}
137
138/// The lifetime 'lib refers to the shared library.
139pub struct Encoder<'lib> {
140    parent: NvEncodeApiFunctionList<'lib>,
141    inner: EncoderPtr,
142    destroyed: bool,
143}
144
145impl<'lib> Encoder<'lib> {
146    pub fn get_encode_guid_count(&self) -> Result<u32, NvencError> {
147        let func = load_func!(self.parent.inner, nvEncGetEncodeGUIDCount)?;
148        let mut value = 0;
149        api_call!(unsafe { func(self.inner.0, &mut value) });
150        Ok(value)
151    }
152
153    pub fn get_encode_preset_config(
154        &self,
155        encode: GUID,
156        preset: GUID,
157    ) -> Result<EncodeConfig, NvencError> {
158        let func = load_func!(self.parent.inner, nvEncGetEncodePresetConfig)?;
159
160        let config = MaybeUninit::zeroed();
161        let mut config: NV_ENC_PRESET_CONFIG = unsafe { config.assume_init() };
162
163        config.presetCfg.version = NV_ENC_CONFIG_VER;
164        config.version = NV_ENC_PRESET_CONFIG_VER;
165
166        api_call!(unsafe { func(self.inner.0, encode, preset, &mut config) });
167        Ok(EncodeConfig {
168            config: config.presetCfg,
169        })
170    }
171
172    // TODO: return an InitializedEncoder type (forces the encoder to be initialized).
173    pub fn initialize(&self, init_params: &InitParams) -> Result<(), NvencError> {
174        // There seem to be under-documented minimum width requirements.
175        // https://forums.developer.nvidia.com/t/minimum-width-in-turing-gpus/155566
176        let func = load_func!(self.parent.inner, nvEncInitializeEncoder)?;
177        // We can safely assume the params won't be changed by the API
178        // according to the API documentation
179        let params = init_params.init_params;
180        let params = &params as *const NV_ENC_INITIALIZE_PARAMS;
181        let params = params as *mut NV_ENC_INITIALIZE_PARAMS;
182
183        api_call!(unsafe { func(self.inner.0, params) });
184
185        Ok(())
186    }
187
188    /// Allocate a new buffer managed by NVIDIA Video SDK
189    pub fn alloc_input_buffer(
190        self_: &Rc<Self>,
191        width: u32,
192        height: u32,
193        format: BufferFormat,
194    ) -> Result<InputBuffer<'lib>, NvencError> {
195        let func = load_func!(self_.parent.inner, nvEncCreateInputBuffer)?;
196
197        let params = MaybeUninit::zeroed();
198        let mut params: NV_ENC_CREATE_INPUT_BUFFER = unsafe { params.assume_init() };
199
200        params.version = NV_ENC_CREATE_INPUT_BUFFER_VER;
201        params.width = width;
202        params.height = height;
203        params.bufferFmt = format as NvInt;
204
205        api_call!(unsafe { func(self_.inner.0, &mut params) });
206
207        Ok(InputBuffer {
208            encoder: self_.clone(),
209            ptr: params.inputBuffer,
210            format,
211            width,
212            height,
213            destroyed: false,
214        })
215    }
216
217    pub fn alloc_output_buffer(self_: &Rc<Self>) -> Result<OutputBuffer<'lib>, NvencError> {
218        let func = load_func!(self_.parent.inner, nvEncCreateBitstreamBuffer)?;
219
220        let params = MaybeUninit::zeroed();
221        let mut params: NV_ENC_CREATE_BITSTREAM_BUFFER = unsafe { params.assume_init() };
222
223        params.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
224        api_call!(unsafe { func(self_.inner.0, &mut params) });
225        Ok(OutputBuffer {
226            encoder: self_.clone(),
227            ptr: params.bitstreamBuffer,
228            destroyed: false,
229        })
230    }
231
232    /// Main entry to encode a video frame with a given presentation time stamp.
233    ///
234    /// Note that since enablePTD is true, this may return
235    /// NV_ENC_ERR_NEED_MORE_INPUT which should not be treated as a fatal error.
236    pub fn encode_picture(
237        &self,
238        input: &InputBuffer,
239        output: &OutputBuffer,
240        pitch: usize,
241        pts: std::time::Duration,
242    ) -> Result<(), NvencError> {
243        let func = load_func!(self.parent.inner, nvEncEncodePicture)?;
244
245        let params = MaybeUninit::zeroed();
246        let mut params: NV_ENC_PIC_PARAMS = unsafe { params.assume_init() };
247
248        params.version = NV_ENC_PIC_PARAMS_VER;
249        params.inputTimeStamp = dur2raw(&pts);
250        params.inputBuffer = input.ptr;
251        params.bufferFmt = input.format as NvInt;
252        params.inputWidth = input.width;
253        params.inputHeight = input.height;
254        params.inputPitch = pitch as u32;
255        params.pictureStruct = _NV_ENC_PIC_STRUCT::NV_ENC_PIC_STRUCT_FRAME;
256        params.outputBitstream = output.ptr;
257
258        api_call!(unsafe { func(self.inner.0, &mut params) });
259        Ok(())
260    }
261
262    /// End the encoder stream
263    ///
264    /// According to the nvenc docs, this can be called multiple times.
265    pub fn end_stream(&self) -> Result<(), NvencError> {
266        let func = load_func!(self.parent.inner, nvEncEncodePicture)?;
267
268        let params = MaybeUninit::zeroed();
269        let mut params: NV_ENC_PIC_PARAMS = unsafe { params.assume_init() };
270
271        params.version = NV_ENC_PIC_PARAMS_VER;
272        params.encodePicFlags = _NV_ENC_PIC_FLAGS::NV_ENC_PIC_FLAG_EOS;
273        api_call!(unsafe { func(self.inner.0, &mut params) });
274        Ok(())
275    }
276
277    pub fn get_sequence_parameter_sets(&self) -> Result<Vec<u8>, NvencError> {
278        let func = load_func!(self.parent.inner, nvEncGetSequenceParams)?;
279
280        let mut buf: Vec<u8> = vec![0; 256];
281        let mut new_len: u32 = 0;
282
283        let mut params = NV_ENC_SEQUENCE_PARAM_PAYLOAD {
284            version: NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER,
285            inBufferSize: buf.len().try_into().unwrap(),
286            spsId: 0,
287            ppsId: 0,
288            spsppsBuffer: buf.as_mut_ptr() as *mut std::ffi::c_void,
289            outSPSPPSPayloadSize: &mut new_len,
290            reserved: [0; 250],
291            reserved2: [std::ptr::null_mut(); 64],
292        };
293
294        api_call!(unsafe { func(self.inner.0, &mut params) });
295
296        buf.truncate(new_len.try_into().unwrap());
297
298        Ok(buf)
299    }
300}
301
302const H264_RATE: u32 = 90_000;
303
304// same as std::time::Duration::from_secs_f64 in rust 1.38
305fn from_secs_f64(secs: f64) -> std::time::Duration {
306    let whole_secs = secs.floor() as u64;
307    let subsec_nanos = ((secs - whole_secs as f64) * 1e9).round() as u32;
308    std::time::Duration::new(whole_secs, subsec_nanos)
309}
310
311fn dur2raw(dur: &std::time::Duration) -> u64 {
312    (dur.as_secs_f64() * H264_RATE as f64).round() as u64
313}
314
315fn raw2dur(raw: u64) -> std::time::Duration {
316    from_secs_f64((raw as f64) / H264_RATE as f64)
317}
318
319#[test]
320fn test_timestamp_conversion() {
321    for expected in &[0, 1, 100, 100_000, 100_000_000] {
322        let dur = raw2dur(*expected);
323        let actual = dur2raw(&dur);
324        assert_eq!(*expected, actual);
325    }
326}
327
328impl Drop for Encoder<'_> {
329    fn drop(&mut self) {
330        if !self.destroyed {
331            let func = if let Some(func) = self.parent.inner.nvEncDestroyEncoder {
332                func
333            } else {
334                panic!("No function 'nvEncDestroyEncoder'");
335            };
336            let status = unsafe { func(self.inner.0) };
337            assert!(
338                status == _NVENCSTATUS::NV_ENC_SUCCESS,
339                "NV_ENC error code: {}",
340                status
341            );
342            self.destroyed = true;
343        }
344    }
345}
346
347/// A simple wrapper of a buffer
348pub struct InputBuffer<'lib> {
349    encoder: Rc<Encoder<'lib>>,
350    ptr: NV_ENC_INPUT_PTR,
351    format: BufferFormat,
352    width: u32,
353    height: u32,
354    destroyed: bool,
355}
356
357/// Acquired by calling `InputBuffer::lock()`
358///
359/// Implements Drop to automatically unlock the InputBuffer.
360pub struct LockedInputBuffer<'lock, 'lib> {
361    inner: &'lock InputBuffer<'lib>,
362    mem: &'lock mut [u8],
363    pitch: usize,
364    dropped: bool,
365}
366
367impl<'lib> InputBuffer<'lib> {
368    pub fn lock<'lock>(&'lock self) -> Result<LockedInputBuffer<'lock, 'lib>, NvencError> {
369        let func = load_func!(self.encoder.parent.inner, nvEncLockInputBuffer)?;
370
371        let params = MaybeUninit::zeroed();
372        let mut params: NV_ENC_LOCK_INPUT_BUFFER = unsafe { params.assume_init() };
373
374        params.version = NV_ENC_LOCK_INPUT_BUFFER_VER;
375        params.inputBuffer = self.ptr;
376
377        api_call!(unsafe { func(self.encoder.inner.0, &mut params) });
378
379        let sz = self.format.calculate_size(params.pitch, self.height)?;
380
381        let mem = unsafe { std::slice::from_raw_parts_mut(params.bufferDataPtr as *mut u8, sz) };
382
383        Ok(LockedInputBuffer {
384            inner: self,
385            mem,
386            pitch: params.pitch as usize,
387            dropped: false,
388        })
389    }
390}
391
392impl Drop for LockedInputBuffer<'_, '_> {
393    fn drop(&mut self) {
394        if !self.dropped {
395            let func = if let Some(func) = self.inner.encoder.parent.inner.nvEncUnlockInputBuffer {
396                func
397            } else {
398                panic!("No function 'nvEncUnlockInputBuffer'");
399            };
400
401            let status = unsafe { func(self.inner.encoder.inner.0, self.inner.ptr) };
402
403            assert!(
404                status == _NVENCSTATUS::NV_ENC_SUCCESS,
405                "NV_ENC error code: {}",
406                status
407            );
408
409            // As far as I understand it, slices (e.g. `self.mem` do not
410            // implement Drop, so we do not need to call `std::mem::forget`
411            // on our slice. Presumably the nvidia driver deallocates the
412            // backing memory in this case.
413
414            self.dropped = true;
415        }
416    }
417}
418
419impl LockedInputBuffer<'_, '_> {
420    pub fn mem(&self) -> &[u8] {
421        self.mem
422    }
423    pub fn mem_mut(&mut self) -> &mut [u8] {
424        self.mem
425    }
426    pub fn pitch(&self) -> usize {
427        self.pitch
428    }
429}
430
431impl Drop for InputBuffer<'_> {
432    fn drop(&mut self) {
433        if !self.destroyed {
434            let func = if let Some(func) = self.encoder.parent.inner.nvEncDestroyInputBuffer {
435                func
436            } else {
437                panic!("No function 'nvEncDestroyInputBuffer'");
438            };
439
440            let status = unsafe { func(self.encoder.inner.0, self.ptr) };
441            assert!(
442                status == _NVENCSTATUS::NV_ENC_SUCCESS,
443                "NV_ENC error code: {}",
444                status
445            );
446
447            self.destroyed = true;
448        }
449    }
450}
451
452pub struct OutputBuffer<'lib> {
453    encoder: Rc<Encoder<'lib>>,
454    ptr: NV_ENC_OUTPUT_PTR,
455    destroyed: bool,
456}
457
458/// Acquired by calling `OutputBuffer::lock()`
459///
460/// Implements Drop to automatically unlock the OutputBuffer.
461pub struct LockedOutputBuffer<'lock, 'lib> {
462    inner: &'lock OutputBuffer<'lib>,
463    mem: &'lock [u8],
464    picture_type: NvInt,
465    /// presentation timestamp (from onset)
466    pts: std::time::Duration,
467    output_time_stamp: u64,
468    output_duration: u64,
469    dropped: bool,
470}
471
472impl Drop for OutputBuffer<'_> {
473    fn drop(&mut self) {
474        if !self.destroyed {
475            let func = if let Some(func) = self.encoder.parent.inner.nvEncDestroyBitstreamBuffer {
476                func
477            } else {
478                panic!("No function 'nvEncDestroyBitstreamBuffer'");
479            };
480
481            let status = unsafe { func(self.encoder.inner.0, self.ptr) };
482            assert!(
483                status == _NVENCSTATUS::NV_ENC_SUCCESS,
484                "NV_ENC error code: {}",
485                status
486            );
487
488            self.destroyed = true;
489        }
490    }
491}
492
493impl LockedOutputBuffer<'_, '_> {
494    pub fn mem(&self) -> &[u8] {
495        self.mem
496    }
497    pub fn pts(&self) -> &std::time::Duration {
498        &self.pts
499    }
500    pub fn output_time_stamp(&self) -> u64 {
501        self.output_time_stamp
502    }
503    pub fn output_duration(&self) -> u64 {
504        self.output_duration
505    }
506    pub fn is_keyframe(&self) -> bool {
507        use crate::ffi::_NV_ENC_PIC_TYPE::*;
508        matches!(self.picture_type, NV_ENC_PIC_TYPE_I | NV_ENC_PIC_TYPE_IDR)
509    }
510}
511
512impl Drop for LockedOutputBuffer<'_, '_> {
513    fn drop(&mut self) {
514        if !self.dropped {
515            let func = if let Some(func) = self.inner.encoder.parent.inner.nvEncUnlockBitstream {
516                func
517            } else {
518                panic!("No function 'nvEncUnlockBitstream'");
519            };
520
521            let status = unsafe { func(self.inner.encoder.inner.0, self.inner.ptr) };
522
523            assert!(
524                status == _NVENCSTATUS::NV_ENC_SUCCESS,
525                "NV_ENC error code: {}",
526                status
527            );
528
529            // As far as I understand it, slices (e.g. `self.mem` do not
530            // implement Drop, so we do not need to call `std::mem::forget`
531            // on our slice. Presumably the nvidia driver deallocates the
532            // backing memory in this case.
533
534            self.dropped = true;
535        }
536    }
537}
538
539impl<'lib> OutputBuffer<'lib> {
540    pub fn lock<'lock>(&'lock self) -> Result<LockedOutputBuffer<'lock, 'lib>, NvencError> {
541        let func = load_func!(self.encoder.parent.inner, nvEncLockBitstream)?;
542
543        let params = MaybeUninit::zeroed();
544        let mut params: NV_ENC_LOCK_BITSTREAM = unsafe { params.assume_init() };
545
546        params.version = NV_ENC_LOCK_BITSTREAM_VER;
547        params.outputBitstream = self.ptr;
548
549        api_call!(unsafe { func(self.encoder.inner.0, &mut params) });
550
551        let output_time_stamp = params.outputTimeStamp;
552        let output_duration = params.outputDuration;
553        let pts = raw2dur(output_time_stamp);
554        let picture_type = params.pictureType;
555
556        let mem = unsafe {
557            std::slice::from_raw_parts(
558                params.bitstreamBufferPtr as *mut u8,
559                params.bitstreamSizeInBytes as usize,
560            )
561        };
562
563        Ok(LockedOutputBuffer {
564            inner: self,
565            mem,
566            pts,
567            output_time_stamp,
568            output_duration,
569            picture_type,
570            dropped: false,
571        })
572    }
573}
574
575/// Data format of input and output buffer
576#[repr(u32)]
577#[derive(Copy, Clone, Debug)]
578pub enum BufferFormat {
579    Undefined = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_UNDEFINED,
580    NV12 = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_NV12,
581    YV12 = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_YV12,
582    IYUV = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_IYUV,
583    YUV444 = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_YUV444,
584    YUV444_10Bit = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_YUV444_10BIT,
585    YUV420_10Bit = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_YUV420_10BIT,
586    ARGB = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_ARGB,
587    ARGB10 = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_ARGB10,
588    ABGR = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_ABGR,
589    AYUV = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_AYUV,
590    ABGR10 = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_ABGR10,
591    // U8 = _NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_U8,
592}
593
594impl BufferFormat {
595    fn calculate_size(&self, stride: u32, height: u32) -> Result<usize, NvencError> {
596        match self {
597            &BufferFormat::NV12 | &BufferFormat::YV12 | &BufferFormat::IYUV => {
598                Ok((stride as usize) * (height as usize) * 3 / 2)
599            }
600            &BufferFormat::ARGB => Ok((stride as usize) * (height as usize) * 4),
601            _ => Err(NvencError::UnableToComputeSize {}),
602        }
603    }
604}
605
606/// Parameters used to initialize the encoder
607pub struct InitParams {
608    init_params: NV_ENC_INITIALIZE_PARAMS,
609    encode_config: EncodeConfig,
610}
611
612impl Debug for InitParams {
613    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
614        let p = &self.init_params;
615        write!(
616            f,
617            "InitParams{{version: {}, encodeGUID: {}, encodeWidth: {}, encodeHeight: {}, darWidth: {}, darHeight: {}, enablePTD: {}, presetGUID: {}, encodeConfig: {:?}, frameRateNum: {}, frameRateDen: {}, maxEncodeWidth: {}, maxEncodeHeight: {} }}",
618            p.version,
619            guid_string(&p.encodeGUID),
620            p.encodeWidth,
621            p.encodeHeight,
622            p.darWidth,
623            p.darHeight,
624            p.enablePTD,
625            guid_string(&p.presetGUID),
626            self.encode_config,
627            p.frameRateNum,
628            p.frameRateDen,
629            p.maxEncodeWidth,
630            p.maxEncodeHeight,
631        )
632    }
633}
634
635fn guid_string(guid: &GUID) -> String {
636    format!(
637        "{{ 0x{:x}, 0x{:x}, 0x{:x}, {} }}",
638        guid.Data1,
639        guid.Data2,
640        guid.Data3,
641        arr_string(&guid.Data4)
642    )
643}
644
645fn arr_string(arr: &[u8; 8]) -> String {
646    format!(
647        "{{ 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x} }}",
648        arr[0], arr[1], arr[2], arr[3], arr[4], arr[5], arr[6], arr[7]
649    )
650}
651
652pub struct InitParamsBuilder {
653    init_params: std::mem::MaybeUninit<NV_ENC_INITIALIZE_PARAMS>,
654    encode_config: Option<EncodeConfig>,
655}
656
657impl InitParamsBuilder {
658    pub fn new(encode: GUID, width: u32, height: u32) -> Self {
659        let mut uninit = std::mem::MaybeUninit::<NV_ENC_INITIALIZE_PARAMS>::zeroed();
660
661        let ptr = uninit.as_mut_ptr();
662
663        unsafe {
664            addr_of_mut!((*ptr).version).write(NV_ENC_INITIALIZE_PARAMS_VER);
665            addr_of_mut!((*ptr).encodeGUID).write(encode);
666            addr_of_mut!((*ptr).encodeWidth).write(width);
667            addr_of_mut!((*ptr).encodeHeight).write(height);
668            addr_of_mut!((*ptr).darWidth).write(width);
669            addr_of_mut!((*ptr).darHeight).write(height);
670            addr_of_mut!((*ptr).enablePTD).write(1);
671            addr_of_mut!((*ptr).maxEncodeWidth).write(width);
672            addr_of_mut!((*ptr).maxEncodeHeight).write(height);
673        }
674        Self {
675            init_params: uninit,
676            encode_config: None,
677        }
678    }
679
680    // display aspect ratio width
681    pub fn dar_width(mut self, width: u32) -> Self {
682        let ptr = self.init_params.as_mut_ptr();
683        unsafe {
684            addr_of_mut!((*ptr).darWidth).write(width);
685        }
686        self
687    }
688
689    // display aspect ratio height
690    pub fn dar_height(mut self, height: u32) -> Self {
691        let ptr = self.init_params.as_mut_ptr();
692        unsafe {
693            addr_of_mut!((*ptr).darHeight).write(height);
694        }
695        self
696    }
697
698    pub fn preset_guid(mut self, preset: GUID) -> Self {
699        let ptr = self.init_params.as_mut_ptr();
700        unsafe {
701            addr_of_mut!((*ptr).presetGUID).write(preset);
702        }
703        self
704    }
705
706    pub fn set_encode_config(mut self, config: EncodeConfig) -> Self {
707        self.encode_config = Some(config);
708        // We will set the `(*ptr).encodeConfig` when `Self::build` is called.
709        self
710    }
711
712    /// Set the frame rate (numerator and denominator)
713    ///
714    /// Note: "The frame rate has no meaning in NVENC other than deciding rate
715    /// control parameters." <https://devtalk.nvidia.com/default/topic/1023473>
716    pub fn set_framerate(mut self, num: u32, den: u32) -> Self {
717        let ptr = self.init_params.as_mut_ptr();
718        unsafe {
719            addr_of_mut!((*ptr).frameRateNum).write(num);
720            addr_of_mut!((*ptr).frameRateDen).write(den);
721        }
722        self
723    }
724
725    pub fn build(self) -> Result<Pin<Box<InitParams>>, NvencError> {
726        let encode_config = match self.encode_config {
727            Some(c) => c,
728            None => {
729                return Err(NvencError::EncodeConfigRequired {});
730            }
731        };
732        let params = InitParams {
733            init_params: unsafe { self.init_params.assume_init() },
734            encode_config,
735        };
736        let mut boxed = Box::pin(params);
737
738        let ptr: *mut NV_ENC_CONFIG = &mut boxed.encode_config.config;
739        // we know this is safe because modifying a field doesn't move the whole struct
740        unsafe {
741            let mut_ref: Pin<&mut InitParams> = Pin::as_mut(&mut boxed);
742            Pin::get_unchecked_mut(mut_ref).init_params.encodeConfig = ptr;
743        }
744        Ok(boxed)
745    }
746}
747
748/// Encoder configuration for a encode session
749pub struct EncodeConfig {
750    config: NV_ENC_CONFIG,
751}
752
753impl Debug for EncodeConfig {
754    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
755        write!(
756            f,
757            "{{rcParams.rateControlMode: {}, rcParams.averageBitRate: {}, rcParams.maxBitRate: {} }}",
758            self.config.rcParams.rateControlMode,
759            self.config.rcParams.averageBitRate,
760            self.config.rcParams.maxBitRate,
761        )
762    }
763}
764
765impl EncodeConfig {
766    pub fn set_rate_control_mode(&mut self, mode: RateControlMode) {
767        self.config.rcParams.rateControlMode = mode.to_c();
768    }
769    pub fn set_average_bit_rate(&mut self, value: u32) {
770        self.config.rcParams.averageBitRate = value;
771    }
772    pub fn set_max_bit_rate(&mut self, value: u32) {
773        self.config.rcParams.maxBitRate = value;
774    }
775}
776
777#[derive(Clone, Copy, Debug)]
778pub enum RateControlMode {
779    /// Constant QP mode
780    Constqp,
781    /// Variable bitrate mode
782    Vbr,
783    /// Constant bitrate mode
784    Cbr,
785    /// low-delay CBR, high quality
786    LowdelayHq,
787    /// CBR, high quality (slower)
788    CbrHq,
789    /// VBR, high quality (slower)
790    VbrHq,
791}
792
793impl RateControlMode {
794    fn to_c(self) -> NvInt {
795        use RateControlMode::*;
796        match self {
797            Constqp => _NV_ENC_PARAMS_RC_MODE::NV_ENC_PARAMS_RC_CONSTQP,
798            Vbr => _NV_ENC_PARAMS_RC_MODE::NV_ENC_PARAMS_RC_VBR,
799            Cbr => _NV_ENC_PARAMS_RC_MODE::NV_ENC_PARAMS_RC_CBR,
800            LowdelayHq => _NV_ENC_PARAMS_RC_MODE::NV_ENC_PARAMS_RC_CBR_LOWDELAY_HQ,
801            CbrHq => _NV_ENC_PARAMS_RC_MODE::NV_ENC_PARAMS_RC_CBR_HQ,
802            VbrHq => _NV_ENC_PARAMS_RC_MODE::NV_ENC_PARAMS_RC_VBR_HQ,
803        }
804    }
805}
dynlink_nvidia_encode/api.rs

dynlink_nvidia_encode/
api.rs
