Add VpxTplGopStats

Contains the size of GOP - also the size of the list of TPL stats for
each frame in this GOP.

VpxTplGopStats will be the unit for VP9E_GET_TPL_STATS control to return
TPL stats from the encoder.

Bug: b/273736974
Change-Id: I1682242fc6db4aafcd6314af023aa0d704976585

7 files changed, 68 insertions, 36 deletions

diff --git a/test/encode_api_test.cc b/test/encode_api_test.cc
index e435ed872..2b0aa1fdf 100644
--- a/test/encode_api_test.cc
+++ b/test/encode_api_test.cc
@@ -384,12 +384,15 @@ class EncodeApiGetTplStatsTest
     }
   }
 
-  vpx_codec_err_t AllocateTplList(VpxTplFrameStats **data) {
-    // Allocate MAX_ARF_GOP_SIZE * sizeof(VpxTplFrameStats) that will be filled
-    // by VP9E_GET_TPL_STATS
-    *data =
+  vpx_codec_err_t AllocateTplList(VpxTplGopStats *data) {
+    // Allocate MAX_ARF_GOP_SIZE (50) * sizeof(VpxTplFrameStats) that will be
+    // filled by VP9E_GET_TPL_STATS.
+    // MAX_ARF_GOP_SIZE is used here because the test doesn't know the size of
+    // each GOP before getting TPL stats from the encoder.
+    data->size = 50;
+    data->frame_stats_list =
         static_cast<VpxTplFrameStats *>(calloc(50, sizeof(VpxTplFrameStats)));
-    if (*data == nullptr) return VPX_CODEC_MEM_ERROR;
+    if (data->frame_stats_list == nullptr) return VPX_CODEC_MEM_ERROR;
     return VPX_CODEC_OK;
   }
 
@@ -398,22 +401,23 @@ class EncodeApiGetTplStatsTest
     while (const vpx_codec_cx_pkt_t *pkt = iter.Next()) {
       switch (pkt->kind) {
         case VPX_CODEC_CX_FRAME_PKT: {
-          VpxTplFrameStats *tpl_stats = NULL;
+          VpxTplGopStats tpl_stats;
           EXPECT_EQ(AllocateTplList(&tpl_stats), VPX_CODEC_OK);
-          encoder->Control(VP9E_GET_TPL_STATS, tpl_stats);
+          encoder->Control(VP9E_GET_TPL_STATS, &tpl_stats);
           bool stats_not_all_zero = false;
-          for (unsigned int i = 0; i < cfg_.g_lag_in_frames; i++) {
-            if (tpl_stats[i].frame_width != 0) {
-              ASSERT_EQ(tpl_stats[i].frame_width, width_);
-              ASSERT_EQ(tpl_stats[i].frame_height, height_);
-              ASSERT_GT(tpl_stats[i].num_blocks, 0);
-              ASSERT_NE(tpl_stats[i].block_stats_list, nullptr);
+          for (int i = 0; i < tpl_stats.size; i++) {
+            VpxTplFrameStats *frame_stats_list = tpl_stats.frame_stats_list;
+            if (frame_stats_list[i].frame_width != 0) {
+              ASSERT_EQ(frame_stats_list[i].frame_width, width_);
+              ASSERT_EQ(frame_stats_list[i].frame_height, height_);
+              ASSERT_GT(frame_stats_list[i].num_blocks, 0);
+              ASSERT_NE(frame_stats_list[i].block_stats_list, nullptr);
               stats_not_all_zero = true;
             }
           }
           ASSERT_TRUE(stats_not_all_zero);
           // Free the memory right away now as this is only a test.
-          free(tpl_stats);
+          free(tpl_stats.frame_stats_list);
           break;
         }
         default: break;
@@ -430,7 +434,7 @@ TEST_P(EncodeApiGetTplStatsTest, GetTplStats) {
   width_ = 352;
   height_ = 288;
   ::libvpx_test::I420VideoSource video("hantro_collage_w352h288.yuv", width_,
-                                       height_, 30, 1, 0, 150);
+                                       height_, 30, 1, 0, 50);
   ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
 }
 
diff --git a/test/encode_test_driver.h b/test/encode_test_driver.h
index a5cd8306e..922c49f42 100644
--- a/test/encode_test_driver.h
+++ b/test/encode_test_driver.h
@@ -154,7 +154,7 @@ class Encoder {
     ASSERT_EQ(VPX_CODEC_OK, res) << EncoderError();
   }
 
-  void Control(int ctrl_id, VpxTplFrameStats *arg) {
+  void Control(int ctrl_id, VpxTplGopStats *arg) {
     const vpx_codec_err_t res = vpx_codec_control_(&encoder_, ctrl_id, arg);
     ASSERT_EQ(VPX_CODEC_OK, res) << EncoderError();
   }
diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index db21509ce..5126a971a 100644
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -2628,7 +2628,6 @@ VP9_COMP *vp9_create_compressor(const VP9EncoderConfig *oxcf,
 #endif  // CONFIG_NON_GREEDY_MV
   for (i = 0; i < MAX_ARF_GOP_SIZE; ++i) {
     cpi->tpl_stats[i].tpl_stats_ptr = NULL;
-    cpi->tpl_frame_stats[i].block_stats_list = NULL;
   }
 
   // Allocate memory to store variances for a frame.
diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index 230a8315b..2528bc231 100644
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -745,7 +745,7 @@ typedef struct VP9_COMP {
   BLOCK_SIZE tpl_bsize;
   TplDepFrame tpl_stats[MAX_ARF_GOP_SIZE];
   // Used to store TPL stats before propagation
-  VpxTplFrameStats tpl_frame_stats[MAX_ARF_GOP_SIZE];
+  VpxTplGopStats tpl_gop_stats;
   YV12_BUFFER_CONFIG *tpl_recon_frames[REF_FRAMES];
   EncFrameBuf enc_frame_buf[REF_FRAMES];
 #if CONFIG_MULTITHREAD
diff --git a/vp9/encoder/vp9_tpl_model.c b/vp9/encoder/vp9_tpl_model.c
index de3783f9a..9f4bafdf8 100644
--- a/vp9/encoder/vp9_tpl_model.c
+++ b/vp9/encoder/vp9_tpl_model.c
@@ -154,17 +154,43 @@ static void init_tpl_stats(VP9_COMP *cpi) {
   int frame_idx;
   for (frame_idx = 0; frame_idx < MAX_ARF_GOP_SIZE; ++frame_idx) {
     TplDepFrame *tpl_frame = &cpi->tpl_stats[frame_idx];
-    VpxTplFrameStats *tpl_frame_stats = &cpi->tpl_frame_stats[frame_idx];
     memset(tpl_frame->tpl_stats_ptr, 0,
            tpl_frame->height * tpl_frame->width *
                sizeof(*tpl_frame->tpl_stats_ptr));
-    memset(tpl_frame_stats->block_stats_list, 0,
-           tpl_frame->height * tpl_frame->width *
-               sizeof(*tpl_frame_stats->block_stats_list));
     tpl_frame->is_valid = 0;
   }
 }
 
+static void free_tpl_frame_stats_list(VpxTplGopStats *tpl_gop_stats) {
+  int frame_idx;
+  for (frame_idx = 0; frame_idx < tpl_gop_stats->size; ++frame_idx) {
+    vpx_free(tpl_gop_stats->frame_stats_list[frame_idx].block_stats_list);
+  }
+  vpx_free(tpl_gop_stats->frame_stats_list);
+}
+
+static void init_tpl_stats_before_propagation(
+    struct vpx_internal_error_info *error_info, VpxTplGopStats *tpl_gop_stats,
+    TplDepFrame *tpl_stats, int tpl_gop_frames) {
+  int frame_idx;
+  free_tpl_frame_stats_list(tpl_gop_stats);
+  CHECK_MEM_ERROR(
+      error_info, tpl_gop_stats->frame_stats_list,
+      vpx_calloc(tpl_gop_frames, sizeof(*tpl_gop_stats->frame_stats_list)));
+  tpl_gop_stats->size = tpl_gop_frames;
+  for (frame_idx = 0; frame_idx < tpl_gop_frames; ++frame_idx) {
+    const int mi_rows = tpl_stats[frame_idx].height;
+    const int mi_cols = tpl_stats[frame_idx].width;
+    CHECK_MEM_ERROR(
+        error_info, tpl_gop_stats->frame_stats_list[frame_idx].block_stats_list,
+        vpx_calloc(
+            mi_rows * mi_cols,
+            sizeof(
+                *tpl_gop_stats->frame_stats_list[frame_idx].block_stats_list)));
+    tpl_gop_stats->frame_stats_list[frame_idx].num_blocks = mi_rows * mi_cols;
+  }
+}
+
 #if CONFIG_NON_GREEDY_MV
 static uint32_t full_pixel_motion_search(VP9_COMP *cpi, ThreadData *td,
                                          MotionField *motion_field,
@@ -1106,7 +1132,7 @@ static void mc_flow_dispenser(VP9_COMP *cpi, GF_PICTURE *gf_picture,
                               int frame_idx, BLOCK_SIZE bsize) {
   TplDepFrame *tpl_frame = &cpi->tpl_stats[frame_idx];
   VpxTplFrameStats *tpl_frame_stats_before_propagation =
-      &cpi->tpl_frame_stats[frame_idx];
+      &cpi->tpl_gop_stats.frame_stats_list[frame_idx];
   YV12_BUFFER_CONFIG *this_frame = gf_picture[frame_idx].frame;
   YV12_BUFFER_CONFIG *ref_frame[MAX_INTER_REF_FRAMES] = { NULL, NULL, NULL };
 
@@ -1349,12 +1375,6 @@ void vp9_init_tpl_buffer(VP9_COMP *cpi) {
     CHECK_MEM_ERROR(&cm->error, cpi->tpl_stats[frame].tpl_stats_ptr,
                     vpx_calloc(mi_rows * mi_cols,
                                sizeof(*cpi->tpl_stats[frame].tpl_stats_ptr)));
-    vpx_free(cpi->tpl_frame_stats[frame].block_stats_list);
-    CHECK_MEM_ERROR(
-        &cm->error, cpi->tpl_frame_stats[frame].block_stats_list,
-        vpx_calloc(mi_rows * mi_cols,
-                   sizeof(*cpi->tpl_frame_stats[frame].block_stats_list)));
-    cpi->tpl_frame_stats[frame].num_blocks = mi_rows * mi_cols;
     cpi->tpl_stats[frame].is_valid = 0;
     cpi->tpl_stats[frame].width = mi_cols;
     cpi->tpl_stats[frame].height = mi_rows;
@@ -1385,8 +1405,8 @@ void vp9_free_tpl_buffer(VP9_COMP *cpi) {
 #endif
     vpx_free(cpi->tpl_stats[frame].tpl_stats_ptr);
     cpi->tpl_stats[frame].is_valid = 0;
-    vpx_free(cpi->tpl_frame_stats[frame].block_stats_list);
   }
+  free_tpl_frame_stats_list(&cpi->tpl_gop_stats);
 }
 
 #if CONFIG_RATE_CTRL
@@ -1442,6 +1462,9 @@ void vp9_setup_tpl_stats(VP9_COMP *cpi) {
 
   init_tpl_stats(cpi);
 
+  init_tpl_stats_before_propagation(&cpi->common.error, &cpi->tpl_gop_stats,
+                                    cpi->tpl_stats, tpl_group_frames);
+
   // Backward propagation from tpl_group_frames to 1.
   for (frame_idx = tpl_group_frames - 1; frame_idx > 0; --frame_idx) {
     if (gf_picture[frame_idx].update_type == USE_BUF_FRAME) continue;
diff --git a/vp9/vp9_cx_iface.c b/vp9/vp9_cx_iface.c
index 7150f7428..62128ff28 100644
--- a/vp9/vp9_cx_iface.c
+++ b/vp9/vp9_cx_iface.c
@@ -1789,16 +1789,16 @@ static vpx_codec_err_t ctrl_get_svc_ref_frame_config(vpx_codec_alg_priv_t *ctx,
 static vpx_codec_err_t ctrl_get_tpl_stats(vpx_codec_alg_priv_t *ctx,
                                           va_list args) {
   VP9_COMP *const cpi = ctx->cpi;
-  VpxTplFrameStats *data = va_arg(args, VpxTplFrameStats *);
+  VpxTplGopStats *data = va_arg(args, VpxTplGopStats *);
+  VpxTplFrameStats *frame_stats_list = cpi->tpl_gop_stats.frame_stats_list;
   int i;
   if (data == NULL) {
     return VPX_CODEC_INVALID_PARAM;
   }
-  for (i = 0; i < MAX_ARF_GOP_SIZE; i++) {
-    data[i].frame_width = cpi->tpl_frame_stats[i].frame_width;
-    data[i].frame_height = cpi->tpl_frame_stats[i].frame_height;
-    data[i].num_blocks = cpi->tpl_frame_stats[i].num_blocks;
-    data[i].block_stats_list = cpi->tpl_frame_stats[i].block_stats_list;
+  data->size = cpi->tpl_gop_stats.size;
+
+  for (i = 0; i < data->size; i++) {
+    data->frame_stats_list[i] = frame_stats_list[i];
   }
 
   return VPX_CODEC_OK;
diff --git a/vpx/vpx_encoder.h b/vpx/vpx_encoder.h
index 2de808973..fb95723dd 100644
--- a/vpx/vpx_encoder.h
+++ b/vpx/vpx_encoder.h
@@ -271,6 +271,12 @@ typedef struct VpxTplFrameStats {
   VpxTplBlockStats *block_stats_list; /**< List of tpl stats for each block */
 } VpxTplFrameStats;
 
+/*!\brief Temporal dependency model stats for each GOP before propagation */
+typedef struct VpxTplGopStats {
+  int size; /**< GOP size, also the size of frame_stats_list. */
+  VpxTplFrameStats *frame_stats_list; /**< List of tpl stats for each frame */
+} VpxTplGopStats;
+
 /*!\brief Encoded Frame Flags
  *
  * This type indicates a bitfield to be passed to vpx_codec_encode(), defining