You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1745 lines
58KB

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