Gate partition ML/pruning in fast two pass partitions

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;

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;

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

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

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 &&

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 &&

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;