From a51fc8023446a8ec3fe2b14d9f2725100a76ce7b Mon Sep 17 00:00:00 2001 From: Yeqing Wu Date: Mon, 24 Aug 2026 10:25:52 -0700 Subject: [PATCH] Gate partition ML/pruning in fast two pass partitions This newly activates fast two pass partition for speed=2, while improving speed-quality ratio for speed=3/4. At speed=2 it is turned on only for class A1 (2160p and above). Class A2 (2K) lost coding efficiency at speed=2, so we leave it off there. Speeds 3 and 4 are not changed and stay on for every resolution. Some of the partition ML/pruning schemes interact poorly with the fast two-pass partition search, so skip them in the dry pass and on trusted wet-pass blocks. The full policy across regimes: * FAST dry pass: disabled. Shape is picked from reduced-tool RD, and the dry pass is already fast. * FAST wet pass, forced: disabled. Partition is already fixed by the dry-pass template; the decision is discarded. * FAST wet pass, reopened: enabled. Block is genuinely searched. * CONSERVATIVE / one-pass: enabled. Predicate is inert on this path. Anchor: commit 7361d54 Speed 4 (cpu-used=4): FG16 CTC (33 frames, class A1 and A2, RA) Speed 3 (cpu-used=3): FG16 CTC (33 frames, class A1 and A2, RA) Speed 2 (cpu-used=2): FG16 CTC (33 frames, class A1 and A2, RA) ``` 1) Speed 4 (cpu-used=4) +------------+-------+-------+-------+-------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+-------+-------+-------+-------+------+------+ | A1 | -0.73 | 0.26 | -0.05 | -0.66 | 106 | 101 | | A2 | -0.49 | 0.22 | 0.20 | -0.43 | 105 | 101 | | Avg w/o B2 | -0.56 | 0.23 | 0.12 | -0.50 | 105 | 101 | +------------+-------+-------+-------+-------+------+------+ 2) Speed 3 (cpu-used=3) +------------+-------+-------+-------+-------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+-------+-------+-------+-------+------+------+ | A1 | -0.75 | 0.08 | 0.17 | -0.68 | 106 | 101 | | A2 | -0.47 | 0.18 | -0.15 | -0.43 | 104 | 100 | | Avg w/o B2 | -0.55 | 0.15 | -0.05 | -0.51 | 105 | 100 | +------------+-------+-------+-------+-------+------+------+ 3) Speed 2 (cpu-used=2) +------------+-------+-------+-------+-------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+-------+-------+-------+-------+------+------+ | A1 | -0.50 | -2.10 | -2.72 | -0.65 | 84 | 87 | | A2 | 0.46 | 1.97 | 1.60 | 0.56 | 83 | 95 | | Avg w/o B2 | 0.17 | 0.76 | 0.32 | 0.20 | 84 | 93 | +------------+-------+-------+-------+-------+------+------+ ``` STATS_CHANGED --- av2/encoder/block.h | 6 +-- av2/encoder/encodeframe.c | 76 ++++++++++++++++++++++---------- av2/encoder/encodeframe_utils.h | 19 +++++++- av2/encoder/encoder_utils.h | 14 +++--- av2/encoder/partition_search.c | 4 +- av2/encoder/partition_strategy.c | 6 ++- av2/encoder/speed_features.c | 6 --- 7 files changed, 90 insertions(+), 41 deletions(-) diff --git a/av2/encoder/block.h b/av2/encoder/block.h index ace2f6b48f..8fac5dee67 100644 --- a/av2/encoder/block.h +++ b/av2/encoder/block.h @@ -1678,9 +1678,9 @@ typedef struct macroblock { * dry_pass_rd_slot(). Reset per superblock. */ int64_t unit_dry_rd[TWO_PASS_DRY_RD_SLOTS]; - /*! \brief True during the wet pass of the SB-level two-pass partition - * search; gates consumption of any dry-pass side information. */ - bool consume_dry_pass_info; + /*! \brief True during the wet pass of the fast SB-level two-pass partition + * search; a block on this path may consume the dry-pass template. */ + bool may_consume_dry_pass_info; /*! \brief Cap on top intra Y candidates for full RD. */ int intra_mode_prune_top; diff --git a/av2/encoder/encodeframe.c b/av2/encoder/encodeframe.c index 267e92fc14..8984de2857 100644 --- a/av2/encoder/encodeframe.c +++ b/av2/encoder/encodeframe.c @@ -598,13 +598,13 @@ static AVM_INLINE void perform_two_partition_passes( * * \ingroup partition_search * - * The wet pass trusts the dry-pass shape for large blocks and searches the - * small ones again. This walks the dry-pass tree and clears the shape of every - * node smaller than min_bsize so the wet pass is free to choose its own. + * Clears the shape of every node smaller than min_bsize so the wet pass is + * free to choose its own. If resplit_max_side_px > 0, also clears unsplit + * blocks up to that pixel size; 0 disables the re-split step. */ static AVM_INLINE void set_min_none_to_invalid(PARTITION_TREE *part_tree, BLOCK_SIZE min_bsize, - bool allow_none_resplit) { + int resplit_max_side_px) { if (!part_tree) return; const BLOCK_SIZE bsize = part_tree->bsize; const PARTITION_TYPE part_type = part_tree->partition; @@ -617,14 +617,11 @@ static AVM_INLINE void set_min_none_to_invalid(PARTITION_TREE *part_tree, return; } - // Large blocks the dry pass left unsplit: re-split them in the wet pass - // (only on the fast level, guarded by allow_none_resplit). The conservative - // level keeps its pre-PR semantics. - if (allow_none_resplit && part_type == PARTITION_NONE) { - // Only do this for blocks up to 128 px. Bigger unsplit blocks are usually - // genuinely flat, so searching them again costs a lot of time for little - // gain. - if (AVMMAX(block_size_wide[bsize], block_size_high[bsize]) <= 128) { + // Re-split unsplit blocks up to the cutoff. Bigger unsplit blocks are + // usually flat, so re-searching them rarely pays. + if (resplit_max_side_px > 0 && part_type == PARTITION_NONE) { + if (AVMMAX(block_size_wide[bsize], block_size_high[bsize]) <= + resplit_max_side_px) { part_tree->partition = PARTITION_INVALID; } return; @@ -649,17 +646,41 @@ static AVM_INLINE void set_min_none_to_invalid(PARTITION_TREE *part_tree, for (int idx = 0; idx < num_subtrees; idx++) { set_min_none_to_invalid(part_tree->sub_tree[idx], min_bsize, - allow_none_resplit); + resplit_max_side_px); } } +// 8-bit-equivalent qindex: strips the bit-depth offset so the threshold +// below matches the command-line --qp regardless of input bit depth. +static AVM_INLINE int two_pass_qindex_8bit_equiv(int base_qindex, + int bit_depth) { + const int qindex_8b = base_qindex - (bit_depth - AVM_BITS_8) * MAXQ_OFFSET; + assert(qindex_8b >= 0); + return qindex_8b; +} + +// Above this qindex the schedule uses a larger dry-pass floor and a +// smaller wet-pass re-split cutoff. +#define TWO_PASS_QP_SCHEDULE_HIGH_QP_THRESH 200 + +// Dry-pass floor. 32x32 at high QP, 16x16 otherwise. +static AVM_INLINE BLOCK_SIZE two_pass_dry_pass_floor(int qindex_8b) { + return qindex_8b >= TWO_PASS_QP_SCHEDULE_HIGH_QP_THRESH ? BLOCK_32X32 + : BLOCK_16X16; +} + +// Wet-pass re-split cutoff in pixels. 64 at high QP, 128 otherwise. +static AVM_INLINE int two_pass_resplit_max_side_px(int qindex_8b) { + return qindex_8b >= TWO_PASS_QP_SCHEDULE_HIGH_QP_THRESH ? 64 : 128; +} + /*!\brief Performs partition search in two passes. * * \ingroup partition_search - * In the first pass, partition search is performed with the - * minimum bsize set to BLOCK_16X16. In the second pass, partition search is - * performed with the same partition tree from the first pass, but partition - * search is allowed to search recursively starting from BLOCK_32X32. + * First pass picks a rough shape with a reduced-tool RD; second pass + * refines it. Minimum bsize is BLOCK_16X16 (BLOCK_32X32 at high QP on the + * fast level); the second pass recurses from BLOCK_32X32 (BLOCK_64X64 at + * high QP on the fast level). Conservative level uses the low-QP values. */ static AVM_INLINE void perform_two_pass_partition_search( AV2_COMP *cpi, ThreadData *td, TileDataEnc *tile_data, TokenExtra **tp, @@ -673,12 +694,20 @@ static AVM_INLINE void perform_two_pass_partition_search( const BLOCK_SIZE sb_size = cm->sb_size; assert(!frame_is_intra_only(cm)); + // QP schedule is fast-only. Conservative uses 16x16 floor and 0 (skip + // re-split), leaving its behavior unchanged. + const int qindex_8b = two_pass_qindex_8bit_equiv(cm->quant_params.base_qindex, + cm->seq_params.bit_depth); + const BLOCK_SIZE dry_floor = + fast_two_pass ? two_pass_dry_pass_floor(qindex_8b) : BLOCK_16X16; + const int resplit_max_side = + fast_two_pass ? two_pass_resplit_max_side_px(qindex_8b) : 0; + // First pass to estimate partition structures SB_FIRST_PASS_STATS sb_fp_stats; av2_backup_sb_state(&sb_fp_stats, cpi, td, tile_data, mi_row, mi_col); - // The dry pass does not go below 16x16: it only needs a rough shape, and it - // scores blocks with a reduced set of tools. - x->sb_enc.min_partition_size = BLOCK_16X16; + // Dry pass: rough shape only, floor at dry_floor. + x->sb_enc.min_partition_size = dry_floor; // Drop the previous superblock's dry-pass rd records. The recorded rdcosts // are in the dry pass's rdmult units, which the wet pass must not recompute. if (fast_two_pass) av2_zero(x->unit_dry_rd); @@ -687,9 +716,10 @@ static AVM_INLINE void perform_two_pass_partition_search( PARTITION_TREE *part_ref = xd->sbi->ptree_root[0]; // Set this to NULL otherwise part_ref will get freed in the second pass. xd->sbi->ptree_root[0] = NULL; - // Trust the dry-pass shape for >=32x32 blocks; re-search smaller ones. - set_min_none_to_invalid(part_ref, get_larger_sqr_bsize(BLOCK_16X16), - fast_two_pass); + // Trust the dry-pass shape for >= (dry_floor << 1) blocks; re-search + // smaller ones. The trust boundary moves with the dry-pass floor. + set_min_none_to_invalid(part_ref, get_larger_sqr_bsize(dry_floor), + resplit_max_side); // Second pass RD_STATS dummy_rdc; diff --git a/av2/encoder/encodeframe_utils.h b/av2/encoder/encodeframe_utils.h index 859ce83a7d..f6531a7720 100644 --- a/av2/encoder/encodeframe_utils.h +++ b/av2/encoder/encodeframe_utils.h @@ -185,7 +185,8 @@ static AVM_INLINE void av2_set_two_pass_flags( const bool fast_two_pass = av2_two_pass_part_is_fast(&cpi->sf.part_sf); x->apply_dry_pass_shortcuts = fast_two_pass && (multi_pass_mode == SB_DRY_PASS); - x->consume_dry_pass_info = fast_two_pass && (multi_pass_mode == SB_WET_PASS); + x->may_consume_dry_pass_info = + fast_two_pass && (multi_pass_mode == SB_WET_PASS); const bool is_wet_pass_reuse = (fast_two_pass && multi_pass_mode == SB_WET_PASS && part_search_state && part_search_state->forced_partition != PARTITION_INVALID); @@ -206,6 +207,22 @@ static AVM_INLINE void av2_set_two_pass_flags( } } +// True when partition-level ML and fast-pruning heuristics should run. +// * FAST dry pass: false -- dry pass is already fast, and +// the models expect full-tool RD inputs. +// * FAST wet pass, forced shape: false -- shape is fixed by the dry pass, +// so the model output is discarded. +// * FAST wet pass, reopened: true -- block is genuinely searched. +// * CONSERVATIVE / one-pass: true -- both flags stay false here. +static AVM_INLINE bool av2_partition_ml_pruning_active( + const MACROBLOCK *x, PARTITION_TYPE forced_partition) { + if (x->apply_dry_pass_shortcuts) return false; + if (x->may_consume_dry_pass_info && forced_partition != PARTITION_INVALID) { + return false; + } + return true; +} + static AVM_INLINE void update_wedge_mode_cdf(FRAME_CONTEXT *fc, const BLOCK_SIZE bsize, const int8_t wedge_index diff --git a/av2/encoder/encoder_utils.h b/av2/encoder/encoder_utils.h index f734d564e9..7d85f6bb8a 100644 --- a/av2/encoder/encoder_utils.h +++ b/av2/encoder/encoder_utils.h @@ -959,15 +959,19 @@ void av2_scale_references(AV2_COMP *cpi, const InterpFilter filter, void av2_setup_frame(AV2_COMP *cpi); -// Config-time part of the TWO_PASS_PART_FAST predicate: GOOD mode at speed >= 3 -// on configs that produce inter frames. Kept config-only (no cm state) because -// av2_select_sb_size() also keys off it, and the superblock size must not -// change between the two passes. +// Config-time part of the TWO_PASS_PART_FAST predicate. // TODO(Yeqing): Extend to intra/key frames and drop the all_intra guard. static INLINE bool av2_wants_two_pass_partition(const AV2EncoderConfig *oxcf) { const int all_intra = oxcf->kf_cfg.key_freq_max == 0 && oxcf->kf_cfg.key_freq_min == 0; - return oxcf->mode == GOOD && oxcf->speed >= 3 && !all_intra; + if (oxcf->mode != GOOD || all_intra) return false; + if (oxcf->speed >= 3) return true; + if (oxcf->speed == 2) { + const int is_2160p_or_larger = + AVMMIN(oxcf->frm_dim_cfg.width, oxcf->frm_dim_cfg.height) >= 2160; + return is_2160p_or_larger; + } + return false; } BLOCK_SIZE av2_select_sb_size(const AV2_COMP *const cpi); diff --git a/av2/encoder/partition_search.c b/av2/encoder/partition_search.c index 09fe661f87..ad610463bc 100644 --- a/av2/encoder/partition_search.c +++ b/av2/encoder/partition_search.c @@ -590,7 +590,7 @@ static AVM_INLINE void record_dry_pass_rd(MACROBLOCK *x, BLOCK_SIZE sb_size, static AVM_INLINE int64_t dry_pass_rd_bound(const MACROBLOCK *x, BLOCK_SIZE sb_size, int mi_row, int mi_col, BLOCK_SIZE bsize) { - if (!x->consume_dry_pass_info) return 0; + if (!x->may_consume_dry_pass_info) return 0; const int slot_idx = dry_pass_rd_slot(sb_size, mi_row, mi_col, bsize); if (slot_idx < 0) return 0; const int64_t rd = x->unit_dry_rd[slot_idx]; @@ -3664,6 +3664,7 @@ static void prune_rect_partitions(AV2_COMP *const cpi, ThreadData *td, const BLOCK_SIZE bsize = blk_params->bsize; if (cpi->sf.part_sf.partition_pruning_with_mlp && + av2_partition_ml_pruning_active(x, part_search_state->forced_partition) && part_search_state->partition_allowed[PARTITION_NONE] && part_search_state->forced_partition == PARTITION_INVALID && block_size_wide[bsize] >= 32 && block_size_high[bsize] >= 32 && @@ -3957,6 +3958,7 @@ static void prune_partitions_after_none(AV2_COMP *const cpi, MACROBLOCK *x, // decision on early terminating at PARTITION_NONE. bool is_early_term_allowed = cpi->sf.part_sf.simple_motion_search_early_term_none && + av2_partition_ml_pruning_active(x, part_search_state->forced_partition) && !frame_is_intra_only(cm) && bsize >= BLOCK_16X16 && blk_params->mi_row_edge < mi_params->mi_rows && blk_params->mi_col_edge < mi_params->mi_cols && diff --git a/av2/encoder/partition_strategy.c b/av2/encoder/partition_strategy.c index efbf5ddf9a..d3aefa2220 100644 --- a/av2/encoder/partition_strategy.c +++ b/av2/encoder/partition_strategy.c @@ -881,8 +881,10 @@ void av2_prune_partitions_before_search( // for NONE partition. const int try_split_only = !cpi->is_screen_content_type && - cpi->sf.part_sf.simple_motion_search_split && *do_square_split && - bsize >= BLOCK_8X8 && + cpi->sf.part_sf.simple_motion_search_split && + av2_partition_ml_pruning_active( + x, partition_search_state->forced_partition) && + *do_square_split && bsize >= BLOCK_8X8 && mi_row + mi_size_high[bsize] <= mi_params->mi_rows && bsize < BLOCK_256X256 && mi_col + mi_size_wide[bsize] <= mi_params->mi_cols && diff --git a/av2/encoder/speed_features.c b/av2/encoder/speed_features.c index 39dd6f2b6c..7ecea411f0 100644 --- a/av2/encoder/speed_features.c +++ b/av2/encoder/speed_features.c @@ -1504,12 +1504,6 @@ static AVM_INLINE void set_two_pass_partition_level(AV2_COMP *cpi) { // TODO(Yeqing): Extend to intra/key frames. if (frame_is_intra_only(cm)) return; - // The fast level takes precedence: it is the speed >= 3 regime, and - // av2_select_sb_size() keys the superblock size off the same config-time - // predicate, so the size and the level have to agree. Note that this - // predicate is GOOD-only, which is what keeps the fast level out of REALTIME - // mode even though set_rt_speed_features_framesize_independent() runs the - // GOOD setters. if (av2_wants_two_pass_partition(&cpi->oxcf)) { *level = TWO_PASS_PART_FAST; return;