From 377700312523135286f9c07ca0bd68e1f48b3694 Mon Sep 17 00:00:00 2001 From: Yeqing Wu Date: Wed, 23 Sep 2026 09:43:05 -0700 Subject: [PATCH] Fast two-pass: additional dry-pass simplifications Three simplifications to the dry pass of fast two-pass partitions. 1. Skip deblock-aware distortion. The dry pass ranks shapes only, so reconstructing and filtering to nudge distortion by the deblocking delta is wasted precision. 2. Reuse the dry pass's simple motion searches in the wet pass. MVs, sse/var, rate/dist/rdcost and start MVs depend only on source, references and geometry, so recomputing them is redundant. A new preserve_sms_cache flag drives av2_reset_sms_cross_pass_state(), which clears only what must not cross the pass boundary. 3. Cap two transform-side inner loops in search_tx_type, select_tx_partition_type and choose_tx_size_type_from_rd on x->apply_dry_pass_shortcuts: force the (set_idx=0, stx=0) IST/stx baseline, and restrict the tx-partition search to NONE and SPLIT (2 of 8 variants). STATS_CHANGED Anchor: commit c5df831 Speed 4 (cpu-used=4): FG16 CTC (33 frames, class A1 and A2, RA) Speed 5 (cpu-used=5): FG16 CTC (33 frames, class A1 and A2, RA) ``` 1) Speed 4 (cpu-used=4) +------------+------+-------+-------+------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+------+-------+-------+------+------+------+ | A1 | 0.27 | -0.23 | -0.21 | 0.23 | 95 | 100 | | A2 | 0.31 | -0.42 | -0.46 | 0.25 | 95 | 100 | | Avg w/o B2 | 0.30 | -0.36 | -0.39 | 0.25 | 95 | 100 | +------------+------+-------+-------+------+------+------+ 2) Speed 5 (cpu-used=5) +------------+------+-------+-------+------+------+------+ | Class | Y | Cb | Cr | wAvg | Enc% | Dec% | +------------+------+-------+-------+------+------+------+ | A1 | 0.39 | 0.01 | 0.13 | 0.37 | 94 | 100 | | A2 | 0.26 | -0.39 | -0.57 | 0.20 | 94 | 100 | | Avg w/o B2 | 0.30 | -0.27 | -0.37 | 0.25 | 94 | 100 | +------------+------+-------+-------+------+------+------+ ``` --- av2/encoder/encodeframe.c | 56 ++++++++++++++++---------- av2/encoder/partition_search.c | 10 ++++- av2/encoder/partition_strategy.h | 21 ++++++++++ av2/encoder/tx_search.c | 68 +++++++++++++++++++++++++------- 4 files changed, 118 insertions(+), 37 deletions(-) diff --git a/av2/encoder/encodeframe.c b/av2/encoder/encodeframe.c index 37f74c1118..51769cd2ad 100644 --- a/av2/encoder/encodeframe.c +++ b/av2/encoder/encodeframe.c @@ -495,11 +495,16 @@ static void fill_sms_buf(SimpleMotionDataBufs *data_buf, } // This function initializes the stats for encode_rd_sb. +// +// preserve_sms_cache is set for the wet pass of a fast two-pass superblock, to +// keep the dry pass's simple-motion-search results. See +// av2_reset_sms_cross_pass_state(). static INLINE void init_encode_rd_sb(AV2_COMP *cpi, ThreadData *td, const TileDataEnc *tile_data, SIMPLE_MOTION_DATA_TREE *sms_root, RD_STATS *rd_cost, int mi_row, int mi_col, - int gather_tpl_data) { + int gather_tpl_data, + int preserve_sms_cache) { const AV2_COMMON *cm = &cpi->common; const TileInfo *tile_info = &tile_data->tile_info; MACROBLOCK *x = &td->mb; @@ -513,10 +518,11 @@ static INLINE void init_encode_rd_sb(AV2_COMP *cpi, ThreadData *td, sf->part_sf.simple_motion_search_early_term_none || sf->part_sf.ml_early_term_after_part_split_level) && !frame_is_intra_only(cm); - if (use_simple_motion_search) { + // Tree-level valid flags gate the motion search, so they must reset with the + // buffers. + if (use_simple_motion_search && !preserve_sms_cache) { init_simple_motion_search_mvs(sms_root); } - (void)sbi; init_ref_frame_space(cpi, td, mi_row, mi_col); x->sb_energy_level = 0; @@ -540,7 +546,13 @@ static INLINE void init_encode_rd_sb(AV2_COMP *cpi, ThreadData *td, if (sf->part_sf.partition_search_type != VAR_BASED_PARTITION && sf->part_sf.partition_search_type != FIXED_PARTITION) { SimpleMotionDataBufs *data_bufs = x->sms_bufs; - av2_init_sms_data_bufs(data_bufs); + // Full reset also separates one superblock from the next, so narrow rather + // than drop when carrying state across passes. + if (preserve_sms_cache) { + av2_reset_sms_cross_pass_state(data_bufs); + } else { + av2_init_sms_data_bufs(data_bufs); + } fill_sms_buf(data_bufs, sms_root, mi_row, mi_col, cm->sb_size, cm->sb_size, 0); fill_sms_buf(data_bufs, sms_root, mi_row, mi_col, cm->sb_size, cm->sb_size, @@ -571,7 +583,7 @@ static AVM_INLINE void perform_one_partition_pass( AV2_COMP *cpi, ThreadData *td, TileDataEnc *tile_data, TokenExtra **tp, TokenExtra **tp_chroma, const int mi_row, const int mi_col, const SB_MULTI_PASS_MODE multi_pass_mode, - const SbMultiPassParams *multi_pass_params) { + const SbMultiPassParams *multi_pass_params, const int preserve_sms_cache) { const AV2_COMMON *const cm = &cpi->common; MACROBLOCK *const x = &td->mb; MACROBLOCKD *const xd = &x->e_mbd; @@ -592,7 +604,7 @@ static AVM_INLINE void perform_one_partition_pass( (total_loop_num == 1 ? SHARED_PART : (loop_idx == 0 ? LUMA_PART : CHROMA_PART)); init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, - 1); + 1, preserve_sms_cache); PC_TREE *const pc_root = av2_alloc_pc_tree_node(xd->tree_type, mi_row, mi_col, cm->sb_size, sb_size, NULL, PARTITION_NONE, 0, 1, ss_x, ss_y); @@ -647,11 +659,12 @@ static AVM_INLINE void perform_two_partition_passes( REF_MV_BANK stored_mv_bank = td->mb.e_mbd.ref_mv_bank; WARP_PARAM_BANK stored_warp_bank = td->mb.e_mbd.warp_param_bank; perform_one_partition_pass(cpi, td, tile_data, tp, tp_chroma, mi_row, mi_col, - SB_DRY_PASS, NULL); + SB_DRY_PASS, NULL, 0); - // Second pass + // Second pass -- clean re-run (this helper proves SB_FIRST_PASS_STATS can + // reproduce a superblock). RD_STATS dummy_rdc; - init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, + init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, 0, 0); av2_reset_mbmi(&cm->mi_params, sb_size, mi_row, mi_col); av2_reset_simple_motion_tree_partition(sms_root, sb_size); @@ -660,7 +673,7 @@ static AVM_INLINE void perform_two_partition_passes( td->mb.e_mbd.ref_mv_bank = stored_mv_bank; td->mb.e_mbd.warp_param_bank = stored_warp_bank; perform_one_partition_pass(cpi, td, tile_data, tp, tp_chroma, mi_row, mi_col, - SB_WET_PASS, NULL); + SB_WET_PASS, NULL, 0); } /*!\brief Mark the small nodes of the dry-pass tree as "search again". @@ -781,7 +794,7 @@ static AVM_INLINE void perform_two_pass_partition_search( // 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); perform_one_partition_pass(cpi, td, tile_data, tp, tp_chroma, mi_row, mi_col, - SB_DRY_PASS, NULL); + SB_DRY_PASS, NULL, 0); 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; @@ -790,17 +803,18 @@ static AVM_INLINE void perform_two_pass_partition_search( set_min_none_to_invalid(part_ref, get_larger_sqr_bsize(dry_floor), resplit_max_side); - // Second pass + // Second pass. Fast reuses the dry pass's simple motion searches; + // conservative starts clean. RD_STATS dummy_rdc; - init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, - 0); + init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, 0, + fast_two_pass); av2_reset_mbmi(&cm->mi_params, sb_size, mi_row, mi_col); av2_reset_simple_motion_tree_partition(sms_root, sb_size); SbMultiPassParams multi_pass_params = { part_ref }; av2_restore_sb_state(&sb_fp_stats, cpi, td, tile_data, mi_row, mi_col); perform_one_partition_pass(cpi, td, tile_data, tp, tp_chroma, mi_row, mi_col, - SB_WET_PASS, &multi_pass_params); + SB_WET_PASS, &multi_pass_params, fast_two_pass); av2_free_ptree_recursive(part_ref); } @@ -839,8 +853,8 @@ static AVM_INLINE void encode_rd_sb(AV2_COMP *cpi, ThreadData *td, x->sms_bufs = td->sms_bufs; x->reuse_inter_mode_cache_type = cpi->sf.inter_sf.reuse_erp_mode_flag; - init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, - 1); + init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, mi_col, 1, + 0); const int intra_sdp_enabled = is_sdp_enabled_in_keyframe(cm); // Encode the superblock @@ -856,7 +870,7 @@ static AVM_INLINE void encode_rd_sb(AV2_COMP *cpi, ThreadData *td, (total_loop_num == 1 ? SHARED_PART : (loop_idx == 0 ? LUMA_PART : CHROMA_PART)); init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, - mi_col, 1); + mi_col, 1, 0); av2_reset_ptree_in_sbi(xd->sbi, xd->tree_type); av2_build_partition_tree_fixed_partitioning( cm, xd->tree_type, mi_row, mi_col, @@ -888,7 +902,7 @@ static AVM_INLINE void encode_rd_sb(AV2_COMP *cpi, ThreadData *td, (total_loop_num == 1 ? SHARED_PART : (loop_idx == 0 ? LUMA_PART : CHROMA_PART)); init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, - mi_col, 1); + mi_col, 1, 0); PC_TREE *const pc_root = av2_alloc_pc_tree_node( xd->tree_type, mi_row, mi_col, cm->sb_size, sb_size, NULL, PARTITION_NONE, 0, 1, ss_x, ss_y); @@ -925,7 +939,7 @@ static AVM_INLINE void encode_rd_sb(AV2_COMP *cpi, ThreadData *td, const int plane_end = get_partition_plane_end(xd->tree_type, num_planes); const BLOCK_SIZE min_partition_size = x->sb_enc.min_partition_size; init_encode_rd_sb(cpi, td, tile_data, sms_root, &dummy_rdc, mi_row, - mi_col, 1); + mi_col, 1, 0); PC_TREE *pc_root; if (cpi->sf.rt_sf.use_nonrd_partition) { if (!td->pc_root) { @@ -1001,7 +1015,7 @@ static AVM_INLINE void encode_rd_sb(AV2_COMP *cpi, ThreadData *td, av2_two_pass_part_is_fast(&sf->part_sf)); } else { perform_one_partition_pass(cpi, td, tile_data, tp, tp_chroma, mi_row, - mi_col, SB_SINGLE_PASS, NULL); + mi_col, SB_SINGLE_PASS, NULL, 0); } // Reset to 0 so that it wouldn't be used elsewhere mistakenly. diff --git a/av2/encoder/partition_search.c b/av2/encoder/partition_search.c index a20f887f70..719d906389 100644 --- a/av2/encoder/partition_search.c +++ b/av2/encoder/partition_search.c @@ -4263,6 +4263,12 @@ static void none_partition_search( const REGION_TYPE region_type = pc_tree->region_type; assert(bsize < BLOCK_SIZES_ALL); + // Skip the deblock-aware distortion refinement in the fast dry pass, which + // only ranks shapes. Derived from multi_pass_mode because + // x->apply_dry_pass_shortcuts is not yet set at the rd-bound site below. + const bool is_fast_dry_pass = av2_two_pass_part_is_fast(&cpi->sf.part_sf) && + multi_pass_mode == SB_DRY_PASS; + // Check if partition none is allowed. const int bw = block_size_wide[bsize]; const int bh = block_size_high[bsize]; @@ -4311,7 +4317,7 @@ static void none_partition_search( RD_STATS best_remain_rdcost; av2_rd_stats_subtraction(x->rdmult, best_rdc, &partition_rdcost, &best_remain_rdcost); - if (best_rdc->rdcost != INT64_MAX && + if (best_rdc->rdcost != INT64_MAX && !is_fast_dry_pass && cpi->sf.lpf_sf.enable_deblock_for_partition_search && !cpi->is_screen_content_type) { // increase the remaining best cost which could have be reduced by deblock @@ -4365,7 +4371,7 @@ static void none_partition_search( av2_add_mode_search_context_to_cache(sms_data, *ctx_none); } - if (cpi->sf.lpf_sf.enable_deblock_for_partition_search && + if (!is_fast_dry_pass && cpi->sf.lpf_sf.enable_deblock_for_partition_search && cm->lf.apply_deblocking_filter[0] && this_rdc->rate != INT_MAX && xd->tree_type != CHROMA_PART && !cpi->is_screen_content_type) { const int64_t distortion_offset = get_dist_offset_by_deblock( diff --git a/av2/encoder/partition_strategy.h b/av2/encoder/partition_strategy.h index c59eef5047..4c58e6dbfe 100644 --- a/av2/encoder/partition_strategy.h +++ b/av2/encoder/partition_strategy.h @@ -229,6 +229,27 @@ static INLINE void av2_init_sms_data_bufs(SimpleMotionDataBufs *data_bufs) { memset(data_bufs, 0, sizeof(*data_bufs)); } +// Boundary reset between the two passes of the fast two-pass partition search. +// Preserves the motion-search results (MVs, sse/var, rate/dist/rdcost, +// start_mv_list, ref_frame, rdmult, and ML residual stats -- all pass-invariant +// functions of source + reference + geometry) so the wet pass reuses them, and +// clears only what must not cross the pass boundary (mode_cache, old_sms, +// prev_partition). +static INLINE void av2_reset_sms_cross_pass_state( + SimpleMotionDataBufs *data_bufs) { + static_assert(sizeof(SimpleMotionDataBufs) % sizeof(SimpleMotionData) == 0, + "SimpleMotionDataBufs must hold a whole number of entries"); + SimpleMotionData *const entries = (SimpleMotionData *)data_bufs; + const size_t num_entries = sizeof(*data_bufs) / sizeof(SimpleMotionData); + for (size_t i = 0; i < num_entries; i++) { + SimpleMotionData *const entry = &entries[i]; + av2_zero(entry->mode_cache); + entry->old_sms = NULL; + entry->has_prev_partition = 0; + entry->prev_partition = (PARTITION_TYPE)0; + } +} + struct PartitionSearchState; void av2_gather_erp_rect_features( float *ml_features, AV2_COMP *cpi, MACROBLOCK *x, const TileInfo *tile_info, diff --git a/av2/encoder/tx_search.c b/av2/encoder/tx_search.c index bb9eb107b8..ce0754894a 100644 --- a/av2/encoder/tx_search.c +++ b/av2/encoder/tx_search.c @@ -2341,6 +2341,28 @@ static INLINE bool prune_tx_type_rd_calc_using_tx_domain_dist( return false; } +// Cap on tx-partition-type sweep in the fast two-pass dry pass. Value 2 = +// TX_PARTITION_NONE + TX_PARTITION_SPLIT. +#define DRY_PASS_ALLOWED_MAX_TX_PARTITION_TYPES 2 + +// True when the fast two-pass dry pass should cap its inner tx-search loops +// (IST set / stx sweep and tx-partition-type sweep). Gated on speed >= 4. +static AVM_INLINE bool dry_pass_caps_tx_search(const MACROBLOCK *x, + const AV2_COMP *cpi) { + return x->apply_dry_pass_shortcuts && cpi->oxcf.speed >= 4; +} + +// True when the caller's winner-mode evaluation has already picked a +// tx-partition in mbmi->tx_partition_type[] for the current block, so a +// tx-partition search should defer to it instead of re-evaluating from scratch. +static AVM_INLINE bool winner_mode_tx_partition_preselected( + const TxfmSearchParams *txfm_params, const MB_MODE_INFO *mbmi, + const AV2_COMP *cpi) { + return txfm_params->eval_mode_type == WINNER_MODE_EVAL && + mbmi->region_type == MIXED_INTER_INTRA_REGION && + cpi->sf.winner_mode_sf.disable_multiway_tx_part_in_rough_mode; +} + // Search for the best transform type for a given transform block. // This function can be used for both inter and intra, both luma and chroma. static void search_tx_type(const AV2_COMP *cpi, MACROBLOCK *x, int plane, @@ -2594,8 +2616,11 @@ static void search_tx_type(const AV2_COMP *cpi, MACROBLOCK *x, int plane, dc_only_blk || (eob_found) || !xd->enable_ist); bool skip_idx = false; + // Fast dry pass ranks shapes only; force the (set_idx=0, stx=0) baseline. const int max_set_id = - get_ist_max_set_id(skip_stx, is_inter, txw, txh, primary_tx_type); + dry_pass_caps_tx_search(x, cpi) + ? 1 + : get_ist_max_set_id(skip_stx, is_inter, txw, txh, primary_tx_type); assert(max_set_id < IST_SET_SIZE); for (int set_idx = 0; set_idx < max_set_id; ++set_idx) { @@ -2603,7 +2628,10 @@ static void search_tx_type(const AV2_COMP *cpi, MACROBLOCK *x, int plane, const uint8_t set_id = get_ist_set_id(set_idx, is_inter, intra_mode, txw, txh, primary_tx_type); - const int max_stx = xd->enable_ist && !(eob_found) ? STX_TYPES : 1; + const int max_stx = + (xd->enable_ist && !(eob_found) && !dry_pass_caps_tx_search(x, cpi)) + ? STX_TYPES + : 1; const int init_stx = (set_idx > 0) ? 1 : 0; for (int stx = init_stx; stx < max_stx; ++stx) { if (eob_found) skip_stx = true; @@ -3373,12 +3401,11 @@ static void select_tx_partition_type( is_inter_block(mbmi, xd->tree_type) ? av2_get_txb_size_index(plane_bsize, blk_row, blk_col) : 0; + const bool tx_partition_preselected = + winner_mode_tx_partition_preselected(txfm_params, mbmi, cpi); TX_PARTITION_TYPE best_tx_partition = - txfm_params->eval_mode_type == WINNER_MODE_EVAL && - mbmi->region_type == MIXED_INTER_INTRA_REGION && - cpi->sf.winner_mode_sf.disable_multiway_tx_part_in_rough_mode - ? mbmi->tx_partition_type[txb_size_index] - : TX_PARTITION_INVALID; + tx_partition_preselected ? mbmi->tx_partition_type[txb_size_index] + : TX_PARTITION_INVALID; uint8_t best_partition_entropy_ctxs[MAX_TX_PARTITIONS] = { 0 }; TX_TYPE best_partition_tx_types[MAX_TX_PARTITIONS] = { 0 }; uint8_t full_blk_skip[MAX_TX_PARTITIONS] = { 0 }; @@ -3389,7 +3416,16 @@ static void select_tx_partition_type( int stationary_cols = 0; int stationarity_valid = 0; - for (TX_PARTITION_TYPE type = 0; type < TX_PARTITION_TYPES; ++type) { + // Fast dry pass ranks shapes only; cap to NONE + SPLIT (NONE alone + // under-ranks blocks that want a split). Skip the cap when the caller has + // pre-selected a multiway partition via WINNER_MODE_EVAL, so we do not + // silently discard it. + const int max_tx_part_type = + (dry_pass_caps_tx_search(x, cpi) && !tx_partition_preselected) + ? DRY_PASS_ALLOWED_MAX_TX_PARTITION_TYPES + : TX_PARTITION_TYPES; + + for (TX_PARTITION_TYPE type = 0; type < max_tx_part_type; ++type) { if (cpi->oxcf.txfm_cfg.reduced_tx_part_set && type > TX_PARTITION_VERT) { break; } @@ -3728,17 +3764,21 @@ static void choose_tx_size_type_from_rd(const AV2_COMP *const cpi, int is_wide_angle_mapped[MAX_TX_PARTITIONS] = { 0 }; int mapped_wide_angle[MAX_TX_PARTITIONS] = { 0 }; assert(!is_inter_block(mbmi, xd->tree_type)); + const bool tx_partition_preselected = + winner_mode_tx_partition_preselected(txfm_params, mbmi, cpi); TX_PARTITION_TYPE best_tx_partition_type = - txfm_params->eval_mode_type == WINNER_MODE_EVAL && - mbmi->region_type == MIXED_INTER_INTRA_REGION && - cpi->sf.winner_mode_sf.disable_multiway_tx_part_in_rough_mode - ? mbmi->tx_partition_type[0] - : TX_PARTITION_NONE; + tx_partition_preselected ? mbmi->tx_partition_type[0] : TX_PARTITION_NONE; int64_t best_rd = INT64_MAX; x->rd_model = FULL_TXFM_RD; int64_t cur_rd = INT64_MAX; const bool is_rect = is_rect_tx(max_tx_size); - for (TX_PARTITION_TYPE type = 0; type < TX_PARTITION_TYPES; ++type) { + // Same NONE+SPLIT cap (and WINNER_MODE_EVAL escape) as in + // select_tx_partition_type(). + const int max_tx_part_type = + (dry_pass_caps_tx_search(x, cpi) && !tx_partition_preselected) + ? DRY_PASS_ALLOWED_MAX_TX_PARTITION_TYPES + : TX_PARTITION_TYPES; + for (TX_PARTITION_TYPE type = 0; type < max_tx_part_type; ++type) { if (cpi->oxcf.txfm_cfg.reduced_tx_part_set && type > TX_PARTITION_VERT) { break; }