Coverage Report

Created: 2022-08-24 06:17

/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
0
                                         int mi_col) {
210
0
  assert(IMPLIES(cpi->ppi->gf_group.size > 0,
211
0
                 cpi->gf_frame_index < cpi->ppi->gf_group.size));
212
0
  const AV1_COMMON *const cm = &cpi->common;
213
0
  return !frame_is_intra_only(cm) && !cpi->use_screen_content_tools &&
214
0
         cpi->sf.part_sf.auto_max_partition_based_on_simple_motion !=
215
0
             NOT_IN_USE &&
216
0
         sb_size == BLOCK_128X128 &&
217
0
         is_full_sb(&cm->mi_params, mi_row, mi_col, sb_size) &&
218
0
         cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
219
0
             OVERLAY_UPDATE &&
220
0
         cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
221
0
             INTNL_OVERLAY_UPDATE;
222
0
}
Unexecuted instantiation: encodeframe.c:use_auto_max_partition
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
0
static BLOCK_SIZE dim_to_size(int dim) {
225
0
  switch (dim) {
226
0
    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
0
    case 128: return BLOCK_128X128;
232
0
    default: assert(0); return 0;
233
0
  }
234
0
}
Unexecuted instantiation: encodeframe.c:dim_to_size
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
0
                                                  int mi_row, int mi_col) {
241
0
  const AV1_COMMON *cm = &cpi->common;
242
243
0
  sb_enc->max_partition_size =
244
0
      AOMMIN(sf->part_sf.default_max_partition_size,
245
0
             dim_to_size(cpi->oxcf.part_cfg.max_partition_size));
246
0
  sb_enc->min_partition_size =
247
0
      AOMMAX(sf->part_sf.default_min_partition_size,
248
0
             dim_to_size(cpi->oxcf.part_cfg.min_partition_size));
249
0
  sb_enc->max_partition_size =
250
0
      AOMMIN(sb_enc->max_partition_size, cm->seq_params->sb_size);
251
0
  sb_enc->min_partition_size =
252
0
      AOMMIN(sb_enc->min_partition_size, cm->seq_params->sb_size);
253
254
0
  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
0
}
Unexecuted instantiation: encodeframe.c:set_max_min_partition_size
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_