nvenc: H264 and HEVC encoders
[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 & CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
482     h264->repeatSPSPPS  = (avctx->flags & 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 & CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
525     hevc->repeatSPSPPS  = (avctx->flags & 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     int ret;
563
564     ctx->params.version = NV_ENC_INITIALIZE_PARAMS_VER;
565
566     ctx->params.encodeHeight = avctx->height;
567     ctx->params.encodeWidth  = avctx->width;
568
569     if (avctx->sample_aspect_ratio.num &&
570         avctx->sample_aspect_ratio.den &&
571         (avctx->sample_aspect_ratio.num != 1 ||
572          avctx->sample_aspect_ratio.den != 1)) {
573         av_reduce(&ctx->params.darWidth,
574                   &ctx->params.darHeight,
575                   avctx->width * avctx->sample_aspect_ratio.num,
576                   avctx->height * avctx->sample_aspect_ratio.den,
577                   INT_MAX / 8);
578     } else {
579         ctx->params.darHeight = avctx->height;
580         ctx->params.darWidth  = avctx->width;
581     }
582
583     ctx->params.frameRateNum = avctx->time_base.den;
584     ctx->params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
585
586     ctx->params.enableEncodeAsync = 0;
587     ctx->params.enablePTD         = 1;
588
589     ctx->params.encodeConfig = &ctx->config;
590
591     nvec_map_preset(ctx);
592
593     preset_cfg.version           = NV_ENC_PRESET_CONFIG_VER;
594     preset_cfg.presetCfg.version = NV_ENC_CONFIG_VER;
595
596     ret = nv->nvEncGetEncodePresetConfig(ctx->nvenc_ctx,
597                                          ctx->params.encodeGUID,
598                                          ctx->params.presetGUID,
599                                          &preset_cfg);
600     if (ret != NV_ENC_SUCCESS) {
601         av_log(avctx, AV_LOG_ERROR,
602                "Cannot get the preset configuration\n");
603         return AVERROR_UNKNOWN;
604     }
605
606     memcpy(&ctx->config, &preset_cfg.presetCfg, sizeof(ctx->config));
607
608     ctx->config.version = NV_ENC_CONFIG_VER;
609
610     if (avctx->gop_size > 0) {
611         if (avctx->max_b_frames > 0) {
612             ctx->last_dts = -2;
613             /* 0 is intra-only,
614              * 1 is I/P only,
615              * 2 is one B Frame,
616              * 3 two B frames, and so on. */
617             ctx->config.frameIntervalP = avctx->max_b_frames + 1;
618         } else if (avctx->max_b_frames == 0) {
619             ctx->config.frameIntervalP = 1;
620         }
621         ctx->config.gopLength = avctx->gop_size;
622     } else if (avctx->gop_size == 0) {
623         ctx->config.frameIntervalP = 0;
624         ctx->config.gopLength      = 1;
625     }
626
627     if (ctx->config.frameIntervalP > 1)
628         avctx->max_b_frames = ctx->config.frameIntervalP - 1;
629
630     nvenc_setup_rate_control(avctx);
631
632     if (avctx->flags & CODEC_FLAG_INTERLACED_DCT) {
633         ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
634     } else {
635         ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;
636     }
637
638     if ((ret = nvenc_setup_codec_config(avctx)) < 0)
639         return ret;
640
641     ret = nv->nvEncInitializeEncoder(ctx->nvenc_ctx, &ctx->params);
642     if (ret != NV_ENC_SUCCESS) {
643         av_log(avctx, AV_LOG_ERROR, "Cannot initialize the decoder");
644         return AVERROR_UNKNOWN;
645     }
646
647     return 0;
648 }
649
650 static int nvenc_alloc_surface(AVCodecContext *avctx, int idx)
651 {
652     NVENCContext *ctx               = avctx->priv_data;
653     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
654     int ret;
655     NV_ENC_CREATE_INPUT_BUFFER in_buffer      = { 0 };
656     NV_ENC_CREATE_BITSTREAM_BUFFER out_buffer = { 0 };
657
658     in_buffer.version  = NV_ENC_CREATE_INPUT_BUFFER_VER;
659     out_buffer.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
660
661     in_buffer.width  = avctx->width;
662     in_buffer.height = avctx->height;
663
664     in_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED;
665
666     switch (avctx->pix_fmt) {
667     case AV_PIX_FMT_YUV420P:
668         in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_YV12_PL;
669         break;
670     case AV_PIX_FMT_NV12:
671         in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_NV12_PL;
672         break;
673     case AV_PIX_FMT_YUV444P:
674         in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_YUV444_PL;
675         break;
676     default:
677         return AVERROR_BUG;
678     }
679
680     ret = nv->nvEncCreateInputBuffer(ctx->nvenc_ctx, &in_buffer);
681     if (ret != NV_ENC_SUCCESS) {
682         av_log(avctx, AV_LOG_ERROR, "CreateInputBuffer failed\n");
683         return AVERROR_UNKNOWN;
684     }
685
686     ctx->in[idx].in        = in_buffer.inputBuffer;
687     ctx->in[idx].format    = in_buffer.bufferFmt;
688
689     /* 1MB is large enough to hold most output frames.
690      * NVENC increases this automaticaly if it's not enough. */
691     out_buffer.size = BITSTREAM_BUFFER_SIZE;
692
693     out_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED;
694
695     ret = nv->nvEncCreateBitstreamBuffer(ctx->nvenc_ctx, &out_buffer);
696     if (ret != NV_ENC_SUCCESS) {
697         av_log(avctx, AV_LOG_ERROR, "CreateBitstreamBuffer failed\n");
698         return AVERROR_UNKNOWN;
699     }
700
701     ctx->out[idx].out  = out_buffer.bitstreamBuffer;
702     ctx->out[idx].busy = 0;
703
704     return 0;
705 }
706
707 static int nvenc_setup_surfaces(AVCodecContext *avctx)
708 {
709     NVENCContext *ctx = avctx->priv_data;
710     int i, ret;
711
712     ctx->nb_surfaces = FFMAX(4 + avctx->max_b_frames,
713                              ctx->nb_surfaces);
714
715     ctx->in = av_mallocz(ctx->nb_surfaces * sizeof(*ctx->in));
716     if (!ctx->in)
717         return AVERROR(ENOMEM);
718
719     ctx->out = av_mallocz(ctx->nb_surfaces * sizeof(*ctx->out));
720     if (!ctx->out)
721         return AVERROR(ENOMEM);
722
723     ctx->timestamps = av_fifo_alloc(ctx->nb_surfaces * sizeof(int64_t));
724     if (!ctx->timestamps)
725         return AVERROR(ENOMEM);
726     ctx->pending = av_fifo_alloc(ctx->nb_surfaces * sizeof(ctx->out));
727     if (!ctx->pending)
728         return AVERROR(ENOMEM);
729     ctx->ready = av_fifo_alloc(ctx->nb_surfaces * sizeof(ctx->out));
730     if (!ctx->ready)
731         return AVERROR(ENOMEM);
732
733     for (i = 0; i < ctx->nb_surfaces; i++) {
734         if ((ret = nvenc_alloc_surface(avctx, i)) < 0)
735             return ret;
736     }
737
738     return 0;
739 }
740
741 #define EXTRADATA_SIZE 512
742
743 static int nvenc_setup_extradata(AVCodecContext *avctx)
744 {
745     NVENCContext *ctx                     = avctx->priv_data;
746     NV_ENCODE_API_FUNCTION_LIST *nv       = &ctx->nvel.nvenc_funcs;
747     NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };
748     int ret;
749
750     avctx->extradata = av_mallocz(EXTRADATA_SIZE + FF_INPUT_BUFFER_PADDING_SIZE);
751     if (!avctx->extradata)
752         return AVERROR(ENOMEM);
753
754     payload.version              = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;
755     payload.spsppsBuffer         = avctx->extradata;
756     payload.inBufferSize         = EXTRADATA_SIZE;
757     payload.outSPSPPSPayloadSize = &avctx->extradata_size;
758
759     ret = nv->nvEncGetSequenceParams(ctx->nvenc_ctx, &payload);
760     if (ret != NV_ENC_SUCCESS) {
761         av_log(avctx, AV_LOG_ERROR, "Cannot get the extradata\n");
762         return AVERROR_UNKNOWN;
763     }
764
765     return 0;
766 }
767
768 av_cold int ff_nvenc_encode_close(AVCodecContext *avctx)
769 {
770     NVENCContext *ctx               = avctx->priv_data;
771     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
772     int i;
773
774     av_frame_free(&avctx->coded_frame);
775
776     if (ctx->in) {
777         for (i = 0; i < ctx->nb_surfaces; ++i) {
778             nv->nvEncDestroyInputBuffer(ctx->nvenc_ctx, ctx->in[i].in);
779             nv->nvEncDestroyBitstreamBuffer(ctx->nvenc_ctx, ctx->out[i].out);
780         }
781     }
782
783     av_freep(&ctx->in);
784     av_freep(&ctx->out);
785
786     if (ctx->nvenc_ctx)
787         nv->nvEncDestroyEncoder(ctx->nvenc_ctx);
788
789     if (ctx->cu_context)
790         ctx->nvel.cu_ctx_destroy(ctx->cu_context);
791
792     if (ctx->nvel.nvenc)
793         dlclose(ctx->nvel.nvenc);
794
795     if (ctx->nvel.cuda)
796         dlclose(ctx->nvel.cuda);
797
798     return 0;
799 }
800
801 av_cold int ff_nvenc_encode_init(AVCodecContext *avctx)
802 {
803     int ret;
804
805     if ((ret = nvenc_load_libraries(avctx)) < 0)
806         return ret;
807
808     if ((ret = nvenc_setup_device(avctx)) < 0)
809         return ret;
810
811     if ((ret = nvenc_setup_encoder(avctx)) < 0)
812         return ret;
813
814     if ((ret = nvenc_setup_surfaces(avctx)) < 0)
815         return ret;
816
817     if (avctx->flags & CODEC_FLAG_GLOBAL_HEADER) {
818         if ((ret = nvenc_setup_extradata(avctx)) < 0)
819             return ret;
820     }
821
822     avctx->coded_frame = av_frame_alloc();
823     if (!avctx->coded_frame)
824         return AVERROR(ENOMEM);
825
826     return 0;
827 }
828
829 static NVENCInputSurface *get_input_surface(NVENCContext *ctx)
830 {
831     int i;
832
833     for (i = 0; i < ctx->nb_surfaces; i++) {
834         if (!ctx->in[i].locked) {
835             ctx->in[i].locked = 1;
836             return &ctx->in[i];
837         }
838     }
839
840     return NULL;
841 }
842
843 static NVENCOutputSurface *get_output_surface(NVENCContext *ctx)
844 {
845     int i;
846
847     for (i = 0; i < ctx->nb_surfaces; i++) {
848         if (!ctx->out[i].busy) {
849             return &ctx->out[i];
850         }
851     }
852
853     return NULL;
854 }
855
856 static int nvenc_copy_frame(NV_ENC_LOCK_INPUT_BUFFER *in, const AVFrame *frame)
857 {
858     uint8_t *buf = in->bufferDataPtr;
859     int off      = frame->height * in->pitch;
860
861     switch (frame->format) {
862     case AV_PIX_FMT_YUV420P:
863         av_image_copy_plane(buf, in->pitch,
864                             frame->data[0], frame->linesize[0],
865                             frame->width, frame->height);
866         buf += off;
867
868         av_image_copy_plane(buf, in->pitch >> 1,
869                             frame->data[2], frame->linesize[2],
870                             frame->width >> 1, frame->height >> 1);
871
872         buf += off >> 2;
873
874         av_image_copy_plane(buf, in->pitch >> 1,
875                             frame->data[1], frame->linesize[1],
876                             frame->width >> 1, frame->height >> 1);
877         break;
878     case AV_PIX_FMT_NV12:
879         av_image_copy_plane(buf, in->pitch,
880                             frame->data[0], frame->linesize[0],
881                             frame->width, frame->height);
882         buf += off;
883
884         av_image_copy_plane(buf, in->pitch >> 1,
885                             frame->data[1], frame->linesize[1],
886                             frame->width >> 1, frame->height >> 1);
887         break;
888     case AV_PIX_FMT_YUV444P:
889         av_image_copy_plane(buf, in->pitch,
890                             frame->data[0], frame->linesize[0],
891                             frame->width, frame->height);
892         buf += off;
893
894         av_image_copy_plane(buf, in->pitch,
895                             frame->data[1], frame->linesize[1],
896                             frame->width, frame->height);
897         buf += off;
898
899         av_image_copy_plane(buf, in->pitch,
900                             frame->data[2], frame->linesize[2],
901                             frame->width, frame->height);
902         break;
903     default:
904         return AVERROR_BUG;
905     }
906
907     return 0;
908 }
909
910 static int nvenc_enqueue_frame(AVCodecContext *avctx, const AVFrame *frame,
911                                NVENCInputSurface **in_surf)
912 {
913     NVENCContext *ctx               = avctx->priv_data;
914     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
915     NV_ENC_LOCK_INPUT_BUFFER params = { 0 };
916     NVENCInputSurface *in           = get_input_surface(ctx);
917     int ret;
918
919     if (!in)
920         return AVERROR_BUG;
921
922     params.version     = NV_ENC_LOCK_INPUT_BUFFER_VER;
923     params.inputBuffer = in->in;
924
925
926     ret = nv->nvEncLockInputBuffer(ctx->nvenc_ctx, &params);
927     if (ret != NV_ENC_SUCCESS) {
928         av_log(avctx, AV_LOG_ERROR, "Cannot lock the buffer %p.\n",
929                in);
930         return AVERROR_UNKNOWN;
931     }
932
933     ret = nvenc_copy_frame(&params, frame);
934     if (ret < 0)
935         goto fail;
936
937     ret = nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in);
938     if (ret != NV_ENC_SUCCESS) {
939         av_log(avctx, AV_LOG_ERROR, "Cannot unlock the buffer %p.\n",
940                in);
941         return AVERROR_UNKNOWN;
942     }
943
944     *in_surf = in;
945
946     return 0;
947
948 fail:
949     nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in);
950
951     return ret;
952 }
953
954 static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,
955                                             NV_ENC_PIC_PARAMS *params)
956 {
957     NVENCContext *ctx = avctx->priv_data;
958
959     switch (avctx->codec->id) {
960     case AV_CODEC_ID_H264:
961         params->codecPicParams.h264PicParams.sliceMode =
962             ctx->config.encodeCodecConfig.h264Config.sliceMode;
963         params->codecPicParams.h264PicParams.sliceModeData =
964             ctx->config.encodeCodecConfig.h264Config.sliceModeData;
965         break;
966     case AV_CODEC_ID_HEVC:
967         params->codecPicParams.hevcPicParams.sliceMode =
968             ctx->config.encodeCodecConfig.hevcConfig.sliceMode;
969         params->codecPicParams.hevcPicParams.sliceModeData =
970             ctx->config.encodeCodecConfig.hevcConfig.sliceModeData;
971         break;
972     }
973 }
974
975 static inline int nvenc_enqueue_timestamp(AVFifoBuffer *f, int64_t pts)
976 {
977     return av_fifo_generic_write(f, &pts, sizeof(pts), NULL);
978 }
979
980 static inline int nvenc_dequeue_timestamp(AVFifoBuffer *f, int64_t *pts)
981 {
982     return av_fifo_generic_read(f, pts, sizeof(*pts), NULL);
983 }
984
985 static inline int nvenc_enqueue_surface(AVFifoBuffer *f,
986                                         NVENCOutputSurface *surf)
987 {
988     surf->busy = 1;
989     return av_fifo_generic_write(f, &surf, sizeof(surf), NULL);
990 }
991
992 static inline int nvenc_dequeue_surface(AVFifoBuffer *f,
993                                         NVENCOutputSurface **surf)
994 {
995     return av_fifo_generic_read(f, surf, sizeof(*surf), NULL);
996 }
997
998 static int nvenc_set_timestamp(NVENCContext *ctx,
999                                NV_ENC_LOCK_BITSTREAM *params,
1000                                AVPacket *pkt)
1001 {
1002     pkt->pts      = params->outputTimeStamp;
1003     pkt->duration = params->outputDuration;
1004
1005     return nvenc_dequeue_timestamp(ctx->timestamps, &pkt->dts);
1006 }
1007
1008 static int nvenc_get_frame(AVCodecContext *avctx, AVPacket *pkt)
1009 {
1010     NVENCContext *ctx               = avctx->priv_data;
1011     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1012     NV_ENC_LOCK_BITSTREAM params    = { 0 };
1013     NVENCOutputSurface *out         = NULL;
1014     int ret;
1015
1016     ret = nvenc_dequeue_surface(ctx->pending, &out);
1017     if (ret)
1018         return ret;
1019
1020     params.version         = NV_ENC_LOCK_BITSTREAM_VER;
1021     params.outputBitstream = out->out;
1022
1023     ret = nv->nvEncLockBitstream(ctx->nvenc_ctx, &params);
1024     if (ret < 0)
1025         return AVERROR_UNKNOWN;
1026
1027     ret = ff_alloc_packet(pkt, params.bitstreamSizeInBytes);
1028     if (ret < 0)
1029         return ret;
1030
1031     memcpy(pkt->data, params.bitstreamBufferPtr, pkt->size);
1032
1033     ret = nv->nvEncUnlockBitstream(ctx->nvenc_ctx, out->out);
1034     if (ret < 0)
1035         return AVERROR_UNKNOWN;
1036
1037     out->busy = out->in->locked = 0;
1038
1039     ret = nvenc_set_timestamp(ctx, &params, pkt);
1040     if (ret < 0)
1041         return ret;
1042
1043     switch (params.pictureType) {
1044     case NV_ENC_PIC_TYPE_IDR:
1045         pkt->flags |= AV_PKT_FLAG_KEY;
1046     case NV_ENC_PIC_TYPE_INTRA_REFRESH:
1047     case NV_ENC_PIC_TYPE_I:
1048         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
1049         break;
1050     case NV_ENC_PIC_TYPE_P:
1051         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_P;
1052         break;
1053     case NV_ENC_PIC_TYPE_B:
1054         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_B;
1055         break;
1056     case NV_ENC_PIC_TYPE_BI:
1057         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_BI;
1058         break;
1059     }
1060
1061     return 0;
1062 }
1063
1064 int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
1065                           const AVFrame *frame, int *got_packet)
1066 {
1067     NVENCContext *ctx               = avctx->priv_data;
1068     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1069     NV_ENC_PIC_PARAMS params        = { 0 };
1070     NVENCInputSurface *in           = NULL;
1071     NVENCOutputSurface *out         = NULL;
1072     int ret;
1073
1074     params.version = NV_ENC_PIC_PARAMS_VER;
1075
1076     if (frame) {
1077         ret = nvenc_enqueue_frame(avctx, frame, &in);
1078         if (ret < 0)
1079             return ret;
1080         out = get_output_surface(ctx);
1081         if (!out)
1082             return AVERROR_BUG;
1083
1084         out->in = in;
1085
1086         params.inputBuffer     = in->in;
1087         params.bufferFmt       = in->format;
1088         params.inputWidth      = frame->width;
1089         params.inputHeight     = frame->height;
1090         params.outputBitstream = out->out;
1091         params.inputTimeStamp  = frame->pts;
1092
1093         if (avctx->flags & CODEC_FLAG_INTERLACED_DCT) {
1094             if (frame->top_field_first)
1095                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
1096             else
1097                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
1098         } else {
1099             params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
1100         }
1101
1102         nvenc_codec_specific_pic_params(avctx, &params);
1103
1104         ret = nvenc_enqueue_timestamp(ctx->timestamps, frame->pts);
1105         if (ret < 0)
1106             return ret;
1107     } else {
1108         params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
1109     }
1110
1111     ret = nv->nvEncEncodePicture(ctx->nvenc_ctx, &params);
1112
1113     if (ret != NV_ENC_SUCCESS &&
1114         ret != NV_ENC_ERR_NEED_MORE_INPUT) {
1115
1116         return AVERROR_UNKNOWN;
1117     }
1118
1119     if (out) {
1120         ret = nvenc_enqueue_surface(ctx->pending, out);
1121         if (ret < 0)
1122             return ret;
1123     }
1124
1125     if (ret != NV_ENC_ERR_NEED_MORE_INPUT &&
1126         av_fifo_size(ctx->pending)) {
1127         ret = nvenc_get_frame(avctx, pkt);
1128         if (ret < 0)
1129             return ret;
1130         *got_packet = 1;
1131     } else {
1132         *got_packet = 0;
1133     }
1134
1135     return 0;
1136 }