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