Enable bit-stream support to SB8X8

This commit enables bit-stream writing and reading for recursive
partition down to block 8x8.

Change-Id: I163cd48d191cc94ead49cbb7fc91374f6bf204e2

4 files changed, 89 insertions, 26 deletions

diff --git a/vp9/common/vp9_blockd.h b/vp9/common/vp9_blockd.h
index 426699e31..2a72034cf 100644
--- a/vp9/common/vp9_blockd.h
+++ b/vp9/common/vp9_blockd.h
@@ -511,6 +511,55 @@ static INLINE int partition_plane_context(MACROBLOCKD *xd,
   return (left * 2 + above) + (bsl - 1) * PARTITION_PLOFFSET;
 }
 
+static BLOCK_SIZE_TYPE get_subsize(BLOCK_SIZE_TYPE bsize,
+                                   PARTITION_TYPE partition) {
+  BLOCK_SIZE_TYPE subsize;
+  switch (partition) {
+    case PARTITION_NONE:
+      subsize = bsize;
+      break;
+    case PARTITION_HORZ:
+      if (bsize == BLOCK_SIZE_SB64X64)
+        subsize = BLOCK_SIZE_SB64X32;
+      else if (bsize == BLOCK_SIZE_SB32X32)
+        subsize = BLOCK_SIZE_SB32X16;
+#if CONFIG_SB8X8
+      else if (bsize == BLOCK_SIZE_MB16X16)
+        subsize = BLOCK_SIZE_SB16X8;
+#endif
+      else
+        assert(0);
+      break;
+    case PARTITION_VERT:
+      if (bsize == BLOCK_SIZE_SB64X64)
+        subsize = BLOCK_SIZE_SB32X64;
+      else if (bsize == BLOCK_SIZE_SB32X32)
+        subsize = BLOCK_SIZE_SB16X32;
+#if CONFIG_SB8X8
+      else if (bsize == BLOCK_SIZE_MB16X16)
+        subsize = BLOCK_SIZE_SB8X16;
+#endif
+      else
+        assert(0);
+      break;
+    case PARTITION_SPLIT:
+      if (bsize == BLOCK_SIZE_SB64X64)
+        subsize = BLOCK_SIZE_SB32X32;
+      else if (bsize == BLOCK_SIZE_SB32X32)
+        subsize = BLOCK_SIZE_MB16X16;
+#if CONFIG_SB8X8
+      else if (bsize == BLOCK_SIZE_MB16X16)
+        subsize = BLOCK_SIZE_SB8X8;
+#endif
+      else
+        assert(0);
+      break;
+    default:
+      assert(0);
+  }
+  return subsize;
+}
+
 #define ACTIVE_HT   110                // quantization stepsize threshold
 
 #define ACTIVE_HT8  300
diff --git a/vp9/decoder/vp9_decodemv.c b/vp9/decoder/vp9_decodemv.c
index 474250cf7..0e55657f8 100644
--- a/vp9/decoder/vp9_decodemv.c
+++ b/vp9/decoder/vp9_decodemv.c
@@ -146,7 +146,11 @@ static void kfread_modes(VP9D_COMP *pbi, MODE_INFO *m,
     m->mbmi.mb_skip_coeff = vp9_read(r, vp9_get_pred_prob(cm, xd, PRED_MBSKIP));
 
   // luma mode
+#if CONFIG_SB8X8
+  m->mbmi.mode = m->mbmi.sb_type > BLOCK_SIZE_SB8X8 ?
+#else
   m->mbmi.mode = m->mbmi.sb_type > BLOCK_SIZE_MB16X16 ?
+#endif
       read_kf_sb_ymode(r, cm->sb_kf_ymode_prob[cm->kf_ymode_probs_index]):
       read_kf_mb_ymode(r, cm->kf_ymode_prob[cm->kf_ymode_probs_index]);
 
@@ -154,11 +158,10 @@ static void kfread_modes(VP9D_COMP *pbi, MODE_INFO *m,
 
   if (m->mbmi.mode == I4X4_PRED) {
     int i;
-    for (i = 0; i < 16; ++i) {
+    for (i = 0; i < (16 >> (2 * CONFIG_SB8X8)); ++i) {
       const B_PREDICTION_MODE a = above_block_mode(m, i, mis);
       const B_PREDICTION_MODE l = xd->left_available || (i & 3) ?
                                   left_block_mode(m, i) : B_DC_PRED;
-
       m->bmi[i].as_mode.first = read_kf_bmode(r, cm->kf_bmode_prob[a][l]);
     }
   }
diff --git a/vp9/decoder/vp9_decodframe.c b/vp9/decoder/vp9_decodframe.c
index f561c4850..724496492 100644
--- a/vp9/decoder/vp9_decodframe.c
+++ b/vp9/decoder/vp9_decodframe.c
@@ -576,10 +576,12 @@ static void decode_modes_b(VP9D_COMP *pbi, int mi_row, int mi_col,
   set_refs(pbi, mi_row, mi_col);
 
 #if CONFIG_SB8X8
-  if (bsize >= BLOCK_SIZE_SB8X8)
-    decode_sb(pbi, xd, mi_row, mi_col, r, bsize);
+  if (bsize == BLOCK_SIZE_SB8X8 &&
+      (xd->mode_info_context->mbmi.mode == SPLITMV ||
+       xd->mode_info_context->mbmi.mode == I4X4_PRED))
+    decode_atom(pbi, xd, mi_row, mi_col, r, bsize);
   else
-    decode_atom(pbi, xd, mi_row, mi_col, r, BLOCK_SIZE_SB8X8);
+    decode_sb(pbi, xd, mi_row, mi_col, r, bsize);
 #else
   // TODO(jingning): merge decode_sb_ and decode_mb_
   if (bsize > BLOCK_SIZE_MB16X16) {
@@ -616,7 +618,11 @@ static void decode_modes_sb(VP9D_COMP *pbi, int mi_row, int mi_col,
   if (mi_row >= pc->mi_rows || mi_col >= pc->mi_cols)
     return;
 
+#if CONFIG_SB8X8
+  if (bsize > BLOCK_SIZE_SB8X8) {
+#else
   if (bsize > BLOCK_SIZE_MB16X16) {
+#endif
     int pl;
     // read the partition information
     xd->left_seg_context =
@@ -628,34 +634,35 @@ static void decode_modes_sb(VP9D_COMP *pbi, int mi_row, int mi_col,
     pc->fc.partition_counts[pl][partition]++;
   }
 
+  subsize = get_subsize(bsize, partition);
   switch (partition) {
     case PARTITION_NONE:
-      subsize = bsize;
       decode_modes_b(pbi, mi_row, mi_col, r, subsize);
       break;
     case PARTITION_HORZ:
-      subsize = (bsize == BLOCK_SIZE_SB64X64) ? BLOCK_SIZE_SB64X32 :
-                                                BLOCK_SIZE_SB32X16;
       decode_modes_b(pbi, mi_row, mi_col, r, subsize);
       if ((mi_row + bs) < pc->mi_rows)
         decode_modes_b(pbi, mi_row + bs, mi_col, r, subsize);
       break;
     case PARTITION_VERT:
-      subsize = (bsize == BLOCK_SIZE_SB64X64) ? BLOCK_SIZE_SB32X64 :
-                                                BLOCK_SIZE_SB16X32;
       decode_modes_b(pbi, mi_row, mi_col, r, subsize);
       if ((mi_col + bs) < pc->mi_cols)
         decode_modes_b(pbi, mi_row, mi_col + bs, r, subsize);
       break;
     case PARTITION_SPLIT:
-      subsize = (bsize == BLOCK_SIZE_SB64X64) ? BLOCK_SIZE_SB32X32 :
-                                                BLOCK_SIZE_MB16X16;
       for (n = 0; n < 4; n++) {
         int j = n >> 1, i = n & 0x01;
         if (subsize == BLOCK_SIZE_SB32X32)
           xd->sb_index = n;
+#if CONFIG_SB8X8
+        else if (subsize == BLOCK_SIZE_MB16X16)
+          xd->mb_index = n;
+        else
+          xd->b_index = n;
+#else
         else
           xd->mb_index = n;
+#endif
         decode_modes_sb(pbi, mi_row + j * bs, mi_col + i * bs, r, subsize);
       }
       break;
@@ -663,7 +670,11 @@ static void decode_modes_sb(VP9D_COMP *pbi, int mi_row, int mi_col,
       assert(0);
   }
   // update partition context
+#if CONFIG_SB8X8
+  if ((partition == PARTITION_SPLIT) && (bsize > BLOCK_SIZE_MB16X16))
+#else
   if ((partition == PARTITION_SPLIT) && (bsize > BLOCK_SIZE_SB32X32))
+#endif
     return;
 
   xd->left_seg_context = pc->left_seg_context + ((mi_row >> CONFIG_SB8X8) & 3);
diff --git a/vp9/encoder/vp9_bitstream.c b/vp9/encoder/vp9_bitstream.c
index 3c0bab2ce..f8e273364 100644
--- a/vp9/encoder/vp9_bitstream.c
+++ b/vp9/encoder/vp9_bitstream.c
@@ -944,7 +944,11 @@ static void write_mb_modes_kf(const VP9_COMP *cpi,
     vp9_write(bc, skip_coeff, vp9_get_pred_prob(c, xd, PRED_MBSKIP));
   }
 
+#if CONFIG_SB8X8
+  if (m->mbmi.sb_type > BLOCK_SIZE_SB8X8)
+#else
   if (m->mbmi.sb_type > BLOCK_SIZE_MB16X16)
+#endif
     sb_kfwrite_ymode(bc, ym, c->sb_kf_ymode_prob[c->kf_ymode_probs_index]);
   else
     kfwrite_ymode(bc, ym, c->kf_ymode_prob[c->kf_ymode_probs_index]);
@@ -960,7 +964,6 @@ static void write_mb_modes_kf(const VP9_COMP *cpi,
 #ifdef ENTROPY_STATS
       ++intra_mode_stats [A] [L] [bm];
 #endif
-
       write_kf_bmode(bc, bm, c->kf_bmode_prob[a][l]);
     } while (++i < (16 >> (CONFIG_SB8X8 * 2)));
   }
@@ -1210,7 +1213,11 @@ static void write_modes_sb(VP9_COMP *cpi, MODE_INFO *m, vp9_writer *bc,
   else
     assert(0);
 
+#if CONFIG_SB8X8
+  if (bsize > BLOCK_SIZE_SB8X8) {
+#else
   if (bsize > BLOCK_SIZE_MB16X16) {
+#endif
     int pl;
     xd->left_seg_context =
         cm->left_seg_context + ((mi_row >> CONFIG_SB8X8) & 3);
@@ -1221,34 +1228,23 @@ static void write_modes_sb(VP9_COMP *cpi, MODE_INFO *m, vp9_writer *bc,
                 vp9_partition_encodings + partition);
   }
 
+  subsize = get_subsize(bsize, partition);
+
   switch (partition) {
     case PARTITION_NONE:
-      subsize = bsize;
       write_modes_b(cpi, m, bc, tok, tok_end, mi_row, mi_col);
       break;
     case PARTITION_HORZ:
-      subsize = (bsize == BLOCK_SIZE_SB64X64) ? BLOCK_SIZE_SB64X32 :
-                                                BLOCK_SIZE_SB32X16;
       write_modes_b(cpi, m, bc, tok, tok_end, mi_row, mi_col);
       if ((mi_row + bh) < cm->mi_rows)
         write_modes_b(cpi, m + bh * mis, bc, tok, tok_end, mi_row + bh, mi_col);
       break;
     case PARTITION_VERT:
-      subsize = (bsize == BLOCK_SIZE_SB64X64) ? BLOCK_SIZE_SB32X64 :
-                                                BLOCK_SIZE_SB16X32;
       write_modes_b(cpi, m, bc, tok, tok_end, mi_row, mi_col);
       if ((mi_col + bw) < cm->mi_cols)
         write_modes_b(cpi, m + bw, bc, tok, tok_end, mi_row, mi_col + bw);
       break;
     case PARTITION_SPLIT:
-      // TODO(jingning): support recursive partitioning down to 16x16 as for
-      // now. need to merge in 16x8, 8x16, 8x8, and smaller partitions.
-      if (bsize == BLOCK_SIZE_SB64X64)
-        subsize = BLOCK_SIZE_SB32X32;
-      else if (bsize == BLOCK_SIZE_SB32X32)
-        subsize = BLOCK_SIZE_MB16X16;
-      else
-        assert(0);
       for (n = 0; n < 4; n++) {
         int j = n >> 1, i = n & 0x01;
         write_modes_sb(cpi, m + j * bs * mis + i * bs, bc, tok, tok_end,
@@ -1260,7 +1256,11 @@ static void write_modes_sb(VP9_COMP *cpi, MODE_INFO *m, vp9_writer *bc,
   }
 
   // update partition context
+#if CONFIG_SB8X8
+  if ((partition == PARTITION_SPLIT) && (bsize > BLOCK_SIZE_MB16X16))
+#else
   if ((partition == PARTITION_SPLIT) && (bsize > BLOCK_SIZE_SB32X32))
+#endif
     return;
 
   xd->left_seg_context = cm->left_seg_context + ((mi_row >> CONFIG_SB8X8) & 3);