Move all partition models in encodeframe.c to partition_strategy
BUG=aomedia:2343
Change-Id: I6189fe7887b89d20551af37bdf6b8fe1d399627e
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 2bf8193..e629b27 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -68,10 +68,6 @@
ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,
int mi_row, int mi_col, BLOCK_SIZE bsize,
int *rate);
-static int ml_predict_breakout(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
- const MACROBLOCK *const x,
- const RD_STATS *const rd_stats,
- unsigned int pb_source_variance);
// This is used as a reference when computing the source variance for the
// purposes of activity masking.
@@ -2363,8 +2359,8 @@
// values as in rd_pick_partition. Retraining the model and tuning the
// threshold values might be helpful to improve the speed.
if (use_ml_based_breakout) {
- if (ml_predict_breakout(cpi, bsize, x, &this_rdc,
- x->source_variance)) {
+ if (av1_ml_predict_breakout(cpi, bsize, x, &this_rdc,
+ x->source_variance)) {
do_square_split = 0;
}
}
@@ -2509,488 +2505,6 @@
}
}
-// split_score indicates confidence of picking split partition;
-// none_score indicates confidence of picking none partition;
-#define FEATURE_SIZE 19
-static int ml_prune_2pass_split_partition(const PC_TREE_STATS *pc_tree_stats,
- BLOCK_SIZE bsize, int *split_score,
- int *none_score) {
- if (!pc_tree_stats->valid) return 0;
- const float *split_weights = NULL;
- const float *none_weights = NULL;
- switch (bsize) {
- case BLOCK_4X4: break;
- case BLOCK_8X8:
- split_weights = av1_2pass_split_partition_weights_8;
- none_weights = av1_2pass_none_partition_weights_8;
- break;
- case BLOCK_16X16:
- split_weights = av1_2pass_split_partition_weights_16;
- none_weights = av1_2pass_none_partition_weights_16;
- break;
- case BLOCK_32X32:
- split_weights = av1_2pass_split_partition_weights_32;
- none_weights = av1_2pass_none_partition_weights_32;
- break;
- case BLOCK_64X64:
- split_weights = av1_2pass_split_partition_weights_64;
- none_weights = av1_2pass_none_partition_weights_64;
- break;
- case BLOCK_128X128:
- split_weights = av1_2pass_split_partition_weights_128;
- none_weights = av1_2pass_none_partition_weights_128;
- break;
- default: assert(0 && "Unexpected bsize.");
- }
- if (!split_weights || !none_weights) return 0;
-
- aom_clear_system_state();
-
- float features[FEATURE_SIZE];
- int feature_index = 0;
- features[feature_index++] = (float)pc_tree_stats->split;
- features[feature_index++] = (float)pc_tree_stats->skip;
- const int rdcost = (int)AOMMIN(INT_MAX, pc_tree_stats->rdcost);
- const int rd_valid = rdcost > 0 && rdcost < 1000000000;
- features[feature_index++] = (float)rd_valid;
- for (int i = 0; i < 4; ++i) {
- features[feature_index++] = (float)pc_tree_stats->sub_block_split[i];
- features[feature_index++] = (float)pc_tree_stats->sub_block_skip[i];
- const int sub_rdcost =
- (int)AOMMIN(INT_MAX, pc_tree_stats->sub_block_rdcost[i]);
- const int sub_rd_valid = sub_rdcost > 0 && sub_rdcost < 1000000000;
- features[feature_index++] = (float)sub_rd_valid;
- // Ratio between the sub-block RD and the whole-block RD.
- float rd_ratio = 1.0f;
- if (rd_valid && sub_rd_valid && sub_rdcost < rdcost)
- rd_ratio = (float)sub_rdcost / (float)rdcost;
- features[feature_index++] = rd_ratio;
- }
- assert(feature_index == FEATURE_SIZE);
-
- float score_1 = split_weights[FEATURE_SIZE];
- float score_2 = none_weights[FEATURE_SIZE];
- for (int i = 0; i < FEATURE_SIZE; ++i) {
- score_1 += features[i] * split_weights[i];
- score_2 += features[i] * none_weights[i];
- }
- *split_score = (int)(score_1 * 100);
- *none_score = (int)(score_2 * 100);
- return 1;
-}
-#undef FEATURE_SIZE
-
-static void ml_prune_rect_partition(const AV1_COMP *const cpi,
- const MACROBLOCK *const x, BLOCK_SIZE bsize,
- int64_t best_rd, int64_t none_rd,
- int64_t *split_rd,
- int *const dst_prune_horz,
- int *const dst_prune_vert) {
- if (bsize < BLOCK_8X8 || best_rd >= 1000000000) return;
- best_rd = AOMMAX(best_rd, 1);
- const NN_CONFIG *nn_config = NULL;
- const float prob_thresholds[5] = { 0.01f, 0.01f, 0.004f, 0.002f, 0.002f };
- float cur_thresh = 0.0f;
- switch (bsize) {
- case BLOCK_8X8:
- nn_config = &av1_rect_partition_nnconfig_8;
- cur_thresh = prob_thresholds[0];
- break;
- case BLOCK_16X16:
- nn_config = &av1_rect_partition_nnconfig_16;
- cur_thresh = prob_thresholds[1];
- break;
- case BLOCK_32X32:
- nn_config = &av1_rect_partition_nnconfig_32;
- cur_thresh = prob_thresholds[2];
- break;
- case BLOCK_64X64:
- nn_config = &av1_rect_partition_nnconfig_64;
- cur_thresh = prob_thresholds[3];
- break;
- case BLOCK_128X128:
- nn_config = &av1_rect_partition_nnconfig_128;
- cur_thresh = prob_thresholds[4];
- break;
- default: assert(0 && "Unexpected bsize.");
- }
- if (!nn_config) return;
- aom_clear_system_state();
-
- // 1. Compute input features
- float features[9];
-
- // RD cost ratios
- for (int i = 0; i < 5; i++) features[i] = 1.0f;
- if (none_rd > 0 && none_rd < 1000000000)
- features[0] = (float)none_rd / (float)best_rd;
- for (int i = 0; i < 4; i++) {
- if (split_rd[i] > 0 && split_rd[i] < 1000000000)
- features[1 + i] = (float)split_rd[i] / (float)best_rd;
- }
-
- // Variance ratios
- const MACROBLOCKD *const xd = &x->e_mbd;
- int whole_block_variance;
- if (is_cur_buf_hbd(xd)) {
- whole_block_variance = av1_high_get_sby_perpixel_variance(
- cpi, &x->plane[0].src, bsize, xd->bd);
- } else {
- whole_block_variance =
- av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize);
- }
- whole_block_variance = AOMMAX(whole_block_variance, 1);
-
- int split_variance[4];
- const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
- struct buf_2d buf;
- buf.stride = x->plane[0].src.stride;
- const int bw = block_size_wide[bsize];
- for (int i = 0; i < 4; ++i) {
- const int x_idx = (i & 1) * bw / 2;
- const int y_idx = (i >> 1) * bw / 2;
- buf.buf = x->plane[0].src.buf + x_idx + y_idx * buf.stride;
- if (is_cur_buf_hbd(xd)) {
- split_variance[i] =
- av1_high_get_sby_perpixel_variance(cpi, &buf, subsize, xd->bd);
- } else {
- split_variance[i] = av1_get_sby_perpixel_variance(cpi, &buf, subsize);
- }
- }
-
- for (int i = 0; i < 4; i++)
- features[5 + i] = (float)split_variance[i] / (float)whole_block_variance;
-
- // 2. Do the prediction and prune 0-2 partitions based on their probabilities
- float raw_scores[3] = { 0.0f };
- av1_nn_predict(features, nn_config, raw_scores);
- aom_clear_system_state();
- float probs[3] = { 0.0f };
- av1_nn_softmax(raw_scores, probs, 3);
-
- // probs[0] is the probability of the fact that both rectangular partitions
- // are worse than current best_rd
- if (probs[1] <= cur_thresh) (*dst_prune_horz) = 1;
- if (probs[2] <= cur_thresh) (*dst_prune_vert) = 1;
-}
-
-// Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
-// considered.
-static void ml_prune_ab_partition(BLOCK_SIZE bsize, int part_ctx, int var_ctx,
- int64_t best_rd, int64_t horz_rd[2],
- int64_t vert_rd[2], int64_t split_rd[4],
- int *const horza_partition_allowed,
- int *const horzb_partition_allowed,
- int *const verta_partition_allowed,
- int *const vertb_partition_allowed) {
- if (bsize < BLOCK_8X8 || best_rd >= 1000000000) return;
- const NN_CONFIG *nn_config = NULL;
- switch (bsize) {
- case BLOCK_8X8: nn_config = NULL; break;
- case BLOCK_16X16: nn_config = &av1_ab_partition_nnconfig_16; break;
- case BLOCK_32X32: nn_config = &av1_ab_partition_nnconfig_32; break;
- case BLOCK_64X64: nn_config = &av1_ab_partition_nnconfig_64; break;
- case BLOCK_128X128: nn_config = &av1_ab_partition_nnconfig_128; break;
- default: assert(0 && "Unexpected bsize.");
- }
- if (!nn_config) return;
-
- aom_clear_system_state();
-
- // Generate features.
- float features[10];
- int feature_index = 0;
- features[feature_index++] = (float)part_ctx;
- features[feature_index++] = (float)var_ctx;
- const int rdcost = (int)AOMMIN(INT_MAX, best_rd);
- int sub_block_rdcost[8] = { 0 };
- int rd_index = 0;
- for (int i = 0; i < 2; ++i) {
- if (horz_rd[i] > 0 && horz_rd[i] < 1000000000)
- sub_block_rdcost[rd_index] = (int)horz_rd[i];
- ++rd_index;
- }
- for (int i = 0; i < 2; ++i) {
- if (vert_rd[i] > 0 && vert_rd[i] < 1000000000)
- sub_block_rdcost[rd_index] = (int)vert_rd[i];
- ++rd_index;
- }
- for (int i = 0; i < 4; ++i) {
- if (split_rd[i] > 0 && split_rd[i] < 1000000000)
- sub_block_rdcost[rd_index] = (int)split_rd[i];
- ++rd_index;
- }
- for (int i = 0; i < 8; ++i) {
- // Ratio between the sub-block RD and the whole-block RD.
- float rd_ratio = 1.0f;
- if (sub_block_rdcost[i] > 0 && sub_block_rdcost[i] < rdcost)
- rd_ratio = (float)sub_block_rdcost[i] / (float)rdcost;
- features[feature_index++] = rd_ratio;
- }
- assert(feature_index == 10);
-
- // Calculate scores using the NN model.
- float score[16] = { 0.0f };
- av1_nn_predict(features, nn_config, score);
- aom_clear_system_state();
- int int_score[16];
- int max_score = -1000;
- for (int i = 0; i < 16; ++i) {
- int_score[i] = (int)(100 * score[i]);
- max_score = AOMMAX(int_score[i], max_score);
- }
-
- // Make decisions based on the model scores.
- int thresh = max_score;
- switch (bsize) {
- case BLOCK_16X16: thresh -= 150; break;
- case BLOCK_32X32: thresh -= 100; break;
- default: break;
- }
- *horza_partition_allowed = 0;
- *horzb_partition_allowed = 0;
- *verta_partition_allowed = 0;
- *vertb_partition_allowed = 0;
- for (int i = 0; i < 16; ++i) {
- if (int_score[i] >= thresh) {
- if ((i >> 0) & 1) *horza_partition_allowed = 1;
- if ((i >> 1) & 1) *horzb_partition_allowed = 1;
- if ((i >> 2) & 1) *verta_partition_allowed = 1;
- if ((i >> 3) & 1) *vertb_partition_allowed = 1;
- }
- }
-}
-
-#define FEATURES 18
-#define LABELS 4
-// Use a ML model to predict if horz4 and vert4 should be considered.
-static void ml_prune_4_partition(const AV1_COMP *const cpi, MACROBLOCK *const x,
- BLOCK_SIZE bsize, int part_ctx,
- int64_t best_rd, int64_t horz_rd[2],
- int64_t vert_rd[2], int64_t split_rd[4],
- int *const partition_horz4_allowed,
- int *const partition_vert4_allowed,
- unsigned int pb_source_variance, int mi_row,
- int mi_col) {
- if (best_rd >= 1000000000) return;
- const NN_CONFIG *nn_config = NULL;
- switch (bsize) {
- case BLOCK_16X16: nn_config = &av1_4_partition_nnconfig_16; break;
- case BLOCK_32X32: nn_config = &av1_4_partition_nnconfig_32; break;
- case BLOCK_64X64: nn_config = &av1_4_partition_nnconfig_64; break;
- default: assert(0 && "Unexpected bsize.");
- }
- if (!nn_config) return;
-
- aom_clear_system_state();
-
- // Generate features.
- float features[FEATURES];
- int feature_index = 0;
- features[feature_index++] = (float)part_ctx;
- features[feature_index++] = (float)get_unsigned_bits(pb_source_variance);
-
- const int rdcost = (int)AOMMIN(INT_MAX, best_rd);
- int sub_block_rdcost[8] = { 0 };
- int rd_index = 0;
- for (int i = 0; i < 2; ++i) {
- if (horz_rd[i] > 0 && horz_rd[i] < 1000000000)
- sub_block_rdcost[rd_index] = (int)horz_rd[i];
- ++rd_index;
- }
- for (int i = 0; i < 2; ++i) {
- if (vert_rd[i] > 0 && vert_rd[i] < 1000000000)
- sub_block_rdcost[rd_index] = (int)vert_rd[i];
- ++rd_index;
- }
- for (int i = 0; i < 4; ++i) {
- if (split_rd[i] > 0 && split_rd[i] < 1000000000)
- sub_block_rdcost[rd_index] = (int)split_rd[i];
- ++rd_index;
- }
- for (int i = 0; i < 8; ++i) {
- // Ratio between the sub-block RD and the whole-block RD.
- float rd_ratio = 1.0f;
- if (sub_block_rdcost[i] > 0 && sub_block_rdcost[i] < rdcost)
- rd_ratio = (float)sub_block_rdcost[i] / (float)rdcost;
- features[feature_index++] = rd_ratio;
- }
-
- // Get variance of the 1:4 and 4:1 sub-blocks.
- unsigned int horz_4_source_var[4] = { 0 };
- unsigned int vert_4_source_var[4] = { 0 };
- {
- BLOCK_SIZE horz_4_bs = get_partition_subsize(bsize, PARTITION_HORZ_4);
- BLOCK_SIZE vert_4_bs = get_partition_subsize(bsize, PARTITION_VERT_4);
- av1_setup_src_planes(x, cpi->source, mi_row, mi_col,
- av1_num_planes(&cpi->common), bsize);
- const int src_stride = x->plane[0].src.stride;
- const uint8_t *src = x->plane[0].src.buf;
- const MACROBLOCKD *const xd = &x->e_mbd;
- for (int i = 0; i < 4; ++i) {
- const uint8_t *horz_src =
- src + i * block_size_high[horz_4_bs] * src_stride;
- const uint8_t *vert_src = src + i * block_size_wide[vert_4_bs];
- unsigned int horz_var, vert_var, sse;
- if (is_cur_buf_hbd(xd)) {
- switch (xd->bd) {
- case 10:
- horz_var = cpi->fn_ptr[horz_4_bs].vf(
- horz_src, src_stride, CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_10),
- 0, &sse);
- vert_var = cpi->fn_ptr[vert_4_bs].vf(
- vert_src, src_stride, CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_10),
- 0, &sse);
- break;
- case 12:
- horz_var = cpi->fn_ptr[horz_4_bs].vf(
- horz_src, src_stride, CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_12),
- 0, &sse);
- vert_var = cpi->fn_ptr[vert_4_bs].vf(
- vert_src, src_stride, CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_12),
- 0, &sse);
- break;
- case 8:
- default:
- horz_var = cpi->fn_ptr[horz_4_bs].vf(
- horz_src, src_stride, CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8),
- 0, &sse);
- vert_var = cpi->fn_ptr[vert_4_bs].vf(
- vert_src, src_stride, CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8),
- 0, &sse);
- break;
- }
- horz_4_source_var[i] =
- ROUND_POWER_OF_TWO(horz_var, num_pels_log2_lookup[horz_4_bs]);
- vert_4_source_var[i] =
- ROUND_POWER_OF_TWO(vert_var, num_pels_log2_lookup[vert_4_bs]);
- } else {
- horz_var = cpi->fn_ptr[horz_4_bs].vf(horz_src, src_stride, AV1_VAR_OFFS,
- 0, &sse);
- vert_var = cpi->fn_ptr[vert_4_bs].vf(vert_src, src_stride, AV1_VAR_OFFS,
- 0, &sse);
- horz_4_source_var[i] =
- ROUND_POWER_OF_TWO(horz_var, num_pels_log2_lookup[horz_4_bs]);
- vert_4_source_var[i] =
- ROUND_POWER_OF_TWO(vert_var, num_pels_log2_lookup[vert_4_bs]);
- }
- }
- }
-
- const float denom = (float)(pb_source_variance + 1);
- const float low_b = 0.1f;
- const float high_b = 10.0f;
- for (int i = 0; i < 4; ++i) {
- // Ratio between the 4:1 sub-block variance and the whole-block variance.
- float var_ratio = (float)(horz_4_source_var[i] + 1) / denom;
- if (var_ratio < low_b) var_ratio = low_b;
- if (var_ratio > high_b) var_ratio = high_b;
- features[feature_index++] = var_ratio;
- }
- for (int i = 0; i < 4; ++i) {
- // Ratio between the 1:4 sub-block RD and the whole-block RD.
- float var_ratio = (float)(vert_4_source_var[i] + 1) / denom;
- if (var_ratio < low_b) var_ratio = low_b;
- if (var_ratio > high_b) var_ratio = high_b;
- features[feature_index++] = var_ratio;
- }
- assert(feature_index == FEATURES);
-
- // Calculate scores using the NN model.
- float score[LABELS] = { 0.0f };
- av1_nn_predict(features, nn_config, score);
- aom_clear_system_state();
- int int_score[LABELS];
- int max_score = -1000;
- for (int i = 0; i < LABELS; ++i) {
- int_score[i] = (int)(100 * score[i]);
- max_score = AOMMAX(int_score[i], max_score);
- }
-
- // Make decisions based on the model scores.
- int thresh = max_score;
- switch (bsize) {
- case BLOCK_16X16: thresh -= 500; break;
- case BLOCK_32X32: thresh -= 500; break;
- case BLOCK_64X64: thresh -= 200; break;
- default: break;
- }
- *partition_horz4_allowed = 0;
- *partition_vert4_allowed = 0;
- for (int i = 0; i < LABELS; ++i) {
- if (int_score[i] >= thresh) {
- if ((i >> 0) & 1) *partition_horz4_allowed = 1;
- if ((i >> 1) & 1) *partition_vert4_allowed = 1;
- }
- }
-}
-#undef FEATURES
-#undef LABELS
-
-#define FEATURES 4
-// ML-based partition search breakout.
-static int ml_predict_breakout(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
- const MACROBLOCK *const x,
- const RD_STATS *const rd_stats,
- unsigned int pb_source_variance) {
- const NN_CONFIG *nn_config = NULL;
- int thresh = 0;
- switch (bsize) {
- case BLOCK_8X8:
- nn_config = &av1_partition_breakout_nnconfig_8;
- thresh = cpi->sf.ml_partition_search_breakout_thresh[0];
- break;
- case BLOCK_16X16:
- nn_config = &av1_partition_breakout_nnconfig_16;
- thresh = cpi->sf.ml_partition_search_breakout_thresh[1];
- break;
- case BLOCK_32X32:
- nn_config = &av1_partition_breakout_nnconfig_32;
- thresh = cpi->sf.ml_partition_search_breakout_thresh[2];
- break;
- case BLOCK_64X64:
- nn_config = &av1_partition_breakout_nnconfig_64;
- thresh = cpi->sf.ml_partition_search_breakout_thresh[3];
- break;
- case BLOCK_128X128:
- nn_config = &av1_partition_breakout_nnconfig_128;
- thresh = cpi->sf.ml_partition_search_breakout_thresh[4];
- break;
- default: assert(0 && "Unexpected bsize.");
- }
- if (!nn_config || thresh < 0) return 0;
-
- // Generate feature values.
- float features[FEATURES];
- int feature_index = 0;
- aom_clear_system_state();
-
- const int num_pels_log2 = num_pels_log2_lookup[bsize];
- float rate_f = (float)AOMMIN(rd_stats->rate, INT_MAX);
- rate_f = ((float)x->rdmult / 128.0f / 512.0f / (float)(1 << num_pels_log2)) *
- rate_f;
- features[feature_index++] = rate_f;
-
- const float dist_f =
- (float)(AOMMIN(rd_stats->dist, INT_MAX) >> num_pels_log2);
- features[feature_index++] = dist_f;
-
- features[feature_index++] = (float)pb_source_variance;
-
- const int dc_q = (int)x->plane[0].dequant_QTX[0];
- features[feature_index++] = (float)(dc_q * dc_q) / 256.0f;
- assert(feature_index == FEATURES);
-
- // Calculate score using the NN model.
- float score = 0.0f;
- av1_nn_predict(features, nn_config, &score);
- aom_clear_system_state();
-
- // Make decision.
- return (int)(score * 100) >= thresh;
-}
-#undef FEATURES
-
// Record the ref frames that have been selected by square partition blocks.
static void update_picked_ref_frames_mask(MACROBLOCK *const x, int ref_type,
BLOCK_SIZE bsize, int mib_size,
@@ -3139,7 +2653,7 @@
if (bsize > BLOCK_4X4 && x->use_cb_search_range) {
int split_score = 0;
int none_score = 0;
- const int score_valid = ml_prune_2pass_split_partition(
+ const int score_valid = av1_ml_prune_2pass_split_partition(
&pc_tree->pc_tree_stats, bsize, &split_score, &none_score);
if (score_valid) {
{
@@ -3350,8 +2864,8 @@
bsize <= cpi->sf.use_square_partition_only_threshold &&
bsize > BLOCK_4X4 && xd->bd == 8;
if (use_ml_based_breakout) {
- if (ml_predict_breakout(cpi, bsize, x, &this_rdc,
- pb_source_variance)) {
+ if (av1_ml_predict_breakout(cpi, bsize, x, &this_rdc,
+ pb_source_variance)) {
do_square_split = 0;
do_rectangular_split = 0;
}
@@ -3480,8 +2994,8 @@
(partition_horz_allowed || partition_vert_allowed) &&
!(prune_horz || prune_vert) && !terminate_partition_search) {
av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes, bsize);
- ml_prune_rect_partition(cpi, x, bsize, best_rdc.rdcost, cur_none_rd,
- split_rd, &prune_horz, &prune_vert);
+ av1_ml_prune_rect_partition(cpi, x, bsize, best_rdc.rdcost, cur_none_rd,
+ split_rd, &prune_horz, &prune_vert);
}
// PARTITION_HORZ
@@ -3774,11 +3288,11 @@
// 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
// re-train the model to fix it.
- ml_prune_ab_partition(bsize, pc_tree->partitioning,
- get_unsigned_bits(x->source_variance),
- best_rdc.rdcost, horz_rd, vert_rd, split_rd,
- &horza_partition_allowed, &horzb_partition_allowed,
- &verta_partition_allowed, &vertb_partition_allowed);
+ av1_ml_prune_ab_partition(
+ bsize, pc_tree->partitioning, get_unsigned_bits(x->source_variance),
+ best_rdc.rdcost, horz_rd, vert_rd, split_rd, &horza_partition_allowed,
+ &horzb_partition_allowed, &verta_partition_allowed,
+ &vertb_partition_allowed);
}
horza_partition_allowed &= cpi->oxcf.enable_ab_partitions;
@@ -3980,10 +3494,10 @@
}
if (cpi->sf.ml_prune_4_partition && partition4_allowed &&
partition_horz_allowed && partition_vert_allowed) {
- ml_prune_4_partition(cpi, x, bsize, pc_tree->partitioning, best_rdc.rdcost,
- horz_rd, vert_rd, split_rd, &partition_horz4_allowed,
- &partition_vert4_allowed, pb_source_variance, mi_row,
- mi_col);
+ av1_ml_prune_4_partition(cpi, x, bsize, pc_tree->partitioning,
+ best_rdc.rdcost, horz_rd, vert_rd, split_rd,
+ &partition_horz4_allowed, &partition_vert4_allowed,
+ pb_source_variance, mi_row, mi_col);
}
#if CONFIG_DIST_8X8
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index e979a49..0a96914 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -949,3 +949,454 @@
if (score < thresh) *terminate_partition_search = 1;
}
#undef FEATURES
+
+#define FEATURE_SIZE 19
+int av1_ml_prune_2pass_split_partition(const PC_TREE_STATS *pc_tree_stats,
+ BLOCK_SIZE bsize, int *split_score,
+ int *none_score) {
+ if (!pc_tree_stats->valid) return 0;
+ const float *split_weights = NULL;
+ const float *none_weights = NULL;
+ switch (bsize) {
+ case BLOCK_4X4: break;
+ case BLOCK_8X8:
+ split_weights = av1_2pass_split_partition_weights_8;
+ none_weights = av1_2pass_none_partition_weights_8;
+ break;
+ case BLOCK_16X16:
+ split_weights = av1_2pass_split_partition_weights_16;
+ none_weights = av1_2pass_none_partition_weights_16;
+ break;
+ case BLOCK_32X32:
+ split_weights = av1_2pass_split_partition_weights_32;
+ none_weights = av1_2pass_none_partition_weights_32;
+ break;
+ case BLOCK_64X64:
+ split_weights = av1_2pass_split_partition_weights_64;
+ none_weights = av1_2pass_none_partition_weights_64;
+ break;
+ case BLOCK_128X128:
+ split_weights = av1_2pass_split_partition_weights_128;
+ none_weights = av1_2pass_none_partition_weights_128;
+ break;
+ default: assert(0 && "Unexpected bsize.");
+ }
+ if (!split_weights || !none_weights) return 0;
+
+ aom_clear_system_state();
+
+ float features[FEATURE_SIZE];
+ int feature_index = 0;
+ features[feature_index++] = (float)pc_tree_stats->split;
+ features[feature_index++] = (float)pc_tree_stats->skip;
+ const int rdcost = (int)AOMMIN(INT_MAX, pc_tree_stats->rdcost);
+ const int rd_valid = rdcost > 0 && rdcost < 1000000000;
+ features[feature_index++] = (float)rd_valid;
+ for (int i = 0; i < 4; ++i) {
+ features[feature_index++] = (float)pc_tree_stats->sub_block_split[i];
+ features[feature_index++] = (float)pc_tree_stats->sub_block_skip[i];
+ const int sub_rdcost =
+ (int)AOMMIN(INT_MAX, pc_tree_stats->sub_block_rdcost[i]);
+ const int sub_rd_valid = sub_rdcost > 0 && sub_rdcost < 1000000000;
+ features[feature_index++] = (float)sub_rd_valid;
+ // Ratio between the sub-block RD and the whole-block RD.
+ float rd_ratio = 1.0f;
+ if (rd_valid && sub_rd_valid && sub_rdcost < rdcost)
+ rd_ratio = (float)sub_rdcost / (float)rdcost;
+ features[feature_index++] = rd_ratio;
+ }
+ assert(feature_index == FEATURE_SIZE);
+
+ float score_1 = split_weights[FEATURE_SIZE];
+ float score_2 = none_weights[FEATURE_SIZE];
+ for (int i = 0; i < FEATURE_SIZE; ++i) {
+ score_1 += features[i] * split_weights[i];
+ score_2 += features[i] * none_weights[i];
+ }
+ *split_score = (int)(score_1 * 100);
+ *none_score = (int)(score_2 * 100);
+ return 1;
+}
+#undef FEATURE_SIZE
+
+void av1_ml_prune_rect_partition(const AV1_COMP *const cpi,
+ const MACROBLOCK *const x, BLOCK_SIZE bsize,
+ int64_t best_rd, int64_t none_rd,
+ int64_t *split_rd, int *const dst_prune_horz,
+ int *const dst_prune_vert) {
+ if (bsize < BLOCK_8X8 || best_rd >= 1000000000) return;
+ best_rd = AOMMAX(best_rd, 1);
+ const NN_CONFIG *nn_config = NULL;
+ const float prob_thresholds[5] = { 0.01f, 0.01f, 0.004f, 0.002f, 0.002f };
+ float cur_thresh = 0.0f;
+ switch (bsize) {
+ case BLOCK_8X8:
+ nn_config = &av1_rect_partition_nnconfig_8;
+ cur_thresh = prob_thresholds[0];
+ break;
+ case BLOCK_16X16:
+ nn_config = &av1_rect_partition_nnconfig_16;
+ cur_thresh = prob_thresholds[1];
+ break;
+ case BLOCK_32X32:
+ nn_config = &av1_rect_partition_nnconfig_32;
+ cur_thresh = prob_thresholds[2];
+ break;
+ case BLOCK_64X64:
+ nn_config = &av1_rect_partition_nnconfig_64;
+ cur_thresh = prob_thresholds[3];
+ break;
+ case BLOCK_128X128:
+ nn_config = &av1_rect_partition_nnconfig_128;
+ cur_thresh = prob_thresholds[4];
+ break;
+ default: assert(0 && "Unexpected bsize.");
+ }
+ if (!nn_config) return;
+ aom_clear_system_state();
+
+ // 1. Compute input features
+ float features[9];
+
+ // RD cost ratios
+ for (int i = 0; i < 5; i++) features[i] = 1.0f;
+ if (none_rd > 0 && none_rd < 1000000000)
+ features[0] = (float)none_rd / (float)best_rd;
+ for (int i = 0; i < 4; i++) {
+ if (split_rd[i] > 0 && split_rd[i] < 1000000000)
+ features[1 + i] = (float)split_rd[i] / (float)best_rd;
+ }
+
+ // Variance ratios
+ const MACROBLOCKD *const xd = &x->e_mbd;
+ int whole_block_variance;
+ if (is_cur_buf_hbd(xd)) {
+ whole_block_variance = av1_high_get_sby_perpixel_variance(
+ cpi, &x->plane[0].src, bsize, xd->bd);
+ } else {
+ whole_block_variance =
+ av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize);
+ }
+ whole_block_variance = AOMMAX(whole_block_variance, 1);
+
+ int split_variance[4];
+ const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
+ struct buf_2d buf;
+ buf.stride = x->plane[0].src.stride;
+ const int bw = block_size_wide[bsize];
+ for (int i = 0; i < 4; ++i) {
+ const int x_idx = (i & 1) * bw / 2;
+ const int y_idx = (i >> 1) * bw / 2;
+ buf.buf = x->plane[0].src.buf + x_idx + y_idx * buf.stride;
+ if (is_cur_buf_hbd(xd)) {
+ split_variance[i] =
+ av1_high_get_sby_perpixel_variance(cpi, &buf, subsize, xd->bd);
+ } else {
+ split_variance[i] = av1_get_sby_perpixel_variance(cpi, &buf, subsize);
+ }
+ }
+
+ for (int i = 0; i < 4; i++)
+ features[5 + i] = (float)split_variance[i] / (float)whole_block_variance;
+
+ // 2. Do the prediction and prune 0-2 partitions based on their probabilities
+ float raw_scores[3] = { 0.0f };
+ av1_nn_predict(features, nn_config, raw_scores);
+ aom_clear_system_state();
+ float probs[3] = { 0.0f };
+ av1_nn_softmax(raw_scores, probs, 3);
+
+ // probs[0] is the probability of the fact that both rectangular partitions
+ // are worse than current best_rd
+ if (probs[1] <= cur_thresh) (*dst_prune_horz) = 1;
+ if (probs[2] <= cur_thresh) (*dst_prune_vert) = 1;
+}
+
+// Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
+// considered.
+void av1_ml_prune_ab_partition(BLOCK_SIZE bsize, int part_ctx, int var_ctx,
+ int64_t best_rd, int64_t horz_rd[2],
+ int64_t vert_rd[2], int64_t split_rd[4],
+ int *const horza_partition_allowed,
+ int *const horzb_partition_allowed,
+ int *const verta_partition_allowed,
+ int *const vertb_partition_allowed) {
+ if (bsize < BLOCK_8X8 || best_rd >= 1000000000) return;
+ const NN_CONFIG *nn_config = NULL;
+ switch (bsize) {
+ case BLOCK_8X8: nn_config = NULL; break;
+ case BLOCK_16X16: nn_config = &av1_ab_partition_nnconfig_16; break;
+ case BLOCK_32X32: nn_config = &av1_ab_partition_nnconfig_32; break;
+ case BLOCK_64X64: nn_config = &av1_ab_partition_nnconfig_64; break;
+ case BLOCK_128X128: nn_config = &av1_ab_partition_nnconfig_128; break;
+ default: assert(0 && "Unexpected bsize.");
+ }
+ if (!nn_config) return;
+
+ aom_clear_system_state();
+
+ // Generate features.
+ float features[10];
+ int feature_index = 0;
+ features[feature_index++] = (float)part_ctx;
+ features[feature_index++] = (float)var_ctx;
+ const int rdcost = (int)AOMMIN(INT_MAX, best_rd);
+ int sub_block_rdcost[8] = { 0 };
+ int rd_index = 0;
+ for (int i = 0; i < 2; ++i) {
+ if (horz_rd[i] > 0 && horz_rd[i] < 1000000000)
+ sub_block_rdcost[rd_index] = (int)horz_rd[i];
+ ++rd_index;
+ }
+ for (int i = 0; i < 2; ++i) {
+ if (vert_rd[i] > 0 && vert_rd[i] < 1000000000)
+ sub_block_rdcost[rd_index] = (int)vert_rd[i];
+ ++rd_index;
+ }
+ for (int i = 0; i < 4; ++i) {
+ if (split_rd[i] > 0 && split_rd[i] < 1000000000)
+ sub_block_rdcost[rd_index] = (int)split_rd[i];
+ ++rd_index;
+ }
+ for (int i = 0; i < 8; ++i) {
+ // Ratio between the sub-block RD and the whole-block RD.
+ float rd_ratio = 1.0f;
+ if (sub_block_rdcost[i] > 0 && sub_block_rdcost[i] < rdcost)
+ rd_ratio = (float)sub_block_rdcost[i] / (float)rdcost;
+ features[feature_index++] = rd_ratio;
+ }
+ assert(feature_index == 10);
+
+ // Calculate scores using the NN model.
+ float score[16] = { 0.0f };
+ av1_nn_predict(features, nn_config, score);
+ aom_clear_system_state();
+ int int_score[16];
+ int max_score = -1000;
+ for (int i = 0; i < 16; ++i) {
+ int_score[i] = (int)(100 * score[i]);
+ max_score = AOMMAX(int_score[i], max_score);
+ }
+
+ // Make decisions based on the model scores.
+ int thresh = max_score;
+ switch (bsize) {
+ case BLOCK_16X16: thresh -= 150; break;
+ case BLOCK_32X32: thresh -= 100; break;
+ default: break;
+ }
+ *horza_partition_allowed = 0;
+ *horzb_partition_allowed = 0;
+ *verta_partition_allowed = 0;
+ *vertb_partition_allowed = 0;
+ for (int i = 0; i < 16; ++i) {
+ if (int_score[i] >= thresh) {
+ if ((i >> 0) & 1) *horza_partition_allowed = 1;
+ if ((i >> 1) & 1) *horzb_partition_allowed = 1;
+ if ((i >> 2) & 1) *verta_partition_allowed = 1;
+ if ((i >> 3) & 1) *vertb_partition_allowed = 1;
+ }
+ }
+}
+
+#define FEATURES 18
+#define LABELS 4
+// Use a ML model to predict if horz4 and vert4 should be considered.
+void av1_ml_prune_4_partition(const AV1_COMP *const cpi, MACROBLOCK *const x,
+ BLOCK_SIZE bsize, int part_ctx, int64_t best_rd,
+ int64_t horz_rd[2], int64_t vert_rd[2],
+ int64_t split_rd[4],
+ int *const partition_horz4_allowed,
+ int *const partition_vert4_allowed,
+ unsigned int pb_source_variance, int mi_row,
+ int mi_col) {
+ if (best_rd >= 1000000000) return;
+ const NN_CONFIG *nn_config = NULL;
+ switch (bsize) {
+ case BLOCK_16X16: nn_config = &av1_4_partition_nnconfig_16; break;
+ case BLOCK_32X32: nn_config = &av1_4_partition_nnconfig_32; break;
+ case BLOCK_64X64: nn_config = &av1_4_partition_nnconfig_64; break;
+ default: assert(0 && "Unexpected bsize.");
+ }
+ if (!nn_config) return;
+
+ aom_clear_system_state();
+
+ // Generate features.
+ float features[FEATURES];
+ int feature_index = 0;
+ features[feature_index++] = (float)part_ctx;
+ features[feature_index++] = (float)get_unsigned_bits(pb_source_variance);
+
+ const int rdcost = (int)AOMMIN(INT_MAX, best_rd);
+ int sub_block_rdcost[8] = { 0 };
+ int rd_index = 0;
+ for (int i = 0; i < 2; ++i) {
+ if (horz_rd[i] > 0 && horz_rd[i] < 1000000000)
+ sub_block_rdcost[rd_index] = (int)horz_rd[i];
+ ++rd_index;
+ }
+ for (int i = 0; i < 2; ++i) {
+ if (vert_rd[i] > 0 && vert_rd[i] < 1000000000)
+ sub_block_rdcost[rd_index] = (int)vert_rd[i];
+ ++rd_index;
+ }
+ for (int i = 0; i < 4; ++i) {
+ if (split_rd[i] > 0 && split_rd[i] < 1000000000)
+ sub_block_rdcost[rd_index] = (int)split_rd[i];
+ ++rd_index;
+ }
+ for (int i = 0; i < 8; ++i) {
+ // Ratio between the sub-block RD and the whole-block RD.
+ float rd_ratio = 1.0f;
+ if (sub_block_rdcost[i] > 0 && sub_block_rdcost[i] < rdcost)
+ rd_ratio = (float)sub_block_rdcost[i] / (float)rdcost;
+ features[feature_index++] = rd_ratio;
+ }
+
+ // Get variance of the 1:4 and 4:1 sub-blocks.
+ unsigned int horz_4_source_var[4] = { 0 };
+ unsigned int vert_4_source_var[4] = { 0 };
+ {
+ BLOCK_SIZE horz_4_bs = get_partition_subsize(bsize, PARTITION_HORZ_4);
+ BLOCK_SIZE vert_4_bs = get_partition_subsize(bsize, PARTITION_VERT_4);
+ av1_setup_src_planes(x, cpi->source, mi_row, mi_col,
+ av1_num_planes(&cpi->common), bsize);
+ const int src_stride = x->plane[0].src.stride;
+ uint8_t *src = x->plane[0].src.buf;
+ const MACROBLOCKD *const xd = &x->e_mbd;
+
+ struct buf_2d horz_4_src, vert_4_src;
+ horz_4_src.stride = src_stride;
+ vert_4_src.stride = src_stride;
+
+ for (int i = 0; i < 4; ++i) {
+ horz_4_src.buf = src + i * block_size_high[horz_4_bs] * src_stride;
+ vert_4_src.buf = src + i * block_size_wide[vert_4_bs];
+
+ if (is_cur_buf_hbd(xd)) {
+ horz_4_source_var[i] = av1_high_get_sby_perpixel_variance(
+ cpi, &horz_4_src, horz_4_bs, xd->bd);
+ vert_4_source_var[i] = av1_high_get_sby_perpixel_variance(
+ cpi, &vert_4_src, vert_4_bs, xd->bd);
+ } else {
+ horz_4_source_var[i] =
+ av1_get_sby_perpixel_variance(cpi, &horz_4_src, horz_4_bs);
+ vert_4_source_var[i] =
+ av1_get_sby_perpixel_variance(cpi, &vert_4_src, vert_4_bs);
+ }
+ }
+ }
+
+ const float denom = (float)(pb_source_variance + 1);
+ const float low_b = 0.1f;
+ const float high_b = 10.0f;
+ for (int i = 0; i < 4; ++i) {
+ // Ratio between the 4:1 sub-block variance and the whole-block variance.
+ float var_ratio = (float)(horz_4_source_var[i] + 1) / denom;
+ if (var_ratio < low_b) var_ratio = low_b;
+ if (var_ratio > high_b) var_ratio = high_b;
+ features[feature_index++] = var_ratio;
+ }
+ for (int i = 0; i < 4; ++i) {
+ // Ratio between the 1:4 sub-block RD and the whole-block RD.
+ float var_ratio = (float)(vert_4_source_var[i] + 1) / denom;
+ if (var_ratio < low_b) var_ratio = low_b;
+ if (var_ratio > high_b) var_ratio = high_b;
+ features[feature_index++] = var_ratio;
+ }
+ assert(feature_index == FEATURES);
+
+ // Calculate scores using the NN model.
+ float score[LABELS] = { 0.0f };
+ av1_nn_predict(features, nn_config, score);
+ aom_clear_system_state();
+ int int_score[LABELS];
+ int max_score = -1000;
+ for (int i = 0; i < LABELS; ++i) {
+ int_score[i] = (int)(100 * score[i]);
+ max_score = AOMMAX(int_score[i], max_score);
+ }
+
+ // Make decisions based on the model scores.
+ int thresh = max_score;
+ switch (bsize) {
+ case BLOCK_16X16: thresh -= 500; break;
+ case BLOCK_32X32: thresh -= 500; break;
+ case BLOCK_64X64: thresh -= 200; break;
+ default: break;
+ }
+ *partition_horz4_allowed = 0;
+ *partition_vert4_allowed = 0;
+ for (int i = 0; i < LABELS; ++i) {
+ if (int_score[i] >= thresh) {
+ if ((i >> 0) & 1) *partition_horz4_allowed = 1;
+ if ((i >> 1) & 1) *partition_vert4_allowed = 1;
+ }
+ }
+}
+#undef FEATURES
+#undef LABELS
+
+#define FEATURES 4
+int av1_ml_predict_breakout(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
+ const MACROBLOCK *const x,
+ const RD_STATS *const rd_stats,
+ unsigned int pb_source_variance) {
+ const NN_CONFIG *nn_config = NULL;
+ int thresh = 0;
+ switch (bsize) {
+ case BLOCK_8X8:
+ nn_config = &av1_partition_breakout_nnconfig_8;
+ thresh = cpi->sf.ml_partition_search_breakout_thresh[0];
+ break;
+ case BLOCK_16X16:
+ nn_config = &av1_partition_breakout_nnconfig_16;
+ thresh = cpi->sf.ml_partition_search_breakout_thresh[1];
+ break;
+ case BLOCK_32X32:
+ nn_config = &av1_partition_breakout_nnconfig_32;
+ thresh = cpi->sf.ml_partition_search_breakout_thresh[2];
+ break;
+ case BLOCK_64X64:
+ nn_config = &av1_partition_breakout_nnconfig_64;
+ thresh = cpi->sf.ml_partition_search_breakout_thresh[3];
+ break;
+ case BLOCK_128X128:
+ nn_config = &av1_partition_breakout_nnconfig_128;
+ thresh = cpi->sf.ml_partition_search_breakout_thresh[4];
+ break;
+ default: assert(0 && "Unexpected bsize.");
+ }
+ if (!nn_config || thresh < 0) return 0;
+
+ // Generate feature values.
+ float features[FEATURES];
+ int feature_index = 0;
+ aom_clear_system_state();
+
+ const int num_pels_log2 = num_pels_log2_lookup[bsize];
+ float rate_f = (float)AOMMIN(rd_stats->rate, INT_MAX);
+ rate_f = ((float)x->rdmult / 128.0f / 512.0f / (float)(1 << num_pels_log2)) *
+ rate_f;
+ features[feature_index++] = rate_f;
+
+ const float dist_f =
+ (float)(AOMMIN(rd_stats->dist, INT_MAX) >> num_pels_log2);
+ features[feature_index++] = dist_f;
+
+ features[feature_index++] = (float)pb_source_variance;
+
+ const int dc_q = (int)x->plane[0].dequant_QTX[0];
+ features[feature_index++] = (float)(dc_q * dc_q) / 256.0f;
+ assert(feature_index == FEATURES);
+
+ // Calculate score using the NN model.
+ float score = 0.0f;
+ av1_nn_predict(features, nn_config, &score);
+ aom_clear_system_state();
+
+ // Make decision.
+ return (int)(score * 100) >= thresh;
+}
+#undef FEATURES
diff --git a/av1/encoder/partition_strategy.h b/av1/encoder/partition_strategy.h
index 064f530..fbe832d 100644
--- a/av1/encoder/partition_strategy.h
+++ b/av1/encoder/partition_strategy.h
@@ -85,6 +85,61 @@
BLOCK_SIZE av1_predict_max_partition(AV1_COMP *const cpi, MACROBLOCK *const x,
const float *features);
+// Attempts an early termination after PARTITION_SPLIT.
+void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
+ PC_TREE *const pc_tree, BLOCK_SIZE bsize,
+ int64_t best_rd, int64_t part_none_rd,
+ int64_t part_split_rd,
+ int64_t *split_block_rd, int mi_row,
+ int mi_col,
+ int *const terminate_partition_search);
+
+// Use data from first partition pass to emit split_scores and none_scores.
+// Returns 0 if the firstpass data is not valid, 1 otherwise.
+// split_score indicates confidence of picking split partition;
+// none_score indicates confidence of picking none partition;
+int av1_ml_prune_2pass_split_partition(const PC_TREE_STATS *pc_tree_stats,
+ BLOCK_SIZE bsize, int *split_score,
+ int *none_score);
+
+// Use the rdcost ratio and source var ratio to prune PARTITION_HORZ and
+// PARTITION_VERT.
+// TODO(chiyotsai@google.com): Currently this model does not use q value and has
+// no information about rectangular partitions. Preliminary experiments suggest
+// that we can get better performance by adding in q_index and rectangular
+// sse/var from SMS. We should retrain and tune this model later.
+void av1_ml_prune_rect_partition(const AV1_COMP *const cpi,
+ const MACROBLOCK *const x, BLOCK_SIZE bsize,
+ int64_t best_rd, int64_t none_rd,
+ int64_t *split_rd, int *const dst_prune_horz,
+ int *const dst_prune_vert);
+
+// Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
+// considered.
+void av1_ml_prune_ab_partition(BLOCK_SIZE bsize, int part_ctx, int var_ctx,
+ int64_t best_rd, int64_t horz_rd[2],
+ int64_t vert_rd[2], int64_t split_rd[4],
+ int *const horza_partition_allowed,
+ int *const horzb_partition_allowed,
+ int *const verta_partition_allowed,
+ int *const vertb_partition_allowed);
+
+// Use a ML model to predict if horz4 and vert4 should be considered.
+void av1_ml_prune_4_partition(const AV1_COMP *const cpi, MACROBLOCK *const x,
+ BLOCK_SIZE bsize, int part_ctx, int64_t best_rd,
+ int64_t horz_rd[2], int64_t vert_rd[2],
+ int64_t split_rd[4],
+ int *const partition_horz4_allowed,
+ int *const partition_vert4_allowed,
+ unsigned int pb_source_variance, int mi_row,
+ int mi_col);
+
+// ML-based partition search breakout after PARTITION_NONE
+int av1_ml_predict_breakout(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
+ const MACROBLOCK *const x,
+ const RD_STATS *const rd_stats,
+ unsigned int pb_source_variance);
+
// A simplified version of set_offsets meant to be used for
// simple_motion_search.
static INLINE void set_offsets_for_motion_search(const AV1_COMP *const cpi,
@@ -168,12 +223,4 @@
INTNL_OVERLAY_UPDATE;
}
-void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
- PC_TREE *const pc_tree, BLOCK_SIZE bsize,
- int64_t best_rd, int64_t part_none_rd,
- int64_t part_split_rd,
- int64_t *split_block_rd, int mi_row,
- int mi_col,
- int *const terminate_partition_search);
-
#endif // AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
