Make max and min partition more well-defined
We redefine the max and min partition in the recursive rd loop as
the max and min square nodes rd_pick_partition() will reach.
Previous implementation only searches NONE when we reach min_bsize.
Now all modes except SPLIT are allowed when bsize reaches the min.
Change-Id: I8a86fed35aa7c057824a8b9527c49142548dbba0
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index a35c136..8ce3119 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -3535,7 +3535,7 @@
static void rd_pick_partition(AV1_COMP *const cpi, ThreadData *td,
TileDataEnc *tile_data, TOKENEXTRA **tp,
int mi_row, int mi_col, BLOCK_SIZE bsize,
- BLOCK_SIZE max_bsize, BLOCK_SIZE min_bsize,
+ BLOCK_SIZE max_sq_part, BLOCK_SIZE min_sq_part,
RD_STATS *rd_cost, int64_t best_rd,
PC_TREE *pc_tree, int64_t *none_rd) {
const AV1_COMMON *const cm = &cpi->common;
@@ -3569,14 +3569,15 @@
int vert_ctx_is_ready = 0;
BLOCK_SIZE bsize2 = get_partition_subsize(bsize, PARTITION_SPLIT);
- // If max and min partition sizes (squares only) are enforced,
- // only PARTITION_NONE is allowed if the current node equals min_bsize,
- // only PARTITION_SPLIT is allowed if the current node exceeds max_bsize.
- assert(block_size_wide[min_bsize] == block_size_high[min_bsize]);
- assert(block_size_wide[max_bsize] == block_size_high[max_bsize]);
- assert(min_bsize <= max_bsize);
- int is_eq_min_bsize = bsize == min_bsize;
- int is_gt_max_bsize = bsize > max_bsize;
+ // Max and min square partition levels are defined as the partition nodes that
+ // the recursive function rd_pick_partition() can reach. To implement this:
+ // only PARTITION_SPLIT is NOT allowed if the current node equals min_sq_part,
+ // only PARTITION_SPLIT is allowed if the current node exceeds max_sq_part.
+ assert(block_size_wide[min_sq_part] == block_size_high[min_sq_part]);
+ assert(block_size_wide[max_sq_part] == block_size_high[max_sq_part]);
+ assert(min_sq_part <= max_sq_part);
+ int is_eq_min_sq_part = bsize == min_sq_part;
+ int is_gt_max_sq_part = bsize > max_sq_part;
if (best_rd < 0) {
pc_tree->none.rdcost = INT64_MAX;
@@ -3868,7 +3869,8 @@
#endif
// PARTITION_NONE
- if ((partition_none_allowed && !is_gt_max_bsize) || is_eq_min_bsize) {
+ if (is_eq_min_sq_part) partition_none_allowed = 1;
+ if (partition_none_allowed && !is_gt_max_sq_part) {
int pt_cost = 0;
if (bsize_at_least_8x8) {
pt_cost = partition_cost[PARTITION_NONE] < INT_MAX
@@ -4008,7 +4010,8 @@
if (cpi->sf.adaptive_motion_search) store_pred_mv(x, ctx_none);
// PARTITION_SPLIT
- if ((do_square_split && !is_eq_min_bsize) || is_gt_max_bsize) {
+ if (is_eq_min_sq_part) do_square_split = 0;
+ if (do_square_split || is_gt_max_sq_part) {
av1_init_rd_stats(&sum_rdc);
subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
sum_rdc.rate = partition_cost[PARTITION_SPLIT];
@@ -4032,7 +4035,7 @@
if (cpi->sf.prune_ref_frame_for_rect_partitions)
pc_tree->split[idx]->none.rate = INT_MAX;
rd_pick_partition(cpi, td, tile_data, tp, mi_row + y_idx, mi_col + x_idx,
- subsize, max_bsize, min_bsize, &this_rdc,
+ subsize, max_sq_part, min_sq_part, &this_rdc,
best_remain_rdcost, pc_tree->split[idx], p_split_rd);
if (this_rdc.rate == INT_MAX) {
@@ -4184,7 +4187,7 @@
// PARTITION_HORZ
if (partition_horz_allowed && !prune_horz &&
(do_rectangular_split || active_h_edge(cpi, mi_row, mi_step)) &&
- !is_eq_min_bsize && !is_gt_max_bsize) {
+ !is_gt_max_sq_part) {
av1_init_rd_stats(&sum_rdc);
subsize = get_partition_subsize(bsize, PARTITION_HORZ);
if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_none);
@@ -4258,7 +4261,7 @@
// PARTITION_VERT
if (partition_vert_allowed && !prune_vert &&
(do_rectangular_split || active_v_edge(cpi, mi_col, mi_step)) &&
- !is_eq_min_bsize && !is_gt_max_bsize) {
+ !is_gt_max_sq_part) {
av1_init_rd_stats(&sum_rdc);
subsize = get_partition_subsize(bsize, PARTITION_VERT);
@@ -4434,8 +4437,7 @@
}
// PARTITION_HORZ_A
- if (partition_horz_allowed && horza_partition_allowed && !is_eq_min_bsize &&
- !is_gt_max_bsize) {
+ if (partition_horz_allowed && horza_partition_allowed && !is_gt_max_sq_part) {
subsize = get_partition_subsize(bsize, PARTITION_HORZ_A);
pc_tree->horizontala[0].rd_mode_is_ready = 0;
pc_tree->horizontala[1].rd_mode_is_ready = 0;
@@ -4500,8 +4502,7 @@
restore_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);
}
// PARTITION_HORZ_B
- if (partition_horz_allowed && horzb_partition_allowed && !is_eq_min_bsize &&
- !is_gt_max_bsize) {
+ if (partition_horz_allowed && horzb_partition_allowed && !is_gt_max_sq_part) {
subsize = get_partition_subsize(bsize, PARTITION_HORZ_B);
pc_tree->horizontalb[0].rd_mode_is_ready = 0;
pc_tree->horizontalb[1].rd_mode_is_ready = 0;
@@ -4561,8 +4562,7 @@
}
// PARTITION_VERT_A
- if (partition_vert_allowed && verta_partition_allowed && !is_eq_min_bsize &&
- !is_gt_max_bsize) {
+ if (partition_vert_allowed && verta_partition_allowed && !is_gt_max_sq_part) {
subsize = get_partition_subsize(bsize, PARTITION_VERT_A);
pc_tree->verticala[0].rd_mode_is_ready = 0;
pc_tree->verticala[1].rd_mode_is_ready = 0;
@@ -4618,8 +4618,7 @@
restore_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);
}
// PARTITION_VERT_B
- if (partition_vert_allowed && vertb_partition_allowed && !is_eq_min_bsize &&
- !is_gt_max_bsize) {
+ if (partition_vert_allowed && vertb_partition_allowed && !is_gt_max_sq_part) {
subsize = get_partition_subsize(bsize, PARTITION_VERT_B);
pc_tree->verticalb[0].rd_mode_is_ready = 0;
pc_tree->verticalb[1].rd_mode_is_ready = 0;
@@ -4715,7 +4714,7 @@
// PARTITION_HORZ_4
if (partition_horz4_allowed && has_rows &&
(do_rectangular_split || active_h_edge(cpi, mi_row, mi_step)) &&
- !is_eq_min_bsize && !is_gt_max_bsize) {
+ !is_gt_max_sq_part) {
av1_init_rd_stats(&sum_rdc);
const int quarter_step = mi_size_high[bsize] / 4;
PICK_MODE_CONTEXT *ctx_prev = ctx_none;
@@ -4760,7 +4759,7 @@
// PARTITION_VERT_4
if (partition_vert4_allowed && has_cols &&
(do_rectangular_split || active_v_edge(cpi, mi_row, mi_step)) &&
- !is_eq_min_bsize && !is_gt_max_bsize) {
+ !is_gt_max_sq_part) {
av1_init_rd_stats(&sum_rdc);
const int quarter_step = mi_size_wide[bsize] / 4;
PICK_MODE_CONTEXT *ctx_prev = ctx_none;