nvenc: Properly free the fifos
[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     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 & AV_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 + AV_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_fifo_free(ctx->timestamps);
775     av_fifo_free(ctx->pending);
776     av_fifo_free(ctx->ready);
777
778     if (ctx->in) {
779         for (i = 0; i < ctx->nb_surfaces; ++i) {
780             nv->nvEncDestroyInputBuffer(ctx->nvenc_ctx, ctx->in[i].in);
781             nv->nvEncDestroyBitstreamBuffer(ctx->nvenc_ctx, ctx->out[i].out);
782         }
783     }
784
785     av_freep(&ctx->in);
786     av_freep(&ctx->out);
787
788     if (ctx->nvenc_ctx)
789         nv->nvEncDestroyEncoder(ctx->nvenc_ctx);
790
791     if (ctx->cu_context)
792         ctx->nvel.cu_ctx_destroy(ctx->cu_context);
793
794     if (ctx->nvel.nvenc)
795         dlclose(ctx->nvel.nvenc);
796
797     if (ctx->nvel.cuda)
798         dlclose(ctx->nvel.cuda);
799
800     return 0;
801 }
802
803 av_cold int ff_nvenc_encode_init(AVCodecContext *avctx)
804 {
805     int ret;
806
807     if ((ret = nvenc_load_libraries(avctx)) < 0)
808         return ret;
809
810     if ((ret = nvenc_setup_device(avctx)) < 0)
811         return ret;
812
813     if ((ret = nvenc_setup_encoder(avctx)) < 0)
814         return ret;
815
816     if ((ret = nvenc_setup_surfaces(avctx)) < 0)
817         return ret;
818
819     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
820         if ((ret = nvenc_setup_extradata(avctx)) < 0)
821             return ret;
822     }
823
824     return 0;
825 }
826
827 static NVENCInputSurface *get_input_surface(NVENCContext *ctx)
828 {
829     int i;
830
831     for (i = 0; i < ctx->nb_surfaces; i++) {
832         if (!ctx->in[i].locked) {
833             ctx->in[i].locked = 1;
834             return &ctx->in[i];
835         }
836     }
837
838     return NULL;
839 }
840
841 static NVENCOutputSurface *get_output_surface(NVENCContext *ctx)
842 {
843     int i;
844
845     for (i = 0; i < ctx->nb_surfaces; i++) {
846         if (!ctx->out[i].busy) {
847             return &ctx->out[i];
848         }
849     }
850
851     return NULL;
852 }
853
854 static int nvenc_copy_frame(NV_ENC_LOCK_INPUT_BUFFER *in, const AVFrame *frame)
855 {
856     uint8_t *buf = in->bufferDataPtr;
857     int off      = frame->height * in->pitch;
858
859     switch (frame->format) {
860     case AV_PIX_FMT_YUV420P:
861         av_image_copy_plane(buf, in->pitch,
862                             frame->data[0], frame->linesize[0],
863                             frame->width, frame->height);
864         buf += off;
865
866         av_image_copy_plane(buf, in->pitch >> 1,
867                             frame->data[2], frame->linesize[2],
868                             frame->width >> 1, frame->height >> 1);
869
870         buf += off >> 2;
871
872         av_image_copy_plane(buf, in->pitch >> 1,
873                             frame->data[1], frame->linesize[1],
874                             frame->width >> 1, frame->height >> 1);
875         break;
876     case AV_PIX_FMT_NV12:
877         av_image_copy_plane(buf, in->pitch,
878                             frame->data[0], frame->linesize[0],
879                             frame->width, frame->height);
880         buf += off;
881
882         av_image_copy_plane(buf, in->pitch,
883                             frame->data[1], frame->linesize[1],
884                             frame->width, frame->height >> 1);
885         break;
886     case AV_PIX_FMT_YUV444P:
887         av_image_copy_plane(buf, in->pitch,
888                             frame->data[0], frame->linesize[0],
889                             frame->width, frame->height);
890         buf += off;
891
892         av_image_copy_plane(buf, in->pitch,
893                             frame->data[1], frame->linesize[1],
894                             frame->width, frame->height);
895         buf += off;
896
897         av_image_copy_plane(buf, in->pitch,
898                             frame->data[2], frame->linesize[2],
899                             frame->width, frame->height);
900         break;
901     default:
902         return AVERROR_BUG;
903     }
904
905     return 0;
906 }
907
908 static int nvenc_enqueue_frame(AVCodecContext *avctx, const AVFrame *frame,
909                                NVENCInputSurface **in_surf)
910 {
911     NVENCContext *ctx               = avctx->priv_data;
912     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
913     NV_ENC_LOCK_INPUT_BUFFER params = { 0 };
914     NVENCInputSurface *in           = get_input_surface(ctx);
915     int ret;
916
917     if (!in)
918         return AVERROR_BUG;
919
920     params.version     = NV_ENC_LOCK_INPUT_BUFFER_VER;
921     params.inputBuffer = in->in;
922
923
924     ret = nv->nvEncLockInputBuffer(ctx->nvenc_ctx, &params);
925     if (ret != NV_ENC_SUCCESS) {
926         av_log(avctx, AV_LOG_ERROR, "Cannot lock the buffer %p.\n",
927                in);
928         return AVERROR_UNKNOWN;
929     }
930
931     ret = nvenc_copy_frame(&params, frame);
932     if (ret < 0)
933         goto fail;
934
935     ret = nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in);
936     if (ret != NV_ENC_SUCCESS) {
937         av_log(avctx, AV_LOG_ERROR, "Cannot unlock the buffer %p.\n",
938                in);
939         return AVERROR_UNKNOWN;
940     }
941
942     *in_surf = in;
943
944     return 0;
945
946 fail:
947     nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in);
948
949     return ret;
950 }
951
952 static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,
953                                             NV_ENC_PIC_PARAMS *params)
954 {
955     NVENCContext *ctx = avctx->priv_data;
956
957     switch (avctx->codec->id) {
958     case AV_CODEC_ID_H264:
959         params->codecPicParams.h264PicParams.sliceMode =
960             ctx->config.encodeCodecConfig.h264Config.sliceMode;
961         params->codecPicParams.h264PicParams.sliceModeData =
962             ctx->config.encodeCodecConfig.h264Config.sliceModeData;
963         break;
964     case AV_CODEC_ID_HEVC:
965         params->codecPicParams.hevcPicParams.sliceMode =
966             ctx->config.encodeCodecConfig.hevcConfig.sliceMode;
967         params->codecPicParams.hevcPicParams.sliceModeData =
968             ctx->config.encodeCodecConfig.hevcConfig.sliceModeData;
969         break;
970     }
971 }
972
973 static inline int nvenc_enqueue_timestamp(AVFifoBuffer *f, int64_t pts)
974 {
975     return av_fifo_generic_write(f, &pts, sizeof(pts), NULL);
976 }
977
978 static inline int nvenc_dequeue_timestamp(AVFifoBuffer *f, int64_t *pts)
979 {
980     return av_fifo_generic_read(f, pts, sizeof(*pts), NULL);
981 }
982
983 static inline int nvenc_enqueue_surface(AVFifoBuffer *f,
984                                         NVENCOutputSurface *surf)
985 {
986     surf->busy = 1;
987     return av_fifo_generic_write(f, &surf, sizeof(surf), NULL);
988 }
989
990 static inline int nvenc_dequeue_surface(AVFifoBuffer *f,
991                                         NVENCOutputSurface **surf)
992 {
993     return av_fifo_generic_read(f, surf, sizeof(*surf), NULL);
994 }
995
996 static int nvenc_set_timestamp(NVENCContext *ctx,
997                                NV_ENC_LOCK_BITSTREAM *params,
998                                AVPacket *pkt)
999 {
1000     pkt->pts      = params->outputTimeStamp;
1001     pkt->duration = params->outputDuration;
1002
1003     return nvenc_dequeue_timestamp(ctx->timestamps, &pkt->dts);
1004 }
1005
1006 static int nvenc_get_frame(AVCodecContext *avctx, AVPacket *pkt)
1007 {
1008     NVENCContext *ctx               = avctx->priv_data;
1009     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1010     NV_ENC_LOCK_BITSTREAM params    = { 0 };
1011     NVENCOutputSurface *out         = NULL;
1012     int ret;
1013
1014     ret = nvenc_dequeue_surface(ctx->pending, &out);
1015     if (ret)
1016         return ret;
1017
1018     params.version         = NV_ENC_LOCK_BITSTREAM_VER;
1019     params.outputBitstream = out->out;
1020
1021     ret = nv->nvEncLockBitstream(ctx->nvenc_ctx, &params);
1022     if (ret < 0)
1023         return AVERROR_UNKNOWN;
1024
1025     ret = ff_alloc_packet(pkt, params.bitstreamSizeInBytes);
1026     if (ret < 0)
1027         return ret;
1028
1029     memcpy(pkt->data, params.bitstreamBufferPtr, pkt->size);
1030
1031     ret = nv->nvEncUnlockBitstream(ctx->nvenc_ctx, out->out);
1032     if (ret < 0)
1033         return AVERROR_UNKNOWN;
1034
1035     out->busy = out->in->locked = 0;
1036
1037     ret = nvenc_set_timestamp(ctx, &params, pkt);
1038     if (ret < 0)
1039         return ret;
1040
1041     switch (params.pictureType) {
1042     case NV_ENC_PIC_TYPE_IDR:
1043         pkt->flags |= AV_PKT_FLAG_KEY;
1044 #if FF_API_CODED_FRAME
1045 FF_DISABLE_DEPRECATION_WARNINGS
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 FF_ENABLE_DEPRECATION_WARNINGS
1060 #endif
1061     }
1062
1063     return 0;
1064 }
1065
1066 int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
1067                           const AVFrame *frame, int *got_packet)
1068 {
1069     NVENCContext *ctx               = avctx->priv_data;
1070     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1071     NV_ENC_PIC_PARAMS params        = { 0 };
1072     NVENCInputSurface *in           = NULL;
1073     NVENCOutputSurface *out         = NULL;
1074     int ret;
1075
1076     params.version = NV_ENC_PIC_PARAMS_VER;
1077
1078     if (frame) {
1079         ret = nvenc_enqueue_frame(avctx, frame, &in);
1080         if (ret < 0)
1081             return ret;
1082         out = get_output_surface(ctx);
1083         if (!out)
1084             return AVERROR_BUG;
1085
1086         out->in = in;
1087
1088         params.inputBuffer     = in->in;
1089         params.bufferFmt       = in->format;
1090         params.inputWidth      = frame->width;
1091         params.inputHeight     = frame->height;
1092         params.outputBitstream = out->out;
1093         params.inputTimeStamp  = frame->pts;
1094
1095         if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
1096             if (frame->top_field_first)
1097                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
1098             else
1099                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
1100         } else {
1101             params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
1102         }
1103
1104         nvenc_codec_specific_pic_params(avctx, &params);
1105
1106         ret = nvenc_enqueue_timestamp(ctx->timestamps, frame->pts);
1107         if (ret < 0)
1108             return ret;
1109     } else {
1110         params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
1111     }
1112
1113     ret = nv->nvEncEncodePicture(ctx->nvenc_ctx, &params);
1114
1115     if (ret != NV_ENC_SUCCESS &&
1116         ret != NV_ENC_ERR_NEED_MORE_INPUT) {
1117
1118         return AVERROR_UNKNOWN;
1119     }
1120
1121     if (out) {
1122         ret = nvenc_enqueue_surface(ctx->pending, out);
1123         if (ret < 0)
1124             return ret;
1125     }
1126
1127     if (ret != NV_ENC_ERR_NEED_MORE_INPUT &&
1128         av_fifo_size(ctx->pending)) {
1129         ret = nvenc_get_frame(avctx, pkt);
1130         if (ret < 0)
1131             return ret;
1132         *got_packet = 1;
1133     } else {
1134         *got_packet = 0;
1135     }
1136
1137     return 0;
1138 }