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