AV1 RT: Move ML-based partitioning to experiments
This will allow to reduce MACROBLOCK structure for ~1.5K to make MT
initialization lighter
Change-Id: I405f676d748f108c57adbd00677cfb46748659e9
diff --git a/av1/encoder/block.h b/av1/encoder/block.h
index 099ba2e..db042d9 100644
--- a/av1/encoder/block.h
+++ b/av1/encoder/block.h
@@ -1172,7 +1172,9 @@
//! SSE of the current predictor.
unsigned int pred_sse[REF_FRAMES];
//! Prediction for ML based partition.
+#if CONFIG_RT_ML_PARTITIONING
DECLARE_ALIGNED(16, uint8_t, est_pred[128 * 128]);
+#endif
/**@}*/
} MACROBLOCK;
#undef SINGLE_REF_MODES
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 1b2174c..981dc95 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -442,6 +442,7 @@
}
#endif // !CONFIG_REALTIME_ONLY
+#if CONFIG_RT_ML_PARTITIONING
// Get a prediction(stored in x->est_pred) for the whole superblock.
static void get_estimated_pred(AV1_COMP *cpi, const TileInfo *const tile,
MACROBLOCK *x, int mi_row, int mi_col) {
@@ -489,6 +490,7 @@
#endif // CONFIG_VP9_HIGHBITDEPTH
}
}
+#endif // CONFIG_RT_ML_PARTITIONING
#define AVG_CDF_WEIGHT_LEFT 3
#define AVG_CDF_WEIGHT_TOP_RIGHT 1
@@ -520,7 +522,7 @@
int offset = cpi->source->y_stride * (mi_row << 2) + (mi_col << 2);
av1_source_content_sb(cpi, x, offset);
}
-
+#if CONFIG_RT_ML_PARTITIONING
if (sf->part_sf.partition_search_type == ML_BASED_PARTITION) {
PC_TREE *const pc_root = av1_alloc_pc_tree_node(sb_size);
RD_STATS dummy_rdc;
@@ -530,7 +532,7 @@
av1_free_pc_tree_recursive(pc_root, av1_num_planes(cm), 0, 0);
return;
}
-
+#endif
// Set the partition
if (sf->part_sf.partition_search_type == FIXED_PARTITION || seg_skip) {
// set a fixed-size partition
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index 8d0bf74..ac3c516 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -3827,6 +3827,7 @@
}
#endif // !CONFIG_REALTIME_ONLY
+#if CONFIG_RT_ML_PARTITIONING
#define FEATURES 6
#define LABELS 2
static int ml_predict_var_paritioning(AV1_COMP *cpi, MACROBLOCK *x,
@@ -4220,3 +4221,4 @@
assert(tp_orig == *tp);
}
}
+#endif // CONFIG_RT_ML_PARTITIONING
diff --git a/av1/encoder/partition_search.h b/av1/encoder/partition_search.h
index b77d4a2..136548e 100644
--- a/av1/encoder/partition_search.h
+++ b/av1/encoder/partition_search.h
@@ -32,11 +32,13 @@
TileDataEnc *tile_data, MB_MODE_INFO **mib,
TokenExtra **tp, int mi_row, int mi_col,
BLOCK_SIZE bsize, PC_TREE *pc_tree);
+#if CONFIG_RT_ML_PARTITIONING
void av1_nonrd_pick_partition(AV1_COMP *cpi, ThreadData *td,
TileDataEnc *tile_data, TokenExtra **tp,
int mi_row, int mi_col, BLOCK_SIZE bsize,
RD_STATS *rd_cost, int do_recon, int64_t best_rd,
PC_TREE *pc_tree);
+#endif
bool av1_rd_pick_partition(AV1_COMP *const cpi, ThreadData *td,
TileDataEnc *tile_data, TokenExtra **tp, int mi_row,
int mi_col, BLOCK_SIZE bsize, RD_STATS *rd_cost,
diff --git a/av1/encoder/speed_features.c b/av1/encoder/speed_features.c
index f0d84d4..0ab71e7 100644
--- a/av1/encoder/speed_features.c
+++ b/av1/encoder/speed_features.c
@@ -334,11 +334,13 @@
sf->rt_sf.nonrd_agressive_skip = 1;
// TODO(kyslov) Re-enable when AV1 models are trained
#if 0
+#if CONFIG_RT_ML_PARTITIONING
if (!frame_is_intra_only(cm)) {
sf->part_sf.partition_search_type = ML_BASED_PARTITION;
sf->rt_sf.reuse_inter_pred_nonrd = 0;
}
#endif
+#endif
}
} else {
if (speed == 8 && !cpi->use_svc) {
diff --git a/av1/encoder/speed_features.h b/av1/encoder/speed_features.h
index 19a4229..7defacf 100644
--- a/av1/encoder/speed_features.h
+++ b/av1/encoder/speed_features.h
@@ -262,8 +262,10 @@
// Partition using source variance
VAR_BASED_PARTITION,
+#if CONFIG_RT_ML_PARTITIONING
// Partition using ML model
ML_BASED_PARTITION
+#endif
} UENUM1BYTE(PARTITION_SEARCH_TYPE);
enum {
diff --git a/build/cmake/aom_config_defaults.cmake b/build/cmake/aom_config_defaults.cmake
index 24020a8..7fdad19 100644
--- a/build/cmake/aom_config_defaults.cmake
+++ b/build/cmake/aom_config_defaults.cmake
@@ -136,6 +136,9 @@
set_aom_config_var(CONFIG_NN_V2 0 "Fully-connected neural nets ver.2.")
set_aom_config_var(CONFIG_OPTICAL_FLOW_API 0
"AV1 experiment flag for optical flow API.")
+set_aom_config_var(CONFIG_RT_ML_PARTITIONING 0
+ "Build with ML-based partitioning for Real Time.")
+
#
# Variables in this section control optional features of the build system.
#
