vaapi_encode: Maintain a pool of bitstream output buffers

Previously we would allocate a new one for every frame.  This instead
maintains an AVBufferPool of them to use as-needed.

Also makes the maximum size of an output buffer adapt to the frame
size - the fixed upper bound was a bit too easy to hit when encoding
large pictures at high quality.

2 files changed, 75 insertions, 19 deletions

diff --git a/libavcodec/vaapi_encode.c b/libavcodec/vaapi_encode.c
index c3f4e44452..7d5614575f 100644
--- a/libavcodec/vaapi_encode.c
+++ b/libavcodec/vaapi_encode.c
@@ -178,16 +178,12 @@ static int vaapi_encode_issue(AVCodecContext *avctx,
     pic->recon_surface = (VASurfaceID)(uintptr_t)pic->recon_image->data[3];
     av_log(avctx, AV_LOG_DEBUG, "Recon surface is %#x.\n", pic->recon_surface);
 
-    vas = vaCreateBuffer(ctx->hwctx->display, ctx->va_context,
-                         VAEncCodedBufferType,
-                         MAX_OUTPUT_BUFFER_SIZE, 1, 0,
-                         &pic->output_buffer);
-    if (vas != VA_STATUS_SUCCESS) {
-        av_log(avctx, AV_LOG_ERROR, "Failed to create bitstream "
-               "output buffer: %d (%s).\n", vas, vaErrorStr(vas));
+    pic->output_buffer_ref = av_buffer_pool_get(ctx->output_buffer_pool);
+    if (!pic->output_buffer_ref) {
         err = AVERROR(ENOMEM);
         goto fail;
     }
+    pic->output_buffer = (VABufferID)(uintptr_t)pic->output_buffer_ref->data;
     av_log(avctx, AV_LOG_DEBUG, "Output buffer is %#x.\n",
            pic->output_buffer);
 
@@ -438,7 +434,7 @@ static int vaapi_encode_output(AVCodecContext *avctx,
 
         err = av_new_packet(pkt, buf->size);
         if (err < 0)
-            goto fail;
+            goto fail_mapped;
 
         memcpy(pkt->data, buf->buf, buf->size);
     }
@@ -456,35 +452,32 @@ static int vaapi_encode_output(AVCodecContext *avctx,
         goto fail;
     }
 
-    vaDestroyBuffer(ctx->hwctx->display, pic->output_buffer);
+    av_buffer_unref(&pic->output_buffer_ref);
     pic->output_buffer = VA_INVALID_ID;
 
     av_log(avctx, AV_LOG_DEBUG, "Output read for pic %"PRId64"/%"PRId64".\n",
            pic->display_order, pic->encode_order);
     return 0;
 
+fail_mapped:
+    vaUnmapBuffer(ctx->hwctx->display, pic->output_buffer);
 fail:
-    if (pic->output_buffer != VA_INVALID_ID) {
-        vaUnmapBuffer(ctx->hwctx->display, pic->output_buffer);
-        vaDestroyBuffer(ctx->hwctx->display, pic->output_buffer);
-        pic->output_buffer = VA_INVALID_ID;
-    }
+    av_buffer_unref(&pic->output_buffer_ref);
+    pic->output_buffer = VA_INVALID_ID;
     return err;
 }
 
 static int vaapi_encode_discard(AVCodecContext *avctx,
                                 VAAPIEncodePicture *pic)
 {
-    VAAPIEncodeContext *ctx = avctx->priv_data;
-
     vaapi_encode_wait(avctx, pic);
 
-    if (pic->output_buffer != VA_INVALID_ID) {
+    if (pic->output_buffer_ref) {
         av_log(avctx, AV_LOG_DEBUG, "Discard output for pic "
                "%"PRId64"/%"PRId64".\n",
                pic->display_order, pic->encode_order);
 
-        vaDestroyBuffer(ctx->hwctx->display, pic->output_buffer);
+        av_buffer_unref(&pic->output_buffer_ref);
         pic->output_buffer = VA_INVALID_ID;
     }
 
@@ -1025,6 +1018,57 @@ fail:
     return err;
 }
 
+static void vaapi_encode_free_output_buffer(void *opaque,
+                                            uint8_t *data)
+{
+    AVCodecContext   *avctx = opaque;
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    VABufferID buffer_id;
+
+    buffer_id = (VABufferID)(uintptr_t)data;
+
+    vaDestroyBuffer(ctx->hwctx->display, buffer_id);
+
+    av_log(avctx, AV_LOG_DEBUG, "Freed output buffer %#x\n", buffer_id);
+}
+
+static AVBufferRef *vaapi_encode_alloc_output_buffer(void *opaque,
+                                                     int size)
+{
+    AVCodecContext   *avctx = opaque;
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    VABufferID buffer_id;
+    VAStatus vas;
+    AVBufferRef *ref;
+
+    // The output buffer size is fixed, so it needs to be large enough
+    // to hold the largest possible compressed frame.  We assume here
+    // that the uncompressed frame plus some header data is an upper
+    // bound on that.
+    vas = vaCreateBuffer(ctx->hwctx->display, ctx->va_context,
+                         VAEncCodedBufferType,
+                         3 * ctx->aligned_width * ctx->aligned_height +
+                         (1 << 16), 1, 0, &buffer_id);
+    if (vas != VA_STATUS_SUCCESS) {
+        av_log(avctx, AV_LOG_ERROR, "Failed to create bitstream "
+               "output buffer: %d (%s).\n", vas, vaErrorStr(vas));
+        return NULL;
+    }
+
+    av_log(avctx, AV_LOG_DEBUG, "Allocated output buffer %#x\n", buffer_id);
+
+    ref = av_buffer_create((uint8_t*)(uintptr_t)buffer_id,
+                           sizeof(buffer_id),
+                           &vaapi_encode_free_output_buffer,
+                           avctx, AV_BUFFER_FLAG_READONLY);
+    if (!ref) {
+        vaDestroyBuffer(ctx->hwctx->display, buffer_id);
+        return NULL;
+    }
+
+    return ref;
+}
+
 av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
                                  const VAAPIEncodeType *type)
 {
@@ -1207,6 +1251,14 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
         }
     }
 
+    ctx->output_buffer_pool =
+        av_buffer_pool_init2(sizeof(VABufferID), avctx,
+                             &vaapi_encode_alloc_output_buffer, NULL);
+    if (!ctx->output_buffer_pool) {
+        err = AVERROR(ENOMEM);
+        goto fail;
+    }
+
     // All I are IDR for now.
     ctx->i_per_idr = 0;
     ctx->p_per_i = ((avctx->gop_size + avctx->max_b_frames) /
@@ -1249,6 +1301,8 @@ av_cold int ff_vaapi_encode_close(AVCodecContext *avctx)
     if (ctx->codec->close)
         ctx->codec->close(avctx);
 
+    av_buffer_pool_uninit(&ctx->output_buffer_pool);
+
     av_freep(&ctx->codec_sequence_params);
     av_freep(&ctx->codec_picture_params);
 
diff --git a/libavcodec/vaapi_encode.h b/libavcodec/vaapi_encode.h
index 85f3a0c894..eede73ca97 100644
--- a/libavcodec/vaapi_encode.h
+++ b/libavcodec/vaapi_encode.h
@@ -39,7 +39,6 @@ enum {
     MAX_PARAM_BUFFERS      = 16,
     MAX_REORDER_DELAY      = 16,
     MAX_PARAM_BUFFER_SIZE  = 1024,
-    MAX_OUTPUT_BUFFER_SIZE = 1024 * 1024,
 };
 
 enum {
@@ -84,6 +83,7 @@ typedef struct VAAPIEncodePicture {
     int          nb_param_buffers;
     VABufferID      param_buffers[MAX_PARAM_BUFFERS];
 
+    AVBufferRef    *output_buffer_ref;
     VABufferID      output_buffer;
 
     void           *priv_data;
@@ -130,6 +130,8 @@ typedef struct VAAPIEncodeContext {
     AVBufferRef    *recon_frames_ref;
     AVHWFramesContext *recon_frames;
 
+    AVBufferPool   *output_buffer_pool;
+
     VAConfigAttrib  config_attributes[MAX_CONFIG_ATTRIBUTES];
     int          nb_config_attributes;