CUVID: Add hardware acceleration support

Signed-off-by: Luca Barbato <lu_zero@gentoo.org>

10 files changed, 190 insertions, 3 deletions

diff --git a/Makefile b/Makefile
index 98eb3ab1d9..c5575a8476 100644
--- a/Makefile
+++ b/Makefile
@@ -82,6 +82,7 @@ ALLAVPROGS  = $(AVBASENAMES:%=%$(EXESUF))
 $(foreach prog,$(AVBASENAMES),$(eval OBJS-$(prog) += cmdutils.o))
 
 OBJS-avconv                   += avconv_opt.o avconv_filter.o
+OBJS-avconv-$(CONFIG_CUVID)   += avconv_cuvid.o
 OBJS-avconv-$(CONFIG_LIBMFX)  += avconv_qsv.o
 OBJS-avconv-$(CONFIG_VAAPI)   += avconv_vaapi.o
 OBJS-avconv-$(CONFIG_VDA)     += avconv_vda.o
diff --git a/avconv.c b/avconv.c
index 94b6da2a8b..df2327ab3f 100644
--- a/avconv.c
+++ b/avconv.c
@@ -2137,6 +2137,17 @@ static int transcode_init(void)
             for (j = 0; j < ifile->nb_streams; j++)
                 input_streams[j + ifile->ist_index]->start = av_gettime_relative();
     }
+    /* For each output stream, compute the right encoding parameters */
+    for (i = 0; i < nb_output_streams; i++) {
+        ost = output_streams[i];
+        if (!ost->stream_copy) {
+#if CONFIG_CUVID
+            if (cuvid_transcode_init(ost))
+                exit_program(1);
+        }
+#endif
+    }
+
 
     /* init input streams */
     for (i = 0; i < nb_input_streams; i++)
diff --git a/avconv.h b/avconv.h
index 3c3f0ef659..70639f11f1 100644
--- a/avconv.h
+++ b/avconv.h
@@ -51,6 +51,7 @@
 enum HWAccelID {
     HWACCEL_NONE = 0,
     HWACCEL_AUTO,
+    HWACCEL_CUVID,
     HWACCEL_VDPAU,
     HWACCEL_DXVA2,
     HWACCEL_VDA,
@@ -502,6 +503,8 @@ int ifilter_parameters_from_frame(InputFilter *ifilter, const AVFrame *frame);
 
 int avconv_parse_options(int argc, char **argv);
 
+int cuvid_init(AVCodecContext *s);
+int cuvid_transcode_init(OutputStream *ost);
 int vdpau_init(AVCodecContext *s);
 int dxva2_init(AVCodecContext *s);
 int vda_init(AVCodecContext *s);
diff --git a/avconv_cuvid.c b/avconv_cuvid.c
new file mode 100644
index 0000000000..5c4d21c887
--- /dev/null
+++ b/avconv_cuvid.c
@@ -0,0 +1,150 @@
+/*
+ * This file is part of Libav.
+ *
+ * Libav is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * Libav is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with Libav; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#include "libavutil/hwcontext.h"
+
+#include "avconv.h"
+
+typedef struct CUVIDContext {
+    AVBufferRef *hw_frames_ctx;
+} CUVIDContext;
+
+static void cuvid_uninit(AVCodecContext *avctx)
+{
+    InputStream  *ist = avctx->opaque;
+    CUVIDContext *ctx = ist->hwaccel_ctx;
+
+    if (ctx) {
+        av_buffer_unref(&ctx->hw_frames_ctx);
+        av_freep(&ctx);
+    }
+
+    av_buffer_unref(&ist->hw_frames_ctx);
+
+    ist->hwaccel_ctx    = 0;
+    ist->hwaccel_uninit = 0;
+}
+
+int cuvid_init(AVCodecContext *avctx)
+{
+    InputStream  *ist = avctx->opaque;
+    CUVIDContext *ctx = ist->hwaccel_ctx;
+
+    if (!ctx) {
+        av_log(NULL, AV_LOG_ERROR, "CUVID transcoding is not initialized. "
+               "-hwaccel cuvid should only be used for one-to-one CUVID transcoding "
+               "with no (software) filters.\n");
+        return AVERROR(EINVAL);
+    }
+
+    return 0;
+}
+
+int cuvid_transcode_init(OutputStream *ost)
+{
+    InputStream *ist;
+    const enum AVPixelFormat *pix_fmt;
+    AVHWFramesContext *hwframe_ctx;
+    AVBufferRef *device_ref = NULL;
+    CUVIDContext *ctx = NULL;
+    int ret = 0;
+
+    if (ost->source_index < 0)
+        return 0;
+
+    ist = input_streams[ost->source_index];
+
+    /* check if the encoder supports CUVID */
+    if (!ost->enc->pix_fmts)
+        goto cancel;
+    for (pix_fmt = ost->enc->pix_fmts; *pix_fmt != AV_PIX_FMT_NONE; pix_fmt++)
+        if (*pix_fmt == AV_PIX_FMT_CUDA)
+            break;
+    if (*pix_fmt == AV_PIX_FMT_NONE)
+        goto cancel;
+
+    /* check if the decoder supports CUVID */
+    if (ist->hwaccel_id != HWACCEL_CUVID || !ist->dec || !ist->dec->pix_fmts)
+        goto cancel;
+    for (pix_fmt = ist->dec->pix_fmts; *pix_fmt != AV_PIX_FMT_NONE; pix_fmt++)
+        if (*pix_fmt == AV_PIX_FMT_CUDA)
+            break;
+    if (*pix_fmt == AV_PIX_FMT_NONE)
+        goto cancel;
+
+    av_log(NULL, AV_LOG_VERBOSE, "Setting up CUVID transcoding\n");
+
+    if (ist->hwaccel_ctx) {
+        ctx = ist->hwaccel_ctx;
+    } else {
+        ctx = av_mallocz(sizeof(*ctx));
+        if (!ctx) {
+            ret = AVERROR(ENOMEM);
+            goto error;
+        }
+    }
+
+    if (!ctx->hw_frames_ctx) {
+        ret = av_hwdevice_ctx_create(&device_ref, AV_HWDEVICE_TYPE_CUDA,
+                                     ist->hwaccel_device, NULL, 0);
+        if (ret < 0)
+            goto error;
+
+        ctx->hw_frames_ctx = av_hwframe_ctx_alloc(device_ref);
+        if (!ctx->hw_frames_ctx) {
+            av_log(NULL, AV_LOG_ERROR, "av_hwframe_ctx_alloc failed\n");
+            ret = AVERROR(ENOMEM);
+            goto error;
+        }
+        av_buffer_unref(&device_ref);
+
+        ist->hw_frames_ctx = av_buffer_ref(ctx->hw_frames_ctx);
+        if (!ist->hw_frames_ctx) {
+            av_log(NULL, AV_LOG_ERROR, "av_buffer_ref failed\n");
+            ret = AVERROR(ENOMEM);
+            goto error;
+        }
+
+        ist->hwaccel_ctx    = ctx;
+        ist->hwaccel_uninit = cuvid_uninit;
+
+        /* This is a bit hacky, av_hwframe_ctx_init is called by the cuvid decoder
+         * once it has probed the necessary format information. But as filters/nvenc
+         * need to know the format/sw_format, set them here so they are happy.
+         * This is fine as long as CUVID doesn't add another supported pix_fmt.
+         */
+        hwframe_ctx            = (AVHWFramesContext*)ctx->hw_frames_ctx->data;
+        hwframe_ctx->format    = AV_PIX_FMT_CUDA;
+        hwframe_ctx->sw_format = AV_PIX_FMT_NV12;
+    }
+
+    return 0;
+
+error:
+    av_freep(&ctx);
+    av_buffer_unref(&device_ref);
+    return ret;
+
+cancel:
+    if (ist->hwaccel_id == HWACCEL_CUVID) {
+        av_log(NULL, AV_LOG_ERROR, "CUVID hwaccel requested, but impossible to achieve.\n");
+        return AVERROR(EINVAL);
+    }
+
+    return 0;
+}
diff --git a/avconv_opt.c b/avconv_opt.c
index e078a0b89d..8ddfded661 100644
--- a/avconv_opt.c
+++ b/avconv_opt.c
@@ -62,6 +62,9 @@ const HWAccel hwaccels[] = {
 #if HAVE_DXVA2_LIB
     { "dxva2", dxva2_init, HWACCEL_DXVA2, AV_PIX_FMT_DXVA2_VLD },
 #endif
+#if CONFIG_CUVID
+    { "cuvid", cuvid_init, HWACCEL_CUVID, AV_PIX_FMT_CUDA },
+#endif
 #if CONFIG_VDA
     { "vda",   vda_init,   HWACCEL_VDA,   AV_PIX_FMT_VDA },
 #endif
diff --git a/libavcodec/h263dec.c b/libavcodec/h263dec.c
index 921ff5fb98..e901a6f9f6 100644
--- a/libavcodec/h263dec.c
+++ b/libavcodec/h263dec.c
@@ -654,6 +654,9 @@ intrax8_decoded:
 }
 
 const enum AVPixelFormat ff_h263_hwaccel_pixfmt_list_420[] = {
+#if CONFIG_MPEG4_CUVID_HWACCEL
+    AV_PIX_FMT_CUDA,
+#endif
 #if CONFIG_H263_VAAPI_HWACCEL || CONFIG_MPEG4_VAAPI_HWACCEL
     AV_PIX_FMT_VAAPI,
 #endif
diff --git a/libavcodec/h264_slice.c b/libavcodec/h264_slice.c
index f1f5fc05f9..ca26af550b 100644
--- a/libavcodec/h264_slice.c
+++ b/libavcodec/h264_slice.c
@@ -716,7 +716,8 @@ static void init_scan_tables(H264Context *h)
 
 static enum AVPixelFormat get_pixel_format(H264Context *h)
 {
-#define HWACCEL_MAX (CONFIG_H264_DXVA2_HWACCEL + \
+#define HWACCEL_MAX (CONFIG_H264_CUVID_HWACCEL + \
+                     CONFIG_H264_DXVA2_HWACCEL + \
                      CONFIG_H264_D3D11VA_HWACCEL + \
                      CONFIG_H264_VAAPI_HWACCEL + \
                      (CONFIG_H264_VDA_HWACCEL * 2) + \
@@ -764,6 +765,9 @@ static enum AVPixelFormat get_pixel_format(H264Context *h)
             else
                 *fmt++ = AV_PIX_FMT_YUV422P;
         } else {
+#if CONFIG_H264_CUVID_HWACCEL
+            *fmt++ = AV_PIX_FMT_CUDA;
+#endif
 #if CONFIG_H264_DXVA2_HWACCEL
             *fmt++ = AV_PIX_FMT_DXVA2_VLD;
 #endif
diff --git a/libavcodec/hevcdec.c b/libavcodec/hevcdec.c
index e24ce1e3c0..03142501ed 100644
--- a/libavcodec/hevcdec.c
+++ b/libavcodec/hevcdec.c
@@ -383,12 +383,18 @@ static void export_stream_params(AVCodecContext *avctx, const HEVCParamSets *ps,
 
 static enum AVPixelFormat get_format(HEVCContext *s, const HEVCSPS *sps)
 {
-    #define HWACCEL_MAX (CONFIG_HEVC_DXVA2_HWACCEL + CONFIG_HEVC_D3D11VA_HWACCEL + \
-                         CONFIG_HEVC_VAAPI_HWACCEL + CONFIG_HEVC_VDPAU_HWACCEL)
+    #define HWACCEL_MAX (CONFIG_HEVC_CUVID_HWACCEL +   \
+                         CONFIG_HEVC_DXVA2_HWACCEL +   \
+                         CONFIG_HEVC_D3D11VA_HWACCEL + \
+                         CONFIG_HEVC_VAAPI_HWACCEL +   \
+                         CONFIG_HEVC_VDPAU_HWACCEL)
     enum AVPixelFormat pix_fmts[HWACCEL_MAX + 2], *fmt = pix_fmts;
 
     if (sps->pix_fmt == AV_PIX_FMT_YUV420P || sps->pix_fmt == AV_PIX_FMT_YUVJ420P ||
         sps->pix_fmt == AV_PIX_FMT_YUV420P10) {
+#if CONFIG_HEVC_CUVID_HWACCEL
+        *fmt++ = AV_PIX_FMT_CUDA;
+#endif
 #if CONFIG_HEVC_D3D11VA_HWACCEL
         *fmt++ = AV_PIX_FMT_D3D11VA_VLD;
 #endif
diff --git a/libavcodec/mpeg12dec.c b/libavcodec/mpeg12dec.c
index afdd652b6a..d9b9f36d37 100644
--- a/libavcodec/mpeg12dec.c
+++ b/libavcodec/mpeg12dec.c
@@ -1130,6 +1130,9 @@ static const enum AVPixelFormat pixfmt_xvmc_mpg2_420[] = {
 #endif /* FF_API_XVMC */
 
 static const enum AVPixelFormat mpeg12_hwaccel_pixfmt_list_420[] = {
+#if CONFIG_MPEG1_CUVID_HWACCEL | CONFIG_MPEG2_CUVID_HWACCEL
+    AV_PIX_FMT_CUDA,
+#endif
 #if CONFIG_MPEG2_DXVA2_HWACCEL
     AV_PIX_FMT_DXVA2_VLD,
 #endif
diff --git a/libavcodec/vc1dec.c b/libavcodec/vc1dec.c
index b26fbf2eff..2d58594e91 100644
--- a/libavcodec/vc1dec.c
+++ b/libavcodec/vc1dec.c
@@ -957,6 +957,9 @@ err:
 
 
 static const enum AVPixelFormat vc1_hwaccel_pixfmt_list_420[] = {
+#if CONFIG_VC1_CUVID_HWACCEL
+    AV_PIX_FMT_CUDA,
+#endif
 #if CONFIG_VC1_DXVA2_HWACCEL
     AV_PIX_FMT_DXVA2_VLD,
 #endif