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