Recursive partition syntax coding

Enable recursive partition information coding from SB64X64 down to
MB16X16. The bit-stream syntax is now supporting rectangular block
sizes. It starts from SB64X64 and recursively describes the partition
type of the current block. If the partition type is PARTITION_NONE,
the block is coded as a single unit; if it is PARTITION_HORZ or
PARTITION_VERT, the block is segmented into two independently coded
rectangular units, with no further partition needed; otherwise, the
block is segmented into 4 square blocks. i.e., PARTITION_SPLIT case,
each can be potentially further partitioned.

Forward adaptive probability modeling is used for the partition
information coding, conditioned on the current block size.

Change-Id: I499365fb547839d555498e3bcc0387d8a3587d87

5 files changed, 58 insertions, 2 deletions

diff --git a/vp9/common/vp9_blockd.h b/vp9/common/vp9_blockd.h
index 95bfcd36c..e034b2ddc 100644
--- a/vp9/common/vp9_blockd.h
+++ b/vp9/common/vp9_blockd.h
@@ -240,6 +240,11 @@ static INLINE int b_height_log2(BLOCK_SIZE_TYPE sb_type) {
   return mb_height_log2(sb_type) + 2;
 }
 
+static INLINE int partition_plane(BLOCK_SIZE_TYPE sb_type) {
+  assert(mb_width_log2(sb_type) == mb_height_log2(sb_type));
+  return (mb_width_log2(sb_type) - 1);
+}
+
 typedef enum {
   BLOCK_4X4_LG2 = 0,
   BLOCK_8X8_LG2 = 2,
diff --git a/vp9/common/vp9_entropymode.c b/vp9/common/vp9_entropymode.c
index ed3c6a554..f4182443a 100644
--- a/vp9/common/vp9_entropymode.c
+++ b/vp9/common/vp9_entropymode.c
@@ -151,6 +151,17 @@ const int vp9_mbsplit_count [VP9_NUMMBSPLITS] = { 2, 2, 4, 16};
 
 const vp9_prob vp9_mbsplit_probs [VP9_NUMMBSPLITS - 1] = { 110, 111, 150};
 
+#if CONFIG_SBSEGMENT
+const vp9_prob vp9_partition_probs[PARTITION_PLANES][PARTITION_TYPES - 1] = {
+  {110, 111, 150},
+  {110, 111, 150},
+};
+#else
+const vp9_prob vp9_partition_probs[PARTITION_PLANES][PARTITION_TYPES - 1] = {
+  {200}, {200},
+};
+#endif
+
 /* Array indices are identical to previously-existing INTRAMODECONTEXTNODES. */
 
 const vp9_tree_index vp9_kf_bmode_tree[VP9_KF_BINTRAMODES * 2 - 2] = {
@@ -283,6 +294,18 @@ const vp9_tree_index vp9_sub_mv_ref_tree[6] = {
   -ZERO4X4, -NEW4X4
 };
 
+#if CONFIG_SBSEGMENT
+const vp9_tree_index vp9_partition_tree[6] = {
+  -PARTITION_NONE, 2,
+  -PARTITION_HORZ, 4,
+  -PARTITION_VERT, -PARTITION_SPLIT
+};
+#else
+const vp9_tree_index vp9_partition_tree[2] = {
+  -PARTITION_NONE, -PARTITION_SPLIT
+};
+#endif
+
 struct vp9_token vp9_bmode_encodings[VP9_NKF_BINTRAMODES];
 struct vp9_token vp9_kf_bmode_encodings[VP9_KF_BINTRAMODES];
 struct vp9_token vp9_ymode_encodings[VP9_YMODES];
@@ -297,6 +320,8 @@ struct vp9_token vp9_mv_ref_encoding_array[VP9_MVREFS];
 struct vp9_token vp9_sb_mv_ref_encoding_array[VP9_MVREFS];
 struct vp9_token vp9_sub_mv_ref_encoding_array[VP9_SUBMVREFS];
 
+struct vp9_token vp9_partition_encodings[PARTITION_TYPES];
+
 void vp9_init_mbmode_probs(VP9_COMMON *x) {
   unsigned int bct [VP9_YMODES] [2];      /* num Ymodes > num UV modes */
 
@@ -332,6 +357,10 @@ void vp9_init_mbmode_probs(VP9_COMMON *x) {
   vpx_memcpy(x->fc.mbsplit_prob, vp9_mbsplit_probs, sizeof(vp9_mbsplit_probs));
   vpx_memcpy(x->fc.switchable_interp_prob, vp9_switchable_interp_prob,
              sizeof(vp9_switchable_interp_prob));
+
+  vpx_memcpy(x->fc.partition_prob, vp9_partition_probs,
+             sizeof(vp9_partition_probs));
+
 #if CONFIG_COMP_INTERINTRA_PRED
   x->fc.interintra_prob = VP9_DEF_INTERINTRA_PROB;
 #endif
@@ -433,6 +462,7 @@ void vp9_entropy_mode_init() {
   vp9_tokens_from_tree(vp9_mbsplit_encodings, vp9_mbsplit_tree);
   vp9_tokens_from_tree(vp9_switchable_interp_encodings,
                        vp9_switchable_interp_tree);
+  vp9_tokens_from_tree(vp9_partition_encodings, vp9_partition_tree);
 
   vp9_tokens_from_tree_offset(vp9_mv_ref_encoding_array,
                               vp9_mv_ref_tree, NEARESTMV);
@@ -631,6 +661,10 @@ void vp9_adapt_mode_probs(VP9_COMMON *cm) {
                                            interintra_prob, factor);
   }
 #endif
+  for (i = 0; i < PARTITION_PLANES; i++)
+    update_mode_probs(PARTITION_TYPES, vp9_partition_tree,
+                      cm->fc.partition_counts[i], cm->fc.pre_partition_prob[i],
+                      cm->fc.partition_prob[i], 0);
 }
 
 static void set_default_lf_deltas(MACROBLOCKD *xd) {
diff --git a/vp9/common/vp9_entropymode.h b/vp9/common/vp9_entropymode.h
index fe97f0e1f..665569578 100644
--- a/vp9/common/vp9_entropymode.h
+++ b/vp9/common/vp9_entropymode.h
@@ -70,6 +70,10 @@ extern struct vp9_token vp9_mv_ref_encoding_array[VP9_MVREFS];
 extern struct vp9_token vp9_sb_mv_ref_encoding_array[VP9_MVREFS];
 extern struct vp9_token vp9_sub_mv_ref_encoding_array[VP9_SUBMVREFS];
 
+// probability models for partition information
+extern const vp9_tree_index  vp9_partition_tree[];
+extern struct vp9_token vp9_partition_encodings[PARTITION_TYPES];
+
 void vp9_entropy_mode_init(void);
 
 struct VP9Common;
diff --git a/vp9/common/vp9_enums.h b/vp9/common/vp9_enums.h
index efa84c40f..930a5975f 100644
--- a/vp9/common/vp9_enums.h
+++ b/vp9/common/vp9_enums.h
@@ -27,4 +27,16 @@ typedef enum BLOCK_SIZE_TYPE {
   BLOCK_SIZE_SB64X64,
 } BLOCK_SIZE_TYPE;
 
+typedef enum PARTITION_TYPE {
+  PARTITION_NONE,
+#if CONFIG_SBSEGMENT
+  PARTITION_HORZ,
+  PARTITION_VERT,
+#endif
+  PARTITION_SPLIT,
+  PARTITION_TYPES
+} PARTITION_TYPE;
+
+#define PARTITION_PLANES 2  // number of probability models
+
 #endif  // VP9_COMMON_VP9_ENUMS_H_
diff --git a/vp9/common/vp9_onyxc_int.h b/vp9/common/vp9_onyxc_int.h
index 71a430341..66698f71a 100644
--- a/vp9/common/vp9_onyxc_int.h
+++ b/vp9/common/vp9_onyxc_int.h
@@ -68,6 +68,7 @@ typedef struct frame_contexts {
   vp9_prob i8x8_mode_prob[VP9_I8X8_MODES - 1];
   vp9_prob sub_mv_ref_prob[SUBMVREF_COUNT][VP9_SUBMVREFS - 1];
   vp9_prob mbsplit_prob[VP9_NUMMBSPLITS - 1];
+  vp9_prob partition_prob[PARTITION_PLANES][PARTITION_TYPES - 1];
 
   vp9_coeff_probs coef_probs_4x4[BLOCK_TYPES];
   vp9_coeff_probs coef_probs_8x8[BLOCK_TYPES];
@@ -95,6 +96,7 @@ typedef struct frame_contexts {
   vp9_prob pre_i8x8_mode_prob[VP9_I8X8_MODES - 1];
   vp9_prob pre_sub_mv_ref_prob[SUBMVREF_COUNT][VP9_SUBMVREFS - 1];
   vp9_prob pre_mbsplit_prob[VP9_NUMMBSPLITS - 1];
+  vp9_prob pre_partition_prob[PARTITION_PLANES][PARTITION_TYPES - 1];
   unsigned int bmode_counts[VP9_NKF_BINTRAMODES];
   unsigned int ymode_counts[VP9_YMODES];   /* interframe intra mode probs */
   unsigned int sb_ymode_counts[VP9_I32X32_MODES];
@@ -102,6 +104,7 @@ typedef struct frame_contexts {
   unsigned int i8x8_mode_counts[VP9_I8X8_MODES];   /* interframe intra probs */
   unsigned int sub_mv_ref_counts[SUBMVREF_COUNT][VP9_SUBMVREFS];
   unsigned int mbsplit_counts[VP9_NUMMBSPLITS];
+  unsigned int partition_counts[PARTITION_PLANES][PARTITION_TYPES];
 
   vp9_coeff_probs pre_coef_probs_4x4[BLOCK_TYPES];
   vp9_coeff_probs pre_coef_probs_8x8[BLOCK_TYPES];
@@ -279,8 +282,6 @@ typedef struct VP9Common {
   vp9_prob prob_intra_coded;
   vp9_prob prob_last_coded;
   vp9_prob prob_gf_coded;
-  vp9_prob prob_sb32_coded;
-  vp9_prob prob_sb64_coded;
 
   // Context probabilities when using predictive coding of segment id
   vp9_prob segment_pred_probs[PREDICTION_PROBS];