Abstract common code in partition4 search
BUG=aomedia:2687
Change-Id: Ice476e0e579f58df2220741c1ad38757c4068bf9
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 5576885..2f84590 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -3408,6 +3408,95 @@
}
}
+// Set mi positions for HORZ4 / VERT4 sub-block partitions.
+static AOM_INLINE void set_mi_pos_partition4(
+ const int inc_step[NUM_PART4_TYPES], int mi_pos[PART4_SUB_PARTITIONS][2],
+ const int mi_row, const int mi_col) {
+ for (int i = 0; i < PART4_SUB_PARTITIONS; i++) {
+ mi_pos[i][0] = mi_row + i * inc_step[HORZ4];
+ mi_pos[i][1] = mi_col + i * inc_step[VERT4];
+ }
+}
+
+// Set context and RD cost for HORZ4 / VERT4 partition types.
+static AOM_INLINE void set_4_part_ctx_and_rdcost(
+ MACROBLOCK *x, const AV1_COMMON *const cm, ThreadData *td,
+ PICK_MODE_CONTEXT *cur_part_ctx[PART4_SUB_PARTITIONS],
+ PartitionSearchState *part_search_state, PARTITION_TYPE partition_type,
+ BLOCK_SIZE bsize) {
+ // Initialize sum_rdc RD cost structure.
+ av1_init_rd_stats(&part_search_state->sum_rdc);
+ const int subsize = get_partition_subsize(bsize, partition_type);
+ part_search_state->sum_rdc.rate =
+ part_search_state->partition_cost[partition_type];
+ part_search_state->sum_rdc.rdcost =
+ RDCOST(x->rdmult, part_search_state->sum_rdc.rate, 0);
+ for (int i = 0; i < PART4_SUB_PARTITIONS; ++i)
+ cur_part_ctx[i] = av1_alloc_pmc(cm, subsize, &td->shared_coeff_buf);
+}
+
+// Partition search of HORZ4 / VERT4 partition types.
+static AOM_INLINE void rd_pick_4partition(
+ AV1_COMP *const cpi, ThreadData *td, TileDataEnc *tile_data,
+ TokenExtra **tp, MACROBLOCK *x, RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx,
+ PC_TREE *pc_tree, PICK_MODE_CONTEXT *cur_part_ctx[PART4_SUB_PARTITIONS],
+ PartitionSearchState *part_search_state, RD_STATS *best_rdc,
+ const int inc_step[NUM_PART4_TYPES], PARTITION_TYPE partition_type) {
+ const AV1_COMMON *const cm = &cpi->common;
+ PartitionBlkParams blk_params = part_search_state->part_blk_params;
+ // mi positions needed for HORZ4 and VERT4 partition types.
+ int mi_pos_check[NUM_PART4_TYPES] = { cm->mi_params.mi_rows,
+ cm->mi_params.mi_cols };
+ const PART4_TYPES part4_idx = (partition_type != PARTITION_HORZ_4);
+ int mi_pos[PART4_SUB_PARTITIONS][2];
+
+ blk_params.subsize = get_partition_subsize(blk_params.bsize, partition_type);
+ // Set partition context and RD cost.
+ set_4_part_ctx_and_rdcost(x, cm, td, cur_part_ctx, part_search_state,
+ partition_type, blk_params.bsize);
+ // Set mi positions for sub-block sizes.
+ set_mi_pos_partition4(inc_step, mi_pos, blk_params.mi_row, blk_params.mi_col);
+#if CONFIG_COLLECT_PARTITION_STATS
+ if (best_rdc.rdcost - part_search_state->sum_rdc.rdcost >= 0) {
+ partition_attempts[partition_type] += 1;
+ aom_usec_timer_start(&partition_timer);
+ partition_timer_on = 1;
+ }
+#endif
+ // Loop over sub-block partitions.
+ for (int i = 0; i < PART4_SUB_PARTITIONS; ++i) {
+ if (i > 0 && mi_pos[i][part4_idx] >= mi_pos_check[part4_idx]) break;
+
+ // Sub-block evaluation of Horz4 / Vert4 partition type.
+ cur_part_ctx[i]->rd_mode_is_ready = 0;
+ if (!rd_try_subblock(
+ cpi, td, tile_data, tp, (i == PART4_SUB_PARTITIONS - 1),
+ mi_pos[i][0], mi_pos[i][1], blk_params.subsize, *best_rdc,
+ &part_search_state->sum_rdc, partition_type, cur_part_ctx[i])) {
+ av1_invalid_rd_stats(&part_search_state->sum_rdc);
+ break;
+ }
+ }
+
+ // Calculate the total cost and update the best partition.
+ av1_rd_cost_update(x->rdmult, &part_search_state->sum_rdc);
+ if (part_search_state->sum_rdc.rdcost < best_rdc->rdcost) {
+ *best_rdc = part_search_state->sum_rdc;
+ part_search_state->found_best_partition = true;
+ pc_tree->partitioning = partition_type;
+ }
+#if CONFIG_COLLECT_PARTITION_STATS
+ if (partition_timer_on) {
+ aom_usec_timer_mark(&partition_timer);
+ int64_t time = aom_usec_timer_elapsed(&partition_timer);
+ partition_times[partition_type] += time;
+ partition_timer_on = 0;
+ }
+#endif
+ restore_context(x, x_ctx, blk_params.mi_row, blk_params.mi_col,
+ blk_params.bsize, av1_num_planes(cm));
+}
+
/*!\brief AV1 block partition search (full search).
*
* \ingroup partition_search
@@ -3944,7 +4033,7 @@
(partition_horz4_allowed || partition_vert4_allowed)) {
// Count of child blocks in which HORZ or VERT partition has won
int num_child_horz_win = 0, num_child_vert_win = 0;
- for (int idx = 0; idx < 4; idx++) {
+ for (int idx = 0; idx < PART4_SUB_PARTITIONS; idx++) {
num_child_horz_win +=
(part_search_state.split_part_rect_win[idx].rect_part_win[HORZ]) ? 1
: 0;
@@ -3971,61 +4060,12 @@
partition_horz4_allowed && blk_params.has_rows &&
(part_search_state.do_rectangular_split ||
active_h_edge(cpi, mi_row, blk_params.mi_step))) {
- av1_init_rd_stats(&part_search_state.sum_rdc);
- const int quarter_step = mi_size_high[bsize] / 4;
-
- blk_params.subsize = get_partition_subsize(bsize, PARTITION_HORZ_4);
- part_search_state.sum_rdc.rate =
- part_search_state.partition_cost[PARTITION_HORZ_4];
- part_search_state.sum_rdc.rdcost =
- RDCOST(x->rdmult, part_search_state.sum_rdc.rate, 0);
-
- for (int i = 0; i < 4; ++i) {
- pc_tree->horizontal4[i] =
- av1_alloc_pmc(cm, blk_params.subsize, &td->shared_coeff_buf);
- }
-
-#if CONFIG_COLLECT_PARTITION_STATS
- if (best_rdc.rdcost - part_search_state.sum_rdc.rdcost >= 0) {
- partition_attempts[PARTITION_HORZ_4] += 1;
- aom_usec_timer_start(&partition_timer);
- partition_timer_on = 1;
- }
-#endif
- for (int i = 0; i < 4; ++i) {
- const int this_mi_row = mi_row + i * quarter_step;
-
- if (i > 0 && this_mi_row >= mi_params->mi_rows) break;
-
- PICK_MODE_CONTEXT *ctx_this = pc_tree->horizontal4[i];
-
- ctx_this->rd_mode_is_ready = 0;
- if (!rd_try_subblock(cpi, td, tile_data, tp, (i == 3), this_mi_row,
- mi_col, blk_params.subsize, best_rdc,
- &part_search_state.sum_rdc, PARTITION_HORZ_4,
- ctx_this)) {
- av1_invalid_rd_stats(&part_search_state.sum_rdc);
- break;
- }
- }
-
- // Calculate the total cost and update the best partition.
- av1_rd_cost_update(x->rdmult, &part_search_state.sum_rdc);
- if (part_search_state.sum_rdc.rdcost < best_rdc.rdcost) {
- best_rdc = part_search_state.sum_rdc;
- part_search_state.found_best_partition = true;
- pc_tree->partitioning = PARTITION_HORZ_4;
- }
-
-#if CONFIG_COLLECT_PARTITION_STATS
- if (partition_timer_on) {
- aom_usec_timer_mark(&partition_timer);
- int64_t time = aom_usec_timer_elapsed(&partition_timer);
- partition_times[PARTITION_HORZ_4] += time;
- partition_timer_on = 0;
- }
-#endif
- restore_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);
+ const int inc_step[NUM_PART4_TYPES] = { mi_size_high[blk_params.bsize] / 4,
+ 0 };
+ // Evaluation of Horz4 partition type.
+ rd_pick_4partition(cpi, td, tile_data, tp, x, &x_ctx, pc_tree,
+ pc_tree->horizontal4, &part_search_state, &best_rdc,
+ inc_step, PARTITION_HORZ_4);
}
// PARTITION_VERT_4
@@ -4034,59 +4074,12 @@
partition_vert4_allowed && blk_params.has_cols &&
(part_search_state.do_rectangular_split ||
active_v_edge(cpi, mi_row, blk_params.mi_step))) {
- av1_init_rd_stats(&part_search_state.sum_rdc);
- const int quarter_step = mi_size_wide[bsize] / 4;
-
- blk_params.subsize = get_partition_subsize(bsize, PARTITION_VERT_4);
- part_search_state.sum_rdc.rate =
- part_search_state.partition_cost[PARTITION_VERT_4];
- part_search_state.sum_rdc.rdcost =
- RDCOST(x->rdmult, part_search_state.sum_rdc.rate, 0);
-
- for (int i = 0; i < 4; ++i)
- pc_tree->vertical4[i] =
- av1_alloc_pmc(cm, blk_params.subsize, &td->shared_coeff_buf);
-
-#if CONFIG_COLLECT_PARTITION_STATS
- if (best_rdc.rdcost - part_search_state.sum_rdc.rdcost >= 0) {
- partition_attempts[PARTITION_VERT_4] += 1;
- aom_usec_timer_start(&partition_timer);
- partition_timer_on = 1;
- }
-#endif
- for (int i = 0; i < 4; ++i) {
- const int this_mi_col = mi_col + i * quarter_step;
-
- if (i > 0 && this_mi_col >= mi_params->mi_cols) break;
-
- PICK_MODE_CONTEXT *ctx_this = pc_tree->vertical4[i];
-
- ctx_this->rd_mode_is_ready = 0;
- if (!rd_try_subblock(cpi, td, tile_data, tp, (i == 3), mi_row,
- this_mi_col, blk_params.subsize, best_rdc,
- &part_search_state.sum_rdc, PARTITION_VERT_4,
- ctx_this)) {
- av1_invalid_rd_stats(&part_search_state.sum_rdc);
- break;
- }
- }
-
- // Calculate the total cost and update the best partition
- av1_rd_cost_update(x->rdmult, &part_search_state.sum_rdc);
- if (part_search_state.sum_rdc.rdcost < best_rdc.rdcost) {
- best_rdc = part_search_state.sum_rdc;
- part_search_state.found_best_partition = true;
- pc_tree->partitioning = PARTITION_VERT_4;
- }
-#if CONFIG_COLLECT_PARTITION_STATS
- if (partition_timer_on) {
- aom_usec_timer_mark(&partition_timer);
- int64_t time = aom_usec_timer_elapsed(&partition_timer);
- partition_times[PARTITION_VERT_4] += time;
- partition_timer_on = 0;
- }
-#endif
- restore_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);
+ const int inc_step[NUM_PART4_TYPES] = { 0, mi_size_wide[blk_params.bsize] /
+ 4 };
+ // Evaluation of Vert4 partition type.
+ rd_pick_4partition(cpi, td, tile_data, tp, x, &x_ctx, pc_tree,
+ pc_tree->vertical4, &part_search_state, &best_rdc,
+ inc_step, PARTITION_VERT_4);
}
if (bsize == cm->seq_params.sb_size &&
diff --git a/av1/encoder/partition_strategy.h b/av1/encoder/partition_strategy.h
index b651673..9c30fb4 100644
--- a/av1/encoder/partition_strategy.h
+++ b/av1/encoder/partition_strategy.h
@@ -36,6 +36,12 @@
// Number of sub-partitions in AB partition types.
#define SUB_PARTITIONS_AB 3
+// Number of sub-partitions in 4-part partition types.
+#define PART4_SUB_PARTITIONS 4
+
+// 4part parition types.
+enum { HORZ4 = 0, VERT4, NUM_PART4_TYPES } UENUM1BYTE(PART4_TYPES);
+
// AB parition types.
enum {
HORZ_A = 0,
