cuvid: Implement flush to support seeking in media players
[ffmpeg.git] / 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 "libavutil/buffer.h"
23 #include "libavutil/mathematics.h"
24 #include "libavutil/hwcontext.h"
25 #include "libavutil/hwcontext_cuda.h"
26 #include "libavutil/fifo.h"
27 #include "libavutil/log.h"
28
29 #include "avcodec.h"
30 #include "internal.h"
31
32 #include <nvcuvid.h>
33
34 #define MAX_FRAME_COUNT 20
35
36 typedef struct CuvidContext
37 {
38     CUvideodecoder cudecoder;
39     CUvideoparser cuparser;
40
41     AVBufferRef *hwdevice;
42     AVBufferRef *hwframe;
43
44     AVBSFContext *bsf;
45
46     AVFifoBuffer *frame_queue;
47
48     int internal_error;
49     int ever_flushed;
50
51     cudaVideoCodec codec_type;
52     cudaVideoChromaFormat chroma_format;
53
54     CUVIDPARSERPARAMS cuparseinfo;
55     CUVIDEOFORMATEX cuparse_ext;
56 } CuvidContext;
57
58 static int check_cu(AVCodecContext *avctx, CUresult err, const char *func)
59 {
60     const char *err_name;
61     const char *err_string;
62
63     av_log(avctx, AV_LOG_TRACE, "Calling %s\n", func);
64
65     if (err == CUDA_SUCCESS)
66         return 0;
67
68     cuGetErrorName(err, &err_name);
69     cuGetErrorString(err, &err_string);
70
71     av_log(avctx, AV_LOG_ERROR, "%s failed", func);
72     if (err_name && err_string)
73         av_log(avctx, AV_LOG_ERROR, " -> %s: %s", err_name, err_string);
74     av_log(avctx, AV_LOG_ERROR, "\n");
75
76     return AVERROR_EXTERNAL;
77 }
78
79 #define CHECK_CU(x) check_cu(avctx, (x), #x)
80
81 static int CUDAAPI cuvid_handle_video_sequence(void *opaque, CUVIDEOFORMAT* format)
82 {
83     AVCodecContext *avctx = opaque;
84     CuvidContext *ctx = avctx->priv_data;
85     AVHWFramesContext *hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
86     CUVIDDECODECREATEINFO cuinfo;
87
88     av_log(avctx, AV_LOG_TRACE, "pfnSequenceCallback\n");
89
90     ctx->internal_error = 0;
91
92     avctx->width = format->display_area.right;
93     avctx->height = format->display_area.bottom;
94
95     ff_set_sar(avctx, av_div_q(
96         (AVRational){ format->display_aspect_ratio.x, format->display_aspect_ratio.y },
97         (AVRational){ avctx->width, avctx->height }));
98
99     if (!format->progressive_sequence)
100         avctx->flags |= AV_CODEC_FLAG_INTERLACED_DCT;
101     else
102         avctx->flags &= ~AV_CODEC_FLAG_INTERLACED_DCT;
103
104     if (format->video_signal_description.video_full_range_flag)
105         avctx->color_range = AVCOL_RANGE_JPEG;
106     else
107         avctx->color_range = AVCOL_RANGE_MPEG;
108
109     avctx->color_primaries = format->video_signal_description.color_primaries;
110     avctx->color_trc = format->video_signal_description.transfer_characteristics;
111     avctx->colorspace = format->video_signal_description.matrix_coefficients;
112
113     if (format->bitrate)
114         avctx->bit_rate = format->bitrate;
115
116     if (format->frame_rate.numerator && format->frame_rate.denominator) {
117         avctx->framerate.num = format->frame_rate.numerator;
118         avctx->framerate.den = format->frame_rate.denominator;
119     }
120
121     if (ctx->cudecoder
122             && avctx->coded_width == format->coded_width
123             && avctx->coded_height == format->coded_height
124             && ctx->chroma_format == format->chroma_format
125             && ctx->codec_type == format->codec)
126         return 1;
127
128     if (ctx->cudecoder) {
129         av_log(avctx, AV_LOG_ERROR, "re-initializing decoder is not supported\n");
130         ctx->internal_error = AVERROR(EINVAL);
131         return 0;
132     }
133
134     if (hwframe_ctx->pool && !ctx->ever_flushed) {
135         av_log(avctx, AV_LOG_ERROR, "AVHWFramesContext is already initialized\n");
136         ctx->internal_error = AVERROR(EINVAL);
137         return 0;
138     }
139
140     avctx->coded_width = format->coded_width;
141     avctx->coded_height = format->coded_height;
142
143     ctx->chroma_format = format->chroma_format;
144
145     memset(&cuinfo, 0, sizeof(cuinfo));
146
147     cuinfo.CodecType = ctx->codec_type = format->codec;
148     cuinfo.ChromaFormat = format->chroma_format;
149     cuinfo.OutputFormat = cudaVideoSurfaceFormat_NV12;
150
151     cuinfo.ulWidth = avctx->coded_width;
152     cuinfo.ulHeight = avctx->coded_height;
153     cuinfo.ulTargetWidth = cuinfo.ulWidth;
154     cuinfo.ulTargetHeight = cuinfo.ulHeight;
155
156     cuinfo.target_rect.left = 0;
157     cuinfo.target_rect.top = 0;
158     cuinfo.target_rect.right = cuinfo.ulWidth;
159     cuinfo.target_rect.bottom = cuinfo.ulHeight;
160
161     cuinfo.ulNumDecodeSurfaces = MAX_FRAME_COUNT;
162     cuinfo.ulNumOutputSurfaces = 1;
163     cuinfo.ulCreationFlags = cudaVideoCreate_PreferCUVID;
164
165     cuinfo.DeinterlaceMode = cudaVideoDeinterlaceMode_Weave;
166
167     ctx->internal_error = CHECK_CU(cuvidCreateDecoder(&ctx->cudecoder, &cuinfo));
168     if (ctx->internal_error < 0)
169         return 0;
170
171     if (!hwframe_ctx->pool) {
172         hwframe_ctx->format = AV_PIX_FMT_CUDA;
173         hwframe_ctx->sw_format = AV_PIX_FMT_NV12;
174         hwframe_ctx->width = FFALIGN(avctx->coded_width, 32);
175         hwframe_ctx->height = FFALIGN(avctx->coded_height, 32);
176
177         if ((ctx->internal_error = av_hwframe_ctx_init(ctx->hwframe)) < 0) {
178             av_log(avctx, AV_LOG_ERROR, "av_hwframe_ctx_init failed\n");
179             return 0;
180         }
181     }
182
183     return 1;
184 }
185
186 static int CUDAAPI cuvid_handle_picture_decode(void *opaque, CUVIDPICPARAMS* picparams)
187 {
188     AVCodecContext *avctx = opaque;
189     CuvidContext *ctx = avctx->priv_data;
190
191     av_log(avctx, AV_LOG_TRACE, "pfnDecodePicture\n");
192
193     ctx->internal_error = CHECK_CU(cuvidDecodePicture(ctx->cudecoder, picparams));
194     if (ctx->internal_error < 0)
195         return 0;
196
197     return 1;
198 }
199
200 static int CUDAAPI cuvid_handle_picture_display(void *opaque, CUVIDPARSERDISPINFO* dispinfo)
201 {
202     AVCodecContext *avctx = opaque;
203     CuvidContext *ctx = avctx->priv_data;
204
205     av_log(avctx, AV_LOG_TRACE, "pfnDisplayPicture\n");
206
207     ctx->internal_error = 0;
208
209     av_fifo_generic_write(ctx->frame_queue, dispinfo, sizeof(CUVIDPARSERDISPINFO), NULL);
210
211     return 1;
212 }
213
214 static int cuvid_decode_frame(AVCodecContext *avctx, void *data, int *got_frame, AVPacket *avpkt)
215 {
216     CuvidContext *ctx = avctx->priv_data;
217     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
218     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
219     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
220     AVFrame *frame = data;
221     CUVIDSOURCEDATAPACKET cupkt;
222     AVPacket filter_packet = { 0 };
223     AVPacket filtered_packet = { 0 };
224     CUdeviceptr mapped_frame = 0;
225     int ret = 0, eret = 0;
226
227     if (ctx->bsf && avpkt->size) {
228         if ((ret = av_packet_ref(&filter_packet, avpkt)) < 0) {
229             av_log(avctx, AV_LOG_ERROR, "av_packet_ref failed\n");
230             return ret;
231         }
232
233         if ((ret = av_bsf_send_packet(ctx->bsf, &filter_packet)) < 0) {
234             av_log(avctx, AV_LOG_ERROR, "av_bsf_send_packet failed\n");
235             av_packet_unref(&filter_packet);
236             return ret;
237         }
238
239         if ((ret = av_bsf_receive_packet(ctx->bsf, &filtered_packet)) < 0) {
240             av_log(avctx, AV_LOG_ERROR, "av_bsf_receive_packet failed\n");
241             return ret;
242         }
243
244         avpkt = &filtered_packet;
245     }
246
247     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
248     if (ret < 0) {
249         av_packet_unref(&filtered_packet);
250         return ret;
251     }
252
253     memset(&cupkt, 0, sizeof(cupkt));
254
255     if (avpkt->size) {
256         cupkt.payload_size = avpkt->size;
257         cupkt.payload = avpkt->data;
258
259         if (avpkt->pts != AV_NOPTS_VALUE) {
260             cupkt.flags = CUVID_PKT_TIMESTAMP;
261             cupkt.timestamp = av_rescale_q(avpkt->pts, avctx->time_base, (AVRational){1, 10000000});
262         }
263     } else {
264         cupkt.flags = CUVID_PKT_ENDOFSTREAM;
265     }
266
267     ret = CHECK_CU(cuvidParseVideoData(ctx->cuparser, &cupkt));
268
269     av_packet_unref(&filtered_packet);
270
271     if (ret < 0) {
272         if (ctx->internal_error)
273             ret = ctx->internal_error;
274         goto error;
275     }
276
277     if (av_fifo_size(ctx->frame_queue)) {
278         CUVIDPARSERDISPINFO dispinfo;
279         CUVIDPROCPARAMS params;
280         unsigned int pitch = 0;
281         int offset = 0;
282         int i;
283
284         av_fifo_generic_read(ctx->frame_queue, &dispinfo, sizeof(CUVIDPARSERDISPINFO), NULL);
285
286         memset(&params, 0, sizeof(params));
287         params.progressive_frame = dispinfo.progressive_frame;
288         params.second_field = 0;
289         params.top_field_first = dispinfo.top_field_first;
290
291         ret = CHECK_CU(cuvidMapVideoFrame(ctx->cudecoder, dispinfo.picture_index, &mapped_frame, &pitch, &params));
292         if (ret < 0)
293             goto error;
294
295         if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
296             ret = av_hwframe_get_buffer(ctx->hwframe, frame, 0);
297             if (ret < 0) {
298                 av_log(avctx, AV_LOG_ERROR, "av_hwframe_get_buffer failed\n");
299                 goto error;
300             }
301
302             ret = ff_decode_frame_props(avctx, frame);
303             if (ret < 0) {
304                 av_log(avctx, AV_LOG_ERROR, "ff_decode_frame_props failed\n");
305                 goto error;
306             }
307
308             for (i = 0; i < 2; i++) {
309                 CUDA_MEMCPY2D cpy = {
310                     .srcMemoryType = CU_MEMORYTYPE_DEVICE,
311                     .dstMemoryType = CU_MEMORYTYPE_DEVICE,
312                     .srcDevice     = mapped_frame,
313                     .dstDevice     = (CUdeviceptr)frame->data[i],
314                     .srcPitch      = pitch,
315                     .dstPitch      = frame->linesize[i],
316                     .srcY          = offset,
317                     .WidthInBytes  = FFMIN(pitch, frame->linesize[i]),
318                     .Height        = avctx->coded_height >> (i ? 1 : 0),
319                 };
320
321                 ret = CHECK_CU(cuMemcpy2D(&cpy));
322                 if (ret < 0)
323                     goto error;
324
325                 offset += avctx->coded_height;
326             }
327         } else if (avctx->pix_fmt == AV_PIX_FMT_NV12) {
328             AVFrame *tmp_frame = av_frame_alloc();
329             if (!tmp_frame) {
330                 av_log(avctx, AV_LOG_ERROR, "av_frame_alloc failed\n");
331                 ret = AVERROR(ENOMEM);
332                 goto error;
333             }
334
335             tmp_frame->format        = AV_PIX_FMT_CUDA;
336             tmp_frame->hw_frames_ctx = av_buffer_ref(ctx->hwframe);
337             tmp_frame->data[0]       = (uint8_t*)mapped_frame;
338             tmp_frame->linesize[0]   = pitch;
339             tmp_frame->data[1]       = (uint8_t*)(mapped_frame + avctx->coded_height * pitch);
340             tmp_frame->linesize[1]   = pitch;
341             tmp_frame->width         = avctx->width;
342             tmp_frame->height        = avctx->height;
343
344             ret = ff_get_buffer(avctx, frame, 0);
345             if (ret < 0) {
346                 av_log(avctx, AV_LOG_ERROR, "ff_get_buffer failed\n");
347                 av_frame_free(&tmp_frame);
348                 goto error;
349             }
350
351             ret = av_hwframe_transfer_data(frame, tmp_frame, 0);
352             if (ret) {
353                 av_log(avctx, AV_LOG_ERROR, "av_hwframe_transfer_data failed\n");
354                 av_frame_free(&tmp_frame);
355                 goto error;
356             }
357
358             av_frame_free(&tmp_frame);
359         } else {
360             ret = AVERROR_BUG;
361             goto error;
362         }
363
364         frame->width = avctx->width;
365         frame->height = avctx->height;
366         frame->pts = av_rescale_q(dispinfo.timestamp, (AVRational){1, 10000000}, avctx->time_base);
367
368         /* CUVIDs opaque reordering breaks the internal pkt logic.
369          * So set pkt_pts and clear all the other pkt_ fields.
370          */
371         frame->pkt_pts = frame->pts;
372         av_frame_set_pkt_pos(frame, -1);
373         av_frame_set_pkt_duration(frame, 0);
374         av_frame_set_pkt_size(frame, -1);
375
376         frame->interlaced_frame = !dispinfo.progressive_frame;
377
378         if (!dispinfo.progressive_frame)
379             frame->top_field_first = dispinfo.top_field_first;
380
381         *got_frame = 1;
382     } else {
383         *got_frame = 0;
384     }
385
386 error:
387     if (mapped_frame)
388         eret = CHECK_CU(cuvidUnmapVideoFrame(ctx->cudecoder, mapped_frame));
389
390     eret = CHECK_CU(cuCtxPopCurrent(&dummy));
391
392     if (eret < 0)
393         return eret;
394     else
395         return ret;
396 }
397
398 static av_cold int cuvid_decode_end(AVCodecContext *avctx)
399 {
400     CuvidContext *ctx = avctx->priv_data;
401
402     av_fifo_freep(&ctx->frame_queue);
403
404     if (ctx->bsf)
405         av_bsf_free(&ctx->bsf);
406
407     if (ctx->cuparser)
408         cuvidDestroyVideoParser(ctx->cuparser);
409
410     if (ctx->cudecoder)
411         cuvidDestroyDecoder(ctx->cudecoder);
412
413     av_buffer_unref(&ctx->hwframe);
414     av_buffer_unref(&ctx->hwdevice);
415
416     return 0;
417 }
418
419 static void cuvid_ctx_free(AVHWDeviceContext *ctx)
420 {
421     AVCUDADeviceContext *hwctx = ctx->hwctx;
422     cuCtxDestroy(hwctx->cuda_ctx);
423 }
424
425 static int cuvid_test_dummy_decoder(AVCodecContext *avctx, CUVIDPARSERPARAMS *cuparseinfo)
426 {
427     CUVIDDECODECREATEINFO cuinfo;
428     CUvideodecoder cudec = 0;
429     int ret = 0;
430
431     memset(&cuinfo, 0, sizeof(cuinfo));
432
433     cuinfo.CodecType = cuparseinfo->CodecType;
434     cuinfo.ChromaFormat = cudaVideoChromaFormat_420;
435     cuinfo.OutputFormat = cudaVideoSurfaceFormat_NV12;
436
437     cuinfo.ulWidth = 1280;
438     cuinfo.ulHeight = 720;
439     cuinfo.ulTargetWidth = cuinfo.ulWidth;
440     cuinfo.ulTargetHeight = cuinfo.ulHeight;
441
442     cuinfo.target_rect.left = 0;
443     cuinfo.target_rect.top = 0;
444     cuinfo.target_rect.right = cuinfo.ulWidth;
445     cuinfo.target_rect.bottom = cuinfo.ulHeight;
446
447     cuinfo.ulNumDecodeSurfaces = MAX_FRAME_COUNT;
448     cuinfo.ulNumOutputSurfaces = 1;
449     cuinfo.ulCreationFlags = cudaVideoCreate_PreferCUVID;
450
451     cuinfo.DeinterlaceMode = cudaVideoDeinterlaceMode_Weave;
452
453     ret = CHECK_CU(cuvidCreateDecoder(&cudec, &cuinfo));
454     if (ret < 0)
455         return ret;
456
457     ret = CHECK_CU(cuvidDestroyDecoder(cudec));
458     if (ret < 0)
459         return ret;
460
461     return 0;
462 }
463
464 static av_cold int cuvid_decode_init(AVCodecContext *avctx)
465 {
466     CuvidContext *ctx = avctx->priv_data;
467     AVCUDADeviceContext *device_hwctx;
468     AVHWDeviceContext *device_ctx;
469     AVHWFramesContext *hwframe_ctx;
470     CUVIDSOURCEDATAPACKET seq_pkt;
471     CUdevice device;
472     CUcontext cuda_ctx = NULL;
473     CUcontext dummy;
474     const AVBitStreamFilter *bsf;
475     int ret = 0;
476
477     enum AVPixelFormat pix_fmts[3] = { AV_PIX_FMT_CUDA,
478                                        AV_PIX_FMT_NV12,
479                                        AV_PIX_FMT_NONE };
480
481     ret = ff_get_format(avctx, pix_fmts);
482     if (ret < 0) {
483         av_log(avctx, AV_LOG_ERROR, "ff_get_format failed: %d\n", ret);
484         return ret;
485     }
486
487     ctx->frame_queue = av_fifo_alloc(MAX_FRAME_COUNT * sizeof(CUVIDPARSERDISPINFO));
488     if (!ctx->frame_queue) {
489         ret = AVERROR(ENOMEM);
490         goto error;
491     }
492
493     avctx->pix_fmt = ret;
494
495     if (avctx->hw_frames_ctx) {
496         ctx->hwframe = av_buffer_ref(avctx->hw_frames_ctx);
497         if (!ctx->hwframe) {
498             ret = AVERROR(ENOMEM);
499             goto error;
500         }
501
502         hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
503
504         ctx->hwdevice = av_buffer_ref(hwframe_ctx->device_ref);
505         if (!ctx->hwdevice) {
506             ret = AVERROR(ENOMEM);
507             goto error;
508         }
509
510         device_ctx = hwframe_ctx->device_ctx;
511         device_hwctx = device_ctx->hwctx;
512         cuda_ctx = device_hwctx->cuda_ctx;
513     } else {
514         ctx->hwdevice = av_hwdevice_ctx_alloc(AV_HWDEVICE_TYPE_CUDA);
515         if (!ctx->hwdevice) {
516             av_log(avctx, AV_LOG_ERROR, "Error allocating hwdevice\n");
517             ret = AVERROR(ENOMEM);
518             goto error;
519         }
520
521         ret = CHECK_CU(cuInit(0));
522         if (ret < 0)
523             goto error;
524
525         ret = CHECK_CU(cuDeviceGet(&device, 0));
526         if (ret < 0)
527             goto error;
528
529         ret = CHECK_CU(cuCtxCreate(&cuda_ctx, CU_CTX_SCHED_BLOCKING_SYNC, device));
530         if (ret < 0)
531             goto error;
532
533         device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
534         device_ctx->free = cuvid_ctx_free;
535
536         device_hwctx = device_ctx->hwctx;
537         device_hwctx->cuda_ctx = cuda_ctx;
538
539         ret = CHECK_CU(cuCtxPopCurrent(&dummy));
540         if (ret < 0)
541             goto error;
542
543         ret = av_hwdevice_ctx_init(ctx->hwdevice);
544         if (ret < 0) {
545             av_log(avctx, AV_LOG_ERROR, "av_hwdevice_ctx_init failed\n");
546             goto error;
547         }
548
549         ctx->hwframe = av_hwframe_ctx_alloc(ctx->hwdevice);
550         if (!ctx->hwframe) {
551             av_log(avctx, AV_LOG_ERROR, "av_hwframe_ctx_alloc failed\n");
552             ret = AVERROR(ENOMEM);
553             goto error;
554         }
555     }
556
557     memset(&ctx->cuparseinfo, 0, sizeof(ctx->cuparseinfo));
558     memset(&ctx->cuparse_ext, 0, sizeof(ctx->cuparse_ext));
559     memset(&seq_pkt, 0, sizeof(seq_pkt));
560
561     ctx->cuparseinfo.pExtVideoInfo = &ctx->cuparse_ext;
562
563     switch (avctx->codec->id) {
564 #if CONFIG_H263_CUVID_DECODER
565     case AV_CODEC_ID_H263:
566         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
567         break;
568 #endif
569 #if CONFIG_H264_CUVID_DECODER
570     case AV_CODEC_ID_H264:
571         ctx->cuparseinfo.CodecType = cudaVideoCodec_H264;
572         break;
573 #endif
574 #if CONFIG_HEVC_CUVID_DECODER
575     case AV_CODEC_ID_HEVC:
576         ctx->cuparseinfo.CodecType = cudaVideoCodec_HEVC;
577         break;
578 #endif
579 #if CONFIG_MJPEG_CUVID_DECODER
580     case AV_CODEC_ID_MJPEG:
581         ctx->cuparseinfo.CodecType = cudaVideoCodec_JPEG;
582         break;
583 #endif
584 #if CONFIG_MPEG1_CUVID_DECODER
585     case AV_CODEC_ID_MPEG1VIDEO:
586         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG1;
587         break;
588 #endif
589 #if CONFIG_MPEG2_CUVID_DECODER
590     case AV_CODEC_ID_MPEG2VIDEO:
591         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG2;
592         break;
593 #endif
594 #if CONFIG_MPEG4_CUVID_DECODER
595     case AV_CODEC_ID_MPEG4:
596         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
597         break;
598 #endif
599 #if CONFIG_VP8_CUVID_DECODER
600     case AV_CODEC_ID_VP8:
601         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP8;
602         break;
603 #endif
604 #if CONFIG_VP9_CUVID_DECODER
605     case AV_CODEC_ID_VP9:
606         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP9;
607         break;
608 #endif
609 #if CONFIG_VC1_CUVID_DECODER
610     case AV_CODEC_ID_VC1:
611         ctx->cuparseinfo.CodecType = cudaVideoCodec_VC1;
612         break;
613 #endif
614     default:
615         av_log(avctx, AV_LOG_ERROR, "Invalid CUVID codec!\n");
616         return AVERROR_BUG;
617     }
618
619     if (avctx->codec->id == AV_CODEC_ID_H264 || avctx->codec->id == AV_CODEC_ID_HEVC) {
620         if (avctx->codec->id == AV_CODEC_ID_H264)
621             bsf = av_bsf_get_by_name("h264_mp4toannexb");
622         else
623             bsf = av_bsf_get_by_name("hevc_mp4toannexb");
624
625         if (!bsf) {
626             ret = AVERROR_BSF_NOT_FOUND;
627             goto error;
628         }
629         if (ret = av_bsf_alloc(bsf, &ctx->bsf)) {
630             goto error;
631         }
632         if (((ret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx)) < 0) || ((ret = av_bsf_init(ctx->bsf)) < 0)) {
633             av_bsf_free(&ctx->bsf);
634             goto error;
635         }
636
637         ctx->cuparse_ext.format.seqhdr_data_length = ctx->bsf->par_out->extradata_size;
638         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
639                ctx->bsf->par_out->extradata,
640                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), ctx->bsf->par_out->extradata_size));
641     } else if (avctx->extradata_size > 0) {
642         ctx->cuparse_ext.format.seqhdr_data_length = avctx->extradata_size;
643         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
644                avctx->extradata,
645                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), avctx->extradata_size));
646     }
647
648     ctx->cuparseinfo.ulMaxNumDecodeSurfaces = MAX_FRAME_COUNT;
649     ctx->cuparseinfo.ulMaxDisplayDelay = 4;
650     ctx->cuparseinfo.pUserData = avctx;
651     ctx->cuparseinfo.pfnSequenceCallback = cuvid_handle_video_sequence;
652     ctx->cuparseinfo.pfnDecodePicture = cuvid_handle_picture_decode;
653     ctx->cuparseinfo.pfnDisplayPicture = cuvid_handle_picture_display;
654
655     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
656     if (ret < 0)
657         goto error;
658
659     ret = cuvid_test_dummy_decoder(avctx, &ctx->cuparseinfo);
660     if (ret < 0)
661         goto error;
662
663     ret = CHECK_CU(cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
664     if (ret < 0)
665         goto error;
666
667     seq_pkt.payload = ctx->cuparse_ext.raw_seqhdr_data;
668     seq_pkt.payload_size = ctx->cuparse_ext.format.seqhdr_data_length;
669
670     if (seq_pkt.payload && seq_pkt.payload_size) {
671         ret = CHECK_CU(cuvidParseVideoData(ctx->cuparser, &seq_pkt));
672         if (ret < 0)
673             goto error;
674     }
675
676     ret = CHECK_CU(cuCtxPopCurrent(&dummy));
677     if (ret < 0)
678         goto error;
679
680     ctx->ever_flushed = 0;
681
682     return 0;
683
684 error:
685     cuvid_decode_end(avctx);
686     return ret;
687 }
688
689 static void cuvid_flush(AVCodecContext *avctx)
690 {
691     CuvidContext *ctx = avctx->priv_data;
692     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
693     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
694     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
695     int ret;
696
697     ctx->ever_flushed = 1;
698
699     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
700     if (ret < 0)
701         goto error;
702
703     av_fifo_freep(&ctx->frame_queue);
704
705     ctx->frame_queue = av_fifo_alloc(MAX_FRAME_COUNT * sizeof(CUVIDPARSERDISPINFO));
706     if (!ctx->frame_queue) {
707         av_log(avctx, AV_LOG_ERROR, "Failed to recreate frame queue on flush\n");
708         return;
709     }
710
711     if (ctx->cudecoder) {
712         cuvidDestroyDecoder(ctx->cudecoder);
713         ctx->cudecoder = NULL;
714     }
715
716     if (ctx->cuparser) {
717         cuvidDestroyVideoParser(ctx->cuparser);
718         ctx->cuparser = NULL;
719     }
720
721     ret = CHECK_CU(cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
722     if (ret < 0)
723         goto error;
724
725     ret = CHECK_CU(cuCtxPopCurrent(&dummy));
726     if (ret < 0)
727         goto error;
728
729     return;
730  error:
731     av_log(avctx, AV_LOG_ERROR, "CUDA reinit on flush failed\n");
732 }
733
734 #define DEFINE_CUVID_CODEC(x, X) \
735     AVHWAccel ff_##x##_cuvid_hwaccel = { \
736         .name           = #x "_cuvid", \
737         .type           = AVMEDIA_TYPE_VIDEO, \
738         .id             = AV_CODEC_ID_##X, \
739         .pix_fmt        = AV_PIX_FMT_CUDA, \
740     }; \
741     AVCodec ff_##x##_cuvid_decoder = { \
742         .name           = #x "_cuvid", \
743         .long_name      = NULL_IF_CONFIG_SMALL("Nvidia CUVID " #X " decoder"), \
744         .type           = AVMEDIA_TYPE_VIDEO, \
745         .id             = AV_CODEC_ID_##X, \
746         .priv_data_size = sizeof(CuvidContext), \
747         .init           = cuvid_decode_init, \
748         .close          = cuvid_decode_end, \
749         .decode         = cuvid_decode_frame, \
750         .flush          = cuvid_flush, \
751         .capabilities   = AV_CODEC_CAP_DELAY, \
752         .pix_fmts       = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, \
753                                                         AV_PIX_FMT_NV12, \
754                                                         AV_PIX_FMT_NONE }, \
755     };
756
757 #if CONFIG_HEVC_CUVID_DECODER
758 DEFINE_CUVID_CODEC(hevc, HEVC)
759 #endif
760
761 #if CONFIG_H263_CUVID_DECODER
762 DEFINE_CUVID_CODEC(h263, H263)
763 #endif
764
765 #if CONFIG_H264_CUVID_DECODER
766 DEFINE_CUVID_CODEC(h264, H264)
767 #endif
768
769 #if CONFIG_MJPEG_CUVID_DECODER
770 DEFINE_CUVID_CODEC(mjpeg, MJPEG)
771 #endif
772
773 #if CONFIG_MPEG1_CUVID_DECODER
774 DEFINE_CUVID_CODEC(mpeg1, MPEG1VIDEO)
775 #endif
776
777 #if CONFIG_MPEG2_CUVID_DECODER
778 DEFINE_CUVID_CODEC(mpeg2, MPEG2VIDEO)
779 #endif
780
781 #if CONFIG_MPEG4_CUVID_DECODER
782 DEFINE_CUVID_CODEC(mpeg4, MPEG4)
783 #endif
784
785 #if CONFIG_VP8_CUVID_DECODER
786 DEFINE_CUVID_CODEC(vp8, VP8)
787 #endif
788
789 #if CONFIG_VP9_CUVID_DECODER
790 DEFINE_CUVID_CODEC(vp9, VP9)
791 #endif
792
793 #if CONFIG_VC1_CUVID_DECODER
794 DEFINE_CUVID_CODEC(vc1, VC1)
795 #endif