[PATCH RFC 10/11] media: v4l2-core: Add a peak bitrate to the encode rate control

[Arun Iyengar via B4 Relay <devnull+aiyengar1990.gmail.com@kernel.org>]

From: Arun Iyengar <aiyengar1990@gmail.com>

The stateless encode cores read V4L2_CID_MPEG_VIDEO_BITRATE_PEAK for
V4L2_MPEG_VIDEO_BITRATE_MODE_VBR, but their rate control only has a
constant bitrate model.

Give the model a peak bitrate, which makes it a variable bitrate one: the
bitrate is met on average over the stream, and the peak bitrate bounds
every second of it.

- The deviation of the coded sizes from the GOP plan is spread over two
  seconds of pictures instead of 8 pictures, so that the picture sizes
  can follow the content for a while.
- The size target of a picture is limited to what is left of the peak in
  the last second (the sizes of up to 128 pictures are kept for it), and
  its QP is raised by the QP steps between its expected size and that: a
  P picture is expected at the size of the last P picture, an I picture
  at the larger of the last I picture and the last P picture times the
  measured I/P size ratio, each at the QP difference. The QP state counts
  the raise as part of the QP error of the picture.

A picture the history cannot foresee (an I picture at a cut to much
harder content) can still take the second over the peak; re-encoding is
up to the driver.

In a host simulation against a model encoder (sizes halving every 6 QP;
720p and 1080p, 256 kbit/s to 20 Mbit/s, peak 1.5 times the bitrate, IDR
periods 1, 10, 30 and 120, a 3x and 0.33x change of the content), the
bitrate is within 0.3% of the target and every second within the peak
wherever QP 51 can reach the target. With the content alternating
between two complexities 4x apart every 3 seconds, the largest second is
0.79-1.0 times the peak, or 1.24-1.26 times where the cut falls on an I
picture. The constant bitrate mode is not changed.

Signed-off-by: Arun Iyengar <aiyengar1990@gmail.com>
---
 drivers/media/v4l2-core/v4l2-enc-cbr.c | 116 +++++++++++++++++++++++++++++++--
 include/media/v4l2-enc-cbr.h           |  24 ++++++-
 2 files changed, 135 insertions(+), 5 deletions(-)

diff --git a/drivers/media/v4l2-core/v4l2-enc-cbr.c b/drivers/media/v4l2-core/v4l2-enc-cbr.c
index 87db60469182..c9e806bbabcb 100644
--- a/drivers/media/v4l2-core/v4l2-enc-cbr.c
+++ b/drivers/media/v4l2-core/v4l2-enc-cbr.c
@@ -30,6 +30,13 @@
  * The QP to size relation (6 QP steps per halving) is the same for H.264
  * and HEVC, so the model does not depend on the codec. Only the QP delta of
  * the I pictures is a parameter.
+ *
+ * With a peak bitrate (v4l2_enc_cbr_set_peak()), the same model gives a
+ * variable bitrate: the bitrate is the average over the stream and the
+ * peak bitrate a bound over any second of it. The deviation from the plan
+ * is spread over two seconds instead of a few pictures, so that the size of
+ * the pictures can follow the content, and a picture that would take the
+ * last second over the peak gets a lower size target and a higher QP.
  */
 
 #include <linux/export.h>
@@ -44,6 +51,8 @@
 
 /* The accumulated size error is spread over this many pictures. */
 #define CBR_WINDOW		8
+/* With a peak bitrate, over this many seconds of pictures */
+#define VBR_WINDOW_SECONDS	2
 /* Limit of the accumulated size error, in seconds of the bitrate */
 #define CBR_BUFFER_SECONDS	3
 /* I/P size ratio, 1/16 units: initial value (10x) and limits (1.5x .. 30x) */
@@ -181,6 +190,76 @@ static int cbr_qp_initial(u64 bits, unsigned int samples)
 	return clamp(DIV_ROUND_CLOSEST(qp_q8, 256), 15, 45);
 }
 
+/* Sum of the sizes of the last @count pictures */
+static u64 vbr_window_sum(const struct v4l2_enc_cbr *cbr, unsigned int count)
+{
+	unsigned int i, index = cbr->peak_index;
+	u64 sum = 0;
+
+	count = min(count, cbr->peak_count);
+	for (i = 0; i < count; i++) {
+		index = (index + V4L2_ENC_CBR_PEAK_PICS - 1) %
+			V4L2_ENC_CBR_PEAK_PICS;
+		sum += cbr->peak_bits[index];
+	}
+
+	return sum;
+}
+
+/*
+ * QP steps (1/256 units) that bring a picture of @bits at @qp, coded at
+ * @qp_pic instead, down to @room: 6 QP steps per halving.
+ */
+static int vbr_qp_add(u64 bits, int qp, int qp_pic, u64 room)
+{
+	if (!bits)
+		return 0;
+
+	return 6 * (cbr_log2_q8(bits) - cbr_log2_q8(room)) -
+	       (qp_pic - qp) * 256;
+}
+
+/*
+ * Keep the last second (at most V4L2_ENC_CBR_PEAK_PICS pictures) within the
+ * peak bitrate: limit the size target of the picture to what is left of the
+ * peak, and raise its QP by the QP steps between its expected size and
+ * that. A P picture is expected at the size of the last one, an I picture
+ * at the larger of the last I picture and the last P picture times the I/P
+ * size ratio (which follows a change of the content sooner), each at the
+ * QP difference.
+ */
+static void vbr_bound(struct v4l2_enc_cbr *cbr,
+		      const struct v4l2_fract *timeperframe, unsigned int fps,
+		      int qp_max)
+{
+	unsigned int pics = clamp(fps, 1U, V4L2_ENC_CBR_PEAK_PICS);
+	u64 peak = cbr_nominal(max(cbr->bitrate_peak, cbr->bitrate),
+			       timeperframe) * pics;
+	u64 used = vbr_window_sum(cbr, pics - 1);
+	u64 room = max_t(u64, peak > used ? peak - used : 0, 8);
+	int add;
+
+	cbr->target_bits = min(cbr->target_bits, room);
+
+	if (cbr->intra)
+		add = max(vbr_qp_add(cbr->i_bits, cbr->i_qp, cbr->qp_pic, room),
+			  vbr_qp_add(div_u64(cbr->p_last_bits * cbr->ratio, 16),
+				     cbr->p_qp - cbr->qp_delta_intra,
+				     cbr->qp_pic, room));
+	else
+		add = vbr_qp_add(cbr->p_last_bits, cbr->p_qp, cbr->qp_pic,
+				 room);
+
+	add = min(DIV_ROUND_UP(add, 256), qp_max - cbr->qp_pic);
+	if (add <= 0)
+		return;
+
+	cbr->qp_pic += add;
+	cbr->qp_peak_add = add;
+
+	pr_debug("  peak: %llu bits left of %llu, QP +%d", room, peak, add);
+}
+
 /**
  * v4l2_enc_cbr_reset() - reset the rate control state
  * @cbr: the rate control state
@@ -195,6 +274,22 @@ void v4l2_enc_cbr_reset(struct v4l2_enc_cbr *cbr, int qp_delta_intra)
 }
 EXPORT_SYMBOL_GPL(v4l2_enc_cbr_reset);
 
+/**
+ * v4l2_enc_cbr_set_peak() - set the peak bitrate
+ * @cbr: the rate control state
+ * @bitrate_peak: peak bitrate in bits per second, 0 for a constant bitrate
+ *
+ * A non-zero peak bitrate makes the rate control a variable bitrate one:
+ * the bitrate of v4l2_enc_cbr_begin() is met on average, and no second of
+ * the stream goes over the peak bitrate (a peak below the bitrate counts as
+ * the bitrate). Takes effect from the next v4l2_enc_cbr_begin().
+ */
+void v4l2_enc_cbr_set_peak(struct v4l2_enc_cbr *cbr, u32 bitrate_peak)
+{
+	cbr->bitrate_peak = bitrate_peak;
+}
+EXPORT_SYMBOL_GPL(v4l2_enc_cbr_set_peak);
+
 /**
  * v4l2_enc_cbr_begin() - set up the size target of the next picture
  * @cbr: the rate control state
@@ -284,12 +379,17 @@ int v4l2_enc_cbr_begin(struct v4l2_enc_cbr *cbr, u32 bitrate,
 
 	target = (s64)cbr->plan_bits -
 		 div_s64(cbr->buffer_bits - max_t(s64, expected, 0),
-			 CBR_WINDOW);
+			 cbr->bitrate_peak ? VBR_WINDOW_SECONDS * fps :
+					     CBR_WINDOW);
 	cbr->target_bits = clamp_t(s64, target,
 				   max_t(u64, cbr->plan_bits / 4, 8),
 				   cbr->plan_bits * 2);
 	cbr->intra = intra;
 	cbr->qp_reenc = 0;
+	cbr->qp_peak_add = 0;
+
+	if (cbr->bitrate_peak)
+		vbr_bound(cbr, timeperframe, fps, qp_max);
 
 	pr_debug("+ v4l2-enc-cbr: begin");
 	pr_debug("  nominal: %llu bits, plan: %llu bits, target: %llu bits",
@@ -408,12 +508,17 @@ void v4l2_enc_cbr_end(struct v4l2_enc_cbr *cbr, u64 bits,
 			   div64_u64(cbr->window_bits,
 				     CBR_BUFFER_SECONDS * fps);
 
+	cbr->peak_bits[cbr->peak_index] = min_t(u64, bits, U32_MAX);
+	cbr->peak_index = (cbr->peak_index + 1) % V4L2_ENC_CBR_PEAK_PICS;
+	cbr->peak_count = min(cbr->peak_count + 1, V4L2_ENC_CBR_PEAK_PICS);
+
 	/*
-	 * QP error of the picture, 1/256 units. A re-encoded picture was coded
-	 * above the QP of the state: its error at the state counts that too.
+	 * QP error of the picture, 1/256 units. A picture coded above the QP
+	 * of the state (re-encoded, or raised for the peak bitrate) counts
+	 * that too: the state would have given a picture that much larger.
 	 */
 	err = 6 * (cbr_log2_q8(bits) - cbr_log2_q8(cbr->target_bits)) +
-	      cbr->qp_reenc * 256;
+	      (cbr->qp_reenc + cbr->qp_peak_add) * 256;
 
 	step = err / 8;
 	if (err > CBR_ERR_SLOW)
@@ -425,6 +530,8 @@ void v4l2_enc_cbr_end(struct v4l2_enc_cbr *cbr, u64 bits,
 	if (!cbr->intra) {
 		cbr->qp_q8 += step;
 		cbr->p_bits = cbr->p_bits ? (7 * cbr->p_bits + bits) / 8 : bits;
+		cbr->p_last_bits = bits;
+		cbr->p_qp = clamp(cbr->qp_pic, qp_min, qp_max);
 		cbr->since_i++;
 
 		/* I/P size ratio, once a few P pictures were coded */
@@ -439,6 +546,7 @@ void v4l2_enc_cbr_end(struct v4l2_enc_cbr *cbr, u64 bits,
 		 * estimate: only all-intra streams steer the QP fully with it.
 		 */
 		cbr->i_bits = bits;
+		cbr->i_qp = clamp(cbr->qp_pic, qp_min, qp_max);
 
 		/*
 		 * Above its share of the GOP at the largest I/P ratio, an I
diff --git a/include/media/v4l2-enc-cbr.h b/include/media/v4l2-enc-cbr.h
index 867e176ecb36..4f4874c41d2a 100644
--- a/include/media/v4l2-enc-cbr.h
+++ b/include/media/v4l2-enc-cbr.h
@@ -12,6 +12,9 @@
 #include <linux/types.h>
 #include <linux/videodev2.h>
 
+/* Pictures of the peak bitrate window: one second, up to 128 pictures. */
+#define V4L2_ENC_CBR_PEAK_PICS	128U
+
 /**
  * struct v4l2_enc_cbr - constant bitrate rate control state
  * @qp_delta_intra: QP of the I pictures below the QP of the P pictures
@@ -37,8 +40,17 @@
  *		 decays by 1/(3 * fps) per picture)
  * @plan_bits: planned size of the current picture
  * @target_bits: size target of the current picture
+ * @bitrate_peak: peak bitrate in bits per second, 0 for a constant bitrate
+ * @qp_peak_add: QP increase of the current picture for the peak bitrate
+ * @i_qp: QP of the last I picture
+ * @p_last_bits: size of the last P picture
+ * @p_qp: QP of the last P picture
+ * @peak_bits: sizes of the last pictures, for the peak bitrate
+ * @peak_index: next entry of @peak_bits
+ * @peak_count: valid entries of @peak_bits
  *
- * Only @qp_delta_intra is set by the user, after v4l2_enc_cbr_reset().
+ * Only @qp_delta_intra is set by the user, after v4l2_enc_cbr_reset(), and
+ * @bitrate_peak by v4l2_enc_cbr_set_peak().
  * @target_bits is valid from v4l2_enc_cbr_begin() of a picture until its
  * v4l2_enc_cbr_end().
  */
@@ -63,9 +75,19 @@ struct v4l2_enc_cbr {
 	u64 window_bits;
 	u64 plan_bits;
 	u64 target_bits;
+
+	u32 bitrate_peak;
+	int qp_peak_add;
+	int i_qp;
+	u64 p_last_bits;
+	int p_qp;
+	u32 peak_bits[V4L2_ENC_CBR_PEAK_PICS];
+	unsigned int peak_index;
+	unsigned int peak_count;
 };
 
 void v4l2_enc_cbr_reset(struct v4l2_enc_cbr *cbr, int qp_delta_intra);
+void v4l2_enc_cbr_set_peak(struct v4l2_enc_cbr *cbr, u32 bitrate_peak);
 int v4l2_enc_cbr_begin(struct v4l2_enc_cbr *cbr, u32 bitrate,
 		       const struct v4l2_fract *timeperframe,
 		       unsigned int samples, bool intra, int qp_min,

-- 
2.55.0