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