LCOV - code coverage report
Current view: top level - libavcodec - nvenc.c (source / functions) Hit Total Coverage
Test: coverage.info Lines: 0 1163 0.0 %
Date: 2017-12-15 11:05:35 Functions: 0 43 0.0 %

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

Generated by: LCOV version 1.13