avcodec/cuvid: check for and warn about invalid pkt_timebase
[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             if (avctx->pkt_timebase.num && avctx->pkt_timebase.den)
262                 cupkt.timestamp = av_rescale_q(avpkt->pts, avctx->pkt_timebase, (AVRational){1, 10000000});
263             else
264                 cupkt.timestamp = avpkt->pts;
265         }
266     } else {
267         cupkt.flags = CUVID_PKT_ENDOFSTREAM;
268     }
269
270     ret = CHECK_CU(cuvidParseVideoData(ctx->cuparser, &cupkt));
271
272     av_packet_unref(&filtered_packet);
273
274     if (ret < 0) {
275         if (ctx->internal_error)
276             ret = ctx->internal_error;
277         goto error;
278     }
279
280     if (av_fifo_size(ctx->frame_queue)) {
281         CUVIDPARSERDISPINFO dispinfo;
282         CUVIDPROCPARAMS params;
283         unsigned int pitch = 0;
284         int offset = 0;
285         int i;
286
287         av_fifo_generic_read(ctx->frame_queue, &dispinfo, sizeof(CUVIDPARSERDISPINFO), NULL);
288
289         memset(&params, 0, sizeof(params));
290         params.progressive_frame = dispinfo.progressive_frame;
291         params.second_field = 0;
292         params.top_field_first = dispinfo.top_field_first;
293
294         ret = CHECK_CU(cuvidMapVideoFrame(ctx->cudecoder, dispinfo.picture_index, &mapped_frame, &pitch, &params));
295         if (ret < 0)
296             goto error;
297
298         if (avctx->pix_fmt == AV_PIX_FMT_CUDA) {
299             ret = av_hwframe_get_buffer(ctx->hwframe, frame, 0);
300             if (ret < 0) {
301                 av_log(avctx, AV_LOG_ERROR, "av_hwframe_get_buffer failed\n");
302                 goto error;
303             }
304
305             ret = ff_decode_frame_props(avctx, frame);
306             if (ret < 0) {
307                 av_log(avctx, AV_LOG_ERROR, "ff_decode_frame_props failed\n");
308                 goto error;
309             }
310
311             for (i = 0; i < 2; i++) {
312                 CUDA_MEMCPY2D cpy = {
313                     .srcMemoryType = CU_MEMORYTYPE_DEVICE,
314                     .dstMemoryType = CU_MEMORYTYPE_DEVICE,
315                     .srcDevice     = mapped_frame,
316                     .dstDevice     = (CUdeviceptr)frame->data[i],
317                     .srcPitch      = pitch,
318                     .dstPitch      = frame->linesize[i],
319                     .srcY          = offset,
320                     .WidthInBytes  = FFMIN(pitch, frame->linesize[i]),
321                     .Height        = avctx->coded_height >> (i ? 1 : 0),
322                 };
323
324                 ret = CHECK_CU(cuMemcpy2D(&cpy));
325                 if (ret < 0)
326                     goto error;
327
328                 offset += avctx->coded_height;
329             }
330         } else if (avctx->pix_fmt == AV_PIX_FMT_NV12) {
331             AVFrame *tmp_frame = av_frame_alloc();
332             if (!tmp_frame) {
333                 av_log(avctx, AV_LOG_ERROR, "av_frame_alloc failed\n");
334                 ret = AVERROR(ENOMEM);
335                 goto error;
336             }
337
338             tmp_frame->format        = AV_PIX_FMT_CUDA;
339             tmp_frame->hw_frames_ctx = av_buffer_ref(ctx->hwframe);
340             tmp_frame->data[0]       = (uint8_t*)mapped_frame;
341             tmp_frame->linesize[0]   = pitch;
342             tmp_frame->data[1]       = (uint8_t*)(mapped_frame + avctx->coded_height * pitch);
343             tmp_frame->linesize[1]   = pitch;
344             tmp_frame->width         = avctx->width;
345             tmp_frame->height        = avctx->height;
346
347             ret = ff_get_buffer(avctx, frame, 0);
348             if (ret < 0) {
349                 av_log(avctx, AV_LOG_ERROR, "ff_get_buffer failed\n");
350                 av_frame_free(&tmp_frame);
351                 goto error;
352             }
353
354             ret = av_hwframe_transfer_data(frame, tmp_frame, 0);
355             if (ret) {
356                 av_log(avctx, AV_LOG_ERROR, "av_hwframe_transfer_data failed\n");
357                 av_frame_free(&tmp_frame);
358                 goto error;
359             }
360
361             av_frame_free(&tmp_frame);
362         } else {
363             ret = AVERROR_BUG;
364             goto error;
365         }
366
367         frame->width = avctx->width;
368         frame->height = avctx->height;
369         if (avctx->pkt_timebase.num && avctx->pkt_timebase.den)
370             frame->pts = av_rescale_q(dispinfo.timestamp, (AVRational){1, 10000000}, avctx->pkt_timebase);
371         else
372             frame->pts = dispinfo.timestamp;
373
374         /* CUVIDs opaque reordering breaks the internal pkt logic.
375          * So set pkt_pts and clear all the other pkt_ fields.
376          */
377         frame->pkt_pts = frame->pts;
378         av_frame_set_pkt_pos(frame, -1);
379         av_frame_set_pkt_duration(frame, 0);
380         av_frame_set_pkt_size(frame, -1);
381
382         frame->interlaced_frame = !dispinfo.progressive_frame;
383
384         if (!dispinfo.progressive_frame)
385             frame->top_field_first = dispinfo.top_field_first;
386
387         *got_frame = 1;
388     } else {
389         *got_frame = 0;
390     }
391
392 error:
393     if (mapped_frame)
394         eret = CHECK_CU(cuvidUnmapVideoFrame(ctx->cudecoder, mapped_frame));
395
396     eret = CHECK_CU(cuCtxPopCurrent(&dummy));
397
398     if (eret < 0)
399         return eret;
400     else
401         return ret;
402 }
403
404 static av_cold int cuvid_decode_end(AVCodecContext *avctx)
405 {
406     CuvidContext *ctx = avctx->priv_data;
407
408     av_fifo_freep(&ctx->frame_queue);
409
410     if (ctx->bsf)
411         av_bsf_free(&ctx->bsf);
412
413     if (ctx->cuparser)
414         cuvidDestroyVideoParser(ctx->cuparser);
415
416     if (ctx->cudecoder)
417         cuvidDestroyDecoder(ctx->cudecoder);
418
419     av_buffer_unref(&ctx->hwframe);
420     av_buffer_unref(&ctx->hwdevice);
421
422     return 0;
423 }
424
425 static void cuvid_ctx_free(AVHWDeviceContext *ctx)
426 {
427     AVCUDADeviceContext *hwctx = ctx->hwctx;
428     cuCtxDestroy(hwctx->cuda_ctx);
429 }
430
431 static int cuvid_test_dummy_decoder(AVCodecContext *avctx, CUVIDPARSERPARAMS *cuparseinfo)
432 {
433     CUVIDDECODECREATEINFO cuinfo;
434     CUvideodecoder cudec = 0;
435     int ret = 0;
436
437     memset(&cuinfo, 0, sizeof(cuinfo));
438
439     cuinfo.CodecType = cuparseinfo->CodecType;
440     cuinfo.ChromaFormat = cudaVideoChromaFormat_420;
441     cuinfo.OutputFormat = cudaVideoSurfaceFormat_NV12;
442
443     cuinfo.ulWidth = 1280;
444     cuinfo.ulHeight = 720;
445     cuinfo.ulTargetWidth = cuinfo.ulWidth;
446     cuinfo.ulTargetHeight = cuinfo.ulHeight;
447
448     cuinfo.target_rect.left = 0;
449     cuinfo.target_rect.top = 0;
450     cuinfo.target_rect.right = cuinfo.ulWidth;
451     cuinfo.target_rect.bottom = cuinfo.ulHeight;
452
453     cuinfo.ulNumDecodeSurfaces = MAX_FRAME_COUNT;
454     cuinfo.ulNumOutputSurfaces = 1;
455     cuinfo.ulCreationFlags = cudaVideoCreate_PreferCUVID;
456
457     cuinfo.DeinterlaceMode = cudaVideoDeinterlaceMode_Weave;
458
459     ret = CHECK_CU(cuvidCreateDecoder(&cudec, &cuinfo));
460     if (ret < 0)
461         return ret;
462
463     ret = CHECK_CU(cuvidDestroyDecoder(cudec));
464     if (ret < 0)
465         return ret;
466
467     return 0;
468 }
469
470 static av_cold int cuvid_decode_init(AVCodecContext *avctx)
471 {
472     CuvidContext *ctx = avctx->priv_data;
473     AVCUDADeviceContext *device_hwctx;
474     AVHWDeviceContext *device_ctx;
475     AVHWFramesContext *hwframe_ctx;
476     CUVIDSOURCEDATAPACKET seq_pkt;
477     CUdevice device;
478     CUcontext cuda_ctx = NULL;
479     CUcontext dummy;
480     const AVBitStreamFilter *bsf;
481     int ret = 0;
482
483     enum AVPixelFormat pix_fmts[3] = { AV_PIX_FMT_CUDA,
484                                        AV_PIX_FMT_NV12,
485                                        AV_PIX_FMT_NONE };
486
487     ret = ff_get_format(avctx, pix_fmts);
488     if (ret < 0) {
489         av_log(avctx, AV_LOG_ERROR, "ff_get_format failed: %d\n", ret);
490         return ret;
491     }
492
493     ctx->frame_queue = av_fifo_alloc(MAX_FRAME_COUNT * sizeof(CUVIDPARSERDISPINFO));
494     if (!ctx->frame_queue) {
495         ret = AVERROR(ENOMEM);
496         goto error;
497     }
498
499     avctx->pix_fmt = ret;
500
501     if (avctx->hw_frames_ctx) {
502         ctx->hwframe = av_buffer_ref(avctx->hw_frames_ctx);
503         if (!ctx->hwframe) {
504             ret = AVERROR(ENOMEM);
505             goto error;
506         }
507
508         hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
509
510         ctx->hwdevice = av_buffer_ref(hwframe_ctx->device_ref);
511         if (!ctx->hwdevice) {
512             ret = AVERROR(ENOMEM);
513             goto error;
514         }
515
516         device_ctx = hwframe_ctx->device_ctx;
517         device_hwctx = device_ctx->hwctx;
518         cuda_ctx = device_hwctx->cuda_ctx;
519     } else {
520         ctx->hwdevice = av_hwdevice_ctx_alloc(AV_HWDEVICE_TYPE_CUDA);
521         if (!ctx->hwdevice) {
522             av_log(avctx, AV_LOG_ERROR, "Error allocating hwdevice\n");
523             ret = AVERROR(ENOMEM);
524             goto error;
525         }
526
527         ret = CHECK_CU(cuInit(0));
528         if (ret < 0)
529             goto error;
530
531         ret = CHECK_CU(cuDeviceGet(&device, 0));
532         if (ret < 0)
533             goto error;
534
535         ret = CHECK_CU(cuCtxCreate(&cuda_ctx, CU_CTX_SCHED_BLOCKING_SYNC, device));
536         if (ret < 0)
537             goto error;
538
539         device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
540         device_ctx->free = cuvid_ctx_free;
541
542         device_hwctx = device_ctx->hwctx;
543         device_hwctx->cuda_ctx = cuda_ctx;
544
545         ret = CHECK_CU(cuCtxPopCurrent(&dummy));
546         if (ret < 0)
547             goto error;
548
549         ret = av_hwdevice_ctx_init(ctx->hwdevice);
550         if (ret < 0) {
551             av_log(avctx, AV_LOG_ERROR, "av_hwdevice_ctx_init failed\n");
552             goto error;
553         }
554
555         ctx->hwframe = av_hwframe_ctx_alloc(ctx->hwdevice);
556         if (!ctx->hwframe) {
557             av_log(avctx, AV_LOG_ERROR, "av_hwframe_ctx_alloc failed\n");
558             ret = AVERROR(ENOMEM);
559             goto error;
560         }
561     }
562
563     memset(&ctx->cuparseinfo, 0, sizeof(ctx->cuparseinfo));
564     memset(&ctx->cuparse_ext, 0, sizeof(ctx->cuparse_ext));
565     memset(&seq_pkt, 0, sizeof(seq_pkt));
566
567     ctx->cuparseinfo.pExtVideoInfo = &ctx->cuparse_ext;
568
569     switch (avctx->codec->id) {
570 #if CONFIG_H263_CUVID_DECODER
571     case AV_CODEC_ID_H263:
572         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
573         break;
574 #endif
575 #if CONFIG_H264_CUVID_DECODER
576     case AV_CODEC_ID_H264:
577         ctx->cuparseinfo.CodecType = cudaVideoCodec_H264;
578         break;
579 #endif
580 #if CONFIG_HEVC_CUVID_DECODER
581     case AV_CODEC_ID_HEVC:
582         ctx->cuparseinfo.CodecType = cudaVideoCodec_HEVC;
583         break;
584 #endif
585 #if CONFIG_MJPEG_CUVID_DECODER
586     case AV_CODEC_ID_MJPEG:
587         ctx->cuparseinfo.CodecType = cudaVideoCodec_JPEG;
588         break;
589 #endif
590 #if CONFIG_MPEG1_CUVID_DECODER
591     case AV_CODEC_ID_MPEG1VIDEO:
592         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG1;
593         break;
594 #endif
595 #if CONFIG_MPEG2_CUVID_DECODER
596     case AV_CODEC_ID_MPEG2VIDEO:
597         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG2;
598         break;
599 #endif
600 #if CONFIG_MPEG4_CUVID_DECODER
601     case AV_CODEC_ID_MPEG4:
602         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
603         break;
604 #endif
605 #if CONFIG_VP8_CUVID_DECODER
606     case AV_CODEC_ID_VP8:
607         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP8;
608         break;
609 #endif
610 #if CONFIG_VP9_CUVID_DECODER
611     case AV_CODEC_ID_VP9:
612         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP9;
613         break;
614 #endif
615 #if CONFIG_VC1_CUVID_DECODER
616     case AV_CODEC_ID_VC1:
617         ctx->cuparseinfo.CodecType = cudaVideoCodec_VC1;
618         break;
619 #endif
620     default:
621         av_log(avctx, AV_LOG_ERROR, "Invalid CUVID codec!\n");
622         return AVERROR_BUG;
623     }
624
625     if (avctx->codec->id == AV_CODEC_ID_H264 || avctx->codec->id == AV_CODEC_ID_HEVC) {
626         if (avctx->codec->id == AV_CODEC_ID_H264)
627             bsf = av_bsf_get_by_name("h264_mp4toannexb");
628         else
629             bsf = av_bsf_get_by_name("hevc_mp4toannexb");
630
631         if (!bsf) {
632             ret = AVERROR_BSF_NOT_FOUND;
633             goto error;
634         }
635         if (ret = av_bsf_alloc(bsf, &ctx->bsf)) {
636             goto error;
637         }
638         if (((ret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx)) < 0) || ((ret = av_bsf_init(ctx->bsf)) < 0)) {
639             av_bsf_free(&ctx->bsf);
640             goto error;
641         }
642
643         ctx->cuparse_ext.format.seqhdr_data_length = ctx->bsf->par_out->extradata_size;
644         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
645                ctx->bsf->par_out->extradata,
646                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), ctx->bsf->par_out->extradata_size));
647     } else if (avctx->extradata_size > 0) {
648         ctx->cuparse_ext.format.seqhdr_data_length = avctx->extradata_size;
649         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
650                avctx->extradata,
651                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), avctx->extradata_size));
652     }
653
654     ctx->cuparseinfo.ulMaxNumDecodeSurfaces = MAX_FRAME_COUNT;
655     ctx->cuparseinfo.ulMaxDisplayDelay = 4;
656     ctx->cuparseinfo.pUserData = avctx;
657     ctx->cuparseinfo.pfnSequenceCallback = cuvid_handle_video_sequence;
658     ctx->cuparseinfo.pfnDecodePicture = cuvid_handle_picture_decode;
659     ctx->cuparseinfo.pfnDisplayPicture = cuvid_handle_picture_display;
660
661     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
662     if (ret < 0)
663         goto error;
664
665     ret = cuvid_test_dummy_decoder(avctx, &ctx->cuparseinfo);
666     if (ret < 0)
667         goto error;
668
669     ret = CHECK_CU(cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
670     if (ret < 0)
671         goto error;
672
673     seq_pkt.payload = ctx->cuparse_ext.raw_seqhdr_data;
674     seq_pkt.payload_size = ctx->cuparse_ext.format.seqhdr_data_length;
675
676     if (seq_pkt.payload && seq_pkt.payload_size) {
677         ret = CHECK_CU(cuvidParseVideoData(ctx->cuparser, &seq_pkt));
678         if (ret < 0)
679             goto error;
680     }
681
682     ret = CHECK_CU(cuCtxPopCurrent(&dummy));
683     if (ret < 0)
684         goto error;
685
686     ctx->ever_flushed = 0;
687
688     if (!avctx->pkt_timebase.num || !avctx->pkt_timebase.den)
689         av_log(avctx, AV_LOG_WARNING, "Invalid pkt_timebase, passing timestamps as-is.\n");
690
691     return 0;
692
693 error:
694     cuvid_decode_end(avctx);
695     return ret;
696 }
697
698 static void cuvid_flush(AVCodecContext *avctx)
699 {
700     CuvidContext *ctx = avctx->priv_data;
701     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
702     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
703     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
704     int ret;
705
706     ctx->ever_flushed = 1;
707
708     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
709     if (ret < 0)
710         goto error;
711
712     av_fifo_freep(&ctx->frame_queue);
713
714     ctx->frame_queue = av_fifo_alloc(MAX_FRAME_COUNT * sizeof(CUVIDPARSERDISPINFO));
715     if (!ctx->frame_queue) {
716         av_log(avctx, AV_LOG_ERROR, "Failed to recreate frame queue on flush\n");
717         return;
718     }
719
720     if (ctx->cudecoder) {
721         cuvidDestroyDecoder(ctx->cudecoder);
722         ctx->cudecoder = NULL;
723     }
724
725     if (ctx->cuparser) {
726         cuvidDestroyVideoParser(ctx->cuparser);
727         ctx->cuparser = NULL;
728     }
729
730     ret = CHECK_CU(cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
731     if (ret < 0)
732         goto error;
733
734     ret = CHECK_CU(cuCtxPopCurrent(&dummy));
735     if (ret < 0)
736         goto error;
737
738     return;
739  error:
740     av_log(avctx, AV_LOG_ERROR, "CUDA reinit on flush failed\n");
741 }
742
743 #define DEFINE_CUVID_CODEC(x, X) \
744     AVHWAccel ff_##x##_cuvid_hwaccel = { \
745         .name           = #x "_cuvid", \
746         .type           = AVMEDIA_TYPE_VIDEO, \
747         .id             = AV_CODEC_ID_##X, \
748         .pix_fmt        = AV_PIX_FMT_CUDA, \
749     }; \
750     AVCodec ff_##x##_cuvid_decoder = { \
751         .name           = #x "_cuvid", \
752         .long_name      = NULL_IF_CONFIG_SMALL("Nvidia CUVID " #X " decoder"), \
753         .type           = AVMEDIA_TYPE_VIDEO, \
754         .id             = AV_CODEC_ID_##X, \
755         .priv_data_size = sizeof(CuvidContext), \
756         .init           = cuvid_decode_init, \
757         .close          = cuvid_decode_end, \
758         .decode         = cuvid_decode_frame, \
759         .flush          = cuvid_flush, \
760         .capabilities   = AV_CODEC_CAP_DELAY, \
761         .pix_fmts       = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, \
762                                                         AV_PIX_FMT_NV12, \
763                                                         AV_PIX_FMT_NONE }, \
764     };
765
766 #if CONFIG_HEVC_CUVID_DECODER
767 DEFINE_CUVID_CODEC(hevc, HEVC)
768 #endif
769
770 #if CONFIG_H263_CUVID_DECODER
771 DEFINE_CUVID_CODEC(h263, H263)
772 #endif
773
774 #if CONFIG_H264_CUVID_DECODER
775 DEFINE_CUVID_CODEC(h264, H264)
776 #endif
777
778 #if CONFIG_MJPEG_CUVID_DECODER
779 DEFINE_CUVID_CODEC(mjpeg, MJPEG)
780 #endif
781
782 #if CONFIG_MPEG1_CUVID_DECODER
783 DEFINE_CUVID_CODEC(mpeg1, MPEG1VIDEO)
784 #endif
785
786 #if CONFIG_MPEG2_CUVID_DECODER
787 DEFINE_CUVID_CODEC(mpeg2, MPEG2VIDEO)
788 #endif
789
790 #if CONFIG_MPEG4_CUVID_DECODER
791 DEFINE_CUVID_CODEC(mpeg4, MPEG4)
792 #endif
793
794 #if CONFIG_VP8_CUVID_DECODER
795 DEFINE_CUVID_CODEC(vp8, VP8)
796 #endif
797
798 #if CONFIG_VP9_CUVID_DECODER
799 DEFINE_CUVID_CODEC(vp9, VP9)
800 #endif
801
802 #if CONFIG_VC1_CUVID_DECODER
803 DEFINE_CUVID_CODEC(vc1, VC1)
804 #endif