Skip to content

Fast two-pass: additional dry-pass simplifications #5435

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leolzhao merged 1 commit into from
Sep 25, 2026
+118 −37
Merged

Fast two-pass: additional dry-pass simplifications #5435

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 35 additions & 21 deletions av2/encoder/encodeframe.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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,
Expand Down Expand Up @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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".
Expand Down Expand Up @@ -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;
Expand All @@ -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);
}
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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.
Expand Down
10 changes: 8 additions & 2 deletions av2/encoder/partition_search.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
21 changes: 21 additions & 0 deletions av2/encoder/partition_strategy.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
68 changes: 54 additions & 14 deletions av2/encoder/tx_search.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -2594,16 +2616,22 @@ 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) {
txfm_param.sec_tx_set_idx = set_idx;
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;
Expand Down Expand Up @@ -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 };
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand Down
Loading