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

   1 /*
   2  * NVIDIA NVENC Support
   3  * Copyright (C) 2015 Luca Barbato
   4  * Copyright (C) 2015 Philip Langdale <philipl@overt.org>
   5  * Copyright (C) 2014 Timo Rothenpieler <timo@rothenpieler.org>
   6  *
   7  * This file is part of Libav.
   8  *
   9  * Libav is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public
  11  * License as published by the Free Software Foundation; either
  12  * version 2.1 of the License, or (at your option) any later version.
  13  *
  14  * Libav is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with Libav; if not, write to the Free Software
  21  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  22  */
  23
  24 #include "config.h"
  25
  26 #include <nvEncodeAPI.h>
  27 #include <string.h>
  28
  29 #define CUDA_LIBNAME "libcuda.so"
  30
  31 #if HAVE_DLFCN_H
  32 #include <dlfcn.h>
  33
  34 #define NVENC_LIBNAME "libnvidia-encode.so"
  35
  36 #elif HAVE_WINDOWS_H
  37 #include <windows.h>
  38
  39 #if ARCH_X86_64
  40 #define NVENC_LIBNAME "nvEncodeAPI64.dll"
  41 #else
  42 #define NVENC_LIBNAME "nvEncodeAPI.dll"
  43 #endif
  44
  45 #define dlopen(filename, flags) LoadLibrary((filename))
  46 #define dlsym(handle, symbol)   GetProcAddress(handle, symbol)
  47 #define dlclose(handle)         FreeLibrary(handle)
  48 #endif
  49
  50 #include "libavutil/common.h"
  51 #include "libavutil/hwcontext.h"
  52 #include "libavutil/imgutils.h"
  53 #include "libavutil/mem.h"
  54 #include "avcodec.h"
  55 #include "internal.h"
  56 #include "nvenc.h"
  57
  58 #if CONFIG_CUDA
  59 #include "libavutil/hwcontext_cuda.h"
  60 #endif
  61
  62 #define NVENC_CAP 0x30
  63 #define BITSTREAM_BUFFER_SIZE 1024 * 1024
  64 #define IS_CBR(rc) (rc == NV_ENC_PARAMS_RC_CBR ||               \
  65                     rc == NV_ENC_PARAMS_RC_2_PASS_QUALITY ||    \
  66                     rc == NV_ENC_PARAMS_RC_2_PASS_FRAMESIZE_CAP)
  67
  68 #define LOAD_LIBRARY(l, path)                   \
  69     do {                                        \
  70         if (!((l) = dlopen(path, RTLD_LAZY))) { \
  71             av_log(avctx, AV_LOG_ERROR,         \
  72                    "Cannot load %s\n",          \
  73                    path);                       \
  74             return AVERROR_UNKNOWN;             \
  75         }                                       \
  76     } while (0)
  77
  78 #define LOAD_SYMBOL(fun, lib, symbol)        \
  79     do {                                     \
  80         if (!((fun) = dlsym(lib, symbol))) { \
  81             av_log(avctx, AV_LOG_ERROR,      \
  82                    "Cannot load %s\n",       \
  83                    symbol);                  \
  84             return AVERROR_UNKNOWN;          \
  85         }                                    \
  86     } while (0)
  87
  88 const enum AVPixelFormat ff_nvenc_pix_fmts[] = {
  89     AV_PIX_FMT_NV12,
  90     AV_PIX_FMT_YUV420P,
  91     AV_PIX_FMT_YUV444P,
  92 #if NVENCAPI_MAJOR_VERSION >= 7
  93     AV_PIX_FMT_P010,
  94     AV_PIX_FMT_YUV444P16,
  95 #endif
  96 #if CONFIG_CUDA
  97     AV_PIX_FMT_CUDA,
  98 #endif
  99     AV_PIX_FMT_NONE
 100 };
 101
 102 #define IS_10BIT(pix_fmt)  (pix_fmt == AV_PIX_FMT_P010    || \
 103                             pix_fmt == AV_PIX_FMT_YUV444P16)
 104
 105 #define IS_YUV444(pix_fmt) (pix_fmt == AV_PIX_FMT_YUV444P || \
 106                             pix_fmt == AV_PIX_FMT_YUV444P16)
 107
 108 static const struct {
 109     NVENCSTATUS nverr;
 110     int         averr;
 111     const char *desc;
 112 } nvenc_errors[] = {
 113     { NV_ENC_SUCCESS,                      0,                "success"                  },
 114     { NV_ENC_ERR_NO_ENCODE_DEVICE,         AVERROR(ENOENT),  "no encode device"         },
 115     { NV_ENC_ERR_UNSUPPORTED_DEVICE,       AVERROR(ENOSYS),  "unsupported device"       },
 116     { NV_ENC_ERR_INVALID_ENCODERDEVICE,    AVERROR(EINVAL),  "invalid encoder device"   },
 117     { NV_ENC_ERR_INVALID_DEVICE,           AVERROR(EINVAL),  "invalid device"           },
 118     { NV_ENC_ERR_DEVICE_NOT_EXIST,         AVERROR(EIO),     "device does not exist"    },
 119     { NV_ENC_ERR_INVALID_PTR,              AVERROR(EFAULT),  "invalid ptr"              },
 120     { NV_ENC_ERR_INVALID_EVENT,            AVERROR(EINVAL),  "invalid event"            },
 121     { NV_ENC_ERR_INVALID_PARAM,            AVERROR(EINVAL),  "invalid param"            },
 122     { NV_ENC_ERR_INVALID_CALL,             AVERROR(EINVAL),  "invalid call"             },
 123     { NV_ENC_ERR_OUT_OF_MEMORY,            AVERROR(ENOMEM),  "out of memory"            },
 124     { NV_ENC_ERR_ENCODER_NOT_INITIALIZED,  AVERROR(EINVAL),  "encoder not initialized"  },
 125     { NV_ENC_ERR_UNSUPPORTED_PARAM,        AVERROR(ENOSYS),  "unsupported param"        },
 126     { NV_ENC_ERR_LOCK_BUSY,                AVERROR(EBUSY),   "lock busy"                },
 127     { NV_ENC_ERR_NOT_ENOUGH_BUFFER,        AVERROR(ENOBUFS), "not enough buffer"        },
 128     { NV_ENC_ERR_INVALID_VERSION,          AVERROR(EINVAL),  "invalid version"          },
 129     { NV_ENC_ERR_MAP_FAILED,               AVERROR(EIO),     "map failed"               },
 130     /* this is error should always be treated specially, so this "mapping"
 131      * is for completeness only */
 132     { NV_ENC_ERR_NEED_MORE_INPUT,          AVERROR_UNKNOWN,  "need more input"          },
 133     { NV_ENC_ERR_ENCODER_BUSY,             AVERROR(EBUSY),   "encoder busy"             },
 134     { NV_ENC_ERR_EVENT_NOT_REGISTERD,      AVERROR(EBADF),   "event not registered"     },
 135     { NV_ENC_ERR_GENERIC,                  AVERROR_UNKNOWN,  "generic error"            },
 136     { NV_ENC_ERR_INCOMPATIBLE_CLIENT_KEY,  AVERROR(EINVAL),  "incompatible client key"  },
 137     { NV_ENC_ERR_UNIMPLEMENTED,            AVERROR(ENOSYS),  "unimplemented"            },
 138     { NV_ENC_ERR_RESOURCE_REGISTER_FAILED, AVERROR(EIO),     "resource register failed" },
 139     { NV_ENC_ERR_RESOURCE_NOT_REGISTERED,  AVERROR(EBADF),   "resource not registered"  },
 140     { NV_ENC_ERR_RESOURCE_NOT_MAPPED,      AVERROR(EBADF),   "resource not mapped"      },
 141 };
 142
 143 static int nvenc_map_error(NVENCSTATUS err, const char **desc)
 144 {
 145     int i;
 146     for (i = 0; i < FF_ARRAY_ELEMS(nvenc_errors); i++) {
 147         if (nvenc_errors[i].nverr == err) {
 148             if (desc)
 149                 *desc = nvenc_errors[i].desc;
 150             return nvenc_errors[i].averr;
 151         }
 152     }
 153     if (desc)
 154         *desc = "unknown error";
 155     return AVERROR_UNKNOWN;
 156 }
 157
 158 static int nvenc_print_error(void *log_ctx, NVENCSTATUS err,
 159                              const char *error_string)
 160 {
 161     const char *desc;
 162     int ret;
 163     ret = nvenc_map_error(err, &desc);
 164     av_log(log_ctx, AV_LOG_ERROR, "%s: %s (%d)\n", error_string, desc, err);
 165     return ret;
 166 }
 167
 168 static av_cold int nvenc_load_libraries(AVCodecContext *avctx)
 169 {
 170     NVENCContext *ctx         = avctx->priv_data;
 171     NVENCLibraryContext *nvel = &ctx->nvel;
 172     PNVENCODEAPICREATEINSTANCE nvenc_create_instance;
 173     NVENCSTATUS err;
 174
 175 #if CONFIG_CUDA
 176     nvel->cu_init                      = cuInit;
 177     nvel->cu_device_get_count          = cuDeviceGetCount;
 178     nvel->cu_device_get                = cuDeviceGet;
 179     nvel->cu_device_get_name           = cuDeviceGetName;
 180     nvel->cu_device_compute_capability = cuDeviceComputeCapability;
 181     nvel->cu_ctx_create                = cuCtxCreate_v2;
 182     nvel->cu_ctx_pop_current           = cuCtxPopCurrent_v2;
 183     nvel->cu_ctx_push_current          = cuCtxPushCurrent_v2;
 184     nvel->cu_ctx_destroy               = cuCtxDestroy_v2;
 185 #else
 186     LOAD_LIBRARY(nvel->cuda, CUDA_LIBNAME);
 187
 188     LOAD_SYMBOL(nvel->cu_init, nvel->cuda, "cuInit");
 189     LOAD_SYMBOL(nvel->cu_device_get_count, nvel->cuda, "cuDeviceGetCount");
 190     LOAD_SYMBOL(nvel->cu_device_get, nvel->cuda, "cuDeviceGet");
 191     LOAD_SYMBOL(nvel->cu_device_get_name, nvel->cuda, "cuDeviceGetName");
 192     LOAD_SYMBOL(nvel->cu_device_compute_capability, nvel->cuda,
 193                 "cuDeviceComputeCapability");
 194     LOAD_SYMBOL(nvel->cu_ctx_create, nvel->cuda, "cuCtxCreate_v2");
 195     LOAD_SYMBOL(nvel->cu_ctx_pop_current, nvel->cuda, "cuCtxPopCurrent_v2");
 196     LOAD_SYMBOL(nvel->cu_ctx_push_current, nvel->cuda, "cuCtxPushCurrent_v2");
 197     LOAD_SYMBOL(nvel->cu_ctx_destroy, nvel->cuda, "cuCtxDestroy_v2");
 198 #endif
 199
 200     LOAD_LIBRARY(nvel->nvenc, NVENC_LIBNAME);
 201
 202     LOAD_SYMBOL(nvenc_create_instance, nvel->nvenc,
 203                 "NvEncodeAPICreateInstance");
 204
 205     nvel->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER;
 206
 207     err = nvenc_create_instance(&nvel->nvenc_funcs);
 208     if (err != NV_ENC_SUCCESS)
 209         return nvenc_print_error(avctx, err, "Cannot create the NVENC instance");
 210
 211     return 0;
 212 }
 213
 214 static int nvenc_open_session(AVCodecContext *avctx)
 215 {
 216     NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS params = { 0 };
 217     NVENCContext *ctx                           = avctx->priv_data;
 218     NV_ENCODE_API_FUNCTION_LIST *nv             = &ctx->nvel.nvenc_funcs;
 219     int ret;
 220
 221     params.version    = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;
 222     params.apiVersion = NVENCAPI_VERSION;
 223     params.device     = ctx->cu_context;
 224     params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
 225
 226     ret = nv->nvEncOpenEncodeSessionEx(&params, &ctx->nvenc_ctx);
 227     if (ret != NV_ENC_SUCCESS) {
 228         ctx->nvenc_ctx = NULL;
 229         return nvenc_print_error(avctx, ret, "Cannot open the NVENC Session");
 230     }
 231
 232     return 0;
 233 }
 234
 235 static int nvenc_check_codec_support(AVCodecContext *avctx)
 236 {
 237     NVENCContext *ctx               = avctx->priv_data;
 238     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
 239     int i, ret, count = 0;
 240     GUID *guids = NULL;
 241
 242     ret = nv->nvEncGetEncodeGUIDCount(ctx->nvenc_ctx, &count);
 243
 244     if (ret != NV_ENC_SUCCESS || !count)
 245         return AVERROR(ENOSYS);
 246
 247     guids = av_malloc(count * sizeof(GUID));
 248     if (!guids)
 249         return AVERROR(ENOMEM);
 250
 251     ret = nv->nvEncGetEncodeGUIDs(ctx->nvenc_ctx, guids, count, &count);
 252     if (ret != NV_ENC_SUCCESS) {
 253         ret = AVERROR(ENOSYS);
 254         goto fail;
 255     }
 256
 257     ret = AVERROR(ENOSYS);
 258     for (i = 0; i < count; i++) {
 259         if (!memcmp(&guids[i], &ctx->params.encodeGUID, sizeof(*guids))) {
 260             ret = 0;
 261             break;
 262         }
 263     }
 264
 265 fail:
 266     av_free(guids);
 267
 268     return ret;
 269 }
 270
 271 static int nvenc_check_cap(AVCodecContext *avctx, NV_ENC_CAPS cap)
 272 {
 273     NVENCContext *ctx               = avctx->priv_data;
 274     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
 275     NV_ENC_CAPS_PARAM params        = { 0 };
 276     int ret, val = 0;
 277
 278     params.version     = NV_ENC_CAPS_PARAM_VER;
 279     params.capsToQuery = cap;
 280
 281     ret = nv->nvEncGetEncodeCaps(ctx->nvenc_ctx, ctx->params.encodeGUID, &params, &val);
 282
 283     if (ret == NV_ENC_SUCCESS)
 284         return val;
 285     return 0;
 286 }
 287
 288 static int nvenc_check_capabilities(AVCodecContext *avctx)
 289 {
 290     NVENCContext *ctx = avctx->priv_data;
 291     int ret;
 292
 293     ret = nvenc_check_codec_support(avctx);
 294     if (ret < 0) {
 295         av_log(avctx, AV_LOG_VERBOSE, "Codec not supported\n");
 296         return ret;
 297     }
 298
 299     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_YUV444_ENCODE);
 300     if (ctx->data_pix_fmt == AV_PIX_FMT_YUV444P && ret <= 0) {
 301         av_log(avctx, AV_LOG_VERBOSE, "YUV444P not supported\n");
 302         return AVERROR(ENOSYS);
 303     }
 304
 305     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_WIDTH_MAX);
 306     if (ret < avctx->width) {
 307         av_log(avctx, AV_LOG_VERBOSE, "Width %d exceeds %d\n",
 308                avctx->width, ret);
 309         return AVERROR(ENOSYS);
 310     }
 311
 312     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_HEIGHT_MAX);
 313     if (ret < avctx->height) {
 314         av_log(avctx, AV_LOG_VERBOSE, "Height %d exceeds %d\n",
 315                avctx->height, ret);
 316         return AVERROR(ENOSYS);
 317     }
 318
 319     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_NUM_MAX_BFRAMES);
 320     if (ret < avctx->max_b_frames) {
 321         av_log(avctx, AV_LOG_VERBOSE, "Max B-frames %d exceed %d\n",
 322                avctx->max_b_frames, ret);
 323
 324         return AVERROR(ENOSYS);
 325     }
 326
 327     return 0;
 328 }
 329
 330 static int nvenc_check_device(AVCodecContext *avctx, int idx)
 331 {
 332     NVENCContext *ctx               = avctx->priv_data;
 333     NVENCLibraryContext *nvel       = &ctx->nvel;
 334     char name[128]                  = { 0 };
 335     int major, minor, ret;
 336     CUdevice cu_device;
 337     CUcontext dummy;
 338     int loglevel = AV_LOG_VERBOSE;
 339
 340     if (ctx->device == LIST_DEVICES)
 341         loglevel = AV_LOG_INFO;
 342
 343     ret = nvel->cu_device_get(&cu_device, idx);
 344     if (ret != CUDA_SUCCESS) {
 345         av_log(avctx, AV_LOG_ERROR,
 346                "Cannot access the CUDA device %d\n",
 347                idx);
 348         return -1;
 349     }
 350
 351     ret = nvel->cu_device_get_name(name, sizeof(name), cu_device);
 352     if (ret != CUDA_SUCCESS)
 353         return -1;
 354
 355     ret = nvel->cu_device_compute_capability(&major, &minor, cu_device);
 356     if (ret != CUDA_SUCCESS)
 357         return -1;
 358
 359     av_log(avctx, loglevel, "Device %d [%s] ", cu_device, name);
 360
 361     if (((major << 4) | minor) < NVENC_CAP)
 362         goto fail;
 363
 364     if (ctx->device != idx && ctx->device != ANY_DEVICE)
 365         return -1;
 366
 367     ret = nvel->cu_ctx_create(&ctx->cu_context_internal, 0, cu_device);
 368     if (ret != CUDA_SUCCESS)
 369         goto fail;
 370
 371     ctx->cu_context = ctx->cu_context_internal;
 372
 373     ret = nvel->cu_ctx_pop_current(&dummy);
 374     if (ret != CUDA_SUCCESS)
 375         goto fail2;
 376
 377     if ((ret = nvenc_open_session(avctx)) < 0)
 378         goto fail2;
 379
 380     if ((ret = nvenc_check_capabilities(avctx)) < 0)
 381         goto fail3;
 382
 383     av_log(avctx, loglevel, "supports NVENC\n");
 384
 385     if (ctx->device == idx || ctx->device == ANY_DEVICE)
 386         return 0;
 387
 388 fail3:
 389     nvel->nvenc_funcs.nvEncDestroyEncoder(ctx->nvenc_ctx);
 390     ctx->nvenc_ctx = NULL;
 391
 392 fail2:
 393     nvel->cu_ctx_destroy(ctx->cu_context_internal);
 394     ctx->cu_context_internal = NULL;
 395
 396 fail:
 397     if (ret != 0)
 398         av_log(avctx, loglevel, "does not support NVENC (major %d minor %d)\n",
 399                major, minor);
 400
 401     return AVERROR(ENOSYS);
 402 }
 403
 404 static int nvenc_setup_device(AVCodecContext *avctx)
 405 {
 406     NVENCContext *ctx         = avctx->priv_data;
 407     NVENCLibraryContext *nvel = &ctx->nvel;
 408
 409     switch (avctx->codec->id) {
 410     case AV_CODEC_ID_H264:
 411         ctx->params.encodeGUID = NV_ENC_CODEC_H264_GUID;
 412         break;
 413     case AV_CODEC_ID_HEVC:
 414         ctx->params.encodeGUID = NV_ENC_CODEC_HEVC_GUID;
 415         break;
 416     default:
 417         return AVERROR_BUG;
 418     }
 419
 420     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
 421 #if CONFIG_CUDA
 422         AVHWFramesContext   *frames_ctx;
 423         AVCUDADeviceContext *device_hwctx;
 424         int ret;
 425
 426         if (!avctx->hw_frames_ctx)
 427             return AVERROR(EINVAL);
 428
 429         frames_ctx   = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
 430         device_hwctx = frames_ctx->device_ctx->hwctx;
 431
 432         ctx->cu_context = device_hwctx->cuda_ctx;
 433
 434         ret = nvenc_open_session(avctx);
 435         if (ret < 0)
 436             return ret;
 437
 438         ret = nvenc_check_capabilities(avctx);
 439         if (ret < 0)
 440             return ret;
 441 #else
 442         return AVERROR_BUG;
 443 #endif
 444     } else {
 445         int i, nb_devices = 0;
 446
 447         if ((nvel->cu_init(0)) != CUDA_SUCCESS) {
 448             av_log(avctx, AV_LOG_ERROR,
 449                    "Cannot init CUDA\n");
 450             return AVERROR_UNKNOWN;
 451         }
 452
 453         if ((nvel->cu_device_get_count(&nb_devices)) != CUDA_SUCCESS) {
 454             av_log(avctx, AV_LOG_ERROR,
 455                    "Cannot enumerate the CUDA devices\n");
 456             return AVERROR_UNKNOWN;
 457         }
 458
 459
 460         for (i = 0; i < nb_devices; ++i) {
 461             if ((nvenc_check_device(avctx, i)) >= 0 && ctx->device != LIST_DEVICES)
 462                 return 0;
 463         }
 464
 465         if (ctx->device == LIST_DEVICES)
 466             return AVERROR_EXIT;
 467
 468         return AVERROR(ENOSYS);
 469     }
 470
 471     return 0;
 472 }
 473
 474 typedef struct GUIDTuple {
 475     const GUID guid;
 476     int flags;
 477 } GUIDTuple;
 478
 479 #define PRESET_ALIAS(alias, name, ...) \
 480     [PRESET_ ## alias] = { NV_ENC_PRESET_ ## name ## _GUID, __VA_ARGS__ }
 481
 482 #define PRESET(name, ...) PRESET_ALIAS(name, name, __VA_ARGS__)
 483
 484 static int nvenc_map_preset(NVENCContext *ctx)
 485 {
 486     GUIDTuple presets[] = {
 487         PRESET(DEFAULT),
 488         PRESET(HP),
 489         PRESET(HQ),
 490         PRESET(BD),
 491         PRESET(LOW_LATENCY_DEFAULT, NVENC_LOWLATENCY),
 492         PRESET(LOW_LATENCY_HP,      NVENC_LOWLATENCY),
 493         PRESET(LOW_LATENCY_HQ,      NVENC_LOWLATENCY),
 494         PRESET(LOSSLESS_DEFAULT,    NVENC_LOSSLESS),
 495         PRESET(LOSSLESS_HP,         NVENC_LOSSLESS),
 496         PRESET_ALIAS(SLOW, HQ,      NVENC_TWO_PASSES),
 497         PRESET_ALIAS(MEDIUM, HQ,    NVENC_ONE_PASS),
 498         PRESET_ALIAS(FAST, HP,      NVENC_ONE_PASS)
 499     };
 500
 501     GUIDTuple *t = &presets[ctx->preset];
 502
 503     ctx->params.presetGUID = t->guid;
 504     ctx->flags             = t->flags;
 505
 506     return AVERROR(EINVAL);
 507 }
 508
 509 #undef PRESET
 510 #undef PRESET_ALIAS
 511
 512 static void set_constqp(AVCodecContext *avctx, NV_ENC_RC_PARAMS *rc)
 513 {
 514     NVENCContext *ctx = avctx->priv_data;
 515     rc->rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
 516
 517     if (ctx->init_qp_p >= 0) {
 518         rc->constQP.qpInterP = ctx->init_qp_p;
 519         if (ctx->init_qp_i >= 0 && ctx->init_qp_b >= 0) {
 520             rc->constQP.qpIntra  = ctx->init_qp_i;
 521             rc->constQP.qpInterB = ctx->init_qp_b;
 522         } else if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
 523             rc->constQP.qpIntra  = av_clip(rc->constQP.qpInterP * fabs(avctx->i_quant_factor) + avctx->i_quant_offset + 0.5, 0, 51);
 524             rc->constQP.qpInterB = av_clip(rc->constQP.qpInterP * fabs(avctx->b_quant_factor) + avctx->b_quant_offset + 0.5, 0, 51);
 525         } else {
 526             rc->constQP.qpIntra  = rc->constQP.qpInterP;
 527             rc->constQP.qpInterB = rc->constQP.qpInterP;
 528         }
 529     } else if (avctx->global_quality >= 0) {
 530         rc->constQP.qpInterP = avctx->global_quality;
 531         rc->constQP.qpInterB = avctx->global_quality;
 532         rc->constQP.qpIntra  = avctx->global_quality;
 533     }
 534 }
 535
 536 static void set_vbr(AVCodecContext *avctx, NV_ENC_RC_PARAMS *rc)
 537 {
 538     NVENCContext *ctx    = avctx->priv_data;
 539
 540     if (avctx->qmin >= 0) {
 541         rc->enableMinQP    = 1;
 542         rc->minQP.qpInterB = avctx->qmin;
 543         rc->minQP.qpInterP = avctx->qmin;
 544         rc->minQP.qpIntra  = avctx->qmin;
 545     }
 546
 547     if (avctx->qmax >= 0) {
 548         rc->enableMaxQP = 1;
 549         rc->maxQP.qpInterB = avctx->qmax;
 550         rc->maxQP.qpInterP = avctx->qmax;
 551         rc->maxQP.qpIntra  = avctx->qmax;
 552     }
 553
 554     if (ctx->init_qp_p >= 0) {
 555         rc->enableInitialRCQP = 1;
 556         rc->initialRCQP.qpInterP = ctx->init_qp_p;
 557         if (ctx->init_qp_i < 0) {
 558             if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
 559                 rc->initialRCQP.qpIntra = av_clip(rc->initialRCQP.qpInterP * fabs(avctx->i_quant_factor) + avctx->i_quant_offset + 0.5, 0, 51);
 560             } else {
 561                 rc->initialRCQP.qpIntra = rc->initialRCQP.qpInterP;
 562             }
 563         } else {
 564             rc->initialRCQP.qpIntra = ctx->init_qp_i;
 565         }
 566
 567         if (ctx->init_qp_b < 0) {
 568             if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {
 569                 rc->initialRCQP.qpInterB = av_clip(rc->initialRCQP.qpInterP * fabs(avctx->b_quant_factor) + avctx->b_quant_offset + 0.5, 0, 51);
 570             } else {
 571                 rc->initialRCQP.qpInterB = rc->initialRCQP.qpInterP;
 572             }
 573         } else {
 574             rc->initialRCQP.qpInterB = ctx->init_qp_b;
 575         }
 576     }
 577 }
 578
 579 static void set_lossless(AVCodecContext *avctx, NV_ENC_RC_PARAMS *rc)
 580 {
 581     rc->rateControlMode  = NV_ENC_PARAMS_RC_CONSTQP;
 582     rc->constQP.qpInterB = 0;
 583     rc->constQP.qpInterP = 0;
 584     rc->constQP.qpIntra  = 0;
 585 }
 586
 587 static void nvenc_override_rate_control(AVCodecContext *avctx,
 588                                         NV_ENC_RC_PARAMS *rc)
 589 {
 590     NVENCContext *ctx    = avctx->priv_data;
 591
 592     switch (ctx->rc) {
 593     case NV_ENC_PARAMS_RC_CONSTQP:
 594         set_constqp(avctx, rc);
 595         return;
 596     case NV_ENC_PARAMS_RC_2_PASS_VBR:
 597     case NV_ENC_PARAMS_RC_VBR:
 598         set_vbr(avctx, rc);
 599         break;
 600     case NV_ENC_PARAMS_RC_VBR_MINQP:
 601         if (avctx->qmin < 0) {
 602             av_log(avctx, AV_LOG_WARNING,
 603                    "The variable bitrate rate-control requires "
 604                    "the 'qmin' option set.\n");
 605             return;
 606         }
 607         set_vbr(avctx, rc);
 608         break;
 609     case NV_ENC_PARAMS_RC_CBR:
 610         break;
 611     case NV_ENC_PARAMS_RC_2_PASS_QUALITY:
 612     case NV_ENC_PARAMS_RC_2_PASS_FRAMESIZE_CAP:
 613         if (!(ctx->flags & NVENC_LOWLATENCY)) {
 614             av_log(avctx, AV_LOG_WARNING,
 615                    "The multipass rate-control requires "
 616                    "a low-latency preset.\n");
 617             return;
 618         }
 619     }
 620
 621     rc->rateControlMode = ctx->rc;
 622 }
 623
 624 static void nvenc_setup_rate_control(AVCodecContext *avctx)
 625 {
 626     NVENCContext *ctx    = avctx->priv_data;
 627     NV_ENC_RC_PARAMS *rc = &ctx->config.rcParams;
 628
 629     if (avctx->bit_rate > 0)
 630         rc->averageBitRate = avctx->bit_rate;
 631
 632     if (avctx->rc_max_rate > 0)
 633         rc->maxBitRate = avctx->rc_max_rate;
 634
 635     if (ctx->rc > 0) {
 636         nvenc_override_rate_control(avctx, rc);
 637     } else if (ctx->flags & NVENC_LOSSLESS) {
 638         set_lossless(avctx, rc);
 639     } else if (avctx->global_quality > 0) {
 640         set_constqp(avctx, rc);
 641     } else {
 642         if (ctx->flags & NVENC_TWO_PASSES)
 643             rc->rateControlMode = NV_ENC_PARAMS_RC_2_PASS_VBR;
 644         else
 645             rc->rateControlMode = NV_ENC_PARAMS_RC_VBR;
 646         set_vbr(avctx, rc);
 647     }
 648
 649     if (avctx->rc_buffer_size > 0)
 650         rc->vbvBufferSize = avctx->rc_buffer_size;
 651
 652     if (rc->averageBitRate > 0)
 653         avctx->bit_rate = rc->averageBitRate;
 654
 655 #if NVENCAPI_MAJOR_VERSION >= 7
 656     if (ctx->aq) {
 657         ctx->config.rcParams.enableAQ   = 1;
 658         ctx->config.rcParams.aqStrength = ctx->aq_strength;
 659         av_log(avctx, AV_LOG_VERBOSE, "AQ enabled.\n");
 660     }
 661
 662     if (ctx->temporal_aq) {
 663         ctx->config.rcParams.enableTemporalAQ = 1;
 664         av_log(avctx, AV_LOG_VERBOSE, "Temporal AQ enabled.\n");
 665     }
 666
 667     if (ctx->rc_lookahead > 0) {
 668         int lkd_bound = FFMIN(ctx->nb_surfaces, ctx->async_depth) -
 669                         ctx->config.frameIntervalP - 4;
 670
 671         if (lkd_bound < 0) {
 672             av_log(avctx, AV_LOG_WARNING,
 673                    "Lookahead not enabled. Increase buffer delay (-delay).\n");
 674         } else {
 675             ctx->config.rcParams.enableLookahead = 1;
 676             ctx->config.rcParams.lookaheadDepth  = av_clip(ctx->rc_lookahead, 0, lkd_bound);
 677             ctx->config.rcParams.disableIadapt   = ctx->no_scenecut;
 678             ctx->config.rcParams.disableBadapt   = !ctx->b_adapt;
 679             av_log(avctx, AV_LOG_VERBOSE,
 680                    "Lookahead enabled: depth %d, scenecut %s, B-adapt %s.\n",
 681                    ctx->config.rcParams.lookaheadDepth,
 682                    ctx->config.rcParams.disableIadapt ? "disabled" : "enabled",
 683                    ctx->config.rcParams.disableBadapt ? "disabled" : "enabled");
 684         }
 685     }
 686
 687     if (ctx->strict_gop) {
 688         ctx->config.rcParams.strictGOPTarget = 1;
 689         av_log(avctx, AV_LOG_VERBOSE, "Strict GOP target enabled.\n");
 690     }
 691
 692     if (ctx->nonref_p)
 693         ctx->config.rcParams.enableNonRefP = 1;
 694
 695     if (ctx->zerolatency)
 696         ctx->config.rcParams.zeroReorderDelay = 1;
 697
 698     if (ctx->quality)
 699         ctx->config.rcParams.targetQuality = ctx->quality;
 700 #endif /* NVENCAPI_MAJOR_VERSION >= 7 */
 701 }
 702
 703 static int nvenc_setup_h264_config(AVCodecContext *avctx)
 704 {
 705     NVENCContext *ctx                      = avctx->priv_data;
 706     NV_ENC_CONFIG *cc                      = &ctx->config;
 707     NV_ENC_CONFIG_H264 *h264               = &cc->encodeCodecConfig.h264Config;
 708     NV_ENC_CONFIG_H264_VUI_PARAMETERS *vui = &h264->h264VUIParameters;
 709
 710     vui->colourDescriptionPresentFlag = avctx->colorspace      != AVCOL_SPC_UNSPECIFIED ||
 711                                         avctx->color_primaries != AVCOL_PRI_UNSPECIFIED ||
 712                                         avctx->color_trc       != AVCOL_TRC_UNSPECIFIED;
 713
 714     vui->colourMatrix            = avctx->colorspace;
 715     vui->colourPrimaries         = avctx->color_primaries;
 716     vui->transferCharacteristics = avctx->color_trc;
 717
 718     vui->videoFullRangeFlag = avctx->color_range == AVCOL_RANGE_JPEG;
 719
 720     vui->videoSignalTypePresentFlag = vui->colourDescriptionPresentFlag ||
 721                                       vui->videoFullRangeFlag;
 722
 723     h264->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
 724     h264->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
 725     h264->outputAUD     = 1;
 726
 727     h264->maxNumRefFrames = avctx->refs;
 728     h264->idrPeriod       = cc->gopLength;
 729
 730     h264->sliceMode     = 3;
 731     h264->sliceModeData = FFMAX(avctx->slices, 1);
 732
 733     if (ctx->flags & NVENC_LOSSLESS)
 734         h264->qpPrimeYZeroTransformBypassFlag = 1;
 735
 736     if (IS_CBR(cc->rcParams.rateControlMode)) {
 737         h264->outputBufferingPeriodSEI = 1;
 738         h264->outputPictureTimingSEI   = 1;
 739     }
 740
 741     if (ctx->profile)
 742         avctx->profile = ctx->profile;
 743
 744     if (ctx->data_pix_fmt == AV_PIX_FMT_YUV444P)
 745         h264->chromaFormatIDC = 3;
 746     else
 747         h264->chromaFormatIDC = 1;
 748
 749     switch (ctx->profile) {
 750     case NV_ENC_H264_PROFILE_BASELINE:
 751         cc->profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;
 752         break;
 753     case NV_ENC_H264_PROFILE_MAIN:
 754         cc->profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;
 755         break;
 756     case NV_ENC_H264_PROFILE_HIGH:
 757         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
 758         break;
 759     case NV_ENC_H264_PROFILE_HIGH_444:
 760         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
 761         break;
 762     case NV_ENC_H264_PROFILE_CONSTRAINED_HIGH:
 763         cc->profileGUID = NV_ENC_H264_PROFILE_CONSTRAINED_HIGH_GUID;
 764         break;
 765     }
 766
 767     if (ctx->data_pix_fmt == AV_PIX_FMT_YUV444P) {
 768         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
 769         avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;
 770     }
 771
 772     h264->level = ctx->level;
 773
 774     return 0;
 775 }
 776
 777 static int nvenc_setup_hevc_config(AVCodecContext *avctx)
 778 {
 779     NVENCContext *ctx                      = avctx->priv_data;
 780     NV_ENC_CONFIG *cc                      = &ctx->config;
 781     NV_ENC_CONFIG_HEVC *hevc               = &cc->encodeCodecConfig.hevcConfig;
 782     NV_ENC_CONFIG_HEVC_VUI_PARAMETERS *vui = &hevc->hevcVUIParameters;
 783
 784     vui->colourDescriptionPresentFlag = avctx->colorspace      != AVCOL_SPC_UNSPECIFIED ||
 785                                         avctx->color_primaries != AVCOL_PRI_UNSPECIFIED ||
 786                                         avctx->color_trc       != AVCOL_TRC_UNSPECIFIED;
 787
 788     vui->colourMatrix            = avctx->colorspace;
 789     vui->colourPrimaries         = avctx->color_primaries;
 790     vui->transferCharacteristics = avctx->color_trc;
 791
 792     vui->videoFullRangeFlag = avctx->color_range == AVCOL_RANGE_JPEG;
 793
 794     vui->videoSignalTypePresentFlag = vui->colourDescriptionPresentFlag ||
 795                                       vui->videoFullRangeFlag;
 796
 797     hevc->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
 798     hevc->repeatSPSPPS  = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;
 799     hevc->outputAUD     = 1;
 800
 801     hevc->maxNumRefFramesInDPB = avctx->refs;
 802     hevc->idrPeriod            = cc->gopLength;
 803
 804     if (IS_CBR(cc->rcParams.rateControlMode)) {
 805         hevc->outputBufferingPeriodSEI = 1;
 806         hevc->outputPictureTimingSEI   = 1;
 807     }
 808
 809     switch (ctx->profile) {
 810     case NV_ENC_HEVC_PROFILE_MAIN:
 811         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID;
 812         avctx->profile  = FF_PROFILE_HEVC_MAIN;
 813         break;
 814 #if NVENCAPI_MAJOR_VERSION >= 7
 815     case NV_ENC_HEVC_PROFILE_MAIN_10:
 816         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN10_GUID;
 817         avctx->profile  = FF_PROFILE_HEVC_MAIN_10;
 818         break;
 819     case NV_ENC_HEVC_PROFILE_REXT:
 820         cc->profileGUID = NV_ENC_HEVC_PROFILE_FREXT_GUID;
 821         avctx->profile  = FF_PROFILE_HEVC_REXT;
 822         break;
 823 #endif /* NVENCAPI_MAJOR_VERSION >= 7 */
 824     }
 825
 826     // force setting profile for various input formats
 827     switch (ctx->data_pix_fmt) {
 828     case AV_PIX_FMT_YUV420P:
 829     case AV_PIX_FMT_NV12:
 830         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID;
 831         avctx->profile  = FF_PROFILE_HEVC_MAIN;
 832         break;
 833 #if NVENCAPI_MAJOR_VERSION >= 7
 834     case AV_PIX_FMT_P010:
 835         cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN10_GUID;
 836         avctx->profile  = FF_PROFILE_HEVC_MAIN_10;
 837         break;
 838     case AV_PIX_FMT_YUV444P:
 839     case AV_PIX_FMT_YUV444P16:
 840         cc->profileGUID = NV_ENC_HEVC_PROFILE_FREXT_GUID;
 841         avctx->profile  = FF_PROFILE_HEVC_REXT;
 842         break;
 843 #endif /* NVENCAPI_MAJOR_VERSION >= 7 */
 844     }
 845
 846 #if NVENCAPI_MAJOR_VERSION >= 7
 847     hevc->chromaFormatIDC     = IS_YUV444(ctx->data_pix_fmt) ? 3 : 1;
 848     hevc->pixelBitDepthMinus8 = IS_10BIT(ctx->data_pix_fmt)  ? 2 : 0;
 849 #endif /* NVENCAPI_MAJOR_VERSION >= 7 */
 850
 851     hevc->sliceMode     = 3;
 852     hevc->sliceModeData = FFMAX(avctx->slices, 1);
 853
 854     if (ctx->level) {
 855         hevc->level = ctx->level;
 856     } else {
 857         hevc->level = NV_ENC_LEVEL_AUTOSELECT;
 858     }
 859
 860     if (ctx->tier) {
 861         hevc->tier = ctx->tier;
 862     }
 863
 864     return 0;
 865 }
 866 static int nvenc_setup_codec_config(AVCodecContext *avctx)
 867 {
 868     switch (avctx->codec->id) {
 869     case AV_CODEC_ID_H264:
 870         return nvenc_setup_h264_config(avctx);
 871     case AV_CODEC_ID_HEVC:
 872         return nvenc_setup_hevc_config(avctx);
 873     }
 874     return 0;
 875 }
 876
 877 static int nvenc_recalc_surfaces(AVCodecContext *avctx)
 878 {
 879     NVENCContext *ctx = avctx->priv_data;
 880     // default minimum of 4 surfaces
 881     // multiply by 2 for number of NVENCs on gpu (hardcode to 2)
 882     // another multiply by 2 to avoid blocking next PBB group
 883     int nb_surfaces = FFMAX(4, ctx->config.frameIntervalP * 2 * 2);
 884
 885     // lookahead enabled
 886     if (ctx->rc_lookahead > 0) {
 887         // +1 is to account for lkd_bound calculation later
 888         // +4 is to allow sufficient pipelining with lookahead
 889         nb_surfaces = FFMAX(1, FFMAX(nb_surfaces, ctx->rc_lookahead + ctx->config.frameIntervalP + 1 + 4));
 890         if (nb_surfaces > ctx->nb_surfaces && ctx->nb_surfaces > 0) {
 891             av_log(avctx, AV_LOG_WARNING,
 892                 "Defined rc_lookahead requires more surfaces, "
 893                 "increasing used surfaces %d -> %d\n",
 894                 ctx->nb_surfaces, nb_surfaces);
 895         }
 896         ctx->nb_surfaces = FFMAX(nb_surfaces, ctx->nb_surfaces);
 897     } else {
 898         if (ctx->config.frameIntervalP > 1 &&
 899             ctx->nb_surfaces < nb_surfaces && ctx->nb_surfaces > 0) {
 900             av_log(avctx, AV_LOG_WARNING,
 901                 "Defined b-frame requires more surfaces, "
 902                 "increasing used surfaces %d -> %d\n",
 903                 ctx->nb_surfaces, nb_surfaces);
 904             ctx->nb_surfaces = FFMAX(ctx->nb_surfaces, nb_surfaces);
 905         } else if (ctx->nb_surfaces <= 0)
 906             ctx->nb_surfaces = nb_surfaces;
 907         // otherwise use user specified value
 908     }
 909
 910     ctx->nb_surfaces = FFMAX(1, FFMIN(MAX_REGISTERED_FRAMES, ctx->nb_surfaces));
 911     ctx->async_depth = FFMIN(ctx->async_depth, ctx->nb_surfaces - 1);
 912     return 0;
 913 }
 914
 915 static int nvenc_setup_encoder(AVCodecContext *avctx)
 916 {
 917     NVENCContext *ctx               = avctx->priv_data;
 918     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
 919     NV_ENC_PRESET_CONFIG preset_cfg = { 0 };
 920     AVCPBProperties *cpb_props;
 921     int ret;
 922
 923     ctx->params.version = NV_ENC_INITIALIZE_PARAMS_VER;
 924
 925     ctx->params.encodeHeight = avctx->height;
 926     ctx->params.encodeWidth  = avctx->width;
 927
 928     if (avctx->sample_aspect_ratio.num &&
 929         avctx->sample_aspect_ratio.den &&
 930         (avctx->sample_aspect_ratio.num != 1 ||
 931          avctx->sample_aspect_ratio.den != 1)) {
 932         av_reduce(&ctx->params.darWidth,
 933                   &ctx->params.darHeight,
 934                   avctx->width * avctx->sample_aspect_ratio.num,
 935                   avctx->height * avctx->sample_aspect_ratio.den,
 936                   INT_MAX / 8);
 937     } else {
 938         ctx->params.darHeight = avctx->height;
 939         ctx->params.darWidth  = avctx->width;
 940     }
 941
 942     // De-compensate for hardware, dubiously, trying to compensate for
 943     // playback at 704 pixel width.
 944     if (avctx->width == 720 && (avctx->height == 480 || avctx->height == 576)) {
 945         av_reduce(&ctx->params.darWidth, &ctx->params.darHeight,
 946                   ctx->params.darWidth * 44,
 947                   ctx->params.darHeight * 45,
 948                   1024 * 1024);
 949     }
 950
 951     ctx->params.frameRateNum = avctx->time_base.den;
 952     ctx->params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
 953
 954     ctx->params.enableEncodeAsync = 0;
 955     ctx->params.enablePTD         = 1;
 956
 957     ctx->params.encodeConfig = &ctx->config;
 958
 959     nvenc_map_preset(ctx);
 960
 961     preset_cfg.version           = NV_ENC_PRESET_CONFIG_VER;
 962     preset_cfg.presetCfg.version = NV_ENC_CONFIG_VER;
 963
 964     ret = nv->nvEncGetEncodePresetConfig(ctx->nvenc_ctx,
 965                                          ctx->params.encodeGUID,
 966                                          ctx->params.presetGUID,
 967                                          &preset_cfg);
 968     if (ret != NV_ENC_SUCCESS)
 969         return nvenc_print_error(avctx, ret, "Cannot get the preset configuration");
 970
 971     memcpy(&ctx->config, &preset_cfg.presetCfg, sizeof(ctx->config));
 972
 973     ctx->config.version = NV_ENC_CONFIG_VER;
 974
 975     if (avctx->gop_size > 0) {
 976         if (avctx->max_b_frames > 0) {
 977             /* 0 is intra-only,
 978              * 1 is I/P only,
 979              * 2 is one B-Frame,
 980              * 3 two B-frames, and so on. */
 981             ctx->config.frameIntervalP = avctx->max_b_frames + 1;
 982         } else if (avctx->max_b_frames == 0) {
 983             ctx->config.frameIntervalP = 1;
 984         }
 985         ctx->config.gopLength = avctx->gop_size;
 986     } else if (avctx->gop_size == 0) {
 987         ctx->config.frameIntervalP = 0;
 988         ctx->config.gopLength      = 1;
 989     }
 990
 991     if (ctx->config.frameIntervalP > 1)
 992         avctx->max_b_frames = ctx->config.frameIntervalP - 1;
 993
 994     ctx->initial_pts[0] = AV_NOPTS_VALUE;
 995     ctx->initial_pts[1] = AV_NOPTS_VALUE;
 996
 997     nvenc_recalc_surfaces(avctx);
 998
 999     nvenc_setup_rate_control(avctx);
1000
1001     if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
1002         ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
1003     } else {
1004         ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;
1005     }
1006
1007     if ((ret = nvenc_setup_codec_config(avctx)) < 0)
1008         return ret;
1009
1010     ret = nv->nvEncInitializeEncoder(ctx->nvenc_ctx, &ctx->params);
1011     if (ret != NV_ENC_SUCCESS)
1012         return nvenc_print_error(avctx, ret, "InitializeEncoder failed");
1013
1014     cpb_props = ff_add_cpb_side_data(avctx);
1015     if (!cpb_props)
1016         return AVERROR(ENOMEM);
1017     cpb_props->max_bitrate = avctx->rc_max_rate;
1018     cpb_props->min_bitrate = avctx->rc_min_rate;
1019     cpb_props->avg_bitrate = avctx->bit_rate;
1020     cpb_props->buffer_size = avctx->rc_buffer_size;
1021
1022     return 0;
1023 }
1024
1025 static int nvenc_alloc_surface(AVCodecContext *avctx, int idx)
1026 {
1027     NVENCContext *ctx               = avctx->priv_data;
1028     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1029     NVENCFrame *tmp_surface         = &ctx->frames[idx];
1030     int ret;
1031     NV_ENC_CREATE_BITSTREAM_BUFFER out_buffer = { 0 };
1032
1033     switch (ctx->data_pix_fmt) {
1034     case AV_PIX_FMT_YUV420P:
1035         ctx->frames[idx].format = NV_ENC_BUFFER_FORMAT_YV12_PL;
1036         break;
1037     case AV_PIX_FMT_NV12:
1038         ctx->frames[idx].format = NV_ENC_BUFFER_FORMAT_NV12_PL;
1039         break;
1040     case AV_PIX_FMT_YUV444P:
1041         ctx->frames[idx].format = NV_ENC_BUFFER_FORMAT_YUV444_PL;
1042         break;
1043 #if NVENCAPI_MAJOR_VERSION >= 7
1044     case AV_PIX_FMT_P010:
1045         ctx->frames[idx].format = NV_ENC_BUFFER_FORMAT_YUV420_10BIT;
1046         break;
1047     case AV_PIX_FMT_YUV444P16:
1048         ctx->frames[idx].format = NV_ENC_BUFFER_FORMAT_YUV444_10BIT;
1049         break;
1050 #endif /* NVENCAPI_MAJOR_VERSION >= 7 */
1051     default:
1052         return AVERROR_BUG;
1053     }
1054
1055     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1056         ctx->frames[idx].in_ref = av_frame_alloc();
1057         if (!ctx->frames[idx].in_ref)
1058             return AVERROR(ENOMEM);
1059     } else {
1060         NV_ENC_CREATE_INPUT_BUFFER in_buffer      = { 0 };
1061
1062         in_buffer.version  = NV_ENC_CREATE_INPUT_BUFFER_VER;
1063
1064         in_buffer.width  = avctx->width;
1065         in_buffer.height = avctx->height;
1066
1067         in_buffer.bufferFmt  = ctx->frames[idx].format;
1068         in_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED;
1069
1070         ret = nv->nvEncCreateInputBuffer(ctx->nvenc_ctx, &in_buffer);
1071         if (ret != NV_ENC_SUCCESS)
1072             return nvenc_print_error(avctx, ret, "CreateInputBuffer failed");
1073
1074         ctx->frames[idx].in     = in_buffer.inputBuffer;
1075     }
1076
1077     out_buffer.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
1078     /* 1MB is large enough to hold most output frames.
1079      * NVENC increases this automatically if it is not enough. */
1080     out_buffer.size = BITSTREAM_BUFFER_SIZE;
1081
1082     out_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED;
1083
1084     ret = nv->nvEncCreateBitstreamBuffer(ctx->nvenc_ctx, &out_buffer);
1085     if (ret != NV_ENC_SUCCESS)
1086         return nvenc_print_error(avctx, ret, "CreateBitstreamBuffer failed");
1087
1088     ctx->frames[idx].out  = out_buffer.bitstreamBuffer;
1089
1090     av_fifo_generic_write(ctx->unused_surface_queue, &tmp_surface, sizeof(tmp_surface), NULL);
1091
1092     return 0;
1093 }
1094
1095 static int nvenc_setup_surfaces(AVCodecContext *avctx)
1096 {
1097     NVENCContext *ctx = avctx->priv_data;
1098     int i, ret;
1099
1100     ctx->frames = av_mallocz_array(ctx->nb_surfaces, sizeof(*ctx->frames));
1101     if (!ctx->frames)
1102         return AVERROR(ENOMEM);
1103
1104     ctx->timestamps = av_fifo_alloc(ctx->nb_surfaces * sizeof(int64_t));
1105     if (!ctx->timestamps)
1106         return AVERROR(ENOMEM);
1107     ctx->unused_surface_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(NVENCFrame*));
1108     if (!ctx->unused_surface_queue)
1109         return AVERROR(ENOMEM);
1110     ctx->pending = av_fifo_alloc(ctx->nb_surfaces * sizeof(*ctx->frames));
1111     if (!ctx->pending)
1112         return AVERROR(ENOMEM);
1113     ctx->ready = av_fifo_alloc(ctx->nb_surfaces * sizeof(*ctx->frames));
1114     if (!ctx->ready)
1115         return AVERROR(ENOMEM);
1116
1117     for (i = 0; i < ctx->nb_surfaces; i++) {
1118         if ((ret = nvenc_alloc_surface(avctx, i)) < 0)
1119             return ret;
1120     }
1121
1122     return 0;
1123 }
1124
1125 #define EXTRADATA_SIZE 512
1126
1127 static int nvenc_setup_extradata(AVCodecContext *avctx)
1128 {
1129     NVENCContext *ctx                     = avctx->priv_data;
1130     NV_ENCODE_API_FUNCTION_LIST *nv       = &ctx->nvel.nvenc_funcs;
1131     NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };
1132     int ret;
1133
1134     avctx->extradata = av_mallocz(EXTRADATA_SIZE + AV_INPUT_BUFFER_PADDING_SIZE);
1135     if (!avctx->extradata)
1136         return AVERROR(ENOMEM);
1137
1138     payload.version              = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;
1139     payload.spsppsBuffer         = avctx->extradata;
1140     payload.inBufferSize         = EXTRADATA_SIZE;
1141     payload.outSPSPPSPayloadSize = &avctx->extradata_size;
1142
1143     ret = nv->nvEncGetSequenceParams(ctx->nvenc_ctx, &payload);
1144     if (ret != NV_ENC_SUCCESS)
1145         return nvenc_print_error(avctx, ret, "Cannot get the extradata");
1146
1147     return 0;
1148 }
1149
1150 av_cold int ff_nvenc_encode_close(AVCodecContext *avctx)
1151 {
1152     NVENCContext *ctx               = avctx->priv_data;
1153     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1154     int i;
1155
1156     /* the encoder has to be flushed before it can be closed */
1157     if (ctx->nvenc_ctx) {
1158         NV_ENC_PIC_PARAMS params        = { .version        = NV_ENC_PIC_PARAMS_VER,
1159                                             .encodePicFlags = NV_ENC_PIC_FLAG_EOS };
1160
1161         nv->nvEncEncodePicture(ctx->nvenc_ctx, &params);
1162     }
1163
1164     av_fifo_free(ctx->timestamps);
1165     av_fifo_free(ctx->pending);
1166     av_fifo_free(ctx->ready);
1167     av_fifo_free(ctx->unused_surface_queue);
1168
1169     if (ctx->frames) {
1170         for (i = 0; i < ctx->nb_surfaces; ++i) {
1171             if (avctx->pix_fmt != AV_PIX_FMT_CUDA) {
1172                 nv->nvEncDestroyInputBuffer(ctx->nvenc_ctx, ctx->frames[i].in);
1173             } else if (ctx->frames[i].in) {
1174                 nv->nvEncUnmapInputResource(ctx->nvenc_ctx, ctx->frames[i].in_map.mappedResource);
1175             }
1176
1177             av_frame_free(&ctx->frames[i].in_ref);
1178             nv->nvEncDestroyBitstreamBuffer(ctx->nvenc_ctx, ctx->frames[i].out);
1179         }
1180     }
1181     for (i = 0; i < ctx->nb_registered_frames; i++) {
1182         if (ctx->registered_frames[i].regptr)
1183             nv->nvEncUnregisterResource(ctx->nvenc_ctx, ctx->registered_frames[i].regptr);
1184     }
1185     ctx->nb_registered_frames = 0;
1186
1187     av_freep(&ctx->frames);
1188
1189     if (ctx->nvenc_ctx)
1190         nv->nvEncDestroyEncoder(ctx->nvenc_ctx);
1191
1192     if (ctx->cu_context_internal)
1193         ctx->nvel.cu_ctx_destroy(ctx->cu_context_internal);
1194
1195     if (ctx->nvel.nvenc)
1196         dlclose(ctx->nvel.nvenc);
1197
1198 #if !CONFIG_CUDA
1199     if (ctx->nvel.cuda)
1200         dlclose(ctx->nvel.cuda);
1201 #endif
1202
1203     return 0;
1204 }
1205
1206 av_cold int ff_nvenc_encode_init(AVCodecContext *avctx)
1207 {
1208     NVENCContext *ctx = avctx->priv_data;
1209     int ret;
1210
1211     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1212         AVHWFramesContext *frames_ctx;
1213         if (!avctx->hw_frames_ctx) {
1214             av_log(avctx, AV_LOG_ERROR,
1215                    "hw_frames_ctx must be set when using GPU frames as input\n");
1216             return AVERROR(EINVAL);
1217         }
1218         frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
1219         ctx->data_pix_fmt = frames_ctx->sw_format;
1220     } else {
1221         ctx->data_pix_fmt = avctx->pix_fmt;
1222     }
1223
1224     if ((ret = nvenc_load_libraries(avctx)) < 0)
1225         return ret;
1226
1227     if ((ret = nvenc_setup_device(avctx)) < 0)
1228         return ret;
1229
1230     if ((ret = nvenc_setup_encoder(avctx)) < 0)
1231         return ret;
1232
1233     if ((ret = nvenc_setup_surfaces(avctx)) < 0)
1234         return ret;
1235
1236     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
1237         if ((ret = nvenc_setup_extradata(avctx)) < 0)
1238             return ret;
1239     }
1240
1241     return 0;
1242 }
1243
1244 static NVENCFrame *get_free_frame(NVENCContext *ctx)
1245 {
1246     NVENCFrame *tmp_surf;
1247
1248     if (!(av_fifo_size(ctx->unused_surface_queue) > 0))
1249         // queue empty
1250         return NULL;
1251
1252     av_fifo_generic_read(ctx->unused_surface_queue, &tmp_surf, sizeof(tmp_surf), NULL);
1253     return tmp_surf;
1254 }
1255
1256 static int nvenc_copy_frame(NV_ENC_LOCK_INPUT_BUFFER *in, const AVFrame *frame)
1257 {
1258     uint8_t *buf = in->bufferDataPtr;
1259     int off      = frame->height * in->pitch;
1260
1261     switch (frame->format) {
1262     case AV_PIX_FMT_YUV420P:
1263         av_image_copy_plane(buf, in->pitch,
1264                             frame->data[0], frame->linesize[0],
1265                             frame->width, frame->height);
1266         buf += off;
1267
1268         av_image_copy_plane(buf, in->pitch >> 1,
1269                             frame->data[2], frame->linesize[2],
1270                             frame->width >> 1, frame->height >> 1);
1271
1272         buf += off >> 2;
1273
1274         av_image_copy_plane(buf, in->pitch >> 1,
1275                             frame->data[1], frame->linesize[1],
1276                             frame->width >> 1, frame->height >> 1);
1277         break;
1278     case AV_PIX_FMT_NV12:
1279         av_image_copy_plane(buf, in->pitch,
1280                             frame->data[0], frame->linesize[0],
1281                             frame->width, frame->height);
1282         buf += off;
1283
1284         av_image_copy_plane(buf, in->pitch,
1285                             frame->data[1], frame->linesize[1],
1286                             frame->width, frame->height >> 1);
1287         break;
1288     case AV_PIX_FMT_P010:
1289         av_image_copy_plane(buf, in->pitch,
1290                             frame->data[0], frame->linesize[0],
1291                             frame->width << 1, frame->height);
1292         buf += off;
1293
1294         av_image_copy_plane(buf, in->pitch,
1295                             frame->data[1], frame->linesize[1],
1296                             frame->width << 1, frame->height >> 1);
1297         break;
1298     case AV_PIX_FMT_YUV444P:
1299         av_image_copy_plane(buf, in->pitch,
1300                             frame->data[0], frame->linesize[0],
1301                             frame->width, frame->height);
1302         buf += off;
1303
1304         av_image_copy_plane(buf, in->pitch,
1305                             frame->data[1], frame->linesize[1],
1306                             frame->width, frame->height);
1307         buf += off;
1308
1309         av_image_copy_plane(buf, in->pitch,
1310                             frame->data[2], frame->linesize[2],
1311                             frame->width, frame->height);
1312         break;
1313     case AV_PIX_FMT_YUV444P16:
1314         av_image_copy_plane(buf, in->pitch,
1315                             frame->data[0], frame->linesize[0],
1316                             frame->width << 1, frame->height);
1317         buf += off;
1318
1319         av_image_copy_plane(buf, in->pitch,
1320                             frame->data[1], frame->linesize[1],
1321                             frame->width << 1, frame->height);
1322         buf += off;
1323
1324         av_image_copy_plane(buf, in->pitch,
1325                             frame->data[2], frame->linesize[2],
1326                             frame->width << 1, frame->height);
1327         break;
1328     default:
1329         return AVERROR_BUG;
1330     }
1331
1332     return 0;
1333 }
1334
1335 static int nvenc_find_free_reg_resource(AVCodecContext *avctx)
1336 {
1337     NVENCContext               *ctx = avctx->priv_data;
1338     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1339     int i;
1340
1341     if (ctx->nb_registered_frames == FF_ARRAY_ELEMS(ctx->registered_frames)) {
1342         for (i = 0; i < ctx->nb_registered_frames; i++) {
1343             if (!ctx->registered_frames[i].mapped) {
1344                 if (ctx->registered_frames[i].regptr) {
1345                     nv->nvEncUnregisterResource(ctx->nvenc_ctx,
1346                                                 ctx->registered_frames[i].regptr);
1347                     ctx->registered_frames[i].regptr = NULL;
1348                 }
1349                 return i;
1350             }
1351         }
1352     } else {
1353         return ctx->nb_registered_frames++;
1354     }
1355
1356     av_log(avctx, AV_LOG_ERROR, "Too many registered CUDA frames\n");
1357     return AVERROR(ENOMEM);
1358 }
1359
1360 static int nvenc_register_frame(AVCodecContext *avctx, const AVFrame *frame)
1361 {
1362     NVENCContext               *ctx = avctx->priv_data;
1363     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1364     AVHWFramesContext   *frames_ctx = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
1365     NV_ENC_REGISTER_RESOURCE reg;
1366     int i, idx, ret;
1367
1368     for (i = 0; i < ctx->nb_registered_frames; i++) {
1369         if (ctx->registered_frames[i].ptr == (CUdeviceptr)frame->data[0])
1370             return i;
1371     }
1372
1373     idx = nvenc_find_free_reg_resource(avctx);
1374     if (idx < 0)
1375         return idx;
1376
1377     reg.version            = NV_ENC_REGISTER_RESOURCE_VER;
1378     reg.resourceType       = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
1379     reg.width              = frames_ctx->width;
1380     reg.height             = frames_ctx->height;
1381     reg.bufferFormat       = ctx->frames[0].format;
1382     reg.pitch              = frame->linesize[0];
1383     reg.resourceToRegister = frame->data[0];
1384
1385     ret = nv->nvEncRegisterResource(ctx->nvenc_ctx, &reg);
1386     if (ret != NV_ENC_SUCCESS) {
1387         nvenc_print_error(avctx, ret, "Error registering an input resource");
1388         return AVERROR_UNKNOWN;
1389     }
1390
1391     ctx->registered_frames[idx].ptr    = (CUdeviceptr)frame->data[0];
1392     ctx->registered_frames[idx].regptr = reg.registeredResource;
1393     return idx;
1394 }
1395
1396 static int nvenc_upload_frame(AVCodecContext *avctx, const AVFrame *frame,
1397                               NVENCFrame *nvenc_frame)
1398 {
1399     NVENCContext *ctx               = avctx->priv_data;
1400     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1401     int ret;
1402
1403     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1404         int reg_idx;
1405
1406         ret = nvenc_register_frame(avctx, frame);
1407         if (ret < 0) {
1408             av_log(avctx, AV_LOG_ERROR, "Could not register an input CUDA frame\n");
1409             return ret;
1410         }
1411         reg_idx = ret;
1412
1413         ret = av_frame_ref(nvenc_frame->in_ref, frame);
1414         if (ret < 0)
1415             return ret;
1416
1417         nvenc_frame->in_map.version            = NV_ENC_MAP_INPUT_RESOURCE_VER;
1418         nvenc_frame->in_map.registeredResource = ctx->registered_frames[reg_idx].regptr;
1419
1420         ret = nv->nvEncMapInputResource(ctx->nvenc_ctx, &nvenc_frame->in_map);
1421         if (ret != NV_ENC_SUCCESS) {
1422             av_frame_unref(nvenc_frame->in_ref);
1423             return nvenc_print_error(avctx, ret, "Error mapping an input resource");
1424         }
1425
1426         ctx->registered_frames[reg_idx].mapped = 1;
1427         nvenc_frame->reg_idx                   = reg_idx;
1428         nvenc_frame->in                        = nvenc_frame->in_map.mappedResource;
1429     } else {
1430         NV_ENC_LOCK_INPUT_BUFFER params = { 0 };
1431
1432         params.version     = NV_ENC_LOCK_INPUT_BUFFER_VER;
1433         params.inputBuffer = nvenc_frame->in;
1434
1435         ret = nv->nvEncLockInputBuffer(ctx->nvenc_ctx, &params);
1436         if (ret != NV_ENC_SUCCESS)
1437             return nvenc_print_error(avctx, ret, "Cannot lock the buffer");
1438
1439         ret = nvenc_copy_frame(&params, frame);
1440         if (ret < 0) {
1441             nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, nvenc_frame->in);
1442             return ret;
1443         }
1444
1445         ret = nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, nvenc_frame->in);
1446         if (ret != NV_ENC_SUCCESS)
1447             return nvenc_print_error(avctx, ret, "Cannot unlock the buffer");
1448     }
1449
1450     return 0;
1451 }
1452
1453 static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,
1454                                             NV_ENC_PIC_PARAMS *params)
1455 {
1456     NVENCContext *ctx = avctx->priv_data;
1457
1458     switch (avctx->codec->id) {
1459     case AV_CODEC_ID_H264:
1460         params->codecPicParams.h264PicParams.sliceMode =
1461             ctx->config.encodeCodecConfig.h264Config.sliceMode;
1462         params->codecPicParams.h264PicParams.sliceModeData =
1463             ctx->config.encodeCodecConfig.h264Config.sliceModeData;
1464         break;
1465     case AV_CODEC_ID_HEVC:
1466         params->codecPicParams.hevcPicParams.sliceMode =
1467             ctx->config.encodeCodecConfig.hevcConfig.sliceMode;
1468         params->codecPicParams.hevcPicParams.sliceModeData =
1469             ctx->config.encodeCodecConfig.hevcConfig.sliceModeData;
1470         break;
1471     }
1472 }
1473
1474 static inline int nvenc_enqueue_timestamp(AVFifoBuffer *f, int64_t pts)
1475 {
1476     return av_fifo_generic_write(f, &pts, sizeof(pts), NULL);
1477 }
1478
1479 static inline int nvenc_dequeue_timestamp(AVFifoBuffer *f, int64_t *pts)
1480 {
1481     return av_fifo_generic_read(f, pts, sizeof(*pts), NULL);
1482 }
1483
1484 static int nvenc_set_timestamp(AVCodecContext *avctx,
1485                                NV_ENC_LOCK_BITSTREAM *params,
1486                                AVPacket *pkt)
1487 {
1488     NVENCContext *ctx = avctx->priv_data;
1489
1490     pkt->pts      = params->outputTimeStamp;
1491     pkt->duration = params->outputDuration;
1492
1493     /* generate the first dts by linearly extrapolating the
1494      * first two pts values to the past */
1495     if (avctx->max_b_frames > 0 && !ctx->first_packet_output &&
1496         ctx->initial_pts[1] != AV_NOPTS_VALUE) {
1497         int64_t ts0 = ctx->initial_pts[0], ts1 = ctx->initial_pts[1];
1498         int64_t delta;
1499
1500         if ((ts0 < 0 && ts1 > INT64_MAX + ts0) ||
1501             (ts0 > 0 && ts1 < INT64_MIN + ts0))
1502             return AVERROR(ERANGE);
1503         delta = ts1 - ts0;
1504
1505         if ((delta < 0 && ts0 > INT64_MAX + delta) ||
1506             (delta > 0 && ts0 < INT64_MIN + delta))
1507             return AVERROR(ERANGE);
1508         pkt->dts = ts0 - delta;
1509
1510         ctx->first_packet_output = 1;
1511         return 0;
1512     }
1513     return nvenc_dequeue_timestamp(ctx->timestamps, &pkt->dts);
1514 }
1515
1516 static int nvenc_get_output(AVCodecContext *avctx, AVPacket *pkt)
1517 {
1518     NVENCContext *ctx               = avctx->priv_data;
1519     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1520     NV_ENC_LOCK_BITSTREAM params    = { 0 };
1521     NVENCFrame *frame;
1522     int ret;
1523
1524     ret = av_fifo_generic_read(ctx->ready, &frame, sizeof(frame), NULL);
1525     if (ret)
1526         return ret;
1527
1528     params.version         = NV_ENC_LOCK_BITSTREAM_VER;
1529     params.outputBitstream = frame->out;
1530
1531     ret = nv->nvEncLockBitstream(ctx->nvenc_ctx, &params);
1532     if (ret < 0)
1533         return nvenc_print_error(avctx, ret, "Cannot lock the bitstream");
1534
1535     ret = ff_alloc_packet(pkt, params.bitstreamSizeInBytes);
1536     if (ret < 0)
1537         return ret;
1538
1539     memcpy(pkt->data, params.bitstreamBufferPtr, pkt->size);
1540
1541     ret = nv->nvEncUnlockBitstream(ctx->nvenc_ctx, frame->out);
1542     if (ret < 0)
1543         return nvenc_print_error(avctx, ret, "Cannot unlock the bitstream");
1544
1545     if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
1546         nv->nvEncUnmapInputResource(ctx->nvenc_ctx, frame->in_map.mappedResource);
1547         av_frame_unref(frame->in_ref);
1548         ctx->registered_frames[frame->reg_idx].mapped = 0;
1549
1550         frame->in = NULL;
1551     }
1552
1553     av_fifo_generic_write(ctx->unused_surface_queue, &frame, sizeof(frame), NULL);
1554
1555     ret = nvenc_set_timestamp(avctx, &params, pkt);
1556     if (ret < 0)
1557         return ret;
1558
1559     switch (params.pictureType) {
1560     case NV_ENC_PIC_TYPE_IDR:
1561         pkt->flags |= AV_PKT_FLAG_KEY;
1562 #if FF_API_CODED_FRAME
1563 FF_DISABLE_DEPRECATION_WARNINGS
1564     case NV_ENC_PIC_TYPE_INTRA_REFRESH:
1565     case NV_ENC_PIC_TYPE_I:
1566         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
1567         break;
1568     case NV_ENC_PIC_TYPE_P:
1569         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_P;
1570         break;
1571     case NV_ENC_PIC_TYPE_B:
1572         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_B;
1573         break;
1574     case NV_ENC_PIC_TYPE_BI:
1575         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_BI;
1576         break;
1577 FF_ENABLE_DEPRECATION_WARNINGS
1578 #endif
1579     }
1580
1581     return 0;
1582 }
1583
1584 static int output_ready(AVCodecContext *avctx, int flush)
1585 {
1586     NVENCContext *ctx = avctx->priv_data;
1587     int nb_ready, nb_pending;
1588
1589     /* when B-frames are enabled, we wait for two initial timestamps to
1590      * calculate the first dts */
1591     if (!flush && avctx->max_b_frames > 0 &&
1592         (ctx->initial_pts[0] == AV_NOPTS_VALUE || ctx->initial_pts[1] == AV_NOPTS_VALUE))
1593         return 0;
1594
1595     nb_ready   = av_fifo_size(ctx->ready)   / sizeof(NVENCFrame*);
1596     nb_pending = av_fifo_size(ctx->pending) / sizeof(NVENCFrame*);
1597     if (flush)
1598         return nb_ready > 0;
1599     return (nb_ready > 0) && (nb_ready + nb_pending >= ctx->async_depth);
1600 }
1601
1602 int ff_nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
1603                           const AVFrame *frame, int *got_packet)
1604 {
1605     NVENCContext *ctx               = avctx->priv_data;
1606     NVENCLibraryContext *nvel       = &ctx->nvel;
1607     NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs;
1608     NV_ENC_PIC_PARAMS params        = { 0 };
1609     NVENCFrame         *nvenc_frame = NULL;
1610     CUcontext dummy;
1611     int enc_ret, ret;
1612
1613     params.version = NV_ENC_PIC_PARAMS_VER;
1614
1615     if (frame) {
1616         nvenc_frame = get_free_frame(ctx);
1617         if (!nvenc_frame) {
1618             av_log(avctx, AV_LOG_ERROR, "No free surfaces\n");
1619             return AVERROR_BUG;
1620         }
1621
1622         ret = nvenc_upload_frame(avctx, frame, nvenc_frame);
1623         if (ret < 0)
1624             return ret;
1625
1626         params.inputBuffer     = nvenc_frame->in;
1627         params.bufferFmt       = nvenc_frame->format;
1628         params.inputWidth      = frame->width;
1629         params.inputHeight     = frame->height;
1630         params.outputBitstream = nvenc_frame->out;
1631         params.inputTimeStamp  = frame->pts;
1632
1633         if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {
1634             if (frame->top_field_first)
1635                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
1636             else
1637                 params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
1638         } else {
1639             params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
1640         }
1641
1642         nvenc_codec_specific_pic_params(avctx, &params);
1643
1644         ret = nvenc_enqueue_timestamp(ctx->timestamps, frame->pts);
1645         if (ret < 0)
1646             return ret;
1647
1648         if (ctx->initial_pts[0] == AV_NOPTS_VALUE)
1649             ctx->initial_pts[0] = frame->pts;
1650         else if (ctx->initial_pts[1] == AV_NOPTS_VALUE)
1651             ctx->initial_pts[1] = frame->pts;
1652     } else {
1653         params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
1654     }
1655
1656     nvel->cu_ctx_push_current(ctx->cu_context);
1657     enc_ret = nv->nvEncEncodePicture(ctx->nvenc_ctx, &params);
1658     nvel->cu_ctx_pop_current(&dummy);
1659
1660     if (enc_ret != NV_ENC_SUCCESS &&
1661         enc_ret != NV_ENC_ERR_NEED_MORE_INPUT)
1662         return nvenc_print_error(avctx, enc_ret, "Error encoding the frame");
1663
1664     if (nvenc_frame) {
1665         ret = av_fifo_generic_write(ctx->pending, &nvenc_frame, sizeof(nvenc_frame), NULL);
1666         if (ret < 0)
1667             return ret;
1668     }
1669
1670     /* all the pending buffers are now ready for output */
1671     if (enc_ret == NV_ENC_SUCCESS) {
1672         while (av_fifo_size(ctx->pending) > 0) {
1673             av_fifo_generic_read(ctx->pending, &nvenc_frame, sizeof(nvenc_frame), NULL);
1674             av_fifo_generic_write(ctx->ready,  &nvenc_frame, sizeof(nvenc_frame), NULL);
1675         }
1676     }
1677
1678     if (output_ready(avctx, !frame)) {
1679         ret = nvenc_get_output(avctx, pkt);
1680         if (ret < 0)
1681             return ret;
1682         *got_packet = 1;
1683     } else {
1684         *got_packet = 0;
1685     }
1686
1687     return 0;
1688 }