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) = ¶ms.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}