nvenc: export CPB props side data
[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/imgutils.h"
51 #include "libavutil/mem.h"
52 #include "avcodec.h"
53 #include "internal.h"
54 #include "nvenc.h"
55
56 #define NVENC_CAP 0x30
57 #define BITSTREAM_BUFFER_SIZE 1024 * 1024
58
59 #define LOAD_LIBRARY(l, path)                   \
60     do {                                        \
61         if (!((l) = dlopen(path, RTLD_LAZY))) { \
62             av_log(avctx, AV_LOG_ERROR,         \
63                    "Cannot load %s\n",          \
64                    path);                       \
65             return AVERROR_UNKNOWN;             \
66         }                                       \
67     } while (0)
68
69 #define LOAD_SYMBOL(fun, lib, symbol)        \
70     do {                                     \
71         if (!((fun) = dlsym(lib, symbol))) { \
72             av_log(avctx, AV_LOG_ERROR,      \
73                    "Cannot load %s\n",       \
74                    symbol);                  \
75             return AVERROR_UNKNOWN;          \
76         }                                    \
77     } while (0)
78
79 static av_cold int nvenc_load_libraries(AVCodecContext *avctx)
80 {
81     NVENCContext *ctx         = avctx->priv_data;
82     NVENCLibraryContext *nvel = &ctx->nvel;
83     PNVENCODEAPICREATEINSTANCE nvenc_create_instance;
84
85     LOAD_LIBRARY(nvel->cuda, CUDA_LIBNAME);
86
87     LOAD_SYMBOL(nvel->cu_init, nvel->cuda, "cuInit");
88     LOAD_SYMBOL(nvel->cu_device_get_count, nvel->cuda, "cuDeviceGetCount");
89     LOAD_SYMBOL(nvel->cu_device_get, nvel->cuda, "cuDeviceGet");
90     LOAD_SYMBOL(nvel->cu_device_get_name, nvel->cuda, "cuDeviceGetName");
91     LOAD_SYMBOL(nvel->cu_device_compute_capability, nvel->cuda,
92                 "cuDeviceComputeCapability");
93     LOAD_SYMBOL(nvel->cu_ctx_create, nvel->cuda, "cuCtxCreate_v2");
94     LOAD_SYMBOL(nvel->cu_ctx_pop_current, nvel->cuda, "cuCtxPopCurrent_v2");
95     LOAD_SYMBOL(nvel->cu_ctx_destroy, nvel->cuda, "cuCtxDestroy_v2");
96
97     LOAD_LIBRARY(nvel->nvenc, NVENC_LIBNAME);
98
99     LOAD_SYMBOL(nvenc_create_instance, nvel->nvenc,
100                 "NvEncodeAPICreateInstance");
101
102     nvel->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER;
103
104     if ((nvenc_create_instance(&nvel->nvenc_funcs)) != NV_ENC_SUCCESS) {
105         av_log(avctx, AV_LOG_ERROR, "Cannot create the NVENC instance");
106         return AVERROR_UNKNOWN;
107     }
108
109     return 0;
110 }
111
112 static int nvenc_open_session(AVCodecContext *avctx)
113 {
114     NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS params = { 0 };
115     NVENCContext *ctx                           = avctx->priv_data;
116     NV_ENCODE_API_FUNCTION_LIST *nv             = &ctx->nvel.nvenc_funcs;
117     int ret;
118
119     params.version    = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;
120     params.apiVersion = NVENCAPI_VERSION;
121     params.device     = ctx->cu_context;
122     params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
123
124     ret = nv->nvEncOpenEncodeSessionEx(&params, &ctx->nvenc_ctx);
125     if (ret != NV_ENC_SUCCESS) {
126         ctx->nvenc_ctx = NULL;
127         av_log(avctx, AV_LOG_ERROR,
128                "Cannot open the NVENC Session\n");
129         return AVERROR_UNKNOWN;
130     }
131
132     return 0;
133 }
134
135 static int nvenc_check_codec_support(AVCodecContext *avctx)
136 {
137     NVENCContext *ctx               = avctx->priv_data;
138     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
139     int i, ret, count = 0;
140     GUID *guids = NULL;
141
142     ret = nv->nvEncGetEncodeGUIDCount(ctx->nvenc_ctx, &count);
143
144     if (ret != NV_ENC_SUCCESS || !count)
145         return AVERROR(ENOSYS);
146
147     guids = av_malloc(count * sizeof(GUID));
148     if (!guids)
149         return AVERROR(ENOMEM);
150
151     ret = nv->nvEncGetEncodeGUIDs(ctx->nvenc_ctx, guids, count, &count);
152     if (ret != NV_ENC_SUCCESS) {
153         ret = AVERROR(ENOSYS);
154         goto fail;
155     }
156
157     ret = AVERROR(ENOSYS);
158     for (i = 0; i < count; i++) {
159         if (!memcmp(&guids[i], &ctx->params.encodeGUID, sizeof(*guids))) {
160             ret = 0;
161             break;
162         }
163     }
164
165 fail:
166     av_free(guids);
167
168     return ret;
169 }
170
171 static int nvenc_check_cap(AVCodecContext *avctx, NV_ENC_CAPS cap)
172 {
173     NVENCContext *ctx               = avctx->priv_data;
174     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
175     NV_ENC_CAPS_PARAM params        = { 0 };
176     int ret, val = 0;
177
178     params.version     = NV_ENC_CAPS_PARAM_VER;
179     params.capsToQuery = cap;
180
181     ret = nv->nvEncGetEncodeCaps(ctx->nvenc_ctx, ctx->params.encodeGUID, &params, &val);
182
183     if (ret == NV_ENC_SUCCESS)
184         return val;
185     return 0;
186 }
187
188 static int nvenc_check_capabilities(AVCodecContext *avctx)
189 {
190     int ret;
191
192     ret = nvenc_check_codec_support(avctx);
193     if (ret < 0) {
194         av_log(avctx, AV_LOG_VERBOSE, "Codec not supported\n");
195         return ret;
196     }
197
198     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_YUV444_ENCODE);
199     if (avctx->pix_fmt == AV_PIX_FMT_YUV444P && ret <= 0) {
200         av_log(avctx, AV_LOG_VERBOSE, "YUV444P not supported\n");
201         return AVERROR(ENOSYS);
202     }
203
204     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_WIDTH_MAX);
205     if (ret < avctx->width) {
206         av_log(avctx, AV_LOG_VERBOSE, "Width %d exceeds %d\n",
207                avctx->width, ret);
208         return AVERROR(ENOSYS);
209     }
210
211     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_HEIGHT_MAX);
212     if (ret < avctx->height) {
213         av_log(avctx, AV_LOG_VERBOSE, "Height %d exceeds %d\n",
214                avctx->height, ret);
215         return AVERROR(ENOSYS);
216     }
217
218     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_NUM_MAX_BFRAMES);
219     if (ret < avctx->max_b_frames) {
220         av_log(avctx, AV_LOG_VERBOSE, "Max b-frames %d exceed %d\n",
221                avctx->max_b_frames, ret);
222
223         return AVERROR(ENOSYS);
224     }
225
226     return 0;
227 }
228
229 static int nvenc_check_device(AVCodecContext *avctx, int idx)
230 {
231     NVENCContext *ctx               = avctx->priv_data;
232     NVENCLibraryContext *nvel       = &ctx->nvel;
233     char name[128]                  = { 0 };
234     int major, minor, ret;
235     CUdevice cu_device;
236     CUcontext dummy;
237     int loglevel = AV_LOG_VERBOSE;
238
239     if (ctx->device == LIST_DEVICES)
240         loglevel = AV_LOG_INFO;
241
242     ret = nvel->cu_device_get(&cu_device, idx);
243     if (ret != CUDA_SUCCESS) {
244         av_log(avctx, AV_LOG_ERROR,
245                "Cannot access the CUDA device %d\n",
246                idx);
247         return -1;
248     }
249
250     ret = nvel->cu_device_get_name(name, sizeof(name), cu_device);
251     if (ret != CUDA_SUCCESS)
252         return -1;
253
254     ret = nvel->cu_device_compute_capability(&major, &minor, cu_device);
255     if (ret != CUDA_SUCCESS)
256         return -1;
257
258     av_log(avctx, loglevel, "Device %d [%s] ", cu_device, name);
259
260     if (((major << 4) | minor) < NVENC_CAP)
261         goto fail;
262
263     ret = nvel->cu_ctx_create(&ctx->cu_context, 0, cu_device);
264     if (ret != CUDA_SUCCESS)
265         goto fail;
266
267     ret = nvel->cu_ctx_pop_current(&dummy);
268     if (ret != CUDA_SUCCESS)
269         goto fail2;
270
271     if ((ret = nvenc_open_session(avctx)) < 0)
272         goto fail2;
273
274     if ((ret = nvenc_check_capabilities(avctx)) < 0)
275         goto fail3;
276
277     av_log(avctx, loglevel, "supports NVENC\n");
278
279     if (ctx->device == cu_device || ctx->device == ANY_DEVICE)
280         return 0;
281
282 fail3:
283     nvel->nvenc_funcs.nvEncDestroyEncoder(ctx->nvenc_ctx);
284     ctx->nvenc_ctx = NULL;
285
286 fail2:
287     nvel->cu_ctx_destroy(ctx->cu_context);
288     ctx->cu_context = NULL;
289
290 fail:
291     if (ret != 0)
292         av_log(avctx, loglevel, "does not support NVENC (major %d minor %d)\n",
293                major, minor);
294
295     return AVERROR(ENOSYS);
296 }
297
298 static int nvenc_setup_device(AVCodecContext *avctx)
299 {
300     NVENCContext *ctx         = avctx->priv_data;
301     NVENCLibraryContext *nvel = &ctx->nvel;
302     int i, nb_devices = 0;
303
304     if ((nvel->cu_init(0)) != CUDA_SUCCESS) {
305         av_log(avctx, AV_LOG_ERROR,
306                "Cannot init CUDA\n");
307         return AVERROR_UNKNOWN;
308     }
309
310     if ((nvel->cu_device_get_count(&nb_devices)) != CUDA_SUCCESS) {
311         av_log(avctx, AV_LOG_ERROR,
312                "Cannot enumerate the CUDA devices\n");
313         return AVERROR_UNKNOWN;
314     }
315
316     switch (avctx->codec->id) {
317     case AV_CODEC_ID_H264:
318         ctx->params.encodeGUID = NV_ENC_CODEC_H264_GUID;
319         break;
320     case AV_CODEC_ID_HEVC:
321         ctx->params.encodeGUID = NV_ENC_CODEC_HEVC_GUID;
322         break;
323     default:
324         return AVERROR_BUG;
325     }
326
327     for (i = 0; i < nb_devices; ++i) {
328         if ((nvenc_check_device(avctx, i)) >= 0 && ctx->device != LIST_DEVICES)
329             return 0;
330     }
331
332     if (ctx->device == LIST_DEVICES)
333         return AVERROR_EXIT;
334
335     return AVERROR(ENOSYS);
336 }
337
338 typedef struct GUIDTuple {
339     const GUID guid;
340     int flags;
341 } GUIDTuple;
342
343 static int nvec_map_preset(NVENCContext *ctx)
344 {
345     GUIDTuple presets[] = {
346         { NV_ENC_PRESET_DEFAULT_GUID },
347         { NV_ENC_PRESET_HP_GUID },
348         { NV_ENC_PRESET_HQ_GUID },
349         { NV_ENC_PRESET_BD_GUID },
350         { NV_ENC_PRESET_LOW_LATENCY_DEFAULT_GUID, NVENC_LOWLATENCY },
351         { NV_ENC_PRESET_LOW_LATENCY_HP_GUID,      NVENC_LOWLATENCY },
352         { NV_ENC_PRESET_LOW_LATENCY_HQ_GUID,      NVENC_LOWLATENCY },
353         { NV_ENC_PRESET_LOSSLESS_DEFAULT_GUID,    NVENC_LOSSLESS },
354         { NV_ENC_PRESET_LOSSLESS_HP_GUID,         NVENC_LOSSLESS },
355         { { 0 } }
356     };
357
358     GUIDTuple *t = &presets[ctx->preset];
359
360     ctx->params.presetGUID = t->guid;
361     ctx->flags             = t->flags;
362
363     return AVERROR(EINVAL);
364 }
365
366 static void set_constqp(AVCodecContext *avctx, NV_ENC_RC_PARAMS *rc)
367 {
368     rc->rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
369     rc->constQP.qpInterB = avctx->global_quality;
370     rc->constQP.qpInterP = avctx->global_quality;
371     rc->constQP.qpIntra  = avctx->global_quality;
372 }
373
374 static void set_vbr(AVCodecContext *avctx, NV_ENC_RC_PARAMS *rc)
375 {
376     if (avctx->qmin >= 0) {
377         rc->enableMinQP    = 1;
378         rc->minQP.qpInterB = avctx->qmin;
379         rc->minQP.qpInterP = avctx->qmin;
380         rc->minQP.qpIntra  = avctx->qmin;
381     }
382
383     if (avctx->qmax >= 0) {
384         rc->enableMaxQP = 1;
385         rc->maxQP.qpInterB = avctx->qmax;
386         rc->maxQP.qpInterP = avctx->qmax;
387         rc->maxQP.qpIntra  = avctx->qmax;
388     }
389 }
390
391 static void nvenc_override_rate_control(AVCodecContext *avctx,
392                                         NV_ENC_RC_PARAMS *rc)
393 {
394     NVENCContext *ctx    = avctx->priv_data;
395
396     switch (ctx->rc) {
397     case NV_ENC_PARAMS_RC_CONSTQP:
398         if (avctx->global_quality < 0) {
399             av_log(avctx, AV_LOG_WARNING,
400                    "The constant quality rate-control requires "
401                    "the 'global_quality' option set.\n");
402             return;
403         }
404         set_constqp(avctx, rc);
405         return;
406     case NV_ENC_PARAMS_RC_2_PASS_VBR:
407     case NV_ENC_PARAMS_RC_VBR:
408         if (avctx->qmin < 0 && avctx->qmax < 0) {
409             av_log(avctx, AV_LOG_WARNING,
410                    "The variable bitrate rate-control requires "
411                    "the 'qmin' and/or 'qmax' option set.\n");
412             return;
413         }
414     case NV_ENC_PARAMS_RC_VBR_MINQP:
415         if (avctx->qmin < 0) {
416             av_log(avctx, AV_LOG_WARNING,
417                    "The variable bitrate rate-control requires "
418                    "the 'qmin' option set.\n");
419             return;
420         }
421         set_vbr(avctx, rc);
422         break;
423     case NV_ENC_PARAMS_RC_CBR:
424         break;
425     case NV_ENC_PARAMS_RC_2_PASS_QUALITY:
426     case NV_ENC_PARAMS_RC_2_PASS_FRAMESIZE_CAP:
427         if (!(ctx->flags & NVENC_LOWLATENCY)) {
428             av_log(avctx, AV_LOG_WARNING,
429                    "The multipass rate-control requires "
430                    "a low-latency preset.\n");
431             return;
432         }
433     }
434
435     rc->rateControlMode = ctx->rc;
436 }
437
438 static void nvenc_setup_rate_control(AVCodecContext *avctx)
439 {
440     NVENCContext *ctx    = avctx->priv_data;
441     NV_ENC_RC_PARAMS *rc = &ctx->config.rcParams;
442
443     if (avctx->bit_rate > 0)
444         rc->averageBitRate = avctx->bit_rate;
445
446     if (avctx->rc_max_rate > 0)
447         rc->maxBitRate = avctx->rc_max_rate;
448
449     if (ctx->rc > 0) {
450         nvenc_override_rate_control(avctx, rc);
451     } else if (avctx->global_quality > 0) {
452         set_constqp(avctx, rc);
453     } else if (avctx->qmin >= 0 && avctx->qmax >= 0) {
454         rc->rateControlMode = NV_ENC_PARAMS_RC_VBR;
455         set_vbr(avctx, rc);
456     }
457
458     if (avctx->rc_buffer_size > 0)
459         rc->vbvBufferSize = avctx->rc_buffer_size;
460
461     if (rc->averageBitRate > 0)
462         avctx->bit_rate = rc->averageBitRate;
463 }
464
465 static int nvenc_setup_h264_config(AVCodecContext *avctx)
466 {
467     NVENCContext *ctx                      = avctx->priv_data;
468     NV_ENC_CONFIG *cc                      = &ctx->config;
469     NV_ENC_CONFIG_H264 *h264               = &cc->encodeCodecConfig.h264Config;
470     NV_ENC_CONFIG_H264_VUI_PARAMETERS *vui = &h264->h264VUIParameters;
471
472     vui->colourDescriptionPresentFlag = 1;
473     vui->videoSignalTypePresentFlag   = 1;
474
475     vui->colourMatrix            = avctx->colorspace;
476     vui->colourPrimaries         = avctx->color_primaries;
477     vui->transferCharacteristics = avctx->color_trc;
478
479     vui->videoFullRangeFlag = avctx->color_range == AVCOL_RANGE_JPEG;
480
481     h264->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
482     h264->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
483
484     h264->maxNumRefFrames = avctx->refs;
485     h264->idrPeriod       = cc->gopLength;
486
487     if (ctx->profile)
488         avctx->profile = ctx->profile;
489
490     if (avctx->pix_fmt == AV_PIX_FMT_YUV444P)
491         h264->chromaFormatIDC = 3;
492     else
493         h264->chromaFormatIDC = 1;
494
495     switch (ctx->profile) {
496     case NV_ENC_H264_PROFILE_BASELINE:
497         cc->profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;
498         break;
499     case NV_ENC_H264_PROFILE_MAIN:
500         cc->profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;
501         break;
502     case NV_ENC_H264_PROFILE_HIGH:
503         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
504         break;
505     case NV_ENC_H264_PROFILE_HIGH_444:
506         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
507         break;
508     case NV_ENC_H264_PROFILE_CONSTRAINED_HIGH:
509         cc->profileGUID = NV_ENC_H264_PROFILE_CONSTRAINED_HIGH_GUID;
510         break;
511     }
512
513     h264->level = ctx->level;
514
515     return 0;
516 }
517
518 static int nvenc_setup_hevc_config(AVCodecContext *avctx)
519 {
520     NVENCContext *ctx                      = avctx->priv_data;
521     NV_ENC_CONFIG *cc                      = &ctx->config;
522     NV_ENC_CONFIG_HEVC *hevc               = &cc->encodeCodecConfig.hevcConfig;
523
524     hevc->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
525     hevc->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
526
527     hevc->maxNumRefFramesInDPB = avctx->refs;
528     hevc->idrPeriod            = cc->gopLength;
529
530     /* No other profile is supported in the current SDK version 5 */
531     cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID;
532     avctx->profile  = FF_PROFILE_HEVC_MAIN;
533
534     if (ctx->level) {
535         hevc->level = ctx->level;
536     } else {
537         hevc->level = NV_ENC_LEVEL_AUTOSELECT;
538     }
539
540     if (ctx->tier) {
541         hevc->tier = ctx->tier;
542     }
543
544     return 0;
545 }
546 static int nvenc_setup_codec_config(AVCodecContext *avctx)
547 {
548     switch (avctx->codec->id) {
549     case AV_CODEC_ID_H264:
550         return nvenc_setup_h264_config(avctx);
551     case AV_CODEC_ID_HEVC:
552         return nvenc_setup_hevc_config(avctx);
553     }
554     return 0;
555 }
556
557 static int nvenc_setup_encoder(AVCodecContext *avctx)
558 {
559     NVENCContext *ctx               = avctx->priv_data;
560     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
561     NV_ENC_PRESET_CONFIG preset_cfg = { 0 };
562     AVCPBProperties *cpb_props;
563     int ret;
564
565     ctx->params.version = NV_ENC_INITIALIZE_PARAMS_VER;
566
567     ctx->params.encodeHeight = avctx->height;
568     ctx->params.encodeWidth  = avctx->width;
569
570     if (avctx->sample_aspect_ratio.num &&
571         avctx->sample_aspect_ratio.den &&
572         (avctx->sample_aspect_ratio.num != 1 ||
573          avctx->sample_aspect_ratio.den != 1)) {
574         av_reduce(&ctx->params.darWidth,
575                   &ctx->params.darHeight,
576                   avctx->width * avctx->sample_aspect_ratio.num,
577                   avctx->height * avctx->sample_aspect_ratio.den,
578                   INT_MAX / 8);
579     } else {
580         ctx->params.darHeight = avctx->height;
581         ctx->params.darWidth  = avctx->width;
582     }
583
584     ctx->params.frameRateNum = avctx->time_base.den;
585     ctx->params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
586
587     ctx->params.enableEncodeAsync = 0;
588     ctx->params.enablePTD         = 1;
589
590     ctx->params.encodeConfig = &ctx->config;
591
592     nvec_map_preset(ctx);
593
594     preset_cfg.version           = NV_ENC_PRESET_CONFIG_VER;
595     preset_cfg.presetCfg.version = NV_ENC_CONFIG_VER;
596
597     ret = nv->nvEncGetEncodePresetConfig(ctx->nvenc_ctx,
598                                          ctx->params.encodeGUID,
599                                          ctx->params.presetGUID,
600                                          &preset_cfg);
601     if (ret != NV_ENC_SUCCESS) {
602         av_log(avctx, AV_LOG_ERROR,
603                "Cannot get the preset configuration\n");
604         return AVERROR_UNKNOWN;
605     }
606
607     memcpy(&ctx->config, &preset_cfg.presetCfg, sizeof(ctx->config));
608
609     ctx->config.version = NV_ENC_CONFIG_VER;
610
611     if (avctx->gop_size > 0) {
612         if (avctx->max_b_frames > 0) {
613             ctx->last_dts = -2;
614             /* 0 is intra-only,
615              * 1 is I/P only,
616              * 2 is one B Frame,
617              * 3 two B frames, and so on. */
618             ctx->config.frameIntervalP = avctx->max_b_frames + 1;
619         } else if (avctx->max_b_frames == 0) {
620             ctx->config.frameIntervalP = 1;
621         }
622         ctx->config.gopLength = avctx->gop_size;
623     } else if (avctx->gop_size == 0) {
624         ctx->config.frameIntervalP = 0;
625         ctx->config.gopLength      = 1;
626     }
627
628     if (ctx->config.frameIntervalP > 1)
629         avctx->max_b_frames = ctx->config.frameIntervalP - 1;
630
631     nvenc_setup_rate_control(avctx);
632
633     if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
634         ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
635     } else {
636         ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;
637     }
638
639     if ((ret = nvenc_setup_codec_config(avctx)) < 0)
640         return ret;
641
642     ret = nv->nvEncInitializeEncoder(ctx->nvenc_ctx, &ctx->params);
643     if (ret != NV_ENC_SUCCESS) {
644         av_log(avctx, AV_LOG_ERROR, "Cannot initialize the decoder");
645         return AVERROR_UNKNOWN;
646     }
647
648     cpb_props = ff_add_cpb_side_data(avctx);
649     if (!cpb_props)
650         return AVERROR(ENOMEM);
651     cpb_props->max_bitrate = avctx->rc_max_rate;
652     cpb_props->min_bitrate = avctx->rc_min_rate;
653     cpb_props->avg_bitrate = avctx->bit_rate;
654     cpb_props->buffer_size = avctx->rc_buffer_size;
655
656     return 0;
657 }
658
659 static int nvenc_alloc_surface(AVCodecContext *avctx, int idx)
660 {
661     NVENCContext *ctx               = avctx->priv_data;
662     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
663     int ret;
664     NV_ENC_CREATE_INPUT_BUFFER in_buffer      = { 0 };
665     NV_ENC_CREATE_BITSTREAM_BUFFER out_buffer = { 0 };
666
667     in_buffer.version  = NV_ENC_CREATE_INPUT_BUFFER_VER;
668     out_buffer.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
669
670     in_buffer.width  = avctx->width;
671     in_buffer.height = avctx->height;
672
673     in_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED;
674
675     switch (avctx->pix_fmt) {
676     case AV_PIX_FMT_YUV420P:
677         in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_YV12_PL;
678         break;
679     case AV_PIX_FMT_NV12:
680         in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_NV12_PL;
681         break;
682     case AV_PIX_FMT_YUV444P:
683         in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_YUV444_PL;
684         break;
685     default:
686         return AVERROR_BUG;
687     }
688
689     ret = nv->nvEncCreateInputBuffer(ctx->nvenc_ctx, &in_buffer);
690     if (ret != NV_ENC_SUCCESS) {
691         av_log(avctx, AV_LOG_ERROR, "CreateInputBuffer failed\n");
692         return AVERROR_UNKNOWN;
693     }
694
695     ctx->in[idx].in        = in_buffer.inputBuffer;
696     ctx->in[idx].format    = in_buffer.bufferFmt;
697
698     /* 1MB is large enough to hold most output frames.
699      * NVENC increases this automaticaly if it's not enough. */
700     out_buffer.size = BITSTREAM_BUFFER_SIZE;
701
702     out_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED;
703
704     ret = nv->nvEncCreateBitstreamBuffer(ctx->nvenc_ctx, &out_buffer);
705     if (ret != NV_ENC_SUCCESS) {
706         av_log(avctx, AV_LOG_ERROR, "CreateBitstreamBuffer failed\n");
707         return AVERROR_UNKNOWN;
708     }
709
710     ctx->out[idx].out  = out_buffer.bitstreamBuffer;
711     ctx->out[idx].busy = 0;
712
713     return 0;
714 }
715
716 static int nvenc_setup_surfaces(AVCodecContext *avctx)
717 {
718     NVENCContext *ctx = avctx->priv_data;
719     int i, ret;
720
721     ctx->nb_surfaces = FFMAX(4 + avctx->max_b_frames,
722                              ctx->nb_surfaces);
723
724     ctx->in = av_mallocz(ctx->nb_surfaces * sizeof(*ctx->in));
725     if (!ctx->in)
726         return AVERROR(ENOMEM);
727
728     ctx->out = av_mallocz(ctx->nb_surfaces * sizeof(*ctx->out));
729     if (!ctx->out)
730         return AVERROR(ENOMEM);
731
732     ctx->timestamps = av_fifo_alloc(ctx->nb_surfaces * sizeof(int64_t));
733     if (!ctx->timestamps)
734         return AVERROR(ENOMEM);
735     ctx->pending = av_fifo_alloc(ctx->nb_surfaces * sizeof(ctx->out));
736     if (!ctx->pending)
737         return AVERROR(ENOMEM);
738     ctx->ready = av_fifo_alloc(ctx->nb_surfaces * sizeof(ctx->out));
739     if (!ctx->ready)
740         return AVERROR(ENOMEM);
741
742     for (i = 0; i < ctx->nb_surfaces; i++) {
743         if ((ret = nvenc_alloc_surface(avctx, i)) < 0)
744             return ret;
745     }
746
747     return 0;
748 }
749
750 #define EXTRADATA_SIZE 512
751
752 static int nvenc_setup_extradata(AVCodecContext *avctx)
753 {
754     NVENCContext *ctx                     = avctx->priv_data;
755     NV_ENCODE_API_FUNCTION_LIST *nv       = &ctx->nvel.nvenc_funcs;
756     NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };
757     int ret;
758
759     avctx->extradata = av_mallocz(EXTRADATA_SIZE + AV_INPUT_BUFFER_PADDING_SIZE);
760     if (!avctx->extradata)
761         return AVERROR(ENOMEM);
762
763     payload.version              = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;
764     payload.spsppsBuffer         = avctx->extradata;
765     payload.inBufferSize         = EXTRADATA_SIZE;
766     payload.outSPSPPSPayloadSize = &avctx->extradata_size;
767
768     ret = nv->nvEncGetSequenceParams(ctx->nvenc_ctx, &payload);
769     if (ret != NV_ENC_SUCCESS) {
770         av_log(avctx, AV_LOG_ERROR, "Cannot get the extradata\n");
771         return AVERROR_UNKNOWN;
772     }
773
774     return 0;
775 }
776
777 av_cold int ff_nvenc_encode_close(AVCodecContext *avctx)
778 {
779     NVENCContext *ctx               = avctx->priv_data;
780     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
781     int i;
782
783     av_fifo_free(ctx->timestamps);
784     av_fifo_free(ctx->pending);
785     av_fifo_free(ctx->ready);
786
787     if (ctx->in) {
788         for (i = 0; i < ctx->nb_surfaces; ++i) {
789             nv->nvEncDestroyInputBuffer(ctx->nvenc_ctx, ctx->in[i].in);
790             nv->nvEncDestroyBitstreamBuffer(ctx->nvenc_ctx, ctx->out[i].out);
791         }
792     }
793
794     av_freep(&ctx->in);
795     av_freep(&ctx->out);
796
797     if (ctx->nvenc_ctx)
798         nv->nvEncDestroyEncoder(ctx->nvenc_ctx);
799
800     if (ctx->cu_context)
801         ctx->nvel.cu_ctx_destroy(ctx->cu_context);
802
803     if (ctx->nvel.nvenc)
804         dlclose(ctx->nvel.nvenc);
805
806     if (ctx->nvel.cuda)
807         dlclose(ctx->nvel.cuda);
808
809     return 0;
810 }
811
812 av_cold int ff_nvenc_encode_init(AVCodecContext *avctx)
813 {
814     int ret;
815
816     if ((ret = nvenc_load_libraries(avctx)) < 0)
817         return ret;
818
819     if ((ret = nvenc_setup_device(avctx)) < 0)
820         return ret;
821
822     if ((ret = nvenc_setup_encoder(avctx)) < 0)
823         return ret;
824
825     if ((ret = nvenc_setup_surfaces(avctx)) < 0)
826         return ret;
827
828     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
829         if ((ret = nvenc_setup_extradata(avctx)) < 0)
830             return ret;
831     }
832
833     return 0;
834 }
835
836 static NVENCInputSurface *get_input_surface(NVENCContext *ctx)
837 {
838     int i;
839
840     for (i = 0; i < ctx->nb_surfaces; i++) {
841         if (!ctx->in[i].locked) {
842             ctx->in[i].locked = 1;
843             return &ctx->in[i];
844         }
845     }
846
847     return NULL;
848 }
849
850 static NVENCOutputSurface *get_output_surface(NVENCContext *ctx)
851 {
852     int i;
853
854     for (i = 0; i < ctx->nb_surfaces; i++) {
855         if (!ctx->out[i].busy) {
856             return &ctx->out[i];
857         }
858     }
859
860     return NULL;
861 }
862
863 static int nvenc_copy_frame(NV_ENC_LOCK_INPUT_BUFFER *in, const AVFrame *frame)
864 {
865     uint8_t *buf = in->bufferDataPtr;
866     int off      = frame->height * in->pitch;
867
868     switch (frame->format) {
869     case AV_PIX_FMT_YUV420P:
870         av_image_copy_plane(buf, in->pitch,
871                             frame->data[0], frame->linesize[0],
872                             frame->width, frame->height);
873         buf += off;
874
875         av_image_copy_plane(buf, in->pitch >> 1,
876                             frame->data[2], frame->linesize[2],
877                             frame->width >> 1, frame->height >> 1);
878
879         buf += off >> 2;
880
881         av_image_copy_plane(buf, in->pitch >> 1,
882                             frame->data[1], frame->linesize[1],
883                             frame->width >> 1, frame->height >> 1);
884         break;
885     case AV_PIX_FMT_NV12:
886         av_image_copy_plane(buf, in->pitch,
887                             frame->data[0], frame->linesize[0],
888                             frame->width, frame->height);
889         buf += off;
890
891         av_image_copy_plane(buf, in->pitch,
892                             frame->data[1], frame->linesize[1],
893                             frame->width, frame->height >> 1);
894         break;
895     case AV_PIX_FMT_YUV444P:
896         av_image_copy_plane(buf, in->pitch,
897                             frame->data[0], frame->linesize[0],
898                             frame->width, frame->height);
899         buf += off;
900
901         av_image_copy_plane(buf, in->pitch,
902                             frame->data[1], frame->linesize[1],
903                             frame->width, frame->height);
904         buf += off;
905
906         av_image_copy_plane(buf, in->pitch,
907                             frame->data[2], frame->linesize[2],
908                             frame->width, frame->height);
909         break;
910     default:
911         return AVERROR_BUG;
912     }
913
914     return 0;
915 }
916
917 static int nvenc_enqueue_frame(AVCodecContext *avctx, const AVFrame *frame,
918                                NVENCInputSurface **in_surf)
919 {
920     NVENCContext *ctx               = avctx->priv_data;
921     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
922     NV_ENC_LOCK_INPUT_BUFFER params = { 0 };
923     NVENCInputSurface *in           = get_input_surface(ctx);
924     int ret;
925
926     if (!in)
927         return AVERROR_BUG;
928
929     params.version     = NV_ENC_LOCK_INPUT_BUFFER_VER;
930     params.inputBuffer = in->in;
931
932
933     ret = nv->nvEncLockInputBuffer(ctx->nvenc_ctx, &params);
934     if (ret != NV_ENC_SUCCESS) {
935         av_log(avctx, AV_LOG_ERROR, "Cannot lock the buffer %p.\n",
936                in);
937         return AVERROR_UNKNOWN;
938     }
939
940     ret = nvenc_copy_frame(&params, frame);
941     if (ret < 0)
942         goto fail;
943
944     ret = nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in);
945     if (ret != NV_ENC_SUCCESS) {
946         av_log(avctx, AV_LOG_ERROR, "Cannot unlock the buffer %p.\n",
947                in);
948         return AVERROR_UNKNOWN;
949     }
950
951     *in_surf = in;
952
953     return 0;
954
955 fail:
956     nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in);
957
958     return ret;
959 }
960
961 static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,
962                                             NV_ENC_PIC_PARAMS *params)
963 {
964     NVENCContext *ctx = avctx->priv_data;
965
966     switch (avctx->codec->id) {
967     case AV_CODEC_ID_H264:
968         params->codecPicParams.h264PicParams.sliceMode =
969             ctx->config.encodeCodecConfig.h264Config.sliceMode;
970         params->codecPicParams.h264PicParams.sliceModeData =
971             ctx->config.encodeCodecConfig.h264Config.sliceModeData;
972         break;
973     case AV_CODEC_ID_HEVC:
974         params->codecPicParams.hevcPicParams.sliceMode =
975             ctx->config.encodeCodecConfig.hevcConfig.sliceMode;
976         params->codecPicParams.hevcPicParams.sliceModeData =
977             ctx->config.encodeCodecConfig.hevcConfig.sliceModeData;
978         break;
979     }
980 }
981
982 static inline int nvenc_enqueue_timestamp(AVFifoBuffer *f, int64_t pts)
983 {
984     return av_fifo_generic_write(f, &pts, sizeof(pts), NULL);
985 }
986
987 static inline int nvenc_dequeue_timestamp(AVFifoBuffer *f, int64_t *pts)
988 {
989     return av_fifo_generic_read(f, pts, sizeof(*pts), NULL);
990 }
991
992 static inline int nvenc_enqueue_surface(AVFifoBuffer *f,
993                                         NVENCOutputSurface *surf)
994 {
995     surf->busy = 1;
996     return av_fifo_generic_write(f, &surf, sizeof(surf), NULL);
997 }
998
999 static inline int nvenc_dequeue_surface(AVFifoBuffer *f,
1000                                         NVENCOutputSurface **surf)
1001 {
1002     return av_fifo_generic_read(f, surf, sizeof(*surf), NULL);
1003 }
1004
1005 static int nvenc_set_timestamp(NVENCContext *ctx,
1006                                NV_ENC_LOCK_BITSTREAM *params,
1007                                AVPacket *pkt)
1008 {
1009     pkt->pts      = params->outputTimeStamp;
1010     pkt->duration = params->outputDuration;
1011
1012     return nvenc_dequeue_timestamp(ctx->timestamps, &pkt->dts);
1013 }
1014
1015 static int nvenc_get_frame(AVCodecContext *avctx, AVPacket *pkt)
1016 {
1017     NVENCContext *ctx               = avctx->priv_data;
1018     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1019     NV_ENC_LOCK_BITSTREAM params    = { 0 };
1020     NVENCOutputSurface *out         = NULL;
1021     int ret;
1022
1023     ret = nvenc_dequeue_surface(ctx->pending, &out);
1024     if (ret)
1025         return ret;
1026
1027     params.version         = NV_ENC_LOCK_BITSTREAM_VER;
1028     params.outputBitstream = out->out;
1029
1030     ret = nv->nvEncLockBitstream(ctx->nvenc_ctx, &params);
1031     if (ret < 0)
1032         return AVERROR_UNKNOWN;
1033
1034     ret = ff_alloc_packet(pkt, params.bitstreamSizeInBytes);
1035     if (ret < 0)
1036         return ret;
1037
1038     memcpy(pkt->data, params.bitstreamBufferPtr, pkt->size);
1039
1040     ret = nv->nvEncUnlockBitstream(ctx->nvenc_ctx, out->out);
1041     if (ret < 0)
1042         return AVERROR_UNKNOWN;
1043
1044     out->busy = out->in->locked = 0;
1045
1046     ret = nvenc_set_timestamp(ctx, &params, pkt);
1047     if (ret < 0)
1048         return ret;
1049
1050     switch (params.pictureType) {
1051     case NV_ENC_PIC_TYPE_IDR:
1052         pkt->flags |= AV_PKT_FLAG_KEY;
1053 #if FF_API_CODED_FRAME
1054 FF_DISABLE_DEPRECATION_WARNINGS
1055     case NV_ENC_PIC_TYPE_INTRA_REFRESH:
1056     case NV_ENC_PIC_TYPE_I:
1057         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
1058         break;
1059     case NV_ENC_PIC_TYPE_P:
1060         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_P;
1061         break;
1062     case NV_ENC_PIC_TYPE_B:
1063         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_B;
1064         break;
1065     case NV_ENC_PIC_TYPE_BI:
1066         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_BI;
1067         break;
1068 FF_ENABLE_DEPRECATION_WARNINGS
1069 #endif
1070     }
1071
1072     return 0;
1073 }
1074
1075 int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
1076                           const AVFrame *frame, int *got_packet)
1077 {
1078     NVENCContext *ctx               = avctx->priv_data;
1079     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1080     NV_ENC_PIC_PARAMS params        = { 0 };
1081     NVENCInputSurface *in           = NULL;
1082     NVENCOutputSurface *out         = NULL;
1083     int ret;
1084
1085     params.version = NV_ENC_PIC_PARAMS_VER;
1086
1087     if (frame) {
1088         ret = nvenc_enqueue_frame(avctx, frame, &in);
1089         if (ret < 0)
1090             return ret;
1091         out = get_output_surface(ctx);
1092         if (!out)
1093             return AVERROR_BUG;
1094
1095         out->in = in;
1096
1097         params.inputBuffer     = in->in;
1098         params.bufferFmt       = in->format;
1099         params.inputWidth      = frame->width;
1100         params.inputHeight     = frame->height;
1101         params.outputBitstream = out->out;
1102         params.inputTimeStamp  = frame->pts;
1103
1104         if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
1105             if (frame->top_field_first)
1106                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
1107             else
1108                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
1109         } else {
1110             params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
1111         }
1112
1113         nvenc_codec_specific_pic_params(avctx, &params);
1114
1115         ret = nvenc_enqueue_timestamp(ctx->timestamps, frame->pts);
1116         if (ret < 0)
1117             return ret;
1118     } else {
1119         params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
1120     }
1121
1122     ret = nv->nvEncEncodePicture(ctx->nvenc_ctx, &params);
1123
1124     if (ret != NV_ENC_SUCCESS &&
1125         ret != NV_ENC_ERR_NEED_MORE_INPUT) {
1126
1127         return AVERROR_UNKNOWN;
1128     }
1129
1130     if (out) {
1131         ret = nvenc_enqueue_surface(ctx->pending, out);
1132         if (ret < 0)
1133             return ret;
1134     }
1135
1136     if (ret != NV_ENC_ERR_NEED_MORE_INPUT &&
1137         av_fifo_size(ctx->pending)) {
1138         ret = nvenc_get_frame(avctx, pkt);
1139         if (ret < 0)
1140             return ret;
1141         *got_packet = 1;
1142     } else {
1143         *got_packet = 0;
1144     }
1145
1146     return 0;
1147 }