Consolidate ML based prune partition speed features
This patch combines rectangular, ab and 4-way
prune partition speed features.
Change-Id: I3d701d85aaec94972be824a1087520292e65aa18
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index 16a7332..8e507bf 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -3081,7 +3081,7 @@
// Pruning: pruning out some 4-way partitions using a DNN taking rd costs of
// sub-blocks from basic partition types.
- if (cpi->sf.part_sf.ml_prune_4_partition && partition4_allowed &&
+ if (cpi->sf.part_sf.ml_prune_partition && partition4_allowed &&
part_search_state->partition_rect_allowed[HORZ] &&
part_search_state->partition_rect_allowed[VERT]) {
av1_ml_prune_4_partition(
@@ -3239,7 +3239,7 @@
// Use the rd costs of PARTITION_NONE and subblocks from PARTITION_SPLIT
// to prune out rectangular partitions in some directions.
if (!cpi->sf.part_sf.ml_early_term_after_part_split_level &&
- cpi->sf.part_sf.ml_prune_rect_partition && !frame_is_intra_only(cm) &&
+ cpi->sf.part_sf.ml_prune_partition && !frame_is_intra_only(cm) &&
(part_search_state->partition_rect_allowed[HORZ] ||
part_search_state->partition_rect_allowed[VERT]) &&
!(part_search_state->prune_rect_part[HORZ] ||
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index e3eceb9..706df56 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -1564,7 +1564,7 @@
// Pruning: pruning out some ab partitions using a DNN taking rd costs of
// sub-blocks from previous basic partition types.
- if (cpi->sf.part_sf.ml_prune_ab_partition && ext_partition_allowed &&
+ if (cpi->sf.part_sf.ml_prune_partition && ext_partition_allowed &&
partition_horz_allowed && partition_vert_allowed) {
// TODO(huisu@google.com): x->source_variance may not be the current
// block's variance. The correct one to use is pb_source_variance. Need to
diff --git a/av1/encoder/speed_features.c b/av1/encoder/speed_features.c
index db1ddc4..5344dc4 100644
--- a/av1/encoder/speed_features.c
+++ b/av1/encoder/speed_features.c
@@ -383,9 +383,7 @@
sf->gm_sf.gm_search_type = GM_REDUCED_REF_SEARCH_SKIP_L2_L3;
sf->part_sf.less_rectangular_check_level = 1;
- sf->part_sf.ml_prune_4_partition = 1;
- sf->part_sf.ml_prune_ab_partition = 1;
- sf->part_sf.ml_prune_rect_partition = 1;
+ sf->part_sf.ml_prune_partition = 1;
sf->part_sf.prune_ext_partition_types_search_level = 1;
sf->part_sf.simple_motion_search_prune_rect = 1;
sf->part_sf.ml_predict_breakout_level = use_hbd ? 1 : 3;
@@ -750,9 +748,7 @@
sf->gm_sf.gm_search_type = GM_REDUCED_REF_SEARCH_SKIP_L2_L3;
sf->part_sf.less_rectangular_check_level = 1;
- sf->part_sf.ml_prune_4_partition = 1;
- sf->part_sf.ml_prune_ab_partition = 1;
- sf->part_sf.ml_prune_rect_partition = 1;
+ sf->part_sf.ml_prune_partition = 1;
sf->part_sf.prune_ext_partition_types_search_level = 1;
// TODO(debargha): Test, tweak and turn on either 1 or 2
@@ -1095,9 +1091,7 @@
part_sf->partition_search_breakout_dist_thr = 0;
part_sf->partition_search_breakout_rate_thr = 0;
part_sf->prune_ext_partition_types_search_level = 0;
- part_sf->ml_prune_rect_partition = 0;
- part_sf->ml_prune_ab_partition = 0;
- part_sf->ml_prune_4_partition = 0;
+ part_sf->ml_prune_partition = 0;
part_sf->ml_early_term_after_part_split_level = 0;
for (int i = 0; i < PARTITION_BLOCK_SIZES; ++i) {
part_sf->ml_partition_search_breakout_thresh[i] =
diff --git a/av1/encoder/speed_features.h b/av1/encoder/speed_features.h
index 6c93d2c..4d11f14 100644
--- a/av1/encoder/speed_features.h
+++ b/av1/encoder/speed_features.h
@@ -399,14 +399,9 @@
// aggressiveness of pruning in order.
int prune_ext_partition_types_search_level;
- // Use a ML model to prune horz and vert partitions
- int ml_prune_rect_partition;
-
- // Use a ML model to prune horz_a, horz_b, vert_a and vert_b partitions.
- int ml_prune_ab_partition;
-
- // Use a ML model to prune horz4 and vert4 partitions.
- int ml_prune_4_partition;
+ // Use a ML model to prune rectangular, ab and 4-way horz
+ // and vert partitions
+ int ml_prune_partition;
// Use a ML model to adaptively terminate partition search after trying
// PARTITION_SPLIT. Can take values 0 - 2, 0 meaning not being enabled, and
