falkTX
/
FFmpeg
mirror of https://github.com/falkTX/FFmpeg.git
1
0
Fork 0
Code Issues 0 Releases 338 Wiki Activity
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.
80089 Commits
32 Branches
646MB
Tree: cfb49fc6f7
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'cfb49fc6f7'
${ noResults }
FFmpeg/libavcodec/nvenc.c

1701 lines
58KB
Raw Blame History 

  1. /*

  2.  * H.264 hardware encoding using nvidia nvenc

  3.  * Copyright (c) 2014 Timo Rothenpieler <timo@rothenpieler.org>

  4.  *

  5.  * This file is part of FFmpeg.

  6.  *

  7.  * FFmpeg is free software; you can redistribute it and/or

  8.  * modify it under the terms of the GNU Lesser General Public

  9.  * License as published by the Free Software Foundation; either

  10.  * version 2.1 of the License, or (at your option) any later version.

  11.  *

  12.  * FFmpeg is distributed in the hope that it will be useful,

  13.  * but WITHOUT ANY WARRANTY; without even the implied warranty of

  14.  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

  15.  * Lesser General Public License for more details.

  16.  *

  17.  * You should have received a copy of the GNU Lesser General Public

  18.  * License along with FFmpeg; if not, write to the Free Software

  19.  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

  20.  */

  21. 

  22. #if defined(_WIN32)

  23. #include <windows.h>

  24. #else

  25. #include <dlfcn.h>

  26. #endif

  27. 

  28. #include <nvEncodeAPI.h>

  29. 

  30. #include "libavutil/fifo.h"

  31. #include "libavutil/internal.h"

  32. #include "libavutil/imgutils.h"

  33. #include "libavutil/avassert.h"

  34. #include "libavutil/opt.h"

  35. #include "libavutil/mem.h"

  36. #include "avcodec.h"

  37. #include "internal.h"

  38. #include "thread.h"

  39. 

  40. #if defined(_WIN32)

  41. #define CUDAAPI __stdcall

  42. #else

  43. #define CUDAAPI

  44. #endif

  45. 

  46. #if defined(_WIN32)

  47. #define LOAD_FUNC(l, s) GetProcAddress(l, s)

  48. #define DL_CLOSE_FUNC(l) FreeLibrary(l)

  49. #else

  50. #define LOAD_FUNC(l, s) dlsym(l, s)

  51. #define DL_CLOSE_FUNC(l) dlclose(l)

  52. #endif

  53. 

  54. typedef enum cudaError_enum {

  55.     CUDA_SUCCESS = 0

  56. } CUresult;

  57. typedef int CUdevice;

  58. typedef void* CUcontext;

  59. 

  60. typedef CUresult(CUDAAPI *PCUINIT)(unsigned int Flags);

  61. typedef CUresult(CUDAAPI *PCUDEVICEGETCOUNT)(int *count);

  62. typedef CUresult(CUDAAPI *PCUDEVICEGET)(CUdevice *device, int ordinal);

  63. typedef CUresult(CUDAAPI *PCUDEVICEGETNAME)(char *name, int len, CUdevice dev);

  64. typedef CUresult(CUDAAPI *PCUDEVICECOMPUTECAPABILITY)(int *major, int *minor, CUdevice dev);

  65. typedef CUresult(CUDAAPI *PCUCTXCREATE)(CUcontext *pctx, unsigned int flags, CUdevice dev);

  66. typedef CUresult(CUDAAPI *PCUCTXPOPCURRENT)(CUcontext *pctx);

  67. typedef CUresult(CUDAAPI *PCUCTXDESTROY)(CUcontext ctx);

  68. 

  69. typedef NVENCSTATUS (NVENCAPI* PNVENCODEAPICREATEINSTANCE)(NV_ENCODE_API_FUNCTION_LIST *functionList);

  70. 

  71. typedef struct NvencSurface

  72. {

  73.     NV_ENC_INPUT_PTR input_surface;

  74.     int width;

  75.     int height;

  76. 

  77.     int lockCount;

  78. 

  79.     NV_ENC_BUFFER_FORMAT format;

  80. 

  81.     NV_ENC_OUTPUT_PTR output_surface;

  82.     int size;

  83. } NvencSurface;

  84. 

  85. typedef struct NvencData

  86. {

  87.     union {

  88.         int64_t timestamp;

  89.         NvencSurface *surface;

  90.     } u;

  91. } NvencData;

  92. 

  93. typedef struct NvencDynLoadFunctions

  94. {

  95.     PCUINIT cu_init;

  96.     PCUDEVICEGETCOUNT cu_device_get_count;

  97.     PCUDEVICEGET cu_device_get;

  98.     PCUDEVICEGETNAME cu_device_get_name;

  99.     PCUDEVICECOMPUTECAPABILITY cu_device_compute_capability;

  100.     PCUCTXCREATE cu_ctx_create;

  101.     PCUCTXPOPCURRENT cu_ctx_pop_current;

  102.     PCUCTXDESTROY cu_ctx_destroy;

  103. 

  104.     NV_ENCODE_API_FUNCTION_LIST nvenc_funcs;

  105.     int nvenc_device_count;

  106.     CUdevice nvenc_devices[16];

  107. 

  108. #if defined(_WIN32)

  109.     HMODULE cuda_lib;

  110.     HMODULE nvenc_lib;

  111. #else

  112.     void* cuda_lib;

  113.     void* nvenc_lib;

  114. #endif

  115. } NvencDynLoadFunctions;

  116. 

  117. typedef struct NvencValuePair

  118. {

  119.     const char *str;

  120.     uint32_t num;

  121. } NvencValuePair;

  122. 

  123. typedef struct NvencContext

  124. {

  125.     AVClass *avclass;

  126. 

  127.     NvencDynLoadFunctions nvenc_dload_funcs;

  128. 

  129.     NV_ENC_INITIALIZE_PARAMS init_encode_params;

  130.     NV_ENC_CONFIG encode_config;

  131.     CUcontext cu_context;

  132. 

  133.     int max_surface_count;

  134.     NvencSurface *surfaces;

  135. 

  136.     AVFifoBuffer *output_surface_queue;

  137.     AVFifoBuffer *output_surface_ready_queue;

  138.     AVFifoBuffer *timestamp_list;

  139.     int64_t last_dts;

  140. 

  141.     void *nvencoder;

  142. 

  143.     char *preset;

  144.     char *profile;

  145.     char *level;

  146.     char *tier;

  147.     int cbr;

  148.     int twopass;

  149.     int gpu;

  150.     int buffer_delay;

  151. } NvencContext;

  152. 

  153. static const NvencValuePair nvenc_h264_level_pairs[] = {

  154.     { "auto", NV_ENC_LEVEL_AUTOSELECT },

  155.     { "1"   , NV_ENC_LEVEL_H264_1     },

  156.     { "1.0" , NV_ENC_LEVEL_H264_1     },

  157.     { "1b"  , NV_ENC_LEVEL_H264_1b    },

  158.     { "1.0b", NV_ENC_LEVEL_H264_1b    },

  159.     { "1.1" , NV_ENC_LEVEL_H264_11    },

  160.     { "1.2" , NV_ENC_LEVEL_H264_12    },

  161.     { "1.3" , NV_ENC_LEVEL_H264_13    },

  162.     { "2"   , NV_ENC_LEVEL_H264_2     },

  163.     { "2.0" , NV_ENC_LEVEL_H264_2     },

  164.     { "2.1" , NV_ENC_LEVEL_H264_21    },

  165.     { "2.2" , NV_ENC_LEVEL_H264_22    },

  166.     { "3"   , NV_ENC_LEVEL_H264_3     },

  167.     { "3.0" , NV_ENC_LEVEL_H264_3     },

  168.     { "3.1" , NV_ENC_LEVEL_H264_31    },

  169.     { "3.2" , NV_ENC_LEVEL_H264_32    },

  170.     { "4"   , NV_ENC_LEVEL_H264_4     },

  171.     { "4.0" , NV_ENC_LEVEL_H264_4     },

  172.     { "4.1" , NV_ENC_LEVEL_H264_41    },

  173.     { "4.2" , NV_ENC_LEVEL_H264_42    },

  174.     { "5"   , NV_ENC_LEVEL_H264_5     },

  175.     { "5.0" , NV_ENC_LEVEL_H264_5     },

  176.     { "5.1" , NV_ENC_LEVEL_H264_51    },

  177.     { NULL }

  178. };

  179. 

  180. static const NvencValuePair nvenc_hevc_level_pairs[] = {

  181.     { "auto", NV_ENC_LEVEL_AUTOSELECT },

  182.     { "1"   , NV_ENC_LEVEL_HEVC_1     },

  183.     { "1.0" , NV_ENC_LEVEL_HEVC_1     },

  184.     { "2"   , NV_ENC_LEVEL_HEVC_2     },

  185.     { "2.0" , NV_ENC_LEVEL_HEVC_2     },

  186.     { "2.1" , NV_ENC_LEVEL_HEVC_21    },

  187.     { "3"   , NV_ENC_LEVEL_HEVC_3     },

  188.     { "3.0" , NV_ENC_LEVEL_HEVC_3     },

  189.     { "3.1" , NV_ENC_LEVEL_HEVC_31    },

  190.     { "4"   , NV_ENC_LEVEL_HEVC_4     },

  191.     { "4.0" , NV_ENC_LEVEL_HEVC_4     },

  192.     { "4.1" , NV_ENC_LEVEL_HEVC_41    },

  193.     { "5"   , NV_ENC_LEVEL_HEVC_5     },

  194.     { "5.0" , NV_ENC_LEVEL_HEVC_5     },

  195.     { "5.1" , NV_ENC_LEVEL_HEVC_51    },

  196.     { "5.2" , NV_ENC_LEVEL_HEVC_52    },

  197.     { "6"   , NV_ENC_LEVEL_HEVC_6     },

  198.     { "6.0" , NV_ENC_LEVEL_HEVC_6     },

  199.     { "6.1" , NV_ENC_LEVEL_HEVC_61    },

  200.     { "6.2" , NV_ENC_LEVEL_HEVC_62    },

  201.     { NULL }

  202. };

  203. 

  204. static const struct {

  205.     NVENCSTATUS nverr;

  206.     int         averr;

  207.     const char *desc;

  208. } nvenc_errors[] = {

  209.     { NV_ENC_SUCCESS,                      0,                "success"                  },

  210.     { NV_ENC_ERR_NO_ENCODE_DEVICE,         AVERROR(ENOENT),  "no encode device"         },

  211.     { NV_ENC_ERR_UNSUPPORTED_DEVICE,       AVERROR(ENOSYS),  "unsupported device"       },

  212.     { NV_ENC_ERR_INVALID_ENCODERDEVICE,    AVERROR(EINVAL),  "invalid encoder device"   },

  213.     { NV_ENC_ERR_INVALID_DEVICE,           AVERROR(EINVAL),  "invalid device"           },

  214.     { NV_ENC_ERR_DEVICE_NOT_EXIST,         AVERROR(EIO),     "device does not exist"    },

  215.     { NV_ENC_ERR_INVALID_PTR,              AVERROR(EFAULT),  "invalid ptr"              },

  216.     { NV_ENC_ERR_INVALID_EVENT,            AVERROR(EINVAL),  "invalid event"            },

  217.     { NV_ENC_ERR_INVALID_PARAM,            AVERROR(EINVAL),  "invalid param"            },

  218.     { NV_ENC_ERR_INVALID_CALL,             AVERROR(EINVAL),  "invalid call"             },

  219.     { NV_ENC_ERR_OUT_OF_MEMORY,            AVERROR(ENOMEM),  "out of memory"            },

  220.     { NV_ENC_ERR_ENCODER_NOT_INITIALIZED,  AVERROR(EINVAL),  "encoder not initialized"  },

  221.     { NV_ENC_ERR_UNSUPPORTED_PARAM,        AVERROR(ENOSYS),  "unsupported param"        },

  222.     { NV_ENC_ERR_LOCK_BUSY,                AVERROR(EAGAIN),  "lock busy"                },

  223.     { NV_ENC_ERR_NOT_ENOUGH_BUFFER,        AVERROR(ENOBUFS), "not enough buffer"        },

  224.     { NV_ENC_ERR_INVALID_VERSION,          AVERROR(EINVAL),  "invalid version"          },

  225.     { NV_ENC_ERR_MAP_FAILED,               AVERROR(EIO),     "map failed"               },

  226.     { NV_ENC_ERR_NEED_MORE_INPUT,          AVERROR(EAGAIN),  "need more input"          },

  227.     { NV_ENC_ERR_ENCODER_BUSY,             AVERROR(EAGAIN),  "encoder busy"             },

  228.     { NV_ENC_ERR_EVENT_NOT_REGISTERD,      AVERROR(EBADF),   "event not registered"     },

  229.     { NV_ENC_ERR_GENERIC,                  AVERROR_UNKNOWN,  "generic error"            },

  230.     { NV_ENC_ERR_INCOMPATIBLE_CLIENT_KEY,  AVERROR(EINVAL),  "incompatible client key"  },

  231.     { NV_ENC_ERR_UNIMPLEMENTED,            AVERROR(ENOSYS),  "unimplemented"            },

  232.     { NV_ENC_ERR_RESOURCE_REGISTER_FAILED, AVERROR(EIO),     "resource register failed" },

  233.     { NV_ENC_ERR_RESOURCE_NOT_REGISTERED,  AVERROR(EBADF),   "resource not registered"  },

  234.     { NV_ENC_ERR_RESOURCE_NOT_MAPPED,      AVERROR(EBADF),   "resource not mapped"      },

  235. };

  236. 

  237. static int nvenc_map_error(NVENCSTATUS err, const char **desc)

  238. {

  239.     int i;

  240.     for (i = 0; i < FF_ARRAY_ELEMS(nvenc_errors); i++) {

  241.         if (nvenc_errors[i].nverr == err) {

  242.             if (desc)

  243.                 *desc = nvenc_errors[i].desc;

  244.             return nvenc_errors[i].averr;

  245.         }

  246.     }

  247.     if (desc)

  248.         *desc = "unknown error";

  249.     return AVERROR_UNKNOWN;

  250. }

  251. 

  252. static int nvenc_print_error(void *log_ctx, NVENCSTATUS err,

  253.                                      const char *error_string)

  254. {

  255.     const char *desc;

  256.     int ret;

  257.     ret = nvenc_map_error(err, &desc);

  258.     av_log(log_ctx, AV_LOG_ERROR, "%s: %s (%d)\n", error_string, desc, err);

  259.     return ret;

  260. }

  261. 

  262. static int input_string_to_uint32(AVCodecContext *avctx, const NvencValuePair *pair, const char *input, uint32_t *output)

  263. {

  264.     for (; pair->str; ++pair) {

  265.         if (!strcmp(input, pair->str)) {

  266.             *output = pair->num;

  267.             return 0;

  268.         }

  269.     }

  270. 

  271.     return AVERROR(EINVAL);

  272. }

  273. 

  274. static void timestamp_queue_enqueue(AVFifoBuffer* queue, int64_t timestamp)

  275. {

  276.     av_fifo_generic_write(queue, &timestamp, sizeof(timestamp), NULL);

  277. }

  278. 

  279. static int64_t timestamp_queue_dequeue(AVFifoBuffer* queue)

  280. {

  281.     int64_t timestamp = AV_NOPTS_VALUE;

  282.     if (av_fifo_size(queue) > 0)

  283.         av_fifo_generic_read(queue, &timestamp, sizeof(timestamp), NULL);

  284. 

  285.     return timestamp;

  286. }

  287. 

  288. #define CHECK_LOAD_FUNC(t, f, s) \

  289. do { \

  290.     (f) = (t)LOAD_FUNC(dl_fn->cuda_lib, s); \

  291.     if (!(f)) { \

  292.         av_log(avctx, AV_LOG_FATAL, "Failed loading %s from CUDA library\n", s); \

  293.         goto error; \

  294.     } \

  295. } while (0)

  296. 

  297. static av_cold int nvenc_dyload_cuda(AVCodecContext *avctx)

  298. {

  299.     NvencContext *ctx = avctx->priv_data;

  300.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  301. 

  302.     if (dl_fn->cuda_lib)

  303.         return 1;

  304. 

  305. #if defined(_WIN32)

  306.     dl_fn->cuda_lib = LoadLibrary(TEXT("nvcuda.dll"));

  307. #else

  308.     dl_fn->cuda_lib = dlopen("libcuda.so", RTLD_LAZY);

  309. #endif

  310. 

  311.     if (!dl_fn->cuda_lib) {

  312.         av_log(avctx, AV_LOG_FATAL, "Failed loading CUDA library\n");

  313.         goto error;

  314.     }

  315. 

  316.     CHECK_LOAD_FUNC(PCUINIT, dl_fn->cu_init, "cuInit");

  317.     CHECK_LOAD_FUNC(PCUDEVICEGETCOUNT, dl_fn->cu_device_get_count, "cuDeviceGetCount");

  318.     CHECK_LOAD_FUNC(PCUDEVICEGET, dl_fn->cu_device_get, "cuDeviceGet");

  319.     CHECK_LOAD_FUNC(PCUDEVICEGETNAME, dl_fn->cu_device_get_name, "cuDeviceGetName");

  320.     CHECK_LOAD_FUNC(PCUDEVICECOMPUTECAPABILITY, dl_fn->cu_device_compute_capability, "cuDeviceComputeCapability");

  321.     CHECK_LOAD_FUNC(PCUCTXCREATE, dl_fn->cu_ctx_create, "cuCtxCreate_v2");

  322.     CHECK_LOAD_FUNC(PCUCTXPOPCURRENT, dl_fn->cu_ctx_pop_current, "cuCtxPopCurrent_v2");

  323.     CHECK_LOAD_FUNC(PCUCTXDESTROY, dl_fn->cu_ctx_destroy, "cuCtxDestroy_v2");

  324. 

  325.     return 1;

  326. 

  327. error:

  328. 

  329.     if (dl_fn->cuda_lib)

  330.         DL_CLOSE_FUNC(dl_fn->cuda_lib);

  331. 

  332.     dl_fn->cuda_lib = NULL;

  333. 

  334.     return 0;

  335. }

  336. 

  337. static av_cold int check_cuda_errors(AVCodecContext *avctx, CUresult err, const char *func)

  338. {

  339.     if (err != CUDA_SUCCESS) {

  340.         av_log(avctx, AV_LOG_FATAL, ">> %s - failed with error code 0x%x\n", func, err);

  341.         return 0;

  342.     }

  343.     return 1;

  344. }

  345. #define check_cuda_errors(f) if (!check_cuda_errors(avctx, f, #f)) goto error

  346. 

  347. static av_cold int nvenc_check_cuda(AVCodecContext *avctx)

  348. {

  349.     int device_count = 0;

  350.     CUdevice cu_device = 0;

  351.     char gpu_name[128];

  352.     int smminor = 0, smmajor = 0;

  353.     int i, smver, target_smver;

  354. 

  355.     NvencContext *ctx = avctx->priv_data;

  356.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  357. 

  358.     switch (avctx->codec->id) {

  359.     case AV_CODEC_ID_H264:

  360.         target_smver = avctx->pix_fmt == AV_PIX_FMT_YUV444P ? 0x52 : 0x30;

  361.         break;

  362.     case AV_CODEC_ID_H265:

  363.         target_smver = 0x52;

  364.         break;

  365.     default:

  366.         av_log(avctx, AV_LOG_FATAL, "Unknown codec name\n");

  367.         goto error;

  368.     }

  369. 

  370.     if (!nvenc_dyload_cuda(avctx))

  371.         return 0;

  372. 

  373.     if (dl_fn->nvenc_device_count > 0)

  374.         return 1;

  375. 

  376.     check_cuda_errors(dl_fn->cu_init(0));

  377. 

  378.     check_cuda_errors(dl_fn->cu_device_get_count(&device_count));

  379. 

  380.     if (!device_count) {

  381.         av_log(avctx, AV_LOG_FATAL, "No CUDA capable devices found\n");

  382.         goto error;

  383.     }

  384. 

  385.     av_log(avctx, AV_LOG_VERBOSE, "%d CUDA capable devices found\n", device_count);

  386. 

  387.     dl_fn->nvenc_device_count = 0;

  388. 

  389.     for (i = 0; i < device_count; ++i) {

  390.         check_cuda_errors(dl_fn->cu_device_get(&cu_device, i));

  391.         check_cuda_errors(dl_fn->cu_device_get_name(gpu_name, sizeof(gpu_name), cu_device));

  392.         check_cuda_errors(dl_fn->cu_device_compute_capability(&smmajor, &smminor, cu_device));

  393. 

  394.         smver = (smmajor << 4) | smminor;

  395. 

  396.         av_log(avctx, AV_LOG_VERBOSE, "[ GPU #%d - < %s > has Compute SM %d.%d, NVENC %s ]\n", i, gpu_name, smmajor, smminor, (smver >= target_smver) ? "Available" : "Not Available");

  397. 

  398.         if (smver >= target_smver)

  399.             dl_fn->nvenc_devices[dl_fn->nvenc_device_count++] = cu_device;

  400.     }

  401. 

  402.     if (!dl_fn->nvenc_device_count) {

  403.         av_log(avctx, AV_LOG_FATAL, "No NVENC capable devices found\n");

  404.         goto error;

  405.     }

  406. 

  407.     return 1;

  408. 

  409. error:

  410. 

  411.     dl_fn->nvenc_device_count = 0;

  412. 

  413.     return 0;

  414. }

  415. 

  416. static av_cold int nvenc_dyload_nvenc(AVCodecContext *avctx)

  417. {

  418.     PNVENCODEAPICREATEINSTANCE nvEncodeAPICreateInstance = 0;

  419.     NVENCSTATUS nvstatus;

  420. 

  421.     NvencContext *ctx = avctx->priv_data;

  422.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  423. 

  424.     if (!nvenc_check_cuda(avctx))

  425.         return 0;

  426. 

  427.     if (dl_fn->nvenc_lib)

  428.         return 1;

  429. 

  430. #if defined(_WIN32)

  431.     if (sizeof(void*) == 8) {

  432.         dl_fn->nvenc_lib = LoadLibrary(TEXT("nvEncodeAPI64.dll"));

  433.     } else {

  434.         dl_fn->nvenc_lib = LoadLibrary(TEXT("nvEncodeAPI.dll"));

  435.     }

  436. #else

  437.     dl_fn->nvenc_lib = dlopen("libnvidia-encode.so.1", RTLD_LAZY);

  438. #endif

  439. 

  440.     if (!dl_fn->nvenc_lib) {

  441.         av_log(avctx, AV_LOG_FATAL, "Failed loading the nvenc library\n");

  442.         goto error;

  443.     }

  444. 

  445.     nvEncodeAPICreateInstance = (PNVENCODEAPICREATEINSTANCE)LOAD_FUNC(dl_fn->nvenc_lib, "NvEncodeAPICreateInstance");

  446. 

  447.     if (!nvEncodeAPICreateInstance) {

  448.         av_log(avctx, AV_LOG_FATAL, "Failed to load nvenc entrypoint\n");

  449.         goto error;

  450.     }

  451. 

  452.     dl_fn->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER;

  453. 

  454.     nvstatus = nvEncodeAPICreateInstance(&dl_fn->nvenc_funcs);

  455. 

  456.     if (nvstatus != NV_ENC_SUCCESS) {

  457.         nvenc_print_error(avctx, nvstatus, "Failed to create nvenc instance");

  458.         goto error;

  459.     }

  460. 

  461.     av_log(avctx, AV_LOG_VERBOSE, "Nvenc initialized successfully\n");

  462. 

  463.     return 1;

  464. 

  465. error:

  466.     if (dl_fn->nvenc_lib)

  467.         DL_CLOSE_FUNC(dl_fn->nvenc_lib);

  468. 

  469.     dl_fn->nvenc_lib = NULL;

  470. 

  471.     return 0;

  472. }

  473. 

  474. static av_cold void nvenc_unload_nvenc(AVCodecContext *avctx)

  475. {

  476.     NvencContext *ctx = avctx->priv_data;

  477.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  478. 

  479.     DL_CLOSE_FUNC(dl_fn->nvenc_lib);

  480.     dl_fn->nvenc_lib = NULL;

  481. 

  482.     dl_fn->nvenc_device_count = 0;

  483. 

  484.     DL_CLOSE_FUNC(dl_fn->cuda_lib);

  485.     dl_fn->cuda_lib = NULL;

  486. 

  487.     dl_fn->cu_init = NULL;

  488.     dl_fn->cu_device_get_count = NULL;

  489.     dl_fn->cu_device_get = NULL;

  490.     dl_fn->cu_device_get_name = NULL;

  491.     dl_fn->cu_device_compute_capability = NULL;

  492.     dl_fn->cu_ctx_create = NULL;

  493.     dl_fn->cu_ctx_pop_current = NULL;

  494.     dl_fn->cu_ctx_destroy = NULL;

  495. 

  496.     av_log(avctx, AV_LOG_VERBOSE, "Nvenc unloaded\n");

  497. }

  498. 

  499. static av_cold int nvenc_setup_device(AVCodecContext *avctx)

  500. {

  501.     NvencContext *ctx = avctx->priv_data;

  502.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  503. 

  504.     CUresult cu_res;

  505.     CUcontext cu_context_curr;

  506. 

  507.     if (ctx->gpu >= dl_fn->nvenc_device_count) {

  508.         av_log(avctx, AV_LOG_FATAL, "Requested GPU %d, but only %d GPUs are available!\n", ctx->gpu, dl_fn->nvenc_device_count);

  509.         return AVERROR(EINVAL);

  510.     }

  511. 

  512.     ctx->cu_context = NULL;

  513.     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

  514. 

  515.     if (cu_res != CUDA_SUCCESS) {

  516.         av_log(avctx, AV_LOG_FATAL, "Failed creating CUDA context for NVENC: 0x%x\n", (int)cu_res);

  517.         return AVERROR_EXTERNAL;

  518.     }

  519. 

  520.     cu_res = dl_fn->cu_ctx_pop_current(&cu_context_curr);

  521. 

  522.     if (cu_res != CUDA_SUCCESS) {

  523.         av_log(avctx, AV_LOG_FATAL, "Failed popping CUDA context: 0x%x\n", (int)cu_res);

  524.         return AVERROR_EXTERNAL;

  525.     }

  526. 

  527.     return 0;

  528. }

  529. 

  530. static av_cold int nvenc_open_session(AVCodecContext *avctx)

  531. {

  532.     NvencContext *ctx = avctx->priv_data;

  533.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  534.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  535. 

  536.     NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS encode_session_params = { 0 };

  537.     NVENCSTATUS nv_status;

  538. 

  539.     encode_session_params.version = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;

  540.     encode_session_params.apiVersion = NVENCAPI_VERSION;

  541.     encode_session_params.device = ctx->cu_context;

  542.     encode_session_params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;

  543. 

  544.     nv_status = p_nvenc->nvEncOpenEncodeSessionEx(&encode_session_params, &ctx->nvencoder);

  545.     if (nv_status != NV_ENC_SUCCESS) {

  546.         ctx->nvencoder = NULL;

  547.         return nvenc_print_error(avctx, nv_status, "OpenEncodeSessionEx failed");

  548.     }

  549. 

  550.     return 0;

  551. }

  552. 

  553. static av_cold void set_constqp(AVCodecContext *avctx)

  554. {

  555.     NvencContext *ctx = avctx->priv_data;

  556. 

  557.     ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;

  558.     ctx->encode_config.rcParams.constQP.qpInterB = avctx->global_quality;

  559.     ctx->encode_config.rcParams.constQP.qpInterP = avctx->global_quality;

  560.     ctx->encode_config.rcParams.constQP.qpIntra = avctx->global_quality;

  561. }

  562. 

  563. static av_cold void set_vbr(AVCodecContext *avctx)

  564. {

  565.     NvencContext *ctx = avctx->priv_data;

  566. 

  567.     ctx->encode_config.rcParams.enableMinQP = 1;

  568.     ctx->encode_config.rcParams.enableMaxQP = 1;

  569. 

  570.     ctx->encode_config.rcParams.minQP.qpInterB = avctx->qmin;

  571.     ctx->encode_config.rcParams.minQP.qpInterP = avctx->qmin;

  572.     ctx->encode_config.rcParams.minQP.qpIntra = avctx->qmin;

  573. 

  574.     ctx->encode_config.rcParams.maxQP.qpInterB = avctx->qmax;

  575.     ctx->encode_config.rcParams.maxQP.qpInterP = avctx->qmax;

  576.     ctx->encode_config.rcParams.maxQP.qpIntra = avctx->qmax;

  577. }

  578. 

  579. static av_cold void set_lossless(AVCodecContext *avctx)

  580. {

  581.     NvencContext *ctx = avctx->priv_data;

  582. 

  583.     ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;

  584.     ctx->encode_config.rcParams.constQP.qpInterB = 0;

  585.     ctx->encode_config.rcParams.constQP.qpInterP = 0;

  586.     ctx->encode_config.rcParams.constQP.qpIntra = 0;

  587. }

  588. 

  589. static av_cold void nvenc_setup_rate_control(AVCodecContext *avctx, int lossless)

  590. {

  591.     NvencContext *ctx = avctx->priv_data;

  592. 

  593.     int qp_inter_p;

  594. 

  595.     if (avctx->bit_rate > 0) {

  596.         ctx->encode_config.rcParams.averageBitRate = avctx->bit_rate;

  597.     } else if (ctx->encode_config.rcParams.averageBitRate > 0) {

  598.         ctx->encode_config.rcParams.maxBitRate = ctx->encode_config.rcParams.averageBitRate;

  599.     }

  600. 

  601.     if (avctx->rc_max_rate > 0)

  602.         ctx->encode_config.rcParams.maxBitRate = avctx->rc_max_rate;

  603. 

  604.     if (lossless) {

  605.         if (avctx->codec->id == AV_CODEC_ID_H264)

  606.             ctx->encode_config.encodeCodecConfig.h264Config.qpPrimeYZeroTransformBypassFlag = 1;

  607. 

  608.         set_lossless(avctx);

  609. 

  610.         avctx->qmin = -1;

  611.         avctx->qmax = -1;

  612.     } else if (ctx->cbr) {

  613.         if (!ctx->twopass) {

  614.             ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CBR;

  615.         } else {

  616.             ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_QUALITY;

  617. 

  618.             if (avctx->codec->id == AV_CODEC_ID_H264) {

  619.                 ctx->encode_config.encodeCodecConfig.h264Config.adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE;

  620.                 ctx->encode_config.encodeCodecConfig.h264Config.fmoMode = NV_ENC_H264_FMO_DISABLE;

  621.             }

  622.         }

  623. 

  624.         if (avctx->codec->id == AV_CODEC_ID_H264) {

  625.             ctx->encode_config.encodeCodecConfig.h264Config.outputBufferingPeriodSEI = 1;

  626.             ctx->encode_config.encodeCodecConfig.h264Config.outputPictureTimingSEI = 1;

  627.         } else if (avctx->codec->id == AV_CODEC_ID_H265) {

  628.             ctx->encode_config.encodeCodecConfig.hevcConfig.outputBufferingPeriodSEI = 1;

  629.             ctx->encode_config.encodeCodecConfig.hevcConfig.outputPictureTimingSEI = 1;

  630.         }

  631.     } else if (avctx->global_quality > 0) {

  632.         set_constqp(avctx);

  633. 

  634.         avctx->qmin = -1;

  635.         avctx->qmax = -1;

  636.     } else {

  637.         if (avctx->qmin >= 0 && avctx->qmax >= 0) {

  638.             set_vbr(avctx);

  639. 

  640.             qp_inter_p = (avctx->qmax + 3 * avctx->qmin) / 4; // biased towards Qmin

  641. 

  642.             if (ctx->twopass) {

  643.                 ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_VBR;

  644.                 if (avctx->codec->id == AV_CODEC_ID_H264) {

  645.                     ctx->encode_config.encodeCodecConfig.h264Config.adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE;

  646.                     ctx->encode_config.encodeCodecConfig.h264Config.fmoMode = NV_ENC_H264_FMO_DISABLE;

  647.                 }

  648.             } else {

  649.                 ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR_MINQP;

  650.             }

  651.         } else {

  652.             qp_inter_p = 26; // default to 26

  653. 

  654.             if (ctx->twopass) {

  655.                 ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_VBR;

  656.             } else {

  657.                 ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR;

  658.             }

  659.         }

  660. 

  661.         ctx->encode_config.rcParams.enableInitialRCQP = 1;

  662.         ctx->encode_config.rcParams.initialRCQP.qpInterP  = qp_inter_p;

  663. 

  664.         if (avctx->i_quant_factor != 0.0 && avctx->b_quant_factor != 0.0) {

  665.             ctx->encode_config.rcParams.initialRCQP.qpIntra = av_clip(

  666.                 qp_inter_p * fabs(avctx->i_quant_factor) + avctx->i_quant_offset, 0, 51);

  667.             ctx->encode_config.rcParams.initialRCQP.qpInterB = av_clip(

  668.                 qp_inter_p * fabs(avctx->b_quant_factor) + avctx->b_quant_offset, 0, 51);

  669.         } else {

  670.             ctx->encode_config.rcParams.initialRCQP.qpIntra = qp_inter_p;

  671.             ctx->encode_config.rcParams.initialRCQP.qpInterB = qp_inter_p;

  672.         }

  673.     }

  674. 

  675.     if (avctx->rc_buffer_size > 0) {

  676.         ctx->encode_config.rcParams.vbvBufferSize = avctx->rc_buffer_size;

  677.     } else if (ctx->encode_config.rcParams.averageBitRate > 0) {

  678.         ctx->encode_config.rcParams.vbvBufferSize = 2 * ctx->encode_config.rcParams.averageBitRate;

  679.     }

  680. }

  681. 

  682. static av_cold int nvenc_setup_h264_config(AVCodecContext *avctx, int lossless)

  683. {

  684.     NvencContext *ctx = avctx->priv_data;

  685.     int res;

  686. 

  687.     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourMatrix = avctx->colorspace;

  688.     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourPrimaries = avctx->color_primaries;

  689.     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.transferCharacteristics = avctx->color_trc;

  690.     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG

  691.         || avctx->pix_fmt == AV_PIX_FMT_YUVJ420P || avctx->pix_fmt == AV_PIX_FMT_YUVJ422P || avctx->pix_fmt == AV_PIX_FMT_YUVJ444P);

  692. 

  693.     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourDescriptionPresentFlag =

  694.         (avctx->colorspace != 2 || avctx->color_primaries != 2 || avctx->color_trc != 2);

  695. 

  696.     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoSignalTypePresentFlag =

  697.         (ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourDescriptionPresentFlag

  698.         || ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoFormat != 5

  699.         || ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoFullRangeFlag != 0);

  700. 

  701.     ctx->encode_config.encodeCodecConfig.h264Config.sliceMode = 3;

  702.     ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData = 1;

  703. 

  704.     ctx->encode_config.encodeCodecConfig.h264Config.disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;

  705.     ctx->encode_config.encodeCodecConfig.h264Config.repeatSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;

  706. 

  707.     ctx->encode_config.encodeCodecConfig.h264Config.outputAUD = 1;

  708. 

  709.     if (!ctx->profile && !lossless) {

  710.         switch (avctx->profile) {

  711.         case FF_PROFILE_H264_HIGH_444_PREDICTIVE:

  712.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;

  713.             break;

  714.         case FF_PROFILE_H264_BASELINE:

  715.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;

  716.             break;

  717.         case FF_PROFILE_H264_MAIN:

  718.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;

  719.             break;

  720.         case FF_PROFILE_H264_HIGH:

  721.         case FF_PROFILE_UNKNOWN:

  722.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;

  723.             break;

  724.         default:

  725.             av_log(avctx, AV_LOG_WARNING, "Unsupported profile requested, falling back to high\n");

  726.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;

  727.             break;

  728.         }

  729.     } else if (!lossless) {

  730.         if (!strcmp(ctx->profile, "high")) {

  731.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;

  732.             avctx->profile = FF_PROFILE_H264_HIGH;

  733.         } else if (!strcmp(ctx->profile, "main")) {

  734.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;

  735.             avctx->profile = FF_PROFILE_H264_MAIN;

  736.         } else if (!strcmp(ctx->profile, "baseline")) {

  737.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;

  738.             avctx->profile = FF_PROFILE_H264_BASELINE;

  739.         } else if (!strcmp(ctx->profile, "high444p")) {

  740.             ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;

  741.             avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;

  742.         } else {

  743.             av_log(avctx, AV_LOG_FATAL, "Profile \"%s\" is unknown! Supported profiles: high, main, baseline\n", ctx->profile);

  744.             return AVERROR(EINVAL);

  745.         }

  746.     }

  747. 

  748.     // force setting profile as high444p if input is AV_PIX_FMT_YUV444P

  749.     if (avctx->pix_fmt == AV_PIX_FMT_YUV444P) {

  750.         ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;

  751.         avctx->profile = FF_PROFILE_H264_HIGH_444_PREDICTIVE;

  752.     }

  753. 

  754.     ctx->encode_config.encodeCodecConfig.h264Config.chromaFormatIDC = avctx->profile == FF_PROFILE_H264_HIGH_444_PREDICTIVE ? 3 : 1;

  755. 

  756.     if (ctx->level) {

  757.         res = input_string_to_uint32(avctx, nvenc_h264_level_pairs, ctx->level, &ctx->encode_config.encodeCodecConfig.h264Config.level);

  758. 

  759.         if (res) {

  760.             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);

  761.             return res;

  762.         }

  763.     } else {

  764.         ctx->encode_config.encodeCodecConfig.h264Config.level = NV_ENC_LEVEL_AUTOSELECT;

  765.     }

  766. 

  767.     return 0;

  768. }

  769. 

  770. static av_cold int nvenc_setup_hevc_config(AVCodecContext *avctx)

  771. {

  772.     NvencContext *ctx = avctx->priv_data;

  773.     int res;

  774. 

  775.     ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.colourMatrix = avctx->colorspace;

  776.     ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.colourPrimaries = avctx->color_primaries;

  777.     ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.transferCharacteristics = avctx->color_trc;

  778.     ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.videoFullRangeFlag = (avctx->color_range == AVCOL_RANGE_JPEG

  779.         || avctx->pix_fmt == AV_PIX_FMT_YUVJ420P || avctx->pix_fmt == AV_PIX_FMT_YUVJ422P || avctx->pix_fmt == AV_PIX_FMT_YUVJ444P);

  780. 

  781.     ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.colourDescriptionPresentFlag =

  782.         (avctx->colorspace != 2 || avctx->color_primaries != 2 || avctx->color_trc != 2);

  783. 

  784.     ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.videoSignalTypePresentFlag =

  785.         (ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.colourDescriptionPresentFlag

  786.         || ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.videoFormat != 5

  787.         || ctx->encode_config.encodeCodecConfig.hevcConfig.hevcVUIParameters.videoFullRangeFlag != 0);

  788. 

  789.     ctx->encode_config.encodeCodecConfig.hevcConfig.sliceMode = 3;

  790.     ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData = 1;

  791. 

  792.     ctx->encode_config.encodeCodecConfig.hevcConfig.disableSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;

  793.     ctx->encode_config.encodeCodecConfig.hevcConfig.repeatSPSPPS = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) ? 0 : 1;

  794. 

  795.     ctx->encode_config.encodeCodecConfig.hevcConfig.outputAUD = 1;

  796. 

  797.     /* No other profile is supported in the current SDK version 5 */

  798.     ctx->encode_config.profileGUID = NV_ENC_HEVC_PROFILE_MAIN_GUID;

  799.     avctx->profile = FF_PROFILE_HEVC_MAIN;

  800. 

  801.     if (ctx->level) {

  802.         res = input_string_to_uint32(avctx, nvenc_hevc_level_pairs, ctx->level, &ctx->encode_config.encodeCodecConfig.hevcConfig.level);

  803. 

  804.         if (res) {

  805.             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);

  806.             return res;

  807.         }

  808.     } else {

  809.         ctx->encode_config.encodeCodecConfig.hevcConfig.level = NV_ENC_LEVEL_AUTOSELECT;

  810.     }

  811. 

  812.     if (ctx->tier) {

  813.         if (!strcmp(ctx->tier, "main")) {

  814.             ctx->encode_config.encodeCodecConfig.hevcConfig.tier = NV_ENC_TIER_HEVC_MAIN;

  815.         } else if (!strcmp(ctx->tier, "high")) {

  816.             ctx->encode_config.encodeCodecConfig.hevcConfig.tier = NV_ENC_TIER_HEVC_HIGH;

  817.         } else {

  818.             av_log(avctx, AV_LOG_FATAL, "Tier \"%s\" is unknown! Supported tiers: main, high\n", ctx->tier);

  819.             return AVERROR(EINVAL);

  820.         }

  821.     }

  822. 

  823.     return 0;

  824. }

  825. 

  826. static av_cold int nvenc_setup_codec_config(AVCodecContext *avctx, int lossless)

  827. {

  828.     switch (avctx->codec->id) {

  829.     case AV_CODEC_ID_H264:

  830.         return nvenc_setup_h264_config(avctx, lossless);

  831.     case AV_CODEC_ID_H265:

  832.         return nvenc_setup_hevc_config(avctx);

  833.     /* Earlier switch/case will return if unknown codec is passed. */

  834.     }

  835. 

  836.     return 0;

  837. }

  838. 

  839. static av_cold int nvenc_setup_encoder(AVCodecContext *avctx)

  840. {

  841.     NvencContext *ctx = avctx->priv_data;

  842.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  843.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  844. 

  845.     NV_ENC_PRESET_CONFIG preset_config = { 0 };

  846.     GUID encoder_preset = NV_ENC_PRESET_HQ_GUID;

  847.     GUID codec;

  848.     NVENCSTATUS nv_status = NV_ENC_SUCCESS;

  849.     AVCPBProperties *cpb_props;

  850.     int num_mbs;

  851.     int isLL = 0;

  852.     int lossless = 0;

  853.     int res = 0;

  854.     int dw, dh;

  855. 

  856.     ctx->last_dts = AV_NOPTS_VALUE;

  857. 

  858.     ctx->encode_config.version = NV_ENC_CONFIG_VER;

  859.     ctx->init_encode_params.version = NV_ENC_INITIALIZE_PARAMS_VER;

  860.     preset_config.version = NV_ENC_PRESET_CONFIG_VER;

  861.     preset_config.presetCfg.version = NV_ENC_CONFIG_VER;

  862. 

  863.     if (ctx->preset) {

  864.         if (!strcmp(ctx->preset, "slow")) {

  865.             encoder_preset = NV_ENC_PRESET_HQ_GUID;

  866.             ctx->twopass = 1;

  867.         } else if (!strcmp(ctx->preset, "medium")) {

  868.             encoder_preset = NV_ENC_PRESET_HQ_GUID;

  869.             ctx->twopass = 0;

  870.         } else if (!strcmp(ctx->preset, "fast")) {

  871.             encoder_preset = NV_ENC_PRESET_HP_GUID;

  872.             ctx->twopass = 0;

  873.         } else if (!strcmp(ctx->preset, "hq")) {

  874.             encoder_preset = NV_ENC_PRESET_HQ_GUID;

  875.         } else if (!strcmp(ctx->preset, "hp")) {

  876.             encoder_preset = NV_ENC_PRESET_HP_GUID;

  877.         } else if (!strcmp(ctx->preset, "bd")) {

  878.             encoder_preset = NV_ENC_PRESET_BD_GUID;

  879.         } else if (!strcmp(ctx->preset, "ll")) {

  880.             encoder_preset = NV_ENC_PRESET_LOW_LATENCY_DEFAULT_GUID;

  881.             isLL = 1;

  882.         } else if (!strcmp(ctx->preset, "llhp")) {

  883.             encoder_preset = NV_ENC_PRESET_LOW_LATENCY_HP_GUID;

  884.             isLL = 1;

  885.         } else if (!strcmp(ctx->preset, "llhq")) {

  886.             encoder_preset = NV_ENC_PRESET_LOW_LATENCY_HQ_GUID;

  887.             isLL = 1;

  888.         } else if (!strcmp(ctx->preset, "lossless")) {

  889.             encoder_preset = NV_ENC_PRESET_LOSSLESS_DEFAULT_GUID;

  890.             lossless = 1;

  891.         } else if (!strcmp(ctx->preset, "losslesshp")) {

  892.             encoder_preset = NV_ENC_PRESET_LOSSLESS_HP_GUID;

  893.             lossless = 1;

  894.         } else if (!strcmp(ctx->preset, "default")) {

  895.             encoder_preset = NV_ENC_PRESET_DEFAULT_GUID;

  896.         } else {

  897.             av_log(avctx, AV_LOG_FATAL, "Preset \"%s\" is unknown! Supported presets: slow, medium, fast, hp, hq, bd, ll, llhp, llhq, lossless, losslesshp, default\n", ctx->preset);

  898.             return AVERROR(EINVAL);

  899.         }

  900.     }

  901. 

  902.     if (ctx->twopass < 0) {

  903.         ctx->twopass = isLL;

  904.     }

  905. 

  906.     switch (avctx->codec->id) {

  907.     case AV_CODEC_ID_H264:

  908.         codec = NV_ENC_CODEC_H264_GUID;

  909.         break;

  910.     case AV_CODEC_ID_H265:

  911.         codec = NV_ENC_CODEC_HEVC_GUID;

  912.         break;

  913.     default:

  914.         av_log(avctx, AV_LOG_ERROR, "Unknown codec name\n");

  915.         return AVERROR(EINVAL);

  916.     }

  917. 

  918.     nv_status = p_nvenc->nvEncGetEncodePresetConfig(ctx->nvencoder, codec, encoder_preset, &preset_config);

  919.     if (nv_status != NV_ENC_SUCCESS) {

  920.         return nvenc_print_error(avctx, nv_status, "GetEncodePresetConfig failed");

  921.     }

  922. 

  923.     ctx->init_encode_params.encodeGUID = codec;

  924.     ctx->init_encode_params.encodeHeight = avctx->height;

  925.     ctx->init_encode_params.encodeWidth = avctx->width;

  926. 

  927.     if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den &&

  928.         (avctx->sample_aspect_ratio.num != 1 || avctx->sample_aspect_ratio.num != 1)) {

  929.         av_reduce(&dw, &dh,

  930.                   avctx->width * avctx->sample_aspect_ratio.num,

  931.                   avctx->height * avctx->sample_aspect_ratio.den,

  932.                   1024 * 1024);

  933.         ctx->init_encode_params.darHeight = dh;

  934.         ctx->init_encode_params.darWidth = dw;

  935.     } else {

  936.         ctx->init_encode_params.darHeight = avctx->height;

  937.         ctx->init_encode_params.darWidth = avctx->width;

  938.     }

  939. 

  940.     // De-compensate for hardware, dubiously, trying to compensate for

  941.     // playback at 704 pixel width.

  942.     if (avctx->width == 720 &&

  943.         (avctx->height == 480 || avctx->height == 576)) {

  944.         av_reduce(&dw, &dh,

  945.                   ctx->init_encode_params.darWidth * 44,

  946.                   ctx->init_encode_params.darHeight * 45,

  947.                   1024 * 1024);

  948.         ctx->init_encode_params.darHeight = dh;

  949.         ctx->init_encode_params.darWidth = dw;

  950.     }

  951. 

  952.     ctx->init_encode_params.frameRateNum = avctx->time_base.den;

  953.     ctx->init_encode_params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;

  954. 

  955.     num_mbs = ((avctx->width + 15) >> 4) * ((avctx->height + 15) >> 4);

  956.     ctx->max_surface_count = (num_mbs >= 8160) ? 32 : 48;

  957. 

  958.     if (ctx->buffer_delay >= ctx->max_surface_count)

  959.         ctx->buffer_delay = ctx->max_surface_count - 1;

  960. 

  961.     ctx->init_encode_params.enableEncodeAsync = 0;

  962.     ctx->init_encode_params.enablePTD = 1;

  963. 

  964.     ctx->init_encode_params.presetGUID = encoder_preset;

  965. 

  966.     ctx->init_encode_params.encodeConfig = &ctx->encode_config;

  967.     memcpy(&ctx->encode_config, &preset_config.presetCfg, sizeof(ctx->encode_config));

  968.     ctx->encode_config.version = NV_ENC_CONFIG_VER;

  969. 

  970.     if (avctx->refs >= 0) {

  971.         /* 0 means "let the hardware decide" */

  972.         switch (avctx->codec->id) {

  973.         case AV_CODEC_ID_H264:

  974.             ctx->encode_config.encodeCodecConfig.h264Config.maxNumRefFrames = avctx->refs;

  975.             break;

  976.         case AV_CODEC_ID_H265:

  977.             ctx->encode_config.encodeCodecConfig.hevcConfig.maxNumRefFramesInDPB = avctx->refs;

  978.             break;

  979.         /* Earlier switch/case will return if unknown codec is passed. */

  980.         }

  981.     }

  982. 

  983.     if (avctx->gop_size > 0) {

  984.         if (avctx->max_b_frames >= 0) {

  985.             /* 0 is intra-only, 1 is I/P only, 2 is one B Frame, 3 two B frames, and so on. */

  986.             ctx->encode_config.frameIntervalP = avctx->max_b_frames + 1;

  987.         }

  988. 

  989.         ctx->encode_config.gopLength = avctx->gop_size;

  990.         switch (avctx->codec->id) {

  991.         case AV_CODEC_ID_H264:

  992.             ctx->encode_config.encodeCodecConfig.h264Config.idrPeriod = avctx->gop_size;

  993.             break;

  994.         case AV_CODEC_ID_H265:

  995.             ctx->encode_config.encodeCodecConfig.hevcConfig.idrPeriod = avctx->gop_size;

  996.             break;

  997.         /* Earlier switch/case will return if unknown codec is passed. */

  998.         }

  999.     } else if (avctx->gop_size == 0) {

  1000.         ctx->encode_config.frameIntervalP = 0;

  1001.         ctx->encode_config.gopLength = 1;

  1002.         switch (avctx->codec->id) {

  1003.         case AV_CODEC_ID_H264:

  1004.             ctx->encode_config.encodeCodecConfig.h264Config.idrPeriod = 1;

  1005.             break;

  1006.         case AV_CODEC_ID_H265:

  1007.             ctx->encode_config.encodeCodecConfig.hevcConfig.idrPeriod = 1;

  1008.             break;

  1009.         /* Earlier switch/case will return if unknown codec is passed. */

  1010.         }

  1011.     }

  1012. 

  1013.     /* when there're b frames, set dts offset */

  1014.     if (ctx->encode_config.frameIntervalP >= 2)

  1015.         ctx->last_dts = -2;

  1016. 

  1017.     nvenc_setup_rate_control(avctx, lossless);

  1018. 

  1019.     if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {

  1020.         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;

  1021.     } else {

  1022.         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;

  1023.     }

  1024. 

  1025.     res = nvenc_setup_codec_config(avctx, lossless);

  1026.     if (res)

  1027.         return res;

  1028. 

  1029.     nv_status = p_nvenc->nvEncInitializeEncoder(ctx->nvencoder, &ctx->init_encode_params);

  1030.     if (nv_status != NV_ENC_SUCCESS) {

  1031.         return nvenc_print_error(avctx, nv_status, "InitializeEncoder failed");

  1032.     }

  1033. 

  1034.     if (ctx->encode_config.frameIntervalP > 1)

  1035.         avctx->has_b_frames = 2;

  1036. 

  1037.     if (ctx->encode_config.rcParams.averageBitRate > 0)

  1038.         avctx->bit_rate = ctx->encode_config.rcParams.averageBitRate;

  1039. 

  1040.     cpb_props = ff_add_cpb_side_data(avctx);

  1041.     if (!cpb_props)

  1042.         return AVERROR(ENOMEM);

  1043.     cpb_props->max_bitrate = ctx->encode_config.rcParams.maxBitRate;

  1044.     cpb_props->avg_bitrate = avctx->bit_rate;

  1045.     cpb_props->buffer_size = ctx->encode_config.rcParams.vbvBufferSize;

  1046. 

  1047.     return 0;

  1048. }

  1049. 

  1050. static av_cold int nvenc_alloc_surface(AVCodecContext *avctx, int idx)

  1051. {

  1052.     NvencContext *ctx = avctx->priv_data;

  1053.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1054.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1055. 

  1056.     NVENCSTATUS nv_status;

  1057.     NV_ENC_CREATE_INPUT_BUFFER allocSurf = { 0 };

  1058.     NV_ENC_CREATE_BITSTREAM_BUFFER allocOut = { 0 };

  1059. 

  1060.     allocSurf.version = NV_ENC_CREATE_INPUT_BUFFER_VER;

  1061.     allocOut.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;

  1062. 

  1063.     allocSurf.width = (avctx->width + 31) & ~31;

  1064.     allocSurf.height = (avctx->height + 31) & ~31;

  1065. 

  1066.     allocSurf.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;

  1067. 

  1068.     switch (avctx->pix_fmt) {

  1069.     case AV_PIX_FMT_YUV420P:

  1070.         allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_YV12_PL;

  1071.         break;

  1072. 

  1073.     case AV_PIX_FMT_NV12:

  1074.         allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_NV12_PL;

  1075.         break;

  1076. 

  1077.     case AV_PIX_FMT_YUV444P:

  1078.         allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_YUV444_PL;

  1079.         break;

  1080. 

  1081.     default:

  1082.         av_log(avctx, AV_LOG_FATAL, "Invalid input pixel format\n");

  1083.         return AVERROR(EINVAL);

  1084.     }

  1085. 

  1086.     nv_status = p_nvenc->nvEncCreateInputBuffer(ctx->nvencoder, &allocSurf);

  1087.     if (nv_status != NV_ENC_SUCCESS) {

  1088.         return nvenc_print_error(avctx, nv_status, "CreateInputBuffer failed");

  1089.     }

  1090. 

  1091.     ctx->surfaces[idx].lockCount = 0;

  1092.     ctx->surfaces[idx].input_surface = allocSurf.inputBuffer;

  1093.     ctx->surfaces[idx].format = allocSurf.bufferFmt;

  1094.     ctx->surfaces[idx].width = allocSurf.width;

  1095.     ctx->surfaces[idx].height = allocSurf.height;

  1096. 

  1097.     /* 1MB is large enough to hold most output frames. NVENC increases this automaticaly if it's not enough. */

  1098.     allocOut.size = 1024 * 1024;

  1099. 

  1100.     allocOut.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;

  1101. 

  1102.     nv_status = p_nvenc->nvEncCreateBitstreamBuffer(ctx->nvencoder, &allocOut);

  1103.     if (nv_status != NV_ENC_SUCCESS) {

  1104.         int err = nvenc_print_error(avctx, nv_status, "CreateBitstreamBuffer failed");

  1105.         p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[idx].input_surface);

  1106.         return err;

  1107.     }

  1108. 

  1109.     ctx->surfaces[idx].output_surface = allocOut.bitstreamBuffer;

  1110.     ctx->surfaces[idx].size = allocOut.size;

  1111. 

  1112.     return 0;

  1113. }

  1114. 

  1115. static av_cold int nvenc_setup_surfaces(AVCodecContext *avctx, int* surfaceCount)

  1116. {

  1117.     int res;

  1118.     NvencContext *ctx = avctx->priv_data;

  1119. 

  1120.     ctx->surfaces = av_malloc(ctx->max_surface_count * sizeof(*ctx->surfaces));

  1121. 

  1122.     if (!ctx->surfaces) {

  1123.         return AVERROR(ENOMEM);

  1124.     }

  1125. 

  1126.     ctx->timestamp_list = av_fifo_alloc(ctx->max_surface_count * sizeof(int64_t));

  1127.     if (!ctx->timestamp_list)

  1128.         return AVERROR(ENOMEM);

  1129.     ctx->output_surface_queue = av_fifo_alloc(ctx->max_surface_count * sizeof(NvencSurface*));

  1130.     if (!ctx->output_surface_queue)

  1131.         return AVERROR(ENOMEM);

  1132.     ctx->output_surface_ready_queue = av_fifo_alloc(ctx->max_surface_count * sizeof(NvencSurface*));

  1133.     if (!ctx->output_surface_ready_queue)

  1134.         return AVERROR(ENOMEM);

  1135. 

  1136.     for (*surfaceCount = 0; *surfaceCount < ctx->max_surface_count; ++*surfaceCount) {

  1137.         res = nvenc_alloc_surface(avctx, *surfaceCount);

  1138.         if (res)

  1139.             return res;

  1140.     }

  1141. 

  1142.     return 0;

  1143. }

  1144. 

  1145. static av_cold int nvenc_setup_extradata(AVCodecContext *avctx)

  1146. {

  1147.     NvencContext *ctx = avctx->priv_data;

  1148.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1149.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1150. 

  1151.     NVENCSTATUS nv_status;

  1152.     uint32_t outSize = 0;

  1153.     char tmpHeader[256];

  1154.     NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };

  1155.     payload.version = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;

  1156. 

  1157.     payload.spsppsBuffer = tmpHeader;

  1158.     payload.inBufferSize = sizeof(tmpHeader);

  1159.     payload.outSPSPPSPayloadSize = &outSize;

  1160. 

  1161.     nv_status = p_nvenc->nvEncGetSequenceParams(ctx->nvencoder, &payload);

  1162.     if (nv_status != NV_ENC_SUCCESS) {

  1163.         return nvenc_print_error(avctx, nv_status, "GetSequenceParams failed");

  1164.     }

  1165. 

  1166.     avctx->extradata_size = outSize;

  1167.     avctx->extradata = av_mallocz(outSize + AV_INPUT_BUFFER_PADDING_SIZE);

  1168. 

  1169.     if (!avctx->extradata) {

  1170.         return AVERROR(ENOMEM);

  1171.     }

  1172. 

  1173.     memcpy(avctx->extradata, tmpHeader, outSize);

  1174. 

  1175.     return 0;

  1176. }

  1177. 

  1178. static av_cold int nvenc_encode_init(AVCodecContext *avctx)

  1179. {

  1180.     NvencContext *ctx = avctx->priv_data;

  1181.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1182.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1183. 

  1184.     int res;

  1185.     int i;

  1186.     int surfaceCount = 0;

  1187. 

  1188.     if (!nvenc_dyload_nvenc(avctx))

  1189.         return AVERROR_EXTERNAL;

  1190. 

  1191.     res = nvenc_setup_device(avctx);

  1192.     if (res)

  1193.         goto error;

  1194. 

  1195.     res = nvenc_open_session(avctx);

  1196.     if (res)

  1197.         goto error;

  1198. 

  1199.     res = nvenc_setup_encoder(avctx);

  1200.     if (res)

  1201.         goto error;

  1202. 

  1203.     res = nvenc_setup_surfaces(avctx, &surfaceCount);

  1204.     if (res)

  1205.         goto error;

  1206. 

  1207.     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {

  1208.         res = nvenc_setup_extradata(avctx);

  1209.         if (res)

  1210.             goto error;

  1211.     }

  1212. 

  1213.     return 0;

  1214. 

  1215. error:

  1216.     av_fifo_freep(&ctx->timestamp_list);

  1217.     av_fifo_freep(&ctx->output_surface_ready_queue);

  1218.     av_fifo_freep(&ctx->output_surface_queue);

  1219. 

  1220.     for (i = 0; i < surfaceCount; ++i) {

  1221.         p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[i].input_surface);

  1222.         p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->surfaces[i].output_surface);

  1223.     }

  1224.     av_freep(&ctx->surfaces);

  1225. 

  1226.     if (ctx->nvencoder)

  1227.         p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);

  1228. 

  1229.     if (ctx->cu_context)

  1230.         dl_fn->cu_ctx_destroy(ctx->cu_context);

  1231. 

  1232.     nvenc_unload_nvenc(avctx);

  1233. 

  1234.     ctx->nvencoder = NULL;

  1235.     ctx->cu_context = NULL;

  1236. 

  1237.     return res;

  1238. }

  1239. 

  1240. static av_cold int nvenc_encode_close(AVCodecContext *avctx)

  1241. {

  1242.     NvencContext *ctx = avctx->priv_data;

  1243.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1244.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1245.     int i;

  1246. 

  1247.     av_fifo_freep(&ctx->timestamp_list);

  1248.     av_fifo_freep(&ctx->output_surface_ready_queue);

  1249.     av_fifo_freep(&ctx->output_surface_queue);

  1250. 

  1251.     for (i = 0; i < ctx->max_surface_count; ++i) {

  1252.         p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->surfaces[i].input_surface);

  1253.         p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->surfaces[i].output_surface);

  1254.     }

  1255.     av_freep(&ctx->surfaces);

  1256.     ctx->max_surface_count = 0;

  1257. 

  1258.     p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);

  1259.     ctx->nvencoder = NULL;

  1260. 

  1261.     dl_fn->cu_ctx_destroy(ctx->cu_context);

  1262.     ctx->cu_context = NULL;

  1263. 

  1264.     nvenc_unload_nvenc(avctx);

  1265. 

  1266.     return 0;

  1267. }

  1268. 

  1269. static NvencSurface *get_free_frame(NvencContext *ctx)

  1270. {

  1271.     int i;

  1272. 

  1273.     for (i = 0; i < ctx->max_surface_count; ++i) {

  1274.         if (!ctx->surfaces[i].lockCount) {

  1275.             ctx->surfaces[i].lockCount = 1;

  1276.             return &ctx->surfaces[i];

  1277.         }

  1278.     }

  1279. 

  1280.     return NULL;

  1281. }

  1282. 

  1283. static int nvenc_copy_frame(AVCodecContext *avctx, NvencSurface *inSurf,

  1284.             NV_ENC_LOCK_INPUT_BUFFER *lockBufferParams, const AVFrame *frame)

  1285. {

  1286.     uint8_t *buf = lockBufferParams->bufferDataPtr;

  1287.     int off = inSurf->height * lockBufferParams->pitch;

  1288. 

  1289.     if (avctx->pix_fmt == AV_PIX_FMT_YUV420P) {

  1290.         av_image_copy_plane(buf, lockBufferParams->pitch,

  1291.             frame->data[0], frame->linesize[0],

  1292.             avctx->width, avctx->height);

  1293. 

  1294.         buf += off;

  1295. 

  1296.         av_image_copy_plane(buf, lockBufferParams->pitch >> 1,

  1297.             frame->data[2], frame->linesize[2],

  1298.             avctx->width >> 1, avctx->height >> 1);

  1299. 

  1300.         buf += off >> 2;

  1301. 

  1302.         av_image_copy_plane(buf, lockBufferParams->pitch >> 1,

  1303.             frame->data[1], frame->linesize[1],

  1304.             avctx->width >> 1, avctx->height >> 1);

  1305.     } else if (avctx->pix_fmt == AV_PIX_FMT_NV12) {

  1306.         av_image_copy_plane(buf, lockBufferParams->pitch,

  1307.             frame->data[0], frame->linesize[0],

  1308.             avctx->width, avctx->height);

  1309. 

  1310.         buf += off;

  1311. 

  1312.         av_image_copy_plane(buf, lockBufferParams->pitch,

  1313.             frame->data[1], frame->linesize[1],

  1314.             avctx->width, avctx->height >> 1);

  1315.     } else if (avctx->pix_fmt == AV_PIX_FMT_YUV444P) {

  1316.         av_image_copy_plane(buf, lockBufferParams->pitch,

  1317.             frame->data[0], frame->linesize[0],

  1318.             avctx->width, avctx->height);

  1319. 

  1320.         buf += off;

  1321. 

  1322.         av_image_copy_plane(buf, lockBufferParams->pitch,

  1323.             frame->data[1], frame->linesize[1],

  1324.             avctx->width, avctx->height);

  1325. 

  1326.         buf += off;

  1327. 

  1328.         av_image_copy_plane(buf, lockBufferParams->pitch,

  1329.             frame->data[2], frame->linesize[2],

  1330.             avctx->width, avctx->height);

  1331.     } else {

  1332.         av_log(avctx, AV_LOG_FATAL, "Invalid pixel format!\n");

  1333.         return AVERROR(EINVAL);

  1334.     }

  1335. 

  1336.     return 0;

  1337. }

  1338. 

  1339. static int nvenc_upload_frame(AVCodecContext *avctx, const AVFrame *frame,

  1340.                                       NvencSurface *nvenc_frame)

  1341. {

  1342.     NvencContext *ctx = avctx->priv_data;

  1343.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1344.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1345. 

  1346.     int res;

  1347.     NVENCSTATUS nv_status;

  1348.     NV_ENC_LOCK_INPUT_BUFFER lockBufferParams = { 0 };

  1349. 

  1350.     lockBufferParams.version = NV_ENC_LOCK_INPUT_BUFFER_VER;

  1351.     lockBufferParams.inputBuffer = nvenc_frame->input_surface;

  1352. 

  1353.     nv_status = p_nvenc->nvEncLockInputBuffer(ctx->nvencoder, &lockBufferParams);

  1354.     if (nv_status != NV_ENC_SUCCESS) {

  1355.         return nvenc_print_error(avctx, nv_status, "Failed locking nvenc input buffer");

  1356.     }

  1357. 

  1358.     res = nvenc_copy_frame(avctx, nvenc_frame, &lockBufferParams, frame);

  1359. 

  1360.     nv_status = p_nvenc->nvEncUnlockInputBuffer(ctx->nvencoder, nvenc_frame->input_surface);

  1361.     if (nv_status != NV_ENC_SUCCESS) {

  1362.         return nvenc_print_error(avctx, nv_status, "Failed unlocking input buffer!");

  1363.     }

  1364. 

  1365.     return res;

  1366. }

  1367. 

  1368. static void nvenc_codec_specific_pic_params(AVCodecContext *avctx,

  1369.                                                     NV_ENC_PIC_PARAMS *params)

  1370. {

  1371.     NvencContext *ctx = avctx->priv_data;

  1372. 

  1373.     switch (avctx->codec->id) {

  1374.     case AV_CODEC_ID_H264:

  1375.       params->codecPicParams.h264PicParams.sliceMode = ctx->encode_config.encodeCodecConfig.h264Config.sliceMode;

  1376.       params->codecPicParams.h264PicParams.sliceModeData = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;

  1377.       break;

  1378.     case AV_CODEC_ID_H265:

  1379.       params->codecPicParams.hevcPicParams.sliceMode = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceMode;

  1380.       params->codecPicParams.hevcPicParams.sliceModeData = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData;

  1381.       break;

  1382.     }

  1383. }

  1384. 

  1385. static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, NvencSurface *tmpoutsurf)

  1386. {

  1387.     NvencContext *ctx = avctx->priv_data;

  1388.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1389.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1390. 

  1391.     uint32_t slice_mode_data;

  1392.     uint32_t *slice_offsets;

  1393.     NV_ENC_LOCK_BITSTREAM lock_params = { 0 };

  1394.     NVENCSTATUS nv_status;

  1395.     int res = 0;

  1396. 

  1397.     enum AVPictureType pict_type;

  1398. 

  1399.     switch (avctx->codec->id) {

  1400.     case AV_CODEC_ID_H264:

  1401.       slice_mode_data = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;

  1402.       break;

  1403.     case AV_CODEC_ID_H265:

  1404.       slice_mode_data = ctx->encode_config.encodeCodecConfig.hevcConfig.sliceModeData;

  1405.       break;

  1406.     default:

  1407.       av_log(avctx, AV_LOG_ERROR, "Unknown codec name\n");

  1408.       res = AVERROR(EINVAL);

  1409.       goto error;

  1410.     }

  1411.     slice_offsets = av_mallocz(slice_mode_data * sizeof(*slice_offsets));

  1412. 

  1413.     if (!slice_offsets)

  1414.         return AVERROR(ENOMEM);

  1415. 

  1416.     lock_params.version = NV_ENC_LOCK_BITSTREAM_VER;

  1417. 

  1418.     lock_params.doNotWait = 0;

  1419.     lock_params.outputBitstream = tmpoutsurf->output_surface;

  1420.     lock_params.sliceOffsets = slice_offsets;

  1421. 

  1422.     nv_status = p_nvenc->nvEncLockBitstream(ctx->nvencoder, &lock_params);

  1423.     if (nv_status != NV_ENC_SUCCESS) {

  1424.         res = nvenc_print_error(avctx, nv_status, "Failed locking bitstream buffer");

  1425.         goto error;

  1426.     }

  1427. 

  1428.     if (res = ff_alloc_packet2(avctx, pkt, lock_params.bitstreamSizeInBytes,0)) {

  1429.         p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);

  1430.         goto error;

  1431.     }

  1432. 

  1433.     memcpy(pkt->data, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes);

  1434. 

  1435.     nv_status = p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);

  1436.     if (nv_status != NV_ENC_SUCCESS)

  1437.         nvenc_print_error(avctx, nv_status, "Failed unlocking bitstream buffer, expect the gates of mordor to open");

  1438. 

  1439.     switch (lock_params.pictureType) {

  1440.     case NV_ENC_PIC_TYPE_IDR:

  1441.         pkt->flags |= AV_PKT_FLAG_KEY;

  1442.     case NV_ENC_PIC_TYPE_I:

  1443.         pict_type = AV_PICTURE_TYPE_I;

  1444.         break;

  1445.     case NV_ENC_PIC_TYPE_P:

  1446.         pict_type = AV_PICTURE_TYPE_P;

  1447.         break;

  1448.     case NV_ENC_PIC_TYPE_B:

  1449.         pict_type = AV_PICTURE_TYPE_B;

  1450.         break;

  1451.     case NV_ENC_PIC_TYPE_BI:

  1452.         pict_type = AV_PICTURE_TYPE_BI;

  1453.         break;

  1454.     default:

  1455.         av_log(avctx, AV_LOG_ERROR, "Unknown picture type encountered, expect the output to be broken.\n");

  1456.         av_log(avctx, AV_LOG_ERROR, "Please report this error and include as much information on how to reproduce it as possible.\n");

  1457.         res = AVERROR_EXTERNAL;

  1458.         goto error;

  1459.     }

  1460. 

  1461. #if FF_API_CODED_FRAME

  1462. FF_DISABLE_DEPRECATION_WARNINGS

  1463.     avctx->coded_frame->pict_type = pict_type;

  1464. FF_ENABLE_DEPRECATION_WARNINGS

  1465. #endif

  1466. 

  1467.     ff_side_data_set_encoder_stats(pkt,

  1468.         (lock_params.frameAvgQP - 1) * FF_QP2LAMBDA, NULL, 0, pict_type);

  1469. 

  1470.     pkt->pts = lock_params.outputTimeStamp;

  1471.     pkt->dts = timestamp_queue_dequeue(ctx->timestamp_list);

  1472. 

  1473.     /* when there're b frame(s), set dts offset */

  1474.     if (ctx->encode_config.frameIntervalP >= 2)

  1475.         pkt->dts -= 1;

  1476. 

  1477.     if (pkt->dts > pkt->pts)

  1478.         pkt->dts = pkt->pts;

  1479. 

  1480.     if (ctx->last_dts != AV_NOPTS_VALUE && pkt->dts <= ctx->last_dts)

  1481.         pkt->dts = ctx->last_dts + 1;

  1482. 

  1483.     ctx->last_dts = pkt->dts;

  1484. 

  1485.     av_free(slice_offsets);

  1486. 

  1487.     return 0;

  1488. 

  1489. error:

  1490. 

  1491.     av_free(slice_offsets);

  1492.     timestamp_queue_dequeue(ctx->timestamp_list);

  1493. 

  1494.     return res;

  1495. }

  1496. 

  1497. static int output_ready(NvencContext *ctx, int flush)

  1498. {

  1499.     int nb_ready, nb_pending;

  1500. 

  1501.     nb_ready   = av_fifo_size(ctx->output_surface_ready_queue)   / sizeof(NvencSurface*);

  1502.     nb_pending = av_fifo_size(ctx->output_surface_queue)         / sizeof(NvencSurface*);

  1503.     return nb_ready > 0 && (flush || nb_ready + nb_pending >= ctx->buffer_delay);

  1504. }

  1505. 

  1506. static int nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,

  1507.     const AVFrame *frame, int *got_packet)

  1508. {

  1509.     NVENCSTATUS nv_status;

  1510.     NvencSurface *tmpoutsurf, *inSurf;

  1511.     int res;

  1512. 

  1513.     NvencContext *ctx = avctx->priv_data;

  1514.     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;

  1515.     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;

  1516. 

  1517.     NV_ENC_PIC_PARAMS pic_params = { 0 };

  1518.     pic_params.version = NV_ENC_PIC_PARAMS_VER;

  1519. 

  1520.     if (frame) {

  1521.         inSurf = get_free_frame(ctx);

  1522.         av_assert0(inSurf);

  1523. 

  1524.         res = nvenc_upload_frame(avctx, frame, inSurf);

  1525.         if (res) {

  1526.             inSurf->lockCount = 0;

  1527.             return res;

  1528.         }

  1529. 

  1530.         pic_params.inputBuffer = inSurf->input_surface;

  1531.         pic_params.bufferFmt = inSurf->format;

  1532.         pic_params.inputWidth = avctx->width;

  1533.         pic_params.inputHeight = avctx->height;

  1534.         pic_params.outputBitstream = inSurf->output_surface;

  1535.         pic_params.completionEvent = 0;

  1536. 

  1537.         if (avctx->flags & AV_CODEC_FLAG_INTERLACED_DCT) {

  1538.             if (frame->top_field_first) {

  1539.                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;

  1540.             } else {

  1541.                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;

  1542.             }

  1543.         } else {

  1544.             pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;

  1545.         }

  1546. 

  1547.         pic_params.encodePicFlags = 0;

  1548.         pic_params.inputTimeStamp = frame->pts;

  1549.         pic_params.inputDuration = 0;

  1550. 

  1551.         nvenc_codec_specific_pic_params(avctx, &pic_params);

  1552. 

  1553.         timestamp_queue_enqueue(ctx->timestamp_list, frame->pts);

  1554.     } else {

  1555.         pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;

  1556.     }

  1557. 

  1558.     nv_status = p_nvenc->nvEncEncodePicture(ctx->nvencoder, &pic_params);

  1559. 

  1560.     if (frame && nv_status == NV_ENC_ERR_NEED_MORE_INPUT)

  1561.         av_fifo_generic_write(ctx->output_surface_queue, &inSurf, sizeof(inSurf), NULL);

  1562. 

  1563.     if (nv_status != NV_ENC_SUCCESS && nv_status != NV_ENC_ERR_NEED_MORE_INPUT) {

  1564.         return nvenc_print_error(avctx, nv_status, "EncodePicture failed!");

  1565.     }

  1566. 

  1567.     if (nv_status != NV_ENC_ERR_NEED_MORE_INPUT) {

  1568.         while (av_fifo_size(ctx->output_surface_queue) > 0) {

  1569.             av_fifo_generic_read(ctx->output_surface_queue, &tmpoutsurf, sizeof(tmpoutsurf), NULL);

  1570.             av_fifo_generic_write(ctx->output_surface_ready_queue, &tmpoutsurf, sizeof(tmpoutsurf), NULL);

  1571.         }

  1572. 

  1573.         if (frame)

  1574.             av_fifo_generic_write(ctx->output_surface_ready_queue, &inSurf, sizeof(inSurf), NULL);

  1575.     }

  1576. 

  1577.     if (output_ready(ctx, !frame)) {

  1578.         av_fifo_generic_read(ctx->output_surface_ready_queue, &tmpoutsurf, sizeof(tmpoutsurf), NULL);

  1579. 

  1580.         res = process_output_surface(avctx, pkt, tmpoutsurf);

  1581. 

  1582.         if (res)

  1583.             return res;

  1584. 

  1585.         av_assert0(tmpoutsurf->lockCount);

  1586.         tmpoutsurf->lockCount--;

  1587. 

  1588.         *got_packet = 1;

  1589.     } else {

  1590.         *got_packet = 0;

  1591.     }

  1592. 

  1593.     return 0;

  1594. }

  1595. 

  1596. static const enum AVPixelFormat pix_fmts_nvenc[] = {

  1597.     AV_PIX_FMT_YUV420P,

  1598.     AV_PIX_FMT_NV12,

  1599.     AV_PIX_FMT_YUV444P,

  1600.     AV_PIX_FMT_NONE

  1601. };

  1602. 

  1603. #define OFFSET(x) offsetof(NvencContext, x)

  1604. #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM

  1605. static const AVOption options[] = {

  1606.     { "preset", "Set the encoding preset (one of slow = hq 2pass, medium = hq, fast = hp, hq, hp, bd, ll, llhq, llhp, lossless, losslesshp, default)", OFFSET(preset), AV_OPT_TYPE_STRING, { .str = "medium" }, 0, 0, VE },

  1607.     { "profile", "Set the encoding profile (high, main, baseline or high444p)", OFFSET(profile), AV_OPT_TYPE_STRING, { .str = "main" }, 0, 0, VE },

  1608.     { "level", "Set the encoding level restriction (auto, 1.0, 1.0b, 1.1, 1.2, ..., 4.2, 5.0, 5.1)", OFFSET(level), AV_OPT_TYPE_STRING, { .str = "auto" }, 0, 0, VE },

  1609.     { "tier", "Set the encoding tier (main or high)", OFFSET(tier), AV_OPT_TYPE_STRING, { .str = "main" }, 0, 0, VE },

  1610.     { "cbr", "Use cbr encoding mode", OFFSET(cbr), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },

  1611.     { "2pass", "Use 2pass encoding mode", OFFSET(twopass), AV_OPT_TYPE_BOOL, { .i64 = -1 }, -1, 1, VE },

  1612.     { "gpu", "Selects which NVENC capable GPU to use. First GPU is 0, second is 1, and so on.", OFFSET(gpu), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INT_MAX, VE },

  1613.     { "delay", "Delays frame output by the given amount of frames.", OFFSET(buffer_delay), AV_OPT_TYPE_INT, { .i64 = INT_MAX }, 0, INT_MAX, VE },

  1614.     { NULL }

  1615. };

  1616. 

  1617. static const AVCodecDefault nvenc_defaults[] = {

  1618.     { "b", "2M" },

  1619.     { "qmin", "-1" },

  1620.     { "qmax", "-1" },

  1621.     { "qdiff", "-1" },

  1622.     { "qblur", "-1" },

  1623.     { "qcomp", "-1" },

  1624.     { "g", "250" },

  1625.     { "bf", "0" },

  1626.     { NULL },

  1627. };

  1628. 

  1629. #if CONFIG_NVENC_ENCODER

  1630. static const AVClass nvenc_class = {

  1631.     .class_name = "nvenc",

  1632.     .item_name = av_default_item_name,

  1633.     .option = options,

  1634.     .version = LIBAVUTIL_VERSION_INT,

  1635. };

  1636. 

  1637. AVCodec ff_nvenc_encoder = {

  1638.     .name = "nvenc",

  1639.     .long_name = NULL_IF_CONFIG_SMALL("NVIDIA NVENC h264 encoder"),

  1640.     .type = AVMEDIA_TYPE_VIDEO,

  1641.     .id = AV_CODEC_ID_H264,

  1642.     .priv_data_size = sizeof(NvencContext),

  1643.     .init = nvenc_encode_init,

  1644.     .encode2 = nvenc_encode_frame,

  1645.     .close = nvenc_encode_close,

  1646.     .capabilities = AV_CODEC_CAP_DELAY,

  1647.     .priv_class = &nvenc_class,

  1648.     .defaults = nvenc_defaults,

  1649.     .pix_fmts = pix_fmts_nvenc,

  1650. };

  1651. #endif

  1652. 

  1653. /* Add an alias for nvenc_h264 */

  1654. #if CONFIG_NVENC_H264_ENCODER

  1655. static const AVClass nvenc_h264_class = {

  1656.     .class_name = "nvenc_h264",

  1657.     .item_name = av_default_item_name,

  1658.     .option = options,

  1659.     .version = LIBAVUTIL_VERSION_INT,

  1660. };

  1661. 

  1662. AVCodec ff_nvenc_h264_encoder = {

  1663.     .name = "nvenc_h264",

  1664.     .long_name = NULL_IF_CONFIG_SMALL("NVIDIA NVENC h264 encoder"),

  1665.     .type = AVMEDIA_TYPE_VIDEO,

  1666.     .id = AV_CODEC_ID_H264,

  1667.     .priv_data_size = sizeof(NvencContext),

  1668.     .init = nvenc_encode_init,

  1669.     .encode2 = nvenc_encode_frame,

  1670.     .close = nvenc_encode_close,

  1671.     .capabilities = AV_CODEC_CAP_DELAY,

  1672.     .priv_class = &nvenc_h264_class,

  1673.     .defaults = nvenc_defaults,

  1674.     .pix_fmts = pix_fmts_nvenc,

  1675. };

  1676. #endif

  1677. 

  1678. #if CONFIG_NVENC_HEVC_ENCODER

  1679. static const AVClass nvenc_hevc_class = {

  1680.     .class_name = "nvenc_hevc",

  1681.     .item_name = av_default_item_name,

  1682.     .option = options,

  1683.     .version = LIBAVUTIL_VERSION_INT,

  1684. };

  1685. 

  1686. AVCodec ff_nvenc_hevc_encoder = {

  1687.     .name = "nvenc_hevc",

  1688.     .long_name = NULL_IF_CONFIG_SMALL("NVIDIA NVENC hevc encoder"),

  1689.     .type = AVMEDIA_TYPE_VIDEO,

  1690.     .id = AV_CODEC_ID_H265,

  1691.     .priv_data_size = sizeof(NvencContext),

  1692.     .init = nvenc_encode_init,

  1693.     .encode2 = nvenc_encode_frame,

  1694.     .close = nvenc_encode_close,

  1695.     .capabilities = AV_CODEC_CAP_DELAY,

  1696.     .priv_class = &nvenc_hevc_class,

  1697.     .defaults = nvenc_defaults,

  1698.     .pix_fmts = pix_fmts_nvenc,

  1699. };

  1700. #endif