nvidia_video_codec_sdk/safe/
session.rs

1//! Defines [`Session`] which represents an ongoing encoder session.
2//!
3//! You need to start a session using [`Encoder::start_session`] before
4//! you can initialize input or output buffers, and before you can encode
5//! frames. The [`Session`] also stores some information such as the encode
6//! width and height so that you do not have to keep repeating it each time.
7
8use std::fmt::Debug;
9
10use super::{api::ENCODE_API, encoder::Encoder, result::EncodeError};
11use crate::{
12    sys::nvEncodeAPI::{
13        GUID,
14        NV_ENC_BUFFER_FORMAT,
15        NV_ENC_CODEC_AV1_GUID,
16        NV_ENC_CODEC_H264_GUID,
17        NV_ENC_CODEC_HEVC_GUID,
18        NV_ENC_CODEC_PIC_PARAMS,
19        NV_ENC_PIC_PARAMS,
20        NV_ENC_PIC_PARAMS_AV1,
21        NV_ENC_PIC_PARAMS_H264,
22        NV_ENC_PIC_PARAMS_HEVC,
23        NV_ENC_PIC_PARAMS_VER,
24        NV_ENC_PIC_STRUCT,
25        NV_ENC_PIC_TYPE,
26    },
27    EncoderInput,
28    EncoderOutput,
29};
30
31/// An encoding session to create input/output buffers and encode frames.
32///
33/// You need to call [`Encoder::start_session`] before you can
34/// encode frames using the session. On drop, the session will automatically
35/// send an empty EOS frame to flush the encoder.
36#[derive(Debug)]
37pub struct Session {
38    pub(crate) encoder: Encoder,
39    pub(crate) width: u32,
40    pub(crate) height: u32,
41    pub(crate) buffer_format: NV_ENC_BUFFER_FORMAT,
42    pub(crate) encode_guid: GUID,
43}
44
45impl Session {
46    /// Get the encoder used for this session.
47    ///
48    /// This might be useful if you want to use some of
49    /// the functions on [`Encoder`].
50    ///
51    /// # Examples
52    ///
53    /// ```
54    /// # use cudarc::driver::CudaContext;
55    /// # use nvidia_video_codec_sdk::{
56    /// #     sys::nvEncodeAPI::{
57    /// #         NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_ARGB,
58    /// #         NV_ENC_CODEC_H264_GUID,
59    /// #     },
60    /// #     Encoder, EncoderInitParams
61    /// # };
62    /// //* Create encoder. *//
63    /// # let cuda_ctx = CudaContext::new(0).unwrap();
64    /// # let encoder = Encoder::initialize_with_cuda(cuda_ctx).unwrap();
65    ///
66    /// //* Set `encode_guid` and check that H.264 encoding is supported. *//
67    /// # let encode_guid = NV_ENC_CODEC_H264_GUID;
68    /// # let encode_guids = encoder.get_encode_guids().unwrap();
69    /// # assert!(encode_guids.contains(&encode_guid));
70    ///
71    /// let session = encoder
72    ///     .start_session(
73    ///         NV_ENC_BUFFER_FORMAT_ARGB,
74    ///         EncoderInitParams::new(encode_guid, 1920, 1080),
75    ///     )
76    ///     .unwrap();
77    /// // We can still use the encoder like this:
78    /// let _input_formats = session
79    ///     .get_encoder()
80    ///     .get_supported_input_formats(encode_guid);
81    /// ```
82    #[must_use]
83    pub fn get_encoder(&self) -> &Encoder {
84        &self.encoder
85    }
86
87    /// Encode a frame.
88    ///
89    /// See [NVIDIA docs](https://docs.nvidia.com/video-technologies/video-codec-sdk/12.0/nvenc-video-encoder-api-prog-guide/index.html#submitting-input-frame-for-encoding).
90    ///
91    /// # Errors
92    ///
93    /// Could error if the encode picture parameters were invalid or otherwise
94    /// incorrect, or if we run out memory.
95    ///
96    /// There are two recoverable errors:
97    /// - If this returns an error with
98    ///   [`ErrorKind::EncoderBusy`](super::ErrorKind::EncoderBusy) then you
99    ///   should retry after a few milliseconds.
100    /// - If this returns an error with
101    ///   [`ErrorKind::NeedMoreInput`](super::ErrorKind::NeedMoreInput), the
102    ///   client should not lock the output bitstream yet. They should continue
103    ///   encoding until this function returns `Ok`, and then lock the
104    ///   bitstreams in the order in which they were originally used.
105    ///
106    /// # Panics
107    ///
108    /// Panics if codec specific parameters are provided for a different codec
109    /// than the one used in the session.
110    ///
111    /// # Examples
112    ///
113    /// ```
114    /// # use cudarc::driver::CudaContext;
115    /// # use nvidia_video_codec_sdk::{
116    /// #     sys::nvEncodeAPI::{
117    /// #         NV_ENC_BUFFER_FORMAT::NV_ENC_BUFFER_FORMAT_ARGB,
118    /// #         NV_ENC_CODEC_H264_GUID,
119    /// #         NV_ENC_PIC_PARAMS,
120    /// #         NV_ENC_PIC_STRUCT,
121    /// #     },
122    /// #     Encoder, EncoderInitParams,
123    /// #     EncodePictureParams
124    /// # };
125    /// # const WIDTH: u32 = 1920;
126    /// # const HEIGHT: u32 = 1080;
127    /// # const DATA_LEN: usize = (WIDTH * HEIGHT * 4) as usize;
128    /// //* Create encoder. *//
129    /// # let cuda_ctx = CudaContext::new(0).unwrap();
130    /// # let encoder = Encoder::initialize_with_cuda(cuda_ctx).unwrap();
131    ///
132    /// //* Set `encode_guid` and `buffer_format`, and check that H.264 encoding and the ARGB format are supported. *//
133    /// # let encode_guid = NV_ENC_CODEC_H264_GUID;
134    /// # let encode_guids = encoder.get_encode_guids().unwrap();
135    /// # assert!(encode_guids.contains(&encode_guid));
136    /// # let buffer_format = NV_ENC_BUFFER_FORMAT_ARGB;
137    /// # let input_formats = encoder.get_supported_input_formats(encode_guid).unwrap();
138    /// # assert!(input_formats.contains(&buffer_format));
139    ///
140    /// // Begin encoder session.
141    /// let mut initialize_params = EncoderInitParams::new(encode_guid, WIDTH, HEIGHT);
142    /// initialize_params.display_aspect_ratio(16, 9)
143    ///     .framerate(30, 1)
144    ///     .enable_picture_type_decision();
145    /// let session = encoder.start_session(
146    ///     buffer_format,
147    ///     initialize_params,
148    /// ).unwrap();
149    ///
150    /// //* Create input and output buffers. *//
151    /// # let mut input_buffer = session
152    /// #     .create_input_buffer()
153    /// #     .unwrap();
154    /// # let mut output_bitstream = session.create_output_bitstream().unwrap();
155    ///
156    /// // Encode frame.
157    /// unsafe { input_buffer.lock().unwrap().write(&[0; DATA_LEN]) };
158    /// session
159    ///     .encode_picture(
160    ///         &mut input_buffer,
161    ///         &mut output_bitstream,
162    ///         // Optional picture parameters
163    ///         EncodePictureParams {
164    ///             input_timestamp: 42,
165    ///             ..Default::default()
166    ///         }
167    ///     )
168    ///     .unwrap();
169    /// # // TODO: check that output is correct.
170    /// let _data = output_bitstream.lock().unwrap().data();
171    /// ```
172    pub fn encode_picture<I: EncoderInput, O: EncoderOutput>(
173        &self,
174        input_buffer: &mut I,
175        output_bitstream: &mut O,
176        params: EncodePictureParams,
177    ) -> Result<(), EncodeError> {
178        if let Some(codec_params) = &params.codec_params {
179            assert_eq!(
180                codec_params.get_codec_guid(),
181                self.encode_guid,
182                "The provided codec specific params must match the codec used"
183            );
184        }
185        let mut encode_pic_params = NV_ENC_PIC_PARAMS {
186            version: NV_ENC_PIC_PARAMS_VER,
187            inputWidth: self.width,
188            inputHeight: self.height,
189            inputPitch: input_buffer.pitch(),
190            inputBuffer: input_buffer.handle(),
191            outputBitstream: output_bitstream.handle(),
192            bufferFmt: self.buffer_format,
193            pictureStruct: NV_ENC_PIC_STRUCT::NV_ENC_PIC_STRUCT_FRAME,
194            inputTimeStamp: params.input_timestamp,
195            codecPicParams: params.codec_params.map(Into::into).unwrap_or_default(),
196            pictureType: params.picture_type,
197            ..Default::default()
198        };
199        unsafe { (ENCODE_API.encode_picture)(self.encoder.ptr, &mut encode_pic_params) }
200            .result(&self.encoder)
201    }
202
203    /// Send an EOS notifications to flush the encoder.
204    ///
205    /// This function is called automatically on drop, but if you wish to
206    /// get the data after flushing, you should call this function yourself.
207    ///
208    /// # Errors
209    ///
210    /// Could error if we run out of memory.
211    ///
212    /// If this returns an error with
213    /// [`ErrorKind::EncoderBusy`](super::ErrorKind::EncoderBusy) then you
214    /// should retry after a few milliseconds.
215    pub fn end_of_stream(&self) -> Result<(), EncodeError> {
216        let mut encode_pic_params = NV_ENC_PIC_PARAMS::end_of_stream();
217        unsafe { (ENCODE_API.encode_picture)(self.encoder.ptr, &mut encode_pic_params) }
218            .result(&self.encoder)
219    }
220}
221
222/// Send an EOS notifications on drop to flush the encoder.
223impl Drop for Session {
224    fn drop(&mut self) {
225        if !std::thread::panicking() {
226            self.end_of_stream()
227                .expect("The encoder should not be busy.");
228        }
229    }
230}
231
232/// Optional parameters for [`Session::encode_picture`].
233#[allow(missing_debug_implementations)] // CodecPictureParams doesn't implement Debug
234pub struct EncodePictureParams {
235    /// Opaque data used for identifying the corresponding encoded frame
236    pub input_timestamp: u64,
237    /// The picture type to use, if picture type decision is disabled in the
238    /// encoder
239    pub picture_type: NV_ENC_PIC_TYPE,
240    /// Codec-specific parameters
241    pub codec_params: Option<CodecPictureParams>,
242}
243
244impl Default for EncodePictureParams {
245    fn default() -> Self {
246        Self {
247            input_timestamp: 0,
248            picture_type: NV_ENC_PIC_TYPE::NV_ENC_PIC_TYPE_UNKNOWN,
249            codec_params: None,
250        }
251    }
252}
253
254/// Codec specific picture parameters
255#[allow(missing_debug_implementations)] // NV_ENC_PIC_PARAMS_H264 contains a union, thus doesn't derive Debug
256pub enum CodecPictureParams {
257    /// Parameters for H.264
258    H264(NV_ENC_PIC_PARAMS_H264),
259    /// Parameters for HEVC or H.265
260    Hevc(NV_ENC_PIC_PARAMS_HEVC),
261    /// Parameters for AV1
262    Av1(NV_ENC_PIC_PARAMS_AV1),
263}
264
265impl CodecPictureParams {
266    /// Returns the GUID representing the codec for which the parameters are
267    /// specified.
268    #[must_use]
269    pub fn get_codec_guid(&self) -> GUID {
270        match self {
271            Self::H264(_) => NV_ENC_CODEC_H264_GUID,
272            Self::Hevc(_) => NV_ENC_CODEC_HEVC_GUID,
273            Self::Av1(_) => NV_ENC_CODEC_AV1_GUID,
274        }
275    }
276}
277
278impl From<CodecPictureParams> for NV_ENC_CODEC_PIC_PARAMS {
279    fn from(value: CodecPictureParams) -> Self {
280        match value {
281            CodecPictureParams::H264(params) => Self {
282                h264PicParams: params,
283            },
284            CodecPictureParams::Hevc(params) => Self {
285                hevcPicParams: params,
286            },
287            CodecPictureParams::Av1(params) => Self {
288                av1PicParams: params,
289            },
290        }
291    }
292}