Add kf_max_pyr_height API.

Bug: b/242914876
Change-Id: I59b1b1c91ac0b3190e15d39204d4c50751ba741d
diff --git a/apps/aomenc.c b/apps/aomenc.c
index 246bb69..60bc01a 100644
--- a/apps/aomenc.c
+++ b/apps/aomenc.c
@@ -451,6 +451,7 @@
   &g_av1_codec_arg_defs.fwd_kf_dist,
   &g_av1_codec_arg_defs.strict_level_conformance,
   &g_av1_codec_arg_defs.dist_metric,
+  &g_av1_codec_arg_defs.kf_max_pyr_height,
   NULL,
 };
 
diff --git a/av1/arg_defs.c b/av1/arg_defs.c
index 58e6f4e..7d1bbab 100644
--- a/av1/arg_defs.c
+++ b/av1/arg_defs.c
@@ -672,5 +672,10 @@
       ARG_DEF(NULL, "strict-level-conformance", 1,
               "When set to 1, exit the encoder when it fails to encode "
               "to a given target level"),
+  .kf_max_pyr_height = ARG_DEF(
+      NULL, "kf-max-pyr-height", 1,
+      "Maximum height of pyrimid structure used for the GOP starting with a "
+      "key frame (-1 to 5). When set to -1 (default), it does not have any "
+      "effects. The height also obeys the value set to gf_max_pyr_height."),
 #endif  // CONFIG_AV1_ENCODER
 };
diff --git a/av1/arg_defs.h b/av1/arg_defs.h
index 5c30f0d..812df27 100644
--- a/av1/arg_defs.h
+++ b/av1/arg_defs.h
@@ -231,6 +231,7 @@
   arg_def_t second_pass_log;
   arg_def_t auto_intra_tools_off;
   arg_def_t strict_level_conformance;
+  arg_def_t kf_max_pyr_height;
 #endif  // CONFIG_AV1_ENCODER
 } av1_codec_arg_definitions_t;
 
diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c
index 8514914..a63f67f 100644
--- a/av1/av1_cx_iface.c
+++ b/av1/av1_cx_iface.c
@@ -188,6 +188,7 @@
   // "--enable_diagonal_intra".
   int auto_intra_tools_off;
   int strict_level_conformance;
+  int kf_max_pyr_height;
 };
 
 #if CONFIG_REALTIME_ONLY
@@ -348,6 +349,7 @@
   NULL,            // second_pass_log
   0,               // auto_intra_tools_off
   0,               // strict_level_conformance
+  -1,              // kf_max_pyr_height
 };
 #else
 static const struct av1_extracfg default_extra_cfg = {
@@ -494,6 +496,7 @@
   NULL,            // second_pass_log
   0,               // auto_intra_tools_off
   0,               // strict_level_conformance
+  -1,              // kf_max_pyr_height
 };
 #endif
 
@@ -840,6 +843,14 @@
   RANGE_CHECK_BOOL(extra_cfg, auto_intra_tools_off);
   RANGE_CHECK_BOOL(extra_cfg, strict_level_conformance);
 
+  RANGE_CHECK(extra_cfg, kf_max_pyr_height, -1, 5);
+  if (extra_cfg->kf_max_pyr_height != -1 &&
+      extra_cfg->kf_max_pyr_height < (int)extra_cfg->gf_min_pyr_height) {
+    ERROR(
+        "The value of kf-max-pyr-height should not be smaller than "
+        "gf-min-pyr-height");
+  }
+
   return AOM_CODEC_OK;
 }
 
@@ -1418,6 +1429,8 @@
 
   oxcf->strict_level_conformance = extra_cfg->strict_level_conformance;
 
+  oxcf->kf_max_pyr_height = extra_cfg->kf_max_pyr_height;
+
   return AOM_CODEC_OK;
 }
 
@@ -3916,6 +3929,9 @@
                               &g_av1_codec_arg_defs.strict_level_conformance,
                               argv, err_string)) {
     extra_cfg.strict_level_conformance = arg_parse_int_helper(&arg, err_string);
+  } else if (arg_match_helper(&arg, &g_av1_codec_arg_defs.kf_max_pyr_height,
+                              argv, err_string)) {
+    extra_cfg.kf_max_pyr_height = arg_parse_int_helper(&arg, err_string);
   } else {
     match = 0;
     snprintf(err_string, ARG_ERR_MSG_MAX_LEN, "Cannot find aom option %s",
diff --git a/av1/encoder/encoder.h b/av1/encoder/encoder.h
index bd6c7a2..c8cf8ad 100644
--- a/av1/encoder/encoder.h
+++ b/av1/encoder/encoder.h
@@ -1067,6 +1067,9 @@
 
   // Exit the encoder when it fails to encode to a given level.
   int strict_level_conformance;
+
+  // Max depth for the GOP after a key frame
+  int kf_max_pyr_height;
   /*!\endcond */
 } AV1EncoderConfig;
 
diff --git a/av1/encoder/gop_structure.c b/av1/encoder/gop_structure.c
index ab20083..33e2074 100644
--- a/av1/encoder/gop_structure.c
+++ b/av1/encoder/gop_structure.c
@@ -824,10 +824,15 @@
   const int key_frame = rc->frames_since_key == 0;
   FRAME_UPDATE_TYPE first_frame_update_type = ARF_UPDATE;
 
-  if (key_frame)
+  if (key_frame) {
     first_frame_update_type = KF_UPDATE;
-  else if (!cpi->ppi->gf_state.arf_gf_boost_lst)
+    if (cpi->oxcf.kf_max_pyr_height != -1) {
+      gf_group->max_layer_depth_allowed = AOMMIN(
+          cpi->oxcf.kf_max_pyr_height, gf_group->max_layer_depth_allowed);
+    }
+  } else if (!cpi->ppi->gf_state.arf_gf_boost_lst) {
     first_frame_update_type = GF_UPDATE;
+  }
 
   gf_group->size = construct_multi_layer_gf_structure(
       cpi, twopass, gf_group, rc, frame_info, p_rc->baseline_gf_interval,