vp9: ML var partition as speed feature & cleanup.

Remove it from runtime flag.

Add new struct for rd ml partition.

BUG=webm:1599

Change-Id: I883edbba83c65b7e557b8832419e212cffc85997

6 files changed, 89 insertions, 116 deletions

diff --git a/configure b/configure
index e91b57a86..71952351e 100755
--- a/configure
+++ b/configure
@@ -272,7 +272,6 @@ EXPERIMENT_LIST="
     fp_mb_stats
     emulate_hardware
     non_greedy_mv
-    ml_var_partition
 "
 CONFIG_LIST="
     dependency_tracking
diff --git a/vp9/encoder/vp9_block.h b/vp9/encoder/vp9_block.h
index 563fdbbde..11ec035e9 100644
--- a/vp9/encoder/vp9_block.h
+++ b/vp9/encoder/vp9_block.h
@@ -208,9 +208,7 @@ struct macroblock {
   void (*highbd_inv_txfm_add)(const tran_low_t *input, uint16_t *dest,
                               int stride, int eob, int bd);
 #endif
-#if CONFIG_ML_VAR_PARTITION
   DECLARE_ALIGNED(16, uint8_t, est_pred[64 * 64]);
-#endif  // CONFIG_ML_VAR_PARTITION
 
   struct scale_factors *me_sf;
 };
diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index f780b2802..b2ad686dc 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -3288,7 +3288,7 @@ static int ml_predict_breakout(VP9_COMP *const cpi, BLOCK_SIZE bsize,
       linear_score += linear_weights[i] * features[i];
   }
 
-  return linear_score >= cpi->sf.ml_partition_search_breakout_thresh[q_ctx];
+  return linear_score >= cpi->sf.rd_ml_partition.search_breakout_thresh[q_ctx];
 }
 #undef FEATURES
 
@@ -3313,15 +3313,15 @@ static void ml_prune_rect_partition(VP9_COMP *const cpi, MACROBLOCK *const x,
     case BLOCK_8X8: break;
     case BLOCK_16X16:
       nn_config = &vp9_rect_part_nnconfig_16;
-      thresh = cpi->sf.ml_prune_rect_partition_threhold[1];
+      thresh = cpi->sf.rd_ml_partition.prune_rect_thresh[1];
       break;
     case BLOCK_32X32:
       nn_config = &vp9_rect_part_nnconfig_32;
-      thresh = cpi->sf.ml_prune_rect_partition_threhold[2];
+      thresh = cpi->sf.rd_ml_partition.prune_rect_thresh[2];
       break;
     case BLOCK_64X64:
       nn_config = &vp9_rect_part_nnconfig_64;
-      thresh = cpi->sf.ml_prune_rect_partition_threhold[3];
+      thresh = cpi->sf.rd_ml_partition.prune_rect_thresh[3];
       break;
     default: assert(0 && "Unexpected block size."); return;
   }
@@ -3798,12 +3798,12 @@ static void rd_pick_partition(VP9_COMP *cpi, ThreadData *td,
 
   pc_tree->partitioning = PARTITION_NONE;
 
-  if (cpi->sf.ml_var_partition_pruning && !frame_is_intra_only(cm)) {
-    const int do_ml_var_partition_pruning =
+  if (cpi->sf.rd_ml_partition.var_pruning && !frame_is_intra_only(cm)) {
+    const int do_rd_ml_partition_var_pruning =
         partition_none_allowed && do_split &&
         mi_row + num_8x8_blocks_high_lookup[bsize] <= cm->mi_rows &&
         mi_col + num_8x8_blocks_wide_lookup[bsize] <= cm->mi_cols;
-    if (do_ml_var_partition_pruning) {
+    if (do_rd_ml_partition_var_pruning) {
       ml_predict_var_rd_paritioning(cpi, x, pc_tree, bsize, mi_row, mi_col,
                                     &partition_none_allowed, &do_split);
     } else {
@@ -3840,7 +3840,7 @@ static void rd_pick_partition(VP9_COMP *cpi, ThreadData *td,
         best_rdc = this_rdc;
         if (bsize >= BLOCK_8X8) pc_tree->partitioning = PARTITION_NONE;
 
-        if (cpi->sf.ml_partition_search_early_termination) {
+        if (cpi->sf.rd_ml_partition.search_early_termination) {
           // Currently, the machine-learning based partition search early
           // termination is only used while bsize is 16x16, 32x32 or 64x64,
           // VPXMIN(cm->width, cm->height) >= 480, and speed = 0.
@@ -3856,15 +3856,14 @@ static void rd_pick_partition(VP9_COMP *cpi, ThreadData *td,
 
         if ((do_split || do_rect) && !x->e_mbd.lossless && ctx->skippable) {
           const int use_ml_based_breakout =
-              cpi->sf.use_ml_partition_search_breakout &&
-              cm->base_qindex >= 100;
+              cpi->sf.rd_ml_partition.search_breakout && cm->base_qindex >= 100;
           if (use_ml_based_breakout) {
             if (ml_predict_breakout(cpi, bsize, x, &this_rdc)) {
               do_split = 0;
               do_rect = 0;
             }
           } else {
-            if (!cpi->sf.ml_partition_search_early_termination) {
+            if (!cpi->sf.rd_ml_partition.search_early_termination) {
               if ((best_rdc.dist < (dist_breakout_thr >> 2)) ||
                   (best_rdc.dist < dist_breakout_thr &&
                    best_rdc.rate < rate_breakout_thr)) {
@@ -4020,7 +4019,7 @@ static void rd_pick_partition(VP9_COMP *cpi, ThreadData *td,
         pc_tree->partitioning = PARTITION_SPLIT;
 
         // Rate and distortion based partition search termination clause.
-        if (!cpi->sf.ml_partition_search_early_termination &&
+        if (!cpi->sf.rd_ml_partition.search_early_termination &&
             !x->e_mbd.lossless &&
             ((best_rdc.dist < (dist_breakout_thr >> 2)) ||
              (best_rdc.dist < dist_breakout_thr &&
@@ -4532,7 +4531,6 @@ static void pred_pixel_ready_reset(PC_TREE *pc_tree, BLOCK_SIZE bsize) {
   }
 }
 
-#if CONFIG_ML_VAR_PARTITION
 #define FEATURES 6
 #define LABELS 2
 static int ml_predict_var_paritioning(VP9_COMP *cpi, MACROBLOCK *x,
@@ -4602,7 +4600,6 @@ static int ml_predict_var_paritioning(VP9_COMP *cpi, MACROBLOCK *x,
 }
 #undef FEATURES
 #undef LABELS
-#endif  // CONFIG_ML_VAR_PARTITION
 
 static void nonrd_pick_partition(VP9_COMP *cpi, ThreadData *td,
                                  TileDataEnc *tile_data, TOKENEXTRA **tp,
@@ -4633,10 +4630,8 @@ static void nonrd_pick_partition(VP9_COMP *cpi, ThreadData *td,
       !force_vert_split && yss <= xss && bsize >= BLOCK_8X8;
   int partition_vert_allowed =
       !force_horz_split && xss <= yss && bsize >= BLOCK_8X8;
-#if CONFIG_ML_VAR_PARTITION
   const int use_ml_based_partitioning =
       sf->partition_search_type == ML_BASED_PARTITION;
-#endif  // CONFIG_ML_VAR_PARTITION
 
   (void)*tp_orig;
 
@@ -4668,7 +4663,6 @@ static void nonrd_pick_partition(VP9_COMP *cpi, ThreadData *td,
     partition_vert_allowed &= force_vert_split;
   }
 
-#if CONFIG_ML_VAR_PARTITION
   if (use_ml_based_partitioning) {
     if (partition_none_allowed || do_split) do_rect = 0;
     if (partition_none_allowed && do_split) {
@@ -4678,7 +4672,6 @@ static void nonrd_pick_partition(VP9_COMP *cpi, ThreadData *td,
       if (ml_predicted_partition == PARTITION_SPLIT) partition_none_allowed = 0;
     }
   }
-#endif  // CONFIG_ML_VAR_PARTITION
 
   if (!partition_none_allowed && !do_split) do_rect = 1;
 
@@ -4703,10 +4696,7 @@ static void nonrd_pick_partition(VP9_COMP *cpi, ThreadData *td,
         best_rdc = this_rdc;
         if (bsize >= BLOCK_8X8) pc_tree->partitioning = PARTITION_NONE;
 
-#if CONFIG_ML_VAR_PARTITION
-        if (!use_ml_based_partitioning)
-#endif  // CONFIG_ML_VAR_PARTITION
-        {
+        if (!use_ml_based_partitioning) {
           int64_t dist_breakout_thr = sf->partition_search_breakout_thr.dist;
           int64_t rate_breakout_thr = sf->partition_search_breakout_thr.rate;
           dist_breakout_thr >>=
@@ -5115,31 +5105,26 @@ static void nonrd_use_partition(VP9_COMP *cpi, ThreadData *td,
     update_partition_context(xd, mi_row, mi_col, subsize, bsize);
 }
 
-#if CONFIG_ML_VAR_PARTITION
 // Get a prediction(stored in x->est_pred) for the whole 64x64 superblock.
 static void get_estimated_pred(VP9_COMP *cpi, const TileInfo *const tile,
                                MACROBLOCK *x, int mi_row, int mi_col) {
   VP9_COMMON *const cm = &cpi->common;
   const int is_key_frame = frame_is_intra_only(cm);
+  MACROBLOCKD *xd = &x->e_mbd;
 
   set_offsets(cpi, tile, x, mi_row, mi_col, BLOCK_64X64);
 
   if (!is_key_frame) {
-    MACROBLOCKD *xd = &x->e_mbd;
     MODE_INFO *mi = xd->mi[0];
     YV12_BUFFER_CONFIG *yv12 = get_ref_frame_buffer(cpi, LAST_FRAME);
     const YV12_BUFFER_CONFIG *yv12_g = NULL;
     const BLOCK_SIZE bsize = BLOCK_32X32 + (mi_col + 4 < cm->mi_cols) * 2 +
                              (mi_row + 4 < cm->mi_rows);
-    int pixels_wide = 64, pixels_high = 64;
     unsigned int y_sad_g, y_sad_thr;
     unsigned int y_sad = UINT_MAX;
 
     assert(yv12 != NULL);
 
-    if (xd->mb_to_right_edge < 0) pixels_wide += (xd->mb_to_right_edge >> 3);
-    if (xd->mb_to_bottom_edge < 0) pixels_high += (xd->mb_to_bottom_edge >> 3);
-
     if (!(is_one_pass_cbr_svc(cpi) && cpi->svc.spatial_layer_id) ||
         cpi->svc.use_gf_temporal_ref_current_layer) {
       // For now, GOLDEN will not be used for non-zero spatial layers, since
@@ -5193,7 +5178,6 @@ static void get_estimated_pred(VP9_COMP *cpi, const TileInfo *const tile,
                            &cm->frame_refs[GOLDEN_FRAME - 1].sf);
       mi->ref_frame[0] = GOLDEN_FRAME;
       mi->mv[0].as_int = 0;
-      y_sad = y_sad_g;
     } else {
       x->pred_mv[LAST_FRAME] = mi->mv[0].as_mv;
     }
@@ -5218,7 +5202,6 @@ static void get_estimated_pred(VP9_COMP *cpi, const TileInfo *const tile,
 #endif  // CONFIG_VP9_HIGHBITDEPTH
   }
 }
-#endif  // CONFIG_ML_VAR_PARTITION
 
 static void encode_nonrd_sb_row(VP9_COMP *cpi, ThreadData *td,
                                 TileDataEnc *tile_data, int mi_row,
@@ -5311,7 +5294,6 @@ static void encode_nonrd_sb_row(VP9_COMP *cpi, ThreadData *td,
         nonrd_use_partition(cpi, td, tile_data, mi, tp, mi_row, mi_col,
                             BLOCK_64X64, 1, &dummy_rdc, td->pc_root);
         break;
-#if CONFIG_ML_VAR_PARTITION
       case ML_BASED_PARTITION:
         get_estimated_pred(cpi, tile_info, x, mi_row, mi_col);
         x->max_partition_size = BLOCK_64X64;
@@ -5321,7 +5303,6 @@ static void encode_nonrd_sb_row(VP9_COMP *cpi, ThreadData *td,
                              BLOCK_64X64, &dummy_rdc, 1, INT64_MAX,
                              td->pc_root);
         break;
-#endif  // CONFIG_ML_VAR_PARTITION
       case SOURCE_VAR_BASED_PARTITION:
         set_source_var_based_partition(cpi, tile_info, x, mi, mi_row, mi_col);
         nonrd_use_partition(cpi, td, tile_data, mi, tp, mi_row, mi_col,
diff --git a/vp9/encoder/vp9_partition_models.h b/vp9/encoder/vp9_partition_models.h
index f2ee72fa7..09c0e30a4 100644
--- a/vp9/encoder/vp9_partition_models.h
+++ b/vp9/encoder/vp9_partition_models.h
@@ -607,7 +607,6 @@ static const NN_CONFIG vp9_partition_nnconfig_16x16 = {
 };
 #undef FEATURES
 
-#if CONFIG_ML_VAR_PARTITION
 #define FEATURES 6
 static const float vp9_var_part_nn_weights_64_layer0[FEATURES * 8] = {
   -0.249572f, 0.205532f,  -2.175608f, 1.094836f,  -2.986370f, 0.193160f,
@@ -735,7 +734,6 @@ static const NN_CONFIG vp9_var_part_nnconfig_16 = {
   },
 };
 #undef FEATURES
-#endif  // CONFIG_ML_VAR_PARTITION
 
 #define FEATURES 12
 #define LABELS 1
diff --git a/vp9/encoder/vp9_speed_features.c b/vp9/encoder/vp9_speed_features.c
index 0ba3edf7f..7c567a14c 100644
--- a/vp9/encoder/vp9_speed_features.c
+++ b/vp9/encoder/vp9_speed_features.c
@@ -76,27 +76,27 @@ static void set_good_speed_feature_framesize_dependent(VP9_COMP *cpi,
   if (is_480p_or_larger) {
     // Currently, the machine-learning based partition search early termination
     // is only used while VPXMIN(cm->width, cm->height) >= 480 and speed = 0.
-    sf->ml_partition_search_early_termination = 1;
+    sf->rd_ml_partition.search_early_termination = 1;
   } else {
     sf->use_square_only_thresh_high = BLOCK_32X32;
   }
 
   if (!is_1080p_or_larger) {
-    sf->use_ml_partition_search_breakout = 1;
+    sf->rd_ml_partition.search_breakout = 1;
     if (is_720p_or_larger) {
-      sf->ml_partition_search_breakout_thresh[0] = 0.0f;
-      sf->ml_partition_search_breakout_thresh[1] = 0.0f;
-      sf->ml_partition_search_breakout_thresh[2] = 0.0f;
+      sf->rd_ml_partition.search_breakout_thresh[0] = 0.0f;
+      sf->rd_ml_partition.search_breakout_thresh[1] = 0.0f;
+      sf->rd_ml_partition.search_breakout_thresh[2] = 0.0f;
     } else {
-      sf->ml_partition_search_breakout_thresh[0] = 2.5f;
-      sf->ml_partition_search_breakout_thresh[1] = 1.5f;
-      sf->ml_partition_search_breakout_thresh[2] = 1.5f;
+      sf->rd_ml_partition.search_breakout_thresh[0] = 2.5f;
+      sf->rd_ml_partition.search_breakout_thresh[1] = 1.5f;
+      sf->rd_ml_partition.search_breakout_thresh[2] = 1.5f;
     }
   }
 
   if (speed >= 1) {
-    sf->ml_partition_search_early_termination = 0;
-    sf->use_ml_partition_search_breakout = 1;
+    sf->rd_ml_partition.search_early_termination = 0;
+    sf->rd_ml_partition.search_breakout = 1;
     if (is_480p_or_larger)
       sf->use_square_only_thresh_high = BLOCK_64X64;
     else
@@ -106,21 +106,21 @@ static void set_good_speed_feature_framesize_dependent(VP9_COMP *cpi,
       sf->disable_split_mask =
           cm->show_frame ? DISABLE_ALL_SPLIT : DISABLE_ALL_INTER_SPLIT;
       sf->partition_search_breakout_thr.dist = (1 << 22);
-      sf->ml_partition_search_breakout_thresh[0] = -5.0f;
-      sf->ml_partition_search_breakout_thresh[1] = -5.0f;
-      sf->ml_partition_search_breakout_thresh[2] = -9.0f;
+      sf->rd_ml_partition.search_breakout_thresh[0] = -5.0f;
+      sf->rd_ml_partition.search_breakout_thresh[1] = -5.0f;
+      sf->rd_ml_partition.search_breakout_thresh[2] = -9.0f;
     } else {
       sf->disable_split_mask = DISABLE_COMPOUND_SPLIT;
       sf->partition_search_breakout_thr.dist = (1 << 21);
-      sf->ml_partition_search_breakout_thresh[0] = -1.0f;
-      sf->ml_partition_search_breakout_thresh[1] = -1.0f;
-      sf->ml_partition_search_breakout_thresh[2] = -1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[0] = -1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[1] = -1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[2] = -1.0f;
     }
 #if CONFIG_VP9_HIGHBITDEPTH
     if (cpi->Source->flags & YV12_FLAG_HIGHBITDEPTH) {
-      sf->ml_partition_search_breakout_thresh[0] -= 1.0f;
-      sf->ml_partition_search_breakout_thresh[1] -= 1.0f;
-      sf->ml_partition_search_breakout_thresh[2] -= 1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[0] -= 1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[1] -= 1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[2] -= 1.0f;
     }
 #endif  // CONFIG_VP9_HIGHBITDEPTH
   }
@@ -134,14 +134,14 @@ static void set_good_speed_feature_framesize_dependent(VP9_COMP *cpi,
       sf->adaptive_pred_interp_filter = 0;
       sf->partition_search_breakout_thr.dist = (1 << 24);
       sf->partition_search_breakout_thr.rate = 120;
-      sf->use_ml_partition_search_breakout = 0;
+      sf->rd_ml_partition.search_breakout = 0;
     } else {
       sf->disable_split_mask = LAST_AND_INTRA_SPLIT_ONLY;
       sf->partition_search_breakout_thr.dist = (1 << 22);
       sf->partition_search_breakout_thr.rate = 100;
-      sf->ml_partition_search_breakout_thresh[0] = 0.0f;
-      sf->ml_partition_search_breakout_thresh[1] = -1.0f;
-      sf->ml_partition_search_breakout_thresh[2] = -4.0f;
+      sf->rd_ml_partition.search_breakout_thresh[0] = 0.0f;
+      sf->rd_ml_partition.search_breakout_thresh[1] = -1.0f;
+      sf->rd_ml_partition.search_breakout_thresh[2] = -4.0f;
     }
     sf->rd_auto_partition_min_limit = set_partition_min_limit(cm);
 
@@ -158,7 +158,7 @@ static void set_good_speed_feature_framesize_dependent(VP9_COMP *cpi,
   }
 
   if (speed >= 3) {
-    sf->use_ml_partition_search_breakout = 0;
+    sf->rd_ml_partition.search_breakout = 0;
     if (is_720p_or_larger) {
       sf->disable_split_mask = DISABLE_ALL_SPLIT;
       sf->schedule_mode_search = cm->base_qindex < 220 ? 1 : 0;
@@ -215,12 +215,12 @@ static void set_good_speed_feature_framesize_independent(VP9_COMP *cpi,
   sf->less_rectangular_check = 1;
   sf->use_square_partition_only = !boosted;
   sf->prune_ref_frame_for_rect_partitions = 1;
-  sf->ml_var_partition_pruning = 1;
+  sf->rd_ml_partition.var_pruning = 1;
 
-  sf->ml_prune_rect_partition_threhold[0] = -1;
-  sf->ml_prune_rect_partition_threhold[1] = 350;
-  sf->ml_prune_rect_partition_threhold[2] = 325;
-  sf->ml_prune_rect_partition_threhold[3] = 250;
+  sf->rd_ml_partition.prune_rect_thresh[0] = -1;
+  sf->rd_ml_partition.prune_rect_thresh[1] = 350;
+  sf->rd_ml_partition.prune_rect_thresh[2] = 325;
+  sf->rd_ml_partition.prune_rect_thresh[3] = 250;
 
   if (cpi->twopass.fr_content_type == FC_GRAPHICS_ANIMATION) {
     sf->exhaustive_searches_thresh = (1 << 22);
@@ -238,10 +238,10 @@ static void set_good_speed_feature_framesize_independent(VP9_COMP *cpi,
 
   if (speed >= 1) {
     sf->temporal_filter_search_method = NSTEP;
-    sf->ml_var_partition_pruning = !boosted;
-    sf->ml_prune_rect_partition_threhold[1] = 225;
-    sf->ml_prune_rect_partition_threhold[2] = 225;
-    sf->ml_prune_rect_partition_threhold[3] = 225;
+    sf->rd_ml_partition.var_pruning = !boosted;
+    sf->rd_ml_partition.prune_rect_thresh[1] = 225;
+    sf->rd_ml_partition.prune_rect_thresh[2] = 225;
+    sf->rd_ml_partition.prune_rect_thresh[3] = 225;
 
     if (oxcf->pass == 2) {
       TWO_PASS *const twopass = &cpi->twopass;
@@ -284,7 +284,7 @@ static void set_good_speed_feature_framesize_independent(VP9_COMP *cpi,
   }
 
   if (speed >= 2) {
-    sf->ml_var_partition_pruning = 0;
+    sf->rd_ml_partition.var_pruning = 0;
     if (oxcf->vbr_corpus_complexity)
       sf->recode_loop = ALLOW_RECODE_FIRST;
     else
@@ -308,9 +308,9 @@ static void set_good_speed_feature_framesize_independent(VP9_COMP *cpi,
     sf->recode_tolerance_high = 45;
     sf->enhanced_full_pixel_motion_search = 0;
     sf->prune_ref_frame_for_rect_partitions = 0;
-    sf->ml_prune_rect_partition_threhold[1] = -1;
-    sf->ml_prune_rect_partition_threhold[2] = -1;
-    sf->ml_prune_rect_partition_threhold[3] = -1;
+    sf->rd_ml_partition.prune_rect_thresh[1] = -1;
+    sf->rd_ml_partition.prune_rect_thresh[2] = -1;
+    sf->rd_ml_partition.prune_rect_thresh[3] = -1;
     sf->mv.subpel_search_level = 0;
 
     if (cpi->twopass.fr_content_type == FC_GRAPHICS_ANIMATION) {
@@ -448,6 +448,7 @@ static void set_rt_speed_feature_framesize_independent(
   sf->enable_tpl_model = 0;
   sf->enhanced_full_pixel_motion_search = 0;
   sf->use_accurate_subpel_search = USE_2_TAPS;
+  sf->nonrd_use_ml_partition = 0;
 
   if (speed >= 1) {
     sf->allow_txfm_domain_distortion = 1;
@@ -563,16 +564,14 @@ static void set_rt_speed_feature_framesize_independent(
         (frames_since_key % (sf->last_partitioning_redo_frequency << 1) == 1);
     sf->max_delta_qindex = is_keyframe ? 20 : 15;
     sf->partition_search_type = REFERENCE_PARTITION;
-#if CONFIG_ML_VAR_PARTITION
-    if (!frame_is_intra_only(cm) && cm->width >= 360 && cm->height >= 360)
-      sf->partition_search_type = ML_BASED_PARTITION;
-    else
-      sf->partition_search_type = REFERENCE_PARTITION;
+    if (sf->nonrd_use_ml_partition) {
+      if (!frame_is_intra_only(cm) && cm->width >= 360 && cm->height >= 360)
+        sf->partition_search_type = ML_BASED_PARTITION;
 #if CONFIG_VP9_HIGHBITDEPTH
-    if (cpi->Source->flags & YV12_FLAG_HIGHBITDEPTH)
-      sf->partition_search_type = REFERENCE_PARTITION;
+      if (cpi->Source->flags & YV12_FLAG_HIGHBITDEPTH)
+        sf->partition_search_type = REFERENCE_PARTITION;
 #endif  // CONFIG_VP9_HIGHBITDEPTH
-#endif  // CONFIG_ML_VAR_PARTITION
+    }
     if (cpi->oxcf.rc_mode == VPX_VBR && cpi->oxcf.lag_in_frames > 0 &&
         cpi->rc.is_src_frame_alt_ref) {
       sf->partition_search_type = VAR_BASED_PARTITION;
@@ -633,16 +632,10 @@ static void set_rt_speed_feature_framesize_independent(
       sf->use_altref_onepass = 1;
       sf->use_compound_nonrd_pickmode = 1;
     }
-#if CONFIG_ML_VAR_PARTITION
-    if (frame_is_intra_only(cm) || cm->width < 360 || cm->height < 360)
-      sf->partition_search_type = VAR_BASED_PARTITION;
-#if CONFIG_VP9_HIGHBITDEPTH
-    if (cpi->Source->flags & YV12_FLAG_HIGHBITDEPTH)
+
+    if (!sf->nonrd_use_ml_partition)
       sf->partition_search_type = VAR_BASED_PARTITION;
-#endif  // CONFIG_VP9_HIGHBITDEPTH
-#else
-    sf->partition_search_type = VAR_BASED_PARTITION;
-#endif  // CONFIG_ML_VAR_PARTITION
+
     sf->mv.search_method = NSTEP;
     sf->mv.reduce_first_step_size = 1;
     sf->skip_encode_sb = 0;
@@ -815,8 +808,8 @@ void vp9_set_speed_features_framesize_dependent(VP9_COMP *cpi) {
   // Some speed-up features even for best quality as minimal impact on quality.
   sf->partition_search_breakout_thr.dist = (1 << 19);
   sf->partition_search_breakout_thr.rate = 80;
-  sf->ml_partition_search_early_termination = 0;
-  sf->use_ml_partition_search_breakout = 0;
+  sf->rd_ml_partition.search_early_termination = 0;
+  sf->rd_ml_partition.search_breakout = 0;
 
   if (oxcf->mode == REALTIME) {
     set_rt_speed_feature_framesize_dependent(cpi, sf, oxcf->speed);
@@ -941,11 +934,11 @@ void vp9_set_speed_features_framesize_independent(VP9_COMP *cpi) {
   sf->limit_newmv_early_exit = 0;
   sf->bias_golden = 0;
   sf->base_mv_aggressive = 0;
-  sf->ml_prune_rect_partition_threhold[0] = -1;
-  sf->ml_prune_rect_partition_threhold[1] = -1;
-  sf->ml_prune_rect_partition_threhold[2] = -1;
-  sf->ml_prune_rect_partition_threhold[3] = -1;
-  sf->ml_var_partition_pruning = 0;
+  sf->rd_ml_partition.prune_rect_thresh[0] = -1;
+  sf->rd_ml_partition.prune_rect_thresh[1] = -1;
+  sf->rd_ml_partition.prune_rect_thresh[2] = -1;
+  sf->rd_ml_partition.prune_rect_thresh[3] = -1;
+  sf->rd_ml_partition.var_pruning = 0;
   sf->use_accurate_subpel_search = USE_8_TAPS;
 
   // Some speed-up features even for best quality as minimal impact on quality.
diff --git a/vp9/encoder/vp9_speed_features.h b/vp9/encoder/vp9_speed_features.h
index 98f8de81d..8609c9837 100644
--- a/vp9/encoder/vp9_speed_features.h
+++ b/vp9/encoder/vp9_speed_features.h
@@ -151,10 +151,8 @@ typedef enum {
   // Use non-fixed partitions based on source variance.
   SOURCE_VAR_BASED_PARTITION,
 
-#if CONFIG_ML_VAR_PARTITION
   // Make partition decisions with machine learning models.
   ML_BASED_PARTITION
-#endif  // CONFIG_ML_VAR_PARTITION
 } PARTITION_SEARCH_TYPE;
 
 typedef enum {
@@ -351,12 +349,6 @@ typedef struct SPEED_FEATURES {
   // Prune reference frames for rectangular partitions.
   int prune_ref_frame_for_rect_partitions;
 
-  // Threshold values used for ML based rectangular partition search pruning.
-  // If < 0, the feature is turned off.
-  // Higher values mean more aggressiveness to skip rectangular partition
-  // search that results in better encoding speed but worse coding performance.
-  int ml_prune_rect_partition_threhold[4];
-
   // Sets min and max partition sizes for this 64x64 region based on the
   // same 64x64 in last encoded frame, and the left and above neighbor.
   AUTO_MIN_MAX_MODE auto_min_max_partition_size;
@@ -511,18 +503,27 @@ typedef struct SPEED_FEATURES {
   // Partition search early breakout thresholds.
   PARTITION_SEARCH_BREAKOUT_THR partition_search_breakout_thr;
 
-  // Use ML-based partition search early breakout.
-  int use_ml_partition_search_breakout;
-  // Higher values mean more aggressiveness for partition search breakout that
-  // results in better encoding  speed but worse compression performance.
-  float ml_partition_search_breakout_thresh[3];
+  struct {
+    // Use ML-based partition search early breakout.
+    int search_breakout;
+    // Higher values mean more aggressiveness for partition search breakout that
+    // results in better encoding  speed but worse compression performance.
+    float search_breakout_thresh[3];
 
-  // Machine-learning based partition search early termination
-  int ml_partition_search_early_termination;
+    // Machine-learning based partition search early termination
+    int search_early_termination;
 
-  // Machine-learning based partition search pruning using prediction residue
-  // variance.
-  int ml_var_partition_pruning;
+    // Machine-learning based partition search pruning using prediction residue
+    // variance.
+    int var_pruning;
+
+    // Threshold values used for ML based rectangular partition search pruning.
+    // If < 0, the feature is turned off.
+    // Higher values mean more aggressiveness to skip rectangular partition
+    // search that results in better encoding speed but worse coding
+    // performance.
+    int prune_rect_thresh[4];
+  } rd_ml_partition;
 
   // Allow skipping partition search for still image frame
   int allow_partition_search_skip;
@@ -598,6 +599,9 @@ typedef struct SPEED_FEATURES {
 
   // Search method used by temporal filtering in full_pixel_motion_search.
   SEARCH_METHODS temporal_filter_search_method;
+
+  // Use machine learning based partition search.
+  int nonrd_use_ml_partition;
 } SPEED_FEATURES;
 
 struct VP9_COMP;