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