CWG-E251 Encoder improvement for forced partitioning regarding SDP

This MR (CWG-E251) is an encoder side change only. In this MR, the forced partition type for partition search prune is now stored for each region type separately.
diff --git a/av1/encoder/block.h b/av1/encoder/block.h
index b4f58fc..f60068f 100644
--- a/av1/encoder/block.h
+++ b/av1/encoder/block.h
@@ -645,9 +645,9 @@
 #define BLOCK_8_COUNT 64
 #define BLOCK_4_COUNT 64
 
-#define MAKE_SM_DATA_BUF(width, height) \
-  SimpleMotionData                      \
-      b_##width##x##height[BLOCK_##width##_COUNT * BLOCK_##height##_COUNT]
+#define MAKE_SM_DATA_BUF(width, height, sdp_flag)                            \
+  SimpleMotionData b_##width##x##height##_##sdp_flag[BLOCK_##width##_COUNT * \
+                                                     BLOCK_##height##_COUNT]
 /*!\endcond */
 
 /*! \brief Simple motion data buffers
@@ -655,53 +655,102 @@
 typedef struct SimpleMotionDataBufs {
   /*!\cond */
   // Square blocks
-  MAKE_SM_DATA_BUF(256, 256);
-  MAKE_SM_DATA_BUF(128, 128);
-  MAKE_SM_DATA_BUF(64, 64);
-  MAKE_SM_DATA_BUF(32, 32);
-  MAKE_SM_DATA_BUF(16, 16);
-  MAKE_SM_DATA_BUF(8, 8);
-  MAKE_SM_DATA_BUF(4, 4);
+  MAKE_SM_DATA_BUF(256, 256, 0);
+  MAKE_SM_DATA_BUF(128, 128, 0);
+  MAKE_SM_DATA_BUF(64, 64, 0);
+  MAKE_SM_DATA_BUF(32, 32, 0);
+  MAKE_SM_DATA_BUF(16, 16, 0);
+  MAKE_SM_DATA_BUF(8, 8, 0);
+  MAKE_SM_DATA_BUF(4, 4, 0);
 
   // 1:2 blocks
-  MAKE_SM_DATA_BUF(128, 256);
-  MAKE_SM_DATA_BUF(64, 128);
-  MAKE_SM_DATA_BUF(32, 64);
-  MAKE_SM_DATA_BUF(16, 32);
-  MAKE_SM_DATA_BUF(8, 16);
-  MAKE_SM_DATA_BUF(4, 8);
+  MAKE_SM_DATA_BUF(128, 256, 0);
+  MAKE_SM_DATA_BUF(64, 128, 0);
+  MAKE_SM_DATA_BUF(32, 64, 0);
+  MAKE_SM_DATA_BUF(16, 32, 0);
+  MAKE_SM_DATA_BUF(8, 16, 0);
+  MAKE_SM_DATA_BUF(4, 8, 0);
 
   // 2:1 blocks
-  MAKE_SM_DATA_BUF(256, 128);
-  MAKE_SM_DATA_BUF(128, 64);
-  MAKE_SM_DATA_BUF(64, 32);
-  MAKE_SM_DATA_BUF(32, 16);
-  MAKE_SM_DATA_BUF(16, 8);
-  MAKE_SM_DATA_BUF(8, 4);
+  MAKE_SM_DATA_BUF(256, 128, 0);
+  MAKE_SM_DATA_BUF(128, 64, 0);
+  MAKE_SM_DATA_BUF(64, 32, 0);
+  MAKE_SM_DATA_BUF(32, 16, 0);
+  MAKE_SM_DATA_BUF(16, 8, 0);
+  MAKE_SM_DATA_BUF(8, 4, 0);
 
   // 1:4 blocks
-  MAKE_SM_DATA_BUF(16, 64);
-  MAKE_SM_DATA_BUF(8, 32);
-  MAKE_SM_DATA_BUF(4, 16);
+  MAKE_SM_DATA_BUF(16, 64, 0);
+  MAKE_SM_DATA_BUF(8, 32, 0);
+  MAKE_SM_DATA_BUF(4, 16, 0);
 
   // 4:1 blocks
-  MAKE_SM_DATA_BUF(64, 16);
-  MAKE_SM_DATA_BUF(32, 8);
-  MAKE_SM_DATA_BUF(16, 4);
+  MAKE_SM_DATA_BUF(64, 16, 0);
+  MAKE_SM_DATA_BUF(32, 8, 0);
+  MAKE_SM_DATA_BUF(16, 4, 0);
 
   // 1:8 blocks
-  MAKE_SM_DATA_BUF(8, 64);
-  MAKE_SM_DATA_BUF(4, 32);
+  MAKE_SM_DATA_BUF(8, 64, 0);
+  MAKE_SM_DATA_BUF(4, 32, 0);
 
   // 8:1 blocks
-  MAKE_SM_DATA_BUF(64, 8);
-  MAKE_SM_DATA_BUF(32, 4);
+  MAKE_SM_DATA_BUF(64, 8, 0);
+  MAKE_SM_DATA_BUF(32, 4, 0);
 
   // 1:16 blocks
-  MAKE_SM_DATA_BUF(4, 64);
+  MAKE_SM_DATA_BUF(4, 64, 0);
 
   // 16:1 blocks
-  MAKE_SM_DATA_BUF(64, 4);
+  MAKE_SM_DATA_BUF(64, 4, 0);
+
+  // Square blocks
+  MAKE_SM_DATA_BUF(256, 256, 1);
+  MAKE_SM_DATA_BUF(128, 128, 1);
+  MAKE_SM_DATA_BUF(64, 64, 1);
+  MAKE_SM_DATA_BUF(32, 32, 1);
+  MAKE_SM_DATA_BUF(16, 16, 1);
+  MAKE_SM_DATA_BUF(8, 8, 1);
+  MAKE_SM_DATA_BUF(4, 4, 1);
+
+  // 1:2 blocks
+  MAKE_SM_DATA_BUF(128, 256, 1);
+  MAKE_SM_DATA_BUF(64, 128, 1);
+  MAKE_SM_DATA_BUF(32, 64, 1);
+  MAKE_SM_DATA_BUF(16, 32, 1);
+  MAKE_SM_DATA_BUF(8, 16, 1);
+  MAKE_SM_DATA_BUF(4, 8, 1);
+
+  // 2:1 blocks
+  MAKE_SM_DATA_BUF(256, 128, 1);
+  MAKE_SM_DATA_BUF(128, 64, 1);
+  MAKE_SM_DATA_BUF(64, 32, 1);
+  MAKE_SM_DATA_BUF(32, 16, 1);
+  MAKE_SM_DATA_BUF(16, 8, 1);
+  MAKE_SM_DATA_BUF(8, 4, 1);
+
+  // 1:4 blocks
+  MAKE_SM_DATA_BUF(16, 64, 1);
+  MAKE_SM_DATA_BUF(8, 32, 1);
+  MAKE_SM_DATA_BUF(4, 16, 1);
+
+  // 4:1 blocks
+  MAKE_SM_DATA_BUF(64, 16, 1);
+  MAKE_SM_DATA_BUF(32, 8, 1);
+  MAKE_SM_DATA_BUF(16, 4, 1);
+
+  // 1:8 blocks
+  MAKE_SM_DATA_BUF(8, 64, 1);
+  MAKE_SM_DATA_BUF(4, 32, 1);
+
+  // 8:1 blocks
+  MAKE_SM_DATA_BUF(64, 8, 1);
+  MAKE_SM_DATA_BUF(32, 4, 1);
+
+  // 1:16 blocks
+  MAKE_SM_DATA_BUF(4, 64, 1);
+
+  // 16:1 blocks
+  MAKE_SM_DATA_BUF(64, 4, 1);
   /*!\endcond */
 } SimpleMotionDataBufs;
 
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 9d082d9..54dde9d 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -438,9 +438,10 @@
 #if CONFIG_EXT_RECUR_PARTITIONS
 static void fill_sms_buf(SimpleMotionDataBufs *data_buf,
                          SIMPLE_MOTION_DATA_TREE *sms_node, int mi_row,
-                         int mi_col, BLOCK_SIZE bsize, BLOCK_SIZE sb_size) {
-  SimpleMotionData *sms_data =
-      av1_get_sms_data_entry(data_buf, mi_row, mi_col, bsize, sb_size);
+                         int mi_col, BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
+                         int8_t sdp_flag) {
+  SimpleMotionData *sms_data = av1_get_sms_data_entry(data_buf, mi_row, mi_col,
+                                                      bsize, sb_size, sdp_flag);
   sms_data->old_sms = sms_node;
   if (bsize >= BLOCK_8X8) {
     const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
@@ -452,8 +453,8 @@
       const int sub_mi_row = mi_row + (r_idx >> 1) * h_mi / 2;
       SIMPLE_MOTION_DATA_TREE *sub_tree = sms_node->split[r_idx];
 
-      fill_sms_buf(data_buf, sub_tree, sub_mi_row, sub_mi_col, subsize,
-                   sb_size);
+      fill_sms_buf(data_buf, sub_tree, sub_mi_row, sub_mi_col, subsize, sb_size,
+                   sdp_flag);
     }
   }
 }
@@ -505,8 +506,10 @@
 #if CONFIG_EXT_RECUR_PARTITIONS
   SimpleMotionDataBufs *data_bufs = x->sms_bufs;
   av1_init_sms_data_bufs(data_bufs);
-  fill_sms_buf(data_bufs, sms_root, mi_row, mi_col, cm->sb_size, cm->sb_size);
-
+  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,
+               1);
 #endif  // CONFIG_EXT_RECUR_PARTITIONS
   if (x->e_mbd.tree_type == CHROMA_PART) {
     assert(is_bsize_square(x->sb_enc.min_partition_size));
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c
index 5eaf7b2..905300a 100644
--- a/av1/encoder/partition_search.c
+++ b/av1/encoder/partition_search.c
@@ -4230,7 +4230,8 @@
       && !is_inter_sdp_chroma(cm, cur_region_type, xd->tree_type)
 #endif  // CONFIG_EXTENDED_SDP
   ) {
-    return av1_get_prev_partition(x, mi_row, mi_col, bsize, cm->sb_size);
+    return av1_get_prev_partition(x, mi_row, mi_col, bsize, cm->sb_size,
+                                  (int8_t)cur_region_type);
   }
   return PARTITION_INVALID;
 }
@@ -5815,8 +5816,8 @@
   }
 #endif
 #if CONFIG_EXT_RECUR_PARTITIONS
-  SimpleMotionData *sms_data =
-      av1_get_sms_data_entry(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size);
+  SimpleMotionData *sms_data = av1_get_sms_data_entry(
+      x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size, (int8_t)cur_region_type);
   av1_set_best_mode_cache(x, sms_data->mode_cache);
 #endif  // CONFIG_EXT_RECUR_PARTITIONS
 
@@ -8750,7 +8751,8 @@
 #if CONFIG_EXT_RECUR_PARTITIONS
   {
     SimpleMotionData *sms_data =
-        av1_get_sms_data_entry(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size);
+        av1_get_sms_data_entry(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size,
+                               (int8_t)pc_tree->region_type);
     sms_tree = sms_data->old_sms;
   }
 #endif  // CONFIG_EXT_RECUR_PARTITIONS
@@ -9288,6 +9290,15 @@
         &level_banks,
 #endif  // CONFIG_MVP_IMPROVEMENT || WARP_CU_BANK
         multi_pass_mode, ext_recur_depth, parent_partition);
+
+    if (part_search_state.found_best_partition) {
+      av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size,
+                               pc_tree->partitioning,
+                               (int8_t)pc_tree->region_type);
+      av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size,
+                               pc_tree->partitioning,
+                               (int8_t)(1 - pc_tree->region_type));
+    }
   }
 #endif  // CONFIG_EXTENDED_SDP
 
@@ -9304,8 +9315,10 @@
     av1_invalid_rd_stats(&pc_tree->rd_cost);
   } else {
 #if CONFIG_EXT_RECUR_PARTITIONS
-    av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size,
-                             pc_tree->partitioning);
+    if (xd->tree_type != CHROMA_PART)
+      av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size,
+                               pc_tree->partitioning,
+                               (int8_t)pc_tree->region_type);
 #endif  // CONFIG_EXT_RECUR_PARTITIONS
   }
 
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 0db8632..4394912 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -1369,13 +1369,15 @@
 #if CONFIG_EXT_RECUR_PARTITIONS
     if (!*partition_none_allowed) {
       av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, bsize, cm->sb_size,
-                               PARTITION_HORZ);
+                               PARTITION_HORZ, (int8_t)pc_tree->region_type);
       const int mi_step = block_size_high[bsize] / 2;
       BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ);
       av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, subsize,
-                               cm->sb_size, PARTITION_VERT);
+                               cm->sb_size, PARTITION_VERT,
+                               (int8_t)pc_tree->region_type);
       av1_cache_best_partition(x->sms_bufs, mi_row + mi_step, mi_col, subsize,
-                               cm->sb_size, PARTITION_VERT);
+                               cm->sb_size, PARTITION_VERT,
+                               (int8_t)pc_tree->region_type);
     }
     (void)pc_tree;
 #endif  // CONFIG_EXT_RECUR_PARTITIONS
@@ -1709,72 +1711,129 @@
   return idx;
 }
 
-#define MAKE_SMS_ARR_SWITCH_CASE(width, height) \
-  case BLOCK_##width##X##height: {              \
-    return sms_bufs->b_##width##x##height;      \
+#define MAKE_SMS_ARR_SWITCH_CASE(width, height, sdp_flag) \
+  case BLOCK_##width##X##height: {                        \
+    return sms_bufs->b_##width##x##height##_##sdp_flag;   \
   }
 
 // Returns the buffer in SimpleMotionDataBufs that correspond to bsize.
 static INLINE SimpleMotionData *get_sms_arr(SimpleMotionDataBufs *sms_bufs,
-                                            BLOCK_SIZE bsize) {
-  switch (bsize) {
-    // Square blocks
-    MAKE_SMS_ARR_SWITCH_CASE(256, 256);
-    MAKE_SMS_ARR_SWITCH_CASE(128, 128);
-    MAKE_SMS_ARR_SWITCH_CASE(64, 64);
-    MAKE_SMS_ARR_SWITCH_CASE(32, 32);
-    MAKE_SMS_ARR_SWITCH_CASE(16, 16);
-    MAKE_SMS_ARR_SWITCH_CASE(8, 8);
-    MAKE_SMS_ARR_SWITCH_CASE(4, 4);
+                                            BLOCK_SIZE bsize,
+                                            int8_t region_type) {
+  if (region_type == 1) {
+    switch (bsize) {
+      // Square blocks
+      MAKE_SMS_ARR_SWITCH_CASE(256, 256, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(128, 128, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(64, 64, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 32, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 16, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 8, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 4, 1);
 
-    // 1:2 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(128, 256);
-    MAKE_SMS_ARR_SWITCH_CASE(64, 128);
-    MAKE_SMS_ARR_SWITCH_CASE(32, 64);
-    MAKE_SMS_ARR_SWITCH_CASE(16, 32);
-    MAKE_SMS_ARR_SWITCH_CASE(8, 16);
-    MAKE_SMS_ARR_SWITCH_CASE(4, 8);
+      // 1:2 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(128, 256, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(64, 128, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 64, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 32, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 16, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 8, 1);
 
-    // 2:1 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(256, 128);
-    MAKE_SMS_ARR_SWITCH_CASE(128, 64);
-    MAKE_SMS_ARR_SWITCH_CASE(64, 32);
-    MAKE_SMS_ARR_SWITCH_CASE(32, 16);
-    MAKE_SMS_ARR_SWITCH_CASE(16, 8);
-    MAKE_SMS_ARR_SWITCH_CASE(8, 4);
+      // 2:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(256, 128, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(128, 64, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(64, 32, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 16, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 8, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 4, 1);
 
-    // 1:4 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(16, 64);
-    MAKE_SMS_ARR_SWITCH_CASE(8, 32);
-    MAKE_SMS_ARR_SWITCH_CASE(4, 16);
+      // 1:4 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(16, 64, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 32, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 16, 1);
 
-    // 4:1 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(64, 16);
-    MAKE_SMS_ARR_SWITCH_CASE(32, 8);
-    MAKE_SMS_ARR_SWITCH_CASE(16, 4);
+      // 4:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(64, 16, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 8, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 4, 1);
 
-    // 1:8 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(8, 64);
-    MAKE_SMS_ARR_SWITCH_CASE(4, 32);
+      // 1:8 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(8, 64, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 32, 1);
 
-    // 8:1 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(64, 8);
-    MAKE_SMS_ARR_SWITCH_CASE(32, 4);
+      // 8:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(64, 8, 1);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 4, 1);
 
-    // 16:1 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(64, 4);
+      // 16:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(64, 4, 1);
 
-    // 1:16 blocks
-    MAKE_SMS_ARR_SWITCH_CASE(4, 64);
+      // 1:16 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(4, 64, 1);
 
-    default: assert(0 && "Invalid bsize"); return NULL;
+      default: assert(0 && "Invalid bsize"); return NULL;
+    }
+  } else {  // region_type = 0
+    switch (bsize) {
+      // Square blocks
+      MAKE_SMS_ARR_SWITCH_CASE(256, 256, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(128, 128, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(64, 64, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 32, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 16, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 8, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 4, 0);
+
+      // 1:2 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(128, 256, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(64, 128, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 64, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 32, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 16, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 8, 0);
+
+      // 2:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(256, 128, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(128, 64, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(64, 32, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 16, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 8, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 4, 0);
+
+      // 1:4 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(16, 64, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(8, 32, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 16, 0);
+
+      // 4:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(64, 16, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 8, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(16, 4, 0);
+
+      // 1:8 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(8, 64, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(4, 32, 0);
+
+      // 8:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(64, 8, 0);
+      MAKE_SMS_ARR_SWITCH_CASE(32, 4, 0);
+
+      // 16:1 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(64, 4, 0);
+
+      // 1:16 blocks
+      MAKE_SMS_ARR_SWITCH_CASE(4, 64, 0);
+
+      default: assert(0 && "Invalid bsize"); return NULL;
+    }
   }
 }
 #undef MAKE_SMS_ARR_SWITCH_CASE
 
-void av1_reset_prev_partition(SimpleMotionDataBufs *sms_bufs) {
+void av1_reset_prev_partition(SimpleMotionDataBufs *sms_bufs,
+                              int8_t region_type) {
   for (BLOCK_SIZE bsize = BLOCK_4X4; bsize < BLOCK_SIZES_ALL; bsize++) {
-    SimpleMotionData *sms_arr = get_sms_arr(sms_bufs, bsize);
+    SimpleMotionData *sms_arr = get_sms_arr(sms_bufs, bsize, region_type);
     const int mi_wide = mi_size_wide[bsize];
     const int mi_high = mi_size_high[bsize];
     const int sms_wide = get_sms_count_from_length(mi_wide);
@@ -1788,7 +1847,8 @@
 // Retrieves the SimpleMotionData from SimpleMotionDataBufs
 SimpleMotionData *av1_get_sms_data_entry(SimpleMotionDataBufs *sms_bufs,
                                          int mi_row, int mi_col,
-                                         BLOCK_SIZE bsize, BLOCK_SIZE sb_size) {
+                                         BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
+                                         int8_t region_type) {
   assert(mi_size_high[sb_size] == mi_size_wide[sb_size]);
   assert(bsize < BLOCK_SIZES_ALL);
   const int mi_in_sb = mi_size_high[sb_size];
@@ -1801,15 +1861,15 @@
   const int idx_col_in_sb = get_sms_arr_1d_idx(mi_wide, mi_col_in_sb);
   if (idx_col_in_sb == -1) return NULL;
   const int arr_stride = get_sms_count_from_length(mi_wide);
-  SimpleMotionData *sms_arr = get_sms_arr(sms_bufs, bsize);
+  SimpleMotionData *sms_arr = get_sms_arr(sms_bufs, bsize, region_type);
   return &sms_arr[idx_row_in_sb * arr_stride + idx_col_in_sb];
 }
 
 void av1_cache_best_partition(SimpleMotionDataBufs *sms_bufs, int mi_row,
                               int mi_col, BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
-                              PARTITION_TYPE partition) {
-  SimpleMotionData *cur_block =
-      av1_get_sms_data_entry(sms_bufs, mi_row, mi_col, bsize, sb_size);
+                              PARTITION_TYPE partition, int8_t region_type) {
+  SimpleMotionData *cur_block = av1_get_sms_data_entry(
+      sms_bufs, mi_row, mi_col, bsize, sb_size, region_type);
   cur_block->has_prev_partition = 1;
   cur_block->prev_partition = partition;
 }
@@ -1990,7 +2050,7 @@
     const int sub_col =
         mi_col + step_multiplier[partition][idx][1] * eighth_step_w / 4;
     SimpleMotionData *subblock = av1_get_sms_data_entry(
-        sms_bufs, sub_row, sub_col, subsizes[idx], sb_size);
+        sms_bufs, sub_row, sub_col, subsizes[idx], sb_size, 1);
     add_start_mv_to_block(subblock, start_mv);
   }
 }
@@ -2004,12 +2064,14 @@
                                    ,
                                    ThreadData *td, bool need_residual_stats
 #endif  // CONFIG_ML_PART_SPLIT
-) {
+                                   ,
+                                   int8_t region_type) {
+  assert(region_type == 1);
   const AV1_COMMON *const cm = &cpi->common;
   const BLOCK_SIZE sb_size = cm->sb_size;
   SimpleMotionDataBufs *sms_bufs = x->sms_bufs;
-  SimpleMotionData *cur_block =
-      av1_get_sms_data_entry(sms_bufs, mi_row, mi_col, bsize, sb_size);
+  SimpleMotionData *cur_block = av1_get_sms_data_entry(
+      sms_bufs, mi_row, mi_col, bsize, sb_size, region_type);
   if (!cur_block->valid
 #if CONFIG_ML_PART_SPLIT
       || (need_residual_stats && !cur_block->residual_stats_valid)
@@ -2031,10 +2093,11 @@
 }
 
 PARTITION_TYPE av1_get_prev_partition(MACROBLOCK *x, int mi_row, int mi_col,
-                                      BLOCK_SIZE bsize, BLOCK_SIZE sb_size) {
+                                      BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
+                                      int8_t region_type) {
   SimpleMotionDataBufs *sms_bufs = x->sms_bufs;
-  const SimpleMotionData *cur_block =
-      av1_get_sms_data_entry(sms_bufs, mi_row, mi_col, bsize, sb_size);
+  const SimpleMotionData *cur_block = av1_get_sms_data_entry(
+      sms_bufs, mi_row, mi_col, bsize, sb_size, region_type);
   if (cur_block && cur_block->has_prev_partition) {
     return cur_block->prev_partition;
   } else {
@@ -2072,7 +2135,8 @@
       ,
       NULL, false
 #endif  // CONFIG_ML_PART_SPLIT
-  );
+      ,
+      1);
 
   const BLOCK_SIZE h_size = get_partition_subsize(bsize, PARTITION_HORZ);
   const SimpleMotionData *blk_h1 =
@@ -2083,7 +2147,8 @@
                              ,
                              NULL, false
 #endif  // CONFIG_ML_PART_SPLIT
-                             )
+                             ,
+                             1)
           : NULL;
   const SimpleMotionData *blk_h2 =
       h_size != BLOCK_INVALID
@@ -2093,7 +2158,8 @@
                              ,
                              NULL, false
 #endif  // CONFIG_ML_PART_SPLIT
-                             )
+                             ,
+                             1)
           : NULL;
 
   const BLOCK_SIZE v_size = get_partition_subsize(bsize, PARTITION_VERT);
@@ -2105,7 +2171,8 @@
                              ,
                              NULL, false
 #endif  // CONFIG_ML_PART_SPLIT
-                             )
+                             ,
+                             1)
           : NULL;
   const SimpleMotionData *blk_v2 =
       v_size != BLOCK_INVALID
@@ -2115,7 +2182,8 @@
                              ,
                              NULL, false
 #endif  // CONFIG_ML_PART_SPLIT
-                             )
+                             ,
+                             1)
           : NULL;
 
   // Results of SMS on the subblocks
@@ -2646,7 +2714,7 @@
   if (cpi->common.current_frame.frame_type != INTER_FRAME) return;
 
   SimpleMotionData *blk_none =
-      av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, bsize, td, true);
+      av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, bsize, td, true, 1);
 
   BLOCK_SIZE subsize_sq = get_partition_subsize(
       get_partition_subsize(bsize, PARTITION_HORZ), PARTITION_VERT);
@@ -2658,15 +2726,16 @@
   if (subsize_sq != BLOCK_INVALID) {
     int w_sub_mi = mi_size_wide[subsize_sq];
     int h_sub_mi = mi_size_high[subsize_sq];
-    SimpleMotionData *blk_sq_0 = av1_get_sms_data(cpi, tile_info, x, mi_row,
-                                                  mi_col, subsize_sq, td, true);
+
+    SimpleMotionData *blk_sq_0 = av1_get_sms_data(
+        cpi, tile_info, x, mi_row, mi_col, subsize_sq, td, true, 1);
     SimpleMotionData *blk_sq_1 = av1_get_sms_data(
-        cpi, tile_info, x, mi_row, mi_col + w_sub_mi, subsize_sq, td, true);
+        cpi, tile_info, x, mi_row, mi_col + w_sub_mi, subsize_sq, td, true, 1);
     SimpleMotionData *blk_sq_2 = av1_get_sms_data(
-        cpi, tile_info, x, mi_row + h_sub_mi, mi_col, subsize_sq, td, true);
+        cpi, tile_info, x, mi_row + h_sub_mi, mi_col, subsize_sq, td, true, 1);
     SimpleMotionData *blk_sq_3 =
         av1_get_sms_data(cpi, tile_info, x, mi_row + h_sub_mi,
-                         mi_col + w_sub_mi, subsize_sq, td, true);
+                         mi_col + w_sub_mi, subsize_sq, td, true, 1);
 
     if (out_features) {
       int blk_area = block_size_wide[bsize] * block_size_high[bsize];
diff --git a/av1/encoder/partition_strategy.h b/av1/encoder/partition_strategy.h
index 2c6b880..80342be 100644
--- a/av1/encoder/partition_strategy.h
+++ b/av1/encoder/partition_strategy.h
@@ -193,7 +193,8 @@
 #if CONFIG_EXT_RECUR_PARTITIONS
 SimpleMotionData *av1_get_sms_data_entry(SimpleMotionDataBufs *sms_bufs,
                                          int mi_row, int mi_col,
-                                         BLOCK_SIZE bsize, BLOCK_SIZE sb_size);
+                                         BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
+                                         int8_t region_type);
 SimpleMotionData *av1_get_sms_data(AV1_COMP *const cpi,
                                    const TileInfo *const tile, MACROBLOCK *x,
                                    int mi_row, int mi_col, BLOCK_SIZE bsize
@@ -202,8 +203,10 @@
                                    ,
                                    ThreadData *td, bool need_residual_stats
 #endif  // CONFIG_ML_PART_SPLIT
-);
-void av1_reset_prev_partition(SimpleMotionDataBufs *sms_bufs);
+                                   ,
+                                   int8_t region_type);
+void av1_reset_prev_partition(SimpleMotionDataBufs *sms_bufs,
+                              int8_t region_type);
 
 static AOM_INLINE void av1_add_mode_search_context_to_cache(
     SimpleMotionData *sms_data, PICK_MODE_CONTEXT *ctx) {
@@ -236,7 +239,7 @@
 
 void av1_cache_best_partition(SimpleMotionDataBufs *sms_bufs, int mi_row,
                               int mi_col, BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
-                              PARTITION_TYPE partition);
+                              PARTITION_TYPE partition, int8_t region_type);
 #endif  // CONFIG_EXT_RECUR_PARTITIONS
 
 // A simplified version of set_offsets meant to be used for
@@ -303,7 +306,8 @@
 }
 
 PARTITION_TYPE av1_get_prev_partition(MACROBLOCK *x, int mi_row, int mi_col,
-                                      BLOCK_SIZE bsize, BLOCK_SIZE sb_size);
+                                      BLOCK_SIZE bsize, BLOCK_SIZE sb_size,
+                                      int8_t region_type);
 
 #if CONFIG_EXT_RECUR_PARTITIONS
 static INLINE void av1_init_sms_data_bufs(SimpleMotionDataBufs *data_bufs) {
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index 7f40c14..505ed21 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -10576,8 +10576,9 @@
       BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ);
       if (subsize != BLOCK_INVALID) {
         for (int r = 0; r <= mi_size_high[bsize] / 2; ++r) {
-          const PARTITION_TYPE prev_part = av1_get_prev_partition(
-              x, xd->mi_row + r, xd->mi_col, subsize, cm->sb_size);
+          const PARTITION_TYPE prev_part =
+              av1_get_prev_partition(x, xd->mi_row + r, xd->mi_col, subsize,
+                                     cm->sb_size, (int8_t)mbmi->region_type);
           if (prev_part != PARTITION_INVALID) {
             prune_ref_frames = true;
             break;
@@ -10588,8 +10589,9 @@
       subsize = get_partition_subsize(bsize, PARTITION_VERT);
       if (subsize != BLOCK_INVALID) {
         for (int c = 0; c <= mi_size_wide[bsize] / 2; ++c) {
-          const PARTITION_TYPE prev_part = av1_get_prev_partition(
-              x, xd->mi_row, xd->mi_col + c, subsize, cm->sb_size);
+          const PARTITION_TYPE prev_part =
+              av1_get_prev_partition(x, xd->mi_row, xd->mi_col + c, subsize,
+                                     cm->sb_size, (int8_t)mbmi->region_type);
           if (prev_part != PARTITION_INVALID) {
             prune_ref_frames = true;
             break;
@@ -10615,7 +10617,8 @@
         for (int sub_idx = 0; sub_idx < 2; sub_idx++) {
           const PARTITION_TYPE prev_part = av1_get_prev_partition(
               x, mi_pos_rect[rect_type][sub_idx][0],
-              mi_pos_rect[rect_type][sub_idx][1], subsize, cm->sb_size);
+              mi_pos_rect[rect_type][sub_idx][1], subsize, cm->sb_size,
+              (int8_t)mbmi->region_type);
           if (prev_part != PARTITION_INVALID) {
             prune_ref_frames = true;
             break;