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