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