]> git.sesse.net Git - ffmpeg/blob - libavcodec/cuvid.c
Merge commit 'a0524d9b1e1bb0012207584f067096df7792df6c'
[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 = FFALIGN(avctx->width, 32);
208         hwframe_ctx->height = FFALIGN(avctx->height, 32);
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         frame->pkt_pts = frame->pts;
473         av_frame_set_pkt_pos(frame, -1);
474         av_frame_set_pkt_duration(frame, 0);
475         av_frame_set_pkt_size(frame, -1);
476
477         frame->interlaced_frame = !parsed_frame.is_deinterlacing && !parsed_frame.dispinfo.progressive_frame;
478
479         if (frame->interlaced_frame)
480             frame->top_field_first = parsed_frame.dispinfo.top_field_first;
481     } else if (ctx->decoder_flushing) {
482         ret = AVERROR_EOF;
483     } else {
484         ret = AVERROR(EAGAIN);
485     }
486
487 error:
488     if (mapped_frame)
489         eret = CHECK_CU(cuvidUnmapVideoFrame(ctx->cudecoder, mapped_frame));
490
491     eret = CHECK_CU(cuCtxPopCurrent(&dummy));
492
493     if (eret < 0)
494         return eret;
495     else
496         return ret;
497 }
498
499 static int cuvid_decode_frame(AVCodecContext *avctx, void *data, int *got_frame, AVPacket *avpkt)
500 {
501     CuvidContext *ctx = avctx->priv_data;
502     AVFrame *frame = data;
503     int ret = 0;
504
505     av_log(avctx, AV_LOG_TRACE, "cuvid_decode_frame\n");
506
507     if (ctx->deint_mode != cudaVideoDeinterlaceMode_Weave) {
508         av_log(avctx, AV_LOG_ERROR, "Deinterlacing is not supported via the old API\n");
509         return AVERROR(EINVAL);
510     }
511
512     if (!ctx->decoder_flushing) {
513         ret = cuvid_decode_packet(avctx, avpkt);
514         if (ret < 0)
515             return ret;
516     }
517
518     ret = cuvid_output_frame(avctx, frame);
519     if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
520         *got_frame = 0;
521     } else if (ret < 0) {
522         return ret;
523     } else {
524         *got_frame = 1;
525     }
526
527     return 0;
528 }
529
530 static av_cold int cuvid_decode_end(AVCodecContext *avctx)
531 {
532     CuvidContext *ctx = avctx->priv_data;
533
534     av_fifo_freep(&ctx->frame_queue);
535
536     if (ctx->bsf)
537         av_bsf_free(&ctx->bsf);
538
539     if (ctx->cuparser)
540         cuvidDestroyVideoParser(ctx->cuparser);
541
542     if (ctx->cudecoder)
543         cuvidDestroyDecoder(ctx->cudecoder);
544
545     av_buffer_unref(&ctx->hwframe);
546     av_buffer_unref(&ctx->hwdevice);
547
548     return 0;
549 }
550
551 static int cuvid_test_dummy_decoder(AVCodecContext *avctx, CUVIDPARSERPARAMS *cuparseinfo)
552 {
553     CUVIDDECODECREATEINFO cuinfo;
554     CUvideodecoder cudec = 0;
555     int ret = 0;
556
557     memset(&cuinfo, 0, sizeof(cuinfo));
558
559     cuinfo.CodecType = cuparseinfo->CodecType;
560     cuinfo.ChromaFormat = cudaVideoChromaFormat_420;
561     cuinfo.OutputFormat = cudaVideoSurfaceFormat_NV12;
562
563     cuinfo.ulWidth = 1280;
564     cuinfo.ulHeight = 720;
565     cuinfo.ulTargetWidth = cuinfo.ulWidth;
566     cuinfo.ulTargetHeight = cuinfo.ulHeight;
567
568     cuinfo.target_rect.left = 0;
569     cuinfo.target_rect.top = 0;
570     cuinfo.target_rect.right = cuinfo.ulWidth;
571     cuinfo.target_rect.bottom = cuinfo.ulHeight;
572
573     cuinfo.ulNumDecodeSurfaces = MAX_FRAME_COUNT;
574     cuinfo.ulNumOutputSurfaces = 1;
575     cuinfo.ulCreationFlags = cudaVideoCreate_PreferCUVID;
576     cuinfo.bitDepthMinus8 = 0;
577
578     cuinfo.DeinterlaceMode = cudaVideoDeinterlaceMode_Weave;
579
580     ret = CHECK_CU(cuvidCreateDecoder(&cudec, &cuinfo));
581     if (ret < 0)
582         return ret;
583
584     ret = CHECK_CU(cuvidDestroyDecoder(cudec));
585     if (ret < 0)
586         return ret;
587
588     return 0;
589 }
590
591 static av_cold int cuvid_decode_init(AVCodecContext *avctx)
592 {
593     CuvidContext *ctx = avctx->priv_data;
594     AVCUDADeviceContext *device_hwctx;
595     AVHWDeviceContext *device_ctx;
596     AVHWFramesContext *hwframe_ctx;
597     CUVIDSOURCEDATAPACKET seq_pkt;
598     CUcontext cuda_ctx = NULL;
599     CUcontext dummy;
600     const AVBitStreamFilter *bsf;
601     int ret = 0;
602
603     enum AVPixelFormat pix_fmts[3] = { AV_PIX_FMT_CUDA,
604                                        AV_PIX_FMT_NV12,
605                                        AV_PIX_FMT_NONE };
606
607     ret = ff_get_format(avctx, pix_fmts);
608     if (ret < 0) {
609         av_log(avctx, AV_LOG_ERROR, "ff_get_format failed: %d\n", ret);
610         return ret;
611     }
612
613     ctx->frame_queue = av_fifo_alloc(MAX_FRAME_COUNT * sizeof(CuvidParsedFrame));
614     if (!ctx->frame_queue) {
615         ret = AVERROR(ENOMEM);
616         goto error;
617     }
618
619     avctx->pix_fmt = ret;
620
621     if (avctx->hw_frames_ctx) {
622         ctx->hwframe = av_buffer_ref(avctx->hw_frames_ctx);
623         if (!ctx->hwframe) {
624             ret = AVERROR(ENOMEM);
625             goto error;
626         }
627
628         hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
629
630         ctx->hwdevice = av_buffer_ref(hwframe_ctx->device_ref);
631         if (!ctx->hwdevice) {
632             ret = AVERROR(ENOMEM);
633             goto error;
634         }
635     } else {
636         ret = av_hwdevice_ctx_create(&ctx->hwdevice, AV_HWDEVICE_TYPE_CUDA, ctx->cu_gpu, NULL, 0);
637         if (ret < 0)
638             goto error;
639
640         ctx->hwframe = av_hwframe_ctx_alloc(ctx->hwdevice);
641         if (!ctx->hwframe) {
642             av_log(avctx, AV_LOG_ERROR, "av_hwframe_ctx_alloc failed\n");
643             ret = AVERROR(ENOMEM);
644             goto error;
645         }
646
647         hwframe_ctx = (AVHWFramesContext*)ctx->hwframe->data;
648     }
649
650     device_ctx = hwframe_ctx->device_ctx;
651     device_hwctx = device_ctx->hwctx;
652     cuda_ctx = device_hwctx->cuda_ctx;
653
654     memset(&ctx->cuparseinfo, 0, sizeof(ctx->cuparseinfo));
655     memset(&ctx->cuparse_ext, 0, sizeof(ctx->cuparse_ext));
656     memset(&seq_pkt, 0, sizeof(seq_pkt));
657
658     ctx->cuparseinfo.pExtVideoInfo = &ctx->cuparse_ext;
659
660     switch (avctx->codec->id) {
661 #if CONFIG_H263_CUVID_DECODER
662     case AV_CODEC_ID_H263:
663         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
664         break;
665 #endif
666 #if CONFIG_H264_CUVID_DECODER
667     case AV_CODEC_ID_H264:
668         ctx->cuparseinfo.CodecType = cudaVideoCodec_H264;
669         break;
670 #endif
671 #if CONFIG_HEVC_CUVID_DECODER
672     case AV_CODEC_ID_HEVC:
673         ctx->cuparseinfo.CodecType = cudaVideoCodec_HEVC;
674         break;
675 #endif
676 #if CONFIG_MJPEG_CUVID_DECODER
677     case AV_CODEC_ID_MJPEG:
678         ctx->cuparseinfo.CodecType = cudaVideoCodec_JPEG;
679         break;
680 #endif
681 #if CONFIG_MPEG1_CUVID_DECODER
682     case AV_CODEC_ID_MPEG1VIDEO:
683         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG1;
684         break;
685 #endif
686 #if CONFIG_MPEG2_CUVID_DECODER
687     case AV_CODEC_ID_MPEG2VIDEO:
688         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG2;
689         break;
690 #endif
691 #if CONFIG_MPEG4_CUVID_DECODER
692     case AV_CODEC_ID_MPEG4:
693         ctx->cuparseinfo.CodecType = cudaVideoCodec_MPEG4;
694         break;
695 #endif
696 #if CONFIG_VP8_CUVID_DECODER
697     case AV_CODEC_ID_VP8:
698         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP8;
699         break;
700 #endif
701 #if CONFIG_VP9_CUVID_DECODER
702     case AV_CODEC_ID_VP9:
703         ctx->cuparseinfo.CodecType = cudaVideoCodec_VP9;
704         break;
705 #endif
706 #if CONFIG_VC1_CUVID_DECODER
707     case AV_CODEC_ID_VC1:
708         ctx->cuparseinfo.CodecType = cudaVideoCodec_VC1;
709         break;
710 #endif
711     default:
712         av_log(avctx, AV_LOG_ERROR, "Invalid CUVID codec!\n");
713         return AVERROR_BUG;
714     }
715
716     if (avctx->codec->id == AV_CODEC_ID_H264 || avctx->codec->id == AV_CODEC_ID_HEVC) {
717         if (avctx->codec->id == AV_CODEC_ID_H264)
718             bsf = av_bsf_get_by_name("h264_mp4toannexb");
719         else
720             bsf = av_bsf_get_by_name("hevc_mp4toannexb");
721
722         if (!bsf) {
723             ret = AVERROR_BSF_NOT_FOUND;
724             goto error;
725         }
726         if (ret = av_bsf_alloc(bsf, &ctx->bsf)) {
727             goto error;
728         }
729         if (((ret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx)) < 0) || ((ret = av_bsf_init(ctx->bsf)) < 0)) {
730             av_bsf_free(&ctx->bsf);
731             goto error;
732         }
733
734         ctx->cuparse_ext.format.seqhdr_data_length = ctx->bsf->par_out->extradata_size;
735         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
736                ctx->bsf->par_out->extradata,
737                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), ctx->bsf->par_out->extradata_size));
738     } else if (avctx->extradata_size > 0) {
739         ctx->cuparse_ext.format.seqhdr_data_length = avctx->extradata_size;
740         memcpy(ctx->cuparse_ext.raw_seqhdr_data,
741                avctx->extradata,
742                FFMIN(sizeof(ctx->cuparse_ext.raw_seqhdr_data), avctx->extradata_size));
743     }
744
745     ctx->cuparseinfo.ulMaxNumDecodeSurfaces = MAX_FRAME_COUNT;
746     ctx->cuparseinfo.ulMaxDisplayDelay = 4;
747     ctx->cuparseinfo.pUserData = avctx;
748     ctx->cuparseinfo.pfnSequenceCallback = cuvid_handle_video_sequence;
749     ctx->cuparseinfo.pfnDecodePicture = cuvid_handle_picture_decode;
750     ctx->cuparseinfo.pfnDisplayPicture = cuvid_handle_picture_display;
751
752     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
753     if (ret < 0)
754         goto error;
755
756     ret = cuvid_test_dummy_decoder(avctx, &ctx->cuparseinfo);
757     if (ret < 0)
758         goto error;
759
760     ret = CHECK_CU(cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
761     if (ret < 0)
762         goto error;
763
764     seq_pkt.payload = ctx->cuparse_ext.raw_seqhdr_data;
765     seq_pkt.payload_size = ctx->cuparse_ext.format.seqhdr_data_length;
766
767     if (seq_pkt.payload && seq_pkt.payload_size) {
768         ret = CHECK_CU(cuvidParseVideoData(ctx->cuparser, &seq_pkt));
769         if (ret < 0)
770             goto error;
771     }
772
773     ret = CHECK_CU(cuCtxPopCurrent(&dummy));
774     if (ret < 0)
775         goto error;
776
777     ctx->prev_pts = INT64_MIN;
778
779     if (!avctx->pkt_timebase.num || !avctx->pkt_timebase.den)
780         av_log(avctx, AV_LOG_WARNING, "Invalid pkt_timebase, passing timestamps as-is.\n");
781
782     return 0;
783
784 error:
785     cuvid_decode_end(avctx);
786     return ret;
787 }
788
789 static void cuvid_flush(AVCodecContext *avctx)
790 {
791     CuvidContext *ctx = avctx->priv_data;
792     AVHWDeviceContext *device_ctx = (AVHWDeviceContext*)ctx->hwdevice->data;
793     AVCUDADeviceContext *device_hwctx = device_ctx->hwctx;
794     CUcontext dummy, cuda_ctx = device_hwctx->cuda_ctx;
795     CUVIDSOURCEDATAPACKET seq_pkt = { 0 };
796     int ret;
797
798     ret = CHECK_CU(cuCtxPushCurrent(cuda_ctx));
799     if (ret < 0)
800         goto error;
801
802     av_fifo_freep(&ctx->frame_queue);
803
804     ctx->frame_queue = av_fifo_alloc(MAX_FRAME_COUNT * sizeof(CuvidParsedFrame));
805     if (!ctx->frame_queue) {
806         av_log(avctx, AV_LOG_ERROR, "Failed to recreate frame queue on flush\n");
807         return;
808     }
809
810     if (ctx->cudecoder) {
811         cuvidDestroyDecoder(ctx->cudecoder);
812         ctx->cudecoder = NULL;
813     }
814
815     if (ctx->cuparser) {
816         cuvidDestroyVideoParser(ctx->cuparser);
817         ctx->cuparser = NULL;
818     }
819
820     ret = CHECK_CU(cuvidCreateVideoParser(&ctx->cuparser, &ctx->cuparseinfo));
821     if (ret < 0)
822         goto error;
823
824     seq_pkt.payload = ctx->cuparse_ext.raw_seqhdr_data;
825     seq_pkt.payload_size = ctx->cuparse_ext.format.seqhdr_data_length;
826
827     if (seq_pkt.payload && seq_pkt.payload_size) {
828         ret = CHECK_CU(cuvidParseVideoData(ctx->cuparser, &seq_pkt));
829         if (ret < 0)
830             goto error;
831     }
832
833     ret = CHECK_CU(cuCtxPopCurrent(&dummy));
834     if (ret < 0)
835         goto error;
836
837     ctx->prev_pts = INT64_MIN;
838     ctx->decoder_flushing = 0;
839
840     return;
841  error:
842     av_log(avctx, AV_LOG_ERROR, "CUDA reinit on flush failed\n");
843 }
844
845 #define OFFSET(x) offsetof(CuvidContext, x)
846 #define VD AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_DECODING_PARAM
847 static const AVOption options[] = {
848     { "deint",    "Set deinterlacing mode", OFFSET(deint_mode), AV_OPT_TYPE_INT,   { .i64 = cudaVideoDeinterlaceMode_Weave    }, cudaVideoDeinterlaceMode_Weave, cudaVideoDeinterlaceMode_Adaptive, VD, "deint" },
849     { "weave",    "Weave deinterlacing (do nothing)",        0, AV_OPT_TYPE_CONST, { .i64 = cudaVideoDeinterlaceMode_Weave    }, 0, 0, VD, "deint" },
850     { "bob",      "Bob deinterlacing",                       0, AV_OPT_TYPE_CONST, { .i64 = cudaVideoDeinterlaceMode_Bob      }, 0, 0, VD, "deint" },
851     { "adaptive", "Adaptive deinterlacing",                  0, AV_OPT_TYPE_CONST, { .i64 = cudaVideoDeinterlaceMode_Adaptive }, 0, 0, VD, "deint" },
852     { "gpu",      "GPU to be used for decoding", OFFSET(cu_gpu), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, VD },
853     { NULL }
854 };
855
856 #define DEFINE_CUVID_CODEC(x, X) \
857     static const AVClass x##_cuvid_class = { \
858         .class_name = #x "_cuvid", \
859         .item_name = av_default_item_name, \
860         .option = options, \
861         .version = LIBAVUTIL_VERSION_INT, \
862     }; \
863     AVHWAccel ff_##x##_cuvid_hwaccel = { \
864         .name           = #x "_cuvid", \
865         .type           = AVMEDIA_TYPE_VIDEO, \
866         .id             = AV_CODEC_ID_##X, \
867         .pix_fmt        = AV_PIX_FMT_CUDA, \
868     }; \
869     AVCodec ff_##x##_cuvid_decoder = { \
870         .name           = #x "_cuvid", \
871         .long_name      = NULL_IF_CONFIG_SMALL("Nvidia CUVID " #X " decoder"), \
872         .type           = AVMEDIA_TYPE_VIDEO, \
873         .id             = AV_CODEC_ID_##X, \
874         .priv_data_size = sizeof(CuvidContext), \
875         .priv_class     = &x##_cuvid_class, \
876         .init           = cuvid_decode_init, \
877         .close          = cuvid_decode_end, \
878         .decode         = cuvid_decode_frame, \
879         .send_packet    = cuvid_decode_packet, \
880         .receive_frame  = cuvid_output_frame, \
881         .flush          = cuvid_flush, \
882         .capabilities   = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AVOID_PROBING, \
883         .pix_fmts       = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, \
884                                                         AV_PIX_FMT_NV12, \
885                                                         AV_PIX_FMT_NONE }, \
886     };
887
888 #if CONFIG_HEVC_CUVID_DECODER
889 DEFINE_CUVID_CODEC(hevc, HEVC)
890 #endif
891
892 #if CONFIG_H263_CUVID_DECODER
893 DEFINE_CUVID_CODEC(h263, H263)
894 #endif
895
896 #if CONFIG_H264_CUVID_DECODER
897 DEFINE_CUVID_CODEC(h264, H264)
898 #endif
899
900 #if CONFIG_MJPEG_CUVID_DECODER
901 DEFINE_CUVID_CODEC(mjpeg, MJPEG)
902 #endif
903
904 #if CONFIG_MPEG1_CUVID_DECODER
905 DEFINE_CUVID_CODEC(mpeg1, MPEG1VIDEO)
906 #endif
907
908 #if CONFIG_MPEG2_CUVID_DECODER
909 DEFINE_CUVID_CODEC(mpeg2, MPEG2VIDEO)
910 #endif
911
912 #if CONFIG_MPEG4_CUVID_DECODER
913 DEFINE_CUVID_CODEC(mpeg4, MPEG4)
914 #endif
915
916 #if CONFIG_VP8_CUVID_DECODER
917 DEFINE_CUVID_CODEC(vp8, VP8)
918 #endif
919
920 #if CONFIG_VP9_CUVID_DECODER
921 DEFINE_CUVID_CODEC(vp9, VP9)
922 #endif
923
924 #if CONFIG_VC1_CUVID_DECODER
925 DEFINE_CUVID_CODEC(vc1, VC1)
926 #endif