Record info about coded LCUs

Adds field lcu_stats to encoder_state_config_frame_t. The following data
is recorded for each LCU:
    - number of bits
    - squared cost
    - used lambda value
    - alpha parameter used for rate control
    - beta parameter used for rate control

src/encoder_state-ctors_dtors.c

@@ -53,11 +53,16 @@ static int encoder_state_config_frame_init(encoder_state_t * const state) {
   state->frame->rc_alpha = 3.2003;
   state->frame->rc_beta = -1.367;
 
+  const encoder_control_t * const encoder = state->encoder_control;
+  const int num_lcus = encoder->in.width_in_lcu * encoder->in.height_in_lcu;
+  state->frame->lcu_stats = MALLOC(lcu_stats_t, num_lcus);
+
   return 1;
 }
 
 static void encoder_state_config_frame_finalize(encoder_state_t * const state) {
   kvz_image_list_destroy(state->frame->ref);
+  FREE_POINTER(state->frame->lcu_stats);
 }
 
 static int encoder_state_config_tile_init(encoder_state_t * const state, 

src/encoderstate.c

@@ -202,7 +202,7 @@ static void encoder_state_worker_encode_lcu(void * opaque) {
   const encoder_control_t * const encoder = state->encoder_control;
   videoframe_t* const frame = state->tile->frame;
 
-  kvz_set_lcu_lambda_and_qp(state);
+  kvz_set_lcu_lambda_and_qp(state, lcu->position);
 
   //This part doesn't write to bitstream, it's only search, deblock and sao
   
@@ -241,6 +241,7 @@ static void encoder_state_worker_encode_lcu(void * opaque) {
   }
   
   //Now write data to bitstream (required to have a correct CABAC state)
+  const uint64_t existing_bits = kvz_bitstream_tell(&state->stream);
   
   //First LCU, and we are in a slice. We need a slice header
   if (state->type == ENCODER_STATE_TYPE_SLICE && lcu->index == 0) {
@@ -267,6 +268,9 @@ static void encoder_state_worker_encode_lcu(void * opaque) {
     kvz_cabac_encode_bin_trm(&state->cabac, 0);  // end_of_slice_segment_flag
   }
 
+  const uint32_t bits = kvz_bitstream_tell(&state->stream) - existing_bits;
+  kvz_get_lcu_stats(state, lcu->position.x, lcu->position.y)->bits = bits;
+
   //Wavefronts need the context to be copied to the next row
   if (state->type == ENCODER_STATE_TYPE_WAVEFRONT_ROW && lcu->index == 1) {
     int j;
@@ -860,6 +864,22 @@ static void encoder_state_init_children(encoder_state_t * const state) {
   }
 }
 
+static void normalize_lcu_weights(encoder_state_t * const state)
+{
+  if (state->frame->num == 0) return;
+
+  const uint32_t num_lcus = state->encoder_control->in.width_in_lcu *
+                            state->encoder_control->in.height_in_lcu;
+  double sum = 0.0;
+  for (uint32_t i = 0; i < num_lcus; i++) {
+    sum += state->frame->lcu_stats[i].weight;
+  }
+
+  for (uint32_t i = 0; i < num_lcus; i++) {
+    state->frame->lcu_stats[i].weight /= sum;
+  }
+}
+
 static void encoder_state_init_new_frame(encoder_state_t * const state, kvz_picture* frame) {
   assert(state->type == ENCODER_STATE_TYPE_MAIN);
 
@@ -908,6 +928,7 @@ static void encoder_state_init_new_frame(encoder_state_t * const state, kvz_pict
   encoder_state_remove_refs(state);
   encoder_state_ref_sort(state);
 
+  normalize_lcu_weights(state);
   kvz_set_picture_lambda_and_qp(state);
 
   encoder_state_init_children(state);
@@ -1045,3 +1066,11 @@ coeff_scan_order_t kvz_get_scan_order(int8_t cu_type, int intra_mode, int depth)
 
   return SCAN_DIAG;
 }
+
+lcu_stats_t* kvz_get_lcu_stats(encoder_state_t *state, int lcu_x, int lcu_y)
+{
+  const int index = lcu_x + state->tile->lcu_offset_x +
+                    (lcu_y + state->tile->lcu_offset_y) *
+                    state->encoder_control->in.width_in_lcu;
+  return &state->frame->lcu_stats[index];
+}

src/encoderstate.h

@@ -49,6 +49,23 @@ typedef enum {
 } encoder_state_type;
 
 
+typedef struct lcu_stats_t {
+  //! \brief Number of bits that were spent
+  uint32_t bits;
+
+  //! \brief Weight of the LCU for rate control
+  double weight;
+
+  //! \brief Lambda value which was used for this LCU
+  double lambda;
+
+  //! \brief Rate control alpha parameter
+  double rc_alpha;
+
+  //! \brief Rate control beta parameter
+  double rc_beta;
+} lcu_stats_t;
+
 
 typedef struct encoder_state_config_frame_t {
   /**
@@ -114,6 +131,13 @@ typedef struct encoder_state_config_frame_t {
    */
   bool done;
 
+  /**
+   * \brief Information about the coded LCUs.
+   *
+   * Used for rate control.
+   */
+  lcu_stats_t *lcu_stats;
+
 } encoder_state_config_frame_t;
 
 typedef struct encoder_state_config_tile_t {
@@ -249,6 +273,8 @@ void kvz_encoder_get_ref_lists(const encoder_state_t *const state,
                                int ref_list_len_out[2],
                                int ref_list_poc_out[2][16]);
 
+lcu_stats_t* kvz_get_lcu_stats(encoder_state_t *state, int lcu_x, int lcu_y);
+
 static const uint8_t g_group_idx[32] = {
   0, 1, 2, 3, 4, 4, 5, 5, 6, 6,
   6, 6, 7, 7, 7, 7, 8, 8, 8, 8,

src/rate_control.c

@@ -190,9 +190,15 @@ void kvz_set_picture_lambda_and_qp(encoder_state_t * const state)
   }
 }
 
-void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state)
+void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state,
+                               vector2d_t pos)
 {
   state->lambda      = state->frame->lambda;
   state->lambda_sqrt = sqrt(state->frame->lambda);
   state->qp          = state->frame->QP;
+
+  lcu_stats_t *lcu_stats = kvz_get_lcu_stats(state, pos.x, pos.y);
+  lcu_stats->lambda      = state->lambda;
+  lcu_stats->rc_alpha    = state->frame->rc_alpha;
+  lcu_stats->rc_beta     = state->frame->rc_beta;
 }

src/rate_control.h

@@ -32,6 +32,7 @@
 
 void kvz_set_picture_lambda_and_qp(encoder_state_t * const state);
 
-void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state);
+void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state,
+                               vector2d_t pos);
 
 #endif // RATE_CONTROL_H_

src/search.c

@@ -949,7 +949,10 @@ void kvz_search_lcu(encoder_state_t * const state, const int x, const int y, con
   }
 
   // Start search from depth 0.
-  search_cu(state, x, y, 0, work_tree);
+  double cost = search_cu(state, x, y, 0, work_tree);
+
+  // Save squared cost for rate control.
+  kvz_get_lcu_stats(state, x / LCU_WIDTH, y / LCU_WIDTH)->weight = cost * cost;
 
   // The best decisions through out the LCU got propagated back to depth 0,
   // so copy those back to the frame.