vaapi_encode: Move block size calculation after entrypoint selection

The block size can be dependent on the profile and entrypoint selected.
It defaults to 16x16, with codecs able to override this choice with their
own function.

Signed-off-by: Fei Wang <fei.w.wang@intel.com>

7 files changed, 76 insertions, 19 deletions

diff --git a/libavcodec/vaapi_encode.c b/libavcodec/vaapi_encode.c
index 21a0ed0827..0e2f5ed447 100644
--- a/libavcodec/vaapi_encode.c
+++ b/libavcodec/vaapi_encode.c
@@ -2117,6 +2117,8 @@ static av_cold int vaapi_encode_init_slice_structure(AVCodecContext *avctx)
         return 0;
     }
 
+    av_assert0(ctx->slice_block_height > 0 && ctx->slice_block_width > 0);
+
     ctx->slice_block_rows = (avctx->height + ctx->slice_block_height - 1) /
                              ctx->slice_block_height;
     ctx->slice_block_cols = (avctx->width  + ctx->slice_block_width  - 1) /
@@ -2506,6 +2508,20 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx)
     if (err < 0)
         goto fail;
 
+    if (ctx->codec->get_encoder_caps) {
+        err = ctx->codec->get_encoder_caps(avctx);
+        if (err < 0)
+            goto fail;
+    } else {
+        // Assume 16x16 blocks.
+        ctx->surface_width  = FFALIGN(avctx->width,  16);
+        ctx->surface_height = FFALIGN(avctx->height, 16);
+        if (ctx->codec->flags & FLAG_SLICE_CONTROL) {
+            ctx->slice_block_width  = 16;
+            ctx->slice_block_height = 16;
+        }
+    }
+
     err = vaapi_encode_init_rate_control(avctx);
     if (err < 0)
         goto fail;
diff --git a/libavcodec/vaapi_encode.h b/libavcodec/vaapi_encode.h
index 4596e49ba9..af0588c30b 100644
--- a/libavcodec/vaapi_encode.h
+++ b/libavcodec/vaapi_encode.h
@@ -386,6 +386,13 @@ typedef struct VAAPIEncodeType {
     // factor depending on RC mode.
     int default_quality;
 
+    // Determine encode parameters like block sizes for surface alignment
+    // and slices. This may need to query the profile and entrypoint,
+    // which will be available when this function is called. If not set,
+    // assume that all blocks are 16x16 and that surfaces should be
+    // aligned to match this.
+    int (*get_encoder_caps)(AVCodecContext *avctx);
+
     // Perform any extra codec-specific configuration after the
     // codec context is initialised (set up the private data and
     // add any necessary global parameters).
diff --git a/libavcodec/vaapi_encode_h265.c b/libavcodec/vaapi_encode_h265.c
index 76d67237d2..3fa1336bb0 100644
--- a/libavcodec/vaapi_encode_h265.c
+++ b/libavcodec/vaapi_encode_h265.c
@@ -55,6 +55,10 @@ typedef struct VAAPIEncodeH265Picture {
 typedef struct VAAPIEncodeH265Context {
     VAAPIEncodeContext common;
 
+    // Encoder features.
+    uint32_t ctu_size;
+    uint32_t min_cb_size;
+
     // User options.
     int qp;
     int aud;
@@ -1091,6 +1095,27 @@ static int vaapi_encode_h265_init_slice_params(AVCodecContext *avctx,
     return 0;
 }
 
+static av_cold int vaapi_encode_h265_get_encoder_caps(AVCodecContext *avctx)
+{
+    VAAPIEncodeContext      *ctx = avctx->priv_data;
+    VAAPIEncodeH265Context *priv = avctx->priv_data;
+
+    if (!priv->ctu_size) {
+        priv->ctu_size     = 32;
+        priv->min_cb_size  = 16;
+    }
+    av_log(avctx, AV_LOG_VERBOSE, "Using CTU size %dx%d, "
+           "min CB size %dx%d.\n", priv->ctu_size, priv->ctu_size,
+           priv->min_cb_size, priv->min_cb_size);
+
+    ctx->surface_width  = FFALIGN(avctx->width,  priv->min_cb_size);
+    ctx->surface_height = FFALIGN(avctx->height, priv->min_cb_size);
+
+    ctx->slice_block_width = ctx->slice_block_height = priv->ctu_size;
+
+    return 0;
+}
+
 static av_cold int vaapi_encode_h265_configure(AVCodecContext *avctx)
 {
     VAAPIEncodeContext      *ctx = avctx->priv_data;
@@ -1160,6 +1185,7 @@ static const VAAPIEncodeType vaapi_encode_type_h265 = {
 
     .default_quality       = 25,
 
+    .get_encoder_caps      = &vaapi_encode_h265_get_encoder_caps,
     .configure             = &vaapi_encode_h265_configure,
 
     .picture_priv_data_size = sizeof(VAAPIEncodeH265Picture),
@@ -1205,12 +1231,6 @@ static av_cold int vaapi_encode_h265_init(AVCodecContext *avctx)
         VA_ENC_PACKED_HEADER_SLICE    | // Slice headers.
         VA_ENC_PACKED_HEADER_MISC;      // SEI
 
-    ctx->surface_width  = FFALIGN(avctx->width,  16);
-    ctx->surface_height = FFALIGN(avctx->height, 16);
-
-    // CTU size is currently hard-coded to 32.
-    ctx->slice_block_width = ctx->slice_block_height = 32;
-
     if (priv->qp > 0)
         ctx->explicit_qp = priv->qp;
 
diff --git a/libavcodec/vaapi_encode_mjpeg.c b/libavcodec/vaapi_encode_mjpeg.c
index ddc1d0fc8d..92960145ac 100644
--- a/libavcodec/vaapi_encode_mjpeg.c
+++ b/libavcodec/vaapi_encode_mjpeg.c
@@ -434,6 +434,20 @@ static int vaapi_encode_mjpeg_init_slice_params(AVCodecContext *avctx,
     return 0;
 }
 
+static av_cold int vaapi_encode_mjpeg_get_encoder_caps(AVCodecContext *avctx)
+{
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    const AVPixFmtDescriptor *desc;
+
+    desc = av_pix_fmt_desc_get(ctx->input_frames->sw_format);
+    av_assert0(desc);
+
+    ctx->surface_width  = FFALIGN(avctx->width,  8 << desc->log2_chroma_w);
+    ctx->surface_height = FFALIGN(avctx->height, 8 << desc->log2_chroma_h);
+
+    return 0;
+}
+
 static av_cold int vaapi_encode_mjpeg_configure(AVCodecContext *avctx)
 {
     VAAPIEncodeContext       *ctx = avctx->priv_data;
@@ -483,6 +497,7 @@ static const VAAPIEncodeType vaapi_encode_type_mjpeg = {
     .flags                 = FLAG_CONSTANT_QUALITY_ONLY |
                              FLAG_INTRA_ONLY,
 
+    .get_encoder_caps      = &vaapi_encode_mjpeg_get_encoder_caps,
     .configure             = &vaapi_encode_mjpeg_configure,
 
     .default_quality       = 80,
@@ -509,9 +524,6 @@ static av_cold int vaapi_encode_mjpeg_init(AVCodecContext *avctx)
     ctx->desired_packed_headers =
         VA_ENC_PACKED_HEADER_RAW_DATA;
 
-    ctx->surface_width  = FFALIGN(avctx->width,  8);
-    ctx->surface_height = FFALIGN(avctx->height, 8);
-
     return ff_vaapi_encode_init(avctx);
 }
 
diff --git a/libavcodec/vaapi_encode_mpeg2.c b/libavcodec/vaapi_encode_mpeg2.c
index 0e30efe119..01d1b24908 100644
--- a/libavcodec/vaapi_encode_mpeg2.c
+++ b/libavcodec/vaapi_encode_mpeg2.c
@@ -623,9 +623,6 @@ static av_cold int vaapi_encode_mpeg2_init(AVCodecContext *avctx)
     ctx->desired_packed_headers = VA_ENC_PACKED_HEADER_SEQUENCE |
                                   VA_ENC_PACKED_HEADER_PICTURE;
 
-    ctx->surface_width  = FFALIGN(avctx->width,  16);
-    ctx->surface_height = FFALIGN(avctx->height, 16);
-
     return ff_vaapi_encode_init(avctx);
 }
 
diff --git a/libavcodec/vaapi_encode_vp8.c b/libavcodec/vaapi_encode_vp8.c
index 0e384d5a58..e1d7379501 100644
--- a/libavcodec/vaapi_encode_vp8.c
+++ b/libavcodec/vaapi_encode_vp8.c
@@ -210,9 +210,6 @@ static av_cold int vaapi_encode_vp8_init(AVCodecContext *avctx)
     // adding them anyway.
     ctx->desired_packed_headers = 0;
 
-    ctx->surface_width  = FFALIGN(avctx->width,  16);
-    ctx->surface_height = FFALIGN(avctx->height, 16);
-
     return ff_vaapi_encode_init(avctx);
 }
 
diff --git a/libavcodec/vaapi_encode_vp9.c b/libavcodec/vaapi_encode_vp9.c
index 9ca4e451a3..52373543d3 100644
--- a/libavcodec/vaapi_encode_vp9.c
+++ b/libavcodec/vaapi_encode_vp9.c
@@ -184,6 +184,17 @@ static int vaapi_encode_vp9_init_picture_params(AVCodecContext *avctx,
     return 0;
 }
 
+static av_cold int vaapi_encode_vp9_get_encoder_caps(AVCodecContext *avctx)
+{
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+
+    // Surfaces must be aligned to 64x64 superblock boundaries.
+    ctx->surface_width  = FFALIGN(avctx->width,  64);
+    ctx->surface_height = FFALIGN(avctx->height, 64);
+
+    return 0;
+}
+
 static av_cold int vaapi_encode_vp9_configure(AVCodecContext *avctx)
 {
     VAAPIEncodeContext     *ctx = avctx->priv_data;
@@ -231,6 +242,7 @@ static const VAAPIEncodeType vaapi_encode_type_vp9 = {
 
     .picture_priv_data_size = sizeof(VAAPIEncodeVP9Picture),
 
+    .get_encoder_caps      = &vaapi_encode_vp9_get_encoder_caps,
     .configure             = &vaapi_encode_vp9_configure,
 
     .sequence_params_size  = sizeof(VAEncSequenceParameterBufferVP9),
@@ -251,10 +263,6 @@ static av_cold int vaapi_encode_vp9_init(AVCodecContext *avctx)
     // can write its own headers and there is no metadata to include.
     ctx->desired_packed_headers = 0;
 
-    // Surfaces must be aligned to superblock boundaries.
-    ctx->surface_width  = FFALIGN(avctx->width,  64);
-    ctx->surface_height = FFALIGN(avctx->height, 64);
-
     return ff_vaapi_encode_init(avctx);
 }