Add a flag to disable ml-based coding block transform features

Runtime flag: --disable-ml-transform-speed-features
OFF(0) by default.

Change-Id: I8f4c6b88c25b0dba562eb64e454f639dfa5ee310

aom/aom_encoder.h

diff --git a/aom/aom_encoder.h b/aom/aom_encoder.h
index bf42dc7..5ec26e4 100644
--- a/aom/aom_encoder.h
+++ b/aom/aom_encoder.h
@@ -238,6 +238,10 @@
    *
    */
   unsigned int disable_ml_partition_speed_features;
+  /*!\brief disable ml-based speed-up for transform search
+   *
+   */
+  unsigned int disable_ml_transform_speed_features;
 #if CONFIG_SDP
   /*!\brief enable Semi-decoupled partitioning
    *

apps/aomenc.c

diff --git a/apps/aomenc.c b/apps/aomenc.c
index 06a1079..fda2159 100644
--- a/apps/aomenc.c
+++ b/apps/aomenc.c
@@ -434,6 +434,7 @@
 };
 
 const arg_def_t *av1_key_val_args[] = {
+  &g_av1_codec_arg_defs.disable_ml_transform_speed_features,
 #if CONFIG_SDP
   &g_av1_codec_arg_defs.enable_sdp,
 #endif
@@ -582,6 +583,7 @@
   config->enable_ab_partitions = 1;
   config->enable_rect_partitions = 1;
   config->enable_1to4_partitions = 1;
+  config->disable_ml_transform_speed_features = 0;
 #if CONFIG_SDP
   config->enable_sdp = 1;
 #endif
@@ -1383,6 +1385,8 @@
   fprintf(
       stdout, "Tool setting (Partition)       : T-Type (%d), 4:1/1:4 (%d)\n",
       encoder_cfg->enable_ab_partitions, encoder_cfg->enable_1to4_partitions);
+  fprintf(stdout, "Disable ml transform speed features          : %d\n",
+          encoder_cfg->disable_ml_transform_speed_features);
 #if CONFIG_SDP
   fprintf(stdout, "                               : SDP (%d)\n",
           encoder_cfg->enable_sdp);

av1/arg_defs.c

diff --git a/av1/arg_defs.c b/av1/arg_defs.c
index c381ac9..ebb2469 100644
--- a/av1/arg_defs.c
+++ b/av1/arg_defs.c
@@ -346,6 +346,10 @@
   .enable_1to4_partitions = ARG_DEF(NULL, "enable-1to4-partitions", 1,
                                     "Enable 1:4 and 4:1 partitions "
                                     "(0: false, 1: true (default))"),
+  .disable_ml_transform_speed_features =
+      ARG_DEF(NULL, "disable-ml-transform-speed-features", 1,
+              "Disable ML transform speed features "
+              "(0: false (default), 1: true)"),
 #if CONFIG_SDP
   .enable_sdp = ARG_DEF(NULL, "enable-sdp", 1,
                         "Enable semi decoupled partitioning for key frame"

av1/arg_defs.h

diff --git a/av1/arg_defs.h b/av1/arg_defs.h
index 60a1ff6..6653b49 100644
--- a/av1/arg_defs.h
+++ b/av1/arg_defs.h
@@ -138,6 +138,7 @@
   arg_def_t enable_rect_partitions;
   arg_def_t enable_ab_partitions;
   arg_def_t enable_1to4_partitions;
+  arg_def_t disable_ml_transform_speed_features;
 #if CONFIG_SDP
   arg_def_t enable_sdp;
 #endif

av1/av1_cx_iface.c

diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c
index 35651e5..a483654 100644
--- a/av1/av1_cx_iface.c
+++ b/av1/av1_cx_iface.c
@@ -113,6 +113,7 @@
   int enable_rect_partitions;  // enable rectangular partitions for sequence
   int enable_ab_partitions;    // enable AB partitions for sequence
   int enable_1to4_partitions;  // enable 1:4 and 4:1 partitions for sequence
+  int disable_ml_transform_speed_features;  // disable all ml transform speedups
 #if CONFIG_SDP
   int enable_sdp;  // enable semi-decoupled partitioning
 #endif             // CONFIG_SDP
@@ -389,6 +390,7 @@
   1,                            // enable rectangular partitions
   1,                            // enable ab shape partitions
   1,                            // enable 1:4 and 4:1 partitions
+  0,                            // disable ml based transform speed features
 #if CONFIG_SDP
   1,    // enable semi-decoupled partitioning
 #endif  // CONFIG_SDP
@@ -841,6 +843,8 @@
   cfg->enable_rect_partitions = extra_cfg->enable_rect_partitions;
   cfg->enable_ab_partitions = extra_cfg->enable_ab_partitions;
   cfg->enable_1to4_partitions = extra_cfg->enable_1to4_partitions;
+  cfg->disable_ml_transform_speed_features =
+      extra_cfg->disable_ml_transform_speed_features;
 #if CONFIG_SDP
   cfg->enable_sdp = extra_cfg->enable_sdp;
 #endif
@@ -907,6 +911,8 @@
   extra_cfg->enable_rect_partitions = cfg->enable_rect_partitions;
   extra_cfg->enable_ab_partitions = cfg->enable_ab_partitions;
   extra_cfg->enable_1to4_partitions = cfg->enable_1to4_partitions;
+  extra_cfg->disable_ml_transform_speed_features =
+      cfg->disable_ml_transform_speed_features;
 #if CONFIG_SDP
   extra_cfg->enable_sdp = cfg->enable_sdp;
 #endif
@@ -1381,6 +1387,8 @@
   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->disable_ml_transform_speed_features =
+      extra_cfg->disable_ml_transform_speed_features;
 #if CONFIG_IST
   txfm_cfg->enable_ist = extra_cfg->enable_ist;
 #endif
@@ -3543,6 +3551,12 @@
                  argv, err_string)) {
     extra_cfg.disable_ml_partition_speed_features =
         arg_parse_int_helper(&arg, err_string);
+  } else if (arg_match_helper(
+                 &arg,
+                 &g_av1_codec_arg_defs.disable_ml_transform_speed_features,
+                 argv, err_string)) {
+    extra_cfg.disable_ml_transform_speed_features =
+        arg_parse_int_helper(&arg, err_string);
 #if CONFIG_SDP
   } else if (arg_match_helper(&arg, &g_av1_codec_arg_defs.enable_sdp, argv,
                               err_string)) {
@@ -3970,7 +3984,7 @@
     0,                           // use_fixed_qp_offsets
     { -1, -1, -1, -1, -1, -1 },  // fixed_qp_offsets
     {
-        0, 128, 128, 4, 1, 1, 1, 0,
+        0, 128, 128, 4, 1, 1, 1, 0, 0,
 #if CONFIG_SDP
         1,
 #endif  // CONFIG_SDP

av1/encoder/encoder.h

diff --git a/av1/encoder/encoder.h b/av1/encoder/encoder.h
index 9c3f4c2..71a1158 100644
--- a/av1/encoder/encoder.h
+++ b/av1/encoder/encoder.h
@@ -285,6 +285,10 @@
  */
 typedef struct {
   /*!
+   * Flag to disable ml based transform speed features.
+   */
+  bool disable_ml_transform_speed_features;
+  /*!
    * Flag to indicate if 64-pt transform should be enabled.
    */
   bool enable_tx64;

av1/encoder/speed_features.c

diff --git a/av1/encoder/speed_features.c b/av1/encoder/speed_features.c
index 067986b..cc42808 100644
--- a/av1/encoder/speed_features.c
+++ b/av1/encoder/speed_features.c
@@ -883,6 +883,14 @@
   lpf_sf->disable_lr_filter = 0;
 }
 
+static void av1_disable_ml_based_transform_sf(TX_SPEED_FEATURES *const tx_sf) {
+  tx_sf->tx_type_search.prune_2d_txfm_mode = TX_TYPE_PRUNE_0;
+  tx_sf->tx_type_search.ml_tx_split_thresh = -1;
+#if CONFIG_NEW_TX_PARTITION
+  tx_sf->tx_type_search.ml_tx_split_horzvert_thresh = -1;
+#endif  // CONFIG_NEW_TX_PARTITION
+}
+
 static void av1_disable_ml_based_partition_sf(
     PARTITION_SPEED_FEATURES *const part_sf) {
   part_sf->ml_prune_4_partition = 0;
@@ -916,6 +924,9 @@
 
   if (oxcf->part_cfg.disable_ml_partition_speed_features)
     av1_disable_ml_based_partition_sf(&sf->part_sf);
+
+  if (oxcf->txfm_cfg.disable_ml_transform_speed_features)
+    av1_disable_ml_based_transform_sf(&sf->tx_sf);
 }
 
 void av1_set_speed_features_framesize_independent(AV1_COMP *cpi, int speed) {

common/args.c

diff --git a/common/args.c b/common/args.c
index 6c8e401..6665645 100644
--- a/common/args.c
+++ b/common/args.c
@@ -81,6 +81,7 @@
     GET_PARAMS(enable_rect_partitions);
     GET_PARAMS(enable_1to4_partitions);
     GET_PARAMS(disable_ml_partition_speed_features);
+    GET_PARAMS(disable_ml_transform_speed_features);
 #if CONFIG_SDP
     GET_PARAMS(enable_sdp);
 #endif