Prune uneven 4A/4B partitions with neighbor block boundaries

av2/common/av2_common_int.h

@@ -6050,6 +6050,22 @@ static INLINE bool av2_skip_reference_buffer_update(
   return clear_multiple_insert_in_one && ref_index != first_ref_index;
 }
 
+// Returns the starting mi location of chroma reference block for the current
+// mbmi, by setting `chroma_mi_row_start` and `chroma_mi_col_start`.
+static INLINE void av2_get_chroma_start_location(const MB_MODE_INFO *mbmi,
+                                                 TREE_TYPE tree_type,
+                                                 int *chroma_mi_row_start,
+                                                 int *chroma_mi_col_start) {
+  assert(tree_type == SHARED_PART || tree_type == CHROMA_PART);
+  if (tree_type == SHARED_PART) {
+    *chroma_mi_row_start = mbmi->chroma_ref_info.mi_row_chroma_base;
+    *chroma_mi_col_start = mbmi->chroma_ref_info.mi_col_chroma_base;
+  } else {
+    *chroma_mi_row_start = mbmi->chroma_mi_row_start;
+    *chroma_mi_col_start = mbmi->chroma_mi_col_start;
+  }
+}
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

av2/common/av2_loopfilter.c

@@ -138,22 +138,6 @@ void av2_loop_filter_frame_init(AV2_COMMON *cm, int plane_start,
   }
 }
 
-// Returns the starting mi location of chroma reference block for the current
-// mbmi, by setting `chroma_mi_row_start` and `chroma_mi_col_start`.
-static void get_chroma_start_location(const MB_MODE_INFO *mbmi,
-                                      TREE_TYPE tree_type,
-                                      int *chroma_mi_row_start,
-                                      int *chroma_mi_col_start) {
-  assert(tree_type == SHARED_PART || tree_type == CHROMA_PART);
-  if (tree_type == SHARED_PART) {
-    *chroma_mi_row_start = mbmi->chroma_ref_info.mi_row_chroma_base;
-    *chroma_mi_col_start = mbmi->chroma_ref_info.mi_col_chroma_base;
-  } else {
-    *chroma_mi_row_start = mbmi->chroma_mi_row_start;
-    *chroma_mi_col_start = mbmi->chroma_mi_col_start;
-  }
-}
-
 // Returns true if we are at the transform boundary.
 static bool is_tu_edge_helper(TX_SIZE tx_size, EDGE_DIR edge_dir,
                               int relative_row, int relative_col) {
@@ -193,7 +177,8 @@ static TX_SIZE get_transform_size(const MACROBLOCKD *const xd,
     int mi_row_start = mbmi->mi_row_start;
     int mi_col_start = mbmi->mi_col_start;
     if (plane != AVM_PLANE_Y) {
-      get_chroma_start_location(mbmi, tree_type, &mi_row_start, &mi_col_start);
+      av2_get_chroma_start_location(mbmi, tree_type, &mi_row_start,
+                                    &mi_col_start);
     }
     *tu_edge = is_tu_edge_helper(
         tx_size, edge_dir, (mi_row - mi_row_start) >> plane_ptr->subsampling_y,
@@ -211,8 +196,8 @@ static TX_SIZE get_transform_size(const MACROBLOCKD *const xd,
                                     plane_ptr->subsampling_y);
     int chroma_mi_row_start;
     int chroma_mi_col_start;
-    get_chroma_start_location(mbmi, tree_type, &chroma_mi_row_start,
-                              &chroma_mi_col_start);
+    av2_get_chroma_start_location(mbmi, tree_type, &chroma_mi_row_start,
+                                  &chroma_mi_col_start);
     *tu_edge = is_tu_edge_helper(
         tx_size, edge_dir,
         (mi_row - chroma_mi_row_start) >> plane_ptr->subsampling_y,
@@ -329,8 +314,8 @@ static uint32_t get_pu_starting_cooord(const MB_MODE_INFO *const mbmi,
   } else {
     int chroma_mi_row_start;
     int chroma_mi_col_start;
-    get_chroma_start_location(mbmi, tree_type, &chroma_mi_row_start,
-                              &chroma_mi_col_start);
+    av2_get_chroma_start_location(mbmi, tree_type, &chroma_mi_row_start,
+                                  &chroma_mi_col_start);
     pu_starting_mi = vert_edge ? chroma_mi_col_start : chroma_mi_row_start;
   }
   const uint32_t pu_stating_coord_luma = pu_starting_mi * MI_SIZE;
@@ -510,8 +495,8 @@ static int get_remaining_mi_size(const MB_MODE_INFO *mbmi,
   } else {
     int mi_row_start_uv;
     int mi_col_start_uv;
-    get_chroma_start_location(mbmi, tree_type, &mi_row_start_uv,
-                              &mi_col_start_uv);
+    av2_get_chroma_start_location(mbmi, tree_type, &mi_row_start_uv,
+                                  &mi_col_start_uv);
     const int mi_pu_start_y = vert_edge ? mi_row_start_uv : mi_col_start_uv;
     const int scale = vert_edge ? ss_y : ss_x;
     mi_pu_start = mi_pu_start_y >> scale;

av2/encoder/partition_search.c

@@ -2896,6 +2896,79 @@ static AVM_INLINE void init_allowed_partitions(
 #endif  // CONFIG_COLLECT_PARTITION_STATS
 }
 
+static AVM_INLINE bool is_same_block_for_tree(const MB_MODE_INFO *m1,
+                                              const MB_MODE_INFO *m2,
+                                              TREE_TYPE tree_type) {
+  if (!m1 || !m2) return false;
+
+  if (tree_type == CHROMA_PART) {
+    int m1_r;
+    int m1_c;
+    av2_get_chroma_start_location(m1, m1->tree_type, &m1_r, &m1_c);
+    int m2_r;
+    int m2_c;
+    av2_get_chroma_start_location(m2, m2->tree_type, &m2_r, &m2_c);
+    return m1_r == m2_r && m1_c == m2_c;
+  } else {
+    return m1->mi_row_start == m2->mi_row_start &&
+           m1->mi_col_start == m2->mi_col_start;
+  }
+}
+
+// Prunes uneven 4-way partitions by checking partition boundary alignment with
+// neighboring top/left blocks. Neighbor boundaries are detected by checking if
+// neighboring top/left MI units share the same block ID.
+static AVM_INLINE void prune_4way_partitions_with_neighbor_boundaries(
+    PartitionSearchState *state, const AV2_COMMON *cm, const MACROBLOCKD *xd,
+    int mi_row, int mi_col, BLOCK_SIZE bsize) {
+  // Check left neighbor for horizontal boundaries using mi array
+  const int mi_height = mi_size_high[bsize];
+  const int available_mi_height =
+      AVMMIN(mi_height, cm->mi_params.mi_rows - mi_row);
+  if (xd->left_available && mi_height >= 8 &&
+      available_mi_height == mi_height) {
+    bool left_horz_boundaries[MAX_MIB_SIZE] = { false };
+    for (int r = 1; r < available_mi_height; r++) {
+      const MB_MODE_INFO *m1 = xd->mi[r * xd->mi_stride - 1];
+      const MB_MODE_INFO *m2 = xd->mi[(r - 1) * xd->mi_stride - 1];
+      if (!is_same_block_for_tree(m1, m2, xd->tree_type)) {
+        left_horz_boundaries[r] = true;
+      }
+    }
+    // Prune HORZ 4A/4B partitions.
+    if (!left_horz_boundaries[mi_height / 8] &&
+        !left_horz_boundaries[3 * mi_height / 8] &&
+        !left_horz_boundaries[5 * mi_height / 8] &&
+        !left_horz_boundaries[7 * mi_height / 8]) {
+      state->prune_partition[PARTITION_HORZ_4A] = true;
+      state->prune_partition[PARTITION_HORZ_4B] = true;
+    }
+  }
+
+  // Check top neighbor for vertical boundaries using mi array
+  const int mi_width = mi_size_wide[bsize];
+  const int available_mi_width =
+      AVMMIN(mi_width, cm->mi_params.mi_cols - mi_col);
+  if (xd->up_available && mi_width >= 8 && available_mi_width == mi_width) {
+    bool top_vert_boundaries[MAX_MIB_SIZE] = { false };
+    for (int c = 1; c < available_mi_width; c++) {
+      const MB_MODE_INFO *m1 = xd->mi[-xd->mi_stride + c];
+      const MB_MODE_INFO *m2 = xd->mi[-xd->mi_stride + c - 1];
+      if (!is_same_block_for_tree(m1, m2, xd->tree_type)) {
+        top_vert_boundaries[c] = true;
+      }
+    }
+    // Prune VERT 4A/4B partitions.
+    if (!top_vert_boundaries[mi_width / 8] &&
+        !top_vert_boundaries[3 * mi_width / 8] &&
+        !top_vert_boundaries[5 * mi_width / 8] &&
+        !top_vert_boundaries[7 * mi_width / 8]) {
+      state->prune_partition[PARTITION_VERT_4A] = true;
+      state->prune_partition[PARTITION_VERT_4B] = true;
+    }
+  }
+}
+
 // Initialize state variables of partition search used in
 // av2_rd_pick_partition().
 static void init_partition_search_state_params(
@@ -5342,6 +5415,12 @@ bool av2_rd_pick_partition(
   av2_set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize,
                   &pc_tree->chroma_ref_info);
 
+  if (cpi->sf.part_sf.prune_part_with_neighbor_boundaries &&
+      !x->must_find_valid_partition && !frame_is_intra_only(cm)) {
+    prune_4way_partitions_with_neighbor_boundaries(&part_search_state, cm, xd,
+                                                   mi_row, mi_col, bsize);
+  }
+
   // Save rdmult before it might be changed, so it can be restored later.
   const int orig_rdmult = x->rdmult;
   setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, NO_AQ, NULL);

av2/encoder/speed_features.c

@@ -750,6 +750,7 @@ static AVM_INLINE void init_part_sf(PARTITION_SPEED_FEATURES *part_sf) {
   part_sf->prune_rect_with_split_depth = 0;
   part_sf->prune_part_h_with_partition_boundary = 0;
   part_sf->inter_sdp_fast_method_level = 0;
+  part_sf->prune_part_with_neighbor_boundaries = 0;
 #if CONFIG_ML_PART_SPLIT
   part_sf->prune_split_with_ml = 0;
   part_sf->prune_none_with_ml = 0;
@@ -983,6 +984,7 @@ static AVM_INLINE void set_erp_speed_features_framesize_dependent(
   const int is_1080p_or_larger = AVMMIN(cm->width, cm->height) >= 1080;
   const unsigned int erp_pruning_level = cpi->oxcf.part_cfg.erp_pruning_level;
   const int is_720p_or_lesser = AVMMIN(cm->width, cm->height) <= 720;
+  const int is_270p_or_lesser = AVMMIN(cm->width, cm->height) <= 270;
 
   switch (erp_pruning_level) {
     case 6: AVM_FALLTHROUGH_INTENDED;
@@ -1031,6 +1033,10 @@ static AVM_INLINE void set_erp_speed_features_framesize_dependent(
       sf->part_sf.remove_qp_restriction_with_ml = 1;
     }
 #endif  // CONFIG_ML_PART_SPLIT
+    if (is_270p_or_lesser) {
+      // For small resolutions, this speed feature has a large coding loss.
+      sf->part_sf.prune_part_with_neighbor_boundaries = 0;
+    }
   }
 
   if (cpi->speed >= 2) {
@@ -1083,6 +1089,7 @@ static AVM_INLINE void set_erp_speed_features(AV2_COMP *cpi) {
       sf->part_sf.ext_recur_depth_level = 2;
       sf->part_sf.simple_motion_search_split = 1;
       sf->part_sf.simple_motion_search_early_term_none = 1;
+      sf->part_sf.prune_part_with_neighbor_boundaries = 1;
       AVM_FALLTHROUGH_INTENDED;
     case 5:
       sf->part_sf.prune_part_h_with_partition_boundary = true;
@@ -1122,6 +1129,7 @@ static AVM_INLINE void set_erp_speed_features(AV2_COMP *cpi) {
     // Emulate erp_pruning_level = 6.
     sf->part_sf.ext_recur_depth_level = 1;
     sf->part_sf.ml_early_term_after_part_split_level = 2;
+    sf->part_sf.prune_part_with_neighbor_boundaries = 1;
   }
 
   if (cpi->speed >= 2) {

av2/encoder/speed_features.h

@@ -441,6 +441,8 @@ typedef struct PARTITION_SPEED_FEATURES {
   //    intra coded block, prunes when inter ratio exceeds 50%, and early skips
   //    when current best partitioning is PARTITION_NONE.
   int inter_sdp_fast_method_level;
+  // Prune partition types if they don't align with neighbor block boundaries.
+  int prune_part_with_neighbor_boundaries;
 #if CONFIG_ML_PART_SPLIT
   int prune_split_with_ml;
   int prune_split_ml_level;