From 9c15fb62b3dfe1c698dc28f9efedb022b0ef8eb8 Mon Sep 17 00:00:00 2001
From: Jerome Jiang <jianj@google.com>
Date: Fri, 17 Mar 2023 14:34:42 -0400
Subject: [PATCH] Add codec control to get tpl stats

Add command line flag to vpxenc to export tpl stats

Bug: b/273736974
Change-Id: I6980096531b0c12fbf7a307fdef4c562d0c29e32
---
 vp9/vp9_cx_iface.c | 23 +++++++++++++++++++++++
 vpx/vp8cx.h        |  9 +++++++++
 vpxenc.c           | 20 ++++++++++++++++++--
 vpxenc.h           |  1 +
 4 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/vp9/vp9_cx_iface.c b/vp9/vp9_cx_iface.c
index 4c7eaed..ec2105b 100644
--- a/vp9/vp9_cx_iface.c
+++ b/vp9/vp9_cx_iface.c
@@ -1788,6 +1788,28 @@ static vpx_codec_err_t ctrl_get_svc_ref_frame_config(vpx_codec_alg_priv_t *ctx,
   return VPX_CODEC_OK;
 }
 
+static vpx_codec_err_t ctrl_get_tpl_stats(vpx_codec_alg_priv_t *ctx,
+                                          va_list args) {
+  VP9_COMP *const cpi = ctx->cpi;
+  VP9_COMMON *const cm = &cpi->common;
+  TplDepFrame **data = va_arg(args, TplDepFrame **);
+  int i;
+  *data = vpx_calloc(MAX_ARF_GOP_SIZE, sizeof(TplDepFrame));
+  for (i = 0; i < MAX_ARF_GOP_SIZE; i++) {
+    const int mi_cols = mi_cols_aligned_to_sb(cm->mi_cols);
+    const int mi_rows = mi_cols_aligned_to_sb(cm->mi_rows);
+    const int copy_size = mi_cols * mi_rows * sizeof(*(*data)[i].tpl_stats_ptr);
+    (*data)[i] = cpi->tpl_stats[i];
+    (*data)[i].tpl_stats_ptr = NULL;
+    (*data)[i].tpl_stats_ptr =
+        vpx_calloc(mi_rows * mi_cols, sizeof(*(*data)[i].tpl_stats_ptr));
+    memcpy((*data)[i].tpl_stats_ptr, cpi->tpl_stats[i].tpl_stats_ptr,
+           copy_size);
+  }
+
+  return VPX_CODEC_OK;
+}
+
 static vpx_codec_err_t ctrl_set_svc_ref_frame_config(vpx_codec_alg_priv_t *ctx,
                                                      va_list args) {
   VP9_COMP *const cpi = ctx->cpi;
@@ -2035,6 +2057,7 @@ static vpx_codec_ctrl_fn_map_t encoder_ctrl_maps[] = {
   { VP9E_GET_ACTIVEMAP, ctrl_get_active_map },
   { VP9E_GET_LEVEL, ctrl_get_level },
   { VP9E_GET_SVC_REF_FRAME_CONFIG, ctrl_get_svc_ref_frame_config },
+  { VP9E_GET_TPL_STATS, ctrl_get_tpl_stats },
 
   { -1, NULL },
 };
diff --git a/vpx/vp8cx.h b/vpx/vp8cx.h
index e0b679f..01c0558 100644
--- a/vpx/vp8cx.h
+++ b/vpx/vp8cx.h
@@ -767,6 +767,13 @@ enum vp8e_enc_control_id {
    *
    */
   VP9E_SET_QUANTIZER_ONE_PASS,
+
+  /*!\brief Codec control to get TPL stats for the current frame.
+   *
+   * Supported in codecs: VP9
+   *
+   */
+  VP9E_GET_TPL_STATS,
 };
 
 /*!\brief vpx 1-D scaling mode
@@ -1097,6 +1104,8 @@ VPX_CTRL_USE_TYPE(VP8E_SET_RTC_EXTERNAL_RATECTRL, int)
 #define VPX_CTRL_VP8E_SET_RTC_EXTERNAL_RATECTRL
 VPX_CTRL_USE_TYPE(VP9E_SET_QUANTIZER_ONE_PASS, int)
 #define VPX_CTRL_VP9E_SET_QUANTIZER_ONE_PASS
+VPX_CTRL_USE_TYPE(VP9E_GET_TPL_STATS, void *)
+#define VPX_CTRL_VP9E_GET_TPL_STATS
 
 /*!\endcond */
 /*! @} - end defgroup vp8_encoder */
diff --git a/vpxenc.c b/vpxenc.c
index 61672ac..9d57708 100644
--- a/vpxenc.c
+++ b/vpxenc.c
@@ -39,6 +39,10 @@
 #include "vpx/vp8dx.h"
 #endif
 
+#if CONFIG_VP9_ENCODER
+#include "vp9/encoder/vp9_encoder.h"
+#endif
+
 #include "vpx/vpx_integer.h"
 #include "vpx_ports/mem_ops.h"
 #include "vpx_ports/vpx_timer.h"
@@ -161,6 +165,8 @@ static const arg_def_t disable_warnings =
 static const arg_def_t disable_warning_prompt =
     ARG_DEF("y", "disable-warning-prompt", 0,
             "Display warnings, but do not prompt user to continue.");
+static const arg_def_t export_tpl_stats =
+    ARG_DEF(NULL, "export-tpl-stats", 0, "Export TPL stats of vp9 encoder");
 
 #if CONFIG_VP9_HIGHBITDEPTH
 static const arg_def_t test16bitinternalarg = ARG_DEF(
@@ -191,6 +197,7 @@ static const arg_def_t *main_args[] = { &help,
                                         &disable_warnings,
                                         &disable_warning_prompt,
                                         &recontest,
+                                        &export_tpl_stats,
                                         NULL };
 
 static const arg_def_t usage =
@@ -531,9 +538,7 @@ static const arg_def_t disable_loopfilter =
             "1: Loopfilter off for non reference frames\n"
             "                                          "
             "2: Loopfilter off for all frames");
-#endif
 
-#if CONFIG_VP9_ENCODER
 static const arg_def_t *vp9_args[] = { &cpu_used_vp9,
                                        &auto_altref_vp9,
                                        &sharpness,
@@ -804,6 +809,8 @@ static void parse_global_config(struct VpxEncoderConfig *global, char **argv) {
       global->disable_warnings = 1;
     else if (arg_match(&arg, &disable_warning_prompt, argi))
       global->disable_warning_prompt = 1;
+    else if (arg_match(&arg, &export_tpl_stats, argi))
+      global->export_tpl_stats = 1;
     else
       argj++;
   }
@@ -1982,6 +1989,15 @@ int main(int argc, const char **argv_) {
 
         if (got_data && global.test_decode != TEST_DECODE_OFF)
           FOREACH_STREAM(test_decode(stream, global.test_decode, global.codec));
+
+#if CONFIG_VP9_ENCODER
+        if (got_data && global.export_tpl_stats) {
+          TplDepFrame *tpl_stats = NULL;
+          FOREACH_STREAM(vpx_codec_control(&stream->encoder, VP9E_GET_TPL_STATS,
+                                           &tpl_stats));
+          vpx_free(tpl_stats);
+        }
+#endif
       }
 
       fflush(stdout);
diff --git a/vpxenc.h b/vpxenc.h
index be54840..f065f08 100644
--- a/vpxenc.h
+++ b/vpxenc.h
@@ -56,6 +56,7 @@ struct VpxEncoderConfig {
   int disable_warnings;
   int disable_warning_prompt;
   int experimental_bitstream;
+  int export_tpl_stats;
 };
 
 #ifdef __cplusplus
-- 
2.7.4