git.sesse.net Git - ffmpeg/blob - libavcodec/nvenc.c

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