Prune reference frames for non-RD partition in realtime mode

av2/encoder/partition_search.c

@@ -2646,6 +2646,14 @@ void av2_nonrd_use_partition(AV2_COMP *cpi, ThreadData *td,
                     bsize, ctx_none, best_rdc, SB_SINGLE_PASS, NULL);
       encode_b(cpi, tile_data, td, tp, mi_row, mi_col, OUTPUT_ENABLED, subsize,
                partition, ctx_none, &rate);
+      // Record picked ref frame in the superblock mask to allow subsequent
+      // sub-blocks to prune unselected reference frames in non-RD mode.
+      if (cpi->sf.inter_sf.prune_ref_frames >= 3 &&
+          xd->tree_type != CHROMA_PART) {
+        const int ref_type = av2_ref_frame_type(ctx_none->mic.ref_frame);
+        av2_update_picked_ref_frames_mask(x, ref_type, bsize, cm->mib_size,
+                                          mi_row, mi_col);
+      }
       break;
     case PARTITION_HORZ:
       if (!pc_tree->horizontal[cur_region_type][0]) {

av2/encoder/rdopt.c

@@ -7906,10 +7906,10 @@ static uint64_t fetch_picked_ref_frames_mask(const MACROBLOCK *const x,
   const MACROBLOCKD *const xd = &x->e_mbd;
   const int mi_row = xd->mi_row;
   const int mi_col = xd->mi_col;
-  const int mi_row_in_sb = mi_row & sb_size_mask;
-  const int mi_col_in_sb = mi_col & sb_size_mask;
   const int mi_w = mi_size_wide[bsize];
   const int mi_h = mi_size_high[bsize];
+  const int mi_row_in_sb = (mi_h >= mib_size) ? 0 : (mi_row & sb_size_mask);
+  const int mi_col_in_sb = (mi_w >= mib_size) ? 0 : (mi_col & sb_size_mask);
   uint64_t picked_ref_frames_mask = 0;
   for (int i = mi_row_in_sb; i < mi_row_in_sb + mi_h; ++i) {
     for (int j = mi_col_in_sb; j < mi_col_in_sb + mi_w; ++j) {
@@ -8114,6 +8114,18 @@ static int inter_mode_search_order_independent_skip(
       skip_motion_mode = (ref_type < INTER_REFS_PER_FRAME &&
                           x->inter_mode_cache[0]->ref_frame[1] != INTRA_FRAME);
     }
+
+    if (cpi->sf.inter_sf.prune_ref_frames >= 3 && ref_type > 0) {
+      if (search_state->best_rd != INT64_MAX) {
+        // Prune secondary references if ref 0 already produced a skippable mode
+        // (zero residual). If ref 0 match was not skippable, do not prune:
+        // allow evaluating secondary references for this block.
+        if (!search_state->best_mode_skippable) {
+          skip_ref = 0;
+        }
+      }
+    }
+
     if (skip_ref) return 1;
   }
 
@@ -9525,15 +9537,28 @@ void av2_rd_pick_inter_mode_sb(struct AV2_COMP *cpi,
   // Ref frames that are selected by square partition blocks.
   uint64_t picked_ref_frames_mask = 0;
   if (inter_sf->prune_ref_frames && !x->inter_mode_cache[0]) {
-    assert(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG));
-
-    // Prune reference frames if we are either a 1:4 block, or if we are a 1:2
-    // block, and we have searched any of the rectangular subblock.
-    if (!is_partition_point(bsize) ||
-        has_searched_rect_subblock(x, xd->mi_row, xd->mi_col, bsize,
-                                   cm->sb_size, (int8_t)mbmi->region_type)) {
-      picked_ref_frames_mask =
-          fetch_picked_ref_frames_mask(x, bsize, cm->mib_size);
+    if (cpi->sf.rt_sf.use_nonrd_partition) {
+      // In non-RD partition mode (realtime variance-based partitioning):
+      // Prune references for sub-blocks (< sb_size) based on reference frames
+      // picked by earlier sub-blocks in this superblock.
+      if (cpi->sf.inter_sf.prune_ref_frames >= 3 && bsize < cm->sb_size) {
+        picked_ref_frames_mask =
+            fetch_picked_ref_frames_mask(x, cm->sb_size, cm->mib_size);
+        if (picked_ref_frames_mask) {
+          // Always ensure the primary reference (ref 0) is tested.
+          picked_ref_frames_mask |= (1ULL << 0);
+        }
+      }
+    } else {
+      assert(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG));
+      // Prune reference frames if we are either a 1:4 block, or if we are a 1:2
+      // block, and we have searched any of the rectangular subblock.
+      if (!is_partition_point(bsize) ||
+          has_searched_rect_subblock(x, xd->mi_row, xd->mi_col, bsize,
+                                     cm->sb_size, (int8_t)mbmi->region_type)) {
+        picked_ref_frames_mask =
+            fetch_picked_ref_frames_mask(x, bsize, cm->mib_size);
+      }
     }
   }
 

av2/encoder/speed_features.c

@@ -754,7 +754,7 @@ static void set_rt_speed_features_framesize_independent(
   // Set this good features as default for now.
   set_good_speed_features_framesize_independent(cpi, sf, speed);
   if (speed >= 6) {
-    sf->inter_sf.prune_ref_frames = 0;
+    sf->inter_sf.prune_ref_frames = 3;
     sf->intra_sf.intra_pruning_with_mlp = 0;
     sf->hl_sf.frame_parameter_update = 0;
     sf->hl_sf.recode_loop = DISALLOW_RECODE;

av2/encoder/speed_features.h

@@ -646,6 +646,8 @@ typedef struct INTER_MODE_SPEED_FEATURES {
   // 0 implies no pruning
   // 1 implies prune for extended partition
   // 2 implies prune horiz, vert and extended partition
+  // 3: prune for all sub-blocks (< sb_size) based on superblock mask gated
+  //    by best_mode_skippable
   int prune_ref_frames;
 
   // When set to N>0, this flag limits the maximum number of compound