blob: 84683f5fd4f8d434d23f077ff9d7892ea17d7f14 [file] [log] [blame]
chiyotsai19a58ee2019-03-18 18:01:05 -07001/*
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"
Cheng Chen6b5305f2021-05-04 12:28:00 -070016#include "av1/encoder/encodeframe_utils.h"
chiyotsai19a58ee2019-03-18 18:01:05 -070017#include "av1/encoder/encodemb.h"
18#include "av1/encoder/encoder.h"
19
chiyotsai36035d12019-06-19 17:39:48 -070020void av1_intra_mode_cnn_partition(const AV1_COMMON *const cm, MACROBLOCK *x,
chiyotsaic3613c82021-06-15 11:51:29 -070021 int label_idx,
venkat sanampudi593c54c2021-06-22 13:49:54 +053022 int intra_cnn_based_part_prune_level,
chiyotsaic3613c82021-06-15 11:51:29 -070023 PartitionSearchState *part_state);
chiyotsai36035d12019-06-19 17:39:48 -070024
chiyotsai19a58ee2019-03-18 18:01:05 -070025// 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
chiyotsaic3613c82021-06-15 11:51:29 -070028void av1_simple_motion_search_based_split(AV1_COMP *const cpi, MACROBLOCK *x,
29 SIMPLE_MOTION_DATA_TREE *sms_tree,
30 PartitionSearchState *part_state);
chiyotsai19a58ee2019-03-18 18:01:05 -070031
chiyotsai200ea982019-03-20 14:39:29 -070032// 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.
chiyotsaic3613c82021-06-15 11:51:29 -070035void av1_simple_motion_search_prune_rect(AV1_COMP *const cpi, MACROBLOCK *x,
36 SIMPLE_MOTION_DATA_TREE *sms_tree,
37 PartitionSearchState *part_state);
chiyotsai200ea982019-03-20 14:39:29 -070038
Jerome Jiangf6af3612019-06-04 12:27:05 -070039#if !CONFIG_REALTIME_ONLY
chiyotsai200ea982019-03-20 14:39:29 -070040// 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
chiyotsaic3613c82021-06-15 11:51:29 -070046void 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);
chiyotsai200ea982019-03-20 14:39:29 -070051
chiyotsai200ea982019-03-20 14:39:29 -070052// Get the features for selecting the max and min partition size. Currently this
chiyotsai23491612019-04-18 12:18:19 -070053// performs simple_motion_search on 16X16 subblocks of the current superblock,
chiyotsai200ea982019-03-20 14:39:29 -070054// and then extract the statistics of sse and motion vectors as features.
55void 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.
chiyotsai68eefbe2020-05-01 15:07:58 -070060BLOCK_SIZE av1_predict_max_partition(const AV1_COMP *const cpi,
61 const MACROBLOCK *const x,
Debargha Mukherjee1c141bc2019-03-21 14:20:56 -070062 const float *features);
chiyotsai200ea982019-03-20 14:39:29 -070063
chiyotsaie10416d2019-04-16 11:16:13 -070064// Attempts an early termination after PARTITION_SPLIT.
65void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
Yue Cheneb628982019-08-29 15:17:13 -070066 SIMPLE_MOTION_DATA_TREE *const sms_tree,
chiyotsaic3613c82021-06-15 11:51:29 -070067 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);
chiyotsaie10416d2019-04-16 11:16:13 -070071
chiyotsaie10416d2019-04-16 11:16:13 -070072// 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.
Cheng Chenc0480112021-05-18 17:54:15 -070078void av1_ml_prune_rect_partition(AV1_COMP *const cpi, const MACROBLOCK *const x,
chiyotsaic3613c82021-06-15 11:51:29 -070079 int64_t best_rd, int64_t none_rd,
80 const int64_t *split_rd,
81 PartitionSearchState *part_state);
chiyotsaie10416d2019-04-16 11:16:13 -070082
83// Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
84// considered.
chiyotsaic3613c82021-06-15 11:51:29 -070085void 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);
chiyotsaie10416d2019-04-16 11:16:13 -070089
90// Use a ML model to predict if horz4 and vert4 should be considered.
chiyotsaic3613c82021-06-15 11:51:29 -070091void 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);
chiyotsaie10416d2019-04-16 11:16:13 -070096
Yue Chen2db09fe2020-05-06 16:23:42 -070097// ML-based partition search breakout after PARTITION_NONE.
chiyotsaic3613c82021-06-15 11:51:29 -070098void av1_ml_predict_breakout(AV1_COMP *const cpi, const MACROBLOCK *const x,
Cheng Chenc0480112021-05-18 17:54:15 -070099 const RD_STATS *const rd_stats,
Cheng Chenc0480112021-05-18 17:54:15 -0700100 unsigned int pb_source_variance, int bit_depth,
chiyotsaic3613c82021-06-15 11:51:29 -0700101 PartitionSearchState *part_state);
Yue Chen2db09fe2020-05-06 16:23:42 -0700102
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.
chiyotsaic3613c82021-06-15 11:51:29 -0700106void 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);
Yue Chen130ad862020-05-06 17:49:04 -0700110
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.
chiyotsaic3613c82021-06-15 11:51:29 -0700117void av1_prune_partitions_by_max_min_bsize(SuperBlockEnc *sb_enc,
118 PartitionSearchState *part_state);
Yue Chen7b09fd02020-05-07 00:54:28 -0700119
120// Prune out AB partitions based on rd decisions made from testing the
121// basic partitions.
chiyotsaic3613c82021-06-15 11:51:29 -0700122void 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);
Cheng Chena92050a2021-06-09 22:34:37 -0700129
130void av1_collect_motion_search_features_sb(AV1_COMP *const cpi, ThreadData *td,
Cheng Chen4eb00772021-09-16 12:23:50 -0700131 TileDataEnc *tile_data,
Cheng Chena92050a2021-06-09 22:34:37 -0700132 const int mi_row, const int mi_col,
Cheng Chen053c49c2021-07-02 13:50:15 -0700133 const BLOCK_SIZE bsize,
134 aom_partition_features_t *features);
Cheng Chenf237a2c2021-10-20 12:41:48 -0700135void 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]);
Jerome Jiangf6af3612019-06-04 12:27:05 -0700143#endif // !CONFIG_REALTIME_ONLY
chiyotsaie10416d2019-04-16 11:16:13 -0700144
chiyotsai19a58ee2019-03-18 18:01:05 -0700145// A simplified version of set_offsets meant to be used for
146// simple_motion_search.
147static INLINE void set_offsets_for_motion_search(const AV1_COMP *const cpi,
148 MACROBLOCK *const x,
149 int mi_row, int mi_col,
150 BLOCK_SIZE bsize) {
151 const AV1_COMMON *const cm = &cpi->common;
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700152 const CommonModeInfoParams *const mi_params = &cm->mi_params;
chiyotsai19a58ee2019-03-18 18:01:05 -0700153 const int num_planes = av1_num_planes(cm);
154 MACROBLOCKD *const xd = &x->e_mbd;
155 const int mi_width = mi_size_wide[bsize];
156 const int mi_height = mi_size_high[bsize];
157
Visheshd1317912020-04-07 14:39:44 +0530158 set_mode_info_offsets(&cpi->common.mi_params, &cpi->mbmi_ext_info, x, xd,
159 mi_row, mi_col);
chiyotsai19a58ee2019-03-18 18:01:05 -0700160
161 // Set up destination pointers.
162 av1_setup_dst_planes(xd->plane, bsize, &cm->cur_frame->buf, mi_row, mi_col, 0,
163 num_planes);
164
165 // Set up limit values for MV components.
166 // Mv beyond the range do not produce new/different prediction block.
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700167 av1_set_mv_limits(mi_params, &x->mv_limits, mi_row, mi_col, mi_height,
168 mi_width, cpi->oxcf.border_in_pixels);
chiyotsai19a58ee2019-03-18 18:01:05 -0700169
170 set_plane_n4(xd, mi_width, mi_height, num_planes);
171
chiyotsai2aac3002020-02-13 17:02:01 -0800172 xd->mi_row = mi_row;
173 xd->mi_col = mi_col;
174
chiyotsai19a58ee2019-03-18 18:01:05 -0700175 // Set up distance of MB to edge of frame in 1/8th pel units.
176 assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));
chiyotsai87bb8052020-02-12 16:56:33 -0800177 xd->mb_to_top_edge = -GET_MV_SUBPEL(mi_row * MI_SIZE);
178 xd->mb_to_bottom_edge =
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700179 GET_MV_SUBPEL((mi_params->mi_rows - mi_height - mi_row) * MI_SIZE);
chiyotsai87bb8052020-02-12 16:56:33 -0800180 xd->mb_to_left_edge = -GET_MV_SUBPEL(mi_col * MI_SIZE);
181 xd->mb_to_right_edge =
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700182 GET_MV_SUBPEL((mi_params->mi_cols - mi_width - mi_col) * MI_SIZE);
chiyotsai19a58ee2019-03-18 18:01:05 -0700183
184 // Set up source buffers.
185 av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes, bsize);
chiyotsai19a58ee2019-03-18 18:01:05 -0700186}
chiyotsai200ea982019-03-20 14:39:29 -0700187
chiyotsaib5faf832021-06-22 14:08:26 -0700188void 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);
chiyotsai200ea982019-03-20 14:39:29 -0700193
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700194static INLINE int is_full_sb(const CommonModeInfoParams *const mi_params,
195 int mi_row, int mi_col, BLOCK_SIZE sb_size) {
chiyotsai200ea982019-03-20 14:39:29 -0700196 const int sb_mi_wide = mi_size_wide[sb_size];
197 const int sb_mi_high = mi_size_high[sb_size];
198
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700199 return (mi_row + sb_mi_high) <= mi_params->mi_rows &&
200 (mi_col + sb_mi_wide) <= mi_params->mi_cols;
chiyotsai200ea982019-03-20 14:39:29 -0700201}
202
Cheng Chen6b5305f2021-05-04 12:28:00 -0700203#if !CONFIG_REALTIME_ONLY
Cheng Chenfa815c62020-03-12 17:22:37 -0700204// 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.
chiyotsai68eefbe2020-05-01 15:07:58 -0700207static INLINE int use_auto_max_partition(const AV1_COMP *const cpi,
chiyotsai200ea982019-03-20 14:39:29 -0700208 BLOCK_SIZE sb_size, int mi_row,
209 int mi_col) {
Mufaddal Chakera8ee04fa2021-03-17 13:33:18 +0530210 assert(IMPLIES(cpi->ppi->gf_group.size > 0,
211 cpi->gf_frame_index < cpi->ppi->gf_group.size));
chiyotsai68eefbe2020-05-01 15:07:58 -0700212 const AV1_COMMON *const cm = &cpi->common;
Cheng Chenfd883dc2020-08-14 15:38:16 -0700213 return !frame_is_intra_only(cm) && !cpi->use_screen_content_tools &&
chiyotsai5c2e42f2019-12-12 12:58:40 -0800214 cpi->sf.part_sf.auto_max_partition_based_on_simple_motion !=
215 NOT_IN_USE &&
Urvang Joshi9dc909d2020-03-23 16:07:02 -0700216 sb_size == BLOCK_128X128 &&
217 is_full_sb(&cm->mi_params, mi_row, mi_col, sb_size) &&
Mufaddal Chakera8ee04fa2021-03-17 13:33:18 +0530218 cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
219 OVERLAY_UPDATE &&
220 cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
221 INTNL_OVERLAY_UPDATE;
chiyotsai200ea982019-03-20 14:39:29 -0700222}
223
Cheng Chen6b5305f2021-05-04 12:28:00 -0700224static BLOCK_SIZE dim_to_size(int dim) {
225 switch (dim) {
226 case 4: return BLOCK_4X4;
227 case 8: return BLOCK_8X8;
228 case 16: return BLOCK_16X16;
229 case 32: return BLOCK_32X32;
230 case 64: return BLOCK_64X64;
231 case 128: return BLOCK_128X128;
232 default: assert(0); return 0;
233 }
234}
235
236static 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 int mi_row, int mi_col) {
241 const AV1_COMMON *cm = &cpi->common;
242
243 sb_enc->max_partition_size =
244 AOMMIN(sf->part_sf.default_max_partition_size,
245 dim_to_size(cpi->oxcf.part_cfg.max_partition_size));
246 sb_enc->min_partition_size =
247 AOMMAX(sf->part_sf.default_min_partition_size,
248 dim_to_size(cpi->oxcf.part_cfg.min_partition_size));
249 sb_enc->max_partition_size =
250 AOMMIN(sb_enc->max_partition_size, cm->seq_params->sb_size);
251 sb_enc->min_partition_size =
252 AOMMIN(sb_enc->min_partition_size, cm->seq_params->sb_size);
253
254 if (use_auto_max_partition(cpi, sb_size, mi_row, mi_col)) {
255 float features[FEATURE_SIZE_MAX_MIN_PART_PRED] = { 0.0f };
256
257 av1_get_max_min_partition_features(cpi, x, mi_row, mi_col, features);
258 sb_enc->max_partition_size =
259 AOMMAX(AOMMIN(av1_predict_max_partition(cpi, x, features),
260 sb_enc->max_partition_size),
261 sb_enc->min_partition_size);
262 }
263}
264#endif // !CONFIG_REALTIME_ONLY
chiyotsai19a58ee2019-03-18 18:01:05 -0700265#endif // AOM_AV1_ENCODER_PARTITION_STRATEGY_H_