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/encodemb.h"
17 #include "av1/encoder/encoder.h"
18
19 #define FEATURE_SIZE_SMS_SPLIT_FAST 6
20 #define FEATURE_SIZE_SMS_SPLIT 17
21 #define FEATURE_SIZE_SMS_PRUNE_PART 25
22 #define FEATURE_SIZE_SMS_TERM_NONE 28
23 #define FEATURE_SIZE_FP_SMS_TERM_NONE 20
24 #define FEATURE_SIZE_MAX_MIN_PART_PRED 13
25 #define MAX_NUM_CLASSES_MAX_MIN_PART_PRED 4
26
27 #define FEATURE_SMS_NONE_FLAG 1
28 #define FEATURE_SMS_SPLIT_FLAG (1 << 1)
29 #define FEATURE_SMS_RECT_FLAG (1 << 2)
30
31 #define FEATURE_SMS_PRUNE_PART_FLAG \
32 (FEATURE_SMS_NONE_FLAG | FEATURE_SMS_SPLIT_FLAG | FEATURE_SMS_RECT_FLAG)
33 #define FEATURE_SMS_SPLIT_MODEL_FLAG \
34 (FEATURE_SMS_NONE_FLAG | FEATURE_SMS_SPLIT_FLAG)
35
36 void av1_intra_mode_cnn_partition(const AV1_COMMON *const cm, MACROBLOCK *x,
37 int bsize, int label_idx,
38 int *partition_none_allowed,
39 int *partition_horz_allowed,
40 int *partition_vert_allowed,
41 int *do_rectangular_split,
42 int *do_square_split);
43
44 // Performs a simple_motion_search with a single reference frame and extract
45 // the variance of residues. Then use the features to determine whether we want
46 // to go straight to splitting without trying PARTITION_NONE
47 void av1_simple_motion_search_based_split(
48 AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
49 int mi_col, BLOCK_SIZE bsize, int *partition_none_allowed,
50 int *partition_horz_allowed, int *partition_vert_allowed,
51 int *do_rectangular_split, int *do_square_split);
52
53 // Performs a simple_motion_search with two reference frames and extract
54 // the variance of residues. Then use the features to determine whether we want
55 // to prune some partitions.
56 void av1_simple_motion_search_prune_rect(AV1_COMP *const cpi, MACROBLOCK *x,
57 PC_TREE *pc_tree, int mi_row,
58 int mi_col, BLOCK_SIZE bsize,
59 int *partition_horz_allowed,
60 int *partition_vert_allowed,
61 int *prune_horz, int *prune_vert);
62
63 #if !CONFIG_REALTIME_ONLY
64 // Early terminates PARTITION_NONE using simple_motion_search features and the
65 // rate, distortion, and rdcost of PARTITION_NONE. This is only called when:
66 // - The frame is a show frame
67 // - The frame is not intra only
68 // - The current bsize is > BLOCK_8X8
69 // - blk_row + blk_height/2 < total_rows and blk_col + blk_width/2 < total_cols
70 void av1_simple_motion_search_early_term_none(AV1_COMP *const cpi,
71 MACROBLOCK *x, PC_TREE *pc_tree,
72 int mi_row, int mi_col,
73 BLOCK_SIZE bsize,
74 const RD_STATS *none_rdc,
75 int *early_terminate);
76
77 // Get the features for selecting the max and min partition size. Currently this
78 // performs simple_motion_search on 16X16 subblocks of the current superblock,
79 // and then extract the statistics of sse and motion vectors as features.
80 void av1_get_max_min_partition_features(AV1_COMP *const cpi, MACROBLOCK *x,
81 int mi_row, int mi_col,
82 float *features);
83
84 // Predict the maximum BLOCK_SIZE to be used to encoder the current superblock.
85 BLOCK_SIZE av1_predict_max_partition(AV1_COMP *const cpi, MACROBLOCK *const x,
86 const float *features);
87
88 // Attempts an early termination after PARTITION_SPLIT.
89 void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
90 PC_TREE *const pc_tree, BLOCK_SIZE bsize,
91 int64_t best_rd, int64_t part_none_rd,
92 int64_t part_split_rd,
93 int64_t *split_block_rd, int mi_row,
94 int mi_col,
95 int *const terminate_partition_search);
96
97 // Use the rdcost ratio and source var ratio to prune PARTITION_HORZ and
98 // PARTITION_VERT.
99 // TODO(chiyotsai@google.com): Currently this model does not use q value and has
100 // no information about rectangular partitions. Preliminary experiments suggest
101 // that we can get better performance by adding in q_index and rectangular
102 // sse/var from SMS. We should retrain and tune this model later.
103 void av1_ml_prune_rect_partition(const AV1_COMP *const cpi,
104 const MACROBLOCK *const x, BLOCK_SIZE bsize,
105 int64_t best_rd, int64_t none_rd,
106 int64_t *split_rd, int *const dst_prune_horz,
107 int *const dst_prune_vert);
108
109 // Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
110 // considered.
111 void av1_ml_prune_ab_partition(BLOCK_SIZE bsize, int part_ctx, int var_ctx,
112 int64_t best_rd, int64_t horz_rd[2],
113 int64_t vert_rd[2], int64_t split_rd[4],
114 int *const horza_partition_allowed,
115 int *const horzb_partition_allowed,
116 int *const verta_partition_allowed,
117 int *const vertb_partition_allowed);
118
119 // Use a ML model to predict if horz4 and vert4 should be considered.
120 void av1_ml_prune_4_partition(const AV1_COMP *const cpi, MACROBLOCK *const x,
121 BLOCK_SIZE bsize, int part_ctx, int64_t best_rd,
122 int64_t horz_rd[2], int64_t vert_rd[2],
123 int64_t split_rd[4],
124 int *const partition_horz4_allowed,
125 int *const partition_vert4_allowed,
126 unsigned int pb_source_variance, int mi_row,
127 int mi_col);
128
129 // ML-based partition search breakout after PARTITION_NONE
130 int av1_ml_predict_breakout(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
131 const MACROBLOCK *const x,
132 const RD_STATS *const rd_stats,
133 unsigned int pb_source_variance);
134 #endif // !CONFIG_REALTIME_ONLY
135
136 // A simplified version of set_offsets meant to be used for
137 // 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)138 static INLINE void set_offsets_for_motion_search(const AV1_COMP *const cpi,
139 MACROBLOCK *const x,
140 int mi_row, int mi_col,
141 BLOCK_SIZE bsize) {
142 const AV1_COMMON *const cm = &cpi->common;
143 const CommonModeInfoParams *const mi_params = &cm->mi_params;
144 const int num_planes = av1_num_planes(cm);
145 MACROBLOCKD *const xd = &x->e_mbd;
146 const int mi_width = mi_size_wide[bsize];
147 const int mi_height = mi_size_high[bsize];
148
149 set_mode_info_offsets(&cpi->common.mi_params, &cpi->mbmi_ext_info, x, xd,
150 mi_row, mi_col);
151
152 // Set up destination pointers.
153 av1_setup_dst_planes(xd->plane, bsize, &cm->cur_frame->buf, mi_row, mi_col, 0,
154 num_planes);
155
156 // Set up limit values for MV components.
157 // Mv beyond the range do not produce new/different prediction block.
158 av1_set_mv_limits(mi_params, &x->mv_limits, mi_row, mi_col, mi_height,
159 mi_width, cpi->oxcf.border_in_pixels);
160
161 set_plane_n4(xd, mi_width, mi_height, num_planes);
162
163 xd->mi_row = mi_row;
164 xd->mi_col = mi_col;
165
166 // Set up distance of MB to edge of frame in 1/8th pel units.
167 assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));
168 xd->mb_to_top_edge = -GET_MV_SUBPEL(mi_row * MI_SIZE);
169 xd->mb_to_bottom_edge =
170 GET_MV_SUBPEL((mi_params->mi_rows - mi_height - mi_row) * MI_SIZE);
171 xd->mb_to_left_edge = -GET_MV_SUBPEL(mi_col * MI_SIZE);
172 xd->mb_to_right_edge =
173 GET_MV_SUBPEL((mi_params->mi_cols - mi_width - mi_col) * MI_SIZE);
174
175 // Set up source buffers.
176 av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes, bsize);
177 }
178
init_simple_motion_search_mvs(PC_TREE * pc_tree)179 static INLINE void init_simple_motion_search_mvs(PC_TREE *pc_tree) {
180 av1_zero(pc_tree->start_mvs);
181
182 av1_zero(pc_tree->sms_none_feat);
183 av1_zero(pc_tree->sms_rect_feat);
184 av1_zero(pc_tree->sms_none_valid);
185 av1_zero(pc_tree->sms_rect_valid);
186
187 if (pc_tree->block_size >= BLOCK_8X8) {
188 init_simple_motion_search_mvs(pc_tree->split[0]);
189 init_simple_motion_search_mvs(pc_tree->split[1]);
190 init_simple_motion_search_mvs(pc_tree->split[2]);
191 init_simple_motion_search_mvs(pc_tree->split[3]);
192 }
193 }
194
is_full_sb(const CommonModeInfoParams * const mi_params,int mi_row,int mi_col,BLOCK_SIZE sb_size)195 static INLINE int is_full_sb(const CommonModeInfoParams *const mi_params,
196 int mi_row, int mi_col, BLOCK_SIZE sb_size) {
197 const int sb_mi_wide = mi_size_wide[sb_size];
198 const int sb_mi_high = mi_size_high[sb_size];
199
200 return (mi_row + sb_mi_high) <= mi_params->mi_rows &&
201 (mi_col + sb_mi_wide) <= mi_params->mi_cols;
202 }
203
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.
use_auto_max_partition(AV1_COMP * const cpi,BLOCK_SIZE sb_size,int mi_row,int mi_col)207 static INLINE int use_auto_max_partition(AV1_COMP *const cpi,
208 BLOCK_SIZE sb_size, int mi_row,
209 int mi_col) {
210 assert(IMPLIES(cpi->gf_group.size > 0,
211 cpi->gf_group.index < cpi->gf_group.size));
212 AV1_COMMON *const cm = &cpi->common;
213 return !frame_is_intra_only(cm) && !cpi->is_screen_content_type &&
214 cpi->sf.part_sf.auto_max_partition_based_on_simple_motion !=
215 NOT_IN_USE &&
216 sb_size == BLOCK_128X128 &&
217 is_full_sb(&cm->mi_params, mi_row, mi_col, sb_size) &&
218 cpi->gf_group.update_type[cpi->gf_group.index] != OVERLAY_UPDATE &&
219 cpi->gf_group.update_type[cpi->gf_group.index] != INTNL_OVERLAY_UPDATE;
220 }
221
222 #endif // AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
223