Add an option to disable transform search size.

Bug: b:191463451

Change-Id: I802f79130b6bd1ce0d7979ddeb293b8f20e5d05b
diff --git a/aom/aomcx.h b/aom/aomcx.h
index 22764c7..621835a 100644
--- a/aom/aomcx.h
+++ b/aom/aomcx.h
@@ -1364,6 +1364,13 @@
    */
   AV1E_SET_ENABLE_DIRECTIONAL_INTRA = 145,
 
+  /*!\brief Control to turn on / off transform size search.
+   *
+   * - 0 = disable, transforms always have the largest possible size
+   * - 1 = enable, search for the best transform size for each block (default)
+   */
+  AV1E_SET_ENABLE_TX_SIZE_SEARCH = 146,
+
   // Any new encoder control IDs should be added above.
   // Maximum allowed encoder control ID is 229.
   // No encoder control ID should be added below.
@@ -1909,6 +1916,9 @@
 AOM_CTRL_USE_TYPE(AV1E_SET_ENABLE_DIRECTIONAL_INTRA, int)
 #define AOM_CTRL_AV1E_SET_ENABLE_DIRECTIONAL_INTRA
 
+AOM_CTRL_USE_TYPE(AV1E_SET_ENABLE_TX_SIZE_SEARCH, int)
+#define AOM_CTRL_AV1E_SET_ENABLE_TX_SIZE_SEARCH
+
 /*!\endcond */
 /*! @} - end defgroup aom_encoder */
 #ifdef __cplusplus
diff --git a/apps/aomenc.c b/apps/aomenc.c
index 16cddd9..02732ec 100644
--- a/apps/aomenc.c
+++ b/apps/aomenc.c
@@ -230,6 +230,7 @@
                                         AV1E_SET_DV_COST_UPD_FREQ,
                                         AV1E_SET_PARTITION_INFO_PATH,
                                         AV1E_SET_ENABLE_DIRECTIONAL_INTRA,
+                                        AV1E_SET_ENABLE_TX_SIZE_SEARCH,
                                         0 };
 
 const arg_def_t *main_args[] = { &g_av1_codec_arg_defs.help,
@@ -428,6 +429,7 @@
   &g_av1_codec_arg_defs.dv_cost_upd_freq,
   &g_av1_codec_arg_defs.partition_info_path,
   &g_av1_codec_arg_defs.enable_directional_intra,
+  &g_av1_codec_arg_defs.enable_tx_size_search,
   NULL,
 };
 
diff --git a/av1/arg_defs.c b/av1/arg_defs.c
index bfb209b..998085e 100644
--- a/av1/arg_defs.c
+++ b/av1/arg_defs.c
@@ -624,6 +624,12 @@
   .fwd_kf_dist =
       ARG_DEF(NULL, "fwd-kf-dist", -1,
               "Set distance between forward keyframes. A value of -1 means no "
-              "repetitive forward keyframes. Default is -1.")
+              "repetitive forward keyframes. Default is -1."),
+
+  .enable_tx_size_search = ARG_DEF(
+      NULL, "enable-tx-size-search", 1,
+      "Enable transform size search to find the best size for each block. "
+      "If false, transforms always have the largest possible size "
+      "(0: false, 1: true (default))")
 #endif  // CONFIG_AV1_ENCODER
 };
diff --git a/av1/arg_defs.h b/av1/arg_defs.h
index 52e0b27..637c6a7 100644
--- a/av1/arg_defs.h
+++ b/av1/arg_defs.h
@@ -218,6 +218,7 @@
   arg_def_t fixed_qp_offsets;
   arg_def_t vbr_corpus_complexity_lap;
   arg_def_t fwd_kf_dist;
+  arg_def_t enable_tx_size_search;
 #endif  // CONFIG_AV1_ENCODER
 } av1_codec_arg_definitions_t;
 
diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c
index 0d7466e..6a0bf16 100644
--- a/av1/av1_cx_iface.c
+++ b/av1/av1_cx_iface.c
@@ -146,6 +146,7 @@
   int use_intra_dct_only;
   int use_inter_dct_only;
   int use_intra_default_tx_only;
+  int enable_tx_size_search;
   int quant_b_adapt;
   unsigned int vbr_corpus_complexity_lap;
   AV1_LEVEL target_seq_level_idx[MAX_NUM_OPERATING_POINTS];
@@ -290,6 +291,7 @@
   0,  // use_intra_dct_only
   0,  // use_inter_dct_only
   0,  // use_intra_default_tx_only
+  1,  // enable_tx_size_search
   0,  // quant_b_adapt
   0,  // vbr_corpus_complexity_lap
   {
@@ -424,6 +426,7 @@
   0,                       // use_intra_dct_only
   0,                       // use_inter_dct_only
   0,                       // use_intra_default_tx_only
+  1,                       // enable_tx_size_search
   0,                       // quant_b_adapt
   0,                       // vbr_corpus_complexity_lap
   {
@@ -1250,6 +1253,7 @@
   txfm_cfg->use_intra_dct_only = extra_cfg->use_intra_dct_only;
   txfm_cfg->use_inter_dct_only = extra_cfg->use_inter_dct_only;
   txfm_cfg->use_intra_default_tx_only = extra_cfg->use_intra_default_tx_only;
+  txfm_cfg->enable_tx_size_search = extra_cfg->enable_tx_size_search;
 
   // Set compound type configuration.
   comp_type_cfg->enable_dist_wtd_comp =
@@ -2027,6 +2031,13 @@
   return update_extra_cfg(ctx, &extra_cfg);
 }
 
+static aom_codec_err_t ctrl_set_enable_tx_size_search(aom_codec_alg_priv_t *ctx,
+                                                      va_list args) {
+  struct av1_extracfg extra_cfg = ctx->extra_cfg;
+  extra_cfg.enable_tx_size_search = CAST(AV1E_SET_ENABLE_TX_SIZE_SEARCH, args);
+  return update_extra_cfg(ctx, &extra_cfg);
+}
+
 static aom_codec_err_t ctrl_set_quant_b_adapt(aom_codec_alg_priv_t *ctx,
                                               va_list args) {
   struct av1_extracfg extra_cfg = ctx->extra_cfg;
@@ -3745,6 +3756,7 @@
   { AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST, ctrl_enable_sb_multipass_unit_test },
   { AV1E_SET_DV_COST_UPD_FREQ, ctrl_set_dv_cost_upd_freq },
   { AV1E_SET_EXTERNAL_PARTITION, ctrl_set_external_partition },
+  { AV1E_SET_ENABLE_TX_SIZE_SEARCH, ctrl_set_enable_tx_size_search },
 
   // Getters
   { AOME_GET_LAST_QUANTIZER, ctrl_get_quantizer },
diff --git a/av1/encoder/encoder.h b/av1/encoder/encoder.h
index 6079a13..1115e907 100644
--- a/av1/encoder/encoder.h
+++ b/av1/encoder/encoder.h
@@ -332,6 +332,10 @@
    * (mode-dependent) only.
    */
   bool use_intra_default_tx_only;
+  /*!
+   * Flag to indicate if transform size search should be enabled.
+   */
+  bool enable_tx_size_search;
 } TxfmSizeTypeCfg;
 
 /*!
diff --git a/av1/encoder/speed_features.c b/av1/encoder/speed_features.c
index 34630e6..aef9ba5 100644
--- a/av1/encoder/speed_features.c
+++ b/av1/encoder/speed_features.c
@@ -92,10 +92,11 @@
 // (Eg : IntraBc) Index 1: Mode evaluation. Index 2: Winner mode evaluation.
 // Index 1 and 2 are applicable when enable_winner_mode_for_tx_size_srch speed
 // feature is ON
-static TX_SIZE_SEARCH_METHOD tx_size_search_methods[3][MODE_EVAL_TYPES] = {
+static TX_SIZE_SEARCH_METHOD tx_size_search_methods[4][MODE_EVAL_TYPES] = {
   { USE_FULL_RD, USE_LARGESTALL, USE_FULL_RD },
   { USE_FAST_RD, USE_LARGESTALL, USE_FULL_RD },
-  { USE_LARGESTALL, USE_LARGESTALL, USE_FULL_RD }
+  { USE_LARGESTALL, USE_LARGESTALL, USE_FULL_RD },
+  { USE_LARGESTALL, USE_LARGESTALL, USE_LARGESTALL }
 };
 
 // Predict transform skip levels to be used for default, mode and winner mode
@@ -1744,7 +1745,7 @@
     WINNER_MODE_SPEED_FEATURES *winner_mode_sf) {
   winner_mode_sf->motion_mode_for_winner_cand = 0;
   // Set this at the appropriate speed levels
-  winner_mode_sf->tx_size_search_level = USE_FULL_RD;
+  winner_mode_sf->tx_size_search_level = 0;
   winner_mode_sf->enable_winner_mode_for_coeff_opt = 0;
   winner_mode_sf->enable_winner_mode_for_tx_size_srch = 0;
   winner_mode_sf->enable_winner_mode_for_use_tx_domain_dist = 0;
@@ -1847,6 +1848,10 @@
       break;
   }
 
+  if (!oxcf->txfm_cfg.enable_tx_size_search) {
+    sf->winner_mode_sf.tx_size_search_level = 3;
+  }
+
   if (!cpi->ppi->seq_params_locked) {
     cpi->common.seq_params->enable_dual_filter &=
         !sf->interp_sf.disable_dual_filter;
diff --git a/av1/encoder/speed_features.h b/av1/encoder/speed_features.h
index 48bf494..96293e3 100644
--- a/av1/encoder/speed_features.h
+++ b/av1/encoder/speed_features.h
@@ -1048,6 +1048,7 @@
   // Level 0  : FULL RD     LARGEST ALL   FULL RD
   // Level 1  : FAST RD     LARGEST ALL   FULL RD
   // Level 2  : LARGEST ALL LARGEST ALL   FULL RD
+  // Level 3 :  LARGEST ALL LARGEST ALL   LARGEST ALL
   int tx_size_search_level;
 
   // Flag used to control the winner mode processing for use transform