1 /*
2 * Copyright (c) 2019, Alliance for Open Media. All rights reserved.
3 *
4 * This source code is subject to the terms of the BSD 2 Clause License and
5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6 * was not distributed with this source code in the LICENSE file, you can
7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8 * Media Patent License 1.0 was not distributed with this source code in the
9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10 */
11
12 #ifndef AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
13 #define AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
14
15 #include "av1/encoder/encodeframe.h"
16 #include "av1/encoder/encodeframe_utils.h"
17 #include "av1/encoder/encodemb.h"
18 #include "av1/encoder/encoder.h"
19
20 #if !CONFIG_REALTIME_ONLY
21 // Early terminates PARTITION_NONE using simple_motion_search features and the
22 // rate, distortion, and rdcost of PARTITION_NONE. This is only called when:
23 // - The frame is a show frame
24 // - The frame is not intra only
25 // - The current bsize is > BLOCK_8X8
26 // - blk_row + blk_height/2 < total_rows and blk_col + blk_width/2 < total_cols
27 void av1_simple_motion_search_early_term_none(AV1_COMP *const cpi,
28 MACROBLOCK *x,
29 SIMPLE_MOTION_DATA_TREE *sms_tree,
30 const RD_STATS *none_rdc,
31 PartitionSearchState *part_state);
32
33 // Get the features for selecting the max and min partition size. Currently this
34 // performs simple_motion_search on 16X16 subblocks of the current superblock,
35 // and then extract the statistics of sse and motion vectors as features.
36 void av1_get_max_min_partition_features(AV1_COMP *const cpi, MACROBLOCK *x,
37 int mi_row, int mi_col,
38 float *features);
39
40 // Predict the maximum BLOCK_SIZE to be used to encoder the current superblock.
41 BLOCK_SIZE av1_predict_max_partition(const AV1_COMP *const cpi,
42 const MACROBLOCK *const x,
43 const float *features);
44
45 // Attempts an early termination after PARTITION_SPLIT.
46 void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
47 SIMPLE_MOTION_DATA_TREE *const sms_tree,
48 int64_t best_rd, int64_t part_none_rd,
49 int64_t part_split_rd,
50 int64_t *split_block_rd,
51 PartitionSearchState *part_state);
52
53 // Use the rdcost ratio and source var ratio to prune PARTITION_HORZ and
54 // PARTITION_VERT.
55 // TODO([email protected]): Currently this model does not use q value and has
56 // no information about rectangular partitions. Preliminary experiments suggest
57 // that we can get better performance by adding in q_index and rectangular
58 // sse/var from SMS. We should retrain and tune this model later.
59 void av1_ml_prune_rect_partition(AV1_COMP *const cpi, const MACROBLOCK *const x,
60 int64_t best_rd, int64_t none_rd,
61 const int64_t *split_rd,
62 PartitionSearchState *part_state);
63
64 // Use a ML model to predict if horz4 and vert4 should be considered.
65 void av1_ml_prune_4_partition(AV1_COMP *const cpi, MACROBLOCK *const x,
66 int part_ctx, int64_t best_rd,
67 PartitionSearchState *part_state,
68 int *part4_allowed,
69 unsigned int pb_source_variance);
70
71 // ML-based partition search breakout after PARTITION_NONE.
72 void av1_ml_predict_breakout(AV1_COMP *const cpi, const MACROBLOCK *const x,
73 const RD_STATS *const rd_stats,
74 unsigned int pb_source_variance, int bit_depth,
75 PartitionSearchState *part_state);
76
77 // The first round of partition pruning determined before any partition
78 // has been tested. The decisions will be updated and passed back
79 // to the partition search function.
80 void av1_prune_partitions_before_search(AV1_COMP *const cpi,
81 MACROBLOCK *const x,
82 SIMPLE_MOTION_DATA_TREE *const sms_tree,
83 PartitionSearchState *part_state);
84
85 // Prune out partitions that lead to coding block sizes outside the min and max
86 // bsizes set by the encoder. Max and min square partition levels are defined as
87 // the partition nodes that the recursive function rd_pick_partition() can
88 // reach. To implement this: only PARTITION_NONE is allowed if the current node
89 // equals max_partition_size, only PARTITION_SPLIT is allowed if the current
90 // node exceeds max_partition_size.
91 void av1_prune_partitions_by_max_min_bsize(SuperBlockEnc *sb_enc,
92 PartitionSearchState *part_state);
93
94 // Prune out AB partitions based on rd decisions made from testing the
95 // basic partitions.
96 void av1_prune_ab_partitions(AV1_COMP *cpi, const MACROBLOCK *x,
97 const PC_TREE *pc_tree, int pb_source_variance,
98 int64_t best_rdcost,
99 const RD_RECT_PART_WIN_INFO *rect_part_win_info,
100 bool ext_partition_allowed,
101 PartitionSearchState *part_state,
102 int *ab_partitions_allowed);
103
104 void av1_collect_motion_search_features_sb(AV1_COMP *const cpi, ThreadData *td,
105 TileDataEnc *tile_data,
106 const int mi_row, const int mi_col,
107 const BLOCK_SIZE bsize,
108 aom_partition_features_t *features);
109 void av1_prepare_motion_search_features_block(
110 AV1_COMP *const cpi, ThreadData *td, TileDataEnc *tile_data,
111 const int mi_row, const int mi_col, const BLOCK_SIZE bsize,
112 const int valid_partition_types, unsigned int *block_sse,
113 unsigned int *block_var, unsigned int sub_block_sse[4],
114 unsigned int sub_block_var[4], unsigned int horz_block_sse[2],
115 unsigned int horz_block_var[2], unsigned int vert_block_sse[2],
116 unsigned int vert_block_var[2]);
117 #endif // !CONFIG_REALTIME_ONLY
118
119 // A simplified version of set_offsets meant to be used for
120 // simple_motion_search.
set_offsets_for_motion_search(const AV1_COMP * const cpi,MACROBLOCK * const x,int mi_row,int mi_col,BLOCK_SIZE bsize)121 static inline void set_offsets_for_motion_search(const AV1_COMP *const cpi,
122 MACROBLOCK *const x,
123 int mi_row, int mi_col,
124 BLOCK_SIZE bsize) {
125 const AV1_COMMON *const cm = &cpi->common;
126 const CommonModeInfoParams *const mi_params = &cm->mi_params;
127 const int num_planes = av1_num_planes(cm);
128 MACROBLOCKD *const xd = &x->e_mbd;
129 const int mi_width = mi_size_wide[bsize];
130 const int mi_height = mi_size_high[bsize];
131
132 set_mode_info_offsets(&cpi->common.mi_params, &cpi->mbmi_ext_info, x, xd,
133 mi_row, mi_col);
134
135 // Set up destination pointers.
136 av1_setup_dst_planes(xd->plane, bsize, &cm->cur_frame->buf, mi_row, mi_col, 0,
137 num_planes);
138
139 // Set up limit values for MV components.
140 // Mv beyond the range do not produce new/different prediction block.
141 av1_set_mv_limits(mi_params, &x->mv_limits, mi_row, mi_col, mi_height,
142 mi_width, cpi->oxcf.border_in_pixels);
143
144 set_plane_n4(xd, mi_width, mi_height, num_planes);
145
146 xd->mi_row = mi_row;
147 xd->mi_col = mi_col;
148
149 // Set up distance of MB to edge of frame in 1/8th pel units.
150 assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));
151 xd->mb_to_top_edge = -GET_MV_SUBPEL(mi_row * MI_SIZE);
152 xd->mb_to_bottom_edge =
153 GET_MV_SUBPEL((mi_params->mi_rows - mi_height - mi_row) * MI_SIZE);
154 xd->mb_to_left_edge = -GET_MV_SUBPEL(mi_col * MI_SIZE);
155 xd->mb_to_right_edge =
156 GET_MV_SUBPEL((mi_params->mi_cols - mi_width - mi_col) * MI_SIZE);
157
158 // Set up source buffers.
159 av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes, bsize);
160 }
161
162 void av1_init_simple_motion_search_mvs_for_sb(const AV1_COMP *cpi,
163 const TileInfo *tile_info,
164 MACROBLOCK *x,
165 SIMPLE_MOTION_DATA_TREE *sms_root,
166 int mi_row, int mi_col);
167
is_full_sb(const CommonModeInfoParams * const mi_params,int mi_row,int mi_col,BLOCK_SIZE sb_size)168 static inline int is_full_sb(const CommonModeInfoParams *const mi_params,
169 int mi_row, int mi_col, BLOCK_SIZE sb_size) {
170 const int sb_mi_wide = mi_size_wide[sb_size];
171 const int sb_mi_high = mi_size_high[sb_size];
172
173 return (mi_row + sb_mi_high) <= mi_params->mi_rows &&
174 (mi_col + sb_mi_wide) <= mi_params->mi_cols;
175 }
176
177 #if !CONFIG_REALTIME_ONLY
178 // Do not use this criteria for screen content videos.
179 // Since screen content videos could often find good predictors and the largest
180 // block size is likely to be used.
use_auto_max_partition(const AV1_COMP * const cpi,BLOCK_SIZE sb_size,int mi_row,int mi_col)181 static inline int use_auto_max_partition(const AV1_COMP *const cpi,
182 BLOCK_SIZE sb_size, int mi_row,
183 int mi_col) {
184 assert(IMPLIES(cpi->ppi->gf_group.size > 0,
185 cpi->gf_frame_index < cpi->ppi->gf_group.size));
186 const AV1_COMMON *const cm = &cpi->common;
187 return !frame_is_intra_only(cm) && !cpi->use_screen_content_tools &&
188 cpi->sf.part_sf.auto_max_partition_based_on_simple_motion !=
189 NOT_IN_USE &&
190 sb_size == BLOCK_128X128 &&
191 is_full_sb(&cm->mi_params, mi_row, mi_col, sb_size) &&
192 cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
193 OVERLAY_UPDATE &&
194 cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
195 INTNL_OVERLAY_UPDATE;
196 }
197
dim_to_size(int dim)198 static BLOCK_SIZE dim_to_size(int dim) {
199 switch (dim) {
200 case 4: return BLOCK_4X4;
201 case 8: return BLOCK_8X8;
202 case 16: return BLOCK_16X16;
203 case 32: return BLOCK_32X32;
204 case 64: return BLOCK_64X64;
205 case 128: return BLOCK_128X128;
206 default: assert(0); return 0;
207 }
208 }
209
set_max_min_partition_size(SuperBlockEnc * sb_enc,AV1_COMP * cpi,MACROBLOCK * x,const SPEED_FEATURES * sf,BLOCK_SIZE sb_size,int mi_row,int mi_col)210 static inline void set_max_min_partition_size(SuperBlockEnc *sb_enc,
211 AV1_COMP *cpi, MACROBLOCK *x,
212 const SPEED_FEATURES *sf,
213 BLOCK_SIZE sb_size, int mi_row,
214 int mi_col) {
215 const AV1_COMMON *cm = &cpi->common;
216
217 sb_enc->max_partition_size =
218 AOMMIN(sf->part_sf.default_max_partition_size,
219 dim_to_size(cpi->oxcf.part_cfg.max_partition_size));
220 sb_enc->min_partition_size =
221 AOMMAX(sf->part_sf.default_min_partition_size,
222 dim_to_size(cpi->oxcf.part_cfg.min_partition_size));
223 sb_enc->max_partition_size =
224 AOMMIN(sb_enc->max_partition_size, cm->seq_params->sb_size);
225 sb_enc->min_partition_size =
226 AOMMIN(sb_enc->min_partition_size, cm->seq_params->sb_size);
227
228 if (use_auto_max_partition(cpi, sb_size, mi_row, mi_col)) {
229 float features[FEATURE_SIZE_MAX_MIN_PART_PRED] = { 0.0f };
230
231 av1_get_max_min_partition_features(cpi, x, mi_row, mi_col, features);
232 sb_enc->max_partition_size =
233 AOMMAX(AOMMIN(av1_predict_max_partition(cpi, x, features),
234 sb_enc->max_partition_size),
235 sb_enc->min_partition_size);
236 }
237 }
238 #endif // !CONFIG_REALTIME_ONLY
239 #endif // AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
240