From: Hendrik Leppkes Date: Thu, 17 Dec 2015 12:41:29 +0000 (+0100) Subject: Merge commit '1520c6ff05d835da4b793318fc88bbbc129c86a1' X-Git-Url: https://git.sesse.net/?a=commitdiff_plain;h=5fc17edc7dcec5ea2989a0d764f3cfb666a7e663;p=ffmpeg Merge commit '1520c6ff05d835da4b793318fc88bbbc129c86a1' * commit '1520c6ff05d835da4b793318fc88bbbc129c86a1': nvenc: export CPB props side data Merged-by: Hendrik Leppkes --- 5fc17edc7dcec5ea2989a0d764f3cfb666a7e663 diff --cc libavcodec/nvenc.c index 31f2dfdd334,92686c32864..a3b02fa99f4 --- a/libavcodec/nvenc.c +++ b/libavcodec/nvenc.c @@@ -458,749 -515,539 +458,757 @@@ error return 0; } -static int nvenc_setup_hevc_config(AVCodecContext *avctx) +static av_cold int nvenc_dyload_nvenc(AVCodecContext *avctx) { - NVENCContext *ctx = avctx->priv_data; - NV_ENC_CONFIG *cc = &ctx->config; - NV_ENC_CONFIG_HEVC *hevc = &cc->encodeCodecConfig.hevcConfig; + PNVENCODEAPICREATEINSTANCE nvEncodeAPICreateInstance = 0; + NVENCSTATUS nvstatus; - hevc->disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0; - hevc->repeatSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1; + NvencContext *ctx = avctx->priv_data; + NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs; - hevc->maxNumRefFramesInDPB = avctx->refs; - hevc->idrPeriod = cc->gopLength; + if (!nvenc_check_cuda(avctx)) + return 0; - /* No other profile is supported in the current SDK version 5 */ - cc->profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID; - avctx->profile = FF_PROFILE_HEVC_MAIN; + if (dl_fn->nvenc_lib) + return 1; - if (ctx->level) { - hevc->level = ctx->level; +#if defined(_WIN32) + if (sizeof(void*) == 8) { + dl_fn->nvenc_lib = LoadLibrary(TEXT("nvEncodeAPI64.dll")); } else { - hevc->level = NV_ENC_LEVEL_AUTOSELECT; + dl_fn->nvenc_lib = LoadLibrary(TEXT("nvEncodeAPI.dll")); } +#else + dl_fn->nvenc_lib = dlopen("libnvidia-encode.so.1", RTLD_LAZY); +#endif - if (ctx->tier) { - hevc->tier = ctx->tier; + if (!dl_fn->nvenc_lib) { + av_log(avctx, AV_LOG_FATAL, "Failed loading the nvenc library\n"); + goto error; } - return 0; -} -static int nvenc_setup_codec_config(AVCodecContext *avctx) -{ - switch (avctx->codec->id) { - case AV_CODEC_ID_H264: - return nvenc_setup_h264_config(avctx); - case AV_CODEC_ID_HEVC: - return nvenc_setup_hevc_config(avctx); - } - return 0; -} + nvEncodeAPICreateInstance = (PNVENCODEAPICREATEINSTANCE)LOAD_FUNC(dl_fn->nvenc_lib, "NvEncodeAPICreateInstance"); -static int nvenc_setup_encoder(AVCodecContext *avctx) -{ - NVENCContext *ctx = avctx->priv_data; - NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs; - NV_ENC_PRESET_CONFIG preset_cfg = { 0 }; - AVCPBProperties *cpb_props; - int ret; + if (!nvEncodeAPICreateInstance) { + av_log(avctx, AV_LOG_FATAL, "Failed to load nvenc entrypoint\n"); + goto error; + } - ctx->params.version = NV_ENC_INITIALIZE_PARAMS_VER; + dl_fn->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER; - ctx->params.encodeHeight = avctx->height; - ctx->params.encodeWidth = avctx->width; + nvstatus = nvEncodeAPICreateInstance(&dl_fn->nvenc_funcs); - if (avctx->sample_aspect_ratio.num && - avctx->sample_aspect_ratio.den && - (avctx->sample_aspect_ratio.num != 1 || - avctx->sample_aspect_ratio.den != 1)) { - av_reduce(&ctx->params.darWidth, - &ctx->params.darHeight, - avctx->width * avctx->sample_aspect_ratio.num, - avctx->height * avctx->sample_aspect_ratio.den, - INT_MAX / 8); - } else { - ctx->params.darHeight = avctx->height; - ctx->params.darWidth = avctx->width; + if (nvstatus != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "Failed to create nvenc instance\n"); + goto error; } - ctx->params.frameRateNum = avctx->time_base.den; - ctx->params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame; + av_log(avctx, AV_LOG_VERBOSE, "Nvenc initialized successfully\n"); - ctx->params.enableEncodeAsync = 0; - ctx->params.enablePTD = 1; + return 1; - ctx->params.encodeConfig = &ctx->config; +error: + if (dl_fn->nvenc_lib) + DL_CLOSE_FUNC(dl_fn->nvenc_lib); - nvec_map_preset(ctx); + dl_fn->nvenc_lib = NULL; - preset_cfg.version = NV_ENC_PRESET_CONFIG_VER; - preset_cfg.presetCfg.version = NV_ENC_CONFIG_VER; + return 0; +} - ret = nv->nvEncGetEncodePresetConfig(ctx->nvenc_ctx, - ctx->params.encodeGUID, - ctx->params.presetGUID, - &preset_cfg); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, - "Cannot get the preset configuration\n"); - return AVERROR_UNKNOWN; - } +static av_cold void nvenc_unload_nvenc(AVCodecContext *avctx) +{ + NvencContext *ctx = avctx->priv_data; + NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs; - memcpy(&ctx->config, &preset_cfg.presetCfg, sizeof(ctx->config)); + DL_CLOSE_FUNC(dl_fn->nvenc_lib); + dl_fn->nvenc_lib = NULL; - ctx->config.version = NV_ENC_CONFIG_VER; + dl_fn->nvenc_device_count = 0; - if (avctx->gop_size > 0) { - if (avctx->max_b_frames > 0) { - ctx->last_dts = -2; - /* 0 is intra-only, - * 1 is I/P only, - * 2 is one B Frame, - * 3 two B frames, and so on. */ - ctx->config.frameIntervalP = avctx->max_b_frames + 1; - } else if (avctx->max_b_frames == 0) { - ctx->config.frameIntervalP = 1; - } - ctx->config.gopLength = avctx->gop_size; - } else if (avctx->gop_size == 0) { - ctx->config.frameIntervalP = 0; - ctx->config.gopLength = 1; - } + DL_CLOSE_FUNC(dl_fn->cuda_lib); + dl_fn->cuda_lib = NULL; - if (ctx->config.frameIntervalP > 1) - avctx->max_b_frames = ctx->config.frameIntervalP - 1; + dl_fn->cu_init = NULL; + dl_fn->cu_device_get_count = NULL; + dl_fn->cu_device_get = NULL; + dl_fn->cu_device_get_name = NULL; + dl_fn->cu_device_compute_capability = NULL; + dl_fn->cu_ctx_create = NULL; + dl_fn->cu_ctx_pop_current = NULL; + dl_fn->cu_ctx_destroy = NULL; - nvenc_setup_rate_control(avctx); + av_log(avctx, AV_LOG_VERBOSE, "Nvenc unloaded\n"); +} - if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) { - ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD; - } else { - ctx->config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME; +static av_cold int nvenc_encode_init(AVCodecContext *avctx) +{ + NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS encode_session_params = { 0 }; + NV_ENC_PRESET_CONFIG preset_config = { 0 }; + CUcontext cu_context_curr; + CUresult cu_res; + GUID encoder_preset = NV_ENC_PRESET_HQ_GUID; + GUID codec; + NVENCSTATUS nv_status = NV_ENC_SUCCESS; ++ AVCPBProperties *cpb_props; + int surfaceCount = 0; + int i, num_mbs; + int isLL = 0; + int lossless = 0; + int res = 0; + int dw, dh; + int qp_inter_p; + + NvencContext *ctx = avctx->priv_data; + NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs; + NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs; + + if (!nvenc_dyload_nvenc(avctx)) + return AVERROR_EXTERNAL; + + ctx->last_dts = AV_NOPTS_VALUE; + + ctx->encode_config.version = NV_ENC_CONFIG_VER; + ctx->init_encode_params.version = NV_ENC_INITIALIZE_PARAMS_VER; + preset_config.version = NV_ENC_PRESET_CONFIG_VER; + preset_config.presetCfg.version = NV_ENC_CONFIG_VER; + encode_session_params.version = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER; + encode_session_params.apiVersion = NVENCAPI_VERSION; + + if (ctx->gpu >= dl_fn->nvenc_device_count) { + av_log(avctx, AV_LOG_FATAL, "Requested GPU %d, but only %d GPUs are available!\n", ctx->gpu, dl_fn->nvenc_device_count); + res = AVERROR(EINVAL); + goto error; } - if ((ret = nvenc_setup_codec_config(avctx)) < 0) - return ret; + ctx->cu_context = NULL; + cu_res = dl_fn->cu_ctx_create(&ctx->cu_context, 4, dl_fn->nvenc_devices[ctx->gpu]); // CU_CTX_SCHED_BLOCKING_SYNC=4, avoid CPU spins - ret = nv->nvEncInitializeEncoder(ctx->nvenc_ctx, &ctx->params); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, "Cannot initialize the decoder"); - return AVERROR_UNKNOWN; + if (cu_res != CUDA_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "Failed creating CUDA context for NVENC: 0x%x\n", (int)cu_res); + res = AVERROR_EXTERNAL; + goto error; } - cpb_props = ff_add_cpb_side_data(avctx); - if (!cpb_props) - return AVERROR(ENOMEM); - cpb_props->max_bitrate = avctx->rc_max_rate; - cpb_props->min_bitrate = avctx->rc_min_rate; - cpb_props->avg_bitrate = avctx->bit_rate; - cpb_props->buffer_size = avctx->rc_buffer_size; + cu_res = dl_fn->cu_ctx_pop_current(&cu_context_curr); - return 0; -} + if (cu_res != CUDA_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "Failed popping CUDA context: 0x%x\n", (int)cu_res); + res = AVERROR_EXTERNAL; + goto error; + } -static int nvenc_alloc_surface(AVCodecContext *avctx, int idx) -{ - NVENCContext *ctx = avctx->priv_data; - NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs; - int ret; - NV_ENC_CREATE_INPUT_BUFFER in_buffer = { 0 }; - NV_ENC_CREATE_BITSTREAM_BUFFER out_buffer = { 0 }; + encode_session_params.device = ctx->cu_context; + encode_session_params.deviceType = NV_ENC_DEVICE_TYPE_CUDA; - in_buffer.version = NV_ENC_CREATE_INPUT_BUFFER_VER; - out_buffer.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER; + nv_status = p_nvenc->nvEncOpenEncodeSessionEx(&encode_session_params, &ctx->nvencoder); + if (nv_status != NV_ENC_SUCCESS) { + ctx->nvencoder = NULL; + av_log(avctx, AV_LOG_FATAL, "OpenEncodeSessionEx failed: 0x%x\n", (int)nv_status); + res = AVERROR_EXTERNAL; + goto error; + } - in_buffer.width = avctx->width; - in_buffer.height = avctx->height; + if (ctx->preset) { + if (!strcmp(ctx->preset, "slow")) { + encoder_preset = NV_ENC_PRESET_HQ_GUID; + ctx->twopass = 1; + } else if (!strcmp(ctx->preset, "medium")) { + encoder_preset = NV_ENC_PRESET_HQ_GUID; + ctx->twopass = 0; + } else if (!strcmp(ctx->preset, "fast")) { + encoder_preset = NV_ENC_PRESET_HP_GUID; + ctx->twopass = 0; + } else if (!strcmp(ctx->preset, "hq")) { + encoder_preset = NV_ENC_PRESET_HQ_GUID; + } else if (!strcmp(ctx->preset, "hp")) { + encoder_preset = NV_ENC_PRESET_HP_GUID; + } else if (!strcmp(ctx->preset, "bd")) { + encoder_preset = NV_ENC_PRESET_BD_GUID; + } else if (!strcmp(ctx->preset, "ll")) { + encoder_preset = NV_ENC_PRESET_LOW_LATENCY_DEFAULT_GUID; + isLL = 1; + } else if (!strcmp(ctx->preset, "llhp")) { + encoder_preset = NV_ENC_PRESET_LOW_LATENCY_HP_GUID; + isLL = 1; + } else if (!strcmp(ctx->preset, "llhq")) { + encoder_preset = NV_ENC_PRESET_LOW_LATENCY_HQ_GUID; + isLL = 1; + } else if (!strcmp(ctx->preset, "lossless")) { + encoder_preset = NV_ENC_PRESET_LOSSLESS_DEFAULT_GUID; + lossless = 1; + } else if (!strcmp(ctx->preset, "losslesshp")) { + encoder_preset = NV_ENC_PRESET_LOSSLESS_HP_GUID; + lossless = 1; + } else if (!strcmp(ctx->preset, "default")) { + encoder_preset = NV_ENC_PRESET_DEFAULT_GUID; + } else { + av_log(avctx, AV_LOG_FATAL, "Preset \"%s\" is unknown! Supported presets: slow, medium, high, hp, hq, bd, ll, llhp, llhq, lossless, losslesshp, default\n", ctx->preset); + res = AVERROR(EINVAL); + goto error; + } + } - in_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED; + if (ctx->twopass < 0) { + ctx->twopass = isLL; + } - switch (avctx->pix_fmt) { - case AV_PIX_FMT_YUV420P: - in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_YV12_PL; - break; - case AV_PIX_FMT_NV12: - in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_NV12_PL; + switch (avctx->codec->id) { + case AV_CODEC_ID_H264: + codec = NV_ENC_CODEC_H264_GUID; break; - case AV_PIX_FMT_YUV444P: - in_buffer.bufferFmt = NV_ENC_BUFFER_FORMAT_YUV444_PL; + case AV_CODEC_ID_H265: + codec = NV_ENC_CODEC_HEVC_GUID; break; default: - return AVERROR_BUG; + av_log(avctx, AV_LOG_ERROR, "Unknown codec name\n"); + res = AVERROR(EINVAL); + goto error; } - ret = nv->nvEncCreateInputBuffer(ctx->nvenc_ctx, &in_buffer); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, "CreateInputBuffer failed\n"); - return AVERROR_UNKNOWN; + nv_status = p_nvenc->nvEncGetEncodePresetConfig(ctx->nvencoder, codec, encoder_preset, &preset_config); + if (nv_status != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "GetEncodePresetConfig failed: 0x%x\n", (int)nv_status); + res = AVERROR_EXTERNAL; + goto error; } - ctx->in[idx].in = in_buffer.inputBuffer; - ctx->in[idx].format = in_buffer.bufferFmt; - - /* 1MB is large enough to hold most output frames. - * NVENC increases this automaticaly if it's not enough. */ - out_buffer.size = BITSTREAM_BUFFER_SIZE; - - out_buffer.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_UNCACHED; + ctx->init_encode_params.encodeGUID = codec; + ctx->init_encode_params.encodeHeight = avctx->height; + ctx->init_encode_params.encodeWidth = avctx->width; - ret = nv->nvEncCreateBitstreamBuffer(ctx->nvenc_ctx, &out_buffer); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, "CreateBitstreamBuffer failed\n"); - return AVERROR_UNKNOWN; + if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den && + (avctx->sample_aspect_ratio.num != 1 || avctx->sample_aspect_ratio.num != 1)) { + av_reduce(&dw, &dh, + avctx->width * avctx->sample_aspect_ratio.num, + avctx->height * avctx->sample_aspect_ratio.den, + 1024 * 1024); + ctx->init_encode_params.darHeight = dh; + ctx->init_encode_params.darWidth = dw; + } else { + ctx->init_encode_params.darHeight = avctx->height; + ctx->init_encode_params.darWidth = avctx->width; } - ctx->out[idx].out = out_buffer.bitstreamBuffer; - ctx->out[idx].busy = 0; + // De-compensate for hardware, dubiously, trying to compensate for + // playback at 704 pixel width. + if (avctx->width == 720 && + (avctx->height == 480 || avctx->height == 576)) { + av_reduce(&dw, &dh, + ctx->init_encode_params.darWidth * 44, + ctx->init_encode_params.darHeight * 45, + 1024 * 1024); + ctx->init_encode_params.darHeight = dh; + ctx->init_encode_params.darWidth = dw; + } - return 0; -} + ctx->init_encode_params.frameRateNum = avctx->time_base.den; + ctx->init_encode_params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame; -static int nvenc_setup_surfaces(AVCodecContext *avctx) -{ - NVENCContext *ctx = avctx->priv_data; - int i, ret; + num_mbs = ((avctx->width + 15) >> 4) * ((avctx->height + 15) >> 4); + ctx->max_surface_count = (num_mbs >= 8160) ? 32 : 48; - ctx->nb_surfaces = FFMAX(4 + avctx->max_b_frames, - ctx->nb_surfaces); + if (ctx->buffer_delay >= ctx->max_surface_count) + ctx->buffer_delay = ctx->max_surface_count - 1; - ctx->in = av_mallocz(ctx->nb_surfaces * sizeof(*ctx->in)); - if (!ctx->in) - return AVERROR(ENOMEM); + ctx->init_encode_params.enableEncodeAsync = 0; + ctx->init_encode_params.enablePTD = 1; - ctx->out = av_mallocz(ctx->nb_surfaces * sizeof(*ctx->out)); - if (!ctx->out) - return AVERROR(ENOMEM); + ctx->init_encode_params.presetGUID = encoder_preset; - ctx->timestamps = av_fifo_alloc(ctx->nb_surfaces * sizeof(int64_t)); - if (!ctx->timestamps) - return AVERROR(ENOMEM); - ctx->pending = av_fifo_alloc(ctx->nb_surfaces * sizeof(ctx->out)); - if (!ctx->pending) - return AVERROR(ENOMEM); - ctx->ready = av_fifo_alloc(ctx->nb_surfaces * sizeof(ctx->out)); - if (!ctx->ready) - return AVERROR(ENOMEM); + ctx->init_encode_params.encodeConfig = &ctx->encode_config; + memcpy(&ctx->encode_config, &preset_config.presetCfg, sizeof(ctx->encode_config)); + ctx->encode_config.version = NV_ENC_CONFIG_VER; - for (i = 0; i < ctx->nb_surfaces; i++) { - if ((ret = nvenc_alloc_surface(avctx, i)) < 0) - return ret; + if (avctx->refs >= 0) { + /* 0 means "let the hardware decide" */ + switch (avctx->codec->id) { + case AV_CODEC_ID_H264: + ctx->encode_config.encodeCodecConfig.h264Config.maxNumRefFrames = avctx->refs; + break; + case AV_CODEC_ID_H265: + ctx->encode_config.encodeCodecConfig.hevcConfig.maxNumRefFramesInDPB = avctx->refs; + break; + /* Earlier switch/case will return if unknown codec is passed. */ + } } - return 0; -} + if (avctx->gop_size > 0) { + if (avctx->max_b_frames >= 0) { + /* 0 is intra-only, 1 is I/P only, 2 is one B Frame, 3 two B frames, and so on. */ + ctx->encode_config.frameIntervalP = avctx->max_b_frames + 1; + } -#define EXTRADATA_SIZE 512 + ctx->encode_config.gopLength = avctx->gop_size; + switch (avctx->codec->id) { + case AV_CODEC_ID_H264: + ctx->encode_config.encodeCodecConfig.h264Config.idrPeriod = avctx->gop_size; + break; + case AV_CODEC_ID_H265: + ctx->encode_config.encodeCodecConfig.hevcConfig.idrPeriod = avctx->gop_size; + break; + /* Earlier switch/case will return if unknown codec is passed. */ + } + } else if (avctx->gop_size == 0) { + ctx->encode_config.frameIntervalP = 0; + ctx->encode_config.gopLength = 1; + switch (avctx->codec->id) { + case AV_CODEC_ID_H264: + ctx->encode_config.encodeCodecConfig.h264Config.idrPeriod = 1; + break; + case AV_CODEC_ID_H265: + ctx->encode_config.encodeCodecConfig.hevcConfig.idrPeriod = 1; + break; + /* Earlier switch/case will return if unknown codec is passed. */ + } + } -static int nvenc_setup_extradata(AVCodecContext *avctx) -{ - NVENCContext *ctx = avctx->priv_data; - NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs; - NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 }; - int ret; + /* when there're b frames, set dts offset */ + if (ctx->encode_config.frameIntervalP >= 2) + ctx->last_dts = -2; - avctx->extradata = av_mallocz(EXTRADATA_SIZE + AV_INPUT_BUFFER_PADDING_SIZE); - if (!avctx->extradata) - return AVERROR(ENOMEM); + if (avctx->bit_rate > 0) { + ctx->encode_config.rcParams.averageBitRate = avctx->bit_rate; + } else if (ctx->encode_config.rcParams.averageBitRate > 0) { + ctx->encode_config.rcParams.maxBitRate = ctx->encode_config.rcParams.averageBitRate; + } - payload.version = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER; - payload.spsppsBuffer = avctx->extradata; - payload.inBufferSize = EXTRADATA_SIZE; - payload.outSPSPPSPayloadSize = &avctx->extradata_size; + if (avctx->rc_max_rate > 0) + ctx->encode_config.rcParams.maxBitRate = avctx->rc_max_rate; + + if (lossless) { + if (avctx->codec->id == AV_CODEC_ID_H264) + ctx->encode_config.encodeCodecConfig.h264Config.qpPrimeYZeroTransformBypassFlag = 1; + + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CONSTQP; + ctx->encode_config.rcParams.constQP.qpInterB = 0; + ctx->encode_config.rcParams.constQP.qpInterP = 0; + ctx->encode_config.rcParams.constQP.qpIntra = 0; + + avctx->qmin = -1; + avctx->qmax = -1; + } else if (ctx->cbr) { + if (!ctx->twopass) { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CBR; + } else { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_QUALITY; - ret = nv->nvEncGetSequenceParams(ctx->nvenc_ctx, &payload); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, "Cannot get the extradata\n"); - return AVERROR_UNKNOWN; - } + if (avctx->codec->id == AV_CODEC_ID_H264) { + ctx->encode_config.encodeCodecConfig.h264Config.adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE; + ctx->encode_config.encodeCodecConfig.h264Config.fmoMode = NV_ENC_H264_FMO_DISABLE; + } + } + } else if (avctx->global_quality > 0) { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CONSTQP; + ctx->encode_config.rcParams.constQP.qpInterB = avctx->global_quality; + ctx->encode_config.rcParams.constQP.qpInterP = avctx->global_quality; + ctx->encode_config.rcParams.constQP.qpIntra = avctx->global_quality; - return 0; -} + avctx->qmin = -1; + avctx->qmax = -1; + } else { + if (avctx->qmin >= 0 && avctx->qmax >= 0) { + ctx->encode_config.rcParams.enableMinQP = 1; + ctx->encode_config.rcParams.enableMaxQP = 1; + + ctx->encode_config.rcParams.minQP.qpInterB = avctx->qmin; + ctx->encode_config.rcParams.minQP.qpInterP = avctx->qmin; + ctx->encode_config.rcParams.minQP.qpIntra = avctx->qmin; + + ctx->encode_config.rcParams.maxQP.qpInterB = avctx->qmax; + ctx->encode_config.rcParams.maxQP.qpInterP = avctx->qmax; + ctx->encode_config.rcParams.maxQP.qpIntra = avctx->qmax; + + qp_inter_p = (avctx->qmax + 3 * avctx->qmin) / 4; // biased towards Qmin + + if (ctx->twopass) { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_VBR; + if (avctx->codec->id == AV_CODEC_ID_H264) { + ctx->encode_config.encodeCodecConfig.h264Config.adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE; + ctx->encode_config.encodeCodecConfig.h264Config.fmoMode = NV_ENC_H264_FMO_DISABLE; + } + } else { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR_MINQP; + } + } else { + qp_inter_p = 26; // default to 26 -av_cold int ff_nvenc_encode_close(AVCodecContext *avctx) -{ - NVENCContext *ctx = avctx->priv_data; - NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs; - int i; + if (ctx->twopass) { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_VBR; + } else { + ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR; + } + } - av_fifo_free(ctx->timestamps); - av_fifo_free(ctx->pending); - av_fifo_free(ctx->ready); + ctx->encode_config.rcParams.enableInitialRCQP = 1; + ctx->encode_config.rcParams.initialRCQP.qpInterP = qp_inter_p; - if (ctx->in) { - for (i = 0; i < ctx->nb_surfaces; ++i) { - nv->nvEncDestroyInputBuffer(ctx->nvenc_ctx, ctx->in[i].in); - nv->nvEncDestroyBitstreamBuffer(ctx->nvenc_ctx, ctx->out[i].out); + if(avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) { + ctx->encode_config.rcParams.initialRCQP.qpIntra = av_clip( + qp_inter_p * fabs(avctx->i_quant_factor) + avctx->i_quant_offset, 0, 51); + ctx->encode_config.rcParams.initialRCQP.qpInterB = av_clip( + qp_inter_p * fabs(avctx->b_quant_factor) + avctx->b_quant_offset, 0, 51); + } else { + ctx->encode_config.rcParams.initialRCQP.qpIntra = qp_inter_p; + ctx->encode_config.rcParams.initialRCQP.qpInterB = qp_inter_p; } } - av_freep(&ctx->in); - av_freep(&ctx->out); + if (avctx->rc_buffer_size > 0) { + ctx->encode_config.rcParams.vbvBufferSize = avctx->rc_buffer_size; + } else if (ctx->encode_config.rcParams.averageBitRate > 0) { + ctx->encode_config.rcParams.vbvBufferSize = 2 * ctx->encode_config.rcParams.averageBitRate; + } - if (ctx->nvenc_ctx) - nv->nvEncDestroyEncoder(ctx->nvenc_ctx); + if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) { + ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD; + } else { + ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME; + } - if (ctx->cu_context) - ctx->nvel.cu_ctx_destroy(ctx->cu_context); + switch (avctx->codec->id) { + case AV_CODEC_ID_H264: + ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourDescriptionPresentFlag = 1; + ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoSignalTypePresentFlag = 1; + + ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourMatrix = avctx->colorspace; + ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourPrimaries = avctx->color_primaries; + ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.transferCharacteristics = avctx->color_trc; + + ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoFullRangeFlag = avctx->color_range == AVCOL_RANGE_JPEG; + + ctx->encode_config.encodeCodecConfig.h264Config.sliceMode = 3; + ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData = 1; + + ctx->encode_config.encodeCodecConfig.h264Config.disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0; + ctx->encode_config.encodeCodecConfig.h264Config.repeatSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1; + + if (!ctx->profile) { + switch (avctx->profile) { + case FF_PROFILE_H264_HIGH_444_PREDICTIVE: + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID; + break; + case FF_PROFILE_H264_BASELINE: + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID; + break; + case FF_PROFILE_H264_MAIN: + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID; + break; + case FF_PROFILE_H264_HIGH: + case FF_PROFILE_UNKNOWN: + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID; + break; + default: + av_log(avctx, AV_LOG_WARNING, "Unsupported profile requested, falling back to high\n"); + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID; + break; + } + } else { + if (!strcmp(ctx->profile, "high")) { + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID; + avctx->profile = FF_PROFILE_H264_HIGH; + } else if (!strcmp(ctx->profile, "main")) { + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID; + avctx->profile = FF_PROFILE_H264_MAIN; + } else if (!strcmp(ctx->profile, "baseline")) { + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID; + avctx->profile = FF_PROFILE_H264_BASELINE; + } else if (!strcmp(ctx->profile, "high444p")) { + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID; + avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE; + } else { + av_log(avctx, AV_LOG_FATAL, "Profile \"%s\" is unknown! Supported profiles: high, main, baseline\n", ctx->profile); + res = AVERROR(EINVAL); + goto error; + } + } - if (ctx->nvel.nvenc) - dlclose(ctx->nvel.nvenc); + // force setting profile as high444p if input is AV_PIX_FMT_YUV444P + if (avctx->pix_fmt == AV_PIX_FMT_YUV444P) { + ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID; + avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE; + } - if (ctx->nvel.cuda) - dlclose(ctx->nvel.cuda); + ctx->encode_config.encodeCodecConfig.h264Config.chromaFormatIDC = avctx->profile == FF_PROFILE_H264_HIGH_444_PREDICTIVE ? 3 : 1; - return 0; -} + if (ctx->level) { + res = input_string_to_uint32(avctx, nvenc_h264_level_pairs, ctx->level, &ctx->encode_config.encodeCodecConfig.h264Config.level); -av_cold int ff_nvenc_encode_init(AVCodecContext *avctx) -{ - int ret; + if (res) { + av_log(avctx, AV_LOG_FATAL, "Level \"%s\" is unknown! Supported levels: auto, 1, 1b, 1.1, 1.2, 1.3, 2, 2.1, 2.2, 3, 3.1, 3.2, 4, 4.1, 4.2, 5, 5.1\n", ctx->level); + goto error; + } + } else { + ctx->encode_config.encodeCodecConfig.h264Config.level = NV_ENC_LEVEL_AUTOSELECT; + } - if ((ret = nvenc_load_libraries(avctx)) < 0) - return ret; + break; + case AV_CODEC_ID_H265: + ctx->encode_config.encodeCodecConfig.hevcConfig.sliceMode = 3; + ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData = 1; - if ((ret = nvenc_setup_device(avctx)) < 0) - return ret; + ctx->encode_config.encodeCodecConfig.hevcConfig.disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0; + ctx->encode_config.encodeCodecConfig.hevcConfig.repeatSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1; - if ((ret = nvenc_setup_encoder(avctx)) < 0) - return ret; + /* No other profile is supported in the current SDK version 5 */ + ctx->encode_config.profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID; + avctx->profile = FF_PROFILE_HEVC_MAIN; - if ((ret = nvenc_setup_surfaces(avctx)) < 0) - return ret; + if (ctx->level) { + res = input_string_to_uint32(avctx, nvenc_hevc_level_pairs, ctx->level, &ctx->encode_config.encodeCodecConfig.hevcConfig.level); - if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) { - if ((ret = nvenc_setup_extradata(avctx)) < 0) - return ret; - } + if (res) { + av_log(avctx, AV_LOG_FATAL, "Level \"%s\" is unknown! Supported levels: auto, 1, 2, 2.1, 3, 3.1, 4, 4.1, 5, 5.1, 5.2, 6, 6.1, 6.2\n", ctx->level); + goto error; + } + } else { + ctx->encode_config.encodeCodecConfig.hevcConfig.level = NV_ENC_LEVEL_AUTOSELECT; + } - return 0; -} + if (ctx->tier) { + if (!strcmp(ctx->tier, "main")) { + ctx->encode_config.encodeCodecConfig.hevcConfig.tier = NV_ENC_TIER_HEVC_MAIN; + } else if (!strcmp(ctx->tier, "high")) { + ctx->encode_config.encodeCodecConfig.hevcConfig.tier = NV_ENC_TIER_HEVC_HIGH; + } else { + av_log(avctx, AV_LOG_FATAL, "Tier \"%s\" is unknown! Supported tiers: main, high\n", ctx->tier); + res = AVERROR(EINVAL); + goto error; + } + } -static NVENCInputSurface *get_input_surface(NVENCContext *ctx) -{ - int i; + break; + /* Earlier switch/case will return if unknown codec is passed. */ + } - for (i = 0; i < ctx->nb_surfaces; i++) { - if (!ctx->in[i].locked) { - ctx->in[i].locked = 1; - return &ctx->in[i]; - } + nv_status = p_nvenc->nvEncInitializeEncoder(ctx->nvencoder, &ctx->init_encode_params); + if (nv_status != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "InitializeEncoder failed: 0x%x\n", (int)nv_status); + res = AVERROR_EXTERNAL; + goto error; } - return NULL; -} + ctx->input_surfaces = av_malloc(ctx->max_surface_count * sizeof(*ctx->input_surfaces)); -static NVENCOutputSurface *get_output_surface(NVENCContext *ctx) -{ - int i; + if (!ctx->input_surfaces) { + res = AVERROR(ENOMEM); + goto error; + } - for (i = 0; i < ctx->nb_surfaces; i++) { - if (!ctx->out[i].busy) { - return &ctx->out[i]; - } + ctx->output_surfaces = av_malloc(ctx->max_surface_count * sizeof(*ctx->output_surfaces)); + + if (!ctx->output_surfaces) { + res = AVERROR(ENOMEM); + goto error; } - return NULL; -} + for (surfaceCount = 0; surfaceCount < ctx->max_surface_count; ++surfaceCount) { + NV_ENC_CREATE_INPUT_BUFFER allocSurf = { 0 }; + NV_ENC_CREATE_BITSTREAM_BUFFER allocOut = { 0 }; + allocSurf.version = NV_ENC_CREATE_INPUT_BUFFER_VER; + allocOut.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER; -static int nvenc_copy_frame(NV_ENC_LOCK_INPUT_BUFFER *in, const AVFrame *frame) -{ - uint8_t *buf = in->bufferDataPtr; - int off = frame->height * in->pitch; + allocSurf.width = (avctx->width + 31) & ~31; + allocSurf.height = (avctx->height + 31) & ~31; - switch (frame->format) { - case AV_PIX_FMT_YUV420P: - av_image_copy_plane(buf, in->pitch, - frame->data[0], frame->linesize[0], - frame->width, frame->height); - buf += off; + allocSurf.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED; - av_image_copy_plane(buf, in->pitch >> 1, - frame->data[2], frame->linesize[2], - frame->width >> 1, frame->height >> 1); + switch (avctx->pix_fmt) { + case AV_PIX_FMT_YUV420P: + allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_YV12_PL; + break; - buf += off >> 2; + case AV_PIX_FMT_NV12: + allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_NV12_PL; + break; - av_image_copy_plane(buf, in->pitch >> 1, - frame->data[1], frame->linesize[1], - frame->width >> 1, frame->height >> 1); - break; - case AV_PIX_FMT_NV12: - av_image_copy_plane(buf, in->pitch, - frame->data[0], frame->linesize[0], - frame->width, frame->height); - buf += off; - - av_image_copy_plane(buf, in->pitch, - frame->data[1], frame->linesize[1], - frame->width, frame->height >> 1); - break; - case AV_PIX_FMT_YUV444P: - av_image_copy_plane(buf, in->pitch, - frame->data[0], frame->linesize[0], - frame->width, frame->height); - buf += off; - - av_image_copy_plane(buf, in->pitch, - frame->data[1], frame->linesize[1], - frame->width, frame->height); - buf += off; - - av_image_copy_plane(buf, in->pitch, - frame->data[2], frame->linesize[2], - frame->width, frame->height); - break; - default: - return AVERROR_BUG; - } + case AV_PIX_FMT_YUV444P: + allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_YUV444_PL; + break; - return 0; -} + default: + av_log(avctx, AV_LOG_FATAL, "Invalid input pixel format\n"); + res = AVERROR(EINVAL); + goto error; + } -static int nvenc_enqueue_frame(AVCodecContext *avctx, const AVFrame *frame, - NVENCInputSurface **in_surf) -{ - NVENCContext *ctx = avctx->priv_data; - NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs; - NV_ENC_LOCK_INPUT_BUFFER params = { 0 }; - NVENCInputSurface *in = get_input_surface(ctx); - int ret; + nv_status = p_nvenc->nvEncCreateInputBuffer(ctx->nvencoder, &allocSurf); + if (nv_status != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "CreateInputBuffer failed\n"); + res = AVERROR_EXTERNAL; + goto error; + } + + ctx->input_surfaces[surfaceCount].lockCount = 0; + ctx->input_surfaces[surfaceCount].input_surface = allocSurf.inputBuffer; + ctx->input_surfaces[surfaceCount].format = allocSurf.bufferFmt; + ctx->input_surfaces[surfaceCount].width = allocSurf.width; + ctx->input_surfaces[surfaceCount].height = allocSurf.height; - if (!in) - return AVERROR_BUG; + /* 1MB is large enough to hold most output frames. NVENC increases this automaticaly if it's not enough. */ + allocOut.size = 1024 * 1024; - params.version = NV_ENC_LOCK_INPUT_BUFFER_VER; - params.inputBuffer = in->in; + allocOut.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED; + nv_status = p_nvenc->nvEncCreateBitstreamBuffer(ctx->nvencoder, &allocOut); + if (nv_status != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "CreateBitstreamBuffer failed\n"); + ctx->output_surfaces[surfaceCount++].output_surface = NULL; + res = AVERROR_EXTERNAL; + goto error; + } - ret = nv->nvEncLockInputBuffer(ctx->nvenc_ctx, ¶ms); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, "Cannot lock the buffer %p.\n", - in); - return AVERROR_UNKNOWN; + ctx->output_surfaces[surfaceCount].output_surface = allocOut.bitstreamBuffer; + ctx->output_surfaces[surfaceCount].size = allocOut.size; + ctx->output_surfaces[surfaceCount].busy = 0; } - ret = nvenc_copy_frame(¶ms, frame); - if (ret < 0) - goto fail; + if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) { + uint32_t outSize = 0; + char tmpHeader[256]; + NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 }; + payload.version = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER; + + payload.spsppsBuffer = tmpHeader; + payload.inBufferSize = sizeof(tmpHeader); + payload.outSPSPPSPayloadSize = &outSize; + + nv_status = p_nvenc->nvEncGetSequenceParams(ctx->nvencoder, &payload); + if (nv_status != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_FATAL, "GetSequenceParams failed\n"); + goto error; + } - ret = nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in); - if (ret != NV_ENC_SUCCESS) { - av_log(avctx, AV_LOG_ERROR, "Cannot unlock the buffer %p.\n", - in); - return AVERROR_UNKNOWN; + avctx->extradata_size = outSize; + avctx->extradata = av_mallocz(outSize + AV_INPUT_BUFFER_PADDING_SIZE); + + if (!avctx->extradata) { + res = AVERROR(ENOMEM); + goto error; + } + + memcpy(avctx->extradata, tmpHeader, outSize); } - *in_surf = in; + if (ctx->encode_config.frameIntervalP > 1) + avctx->has_b_frames = 2; - return 0; + if (ctx->encode_config.rcParams.averageBitRate > 0) + avctx->bit_rate = ctx->encode_config.rcParams.averageBitRate; -fail: - nv->nvEncUnlockInputBuffer(ctx->nvenc_ctx, in->in); ++ cpb_props = ff_add_cpb_side_data(avctx); ++ if (!cpb_props) ++ return AVERROR(ENOMEM); ++ cpb_props->max_bitrate = ctx->encode_config.rcParams.maxBitRate; ++ cpb_props->avg_bitrate = avctx->bit_rate; ++ cpb_props->buffer_size = ctx->encode_config.rcParams.vbvBufferSize; + - return ret; -} + return 0; -static void nvenc_codec_specific_pic_params(AVCodecContext *avctx, - NV_ENC_PIC_PARAMS *params) -{ - NVENCContext *ctx = avctx->priv_data; +error: - switch (avctx->codec->id) { - case AV_CODEC_ID_H264: - params->codecPicParams.h264PicParams.sliceMode = - ctx->config.encodeCodecConfig.h264Config.sliceMode; - params->codecPicParams.h264PicParams.sliceModeData = - ctx->config.encodeCodecConfig.h264Config.sliceModeData; - break; - case AV_CODEC_ID_HEVC: - params->codecPicParams.hevcPicParams.sliceMode = - ctx->config.encodeCodecConfig.hevcConfig.sliceMode; - params->codecPicParams.hevcPicParams.sliceModeData = - ctx->config.encodeCodecConfig.hevcConfig.sliceModeData; - break; + for (i = 0; i < surfaceCount; ++i) { + p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->input_surfaces[i].input_surface); + if (ctx->output_surfaces[i].output_surface) + p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->output_surfaces[i].output_surface); } -} -static inline int nvenc_enqueue_timestamp(AVFifoBuffer *f, int64_t pts) -{ - return av_fifo_generic_write(f, &pts, sizeof(pts), NULL); -} + if (ctx->nvencoder) + p_nvenc->nvEncDestroyEncoder(ctx->nvencoder); -static inline int nvenc_dequeue_timestamp(AVFifoBuffer *f, int64_t *pts) -{ - return av_fifo_generic_read(f, pts, sizeof(*pts), NULL); -} + if (ctx->cu_context) + dl_fn->cu_ctx_destroy(ctx->cu_context); -static inline int nvenc_enqueue_surface(AVFifoBuffer *f, - NVENCOutputSurface *surf) -{ - surf->busy = 1; - return av_fifo_generic_write(f, &surf, sizeof(surf), NULL); -} + nvenc_unload_nvenc(avctx); -static inline int nvenc_dequeue_surface(AVFifoBuffer *f, - NVENCOutputSurface **surf) -{ - return av_fifo_generic_read(f, surf, sizeof(*surf), NULL); + ctx->nvencoder = NULL; + ctx->cu_context = NULL; + + return res; } -static int nvenc_set_timestamp(NVENCContext *ctx, - NV_ENC_LOCK_BITSTREAM *params, - AVPacket *pkt) +static av_cold int nvenc_encode_close(AVCodecContext *avctx) { - pkt->pts = params->outputTimeStamp; - pkt->duration = params->outputDuration; + NvencContext *ctx = avctx->priv_data; + NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs; + NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs; + int i; + + av_freep(&ctx->timestamp_list.data); + av_freep(&ctx->output_surface_ready_queue.data); + av_freep(&ctx->output_surface_queue.data); - return nvenc_dequeue_timestamp(ctx->timestamps, &pkt->dts); + for (i = 0; i < ctx->max_surface_count; ++i) { + p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->input_surfaces[i].input_surface); + p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->output_surfaces[i].output_surface); + } + ctx->max_surface_count = 0; + + p_nvenc->nvEncDestroyEncoder(ctx->nvencoder); + ctx->nvencoder = NULL; + + dl_fn->cu_ctx_destroy(ctx->cu_context); + ctx->cu_context = NULL; + + nvenc_unload_nvenc(avctx); + + return 0; } -static int nvenc_get_frame(AVCodecContext *avctx, AVPacket *pkt) +static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencOutputSurface *tmpoutsurf) { - NVENCContext *ctx = avctx->priv_data; - NV_ENCODE_API_FUNCTION_LIST *nv = &ctx->nvel.nvenc_funcs; - NV_ENC_LOCK_BITSTREAM params = { 0 }; - NVENCOutputSurface *out = NULL; - int ret; + NvencContext *ctx = avctx->priv_data; + NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs; + NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs; - ret = nvenc_dequeue_surface(ctx->pending, &out); - if (ret) - return ret; + uint32_t slice_mode_data; + uint32_t *slice_offsets; + NV_ENC_LOCK_BITSTREAM lock_params = { 0 }; + NVENCSTATUS nv_status; + int res = 0; - params.version = NV_ENC_LOCK_BITSTREAM_VER; - params.outputBitstream = out->out; + switch (avctx->codec->id) { + case AV_CODEC_ID_H264: + slice_mode_data = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData; + break; + case AV_CODEC_ID_H265: + slice_mode_data = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData; + break; + default: + av_log(avctx, AV_LOG_ERROR, "Unknown codec name\n"); + res = AVERROR(EINVAL); + goto error; + } + slice_offsets = av_mallocz(slice_mode_data * sizeof(*slice_offsets)); - ret = nv->nvEncLockBitstream(ctx->nvenc_ctx, ¶ms); - if (ret < 0) - return AVERROR_UNKNOWN; + if (!slice_offsets) + return AVERROR(ENOMEM); + + lock_params.version = NV_ENC_LOCK_BITSTREAM_VER; - ret = ff_alloc_packet(pkt, params.bitstreamSizeInBytes); - if (ret < 0) - return ret; + lock_params.doNotWait = 0; + lock_params.outputBitstream = tmpoutsurf->output_surface; + lock_params.sliceOffsets = slice_offsets; - memcpy(pkt->data, params.bitstreamBufferPtr, pkt->size); + nv_status = p_nvenc->nvEncLockBitstream(ctx->nvencoder, &lock_params); + if (nv_status != NV_ENC_SUCCESS) { + av_log(avctx, AV_LOG_ERROR, "Failed locking bitstream buffer\n"); + res = AVERROR_EXTERNAL; + goto error; + } - ret = nv->nvEncUnlockBitstream(ctx->nvenc_ctx, out->out); - if (ret < 0) - return AVERROR_UNKNOWN; + if (res = ff_alloc_packet2(avctx, pkt, lock_params.bitstreamSizeInBytes,0)) { + p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface); + goto error; + } - out->busy = out->in->locked = 0; + memcpy(pkt->data, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes); - ret = nvenc_set_timestamp(ctx, ¶ms, pkt); - if (ret < 0) - return ret; + nv_status = p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface); + if (nv_status != NV_ENC_SUCCESS) + av_log(avctx, AV_LOG_ERROR, "Failed unlocking bitstream buffer, expect the gates of mordor to open\n"); - switch (params.pictureType) { + switch (lock_params.pictureType) { case NV_ENC_PIC_TYPE_IDR: pkt->flags |= AV_PKT_FLAG_KEY; #if FF_API_CODED_FRAME