[FFmpeg-devel,v3,12/12] vaapi_encode: Support limiting slice size

diff --git a/doc/encoders.texi b/doc/encoders.texi
index 39b8fc1d37..21b6d147fe 100644
--- a/doc/encoders.texi
+++ b/doc/encoders.texi
@@ -2855,6 +2855,10 @@  the stream, and requires driver support.  The minimum interval between frames
 must not be smaller than this, and there may be problems if the maximum
 interval is more than a small multiple of it.
 
+@item max_slice_bytes
+Set the maximum number of bytes allowed in a single slice.  Requires driver
+support.  Not limited by default.
+
 @end table
 
 Each encoder also has its own specific options:
diff --git a/libavcodec/vaapi_encode.c b/libavcodec/vaapi_encode.c
index df8f65e431..44d04ca6a9 100644
--- a/libavcodec/vaapi_encode.c
+++ b/libavcodec/vaapi_encode.c
@@ -316,14 +316,14 @@  static int vaapi_encode_issue(AVCodecContext *avctx,
     if (pic->nb_slices == 0)
         pic->nb_slices = ctx->nb_slices;
     if (pic->nb_slices > 0) {
-        int rounding;
-
         pic->slices = av_mallocz_array(pic->nb_slices, sizeof(*pic->slices));
         if (!pic->slices) {
             err = AVERROR(ENOMEM);
             goto fail;
         }
-
+    }
+    if (pic->nb_slices > 0 && ctx->max_slice_bytes == 0) {
+        int rounding;
         for (i = 0; i < pic->nb_slices; i++)
             pic->slices[i].row_size = ctx->slice_size;
 
@@ -435,6 +435,31 @@  static int vaapi_encode_issue(AVCodecContext *avctx,
     }
 #endif
 
+#if VA_CHECK_VERSION(1, 0, 0)
+    if (ctx->max_slice_bytes > 0) {
+        struct {
+            VAEncMiscParameterBuffer misc;
+            VAEncMiscParameterMaxSliceSize slice;
+        } param = {
+            .misc = {
+                .type = VAEncMiscParameterTypeMaxSliceSize,
+            },
+            .slice = {
+                // VAAPI wants this value in bits.
+                .max_slice_size = 8 * ctx->max_slice_bytes,
+            },
+        };
+
+        err = vaapi_encode_make_param_buffer(avctx, pic,
+                                             VAEncMiscParameterBufferType,
+                                             (char*)&param, sizeof(param));
+        if (err < 0)
+            goto fail;
+    }
+#else
+    av_assert0(ctx->max_slice_bytes == 0);
+#endif
+
     vas = vaBeginPicture(ctx->hwctx->display, ctx->va_context,
                          pic->input_surface);
     if (vas != VA_STATUS_SUCCESS) {
@@ -1812,7 +1837,7 @@  static av_cold int vaapi_encode_init_slice_structure(AVCodecContext *avctx)
     ctx->slice_block_cols = (avctx->width  + ctx->slice_block_width  - 1) /
                              ctx->slice_block_width;
 
-    if (avctx->slices <= 1) {
+    if (avctx->slices <= 1 && ctx->max_slice_bytes == 0) {
         ctx->nb_slices  = 1;
         ctx->slice_size = ctx->slice_block_rows;
         return 0;
@@ -1836,6 +1861,30 @@  static av_cold int vaapi_encode_init_slice_structure(AVCodecContext *avctx)
         return AVERROR(EINVAL);
     }
 
+    if (ctx->max_slice_bytes > 0) {
+#if VA_CHECK_VERSION(1, 0, 0)
+        if (slice_structure & VA_ENC_SLICE_STRUCTURE_MAX_SLICE_SIZE) {
+            av_log(avctx, AV_LOG_VERBOSE, "Encoding pictures using "
+                   "at most %d bytes per slice.\n",
+                   ctx->max_slice_bytes);
+            // In this mode we supply only a single slice structure and
+            // no packed headers - the driver generates the headers for
+            // each slice itself.
+            ctx->nb_slices = 1;
+            ctx->desired_packed_headers &= ~VA_ENC_PACKED_HEADER_SLICE;
+            return 0;
+        } else {
+            av_log(avctx, AV_LOG_ERROR, "Driver does not support "
+                   "encoding pictures with a slice size limit.\n");
+            return AVERROR(EINVAL);
+        }
+#else
+        av_log(avctx, AV_LOG_ERROR, "Slice size limiting is "
+               "not supported with this VAAPI version.\n");
+        return AVERROR(EINVAL);
+#endif
+    }
+
     // For fixed-size slices currently we only support whole rows, making
     // rectangular slices.  This could be extended to arbitrary runs of
     // blocks, but since slices tend to be a conformance requirement and
diff --git a/libavcodec/vaapi_encode.h b/libavcodec/vaapi_encode.h
index 860b67dbe2..ddc947edbe 100644
--- a/libavcodec/vaapi_encode.h
+++ b/libavcodec/vaapi_encode.h
@@ -191,6 +191,9 @@  typedef struct VAAPIEncodeContext {
     // framerate.
     AVRational      vfr_max_fps;
 
+    // Requested maximum slice size in bytes.  Ignored if zero.
+    unsigned int    max_slice_bytes;
+
     // Desired packed headers.
     unsigned int    desired_packed_headers;
 
@@ -441,7 +444,11 @@  int ff_vaapi_encode_close(AVCodecContext *avctx);
     { "b_depth", \
       "Maximum B-frame reference depth", \
       OFFSET(common.desired_b_depth), AV_OPT_TYPE_INT, \
-      { .i64 = 1 }, 1, INT_MAX, FLAGS }
+      { .i64 = 1 }, 1, INT_MAX, FLAGS }, \
+    { "max_slice_bytes", \
+      "Maximum number of bytes allowed in a single slice", \
+      OFFSET(common.max_slice_bytes), AV_OPT_TYPE_INT, \
+      { .i64 = 0 }, 0, INT_MAX, FLAGS }
 
 #define VAAPI_ENCODE_RC_MODE(name, desc) \
     { #name, desc, 0, AV_OPT_TYPE_CONST, { .i64 = RC_MODE_ ## name }, \