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