git.sesse.net Git - ffmpeg/blob - 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 #if defined(_WIN32) && !defined(__CYGWIN__)
  23 #include <windows.h>
  24 #else
  25 #include <dlfcn.h>
  26 #endif
  27
  28 /* NVENC API is stdcall in cygwin, as it's still Windows */
  29 #if defined(__CYGWIN__) && !defined(_WIN32)
  30 #define _WIN32
  31 #define TEMP_WIN32
  32 #endif
  33
  34 #include <nvEncodeAPI.h>
  35
  36 #if defined(TEMP_WIN32)
  37 #undef _WIN32
  38 #endif
  39
  40
  41 #include "libavutil/internal.h"
  42 #include "libavutil/imgutils.h"
  43 #include "libavutil/avassert.h"
  44 #include "libavutil/opt.h"
  45 #include "libavutil/mem.h"
  46 #include "avcodec.h"
  47 #include "internal.h"
  48 #include "thread.h"
  49
  50 #if defined(_WIN32) || defined(__CYGWIN__)
  51 #define CUDAAPI __stdcall
  52 #else
  53 #define CUDAAPI
  54 #endif
  55
  56 #if defined(_WIN32) && !defined(__CYGWIN__)
  57 #define LOAD_FUNC(l, s) GetProcAddress(l, s)
  58 #define DL_CLOSE_FUNC(l) FreeLibrary(l)
  59 #else
  60 #define LOAD_FUNC(l, s) dlsym(l, s)
  61 #define DL_CLOSE_FUNC(l) dlclose(l)
  62 #endif
  63
  64 typedef enum cudaError_enum {
  65     CUDA_SUCCESS = 0
  66 } CUresult;
  67 typedef int CUdevice;
  68 typedef void* CUcontext;
  69
  70 typedef CUresult(CUDAAPI *PCUINIT)(unsigned int Flags);
  71 typedef CUresult(CUDAAPI *PCUDEVICEGETCOUNT)(int *count);
  72 typedef CUresult(CUDAAPI *PCUDEVICEGET)(CUdevice *device, int ordinal);
  73 typedef CUresult(CUDAAPI *PCUDEVICEGETNAME)(char *name, int len, CUdevice dev);
  74 typedef CUresult(CUDAAPI *PCUDEVICECOMPUTECAPABILITY)(int *major, int *minor, CUdevice dev);
  75 typedef CUresult(CUDAAPI *PCUCTXCREATE)(CUcontext *pctx, unsigned int flags, CUdevice dev);
  76 typedef CUresult(CUDAAPI *PCUCTXPOPCURRENT)(CUcontext *pctx);
  77 typedef CUresult(CUDAAPI *PCUCTXDESTROY)(CUcontext ctx);
  78
  79 typedef NVENCSTATUS (NVENCAPI* PNVENCODEAPICREATEINSTANCE)(NV_ENCODE_API_FUNCTION_LIST *functionList);
  80
  81 static const GUID dummy_license = { 0x0, 0x0, 0x0, { 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0 } };
  82
  83 typedef struct NvencInputSurface
  84 {
  85     NV_ENC_INPUT_PTR input_surface;
  86     int width;
  87     int height;
  88
  89     int lockCount;
  90
  91     NV_ENC_BUFFER_FORMAT format;
  92 } NvencInputSurface;
  93
  94 typedef struct NvencOutputSurface
  95 {
  96     NV_ENC_OUTPUT_PTR output_surface;
  97     int size;
  98
  99     NvencInputSurface* input_surface;
 100
 101     int busy;
 102 } NvencOutputSurface;
 103
 104 typedef struct NvencData
 105 {
 106     union {
 107         int64_t timestamp;
 108         NvencOutputSurface *surface;
 109     };
 110 } NvencData;
 111
 112 typedef struct NvencDataList
 113 {
 114     NvencData* data;
 115
 116     uint32_t pos;
 117     uint32_t count;
 118     uint32_t size;
 119 } NvencDataList;
 120
 121 typedef struct NvencDynLoadFunctions
 122 {
 123     PCUINIT cu_init;
 124     PCUDEVICEGETCOUNT cu_device_get_count;
 125     PCUDEVICEGET cu_device_get;
 126     PCUDEVICEGETNAME cu_device_get_name;
 127     PCUDEVICECOMPUTECAPABILITY cu_device_compute_capability;
 128     PCUCTXCREATE cu_ctx_create;
 129     PCUCTXPOPCURRENT cu_ctx_pop_current;
 130     PCUCTXDESTROY cu_ctx_destroy;
 131
 132     NV_ENCODE_API_FUNCTION_LIST nvenc_funcs;
 133     int nvenc_device_count;
 134     CUdevice nvenc_devices[16];
 135
 136 #if defined(_WIN32) && !defined(__CYGWIN__)
 137     HMODULE cuda_lib;
 138     HMODULE nvenc_lib;
 139 #else
 140     void* cuda_lib;
 141     void* nvenc_lib;
 142 #endif
 143 } NvencDynLoadFunctions;
 144
 145 typedef struct NvencContext
 146 {
 147     AVClass *avclass;
 148
 149     NvencDynLoadFunctions nvenc_dload_funcs;
 150
 151     NV_ENC_INITIALIZE_PARAMS init_encode_params;
 152     NV_ENC_CONFIG encode_config;
 153     CUcontext cu_context;
 154
 155     int max_surface_count;
 156     NvencInputSurface *input_surfaces;
 157     NvencOutputSurface *output_surfaces;
 158
 159     NvencDataList output_surface_queue;
 160     NvencDataList output_surface_ready_queue;
 161     NvencDataList timestamp_list;
 162     int64_t last_dts;
 163
 164     void *nvencoder;
 165
 166     char *preset;
 167     int cbr;
 168     int twopass;
 169     int gobpattern;
 170     int gpu;
 171 } NvencContext;
 172
 173 static NvencData* data_queue_dequeue(NvencDataList* queue)
 174 {
 175     uint32_t mask;
 176     uint32_t read_pos;
 177
 178     av_assert0(queue);
 179     av_assert0(queue->size);
 180     av_assert0(queue->data);
 181
 182     if (!queue->count)
 183         return NULL;
 184
 185     /* Size always is a multiple of two */
 186     mask = queue->size - 1;
 187     read_pos = (queue->pos - queue->count) & mask;
 188     queue->count--;
 189
 190     return &queue->data[read_pos];
 191 }
 192
 193 static int data_queue_enqueue(NvencDataList* queue, NvencData *data)
 194 {
 195     NvencDataList new_queue;
 196     NvencData* tmp_data;
 197     uint32_t mask;
 198
 199     if (!queue->size) {
 200         /* size always has to be a multiple of two */
 201         queue->size = 4;
 202         queue->pos = 0;
 203         queue->count = 0;
 204
 205         queue->data = av_malloc(queue->size * sizeof(*(queue->data)));
 206
 207         if (!queue->data) {
 208             queue->size = 0;
 209             return AVERROR(ENOMEM);
 210         }
 211     }
 212
 213     if (queue->count == queue->size) {
 214         new_queue.size = queue->size << 1;
 215         new_queue.pos = 0;
 216         new_queue.count = 0;
 217         new_queue.data = av_malloc(new_queue.size * sizeof(*(queue->data)));
 218
 219         if (!new_queue.data)
 220             return AVERROR(ENOMEM);
 221
 222         while (tmp_data = data_queue_dequeue(queue))
 223             data_queue_enqueue(&new_queue, tmp_data);
 224
 225         av_free(queue->data);
 226         *queue = new_queue;
 227     }
 228
 229     mask = queue->size - 1;
 230
 231     queue->data[queue->pos] = *data;
 232     queue->pos = (queue->pos + 1) & mask;
 233     queue->count++;
 234
 235     return 0;
 236 }
 237
 238 static int out_surf_queue_enqueue(NvencDataList* queue, NvencOutputSurface* surface)
 239 {
 240     NvencData data;
 241     data.surface = surface;
 242
 243     return data_queue_enqueue(queue, &data);
 244 }
 245
 246 static NvencOutputSurface* out_surf_queue_dequeue(NvencDataList* queue)
 247 {
 248     NvencData* res = data_queue_dequeue(queue);
 249
 250     if (!res)
 251         return NULL;
 252
 253     return res->surface;
 254 }
 255
 256 static int timestamp_queue_enqueue(NvencDataList* queue, int64_t timestamp)
 257 {
 258     NvencData data;
 259     data.timestamp = timestamp;
 260
 261     return data_queue_enqueue(queue, &data);
 262 }
 263
 264 static int64_t timestamp_queue_dequeue(NvencDataList* queue)
 265 {
 266     NvencData* res = data_queue_dequeue(queue);
 267
 268     if (!res)
 269         return AV_NOPTS_VALUE;
 270
 271     return res->timestamp;
 272 }
 273
 274 #define CHECK_LOAD_FUNC(t, f, s) \
 275 do { \
 276     (f) = (t)LOAD_FUNC(dl_fn->cuda_lib, s); \
 277     if (!(f)) { \
 278         av_log(avctx, AV_LOG_FATAL, "Failed loading %s from CUDA library\n", s); \
 279         goto error; \
 280     } \
 281 } while (0)
 282
 283 static av_cold int nvenc_dyload_cuda(AVCodecContext *avctx)
 284 {
 285     NvencContext *ctx = avctx->priv_data;
 286     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 287
 288     if (dl_fn->cuda_lib)
 289         return 1;
 290
 291 #if defined(_WIN32)
 292     dl_fn->cuda_lib = LoadLibrary(TEXT("nvcuda.dll"));
 293 #elif defined(__CYGWIN__)
 294     dl_fn->cuda_lib = dlopen("nvcuda.dll", RTLD_LAZY);
 295 #else
 296     dl_fn->cuda_lib = dlopen("libcuda.so", RTLD_LAZY);
 297 #endif
 298
 299     if (!dl_fn->cuda_lib) {
 300         av_log(avctx, AV_LOG_FATAL, "Failed loading CUDA library\n");
 301         goto error;
 302     }
 303
 304     CHECK_LOAD_FUNC(PCUINIT, dl_fn->cu_init, "cuInit");
 305     CHECK_LOAD_FUNC(PCUDEVICEGETCOUNT, dl_fn->cu_device_get_count, "cuDeviceGetCount");
 306     CHECK_LOAD_FUNC(PCUDEVICEGET, dl_fn->cu_device_get, "cuDeviceGet");
 307     CHECK_LOAD_FUNC(PCUDEVICEGETNAME, dl_fn->cu_device_get_name, "cuDeviceGetName");
 308     CHECK_LOAD_FUNC(PCUDEVICECOMPUTECAPABILITY, dl_fn->cu_device_compute_capability, "cuDeviceComputeCapability");
 309     CHECK_LOAD_FUNC(PCUCTXCREATE, dl_fn->cu_ctx_create, "cuCtxCreate_v2");
 310     CHECK_LOAD_FUNC(PCUCTXPOPCURRENT, dl_fn->cu_ctx_pop_current, "cuCtxPopCurrent_v2");
 311     CHECK_LOAD_FUNC(PCUCTXDESTROY, dl_fn->cu_ctx_destroy, "cuCtxDestroy_v2");
 312
 313     return 1;
 314
 315 error:
 316
 317     if (dl_fn->cuda_lib)
 318         DL_CLOSE_FUNC(dl_fn->cuda_lib);
 319
 320     dl_fn->cuda_lib = NULL;
 321
 322     return 0;
 323 }
 324
 325 static av_cold int check_cuda_errors(AVCodecContext *avctx, CUresult err, const char *func)
 326 {
 327     if (err != CUDA_SUCCESS) {
 328         av_log(avctx, AV_LOG_FATAL, ">> %s - failed with error code 0x%x\n", func, err);
 329         return 0;
 330     }
 331     return 1;
 332 }
 333 #define check_cuda_errors(f) if (!check_cuda_errors(avctx, f, #f)) goto error
 334
 335 static av_cold int nvenc_check_cuda(AVCodecContext *avctx)
 336 {
 337     int device_count = 0;
 338     CUdevice cu_device = 0;
 339     char gpu_name[128];
 340     int smminor = 0, smmajor = 0;
 341     int i, smver;
 342
 343     NvencContext *ctx = avctx->priv_data;
 344     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 345
 346     if (!nvenc_dyload_cuda(avctx))
 347         return 0;
 348
 349     if (dl_fn->nvenc_device_count > 0)
 350         return 1;
 351
 352     check_cuda_errors(dl_fn->cu_init(0));
 353
 354     check_cuda_errors(dl_fn->cu_device_get_count(&device_count));
 355
 356     if (!device_count) {
 357         av_log(avctx, AV_LOG_FATAL, "No CUDA capable devices found\n");
 358         goto error;
 359     }
 360
 361     av_log(avctx, AV_LOG_VERBOSE, "%d CUDA capable devices found\n", device_count);
 362
 363     dl_fn->nvenc_device_count = 0;
 364
 365     for (i = 0; i < device_count; ++i) {
 366         check_cuda_errors(dl_fn->cu_device_get(&cu_device, i));
 367         check_cuda_errors(dl_fn->cu_device_get_name(gpu_name, sizeof(gpu_name), cu_device));
 368         check_cuda_errors(dl_fn->cu_device_compute_capability(&smmajor, &smminor, cu_device));
 369
 370         smver = (smmajor << 4) | smminor;
 371
 372         av_log(avctx, AV_LOG_VERBOSE, "[ GPU #%d - < %s > has Compute SM %d.%d, NVENC %s ]\n", i, gpu_name, smmajor, smminor, (smver >= 0x30) ? "Available" : "Not Available");
 373
 374         if (smver >= 0x30)
 375             dl_fn->nvenc_devices[dl_fn->nvenc_device_count++] = cu_device;
 376     }
 377
 378     if (!dl_fn->nvenc_device_count) {
 379         av_log(avctx, AV_LOG_FATAL, "No NVENC capable devices found\n");
 380         goto error;
 381     }
 382
 383     return 1;
 384
 385 error:
 386
 387     dl_fn->nvenc_device_count = 0;
 388
 389     return 0;
 390 }
 391
 392 static av_cold int nvenc_dyload_nvenc(AVCodecContext *avctx)
 393 {
 394     PNVENCODEAPICREATEINSTANCE nvEncodeAPICreateInstance = 0;
 395     NVENCSTATUS nvstatus;
 396
 397     NvencContext *ctx = avctx->priv_data;
 398     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 399
 400     if (!nvenc_check_cuda(avctx))
 401         return 0;
 402
 403     if (dl_fn->nvenc_lib)
 404         return 1;
 405
 406 #if defined(_WIN32)
 407     if (sizeof(void*) == 8) {
 408         dl_fn->nvenc_lib = LoadLibrary(TEXT("nvEncodeAPI64.dll"));
 409     } else {
 410         dl_fn->nvenc_lib = LoadLibrary(TEXT("nvEncodeAPI.dll"));
 411     }
 412 #elif defined(__CYGWIN__)
 413     if (sizeof(void*) == 8) {
 414         dl_fn->nvenc_lib = dlopen("nvEncodeAPI64.dll", RTLD_LAZY);
 415     } else {
 416         dl_fn->nvenc_lib = dlopen("nvEncodeAPI.dll", RTLD_LAZY);
 417     }
 418 #else
 419     dl_fn->nvenc_lib = dlopen("libnvidia-encode.so.1", RTLD_LAZY);
 420 #endif
 421
 422     if (!dl_fn->nvenc_lib) {
 423         av_log(avctx, AV_LOG_FATAL, "Failed loading the nvenc library\n");
 424         goto error;
 425     }
 426
 427     nvEncodeAPICreateInstance = (PNVENCODEAPICREATEINSTANCE)LOAD_FUNC(dl_fn->nvenc_lib, "NvEncodeAPICreateInstance");
 428
 429     if (!nvEncodeAPICreateInstance) {
 430         av_log(avctx, AV_LOG_FATAL, "Failed to load nvenc entrypoint\n");
 431         goto error;
 432     }
 433
 434     dl_fn->nvenc_funcs.version = NV_ENCODE_API_FUNCTION_LIST_VER;
 435
 436     nvstatus = nvEncodeAPICreateInstance(&dl_fn->nvenc_funcs);
 437
 438     if (nvstatus != NV_ENC_SUCCESS) {
 439         av_log(avctx, AV_LOG_FATAL, "Failed to create nvenc instance\n");
 440         goto error;
 441     }
 442
 443     av_log(avctx, AV_LOG_VERBOSE, "Nvenc initialized successfully\n");
 444
 445     return 1;
 446
 447 error:
 448     if (dl_fn->nvenc_lib)
 449         DL_CLOSE_FUNC(dl_fn->nvenc_lib);
 450
 451     dl_fn->nvenc_lib = NULL;
 452
 453     return 0;
 454 }
 455
 456 static av_cold void nvenc_unload_nvenc(AVCodecContext *avctx)
 457 {
 458     NvencContext *ctx = avctx->priv_data;
 459     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 460
 461     DL_CLOSE_FUNC(dl_fn->nvenc_lib);
 462     dl_fn->nvenc_lib = NULL;
 463
 464     dl_fn->nvenc_device_count = 0;
 465
 466     DL_CLOSE_FUNC(dl_fn->cuda_lib);
 467     dl_fn->cuda_lib = NULL;
 468
 469     dl_fn->cu_init = NULL;
 470     dl_fn->cu_device_get_count = NULL;
 471     dl_fn->cu_device_get = NULL;
 472     dl_fn->cu_device_get_name = NULL;
 473     dl_fn->cu_device_compute_capability = NULL;
 474     dl_fn->cu_ctx_create = NULL;
 475     dl_fn->cu_ctx_pop_current = NULL;
 476     dl_fn->cu_ctx_destroy = NULL;
 477
 478     av_log(avctx, AV_LOG_VERBOSE, "Nvenc unloaded\n");
 479 }
 480
 481 static av_cold int nvenc_encode_init(AVCodecContext *avctx)
 482 {
 483     NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS encode_session_params = { 0 };
 484     NV_ENC_PRESET_CONFIG preset_config = { 0 };
 485     CUcontext cu_context_curr;
 486     CUresult cu_res;
 487     GUID encoder_preset = NV_ENC_PRESET_HQ_GUID;
 488     GUID license = dummy_license;
 489     NVENCSTATUS nv_status = NV_ENC_SUCCESS;
 490     int surfaceCount = 0;
 491     int i, num_mbs;
 492     int isLL = 0;
 493     int res = 0;
 494
 495     NvencContext *ctx = avctx->priv_data;
 496     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 497     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
 498
 499     if (!nvenc_dyload_nvenc(avctx))
 500         return AVERROR_EXTERNAL;
 501
 502     avctx->coded_frame = av_frame_alloc();
 503     if (!avctx->coded_frame) {
 504         res = AVERROR(ENOMEM);
 505         goto error;
 506     }
 507
 508     ctx->last_dts = AV_NOPTS_VALUE;
 509
 510     ctx->encode_config.version = NV_ENC_CONFIG_VER;
 511     ctx->init_encode_params.version = NV_ENC_INITIALIZE_PARAMS_VER;
 512     preset_config.version = NV_ENC_PRESET_CONFIG_VER;
 513     preset_config.presetCfg.version = NV_ENC_CONFIG_VER;
 514     encode_session_params.version = NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER;
 515     encode_session_params.apiVersion = NVENCAPI_VERSION;
 516     encode_session_params.clientKeyPtr = &license;
 517
 518     if (ctx->gpu >= dl_fn->nvenc_device_count) {
 519         av_log(avctx, AV_LOG_FATAL, "Requested GPU %d, but only %d GPUs are available!\n", ctx->gpu, dl_fn->nvenc_device_count);
 520         res = AVERROR(EINVAL);
 521         goto error;
 522     }
 523
 524     ctx->cu_context = NULL;
 525     cu_res = dl_fn->cu_ctx_create(&ctx->cu_context, 0, dl_fn->nvenc_devices[ctx->gpu]);
 526
 527     if (cu_res != CUDA_SUCCESS) {
 528         av_log(avctx, AV_LOG_FATAL, "Failed creating CUDA context for NVENC: 0x%x\n", (int)cu_res);
 529         res = AVERROR_EXTERNAL;
 530         goto error;
 531     }
 532
 533     cu_res = dl_fn->cu_ctx_pop_current(&cu_context_curr);
 534
 535     if (cu_res != CUDA_SUCCESS) {
 536         av_log(avctx, AV_LOG_FATAL, "Failed popping CUDA context: 0x%x\n", (int)cu_res);
 537         res = AVERROR_EXTERNAL;
 538         goto error;
 539     }
 540
 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         av_log(avctx, AV_LOG_FATAL, "OpenEncodeSessionEx failed: 0x%x - invalid license key?\n", (int)nv_status);
 548         res = AVERROR_EXTERNAL;
 549         goto error;
 550     }
 551
 552     if (ctx->preset) {
 553         if (!strcmp(ctx->preset, "hp")) {
 554             encoder_preset = NV_ENC_PRESET_HP_GUID;
 555         } else if (!strcmp(ctx->preset, "hq")) {
 556             encoder_preset = NV_ENC_PRESET_HQ_GUID;
 557         } else if (!strcmp(ctx->preset, "bd")) {
 558             encoder_preset = NV_ENC_PRESET_BD_GUID;
 559         } else if (!strcmp(ctx->preset, "ll")) {
 560             encoder_preset = NV_ENC_PRESET_LOW_LATENCY_DEFAULT_GUID;
 561             isLL = 1;
 562         } else if (!strcmp(ctx->preset, "llhp")) {
 563             encoder_preset = NV_ENC_PRESET_LOW_LATENCY_HP_GUID;
 564             isLL = 1;
 565         } else if (!strcmp(ctx->preset, "llhq")) {
 566             encoder_preset = NV_ENC_PRESET_LOW_LATENCY_HQ_GUID;
 567             isLL = 1;
 568         } else if (!strcmp(ctx->preset, "default")) {
 569             encoder_preset = NV_ENC_PRESET_DEFAULT_GUID;
 570         } else {
 571             av_log(avctx, AV_LOG_FATAL, "Preset \"%s\" is unknown! Supported presets: hp, hq, bd, ll, llhp, llhq, default\n", ctx->preset);
 572             res = AVERROR(EINVAL);
 573             goto error;
 574         }
 575     }
 576
 577     nv_status = p_nvenc->nvEncGetEncodePresetConfig(ctx->nvencoder, NV_ENC_CODEC_H264_GUID, encoder_preset, &preset_config);
 578     if (nv_status != NV_ENC_SUCCESS) {
 579         av_log(avctx, AV_LOG_FATAL, "GetEncodePresetConfig failed: 0x%x\n", (int)nv_status);
 580         res = AVERROR_EXTERNAL;
 581         goto error;
 582     }
 583
 584     ctx->init_encode_params.encodeGUID = NV_ENC_CODEC_H264_GUID;
 585     ctx->init_encode_params.encodeHeight = avctx->height;
 586     ctx->init_encode_params.encodeWidth = avctx->width;
 587     ctx->init_encode_params.darHeight = avctx->height;
 588     ctx->init_encode_params.darWidth = avctx->width;
 589     ctx->init_encode_params.frameRateNum = avctx->time_base.den;
 590     ctx->init_encode_params.frameRateDen = avctx->time_base.num * avctx->ticks_per_frame;
 591
 592     num_mbs = ((avctx->width + 15) >> 4) * ((avctx->height + 15) >> 4);
 593     ctx->max_surface_count = (num_mbs >= 8160) ? 32 : 48;
 594
 595     ctx->init_encode_params.enableEncodeAsync = 0;
 596     ctx->init_encode_params.enablePTD = 1;
 597
 598     ctx->init_encode_params.presetGUID = encoder_preset;
 599
 600     ctx->init_encode_params.encodeConfig = &ctx->encode_config;
 601     memcpy(&ctx->encode_config, &preset_config.presetCfg, sizeof(ctx->encode_config));
 602     ctx->encode_config.version = NV_ENC_CONFIG_VER;
 603
 604     if (avctx->gop_size >= 0) {
 605         ctx->encode_config.gopLength = avctx->gop_size;
 606         ctx->encode_config.encodeCodecConfig.h264Config.idrPeriod = avctx->gop_size;
 607     }
 608
 609     if (avctx->bit_rate > 0)
 610         ctx->encode_config.rcParams.averageBitRate = avctx->bit_rate;
 611
 612     if (avctx->rc_max_rate > 0)
 613         ctx->encode_config.rcParams.maxBitRate = avctx->rc_max_rate;
 614
 615     if (ctx->cbr) {
 616         if (!ctx->twopass) {
 617             ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CBR;
 618         } else if (ctx->twopass == 1 || isLL) {
 619             ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_2_PASS_QUALITY;
 620
 621             ctx->encode_config.encodeCodecConfig.h264Config.adaptiveTransformMode = NV_ENC_H264_ADAPTIVE_TRANSFORM_ENABLE;
 622             ctx->encode_config.encodeCodecConfig.h264Config.fmoMode = NV_ENC_H264_FMO_DISABLE;
 623
 624             if (!isLL)
 625                 av_log(avctx, AV_LOG_WARNING, "Twopass mode is only known to work with low latency (ll, llhq, llhp) presets.\n");
 626         } else {
 627             ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CBR;
 628         }
 629     } else if (avctx->global_quality > 0) {
 630         ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_CONSTQP;
 631         ctx->encode_config.rcParams.constQP.qpInterB = avctx->global_quality;
 632         ctx->encode_config.rcParams.constQP.qpInterP = avctx->global_quality;
 633         ctx->encode_config.rcParams.constQP.qpIntra = avctx->global_quality;
 634
 635         avctx->qmin = -1;
 636         avctx->qmax = -1;
 637     } else if (avctx->qmin >= 0 && avctx->qmax >= 0) {
 638         ctx->encode_config.rcParams.rateControlMode = NV_ENC_PARAMS_RC_VBR;
 639
 640         ctx->encode_config.rcParams.enableMinQP = 1;
 641         ctx->encode_config.rcParams.enableMaxQP = 1;
 642
 643         ctx->encode_config.rcParams.minQP.qpInterB = avctx->qmin;
 644         ctx->encode_config.rcParams.minQP.qpInterP = avctx->qmin;
 645         ctx->encode_config.rcParams.minQP.qpIntra = avctx->qmin;
 646
 647         ctx->encode_config.rcParams.maxQP.qpInterB = avctx->qmax;
 648         ctx->encode_config.rcParams.maxQP.qpInterP = avctx->qmax;
 649         ctx->encode_config.rcParams.maxQP.qpIntra = avctx->qmax;
 650     }
 651
 652     if (avctx->rc_buffer_size > 0)
 653         ctx->encode_config.rcParams.vbvBufferSize = avctx->rc_buffer_size;
 654
 655     if (avctx->flags & CODEC_FLAG_INTERLACED_DCT) {
 656         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FIELD;
 657     } else {
 658         ctx->encode_config.frameFieldMode = NV_ENC_PARAMS_FRAME_FIELD_MODE_FRAME;
 659     }
 660
 661     switch (avctx->profile) {
 662     case FF_PROFILE_H264_BASELINE:
 663         ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_BASELINE_GUID;
 664         break;
 665     case FF_PROFILE_H264_MAIN:
 666         ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_MAIN_GUID;
 667         break;
 668     case FF_PROFILE_H264_HIGH:
 669     case FF_PROFILE_UNKNOWN:
 670         ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
 671         break;
 672     default:
 673         av_log(avctx, AV_LOG_WARNING, "Unsupported h264 profile requested, falling back to high\n");
 674         ctx->encode_config.profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
 675         break;
 676     }
 677
 678     if (ctx->gobpattern >= 0) {
 679         ctx->encode_config.frameIntervalP = 1;
 680     }
 681
 682     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourDescriptionPresentFlag = 1;
 683     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoSignalTypePresentFlag = 1;
 684
 685     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourMatrix = avctx->colorspace;
 686     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.colourPrimaries = avctx->color_primaries;
 687     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.transferCharacteristics = avctx->color_trc;
 688
 689     ctx->encode_config.encodeCodecConfig.h264Config.h264VUIParameters.videoFullRangeFlag = avctx->color_range == AVCOL_RANGE_JPEG;
 690
 691     ctx->encode_config.encodeCodecConfig.h264Config.disableSPSPPS = (avctx->flags & CODEC_FLAG_GLOBAL_HEADER) ? 1 : 0;
 692
 693     nv_status = p_nvenc->nvEncInitializeEncoder(ctx->nvencoder, &ctx->init_encode_params);
 694     if (nv_status != NV_ENC_SUCCESS) {
 695         av_log(avctx, AV_LOG_FATAL, "InitializeEncoder failed: 0x%x\n", (int)nv_status);
 696         res = AVERROR_EXTERNAL;
 697         goto error;
 698     }
 699
 700     ctx->input_surfaces = av_malloc(ctx->max_surface_count * sizeof(*ctx->input_surfaces));
 701
 702     if (!ctx->input_surfaces) {
 703         res = AVERROR(ENOMEM);
 704         goto error;
 705     }
 706
 707     ctx->output_surfaces = av_malloc(ctx->max_surface_count * sizeof(*ctx->output_surfaces));
 708
 709     if (!ctx->output_surfaces) {
 710         res = AVERROR(ENOMEM);
 711         goto error;
 712     }
 713
 714     for (surfaceCount = 0; surfaceCount < ctx->max_surface_count; ++surfaceCount) {
 715         NV_ENC_CREATE_INPUT_BUFFER allocSurf = { 0 };
 716         NV_ENC_CREATE_BITSTREAM_BUFFER allocOut = { 0 };
 717         allocSurf.version = NV_ENC_CREATE_INPUT_BUFFER_VER;
 718         allocOut.version = NV_ENC_CREATE_BITSTREAM_BUFFER_VER;
 719
 720         allocSurf.width = (avctx->width + 31) & ~31;
 721         allocSurf.height = (avctx->height + 31) & ~31;
 722
 723         allocSurf.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;
 724
 725         switch (avctx->pix_fmt) {
 726         case AV_PIX_FMT_YUV420P:
 727             allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_YV12_PL;
 728             break;
 729
 730         case AV_PIX_FMT_NV12:
 731             allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_NV12_PL;
 732             break;
 733
 734         case AV_PIX_FMT_YUV444P:
 735             allocSurf.bufferFmt = NV_ENC_BUFFER_FORMAT_YUV444_PL;
 736             break;
 737
 738         default:
 739             av_log(avctx, AV_LOG_FATAL, "Invalid input pixel format\n");
 740             res = AVERROR(EINVAL);
 741             goto error;
 742         }
 743
 744         nv_status = p_nvenc->nvEncCreateInputBuffer(ctx->nvencoder, &allocSurf);
 745         if (nv_status = NV_ENC_SUCCESS){
 746             av_log(avctx, AV_LOG_FATAL, "CreateInputBuffer failed\n");
 747             res = AVERROR_EXTERNAL;
 748             goto error;
 749         }
 750
 751         ctx->input_surfaces[surfaceCount].lockCount = 0;
 752         ctx->input_surfaces[surfaceCount].input_surface = allocSurf.inputBuffer;
 753         ctx->input_surfaces[surfaceCount].format = allocSurf.bufferFmt;
 754         ctx->input_surfaces[surfaceCount].width = allocSurf.width;
 755         ctx->input_surfaces[surfaceCount].height = allocSurf.height;
 756
 757         /* 1MB is large enough to hold most output frames. NVENC increases this automaticaly if it's not enough. */
 758         allocOut.size = 1024 * 1024;
 759
 760         allocOut.memoryHeap = NV_ENC_MEMORY_HEAP_SYSMEM_CACHED;
 761
 762         nv_status = p_nvenc->nvEncCreateBitstreamBuffer(ctx->nvencoder, &allocOut);
 763         if (nv_status = NV_ENC_SUCCESS) {
 764             av_log(avctx, AV_LOG_FATAL, "CreateBitstreamBuffer failed\n");
 765             ctx->output_surfaces[surfaceCount++].output_surface = NULL;
 766             res = AVERROR_EXTERNAL;
 767             goto error;
 768         }
 769
 770         ctx->output_surfaces[surfaceCount].output_surface = allocOut.bitstreamBuffer;
 771         ctx->output_surfaces[surfaceCount].size = allocOut.size;
 772         ctx->output_surfaces[surfaceCount].busy = 0;
 773     }
 774
 775     if (avctx->flags & CODEC_FLAG_GLOBAL_HEADER) {
 776         uint32_t outSize = 0;
 777         char tmpHeader[256];
 778         NV_ENC_SEQUENCE_PARAM_PAYLOAD payload = { 0 };
 779         payload.version = NV_ENC_SEQUENCE_PARAM_PAYLOAD_VER;
 780
 781         payload.spsppsBuffer = tmpHeader;
 782         payload.inBufferSize = sizeof(tmpHeader);
 783         payload.outSPSPPSPayloadSize = &outSize;
 784
 785         nv_status = p_nvenc->nvEncGetSequenceParams(ctx->nvencoder, &payload);
 786         if (nv_status != NV_ENC_SUCCESS) {
 787             av_log(avctx, AV_LOG_FATAL, "GetSequenceParams failed\n");
 788             goto error;
 789         }
 790
 791         avctx->extradata_size = outSize;
 792         avctx->extradata = av_mallocz(outSize + FF_INPUT_BUFFER_PADDING_SIZE);
 793
 794         if (!avctx->extradata) {
 795             res = AVERROR(ENOMEM);
 796             goto error;
 797         }
 798
 799         memcpy(avctx->extradata, tmpHeader, outSize);
 800     }
 801
 802     if (ctx->encode_config.frameIntervalP > 1)
 803         avctx->has_b_frames = 2;
 804
 805     if (ctx->encode_config.rcParams.averageBitRate > 0)
 806         avctx->bit_rate = ctx->encode_config.rcParams.averageBitRate;
 807
 808     return 0;
 809
 810 error:
 811
 812     for (i = 0; i < surfaceCount; ++i) {
 813         p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->input_surfaces[i].input_surface);
 814         if (ctx->output_surfaces[i].output_surface)
 815             p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->output_surfaces[i].output_surface);
 816     }
 817
 818     if (ctx->nvencoder)
 819         p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);
 820
 821     if (ctx->cu_context)
 822         dl_fn->cu_ctx_destroy(ctx->cu_context);
 823
 824     av_frame_free(&avctx->coded_frame);
 825
 826     nvenc_unload_nvenc(avctx);
 827
 828     ctx->nvencoder = NULL;
 829     ctx->cu_context = NULL;
 830
 831     return res;
 832 }
 833
 834 static av_cold int nvenc_encode_close(AVCodecContext *avctx)
 835 {
 836     NvencContext *ctx = avctx->priv_data;
 837     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 838     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
 839     int i;
 840
 841     av_freep(&ctx->timestamp_list.data);
 842     av_freep(&ctx->output_surface_ready_queue.data);
 843     av_freep(&ctx->output_surface_queue.data);
 844
 845     for (i = 0; i < ctx->max_surface_count; ++i) {
 846         p_nvenc->nvEncDestroyInputBuffer(ctx->nvencoder, ctx->input_surfaces[i].input_surface);
 847         p_nvenc->nvEncDestroyBitstreamBuffer(ctx->nvencoder, ctx->output_surfaces[i].output_surface);
 848     }
 849     ctx->max_surface_count = 0;
 850
 851     p_nvenc->nvEncDestroyEncoder(ctx->nvencoder);
 852     ctx->nvencoder = NULL;
 853
 854     dl_fn->cu_ctx_destroy(ctx->cu_context);
 855     ctx->cu_context = NULL;
 856
 857     nvenc_unload_nvenc(avctx);
 858
 859     av_frame_free(&avctx->coded_frame);
 860
 861     return 0;
 862 }
 863
 864 static int process_output_surface(AVCodecContext *avctx, AVPacket *pkt, AVFrame *coded_frame, NvencOutputSurface *tmpoutsurf)
 865 {
 866     NvencContext *ctx = avctx->priv_data;
 867     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 868     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
 869
 870     uint32_t *slice_offsets = av_mallocz(ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData * sizeof(*slice_offsets));
 871     NV_ENC_LOCK_BITSTREAM lock_params = { 0 };
 872     NVENCSTATUS nv_status;
 873     int res = 0;
 874
 875     if (!slice_offsets)
 876         return AVERROR(ENOMEM);
 877
 878     lock_params.version = NV_ENC_LOCK_BITSTREAM_VER;
 879
 880     lock_params.doNotWait = 0;
 881     lock_params.outputBitstream = tmpoutsurf->output_surface;
 882     lock_params.sliceOffsets = slice_offsets;
 883
 884     nv_status = p_nvenc->nvEncLockBitstream(ctx->nvencoder, &lock_params);
 885     if (nv_status != NV_ENC_SUCCESS) {
 886         av_log(avctx, AV_LOG_ERROR, "Failed locking bitstream buffer\n");
 887         res = AVERROR_EXTERNAL;
 888         goto error;
 889     }
 890
 891     if (res = ff_alloc_packet2(avctx, pkt, lock_params.bitstreamSizeInBytes)) {
 892         p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
 893         goto error;
 894     }
 895
 896     memcpy(pkt->data, lock_params.bitstreamBufferPtr, lock_params.bitstreamSizeInBytes);
 897
 898     nv_status = p_nvenc->nvEncUnlockBitstream(ctx->nvencoder, tmpoutsurf->output_surface);
 899     if (nv_status != NV_ENC_SUCCESS)
 900         av_log(avctx, AV_LOG_ERROR, "Failed unlocking bitstream buffer, expect the gates of mordor to open\n");
 901
 902     switch (lock_params.pictureType) {
 903     case NV_ENC_PIC_TYPE_IDR:
 904         pkt->flags |= AV_PKT_FLAG_KEY;
 905     case NV_ENC_PIC_TYPE_I:
 906         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
 907         break;
 908     case NV_ENC_PIC_TYPE_P:
 909         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_P;
 910         break;
 911     case NV_ENC_PIC_TYPE_B:
 912         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_B;
 913         break;
 914     case NV_ENC_PIC_TYPE_BI:
 915         avctx->coded_frame->pict_type = AV_PICTURE_TYPE_BI;
 916         break;
 917     default:
 918         av_log(avctx, AV_LOG_ERROR, "Unknown picture type encountered, expect the output to be broken.\n");
 919         av_log(avctx, AV_LOG_ERROR, "Please report this error and include as much information on how to reproduce it as possible.\n");
 920         res = AVERROR_EXTERNAL;
 921         goto error;
 922     }
 923
 924     pkt->pts = lock_params.outputTimeStamp;
 925     pkt->dts = timestamp_queue_dequeue(&ctx->timestamp_list);
 926
 927     if (pkt->dts > pkt->pts)
 928         pkt->dts = pkt->pts;
 929
 930     if (ctx->last_dts != AV_NOPTS_VALUE && pkt->dts <= ctx->last_dts)
 931         pkt->dts = ctx->last_dts + 1;
 932
 933     ctx->last_dts = pkt->dts;
 934
 935     av_free(slice_offsets);
 936
 937     return 0;
 938
 939 error:
 940
 941     av_free(slice_offsets);
 942     timestamp_queue_dequeue(&ctx->timestamp_list);
 943
 944     return res;
 945 }
 946
 947 static int nvenc_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
 948     const AVFrame *frame, int *got_packet)
 949 {
 950     NVENCSTATUS nv_status;
 951     NvencOutputSurface *tmpoutsurf;
 952     int res, i = 0;
 953
 954     NvencContext *ctx = avctx->priv_data;
 955     NvencDynLoadFunctions *dl_fn = &ctx->nvenc_dload_funcs;
 956     NV_ENCODE_API_FUNCTION_LIST *p_nvenc = &dl_fn->nvenc_funcs;
 957
 958     NV_ENC_PIC_PARAMS pic_params = { 0 };
 959     pic_params.version = NV_ENC_PIC_PARAMS_VER;
 960
 961     if (frame) {
 962         NV_ENC_LOCK_INPUT_BUFFER lockBufferParams = { 0 };
 963         NvencInputSurface *inSurf = NULL;
 964
 965         for (i = 0; i < ctx->max_surface_count; ++i) {
 966             if (!ctx->input_surfaces[i].lockCount) {
 967                 inSurf = &ctx->input_surfaces[i];
 968                 break;
 969             }
 970         }
 971
 972         av_assert0(inSurf);
 973
 974         inSurf->lockCount = 1;
 975
 976         lockBufferParams.version = NV_ENC_LOCK_INPUT_BUFFER_VER;
 977         lockBufferParams.inputBuffer = inSurf->input_surface;
 978
 979         nv_status = p_nvenc->nvEncLockInputBuffer(ctx->nvencoder, &lockBufferParams);
 980         if (nv_status != NV_ENC_SUCCESS) {
 981             av_log(avctx, AV_LOG_ERROR, "Failed locking nvenc input buffer\n");
 982             return 0;
 983         }
 984
 985         if (avctx->pix_fmt == AV_PIX_FMT_YUV420P) {
 986             uint8_t *buf = lockBufferParams.bufferDataPtr;
 987
 988             av_image_copy_plane(buf, lockBufferParams.pitch,
 989                 frame->data[0], frame->linesize[0],
 990                 avctx->width, avctx->height);
 991
 992             buf += inSurf->height * lockBufferParams.pitch;
 993
 994             av_image_copy_plane(buf, lockBufferParams.pitch >> 1,
 995                 frame->data[2], frame->linesize[2],
 996                 avctx->width >> 1, avctx->height >> 1);
 997
 998             buf += (inSurf->height * lockBufferParams.pitch) >> 2;
 999
1000             av_image_copy_plane(buf, lockBufferParams.pitch >> 1,
1001                 frame->data[1], frame->linesize[1],
1002                 avctx->width >> 1, avctx->height >> 1);
1003         } else if (avctx->pix_fmt == AV_PIX_FMT_NV12) {
1004             uint8_t *buf = lockBufferParams.bufferDataPtr;
1005
1006             av_image_copy_plane(buf, lockBufferParams.pitch,
1007                 frame->data[0], frame->linesize[0],
1008                 avctx->width, avctx->height);
1009
1010             buf += inSurf->height * lockBufferParams.pitch;
1011
1012             av_image_copy_plane(buf, lockBufferParams.pitch,
1013                 frame->data[1], frame->linesize[1],
1014                 avctx->width, avctx->height >> 1);
1015         } else if (avctx->pix_fmt == AV_PIX_FMT_YUV444P) {
1016             uint8_t *buf = lockBufferParams.bufferDataPtr;
1017
1018             av_image_copy_plane(buf, lockBufferParams.pitch,
1019                 frame->data[0], frame->linesize[0],
1020                 avctx->width, avctx->height);
1021
1022             buf += inSurf->height * lockBufferParams.pitch;
1023
1024             av_image_copy_plane(buf, lockBufferParams.pitch,
1025                 frame->data[1], frame->linesize[1],
1026                 avctx->width, avctx->height);
1027
1028             buf += inSurf->height * lockBufferParams.pitch;
1029
1030             av_image_copy_plane(buf, lockBufferParams.pitch,
1031                 frame->data[2], frame->linesize[2],
1032                 avctx->width, avctx->height);
1033         } else {
1034             av_log(avctx, AV_LOG_FATAL, "Invalid pixel format!\n");
1035             return AVERROR(EINVAL);
1036         }
1037
1038         nv_status = p_nvenc->nvEncUnlockInputBuffer(ctx->nvencoder, inSurf->input_surface);
1039         if (nv_status != NV_ENC_SUCCESS) {
1040             av_log(avctx, AV_LOG_FATAL, "Failed unlocking input buffer!\n");
1041             return AVERROR_EXTERNAL;
1042         }
1043
1044         for (i = 0; i < ctx->max_surface_count; ++i)
1045             if (!ctx->output_surfaces[i].busy)
1046                 break;
1047
1048         if (i == ctx->max_surface_count) {
1049             inSurf->lockCount = 0;
1050             av_log(avctx, AV_LOG_FATAL, "No free output surface found!\n");
1051             return AVERROR_EXTERNAL;
1052         }
1053
1054         ctx->output_surfaces[i].input_surface = inSurf;
1055
1056         pic_params.inputBuffer = inSurf->input_surface;
1057         pic_params.bufferFmt = inSurf->format;
1058         pic_params.inputWidth = avctx->width;
1059         pic_params.inputHeight = avctx->height;
1060         pic_params.outputBitstream = ctx->output_surfaces[i].output_surface;
1061         pic_params.completionEvent = 0;
1062
1063         if (avctx->flags & CODEC_FLAG_INTERLACED_DCT) {
1064             if (frame->top_field_first) {
1065                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
1066             } else {
1067                 pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
1068             }
1069         } else {
1070             pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
1071         }
1072
1073         pic_params.encodePicFlags = 0;
1074         pic_params.inputTimeStamp = frame->pts;
1075         pic_params.inputDuration = 0;
1076         pic_params.codecPicParams.h264PicParams.sliceMode = ctx->encode_config.encodeCodecConfig.h264Config.sliceMode;
1077         pic_params.codecPicParams.h264PicParams.sliceModeData = ctx->encode_config.encodeCodecConfig.h264Config.sliceModeData;
1078         memcpy(&pic_params.rcParams, &ctx->encode_config.rcParams, sizeof(NV_ENC_RC_PARAMS));
1079
1080         res = timestamp_queue_enqueue(&ctx->timestamp_list, frame->pts);
1081
1082         if (res)
1083             return res;
1084     } else {
1085         pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
1086     }
1087
1088     nv_status = p_nvenc->nvEncEncodePicture(ctx->nvencoder, &pic_params);
1089
1090     if (frame && nv_status == NV_ENC_ERR_NEED_MORE_INPUT) {
1091         res = out_surf_queue_enqueue(&ctx->output_surface_queue, &ctx->output_surfaces[i]);
1092
1093         if (res)
1094             return res;
1095
1096         ctx->output_surfaces[i].busy = 1;
1097     }
1098
1099     if (nv_status != NV_ENC_SUCCESS && nv_status != NV_ENC_ERR_NEED_MORE_INPUT) {
1100         av_log(avctx, AV_LOG_ERROR, "EncodePicture failed!\n");
1101         return AVERROR_EXTERNAL;
1102     }
1103
1104     if (nv_status != NV_ENC_ERR_NEED_MORE_INPUT) {
1105         while (ctx->output_surface_queue.count) {
1106             tmpoutsurf = out_surf_queue_dequeue(&ctx->output_surface_queue);
1107             res = out_surf_queue_enqueue(&ctx->output_surface_ready_queue, tmpoutsurf);
1108
1109             if (res)
1110                 return res;
1111         }
1112
1113         if (frame) {
1114             res = out_surf_queue_enqueue(&ctx->output_surface_ready_queue, &ctx->output_surfaces[i]);
1115
1116             if (res)
1117                 return res;
1118
1119             ctx->output_surfaces[i].busy = 1;
1120         }
1121     }
1122
1123     if (ctx->output_surface_ready_queue.count) {
1124         tmpoutsurf = out_surf_queue_dequeue(&ctx->output_surface_ready_queue);
1125
1126         res = process_output_surface(avctx, pkt, avctx->coded_frame, tmpoutsurf);
1127
1128         if (res)
1129             return res;
1130
1131         tmpoutsurf->busy = 0;
1132         av_assert0(tmpoutsurf->input_surface->lockCount);
1133         tmpoutsurf->input_surface->lockCount--;
1134
1135         *got_packet = 1;
1136     } else {
1137         *got_packet = 0;
1138     }
1139
1140     return 0;
1141 }
1142
1143 static enum AVPixelFormat pix_fmts_nvenc[] = {
1144     AV_PIX_FMT_NV12,
1145     AV_PIX_FMT_NONE
1146 };
1147
1148 #define OFFSET(x) offsetof(NvencContext, x)
1149 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
1150 static const AVOption options[] = {
1151     { "preset", "Set the encoding preset (one of hq, hp, bd, ll, llhq, llhp, default)", OFFSET(preset), AV_OPT_TYPE_STRING, { .str = "hq" }, 0, 0, VE },
1152     { "cbr", "Use cbr encoding mode", OFFSET(cbr), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, VE },
1153     { "2pass", "Use 2pass cbr encoding mode (low latency mode only)", OFFSET(twopass), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 1, VE },
1154     { "goppattern", "Specifies the GOP pattern as follows: 0: I, 1: IPP, 2: IBP, 3: IBBP", OFFSET(gobpattern), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 3, VE },
1155     { "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 },
1156     { NULL }
1157 };
1158
1159 static const AVClass nvenc_class = {
1160     .class_name = "nvenc",
1161     .item_name = av_default_item_name,
1162     .option = options,
1163     .version = LIBAVUTIL_VERSION_INT,
1164 };
1165
1166 static const AVCodecDefault nvenc_defaults[] = {
1167     { "b", "0" },
1168     { "qmin", "-1" },
1169     { "qmax", "-1" },
1170     { "qdiff", "-1" },
1171     { "qblur", "-1" },
1172     { "qcomp", "-1" },
1173     { NULL },
1174 };
1175
1176 AVCodec ff_nvenc_encoder = {
1177     .name = "nvenc",
1178     .long_name = NULL_IF_CONFIG_SMALL("Nvidia NVENC h264 encoder"),
1179     .type = AVMEDIA_TYPE_VIDEO,
1180     .id = AV_CODEC_ID_H264,
1181     .priv_data_size = sizeof(NvencContext),
1182     .init = nvenc_encode_init,
1183     .encode2 = nvenc_encode_frame,
1184     .close = nvenc_encode_close,
1185     .capabilities = CODEC_CAP_DELAY,
1186     .priv_class = &nvenc_class,
1187     .defaults = nvenc_defaults,
1188     .pix_fmts = pix_fmts_nvenc,
1189 };