From 02a6e9a2735a88e212716f3c5e414db0dcf5b597 Mon Sep 17 00:00:00 2001
From: Sreerenj Balachandran <sreerenj.balachandran@intel.com>
Date: Wed, 5 Apr 2017 14:48:46 -0700
Subject: [PATCH] libs: encoder: h264: insert AU delimiter

Insert an AUD as the first NAL of each encoded frame.

Some applications require Access Unit Delimiter for decoding the
stream.

The AU delimeter insertion is done only when the aud parameter is
TRUE (by default is disabled). The reason of this it is because this
header is only available from Intel Gen9 and the VA intel driver
should be 1.8 or superior. Otherwise, the output will be corrupted.

https://bugzilla.gnome.org/show_bug.cgi?id=776712

Signed-off-by: Victor Jaquez <vjaquez@igalia.com>
---
 gst-libs/gst/vaapi/gstvaapiencoder_h264.c | 75 +++++++++++++++++++++++
 gst-libs/gst/vaapi/gstvaapiencoder_h264.h |  1 +
 2 files changed, 76 insertions(+)

diff --git a/gst-libs/gst/vaapi/gstvaapiencoder_h264.c b/gst-libs/gst/vaapi/gstvaapiencoder_h264.c
index fda0c31e72..c5e33fd1a9 100644
--- a/gst-libs/gst/vaapi/gstvaapiencoder_h264.c
+++ b/gst-libs/gst/vaapi/gstvaapiencoder_h264.c
@@ -715,6 +715,8 @@ struct _GstVaapiEncoderH264
   guint16 view_ids[MAX_NUM_VIEWS];
   GstVaapiH264ViewRefPool ref_pools[MAX_NUM_VIEWS];
   GstVaapiH264ViewReorderPool reorder_pools[MAX_NUM_VIEWS];
+
+  gboolean use_aud;
 };
 
 /* Write a SEI buffering period payload */
@@ -1293,6 +1295,52 @@ fill_hrd_params (GstVaapiEncoderH264 * encoder, VAEncMiscParameterHRD * hrd)
   }
 }
 
+static gboolean
+add_packed_au_delimiter (GstVaapiEncoderH264 * encoder,
+    GstVaapiEncPicture * picture)
+{
+  GstVaapiEncPackedHeader *packed_aud;
+  GstBitWriter bs;
+  VAEncPackedHeaderParameterBuffer packed_header_param_buffer = { 0 };
+  guint32 data_bit_size;
+  guint8 *data;
+
+  gst_bit_writer_init (&bs, 128 * 8);
+  WRITE_UINT32 (&bs, 0x00000001, 32);   /* start code */
+  bs_write_nal_header (&bs, GST_H264_NAL_REF_IDC_NONE,
+      GST_H264_NAL_AU_DELIMITER);
+  WRITE_UINT32 (&bs, picture->type - 1, 3);
+  if (!bs_write_trailing_bits (&bs))
+    goto bs_error;
+
+  g_assert (GST_BIT_WRITER_BIT_SIZE (&bs) % 8 == 0);
+  data_bit_size = GST_BIT_WRITER_BIT_SIZE (&bs);
+  data = GST_BIT_WRITER_DATA (&bs);
+
+  packed_header_param_buffer.type = VAEncPackedHeaderRawData;
+  packed_header_param_buffer.bit_length = data_bit_size;
+  packed_header_param_buffer.has_emulation_bytes = 0;
+
+  packed_aud = gst_vaapi_enc_packed_header_new (GST_VAAPI_ENCODER (encoder),
+      &packed_header_param_buffer, sizeof (packed_header_param_buffer),
+      data, (data_bit_size + 7) / 8);
+  g_assert (packed_aud);
+
+  gst_vaapi_enc_picture_add_packed_header (picture, packed_aud);
+  gst_vaapi_codec_object_replace (&packed_aud, NULL);
+
+  gst_bit_writer_clear (&bs, TRUE);
+  return TRUE;
+
+  /* ERRORS */
+bs_error:
+  {
+    GST_WARNING ("failed to write AU Delimiter  NAL unit");
+    gst_bit_writer_clear (&bs, TRUE);
+    return FALSE;
+  }
+}
+
 /* Adds the supplied sequence header (SPS) to the list of packed
    headers to pass down as-is to the encoder */
 static gboolean
@@ -2120,6 +2168,13 @@ ensure_sequence (GstVaapiEncoderH264 * encoder, GstVaapiEncPicture * picture)
 {
   GstVaapiEncSequence *sequence = NULL;
 
+  /* Insert an AU delimiter */
+  if ((GST_VAAPI_ENCODER_PACKED_HEADERS (encoder) &
+          VA_ENC_PACKED_HEADER_RAW_DATA) && encoder->use_aud) {
+    if (!add_packed_au_delimiter (encoder, picture))
+      goto error_create_packed_au_delimiter;
+  }
+
   /* submit an SPS header before every new I-frame, if codec config changed */
   if (!encoder->config_changed || picture->type != GST_VAAPI_PICTURE_TYPE_I)
     return TRUE;
@@ -2157,6 +2212,11 @@ error_create_seq_param:
     gst_vaapi_codec_object_replace (&sequence, NULL);
     return FALSE;
   }
+error_create_packed_au_delimiter:
+  {
+    GST_ERROR ("failed to create AU delimiter");
+    gst_vaapi_codec_object_replace (&sequence, NULL);
+  }
 error_create_packed_seq_hdr:
   {
     GST_ERROR ("failed to create packed sequence header buffer");
@@ -2983,6 +3043,9 @@ gst_vaapi_encoder_h264_set_property (GstVaapiEncoder * base_encoder,
       }
       break;
     }
+    case GST_VAAPI_ENCODER_H264_PROP_AUD:
+      encoder->use_aud = g_value_get_boolean (value);
+      break;
     default:
       return GST_VAAPI_ENCODER_STATUS_ERROR_INVALID_PARAMETER;
   }
@@ -3134,6 +3197,7 @@ gst_vaapi_encoder_h264_get_default_properties (void)
           "Number of Views",
           "Number of Views for MVC encoding",
           1, MAX_NUM_VIEWS, 1, G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS));
+
   /**
    * GstVaapiEncoderH264:view-ids:
    *
@@ -3148,6 +3212,17 @@ gst_vaapi_encoder_h264_get_default_properties (void)
               G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS),
           G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS));
 
+  /**
+   * GstVaapiEncoderH264:aud:
+   *
+   * Use AU (Access Unit) delimeter.
+   */
+  GST_VAAPI_ENCODER_PROPERTIES_APPEND (props,
+      GST_VAAPI_ENCODER_H264_PROP_AUD,
+      g_param_spec_boolean ("aud", "AU delimiter",
+          "Use AU (Access Unit) delimeter", FALSE,
+          G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS));
+
   return props;
 }
 
diff --git a/gst-libs/gst/vaapi/gstvaapiencoder_h264.h b/gst-libs/gst/vaapi/gstvaapiencoder_h264.h
index 14e35efd6e..5019b6cf4e 100644
--- a/gst-libs/gst/vaapi/gstvaapiencoder_h264.h
+++ b/gst-libs/gst/vaapi/gstvaapiencoder_h264.h
@@ -61,6 +61,7 @@ typedef enum {
   GST_VAAPI_ENCODER_H264_PROP_CPB_LENGTH = -7,
   GST_VAAPI_ENCODER_H264_PROP_NUM_VIEWS = -8,
   GST_VAAPI_ENCODER_H264_PROP_VIEW_IDS = -9,
+  GST_VAAPI_ENCODER_H264_PROP_AUD = -10,
 } GstVaapiEncoderH264Prop;
 
 GstVaapiEncoder *