1 files changed, 174 insertions, 85 deletions

diff --git a/libavcodec/vaapi_encode.c b/libavcodec/vaapi_encode.c
index 4a92506b41..ddd0308492 100644
--- a/libavcodec/vaapi_encode.c
+++ b/libavcodec/vaapi_encode.c
@@ -922,7 +922,7 @@ fail:
     return err;
 }
 
-static av_cold int vaapi_encode_check_config(AVCodecContext *avctx)
+static av_cold int vaapi_encode_config_attributes(AVCodecContext *avctx)
 {
     VAAPIEncodeContext *ctx = avctx->priv_data;
     VAStatus vas;
@@ -930,6 +930,7 @@ static av_cold int vaapi_encode_check_config(AVCodecContext *avctx)
     VAProfile    *profiles    = NULL;
     VAEntrypoint *entrypoints = NULL;
     VAConfigAttrib attr[] = {
+        { VAConfigAttribRTFormat        },
         { VAConfigAttribRateControl     },
         { VAConfigAttribEncMaxRefFrames },
     };
@@ -1001,13 +1002,33 @@ static av_cold int vaapi_encode_check_config(AVCodecContext *avctx)
             continue;
         }
         switch (attr[i].type) {
+        case VAConfigAttribRTFormat:
+            if (!(ctx->va_rt_format & attr[i].value)) {
+                av_log(avctx, AV_LOG_ERROR, "Surface RT format %#x "
+                       "is not supported (mask %#x).\n",
+                       ctx->va_rt_format, attr[i].value);
+                err = AVERROR(EINVAL);
+                goto fail;
+            }
+            ctx->config_attributes[ctx->nb_config_attributes++] =
+                (VAConfigAttrib) {
+                .type  = VAConfigAttribRTFormat,
+                .value = ctx->va_rt_format,
+            };
+            break;
         case VAConfigAttribRateControl:
             if (!(ctx->va_rc_mode & attr[i].value)) {
-                av_log(avctx, AV_LOG_ERROR, "Rate control mode is not "
-                       "supported: %x\n", attr[i].value);
+                av_log(avctx, AV_LOG_ERROR, "Rate control mode %#x "
+                       "is not supported (mask: %#x).\n",
+                       ctx->va_rc_mode, attr[i].value);
                 err = AVERROR(EINVAL);
                 goto fail;
             }
+            ctx->config_attributes[ctx->nb_config_attributes++] =
+                (VAConfigAttrib) {
+                .type  = VAConfigAttribRateControl,
+                .value = ctx->va_rc_mode,
+            };
             break;
         case VAConfigAttribEncMaxRefFrames:
         {
@@ -1016,18 +1037,20 @@ static av_cold int vaapi_encode_check_config(AVCodecContext *avctx)
 
             if (avctx->gop_size > 1 && ref_l0 < 1) {
                 av_log(avctx, AV_LOG_ERROR, "P frames are not "
-                       "supported (%x).\n", attr[i].value);
+                       "supported (%#x).\n", attr[i].value);
                 err = AVERROR(EINVAL);
                 goto fail;
             }
             if (avctx->max_b_frames > 0 && ref_l1 < 1) {
                 av_log(avctx, AV_LOG_ERROR, "B frames are not "
-                       "supported (%x).\n", attr[i].value);
+                       "supported (%#x).\n", attr[i].value);
                 err = AVERROR(EINVAL);
                 goto fail;
             }
         }
         break;
+        default:
+            av_assert0(0 && "Unexpected config attribute.");
         }
     }
 
@@ -1038,6 +1061,54 @@ fail:
     return err;
 }
 
+static av_cold int vaapi_encode_init_rate_control(AVCodecContext *avctx)
+{
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    int hrd_buffer_size;
+    int hrd_initial_buffer_fullness;
+
+    if (avctx->bit_rate > INT32_MAX) {
+        av_log(avctx, AV_LOG_ERROR, "Target bitrate of 2^31 bps or "
+               "higher is not supported.\n");
+        return AVERROR(EINVAL);
+    }
+
+    if (avctx->rc_buffer_size)
+        hrd_buffer_size = avctx->rc_buffer_size;
+    else
+        hrd_buffer_size = avctx->bit_rate;
+    if (avctx->rc_initial_buffer_occupancy)
+        hrd_initial_buffer_fullness = avctx->rc_initial_buffer_occupancy;
+    else
+        hrd_initial_buffer_fullness = hrd_buffer_size * 3 / 4;
+
+    ctx->rc_params.misc.type = VAEncMiscParameterTypeRateControl;
+    ctx->rc_params.rc = (VAEncMiscParameterRateControl) {
+        .bits_per_second   = avctx->bit_rate,
+        .target_percentage = 66,
+        .window_size       = 1000,
+        .initial_qp        = (avctx->qmax >= 0 ? avctx->qmax : 40),
+        .min_qp            = (avctx->qmin >= 0 ? avctx->qmin : 18),
+        .basic_unit_size   = 0,
+    };
+    ctx->global_params[ctx->nb_global_params] =
+        &ctx->rc_params.misc;
+    ctx->global_params_size[ctx->nb_global_params++] =
+        sizeof(ctx->rc_params);
+
+    ctx->hrd_params.misc.type = VAEncMiscParameterTypeHRD;
+    ctx->hrd_params.hrd = (VAEncMiscParameterHRD) {
+        .initial_buffer_fullness = hrd_initial_buffer_fullness,
+        .buffer_size             = hrd_buffer_size,
+    };
+    ctx->global_params[ctx->nb_global_params] =
+        &ctx->hrd_params.misc;
+    ctx->global_params_size[ctx->nb_global_params++] =
+        sizeof(ctx->hrd_params);
+
+    return 0;
+}
+
 static void vaapi_encode_free_output_buffer(void *opaque,
                                             uint8_t *data)
 {
@@ -1067,7 +1138,7 @@ static AVBufferRef *vaapi_encode_alloc_output_buffer(void *opaque,
     // bound on that.
     vas = vaCreateBuffer(ctx->hwctx->display, ctx->va_context,
                          VAEncCodedBufferType,
-                         3 * ctx->aligned_width * ctx->aligned_height +
+                         3 * ctx->surface_width * ctx->surface_height +
                          (1 << 16), 1, 0, &buffer_id);
     if (vas != VA_STATUS_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to create bitstream "
@@ -1089,69 +1160,14 @@ static AVBufferRef *vaapi_encode_alloc_output_buffer(void *opaque,
     return ref;
 }
 
-av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
-                                 const VAAPIEncodeType *type)
+static av_cold int vaapi_encode_create_recon_frames(AVCodecContext *avctx)
 {
     VAAPIEncodeContext *ctx = avctx->priv_data;
-    AVVAAPIFramesContext *recon_hwctx = NULL;
     AVVAAPIHWConfig *hwconfig = NULL;
     AVHWFramesConstraints *constraints = NULL;
     enum AVPixelFormat recon_format;
-    VAStatus vas;
     int err, i;
 
-    if (!avctx->hw_frames_ctx) {
-        av_log(avctx, AV_LOG_ERROR, "A hardware frames reference is "
-               "required to associate the encoding device.\n");
-        return AVERROR(EINVAL);
-    }
-
-    ctx->codec = type;
-    ctx->codec_options = ctx->codec_options_data;
-
-    ctx->va_config  = VA_INVALID_ID;
-    ctx->va_context = VA_INVALID_ID;
-
-    ctx->priv_data = av_mallocz(type->priv_data_size);
-    if (!ctx->priv_data) {
-        err = AVERROR(ENOMEM);
-        goto fail;
-    }
-
-    ctx->input_frames_ref = av_buffer_ref(avctx->hw_frames_ctx);
-    if (!ctx->input_frames_ref) {
-        err = AVERROR(ENOMEM);
-        goto fail;
-    }
-    ctx->input_frames = (AVHWFramesContext*)ctx->input_frames_ref->data;
-
-    ctx->device_ref = av_buffer_ref(ctx->input_frames->device_ref);
-    if (!ctx->device_ref) {
-        err = AVERROR(ENOMEM);
-        goto fail;
-    }
-    ctx->device = (AVHWDeviceContext*)ctx->device_ref->data;
-    ctx->hwctx = ctx->device->hwctx;
-
-    err = ctx->codec->init(avctx);
-    if (err < 0)
-        goto fail;
-
-    err = vaapi_encode_check_config(avctx);
-    if (err < 0)
-        goto fail;
-
-    vas = vaCreateConfig(ctx->hwctx->display,
-                         ctx->va_profile, ctx->va_entrypoint,
-                         ctx->config_attributes, ctx->nb_config_attributes,
-                         &ctx->va_config);
-    if (vas != VA_STATUS_SUCCESS) {
-        av_log(avctx, AV_LOG_ERROR, "Failed to create encode pipeline "
-               "configuration: %d (%s).\n", vas, vaErrorStr(vas));
-        err = AVERROR(EIO);
-        goto fail;
-    }
-
     hwconfig = av_hwdevice_hwconfig_alloc(ctx->device_ref);
     if (!hwconfig) {
         err = AVERROR(ENOMEM);
@@ -1190,13 +1206,13 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
     av_log(avctx, AV_LOG_DEBUG, "Using %s as format of "
            "reconstructed frames.\n", av_get_pix_fmt_name(recon_format));
 
-    if (ctx->aligned_width  < constraints->min_width  ||
-        ctx->aligned_height < constraints->min_height ||
-        ctx->aligned_width  > constraints->max_width ||
-        ctx->aligned_height > constraints->max_height) {
+    if (ctx->surface_width  < constraints->min_width  ||
+        ctx->surface_height < constraints->min_height ||
+        ctx->surface_width  > constraints->max_width ||
+        ctx->surface_height > constraints->max_height) {
         av_log(avctx, AV_LOG_ERROR, "Hardware does not support encoding at "
                "size %dx%d (constraints: width %d-%d height %d-%d).\n",
-               ctx->aligned_width, ctx->aligned_height,
+               ctx->surface_width, ctx->surface_height,
                constraints->min_width,  constraints->max_width,
                constraints->min_height, constraints->max_height);
         err = AVERROR(EINVAL);
@@ -1215,9 +1231,10 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
 
     ctx->recon_frames->format    = AV_PIX_FMT_VAAPI;
     ctx->recon_frames->sw_format = recon_format;
-    ctx->recon_frames->width     = ctx->aligned_width;
-    ctx->recon_frames->height    = ctx->aligned_height;
-    ctx->recon_frames->initial_pool_size = ctx->nb_recon_frames;
+    ctx->recon_frames->width     = ctx->surface_width;
+    ctx->recon_frames->height    = ctx->surface_height;
+    ctx->recon_frames->initial_pool_size =
+        avctx->max_b_frames + 3;
 
     err = av_hwframe_ctx_init(ctx->recon_frames_ref);
     if (err < 0) {
@@ -1225,10 +1242,75 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
                "frame context: %d.\n", err);
         goto fail;
     }
-    recon_hwctx = ctx->recon_frames->hwctx;
 
+    err = 0;
+  fail:
+    av_freep(&hwconfig);
+    av_hwframe_constraints_free(&constraints);
+    return err;
+}
+
+av_cold int ff_vaapi_encode_init(AVCodecContext *avctx)
+{
+    VAAPIEncodeContext *ctx = avctx->priv_data;
+    AVVAAPIFramesContext *recon_hwctx = NULL;
+    VAStatus vas;
+    int err;
+
+    if (!avctx->hw_frames_ctx) {
+        av_log(avctx, AV_LOG_ERROR, "A hardware frames reference is "
+               "required to associate the encoding device.\n");
+        return AVERROR(EINVAL);
+    }
+
+    ctx->codec_options = ctx->codec_options_data;
+
+    ctx->va_config  = VA_INVALID_ID;
+    ctx->va_context = VA_INVALID_ID;
+
+    ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
+    if (!ctx->priv_data) {
+        err = AVERROR(ENOMEM);
+        goto fail;
+    }
+
+    ctx->input_frames_ref = av_buffer_ref(avctx->hw_frames_ctx);
+    if (!ctx->input_frames_ref) {
+        err = AVERROR(ENOMEM);
+        goto fail;
+    }
+    ctx->input_frames = (AVHWFramesContext*)ctx->input_frames_ref->data;
+
+    ctx->device_ref = av_buffer_ref(ctx->input_frames->device_ref);
+    if (!ctx->device_ref) {
+        err = AVERROR(ENOMEM);
+        goto fail;
+    }
+    ctx->device = (AVHWDeviceContext*)ctx->device_ref->data;
+    ctx->hwctx = ctx->device->hwctx;
+
+    err = vaapi_encode_config_attributes(avctx);
+    if (err < 0)
+        goto fail;
+
+    vas = vaCreateConfig(ctx->hwctx->display,
+                         ctx->va_profile, ctx->va_entrypoint,
+                         ctx->config_attributes, ctx->nb_config_attributes,
+                         &ctx->va_config);
+    if (vas != VA_STATUS_SUCCESS) {
+        av_log(avctx, AV_LOG_ERROR, "Failed to create encode pipeline "
+               "configuration: %d (%s).\n", vas, vaErrorStr(vas));
+        err = AVERROR(EIO);
+        goto fail;
+    }
+
+    err = vaapi_encode_create_recon_frames(avctx);
+    if (err < 0)
+        goto fail;
+
+    recon_hwctx = ctx->recon_frames->hwctx;
     vas = vaCreateContext(ctx->hwctx->display, ctx->va_config,
-                          ctx->aligned_width, ctx->aligned_height,
+                          ctx->surface_width, ctx->surface_height,
                           VA_PROGRESSIVE,
                           recon_hwctx->surface_ids,
                           recon_hwctx->nb_surfaces,
@@ -1240,6 +1322,26 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
         goto fail;
     }
 
+    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;
+    }
+
+    if (ctx->va_rc_mode & ~VA_RC_CQP) {
+        err = vaapi_encode_init_rate_control(avctx);
+        if (err < 0)
+            goto fail;
+    }
+
+    if (ctx->codec->configure) {
+        err = ctx->codec->configure(avctx);
+        if (err < 0)
+            goto fail;
+    }
+
     ctx->input_order  = 0;
     ctx->output_delay = avctx->max_b_frames;
     ctx->decode_delay = 1;
@@ -1271,14 +1373,6 @@ 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) /
@@ -1292,8 +1386,6 @@ av_cold int ff_vaapi_encode_init(AVCodecContext *avctx,
     return 0;
 
 fail:
-    av_freep(&hwconfig);
-    av_hwframe_constraints_free(&constraints);
     ff_vaapi_encode_close(avctx);
     return err;
 }
@@ -1318,9 +1410,6 @@ av_cold int ff_vaapi_encode_close(AVCodecContext *avctx)
         ctx->va_config = VA_INVALID_ID;
     }
 
-    if (ctx->codec->close)
-        ctx->codec->close(avctx);
-
     av_buffer_pool_uninit(&ctx->output_buffer_pool);
 
     av_freep(&ctx->codec_sequence_params);