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; }