summaryrefslogtreecommitdiffstats
path: root/third_party/aom/av1/encoder/partition_strategy.h
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--third_party/aom/av1/encoder/partition_strategy.h265
1 files changed, 265 insertions, 0 deletions
diff --git a/third_party/aom/av1/encoder/partition_strategy.h b/third_party/aom/av1/encoder/partition_strategy.h
new file mode 100644
index 0000000000..84683f5fd4
--- /dev/null
+++ b/third_party/aom/av1/encoder/partition_strategy.h
@@ -0,0 +1,265 @@
+/*
+ * Copyright (c) 2019, Alliance for Open Media. All rights reserved
+ *
+ * This source code is subject to the terms of the BSD 2 Clause License and
+ * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
+ * was not distributed with this source code in the LICENSE file, you can
+ * obtain it at www.aomedia.org/license/software. If the Alliance for Open
+ * Media Patent License 1.0 was not distributed with this source code in the
+ * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
+ */
+
+#ifndef AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
+#define AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
+
+#include "av1/encoder/encodeframe.h"
+#include "av1/encoder/encodeframe_utils.h"
+#include "av1/encoder/encodemb.h"
+#include "av1/encoder/encoder.h"
+
+void av1_intra_mode_cnn_partition(const AV1_COMMON *const cm, MACROBLOCK *x,
+ int label_idx,
+ int intra_cnn_based_part_prune_level,
+ PartitionSearchState *part_state);
+
+// Performs a simple_motion_search with a single reference frame and extract
+// the variance of residues. Then use the features to determine whether we want
+// to go straight to splitting without trying PARTITION_NONE
+void av1_simple_motion_search_based_split(AV1_COMP *const cpi, MACROBLOCK *x,
+ SIMPLE_MOTION_DATA_TREE *sms_tree,
+ PartitionSearchState *part_state);
+
+// Performs a simple_motion_search with two reference frames and extract
+// the variance of residues. Then use the features to determine whether we want
+// to prune some partitions.
+void av1_simple_motion_search_prune_rect(AV1_COMP *const cpi, MACROBLOCK *x,
+ SIMPLE_MOTION_DATA_TREE *sms_tree,
+ PartitionSearchState *part_state);
+
+#if !CONFIG_REALTIME_ONLY
+// Early terminates PARTITION_NONE using simple_motion_search features and the
+// rate, distortion, and rdcost of PARTITION_NONE. This is only called when:
+// - The frame is a show frame
+// - The frame is not intra only
+// - The current bsize is > BLOCK_8X8
+// - blk_row + blk_height/2 < total_rows and blk_col + blk_width/2 < total_cols
+void av1_simple_motion_search_early_term_none(AV1_COMP *const cpi,
+ MACROBLOCK *x,
+ SIMPLE_MOTION_DATA_TREE *sms_tree,
+ const RD_STATS *none_rdc,
+ PartitionSearchState *part_state);
+
+// Get the features for selecting the max and min partition size. Currently this
+// performs simple_motion_search on 16X16 subblocks of the current superblock,
+// and then extract the statistics of sse and motion vectors as features.
+void av1_get_max_min_partition_features(AV1_COMP *const cpi, MACROBLOCK *x,
+ int mi_row, int mi_col,
+ float *features);
+
+// Predict the maximum BLOCK_SIZE to be used to encoder the current superblock.
+BLOCK_SIZE av1_predict_max_partition(const AV1_COMP *const cpi,
+ const MACROBLOCK *const x,
+ const float *features);
+
+// Attempts an early termination after PARTITION_SPLIT.
+void av1_ml_early_term_after_split(AV1_COMP *const cpi, MACROBLOCK *const x,
+ SIMPLE_MOTION_DATA_TREE *const sms_tree,
+ int64_t best_rd, int64_t part_none_rd,
+ int64_t part_split_rd,
+ int64_t *split_block_rd,
+ PartitionSearchState *part_state);
+
+// Use the rdcost ratio and source var ratio to prune PARTITION_HORZ and
+// PARTITION_VERT.
+// TODO(chiyotsai@google.com): Currently this model does not use q value and has
+// no information about rectangular partitions. Preliminary experiments suggest
+// that we can get better performance by adding in q_index and rectangular
+// sse/var from SMS. We should retrain and tune this model later.
+void av1_ml_prune_rect_partition(AV1_COMP *const cpi, const MACROBLOCK *const x,
+ int64_t best_rd, int64_t none_rd,
+ const int64_t *split_rd,
+ PartitionSearchState *part_state);
+
+// Use a ML model to predict if horz_a, horz_b, vert_a, and vert_b should be
+// considered.
+void av1_ml_prune_ab_partition(AV1_COMP *const cpi, int part_ctx, int var_ctx,
+ int64_t best_rd,
+ PartitionSearchState *part_state,
+ int *ab_partitions_allowed);
+
+// Use a ML model to predict if horz4 and vert4 should be considered.
+void av1_ml_prune_4_partition(AV1_COMP *const cpi, MACROBLOCK *const x,
+ int part_ctx, int64_t best_rd,
+ PartitionSearchState *part_state,
+ int *part4_allowed,
+ unsigned int pb_source_variance);
+
+// ML-based partition search breakout after PARTITION_NONE.
+void av1_ml_predict_breakout(AV1_COMP *const cpi, const MACROBLOCK *const x,
+ const RD_STATS *const rd_stats,
+ unsigned int pb_source_variance, int bit_depth,
+ PartitionSearchState *part_state);
+
+// The first round of partition pruning determined before any partition
+// has been tested. The decisions will be updated and passed back
+// to the partition search function.
+void av1_prune_partitions_before_search(AV1_COMP *const cpi,
+ MACROBLOCK *const x,
+ SIMPLE_MOTION_DATA_TREE *const sms_tree,
+ PartitionSearchState *part_state);
+
+// Prune out partitions that lead to coding block sizes outside the min and max
+// bsizes set by the encoder. Max and min square partition levels are defined as
+// the partition nodes that the recursive function rd_pick_partition() can
+// reach. To implement this: only PARTITION_NONE is allowed if the current node
+// equals max_partition_size, only PARTITION_SPLIT is allowed if the current
+// node exceeds max_partition_size.
+void av1_prune_partitions_by_max_min_bsize(SuperBlockEnc *sb_enc,
+ PartitionSearchState *part_state);
+
+// Prune out AB partitions based on rd decisions made from testing the
+// basic partitions.
+void av1_prune_ab_partitions(AV1_COMP *cpi, const MACROBLOCK *x,
+ const PC_TREE *pc_tree, int pb_source_variance,
+ int64_t best_rdcost,
+ const RD_RECT_PART_WIN_INFO *rect_part_win_info,
+ bool ext_partition_allowed,
+ PartitionSearchState *part_state,
+ int *ab_partitions_allowed);
+
+void av1_collect_motion_search_features_sb(AV1_COMP *const cpi, ThreadData *td,
+ TileDataEnc *tile_data,
+ const int mi_row, const int mi_col,
+ const BLOCK_SIZE bsize,
+ aom_partition_features_t *features);
+void av1_prepare_motion_search_features_block(
+ AV1_COMP *const cpi, ThreadData *td, TileDataEnc *tile_data,
+ const int mi_row, const int mi_col, const BLOCK_SIZE bsize,
+ const int valid_partition_types, unsigned int *block_sse,
+ unsigned int *block_var, unsigned int sub_block_sse[4],
+ unsigned int sub_block_var[4], unsigned int horz_block_sse[2],
+ unsigned int horz_block_var[2], unsigned int vert_block_sse[2],
+ unsigned int vert_block_var[2]);
+#endif // !CONFIG_REALTIME_ONLY
+
+// A simplified version of set_offsets meant to be used for
+// simple_motion_search.
+static INLINE void set_offsets_for_motion_search(const AV1_COMP *const cpi,
+ MACROBLOCK *const x,
+ int mi_row, int mi_col,
+ BLOCK_SIZE bsize) {
+ const AV1_COMMON *const cm = &cpi->common;
+ const CommonModeInfoParams *const mi_params = &cm->mi_params;
+ const int num_planes = av1_num_planes(cm);
+ MACROBLOCKD *const xd = &x->e_mbd;
+ const int mi_width = mi_size_wide[bsize];
+ const int mi_height = mi_size_high[bsize];
+
+ set_mode_info_offsets(&cpi->common.mi_params, &cpi->mbmi_ext_info, x, xd,
+ mi_row, mi_col);
+
+ // Set up destination pointers.
+ av1_setup_dst_planes(xd->plane, bsize, &cm->cur_frame->buf, mi_row, mi_col, 0,
+ num_planes);
+
+ // Set up limit values for MV components.
+ // Mv beyond the range do not produce new/different prediction block.
+ av1_set_mv_limits(mi_params, &x->mv_limits, mi_row, mi_col, mi_height,
+ mi_width, cpi->oxcf.border_in_pixels);
+
+ set_plane_n4(xd, mi_width, mi_height, num_planes);
+
+ xd->mi_row = mi_row;
+ xd->mi_col = mi_col;
+
+ // Set up distance of MB to edge of frame in 1/8th pel units.
+ assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));
+ xd->mb_to_top_edge = -GET_MV_SUBPEL(mi_row * MI_SIZE);
+ xd->mb_to_bottom_edge =
+ GET_MV_SUBPEL((mi_params->mi_rows - mi_height - mi_row) * MI_SIZE);
+ xd->mb_to_left_edge = -GET_MV_SUBPEL(mi_col * MI_SIZE);
+ xd->mb_to_right_edge =
+ GET_MV_SUBPEL((mi_params->mi_cols - mi_width - mi_col) * MI_SIZE);
+
+ // Set up source buffers.
+ av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes, bsize);
+}
+
+void av1_init_simple_motion_search_mvs_for_sb(const AV1_COMP *cpi,
+ const TileInfo *tile_info,
+ MACROBLOCK *x,
+ SIMPLE_MOTION_DATA_TREE *sms_root,
+ int mi_row, int mi_col);
+
+static INLINE int is_full_sb(const CommonModeInfoParams *const mi_params,
+ int mi_row, int mi_col, BLOCK_SIZE sb_size) {
+ const int sb_mi_wide = mi_size_wide[sb_size];
+ const int sb_mi_high = mi_size_high[sb_size];
+
+ return (mi_row + sb_mi_high) <= mi_params->mi_rows &&
+ (mi_col + sb_mi_wide) <= mi_params->mi_cols;
+}
+
+#if !CONFIG_REALTIME_ONLY
+// Do not use this criteria for screen content videos.
+// Since screen content videos could often find good predictors and the largest
+// block size is likely to be used.
+static INLINE int use_auto_max_partition(const AV1_COMP *const cpi,
+ BLOCK_SIZE sb_size, int mi_row,
+ int mi_col) {
+ assert(IMPLIES(cpi->ppi->gf_group.size > 0,
+ cpi->gf_frame_index < cpi->ppi->gf_group.size));
+ const AV1_COMMON *const cm = &cpi->common;
+ return !frame_is_intra_only(cm) && !cpi->use_screen_content_tools &&
+ cpi->sf.part_sf.auto_max_partition_based_on_simple_motion !=
+ NOT_IN_USE &&
+ sb_size == BLOCK_128X128 &&
+ is_full_sb(&cm->mi_params, mi_row, mi_col, sb_size) &&
+ cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
+ OVERLAY_UPDATE &&
+ cpi->ppi->gf_group.update_type[cpi->gf_frame_index] !=
+ INTNL_OVERLAY_UPDATE;
+}
+
+static BLOCK_SIZE dim_to_size(int dim) {
+ switch (dim) {
+ case 4: return BLOCK_4X4;
+ case 8: return BLOCK_8X8;
+ case 16: return BLOCK_16X16;
+ case 32: return BLOCK_32X32;
+ case 64: return BLOCK_64X64;
+ case 128: return BLOCK_128X128;
+ default: assert(0); return 0;
+ }
+}
+
+static AOM_INLINE void set_max_min_partition_size(SuperBlockEnc *sb_enc,
+ AV1_COMP *cpi, MACROBLOCK *x,
+ const SPEED_FEATURES *sf,
+ BLOCK_SIZE sb_size,
+ int mi_row, int mi_col) {
+ const AV1_COMMON *cm = &cpi->common;
+
+ sb_enc->max_partition_size =
+ AOMMIN(sf->part_sf.default_max_partition_size,
+ dim_to_size(cpi->oxcf.part_cfg.max_partition_size));
+ sb_enc->min_partition_size =
+ AOMMAX(sf->part_sf.default_min_partition_size,
+ dim_to_size(cpi->oxcf.part_cfg.min_partition_size));
+ sb_enc->max_partition_size =
+ AOMMIN(sb_enc->max_partition_size, cm->seq_params->sb_size);
+ sb_enc->min_partition_size =
+ AOMMIN(sb_enc->min_partition_size, cm->seq_params->sb_size);
+
+ if (use_auto_max_partition(cpi, sb_size, mi_row, mi_col)) {
+ float features[FEATURE_SIZE_MAX_MIN_PART_PRED] = { 0.0f };
+
+ av1_get_max_min_partition_features(cpi, x, mi_row, mi_col, features);
+ sb_enc->max_partition_size =
+ AOMMAX(AOMMIN(av1_predict_max_partition(cpi, x, features),
+ sb_enc->max_partition_size),
+ sb_enc->min_partition_size);
+ }
+}
+#endif // !CONFIG_REALTIME_ONLY
+#endif // AOM_AV1_ENCODER_PARTITION_STRATEGY_H_
generated by cgit v1.2.3 (git 2.39.1) at 2024-08-27 07:30:13 +0000