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