7a97f7f61a1e655ded27ca58f16102693e4b2657
[platform/upstream/gstreamer.git] / sys / nvcodec / gstnvbaseenc.c
1 /* GStreamer NVENC plugin
2  * Copyright (C) 2015 Centricular Ltd
3  *
4  * This library is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU Library General Public
6  * License as published by the Free Software Foundation; either
7  * version 2 of the License, or (at your option) any later version.
8  *
9  * This library is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12  * Library General Public License for more details.
13  *
14  * You should have received a copy of the GNU Library General Public
15  * License along with this library; if not, write to the
16  * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
17  * Boston, MA 02110-1301, USA.
18  */
19
20 #ifdef HAVE_CONFIG_H
21 #include "config.h"
22 #endif
23
24 #include "gstnvbaseenc.h"
25 #include "gstcudautils.h"
26
27 #include <gst/pbutils/codec-utils.h>
28
29 #include <string.h>
30
31 #define GST_CAT_DEFAULT gst_nvenc_debug
32
33 #if HAVE_NVCODEC_GST_GL
34 #include <gst/gl/gl.h>
35 #endif
36
37 /* This currently supports both 5.x and 6.x versions of the NvEncodeAPI.h
38  * header which are mostly API compatible. */
39
40 #define SUPPORTED_GL_APIS GST_GL_API_OPENGL3
41
42 /* magic pointer value we can put in the async queue to signal shut down */
43 #define SHUTDOWN_COOKIE ((gpointer)GINT_TO_POINTER (1))
44
45 #define parent_class gst_nv_base_enc_parent_class
46 G_DEFINE_ABSTRACT_TYPE (GstNvBaseEnc, gst_nv_base_enc, GST_TYPE_VIDEO_ENCODER);
47
48 #define GST_TYPE_NV_PRESET (gst_nv_preset_get_type())
49 static GType
50 gst_nv_preset_get_type (void)
51 {
52   static GType nv_preset_type = 0;
53
54   static const GEnumValue presets[] = {
55     {GST_NV_PRESET_DEFAULT, "Default", "default"},
56     {GST_NV_PRESET_HP, "High Performance", "hp"},
57     {GST_NV_PRESET_HQ, "High Quality", "hq"},
58 /*    {GST_NV_PRESET_BD, "BD", "bd"}, */
59     {GST_NV_PRESET_LOW_LATENCY_DEFAULT, "Low Latency", "low-latency"},
60     {GST_NV_PRESET_LOW_LATENCY_HQ, "Low Latency, High Quality",
61         "low-latency-hq"},
62     {GST_NV_PRESET_LOW_LATENCY_HP, "Low Latency, High Performance",
63         "low-latency-hp"},
64     {GST_NV_PRESET_LOSSLESS_DEFAULT, "Lossless", "lossless"},
65     {GST_NV_PRESET_LOSSLESS_HP, "Lossless, High Performance", "lossless-hp"},
66     {0, NULL, NULL},
67   };
68
69   if (!nv_preset_type) {
70     nv_preset_type = g_enum_register_static ("GstNvPreset", presets);
71   }
72   return nv_preset_type;
73 }
74
75 static GUID
76 _nv_preset_to_guid (GstNvPreset preset)
77 {
78   GUID null = { 0, };
79
80   switch (preset) {
81 #define CASE(gst,nv) case G_PASTE(GST_NV_PRESET_,gst): return G_PASTE(G_PASTE(NV_ENC_PRESET_,nv),_GUID)
82       CASE (DEFAULT, DEFAULT);
83       CASE (HP, HP);
84       CASE (HQ, HQ);
85 /*    CASE (BD, BD);*/
86       CASE (LOW_LATENCY_DEFAULT, LOW_LATENCY_DEFAULT);
87       CASE (LOW_LATENCY_HQ, LOW_LATENCY_HQ);
88       CASE (LOW_LATENCY_HP, LOW_LATENCY_HQ);
89       CASE (LOSSLESS_DEFAULT, LOSSLESS_DEFAULT);
90       CASE (LOSSLESS_HP, LOSSLESS_HP);
91 #undef CASE
92     default:
93       return null;
94   }
95 }
96
97 #define GST_TYPE_NV_RC_MODE (gst_nv_rc_mode_get_type())
98 static GType
99 gst_nv_rc_mode_get_type (void)
100 {
101   static GType nv_rc_mode_type = 0;
102
103   static const GEnumValue modes[] = {
104     {GST_NV_RC_MODE_DEFAULT, "Default", "default"},
105     {GST_NV_RC_MODE_CONSTQP, "Constant Quantization", "constqp"},
106     {GST_NV_RC_MODE_CBR, "Constant Bit Rate", "cbr"},
107     {GST_NV_RC_MODE_VBR, "Variable Bit Rate", "vbr"},
108     {GST_NV_RC_MODE_VBR_MINQP,
109         "Variable Bit Rate "
110           "(with minimum quantization parameter, DEPRECATED)", "vbr-minqp"},
111     {GST_NV_RC_MODE_CBR_LOWDELAY_HQ,
112         "Low-Delay CBR, High Quality", "cbr-ld-hq"},
113     {GST_NV_RC_MODE_CBR_HQ, "CBR, High Quality (slower)", "cbr-hq"},
114     {GST_NV_RC_MODE_VBR_HQ, "VBR, High Quality (slower)", "vbr-hq"},
115     {0, NULL, NULL},
116   };
117
118   if (!nv_rc_mode_type) {
119     nv_rc_mode_type = g_enum_register_static ("GstNvRCMode", modes);
120   }
121   return nv_rc_mode_type;
122 }
123
124 static NV_ENC_PARAMS_RC_MODE
125 _rc_mode_to_nv (GstNvRCMode mode)
126 {
127   switch (mode) {
128     case GST_NV_RC_MODE_DEFAULT:
129       return NV_ENC_PARAMS_RC_VBR;
130 #define CASE(gst,nv) case G_PASTE(GST_NV_RC_MODE_,gst): return G_PASTE(NV_ENC_PARAMS_RC_,nv)
131       CASE (CONSTQP, CONSTQP);
132       CASE (CBR, CBR);
133       CASE (VBR, VBR);
134       CASE (VBR_MINQP, VBR_MINQP);
135       CASE (CBR_LOWDELAY_HQ, CBR_LOWDELAY_HQ);
136       CASE (CBR_HQ, CBR_HQ);
137       CASE (VBR_HQ, VBR_HQ);
138 #undef CASE
139     default:
140       return NV_ENC_PARAMS_RC_VBR;
141   }
142 }
143
144 enum
145 {
146   PROP_0,
147   PROP_DEVICE_ID,
148   PROP_PRESET,
149   PROP_BITRATE,
150   PROP_RC_MODE,
151   PROP_QP_MIN,
152   PROP_QP_MAX,
153   PROP_QP_CONST,
154   PROP_GOP_SIZE,
155   PROP_MAX_BITRATE,
156   PROP_SPATIAL_AQ,
157   PROP_AQ_STRENGTH,
158   PROP_NON_REF_P,
159   PROP_ZEROLATENCY,
160   PROP_STRICT_GOP,
161   PROP_CONST_QUALITY,
162   PROP_I_ADAPT,
163   PROP_QP_MIN_I,
164   PROP_QP_MIN_P,
165   PROP_QP_MIN_B,
166   PROP_QP_MAX_I,
167   PROP_QP_MAX_P,
168   PROP_QP_MAX_B,
169   PROP_QP_CONST_I,
170   PROP_QP_CONST_P,
171   PROP_QP_CONST_B,
172 };
173
174 #define DEFAULT_PRESET GST_NV_PRESET_DEFAULT
175 #define DEFAULT_BITRATE 0
176 #define DEFAULT_RC_MODE GST_NV_RC_MODE_DEFAULT
177 #define DEFAULT_QP_MIN -1
178 #define DEFAULT_QP_MAX -1
179 #define DEFAULT_QP_CONST -1
180 #define DEFAULT_GOP_SIZE 75
181 #define DEFAULT_MAX_BITRATE 0
182 #define DEFAULT_SPATIAL_AQ FALSE
183 #define DEFAULT_AQ_STRENGTH 0
184 #define DEFAULT_NON_REF_P FALSE
185 #define DEFAULT_ZEROLATENCY FALSE
186 #define DEFAULT_STRICT_GOP FALSE
187 #define DEFAULT_CONST_QUALITY 0
188 #define DEFAULT_I_ADAPT FALSE
189 #define DEFAULT_QP_DETAIL -1
190
191 /* This lock is needed to prevent the situation where multiple encoders are
192  * initialised at the same time which appears to cause excessive CPU usage over
193  * some period of time. */
194 G_LOCK_DEFINE_STATIC (initialization_lock);
195
196 typedef struct
197 {
198   /* Allocated CUDA device memory and registered to NVENC to be used as input
199    * buffer regardless of the input memory type (OpenGL or System memory) */
200   CUdeviceptr cuda_pointer;
201
202   /* The stride of allocated CUDA device memory (CuMemAllocPitch).
203    * This might be different from the stride of GstVideoInfo */
204   gsize cuda_stride;
205
206   /* Registered NVENC resource (cuda_pointer is used for this) */
207   NV_ENC_REGISTER_RESOURCE nv_resource;
208
209   /* Mapped resource of nv_resource */
210   NV_ENC_MAP_INPUT_RESOURCE nv_mapped_resource;
211
212   /* whether nv_mapped_resource was mapped via NvEncMapInputResource()
213    * and therefore should unmap via NvEncUnmapInputResource or not */
214   gboolean mapped;
215 } GstNvEncInputResource;
216
217 /* The pair of GstNvEncInputResource () and NV_ENC_OUTPUT_PTR.
218  * The number of input/output resource are always identical */
219 typedef struct
220 {
221   GstNvEncInputResource *in_buf;
222   NV_ENC_OUTPUT_PTR out_buf;
223 } GstNvEncFrameState;
224
225 static gboolean gst_nv_base_enc_open (GstVideoEncoder * enc);
226 static gboolean gst_nv_base_enc_close (GstVideoEncoder * enc);
227 static gboolean gst_nv_base_enc_start (GstVideoEncoder * enc);
228 static gboolean gst_nv_base_enc_stop (GstVideoEncoder * enc);
229 static void gst_nv_base_enc_set_context (GstElement * element,
230     GstContext * context);
231 static gboolean gst_nv_base_enc_sink_query (GstVideoEncoder * enc,
232     GstQuery * query);
233 static gboolean gst_nv_base_enc_sink_event (GstVideoEncoder * enc,
234     GstEvent * event);
235 static gboolean gst_nv_base_enc_set_format (GstVideoEncoder * enc,
236     GstVideoCodecState * state);
237 static GstFlowReturn gst_nv_base_enc_handle_frame (GstVideoEncoder * enc,
238     GstVideoCodecFrame * frame);
239 static void gst_nv_base_enc_free_buffers (GstNvBaseEnc * nvenc);
240 static GstFlowReturn gst_nv_base_enc_finish (GstVideoEncoder * enc);
241 static void gst_nv_base_enc_set_property (GObject * object, guint prop_id,
242     const GValue * value, GParamSpec * pspec);
243 static void gst_nv_base_enc_get_property (GObject * object, guint prop_id,
244     GValue * value, GParamSpec * pspec);
245 static void gst_nv_base_enc_finalize (GObject * obj);
246 static GstCaps *gst_nv_base_enc_getcaps (GstVideoEncoder * enc,
247     GstCaps * filter);
248 static gboolean gst_nv_base_enc_stop_bitstream_thread (GstNvBaseEnc * nvenc,
249     gboolean force);
250 static gboolean gst_nv_base_enc_drain_encoder (GstNvBaseEnc * nvenc);
251
252 static void
253 gst_nv_base_enc_class_init (GstNvBaseEncClass * klass)
254 {
255   GObjectClass *gobject_class = G_OBJECT_CLASS (klass);
256   GstElementClass *element_class = GST_ELEMENT_CLASS (klass);
257   GstVideoEncoderClass *videoenc_class = GST_VIDEO_ENCODER_CLASS (klass);
258
259   gobject_class->set_property = gst_nv_base_enc_set_property;
260   gobject_class->get_property = gst_nv_base_enc_get_property;
261   gobject_class->finalize = gst_nv_base_enc_finalize;
262
263   element_class->set_context = GST_DEBUG_FUNCPTR (gst_nv_base_enc_set_context);
264
265   videoenc_class->open = GST_DEBUG_FUNCPTR (gst_nv_base_enc_open);
266   videoenc_class->close = GST_DEBUG_FUNCPTR (gst_nv_base_enc_close);
267
268   videoenc_class->start = GST_DEBUG_FUNCPTR (gst_nv_base_enc_start);
269   videoenc_class->stop = GST_DEBUG_FUNCPTR (gst_nv_base_enc_stop);
270
271   videoenc_class->set_format = GST_DEBUG_FUNCPTR (gst_nv_base_enc_set_format);
272   videoenc_class->getcaps = GST_DEBUG_FUNCPTR (gst_nv_base_enc_getcaps);
273   videoenc_class->handle_frame =
274       GST_DEBUG_FUNCPTR (gst_nv_base_enc_handle_frame);
275   videoenc_class->finish = GST_DEBUG_FUNCPTR (gst_nv_base_enc_finish);
276   videoenc_class->sink_query = GST_DEBUG_FUNCPTR (gst_nv_base_enc_sink_query);
277   videoenc_class->sink_event = GST_DEBUG_FUNCPTR (gst_nv_base_enc_sink_event);
278
279   g_object_class_install_property (gobject_class, PROP_DEVICE_ID,
280       g_param_spec_uint ("cuda-device-id",
281           "Cuda Device ID",
282           "Get the GPU device to use for operations",
283           0, G_MAXUINT, 0, G_PARAM_READABLE | G_PARAM_STATIC_STRINGS));
284   g_object_class_install_property (gobject_class, PROP_PRESET,
285       g_param_spec_enum ("preset", "Encoding Preset",
286           "Encoding Preset",
287           GST_TYPE_NV_PRESET, DEFAULT_PRESET,
288           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
289           G_PARAM_STATIC_STRINGS));
290   g_object_class_install_property (gobject_class, PROP_RC_MODE,
291       g_param_spec_enum ("rc-mode", "RC Mode", "Rate Control Mode",
292           GST_TYPE_NV_RC_MODE, DEFAULT_RC_MODE,
293           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
294           G_PARAM_STATIC_STRINGS));
295   g_object_class_install_property (gobject_class, PROP_QP_MIN,
296       g_param_spec_int ("qp-min", "Minimum Quantizer",
297           "Minimum quantizer (-1 = from NVENC preset)", -1, 51, DEFAULT_QP_MIN,
298           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
299           G_PARAM_STATIC_STRINGS));
300   g_object_class_install_property (gobject_class, PROP_QP_MAX,
301       g_param_spec_int ("qp-max", "Maximum Quantizer",
302           "Maximum quantizer (-1 = from NVENC preset)", -1, 51, DEFAULT_QP_MAX,
303           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
304           G_PARAM_STATIC_STRINGS));
305   g_object_class_install_property (gobject_class, PROP_QP_CONST,
306       g_param_spec_int ("qp-const", "Constant Quantizer",
307           "Constant quantizer (-1 = from NVENC preset)", -1, 51,
308           DEFAULT_QP_CONST,
309           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
310           G_PARAM_STATIC_STRINGS));
311   g_object_class_install_property (gobject_class, PROP_GOP_SIZE,
312       g_param_spec_int ("gop-size", "GOP size",
313           "Number of frames between intra frames (-1 = infinite)",
314           -1, G_MAXINT, DEFAULT_GOP_SIZE,
315           (GParamFlags) (G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
316               G_PARAM_STATIC_STRINGS)));
317   g_object_class_install_property (gobject_class, PROP_BITRATE,
318       g_param_spec_uint ("bitrate", "Bitrate",
319           "Bitrate in kbit/sec (0 = from NVENC preset)", 0, 2000 * 1024,
320           DEFAULT_BITRATE,
321           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
322           G_PARAM_STATIC_STRINGS));
323   g_object_class_install_property (gobject_class, PROP_MAX_BITRATE,
324       g_param_spec_uint ("max-bitrate", "Max Bitrate",
325           "Maximum Bitrate in kbit/sec (ignored for CBR mode)", 0, 2000 * 1024,
326           DEFAULT_MAX_BITRATE,
327           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
328           G_PARAM_STATIC_STRINGS));
329   g_object_class_install_property (gobject_class, PROP_SPATIAL_AQ,
330       g_param_spec_boolean ("spatial-aq", "Spatial AQ",
331           "Spatial Adaptive Quantization",
332           DEFAULT_SPATIAL_AQ,
333           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
334           G_PARAM_STATIC_STRINGS));
335   g_object_class_install_property (gobject_class, PROP_AQ_STRENGTH,
336       g_param_spec_uint ("aq-strength", "AQ Strength",
337           "Adaptive Quantization Strength when spatial-aq is enabled"
338           " from 1 (low) to 15 (aggressive), (0 = autoselect)",
339           0, 15, DEFAULT_AQ_STRENGTH,
340           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
341           G_PARAM_STATIC_STRINGS));
342   g_object_class_install_property (gobject_class, PROP_NON_REF_P,
343       g_param_spec_boolean ("nonref-p", "Nonref P",
344           "Automatic insertion of non-reference P-frames", DEFAULT_NON_REF_P,
345           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
346           G_PARAM_STATIC_STRINGS));
347   g_object_class_install_property (gobject_class, PROP_ZEROLATENCY,
348       g_param_spec_boolean ("zerolatency", "Zerolatency",
349           "Zero latency operation (no reordering delay)", DEFAULT_ZEROLATENCY,
350           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
351           G_PARAM_STATIC_STRINGS));
352   g_object_class_install_property (gobject_class, PROP_STRICT_GOP,
353       g_param_spec_boolean ("strict-gop", "Strict GOP",
354           "Minimize GOP-to-GOP rate fluctuations", DEFAULT_STRICT_GOP,
355           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
356           G_PARAM_STATIC_STRINGS));
357   g_object_class_install_property (gobject_class, PROP_CONST_QUALITY,
358       g_param_spec_double ("const-quality", "Constant Quality",
359           "Target Constant Quality level for VBR mode (0 = automatic)",
360           0, 51, DEFAULT_CONST_QUALITY,
361           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
362           G_PARAM_STATIC_STRINGS));
363   g_object_class_install_property (gobject_class, PROP_I_ADAPT,
364       g_param_spec_boolean ("i-adapt", "I Adapt",
365           "Enable adaptive I-frame insert when lookahead is enabled",
366           DEFAULT_I_ADAPT,
367           G_PARAM_READWRITE | GST_PARAM_MUTABLE_READY |
368           G_PARAM_STATIC_STRINGS));
369   g_object_class_install_property (gobject_class, PROP_QP_MIN_I,
370       g_param_spec_int ("qp-min-i", "QP Min I",
371           "Minimum QP value for I frame, When >= 0, \"qp-min-p\" and "
372           "\"qp-min-b\" should be also >= 0. Overwritten by \"qp-min\""
373           " (-1 = from NVENC preset)", -1, 51,
374           DEFAULT_QP_DETAIL,
375           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
376           G_PARAM_STATIC_STRINGS));
377   g_object_class_install_property (gobject_class, PROP_QP_MIN_P,
378       g_param_spec_int ("qp-min-p", "QP Min P",
379           "Minimum QP value for P frame, When >= 0, \"qp-min-i\" and "
380           "\"qp-min-b\" should be also >= 0. Overwritten by \"qp-min\""
381           " (-1 = from NVENC preset)", -1, 51,
382           DEFAULT_QP_DETAIL,
383           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
384           G_PARAM_STATIC_STRINGS));
385   g_object_class_install_property (gobject_class, PROP_QP_MIN_B,
386       g_param_spec_int ("qp-min-b", "QP Min B",
387           "Minimum QP value for B frame, When >= 0, \"qp-min-i\" and "
388           "\"qp-min-p\" should be also >= 0. Overwritten by \"qp-min\""
389           " (-1 = from NVENC preset)", -1, 51,
390           DEFAULT_QP_DETAIL,
391           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
392           G_PARAM_STATIC_STRINGS));
393   g_object_class_install_property (gobject_class, PROP_QP_MAX_I,
394       g_param_spec_int ("qp-max-i", "QP Max I",
395           "Maximum QP value for I frame, When >= 0, \"qp-max-p\" and "
396           "\"qp-max-b\" should be also >= 0. Overwritten by \"qp-max\""
397           " (-1 = from NVENC preset)", -1, 51,
398           DEFAULT_QP_DETAIL,
399           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
400           G_PARAM_STATIC_STRINGS));
401   g_object_class_install_property (gobject_class, PROP_QP_MAX_P,
402       g_param_spec_int ("qp-max-p", "QP Max P",
403           "Maximum QP value for P frame, When >= 0, \"qp-max-i\" and "
404           "\"qp-max-b\" should be also >= 0. Overwritten by \"qp-max\""
405           " (-1 = from NVENC preset)", -1, 51,
406           DEFAULT_QP_DETAIL,
407           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
408           G_PARAM_STATIC_STRINGS));
409   g_object_class_install_property (gobject_class, PROP_QP_MAX_B,
410       g_param_spec_int ("qp-max-b", "QP Max B",
411           "Maximum QP value for B frame, When >= 0, \"qp-max-i\" and "
412           "\"qp-max-p\" should be also >= 0. Overwritten by \"qp-max\""
413           " (-1 = from NVENC preset)", -1, 51,
414           DEFAULT_QP_DETAIL,
415           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
416           G_PARAM_STATIC_STRINGS));
417   g_object_class_install_property (gobject_class, PROP_QP_CONST_I,
418       g_param_spec_int ("qp-const-i", "QP Const I",
419           "Constant QP value for I frame, When >= 0, \"qp-const-p\" and "
420           "\"qp-const-b\" should be also >= 0. Overwritten by \"qp-const\""
421           " (-1 = from NVENC preset)", -1, 51,
422           DEFAULT_QP_DETAIL,
423           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
424           G_PARAM_STATIC_STRINGS));
425   g_object_class_install_property (gobject_class, PROP_QP_CONST_P,
426       g_param_spec_int ("qp-const-p", "QP Const P",
427           "Constant QP value for P frame, When >= 0, \"qp-const-i\" and "
428           "\"qp-const-b\" should be also >= 0. Overwritten by \"qp-const\""
429           " (-1 = from NVENC preset)", -1, 51,
430           DEFAULT_QP_DETAIL,
431           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
432           G_PARAM_STATIC_STRINGS));
433   g_object_class_install_property (gobject_class, PROP_QP_CONST_B,
434       g_param_spec_int ("qp-const-b", "QP Const B",
435           "Constant QP value for B frame, When >= 0, \"qp-const-i\" and "
436           "\"qp-const-p\" should be also >= 0. Overwritten by \"qp-const\""
437           " (-1 = from NVENC preset)", -1, 51,
438           DEFAULT_QP_DETAIL,
439           G_PARAM_READWRITE | GST_PARAM_MUTABLE_PLAYING |
440           G_PARAM_STATIC_STRINGS));
441 }
442
443 static gboolean
444 gst_nv_base_enc_open_encode_session (GstNvBaseEnc * nvenc)
445 {
446   NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS params = { 0, };
447   NVENCSTATUS nv_ret;
448
449   params.version = gst_nvenc_get_open_encode_session_ex_params_version ();
450   params.apiVersion = gst_nvenc_get_api_version ();
451   params.device = gst_cuda_context_get_handle (nvenc->cuda_ctx);
452   params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
453   nv_ret = NvEncOpenEncodeSessionEx (&params, &nvenc->encoder);
454
455   return nv_ret == NV_ENC_SUCCESS;
456 }
457
458 static gboolean
459 gst_nv_base_enc_open (GstVideoEncoder * enc)
460 {
461   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
462   GstNvBaseEncClass *klass = GST_NV_BASE_ENC_GET_CLASS (enc);
463   GValue *formats = NULL;
464   CUresult cuda_ret;
465
466   if (!gst_cuda_ensure_element_context (GST_ELEMENT_CAST (enc),
467           klass->cuda_device_id, &nvenc->cuda_ctx)) {
468     GST_ERROR_OBJECT (nvenc, "failed to create CUDA context");
469     return FALSE;
470   }
471
472   if (gst_cuda_context_push (nvenc->cuda_ctx)) {
473     cuda_ret = CuStreamCreate (&nvenc->cuda_stream, CU_STREAM_DEFAULT);
474     if (!gst_cuda_result (cuda_ret)) {
475       GST_WARNING_OBJECT (nvenc,
476           "Could not create cuda stream, will use default stream");
477       nvenc->cuda_stream = NULL;
478     }
479     gst_cuda_context_pop (NULL);
480   }
481
482   if (!gst_nv_base_enc_open_encode_session (nvenc)) {
483     GST_ERROR ("Failed to create NVENC encoder session");
484     gst_clear_object (&nvenc->cuda_ctx);
485     return FALSE;
486   }
487
488   GST_INFO ("created NVENC encoder %p", nvenc->encoder);
489
490   /* query supported input formats */
491   if (!gst_nvenc_get_supported_input_formats (nvenc->encoder, klass->codec_id,
492           &formats)) {
493     GST_WARNING_OBJECT (nvenc, "No supported input formats");
494     gst_nv_base_enc_close (enc);
495     return FALSE;
496   }
497
498   nvenc->input_formats = formats;
499
500   return TRUE;
501 }
502
503 static void
504 gst_nv_base_enc_set_context (GstElement * element, GstContext * context)
505 {
506   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (element);
507   GstNvBaseEncClass *klass = GST_NV_BASE_ENC_GET_CLASS (nvenc);
508
509   if (gst_cuda_handle_set_context (element, context, klass->cuda_device_id,
510           &nvenc->cuda_ctx)) {
511     goto done;
512   }
513 #if HAVE_NVCODEC_GST_GL
514   gst_gl_handle_set_context (element, context,
515       (GstGLDisplay **) & nvenc->display,
516       (GstGLContext **) & nvenc->other_context);
517   if (nvenc->display)
518     gst_gl_display_filter_gl_api (GST_GL_DISPLAY (nvenc->display),
519         SUPPORTED_GL_APIS);
520 #endif
521
522 done:
523   GST_ELEMENT_CLASS (parent_class)->set_context (element, context);
524 }
525
526 static gboolean
527 gst_nv_base_enc_sink_query (GstVideoEncoder * enc, GstQuery * query)
528 {
529   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
530
531   switch (GST_QUERY_TYPE (query)) {
532     case GST_QUERY_CONTEXT:{
533       if (gst_cuda_handle_context_query (GST_ELEMENT (nvenc),
534               query, nvenc->cuda_ctx))
535         return TRUE;
536
537 #if HAVE_NVCODEC_GST_GL
538       {
539         gboolean ret;
540
541         ret = gst_gl_handle_context_query ((GstElement *) nvenc, query,
542             (GstGLDisplay *) nvenc->display, NULL,
543             (GstGLContext *) nvenc->other_context);
544         if (nvenc->display) {
545           gst_gl_display_filter_gl_api (GST_GL_DISPLAY (nvenc->display),
546               SUPPORTED_GL_APIS);
547         }
548
549         if (ret)
550           return ret;
551       }
552 #endif
553       break;
554     }
555     default:
556       break;
557   }
558
559   return GST_VIDEO_ENCODER_CLASS (parent_class)->sink_query (enc, query);
560 }
561
562 static gboolean
563 gst_nv_base_enc_sink_event (GstVideoEncoder * enc, GstEvent * event)
564 {
565   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
566   gboolean ret;
567
568   ret = GST_VIDEO_ENCODER_CLASS (parent_class)->sink_event (enc, event);
569
570   switch (GST_EVENT_TYPE (event)) {
571     case GST_EVENT_STREAM_START:
572     case GST_EVENT_FLUSH_STOP:
573       nvenc->last_flow = GST_FLOW_OK;
574       break;
575     default:
576       break;
577   }
578
579   return ret;
580 }
581
582 static gboolean
583 gst_nv_base_enc_start (GstVideoEncoder * enc)
584 {
585   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
586
587   nvenc->available_queue = g_async_queue_new ();
588   nvenc->pending_queue = g_async_queue_new ();
589   nvenc->bitstream_queue = g_async_queue_new ();
590   nvenc->items = g_array_new (FALSE, TRUE, sizeof (GstNvEncFrameState));
591
592   nvenc->last_flow = GST_FLOW_OK;
593   memset (&nvenc->init_params, 0, sizeof (NV_ENC_INITIALIZE_PARAMS));
594   memset (&nvenc->config, 0, sizeof (NV_ENC_CONFIG));
595
596 #if HAVE_NVCODEC_GST_GL
597   {
598     gst_gl_ensure_element_data (GST_ELEMENT (nvenc),
599         (GstGLDisplay **) & nvenc->display,
600         (GstGLContext **) & nvenc->other_context);
601     if (nvenc->display)
602       gst_gl_display_filter_gl_api (GST_GL_DISPLAY (nvenc->display),
603           SUPPORTED_GL_APIS);
604   }
605 #endif
606
607   /* DTS can be negative if bframe was enabled */
608   gst_video_encoder_set_min_pts (enc, GST_SECOND * 60 * 60 * 1000);
609
610   return TRUE;
611 }
612
613 static gboolean
614 gst_nv_base_enc_stop (GstVideoEncoder * enc)
615 {
616   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
617
618   gst_nv_base_enc_stop_bitstream_thread (nvenc, TRUE);
619
620   gst_nv_base_enc_free_buffers (nvenc);
621
622   if (nvenc->input_state) {
623     gst_video_codec_state_unref (nvenc->input_state);
624     nvenc->input_state = NULL;
625   }
626
627   if (nvenc->available_queue) {
628     g_async_queue_unref (nvenc->available_queue);
629     nvenc->available_queue = NULL;
630   }
631   if (nvenc->pending_queue) {
632     g_async_queue_unref (nvenc->pending_queue);
633     nvenc->pending_queue = NULL;
634   }
635   if (nvenc->bitstream_queue) {
636     g_async_queue_unref (nvenc->bitstream_queue);
637     nvenc->bitstream_queue = NULL;
638   }
639   if (nvenc->display) {
640     gst_object_unref (nvenc->display);
641     nvenc->display = NULL;
642   }
643   if (nvenc->other_context) {
644     gst_object_unref (nvenc->other_context);
645     nvenc->other_context = NULL;
646   }
647
648   if (nvenc->items) {
649     g_array_free (nvenc->items, TRUE);
650     nvenc->items = NULL;
651   }
652
653   return TRUE;
654 }
655
656 static void
657 check_formats (const gchar * str, guint * max_chroma, guint * max_bit_minus8)
658 {
659   if (!str)
660     return;
661
662   if (g_strrstr (str, "-444") || g_strrstr (str, "-4:4:4"))
663     *max_chroma = 2;
664   else if ((g_strrstr (str, "-4:2:2") || g_strrstr (str, "-422"))
665       && *max_chroma < 1)
666     *max_chroma = 1;
667
668   if (g_strrstr (str, "-12"))
669     *max_bit_minus8 = 4;
670   else if (g_strrstr (str, "-10") && *max_bit_minus8 < 2)
671     *max_bit_minus8 = 2;
672 }
673
674 static gboolean
675 gst_nv_base_enc_set_filtered_input_formats (GstNvBaseEnc * nvenc,
676     GstCaps * caps, const GValue * input_formats, guint max_chroma,
677     guint max_bit_minus8)
678 {
679   gint i;
680   GValue supported_format = G_VALUE_INIT;
681   gint num_format = 0;
682   const GValue *last_format = NULL;
683
684   g_value_init (&supported_format, GST_TYPE_LIST);
685
686   for (i = 0; i < gst_value_list_get_size (input_formats); i++) {
687     const GValue *val;
688     GstVideoFormat format;
689
690     val = gst_value_list_get_value (input_formats, i);
691     format = gst_video_format_from_string (g_value_get_string (val));
692
693     switch (format) {
694       case GST_VIDEO_FORMAT_NV12:
695       case GST_VIDEO_FORMAT_YV12:
696       case GST_VIDEO_FORMAT_I420:
697         /* 8bits 4:2:0 formats are always supported */
698       case GST_VIDEO_FORMAT_BGRA:
699       case GST_VIDEO_FORMAT_RGBA:
700         /* NOTE: RGB formats seems to also supported format, which are
701          * encoded to 4:2:0 formats */
702         gst_value_list_append_value (&supported_format, val);
703         last_format = val;
704         num_format++;
705         break;
706       case GST_VIDEO_FORMAT_Y444:
707       case GST_VIDEO_FORMAT_VUYA:
708         if (max_chroma >= 2) {
709           gst_value_list_append_value (&supported_format, val);
710           last_format = val;
711           num_format++;
712         }
713         break;
714       case GST_VIDEO_FORMAT_P010_10LE:
715       case GST_VIDEO_FORMAT_P010_10BE:
716       case GST_VIDEO_FORMAT_BGR10A2_LE:
717       case GST_VIDEO_FORMAT_RGB10A2_LE:
718       case GST_VIDEO_FORMAT_Y444_16LE:
719       case GST_VIDEO_FORMAT_Y444_16BE:
720         if (max_bit_minus8 >= 2) {
721           gst_value_list_append_value (&supported_format, val);
722           last_format = val;
723           num_format++;
724         }
725         break;
726       default:
727         break;
728     }
729   }
730
731   if (num_format == 0) {
732     g_value_unset (&supported_format);
733     GST_WARNING_OBJECT (nvenc, "Cannot find matching input format");
734     return FALSE;
735   }
736
737   if (num_format > 1)
738     gst_caps_set_value (caps, "format", &supported_format);
739   else
740     gst_caps_set_value (caps, "format", last_format);
741
742   g_value_unset (&supported_format);
743
744   return TRUE;
745 }
746
747 static GstCaps *
748 gst_nv_base_enc_getcaps (GstVideoEncoder * enc, GstCaps * filter)
749 {
750   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
751   GstNvBaseEncClass *klass = GST_NV_BASE_ENC_GET_CLASS (enc);
752   GstCaps *supported_incaps = NULL;
753   GstCaps *template_caps, *caps, *allowed;
754
755   template_caps = gst_pad_get_pad_template_caps (enc->sinkpad);
756   allowed = gst_pad_get_allowed_caps (enc->srcpad);
757
758   GST_LOG_OBJECT (enc, "template caps %" GST_PTR_FORMAT, template_caps);
759   GST_LOG_OBJECT (enc, "allowed caps %" GST_PTR_FORMAT, allowed);
760
761   if (!allowed) {
762     /* no peer */
763     supported_incaps = template_caps;
764     template_caps = NULL;
765     goto done;
766   } else if (gst_caps_is_empty (allowed)) {
767     /* couldn't be negotiated, just return empty caps */
768     gst_caps_unref (template_caps);
769     return allowed;
770   }
771
772   GST_OBJECT_LOCK (nvenc);
773
774   if (nvenc->input_formats != NULL) {
775     GValue *val;
776     gboolean has_profile = FALSE;
777     guint max_chroma_index = 0;
778     guint max_bit_minus8 = 0;
779     gint i, j;
780
781     for (i = 0; i < gst_caps_get_size (allowed); i++) {
782       const GstStructure *allowed_s = gst_caps_get_structure (allowed, i);
783       const GValue *val;
784
785       if ((val = gst_structure_get_value (allowed_s, "profile"))) {
786         if (G_VALUE_HOLDS_STRING (val)) {
787           check_formats (g_value_get_string (val), &max_chroma_index,
788               &max_bit_minus8);
789           has_profile = TRUE;
790         } else if (GST_VALUE_HOLDS_LIST (val)) {
791           for (j = 0; j < gst_value_list_get_size (val); j++) {
792             const GValue *vlist = gst_value_list_get_value (val, j);
793
794             if (G_VALUE_HOLDS_STRING (vlist)) {
795               check_formats (g_value_get_string (vlist), &max_chroma_index,
796                   &max_bit_minus8);
797               has_profile = TRUE;
798             }
799           }
800         }
801       }
802     }
803
804     GST_LOG_OBJECT (enc,
805         "downstream requested profile %d, max bitdepth %d, max chroma %d",
806         has_profile, max_bit_minus8 + 8, max_chroma_index);
807
808     supported_incaps = gst_caps_copy (template_caps);
809     if (!has_profile ||
810         !gst_nv_base_enc_set_filtered_input_formats (nvenc, supported_incaps,
811             nvenc->input_formats, max_chroma_index, max_bit_minus8)) {
812       gst_caps_set_value (supported_incaps, "format", nvenc->input_formats);
813     }
814
815     val = gst_nvenc_get_interlace_modes (nvenc->encoder, klass->codec_id);
816     gst_caps_set_value (supported_incaps, "interlace-mode", val);
817     g_value_unset (val);
818     g_free (val);
819
820     GST_LOG_OBJECT (enc, "codec input caps %" GST_PTR_FORMAT, supported_incaps);
821     GST_LOG_OBJECT (enc, "   template caps %" GST_PTR_FORMAT, template_caps);
822     caps = gst_caps_intersect (template_caps, supported_incaps);
823     gst_caps_unref (supported_incaps);
824     supported_incaps = caps;
825     GST_LOG_OBJECT (enc, "  supported caps %" GST_PTR_FORMAT, supported_incaps);
826   }
827
828   GST_OBJECT_UNLOCK (nvenc);
829
830 done:
831   caps = gst_video_encoder_proxy_getcaps (enc, supported_incaps, filter);
832
833   if (supported_incaps)
834     gst_caps_unref (supported_incaps);
835   gst_clear_caps (&allowed);
836   gst_clear_caps (&template_caps);
837
838   GST_DEBUG_OBJECT (nvenc, "  returning caps %" GST_PTR_FORMAT, caps);
839
840   return caps;
841 }
842
843 static gboolean
844 gst_nv_base_enc_close (GstVideoEncoder * enc)
845 {
846   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
847   gboolean ret = TRUE;
848
849   if (nvenc->encoder) {
850     if (NvEncDestroyEncoder (nvenc->encoder) != NV_ENC_SUCCESS)
851       ret = FALSE;
852
853     nvenc->encoder = NULL;
854   }
855
856   if (nvenc->cuda_ctx && nvenc->cuda_stream) {
857     if (gst_cuda_context_push (nvenc->cuda_ctx)) {
858       gst_cuda_result (CuStreamDestroy (nvenc->cuda_stream));
859       gst_cuda_context_pop (NULL);
860     }
861   }
862
863   gst_clear_object (&nvenc->cuda_ctx);
864   nvenc->cuda_stream = NULL;
865
866   GST_OBJECT_LOCK (nvenc);
867   if (nvenc->input_formats)
868     g_value_unset (nvenc->input_formats);
869   g_free (nvenc->input_formats);
870   nvenc->input_formats = NULL;
871   GST_OBJECT_UNLOCK (nvenc);
872
873   if (nvenc->input_state) {
874     gst_video_codec_state_unref (nvenc->input_state);
875     nvenc->input_state = NULL;
876   }
877
878   return ret;
879 }
880
881 static void
882 gst_nv_base_enc_init (GstNvBaseEnc * nvenc)
883 {
884   GstVideoEncoder *encoder = GST_VIDEO_ENCODER (nvenc);
885   GstNvEncQP qp_detail =
886       { DEFAULT_QP_DETAIL, DEFAULT_QP_DETAIL, DEFAULT_QP_DETAIL };
887
888   nvenc->preset_enum = DEFAULT_PRESET;
889   nvenc->selected_preset = _nv_preset_to_guid (nvenc->preset_enum);
890   nvenc->rate_control_mode = DEFAULT_RC_MODE;
891   nvenc->qp_min = DEFAULT_QP_MIN;
892   nvenc->qp_max = DEFAULT_QP_MAX;
893   nvenc->qp_const = DEFAULT_QP_CONST;
894   nvenc->bitrate = DEFAULT_BITRATE;
895   nvenc->gop_size = DEFAULT_GOP_SIZE;
896   nvenc->max_bitrate = DEFAULT_MAX_BITRATE;
897   nvenc->spatial_aq = DEFAULT_SPATIAL_AQ;
898   nvenc->aq_strength = DEFAULT_AQ_STRENGTH;
899   nvenc->non_refp = DEFAULT_NON_REF_P;
900   nvenc->zerolatency = DEFAULT_ZEROLATENCY;
901   nvenc->strict_gop = DEFAULT_STRICT_GOP;
902   nvenc->const_quality = DEFAULT_CONST_QUALITY;
903   nvenc->i_adapt = DEFAULT_I_ADAPT;
904   nvenc->qp_min_detail = qp_detail;
905   nvenc->qp_max_detail = qp_detail;
906   nvenc->qp_const_detail = qp_detail;
907
908   GST_VIDEO_ENCODER_STREAM_LOCK (encoder);
909   GST_VIDEO_ENCODER_STREAM_UNLOCK (encoder);
910
911   GST_PAD_SET_ACCEPT_INTERSECT (GST_VIDEO_ENCODER_SINK_PAD (encoder));
912 }
913
914 static void
915 gst_nv_base_enc_finalize (GObject * obj)
916 {
917   G_OBJECT_CLASS (gst_nv_base_enc_parent_class)->finalize (obj);
918 }
919
920 static GstVideoCodecFrame *
921 _find_frame_with_output_buffer (GstNvBaseEnc * nvenc, NV_ENC_OUTPUT_PTR out_buf)
922 {
923   GList *l, *walk = gst_video_encoder_get_frames (GST_VIDEO_ENCODER (nvenc));
924   GstVideoCodecFrame *ret = NULL;
925
926   for (l = walk; l; l = l->next) {
927     GstVideoCodecFrame *frame = (GstVideoCodecFrame *) l->data;
928     GstNvEncFrameState *state = gst_video_codec_frame_get_user_data (frame);
929
930     if (!state || !state->out_buf)
931       continue;
932
933     if (state->out_buf == out_buf) {
934       ret = frame;
935       break;
936     }
937   }
938
939   if (ret)
940     gst_video_codec_frame_ref (ret);
941
942   g_list_free_full (walk, (GDestroyNotify) gst_video_codec_frame_unref);
943
944   return ret;
945 }
946
947 static const gchar *
948 picture_type_to_string (NV_ENC_PIC_TYPE type)
949 {
950   switch (type) {
951     case NV_ENC_PIC_TYPE_P:
952       return "P";
953     case NV_ENC_PIC_TYPE_B:
954       return "B";
955     case NV_ENC_PIC_TYPE_I:
956       return "I";
957     case NV_ENC_PIC_TYPE_IDR:
958       return "IDR";
959     case NV_ENC_PIC_TYPE_BI:
960       return "BI";
961     case NV_ENC_PIC_TYPE_SKIPPED:
962       return "SKIPPED";
963     case NV_ENC_PIC_TYPE_INTRA_REFRESH:
964       return "INTRA-REFRESH";
965     case NV_ENC_PIC_TYPE_UNKNOWN:
966     default:
967       break;
968   }
969
970   return "UNKNOWN";
971 }
972
973 static gpointer
974 gst_nv_base_enc_bitstream_thread (gpointer user_data)
975 {
976   GstVideoEncoder *enc = user_data;
977   GstNvBaseEnc *nvenc = user_data;
978   GstFlowReturn flow = GST_FLOW_OK;
979
980   /* overview of operation:
981    * 1. retreive the next buffer submitted to the bitstream pool
982    * 2. wait for that buffer to be ready from nvenc (LockBitsream)
983    * 3. retreive the GstVideoCodecFrame associated with that buffer
984    * 4. for each buffer in the frame
985    * 4.1 (step 2): wait for that buffer to be ready from nvenc (LockBitsream)
986    * 4.2 create an output GstBuffer from the nvenc buffers
987    * 4.3 unlock the nvenc bitstream buffers UnlockBitsream
988    * 5. finish_frame()
989    * 6. cleanup
990    */
991   do {
992     GstBuffer *buffer = NULL;
993     GstNvEncFrameState *state_in_queue = NULL;
994     GstNvEncFrameState *state = NULL;
995     GstVideoCodecFrame *frame = NULL;
996     NVENCSTATUS nv_ret;
997     NV_ENC_LOCK_BITSTREAM lock_bs = { 0, };
998     NV_ENC_OUTPUT_PTR out_buf;
999     GstNvEncInputResource *resource;
1000
1001     GST_LOG_OBJECT (enc, "wait for bitstream buffer..");
1002
1003     state_in_queue = g_async_queue_pop (nvenc->bitstream_queue);
1004     if ((gpointer) state_in_queue == SHUTDOWN_COOKIE)
1005       goto exit_thread;
1006
1007     out_buf = state_in_queue->out_buf;
1008     resource = state_in_queue->in_buf;
1009
1010     GST_LOG_OBJECT (nvenc, "waiting for output buffer %p to be ready", out_buf);
1011
1012     lock_bs.version = gst_nvenc_get_lock_bitstream_version ();
1013     lock_bs.outputBitstream = out_buf;
1014     lock_bs.doNotWait = 0;
1015
1016     /* FIXME: this would need to be updated for other slice modes */
1017     lock_bs.sliceOffsets = NULL;
1018
1019     if (!gst_cuda_context_push (nvenc->cuda_ctx)) {
1020       GST_ELEMENT_ERROR (nvenc, LIBRARY, ENCODE, (NULL),
1021           ("Failed to push current context"));
1022       goto error_shutdown;
1023     }
1024
1025     nv_ret = NvEncLockBitstream (nvenc->encoder, &lock_bs);
1026     if (nv_ret != NV_ENC_SUCCESS) {
1027       gst_cuda_context_pop (NULL);
1028
1029       GST_ELEMENT_ERROR (nvenc, STREAM, ENCODE, (NULL),
1030           ("Failed to lock bitstream buffer %p, ret %d",
1031               lock_bs.outputBitstream, nv_ret));
1032       goto error_shutdown;
1033     }
1034
1035     frame = _find_frame_with_output_buffer (nvenc, out_buf);
1036     state = gst_video_codec_frame_get_user_data (frame);
1037     g_assert (state->out_buf == out_buf);
1038
1039     /* copy into output buffer */
1040     buffer = gst_buffer_new_allocate (NULL, lock_bs.bitstreamSizeInBytes, NULL);
1041     gst_buffer_fill (buffer, 0, lock_bs.bitstreamBufferPtr,
1042         lock_bs.bitstreamSizeInBytes);
1043
1044     if (lock_bs.pictureType == NV_ENC_PIC_TYPE_IDR) {
1045       GST_DEBUG_OBJECT (nvenc, "This is a keyframe");
1046       GST_VIDEO_CODEC_FRAME_SET_SYNC_POINT (frame);
1047     }
1048
1049     nv_ret = NvEncUnlockBitstream (nvenc->encoder, state->out_buf);
1050
1051     if (nv_ret != NV_ENC_SUCCESS) {
1052       gst_cuda_context_pop (NULL);
1053
1054       GST_ELEMENT_ERROR (nvenc, STREAM, ENCODE, (NULL),
1055           ("Failed to unlock bitstream buffer %p, ret %d",
1056               lock_bs.outputBitstream, nv_ret));
1057       gst_buffer_unref (buffer);
1058       gst_video_encoder_finish_frame (enc, frame);
1059
1060       goto error_shutdown;
1061     }
1062
1063     frame->dts = frame->pts;
1064     frame->pts = lock_bs.outputTimeStamp;
1065     frame->duration = lock_bs.outputDuration;
1066
1067     GST_LOG_OBJECT (nvenc, "frame index %" G_GUINT32_FORMAT
1068         ", frame type %s, dts %" GST_TIME_FORMAT
1069         ", pts %" GST_TIME_FORMAT,
1070         lock_bs.frameIdx, picture_type_to_string (lock_bs.pictureType),
1071         GST_TIME_ARGS (frame->dts), GST_TIME_ARGS (frame->pts));
1072
1073     frame->output_buffer = buffer;
1074
1075     nv_ret =
1076         NvEncUnmapInputResource (nvenc->encoder,
1077         resource->nv_mapped_resource.mappedResource);
1078     resource->mapped = FALSE;
1079
1080     if (nv_ret != NV_ENC_SUCCESS) {
1081       GST_ERROR_OBJECT (nvenc, "Failed to unmap input resource %p, ret %d",
1082           resource, nv_ret);
1083     }
1084
1085     gst_cuda_context_pop (NULL);
1086
1087     memset (&resource->nv_mapped_resource, 0,
1088         sizeof (resource->nv_mapped_resource));
1089
1090     g_async_queue_push (nvenc->available_queue, state_in_queue);
1091
1092     /* Ugly but no other way to get DTS offset since nvenc dose not adjust
1093      * dts/pts even if bframe was enabled. So the output PTS can be smaller
1094      * than DTS. The maximum difference between DTS and PTS can be calculated
1095      * using the PTS difference between the first frame and the second frame.
1096      */
1097     if (nvenc->bframes > 0) {
1098       if (nvenc->dts_offset == 0) {
1099         if (!nvenc->first_frame) {
1100           /* store the first frame to get dts offset */
1101           nvenc->first_frame = frame;
1102           continue;
1103         } else {
1104           if (nvenc->first_frame->pts >= frame->pts) {
1105             GstClockTime duration = 0;
1106
1107             GST_WARNING_OBJECT (enc, "Could not calculate DTS offset");
1108
1109             if (nvenc->input_info.fps_n > 0 && nvenc->input_info.fps_d > 0) {
1110               duration =
1111                   gst_util_uint64_scale (GST_SECOND, nvenc->input_info.fps_d,
1112                   nvenc->input_info.fps_n);
1113             } else if (nvenc->first_frame->duration > 0 &&
1114                 GST_CLOCK_TIME_IS_VALID (nvenc->first_frame->duration)) {
1115               duration = nvenc->first_frame->duration;
1116             } else {
1117               GST_WARNING_OBJECT (enc,
1118                   "No way to get frame duration, assuming 30fps");
1119               duration = gst_util_uint64_scale (GST_SECOND, 1, 30);
1120             }
1121
1122             nvenc->dts_offset = duration * nvenc->bframes;
1123           } else {
1124             nvenc->dts_offset = frame->pts - nvenc->first_frame->pts;
1125           }
1126
1127           /* + 1 to dts_offset to adjust fraction */
1128           nvenc->dts_offset++;
1129
1130           GST_DEBUG_OBJECT (enc,
1131               "Calculated DTS offset %" GST_TIME_FORMAT,
1132               GST_TIME_ARGS (nvenc->dts_offset));
1133         }
1134
1135         nvenc->first_frame->dts -= nvenc->dts_offset;
1136         gst_video_encoder_finish_frame (enc, nvenc->first_frame);
1137         nvenc->first_frame = NULL;
1138       }
1139
1140       frame->dts -= nvenc->dts_offset;
1141     }
1142
1143     flow = gst_video_encoder_finish_frame (enc, frame);
1144
1145     if (flow != GST_FLOW_OK) {
1146       GST_INFO_OBJECT (enc, "got flow %s", gst_flow_get_name (flow));
1147       g_atomic_int_set (&nvenc->last_flow, flow);
1148       g_async_queue_push (nvenc->available_queue, SHUTDOWN_COOKIE);
1149       goto exit_thread;
1150     }
1151   }
1152   while (TRUE);
1153
1154 error_shutdown:
1155   {
1156     if (nvenc->first_frame) {
1157       gst_clear_buffer (&nvenc->first_frame->output_buffer);
1158       gst_video_encoder_finish_frame (enc, nvenc->first_frame);
1159       nvenc->first_frame = NULL;
1160     }
1161     g_atomic_int_set (&nvenc->last_flow, GST_FLOW_ERROR);
1162     g_async_queue_push (nvenc->available_queue, SHUTDOWN_COOKIE);
1163
1164     goto exit_thread;
1165   }
1166
1167 exit_thread:
1168   {
1169     if (nvenc->first_frame) {
1170       gst_video_encoder_finish_frame (enc, nvenc->first_frame);
1171       nvenc->first_frame = NULL;
1172     }
1173
1174     GST_INFO_OBJECT (nvenc, "exiting thread");
1175
1176     return NULL;
1177   }
1178 }
1179
1180 static gboolean
1181 gst_nv_base_enc_start_bitstream_thread (GstNvBaseEnc * nvenc)
1182 {
1183   gchar *name = g_strdup_printf ("%s-read-bits", GST_OBJECT_NAME (nvenc));
1184
1185   g_assert (nvenc->bitstream_thread == NULL);
1186
1187   g_assert (g_async_queue_length (nvenc->bitstream_queue) == 0);
1188
1189   nvenc->bitstream_thread =
1190       g_thread_try_new (name, gst_nv_base_enc_bitstream_thread, nvenc, NULL);
1191
1192   g_free (name);
1193
1194   if (nvenc->bitstream_thread == NULL)
1195     return FALSE;
1196
1197   GST_INFO_OBJECT (nvenc, "started thread to read bitstream");
1198   return TRUE;
1199 }
1200
1201 static gboolean
1202 gst_nv_base_enc_stop_bitstream_thread (GstNvBaseEnc * nvenc, gboolean force)
1203 {
1204   GstNvEncFrameState *state;
1205
1206   if (nvenc->bitstream_thread == NULL)
1207     return TRUE;
1208
1209   /* Always send EOS packet to flush GPU. Otherwise, randomly crash happens
1210    * during NvEncDestroyEncoder especially when rc-lookahead or bframe was
1211    * enabled */
1212   gst_nv_base_enc_drain_encoder (nvenc);
1213
1214   if (force) {
1215     g_async_queue_lock (nvenc->available_queue);
1216     g_async_queue_lock (nvenc->pending_queue);
1217     g_async_queue_lock (nvenc->bitstream_queue);
1218     while ((state = g_async_queue_try_pop_unlocked (nvenc->bitstream_queue))) {
1219       GST_INFO_OBJECT (nvenc, "stole bitstream buffer %p from queue", state);
1220       g_async_queue_push_unlocked (nvenc->available_queue, state);
1221     }
1222     g_async_queue_push_unlocked (nvenc->bitstream_queue, SHUTDOWN_COOKIE);
1223     g_async_queue_unlock (nvenc->available_queue);
1224     g_async_queue_unlock (nvenc->pending_queue);
1225     g_async_queue_unlock (nvenc->bitstream_queue);
1226   } else {
1227     /* wait for encoder to drain the remaining buffers */
1228     g_async_queue_push (nvenc->bitstream_queue, SHUTDOWN_COOKIE);
1229   }
1230
1231   if (!force) {
1232     /* temporary unlock during finish, so other thread can find and push frame */
1233     GST_VIDEO_ENCODER_STREAM_UNLOCK (nvenc);
1234   }
1235
1236   g_thread_join (nvenc->bitstream_thread);
1237
1238   if (!force)
1239     GST_VIDEO_ENCODER_STREAM_LOCK (nvenc);
1240
1241   nvenc->bitstream_thread = NULL;
1242   return TRUE;
1243 }
1244
1245 static void
1246 gst_nv_base_enc_reset_queues (GstNvBaseEnc * nvenc, gboolean refill)
1247 {
1248   gpointer ptr;
1249   gint i;
1250
1251   GST_INFO_OBJECT (nvenc, "clearing queues");
1252
1253   while ((ptr = g_async_queue_try_pop (nvenc->available_queue))) {
1254     /* do nothing */
1255   }
1256   while ((ptr = g_async_queue_try_pop (nvenc->pending_queue))) {
1257     /* do nothing */
1258   }
1259   while ((ptr = g_async_queue_try_pop (nvenc->bitstream_queue))) {
1260     /* do nothing */
1261   }
1262
1263   if (refill) {
1264     GST_INFO_OBJECT (nvenc, "refilling buffer pools");
1265     for (i = 0; i < nvenc->n_bufs; ++i) {
1266       g_async_queue_push (nvenc->available_queue,
1267           &g_array_index (nvenc->items, GstNvEncFrameState, i));
1268     }
1269   }
1270 }
1271
1272 static void
1273 gst_nv_base_enc_free_buffers (GstNvBaseEnc * nvenc)
1274 {
1275   NVENCSTATUS nv_ret;
1276   CUresult cuda_ret;
1277   guint i;
1278
1279   if (nvenc->encoder == NULL)
1280     return;
1281
1282   gst_nv_base_enc_reset_queues (nvenc, FALSE);
1283
1284   gst_cuda_context_push (nvenc->cuda_ctx);
1285   for (i = 0; i < nvenc->n_bufs; ++i) {
1286     NV_ENC_OUTPUT_PTR out_buf =
1287         g_array_index (nvenc->items, GstNvEncFrameState, i).out_buf;
1288     GstNvEncInputResource *in_buf =
1289         g_array_index (nvenc->items, GstNvEncFrameState, i).in_buf;
1290
1291     if (in_buf->mapped) {
1292       GST_LOG_OBJECT (nvenc, "Unmap resource %p", in_buf);
1293
1294       nv_ret =
1295           NvEncUnmapInputResource (nvenc->encoder,
1296           in_buf->nv_mapped_resource.mappedResource);
1297
1298       if (nv_ret != NV_ENC_SUCCESS) {
1299         GST_ERROR_OBJECT (nvenc, "Failed to unmap input resource %p, ret %d",
1300             in_buf, nv_ret);
1301       }
1302     }
1303
1304     nv_ret =
1305         NvEncUnregisterResource (nvenc->encoder,
1306         in_buf->nv_resource.registeredResource);
1307     if (nv_ret != NV_ENC_SUCCESS)
1308       GST_ERROR_OBJECT (nvenc, "Failed to unregister resource %p, ret %d",
1309           in_buf, nv_ret);
1310
1311     cuda_ret = CuMemFree (in_buf->cuda_pointer);
1312     if (!gst_cuda_result (cuda_ret)) {
1313       GST_ERROR_OBJECT (nvenc, "Failed to free CUDA device memory, ret %d",
1314           cuda_ret);
1315     }
1316
1317     g_free (in_buf);
1318
1319     GST_DEBUG_OBJECT (nvenc, "Destroying output bitstream buffer %p", out_buf);
1320     nv_ret = NvEncDestroyBitstreamBuffer (nvenc->encoder, out_buf);
1321     if (nv_ret != NV_ENC_SUCCESS) {
1322       GST_ERROR_OBJECT (nvenc, "Failed to destroy output buffer %p, ret %d",
1323           out_buf, nv_ret);
1324     }
1325   }
1326   gst_cuda_context_pop (NULL);
1327   g_array_set_size (nvenc->items, 0);
1328 }
1329
1330 static inline guint
1331 _get_plane_width (GstVideoInfo * info, guint plane)
1332 {
1333   return GST_VIDEO_INFO_COMP_WIDTH (info, plane)
1334       * GST_VIDEO_INFO_COMP_PSTRIDE (info, plane);
1335 }
1336
1337 static inline guint
1338 _get_plane_height (GstVideoInfo * info, guint plane)
1339 {
1340   if (GST_VIDEO_INFO_IS_YUV (info))
1341     /* For now component width and plane width are the same and the
1342      * plane-component mapping matches
1343      */
1344     return GST_VIDEO_INFO_COMP_HEIGHT (info, plane);
1345   else                          /* RGB, GRAY */
1346     return GST_VIDEO_INFO_HEIGHT (info);
1347 }
1348
1349 static inline gsize
1350 _get_frame_data_height (GstVideoInfo * info)
1351 {
1352   gsize ret = 0;
1353   gint i;
1354
1355   for (i = 0; i < GST_VIDEO_INFO_N_PLANES (info); i++) {
1356     ret += _get_plane_height (info, i);
1357   }
1358
1359   return ret;
1360 }
1361
1362 static gboolean
1363 qp_has_values (const GstNvEncQP * qp)
1364 {
1365   return qp->qp_i >= 0 && qp->qp_p >= 0 && qp->qp_b >= 0;
1366 }
1367
1368 static void
1369 gst_nv_base_enc_setup_rate_control (GstNvBaseEnc * nvenc,
1370     NV_ENC_RC_PARAMS * rc_params)
1371 {
1372   GstNvRCMode rc_mode = nvenc->rate_control_mode;
1373   NV_ENC_PARAMS_RC_MODE nv_rcmode;
1374
1375   if (nvenc->bitrate)
1376     rc_params->averageBitRate = nvenc->bitrate * 1024;
1377
1378   if (nvenc->max_bitrate)
1379     rc_params->maxBitRate = nvenc->max_bitrate * 1024;
1380
1381   if (nvenc->vbv_buffersize)
1382     rc_params->vbvBufferSize = nvenc->vbv_buffersize * 1024;
1383
1384   /* Guess the best matching mode */
1385   if (rc_mode == GST_NV_RC_MODE_DEFAULT) {
1386     if (nvenc->qp_const >= 0) {
1387       /* constQP is used only for RC_CONSTQP mode */
1388       rc_mode = GST_NV_RC_MODE_CONSTQP;
1389     }
1390   }
1391
1392   if (nvenc->qp_min >= 0) {
1393     rc_params->enableMinQP = 1;
1394     rc_params->minQP.qpInterB = nvenc->qp_min;
1395     rc_params->minQP.qpInterP = nvenc->qp_min;
1396     rc_params->minQP.qpIntra = nvenc->qp_min;
1397   } else if (qp_has_values (&nvenc->qp_min_detail)) {
1398     rc_params->enableMinQP = 1;
1399     rc_params->minQP.qpInterB = nvenc->qp_min_detail.qp_b;
1400     rc_params->minQP.qpInterP = nvenc->qp_min_detail.qp_p;
1401     rc_params->minQP.qpIntra = nvenc->qp_min_detail.qp_i;
1402   }
1403
1404   if (nvenc->qp_max >= 0) {
1405     rc_params->enableMaxQP = 1;
1406     rc_params->maxQP.qpInterB = nvenc->qp_max;
1407     rc_params->maxQP.qpInterP = nvenc->qp_max;
1408     rc_params->maxQP.qpIntra = nvenc->qp_max;
1409   } else if (qp_has_values (&nvenc->qp_max_detail)) {
1410     rc_params->enableMaxQP = 1;
1411     rc_params->maxQP.qpInterB = nvenc->qp_max_detail.qp_b;
1412     rc_params->maxQP.qpInterP = nvenc->qp_max_detail.qp_p;
1413     rc_params->maxQP.qpIntra = nvenc->qp_max_detail.qp_i;
1414   }
1415
1416   if (nvenc->qp_const >= 0) {
1417     rc_params->constQP.qpInterB = nvenc->qp_const;
1418     rc_params->constQP.qpInterP = nvenc->qp_const;
1419     rc_params->constQP.qpIntra = nvenc->qp_const;
1420   } else if (qp_has_values (&nvenc->qp_const_detail)) {
1421     rc_params->constQP.qpInterB = nvenc->qp_const_detail.qp_b;
1422     rc_params->constQP.qpInterP = nvenc->qp_const_detail.qp_p;
1423     rc_params->constQP.qpIntra = nvenc->qp_const_detail.qp_i;
1424   }
1425
1426   nv_rcmode = _rc_mode_to_nv (rc_mode);
1427   if (nv_rcmode == NV_ENC_PARAMS_RC_VBR_MINQP && nvenc->qp_min < 0) {
1428     GST_WARNING_OBJECT (nvenc, "vbr-minqp was requested without qp-min");
1429     nv_rcmode = NV_ENC_PARAMS_RC_VBR;
1430   }
1431
1432   rc_params->rateControlMode = nv_rcmode;
1433
1434   if (nvenc->spatial_aq) {
1435     rc_params->enableAQ = 1;
1436     rc_params->aqStrength = nvenc->aq_strength;
1437   }
1438
1439   rc_params->enableTemporalAQ = nvenc->temporal_aq;
1440
1441   if (nvenc->rc_lookahead) {
1442     rc_params->enableLookahead = 1;
1443     rc_params->lookaheadDepth = nvenc->rc_lookahead;
1444     rc_params->disableIadapt = !nvenc->i_adapt;
1445     rc_params->disableBadapt = !nvenc->b_adapt;
1446   }
1447
1448   rc_params->strictGOPTarget = nvenc->strict_gop;
1449   rc_params->enableNonRefP = nvenc->non_refp;
1450   rc_params->zeroReorderDelay = nvenc->zerolatency;
1451
1452   if (nvenc->const_quality) {
1453     guint scaled = (gint) (nvenc->const_quality * 256.0);
1454
1455     rc_params->targetQuality = (guint8) (scaled >> 8);
1456     rc_params->targetQualityLSB = (guint8) (scaled & 0xff);
1457   }
1458 }
1459
1460 static gint
1461 gst_nv_base_enc_calculate_num_prealloc_buffers (GstNvBaseEnc * enc,
1462     NV_ENC_CONFIG * config)
1463 {
1464   gint num_buffers;
1465
1466   /* At least 4 surfaces are required as documented by Nvidia Encoder guide */
1467   num_buffers = 4;
1468
1469   /* + lookahead depth */
1470   num_buffers += config->rcParams.lookaheadDepth;
1471
1472   /* + GOP size */
1473   num_buffers += config->frameIntervalP;
1474
1475   /* hardcoded upper bound "48"
1476    * The worst case
1477    *   default num buffers: 4
1478    *   maximum allowed lookahead: 32
1479    *   max bfraems: 4 -> frameIntervalP: 5
1480    * "4 + 32 + 5" < "48" so it seems to sufficiently safe upper bound */
1481   num_buffers = MIN (num_buffers, 48);
1482
1483   GST_DEBUG_OBJECT (enc, "Calculated num buffers: %d "
1484       "(lookahead %d, frameIntervalP %d)",
1485       num_buffers, config->rcParams.lookaheadDepth, config->frameIntervalP);
1486
1487   return num_buffers;
1488 }
1489
1490 /* GstVideoEncoder::set_format or by nvenc self if new properties were set.
1491  *
1492  * NvEncReconfigureEncoder with following conditions are not allowed
1493  * 1) GOP structure change
1494  * 2) sync-Async mode change (Async mode is Windows only and we didn't support it)
1495  * 3) MaxWidth, MaxHeight
1496  * 4) PTDmode (Picture Type Decision mode)
1497  *
1498  * So we will force to re-init the encode session if
1499  * 1) New resolution is larger than previous config
1500  * 2) GOP size changed
1501  * 3) Input pixel format change
1502  *    pre-allocated CUDA memory could not ensure stride, width and height
1503  *
1504  * TODO: bframe also considered as force re-init case
1505  */
1506 static gboolean
1507 gst_nv_base_enc_set_format (GstVideoEncoder * enc, GstVideoCodecState * state)
1508 {
1509   GstNvBaseEncClass *nvenc_class = GST_NV_BASE_ENC_GET_CLASS (enc);
1510   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
1511   GstVideoInfo *info = &state->info;
1512   GstVideoCodecState *old_state = nvenc->input_state;
1513   NV_ENC_RECONFIGURE_PARAMS reconfigure_params = { 0, };
1514   NV_ENC_INITIALIZE_PARAMS *params = &nvenc->init_params;
1515   NV_ENC_PRESET_CONFIG preset_config = { 0, };
1516   NVENCSTATUS nv_ret;
1517   gint dar_n, dar_d;
1518   gboolean reconfigure = FALSE;
1519
1520   g_atomic_int_set (&nvenc->reconfig, FALSE);
1521
1522   if (old_state) {
1523     gboolean larger_resolution;
1524     gboolean format_changed;
1525     gboolean gop_size_changed;
1526
1527     larger_resolution =
1528         (GST_VIDEO_INFO_WIDTH (info) > nvenc->init_params.maxEncodeWidth ||
1529         GST_VIDEO_INFO_HEIGHT (info) > nvenc->init_params.maxEncodeHeight);
1530     format_changed =
1531         GST_VIDEO_INFO_FORMAT (info) !=
1532         GST_VIDEO_INFO_FORMAT (&old_state->info);
1533
1534     if (nvenc->config.gopLength == NVENC_INFINITE_GOPLENGTH
1535         && nvenc->gop_size == -1) {
1536       gop_size_changed = FALSE;
1537     } else if (nvenc->config.gopLength != nvenc->gop_size) {
1538       gop_size_changed = TRUE;
1539     } else {
1540       gop_size_changed = FALSE;
1541     }
1542
1543     if (larger_resolution || format_changed || gop_size_changed) {
1544       GST_DEBUG_OBJECT (nvenc,
1545           "resolution %dx%d -> %dx%d, format %s -> %s, re-init",
1546           nvenc->init_params.maxEncodeWidth, nvenc->init_params.maxEncodeHeight,
1547           GST_VIDEO_INFO_WIDTH (info), GST_VIDEO_INFO_HEIGHT (info),
1548           gst_video_format_to_string (GST_VIDEO_INFO_FORMAT (&old_state->info)),
1549           gst_video_format_to_string (GST_VIDEO_INFO_FORMAT (info)));
1550
1551       gst_nv_base_enc_drain_encoder (nvenc);
1552       gst_nv_base_enc_stop_bitstream_thread (nvenc, FALSE);
1553       gst_nv_base_enc_free_buffers (nvenc);
1554       NvEncDestroyEncoder (nvenc->encoder);
1555       nvenc->encoder = NULL;
1556
1557       if (!gst_nv_base_enc_open_encode_session (nvenc)) {
1558         GST_ERROR_OBJECT (nvenc, "Failed to open encode session");
1559         return FALSE;
1560       }
1561     } else {
1562       reconfigure_params.version = gst_nvenc_get_reconfigure_params_version ();
1563       /* reset rate control state and start from IDR */
1564       reconfigure_params.resetEncoder = TRUE;
1565       reconfigure_params.forceIDR = TRUE;
1566       reconfigure = TRUE;
1567     }
1568   }
1569
1570   params->version = gst_nvenc_get_initialize_params_version ();
1571   params->encodeGUID = nvenc_class->codec_id;
1572   params->encodeWidth = GST_VIDEO_INFO_WIDTH (info);
1573   params->encodeHeight = GST_VIDEO_INFO_HEIGHT (info);
1574
1575   {
1576     guint32 n_presets;
1577     GUID *presets;
1578     guint32 i;
1579
1580     nv_ret =
1581         NvEncGetEncodePresetCount (nvenc->encoder,
1582         params->encodeGUID, &n_presets);
1583     if (nv_ret != NV_ENC_SUCCESS) {
1584       GST_ELEMENT_ERROR (nvenc, LIBRARY, SETTINGS, (NULL),
1585           ("Failed to get encoder presets"));
1586       return FALSE;
1587     }
1588
1589     presets = g_new0 (GUID, n_presets);
1590     nv_ret =
1591         NvEncGetEncodePresetGUIDs (nvenc->encoder,
1592         params->encodeGUID, presets, n_presets, &n_presets);
1593     if (nv_ret != NV_ENC_SUCCESS) {
1594       GST_ELEMENT_ERROR (nvenc, LIBRARY, SETTINGS, (NULL),
1595           ("Failed to get encoder presets"));
1596       g_free (presets);
1597       return FALSE;
1598     }
1599
1600     for (i = 0; i < n_presets; i++) {
1601       if (gst_nvenc_cmp_guid (presets[i], nvenc->selected_preset))
1602         break;
1603     }
1604     g_free (presets);
1605     if (i >= n_presets) {
1606       GST_ELEMENT_ERROR (nvenc, LIBRARY, SETTINGS, (NULL),
1607           ("Selected preset not supported"));
1608       return FALSE;
1609     }
1610
1611     params->presetGUID = nvenc->selected_preset;
1612   }
1613
1614   params->enablePTD = 1;
1615   if (!reconfigure) {
1616     /* this sets the required buffer size and the maximum allowed size on
1617      * subsequent reconfigures */
1618     params->maxEncodeWidth = GST_VIDEO_INFO_WIDTH (info);
1619     params->maxEncodeHeight = GST_VIDEO_INFO_HEIGHT (info);
1620   }
1621
1622   preset_config.version = gst_nvenc_get_preset_config_version ();
1623   preset_config.presetCfg.version = gst_nvenc_get_config_version ();
1624
1625   nv_ret =
1626       NvEncGetEncodePresetConfig (nvenc->encoder,
1627       params->encodeGUID, params->presetGUID, &preset_config);
1628   if (nv_ret != NV_ENC_SUCCESS) {
1629     GST_ELEMENT_ERROR (nvenc, LIBRARY, SETTINGS, (NULL),
1630         ("Failed to get encode preset configuration: %d", nv_ret));
1631     return FALSE;
1632   }
1633
1634   params->encodeConfig = &preset_config.presetCfg;
1635
1636   if (GST_VIDEO_INFO_IS_INTERLACED (info)) {
1637     if (GST_VIDEO_INFO_INTERLACE_MODE (info) ==
1638         GST_VIDEO_INTERLACE_MODE_INTERLEAVED
1639         || GST_VIDEO_INFO_INTERLACE_MODE (info) ==
1640         GST_VIDEO_INTERLACE_MODE_MIXED) {
1641       preset_config.presetCfg.frameFieldMode =
1642           NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
1643     }
1644   }
1645
1646   if (info->fps_d > 0 && info->fps_n > 0) {
1647     params->frameRateNum = info->fps_n;
1648     params->frameRateDen = info->fps_d;
1649   } else {
1650     params->frameRateNum = 0;
1651     params->frameRateDen = 1;
1652   }
1653
1654   if (gst_util_fraction_multiply (GST_VIDEO_INFO_WIDTH (info),
1655           GST_VIDEO_INFO_HEIGHT (info), GST_VIDEO_INFO_PAR_N (info),
1656           GST_VIDEO_INFO_PAR_D (info), &dar_n, &dar_d) && dar_n > 0
1657       && dar_d > 0) {
1658     params->darWidth = dar_n;
1659     params->darHeight = dar_d;
1660   }
1661
1662   gst_nv_base_enc_setup_rate_control (nvenc, &params->encodeConfig->rcParams);
1663
1664   params->enableWeightedPrediction = nvenc->weighted_pred;
1665
1666   if (nvenc->gop_size < 0) {
1667     params->encodeConfig->gopLength = NVENC_INFINITE_GOPLENGTH;
1668     params->encodeConfig->frameIntervalP = 1;
1669   } else if (nvenc->gop_size > 0) {
1670     params->encodeConfig->gopLength = nvenc->gop_size;
1671     /* frameIntervalP
1672      * 0: All Intra frames
1673      * 1: I/P only
1674      * n ( > 1): n - 1 bframes
1675      */
1676     params->encodeConfig->frameIntervalP = nvenc->bframes + 1;
1677   } else {
1678     /* gop size == 0 means all intra frames */
1679     params->encodeConfig->gopLength = 1;
1680     params->encodeConfig->frameIntervalP = 0;
1681   }
1682
1683   g_assert (nvenc_class->set_encoder_config);
1684   if (!nvenc_class->set_encoder_config (nvenc, state, params->encodeConfig)) {
1685     GST_ERROR_OBJECT (enc, "Subclass failed to set encoder configuration");
1686     return FALSE;
1687   }
1688
1689   /* store the last config to reconfig/re-init decision in the next time */
1690   nvenc->config = *params->encodeConfig;
1691
1692   G_LOCK (initialization_lock);
1693   if (reconfigure) {
1694     reconfigure_params.reInitEncodeParams = nvenc->init_params;
1695     nv_ret = NvEncReconfigureEncoder (nvenc->encoder, &reconfigure_params);
1696   } else {
1697     nv_ret = NvEncInitializeEncoder (nvenc->encoder, params);
1698   }
1699   G_UNLOCK (initialization_lock);
1700
1701   if (nv_ret != NV_ENC_SUCCESS) {
1702     GST_ELEMENT_ERROR (nvenc, LIBRARY, SETTINGS, (NULL),
1703         ("Failed to %sinit encoder: %d", reconfigure ? "re" : "", nv_ret));
1704     return FALSE;
1705   }
1706
1707   if (!reconfigure) {
1708     nvenc->input_info = *info;
1709     nvenc->gl_input = FALSE;
1710   }
1711
1712   if (nvenc->input_state)
1713     gst_video_codec_state_unref (nvenc->input_state);
1714   nvenc->input_state = gst_video_codec_state_ref (state);
1715   GST_INFO_OBJECT (nvenc, "%sconfigured encoder", reconfigure ? "re" : "");
1716
1717   /* now allocate some buffers only on first configuration */
1718   if (!reconfigure) {
1719 #if HAVE_NVCODEC_GST_GL
1720     GstCapsFeatures *features;
1721 #endif
1722     guint i;
1723     guint input_width, input_height;
1724
1725     input_width = GST_VIDEO_INFO_WIDTH (info);
1726     input_height = GST_VIDEO_INFO_HEIGHT (info);
1727
1728     nvenc->n_bufs =
1729         gst_nv_base_enc_calculate_num_prealloc_buffers (nvenc,
1730         params->encodeConfig);
1731
1732     /* input buffers */
1733     g_array_set_size (nvenc->items, nvenc->n_bufs);
1734
1735 #if HAVE_NVCODEC_GST_GL
1736     features = gst_caps_get_features (state->caps, 0);
1737     if (gst_caps_features_contains (features,
1738             GST_CAPS_FEATURE_MEMORY_GL_MEMORY)) {
1739       nvenc->gl_input = TRUE;
1740     }
1741 #endif
1742
1743     gst_cuda_context_push (nvenc->cuda_ctx);
1744     for (i = 0; i < nvenc->n_bufs; ++i) {
1745       GstNvEncInputResource *resource = g_new0 (GstNvEncInputResource, 1);
1746       CUresult cu_ret;
1747
1748       memset (&resource->nv_resource, 0, sizeof (resource->nv_resource));
1749       memset (&resource->nv_mapped_resource, 0,
1750           sizeof (resource->nv_mapped_resource));
1751
1752       /* scratch buffer for non-contigious planer into a contigious buffer */
1753       cu_ret =
1754           CuMemAllocPitch (&resource->cuda_pointer,
1755           &resource->cuda_stride, _get_plane_width (info, 0),
1756           _get_frame_data_height (info), 16);
1757       if (!gst_cuda_result (cu_ret)) {
1758         GST_ERROR_OBJECT (nvenc, "failed to alocate cuda scratch buffer "
1759             "ret %d", cu_ret);
1760         g_assert_not_reached ();
1761       }
1762
1763       resource->nv_resource.version =
1764           gst_nvenc_get_registure_resource_version ();
1765       resource->nv_resource.resourceType =
1766           NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
1767       resource->nv_resource.width = input_width;
1768       resource->nv_resource.height = input_height;
1769       resource->nv_resource.pitch = resource->cuda_stride;
1770       resource->nv_resource.bufferFormat =
1771           gst_nvenc_get_nv_buffer_format (GST_VIDEO_INFO_FORMAT (info));
1772       resource->nv_resource.resourceToRegister =
1773           (gpointer) resource->cuda_pointer;
1774
1775       nv_ret = NvEncRegisterResource (nvenc->encoder, &resource->nv_resource);
1776       if (nv_ret != NV_ENC_SUCCESS)
1777         GST_ERROR_OBJECT (nvenc, "Failed to register resource %p, ret %d",
1778             resource, nv_ret);
1779
1780       g_array_index (nvenc->items, GstNvEncFrameState, i).in_buf = resource;
1781     }
1782     gst_cuda_context_pop (NULL);
1783
1784     /* output buffers */
1785     for (i = 0; i < nvenc->n_bufs; ++i) {
1786       NV_ENC_CREATE_BITSTREAM_BUFFER cout_buf = { 0, };
1787
1788       cout_buf.version = gst_nvenc_get_create_bitstream_buffer_version ();
1789
1790       /* 1 MB should be large enough to hold most output frames.
1791        * NVENC will automatically increase this if it's not enough. */
1792       cout_buf.size = 1024 * 1024;
1793       cout_buf.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;
1794
1795       G_LOCK (initialization_lock);
1796       nv_ret = NvEncCreateBitstreamBuffer (nvenc->encoder, &cout_buf);
1797       G_UNLOCK (initialization_lock);
1798
1799       if (nv_ret != NV_ENC_SUCCESS) {
1800         GST_WARNING_OBJECT (enc, "Failed to allocate input buffer: %d", nv_ret);
1801         /* FIXME: clean up */
1802         return FALSE;
1803       }
1804
1805       GST_INFO_OBJECT (nvenc, "allocated output buffer %2d: %p", i,
1806           cout_buf.bitstreamBuffer);
1807
1808       g_array_index (nvenc->items, GstNvEncFrameState, i).out_buf =
1809           cout_buf.bitstreamBuffer;
1810
1811       g_async_queue_push (nvenc->available_queue, &g_array_index (nvenc->items,
1812               GstNvEncFrameState, i));
1813     }
1814
1815 #if 0
1816     /* Get SPS/PPS */
1817     {
1818       NV_ENC_SEQUENCE_PARAM_PAYLOAD seq_param = { 0 };
1819       uint32_t seq_size = 0;
1820
1821       seq_param.version = gst_nvenc_get_sequence_param_payload_version ();
1822       seq_param.spsppsBuffer = g_alloca (1024);
1823       seq_param.inBufferSize = 1024;
1824       seq_param.outSPSPPSPayloadSize = &seq_size;
1825
1826       nv_ret = NvEncGetSequenceParams (nvenc->encoder, &seq_param);
1827       if (nv_ret != NV_ENC_SUCCESS) {
1828         GST_WARNING_OBJECT (enc, "Failed to retrieve SPS/PPS: %d", nv_ret);
1829         return FALSE;
1830       }
1831
1832       /* FIXME: use SPS/PPS */
1833       GST_MEMDUMP_OBJECT (enc, "SPS/PPS", seq_param.spsppsBuffer, seq_size);
1834     }
1835 #endif
1836   }
1837
1838   g_assert (nvenc_class->set_src_caps);
1839   if (!nvenc_class->set_src_caps (nvenc, state)) {
1840     GST_ERROR_OBJECT (nvenc, "Subclass failed to set output caps");
1841     /* FIXME: clean up */
1842     return FALSE;
1843   }
1844
1845   return TRUE;
1846 }
1847
1848 static guint
1849 _get_cuda_device_stride (GstVideoInfo * info, guint plane, gsize cuda_stride)
1850 {
1851   switch (GST_VIDEO_INFO_FORMAT (info)) {
1852     case GST_VIDEO_FORMAT_NV12:
1853     case GST_VIDEO_FORMAT_NV21:
1854     case GST_VIDEO_FORMAT_P010_10LE:
1855     case GST_VIDEO_FORMAT_P010_10BE:
1856     case GST_VIDEO_FORMAT_Y444:
1857     case GST_VIDEO_FORMAT_BGRA:
1858     case GST_VIDEO_FORMAT_RGBA:
1859     case GST_VIDEO_FORMAT_BGR10A2_LE:
1860     case GST_VIDEO_FORMAT_RGB10A2_LE:
1861     case GST_VIDEO_FORMAT_Y444_16LE:
1862     case GST_VIDEO_FORMAT_Y444_16BE:
1863     case GST_VIDEO_FORMAT_VUYA:
1864       return cuda_stride;
1865     case GST_VIDEO_FORMAT_I420:
1866     case GST_VIDEO_FORMAT_YV12:
1867       return plane == 0 ? cuda_stride : (GST_ROUND_UP_2 (cuda_stride) / 2);
1868     default:
1869       g_assert_not_reached ();
1870       return cuda_stride;
1871   }
1872 }
1873
1874 #if HAVE_NVCODEC_GST_GL
1875 typedef struct _GstNvEncRegisterResourceData
1876 {
1877   GstMemory *mem;
1878   GstCudaGraphicsResource *resource;
1879   GstNvBaseEnc *nvenc;
1880   gboolean ret;
1881 } GstNvEncRegisterResourceData;
1882
1883 static void
1884 register_cuda_resource (GstGLContext * context,
1885     GstNvEncRegisterResourceData * data)
1886 {
1887   GstMemory *mem = data->mem;
1888   GstCudaGraphicsResource *resource = data->resource;
1889   GstNvBaseEnc *nvenc = data->nvenc;
1890   GstMapInfo map_info = GST_MAP_INFO_INIT;
1891   GstGLBuffer *gl_buf_obj;
1892
1893   data->ret = FALSE;
1894
1895   if (!gst_cuda_context_push (nvenc->cuda_ctx)) {
1896     GST_WARNING_OBJECT (nvenc, "failed to push CUDA context");
1897     return;
1898   }
1899
1900   if (gst_memory_map (mem, &map_info, GST_MAP_READ | GST_MAP_GL)) {
1901     GstGLMemoryPBO *gl_mem = (GstGLMemoryPBO *) data->mem;
1902     gl_buf_obj = gl_mem->pbo;
1903
1904     GST_LOG_OBJECT (nvenc,
1905         "registure glbuffer %d to CUDA resource", gl_buf_obj->id);
1906
1907     if (gst_cuda_graphics_resource_register_gl_buffer (resource,
1908             gl_buf_obj->id, CU_GRAPHICS_REGISTER_FLAGS_NONE)) {
1909       data->ret = TRUE;
1910     } else {
1911       GST_WARNING_OBJECT (nvenc, "failed to register memory");
1912     }
1913
1914     gst_memory_unmap (mem, &map_info);
1915   } else {
1916     GST_WARNING_OBJECT (nvenc, "failed to map memory");
1917   }
1918
1919   if (!gst_cuda_context_pop (NULL))
1920     GST_WARNING_OBJECT (nvenc, "failed to unlock CUDA context");
1921 }
1922
1923 static GstCudaGraphicsResource *
1924 ensure_cuda_graphics_resource (GstMemory * mem, GstNvBaseEnc * nvenc)
1925 {
1926   GQuark quark;
1927   GstCudaGraphicsResource *cgr_info;
1928   GstNvEncRegisterResourceData data;
1929
1930   if (!gst_is_gl_memory_pbo (mem)) {
1931     GST_WARNING_OBJECT (nvenc, "memory is not GL PBO memory, %s",
1932         mem->allocator->mem_type);
1933     return NULL;
1934   }
1935
1936   quark = gst_cuda_quark_from_id (GST_CUDA_QUARK_GRAPHICS_RESOURCE);
1937
1938   cgr_info = gst_mini_object_get_qdata (GST_MINI_OBJECT (mem), quark);
1939   if (!cgr_info) {
1940     cgr_info = gst_cuda_graphics_resource_new (nvenc->cuda_ctx,
1941         GST_OBJECT (GST_GL_BASE_MEMORY_CAST (mem)->context),
1942         GST_CUDA_GRAPHICS_RESOURCE_GL_BUFFER);
1943     data.mem = mem;
1944     data.resource = cgr_info;
1945     data.nvenc = nvenc;
1946     gst_gl_context_thread_add ((GstGLContext *) cgr_info->graphics_context,
1947         (GstGLContextThreadFunc) register_cuda_resource, &data);
1948     if (!data.ret) {
1949       GST_WARNING_OBJECT (nvenc, "could not register resource");
1950       gst_cuda_graphics_resource_free (cgr_info);
1951
1952       return NULL;
1953     }
1954
1955     gst_mini_object_set_qdata (GST_MINI_OBJECT (mem), quark, cgr_info,
1956         (GDestroyNotify) gst_cuda_graphics_resource_free);
1957   }
1958
1959   return cgr_info;
1960 }
1961
1962 typedef struct _GstNvEncGLMapData
1963 {
1964   GstNvBaseEnc *nvenc;
1965   GstBuffer *buffer;
1966   GstVideoInfo *info;
1967   GstNvEncInputResource *resource;
1968
1969   gboolean ret;
1970 } GstNvEncGLMapData;
1971
1972 static void
1973 _map_gl_input_buffer (GstGLContext * context, GstNvEncGLMapData * data)
1974 {
1975   GstNvBaseEnc *nvenc = data->nvenc;
1976   CUresult cuda_ret;
1977   CUdeviceptr data_pointer;
1978   guint i;
1979   CUDA_MEMCPY2D param;
1980   GstCudaGraphicsResource **resources;
1981   guint num_resources;
1982
1983   data->ret = FALSE;
1984
1985   num_resources = gst_buffer_n_memory (data->buffer);
1986   resources = g_newa (GstCudaGraphicsResource *, num_resources);
1987
1988   for (i = 0; i < num_resources; i++) {
1989     GstMemory *mem;
1990
1991     mem = gst_buffer_peek_memory (data->buffer, i);
1992     resources[i] = ensure_cuda_graphics_resource (mem, nvenc);
1993     if (!resources[i]) {
1994       GST_ERROR_OBJECT (nvenc, "could not register %dth memory", i);
1995       return;
1996     }
1997   }
1998
1999   gst_cuda_context_push (nvenc->cuda_ctx);
2000   data_pointer = data->resource->cuda_pointer;
2001   for (i = 0; i < GST_VIDEO_INFO_N_PLANES (data->info); i++) {
2002     GstGLBuffer *gl_buf_obj;
2003     GstGLMemoryPBO *gl_mem;
2004     guint src_stride, dest_stride;
2005     CUgraphicsResource cuda_resource;
2006     gsize cuda_num_bytes;
2007     CUdeviceptr cuda_plane_pointer;
2008
2009     gl_mem = (GstGLMemoryPBO *) gst_buffer_peek_memory (data->buffer, i);
2010     g_return_if_fail (gst_is_gl_memory_pbo ((GstMemory *) gl_mem));
2011
2012     gl_buf_obj = (GstGLBuffer *) gl_mem->pbo;
2013     g_return_if_fail (gl_buf_obj != NULL);
2014
2015     /* get the texture into the PBO */
2016     gst_gl_memory_pbo_upload_transfer (gl_mem);
2017     gst_gl_memory_pbo_download_transfer (gl_mem);
2018
2019     GST_LOG_OBJECT (nvenc, "attempting to copy texture %u into cuda",
2020         gl_mem->mem.tex_id);
2021
2022     cuda_resource =
2023         gst_cuda_graphics_resource_map (resources[i], nvenc->cuda_stream,
2024         CU_GRAPHICS_MAP_RESOURCE_FLAGS_READ_ONLY);
2025
2026     if (!cuda_resource) {
2027       GST_ERROR_OBJECT (nvenc, "failed to map GL texture %u into cuda",
2028           gl_mem->mem.tex_id);
2029       g_assert_not_reached ();
2030     }
2031
2032     cuda_ret =
2033         CuGraphicsResourceGetMappedPointer (&cuda_plane_pointer,
2034         &cuda_num_bytes, cuda_resource);
2035
2036     if (!gst_cuda_result (cuda_ret)) {
2037       GST_ERROR_OBJECT (nvenc, "failed to get mapped pointer of map GL "
2038           "texture %u in cuda ret :%d", gl_mem->mem.tex_id, cuda_ret);
2039       g_assert_not_reached ();
2040     }
2041
2042     src_stride = GST_VIDEO_INFO_PLANE_STRIDE (data->info, i);
2043     dest_stride = _get_cuda_device_stride (&nvenc->input_info,
2044         i, data->resource->cuda_stride);
2045
2046     /* copy into scratch buffer */
2047     param.srcXInBytes = 0;
2048     param.srcY = 0;
2049     param.srcMemoryType = CU_MEMORYTYPE_DEVICE;
2050     param.srcDevice = cuda_plane_pointer;
2051     param.srcPitch = src_stride;
2052
2053     param.dstXInBytes = 0;
2054     param.dstY = 0;
2055     param.dstMemoryType = CU_MEMORYTYPE_DEVICE;
2056     param.dstDevice = data_pointer;
2057     param.dstPitch = dest_stride;
2058     param.WidthInBytes = _get_plane_width (data->info, i);
2059     param.Height = _get_plane_height (data->info, i);
2060
2061     cuda_ret = CuMemcpy2DAsync (&param, nvenc->cuda_stream);
2062     if (!gst_cuda_result (cuda_ret)) {
2063       GST_ERROR_OBJECT (data->nvenc, "failed to copy GL texture %u into cuda "
2064           "ret :%d", gl_mem->mem.tex_id, cuda_ret);
2065       g_assert_not_reached ();
2066     }
2067
2068     gst_cuda_graphics_resource_unmap (resources[i], nvenc->cuda_stream);
2069
2070     data_pointer += dest_stride * _get_plane_height (&nvenc->input_info, i);
2071   }
2072   gst_cuda_result (CuStreamSynchronize (nvenc->cuda_stream));
2073   gst_cuda_context_pop (NULL);
2074
2075   data->ret = TRUE;
2076 }
2077 #endif
2078
2079 static gboolean
2080 gst_nv_base_enc_upload_frame (GstNvBaseEnc * nvenc, GstVideoFrame * frame,
2081     GstNvEncInputResource * resource)
2082 {
2083   gint i;
2084   CUdeviceptr dst = resource->cuda_pointer;
2085   GstVideoInfo *info = &frame->info;
2086   CUresult cuda_ret;
2087
2088   if (!gst_cuda_context_push (nvenc->cuda_ctx)) {
2089     GST_ERROR_OBJECT (nvenc, "cannot push context");
2090     return FALSE;
2091   }
2092
2093   for (i = 0; i < GST_VIDEO_FRAME_N_PLANES (frame); i++) {
2094     CUDA_MEMCPY2D param = { 0, };
2095     guint dest_stride = _get_cuda_device_stride (&nvenc->input_info, i,
2096         resource->cuda_stride);
2097
2098     param.srcMemoryType = CU_MEMORYTYPE_HOST;
2099     param.srcHost = GST_VIDEO_FRAME_PLANE_DATA (frame, i);
2100     param.srcPitch = GST_VIDEO_FRAME_PLANE_STRIDE (frame, i);
2101
2102     param.dstMemoryType = CU_MEMORYTYPE_DEVICE;
2103     param.dstDevice = dst;
2104     param.dstPitch = dest_stride;
2105     param.WidthInBytes = _get_plane_width (info, i);
2106     param.Height = _get_plane_height (info, i);
2107
2108     cuda_ret = CuMemcpy2DAsync (&param, nvenc->cuda_stream);
2109     if (!gst_cuda_result (cuda_ret)) {
2110       GST_ERROR_OBJECT (nvenc, "cannot copy %dth plane, ret %d", i, cuda_ret);
2111       gst_cuda_context_pop (NULL);
2112
2113       return FALSE;
2114     }
2115
2116     dst += dest_stride * _get_plane_height (&nvenc->input_info, i);
2117   }
2118
2119   gst_cuda_result (CuStreamSynchronize (nvenc->cuda_stream));
2120   gst_cuda_context_pop (NULL);
2121
2122   return TRUE;
2123 }
2124
2125 static GstFlowReturn
2126 _acquire_input_buffer (GstNvBaseEnc * nvenc, GstNvEncFrameState ** input)
2127 {
2128   GST_LOG_OBJECT (nvenc, "acquiring input buffer..");
2129   GST_VIDEO_ENCODER_STREAM_UNLOCK (nvenc);
2130   *input = g_async_queue_pop (nvenc->available_queue);
2131   GST_VIDEO_ENCODER_STREAM_LOCK (nvenc);
2132
2133   if (*input == SHUTDOWN_COOKIE)
2134     return g_atomic_int_get (&nvenc->last_flow);
2135
2136   return GST_FLOW_OK;
2137 }
2138
2139 static GstFlowReturn
2140 _submit_input_buffer (GstNvBaseEnc * nvenc, GstVideoCodecFrame * frame,
2141     GstVideoFrame * vframe, GstNvEncFrameState * state, void *inputBufferPtr,
2142     NV_ENC_BUFFER_FORMAT bufferFormat)
2143 {
2144   GstNvBaseEncClass *nvenc_class = GST_NV_BASE_ENC_GET_CLASS (nvenc);
2145   NV_ENC_PIC_PARAMS pic_params = { 0, };
2146   NVENCSTATUS nv_ret;
2147   gpointer inputBuffer, outputBufferPtr;
2148
2149   inputBuffer = state->in_buf;
2150   outputBufferPtr = state->out_buf;
2151
2152   GST_LOG_OBJECT (nvenc, "%u: input buffer %p, output buffer %p, "
2153       "pts %" GST_TIME_FORMAT, frame->system_frame_number, inputBuffer,
2154       outputBufferPtr, GST_TIME_ARGS (frame->pts));
2155
2156   pic_params.version = gst_nvenc_get_pic_params_version ();
2157   pic_params.inputBuffer = inputBufferPtr;
2158   pic_params.bufferFmt = bufferFormat;
2159
2160   pic_params.inputWidth = GST_VIDEO_FRAME_WIDTH (vframe);
2161   pic_params.inputHeight = GST_VIDEO_FRAME_HEIGHT (vframe);
2162   pic_params.outputBitstream = outputBufferPtr;
2163   pic_params.completionEvent = NULL;
2164   if (GST_VIDEO_FRAME_IS_INTERLACED (vframe)) {
2165     if (GST_VIDEO_FRAME_IS_TFF (vframe))
2166       pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
2167     else
2168       pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
2169   } else {
2170     pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
2171   }
2172   pic_params.inputTimeStamp = frame->pts;
2173   pic_params.inputDuration =
2174       GST_CLOCK_TIME_IS_VALID (frame->duration) ? frame->duration : 0;
2175   pic_params.frameIdx = frame->system_frame_number;
2176
2177   if (GST_VIDEO_CODEC_FRAME_IS_FORCE_KEYFRAME (frame))
2178     pic_params.encodePicFlags = NV_ENC_PIC_FLAG_FORCEIDR;
2179   else
2180     pic_params.encodePicFlags = 0;
2181
2182   if (nvenc_class->set_pic_params
2183       && !nvenc_class->set_pic_params (nvenc, frame, &pic_params)) {
2184     GST_ERROR_OBJECT (nvenc, "Subclass failed to submit buffer");
2185     return GST_FLOW_ERROR;
2186   }
2187
2188   if (!gst_cuda_context_push (nvenc->cuda_ctx)) {
2189     GST_ELEMENT_ERROR (nvenc, LIBRARY, ENCODE, (NULL),
2190         ("Failed to push current context"));
2191     return GST_FLOW_ERROR;
2192   }
2193
2194   nv_ret = NvEncEncodePicture (nvenc->encoder, &pic_params);
2195
2196   gst_cuda_context_pop (NULL);
2197
2198   if (nv_ret == NV_ENC_SUCCESS) {
2199     GST_LOG_OBJECT (nvenc, "Encoded picture");
2200   } else if (nv_ret == NV_ENC_ERR_NEED_MORE_INPUT) {
2201     GST_DEBUG_OBJECT (nvenc, "Encoded picture (encoder needs more input)");
2202   } else {
2203     GST_ERROR_OBJECT (nvenc, "Failed to encode picture: %d", nv_ret);
2204     g_async_queue_push (nvenc->available_queue, state);
2205
2206     return GST_FLOW_ERROR;
2207   }
2208
2209   /* GstNvEncFrameState shouldn't be freed by DestroyNotify */
2210   gst_video_codec_frame_set_user_data (frame, state, NULL);
2211   g_async_queue_push (nvenc->pending_queue, state);
2212
2213   if (nv_ret == NV_ENC_SUCCESS) {
2214     GstNvEncFrameState *pending_state;
2215     gint len, i, end;
2216
2217     /* HACK: NvEncEncodePicture() with returning NV_ENC_SUCCESS means that
2218      * we can pop encoded bitstream from GPU
2219      * (via NvEncLockBitstream and copy to memory then NvEncUnlockBitstream).
2220      * But if we try to pop every buffer from GPU when the rc-lookahead
2221      * was enabled, NvEncLockBitstream returns error NV_ENC_ERR_INVALID_PARAM
2222      * randomly (seemingly it's dependent on how fast the encoding thread
2223      * dequeued the encoded picture).
2224      * So make "pending_queue" having the number of lookahead pictures always,
2225      * so that GPU should be able to reference the lookahead pictures.
2226      *
2227      * This behavior is not documented by Nvidia. The guess here is that
2228      * the lookahead pictures are still used for rate-control by Nvidia driver
2229      * and dequeuing the lookahead picture from GPU seems to be causing the
2230      * problem.
2231      */
2232     end = nvenc->rc_lookahead;
2233
2234     g_async_queue_lock (nvenc->pending_queue);
2235
2236     len = g_async_queue_length_unlocked (nvenc->pending_queue);
2237     for (i = len; i > end; i--) {
2238       pending_state = g_async_queue_pop_unlocked (nvenc->pending_queue);
2239       g_async_queue_push (nvenc->bitstream_queue, pending_state);
2240     }
2241
2242     g_async_queue_unlock (nvenc->pending_queue);
2243   }
2244
2245   return GST_FLOW_OK;
2246 }
2247
2248 static GstFlowReturn
2249 gst_nv_base_enc_handle_frame (GstVideoEncoder * enc, GstVideoCodecFrame * frame)
2250 {
2251   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
2252   NVENCSTATUS nv_ret;
2253   GstVideoFrame vframe;
2254   GstVideoInfo *info = &nvenc->input_state->info;
2255   GstFlowReturn flow = GST_FLOW_OK;
2256   GstMapFlags in_map_flags = GST_MAP_READ;
2257   GstNvEncFrameState *state = NULL;
2258   GstNvEncInputResource *resource = NULL;
2259
2260   g_assert (nvenc->encoder != NULL);
2261
2262   /* check last flow and if it's not OK, just return the last flow,
2263    * non-OK flow means that encoding thread was terminated */
2264   flow = g_atomic_int_get (&nvenc->last_flow);
2265   if (flow != GST_FLOW_OK) {
2266     GST_DEBUG_OBJECT (nvenc, "last flow was %s", gst_flow_get_name (flow));
2267     /* just drop this frame */
2268     gst_video_encoder_finish_frame (enc, frame);
2269
2270     return flow;
2271   }
2272
2273   if (g_atomic_int_compare_and_exchange (&nvenc->reconfig, TRUE, FALSE)) {
2274     if (!gst_nv_base_enc_set_format (enc, nvenc->input_state)) {
2275       flow = GST_FLOW_NOT_NEGOTIATED;
2276       goto drop;
2277     }
2278
2279     /* reconfigured encode session should start from keyframe */
2280     GST_VIDEO_CODEC_FRAME_SET_FORCE_KEYFRAME (frame);
2281   }
2282 #if HAVE_NVCODEC_GST_GL
2283   if (nvenc->gl_input)
2284     in_map_flags |= GST_MAP_GL;
2285 #endif
2286
2287   if (!gst_video_frame_map (&vframe, info, frame->input_buffer, in_map_flags)) {
2288     goto drop;
2289   }
2290
2291   /* make sure our thread that waits for output to be ready is started */
2292   if (nvenc->bitstream_thread == NULL) {
2293     if (!gst_nv_base_enc_start_bitstream_thread (nvenc)) {
2294       gst_video_frame_unmap (&vframe);
2295       goto unmap_and_drop;
2296     }
2297   }
2298
2299   flow = _acquire_input_buffer (nvenc, &state);
2300   if (flow != GST_FLOW_OK || state == SHUTDOWN_COOKIE || !state)
2301     goto unmap_and_drop;
2302
2303   resource = state->in_buf;
2304
2305 #if HAVE_NVCODEC_GST_GL
2306   if (nvenc->gl_input) {
2307     GstGLMemory *gl_mem;
2308     GstNvEncGLMapData data;
2309
2310     gl_mem = (GstGLMemory *) gst_buffer_peek_memory (frame->input_buffer, 0);
2311     g_assert (gst_is_gl_memory ((GstMemory *) gl_mem));
2312
2313     data.nvenc = nvenc;
2314     data.buffer = frame->input_buffer;
2315     data.info = &vframe.info;
2316     data.resource = resource;
2317
2318     gst_gl_context_thread_add (gl_mem->mem.context,
2319         (GstGLContextThreadFunc) _map_gl_input_buffer, &data);
2320     if (!data.ret) {
2321       flow = GST_FLOW_ERROR;
2322       goto unmap_and_drop;
2323     }
2324   } else
2325 #endif
2326   if (!gst_nv_base_enc_upload_frame (nvenc, &vframe, resource)) {
2327     flow = GST_FLOW_ERROR;
2328     goto unmap_and_drop;
2329   }
2330
2331   resource->nv_mapped_resource.version =
2332       gst_nvenc_get_map_input_resource_version ();
2333   resource->nv_mapped_resource.registeredResource =
2334       resource->nv_resource.registeredResource;
2335
2336   if (!gst_cuda_context_push (nvenc->cuda_ctx)) {
2337     GST_ELEMENT_ERROR (nvenc, LIBRARY, ENCODE, (NULL),
2338         ("Failed to push current context"));
2339     flow = GST_FLOW_ERROR;
2340     goto unmap_and_drop;
2341   }
2342
2343   nv_ret =
2344       NvEncMapInputResource (nvenc->encoder, &resource->nv_mapped_resource);
2345   gst_cuda_context_pop (NULL);
2346
2347   if (nv_ret != NV_ENC_SUCCESS) {
2348     GST_ERROR_OBJECT (nvenc, "Failed to map input resource %p, ret %d",
2349         resource, nv_ret);
2350     flow = GST_FLOW_ERROR;
2351     goto unmap_and_drop;
2352   }
2353
2354   resource->mapped = TRUE;
2355
2356   flow =
2357       _submit_input_buffer (nvenc, frame, &vframe, state,
2358       resource->nv_mapped_resource.mappedResource,
2359       resource->nv_mapped_resource.mappedBufferFmt);
2360
2361   if (flow != GST_FLOW_OK) {
2362     GST_DEBUG_OBJECT (nvenc, "return state to pool");
2363     g_async_queue_push (nvenc->available_queue, state);
2364     goto unmap_and_drop;
2365   }
2366
2367   flow = g_atomic_int_get (&nvenc->last_flow);
2368
2369   gst_video_frame_unmap (&vframe);
2370   /* encoder will keep frame in list internally, we'll look it up again later
2371    * in the thread where we get the output buffers and finish it there */
2372   gst_video_codec_frame_unref (frame);
2373
2374   return flow;
2375
2376 /* ERRORS */
2377 unmap_and_drop:
2378   {
2379     gst_video_frame_unmap (&vframe);
2380     goto drop;
2381   }
2382 drop:
2383   {
2384     gst_video_encoder_finish_frame (enc, frame);
2385     return flow;
2386   }
2387 }
2388
2389 static gboolean
2390 gst_nv_base_enc_drain_encoder (GstNvBaseEnc * nvenc)
2391 {
2392   NV_ENC_PIC_PARAMS pic_params = { 0, };
2393   NVENCSTATUS nv_ret;
2394   gboolean ret = TRUE;
2395
2396   GST_INFO_OBJECT (nvenc, "draining encoder");
2397
2398   if (nvenc->input_state == NULL) {
2399     GST_DEBUG_OBJECT (nvenc, "no input state, nothing to do");
2400     return TRUE;
2401   }
2402
2403   pic_params.version = gst_nvenc_get_pic_params_version ();
2404   pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
2405
2406   if (!gst_cuda_context_push (nvenc->cuda_ctx)) {
2407     GST_ERROR_OBJECT (nvenc, "Could not push context");
2408     return GST_FLOW_ERROR;
2409   }
2410
2411   nv_ret = NvEncEncodePicture (nvenc->encoder, &pic_params);
2412
2413   if (nv_ret != NV_ENC_SUCCESS) {
2414     GST_LOG_OBJECT (nvenc, "Failed to drain encoder, ret %d", nv_ret);
2415
2416     ret = FALSE;
2417   } else {
2418     GstNvEncFrameState *pending_state;
2419
2420     g_async_queue_lock (nvenc->pending_queue);
2421     while ((pending_state =
2422             g_async_queue_try_pop_unlocked (nvenc->pending_queue))) {
2423       g_async_queue_push (nvenc->bitstream_queue, pending_state);
2424     }
2425     g_async_queue_unlock (nvenc->pending_queue);
2426   }
2427
2428   gst_cuda_context_pop (NULL);
2429
2430   return ret;
2431 }
2432
2433 static GstFlowReturn
2434 gst_nv_base_enc_finish (GstVideoEncoder * enc)
2435 {
2436   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
2437
2438   gst_nv_base_enc_stop_bitstream_thread (nvenc, FALSE);
2439
2440   return GST_FLOW_OK;
2441 }
2442
2443 #if 0
2444 static gboolean
2445 gst_nv_base_enc_flush (GstVideoEncoder * enc)
2446 {
2447   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (enc);
2448   GST_INFO_OBJECT (nvenc, "done flushing encoder");
2449   return TRUE;
2450 }
2451 #endif
2452
2453 void
2454 gst_nv_base_enc_schedule_reconfig (GstNvBaseEnc * nvenc)
2455 {
2456   g_atomic_int_set (&nvenc->reconfig, TRUE);
2457 }
2458
2459 static void
2460 gst_nv_base_enc_set_property (GObject * object, guint prop_id,
2461     const GValue * value, GParamSpec * pspec)
2462 {
2463   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (object);
2464   GstNvBaseEncClass *klass = GST_NV_BASE_ENC_GET_CLASS (nvenc);
2465   gboolean reconfig = TRUE;
2466
2467   switch (prop_id) {
2468     case PROP_PRESET:
2469       nvenc->preset_enum = g_value_get_enum (value);
2470       nvenc->selected_preset = _nv_preset_to_guid (nvenc->preset_enum);
2471       gst_nv_base_enc_schedule_reconfig (nvenc);
2472       break;
2473     case PROP_RC_MODE:
2474     {
2475       GstNvRCMode rc_mode = g_value_get_enum (value);
2476       NV_ENC_PARAMS_RC_MODE nv_rc_mode = _rc_mode_to_nv (rc_mode);
2477
2478       if ((klass->device_caps.rc_modes & nv_rc_mode) == nv_rc_mode) {
2479         nvenc->rate_control_mode = rc_mode;
2480       } else {
2481         GST_WARNING_OBJECT (nvenc,
2482             "device does not support requested rate control mode %d", rc_mode);
2483         reconfig = FALSE;
2484       }
2485       break;
2486     }
2487     case PROP_QP_MIN:
2488       nvenc->qp_min = g_value_get_int (value);
2489       break;
2490     case PROP_QP_MAX:
2491       nvenc->qp_max = g_value_get_int (value);
2492       break;
2493     case PROP_QP_CONST:
2494       nvenc->qp_const = g_value_get_int (value);
2495       break;
2496     case PROP_BITRATE:
2497       nvenc->bitrate = g_value_get_uint (value);
2498       break;
2499     case PROP_GOP_SIZE:
2500       nvenc->gop_size = g_value_get_int (value);
2501       break;
2502     case PROP_MAX_BITRATE:
2503       nvenc->max_bitrate = g_value_get_uint (value);
2504       break;
2505     case PROP_SPATIAL_AQ:
2506       nvenc->spatial_aq = g_value_get_boolean (value);
2507       break;
2508     case PROP_AQ_STRENGTH:
2509       nvenc->aq_strength = g_value_get_uint (value);
2510       break;
2511     case PROP_NON_REF_P:
2512       nvenc->non_refp = g_value_get_boolean (value);
2513       break;
2514     case PROP_ZEROLATENCY:
2515       nvenc->zerolatency = g_value_get_boolean (value);
2516       break;
2517     case PROP_STRICT_GOP:
2518       nvenc->strict_gop = g_value_get_boolean (value);
2519       break;
2520     case PROP_CONST_QUALITY:
2521       nvenc->const_quality = g_value_get_double (value);
2522       break;
2523     case PROP_I_ADAPT:
2524       nvenc->i_adapt = g_value_get_boolean (value);
2525       break;
2526     case PROP_QP_MIN_I:
2527       nvenc->qp_min_detail.qp_i = g_value_get_int (value);
2528       break;
2529     case PROP_QP_MIN_P:
2530       nvenc->qp_min_detail.qp_p = g_value_get_int (value);
2531       break;
2532     case PROP_QP_MIN_B:
2533       nvenc->qp_min_detail.qp_b = g_value_get_int (value);
2534       break;
2535     case PROP_QP_MAX_I:
2536       nvenc->qp_max_detail.qp_i = g_value_get_int (value);
2537       break;
2538     case PROP_QP_MAX_P:
2539       nvenc->qp_max_detail.qp_p = g_value_get_int (value);
2540       break;
2541     case PROP_QP_MAX_B:
2542       nvenc->qp_max_detail.qp_b = g_value_get_int (value);
2543       break;
2544     case PROP_QP_CONST_I:
2545       nvenc->qp_const_detail.qp_i = g_value_get_int (value);
2546       break;
2547     case PROP_QP_CONST_P:
2548       nvenc->qp_const_detail.qp_p = g_value_get_int (value);
2549       break;
2550     case PROP_QP_CONST_B:
2551       nvenc->qp_const_detail.qp_b = g_value_get_int (value);
2552       break;
2553     default:
2554       G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);
2555       reconfig = FALSE;
2556       break;
2557   }
2558
2559   if (reconfig)
2560     gst_nv_base_enc_schedule_reconfig (nvenc);
2561 }
2562
2563 static void
2564 gst_nv_base_enc_get_property (GObject * object, guint prop_id, GValue * value,
2565     GParamSpec * pspec)
2566 {
2567   GstNvBaseEnc *nvenc = GST_NV_BASE_ENC (object);
2568   GstNvBaseEncClass *nvenc_class = GST_NV_BASE_ENC_GET_CLASS (object);
2569
2570   switch (prop_id) {
2571     case PROP_DEVICE_ID:
2572       g_value_set_uint (value, nvenc_class->cuda_device_id);
2573       break;
2574     case PROP_PRESET:
2575       g_value_set_enum (value, nvenc->preset_enum);
2576       break;
2577     case PROP_RC_MODE:
2578       g_value_set_enum (value, nvenc->rate_control_mode);
2579       break;
2580     case PROP_QP_MIN:
2581       g_value_set_int (value, nvenc->qp_min);
2582       break;
2583     case PROP_QP_MAX:
2584       g_value_set_int (value, nvenc->qp_max);
2585       break;
2586     case PROP_QP_CONST:
2587       g_value_set_int (value, nvenc->qp_const);
2588       break;
2589     case PROP_BITRATE:
2590       g_value_set_uint (value, nvenc->bitrate);
2591       break;
2592     case PROP_GOP_SIZE:
2593       g_value_set_int (value, nvenc->gop_size);
2594       break;
2595     case PROP_MAX_BITRATE:
2596       g_value_set_uint (value, nvenc->max_bitrate);
2597       break;
2598     case PROP_SPATIAL_AQ:
2599       g_value_set_boolean (value, nvenc->spatial_aq);
2600       break;
2601     case PROP_AQ_STRENGTH:
2602       g_value_set_uint (value, nvenc->aq_strength);
2603       break;
2604     case PROP_NON_REF_P:
2605       g_value_set_boolean (value, nvenc->non_refp);
2606       break;
2607     case PROP_ZEROLATENCY:
2608       g_value_set_boolean (value, nvenc->zerolatency);
2609       break;
2610     case PROP_STRICT_GOP:
2611       g_value_set_boolean (value, nvenc->strict_gop);
2612       break;
2613     case PROP_CONST_QUALITY:
2614       g_value_set_double (value, nvenc->const_quality);
2615       break;
2616     case PROP_I_ADAPT:
2617       g_value_set_boolean (value, nvenc->i_adapt);
2618       break;
2619     case PROP_QP_MIN_I:
2620       g_value_set_int (value, nvenc->qp_min_detail.qp_i);
2621       break;
2622     case PROP_QP_MIN_P:
2623       g_value_set_int (value, nvenc->qp_min_detail.qp_p);
2624       break;
2625     case PROP_QP_MIN_B:
2626       g_value_set_int (value, nvenc->qp_min_detail.qp_b);
2627       break;
2628     case PROP_QP_MAX_I:
2629       g_value_set_int (value, nvenc->qp_max_detail.qp_i);
2630       break;
2631     case PROP_QP_MAX_P:
2632       g_value_set_int (value, nvenc->qp_max_detail.qp_p);
2633       break;
2634     case PROP_QP_MAX_B:
2635       g_value_set_int (value, nvenc->qp_max_detail.qp_b);
2636       break;
2637     case PROP_QP_CONST_I:
2638       g_value_set_int (value, nvenc->qp_const_detail.qp_i);
2639       break;
2640     case PROP_QP_CONST_P:
2641       g_value_set_int (value, nvenc->qp_const_detail.qp_p);
2642       break;
2643     case PROP_QP_CONST_B:
2644       g_value_set_int (value, nvenc->qp_const_detail.qp_b);
2645       break;
2646     default:
2647       G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);
2648       break;
2649   }
2650 }
2651
2652 typedef struct
2653 {
2654   guint cuda_device_id;
2655   GstNvEncDeviceCaps device_caps;
2656 } GstNvEncClassData;
2657
2658 static void
2659 gst_nv_base_enc_subclass_init (gpointer g_class, gpointer data)
2660 {
2661   GstNvBaseEncClass *nvbaseenc_class = GST_NV_BASE_ENC_CLASS (g_class);
2662   GstNvEncClassData *cdata = (GstNvEncClassData *) data;
2663
2664   nvbaseenc_class->cuda_device_id = cdata->cuda_device_id;
2665   nvbaseenc_class->device_caps = cdata->device_caps;
2666
2667   g_free (cdata);
2668 }
2669
2670 GType
2671 gst_nv_base_enc_register (const char *codec, guint device_id,
2672     GstNvEncDeviceCaps * device_caps)
2673 {
2674   GTypeQuery type_query;
2675   GTypeInfo type_info = { 0, };
2676   GType subtype;
2677   gchar *type_name;
2678   GstNvEncClassData *cdata;
2679
2680   type_name = g_strdup_printf ("GstNvDevice%d%sEnc", device_id, codec);
2681   subtype = g_type_from_name (type_name);
2682
2683   /* has already registered nvdeviceenc class */
2684   if (subtype)
2685     goto done;
2686
2687   cdata = g_new0 (GstNvEncClassData, 1);
2688   cdata->cuda_device_id = device_id;
2689   cdata->device_caps = *device_caps;
2690
2691   g_type_query (GST_TYPE_NV_BASE_ENC, &type_query);
2692   memset (&type_info, 0, sizeof (type_info));
2693   type_info.class_size = type_query.class_size;
2694   type_info.instance_size = type_query.instance_size;
2695   type_info.class_init = (GClassInitFunc) gst_nv_base_enc_subclass_init;
2696   type_info.class_data = cdata;
2697
2698   subtype = g_type_register_static (GST_TYPE_NV_BASE_ENC,
2699       type_name, &type_info, 0);
2700
2701 done:
2702   g_free (type_name);
2703   return subtype;
2704 }