From d071a1117daff0a84bafe9b437164fdef97af98f Mon Sep 17 00:00:00 2001
From: Zhong Li <zhong.li@intel.com>
Date: Wed, 27 Mar 2019 02:56:24 +0800
Subject: [PATCH] lavc/qsvenc: get vps extradata from MSDK

Signed-off-by: Zhong Li <zhong.li@intel.com>
Reviewed-by: Mark Thompson <sw@jkqxz.net>
---
 libavcodec/qsvenc.c      | 68 +++++++++++++++++++++++++++++-----------
 libavcodec/qsvenc.h      |  3 ++
 libavcodec/qsvenc_hevc.c | 10 +++---
 3 files changed, 59 insertions(+), 22 deletions(-)

diff --git a/libavcodec/qsvenc.c b/libavcodec/qsvenc.c
index 4c03133817..e76c580b4a 100644
--- a/libavcodec/qsvenc.c
+++ b/libavcodec/qsvenc.c
@@ -810,22 +810,37 @@ static int qsv_retrieve_enc_params(AVCodecContext *avctx, QSVEncContext *q)
     };
 #endif
 
-    mfxExtBuffer *ext_buffers[] = {
-        (mfxExtBuffer*)&extradata,
-        (mfxExtBuffer*)&co,
-#if QSV_HAVE_CO2
-        (mfxExtBuffer*)&co2,
-#endif
-#if QSV_HAVE_CO3
-        (mfxExtBuffer*)&co3,
-#endif
+#if QSV_HAVE_CO_VPS
+    uint8_t vps_buf[128];
+    mfxExtCodingOptionVPS extradata_vps = {
+        .Header.BufferId = MFX_EXTBUFF_CODING_OPTION_VPS,
+        .Header.BufferSz = sizeof(extradata_vps),
+        .VPSBuffer       = vps_buf,
+        .VPSBufSize      = sizeof(vps_buf),
     };
+#endif
+
+    mfxExtBuffer *ext_buffers[2 + QSV_HAVE_CO2 + QSV_HAVE_CO3 + QSV_HAVE_CO_VPS];
 
     int need_pps = avctx->codec_id != AV_CODEC_ID_MPEG2VIDEO;
-    int ret;
+    int ret, ext_buf_num = 0, extradata_offset = 0;
+
+    ext_buffers[ext_buf_num++] = (mfxExtBuffer*)&extradata;
+    ext_buffers[ext_buf_num++] = (mfxExtBuffer*)&co;
+#if QSV_HAVE_CO2
+    ext_buffers[ext_buf_num++] = (mfxExtBuffer*)&co2;
+#endif
+#if QSV_HAVE_CO3
+    ext_buffers[ext_buf_num++] = (mfxExtBuffer*)&co3;
+#endif
+#if QSV_HAVE_CO_VPS
+    q->hevc_vps = ((avctx->codec_id == AV_CODEC_ID_HEVC) && QSV_RUNTIME_VERSION_ATLEAST(q->ver, 1, 17));
+    if (q->hevc_vps)
+        ext_buffers[ext_buf_num++] = (mfxExtBuffer*)&extradata_vps;
+#endif
 
     q->param.ExtParam    = ext_buffers;
-    q->param.NumExtParam = FF_ARRAY_ELEMS(ext_buffers);
+    q->param.NumExtParam = ext_buf_num;
 
     ret = MFXVideoENCODE_GetVideoParam(q->session, &q->param);
     if (ret < 0)
@@ -834,20 +849,37 @@ static int qsv_retrieve_enc_params(AVCodecContext *avctx, QSVEncContext *q)
 
     q->packet_size = q->param.mfx.BufferSizeInKB * q->param.mfx.BRCParamMultiplier * 1000;
 
-    if (!extradata.SPSBufSize || (need_pps && !extradata.PPSBufSize)) {
+    if (!extradata.SPSBufSize || (need_pps && !extradata.PPSBufSize)
+#if QSV_HAVE_CO_VPS
+        || (q->hevc_vps && !extradata_vps.VPSBufSize)
+#endif
+    ) {
         av_log(avctx, AV_LOG_ERROR, "No extradata returned from libmfx.\n");
         return AVERROR_UNKNOWN;
     }
 
-    avctx->extradata = av_malloc(extradata.SPSBufSize + need_pps * extradata.PPSBufSize +
-                                 AV_INPUT_BUFFER_PADDING_SIZE);
+    avctx->extradata_size = extradata.SPSBufSize + need_pps * extradata.PPSBufSize;
+#if QSV_HAVE_CO_VPS
+    avctx->extradata_size += q->hevc_vps * extradata_vps.VPSBufSize;
+#endif
+
+    avctx->extradata = av_malloc(avctx->extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
     if (!avctx->extradata)
         return AVERROR(ENOMEM);
 
-    memcpy(avctx->extradata,                        sps_buf, extradata.SPSBufSize);
-    if (need_pps)
-        memcpy(avctx->extradata + extradata.SPSBufSize, pps_buf, extradata.PPSBufSize);
-    avctx->extradata_size = extradata.SPSBufSize + need_pps * extradata.PPSBufSize;
+#if QSV_HAVE_CO_VPS
+    if (q->hevc_vps) {
+        memcpy(avctx->extradata, vps_buf, extradata_vps.VPSBufSize);
+        extradata_offset += extradata_vps.VPSBufSize;
+    }
+#endif
+
+    memcpy(avctx->extradata + extradata_offset, sps_buf, extradata.SPSBufSize);
+    extradata_offset += extradata.SPSBufSize;
+    if (need_pps) {
+        memcpy(avctx->extradata + extradata_offset, pps_buf, extradata.PPSBufSize);
+        extradata_offset += extradata.PPSBufSize;
+    }
     memset(avctx->extradata + avctx->extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 
     cpb_props = ff_add_cpb_side_data(avctx);
diff --git a/libavcodec/qsvenc.h b/libavcodec/qsvenc.h
index 70e0b048ee..f6d53d44cd 100644
--- a/libavcodec/qsvenc.h
+++ b/libavcodec/qsvenc.h
@@ -36,6 +36,7 @@
 
 #define QSV_HAVE_CO2 QSV_VERSION_ATLEAST(1, 6)
 #define QSV_HAVE_CO3 QSV_VERSION_ATLEAST(1, 11)
+#define QSV_HAVE_CO_VPS  QSV_VERSION_ATLEAST(1, 17)
 
 #define QSV_HAVE_TRELLIS QSV_VERSION_ATLEAST(1, 8)
 #define QSV_HAVE_MAX_SLICE_SIZE QSV_VERSION_ATLEAST(1, 9)
@@ -134,6 +135,8 @@ typedef struct QSVEncContext {
 
     mfxVersion          ver;
 
+    int hevc_vps;
+
     // options set by the caller
     int async_depth;
     int idr_interval;
diff --git a/libavcodec/qsvenc_hevc.c b/libavcodec/qsvenc_hevc.c
index 39aa8843c2..4c2592355e 100644
--- a/libavcodec/qsvenc_hevc.c
+++ b/libavcodec/qsvenc_hevc.c
@@ -194,10 +194,12 @@ static av_cold int qsv_enc_init(AVCodecContext *avctx)
     if (ret < 0)
         return ret;
 
-    ret = generate_fake_vps(&q->qsv, avctx);
-    if (ret < 0) {
-        ff_qsv_enc_close(avctx, &q->qsv);
-        return ret;
+    if (!q->qsv.hevc_vps) {
+        ret = generate_fake_vps(&q->qsv, avctx);
+        if (ret < 0) {
+            ff_qsv_enc_close(avctx, &q->qsv);
+            return ret;
+        }
     }
 
     return 0;