Add superblock source SAD speed feature for RTC partitioning

av2/encoder/block.h

@@ -1210,6 +1210,13 @@ typedef struct {
 } CoeffCosts;
 
 /*!\cond */
+typedef enum {
+  kZeroSad = 0,
+  kVeryLowSad = 1,
+  kLowSad = 2,
+  kMedSad = 3,
+  kHighSad = 4
+} SOURCE_SAD;
 #define SINGLE_REF_MODES ((REF_FRAMES - 1) * 4)
 /*!\endcond */
 struct inter_modes_info;
@@ -1537,6 +1544,10 @@ typedef struct macroblock {
   /**@{*/
   //! Variance of the source frame.
   unsigned int source_variance;
+  //! Superblock source SAD.
+  unsigned int source_sad;
+  //! Superblock source SAD level.
+  SOURCE_SAD source_sad_level;
   //! SSE of the current predictor.
   unsigned int pred_sse[SINGLE_REF_FRAMES];
   /*! Simple motion search buffers. */

av2/encoder/encodeframe.c

@@ -179,6 +179,75 @@ static BLOCK_SIZE get_rd_var_based_fixed_partition(AV2_COMP *cpi, MACROBLOCK *x,
     return BLOCK_8X8;
 }
 
+static unsigned int get_sb_source_sad(const AV2_COMP *cpi, const MACROBLOCK *x,
+                                      int mi_row, int mi_col) {
+  if (cpi->last_source == NULL || cpi->source == NULL) return UINT_MAX;
+  if (frame_is_intra_only(&cpi->common)) return UINT_MAX;
+  if (cpi->last_source->y_width != cpi->source->y_width ||
+      cpi->last_source->y_height != cpi->source->y_height)
+    return UINT_MAX;
+
+  const BLOCK_SIZE sb_size = cpi->common.seq_params.sb_size;
+  const int block_width = mi_size_wide[sb_size];
+  const int block_height = mi_size_high[sb_size];
+  const uint16_t *src_y = x->plane[AVM_PLANE_Y].src.buf;
+  const int src_stride = x->plane[AVM_PLANE_Y].src.stride;
+  const int last_src_offset =
+      (mi_row * MI_SIZE) * cpi->last_source->y_stride + (mi_col * MI_SIZE);
+  const uint16_t *last_src_y = cpi->last_source->y_buffer + last_src_offset;
+  const int last_src_stride = cpi->last_source->y_stride;
+
+  if (mi_row + block_height <= cpi->common.mi_params.mi_rows &&
+      mi_col + block_width <= cpi->common.mi_params.mi_cols) {
+    return cpi->fn_ptr[sb_size].sdf(src_y, src_stride, last_src_y,
+                                    last_src_stride);
+  }
+
+  // Border superblock: compute SAD over the valid pixels within the frame
+  // and normalize to the full superblock area.
+  const int sb_pix_w = block_width * MI_SIZE;
+  const int sb_pix_h = block_height * MI_SIZE;
+  const int valid_w =
+      AVMMIN(sb_pix_w, (int)cpi->common.width - mi_col * MI_SIZE);
+  const int valid_h =
+      AVMMIN(sb_pix_h, (int)cpi->common.height - mi_row * MI_SIZE);
+  if (valid_w <= 0 || valid_h <= 0) return UINT_MAX;
+
+  uint64_t sum_sad = 0;
+  for (int r = 0; r < valid_h; ++r) {
+    const uint16_t *s = src_y + r * src_stride;
+    const uint16_t *ls = last_src_y + r * last_src_stride;
+    for (int c = 0; c < valid_w; ++c) {
+      sum_sad += abs((int)s[c] - (int)ls[c]);
+    }
+  }
+
+  const uint64_t total_sb_pix = (uint64_t)sb_pix_w * sb_pix_h;
+  const uint64_t valid_pix = (uint64_t)valid_w * valid_h;
+  uint64_t norm_sad = (sum_sad * total_sb_pix + (valid_pix >> 1)) / valid_pix;
+  if (cpi->common.seq_params.bit_depth > 8) {
+    norm_sad >>= (cpi->common.seq_params.bit_depth - 8);
+  }
+  return (unsigned int)norm_sad;
+}
+
+static SOURCE_SAD get_source_sad_level(unsigned int sb_source_sad,
+                                       BLOCK_SIZE sb_size) {
+  if (sb_source_sad == UINT_MAX) return kMedSad;
+  if (sb_source_sad == 0) return kZeroSad;
+
+  const int num_64x64 = (sb_size == BLOCK_256X256)   ? 16
+                        : (sb_size == BLOCK_128X128) ? 4
+                                                     : 1;
+  const unsigned int avg_64x64_sad =
+      (sb_source_sad + (num_64x64 >> 1)) / num_64x64;
+
+  if (avg_64x64_sad < 5000) return kVeryLowSad;
+  if (avg_64x64_sad < 20000) return kLowSad;
+  if (avg_64x64_sad > 50000) return kHighSad;
+  return kMedSad;
+}
+
 void av2_setup_src_planes(MACROBLOCK *x, const YV12_BUFFER_CONFIG *src,
                           int mi_row, int mi_col, const int num_planes,
                           const CHROMA_REF_INFO *chroma_ref_info) {
@@ -840,7 +909,14 @@ static AVM_INLINE void encode_rd_sb(AV2_COMP *cpi, ThreadData *td,
     xd->tree_type = SHARED_PART;
   } else if (sf->part_sf.partition_search_type == VAR_BASED_PARTITION) {
     av2_set_offsets(cpi, tile_info, x, mi_row, mi_col, sb_size, NULL);
-    av2_choose_var_based_partitioning(cpi, tile_info, td, x, mi_row, mi_col);
+    unsigned int source_sad = UINT_MAX;
+    if (cpi->sf.rt_sf.source_metrics_sb) {
+      source_sad = get_sb_source_sad(cpi, x, mi_row, mi_col);
+    }
+    x->source_sad = source_sad;
+    x->source_sad_level = get_source_sad_level(source_sad, sb_size);
+    av2_choose_var_based_partitioning(cpi, tile_info, td, x, mi_row, mi_col,
+                                      source_sad);
     for (int loop_idx = 0; loop_idx < total_loop_num; loop_idx++) {
       xd->tree_type =
           (total_loop_num == 1 ? SHARED_PART
@@ -1371,6 +1447,8 @@ static AVM_INLINE void encode_sb_row(AV2_COMP *cpi, ThreadData *td,
 
     xd->cur_frame_force_integer_mv = cm->features.cur_frame_force_integer_mv;
     x->source_variance = UINT_MAX;
+    x->source_sad = UINT_MAX;
+    x->source_sad_level = kMedSad;
     td->mb.cb_coef_buff = av2_get_cb_coeff_buffer(cpi, mi_row, mi_col);
 
     av2_reset_refmv_bank(cm, xd, tile_info, mi_row, mi_col);

av2/encoder/speed_features.c

@@ -766,6 +766,7 @@ static void set_rt_speed_features_framesize_independent(
     sf->rd_sf.tx_domain_dist_thres_level = 2;
     sf->rt_sf.use_nonrd_partition = 1;
     sf->rt_sf.use_only_dc_intra_interframe = true;
+    sf->rt_sf.source_metrics_sb = 1;
     sf->winner_mode_sf.tx_size_search_level = USE_FAST_RD;
     sf->tx_sf.restrict_tx_partition_type_search = 3;
     sf->tx_sf.enable_tx_partition = true;
@@ -1098,6 +1099,12 @@ static AVM_INLINE void init_lc_sf(LC_DEC_SPEED_FEATURES *lc_sf) {
   lc_sf->skip_loop_filter_based_on_error = 0;
 }
 
+static AVM_INLINE void init_rt_sf(REALTIME_SPEED_FEATURES *rt_sf) {
+  rt_sf->use_nonrd_partition = 0;
+  rt_sf->use_only_dc_intra_interframe = false;
+  rt_sf->source_metrics_sb = 0;
+}
+
 static AVM_INLINE void set_erp_speed_features_framesize_dependent(
     AV2_COMP *cpi) {
   SPEED_FEATURES *const sf = &cpi->sf;
@@ -1317,6 +1324,7 @@ void av2_set_speed_features_framesize_independent(AV2_COMP *cpi, int speed) {
   init_lpf_sf(&sf->lpf_sf);
   init_flexmv_sf(&sf->flexmv_sf);
   init_lc_sf(&sf->lc_sf);
+  init_rt_sf(&sf->rt_sf);
 
   if (oxcf->mode == GOOD) {
     set_good_speed_features_framesize_independent(cpi, sf, speed);

av2/encoder/speed_features.h

@@ -1110,6 +1110,9 @@ typedef struct REALTIME_SPEED_FEATURES {
 
   // Flag to disable all but DC intra mode for inter frame prediction.
   bool use_only_dc_intra_interframe;
+
+  // Compute source sad metrics for superblock.
+  int source_metrics_sb;
 } REALTIME_SPEED_FEATURES;
 
 typedef struct LC_DEC_SPEED_FEATURES {

av2/encoder/var_based_part.c

@@ -510,26 +510,50 @@ static void set_vbp_thresholds_key_frame(int64_t thresholds[],
  * threshold
  */
 static inline void set_vbp_thresholds(AV2_COMP *cpi, int64_t thresholds[],
-                                      int qindex) {
+                                      int qindex, SOURCE_SAD source_sad_level) {
   AV2_COMMON *const cm = &cpi->common;
   const int is_key_frame = frame_is_intra_only(cm);
   const int threshold_multiplier = is_key_frame ? 120 : 1;
   // In AV2, qindex is linear in dB (exponential step size 2^(q/24)), whereas
-  // AV1 had a linear quantizer table. We compute threshold_base directly as a
-  // function of qindex to account for the exponential range in AV2 vs linear in
-  // AV1.
+  // AV1 had a linear quantizer table. We compute threshold_base as a function
+  // of qindex and source_sad_level to account for the exponential range in AV2.
+  const int num_pixels = cm->width * cm->height;
+  const int32_t ac_q = av2_ac_quant_QTX(qindex, 0, 0, cm->seq_params.bit_depth);
+  const int64_t q_offset = qindex - 65;
+  const int64_t poly_base =
+      (qindex <= 65) ? 12 : (12 + ((q_offset * (qindex + 15)) / 40));
+
+  int64_t quant_base;
+  if (num_pixels <= RESOLUTION_240P) {
+    if (source_sad_level == kZeroSad) {
+      quant_base = ac_q >> 2;
+    } else {
+      quant_base = poly_base;
+    }
+  } else if (num_pixels <= RESOLUTION_288P) {
+    if (source_sad_level <= kVeryLowSad) {
+      quant_base = ac_q >> 2;
+    } else if (source_sad_level <= kLowSad) {
+      quant_base = (3 * ac_q) >> 4;
+    } else {
+      quant_base = poly_base;
+    }
+  } else {
+    quant_base = ac_q >> 2;
+  }
+
   int64_t threshold_base;
-  if (qindex <= 65) {
-    threshold_base = 12;
+  if (qindex <= 80) {
+    threshold_base = poly_base;
+  } else if (qindex >= 120) {
+    threshold_base = quant_base;
   } else {
-    const int64_t q_offset = qindex - 65;
-    threshold_base = 12 + ((q_offset * (qindex + 15)) / 40);
+    threshold_base =
+        interpolate_threshold(poly_base, quant_base, qindex, 80, 120);
   }
   threshold_base *= threshold_multiplier;
 
-  const int current_qindex = cm->quant_params.base_qindex;
   const int threshold_left_shift = 7;
-  const int num_pixels = cm->width * cm->height;
 
   if (is_key_frame) {
     set_vbp_thresholds_key_frame(thresholds, threshold_base, num_pixels);
@@ -544,8 +568,57 @@ static inline void set_vbp_thresholds(AV2_COMP *cpi, int64_t thresholds[],
   thresholds[4] = threshold_base << threshold_left_shift;
   thresholds[5] = INT64_MAX;
 
-  tune_thresh_based_on_resolution(cpi, thresholds, threshold_base,
-                                  current_qindex, num_pixels);
+  tune_thresh_based_on_resolution(cpi, thresholds, threshold_base, qindex,
+                                  num_pixels);
+
+  const bool is_screen = (cpi->oxcf.tune_cfg.content == AVM_CONTENT_SCREEN);
+  if (is_screen) {
+    if (source_sad_level == kZeroSad) {
+      thresholds[1] = (2 * thresholds[1]);
+      thresholds[2] = (2 * thresholds[2]);
+      thresholds[3] = (2 * thresholds[3]);
+      thresholds[4] = INT64_MAX;
+    } else if (source_sad_level <= kVeryLowSad) {
+      thresholds[1] = (3 * thresholds[1]) >> 1;
+      thresholds[2] = (3 * thresholds[2]) >> 1;
+      thresholds[3] = (3 * thresholds[3]) >> 1;
+    } else if (source_sad_level <= kLowSad) {
+      thresholds[1] = (5 * thresholds[1]) >> 2;
+      thresholds[2] = (5 * thresholds[2]) >> 2;
+      thresholds[3] = (5 * thresholds[3]) >> 2;
+    }
+  } else {
+    if (num_pixels <= RESOLUTION_288P) {
+      if (source_sad_level == kZeroSad) {
+        thresholds[2] = thresholds[2] << 2;
+        thresholds[3] = thresholds[3] << 1;
+        thresholds[4] = INT64_MAX;
+      } else if (source_sad_level <= kVeryLowSad && qindex >= 100) {
+        thresholds[2] = (3 * thresholds[2]) >> 1;
+        thresholds[3] = (5 * thresholds[3]) >> 2;
+      } else if (source_sad_level <= kLowSad && qindex >= 110) {
+        thresholds[2] = (5 * thresholds[2]) >> 2;
+        thresholds[3] = (9 * thresholds[3]) >> 3;
+      }
+    } else {
+      if (source_sad_level == kZeroSad && qindex >= 100) {
+        thresholds[1] = (7 * thresholds[1]) >> 2;
+        if (num_pixels > RESOLUTION_360P) {
+          thresholds[2] = (3 * thresholds[2]) >> 1;
+        }
+      } else if (source_sad_level <= kVeryLowSad && qindex >= 100) {
+        thresholds[1] = (3 * thresholds[1]) >> 1;
+        if (num_pixels > RESOLUTION_360P) {
+          thresholds[2] = (5 * thresholds[2]) >> 2;
+        }
+      } else if (source_sad_level <= kLowSad && qindex >= 100) {
+        thresholds[1] = (5 * thresholds[1]) >> 2;
+        if (num_pixels > RESOLUTION_360P) {
+          thresholds[2] = (9 * thresholds[2]) >> 3;
+        }
+      }
+    }
+  }
 }
 
 static inline void force_split_ancestors(PART_EVAL_STATUS *force_split,
@@ -792,7 +865,8 @@ static void set_vt_partitioning_64x64(AV2_COMP *cpi, MACROBLOCKD *xd,
 void av2_choose_var_based_partitioning(AV2_COMP *cpi,
                                        const TileInfo *const tile,
                                        ThreadData *td, MACROBLOCK *x,
-                                       int mi_row, int mi_col) {
+                                       int mi_row, int mi_col,
+                                       unsigned int source_sad) {
   AV2_COMMON *const cm = &cpi->common;
   MACROBLOCKD *xd = &x->e_mbd;
   // Flat array representing quadtree nodes up to 16x16 level for 256x256 SB
@@ -813,6 +887,8 @@ void av2_choose_var_based_partitioning(AV2_COMP *cpi,
   bool is_key_frame = frame_is_intra_only(cm);
   bool scaled_ref_last = false;
   const int is_360p_or_smaller = cm->width * cm->height <= RESOLUTION_360P;
+  const SOURCE_SAD source_sad_level = x->source_sad_level;
+  const bool is_screen = (cpi->oxcf.tune_cfg.content == AVM_CONTENT_SCREEN);
 
   assert(cm->seq_params.sb_size == BLOCK_64X64 ||
          cm->seq_params.sb_size == BLOCK_128X128 ||
@@ -827,7 +903,7 @@ void av2_choose_var_based_partitioning(AV2_COMP *cpi,
   int64_t thresholds[6];
   const int qindex = cm->quant_params.base_qindex;
 
-  set_vbp_thresholds(cpi, thresholds, qindex);
+  set_vbp_thresholds(cpi, thresholds, qindex, source_sad_level);
 
   src_buf = x->plane[AVM_PLANE_Y].src.buf;
   int src_stride = x->plane[AVM_PLANE_Y].src.stride;
@@ -844,6 +920,18 @@ void av2_choose_var_based_partitioning(AV2_COMP *cpi,
       dst_buf = xd->plane[AVM_PLANE_Y].pre[0].buf;
       dst_stride = xd->plane[AVM_PLANE_Y].pre[0].stride;
     }
+
+    const int block_width = mi_size_wide[cm->seq_params.sb_size];
+    const int block_height = mi_size_high[cm->seq_params.sb_size];
+    if (source_sad == 0 && mi_col + block_width <= tile->mi_col_end &&
+        mi_row + block_height <= tile->mi_row_end) {
+      if (y_sad == 0 ||
+          (is_screen &&
+           y_sad < (unsigned int)(block_width * block_height * 2))) {
+        set_block_size(cpi, mi_row, mi_col, cm->seq_params.sb_size);
+        return;
+      }
+    }
   } else {
     dst_buf = NULL;
     dst_stride = 0;
@@ -888,7 +976,6 @@ void av2_choose_var_based_partitioning(AV2_COMP *cpi,
   const int offset_32x32 = offsets.offset_32x32;
   const int offset_16x16 = offsets.offset_16x16;
   const bool is_lowres = (cm->width * cm->height <= RESOLUTION_288P);
-  const bool is_screen = (cpi->oxcf.tune_cfg.content == AVM_CONTENT_SCREEN);
 
   avg_64x64 = 0;
   for (int blk64_idx = 0; blk64_idx < num_64x64_blocks; ++blk64_idx) {

av2/encoder/var_based_part.h

@@ -99,7 +99,8 @@ extern "C" {
 void av2_choose_var_based_partitioning(AV2_COMP *cpi,
                                        const TileInfo *const tile,
                                        ThreadData *td, MACROBLOCK *x,
-                                       int mi_row, int mi_col);
+                                       int mi_row, int mi_col,
+                                       unsigned int source_sad);
 
 #ifdef __cplusplus
 }  // extern "C"