Handle partition type coding of boundary blocks

The partition types of blocks sitting on the frame boundary are
constrained by the block size and the position of each sub-block
relative to the frame. Hence we use truncated probability models
to handle the coding of such information.

100 frames run:
yt 0.138%

Change-Id: I85d9b45665c15280069c0234ea6f778af586d87d

test/borders_test.cc

@@ -34,7 +34,7 @@ class BordersTest : public ::libvpx_test::EncoderTest,
   virtual void PreEncodeFrameHook(::libvpx_test::VideoSource *video,
                                   ::libvpx_test::Encoder *encoder) {
     if ( video->frame() == 1) {
-      encoder->Control(VP8E_SET_CPUUSED, 5);
+      encoder->Control(VP8E_SET_CPUUSED, 0);
       encoder->Control(VP8E_SET_ENABLEAUTOALTREF, 1);
       encoder->Control(VP8E_SET_ARNR_MAXFRAMES, 7);
       encoder->Control(VP8E_SET_ARNR_STRENGTH, 5);

vp9/common/vp9_alloccommon.c

@@ -74,7 +74,7 @@ static void set_mb_mi(VP9_COMMON *cm, int aligned_width, int aligned_height) {
 
   cm->mi_cols = aligned_width >> LOG2_MI_SIZE;
   cm->mi_rows = aligned_height >> LOG2_MI_SIZE;
-  cm->mode_info_stride = cm->mi_cols + 1;
+  cm->mode_info_stride = cm->mi_cols + 64 / MI_SIZE;
 }
 
 static void setup_mi(VP9_COMMON *cm) {
@@ -131,12 +131,13 @@ int vp9_alloc_frame_buffers(VP9_COMMON *oci, int width, int height) {
   set_mb_mi(oci, aligned_width, aligned_height);
 
   // Allocation
-  oci->mip = vpx_calloc(oci->mode_info_stride * (oci->mi_rows + 1),
+  oci->mip = vpx_calloc(oci->mode_info_stride * (oci->mi_rows + 64 / MI_SIZE),
                         sizeof(MODE_INFO));
   if (!oci->mip)
     goto fail;
 
-  oci->prev_mip = vpx_calloc(oci->mode_info_stride * (oci->mi_rows + 1),
+  oci->prev_mip = vpx_calloc(oci->mode_info_stride *
+                             (oci->mi_rows + 64 / MI_SIZE),
                              sizeof(MODE_INFO));
   if (!oci->prev_mip)
     goto fail;

vp9/common/vp9_onyxc_int.h

@@ -301,6 +301,30 @@ static INLINE void set_partition_seg_context(VP9_COMMON *cm,
   xd->left_seg_context  = cm->left_seg_context + (mi_row & MI_MASK);
 }
 
+static int check_bsize_coverage(VP9_COMMON *cm, MACROBLOCKD *xd,
+                                int mi_row, int mi_col,
+                                BLOCK_SIZE_TYPE bsize) {
+  int bsl = mi_width_log2(bsize), bs = 1 << bsl;
+  int ms = bs / 2;
+
+  if ((mi_row + bs <= cm->mi_rows) && (mi_col + ms < cm->mi_cols))
+    return 0;
+  if ((mi_col + bs <= cm->mi_cols) && (mi_row + ms < cm->mi_rows))
+    return 0;
+
+  // frame width/height are multiples of 8, hence 8x8 block should always
+  // pass the above check
+  assert(bsize > BLOCK_SIZE_SB8X8);
+
+  // return the node index in the prob tree for binary coding
+  if ((mi_col + bs <= cm->mi_cols) && (mi_row + ms >= cm->mi_rows))
+    return 1;
+  if ((mi_row + bs <= cm->mi_rows) && (mi_col + ms >= cm->mi_cols))
+    return 2;
+
+  return -1;
+}
+
 static void set_mi_row_col(VP9_COMMON *cm, MACROBLOCKD *xd,
                        int mi_row, int bh,
                        int mi_col, int bw) {

vp9/decoder/vp9_decodframe.c

@@ -469,12 +469,21 @@ static void decode_modes_sb(VP9D_COMP *pbi, int mi_row, int mi_col,
 
   if (bsize >= BLOCK_SIZE_SB8X8) {
     int pl;
+    int idx = check_bsize_coverage(pc, xd, mi_row, mi_col, bsize);
     // read the partition information
     xd->left_seg_context = pc->left_seg_context + (mi_row & MI_MASK);
     xd->above_seg_context = pc->above_seg_context + mi_col;
     pl = partition_plane_context(xd, bsize);
+
+    if (idx == 0)
       partition = treed_read(r, vp9_partition_tree,
                              pc->fc.partition_prob[pc->frame_type][pl]);
+    else if (idx > 0 &&
+        !vp9_read(r, pc->fc.partition_prob[pc->frame_type][pl][idx]))
+      partition = (idx == 1) ? PARTITION_HORZ : PARTITION_VERT;
+    else
+      partition = PARTITION_SPLIT;
+
     pc->fc.partition_counts[pl][partition]++;
   }
 

vp9/encoder/vp9_bitstream.c

@@ -803,13 +803,18 @@ static void write_modes_sb(VP9_COMP *cpi, MODE_INFO *m, vp9_writer *bc,
 
   if (bsize >= BLOCK_SIZE_SB8X8) {
     int pl;
+    int idx = check_bsize_coverage(cm, xd, mi_row, mi_col, bsize);
     xd->left_seg_context = cm->left_seg_context + (mi_row & MI_MASK);
     xd->above_seg_context = cm->above_seg_context + mi_col;
     pl = partition_plane_context(xd, bsize);
     // encode the partition information
+    if (idx == 0)
       write_token(bc, vp9_partition_tree,
                   cm->fc.partition_prob[cm->frame_type][pl],
                   vp9_partition_encodings + partition);
+    else if (idx > 0)
+      vp9_write(bc, partition == PARTITION_SPLIT,
+                cm->fc.partition_prob[cm->frame_type][pl][idx]);
   }
 
   subsize = get_subsize(bsize, partition);

vp9/encoder/vp9_encodeframe.c

@@ -1286,8 +1286,7 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
   }
 
   // PARTITION_HORZ
-  if ((mi_col + ms <= cm->mi_cols) && (mi_row + (ms >> 1) <= cm->mi_rows) &&
+    if ((bsize >= BLOCK_SIZE_SB8X8) && (mi_col + ms <= cm->mi_cols)) {
-      (bsize >= BLOCK_SIZE_SB8X8)) {
     int r2, d2;
     int mb_skip = 0;
     subsize = get_subsize(bsize, PARTITION_HORZ);
@@ -1295,7 +1294,7 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
     pick_sb_modes(cpi, mi_row, mi_col, tp, &r2, &d2, subsize,
                   get_block_context(x, subsize));
 
-    if (mi_row + ms <= cm->mi_rows) {
+    if (mi_row < cm->mi_rows) {
       int r = 0, d = 0;
       update_state(cpi, get_block_context(x, subsize), subsize, 0);
       encode_superblock(cpi, tp, 0, mi_row, mi_col, subsize);
@@ -1322,15 +1321,14 @@ static void rd_pick_partition(VP9_COMP *cpi, TOKENEXTRA **tp,
   }
 
   // PARTITION_VERT
-  if ((mi_row + ms <= cm->mi_rows) && (mi_col + (ms >> 1) <= cm->mi_cols) &&
+  if ((bsize >= BLOCK_SIZE_SB8X8) && (mi_row + ms <= cm->mi_rows)) {
-      (bsize >= BLOCK_SIZE_SB8X8)) {
     int r2, d2;
     int mb_skip = 0;
     subsize = get_subsize(bsize, PARTITION_VERT);
     *(get_sb_index(xd, subsize)) = 0;
     pick_sb_modes(cpi, mi_row, mi_col, tp, &r2, &d2, subsize,
                   get_block_context(x, subsize));
-    if (mi_col + ms <= cm->mi_cols) {
+    if (mi_col < cm->mi_cols) {
       int r = 0, d = 0;
       update_state(cpi, get_block_context(x, subsize), subsize, 0);
       encode_superblock(cpi, tp, 0, mi_row, mi_col, subsize);
@@ -1404,7 +1402,7 @@ static void encode_sb_row(VP9_COMP *cpi, int mi_row,
 
   // Code each SB in the row
   for (mi_col = cm->cur_tile_mi_col_start;
-       mi_col < cm->cur_tile_mi_col_end; mi_col += 8) {
+       mi_col < cm->cur_tile_mi_col_end; mi_col += 64 / MI_SIZE) {
     int dummy_rate, dummy_dist;
     if (cpi->speed < 5) {
       rd_pick_partition(cpi, tp, mi_row, mi_col, BLOCK_SIZE_SB64X64,

vp9/encoder/vp9_onyx_if.c

@@ -815,7 +815,7 @@ static int alloc_partition_data(VP9_COMP *cpi) {
   vpx_free(cpi->mb.pip);
 
   cpi->mb.pip = vpx_calloc((cpi->common.mode_info_stride) *
-                           (cpi->common.mi_rows + 1),
+                           (cpi->common.mi_rows + 64 / MI_SIZE),
                            sizeof(PARTITION_INFO));
   if (!cpi->mb.pip)
     return 1;
@@ -3385,11 +3385,11 @@ static void encode_frame_to_data_rate(VP9_COMP *cpi,
 
   if (cm->show_frame) {
     vpx_memcpy(cm->prev_mip, cm->mip,
-               cm->mode_info_stride * (cm->mi_rows + 1) *
+               cm->mode_info_stride * (cm->mi_rows + 64 / MI_SIZE) *
                sizeof(MODE_INFO));
   } else {
     vpx_memset(cm->prev_mip, 0,
-               cm->mode_info_stride * (cm->mi_rows + 1) *
+               cm->mode_info_stride * (cm->mi_rows + 64 / MI_SIZE) *
                sizeof(MODE_INFO));
   }
   // restore prev_mi