avcodec/nvenc: rework library load and GPU selection
[ffmpeg.git] / libavcodec / nvenc.c
1 /*
2  * H.264 hardware encoding using nvidia nvenc
3  * Copyright (c) 2014 Timo Rothenpieler <timo@rothenpieler.org>
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with FFmpeg; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21
22 #include "config.h"
23
24 #if defined(_WIN32)
25 #include <windows.h>
26
27 #define CUDA_LIBNAME TEXT("nvcuda.dll")
28 #if ARCH_X86_64
29 #define NVENC_LIBNAME TEXT("nvEncodeAPI64.dll")
30 #else
31 #define NVENC_LIBNAME TEXT("nvEncodeAPI.dll")
32 #endif
33
34 #define dlopen(filename, flags) LoadLibrary((filename))
35 #define dlsym(handle, symbol)   GetProcAddress(handle, symbol)
36 #define dlclose(handle)         FreeLibrary(handle)
37 #else
38 #include <dlfcn.h>
39
40 #define CUDA_LIBNAME "libcuda.so"
41 #define NVENC_LIBNAME "libnvidia-encode.so"
42 #endif
43
44 #include "libavutil/hwcontext.h"
45 #include "libavutil/imgutils.h"
46 #include "libavutil/avassert.h"
47 #include "libavutil/mem.h"
48 #include "internal.h"
49 #include "nvenc.h"
50
51 #define NVENC_CAP 0x30
52 #define IS_CBR(rc) (rc == NV_ENC_PARAMS_RC_CBR ||               \
53                     rc == NV_ENC_PARAMS_RC_2_PASS_QUALITY ||    \
54                     rc == NV_ENC_PARAMS_RC_2_PASS_FRAMESIZE_CAP)
55
56 #define LOAD_LIBRARY(l, path)                   \
57     do {                                        \
58         if (!((l) = dlopen(path, RTLD_LAZY))) { \
59             av_log(avctx, AV_LOG_ERROR,         \
60                    "Cannot load %s\n",          \
61                    path);                       \
62             return AVERROR_UNKNOWN;             \
63         }                                       \
64     } while (0)
65
66 #define LOAD_SYMBOL(fun, lib, symbol)        \
67     do {                                     \
68         if (!((fun) = dlsym(lib, symbol))) { \
69             av_log(avctx, AV_LOG_ERROR,      \
70                    "Cannot load %s\n",       \
71                    symbol);                  \
72             return AVERROR_UNKNOWN;          \
73         }                                    \
74     } while (0)
75
76 const enum AVPixelFormat ff_nvenc_pix_fmts[] = {
77     AV_PIX_FMT_YUV420P,
78     AV_PIX_FMT_NV12,
79     AV_PIX_FMT_YUV444P,
80 #if CONFIG_CUDA
81     AV_PIX_FMT_CUDA,
82 #endif
83     AV_PIX_FMT_NONE
84 };
85
86 static const struct {
87     NVENCSTATUS nverr;
88     int         averr;
89     const char *desc;
90 } nvenc_errors[] = {
91     { NV_ENC_SUCCESS,                      0,                "success"                  },
92     { NV_ENC_ERR_NO_ENCODE_DEVICE,         AVERROR(ENOENT),  "no encode device"         },
93     { NV_ENC_ERR_UNSUPPORTED_DEVICE,       AVERROR(ENOSYS),  "unsupported device"       },
94     { NV_ENC_ERR_INVALID_ENCODERDEVICE,    AVERROR(EINVAL),  "invalid encoder device"   },
95     { NV_ENC_ERR_INVALID_DEVICE,           AVERROR(EINVAL),  "invalid device"           },
96     { NV_ENC_ERR_DEVICE_NOT_EXIST,         AVERROR(EIO),     "device does not exist"    },
97     { NV_ENC_ERR_INVALID_PTR,              AVERROR(EFAULT),  "invalid ptr"              },
98     { NV_ENC_ERR_INVALID_EVENT,            AVERROR(EINVAL),  "invalid event"            },
99     { NV_ENC_ERR_INVALID_PARAM,            AVERROR(EINVAL),  "invalid param"            },
100     { NV_ENC_ERR_INVALID_CALL,             AVERROR(EINVAL),  "invalid call"             },
101     { NV_ENC_ERR_OUT_OF_MEMORY,            AVERROR(ENOMEM),  "out of memory"            },
102     { NV_ENC_ERR_ENCODER_NOT_INITIALIZED,  AVERROR(EINVAL),  "encoder not initialized"  },
103     { NV_ENC_ERR_UNSUPPORTED_PARAM,        AVERROR(ENOSYS),  "unsupported param"        },
104     { NV_ENC_ERR_LOCK_BUSY,                AVERROR(EAGAIN),  "lock busy"                },
105     { NV_ENC_ERR_NOT_ENOUGH_BUFFER,        AVERROR(ENOBUFS), "not enough buffer"        },
106     { NV_ENC_ERR_INVALID_VERSION,          AVERROR(EINVAL),  "invalid version"          },
107     { NV_ENC_ERR_MAP_FAILED,               AVERROR(EIO),     "map failed"               },
108     { NV_ENC_ERR_NEED_MORE_INPUT,          AVERROR(EAGAIN),  "need more input"          },
109     { NV_ENC_ERR_ENCODER_BUSY,             AVERROR(EAGAIN),  "encoder busy"             },
110     { NV_ENC_ERR_EVENT_NOT_REGISTERD,      AVERROR(EBADF),   "event not registered"     },
111     { NV_ENC_ERR_GENERIC,                  AVERROR_UNKNOWN,  "generic error"            },
112     { NV_ENC_ERR_INCOMPATIBLE_CLIENT_KEY,  AVERROR(EINVAL),  "incompatible client key"  },
113     { NV_ENC_ERR_UNIMPLEMENTED,            AVERROR(ENOSYS),  "unimplemented"            },
114     { NV_ENC_ERR_RESOURCE_REGISTER_FAILED, AVERROR(EIO),     "resource register failed" },
115     { NV_ENC_ERR_RESOURCE_NOT_REGISTERED,  AVERROR(EBADF),   "resource not registered"  },
116     { NV_ENC_ERR_RESOURCE_NOT_MAPPED,      AVERROR(EBADF),   "resource not mapped"      },
117 };
118
119 static int nvenc_map_error(NVENCSTATUS err, const char **desc)
120 {
121     int i;
122     for (i = 0; i < FF_ARRAY_ELEMS(nvenc_errors); i++) {
123         if (nvenc_errors[i].nverr == err) {
124             if (desc)
125                 *desc = nvenc_errors[i].desc;
126             return nvenc_errors[i].averr;
127         }
128     }
129     if (desc)
130         *desc = "unknown error";
131     return AVERROR_UNKNOWN;
132 }
133
134 static int nvenc_print_error(void *log_ctx, NVENCSTATUS err,
135                                      const char *error_string)
136 {
137     const char *desc;
138     int ret;
139     ret = nvenc_map_error(err, &desc);
140     av_log(log_ctx, AV_LOG_ERROR, "%s: %s (%d)\n", error_string, desc, err);
141     return ret;
142 }
143
144 static void timestamp_queue_enqueue(AVFifoBuffer* queue, int64_t timestamp)
145 {
146     av_fifo_generic_write(queue, &timestamp, sizeof(timestamp), NULL);
147 }
148
149 static int64_t timestamp_queue_dequeue(AVFifoBuffer* queue)
150 {
151     int64_t timestamp = AV_NOPTS_VALUE;
152     if (av_fifo_size(queue) > 0)
153         av_fifo_generic_read(queue, &timestamp, sizeof(timestamp), NULL);
154
155     return timestamp;
156 }
157
158 static av_cold int nvenc_load_libraries(AVCodecContext *avctx)
159 {
160     NvencContext *ctx = avctx->priv_data;
161     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
162     PNVENCODEAPICREATEINSTANCE nvenc_create_instance;
163     NVENCSTATUS err;
164
165 #if CONFIG_CUDA
166     dl_fn->cu_init                      = cuInit;
167     dl_fn->cu_device_get_count          = cuDeviceGetCount;
168     dl_fn->cu_device_get                = cuDeviceGet;
169     dl_fn->cu_device_get_name           = cuDeviceGetName;
170     dl_fn->cu_device_compute_capability = cuDeviceComputeCapability;
171     dl_fn->cu_ctx_create                = cuCtxCreate_v2;
172     dl_fn->cu_ctx_pop_current           = cuCtxPopCurrent_v2;
173     dl_fn->cu_ctx_destroy               = cuCtxDestroy_v2;
174 #else
175     LOAD_LIBRARY(dl_fn->cuda, CUDA_LIBNAME);
176
177     LOAD_SYMBOL(dl_fn->cu_init, dl_fn->cuda, "cuInit");
178     LOAD_SYMBOL(dl_fn->cu_device_get_count, dl_fn->cuda, "cuDeviceGetCount");
179     LOAD_SYMBOL(dl_fn->cu_device_get, dl_fn->cuda, "cuDeviceGet");
180     LOAD_SYMBOL(dl_fn->cu_device_get_name, dl_fn->cuda, "cuDeviceGetName");
181     LOAD_SYMBOL(dl_fn->cu_device_compute_capability, dl_fn->cuda,
182                 "cuDeviceComputeCapability");
183     LOAD_SYMBOL(dl_fn->cu_ctx_create, dl_fn->cuda, "cuCtxCreate_v2");
184     LOAD_SYMBOL(dl_fn->cu_ctx_pop_current, dl_fn->cuda, "cuCtxPopCurrent_v2");
185     LOAD_SYMBOL(dl_fn->cu_ctx_destroy, dl_fn->cuda, "cuCtxDestroy_v2");
186 #endif
187
188     LOAD_LIBRARY(dl_fn->nvenc, NVENC_LIBNAME);
189
190     LOAD_SYMBOL(nvenc_create_instance, dl_fn->nvenc,
191                 "NvEncodeAPICreateInstance");
192
193     dl_fn->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER;
194
195     err = nvenc_create_instance(&dl_fn->nvenc_funcs);
196     if (err != NV_ENC_SUCCESS)
197         return nvenc_print_error(avctx, err, "Failed to create nvenc instance");
198
199     av_log(avctx, AV_LOG_VERBOSE, "Nvenc initialized successfully\n");
200
201     return 0;
202 }
203
204 static av_cold int nvenc_open_session(AVCodecContext *avctx)
205 {
206     NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS params = { 0 };
207     NvencContext *ctx = avctx->priv_data;
208     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
209     NVENCSTATUS ret;
210
211     params.version    = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;
212     params.apiVersion = NVENCAPI_VERSION;
213     params.device     = ctx->cu_context;
214     params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
215
216     ret = p_nvenc->nvEncOpenEncodeSessionEx(&params, &ctx->nvencoder);
217     if (ret != NV_ENC_SUCCESS) {
218         ctx->nvencoder = NULL;
219         return nvenc_print_error(avctx, ret, "OpenEncodeSessionEx failed");
220     }
221
222     return 0;
223 }
224
225 static int nvenc_check_codec_support(AVCodecContext *avctx)
226 {
227     NvencContext *ctx = avctx->priv_data;
228     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
229     int i, ret, count = 0;
230     GUID *guids = NULL;
231
232     ret = p_nvenc->nvEncGetEncodeGUIDCount(ctx->nvencoder, &count);
233
234     if (ret != NV_ENC_SUCCESS || !count)
235         return AVERROR(ENOSYS);
236
237     guids = av_malloc(count * sizeof(GUID));
238     if (!guids)
239         return AVERROR(ENOMEM);
240
241     ret = p_nvenc->nvEncGetEncodeGUIDs(ctx->nvencoder, guids, count, &count);
242     if (ret != NV_ENC_SUCCESS) {
243         ret = AVERROR(ENOSYS);
244         goto fail;
245     }
246
247     ret = AVERROR(ENOSYS);
248     for (i = 0; i < count; i++) {
249         if (!memcmp(&guids[i], &ctx->init_encode_params.encodeGUID, sizeof(*guids))) {
250             ret = 0;
251             break;
252         }
253     }
254
255 fail:
256     av_free(guids);
257
258     return ret;
259 }
260
261 static int nvenc_check_cap(AVCodecContext *avctx, NV_ENC_CAPS cap)
262 {
263     NvencContext *ctx = avctx->priv_data;
264     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
265     NV_ENC_CAPS_PARAM params        = { 0 };
266     int ret, val = 0;
267
268     params.version     = NV_ENC_CAPS_PARAM_VER;
269     params.capsToQuery = cap;
270
271     ret = p_nvenc->nvEncGetEncodeCaps(ctx->nvencoder, ctx->init_encode_params.encodeGUID, &params, &val);
272
273     if (ret == NV_ENC_SUCCESS)
274         return val;
275     return 0;
276 }
277
278 static int nvenc_check_capabilities(AVCodecContext *avctx)
279 {
280     NvencContext *ctx = avctx->priv_data;
281     int ret;
282
283     ret = nvenc_check_codec_support(avctx);
284     if (ret < 0) {
285         av_log(avctx, AV_LOG_VERBOSE, "Codec not supported\n");
286         return ret;
287     }
288
289     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_YUV444_ENCODE);
290     if (ctx->data_pix_fmt == AV_PIX_FMT_YUV444P && ret <= 0) {
291         av_log(avctx, AV_LOG_VERBOSE, "YUV444P not supported\n");
292         return AVERROR(ENOSYS);
293     }
294
295     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_LOSSLESS_ENCODE);
296     if (ctx->preset >= PRESET_LOSSLESS_DEFAULT && ret <= 0) {
297         av_log(avctx, AV_LOG_VERBOSE, "Lossless encoding not supported\n");
298         return AVERROR(ENOSYS);
299     }
300
301     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_WIDTH_MAX);
302     if (ret < avctx->width) {
303         av_log(avctx, AV_LOG_VERBOSE, "Width %d exceeds %d\n",
304                avctx->width, ret);
305         return AVERROR(ENOSYS);
306     }
307
308     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_HEIGHT_MAX);
309     if (ret < avctx->height) {
310         av_log(avctx, AV_LOG_VERBOSE, "Height %d exceeds %d\n",
311                avctx->height, ret);
312         return AVERROR(ENOSYS);
313     }
314
315     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_NUM_MAX_BFRAMES);
316     if (ret < avctx->max_b_frames) {
317         av_log(avctx, AV_LOG_VERBOSE, "Max b-frames %d exceed %d\n",
318                avctx->max_b_frames, ret);
319
320         return AVERROR(ENOSYS);
321     }
322
323     return 0;
324 }
325
326 static av_cold int nvenc_check_device(AVCodecContext *avctx, int idx)
327 {
328     NvencContext *ctx = avctx->priv_data;
329     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
330     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
331     char name[128] = { 0};
332     int major, minor, ret;
333     CUresult cu_res;
334     CUdevice cu_device;
335     CUcontext dummy;
336     int loglevel = AV_LOG_VERBOSE;
337
338     if (ctx->device == LIST_DEVICES)
339         loglevel = AV_LOG_INFO;
340
341     cu_res = dl_fn->cu_device_get(&cu_device, idx);
342     if (cu_res != CUDA_SUCCESS) {
343         av_log(avctx, AV_LOG_ERROR,
344                "Cannot access the CUDA device %d\n",
345                idx);
346         return -1;
347     }
348
349     cu_res = dl_fn->cu_device_get_name(name, sizeof(name), cu_device);
350     if (cu_res != CUDA_SUCCESS)
351         return -1;
352
353     cu_res = dl_fn->cu_device_compute_capability(&major, &minor, cu_device);
354     if (cu_res != CUDA_SUCCESS)
355         return -1;
356
357     av_log(avctx, loglevel, "[ GPU #%d - < %s > has Compute SM %d.%d ]\n", idx, name, major, minor);
358     if (((major << 4) | minor) < NVENC_CAP) {
359         av_log(avctx, loglevel, "does not support NVENC\n");
360         goto fail;
361     }
362
363     cu_res = dl_fn->cu_ctx_create(&ctx->cu_context_internal, 0, cu_device);
364     if (cu_res != CUDA_SUCCESS) {
365         av_log(avctx, AV_LOG_FATAL, "Failed creating CUDA context for NVENC: 0x%x\n", (int)cu_res);
366         goto fail;
367     }
368
369     ctx->cu_context = ctx->cu_context_internal;
370
371     cu_res = dl_fn->cu_ctx_pop_current(&dummy);
372     if (cu_res != CUDA_SUCCESS) {
373         av_log(avctx, AV_LOG_FATAL, "Failed popping CUDA context: 0x%x\n", (int)cu_res);
374         goto fail2;
375     }
376
377     if ((ret = nvenc_open_session(avctx)) < 0)
378         goto fail2;
379
380     if ((ret = nvenc_check_capabilities(avctx)) < 0)
381         goto fail3;
382
383     av_log(avctx, loglevel, "supports NVENC\n");
384
385     dl_fn->nvenc_device_count++;
386
387     if (ctx->device == dl_fn->nvenc_device_count - 1 || ctx->device == ANY_DEVICE)
388         return 0;
389
390 fail3:
391     p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);
392     ctx->nvencoder = NULL;
393
394 fail2:
395     dl_fn->cu_ctx_destroy(ctx->cu_context_internal);
396     ctx->cu_context_internal = NULL;
397
398 fail:
399     return AVERROR(ENOSYS);
400 }
401
402 static av_cold int nvenc_setup_device(AVCodecContext *avctx)
403 {
404     NvencContext *ctx = avctx->priv_data;
405     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
406
407     switch (avctx->codec->id) {
408     case AV_CODEC_ID_H264:
409         ctx->init_encode_params.encodeGUID = NV_ENC_CODEC_H264_GUID;
410         break;
411     case AV_CODEC_ID_HEVC:
412         ctx->init_encode_params.encodeGUID = NV_ENC_CODEC_HEVC_GUID;
413         break;
414     default:
415         return AVERROR_BUG;
416     }
417
418     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
419 #if CONFIG_CUDA
420         AVHWFramesContext   *frames_ctx;
421         AVCUDADeviceContext *device_hwctx;
422         int ret;
423
424         if (!avctx->hw_frames_ctx)
425             return AVERROR(EINVAL);
426
427         frames_ctx   = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
428         device_hwctx = frames_ctx->device_ctx->hwctx;
429
430         ctx->cu_context = device_hwctx->cuda_ctx;
431
432         ret = nvenc_open_session(avctx);
433         if (ret < 0)
434             return ret;
435
436         ret = nvenc_check_capabilities(avctx);
437         if (ret < 0) {
438             av_log(avctx, AV_LOG_FATAL, "Provided device doesn't support required NVENC features\n");
439             return ret;
440         }
441 #else
442         return AVERROR_BUG;
443 #endif
444     } else {
445         int i, nb_devices = 0;
446
447         if ((dl_fn->cu_init(0)) != CUDA_SUCCESS) {
448             av_log(avctx, AV_LOG_ERROR,
449                    "Cannot init CUDA\n");
450             return AVERROR_UNKNOWN;
451         }
452
453         if ((dl_fn->cu_device_get_count(&nb_devices)) != CUDA_SUCCESS) {
454             av_log(avctx, AV_LOG_ERROR,
455                    "Cannot enumerate the CUDA devices\n");
456             return AVERROR_UNKNOWN;
457         }
458
459         if (!nb_devices) {
460             av_log(avctx, AV_LOG_FATAL, "No CUDA capable devices found\n");
461                 return AVERROR_EXTERNAL;
462         }
463
464         av_log(avctx, AV_LOG_VERBOSE, "%d CUDA capable devices found\n", nb_devices);
465
466         dl_fn->nvenc_device_count = 0;
467         for (i = 0; i < nb_devices; ++i) {
468             if ((nvenc_check_device(avctx, i)) >= 0 && ctx->device != LIST_DEVICES)
469                 return 0;
470         }
471
472         if (ctx->device == LIST_DEVICES)
473             return AVERROR_EXIT;
474
475         if (!dl_fn->nvenc_device_count) {
476             av_log(avctx, AV_LOG_FATAL, "No NVENC capable devices found\n");
477             return AVERROR_EXTERNAL;
478         }
479
480         av_log(avctx, AV_LOG_FATAL, "Requested GPU %d, but only %d GPUs are available!\n", ctx->device, dl_fn->nvenc_device_count);
481         return AVERROR(EINVAL);
482     }
483
484     return 0;
485 }
486
487 typedef struct GUIDTuple {
488     const GUID guid;
489     int flags;
490 } GUIDTuple;
491
492 static void nvenc_map_preset(NvencContext *ctx)
493 {
494     GUIDTuple presets[] = {
495         { NV_ENC_PRESET_DEFAULT_GUID },
496         { NV_ENC_PRESET_HQ_GUID,                  NVENC_TWO_PASSES }, /* slow */
497         { NV_ENC_PRESET_HQ_GUID,                  NVENC_ONE_PASS }, /* medium */
498         { NV_ENC_PRESET_HP_GUID,                  NVENC_ONE_PASS }, /* fast */
499         { NV_ENC_PRESET_HP_GUID },
500         { NV_ENC_PRESET_HQ_GUID },
501         { NV_ENC_PRESET_BD_GUID },
502         { NV_ENC_PRESET_LOW_LATENCY_DEFAULT_GUID, NVENC_LOWLATENCY },
503         { NV_ENC_PRESET_LOW_LATENCY_HQ_GUID,      NVENC_LOWLATENCY },
504         { NV_ENC_PRESET_LOW_LATENCY_HP_GUID,      NVENC_LOWLATENCY },
505         { NV_ENC_PRESET_LOSSLESS_DEFAULT_GUID,    NVENC_LOSSLESS },
506         { NV_ENC_PRESET_LOSSLESS_HP_GUID,         NVENC_LOSSLESS },
507     };
508
509     GUIDTuple *t = &presets[ctx->preset];
510
511     ctx->init_encode_params.presetGUID = t->guid;
512     ctx->flags = t->flags;
513 }
514
515 static av_cold void set_constqp(AVCodecContext *avctx)
516 {
517     NvencContext *ctx = avctx->priv_data;
518     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
519
520     rc->rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
521     rc->constQP.qpInterB = avctx->global_quality;
522     rc->constQP.qpInterP = avctx->global_quality;
523     rc->constQP.qpIntra = avctx->global_quality;
524
525     avctx->qmin = -1;
526     avctx->qmax = -1;
527 }
528
529 static av_cold void set_vbr(AVCodecContext *avctx)
530 {
531     NvencContext *ctx = avctx->priv_data;
532     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
533     int qp_inter_p;
534
535     if (avctx->qmin >= 0 && avctx->qmax >= 0) {
536         rc->enableMinQP = 1;
537         rc->enableMaxQP = 1;
538
539         rc->minQP.qpInterB = avctx->qmin;
540         rc->minQP.qpInterP = avctx->qmin;
541         rc->minQP.qpIntra = avctx->qmin;
542
543         rc->maxQP.qpInterB = avctx->qmax;
544         rc->maxQP.qpInterP = avctx->qmax;
545         rc->maxQP.qpIntra = avctx->qmax;
546
547         qp_inter_p = (avctx->qmax + 3 * avctx->qmin) / 4; // biased towards Qmin
548     } else {
549         qp_inter_p = 26; // default to 26
550     }
551
552     rc->enableInitialRCQP = 1;
553     rc->initialRCQP.qpInterP  = qp_inter_p;
554
555     if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
556         rc->initialRCQP.qpIntra = av_clip(
557             qp_inter_p * fabs(avctx->i_quant_factor) + avctx->i_quant_offset, 0, 51);
558         rc->initialRCQP.qpInterB = av_clip(
559             qp_inter_p * fabs(avctx->b_quant_factor) + avctx->b_quant_offset, 0, 51);
560     } else {
561         rc->initialRCQP.qpIntra = qp_inter_p;
562         rc->initialRCQP.qpInterB = qp_inter_p;
563     }
564 }
565
566 static av_cold void set_lossless(AVCodecContext *avctx)
567 {
568     NvencContext *ctx = avctx->priv_data;
569     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
570
571     rc->rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
572     rc->constQP.qpInterB = 0;
573     rc->constQP.qpInterP = 0;
574     rc->constQP.qpIntra = 0;
575
576     avctx->qmin = -1;
577     avctx->qmax = -1;
578 }
579
580 static void nvenc_override_rate_control(AVCodecContext *avctx)
581 {
582     NvencContext *ctx    = avctx->priv_data;
583     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
584
585     switch (ctx->rc) {
586     case NV_ENC_PARAMS_RC_CONSTQP:
587         if (avctx->global_quality <= 0) {
588             av_log(avctx, AV_LOG_WARNING,
589                    "The constant quality rate-control requires "
590                    "the 'global_quality' option set.\n");
591             return;
592         }
593         set_constqp(avctx);
594         return;
595     case NV_ENC_PARAMS_RC_2_PASS_VBR:
596     case NV_ENC_PARAMS_RC_VBR:
597         if (avctx->qmin < 0 && avctx->qmax < 0) {
598             av_log(avctx, AV_LOG_WARNING,
599                    "The variable bitrate rate-control requires "
600                    "the 'qmin' and/or 'qmax' option set.\n");
601             set_vbr(avctx);
602             return;
603         }
604     case NV_ENC_PARAMS_RC_VBR_MINQP:
605         if (avctx->qmin < 0) {
606             av_log(avctx, AV_LOG_WARNING,
607                    "The variable bitrate rate-control requires "
608                    "the 'qmin' option set.\n");
609             set_vbr(avctx);
610             return;
611         }
612         set_vbr(avctx);
613         break;
614     case NV_ENC_PARAMS_RC_CBR:
615         break;
616     case NV_ENC_PARAMS_RC_2_PASS_QUALITY:
617     case NV_ENC_PARAMS_RC_2_PASS_FRAMESIZE_CAP:
618         if (!(ctx->flags & NVENC_LOWLATENCY)) {
619             av_log(avctx, AV_LOG_WARNING,
620                    "The multipass rate-control requires "
621                    "a low-latency preset.\n");
622             return;
623         }
624     }
625
626     rc->rateControlMode = ctx->rc;
627 }
628
629 static av_cold void nvenc_setup_rate_control(AVCodecContext *avctx)
630 {
631     NvencContext *ctx = avctx->priv_data;
632
633     if (avctx->bit_rate > 0) {
634         ctx->encode_config.rcParams.averageBitRate = avctx->bit_rate;
635     } else if (ctx->encode_config.rcParams.averageBitRate > 0) {
636         ctx->encode_config.rcParams.maxBitRate = ctx->encode_config.rcParams.averageBitRate;
637     }
638
639     if (avctx->rc_max_rate > 0)
640         ctx->encode_config.rcParams.maxBitRate = avctx->rc_max_rate;
641
642     if (ctx->rc < 0) {
643         if (ctx->flags & NVENC_ONE_PASS)
644             ctx->twopass = 0;
645         if (ctx->flags & NVENC_TWO_PASSES)
646             ctx->twopass = 1;
647
648         if (ctx->twopass < 0)
649             ctx->twopass = (ctx->flags & NVENC_LOWLATENCY) != 0;
650
651         if (ctx->cbr) {
652             if (ctx->twopass) {
653                 ctx->rc = NV_ENC_PARAMS_RC_2_PASS_QUALITY;
654             } else {
655                 ctx->rc = NV_ENC_PARAMS_RC_CBR;
656             }
657         } else if (avctx->global_quality > 0) {
658             ctx->rc = NV_ENC_PARAMS_RC_CONSTQP;
659         } else if (ctx->twopass) {
660             ctx->rc = NV_ENC_PARAMS_RC_2_PASS_VBR;
661         } else if (avctx->qmin >= 0 && avctx->qmax >= 0) {
662             ctx->rc = NV_ENC_PARAMS_RC_VBR_MINQP;
663         }
664     }
665
666     if (ctx->flags & NVENC_LOSSLESS) {
667         set_lossless(avctx);
668     } else if (ctx->rc > 0) {
669         nvenc_override_rate_control(avctx);
670     } else {
671         ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR;
672         set_vbr(avctx);
673     }
674
675     if (avctx->rc_buffer_size > 0) {
676         ctx->encode_config.rcParams.vbvBufferSize = avctx->rc_buffer_size;
677     } else if (ctx->encode_config.rcParams.averageBitRate > 0) {
678         ctx->encode_config.rcParams.vbvBufferSize = 2 * ctx->encode_config.rcParams.averageBitRate;
679     }
680 }
681
682 static av_cold int nvenc_setup_h264_config(AVCodecContext *avctx)
683 {
684     NvencContext *ctx                      = avctx->priv_data;
685     NV_ENC_CONFIG *cc                      = &ctx->encode_config;
686     NV_ENC_CONFIG_H264 *h264               = &cc->encodeCodecConfig.h264Config;
687     NV_ENC_CONFIG_H264_VUI_PARAMETERS *vui = &h264->h264VUIParameters;
688
689     vui->colourMatrix = avctx->colorspace;
690     vui->colourPrimaries = avctx->color_primaries;
691     vui->transferCharacteristics = avctx->color_trc;
692     vui->videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG
693         || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ420P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ422P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ444P);
694
695     vui->colourDescriptionPresentFlag =
696         (avctx->colorspace != 2 || avctx->color_primaries != 2 || avctx->color_trc != 2);
697
698     vui->videoSignalTypePresentFlag =
699         (vui->colourDescriptionPresentFlag
700         || vui->videoFormat != 5
701         || vui->videoFullRangeFlag != 0);
702
703     h264->sliceMode = 3;
704     h264->sliceModeData = 1;
705
706     h264->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
707     h264->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
708     h264->outputAUD     = 1;
709
710     if (avctx->refs >= 0) {
711         /* 0 means "let the hardware decide" */
712         h264->maxNumRefFrames = avctx->refs;
713     }
714     if (avctx->gop_size >= 0) {
715         h264->idrPeriod = cc->gopLength;
716     }
717
718     if (IS_CBR(cc->rcParams.rateControlMode)) {
719         h264->outputBufferingPeriodSEI = 1;
720         h264->outputPictureTimingSEI   = 1;
721     }
722
723     if (cc->rcParams.rateControlMode == NV_ENC_PARAMS_RC_2_PASS_QUALITY ||
724         cc->rcParams.rateControlMode == NV_ENC_PARAMS_RC_2_PASS_FRAMESIZE_CAP ||
725         cc->rcParams.rateControlMode == NV_ENC_PARAMS_RC_2_PASS_VBR) {
726         h264->adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE;
727         h264->fmoMode = NV_ENC_H264_FMO_DISABLE;
728     }
729
730     if (ctx->flags & NVENC_LOSSLESS) {
731         h264->qpPrimeYZeroTransformBypassFlag = 1;
732     } else {
733         switch(ctx->profile) {
734         case NV_ENC_H264_PROFILE_BASELINE:
735             cc->profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;
736             avctx->profile = FF_PROFILE_H264_BASELINE;
737             break;
738         case NV_ENC_H264_PROFILE_MAIN:
739             cc->profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;
740             avctx->profile = FF_PROFILE_H264_MAIN;
741             break;
742         case NV_ENC_H264_PROFILE_HIGH:
743             cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
744             avctx->profile = FF_PROFILE_H264_HIGH;
745             break;
746         case NV_ENC_H264_PROFILE_HIGH_444P:
747             cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
748             avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;
749             break;
750         }
751     }
752
753     // force setting profile as high444p if input is AV_PIX_FMT_YUV444P
754     if (ctx->data_pix_fmt == AV_PIX_FMT_YUV444P) {
755         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
756         avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;
757     }
758
759     h264->chromaFormatIDC = avctx->profile == FF_PROFILE_H264_HIGH_444_PREDICTIVE ? 3 : 1;
760
761     h264->level = ctx->level;
762
763     return 0;
764 }
765
766 static av_cold int nvenc_setup_hevc_config(AVCodecContext *avctx)
767 {
768     NvencContext *ctx                      = avctx->priv_data;
769     NV_ENC_CONFIG *cc                      = &ctx->encode_config;
770     NV_ENC_CONFIG_HEVC *hevc               = &cc->encodeCodecConfig.hevcConfig;
771     NV_ENC_CONFIG_HEVC_VUI_PARAMETERS *vui = &hevc->hevcVUIParameters;
772
773     vui->colourMatrix = avctx->colorspace;
774     vui->colourPrimaries = avctx->color_primaries;
775     vui->transferCharacteristics = avctx->color_trc;
776     vui->videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG
777         || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ420P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ422P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ444P);
778
779     vui->colourDescriptionPresentFlag =
780         (avctx->colorspace != 2 || avctx->color_primaries != 2 || avctx->color_trc != 2);
781
782     vui->videoSignalTypePresentFlag =
783         (vui->colourDescriptionPresentFlag
784         || vui->videoFormat != 5
785         || vui->videoFullRangeFlag != 0);
786
787     hevc->sliceMode = 3;
788     hevc->sliceModeData = 1;
789
790     hevc->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
791     hevc->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
792     hevc->outputAUD     = 1;
793
794     if (avctx->refs >= 0) {
795         /* 0 means "let the hardware decide" */
796         hevc->maxNumRefFramesInDPB = avctx->refs;
797     }
798     if (avctx->gop_size >= 0) {
799         hevc->idrPeriod = cc->gopLength;
800     }
801
802     if (IS_CBR(cc->rcParams.rateControlMode)) {
803         hevc->outputBufferingPeriodSEI = 1;
804         hevc->outputPictureTimingSEI   = 1;
805     }
806
807     /* No other profile is supported in the current SDK version 5 */
808     cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID;
809     avctx->profile = FF_PROFILE_HEVC_MAIN;
810
811     hevc->level = ctx->level;
812
813     hevc->tier = ctx->tier;
814
815     return 0;
816 }
817
818 static av_cold int nvenc_setup_codec_config(AVCodecContext *avctx)
819 {
820     switch (avctx->codec->id) {
821     case AV_CODEC_ID_H264:
822         return nvenc_setup_h264_config(avctx);
823     case AV_CODEC_ID_HEVC:
824         return nvenc_setup_hevc_config(avctx);
825     /* Earlier switch/case will return if unknown codec is passed. */
826     }
827
828     return 0;
829 }
830
831 static av_cold int nvenc_setup_encoder(AVCodecContext *avctx)
832 {
833     NvencContext *ctx = avctx->priv_data;
834     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
835     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
836
837     NV_ENC_PRESET_CONFIG preset_config = { 0 };
838     NVENCSTATUS nv_status = NV_ENC_SUCCESS;
839     AVCPBProperties *cpb_props;
840     int res = 0;
841     int dw, dh;
842
843     ctx->last_dts = AV_NOPTS_VALUE;
844
845     ctx->encode_config.version = NV_ENC_CONFIG_VER;
846     ctx->init_encode_params.version = NV_ENC_INITIALIZE_PARAMS_VER;
847
848     ctx->init_encode_params.encodeHeight = avctx->height;
849     ctx->init_encode_params.encodeWidth = avctx->width;
850
851     ctx->init_encode_params.encodeConfig = &ctx->encode_config;
852
853     nvenc_map_preset(ctx);
854
855     preset_config.version = NV_ENC_PRESET_CONFIG_VER;
856     preset_config.presetCfg.version = NV_ENC_CONFIG_VER;
857
858     nv_status = p_nvenc->nvEncGetEncodePresetConfig(ctx->nvencoder,
859                                                     ctx->init_encode_params.encodeGUID,
860                                                     ctx->init_encode_params.presetGUID,
861                                                     &preset_config);
862     if (nv_status != NV_ENC_SUCCESS)
863         return nvenc_print_error(avctx, nv_status, "Cannot get the preset configuration");
864
865     memcpy(&ctx->encode_config, &preset_config.presetCfg, sizeof(ctx->encode_config));
866
867     ctx->encode_config.version = NV_ENC_CONFIG_VER;
868
869     if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den &&
870         (avctx->sample_aspect_ratio.num != 1 || avctx->sample_aspect_ratio.num != 1)) {
871         av_reduce(&dw, &dh,
872                   avctx->width * avctx->sample_aspect_ratio.num,
873                   avctx->height * avctx->sample_aspect_ratio.den,
874                   1024 * 1024);
875         ctx->init_encode_params.darHeight = dh;
876         ctx->init_encode_params.darWidth = dw;
877     } else {
878         ctx->init_encode_params.darHeight = avctx->height;
879         ctx->init_encode_params.darWidth = avctx->width;
880     }
881
882     // De-compensate for hardware, dubiously, trying to compensate for
883     // playback at 704 pixel width.
884     if (avctx->width == 720 &&
885         (avctx->height == 480 || avctx->height == 576)) {
886         av_reduce(&dw, &dh,
887                   ctx->init_encode_params.darWidth * 44,
888                   ctx->init_encode_params.darHeight * 45,
889                   1024 * 1024);
890         ctx->init_encode_params.darHeight = dh;
891         ctx->init_encode_params.darWidth = dw;
892     }
893
894     ctx->init_encode_params.frameRateNum = avctx->time_base.den;
895     ctx->init_encode_params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
896
897     ctx->init_encode_params.enableEncodeAsync = 0;
898     ctx->init_encode_params.enablePTD = 1;
899
900     if (avctx->gop_size > 0) {
901         if (avctx->max_b_frames >= 0) {
902             /* 0 is intra-only, 1 is I/P only, 2 is one B Frame, 3 two B frames, and so on. */
903             ctx->encode_config.frameIntervalP = avctx->max_b_frames + 1;
904         }
905
906         ctx->encode_config.gopLength = avctx->gop_size;
907     } else if (avctx->gop_size == 0) {
908         ctx->encode_config.frameIntervalP = 0;
909         ctx->encode_config.gopLength = 1;
910     }
911
912     /* when there're b frames, set dts offset */
913     if (ctx->encode_config.frameIntervalP >= 2)
914         ctx->last_dts = -2;
915
916     nvenc_setup_rate_control(avctx);
917
918     if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
919         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
920     } else {
921         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;
922     }
923
924     res = nvenc_setup_codec_config(avctx);
925     if (res)
926         return res;
927
928     nv_status = p_nvenc->nvEncInitializeEncoder(ctx->nvencoder, &ctx->init_encode_params);
929     if (nv_status != NV_ENC_SUCCESS) {
930         return nvenc_print_error(avctx, nv_status, "InitializeEncoder failed");
931     }
932
933     if (ctx->encode_config.frameIntervalP > 1)
934         avctx->has_b_frames = 2;
935
936     if (ctx->encode_config.rcParams.averageBitRate > 0)
937         avctx->bit_rate = ctx->encode_config.rcParams.averageBitRate;
938
939     cpb_props = ff_add_cpb_side_data(avctx);
940     if (!cpb_props)
941         return AVERROR(ENOMEM);
942     cpb_props->max_bitrate = ctx->encode_config.rcParams.maxBitRate;
943     cpb_props->avg_bitrate = avctx->bit_rate;
944     cpb_props->buffer_size = ctx->encode_config.rcParams.vbvBufferSize;
945
946     return 0;
947 }
948
949 static av_cold int nvenc_alloc_surface(AVCodecContext *avctx, int idx)
950 {
951     NvencContext *ctx = avctx->priv_data;
952     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
953     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
954
955     NVENCSTATUS nv_status;
956     NV_ENC_CREATE_BITSTREAM_BUFFER allocOut = { 0 };
957     allocOut.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
958
959     switch (ctx->data_pix_fmt) {
960     case AV_PIX_FMT_YUV420P:
961         ctx->surfaces[idx].format = NV_ENC_BUFFER_FORMAT_YV12_PL;
962         break;
963
964     case AV_PIX_FMT_NV12:
965         ctx->surfaces[idx].format = NV_ENC_BUFFER_FORMAT_NV12_PL;
966         break;
967
968     case AV_PIX_FMT_YUV444P:
969         ctx->surfaces[idx].format = NV_ENC_BUFFER_FORMAT_YUV444_PL;
970         break;
971
972     default:
973         av_log(avctx, AV_LOG_FATAL, "Invalid input pixel format\n");
974         return AVERROR(EINVAL);
975     }
976
977     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
978         ctx->surfaces[idx].in_ref = av_frame_alloc();
979         if (!ctx->surfaces[idx].in_ref)
980             return AVERROR(ENOMEM);
981     } else {
982         NV_ENC_CREATE_INPUT_BUFFER allocSurf = { 0 };
983         allocSurf.version = NV_ENC_CREATE_INPUT_BUFFER_VER;
984         allocSurf.width = (avctx->width + 31) & ~31;
985         allocSurf.height = (avctx->height + 31) & ~31;
986         allocSurf.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;
987         allocSurf.bufferFmt = ctx->surfaces[idx].format;
988
989         nv_status = p_nvenc->nvEncCreateInputBuffer(ctx->nvencoder, &allocSurf);
990         if (nv_status != NV_ENC_SUCCESS) {
991             return nvenc_print_error(avctx, nv_status, "CreateInputBuffer failed");
992         }
993
994         ctx->surfaces[idx].input_surface = allocSurf.inputBuffer;
995         ctx->surfaces[idx].width = allocSurf.width;
996         ctx->surfaces[idx].height = allocSurf.height;
997     }
998
999     ctx->surfaces[idx].lockCount = 0;
1000
1001     /* 1MB is large enough to hold most output frames. NVENC increases this automaticaly if it's not enough. */
1002     allocOut.size = 1024 * 1024;
1003
1004     allocOut.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;
1005
1006     nv_status = p_nvenc->nvEncCreateBitstreamBuffer(ctx->nvencoder, &allocOut);
1007     if (nv_status != NV_ENC_SUCCESS) {
1008         int err = nvenc_print_error(avctx, nv_status, "CreateBitstreamBuffer failed");
1009         if (avctx->pix_fmt != AV_PIX_FMT_CUDA)
1010             p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[idx].input_surface);
1011         av_frame_free(&ctx->surfaces[idx].in_ref);
1012         return err;
1013     }
1014
1015     ctx->surfaces[idx].output_surface = allocOut.bitstreamBuffer;
1016     ctx->surfaces[idx].size = allocOut.size;
1017
1018     return 0;
1019 }
1020
1021 static av_cold int nvenc_setup_surfaces(AVCodecContext *avctx)
1022 {
1023     NvencContext *ctx = avctx->priv_data;
1024     int i, res;
1025     int num_mbs = ((avctx->width + 15) >> 4) * ((avctx->height + 15) >> 4);
1026     ctx->nb_surfaces = FFMAX((num_mbs >= 8160) ? 32 : 48,
1027                              ctx->nb_surfaces);
1028     ctx->async_depth = FFMIN(ctx->async_depth, ctx->nb_surfaces - 1);
1029
1030
1031     ctx->surfaces = av_mallocz_array(ctx->nb_surfaces, sizeof(*ctx->surfaces));
1032     if (!ctx->surfaces)
1033         return AVERROR(ENOMEM);
1034
1035     ctx->timestamp_list = av_fifo_alloc(ctx->nb_surfaces * sizeof(int64_t));
1036     if (!ctx->timestamp_list)
1037         return AVERROR(ENOMEM);
1038     ctx->output_surface_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(NvencSurface*));
1039     if (!ctx->output_surface_queue)
1040         return AVERROR(ENOMEM);
1041     ctx->output_surface_ready_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(NvencSurface*));
1042     if (!ctx->output_surface_ready_queue)
1043         return AVERROR(ENOMEM);
1044
1045     for (i = 0; i < ctx->nb_surfaces; i++) {
1046         if ((res = nvenc_alloc_surface(avctx, i)) < 0)
1047             return res;
1048     }
1049
1050     return 0;
1051 }
1052
1053 static av_cold int nvenc_setup_extradata(AVCodecContext *avctx)
1054 {
1055     NvencContext *ctx = avctx->priv_data;
1056     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1057     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1058
1059     NVENCSTATUS nv_status;
1060     uint32_t outSize = 0;
1061     char tmpHeader[256];
1062     NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };
1063     payload.version = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;
1064
1065     payload.spsppsBuffer = tmpHeader;
1066     payload.inBufferSize = sizeof(tmpHeader);
1067     payload.outSPSPPSPayloadSize = &outSize;
1068
1069     nv_status = p_nvenc->nvEncGetSequenceParams(ctx->nvencoder, &payload);
1070     if (nv_status != NV_ENC_SUCCESS) {
1071         return nvenc_print_error(avctx, nv_status, "GetSequenceParams failed");
1072     }
1073
1074     avctx->extradata_size = outSize;
1075     avctx->extradata = av_mallocz(outSize + AV_INPUT_BUFFER_PADDING_SIZE);
1076
1077     if (!avctx->extradata) {
1078         return AVERROR(ENOMEM);
1079     }
1080
1081     memcpy(avctx->extradata, tmpHeader, outSize);
1082
1083     return 0;
1084 }
1085
1086 av_cold int ff_nvenc_encode_close(AVCodecContext *avctx)
1087 {
1088     NvencContext *ctx               = avctx->priv_data;
1089     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1090     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1091     int i;
1092
1093     /* the encoder has to be flushed before it can be closed */
1094     if (ctx->nvencoder) {
1095         NV_ENC_PIC_PARAMS params        = { .version        = NV_ENC_PIC_PARAMS_VER,
1096                                             .encodePicFlags = NV_ENC_PIC_FLAG_EOS };
1097
1098         p_nvenc->nvEncEncodePicture(ctx->nvencoder, &params);
1099     }
1100
1101     av_fifo_freep(&ctx->timestamp_list);
1102     av_fifo_freep(&ctx->output_surface_ready_queue);
1103     av_fifo_freep(&ctx->output_surface_queue);
1104
1105     if (ctx->surfaces && avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1106         for (i = 0; i < ctx->nb_surfaces; ++i) {
1107             if (ctx->surfaces[i].input_surface) {
1108                  p_nvenc->nvEncUnmapInputResource(ctx->nvencoder, ctx->surfaces[i].in_map.mappedResource);
1109             }
1110         }
1111         for (i = 0; i < ctx->nb_registered_frames; i++) {
1112             if (ctx->registered_frames[i].regptr)
1113                 p_nvenc->nvEncUnregisterResource(ctx->nvencoder, ctx->registered_frames[i].regptr);
1114         }
1115         ctx->nb_registered_frames = 0;
1116     }
1117
1118     if (ctx->surfaces) {
1119         for (i = 0; i < ctx->nb_surfaces; ++i) {
1120             if (avctx->pix_fmt != AV_PIX_FMT_CUDA)
1121                 p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[i].input_surface);
1122             av_frame_free(&ctx->surfaces[i].in_ref);
1123             p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->surfaces[i].output_surface);
1124         }
1125     }
1126     av_freep(&ctx->surfaces);
1127     ctx->nb_surfaces = 0;
1128
1129     if (ctx->nvencoder)
1130         p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);
1131     ctx->nvencoder = NULL;
1132
1133     if (ctx->cu_context_internal)
1134         dl_fn->cu_ctx_destroy(ctx->cu_context_internal);
1135     ctx->cu_context = ctx->cu_context_internal = NULL;
1136
1137     if (dl_fn->nvenc)
1138         dlclose(dl_fn->nvenc);
1139     dl_fn->nvenc = NULL;
1140
1141     dl_fn->nvenc_device_count = 0;
1142
1143 #if !CONFIG_CUDA
1144     if (dl_fn->cuda)
1145         dlclose(dl_fn->cuda);
1146     dl_fn->cuda = NULL;
1147 #endif
1148
1149     dl_fn->cu_init = NULL;
1150     dl_fn->cu_device_get_count = NULL;
1151     dl_fn->cu_device_get = NULL;
1152     dl_fn->cu_device_get_name = NULL;
1153     dl_fn->cu_device_compute_capability = NULL;
1154     dl_fn->cu_ctx_create = NULL;
1155     dl_fn->cu_ctx_pop_current = NULL;
1156     dl_fn->cu_ctx_destroy = NULL;
1157
1158     av_log(avctx, AV_LOG_VERBOSE, "Nvenc unloaded\n");
1159
1160     return 0;
1161 }
1162
1163 av_cold int ff_nvenc_encode_init(AVCodecContext *avctx)
1164 {
1165     NvencContext *ctx = avctx->priv_data;
1166     int ret;
1167
1168     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1169         AVHWFramesContext *frames_ctx;
1170         if (!avctx->hw_frames_ctx) {
1171             av_log(avctx, AV_LOG_ERROR,
1172                    "hw_frames_ctx must be set when using GPU frames as input\n");
1173             return AVERROR(EINVAL);
1174         }
1175         frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
1176         ctx->data_pix_fmt = frames_ctx->sw_format;
1177     } else {
1178         ctx->data_pix_fmt = avctx->pix_fmt;
1179     }
1180
1181     if ((ret = nvenc_load_libraries(avctx)) < 0)
1182         return ret;
1183
1184     if ((ret = nvenc_setup_device(avctx)) < 0)
1185         return ret;
1186
1187     if ((ret = nvenc_setup_encoder(avctx)) < 0)
1188         return ret;
1189
1190     if ((ret = nvenc_setup_surfaces(avctx)) < 0)
1191         return ret;
1192
1193     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
1194         if ((ret = nvenc_setup_extradata(avctx)) < 0)
1195             return ret;
1196     }
1197
1198     return 0;
1199 }
1200
1201 static NvencSurface *get_free_frame(NvencContext *ctx)
1202 {
1203     int i;
1204
1205     for (i = 0; i < ctx->nb_surfaces; ++i) {
1206         if (!ctx->surfaces[i].lockCount) {
1207             ctx->surfaces[i].lockCount = 1;
1208             return &ctx->surfaces[i];
1209         }
1210     }
1211
1212     return NULL;
1213 }
1214
1215 static int nvenc_copy_frame(AVCodecContext *avctx, NvencSurface *inSurf,
1216             NV_ENC_LOCK_INPUT_BUFFER *lockBufferParams, const AVFrame *frame)
1217 {
1218     uint8_t *buf = lockBufferParams->bufferDataPtr;
1219     int off = inSurf->height * lockBufferParams->pitch;
1220
1221     if (frame->format == AV_PIX_FMT_YUV420P) {
1222         av_image_copy_plane(buf, lockBufferParams->pitch,
1223             frame->data[0], frame->linesize[0],
1224             avctx->width, avctx->height);
1225
1226         buf += off;
1227
1228         av_image_copy_plane(buf, lockBufferParams->pitch >> 1,
1229             frame->data[2], frame->linesize[2],
1230             avctx->width >> 1, avctx->height >> 1);
1231
1232         buf += off >> 2;
1233
1234         av_image_copy_plane(buf, lockBufferParams->pitch >> 1,
1235             frame->data[1], frame->linesize[1],
1236             avctx->width >> 1, avctx->height >> 1);
1237     } else if (frame->format == AV_PIX_FMT_NV12) {
1238         av_image_copy_plane(buf, lockBufferParams->pitch,
1239             frame->data[0], frame->linesize[0],
1240             avctx->width, avctx->height);
1241
1242         buf += off;
1243
1244         av_image_copy_plane(buf, lockBufferParams->pitch,
1245             frame->data[1], frame->linesize[1],
1246             avctx->width, avctx->height >> 1);
1247     } else if (frame->format == AV_PIX_FMT_YUV444P) {
1248         av_image_copy_plane(buf, lockBufferParams->pitch,
1249             frame->data[0], frame->linesize[0],
1250             avctx->width, avctx->height);
1251
1252         buf += off;
1253
1254         av_image_copy_plane(buf, lockBufferParams->pitch,
1255             frame->data[1], frame->linesize[1],
1256             avctx->width, avctx->height);
1257
1258         buf += off;
1259
1260         av_image_copy_plane(buf, lockBufferParams->pitch,
1261             frame->data[2], frame->linesize[2],
1262             avctx->width, avctx->height);
1263     } else {
1264         av_log(avctx, AV_LOG_FATAL, "Invalid pixel format!\n");
1265         return AVERROR(EINVAL);
1266     }
1267
1268     return 0;
1269 }
1270
1271 static int nvenc_find_free_reg_resource(AVCodecContext *avctx)
1272 {
1273     NvencContext *ctx = avctx->priv_data;
1274     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1275     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1276
1277     int i;
1278
1279     if (ctx->nb_registered_frames == FF_ARRAY_ELEMS(ctx->registered_frames)) {
1280         for (i = 0; i < ctx->nb_registered_frames; i++) {
1281             if (!ctx->registered_frames[i].mapped) {
1282                 if (ctx->registered_frames[i].regptr) {
1283                     p_nvenc->nvEncUnregisterResource(ctx->nvencoder,
1284                                                 ctx->registered_frames[i].regptr);
1285                     ctx->registered_frames[i].regptr = NULL;
1286                 }
1287                 return i;
1288             }
1289         }
1290     } else {
1291         return ctx->nb_registered_frames++;
1292     }
1293
1294     av_log(avctx, AV_LOG_ERROR, "Too many registered CUDA frames\n");
1295     return AVERROR(ENOMEM);
1296 }
1297
1298 static int nvenc_register_frame(AVCodecContext *avctx, const AVFrame *frame)
1299 {
1300     NvencContext *ctx = avctx->priv_data;
1301     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1302     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1303
1304     AVHWFramesContext *frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
1305     NV_ENC_REGISTER_RESOURCE reg;
1306     int i, idx, ret;
1307
1308     for (i = 0; i < ctx->nb_registered_frames; i++) {
1309         if (ctx->registered_frames[i].ptr == (CUdeviceptr)frame->data[0])
1310             return i;
1311     }
1312
1313     idx = nvenc_find_free_reg_resource(avctx);
1314     if (idx < 0)
1315         return idx;
1316
1317     reg.version            = NV_ENC_REGISTER_RESOURCE_VER;
1318     reg.resourceType       = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
1319     reg.width              = frames_ctx->width;
1320     reg.height             = frames_ctx->height;
1321     reg.bufferFormat       = ctx->surfaces[0].format;
1322     reg.pitch              = frame->linesize[0];
1323     reg.resourceToRegister = frame->data[0];
1324
1325     ret = p_nvenc->nvEncRegisterResource(ctx->nvencoder, &reg);
1326     if (ret != NV_ENC_SUCCESS) {
1327         nvenc_print_error(avctx, ret, "Error registering an input resource");
1328         return AVERROR_UNKNOWN;
1329     }
1330
1331     ctx->registered_frames[idx].ptr    = (CUdeviceptr)frame->data[0];
1332     ctx->registered_frames[idx].regptr = reg.registeredResource;
1333     return idx;
1334 }
1335
1336 static int nvenc_upload_frame(AVCodecContext *avctx, const AVFrame *frame,
1337                                       NvencSurface *nvenc_frame)
1338 {
1339     NvencContext *ctx = avctx->priv_data;
1340     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1341     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1342
1343     int res;
1344     NVENCSTATUS nv_status;
1345
1346     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1347         int reg_idx = nvenc_register_frame(avctx, frame);
1348         if (reg_idx < 0) {
1349             av_log(avctx, AV_LOG_ERROR, "Could not register an input CUDA frame\n");
1350             return reg_idx;
1351         }
1352
1353         res = av_frame_ref(nvenc_frame->in_ref, frame);
1354         if (res < 0)
1355             return res;
1356
1357         nvenc_frame->in_map.version = NV_ENC_MAP_INPUT_RESOURCE_VER;
1358         nvenc_frame->in_map.registeredResource = ctx->registered_frames[reg_idx].regptr;
1359         nv_status = p_nvenc->nvEncMapInputResource(ctx->nvencoder, &nvenc_frame->in_map);
1360         if (nv_status != NV_ENC_SUCCESS) {
1361             av_frame_unref(nvenc_frame->in_ref);
1362             return nvenc_print_error(avctx, nv_status, "Error mapping an input resource");
1363         }
1364
1365         ctx->registered_frames[reg_idx].mapped = 1;
1366         nvenc_frame->reg_idx                   = reg_idx;
1367         nvenc_frame->input_surface             = nvenc_frame->in_map.mappedResource;
1368         return 0;
1369     } else {
1370         NV_ENC_LOCK_INPUT_BUFFER lockBufferParams = { 0 };
1371
1372         lockBufferParams.version = NV_ENC_LOCK_INPUT_BUFFER_VER;
1373         lockBufferParams.inputBuffer = nvenc_frame->input_surface;
1374
1375         nv_status = p_nvenc->nvEncLockInputBuffer(ctx->nvencoder, &lockBufferParams);
1376         if (nv_status != NV_ENC_SUCCESS) {
1377             return nvenc_print_error(avctx, nv_status, "Failed locking nvenc input buffer");
1378         }
1379
1380         res = nvenc_copy_frame(avctx, nvenc_frame, &lockBufferParams, frame);
1381
1382         nv_status = p_nvenc->nvEncUnlockInputBuffer(ctx->nvencoder, nvenc_frame->input_surface);
1383         if (nv_status != NV_ENC_SUCCESS) {
1384             return nvenc_print_error(avctx, nv_status, "Failed unlocking input buffer!");
1385         }
1386
1387         return res;
1388     }
1389 }
1390
1391 static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,
1392                                                     NV_ENC_PIC_PARAMS *params)
1393 {
1394     NvencContext *ctx = avctx->priv_data;
1395
1396     switch (avctx->codec->id) {
1397     case AV_CODEC_ID_H264:
1398       params->codecPicParams.h264PicParams.sliceMode = ctx->encode_config.encodeCodecConfig.h264Config.sliceMode;
1399       params->codecPicParams.h264PicParams.sliceModeData = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;
1400       break;
1401     case AV_CODEC_ID_H265:
1402       params->codecPicParams.hevcPicParams.sliceMode = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceMode;
1403       params->codecPicParams.hevcPicParams.sliceModeData = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData;
1404       break;
1405     }
1406 }
1407
1408 static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencSurface *tmpoutsurf)
1409 {
1410     NvencContext *ctx = avctx->priv_data;
1411     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1412     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1413
1414     uint32_t slice_mode_data;
1415     uint32_t *slice_offsets;
1416     NV_ENC_LOCK_BITSTREAM lock_params = { 0 };
1417     NVENCSTATUS nv_status;
1418     int res = 0;
1419
1420     enum AVPictureType pict_type;
1421
1422     switch (avctx->codec->id) {
1423     case AV_CODEC_ID_H264:
1424       slice_mode_data = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;
1425       break;
1426     case AV_CODEC_ID_H265:
1427       slice_mode_data = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData;
1428       break;
1429     default:
1430       av_log(avctx, AV_LOG_ERROR, "Unknown codec name\n");
1431       res = AVERROR(EINVAL);
1432       goto error;
1433     }
1434     slice_offsets = av_mallocz(slice_mode_data * sizeof(*slice_offsets));
1435
1436     if (!slice_offsets)
1437         return AVERROR(ENOMEM);
1438
1439     lock_params.version = NV_ENC_LOCK_BITSTREAM_VER;
1440
1441     lock_params.doNotWait = 0;
1442     lock_params.outputBitstream = tmpoutsurf->output_surface;
1443     lock_params.sliceOffsets = slice_offsets;
1444
1445     nv_status = p_nvenc->nvEncLockBitstream(ctx->nvencoder, &lock_params);
1446     if (nv_status != NV_ENC_SUCCESS) {
1447         res = nvenc_print_error(avctx, nv_status, "Failed locking bitstream buffer");
1448         goto error;
1449     }
1450
1451     if (res = ff_alloc_packet2(avctx, pkt, lock_params.bitstreamSizeInBytes,0)) {
1452         p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
1453         goto error;
1454     }
1455
1456     memcpy(pkt->data, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes);
1457
1458     nv_status = p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
1459     if (nv_status != NV_ENC_SUCCESS)
1460         nvenc_print_error(avctx, nv_status, "Failed unlocking bitstream buffer, expect the gates of mordor to open");
1461
1462
1463     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1464         p_nvenc->nvEncUnmapInputResource(ctx->nvencoder, tmpoutsurf->in_map.mappedResource);
1465         av_frame_unref(tmpoutsurf->in_ref);
1466         ctx->registered_frames[tmpoutsurf->reg_idx].mapped = 0;
1467
1468         tmpoutsurf->input_surface = NULL;
1469     }
1470
1471     switch (lock_params.pictureType) {
1472     case NV_ENC_PIC_TYPE_IDR:
1473         pkt->flags |= AV_PKT_FLAG_KEY;
1474     case NV_ENC_PIC_TYPE_I:
1475         pict_type = AV_PICTURE_TYPE_I;
1476         break;
1477     case NV_ENC_PIC_TYPE_P:
1478         pict_type = AV_PICTURE_TYPE_P;
1479         break;
1480     case NV_ENC_PIC_TYPE_B:
1481         pict_type = AV_PICTURE_TYPE_B;
1482         break;
1483     case NV_ENC_PIC_TYPE_BI:
1484         pict_type = AV_PICTURE_TYPE_BI;
1485         break;
1486     default:
1487         av_log(avctx, AV_LOG_ERROR, "Unknown picture type encountered, expect the output to be broken.\n");
1488         av_log(avctx, AV_LOG_ERROR, "Please report this error and include as much information on how to reproduce it as possible.\n");
1489         res = AVERROR_EXTERNAL;
1490         goto error;
1491     }
1492
1493 #if FF_API_CODED_FRAME
1494 FF_DISABLE_DEPRECATION_WARNINGS
1495     avctx->coded_frame->pict_type = pict_type;
1496 FF_ENABLE_DEPRECATION_WARNINGS
1497 #endif
1498
1499     ff_side_data_set_encoder_stats(pkt,
1500         (lock_params.frameAvgQP - 1) * FF_QP2LAMBDA, NULL, 0, pict_type);
1501
1502     pkt->pts = lock_params.outputTimeStamp;
1503     pkt->dts = timestamp_queue_dequeue(ctx->timestamp_list);
1504
1505     /* when there're b frame(s), set dts offset */
1506     if (ctx->encode_config.frameIntervalP >= 2)
1507         pkt->dts -= 1;
1508
1509     if (pkt->dts > pkt->pts)
1510         pkt->dts = pkt->pts;
1511
1512     if (ctx->last_dts != AV_NOPTS_VALUE && pkt->dts <= ctx->last_dts)
1513         pkt->dts = ctx->last_dts + 1;
1514
1515     ctx->last_dts = pkt->dts;
1516
1517     av_free(slice_offsets);
1518
1519     return 0;
1520
1521 error:
1522
1523     av_free(slice_offsets);
1524     timestamp_queue_dequeue(ctx->timestamp_list);
1525
1526     return res;
1527 }
1528
1529 static int output_ready(NvencContext *ctx, int flush)
1530 {
1531     int nb_ready, nb_pending;
1532
1533     nb_ready   = av_fifo_size(ctx->output_surface_ready_queue)   / sizeof(NvencSurface*);
1534     nb_pending = av_fifo_size(ctx->output_surface_queue)         / sizeof(NvencSurface*);
1535     return nb_ready > 0 && (flush || nb_ready + nb_pending >= ctx->async_depth);
1536 }
1537
1538 int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
1539     const AVFrame *frame, int *got_packet)
1540 {
1541     NVENCSTATUS nv_status;
1542     NvencSurface *tmpoutsurf, *inSurf;
1543     int res;
1544
1545     NvencContext *ctx = avctx->priv_data;
1546     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1547     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1548
1549     NV_ENC_PIC_PARAMS pic_params = { 0 };
1550     pic_params.version = NV_ENC_PIC_PARAMS_VER;
1551
1552     if (frame) {
1553         inSurf = get_free_frame(ctx);
1554         av_assert0(inSurf);
1555
1556         res = nvenc_upload_frame(avctx, frame, inSurf);
1557         if (res) {
1558             inSurf->lockCount = 0;
1559             return res;
1560         }
1561
1562         pic_params.inputBuffer = inSurf->input_surface;
1563         pic_params.bufferFmt = inSurf->format;
1564         pic_params.inputWidth = avctx->width;
1565         pic_params.inputHeight = avctx->height;
1566         pic_params.outputBitstream = inSurf->output_surface;
1567         pic_params.completionEvent = 0;
1568
1569         if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
1570             if (frame->top_field_first) {
1571                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
1572             } else {
1573                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
1574             }
1575         } else {
1576             pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
1577         }
1578
1579         pic_params.encodePicFlags = 0;
1580         pic_params.inputTimeStamp = frame->pts;
1581         pic_params.inputDuration = 0;
1582
1583         nvenc_codec_specific_pic_params(avctx, &pic_params);
1584
1585         timestamp_queue_enqueue(ctx->timestamp_list, frame->pts);
1586     } else {
1587         pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
1588     }
1589
1590     nv_status = p_nvenc->nvEncEncodePicture(ctx->nvencoder, &pic_params);
1591
1592     if (frame && nv_status == NV_ENC_ERR_NEED_MORE_INPUT)
1593         av_fifo_generic_write(ctx->output_surface_queue, &inSurf, sizeof(inSurf), NULL);
1594
1595     if (nv_status != NV_ENC_SUCCESS && nv_status != NV_ENC_ERR_NEED_MORE_INPUT) {
1596         return nvenc_print_error(avctx, nv_status, "EncodePicture failed!");
1597     }
1598
1599     if (nv_status != NV_ENC_ERR_NEED_MORE_INPUT) {
1600         while (av_fifo_size(ctx->output_surface_queue) > 0) {
1601             av_fifo_generic_read(ctx->output_surface_queue, &tmpoutsurf, sizeof(tmpoutsurf), NULL);
1602             av_fifo_generic_write(ctx->output_surface_ready_queue, &tmpoutsurf, sizeof(tmpoutsurf), NULL);
1603         }
1604
1605         if (frame)
1606             av_fifo_generic_write(ctx->output_surface_ready_queue, &inSurf, sizeof(inSurf), NULL);
1607     }
1608
1609     if (output_ready(ctx, !frame)) {
1610         av_fifo_generic_read(ctx->output_surface_ready_queue, &tmpoutsurf, sizeof(tmpoutsurf), NULL);
1611
1612         res = process_output_surface(avctx, pkt, tmpoutsurf);
1613
1614         if (res)
1615             return res;
1616
1617         av_assert0(tmpoutsurf->lockCount);
1618         tmpoutsurf->lockCount--;
1619
1620         *got_packet = 1;
1621     } else {
1622         *got_packet = 0;
1623     }
1624
1625     return 0;
1626 }