Add chroma 2x2 block experiment

Support the use of 2x2 prediction and transform block size for
chroma components. This experiment builds on top of cb4x4. It
brings back the loss in performance gains due to 4x4 chroma
component constrain. The coding gains for lowres are up from 1.5%
to 2.5%.

Change-Id: I8e461e230c2d920f6df1af9d935ebe9dfa3820d9

av1/common/blockd.c

@@ -130,7 +130,7 @@ void av1_foreach_transformed_block_in_plane(
   // 4x4=0, 8x8=2, 16x16=4, 32x32=6, 64x64=8
   // transform size varies per plane, look it up in a common way.
   const TX_SIZE tx_size = plane ? get_uv_tx_size(mbmi, pd) : mbmi->tx_size;
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
   const BLOCK_SIZE plane_bsize =
       AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
 #else

av1/common/blockd.h

@@ -924,7 +924,7 @@ static INLINE TX_SIZE get_uv_tx_size(const MB_MODE_INFO *mbmi,
 
   uv_txsize = uv_txsize_lookup[mbmi->sb_type][mbmi->tx_size][pd->subsampling_x]
                               [pd->subsampling_y];
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
   uv_txsize = AOMMAX(uv_txsize, TX_4X4);
 #endif
   assert(uv_txsize != TX_INVALID);

av1/common/onyxc_int.h

@@ -585,8 +585,10 @@ static INLINE void set_plane_n4(MACROBLOCKD *const xd, int bw, int bh) {
     xd->plane[i].width = (bw * MI_SIZE) >> xd->plane[i].subsampling_x;
     xd->plane[i].height = (bh * MI_SIZE) >> xd->plane[i].subsampling_y;
 
+#if !CONFIG_CHROMA_2X2
     xd->plane[i].width = AOMMAX(xd->plane[i].width, 4);
     xd->plane[i].height = AOMMAX(xd->plane[i].height, 4);
+#endif
   }
 }
 
@@ -728,6 +730,9 @@ static INLINE void update_partition_context(MACROBLOCKD *xd, int mi_row,
 
 #if CONFIG_CB4X4
 static INLINE int is_chroma_reference(const int mi_row, const int mi_col) {
+#if CONFIG_CHROMA_2X2
+  return 1;
+#endif
   return !((mi_row & 0x01) || (mi_col & 0x01));
 }
 #endif

av1/common/reconintra.c

@@ -1895,7 +1895,7 @@ static void predict_square_intra_block(const MACROBLOCKD *xd, int wpx, int hpx,
   const int mi_col = -xd->mb_to_left_edge >> (3 + MI_SIZE_LOG2);
   const int txwpx = tx_size_wide[tx_size];
   const int txhpx = tx_size_high[tx_size];
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
   const int xr_chr_offset = (plane && bsize < BLOCK_8X8) ? 2 : 0;
   const int yd_chr_offset = (plane && bsize < BLOCK_8X8) ? 2 : 0;
 #else
@@ -1918,7 +1918,7 @@ static void predict_square_intra_block(const MACROBLOCKD *xd, int wpx, int hpx,
   const PARTITION_TYPE partition = xd->mi[0]->mbmi.partition;
 #endif
 
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
   // force 4x4 chroma component block size.
   if (plane && bsize < BLOCK_8X8) bsize = BLOCK_8X8;
 #endif

av1/decoder/decodeframe.c

@@ -1659,8 +1659,12 @@ static void decode_token_and_recon_block(AV1Decoder *const pbi,
       const int stepr = tx_size_high_unit[tx_size];
       const int stepc = tx_size_wide_unit[tx_size];
 #if CONFIG_CB4X4
+#if CONFIG_CHROMA_2X2
+      const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
+#else
       const BLOCK_SIZE plane_bsize =
           AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
+#endif  // CONFIG_CHROMA_2X2
 #else
       const BLOCK_SIZE plane_bsize =
           get_plane_block_size(AOMMAX(BLOCK_8X8, bsize), pd);
@@ -1744,8 +1748,12 @@ static void decode_token_and_recon_block(AV1Decoder *const pbi,
       for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
         const struct macroblockd_plane *const pd = &xd->plane[plane];
 #if CONFIG_CB4X4
+#if CONFIG_CHROMA_2X2
+        const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
+#else
         const BLOCK_SIZE plane_bsize =
             AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
+#endif  // CONFIG_CHROMA_2X2
 #else
         const BLOCK_SIZE plane_bsize =
             get_plane_block_size(AOMMAX(BLOCK_8X8, bsize), pd);

av1/encoder/encodemb.c

@@ -887,7 +887,7 @@ void av1_encode_sb(AV1_COMMON *cm, MACROBLOCK *x, BLOCK_SIZE bsize,
   if (x->skip) return;
 
   for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
     if (bsize < BLOCK_8X8 && plane && !is_chroma_reference(mi_row, mi_col))
       continue;
     if (plane) bsize = AOMMAX(bsize, BLOCK_8X8);

av1/encoder/rdopt.c

@@ -974,7 +974,7 @@ static void model_rd_for_sb(const AV1_COMP *const cpi, BLOCK_SIZE bsize,
   for (plane = plane_from; plane <= plane_to; ++plane) {
     struct macroblock_plane *const p = &x->plane[plane];
     struct macroblockd_plane *const pd = &xd->plane[plane];
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
     const BLOCK_SIZE bs = AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
 #else
     const BLOCK_SIZE bs = get_plane_block_size(bsize, pd);
@@ -3721,7 +3721,7 @@ static int super_block_uvrd(const AV1_COMP *const cpi, MACROBLOCK *x,
 
   if (ref_best_rd < 0) is_cost_valid = 0;
 
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
   if (x->skip_chroma_rd) return is_cost_valid;
   bsize = AOMMAX(BLOCK_8X8, bsize);
 #endif
@@ -4347,7 +4347,7 @@ static int inter_block_uvrd(const AV1_COMP *cpi, MACROBLOCK *x,
 
   av1_init_rd_stats(rd_stats);
 
-#if CONFIG_CB4X4
+#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
   if (x->skip_chroma_rd) return is_cost_valid;
   bsize = AOMMAX(BLOCK_8X8, bsize);
 #endif
@@ -4845,6 +4845,10 @@ static void choose_intra_uv_mode(const AV1_COMP *const cpi, MACROBLOCK *const x,
   // appropriate speed flag is set.
   (void)ctx;
 #if CONFIG_CB4X4
+#if CONFIG_CHROMA_2X2
+  rd_pick_intra_sbuv_mode(cpi, x, rate_uv, rate_uv_tokenonly, dist_uv, skip_uv,
+                          bsize, max_tx_size);
+#else
   max_tx_size = AOMMAX(max_tx_size, TX_4X4);
   if (x->skip_chroma_rd) {
     *rate_uv = 0;
@@ -4854,10 +4858,13 @@ static void choose_intra_uv_mode(const AV1_COMP *const cpi, MACROBLOCK *const x,
     *mode_uv = DC_PRED;
     return;
   }
-#endif
-
   rd_pick_intra_sbuv_mode(cpi, x, rate_uv, rate_uv_tokenonly, dist_uv, skip_uv,
                           bsize < BLOCK_8X8 ? BLOCK_8X8 : bsize, max_tx_size);
+#endif  // CONFIG_CHROMA_2X2
+#else
+  rd_pick_intra_sbuv_mode(cpi, x, rate_uv, rate_uv_tokenonly, dist_uv, skip_uv,
+                          bsize < BLOCK_8X8 ? BLOCK_8X8 : bsize, max_tx_size);
+#endif  // CONFIG_CB4X4
   *mode_uv = x->e_mbd.mi[0]->mbmi.uv_mode;
 }
 
@@ -9059,7 +9066,9 @@ void av1_rd_pick_intra_mode_sb(const AV1_COMP *cpi, MACROBLOCK *x,
                                    [pd[1].subsampling_x][pd[1].subsampling_y];
 
 #if CONFIG_CB4X4
+#if !CONFIG_CHROMA_2X2
   max_uv_tx_size = AOMMAX(max_uv_tx_size, TX_4X4);
+#endif
   if (!x->skip_chroma_rd)
     rd_pick_intra_sbuv_mode(cpi, x, &rate_uv, &rate_uv_tokenonly, &dist_uv,
                             &uv_skip, bsize, max_uv_tx_size);