Rate-distortion optimization for recursive transform block coding

This commit enables the rate-distortion optimization for recursive
transform block coding scheme.

Change-Id: Id6a8336ca847bb3af1e94cbfb51db1f4da12d38f

vp10/common/pred_common.h

@@ -165,6 +165,70 @@ static INLINE unsigned int *get_tx_counts(TX_SIZE max_tx_size, int ctx,
   }
 }
 
+#if CONFIG_VAR_TX
+static void update_tx_counts(VP10_COMMON *cm, MACROBLOCKD *xd,
+                             MB_MODE_INFO *mbmi, BLOCK_SIZE plane_bsize,
+                             TX_SIZE tx_size, int blk_row, int blk_col,
+                             TX_SIZE max_tx_size, int ctx,
+                             struct tx_counts *tx_counts) {
+  const struct macroblockd_plane *const pd = &xd->plane[0];
+  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
+  int tx_idx = (blk_row >> (1 - pd->subsampling_y)) * 8 +
+               (blk_col >> (1 - pd->subsampling_x));
+  TX_SIZE plane_tx_size = 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];
+
+  if (xd->mb_to_bottom_edge < 0)
+    max_blocks_high += xd->mb_to_bottom_edge >> (5 + pd->subsampling_y);
+  if (xd->mb_to_right_edge < 0)
+    max_blocks_wide += xd->mb_to_right_edge >> (5 + pd->subsampling_x);
+
+  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide)
+    return;
+
+  if (tx_size == plane_tx_size) {
+    ++get_tx_counts(max_tx_size, ctx, tx_counts)[tx_size];
+  } else {
+    int bsl = b_width_log2_lookup[bsize];
+    int i;
+
+    assert(bsl > 0);
+    --bsl;
+
+    for (i = 0; i < 4; ++i) {
+      const int offsetr = blk_row + ((i >> 1) << bsl);
+      const int offsetc = blk_col + ((i & 0x01) << bsl);
+
+      if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide)
+        continue;
+      update_tx_counts(cm, xd, mbmi, plane_bsize,
+                       tx_size - 1, offsetr, offsetc,
+                       max_tx_size, ctx, tx_counts);
+    }
+  }
+}
+
+static INLINE void inter_block_tx_count_update(VP10_COMMON *cm,
+                                               MACROBLOCKD *xd,
+                                               MB_MODE_INFO *mbmi,
+                                               BLOCK_SIZE plane_bsize,
+                                               int ctx,
+                                               struct tx_counts *tx_counts) {
+  const int mi_width = num_4x4_blocks_wide_lookup[plane_bsize];
+  const int mi_height = num_4x4_blocks_high_lookup[plane_bsize];
+  TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
+  BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
+  int bh = num_4x4_blocks_wide_lookup[txb_size];
+  int idx, idy;
+
+  for (idy = 0; idy < mi_height; idy += bh)
+    for (idx = 0; idx < mi_width; idx += bh)
+      update_tx_counts(cm, xd, mbmi, plane_bsize, max_tx_size, idy, idx,
+                       max_tx_size, ctx, tx_counts);
+}
+#endif
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

vp10/decoder/decodeframe.c

@@ -390,11 +390,11 @@ 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];
+  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
   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->inter_tx_size[tx_idx], mbmi->sb_type,
-                          pd->subsampling_x, pd->subsampling_y) :
+      get_uv_tx_size_impl(mbmi->inter_tx_size[tx_idx], bsize, 0, 0) :
       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];
@@ -419,7 +419,6 @@ static void decode_reconstruct_tx(MACROBLOCKD *const xd, vpx_reader *r,
         pd->dst.stride, eob, block);
     *eob_total += eob;
   } else {
-    const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
     int bsl = b_width_log2_lookup[bsize];
     int i;
 

vp10/decoder/decodemv.c

@@ -668,7 +668,8 @@ static void read_inter_frame_mode_info(VP10Decoder *const pbi,
                            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];
+      inter_block_tx_count_update(cm, xd, mbmi, bsize,
+                                  ctx, &xd->counts->tx);
     }
   } else {
     mbmi->tx_size = read_tx_size(cm, xd, !mbmi->skip || !inter_block, r);
@@ -731,7 +732,6 @@ void vp10_read_mode_info(VP10Decoder *const pbi, MACROBLOCKD *xd,
     read_intra_frame_mode_info(cm, xd, mi_row, mi_col, r);
   } else {
     read_inter_frame_mode_info(pbi, xd, mi_row, mi_col, r);
-
     for (h = 0; h < y_mis; ++h) {
       MV_REF *const frame_mv = frame_mvs + h * cm->mi_cols;
       for (w = 0; w < x_mis; ++w) {

vp10/encoder/encodeframe.c

@@ -2974,8 +2974,18 @@ static void encode_superblock(VP10_COMP *cpi, ThreadData *td,
     if (cm->tx_mode == TX_MODE_SELECT &&
         mbmi->sb_type >= BLOCK_8X8  &&
         !(is_inter_block(mbmi) && (mbmi->skip || seg_skip))) {
+#if CONFIG_VAR_TX
+      int tx_size_ctx = get_tx_size_context(xd);
+      if (is_inter_block(mbmi))
+        inter_block_tx_count_update(cm, xd, mbmi, bsize,
+                                    tx_size_ctx, &td->counts->tx);
+      else
+        ++get_tx_counts(max_txsize_lookup[bsize], get_tx_size_context(xd),
+                        &td->counts->tx)[mbmi->tx_size];
+#else
       ++get_tx_counts(max_txsize_lookup[bsize], get_tx_size_context(xd),
                       &td->counts->tx)[mbmi->tx_size];
+#endif
     } else {
       int x, y;
       TX_SIZE tx_size;

vp10/encoder/encodemb.c

@@ -1530,12 +1530,13 @@ static void encode_block_inter(int plane, int block, int blk_row, int blk_col,
   MACROBLOCK *const x = args->x;
   MACROBLOCKD *const xd = &x->e_mbd;
   MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
-
+  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
   const struct macroblockd_plane *const pd = &xd->plane[plane];
-  const int blk_idx = ((blk_row >> 1) << 3) + (blk_col >> 1);
+  int blk_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->inter_tx_size[blk_idx], mbmi->sb_type,
-                          pd->subsampling_x, pd->subsampling_y) :
+      get_uv_tx_size_impl(mbmi->inter_tx_size[blk_idx], bsize,
+                          0, 0) :
       mbmi->inter_tx_size[blk_idx];
 
   int max_blocks_high = num_4x4_blocks_high_lookup[plane_bsize];
@@ -1553,7 +1554,6 @@ static void encode_block_inter(int plane, int block, int blk_row, int blk_col,
     encode_block(plane, block, blk_row, blk_col, plane_bsize,
                  tx_size, arg);
   } else {
-    const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
     int bsl = b_width_log2_lookup[bsize];
     int i;
 

vp10/encoder/tokenize.c

@@ -622,10 +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);
+  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
+  int blk_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->inter_tx_size[blk_idx], mbmi->sb_type,
-                          pd->subsampling_x, pd->subsampling_y) :
+      get_uv_tx_size_impl(mbmi->inter_tx_size[blk_idx], bsize, 0, 0) :
       mbmi->inter_tx_size[blk_idx];
 
   int max_blocks_high = num_4x4_blocks_high_lookup[plane_bsize];
@@ -647,7 +648,6 @@ void tokenize_tx(ThreadData *td, TOKENEXTRA **t,
       set_entropy_context_b(plane, block, blk_row, blk_col,
                             plane_bsize, tx_size, arg);
   } else {
-    const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
     int bsl = b_width_log2_lookup[bsize];
     int i;