Add decoder support to recursive transform block partition

This commit allows the decoder to recursively parse and rebuild
the pixel blocks.

Change-Id: I510f3a30ae7cdad5b70725c66882b00a0594e96f

vp10/decoder/decodeframe.c

@@ -390,9 +390,12 @@ static void decode_reconstruct_tx(MACROBLOCKD *const xd, vpx_reader *r,
                                   int block, int blk_row, int blk_col,
                                   TX_SIZE tx_size, int *eob_total) {
   const struct macroblockd_plane *const pd = &xd->plane[plane];
+  int tx_idx = (blk_row >> (1 - pd->subsampling_y)) * 8 +
+               (blk_col >> (1 - pd->subsampling_x));
   TX_SIZE plane_tx_size = plane ?
-      get_uv_tx_size_impl(mbmi->tx_size, mbmi->sb_type,
-                          pd->subsampling_x, pd->subsampling_y) : mbmi->tx_size;
+      get_uv_tx_size_impl(mbmi->inter_tx_size[tx_idx], mbmi->sb_type,
+                          pd->subsampling_x, pd->subsampling_y) :
+      mbmi->inter_tx_size[tx_idx];
   int max_blocks_high = num_4x4_blocks_high_lookup[plane_bsize];
   int max_blocks_wide = num_4x4_blocks_wide_lookup[plane_bsize];
 
@@ -935,7 +938,7 @@ static void decode_block(VP10Decoder *const pbi, MACROBLOCKD *const xd,
         const BLOCK_SIZE plane_bsize =
             get_plane_block_size(VPXMAX(bsize, BLOCK_8X8), pd);
         const TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
-        const int txb_size = txsize_to_bsize[max_tx_size];
+        const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
         int bw = num_4x4_blocks_wide_lookup[txb_size];
         int block = 0;
         const int step = 1 << (max_tx_size << 1);
@@ -2178,6 +2181,7 @@ static int read_compressed_header(VP10Decoder *pbi, const uint8_t *data,
 #if !CONFIG_MISC_FIXES
   cm->tx_mode = xd->lossless ? ONLY_4X4 : read_tx_mode(&r);
 #endif
+
   if (cm->tx_mode == TX_MODE_SELECT)
     read_tx_mode_probs(&fc->tx_probs, &r);
   read_coef_probs(fc, cm->tx_mode, &r);

vp10/decoder/decodemv.c

@@ -65,31 +65,44 @@ static int read_segment_id(vpx_reader *r, const struct segmentation *seg) {
 }
 
 #if CONFIG_VAR_TX
-static void read_tx_size_inter(VP10_COMMON *cm, MB_MODE_INFO *mbmi,
-                               TX_SIZE tx_size, int mi_row, int mi_col,
+static void read_tx_size_inter(VP10_COMMON *cm, MACROBLOCKD *xd,
+                               MB_MODE_INFO *mbmi,
+                               TX_SIZE tx_size, int blk_row, int blk_col,
                                vpx_reader *r) {
-  int is_split = vpx_read_bit(r);
+  int is_split = 0;
+  const int tx_idx = (blk_row >> 1) * 8 + (blk_col >> 1);
+  int max_blocks_high = num_4x4_blocks_high_lookup[mbmi->sb_type];
+  int max_blocks_wide = num_4x4_blocks_wide_lookup[mbmi->sb_type];
+  if (xd->mb_to_bottom_edge < 0)
+    max_blocks_high += xd->mb_to_bottom_edge >> 5;
+  if (xd->mb_to_right_edge < 0)
+     max_blocks_wide += xd->mb_to_right_edge >> 5;
+
+  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide)
+     return;
+
+  is_split = vpx_read_bit(r);
 
   if (is_split) {
     BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
-    int bsl = mi_width_log2_lookup[bsize];
+    int bsl = b_width_log2_lookup[bsize];
     int i;
     if (tx_size == TX_8X8) {
-      mbmi->tx_size = TX_4X4;
+      mbmi->inter_tx_size[tx_idx] = TX_4X4;
+      mbmi->tx_size = mbmi->inter_tx_size[tx_idx];
       return;
     }
 
     assert(bsl > 0);
     --bsl;
     for (i = 0; i < 4; ++i) {
-      int offsetr = mi_row + ((i >> 1) << bsl);
-      int offsetc = mi_col + ((i & 0x01) << bsl);
-      if (offsetr >= cm->mi_rows || offsetc >= cm->mi_cols)
-        continue;
-      read_tx_size_inter(cm, mbmi, tx_size - 1, offsetr, offsetc, r);
+      int offsetr = blk_row + ((i >> 1) << bsl);
+      int offsetc = blk_col + ((i & 0x01) << bsl);
+      read_tx_size_inter(cm, xd, mbmi, tx_size - 1, offsetr, offsetc, r);
     }
   } else {
-    mbmi->tx_size = tx_size;
+    mbmi->inter_tx_size[tx_idx] = tx_size;
+    mbmi->tx_size = mbmi->inter_tx_size[tx_idx];
   }
 }
 #endif
@@ -644,21 +657,31 @@ static void read_inter_frame_mode_info(VP10Decoder *const pbi,
   if (bsize >= BLOCK_8X8 && cm->tx_mode == TX_MODE_SELECT &&
       !mbmi->skip && inter_block) {
     const TX_SIZE max_tx_size = max_txsize_lookup[bsize];
-    const int txb_size = txsize_to_bsize[max_tx_size];
-    const int bs = num_8x8_blocks_wide_lookup[txb_size];
-    const int width  = num_8x8_blocks_wide_lookup[bsize];
-    const int height = num_8x8_blocks_high_lookup[bsize];
+    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
+    const int bs = num_4x4_blocks_wide_lookup[txb_size];
+    const int width  = num_4x4_blocks_wide_lookup[bsize];
+    const int height = num_4x4_blocks_high_lookup[bsize];
     int idx, idy;
     for (idy = 0; idy < height; idy += bs)
       for (idx = 0; idx < width; idx += bs)
-        read_tx_size_inter(cm, mbmi, max_tx_size,
-                           mi_row + idy, mi_col + idx, r);
+        read_tx_size_inter(cm, xd, mbmi, max_tx_size,
+                           idy, idx, r);
     if (xd->counts) {
       const int ctx = get_tx_size_context(xd);
       ++get_tx_counts(max_tx_size, ctx, &xd->counts->tx)[mbmi->tx_size];
     }
   } else {
     mbmi->tx_size = read_tx_size(cm, xd, !mbmi->skip || !inter_block, r);
+    if (inter_block) {
+      const BLOCK_SIZE txb_size = txsize_to_bsize[mbmi->tx_size];
+      const int bs = num_4x4_blocks_wide_lookup[txb_size];
+      const int width  = num_4x4_blocks_wide_lookup[bsize];
+      const int height = num_4x4_blocks_high_lookup[bsize];
+      int idx, idy;
+      for (idy = 0; idy < height; idy += bs)
+        for (idx = 0; idx < width; idx += bs)
+          mbmi->inter_tx_size[(idy >> 1) * 8 + (idx >> 1)] = mbmi->tx_size;
+    }
   }
 #else
   mbmi->tx_size = read_tx_size(cm, xd, !mbmi->skip || !inter_block, r);

vp10/encoder/bitstream.c

@@ -119,14 +119,26 @@ static int prob_diff_update_savings(const vpx_tree_index *tree,
 
 #if CONFIG_VAR_TX
 static void write_tx_size_inter(const VP10_COMMON *cm,
+                                const MACROBLOCKD *xd,
                                 const MB_MODE_INFO *mbmi,
-                                TX_SIZE tx_size, int mi_row, int mi_col,
+                                TX_SIZE tx_size, int blk_row, int blk_col,
                                 vpx_writer *w) {
-  if (tx_size == mbmi->tx_size) {
+  const int tx_idx = (blk_row >> 1) * 8 + (blk_col >> 1);
+  int max_blocks_high = num_4x4_blocks_high_lookup[mbmi->sb_type];
+  int max_blocks_wide = num_4x4_blocks_wide_lookup[mbmi->sb_type];
+  if (xd->mb_to_bottom_edge < 0)
+    max_blocks_high += xd->mb_to_bottom_edge >> 5;
+  if (xd->mb_to_right_edge < 0)
+     max_blocks_wide += xd->mb_to_right_edge >> 5;
+
+  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide)
+     return;
+
+  if (tx_size == mbmi->inter_tx_size[tx_idx]) {
     vpx_write_bit(w, 0);
   } else {
     const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
-    int bsl = mi_width_log2_lookup[bsize];
+    int bsl = b_width_log2_lookup[bsize];
     int i;
     vpx_write_bit(w, 1);
 
@@ -136,11 +148,9 @@ static void write_tx_size_inter(const VP10_COMMON *cm,
     assert(bsl > 0);
     --bsl;
     for (i = 0; i < 4; ++i) {
-      int offsetr = mi_row + ((i >> 1) << bsl);
-      int offsetc = mi_col + ((i & 0x01) << bsl);
-      if (offsetr >= cm->mi_rows || offsetc >= cm->mi_cols)
-        continue;
-      write_tx_size_inter(cm, mbmi, tx_size - 1, offsetr, offsetc, w);
+      int offsetr = blk_row + ((i >> 1) << bsl);
+      int offsetc = blk_col + ((i & 0x01) << bsl);
+      write_tx_size_inter(cm, xd, mbmi, tx_size - 1, offsetr, offsetc, w);
     }
   }
 }
@@ -345,9 +355,6 @@ static void write_ref_frames(const VP10_COMMON *cm, const MACROBLOCKD *xd,
 }
 
 static void pack_inter_mode_mvs(VP10_COMP *cpi, const MODE_INFO *mi,
-#if CONFIG_VAR_TX
-                                int mi_row, int mi_col,
-#endif
                                 vpx_writer *w) {
   VP10_COMMON *const cm = &cpi->common;
   const nmv_context *nmvc = &cm->fc->nmvc;
@@ -387,14 +394,13 @@ static void pack_inter_mode_mvs(VP10_COMP *cpi, const MODE_INFO *mi,
     if (is_inter) {  // This implies skip flag is 0.
       const TX_SIZE max_tx_size = max_txsize_lookup[bsize];
       const int txb_size = txsize_to_bsize[max_tx_size];
-      const int bs = num_8x8_blocks_wide_lookup[txb_size];
-      const int width  = num_8x8_blocks_wide_lookup[bsize];
-      const int height = num_8x8_blocks_high_lookup[bsize];
+      const int bs = num_4x4_blocks_wide_lookup[txb_size];
+      const int width  = num_4x4_blocks_wide_lookup[bsize];
+      const int height = num_4x4_blocks_high_lookup[bsize];
       int idx, idy;
       for (idy = 0; idy < height; idy += bs)
         for (idx = 0; idx < width; idx += bs)
-          write_tx_size_inter(cm, mbmi, max_tx_size,
-                              mi_row + idy, mi_col + idx, w);
+          write_tx_size_inter(cm, xd, mbmi, max_tx_size, idy, idx, w);
     } else {
       write_selected_tx_size(cm, xd, w);
     }
@@ -549,11 +555,7 @@ static void write_modes_b(VP10_COMP *cpi, const TileInfo *const tile,
   if (frame_is_intra_only(cm)) {
     write_mb_modes_kf(cm, xd, xd->mi, w);
   } else {
-#if CONFIG_VAR_TX
-    pack_inter_mode_mvs(cpi, m, mi_row, mi_col, w);
-#else
     pack_inter_mode_mvs(cpi, m, w);
-#endif
   }
 
   assert(*tok < tok_end);

vp10/encoder/encodemb.c

@@ -1500,9 +1500,11 @@ static void encode_block_inter(int plane, int block, int blk_row, int blk_col,
   MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
 
   const struct macroblockd_plane *const pd = &xd->plane[plane];
+  const int blk_idx = ((blk_row >> 1) << 3) + (blk_col >> 1);
   TX_SIZE plane_tx_size = plane ?
-      get_uv_tx_size_impl(mbmi->tx_size, mbmi->sb_type,
-                          pd->subsampling_x, pd->subsampling_y) : mbmi->tx_size;
+      get_uv_tx_size_impl(mbmi->inter_tx_size[blk_idx], mbmi->sb_type,
+                          pd->subsampling_x, pd->subsampling_y) :
+      mbmi->inter_tx_size[blk_idx];
 
   int max_blocks_high = num_4x4_blocks_high_lookup[plane_bsize];
   int max_blocks_wide = num_4x4_blocks_wide_lookup[plane_bsize];
@@ -1601,8 +1603,8 @@ void vp10_encode_sb(MACROBLOCK *x, BLOCK_SIZE bsize) {
     const int mi_width = num_4x4_blocks_wide_lookup[plane_bsize];
     const int mi_height = num_4x4_blocks_high_lookup[plane_bsize];
     const TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
-    int txb_size = txsize_to_bsize[max_tx_size];
-    int bh = num_4x4_blocks_wide_lookup[txb_size];
+    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
+    const int bh = num_4x4_blocks_wide_lookup[txb_size];
     int idx, idy;
     int block = 0;
     int step = 1 << (max_tx_size * 2);

vp10/encoder/rdopt.c

@@ -2781,6 +2781,11 @@ static int64_t handle_inter_mode(VP10_COMP *cpi, MACROBLOCK *x,
     super_block_yrd(cpi, x, rate_y, &distortion_y, &skippable_y, psse,
                     bsize, ref_best_rd);
 
+#if CONFIG_VAR_TX
+    for (i = 0; i < 64; ++i)
+      mbmi->inter_tx_size[i] = mbmi->tx_size;
+#endif
+
     if (*rate_y == INT_MAX) {
       *rate2 = INT_MAX;
       *distortion = INT64_MAX;
@@ -3932,6 +3937,11 @@ void vp10_rd_pick_inter_mode_sub8x8(VP10_COMP *cpi,
         xd->plane[i].pre[1] = yv12_mb[second_ref_frame][i];
     }
 
+#if CONFIG_VAR_TX
+    for (i = 0; i < 64; ++i)
+      mbmi->inter_tx_size[i] = mbmi->tx_size;
+#endif
+
     if (ref_frame == INTRA_FRAME) {
       int rate;
       if (rd_pick_intra_sub_8x8_y_mode(cpi, x, &rate, &rate_y,

vp10/encoder/tokenize.c

@@ -622,9 +622,11 @@ void tokenize_tx(ThreadData *td, TOKENEXTRA **t,
   MACROBLOCKD *const xd = &x->e_mbd;
   MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
   const struct macroblockd_plane *const pd = &xd->plane[plane];
+  const int blk_idx = ((blk_row >> 1) << 3) + (blk_col >> 1);
   TX_SIZE plane_tx_size = plane ?
-      get_uv_tx_size_impl(mbmi->tx_size, mbmi->sb_type,
-                          pd->subsampling_x, pd->subsampling_y) : mbmi->tx_size;
+      get_uv_tx_size_impl(mbmi->inter_tx_size[blk_idx], mbmi->sb_type,
+                          pd->subsampling_x, pd->subsampling_y) :
+      mbmi->inter_tx_size[blk_idx];
 
   int max_blocks_high = num_4x4_blocks_high_lookup[plane_bsize];
   int max_blocks_wide = num_4x4_blocks_wide_lookup[plane_bsize];
@@ -702,7 +704,7 @@ void vp10_tokenize_sb_inter(VP10_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
     const int mi_width = num_4x4_blocks_wide_lookup[plane_bsize];
     const int mi_height = num_4x4_blocks_high_lookup[plane_bsize];
     const TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
-    int txb_size = txsize_to_bsize[max_tx_size];
+    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
     int bh = num_4x4_blocks_wide_lookup[txb_size];
     int idx, idy;
     int block = 0;