1 /*
2  * H.264/HEVC hardware encoding using nvidia nvenc
3  * Copyright (c) 2016 Timo Rothenpieler <timo@rothenpieler.org>
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with FFmpeg; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21 
22 #include "config.h"
23 
24 #include "nvenc.h"
25 #include "hevc_sei.h"
26 
27 #include "libavutil/hwcontext_cuda.h"
28 #include "libavutil/hwcontext.h"
29 #include "libavutil/cuda_check.h"
30 #include "libavutil/imgutils.h"
31 #include "libavutil/avassert.h"
32 #include "libavutil/mem.h"
33 #include "libavutil/pixdesc.h"
34 #include "atsc_a53.h"
35 #include "encode.h"
36 #include "internal.h"
37 #include "packet_internal.h"
38 
39 #define CHECK_CU(x) FF_CUDA_CHECK_DL(avctx, dl_fn->cuda_dl, x)
40 
41 #define NVENC_CAP 0x30
42 #define IS_CBR(rc) (rc == NV_ENC_PARAMS_RC_CBR ||             \
43                     rc == NV_ENC_PARAMS_RC_CBR_LOWDELAY_HQ || \
44                     rc == NV_ENC_PARAMS_RC_CBR_HQ)
45 
46 const enum AVPixelFormat ff_nvenc_pix_fmts[] = {
47     AV_PIX_FMT_YUV420P,
48     AV_PIX_FMT_NV12,
49     AV_PIX_FMT_P010,
50     AV_PIX_FMT_YUV444P,
51     AV_PIX_FMT_P016,      // Truncated to 10bits
52     AV_PIX_FMT_YUV444P16, // Truncated to 10bits
53     AV_PIX_FMT_0RGB32,
54     AV_PIX_FMT_0BGR32,
55     AV_PIX_FMT_CUDA,
56 #if CONFIG_D3D11VA
57     AV_PIX_FMT_D3D11,
58 #endif
59     AV_PIX_FMT_NONE
60 };
61 
62 const AVCodecHWConfigInternal *const ff_nvenc_hw_configs[] = {
63     HW_CONFIG_ENCODER_FRAMES(CUDA,  CUDA),
64     HW_CONFIG_ENCODER_DEVICE(NONE,  CUDA),
65 #if CONFIG_D3D11VA
66     HW_CONFIG_ENCODER_FRAMES(D3D11, D3D11VA),
67     HW_CONFIG_ENCODER_DEVICE(NONE,  D3D11VA),
68 #endif
69     NULL,
70 };
71 
72 #define IS_10BIT(pix_fmt)  (pix_fmt == AV_PIX_FMT_P010    || \
73                             pix_fmt == AV_PIX_FMT_P016    || \
74                             pix_fmt == AV_PIX_FMT_YUV444P16)
75 
76 #define IS_YUV444(pix_fmt) (pix_fmt == AV_PIX_FMT_YUV444P || \
77                             pix_fmt == AV_PIX_FMT_YUV444P16)
78 
79 static const struct {
80     NVENCSTATUS nverr;
81     int         averr;
82     const char *desc;
83 } nvenc_errors[] = {
84     { NV_ENC_SUCCESS,                      0,                "success"                  },
85     { NV_ENC_ERR_NO_ENCODE_DEVICE,         AVERROR(ENOENT),  "no encode device"         },
86     { NV_ENC_ERR_UNSUPPORTED_DEVICE,       AVERROR(ENOSYS),  "unsupported device"       },
87     { NV_ENC_ERR_INVALID_ENCODERDEVICE,    AVERROR(EINVAL),  "invalid encoder device"   },
88     { NV_ENC_ERR_INVALID_DEVICE,           AVERROR(EINVAL),  "invalid device"           },
89     { NV_ENC_ERR_DEVICE_NOT_EXIST,         AVERROR(EIO),     "device does not exist"    },
90     { NV_ENC_ERR_INVALID_PTR,              AVERROR(EFAULT),  "invalid ptr"              },
91     { NV_ENC_ERR_INVALID_EVENT,            AVERROR(EINVAL),  "invalid event"            },
92     { NV_ENC_ERR_INVALID_PARAM,            AVERROR(EINVAL),  "invalid param"            },
93     { NV_ENC_ERR_INVALID_CALL,             AVERROR(EINVAL),  "invalid call"             },
94     { NV_ENC_ERR_OUT_OF_MEMORY,            AVERROR(ENOMEM),  "out of memory"            },
95     { NV_ENC_ERR_ENCODER_NOT_INITIALIZED,  AVERROR(EINVAL),  "encoder not initialized"  },
96     { NV_ENC_ERR_UNSUPPORTED_PARAM,        AVERROR(ENOSYS),  "unsupported param"        },
97     { NV_ENC_ERR_LOCK_BUSY,                AVERROR(EAGAIN),  "lock busy"                },
98     { NV_ENC_ERR_NOT_ENOUGH_BUFFER,        AVERROR_BUFFER_TOO_SMALL, "not enough buffer"},
99     { NV_ENC_ERR_INVALID_VERSION,          AVERROR(EINVAL),  "invalid version"          },
100     { NV_ENC_ERR_MAP_FAILED,               AVERROR(EIO),     "map failed"               },
101     { NV_ENC_ERR_NEED_MORE_INPUT,          AVERROR(EAGAIN),  "need more input"          },
102     { NV_ENC_ERR_ENCODER_BUSY,             AVERROR(EAGAIN),  "encoder busy"             },
103     { NV_ENC_ERR_EVENT_NOT_REGISTERD,      AVERROR(EBADF),   "event not registered"     },
104     { NV_ENC_ERR_GENERIC,                  AVERROR_UNKNOWN,  "generic error"            },
105     { NV_ENC_ERR_INCOMPATIBLE_CLIENT_KEY,  AVERROR(EINVAL),  "incompatible client key"  },
106     { NV_ENC_ERR_UNIMPLEMENTED,            AVERROR(ENOSYS),  "unimplemented"            },
107     { NV_ENC_ERR_RESOURCE_REGISTER_FAILED, AVERROR(EIO),     "resource register failed" },
108     { NV_ENC_ERR_RESOURCE_NOT_REGISTERED,  AVERROR(EBADF),   "resource not registered"  },
109     { NV_ENC_ERR_RESOURCE_NOT_MAPPED,      AVERROR(EBADF),   "resource not mapped"      },
110 };
111 
nvenc_map_error(NVENCSTATUS err,const char ** desc)112 static int nvenc_map_error(NVENCSTATUS err, const char **desc)
113 {
114     int i;
115     for (i = 0; i < FF_ARRAY_ELEMS(nvenc_errors); i++) {
116         if (nvenc_errors[i].nverr == err) {
117             if (desc)
118                 *desc = nvenc_errors[i].desc;
119             return nvenc_errors[i].averr;
120         }
121     }
122     if (desc)
123         *desc = "unknown error";
124     return AVERROR_UNKNOWN;
125 }
126 
nvenc_print_error(AVCodecContext * avctx,NVENCSTATUS err,const char * error_string)127 static int nvenc_print_error(AVCodecContext *avctx, NVENCSTATUS err,
128                              const char *error_string)
129 {
130     const char *desc;
131     const char *details = "(no details)";
132     int ret = nvenc_map_error(err, &desc);
133 
134 #ifdef NVENC_HAVE_GETLASTERRORSTRING
135     NvencContext *ctx = avctx->priv_data;
136     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
137 
138     if (p_nvenc && ctx->nvencoder)
139         details = p_nvenc->nvEncGetLastErrorString(ctx->nvencoder);
140 #endif
141 
142     av_log(avctx, AV_LOG_ERROR, "%s: %s (%d): %s\n", error_string, desc, err, details);
143 
144     return ret;
145 }
146 
nvenc_print_driver_requirement(AVCodecContext * avctx,int level)147 static void nvenc_print_driver_requirement(AVCodecContext *avctx, int level)
148 {
149 #if NVENCAPI_CHECK_VERSION(11, 1)
150     const char *minver = "(unknown)";
151 #elif NVENCAPI_CHECK_VERSION(11, 0)
152 # if defined(_WIN32) || defined(__CYGWIN__)
153     const char *minver = "456.71";
154 # else
155     const char *minver = "455.28";
156 # endif
157 #elif NVENCAPI_CHECK_VERSION(10, 0)
158 # if defined(_WIN32) || defined(__CYGWIN__)
159     const char *minver = "450.51";
160 # else
161     const char *minver = "445.87";
162 # endif
163 #elif NVENCAPI_CHECK_VERSION(9, 1)
164 # if defined(_WIN32) || defined(__CYGWIN__)
165     const char *minver = "436.15";
166 # else
167     const char *minver = "435.21";
168 # endif
169 #elif NVENCAPI_CHECK_VERSION(9, 0)
170 # if defined(_WIN32) || defined(__CYGWIN__)
171     const char *minver = "418.81";
172 # else
173     const char *minver = "418.30";
174 # endif
175 #elif NVENCAPI_CHECK_VERSION(8, 2)
176 # if defined(_WIN32) || defined(__CYGWIN__)
177     const char *minver = "397.93";
178 # else
179     const char *minver = "396.24";
180 #endif
181 #elif NVENCAPI_CHECK_VERSION(8, 1)
182 # if defined(_WIN32) || defined(__CYGWIN__)
183     const char *minver = "390.77";
184 # else
185     const char *minver = "390.25";
186 # endif
187 #else
188 # if defined(_WIN32) || defined(__CYGWIN__)
189     const char *minver = "378.66";
190 # else
191     const char *minver = "378.13";
192 # endif
193 #endif
194     av_log(avctx, level, "The minimum required Nvidia driver for nvenc is %s or newer\n", minver);
195 }
196 
nvenc_load_libraries(AVCodecContext * avctx)197 static av_cold int nvenc_load_libraries(AVCodecContext *avctx)
198 {
199     NvencContext *ctx            = avctx->priv_data;
200     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
201     NVENCSTATUS err;
202     uint32_t nvenc_max_ver;
203     int ret;
204 
205     ret = cuda_load_functions(&dl_fn->cuda_dl, avctx);
206     if (ret < 0)
207         return ret;
208 
209     ret = nvenc_load_functions(&dl_fn->nvenc_dl, avctx);
210     if (ret < 0) {
211         nvenc_print_driver_requirement(avctx, AV_LOG_ERROR);
212         return ret;
213     }
214 
215     err = dl_fn->nvenc_dl->NvEncodeAPIGetMaxSupportedVersion(&nvenc_max_ver);
216     if (err != NV_ENC_SUCCESS)
217         return nvenc_print_error(avctx, err, "Failed to query nvenc max version");
218 
219     av_log(avctx, AV_LOG_VERBOSE, "Loaded Nvenc version %d.%d\n", nvenc_max_ver >> 4, nvenc_max_ver & 0xf);
220 
221     if ((NVENCAPI_MAJOR_VERSION << 4 | NVENCAPI_MINOR_VERSION) > nvenc_max_ver) {
222         av_log(avctx, AV_LOG_ERROR, "Driver does not support the required nvenc API version. "
223                "Required: %d.%d Found: %d.%d\n",
224                NVENCAPI_MAJOR_VERSION, NVENCAPI_MINOR_VERSION,
225                nvenc_max_ver >> 4, nvenc_max_ver & 0xf);
226         nvenc_print_driver_requirement(avctx, AV_LOG_ERROR);
227         return AVERROR(ENOSYS);
228     }
229 
230     dl_fn->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER;
231 
232     err = dl_fn->nvenc_dl->NvEncodeAPICreateInstance(&dl_fn->nvenc_funcs);
233     if (err != NV_ENC_SUCCESS)
234         return nvenc_print_error(avctx, err, "Failed to create nvenc instance");
235 
236     av_log(avctx, AV_LOG_VERBOSE, "Nvenc initialized successfully\n");
237 
238     return 0;
239 }
240 
nvenc_push_context(AVCodecContext * avctx)241 static int nvenc_push_context(AVCodecContext *avctx)
242 {
243     NvencContext *ctx            = avctx->priv_data;
244     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
245 
246     if (ctx->d3d11_device)
247         return 0;
248 
249     return CHECK_CU(dl_fn->cuda_dl->cuCtxPushCurrent(ctx->cu_context));
250 }
251 
nvenc_pop_context(AVCodecContext * avctx)252 static int nvenc_pop_context(AVCodecContext *avctx)
253 {
254     NvencContext *ctx            = avctx->priv_data;
255     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
256     CUcontext dummy;
257 
258     if (ctx->d3d11_device)
259         return 0;
260 
261     return CHECK_CU(dl_fn->cuda_dl->cuCtxPopCurrent(&dummy));
262 }
263 
nvenc_open_session(AVCodecContext * avctx)264 static av_cold int nvenc_open_session(AVCodecContext *avctx)
265 {
266     NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS params = { 0 };
267     NvencContext *ctx = avctx->priv_data;
268     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
269     NVENCSTATUS ret;
270 
271     params.version    = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;
272     params.apiVersion = NVENCAPI_VERSION;
273     if (ctx->d3d11_device) {
274         params.device     = ctx->d3d11_device;
275         params.deviceType = NV_ENC_DEVICE_TYPE_DIRECTX;
276     } else {
277         params.device     = ctx->cu_context;
278         params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
279     }
280 
281     ret = p_nvenc->nvEncOpenEncodeSessionEx(&params, &ctx->nvencoder);
282     if (ret != NV_ENC_SUCCESS) {
283         ctx->nvencoder = NULL;
284         return nvenc_print_error(avctx, ret, "OpenEncodeSessionEx failed");
285     }
286 
287     return 0;
288 }
289 
nvenc_check_codec_support(AVCodecContext * avctx)290 static int nvenc_check_codec_support(AVCodecContext *avctx)
291 {
292     NvencContext *ctx                    = avctx->priv_data;
293     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
294     int i, ret, count = 0;
295     GUID *guids = NULL;
296 
297     ret = p_nvenc->nvEncGetEncodeGUIDCount(ctx->nvencoder, &count);
298 
299     if (ret != NV_ENC_SUCCESS || !count)
300         return AVERROR(ENOSYS);
301 
302     guids = av_malloc(count * sizeof(GUID));
303     if (!guids)
304         return AVERROR(ENOMEM);
305 
306     ret = p_nvenc->nvEncGetEncodeGUIDs(ctx->nvencoder, guids, count, &count);
307     if (ret != NV_ENC_SUCCESS) {
308         ret = AVERROR(ENOSYS);
309         goto fail;
310     }
311 
312     ret = AVERROR(ENOSYS);
313     for (i = 0; i < count; i++) {
314         if (!memcmp(&guids[i], &ctx->init_encode_params.encodeGUID, sizeof(*guids))) {
315             ret = 0;
316             break;
317         }
318     }
319 
320 fail:
321     av_free(guids);
322 
323     return ret;
324 }
325 
nvenc_check_cap(AVCodecContext * avctx,NV_ENC_CAPS cap)326 static int nvenc_check_cap(AVCodecContext *avctx, NV_ENC_CAPS cap)
327 {
328     NvencContext *ctx = avctx->priv_data;
329     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
330     NV_ENC_CAPS_PARAM params        = { 0 };
331     int ret, val = 0;
332 
333     params.version     = NV_ENC_CAPS_PARAM_VER;
334     params.capsToQuery = cap;
335 
336     ret = p_nvenc->nvEncGetEncodeCaps(ctx->nvencoder, ctx->init_encode_params.encodeGUID, &params, &val);
337 
338     if (ret == NV_ENC_SUCCESS)
339         return val;
340     return 0;
341 }
342 
nvenc_check_capabilities(AVCodecContext * avctx)343 static int nvenc_check_capabilities(AVCodecContext *avctx)
344 {
345     NvencContext *ctx = avctx->priv_data;
346     int ret;
347 
348     ret = nvenc_check_codec_support(avctx);
349     if (ret < 0) {
350         av_log(avctx, AV_LOG_WARNING, "Codec not supported\n");
351         return ret;
352     }
353 
354     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_YUV444_ENCODE);
355     if (IS_YUV444(ctx->data_pix_fmt) && ret <= 0) {
356         av_log(avctx, AV_LOG_WARNING, "YUV444P not supported\n");
357         return AVERROR(ENOSYS);
358     }
359 
360     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_LOSSLESS_ENCODE);
361     if (ctx->preset >= PRESET_LOSSLESS_DEFAULT && ret <= 0) {
362         av_log(avctx, AV_LOG_WARNING, "Lossless encoding not supported\n");
363         return AVERROR(ENOSYS);
364     }
365 
366     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_WIDTH_MAX);
367     if (ret < avctx->width) {
368         av_log(avctx, AV_LOG_WARNING, "Width %d exceeds %d\n",
369                avctx->width, ret);
370         return AVERROR(ENOSYS);
371     }
372 
373     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_HEIGHT_MAX);
374     if (ret < avctx->height) {
375         av_log(avctx, AV_LOG_WARNING, "Height %d exceeds %d\n",
376                avctx->height, ret);
377         return AVERROR(ENOSYS);
378     }
379 
380     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_NUM_MAX_BFRAMES);
381     if (ret < avctx->max_b_frames) {
382         av_log(avctx, AV_LOG_WARNING, "Max B-frames %d exceed %d\n",
383                avctx->max_b_frames, ret);
384 
385         return AVERROR(ENOSYS);
386     }
387 
388     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_FIELD_ENCODING);
389     if (ret < 1 && avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
390         av_log(avctx, AV_LOG_WARNING,
391                "Interlaced encoding is not supported. Supported level: %d\n",
392                ret);
393         return AVERROR(ENOSYS);
394     }
395 
396     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_10BIT_ENCODE);
397     if (IS_10BIT(ctx->data_pix_fmt) && ret <= 0) {
398         av_log(avctx, AV_LOG_WARNING, "10 bit encode not supported\n");
399         return AVERROR(ENOSYS);
400     }
401 
402     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_LOOKAHEAD);
403     if (ctx->rc_lookahead > 0 && ret <= 0) {
404         av_log(avctx, AV_LOG_WARNING, "RC lookahead not supported\n");
405         return AVERROR(ENOSYS);
406     }
407 
408     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_TEMPORAL_AQ);
409     if (ctx->temporal_aq > 0 && ret <= 0) {
410         av_log(avctx, AV_LOG_WARNING, "Temporal AQ not supported\n");
411         return AVERROR(ENOSYS);
412     }
413 
414     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_WEIGHTED_PREDICTION);
415     if (ctx->weighted_pred > 0 && ret <= 0) {
416         av_log (avctx, AV_LOG_WARNING, "Weighted Prediction not supported\n");
417         return AVERROR(ENOSYS);
418     }
419 
420     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_CABAC);
421     if (ctx->coder == NV_ENC_H264_ENTROPY_CODING_MODE_CABAC && ret <= 0) {
422         av_log(avctx, AV_LOG_WARNING, "CABAC entropy coding not supported\n");
423         return AVERROR(ENOSYS);
424     }
425 
426 #ifdef NVENC_HAVE_BFRAME_REF_MODE
427     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_BFRAME_REF_MODE);
428     if (ctx->b_ref_mode == NV_ENC_BFRAME_REF_MODE_EACH && ret != 1 && ret != 3) {
429         av_log(avctx, AV_LOG_WARNING, "Each B frame as reference is not supported\n");
430         return AVERROR(ENOSYS);
431     } else if (ctx->b_ref_mode != NV_ENC_BFRAME_REF_MODE_DISABLED && ret == 0) {
432         av_log(avctx, AV_LOG_WARNING, "B frames as references are not supported\n");
433         return AVERROR(ENOSYS);
434     }
435 #else
436     if (ctx->b_ref_mode != 0) {
437         av_log(avctx, AV_LOG_WARNING, "B frames as references need SDK 8.1 at build time\n");
438         return AVERROR(ENOSYS);
439     }
440 #endif
441 
442 #ifdef NVENC_HAVE_MULTIPLE_REF_FRAMES
443     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_MULTIPLE_REF_FRAMES);
444     if(avctx->refs != NV_ENC_NUM_REF_FRAMES_AUTOSELECT && ret <= 0) {
445         av_log(avctx, AV_LOG_WARNING, "Multiple reference frames are not supported by the device\n");
446         return AVERROR(ENOSYS);
447     }
448 #else
449     if(avctx->refs != 0) {
450         av_log(avctx, AV_LOG_WARNING, "Multiple reference frames need SDK 9.1 at build time\n");
451         return AVERROR(ENOSYS);
452     }
453 #endif
454 
455     ctx->support_dyn_bitrate = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_DYN_BITRATE_CHANGE);
456 
457     return 0;
458 }
459 
nvenc_check_device(AVCodecContext * avctx,int idx)460 static av_cold int nvenc_check_device(AVCodecContext *avctx, int idx)
461 {
462     NvencContext *ctx = avctx->priv_data;
463     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
464     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
465     char name[128] = { 0};
466     int major, minor, ret;
467     CUdevice cu_device;
468     int loglevel = AV_LOG_VERBOSE;
469 
470     if (ctx->device == LIST_DEVICES)
471         loglevel = AV_LOG_INFO;
472 
473     ret = CHECK_CU(dl_fn->cuda_dl->cuDeviceGet(&cu_device, idx));
474     if (ret < 0)
475         return ret;
476 
477     ret = CHECK_CU(dl_fn->cuda_dl->cuDeviceGetName(name, sizeof(name), cu_device));
478     if (ret < 0)
479         return ret;
480 
481     ret = CHECK_CU(dl_fn->cuda_dl->cuDeviceComputeCapability(&major, &minor, cu_device));
482     if (ret < 0)
483         return ret;
484 
485     av_log(avctx, loglevel, "[ GPU #%d - < %s > has Compute SM %d.%d ]\n", idx, name, major, minor);
486     if (((major << 4) | minor) < NVENC_CAP) {
487         av_log(avctx, loglevel, "does not support NVENC\n");
488         goto fail;
489     }
490 
491     if (ctx->device != idx && ctx->device != ANY_DEVICE)
492         return -1;
493 
494     ret = CHECK_CU(dl_fn->cuda_dl->cuCtxCreate(&ctx->cu_context_internal, 0, cu_device));
495     if (ret < 0)
496         goto fail;
497 
498     ctx->cu_context = ctx->cu_context_internal;
499     ctx->cu_stream = NULL;
500 
501     if ((ret = nvenc_pop_context(avctx)) < 0)
502         goto fail2;
503 
504     if ((ret = nvenc_open_session(avctx)) < 0)
505         goto fail2;
506 
507     if ((ret = nvenc_check_capabilities(avctx)) < 0)
508         goto fail3;
509 
510     av_log(avctx, loglevel, "supports NVENC\n");
511 
512     dl_fn->nvenc_device_count++;
513 
514     if (ctx->device == idx || ctx->device == ANY_DEVICE)
515         return 0;
516 
517 fail3:
518     if ((ret = nvenc_push_context(avctx)) < 0)
519         return ret;
520 
521     p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);
522     ctx->nvencoder = NULL;
523 
524     if ((ret = nvenc_pop_context(avctx)) < 0)
525         return ret;
526 
527 fail2:
528     CHECK_CU(dl_fn->cuda_dl->cuCtxDestroy(ctx->cu_context_internal));
529     ctx->cu_context_internal = NULL;
530 
531 fail:
532     return AVERROR(ENOSYS);
533 }
534 
nvenc_setup_device(AVCodecContext * avctx)535 static av_cold int nvenc_setup_device(AVCodecContext *avctx)
536 {
537     NvencContext *ctx            = avctx->priv_data;
538     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
539 
540     switch (avctx->codec->id) {
541     case AV_CODEC_ID_H264:
542         ctx->init_encode_params.encodeGUID = NV_ENC_CODEC_H264_GUID;
543         break;
544     case AV_CODEC_ID_HEVC:
545         ctx->init_encode_params.encodeGUID = NV_ENC_CODEC_HEVC_GUID;
546         break;
547     default:
548         return AVERROR_BUG;
549     }
550 
551     if (avctx->pix_fmt == AV_PIX_FMT_CUDA || avctx->pix_fmt == AV_PIX_FMT_D3D11 || avctx->hw_frames_ctx || avctx->hw_device_ctx) {
552         AVHWFramesContext   *frames_ctx;
553         AVHWDeviceContext   *hwdev_ctx;
554         AVCUDADeviceContext *cuda_device_hwctx = NULL;
555 #if CONFIG_D3D11VA
556         AVD3D11VADeviceContext *d3d11_device_hwctx = NULL;
557 #endif
558         int ret;
559 
560         if (avctx->hw_frames_ctx) {
561             frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
562             if (frames_ctx->format == AV_PIX_FMT_CUDA)
563                 cuda_device_hwctx = frames_ctx->device_ctx->hwctx;
564 #if CONFIG_D3D11VA
565             else if (frames_ctx->format == AV_PIX_FMT_D3D11)
566                 d3d11_device_hwctx = frames_ctx->device_ctx->hwctx;
567 #endif
568             else
569                 return AVERROR(EINVAL);
570         } else if (avctx->hw_device_ctx) {
571             hwdev_ctx = (AVHWDeviceContext*)avctx->hw_device_ctx->data;
572             if (hwdev_ctx->type == AV_HWDEVICE_TYPE_CUDA)
573                 cuda_device_hwctx = hwdev_ctx->hwctx;
574 #if CONFIG_D3D11VA
575             else if (hwdev_ctx->type == AV_HWDEVICE_TYPE_D3D11VA)
576                 d3d11_device_hwctx = hwdev_ctx->hwctx;
577 #endif
578             else
579                 return AVERROR(EINVAL);
580         } else {
581             return AVERROR(EINVAL);
582         }
583 
584         if (cuda_device_hwctx) {
585             ctx->cu_context = cuda_device_hwctx->cuda_ctx;
586             ctx->cu_stream = cuda_device_hwctx->stream;
587         }
588 #if CONFIG_D3D11VA
589         else if (d3d11_device_hwctx) {
590             ctx->d3d11_device = d3d11_device_hwctx->device;
591             ID3D11Device_AddRef(ctx->d3d11_device);
592         }
593 #endif
594 
595         ret = nvenc_open_session(avctx);
596         if (ret < 0)
597             return ret;
598 
599         ret = nvenc_check_capabilities(avctx);
600         if (ret < 0) {
601             av_log(avctx, AV_LOG_FATAL, "Provided device doesn't support required NVENC features\n");
602             return ret;
603         }
604     } else {
605         int i, nb_devices = 0;
606 
607         if (CHECK_CU(dl_fn->cuda_dl->cuInit(0)) < 0)
608             return AVERROR_UNKNOWN;
609 
610         if (CHECK_CU(dl_fn->cuda_dl->cuDeviceGetCount(&nb_devices)) < 0)
611             return AVERROR_UNKNOWN;
612 
613         if (!nb_devices) {
614             av_log(avctx, AV_LOG_FATAL, "No CUDA capable devices found\n");
615                 return AVERROR_EXTERNAL;
616         }
617 
618         av_log(avctx, AV_LOG_VERBOSE, "%d CUDA capable devices found\n", nb_devices);
619 
620         dl_fn->nvenc_device_count = 0;
621         for (i = 0; i < nb_devices; ++i) {
622             if ((nvenc_check_device(avctx, i)) >= 0 && ctx->device != LIST_DEVICES)
623                 return 0;
624         }
625 
626         if (ctx->device == LIST_DEVICES)
627             return AVERROR_EXIT;
628 
629         if (!dl_fn->nvenc_device_count) {
630             av_log(avctx, AV_LOG_FATAL, "No capable devices found\n");
631             return AVERROR_EXTERNAL;
632         }
633 
634         av_log(avctx, AV_LOG_FATAL, "Requested GPU %d, but only %d GPUs are available!\n", ctx->device, nb_devices);
635         return AVERROR(EINVAL);
636     }
637 
638     return 0;
639 }
640 
641 typedef struct GUIDTuple {
642     const GUID guid;
643     int flags;
644 } GUIDTuple;
645 
646 #define PRESET_ALIAS(alias, name, ...) \
647     [PRESET_ ## alias] = { NV_ENC_PRESET_ ## name ## _GUID, __VA_ARGS__ }
648 
649 #define PRESET(name, ...) PRESET_ALIAS(name, name, __VA_ARGS__)
650 
nvenc_map_preset(NvencContext * ctx)651 static void nvenc_map_preset(NvencContext *ctx)
652 {
653     GUIDTuple presets[] = {
654 #ifdef NVENC_HAVE_NEW_PRESETS
655         PRESET(P1),
656         PRESET(P2),
657         PRESET(P3),
658         PRESET(P4),
659         PRESET(P5),
660         PRESET(P6),
661         PRESET(P7),
662         PRESET_ALIAS(SLOW,   P7, NVENC_TWO_PASSES),
663         PRESET_ALIAS(MEDIUM, P4, NVENC_ONE_PASS),
664         PRESET_ALIAS(FAST,   P1, NVENC_ONE_PASS),
665         // Compat aliases
666         PRESET_ALIAS(DEFAULT,             P4, NVENC_DEPRECATED_PRESET),
667         PRESET_ALIAS(HP,                  P1, NVENC_DEPRECATED_PRESET),
668         PRESET_ALIAS(HQ,                  P7, NVENC_DEPRECATED_PRESET),
669         PRESET_ALIAS(BD,                  P5, NVENC_DEPRECATED_PRESET),
670         PRESET_ALIAS(LOW_LATENCY_DEFAULT, P4, NVENC_DEPRECATED_PRESET | NVENC_LOWLATENCY),
671         PRESET_ALIAS(LOW_LATENCY_HP,      P1, NVENC_DEPRECATED_PRESET | NVENC_LOWLATENCY),
672         PRESET_ALIAS(LOW_LATENCY_HQ,      P7, NVENC_DEPRECATED_PRESET | NVENC_LOWLATENCY),
673         PRESET_ALIAS(LOSSLESS_DEFAULT,    P4, NVENC_DEPRECATED_PRESET | NVENC_LOSSLESS),
674         PRESET_ALIAS(LOSSLESS_HP,         P1, NVENC_DEPRECATED_PRESET | NVENC_LOSSLESS),
675 #else
676         PRESET(DEFAULT),
677         PRESET(HP),
678         PRESET(HQ),
679         PRESET(BD),
680         PRESET_ALIAS(SLOW,   HQ,    NVENC_TWO_PASSES),
681         PRESET_ALIAS(MEDIUM, HQ,    NVENC_ONE_PASS),
682         PRESET_ALIAS(FAST,   HP,    NVENC_ONE_PASS),
683         PRESET(LOW_LATENCY_DEFAULT, NVENC_LOWLATENCY),
684         PRESET(LOW_LATENCY_HP,      NVENC_LOWLATENCY),
685         PRESET(LOW_LATENCY_HQ,      NVENC_LOWLATENCY),
686         PRESET(LOSSLESS_DEFAULT,    NVENC_LOSSLESS),
687         PRESET(LOSSLESS_HP,         NVENC_LOSSLESS),
688 #endif
689     };
690 
691     GUIDTuple *t = &presets[ctx->preset];
692 
693     ctx->init_encode_params.presetGUID = t->guid;
694     ctx->flags = t->flags;
695 }
696 
697 #undef PRESET
698 #undef PRESET_ALIAS
699 
set_constqp(AVCodecContext * avctx)700 static av_cold void set_constqp(AVCodecContext *avctx)
701 {
702     NvencContext *ctx = avctx->priv_data;
703     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
704 
705     rc->rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
706 
707     if (ctx->init_qp_p >= 0) {
708         rc->constQP.qpInterP = ctx->init_qp_p;
709         if (ctx->init_qp_i >= 0 && ctx->init_qp_b >= 0) {
710             rc->constQP.qpIntra = ctx->init_qp_i;
711             rc->constQP.qpInterB = ctx->init_qp_b;
712         } else if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
713             rc->constQP.qpIntra = av_clip(
714                 rc->constQP.qpInterP * fabs(avctx->i_quant_factor) + avctx->i_quant_offset + 0.5, 0, 51);
715             rc->constQP.qpInterB = av_clip(
716                 rc->constQP.qpInterP * fabs(avctx->b_quant_factor) + avctx->b_quant_offset + 0.5, 0, 51);
717         } else {
718             rc->constQP.qpIntra = rc->constQP.qpInterP;
719             rc->constQP.qpInterB = rc->constQP.qpInterP;
720         }
721     } else if (ctx->cqp >= 0) {
722         rc->constQP.qpInterP = rc->constQP.qpInterB = rc->constQP.qpIntra = ctx->cqp;
723         if (avctx->b_quant_factor != 0.0)
724             rc->constQP.qpInterB = av_clip(ctx->cqp * fabs(avctx->b_quant_factor) + avctx->b_quant_offset + 0.5, 0, 51);
725         if (avctx->i_quant_factor != 0.0)
726             rc->constQP.qpIntra = av_clip(ctx->cqp * fabs(avctx->i_quant_factor) + avctx->i_quant_offset + 0.5, 0, 51);
727     }
728 
729     avctx->qmin = -1;
730     avctx->qmax = -1;
731 }
732 
set_vbr(AVCodecContext * avctx)733 static av_cold void set_vbr(AVCodecContext *avctx)
734 {
735     NvencContext *ctx = avctx->priv_data;
736     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
737     int qp_inter_p;
738 
739     if (avctx->qmin >= 0 && avctx->qmax >= 0) {
740         rc->enableMinQP = 1;
741         rc->enableMaxQP = 1;
742 
743         rc->minQP.qpInterB = avctx->qmin;
744         rc->minQP.qpInterP = avctx->qmin;
745         rc->minQP.qpIntra  = avctx->qmin;
746 
747         rc->maxQP.qpInterB = avctx->qmax;
748         rc->maxQP.qpInterP = avctx->qmax;
749         rc->maxQP.qpIntra = avctx->qmax;
750 
751         qp_inter_p = (avctx->qmax + 3 * avctx->qmin) / 4; // biased towards Qmin
752     } else if (avctx->qmin >= 0) {
753         rc->enableMinQP = 1;
754 
755         rc->minQP.qpInterB = avctx->qmin;
756         rc->minQP.qpInterP = avctx->qmin;
757         rc->minQP.qpIntra = avctx->qmin;
758 
759         qp_inter_p = avctx->qmin;
760     } else {
761         qp_inter_p = 26; // default to 26
762     }
763 
764     rc->enableInitialRCQP = 1;
765 
766     if (ctx->init_qp_p < 0) {
767         rc->initialRCQP.qpInterP  = qp_inter_p;
768     } else {
769         rc->initialRCQP.qpInterP = ctx->init_qp_p;
770     }
771 
772     if (ctx->init_qp_i < 0) {
773         if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
774             rc->initialRCQP.qpIntra = av_clip(
775                 rc->initialRCQP.qpInterP * fabs(avctx->i_quant_factor) + avctx->i_quant_offset + 0.5, 0, 51);
776         } else {
777             rc->initialRCQP.qpIntra = rc->initialRCQP.qpInterP;
778         }
779     } else {
780         rc->initialRCQP.qpIntra = ctx->init_qp_i;
781     }
782 
783     if (ctx->init_qp_b < 0) {
784         if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
785             rc->initialRCQP.qpInterB = av_clip(
786                 rc->initialRCQP.qpInterP * fabs(avctx->b_quant_factor) + avctx->b_quant_offset + 0.5, 0, 51);
787         } else {
788             rc->initialRCQP.qpInterB = rc->initialRCQP.qpInterP;
789         }
790     } else {
791         rc->initialRCQP.qpInterB = ctx->init_qp_b;
792     }
793 }
794 
set_lossless(AVCodecContext * avctx)795 static av_cold void set_lossless(AVCodecContext *avctx)
796 {
797     NvencContext *ctx = avctx->priv_data;
798     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
799 
800     rc->rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
801     rc->constQP.qpInterB = 0;
802     rc->constQP.qpInterP = 0;
803     rc->constQP.qpIntra  = 0;
804 
805     avctx->qmin = -1;
806     avctx->qmax = -1;
807 }
808 
nvenc_override_rate_control(AVCodecContext * avctx)809 static void nvenc_override_rate_control(AVCodecContext *avctx)
810 {
811     NvencContext *ctx    = avctx->priv_data;
812     NV_ENC_RC_PARAMS *rc = &ctx->encode_config.rcParams;
813 
814     switch (ctx->rc) {
815     case NV_ENC_PARAMS_RC_CONSTQP:
816         set_constqp(avctx);
817         return;
818     case NV_ENC_PARAMS_RC_VBR_MINQP:
819         if (avctx->qmin < 0) {
820             av_log(avctx, AV_LOG_WARNING,
821                    "The variable bitrate rate-control requires "
822                    "the 'qmin' option set.\n");
823             set_vbr(avctx);
824             return;
825         }
826         /* fall through */
827     case NV_ENC_PARAMS_RC_VBR_HQ:
828     case NV_ENC_PARAMS_RC_VBR:
829         set_vbr(avctx);
830         break;
831     case NV_ENC_PARAMS_RC_CBR:
832     case NV_ENC_PARAMS_RC_CBR_HQ:
833     case NV_ENC_PARAMS_RC_CBR_LOWDELAY_HQ:
834         break;
835     }
836 
837     rc->rateControlMode = ctx->rc;
838 }
839 
nvenc_recalc_surfaces(AVCodecContext * avctx)840 static av_cold int nvenc_recalc_surfaces(AVCodecContext *avctx)
841 {
842     NvencContext *ctx = avctx->priv_data;
843     // default minimum of 4 surfaces
844     // multiply by 2 for number of NVENCs on gpu (hardcode to 2)
845     // another multiply by 2 to avoid blocking next PBB group
846     int nb_surfaces = FFMAX(4, ctx->encode_config.frameIntervalP * 2 * 2);
847 
848     // lookahead enabled
849     if (ctx->rc_lookahead > 0) {
850         // +1 is to account for lkd_bound calculation later
851         // +4 is to allow sufficient pipelining with lookahead
852         nb_surfaces = FFMAX(1, FFMAX(nb_surfaces, ctx->rc_lookahead + ctx->encode_config.frameIntervalP + 1 + 4));
853         if (nb_surfaces > ctx->nb_surfaces && ctx->nb_surfaces > 0)
854         {
855             av_log(avctx, AV_LOG_WARNING,
856                    "Defined rc_lookahead requires more surfaces, "
857                    "increasing used surfaces %d -> %d\n", ctx->nb_surfaces, nb_surfaces);
858         }
859         ctx->nb_surfaces = FFMAX(nb_surfaces, ctx->nb_surfaces);
860     } else {
861         if (ctx->encode_config.frameIntervalP > 1 && ctx->nb_surfaces < nb_surfaces && ctx->nb_surfaces > 0)
862         {
863             av_log(avctx, AV_LOG_WARNING,
864                    "Defined b-frame requires more surfaces, "
865                    "increasing used surfaces %d -> %d\n", ctx->nb_surfaces, nb_surfaces);
866             ctx->nb_surfaces = FFMAX(ctx->nb_surfaces, nb_surfaces);
867         }
868         else if (ctx->nb_surfaces <= 0)
869             ctx->nb_surfaces = nb_surfaces;
870         // otherwise use user specified value
871     }
872 
873     ctx->nb_surfaces = FFMAX(1, FFMIN(MAX_REGISTERED_FRAMES, ctx->nb_surfaces));
874     ctx->async_depth = FFMIN(ctx->async_depth, ctx->nb_surfaces - 1);
875 
876     return 0;
877 }
878 
nvenc_setup_rate_control(AVCodecContext * avctx)879 static av_cold void nvenc_setup_rate_control(AVCodecContext *avctx)
880 {
881     NvencContext *ctx = avctx->priv_data;
882 
883     if (avctx->global_quality > 0)
884         av_log(avctx, AV_LOG_WARNING, "Using global_quality with nvenc is deprecated. Use qp instead.\n");
885 
886     if (ctx->cqp < 0 && avctx->global_quality > 0)
887         ctx->cqp = avctx->global_quality;
888 
889     if (avctx->bit_rate > 0) {
890         ctx->encode_config.rcParams.averageBitRate = avctx->bit_rate;
891     } else if (ctx->encode_config.rcParams.averageBitRate > 0) {
892         ctx->encode_config.rcParams.maxBitRate = ctx->encode_config.rcParams.averageBitRate;
893     }
894 
895     if (avctx->rc_max_rate > 0)
896         ctx->encode_config.rcParams.maxBitRate = avctx->rc_max_rate;
897 
898 #ifdef NVENC_HAVE_MULTIPASS
899     ctx->encode_config.rcParams.multiPass = ctx->multipass;
900 
901     if (ctx->flags & NVENC_ONE_PASS)
902         ctx->encode_config.rcParams.multiPass = NV_ENC_MULTI_PASS_DISABLED;
903     if (ctx->flags & NVENC_TWO_PASSES || ctx->twopass > 0)
904         ctx->encode_config.rcParams.multiPass = NV_ENC_TWO_PASS_FULL_RESOLUTION;
905 
906     if (ctx->rc < 0) {
907         if (ctx->cbr) {
908             ctx->rc = NV_ENC_PARAMS_RC_CBR;
909         } else if (ctx->cqp >= 0) {
910             ctx->rc = NV_ENC_PARAMS_RC_CONSTQP;
911         } else if (ctx->quality >= 0.0f) {
912             ctx->rc = NV_ENC_PARAMS_RC_VBR;
913         }
914     }
915 #else
916     if (ctx->rc < 0) {
917         if (ctx->flags & NVENC_ONE_PASS)
918             ctx->twopass = 0;
919         if (ctx->flags & NVENC_TWO_PASSES)
920             ctx->twopass = 1;
921 
922         if (ctx->twopass < 0)
923             ctx->twopass = (ctx->flags & NVENC_LOWLATENCY) != 0;
924 
925         if (ctx->cbr) {
926             if (ctx->twopass) {
927                 ctx->rc = NV_ENC_PARAMS_RC_CBR_LOWDELAY_HQ;
928             } else {
929                 ctx->rc = NV_ENC_PARAMS_RC_CBR;
930             }
931         } else if (ctx->cqp >= 0) {
932             ctx->rc = NV_ENC_PARAMS_RC_CONSTQP;
933         } else if (ctx->twopass) {
934             ctx->rc = NV_ENC_PARAMS_RC_VBR_HQ;
935         } else if (avctx->qmin >= 0 && avctx->qmax >= 0) {
936             ctx->rc = NV_ENC_PARAMS_RC_VBR_MINQP;
937         }
938     }
939 #endif
940 
941     if (ctx->rc >= 0 && ctx->rc & RC_MODE_DEPRECATED) {
942         av_log(avctx, AV_LOG_WARNING, "Specified rc mode is deprecated.\n");
943         av_log(avctx, AV_LOG_WARNING, "Use -rc constqp/cbr/vbr, -tune and -multipass instead.\n");
944 
945         ctx->rc &= ~RC_MODE_DEPRECATED;
946     }
947 
948 #ifdef NVENC_HAVE_LDKFS
949     if (ctx->ldkfs)
950          ctx->encode_config.rcParams.lowDelayKeyFrameScale = ctx->ldkfs;
951 #endif
952 
953     if (ctx->flags & NVENC_LOSSLESS) {
954         set_lossless(avctx);
955     } else if (ctx->rc >= 0) {
956         nvenc_override_rate_control(avctx);
957     } else {
958         ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR;
959         set_vbr(avctx);
960     }
961 
962     if (avctx->rc_buffer_size > 0) {
963         ctx->encode_config.rcParams.vbvBufferSize = avctx->rc_buffer_size;
964     } else if (ctx->encode_config.rcParams.averageBitRate > 0) {
965         avctx->rc_buffer_size = ctx->encode_config.rcParams.vbvBufferSize = 2 * ctx->encode_config.rcParams.averageBitRate;
966     }
967 
968     if (ctx->aq) {
969         ctx->encode_config.rcParams.enableAQ   = 1;
970         ctx->encode_config.rcParams.aqStrength = ctx->aq_strength;
971         av_log(avctx, AV_LOG_VERBOSE, "AQ enabled.\n");
972     }
973 
974     if (ctx->temporal_aq) {
975         ctx->encode_config.rcParams.enableTemporalAQ = 1;
976         av_log(avctx, AV_LOG_VERBOSE, "Temporal AQ enabled.\n");
977     }
978 
979     if (ctx->rc_lookahead > 0) {
980         int lkd_bound = FFMIN(ctx->nb_surfaces, ctx->async_depth) -
981                         ctx->encode_config.frameIntervalP - 4;
982 
983         if (lkd_bound < 0) {
984             av_log(avctx, AV_LOG_WARNING,
985                    "Lookahead not enabled. Increase buffer delay (-delay).\n");
986         } else {
987             ctx->encode_config.rcParams.enableLookahead = 1;
988             ctx->encode_config.rcParams.lookaheadDepth  = av_clip(ctx->rc_lookahead, 0, lkd_bound);
989             ctx->encode_config.rcParams.disableIadapt   = ctx->no_scenecut;
990             ctx->encode_config.rcParams.disableBadapt   = !ctx->b_adapt;
991             av_log(avctx, AV_LOG_VERBOSE,
992                    "Lookahead enabled: depth %d, scenecut %s, B-adapt %s.\n",
993                    ctx->encode_config.rcParams.lookaheadDepth,
994                    ctx->encode_config.rcParams.disableIadapt ? "disabled" : "enabled",
995                    ctx->encode_config.rcParams.disableBadapt ? "disabled" : "enabled");
996         }
997     }
998 
999     if (ctx->strict_gop) {
1000         ctx->encode_config.rcParams.strictGOPTarget = 1;
1001         av_log(avctx, AV_LOG_VERBOSE, "Strict GOP target enabled.\n");
1002     }
1003 
1004     if (ctx->nonref_p)
1005         ctx->encode_config.rcParams.enableNonRefP = 1;
1006 
1007     if (ctx->zerolatency)
1008         ctx->encode_config.rcParams.zeroReorderDelay = 1;
1009 
1010     if (ctx->quality) {
1011         //convert from float to fixed point 8.8
1012         int tmp_quality = (int)(ctx->quality * 256.0f);
1013         ctx->encode_config.rcParams.targetQuality = (uint8_t)(tmp_quality >> 8);
1014         ctx->encode_config.rcParams.targetQualityLSB = (uint8_t)(tmp_quality & 0xff);
1015 
1016         av_log(avctx, AV_LOG_VERBOSE, "CQ(%d) mode enabled.\n", tmp_quality);
1017 
1018         //CQ mode shall discard avg bitrate & honor max bitrate;
1019         ctx->encode_config.rcParams.averageBitRate = avctx->bit_rate = 0;
1020         ctx->encode_config.rcParams.maxBitRate = avctx->rc_max_rate;
1021     }
1022 }
1023 
nvenc_setup_h264_config(AVCodecContext * avctx)1024 static av_cold int nvenc_setup_h264_config(AVCodecContext *avctx)
1025 {
1026     NvencContext *ctx                      = avctx->priv_data;
1027     NV_ENC_CONFIG *cc                      = &ctx->encode_config;
1028     NV_ENC_CONFIG_H264 *h264               = &cc->encodeCodecConfig.h264Config;
1029     NV_ENC_CONFIG_H264_VUI_PARAMETERS *vui = &h264->h264VUIParameters;
1030 
1031     vui->colourMatrix = avctx->colorspace;
1032     vui->colourPrimaries = avctx->color_primaries;
1033     vui->transferCharacteristics = avctx->color_trc;
1034     vui->videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG
1035         || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ420P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ422P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ444P);
1036 
1037     vui->colourDescriptionPresentFlag =
1038         (avctx->colorspace != 2 || avctx->color_primaries != 2 || avctx->color_trc != 2);
1039 
1040     vui->videoSignalTypePresentFlag =
1041         (vui->colourDescriptionPresentFlag
1042         || vui->videoFormat != 5
1043         || vui->videoFullRangeFlag != 0);
1044 
1045     h264->sliceMode = 3;
1046     h264->sliceModeData = 1;
1047 
1048     h264->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
1049     h264->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
1050     h264->outputAUD     = ctx->aud;
1051 
1052     if (ctx->dpb_size >= 0) {
1053         /* 0 means "let the hardware decide" */
1054         h264->maxNumRefFrames = ctx->dpb_size;
1055     }
1056     if (avctx->gop_size >= 0) {
1057         h264->idrPeriod = cc->gopLength;
1058     }
1059 
1060     if (IS_CBR(cc->rcParams.rateControlMode)) {
1061         h264->outputBufferingPeriodSEI = 1;
1062     }
1063 
1064     h264->outputPictureTimingSEI = 1;
1065 
1066     if (cc->rcParams.rateControlMode == NV_ENC_PARAMS_RC_CBR_LOWDELAY_HQ ||
1067         cc->rcParams.rateControlMode == NV_ENC_PARAMS_RC_CBR_HQ ||
1068         cc->rcParams.rateControlMode == NV_ENC_PARAMS_RC_VBR_HQ) {
1069         h264->adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE;
1070         h264->fmoMode = NV_ENC_H264_FMO_DISABLE;
1071     }
1072 
1073     if (ctx->flags & NVENC_LOSSLESS) {
1074         h264->qpPrimeYZeroTransformBypassFlag = 1;
1075     } else {
1076         switch(ctx->profile) {
1077         case NV_ENC_H264_PROFILE_BASELINE:
1078             cc->profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;
1079             avctx->profile = FF_PROFILE_H264_BASELINE;
1080             break;
1081         case NV_ENC_H264_PROFILE_MAIN:
1082             cc->profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;
1083             avctx->profile = FF_PROFILE_H264_MAIN;
1084             break;
1085         case NV_ENC_H264_PROFILE_HIGH:
1086             cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
1087             avctx->profile = FF_PROFILE_H264_HIGH;
1088             break;
1089         case NV_ENC_H264_PROFILE_HIGH_444P:
1090             cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
1091             avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;
1092             break;
1093         }
1094     }
1095 
1096     // force setting profile as high444p if input is AV_PIX_FMT_YUV444P
1097     if (ctx->data_pix_fmt == AV_PIX_FMT_YUV444P) {
1098         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
1099         avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;
1100     }
1101 
1102     h264->chromaFormatIDC = avctx->profile == FF_PROFILE_H264_HIGH_444_PREDICTIVE ? 3 : 1;
1103 
1104     h264->level = ctx->level;
1105 
1106     if (ctx->coder >= 0)
1107         h264->entropyCodingMode = ctx->coder;
1108 
1109 #ifdef NVENC_HAVE_BFRAME_REF_MODE
1110     h264->useBFramesAsRef = ctx->b_ref_mode;
1111 #endif
1112 
1113 #ifdef NVENC_HAVE_MULTIPLE_REF_FRAMES
1114     h264->numRefL0 = avctx->refs;
1115     h264->numRefL1 = avctx->refs;
1116 #endif
1117 
1118     return 0;
1119 }
1120 
nvenc_setup_hevc_config(AVCodecContext * avctx)1121 static av_cold int nvenc_setup_hevc_config(AVCodecContext *avctx)
1122 {
1123     NvencContext *ctx                      = avctx->priv_data;
1124     NV_ENC_CONFIG *cc                      = &ctx->encode_config;
1125     NV_ENC_CONFIG_HEVC *hevc               = &cc->encodeCodecConfig.hevcConfig;
1126     NV_ENC_CONFIG_HEVC_VUI_PARAMETERS *vui = &hevc->hevcVUIParameters;
1127 
1128     vui->colourMatrix = avctx->colorspace;
1129     vui->colourPrimaries = avctx->color_primaries;
1130     vui->transferCharacteristics = avctx->color_trc;
1131     vui->videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG
1132         || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ420P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ422P || ctx->data_pix_fmt == AV_PIX_FMT_YUVJ444P);
1133 
1134     vui->colourDescriptionPresentFlag =
1135         (avctx->colorspace != 2 || avctx->color_primaries != 2 || avctx->color_trc != 2);
1136 
1137     vui->videoSignalTypePresentFlag =
1138         (vui->colourDescriptionPresentFlag
1139         || vui->videoFormat != 5
1140         || vui->videoFullRangeFlag != 0);
1141 
1142     hevc->sliceMode = 3;
1143     hevc->sliceModeData = 1;
1144 
1145     hevc->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
1146     hevc->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
1147     hevc->outputAUD     = ctx->aud;
1148 
1149     if (ctx->dpb_size >= 0) {
1150         /* 0 means "let the hardware decide" */
1151         hevc->maxNumRefFramesInDPB = ctx->dpb_size;
1152     }
1153     if (avctx->gop_size >= 0) {
1154         hevc->idrPeriod = cc->gopLength;
1155     }
1156 
1157     if (IS_CBR(cc->rcParams.rateControlMode)) {
1158         hevc->outputBufferingPeriodSEI = 1;
1159     }
1160 
1161     hevc->outputPictureTimingSEI = 1;
1162 
1163     switch (ctx->profile) {
1164     case NV_ENC_HEVC_PROFILE_MAIN:
1165         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID;
1166         avctx->profile  = FF_PROFILE_HEVC_MAIN;
1167         break;
1168     case NV_ENC_HEVC_PROFILE_MAIN_10:
1169         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN10_GUID;
1170         avctx->profile  = FF_PROFILE_HEVC_MAIN_10;
1171         break;
1172     case NV_ENC_HEVC_PROFILE_REXT:
1173         cc->profileGUID = NV_ENC_HEVC_PROFILE_FREXT_GUID;
1174         avctx->profile  = FF_PROFILE_HEVC_REXT;
1175         break;
1176     }
1177 
1178     // force setting profile as main10 if input is 10 bit
1179     if (IS_10BIT(ctx->data_pix_fmt)) {
1180         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN10_GUID;
1181         avctx->profile = FF_PROFILE_HEVC_MAIN_10;
1182     }
1183 
1184     // force setting profile as rext if input is yuv444
1185     if (IS_YUV444(ctx->data_pix_fmt)) {
1186         cc->profileGUID = NV_ENC_HEVC_PROFILE_FREXT_GUID;
1187         avctx->profile = FF_PROFILE_HEVC_REXT;
1188     }
1189 
1190     hevc->chromaFormatIDC = IS_YUV444(ctx->data_pix_fmt) ? 3 : 1;
1191 
1192     hevc->pixelBitDepthMinus8 = IS_10BIT(ctx->data_pix_fmt) ? 2 : 0;
1193 
1194     hevc->level = ctx->level;
1195 
1196     hevc->tier = ctx->tier;
1197 
1198 #ifdef NVENC_HAVE_HEVC_BFRAME_REF_MODE
1199     hevc->useBFramesAsRef = ctx->b_ref_mode;
1200 #endif
1201 
1202 #ifdef NVENC_HAVE_MULTIPLE_REF_FRAMES
1203     hevc->numRefL0 = avctx->refs;
1204     hevc->numRefL1 = avctx->refs;
1205 #endif
1206 
1207     return 0;
1208 }
1209 
nvenc_setup_codec_config(AVCodecContext * avctx)1210 static av_cold int nvenc_setup_codec_config(AVCodecContext *avctx)
1211 {
1212     switch (avctx->codec->id) {
1213     case AV_CODEC_ID_H264:
1214         return nvenc_setup_h264_config(avctx);
1215     case AV_CODEC_ID_HEVC:
1216         return nvenc_setup_hevc_config(avctx);
1217     /* Earlier switch/case will return if unknown codec is passed. */
1218     }
1219 
1220     return 0;
1221 }
1222 
compute_dar(AVCodecContext * avctx,int * dw,int * dh)1223 static void compute_dar(AVCodecContext *avctx, int *dw, int *dh) {
1224     int sw, sh;
1225 
1226     sw = avctx->width;
1227     sh = avctx->height;
1228 
1229     if (avctx->sample_aspect_ratio.num > 0 && avctx->sample_aspect_ratio.den > 0) {
1230         sw *= avctx->sample_aspect_ratio.num;
1231         sh *= avctx->sample_aspect_ratio.den;
1232     }
1233 
1234     av_reduce(dw, dh, sw, sh, 1024 * 1024);
1235 }
1236 
nvenc_setup_encoder(AVCodecContext * avctx)1237 static av_cold int nvenc_setup_encoder(AVCodecContext *avctx)
1238 {
1239     NvencContext *ctx = avctx->priv_data;
1240     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1241     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1242 
1243     NV_ENC_PRESET_CONFIG preset_config = { 0 };
1244     NVENCSTATUS nv_status = NV_ENC_SUCCESS;
1245     AVCPBProperties *cpb_props;
1246     int res = 0;
1247     int dw, dh;
1248 
1249     ctx->encode_config.version = NV_ENC_CONFIG_VER;
1250     ctx->init_encode_params.version = NV_ENC_INITIALIZE_PARAMS_VER;
1251 
1252     ctx->init_encode_params.encodeHeight = avctx->height;
1253     ctx->init_encode_params.encodeWidth = avctx->width;
1254 
1255     ctx->init_encode_params.encodeConfig = &ctx->encode_config;
1256 
1257     nvenc_map_preset(ctx);
1258 
1259     if (ctx->flags & NVENC_DEPRECATED_PRESET)
1260         av_log(avctx, AV_LOG_WARNING, "The selected preset is deprecated. Use p1 to p7 + -tune or fast/medium/slow.\n");
1261 
1262     preset_config.version = NV_ENC_PRESET_CONFIG_VER;
1263     preset_config.presetCfg.version = NV_ENC_CONFIG_VER;
1264 
1265 #ifdef NVENC_HAVE_NEW_PRESETS
1266     ctx->init_encode_params.tuningInfo = ctx->tuning_info;
1267 
1268     if (ctx->flags & NVENC_LOWLATENCY)
1269         ctx->init_encode_params.tuningInfo = NV_ENC_TUNING_INFO_LOW_LATENCY;
1270 
1271     nv_status = p_nvenc->nvEncGetEncodePresetConfigEx(ctx->nvencoder,
1272         ctx->init_encode_params.encodeGUID,
1273         ctx->init_encode_params.presetGUID,
1274         ctx->init_encode_params.tuningInfo,
1275         &preset_config);
1276 #else
1277     nv_status = p_nvenc->nvEncGetEncodePresetConfig(ctx->nvencoder,
1278         ctx->init_encode_params.encodeGUID,
1279         ctx->init_encode_params.presetGUID,
1280         &preset_config);
1281 #endif
1282     if (nv_status != NV_ENC_SUCCESS)
1283         return nvenc_print_error(avctx, nv_status, "Cannot get the preset configuration");
1284 
1285     memcpy(&ctx->encode_config, &preset_config.presetCfg, sizeof(ctx->encode_config));
1286 
1287     ctx->encode_config.version = NV_ENC_CONFIG_VER;
1288 
1289     compute_dar(avctx, &dw, &dh);
1290     ctx->init_encode_params.darHeight = dh;
1291     ctx->init_encode_params.darWidth = dw;
1292 
1293     if (avctx->framerate.num > 0 && avctx->framerate.den > 0) {
1294         ctx->init_encode_params.frameRateNum = avctx->framerate.num;
1295         ctx->init_encode_params.frameRateDen = avctx->framerate.den;
1296     } else {
1297         ctx->init_encode_params.frameRateNum = avctx->time_base.den;
1298         ctx->init_encode_params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
1299     }
1300 
1301     ctx->init_encode_params.enableEncodeAsync = 0;
1302     ctx->init_encode_params.enablePTD = 1;
1303 
1304 #ifdef NVENC_HAVE_NEW_PRESETS
1305     /* If lookahead isn't set from CLI, use value from preset.
1306      * P6 & P7 presets may enable lookahead for better quality.
1307      * */
1308     if (ctx->rc_lookahead == 0 && ctx->encode_config.rcParams.enableLookahead)
1309         ctx->rc_lookahead = ctx->encode_config.rcParams.lookaheadDepth;
1310 
1311     if (ctx->init_encode_params.tuningInfo == NV_ENC_TUNING_INFO_LOSSLESS)
1312         ctx->flags |= NVENC_LOSSLESS;
1313 #endif
1314 
1315     if (ctx->weighted_pred == 1)
1316         ctx->init_encode_params.enableWeightedPrediction = 1;
1317 
1318     if (ctx->bluray_compat) {
1319         ctx->aud = 1;
1320         ctx->dpb_size = FFMIN(FFMAX(avctx->refs, 0), 6);
1321         avctx->max_b_frames = FFMIN(avctx->max_b_frames, 3);
1322         switch (avctx->codec->id) {
1323         case AV_CODEC_ID_H264:
1324             /* maximum level depends on used resolution */
1325             break;
1326         case AV_CODEC_ID_HEVC:
1327             ctx->level = NV_ENC_LEVEL_HEVC_51;
1328             ctx->tier = NV_ENC_TIER_HEVC_HIGH;
1329             break;
1330         }
1331     }
1332 
1333     if (avctx->gop_size > 0) {
1334         if (avctx->max_b_frames >= 0) {
1335             /* 0 is intra-only, 1 is I/P only, 2 is one B-Frame, 3 two B-frames, and so on. */
1336             ctx->encode_config.frameIntervalP = avctx->max_b_frames + 1;
1337         }
1338 
1339         ctx->encode_config.gopLength = avctx->gop_size;
1340     } else if (avctx->gop_size == 0) {
1341         ctx->encode_config.frameIntervalP = 0;
1342         ctx->encode_config.gopLength = 1;
1343     }
1344 
1345     nvenc_recalc_surfaces(avctx);
1346 
1347     nvenc_setup_rate_control(avctx);
1348 
1349     if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
1350         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
1351     } else {
1352         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;
1353     }
1354 
1355     res = nvenc_setup_codec_config(avctx);
1356     if (res)
1357         return res;
1358 
1359     res = nvenc_push_context(avctx);
1360     if (res < 0)
1361         return res;
1362 
1363     nv_status = p_nvenc->nvEncInitializeEncoder(ctx->nvencoder, &ctx->init_encode_params);
1364     if (nv_status != NV_ENC_SUCCESS) {
1365         nvenc_pop_context(avctx);
1366         return nvenc_print_error(avctx, nv_status, "InitializeEncoder failed");
1367     }
1368 
1369 #ifdef NVENC_HAVE_CUSTREAM_PTR
1370     if (ctx->cu_context) {
1371         nv_status = p_nvenc->nvEncSetIOCudaStreams(ctx->nvencoder, &ctx->cu_stream, &ctx->cu_stream);
1372         if (nv_status != NV_ENC_SUCCESS) {
1373             nvenc_pop_context(avctx);
1374             return nvenc_print_error(avctx, nv_status, "SetIOCudaStreams failed");
1375         }
1376     }
1377 #endif
1378 
1379     res = nvenc_pop_context(avctx);
1380     if (res < 0)
1381         return res;
1382 
1383     if (ctx->encode_config.frameIntervalP > 1)
1384         avctx->has_b_frames = 2;
1385 
1386     if (ctx->encode_config.rcParams.averageBitRate > 0)
1387         avctx->bit_rate = ctx->encode_config.rcParams.averageBitRate;
1388 
1389     cpb_props = ff_add_cpb_side_data(avctx);
1390     if (!cpb_props)
1391         return AVERROR(ENOMEM);
1392     cpb_props->max_bitrate = ctx->encode_config.rcParams.maxBitRate;
1393     cpb_props->avg_bitrate = avctx->bit_rate;
1394     cpb_props->buffer_size = ctx->encode_config.rcParams.vbvBufferSize;
1395 
1396     return 0;
1397 }
1398 
nvenc_map_buffer_format(enum AVPixelFormat pix_fmt)1399 static NV_ENC_BUFFER_FORMAT nvenc_map_buffer_format(enum AVPixelFormat pix_fmt)
1400 {
1401     switch (pix_fmt) {
1402     case AV_PIX_FMT_YUV420P:
1403         return NV_ENC_BUFFER_FORMAT_YV12_PL;
1404     case AV_PIX_FMT_NV12:
1405         return NV_ENC_BUFFER_FORMAT_NV12_PL;
1406     case AV_PIX_FMT_P010:
1407     case AV_PIX_FMT_P016:
1408         return NV_ENC_BUFFER_FORMAT_YUV420_10BIT;
1409     case AV_PIX_FMT_YUV444P:
1410         return NV_ENC_BUFFER_FORMAT_YUV444_PL;
1411     case AV_PIX_FMT_YUV444P16:
1412         return NV_ENC_BUFFER_FORMAT_YUV444_10BIT;
1413     case AV_PIX_FMT_0RGB32:
1414         return NV_ENC_BUFFER_FORMAT_ARGB;
1415     case AV_PIX_FMT_0BGR32:
1416         return NV_ENC_BUFFER_FORMAT_ABGR;
1417     default:
1418         return NV_ENC_BUFFER_FORMAT_UNDEFINED;
1419     }
1420 }
1421 
nvenc_alloc_surface(AVCodecContext * avctx,int idx)1422 static av_cold int nvenc_alloc_surface(AVCodecContext *avctx, int idx)
1423 {
1424     NvencContext *ctx = avctx->priv_data;
1425     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1426     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1427     NvencSurface* tmp_surface = &ctx->surfaces[idx];
1428 
1429     NVENCSTATUS nv_status;
1430     NV_ENC_CREATE_BITSTREAM_BUFFER allocOut = { 0 };
1431     allocOut.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
1432 
1433     if (avctx->pix_fmt == AV_PIX_FMT_CUDA || avctx->pix_fmt == AV_PIX_FMT_D3D11) {
1434         ctx->surfaces[idx].in_ref = av_frame_alloc();
1435         if (!ctx->surfaces[idx].in_ref)
1436             return AVERROR(ENOMEM);
1437     } else {
1438         NV_ENC_CREATE_INPUT_BUFFER allocSurf = { 0 };
1439 
1440         ctx->surfaces[idx].format = nvenc_map_buffer_format(ctx->data_pix_fmt);
1441         if (ctx->surfaces[idx].format == NV_ENC_BUFFER_FORMAT_UNDEFINED) {
1442             av_log(avctx, AV_LOG_FATAL, "Invalid input pixel format: %s\n",
1443                    av_get_pix_fmt_name(ctx->data_pix_fmt));
1444             return AVERROR(EINVAL);
1445         }
1446 
1447         allocSurf.version = NV_ENC_CREATE_INPUT_BUFFER_VER;
1448         allocSurf.width = avctx->width;
1449         allocSurf.height = avctx->height;
1450         allocSurf.bufferFmt = ctx->surfaces[idx].format;
1451 
1452         nv_status = p_nvenc->nvEncCreateInputBuffer(ctx->nvencoder, &allocSurf);
1453         if (nv_status != NV_ENC_SUCCESS) {
1454             return nvenc_print_error(avctx, nv_status, "CreateInputBuffer failed");
1455         }
1456 
1457         ctx->surfaces[idx].input_surface = allocSurf.inputBuffer;
1458         ctx->surfaces[idx].width = allocSurf.width;
1459         ctx->surfaces[idx].height = allocSurf.height;
1460     }
1461 
1462     nv_status = p_nvenc->nvEncCreateBitstreamBuffer(ctx->nvencoder, &allocOut);
1463     if (nv_status != NV_ENC_SUCCESS) {
1464         int err = nvenc_print_error(avctx, nv_status, "CreateBitstreamBuffer failed");
1465         if (avctx->pix_fmt != AV_PIX_FMT_CUDA && avctx->pix_fmt != AV_PIX_FMT_D3D11)
1466             p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[idx].input_surface);
1467         av_frame_free(&ctx->surfaces[idx].in_ref);
1468         return err;
1469     }
1470 
1471     ctx->surfaces[idx].output_surface = allocOut.bitstreamBuffer;
1472 
1473     av_fifo_generic_write(ctx->unused_surface_queue, &tmp_surface, sizeof(tmp_surface), NULL);
1474 
1475     return 0;
1476 }
1477 
nvenc_setup_surfaces(AVCodecContext * avctx)1478 static av_cold int nvenc_setup_surfaces(AVCodecContext *avctx)
1479 {
1480     NvencContext *ctx = avctx->priv_data;
1481     int i, res = 0, res2;
1482 
1483     ctx->surfaces = av_mallocz_array(ctx->nb_surfaces, sizeof(*ctx->surfaces));
1484     if (!ctx->surfaces)
1485         return AVERROR(ENOMEM);
1486 
1487     ctx->timestamp_list = av_fifo_alloc(ctx->nb_surfaces * sizeof(int64_t));
1488     if (!ctx->timestamp_list)
1489         return AVERROR(ENOMEM);
1490 
1491     ctx->unused_surface_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(NvencSurface*));
1492     if (!ctx->unused_surface_queue)
1493         return AVERROR(ENOMEM);
1494 
1495     ctx->output_surface_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(NvencSurface*));
1496     if (!ctx->output_surface_queue)
1497         return AVERROR(ENOMEM);
1498     ctx->output_surface_ready_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(NvencSurface*));
1499     if (!ctx->output_surface_ready_queue)
1500         return AVERROR(ENOMEM);
1501 
1502     res = nvenc_push_context(avctx);
1503     if (res < 0)
1504         return res;
1505 
1506     for (i = 0; i < ctx->nb_surfaces; i++) {
1507         if ((res = nvenc_alloc_surface(avctx, i)) < 0)
1508             goto fail;
1509     }
1510 
1511 fail:
1512     res2 = nvenc_pop_context(avctx);
1513     if (res2 < 0)
1514         return res2;
1515 
1516     return res;
1517 }
1518 
nvenc_setup_extradata(AVCodecContext * avctx)1519 static av_cold int nvenc_setup_extradata(AVCodecContext *avctx)
1520 {
1521     NvencContext *ctx = avctx->priv_data;
1522     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1523     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1524 
1525     NVENCSTATUS nv_status;
1526     uint32_t outSize = 0;
1527     char tmpHeader[256];
1528     NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };
1529     payload.version = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;
1530 
1531     payload.spsppsBuffer = tmpHeader;
1532     payload.inBufferSize = sizeof(tmpHeader);
1533     payload.outSPSPPSPayloadSize = &outSize;
1534 
1535     nv_status = p_nvenc->nvEncGetSequenceParams(ctx->nvencoder, &payload);
1536     if (nv_status != NV_ENC_SUCCESS) {
1537         return nvenc_print_error(avctx, nv_status, "GetSequenceParams failed");
1538     }
1539 
1540     avctx->extradata_size = outSize;
1541     avctx->extradata = av_mallocz(outSize + AV_INPUT_BUFFER_PADDING_SIZE);
1542 
1543     if (!avctx->extradata) {
1544         return AVERROR(ENOMEM);
1545     }
1546 
1547     memcpy(avctx->extradata, tmpHeader, outSize);
1548 
1549     return 0;
1550 }
1551 
ff_nvenc_encode_close(AVCodecContext * avctx)1552 av_cold int ff_nvenc_encode_close(AVCodecContext *avctx)
1553 {
1554     NvencContext *ctx               = avctx->priv_data;
1555     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1556     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1557     int i, res;
1558 
1559     /* the encoder has to be flushed before it can be closed */
1560     if (ctx->nvencoder) {
1561         NV_ENC_PIC_PARAMS params        = { .version        = NV_ENC_PIC_PARAMS_VER,
1562                                             .encodePicFlags = NV_ENC_PIC_FLAG_EOS };
1563 
1564         res = nvenc_push_context(avctx);
1565         if (res < 0)
1566             return res;
1567 
1568         p_nvenc->nvEncEncodePicture(ctx->nvencoder, &params);
1569     }
1570 
1571     av_fifo_freep(&ctx->timestamp_list);
1572     av_fifo_freep(&ctx->output_surface_ready_queue);
1573     av_fifo_freep(&ctx->output_surface_queue);
1574     av_fifo_freep(&ctx->unused_surface_queue);
1575 
1576     if (ctx->surfaces && (avctx->pix_fmt == AV_PIX_FMT_CUDA || avctx->pix_fmt == AV_PIX_FMT_D3D11)) {
1577         for (i = 0; i < ctx->nb_registered_frames; i++) {
1578             if (ctx->registered_frames[i].mapped)
1579                 p_nvenc->nvEncUnmapInputResource(ctx->nvencoder, ctx->registered_frames[i].in_map.mappedResource);
1580             if (ctx->registered_frames[i].regptr)
1581                 p_nvenc->nvEncUnregisterResource(ctx->nvencoder, ctx->registered_frames[i].regptr);
1582         }
1583         ctx->nb_registered_frames = 0;
1584     }
1585 
1586     if (ctx->surfaces) {
1587         for (i = 0; i < ctx->nb_surfaces; ++i) {
1588             if (avctx->pix_fmt != AV_PIX_FMT_CUDA && avctx->pix_fmt != AV_PIX_FMT_D3D11)
1589                 p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[i].input_surface);
1590             av_frame_free(&ctx->surfaces[i].in_ref);
1591             p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->surfaces[i].output_surface);
1592         }
1593     }
1594     av_freep(&ctx->surfaces);
1595     ctx->nb_surfaces = 0;
1596 
1597     av_frame_free(&ctx->frame);
1598 
1599     if (ctx->nvencoder) {
1600         p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);
1601 
1602         res = nvenc_pop_context(avctx);
1603         if (res < 0)
1604             return res;
1605     }
1606     ctx->nvencoder = NULL;
1607 
1608     if (ctx->cu_context_internal)
1609         CHECK_CU(dl_fn->cuda_dl->cuCtxDestroy(ctx->cu_context_internal));
1610     ctx->cu_context = ctx->cu_context_internal = NULL;
1611 
1612 #if CONFIG_D3D11VA
1613     if (ctx->d3d11_device) {
1614         ID3D11Device_Release(ctx->d3d11_device);
1615         ctx->d3d11_device = NULL;
1616     }
1617 #endif
1618 
1619     nvenc_free_functions(&dl_fn->nvenc_dl);
1620     cuda_free_functions(&dl_fn->cuda_dl);
1621 
1622     dl_fn->nvenc_device_count = 0;
1623 
1624     av_log(avctx, AV_LOG_VERBOSE, "Nvenc unloaded\n");
1625 
1626     return 0;
1627 }
1628 
ff_nvenc_encode_init(AVCodecContext * avctx)1629 av_cold int ff_nvenc_encode_init(AVCodecContext *avctx)
1630 {
1631     NvencContext *ctx = avctx->priv_data;
1632     int ret;
1633 
1634     if (avctx->pix_fmt == AV_PIX_FMT_CUDA || avctx->pix_fmt == AV_PIX_FMT_D3D11) {
1635         AVHWFramesContext *frames_ctx;
1636         if (!avctx->hw_frames_ctx) {
1637             av_log(avctx, AV_LOG_ERROR,
1638                    "hw_frames_ctx must be set when using GPU frames as input\n");
1639             return AVERROR(EINVAL);
1640         }
1641         frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
1642         if (frames_ctx->format != avctx->pix_fmt) {
1643             av_log(avctx, AV_LOG_ERROR,
1644                    "hw_frames_ctx must match the GPU frame type\n");
1645             return AVERROR(EINVAL);
1646         }
1647         ctx->data_pix_fmt = frames_ctx->sw_format;
1648     } else {
1649         ctx->data_pix_fmt = avctx->pix_fmt;
1650     }
1651 
1652     ctx->frame = av_frame_alloc();
1653     if (!ctx->frame)
1654         return AVERROR(ENOMEM);
1655 
1656     if ((ret = nvenc_load_libraries(avctx)) < 0)
1657         return ret;
1658 
1659     if ((ret = nvenc_setup_device(avctx)) < 0)
1660         return ret;
1661 
1662     if ((ret = nvenc_setup_encoder(avctx)) < 0)
1663         return ret;
1664 
1665     if ((ret = nvenc_setup_surfaces(avctx)) < 0)
1666         return ret;
1667 
1668     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
1669         if ((ret = nvenc_setup_extradata(avctx)) < 0)
1670             return ret;
1671     }
1672 
1673     return 0;
1674 }
1675 
get_free_frame(NvencContext * ctx)1676 static NvencSurface *get_free_frame(NvencContext *ctx)
1677 {
1678     NvencSurface *tmp_surf;
1679 
1680     if (!(av_fifo_size(ctx->unused_surface_queue) > 0))
1681         // queue empty
1682         return NULL;
1683 
1684     av_fifo_generic_read(ctx->unused_surface_queue, &tmp_surf, sizeof(tmp_surf), NULL);
1685     return tmp_surf;
1686 }
1687 
nvenc_copy_frame(AVCodecContext * avctx,NvencSurface * nv_surface,NV_ENC_LOCK_INPUT_BUFFER * lock_buffer_params,const AVFrame * frame)1688 static int nvenc_copy_frame(AVCodecContext *avctx, NvencSurface *nv_surface,
1689             NV_ENC_LOCK_INPUT_BUFFER *lock_buffer_params, const AVFrame *frame)
1690 {
1691     int dst_linesize[4] = {
1692         lock_buffer_params->pitch,
1693         lock_buffer_params->pitch,
1694         lock_buffer_params->pitch,
1695         lock_buffer_params->pitch
1696     };
1697     uint8_t *dst_data[4];
1698     int ret;
1699 
1700     if (frame->format == AV_PIX_FMT_YUV420P)
1701         dst_linesize[1] = dst_linesize[2] >>= 1;
1702 
1703     ret = av_image_fill_pointers(dst_data, frame->format, nv_surface->height,
1704                                  lock_buffer_params->bufferDataPtr, dst_linesize);
1705     if (ret < 0)
1706         return ret;
1707 
1708     if (frame->format == AV_PIX_FMT_YUV420P)
1709         FFSWAP(uint8_t*, dst_data[1], dst_data[2]);
1710 
1711     av_image_copy(dst_data, dst_linesize,
1712                   (const uint8_t**)frame->data, frame->linesize, frame->format,
1713                   avctx->width, avctx->height);
1714 
1715     return 0;
1716 }
1717 
nvenc_find_free_reg_resource(AVCodecContext * avctx)1718 static int nvenc_find_free_reg_resource(AVCodecContext *avctx)
1719 {
1720     NvencContext *ctx = avctx->priv_data;
1721     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1722     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1723     NVENCSTATUS nv_status;
1724 
1725     int i, first_round;
1726 
1727     if (ctx->nb_registered_frames == FF_ARRAY_ELEMS(ctx->registered_frames)) {
1728         for (first_round = 1; first_round >= 0; first_round--) {
1729             for (i = 0; i < ctx->nb_registered_frames; i++) {
1730                 if (!ctx->registered_frames[i].mapped) {
1731                     if (ctx->registered_frames[i].regptr) {
1732                         if (first_round)
1733                             continue;
1734                         nv_status = p_nvenc->nvEncUnregisterResource(ctx->nvencoder, ctx->registered_frames[i].regptr);
1735                         if (nv_status != NV_ENC_SUCCESS)
1736                             return nvenc_print_error(avctx, nv_status, "Failed unregistering unused input resource");
1737                         ctx->registered_frames[i].ptr = NULL;
1738                         ctx->registered_frames[i].regptr = NULL;
1739                     }
1740                     return i;
1741                 }
1742             }
1743         }
1744     } else {
1745         return ctx->nb_registered_frames++;
1746     }
1747 
1748     av_log(avctx, AV_LOG_ERROR, "Too many registered CUDA frames\n");
1749     return AVERROR(ENOMEM);
1750 }
1751 
nvenc_register_frame(AVCodecContext * avctx,const AVFrame * frame)1752 static int nvenc_register_frame(AVCodecContext *avctx, const AVFrame *frame)
1753 {
1754     NvencContext *ctx = avctx->priv_data;
1755     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1756     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1757 
1758     AVHWFramesContext *frames_ctx = (AVHWFramesContext*)frame->hw_frames_ctx->data;
1759     NV_ENC_REGISTER_RESOURCE reg;
1760     int i, idx, ret;
1761 
1762     for (i = 0; i < ctx->nb_registered_frames; i++) {
1763         if (avctx->pix_fmt == AV_PIX_FMT_CUDA && ctx->registered_frames[i].ptr == frame->data[0])
1764             return i;
1765         else if (avctx->pix_fmt == AV_PIX_FMT_D3D11 && ctx->registered_frames[i].ptr == frame->data[0] && ctx->registered_frames[i].ptr_index == (intptr_t)frame->data[1])
1766             return i;
1767     }
1768 
1769     idx = nvenc_find_free_reg_resource(avctx);
1770     if (idx < 0)
1771         return idx;
1772 
1773     reg.version            = NV_ENC_REGISTER_RESOURCE_VER;
1774     reg.width              = frames_ctx->width;
1775     reg.height             = frames_ctx->height;
1776     reg.pitch              = frame->linesize[0];
1777     reg.resourceToRegister = frame->data[0];
1778 
1779     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1780         reg.resourceType   = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
1781     }
1782     else if (avctx->pix_fmt == AV_PIX_FMT_D3D11) {
1783         reg.resourceType     = NV_ENC_INPUT_RESOURCE_TYPE_DIRECTX;
1784         reg.subResourceIndex = (intptr_t)frame->data[1];
1785     }
1786 
1787     reg.bufferFormat       = nvenc_map_buffer_format(frames_ctx->sw_format);
1788     if (reg.bufferFormat == NV_ENC_BUFFER_FORMAT_UNDEFINED) {
1789         av_log(avctx, AV_LOG_FATAL, "Invalid input pixel format: %s\n",
1790                av_get_pix_fmt_name(frames_ctx->sw_format));
1791         return AVERROR(EINVAL);
1792     }
1793 
1794     ret = p_nvenc->nvEncRegisterResource(ctx->nvencoder, &reg);
1795     if (ret != NV_ENC_SUCCESS) {
1796         nvenc_print_error(avctx, ret, "Error registering an input resource");
1797         return AVERROR_UNKNOWN;
1798     }
1799 
1800     ctx->registered_frames[idx].ptr       = frame->data[0];
1801     ctx->registered_frames[idx].ptr_index = reg.subResourceIndex;
1802     ctx->registered_frames[idx].regptr    = reg.registeredResource;
1803     return idx;
1804 }
1805 
nvenc_upload_frame(AVCodecContext * avctx,const AVFrame * frame,NvencSurface * nvenc_frame)1806 static int nvenc_upload_frame(AVCodecContext *avctx, const AVFrame *frame,
1807                                       NvencSurface *nvenc_frame)
1808 {
1809     NvencContext *ctx = avctx->priv_data;
1810     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1811     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1812 
1813     int res;
1814     NVENCSTATUS nv_status;
1815 
1816     if (avctx->pix_fmt == AV_PIX_FMT_CUDA || avctx->pix_fmt == AV_PIX_FMT_D3D11) {
1817         int reg_idx = nvenc_register_frame(avctx, frame);
1818         if (reg_idx < 0) {
1819             av_log(avctx, AV_LOG_ERROR, "Could not register an input HW frame\n");
1820             return reg_idx;
1821         }
1822 
1823         res = av_frame_ref(nvenc_frame->in_ref, frame);
1824         if (res < 0)
1825             return res;
1826 
1827         if (!ctx->registered_frames[reg_idx].mapped) {
1828             ctx->registered_frames[reg_idx].in_map.version = NV_ENC_MAP_INPUT_RESOURCE_VER;
1829             ctx->registered_frames[reg_idx].in_map.registeredResource = ctx->registered_frames[reg_idx].regptr;
1830             nv_status = p_nvenc->nvEncMapInputResource(ctx->nvencoder, &ctx->registered_frames[reg_idx].in_map);
1831             if (nv_status != NV_ENC_SUCCESS) {
1832                 av_frame_unref(nvenc_frame->in_ref);
1833                 return nvenc_print_error(avctx, nv_status, "Error mapping an input resource");
1834             }
1835         }
1836 
1837         ctx->registered_frames[reg_idx].mapped += 1;
1838 
1839         nvenc_frame->reg_idx                   = reg_idx;
1840         nvenc_frame->input_surface             = ctx->registered_frames[reg_idx].in_map.mappedResource;
1841         nvenc_frame->format                    = ctx->registered_frames[reg_idx].in_map.mappedBufferFmt;
1842         nvenc_frame->pitch                     = frame->linesize[0];
1843 
1844         return 0;
1845     } else {
1846         NV_ENC_LOCK_INPUT_BUFFER lockBufferParams = { 0 };
1847 
1848         lockBufferParams.version = NV_ENC_LOCK_INPUT_BUFFER_VER;
1849         lockBufferParams.inputBuffer = nvenc_frame->input_surface;
1850 
1851         nv_status = p_nvenc->nvEncLockInputBuffer(ctx->nvencoder, &lockBufferParams);
1852         if (nv_status != NV_ENC_SUCCESS) {
1853             return nvenc_print_error(avctx, nv_status, "Failed locking nvenc input buffer");
1854         }
1855 
1856         nvenc_frame->pitch = lockBufferParams.pitch;
1857         res = nvenc_copy_frame(avctx, nvenc_frame, &lockBufferParams, frame);
1858 
1859         nv_status = p_nvenc->nvEncUnlockInputBuffer(ctx->nvencoder, nvenc_frame->input_surface);
1860         if (nv_status != NV_ENC_SUCCESS) {
1861             return nvenc_print_error(avctx, nv_status, "Failed unlocking input buffer!");
1862         }
1863 
1864         return res;
1865     }
1866 }
1867 
nvenc_codec_specific_pic_params(AVCodecContext * avctx,NV_ENC_PIC_PARAMS * params,NV_ENC_SEI_PAYLOAD * sei_data,int sei_count)1868 static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,
1869                                             NV_ENC_PIC_PARAMS *params,
1870                                             NV_ENC_SEI_PAYLOAD *sei_data,
1871                                             int sei_count)
1872 {
1873     NvencContext *ctx = avctx->priv_data;
1874 
1875     switch (avctx->codec->id) {
1876     case AV_CODEC_ID_H264:
1877         params->codecPicParams.h264PicParams.sliceMode =
1878             ctx->encode_config.encodeCodecConfig.h264Config.sliceMode;
1879         params->codecPicParams.h264PicParams.sliceModeData =
1880             ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;
1881         if (sei_count > 0) {
1882             params->codecPicParams.h264PicParams.seiPayloadArray = sei_data;
1883             params->codecPicParams.h264PicParams.seiPayloadArrayCnt = sei_count;
1884         }
1885 
1886       break;
1887     case AV_CODEC_ID_HEVC:
1888         params->codecPicParams.hevcPicParams.sliceMode =
1889             ctx->encode_config.encodeCodecConfig.hevcConfig.sliceMode;
1890         params->codecPicParams.hevcPicParams.sliceModeData =
1891             ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData;
1892         if (sei_count > 0) {
1893             params->codecPicParams.hevcPicParams.seiPayloadArray = sei_data;
1894             params->codecPicParams.hevcPicParams.seiPayloadArrayCnt = sei_count;
1895         }
1896 
1897         break;
1898     }
1899 }
1900 
timestamp_queue_enqueue(AVFifoBuffer * queue,int64_t timestamp)1901 static inline void timestamp_queue_enqueue(AVFifoBuffer* queue, int64_t timestamp)
1902 {
1903     av_fifo_generic_write(queue, &timestamp, sizeof(timestamp), NULL);
1904 }
1905 
timestamp_queue_dequeue(AVFifoBuffer * queue)1906 static inline int64_t timestamp_queue_dequeue(AVFifoBuffer* queue)
1907 {
1908     int64_t timestamp = AV_NOPTS_VALUE;
1909     if (av_fifo_size(queue) > 0)
1910         av_fifo_generic_read(queue, &timestamp, sizeof(timestamp), NULL);
1911 
1912     return timestamp;
1913 }
1914 
nvenc_set_timestamp(AVCodecContext * avctx,NV_ENC_LOCK_BITSTREAM * params,AVPacket * pkt)1915 static int nvenc_set_timestamp(AVCodecContext *avctx,
1916                                NV_ENC_LOCK_BITSTREAM *params,
1917                                AVPacket *pkt)
1918 {
1919     NvencContext *ctx = avctx->priv_data;
1920 
1921     pkt->pts = params->outputTimeStamp;
1922     pkt->dts = timestamp_queue_dequeue(ctx->timestamp_list);
1923 
1924     pkt->dts -= FFMAX(ctx->encode_config.frameIntervalP - 1, 0) * FFMAX(avctx->ticks_per_frame, 1);
1925 
1926     return 0;
1927 }
1928 
process_output_surface(AVCodecContext * avctx,AVPacket * pkt,NvencSurface * tmpoutsurf)1929 static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencSurface *tmpoutsurf)
1930 {
1931     NvencContext *ctx = avctx->priv_data;
1932     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
1933     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
1934 
1935     uint32_t slice_mode_data;
1936     uint32_t *slice_offsets = NULL;
1937     NV_ENC_LOCK_BITSTREAM lock_params = { 0 };
1938     NVENCSTATUS nv_status;
1939     int res = 0;
1940 
1941     enum AVPictureType pict_type;
1942 
1943     switch (avctx->codec->id) {
1944     case AV_CODEC_ID_H264:
1945       slice_mode_data = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;
1946       break;
1947     case AV_CODEC_ID_H265:
1948       slice_mode_data = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData;
1949       break;
1950     default:
1951       av_log(avctx, AV_LOG_ERROR, "Unknown codec name\n");
1952       res = AVERROR(EINVAL);
1953       goto error;
1954     }
1955     slice_offsets = av_mallocz(slice_mode_data * sizeof(*slice_offsets));
1956 
1957     if (!slice_offsets) {
1958         res = AVERROR(ENOMEM);
1959         goto error;
1960     }
1961 
1962     lock_params.version = NV_ENC_LOCK_BITSTREAM_VER;
1963 
1964     lock_params.doNotWait = 0;
1965     lock_params.outputBitstream = tmpoutsurf->output_surface;
1966     lock_params.sliceOffsets = slice_offsets;
1967 
1968     nv_status = p_nvenc->nvEncLockBitstream(ctx->nvencoder, &lock_params);
1969     if (nv_status != NV_ENC_SUCCESS) {
1970         res = nvenc_print_error(avctx, nv_status, "Failed locking bitstream buffer");
1971         goto error;
1972     }
1973 
1974     res = ff_get_encode_buffer(avctx, pkt, lock_params.bitstreamSizeInBytes, 0);
1975 
1976     if (res < 0) {
1977         p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
1978         goto error;
1979     }
1980 
1981     memcpy(pkt->data, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes);
1982 
1983     nv_status = p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
1984     if (nv_status != NV_ENC_SUCCESS) {
1985         res = nvenc_print_error(avctx, nv_status, "Failed unlocking bitstream buffer, expect the gates of mordor to open");
1986         goto error;
1987     }
1988 
1989 
1990     if (avctx->pix_fmt == AV_PIX_FMT_CUDA || avctx->pix_fmt == AV_PIX_FMT_D3D11) {
1991         ctx->registered_frames[tmpoutsurf->reg_idx].mapped -= 1;
1992         if (ctx->registered_frames[tmpoutsurf->reg_idx].mapped == 0) {
1993             nv_status = p_nvenc->nvEncUnmapInputResource(ctx->nvencoder, ctx->registered_frames[tmpoutsurf->reg_idx].in_map.mappedResource);
1994             if (nv_status != NV_ENC_SUCCESS) {
1995                 res = nvenc_print_error(avctx, nv_status, "Failed unmapping input resource");
1996                 goto error;
1997             }
1998         } else if (ctx->registered_frames[tmpoutsurf->reg_idx].mapped < 0) {
1999             res = AVERROR_BUG;
2000             goto error;
2001         }
2002 
2003         av_frame_unref(tmpoutsurf->in_ref);
2004 
2005         tmpoutsurf->input_surface = NULL;
2006     }
2007 
2008     switch (lock_params.pictureType) {
2009     case NV_ENC_PIC_TYPE_IDR:
2010         pkt->flags |= AV_PKT_FLAG_KEY;
2011     case NV_ENC_PIC_TYPE_I:
2012         pict_type = AV_PICTURE_TYPE_I;
2013         break;
2014     case NV_ENC_PIC_TYPE_P:
2015         pict_type = AV_PICTURE_TYPE_P;
2016         break;
2017     case NV_ENC_PIC_TYPE_B:
2018         pict_type = AV_PICTURE_TYPE_B;
2019         break;
2020     case NV_ENC_PIC_TYPE_BI:
2021         pict_type = AV_PICTURE_TYPE_BI;
2022         break;
2023     default:
2024         av_log(avctx, AV_LOG_ERROR, "Unknown picture type encountered, expect the output to be broken.\n");
2025         av_log(avctx, AV_LOG_ERROR, "Please report this error and include as much information on how to reproduce it as possible.\n");
2026         res = AVERROR_EXTERNAL;
2027         goto error;
2028     }
2029 
2030 #if FF_API_CODED_FRAME
2031 FF_DISABLE_DEPRECATION_WARNINGS
2032     avctx->coded_frame->pict_type = pict_type;
2033 FF_ENABLE_DEPRECATION_WARNINGS
2034 #endif
2035 
2036     ff_side_data_set_encoder_stats(pkt,
2037         (lock_params.frameAvgQP - 1) * FF_QP2LAMBDA, NULL, 0, pict_type);
2038 
2039     res = nvenc_set_timestamp(avctx, &lock_params, pkt);
2040     if (res < 0)
2041         goto error2;
2042 
2043     av_free(slice_offsets);
2044 
2045     return 0;
2046 
2047 error:
2048     timestamp_queue_dequeue(ctx->timestamp_list);
2049 
2050 error2:
2051     av_free(slice_offsets);
2052 
2053     return res;
2054 }
2055 
output_ready(AVCodecContext * avctx,int flush)2056 static int output_ready(AVCodecContext *avctx, int flush)
2057 {
2058     NvencContext *ctx = avctx->priv_data;
2059     int nb_ready, nb_pending;
2060 
2061     nb_ready   = av_fifo_size(ctx->output_surface_ready_queue)   / sizeof(NvencSurface*);
2062     nb_pending = av_fifo_size(ctx->output_surface_queue)         / sizeof(NvencSurface*);
2063     if (flush)
2064         return nb_ready > 0;
2065     return (nb_ready > 0) && (nb_ready + nb_pending >= ctx->async_depth);
2066 }
2067 
reconfig_encoder(AVCodecContext * avctx,const AVFrame * frame)2068 static void reconfig_encoder(AVCodecContext *avctx, const AVFrame *frame)
2069 {
2070     NvencContext *ctx = avctx->priv_data;
2071     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &ctx->nvenc_dload_funcs.nvenc_funcs;
2072     NVENCSTATUS ret;
2073 
2074     NV_ENC_RECONFIGURE_PARAMS params = { 0 };
2075     int needs_reconfig = 0;
2076     int needs_encode_config = 0;
2077     int reconfig_bitrate = 0, reconfig_dar = 0;
2078     int dw, dh;
2079 
2080     params.version = NV_ENC_RECONFIGURE_PARAMS_VER;
2081     params.reInitEncodeParams = ctx->init_encode_params;
2082 
2083     compute_dar(avctx, &dw, &dh);
2084     if (dw != ctx->init_encode_params.darWidth || dh != ctx->init_encode_params.darHeight) {
2085         av_log(avctx, AV_LOG_VERBOSE,
2086                "aspect ratio change (DAR): %d:%d -> %d:%d\n",
2087                ctx->init_encode_params.darWidth,
2088                ctx->init_encode_params.darHeight, dw, dh);
2089 
2090         params.reInitEncodeParams.darHeight = dh;
2091         params.reInitEncodeParams.darWidth = dw;
2092 
2093         needs_reconfig = 1;
2094         reconfig_dar = 1;
2095     }
2096 
2097     if (ctx->rc != NV_ENC_PARAMS_RC_CONSTQP && ctx->support_dyn_bitrate) {
2098         if (avctx->bit_rate > 0 && params.reInitEncodeParams.encodeConfig->rcParams.averageBitRate != avctx->bit_rate) {
2099             av_log(avctx, AV_LOG_VERBOSE,
2100                    "avg bitrate change: %d -> %d\n",
2101                    params.reInitEncodeParams.encodeConfig->rcParams.averageBitRate,
2102                    (uint32_t)avctx->bit_rate);
2103 
2104             params.reInitEncodeParams.encodeConfig->rcParams.averageBitRate = avctx->bit_rate;
2105             reconfig_bitrate = 1;
2106         }
2107 
2108         if (avctx->rc_max_rate > 0 && ctx->encode_config.rcParams.maxBitRate != avctx->rc_max_rate) {
2109             av_log(avctx, AV_LOG_VERBOSE,
2110                    "max bitrate change: %d -> %d\n",
2111                    params.reInitEncodeParams.encodeConfig->rcParams.maxBitRate,
2112                    (uint32_t)avctx->rc_max_rate);
2113 
2114             params.reInitEncodeParams.encodeConfig->rcParams.maxBitRate = avctx->rc_max_rate;
2115             reconfig_bitrate = 1;
2116         }
2117 
2118         if (avctx->rc_buffer_size > 0 && ctx->encode_config.rcParams.vbvBufferSize != avctx->rc_buffer_size) {
2119             av_log(avctx, AV_LOG_VERBOSE,
2120                    "vbv buffer size change: %d -> %d\n",
2121                    params.reInitEncodeParams.encodeConfig->rcParams.vbvBufferSize,
2122                    avctx->rc_buffer_size);
2123 
2124             params.reInitEncodeParams.encodeConfig->rcParams.vbvBufferSize = avctx->rc_buffer_size;
2125             reconfig_bitrate = 1;
2126         }
2127 
2128         if (reconfig_bitrate) {
2129             params.resetEncoder = 1;
2130             params.forceIDR = 1;
2131 
2132             needs_encode_config = 1;
2133             needs_reconfig = 1;
2134         }
2135     }
2136 
2137     if (!needs_encode_config)
2138         params.reInitEncodeParams.encodeConfig = NULL;
2139 
2140     if (needs_reconfig) {
2141         ret = p_nvenc->nvEncReconfigureEncoder(ctx->nvencoder, &params);
2142         if (ret != NV_ENC_SUCCESS) {
2143             nvenc_print_error(avctx, ret, "failed to reconfigure nvenc");
2144         } else {
2145             if (reconfig_dar) {
2146                 ctx->init_encode_params.darHeight = dh;
2147                 ctx->init_encode_params.darWidth = dw;
2148             }
2149 
2150             if (reconfig_bitrate) {
2151                 ctx->encode_config.rcParams.averageBitRate = params.reInitEncodeParams.encodeConfig->rcParams.averageBitRate;
2152                 ctx->encode_config.rcParams.maxBitRate = params.reInitEncodeParams.encodeConfig->rcParams.maxBitRate;
2153                 ctx->encode_config.rcParams.vbvBufferSize = params.reInitEncodeParams.encodeConfig->rcParams.vbvBufferSize;
2154             }
2155 
2156         }
2157     }
2158 }
2159 
nvenc_send_frame(AVCodecContext * avctx,const AVFrame * frame)2160 static int nvenc_send_frame(AVCodecContext *avctx, const AVFrame *frame)
2161 {
2162     NVENCSTATUS nv_status;
2163     NvencSurface *tmp_out_surf, *in_surf;
2164     int res, res2;
2165     NV_ENC_SEI_PAYLOAD sei_data[8];
2166     int sei_count = 0;
2167     int i;
2168 
2169     NvencContext *ctx = avctx->priv_data;
2170     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
2171     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
2172 
2173     NV_ENC_PIC_PARAMS pic_params = { 0 };
2174     pic_params.version = NV_ENC_PIC_PARAMS_VER;
2175 
2176     if ((!ctx->cu_context && !ctx->d3d11_device) || !ctx->nvencoder)
2177         return AVERROR(EINVAL);
2178 
2179     if (frame && frame->buf[0]) {
2180         in_surf = get_free_frame(ctx);
2181         if (!in_surf)
2182             return AVERROR(EAGAIN);
2183 
2184         res = nvenc_push_context(avctx);
2185         if (res < 0)
2186             return res;
2187 
2188         reconfig_encoder(avctx, frame);
2189 
2190         res = nvenc_upload_frame(avctx, frame, in_surf);
2191 
2192         res2 = nvenc_pop_context(avctx);
2193         if (res2 < 0)
2194             return res2;
2195 
2196         if (res)
2197             return res;
2198 
2199         pic_params.inputBuffer = in_surf->input_surface;
2200         pic_params.bufferFmt = in_surf->format;
2201         pic_params.inputWidth = in_surf->width;
2202         pic_params.inputHeight = in_surf->height;
2203         pic_params.inputPitch = in_surf->pitch;
2204         pic_params.outputBitstream = in_surf->output_surface;
2205 
2206         if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
2207             if (frame->top_field_first)
2208                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
2209             else
2210                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
2211         } else {
2212             pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
2213         }
2214 
2215         if (ctx->forced_idr >= 0 && frame->pict_type == AV_PICTURE_TYPE_I) {
2216             pic_params.encodePicFlags =
2217                 ctx->forced_idr ? NV_ENC_PIC_FLAG_FORCEIDR : NV_ENC_PIC_FLAG_FORCEINTRA;
2218         } else {
2219             pic_params.encodePicFlags = 0;
2220         }
2221 
2222         pic_params.inputTimeStamp = frame->pts;
2223 
2224         if (ctx->a53_cc && av_frame_get_side_data(frame, AV_FRAME_DATA_A53_CC)) {
2225             void *a53_data = NULL;
2226             size_t a53_size = 0;
2227 
2228             if (ff_alloc_a53_sei(frame, 0, (void**)&a53_data, &a53_size) < 0) {
2229                 av_log(ctx, AV_LOG_ERROR, "Not enough memory for closed captions, skipping\n");
2230             }
2231 
2232             if (a53_data) {
2233                 sei_data[sei_count].payloadSize = (uint32_t)a53_size;
2234                 sei_data[sei_count].payloadType = 4;
2235                 sei_data[sei_count].payload = (uint8_t*)a53_data;
2236                 sei_count ++;
2237             }
2238         }
2239 
2240         if (ctx->s12m_tc && av_frame_get_side_data(frame, AV_FRAME_DATA_S12M_TIMECODE)) {
2241             void *tc_data = NULL;
2242             size_t tc_size = 0;
2243 
2244             if (ff_alloc_timecode_sei(frame, avctx->framerate, 0, (void**)&tc_data, &tc_size) < 0) {
2245                 av_log(ctx, AV_LOG_ERROR, "Not enough memory for timecode sei, skipping\n");
2246             }
2247 
2248             if (tc_data) {
2249                 sei_data[sei_count].payloadSize = (uint32_t)tc_size;
2250                 sei_data[sei_count].payloadType = SEI_TYPE_TIME_CODE;
2251                 sei_data[sei_count].payload = (uint8_t*)tc_data;
2252                 sei_count ++;
2253             }
2254         }
2255 
2256         nvenc_codec_specific_pic_params(avctx, &pic_params, sei_data, sei_count);
2257     } else {
2258         pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
2259     }
2260 
2261     res = nvenc_push_context(avctx);
2262     if (res < 0)
2263         return res;
2264 
2265     nv_status = p_nvenc->nvEncEncodePicture(ctx->nvencoder, &pic_params);
2266 
2267     for ( i = 0; i < sei_count; i++)
2268         av_freep(&sei_data[i].payload);
2269 
2270     res = nvenc_pop_context(avctx);
2271     if (res < 0)
2272         return res;
2273 
2274     if (nv_status != NV_ENC_SUCCESS &&
2275         nv_status != NV_ENC_ERR_NEED_MORE_INPUT)
2276         return nvenc_print_error(avctx, nv_status, "EncodePicture failed!");
2277 
2278     if (frame && frame->buf[0]) {
2279         av_fifo_generic_write(ctx->output_surface_queue, &in_surf, sizeof(in_surf), NULL);
2280         timestamp_queue_enqueue(ctx->timestamp_list, frame->pts);
2281     }
2282 
2283     /* all the pending buffers are now ready for output */
2284     if (nv_status == NV_ENC_SUCCESS) {
2285         while (av_fifo_size(ctx->output_surface_queue) > 0) {
2286             av_fifo_generic_read(ctx->output_surface_queue, &tmp_out_surf, sizeof(tmp_out_surf), NULL);
2287             av_fifo_generic_write(ctx->output_surface_ready_queue, &tmp_out_surf, sizeof(tmp_out_surf), NULL);
2288         }
2289     }
2290 
2291     return 0;
2292 }
2293 
ff_nvenc_receive_packet(AVCodecContext * avctx,AVPacket * pkt)2294 int ff_nvenc_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
2295 {
2296     NvencSurface *tmp_out_surf;
2297     int res, res2;
2298 
2299     NvencContext *ctx = avctx->priv_data;
2300 
2301     AVFrame *frame = ctx->frame;
2302 
2303     if ((!ctx->cu_context && !ctx->d3d11_device) || !ctx->nvencoder)
2304         return AVERROR(EINVAL);
2305 
2306     if (!frame->buf[0]) {
2307         res = ff_encode_get_frame(avctx, frame);
2308         if (res < 0 && res != AVERROR_EOF)
2309             return res;
2310     }
2311 
2312     res = nvenc_send_frame(avctx, frame);
2313     if (res < 0) {
2314         if (res != AVERROR(EAGAIN))
2315             return res;
2316     } else
2317         av_frame_unref(frame);
2318 
2319     if (output_ready(avctx, avctx->internal->draining)) {
2320         av_fifo_generic_read(ctx->output_surface_ready_queue, &tmp_out_surf, sizeof(tmp_out_surf), NULL);
2321 
2322         res = nvenc_push_context(avctx);
2323         if (res < 0)
2324             return res;
2325 
2326         res = process_output_surface(avctx, pkt, tmp_out_surf);
2327 
2328         res2 = nvenc_pop_context(avctx);
2329         if (res2 < 0)
2330             return res2;
2331 
2332         if (res)
2333             return res;
2334 
2335         av_fifo_generic_write(ctx->unused_surface_queue, &tmp_out_surf, sizeof(tmp_out_surf), NULL);
2336     } else if (avctx->internal->draining) {
2337         return AVERROR_EOF;
2338     } else {
2339         return AVERROR(EAGAIN);
2340     }
2341 
2342     return 0;
2343 }
2344 
ff_nvenc_encode_flush(AVCodecContext * avctx)2345 av_cold void ff_nvenc_encode_flush(AVCodecContext *avctx)
2346 {
2347     NvencContext *ctx = avctx->priv_data;
2348 
2349     nvenc_send_frame(avctx, NULL);
2350     av_fifo_reset(ctx->timestamp_list);
2351 }
2352