git.sesse.net Git - ffmpeg/blob - libavcodec/cuvid.c

   1 /*
   2  * Nvidia CUVID decoder
   3  * Copyright (c) 2016 Timo Rothenpieler <timo@rothenpieler.org>
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 #include "compat/cuda/dynlink_loader.h"
  23
  24 #include "libavutil/buffer.h"
  25 #include "libavutil/mathematics.h"
  26 #include "libavutil/hwcontext.h"
  27 #include "libavutil/hwcontext_cuda_internal.h"
  28 #include "libavutil/fifo.h"
  29 #include "libavutil/log.h"
  30 #include "libavutil/opt.h"
  31
  32 #include "avcodec.h"
  33 #include "internal.h"
  34
  35 typedef struct CuvidContext
  36 {
  37     AVClass *avclass;
  38
  39     CUvideodecoder cudecoder;
  40     CUvideoparser cuparser;
  41
  42     char *cu_gpu;
  43     int nb_surfaces;
  44
  45     AVBufferRef *hwdevice;
  46     AVBufferRef *hwframe;
  47
  48     AVBSFContext *bsf;
  49
  50     AVFifoBuffer *frame_queue;
  51
  52     int deint_mode;
  53     int64_t prev_pts;
  54
  55     int internal_error;
  56     int decoder_flushing;
  57
  58     cudaVideoCodec codec_type;
  59     cudaVideoChromaFormat chroma_format;
  60
  61     CUVIDPARSERPARAMS cuparseinfo;
  62     CUVIDEOFORMATEX cuparse_ext;
  63
  64     CudaFunctions *cudl;
  65     CuvidFunctions *cvdl;
  66 } CuvidContext;
  67
  68 typedef struct CuvidParsedFrame
  69 {
  70     CUVIDPARSERDISPINFO dispinfo;
  71     int second_field;
  72     int is_deinterlacing;
  73 } CuvidParsedFrame;
  74
  75 static int check_cu(AVCodecContext *avctx, CUresult err, const char *func)
  76 {
  77     CuvidContext *ctx = avctx->priv_data;
  78     const char *err_name;
  79     const char *err_string;
  80
  81     av_log(avctx, AV_LOG_TRACE, "Calling %s\n", func);
  82
  83     if (err == CUDA_SUCCESS)
  84         return 0;
  85
  86     ctx->cudl->cuGetErrorName(err, &err_name);
  87     ctx->cudl->cuGetErrorString(err, &err_string);
  88
  89     av_log(avctx, AV_LOG_ERROR, "%s failed", func);
  90     if (err_name && err_string)
  91         av_log(avctx, AV_LOG_ERROR, " -> %s: %s", err_name, err_string);
  92     av_log(avctx, AV_LOG_ERROR, "\n");
  93
  94     return AVERROR_EXTERNAL;
  95 }
  96
  97 #define CHECK_CU(x) check_cu(avctx, (x), #x)
  98
  99 static int CUDAAPI cuvid_handle_video_sequence(void *opaque, CUVIDEOFORMAT* format)
 100 {
 101     AVCodecContext *avctx = opaque;
 102     CuvidContext *ctx = avctx->priv_data;
 103     AVHWFramesContext *hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
 104     CUVIDDECODECREATEINFO cuinfo;
 105
 106     av_log(avctx, AV_LOG_TRACE, "pfnSequenceCallback, progressive_sequence=%d\n", format->progressive_sequence);
 107
 108     ctx->internal_error = 0;
 109
 110     avctx->width = format->display_area.right;
 111     avctx->height = format->display_area.bottom;
 112
 113     ff_set_sar(avctx, av_div_q(
 114         (AVRational){ format->display_aspect_ratio.x, format->display_aspect_ratio.y },
 115         (AVRational){ avctx->width, avctx->height }));
 116
 117     if (!format->progressive_sequence && ctx->deint_mode == cudaVideoDeinterlaceMode_Weave)
 118         avctx->flags |= AV_CODEC_FLAG_INTERLACED_DCT;
 119     else
 120         avctx->flags &= ~AV_CODEC_FLAG_INTERLACED_DCT;
 121
 122     if (format->video_signal_description.video_full_range_flag)
 123         avctx->color_range = AVCOL_RANGE_JPEG;
 124     else
 125         avctx->color_range = AVCOL_RANGE_MPEG;
 126
 127     avctx->color_primaries = format->video_signal_description.color_primaries;
 128     avctx->color_trc = format->video_signal_description.transfer_characteristics;
 129     avctx->colorspace = format->video_signal_description.matrix_coefficients;
 130
 131     if (format->bitrate)
 132         avctx->bit_rate = format->bitrate;
 133
 134     if (format->frame_rate.numerator && format->frame_rate.denominator) {
 135         avctx->framerate.num = format->frame_rate.numerator;
 136         avctx->framerate.den = format->frame_rate.denominator;
 137     }
 138
 139     if (ctx->cudecoder
 140             && avctx->coded_width == format->coded_width
 141             && avctx->coded_height == format->coded_height
 142             && ctx->chroma_format == format->chroma_format
 143             && ctx->codec_type == format->codec)
 144         return 1;
 145
 146     if (ctx->cudecoder) {
 147         av_log(avctx, AV_LOG_TRACE, "Re-initializing decoder\n");
 148         ctx->internal_error = CHECK_CU(ctx->cvdl->cuvidDestroyDecoder(ctx->cudecoder));
 149         if (ctx->internal_error < 0)
 150             return 0;
 151         ctx->cudecoder = NULL;
 152     }
 153
 154     if (hwframe_ctx->pool && (
 155             hwframe_ctx->width < avctx->width ||
 156             hwframe_ctx->height < avctx->height ||
 157             hwframe_ctx->format != AV_PIX_FMT_CUDA ||
 158             hwframe_ctx->sw_format != AV_PIX_FMT_NV12)) {
 159         av_log(avctx, AV_LOG_ERROR, "AVHWFramesContext is already initialized with incompatible parameters\n");
 160         ctx->internal_error = AVERROR(EINVAL);
 161         return 0;
 162     }
 163
 164     if (format->chroma_format != cudaVideoChromaFormat_420) {
 165         av_log(avctx, AV_LOG_ERROR, "Chroma formats other than 420 are not supported\n");
 166         ctx->internal_error = AVERROR(EINVAL);
 167         return 0;
 168     }
 169
 170     avctx->coded_width = format->coded_width;
 171     avctx->coded_height = format->coded_height;
 172
 173     ctx->chroma_format = format->chroma_format;
 174
 175     memset(&cuinfo, 0, sizeof(cuinfo));
 176
 177     cuinfo.CodecType = ctx->codec_type = format->codec;
 178     cuinfo.ChromaFormat = format->chroma_format;
 179     cuinfo.OutputFormat = cudaVideoSurfaceFormat_NV12;
 180
 181     cuinfo.ulWidth = avctx->coded_width;
 182     cuinfo.ulHeight = avctx->coded_height;
 183     cuinfo.ulTargetWidth = cuinfo.ulWidth;
 184     cuinfo.ulTargetHeight = cuinfo.ulHeight;
 185
 186     cuinfo.target_rect.left = 0;
 187     cuinfo.target_rect.top = 0;
 188     cuinfo.target_rect.right = cuinfo.ulWidth;
 189     cuinfo.target_rect.bottom = cuinfo.ulHeight;
 190
 191     cuinfo.ulNumDecodeSurfaces = ctx->nb_surfaces;
 192     cuinfo.ulNumOutputSurfaces = 1;
 193     cuinfo.ulCreationFlags = cudaVideoCreate_PreferCUVID;
 194     cuinfo.bitDepthMinus8 = format->bit_depth_luma_minus8;
 195
 196     if (format->progressive_sequence) {
 197         ctx->deint_mode = cuinfo.DeinterlaceMode = cudaVideoDeinterlaceMode_Weave;
 198     } else {
 199         cuinfo.DeinterlaceMode = ctx->deint_mode;
 200     }
 201
 202     if (ctx->deint_mode != cudaVideoDeinterlaceMode_Weave)
 203         avctx->framerate = av_mul_q(avctx->framerate, (AVRational){2, 1});
 204
 205     ctx->internal_error = CHECK_CU(ctx->cvdl->cuvidCreateDecoder(&ctx->cudecoder, &cuinfo));
 206     if (ctx->internal_error < 0)
 207         return 0;
 208
 209     if (!hwframe_ctx->pool) {
 210         hwframe_ctx->format = AV_PIX_FMT_CUDA;
 211         hwframe_ctx->sw_format = AV_PIX_FMT_NV12;
 212         hwframe_ctx->width = avctx->width;
 213         hwframe_ctx->height = avctx->height;
 214
 215         if ((ctx->internal_error = av_hwframe_ctx_init(ctx->hwframe)) < 0) {
 216             av_log(avctx, AV_LOG_ERROR, "av_hwframe_ctx_init failed\n");
 217             return 0;
 218         }
 219     }
 220
 221     return 1;
 222 }
 223
 224 static int CUDAAPI cuvid_handle_picture_decode(void *opaque, CUVIDPICPARAMS* picparams)
 225 {
 226     AVCodecContext *avctx = opaque;
 227     CuvidContext *ctx = avctx->priv_data;
 228
 229     av_log(avctx, AV_LOG_TRACE, "pfnDecodePicture\n");
 230
 231     ctx->internal_error = CHECK_CU(ctx->cvdl->cuvidDecodePicture(ctx->cudecoder, picparams));
 232     if (ctx->internal_error < 0)
 233         return 0;
 234
 235     return 1;
 236 }
 237
 238 static int CUDAAPI cuvid_handle_picture_display(void *opaque, CUVIDPARSERDISPINFO* dispinfo)
 239 {
 240     AVCodecContext *avctx = opaque;
 241     CuvidContext *ctx = avctx->priv_data;
 242     CuvidParsedFrame parsed_frame = { *dispinfo, 0, 0 };
 243
 244     ctx->internal_error = 0;
 245
 246     if (ctx->deint_mode == cudaVideoDeinterlaceMode_Weave) {
 247         av_fifo_generic_write(ctx->frame_queue, &parsed_frame, sizeof(CuvidParsedFrame), NULL);
 248     } else {
 249         parsed_frame.is_deinterlacing = 1;
 250         av_fifo_generic_write(ctx->frame_queue, &parsed_frame, sizeof(CuvidParsedFrame), NULL);
 251         parsed_frame.second_field = 1;
 252         av_fifo_generic_write(ctx->frame_queue, &parsed_frame, sizeof(CuvidParsedFrame), NULL);
 253     }
 254
 255     return 1;
 256 }
 257
 258 static int cuvid_decode_packet(AVCodecContext *avctx, const AVPacket *avpkt)
 259 {
 260     CuvidContext *ctx = avctx->priv_data;
 261     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
 262     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
 263     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
 264     CUVIDSOURCEDATAPACKET cupkt;
 265     AVPacket filter_packet = { 0 };
 266     AVPacket filtered_packet = { 0 };
 267     int ret = 0, eret = 0, is_flush = ctx->decoder_flushing;
 268
 269     av_log(avctx, AV_LOG_TRACE, "cuvid_decode_packet\n");
 270
 271     if (is_flush && avpkt && avpkt->size)
 272         return AVERROR_EOF;
 273
 274     if ((av_fifo_size(ctx->frame_queue) / sizeof(CuvidParsedFrame)) + 2 > ctx->nb_surfaces && avpkt && avpkt->size)
 275         return AVERROR(EAGAIN);
 276
 277     if (ctx->bsf && avpkt && avpkt->size) {
 278         if ((ret = av_packet_ref(&filter_packet, avpkt)) < 0) {
 279             av_log(avctx, AV_LOG_ERROR, "av_packet_ref failed\n");
 280             return ret;
 281         }
 282
 283         if ((ret = av_bsf_send_packet(ctx->bsf, &filter_packet)) < 0) {
 284             av_log(avctx, AV_LOG_ERROR, "av_bsf_send_packet failed\n");
 285             av_packet_unref(&filter_packet);
 286             return ret;
 287         }
 288
 289         if ((ret = av_bsf_receive_packet(ctx->bsf, &filtered_packet)) < 0) {
 290             av_log(avctx, AV_LOG_ERROR, "av_bsf_receive_packet failed\n");
 291             return ret;
 292         }
 293
 294         avpkt = &filtered_packet;
 295     }
 296
 297     ret = CHECK_CU(ctx->cudl->cuCtxPushCurrent(cuda_ctx));
 298     if (ret < 0) {
 299         av_packet_unref(&filtered_packet);
 300         return ret;
 301     }
 302
 303     memset(&cupkt, 0, sizeof(cupkt));
 304
 305     if (avpkt && avpkt->size) {
 306         cupkt.payload_size = avpkt->size;
 307         cupkt.payload = avpkt->data;
 308
 309         if (avpkt->pts != AV_NOPTS_VALUE) {
 310             cupkt.flags = CUVID_PKT_TIMESTAMP;
 311             if (avctx->pkt_timebase.num && avctx->pkt_timebase.den)
 312                 cupkt.timestamp = av_rescale_q(avpkt->pts, avctx->pkt_timebase, (AVRational){1, 10000000});
 313             else
 314                 cupkt.timestamp = avpkt->pts;
 315         }
 316     } else {
 317         cupkt.flags = CUVID_PKT_ENDOFSTREAM;
 318         ctx->decoder_flushing = 1;
 319     }
 320
 321     ret = CHECK_CU(ctx->cvdl->cuvidParseVideoData(ctx->cuparser, &cupkt));
 322
 323     av_packet_unref(&filtered_packet);
 324
 325     if (ret < 0)
 326         goto error;
 327
 328     // cuvidParseVideoData doesn't return an error just because stuff failed...
 329     if (ctx->internal_error) {
 330         av_log(avctx, AV_LOG_ERROR, "cuvid decode callback error\n");
 331         ret = ctx->internal_error;
 332         goto error;
 333     }
 334
 335 error:
 336     eret = CHECK_CU(ctx->cudl->cuCtxPopCurrent(&dummy));
 337
 338     if (eret < 0)
 339         return eret;
 340     else if (ret < 0)
 341         return ret;
 342     else if (is_flush)
 343         return AVERROR_EOF;
 344     else
 345         return 0;
 346 }
 347
 348 static int cuvid_output_frame(AVCodecContext *avctx, AVFrame *frame)
 349 {
 350     CuvidContext *ctx = avctx->priv_data;
 351     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
 352     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
 353     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
 354     CUdeviceptr mapped_frame = 0;
 355     int ret = 0, eret = 0;
 356
 357     av_log(avctx, AV_LOG_TRACE, "cuvid_output_frame\n");
 358
 359     if (ctx->decoder_flushing) {
 360         ret = cuvid_decode_packet(avctx, NULL);
 361         if (ret < 0 && ret != AVERROR_EOF)
 362             return ret;
 363     }
 364
 365     ret = CHECK_CU(ctx->cudl->cuCtxPushCurrent(cuda_ctx));
 366     if (ret < 0)
 367         return ret;
 368
 369     if (av_fifo_size(ctx->frame_queue)) {
 370         CuvidParsedFrame parsed_frame;
 371         CUVIDPROCPARAMS params;
 372         unsigned int pitch = 0;
 373         int offset = 0;
 374         int i;
 375
 376         av_fifo_generic_read(ctx->frame_queue, &parsed_frame, sizeof(CuvidParsedFrame), NULL);
 377
 378         memset(&params, 0, sizeof(params));
 379         params.progressive_frame = parsed_frame.dispinfo.progressive_frame;
 380         params.second_field = parsed_frame.second_field;
 381         params.top_field_first = parsed_frame.dispinfo.top_field_first;
 382
 383         ret = CHECK_CU(ctx->cvdl->cuvidMapVideoFrame(ctx->cudecoder, parsed_frame.dispinfo.picture_index, &mapped_frame, &pitch, &params));
 384         if (ret < 0)
 385             goto error;
 386
 387         if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
 388             ret = av_hwframe_get_buffer(ctx->hwframe, frame, 0);
 389             if (ret < 0) {
 390                 av_log(avctx, AV_LOG_ERROR, "av_hwframe_get_buffer failed\n");
 391                 goto error;
 392             }
 393
 394             ret = ff_decode_frame_props(avctx, frame);
 395             if (ret < 0) {
 396                 av_log(avctx, AV_LOG_ERROR, "ff_decode_frame_props failed\n");
 397                 goto error;
 398             }
 399
 400             for (i = 0; i < 2; i++) {
 401                 CUDA_MEMCPY2D cpy = {
 402                     .srcMemoryType = CU_MEMORYTYPE_DEVICE,
 403                     .dstMemoryType = CU_MEMORYTYPE_DEVICE,
 404                     .srcDevice     = mapped_frame,
 405                     .dstDevice     = (CUdeviceptr)frame->data[i],
 406                     .srcPitch      = pitch,
 407                     .dstPitch      = frame->linesize[i],
 408                     .srcY          = offset,
 409                     .WidthInBytes  = FFMIN(pitch, frame->linesize[i]),
 410                     .Height        = avctx->height >> (i ? 1 : 0),
 411                 };
 412
 413                 ret = CHECK_CU(ctx->cudl->cuMemcpy2D(&cpy));
 414                 if (ret < 0)
 415                     goto error;
 416
 417                 offset += avctx->coded_height;
 418             }
 419         } else if (avctx->pix_fmt == AV_PIX_FMT_NV12) {
 420             AVFrame *tmp_frame = av_frame_alloc();
 421             if (!tmp_frame) {
 422                 av_log(avctx, AV_LOG_ERROR, "av_frame_alloc failed\n");
 423                 ret = AVERROR(ENOMEM);
 424                 goto error;
 425             }
 426
 427             tmp_frame->format        = AV_PIX_FMT_CUDA;
 428             tmp_frame->hw_frames_ctx = av_buffer_ref(ctx->hwframe);
 429             tmp_frame->data[0]       = (uint8_t*)mapped_frame;
 430             tmp_frame->linesize[0]   = pitch;
 431             tmp_frame->data[1]       = (uint8_t*)(mapped_frame + avctx->coded_height * pitch);
 432             tmp_frame->linesize[1]   = pitch;
 433             tmp_frame->width         = avctx->width;
 434             tmp_frame->height        = avctx->height;
 435
 436             ret = ff_get_buffer(avctx, frame, 0);
 437             if (ret < 0) {
 438                 av_log(avctx, AV_LOG_ERROR, "ff_get_buffer failed\n");
 439                 av_frame_free(&tmp_frame);
 440                 goto error;
 441             }
 442
 443             ret = av_hwframe_transfer_data(frame, tmp_frame, 0);
 444             if (ret) {
 445                 av_log(avctx, AV_LOG_ERROR, "av_hwframe_transfer_data failed\n");
 446                 av_frame_free(&tmp_frame);
 447                 goto error;
 448             }
 449
 450             av_frame_free(&tmp_frame);
 451         } else {
 452             ret = AVERROR_BUG;
 453             goto error;
 454         }
 455
 456         frame->width = avctx->width;
 457         frame->height = avctx->height;
 458         if (avctx->pkt_timebase.num && avctx->pkt_timebase.den)
 459             frame->pts = av_rescale_q(parsed_frame.dispinfo.timestamp, (AVRational){1, 10000000}, avctx->pkt_timebase);
 460         else
 461             frame->pts = parsed_frame.dispinfo.timestamp;
 462
 463         if (parsed_frame.second_field) {
 464             if (ctx->prev_pts == INT64_MIN) {
 465                 ctx->prev_pts = frame->pts;
 466                 frame->pts += (avctx->pkt_timebase.den * avctx->framerate.den) / (avctx->pkt_timebase.num * avctx->framerate.num);
 467             } else {
 468                 int pts_diff = (frame->pts - ctx->prev_pts) / 2;
 469                 ctx->prev_pts = frame->pts;
 470                 frame->pts += pts_diff;
 471             }
 472         }
 473
 474         /* CUVIDs opaque reordering breaks the internal pkt logic.
 475          * So set pkt_pts and clear all the other pkt_ fields.
 476          */
 477 #if FF_API_PKT_PTS
 478 FF_DISABLE_DEPRECATION_WARNINGS
 479         frame->pkt_pts = frame->pts;
 480 FF_ENABLE_DEPRECATION_WARNINGS
 481 #endif
 482         av_frame_set_pkt_pos(frame, -1);
 483         av_frame_set_pkt_duration(frame, 0);
 484         av_frame_set_pkt_size(frame, -1);
 485
 486         frame->interlaced_frame = !parsed_frame.is_deinterlacing && !parsed_frame.dispinfo.progressive_frame;
 487
 488         if (frame->interlaced_frame)
 489             frame->top_field_first = parsed_frame.dispinfo.top_field_first;
 490     } else if (ctx->decoder_flushing) {
 491         ret = AVERROR_EOF;
 492     } else {
 493         ret = AVERROR(EAGAIN);
 494     }
 495
 496 error:
 497     if (mapped_frame)
 498         eret = CHECK_CU(ctx->cvdl->cuvidUnmapVideoFrame(ctx->cudecoder, mapped_frame));
 499
 500     eret = CHECK_CU(ctx->cudl->cuCtxPopCurrent(&dummy));
 501
 502     if (eret < 0)
 503         return eret;
 504     else
 505         return ret;
 506 }
 507
 508 static int cuvid_decode_frame(AVCodecContext *avctx, void *data, int *got_frame, AVPacket *avpkt)
 509 {
 510     CuvidContext *ctx = avctx->priv_data;
 511     AVFrame *frame = data;
 512     int ret = 0;
 513
 514     av_log(avctx, AV_LOG_TRACE, "cuvid_decode_frame\n");
 515
 516     if (ctx->deint_mode != cudaVideoDeinterlaceMode_Weave) {
 517         av_log(avctx, AV_LOG_ERROR, "Deinterlacing is not supported via the old API\n");
 518         return AVERROR(EINVAL);
 519     }
 520
 521     if (!ctx->decoder_flushing) {
 522         ret = cuvid_decode_packet(avctx, avpkt);
 523         if (ret < 0)
 524             return ret;
 525     }
 526
 527     ret = cuvid_output_frame(avctx, frame);
 528     if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
 529         *got_frame = 0;
 530     } else if (ret < 0) {
 531         return ret;
 532     } else {
 533         *got_frame = 1;
 534     }
 535
 536     return 0;
 537 }
 538
 539 static av_cold int cuvid_decode_end(AVCodecContext *avctx)
 540 {
 541     CuvidContext *ctx = avctx->priv_data;
 542
 543     av_fifo_freep(&ctx->frame_queue);
 544
 545     if (ctx->bsf)
 546         av_bsf_free(&ctx->bsf);
 547
 548     if (ctx->cuparser)
 549         ctx->cvdl->cuvidDestroyVideoParser(ctx->cuparser);
 550
 551     if (ctx->cudecoder)
 552         ctx->cvdl->cuvidDestroyDecoder(ctx->cudecoder);
 553
 554     ctx->cudl = NULL;
 555
 556     av_buffer_unref(&ctx->hwframe);
 557     av_buffer_unref(&ctx->hwdevice);
 558
 559     cuvid_free_functions(&ctx->cvdl);
 560
 561     return 0;
 562 }
 563
 564 static int cuvid_test_dummy_decoder(AVCodecContext *avctx, CUVIDPARSERPARAMS *cuparseinfo)
 565 {
 566     CuvidContext *ctx = avctx->priv_data;
 567     CUVIDDECODECREATEINFO cuinfo;
 568     CUvideodecoder cudec = 0;
 569     int ret = 0;
 570
 571     memset(&cuinfo, 0, sizeof(cuinfo));
 572
 573     cuinfo.CodecType = cuparseinfo->CodecType;
 574     cuinfo.ChromaFormat = cudaVideoChromaFormat_420;
 575     cuinfo.OutputFormat = cudaVideoSurfaceFormat_NV12;
 576
 577     cuinfo.ulWidth = 1280;
 578     cuinfo.ulHeight = 720;
 579     cuinfo.ulTargetWidth = cuinfo.ulWidth;
 580     cuinfo.ulTargetHeight = cuinfo.ulHeight;
 581
 582     cuinfo.target_rect.left = 0;
 583     cuinfo.target_rect.top = 0;
 584     cuinfo.target_rect.right = cuinfo.ulWidth;
 585     cuinfo.target_rect.bottom = cuinfo.ulHeight;
 586
 587     cuinfo.ulNumDecodeSurfaces = ctx->nb_surfaces;
 588     cuinfo.ulNumOutputSurfaces = 1;
 589     cuinfo.ulCreationFlags = cudaVideoCreate_PreferCUVID;
 590     cuinfo.bitDepthMinus8 = 0;
 591
 592     cuinfo.DeinterlaceMode = cudaVideoDeinterlaceMode_Weave;
 593
 594     ret = CHECK_CU(ctx->cvdl->cuvidCreateDecoder(&cudec, &cuinfo));
 595     if (ret < 0)
 596         return ret;
 597
 598     ret = CHECK_CU(ctx->cvdl->cuvidDestroyDecoder(cudec));
 599     if (ret < 0)
 600         return ret;
 601
 602     return 0;
 603 }
 604
 605 static av_cold int cuvid_decode_init(AVCodecContext *avctx)
 606 {
 607     CuvidContext *ctx = avctx->priv_data;
 608     AVCUDADeviceContext *device_hwctx;
 609     AVHWDeviceContext *device_ctx;
 610     AVHWFramesContext *hwframe_ctx;
 611     CUVIDSOURCEDATAPACKET seq_pkt;
 612     CUcontext cuda_ctx = NULL;
 613     CUcontext dummy;
 614     const AVBitStreamFilter *bsf;
 615     int ret = 0;
 616
 617     enum AVPixelFormat pix_fmts[3] = { AV_PIX_FMT_CUDA,
 618                                        AV_PIX_FMT_NV12,
 619                                        AV_PIX_FMT_NONE };
 620
 621     ret = ff_get_format(avctx, pix_fmts);
 622     if (ret < 0) {
 623         av_log(avctx, AV_LOG_ERROR, "ff_get_format failed: %d\n", ret);
 624         return ret;
 625     }
 626     avctx->pix_fmt = ret;
 627
 628     ret = cuvid_load_functions(&ctx->cvdl);
 629     if (ret < 0) {
 630         av_log(avctx, AV_LOG_ERROR, "Failed loading nvcuvid.\n");
 631         goto error;
 632     }
 633
 634     ctx->frame_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(CuvidParsedFrame));
 635     if (!ctx->frame_queue) {
 636         ret = AVERROR(ENOMEM);
 637         goto error;
 638     }
 639
 640     if (avctx->hw_frames_ctx) {
 641         ctx->hwframe = av_buffer_ref(avctx->hw_frames_ctx);
 642         if (!ctx->hwframe) {
 643             ret = AVERROR(ENOMEM);
 644             goto error;
 645         }
 646
 647         hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
 648
 649         ctx->hwdevice = av_buffer_ref(hwframe_ctx->device_ref);
 650         if (!ctx->hwdevice) {
 651             ret = AVERROR(ENOMEM);
 652             goto error;
 653         }
 654     } else {
 655         ret = av_hwdevice_ctx_create(&ctx->hwdevice, AV_HWDEVICE_TYPE_CUDA, ctx->cu_gpu, NULL, 0);
 656         if (ret < 0)
 657             goto error;
 658
 659         ctx->hwframe = av_hwframe_ctx_alloc(ctx->hwdevice);
 660         if (!ctx->hwframe) {
 661             av_log(avctx, AV_LOG_ERROR, "av_hwframe_ctx_alloc failed\n");
 662             ret = AVERROR(ENOMEM);
 663             goto error;
 664         }
 665
 666         hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
 667     }
 668
 669     device_ctx = hwframe_ctx->device_ctx;
 670     device_hwctx = device_ctx->hwctx;
 671
 672     cuda_ctx = device_hwctx->cuda_ctx;
 673     ctx->cudl = device_hwctx->internal->cuda_dl;
 674
 675     memset(&ctx->cuparseinfo, 0, sizeof(ctx->cuparseinfo));
 676     memset(&ctx->cuparse_ext, 0, sizeof(ctx->cuparse_ext));
 677     memset(&seq_pkt, 0, sizeof(seq_pkt));
 678
 679     ctx->cuparseinfo.pExtVideoInfo = &ctx->cuparse_ext;
 680
 681     switch (avctx->codec->id) {
 682 #if CONFIG_H264_CUVID_DECODER
 683     case AV_CODEC_ID_H264:
 684         ctx->cuparseinfo.CodecType = cudaVideoCodec_H264;
 685         break;
 686 #endif
 687 #if CONFIG_HEVC_CUVID_DECODER
 688     case AV_CODEC_ID_HEVC:
 689         ctx->cuparseinfo.CodecType = cudaVideoCodec_HEVC;
 690         break;
 691 #endif
 692 #if CONFIG_MJPEG_CUVID_DECODER
 693     case AV_CODEC_ID_MJPEG:
 694         ctx->cuparseinfo.CodecType = cudaVideoCodec_JPEG;
 695         break;
 696 #endif
 697 #if CONFIG_MPEG1_CUVID_DECODER
 698     case AV_CODEC_ID_MPEG1VIDEO:
 699         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG1;
 700         break;
 701 #endif
 702 #if CONFIG_MPEG2_CUVID_DECODER
 703     case AV_CODEC_ID_MPEG2VIDEO:
 704         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG2;
 705         break;
 706 #endif
 707 #if CONFIG_MPEG4_CUVID_DECODER
 708     case AV_CODEC_ID_MPEG4:
 709         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
 710         break;
 711 #endif
 712 #if CONFIG_VP8_CUVID_DECODER
 713     case AV_CODEC_ID_VP8:
 714         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP8;
 715         break;
 716 #endif
 717 #if CONFIG_VP9_CUVID_DECODER
 718     case AV_CODEC_ID_VP9:
 719         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP9;
 720         break;
 721 #endif
 722 #if CONFIG_VC1_CUVID_DECODER
 723     case AV_CODEC_ID_VC1:
 724         ctx->cuparseinfo.CodecType = cudaVideoCodec_VC1;
 725         break;
 726 #endif
 727     default:
 728         av_log(avctx, AV_LOG_ERROR, "Invalid CUVID codec!\n");
 729         return AVERROR_BUG;
 730     }
 731
 732     if (avctx->codec->id == AV_CODEC_ID_H264 || avctx->codec->id == AV_CODEC_ID_HEVC) {
 733         if (avctx->codec->id == AV_CODEC_ID_H264)
 734             bsf = av_bsf_get_by_name("h264_mp4toannexb");
 735         else
 736             bsf = av_bsf_get_by_name("hevc_mp4toannexb");
 737
 738         if (!bsf) {
 739             ret = AVERROR_BSF_NOT_FOUND;
 740             goto error;
 741         }
 742         if (ret = av_bsf_alloc(bsf, &ctx->bsf)) {
 743             goto error;
 744         }
 745         if (((ret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx)) < 0) || ((ret = av_bsf_init(ctx->bsf)) < 0)) {
 746             av_bsf_free(&ctx->bsf);
 747             goto error;
 748         }
 749
 750         ctx->cuparse_ext.format.seqhdr_data_length = ctx->bsf->par_out->extradata_size;
 751         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
 752                ctx->bsf->par_out->extradata,
 753                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), ctx->bsf->par_out->extradata_size));
 754     } else if (avctx->extradata_size > 0) {
 755         ctx->cuparse_ext.format.seqhdr_data_length = avctx->extradata_size;
 756         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
 757                avctx->extradata,
 758                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), avctx->extradata_size));
 759     }
 760
 761     ctx->cuparseinfo.ulMaxNumDecodeSurfaces = ctx->nb_surfaces;
 762     ctx->cuparseinfo.ulMaxDisplayDelay = 4;
 763     ctx->cuparseinfo.pUserData = avctx;
 764     ctx->cuparseinfo.pfnSequenceCallback = cuvid_handle_video_sequence;
 765     ctx->cuparseinfo.pfnDecodePicture = cuvid_handle_picture_decode;
 766     ctx->cuparseinfo.pfnDisplayPicture = cuvid_handle_picture_display;
 767
 768     ret = CHECK_CU(ctx->cudl->cuCtxPushCurrent(cuda_ctx));
 769     if (ret < 0)
 770         goto error;
 771
 772     ret = cuvid_test_dummy_decoder(avctx, &ctx->cuparseinfo);
 773     if (ret < 0)
 774         goto error;
 775
 776     ret = CHECK_CU(ctx->cvdl->cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
 777     if (ret < 0)
 778         goto error;
 779
 780     seq_pkt.payload = ctx->cuparse_ext.raw_seqhdr_data;
 781     seq_pkt.payload_size = ctx->cuparse_ext.format.seqhdr_data_length;
 782
 783     if (seq_pkt.payload && seq_pkt.payload_size) {
 784         ret = CHECK_CU(ctx->cvdl->cuvidParseVideoData(ctx->cuparser, &seq_pkt));
 785         if (ret < 0)
 786             goto error;
 787     }
 788
 789     ret = CHECK_CU(ctx->cudl->cuCtxPopCurrent(&dummy));
 790     if (ret < 0)
 791         goto error;
 792
 793     ctx->prev_pts = INT64_MIN;
 794
 795     if (!avctx->pkt_timebase.num || !avctx->pkt_timebase.den)
 796         av_log(avctx, AV_LOG_WARNING, "Invalid pkt_timebase, passing timestamps as-is.\n");
 797
 798     return 0;
 799
 800 error:
 801     cuvid_decode_end(avctx);
 802     return ret;
 803 }
 804
 805 static void cuvid_flush(AVCodecContext *avctx)
 806 {
 807     CuvidContext *ctx = avctx->priv_data;
 808     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
 809     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
 810     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
 811     CUVIDSOURCEDATAPACKET seq_pkt = { 0 };
 812     int ret;
 813
 814     ret = CHECK_CU(ctx->cudl->cuCtxPushCurrent(cuda_ctx));
 815     if (ret < 0)
 816         goto error;
 817
 818     av_fifo_freep(&ctx->frame_queue);
 819
 820     ctx->frame_queue = av_fifo_alloc(ctx->nb_surfaces * sizeof(CuvidParsedFrame));
 821     if (!ctx->frame_queue) {
 822         av_log(avctx, AV_LOG_ERROR, "Failed to recreate frame queue on flush\n");
 823         return;
 824     }
 825
 826     if (ctx->cudecoder) {
 827         ctx->cvdl->cuvidDestroyDecoder(ctx->cudecoder);
 828         ctx->cudecoder = NULL;
 829     }
 830
 831     if (ctx->cuparser) {
 832         ctx->cvdl->cuvidDestroyVideoParser(ctx->cuparser);
 833         ctx->cuparser = NULL;
 834     }
 835
 836     ret = CHECK_CU(ctx->cvdl->cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
 837     if (ret < 0)
 838         goto error;
 839
 840     seq_pkt.payload = ctx->cuparse_ext.raw_seqhdr_data;
 841     seq_pkt.payload_size = ctx->cuparse_ext.format.seqhdr_data_length;
 842
 843     if (seq_pkt.payload && seq_pkt.payload_size) {
 844         ret = CHECK_CU(ctx->cvdl->cuvidParseVideoData(ctx->cuparser, &seq_pkt));
 845         if (ret < 0)
 846             goto error;
 847     }
 848
 849     ret = CHECK_CU(ctx->cudl->cuCtxPopCurrent(&dummy));
 850     if (ret < 0)
 851         goto error;
 852
 853     ctx->prev_pts = INT64_MIN;
 854     ctx->decoder_flushing = 0;
 855
 856     return;
 857  error:
 858     av_log(avctx, AV_LOG_ERROR, "CUDA reinit on flush failed\n");
 859 }
 860
 861 #define OFFSET(x) offsetof(CuvidContext, x)
 862 #define VD AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_DECODING_PARAM
 863 static const AVOption options[] = {
 864     { "deint",    "Set deinterlacing mode", OFFSET(deint_mode), AV_OPT_TYPE_INT,   { .i64 = cudaVideoDeinterlaceMode_Weave    }, cudaVideoDeinterlaceMode_Weave, cudaVideoDeinterlaceMode_Adaptive, VD, "deint" },
 865     { "weave",    "Weave deinterlacing (do nothing)",        0, AV_OPT_TYPE_CONST, { .i64 = cudaVideoDeinterlaceMode_Weave    }, 0, 0, VD, "deint" },
 866     { "bob",      "Bob deinterlacing",                       0, AV_OPT_TYPE_CONST, { .i64 = cudaVideoDeinterlaceMode_Bob      }, 0, 0, VD, "deint" },
 867     { "adaptive", "Adaptive deinterlacing",                  0, AV_OPT_TYPE_CONST, { .i64 = cudaVideoDeinterlaceMode_Adaptive }, 0, 0, VD, "deint" },
 868     { "gpu",      "GPU to be used for decoding", OFFSET(cu_gpu), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, VD },
 869     { "surfaces", "Maximum surfaces to be used for decoding", OFFSET(nb_surfaces), AV_OPT_TYPE_INT, { .i64 = 25 }, 0, INT_MAX, VD },
 870     { NULL }
 871 };
 872
 873 #define DEFINE_CUVID_CODEC(x, X) \
 874     static const AVClass x##_cuvid_class = { \
 875         .class_name = #x "_cuvid", \
 876         .item_name = av_default_item_name, \
 877         .option = options, \
 878         .version = LIBAVUTIL_VERSION_INT, \
 879     }; \
 880     AVHWAccel ff_##x##_cuvid_hwaccel = { \
 881         .name           = #x "_cuvid", \
 882         .type           = AVMEDIA_TYPE_VIDEO, \
 883         .id             = AV_CODEC_ID_##X, \
 884         .pix_fmt        = AV_PIX_FMT_CUDA, \
 885     }; \
 886     AVCodec ff_##x##_cuvid_decoder = { \
 887         .name           = #x "_cuvid", \
 888         .long_name      = NULL_IF_CONFIG_SMALL("Nvidia CUVID " #X " decoder"), \
 889         .type           = AVMEDIA_TYPE_VIDEO, \
 890         .id             = AV_CODEC_ID_##X, \
 891         .priv_data_size = sizeof(CuvidContext), \
 892         .priv_class     = &x##_cuvid_class, \
 893         .init           = cuvid_decode_init, \
 894         .close          = cuvid_decode_end, \
 895         .decode         = cuvid_decode_frame, \
 896         .send_packet    = cuvid_decode_packet, \
 897         .receive_frame  = cuvid_output_frame, \
 898         .flush          = cuvid_flush, \
 899         .capabilities   = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AVOID_PROBING, \
 900         .pix_fmts       = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, \
 901                                                         AV_PIX_FMT_NV12, \
 902                                                         AV_PIX_FMT_NONE }, \
 903     };
 904
 905 #if CONFIG_HEVC_CUVID_DECODER
 906 DEFINE_CUVID_CODEC(hevc, HEVC)
 907 #endif
 908
 909 #if CONFIG_H264_CUVID_DECODER
 910 DEFINE_CUVID_CODEC(h264, H264)
 911 #endif
 912
 913 #if CONFIG_MJPEG_CUVID_DECODER
 914 DEFINE_CUVID_CODEC(mjpeg, MJPEG)
 915 #endif
 916
 917 #if CONFIG_MPEG1_CUVID_DECODER
 918 DEFINE_CUVID_CODEC(mpeg1, MPEG1VIDEO)
 919 #endif
 920
 921 #if CONFIG_MPEG2_CUVID_DECODER
 922 DEFINE_CUVID_CODEC(mpeg2, MPEG2VIDEO)
 923 #endif
 924
 925 #if CONFIG_MPEG4_CUVID_DECODER
 926 DEFINE_CUVID_CODEC(mpeg4, MPEG4)
 927 #endif
 928
 929 #if CONFIG_VP8_CUVID_DECODER
 930 DEFINE_CUVID_CODEC(vp8, VP8)
 931 #endif
 932
 933 #if CONFIG_VP9_CUVID_DECODER
 934 DEFINE_CUVID_CODEC(vp9, VP9)
 935 #endif
 936
 937 #if CONFIG_VC1_CUVID_DECODER
 938 DEFINE_CUVID_CODEC(vc1, VC1)
 939 #endif