Support potential 2x2 transform block unit

Make the codec support 2x2 tranform block unit for chroma components.

Change-Id: If428abeb55fd772bdd74da0d0449b436f60eb0e3

av1/common/common_data.h

@@ -112,6 +112,24 @@ static const BLOCK_SIZE subsize_lookup[PARTITION_TYPES][BLOCK_SIZES] = {
 };
 
 // transform block size in pixels
+#if CONFIG_CB4X4
+static const int tx_size_1d[TX_SIZES] = { 2, 4, 8, 16, 32 };
+
+static const uint8_t tx_size_1d_log2[TX_SIZES] = { 1, 2, 3, 4, 5 };
+
+static const int tx_size_1d_in_unit[TX_SIZES] = { 1, 1, 2, 4, 8 };
+
+// TODO(jingning): Temporary table during the construction.
+static const int tx_size_1d_in_unit_log2[TX_SIZES] = { 0, 0, 1, 2, 3 };
+
+static const BLOCK_SIZE txsize_to_bsize[TX_SIZES] = {
+  BLOCK_4X4,    // TODO(jingning): replace with BLOCK_2X2
+  BLOCK_4X4,    // TX_4X4
+  BLOCK_8X8,    // TX_8X8
+  BLOCK_16X16,  // TX_16X16
+  BLOCK_32X32,  // TX_32X32
+};
+#else
 static const int tx_size_1d[TX_SIZES] = { 4, 8, 16, 32 };
 
 static const uint8_t tx_size_1d_log2[TX_SIZES] = { 2, 3, 4, 5 };
@@ -121,17 +139,18 @@ static const int tx_size_1d_in_unit[TX_SIZES] = { 1, 2, 4, 8 };
 // TODO(jingning): Temporary table during the construction.
 static const int tx_size_1d_in_unit_log2[TX_SIZES] = { 0, 1, 2, 3 };
 
-static const TX_SIZE max_txsize_lookup[BLOCK_SIZES] = {
-  TX_4X4,   TX_4X4,   TX_4X4,   TX_8X8,   TX_8X8,   TX_8X8,  TX_16X16,
-  TX_16X16, TX_16X16, TX_32X32, TX_32X32, TX_32X32, TX_32X32
-};
-
 static const BLOCK_SIZE txsize_to_bsize[TX_SIZES] = {
   BLOCK_4X4,    // TX_4X4
   BLOCK_8X8,    // TX_8X8
   BLOCK_16X16,  // TX_16X16
   BLOCK_32X32,  // TX_32X32
 };
+#endif
+
+static const TX_SIZE max_txsize_lookup[BLOCK_SIZES] = {
+  TX_4X4,   TX_4X4,   TX_4X4,   TX_8X8,   TX_8X8,   TX_8X8,  TX_16X16,
+  TX_16X16, TX_16X16, TX_32X32, TX_32X32, TX_32X32, TX_32X32
+};
 
 static const TX_SIZE tx_mode_to_biggest_tx_size[TX_MODES] = {
   TX_4X4,    // ONLY_4X4

av1/common/entropy.c

@@ -1676,7 +1676,7 @@ void av1_adapt_coef_probs(AV1_COMMON *cm) {
     update_factor = COEF_MAX_UPDATE_FACTOR;
     count_sat = COEF_COUNT_SAT;
   }
-  for (t = TX_4X4; t <= TX_32X32; t++)
+  for (t = 0; t <= TX_32X32; t++)
     adapt_coef_probs(cm, t, count_sat, update_factor);
 #if CONFIG_RANS || CONFIG_DAALA_EC
   av1_coef_pareto_cdfs(cm->fc);

av1/common/entropy.h

@@ -197,6 +197,10 @@ static INLINE int get_entropy_context(TX_SIZE tx_size, const ENTROPY_CONTEXT *a,
                                       const ENTROPY_CONTEXT *l) {
   ENTROPY_CONTEXT above_ec = 0, left_ec = 0;
 
+#if CONFIG_CB4X4
+  if (tx_size == 0) assert(0);
+#endif
+
   switch (tx_size) {
     case TX_4X4:
       above_ec = a[0] != 0;

av1/common/entropymode.c

@@ -308,12 +308,22 @@ static const aom_prob default_single_ref_p[REF_CONTEXTS][SINGLE_REFS - 1] = {
 #endif  // CONFIG_EXT_REFS
 };
 
+#if CONFIG_CB4X4
+static const struct tx_probs default_tx_probs = {
+  { { 1, 3, 136, 37 }, { 1, 5, 52, 13 } },
+
+  { { 1, 20, 152 }, { 1, 15, 101 } },
+
+  { { 1, 100 }, { 1, 66 } }
+};
+#else
 static const struct tx_probs default_tx_probs = { { { 3, 136, 37 },
                                                     { 5, 52, 13 } },
 
                                                   { { 20, 152 }, { 15, 101 } },
 
                                                   { { 100 }, { 66 } } };
+#endif
 
 #if CONFIG_PALETTE
 const aom_tree_index av1_palette_size_tree[TREE_SIZE(PALETTE_SIZES)] = {
@@ -773,27 +783,27 @@ int av1_get_palette_color_context(const uint8_t *color_map, int cols, int r,
 
 void av1_tx_counts_to_branch_counts_32x32(const unsigned int *tx_count_32x32p,
                                           unsigned int (*ct_32x32p)[2]) {
-  ct_32x32p[0][0] = tx_count_32x32p[TX_4X4];
-  ct_32x32p[0][1] = tx_count_32x32p[TX_8X8] + tx_count_32x32p[TX_16X16] +
-                    tx_count_32x32p[TX_32X32];
-  ct_32x32p[1][0] = tx_count_32x32p[TX_8X8];
-  ct_32x32p[1][1] = tx_count_32x32p[TX_16X16] + tx_count_32x32p[TX_32X32];
-  ct_32x32p[2][0] = tx_count_32x32p[TX_16X16];
-  ct_32x32p[2][1] = tx_count_32x32p[TX_32X32];
+  ct_32x32p[TX_4X4][0] = tx_count_32x32p[TX_4X4];
+  ct_32x32p[TX_4X4][1] = tx_count_32x32p[TX_8X8] + tx_count_32x32p[TX_16X16] +
+                         tx_count_32x32p[TX_32X32];
+  ct_32x32p[TX_8X8][0] = tx_count_32x32p[TX_8X8];
+  ct_32x32p[TX_8X8][1] = tx_count_32x32p[TX_16X16] + tx_count_32x32p[TX_32X32];
+  ct_32x32p[TX_16X16][0] = tx_count_32x32p[TX_16X16];
+  ct_32x32p[TX_16X16][1] = tx_count_32x32p[TX_32X32];
 }
 
 void av1_tx_counts_to_branch_counts_16x16(const unsigned int *tx_count_16x16p,
                                           unsigned int (*ct_16x16p)[2]) {
-  ct_16x16p[0][0] = tx_count_16x16p[TX_4X4];
-  ct_16x16p[0][1] = tx_count_16x16p[TX_8X8] + tx_count_16x16p[TX_16X16];
-  ct_16x16p[1][0] = tx_count_16x16p[TX_8X8];
-  ct_16x16p[1][1] = tx_count_16x16p[TX_16X16];
+  ct_16x16p[TX_4X4][0] = tx_count_16x16p[TX_4X4];
+  ct_16x16p[TX_4X4][1] = tx_count_16x16p[TX_8X8] + tx_count_16x16p[TX_16X16];
+  ct_16x16p[TX_8X8][0] = tx_count_16x16p[TX_8X8];
+  ct_16x16p[TX_8X8][1] = tx_count_16x16p[TX_16X16];
 }
 
 void av1_tx_counts_to_branch_counts_8x8(const unsigned int *tx_count_8x8p,
                                         unsigned int (*ct_8x8p)[2]) {
-  ct_8x8p[0][0] = tx_count_8x8p[TX_4X4];
-  ct_8x8p[0][1] = tx_count_8x8p[TX_8X8];
+  ct_8x8p[TX_4X4][0] = tx_count_8x8p[TX_4X4];
+  ct_8x8p[TX_4X4][1] = tx_count_8x8p[TX_8X8];
 }
 
 static const aom_prob default_skip_probs[SKIP_CONTEXTS] = { 192, 128, 64 };
@@ -839,13 +849,21 @@ int av1_ext_tx_inv[TX_TYPES];
 
 static const aom_prob
     default_intra_ext_tx_prob[EXT_TX_SIZES][TX_TYPES][TX_TYPES - 1] = {
+#if CONFIG_CB4X4
+      { { 240, 85, 128 }, { 4, 1, 248 }, { 4, 1, 8 }, { 4, 248, 128 } },
+#endif
       { { 240, 85, 128 }, { 4, 1, 248 }, { 4, 1, 8 }, { 4, 248, 128 } },
       { { 244, 85, 128 }, { 8, 2, 248 }, { 8, 2, 8 }, { 8, 248, 128 } },
       { { 248, 85, 128 }, { 16, 4, 248 }, { 16, 4, 8 }, { 16, 248, 128 } },
     };
 
 static const aom_prob default_inter_ext_tx_prob[EXT_TX_SIZES][TX_TYPES - 1] = {
-  { 160, 85, 128 }, { 176, 85, 128 }, { 192, 85, 128 },
+#if CONFIG_CB4X4
+  { 160, 85, 128 },
+#endif
+  { 160, 85, 128 },
+  { 176, 85, 128 },
+  { 192, 85, 128 },
 };
 
 static void init_mode_probs(FRAME_CONTEXT *fc) {

av1/common/enums.h

@@ -84,11 +84,20 @@ typedef char PARTITION_CONTEXT;
 
 // block transform size
 typedef uint8_t TX_SIZE;
+#if CONFIG_CB4X4
+#define TX_2X2 ((TX_SIZE)0)    // 2x2 transform
+#define TX_4X4 ((TX_SIZE)1)    // 4x4 transform
+#define TX_8X8 ((TX_SIZE)2)    // 8x8 transform
+#define TX_16X16 ((TX_SIZE)3)  // 16x16 transform
+#define TX_32X32 ((TX_SIZE)4)  // 32x32 transform
+#define TX_SIZES ((TX_SIZE)5)
+#else
 #define TX_4X4 ((TX_SIZE)0)    // 4x4 transform
 #define TX_8X8 ((TX_SIZE)1)    // 8x8 transform
 #define TX_16X16 ((TX_SIZE)2)  // 16x16 transform
 #define TX_32X32 ((TX_SIZE)3)  // 32x32 transform
 #define TX_SIZES ((TX_SIZE)4)
+#endif
 
 // frame transform mode
 typedef enum {
@@ -123,7 +132,11 @@ typedef enum {
   TX_TYPES = 4,
 } TX_TYPE;
 
+#if CONFIG_CB4X4
+#define EXT_TX_SIZES 4  // number of sizes that use extended transforms
+#else
 #define EXT_TX_SIZES 3  // number of sizes that use extended transforms
+#endif
 
 typedef enum {
   AOM_LAST_FLAG = 1 << 0,

av1/common/loopfilter.c

@@ -38,6 +38,9 @@
 //
 // A loopfilter should be applied to every other 8x8 horizontally.
 static const uint64_t left_64x64_txform_mask[TX_SIZES] = {
+#if CONFIG_CB4X4
+  0xffffffffffffffffULL,  // TX_2X2
+#endif
   0xffffffffffffffffULL,  // TX_4X4
   0xffffffffffffffffULL,  // TX_8x8
   0x5555555555555555ULL,  // TX_16x16
@@ -62,6 +65,9 @@ static const uint64_t left_64x64_txform_mask[TX_SIZES] = {
 //
 // A loopfilter should be applied to every other 4 the row vertically.
 static const uint64_t above_64x64_txform_mask[TX_SIZES] = {
+#if CONFIG_CB4X4
+  0xffffffffffffffffULL,  // TX_2X2
+#endif
   0xffffffffffffffffULL,  // TX_4X4
   0xffffffffffffffffULL,  // TX_8x8
   0x00ff00ff00ff00ffULL,  // TX_16x16
@@ -140,6 +146,9 @@ static const uint64_t above_border = 0x000000ff000000ffULL;
 
 // 16 bit masks for uv transform sizes.
 static const uint16_t left_64x64_txform_mask_uv[TX_SIZES] = {
+#if CONFIG_CB4X4
+  0xffff,  // TX_2X2
+#endif
   0xffff,  // TX_4X4
   0xffff,  // TX_8x8
   0x5555,  // TX_16x16
@@ -147,6 +156,9 @@ static const uint16_t left_64x64_txform_mask_uv[TX_SIZES] = {
 };
 
 static const uint16_t above_64x64_txform_mask_uv[TX_SIZES] = {
+#if CONFIG_CB4X4
+  0xffff,  // TX_2X2
+#endif
   0xffff,  // TX_4X4
   0xffff,  // TX_8x8
   0x0f0f,  // TX_16x16

av1/common/pred_common.h

@@ -187,7 +187,13 @@ static INLINE int get_tx_size_context(const MACROBLOCKD *xd) {
 
   if (!has_above) above_ctx = left_ctx;
 
+#if CONFIG_CB4X4
+  // TODO(jingning): Temporary setup. Will rework this after the cb4x4
+  // framework is up running.
+  return (above_ctx + left_ctx) > max_tx_size + 1;
+#else
   return (above_ctx + left_ctx) > max_tx_size;
+#endif
 }
 
 static INLINE const aom_prob *get_tx_probs(TX_SIZE max_tx_size, int ctx,

av1/common/scan.c

@@ -691,6 +691,10 @@ DECLARE_ALIGNED(16, static const int16_t, av1_default_iscan_32x32[1024]) = {
 };
 
 const SCAN_ORDER av1_default_scan_orders[TX_SIZES] = {
+#if CONFIG_CB4X4
+  // TODO(jingning): use 2x2 scan order
+  { default_scan_4x4, av1_default_iscan_4x4, default_scan_4x4_neighbors },
+#endif
   { default_scan_4x4, av1_default_iscan_4x4, default_scan_4x4_neighbors },
   { default_scan_8x8, av1_default_iscan_8x8, default_scan_8x8_neighbors },
   { default_scan_16x16, av1_default_iscan_16x16, default_scan_16x16_neighbors },
@@ -698,6 +702,13 @@ const SCAN_ORDER av1_default_scan_orders[TX_SIZES] = {
 };
 
 const SCAN_ORDER av1_scan_orders[TX_SIZES][TX_TYPES] = {
+#if CONFIG_CB4X4
+  { // TX_2X2
+    { default_scan_4x4, av1_default_iscan_4x4, default_scan_4x4_neighbors },
+    { row_scan_4x4, av1_row_iscan_4x4, row_scan_4x4_neighbors },
+    { col_scan_4x4, av1_col_iscan_4x4, col_scan_4x4_neighbors },
+    { default_scan_4x4, av1_default_iscan_4x4, default_scan_4x4_neighbors } },
+#endif
   { // TX_4X4
     { default_scan_4x4, av1_default_iscan_4x4, default_scan_4x4_neighbors },
     { row_scan_4x4, av1_row_iscan_4x4, row_scan_4x4_neighbors },

av1/common/thread_common.c

@@ -489,7 +489,7 @@ void av1_accumulate_frame_counts(AV1_COMMON *cm, FRAME_COUNTS *counts,
 
   if (is_dec) {
     int n;
-    for (i = TX_4X4; i < TX_SIZES; i++)
+    for (i = 0; i < TX_SIZES; i++)
       for (j = 0; j < PLANE_TYPES; j++)
         for (k = 0; k < REF_TYPES; k++)
           for (l = 0; l < COEF_BANDS; l++)
@@ -501,7 +501,7 @@ void av1_accumulate_frame_counts(AV1_COMMON *cm, FRAME_COUNTS *counts,
                     counts->coef[i][j][k][l][m][n];
             }
   } else {
-    for (i = TX_4X4; i < TX_SIZES; i++)
+    for (i = 0; i < TX_SIZES; i++)
       for (j = 0; j < PLANE_TYPES; j++)
         for (k = 0; k < REF_TYPES; k++)
           for (l = 0; l < COEF_BANDS; l++)

av1/decoder/decodeframe.c

@@ -762,7 +762,7 @@ static void read_coef_probs_common(av1_coeff_probs_model *coef_probs,
 static void read_coef_probs(FRAME_CONTEXT *fc, TX_MODE tx_mode, aom_reader *r) {
   const TX_SIZE max_tx_size = tx_mode_to_biggest_tx_size[tx_mode];
   TX_SIZE tx_size;
-  for (tx_size = 0; tx_size <= max_tx_size; ++tx_size)
+  for (tx_size = TX_4X4; tx_size <= max_tx_size; ++tx_size)
     read_coef_probs_common(fc->coef_probs[tx_size], r);
 #if CONFIG_RANS || CONFIG_DAALA_EC
   av1_coef_pareto_cdfs(fc);

av1/decoder/decodemv.c

@@ -234,11 +234,11 @@ static TX_SIZE read_selected_tx_size(AV1_COMMON *cm, MACROBLOCKD *xd,
   FRAME_COUNTS *counts = xd->counts;
   const int ctx = get_tx_size_context(xd);
   const aom_prob *tx_probs = get_tx_probs(max_tx_size, ctx, &cm->fc->tx_probs);
-  TX_SIZE tx_size = aom_read(r, tx_probs[0], ACCT_STR) ? TX_8X8 : TX_4X4;
+  TX_SIZE tx_size = aom_read(r, tx_probs[TX_4X4], ACCT_STR) ? TX_8X8 : TX_4X4;
   if (tx_size != TX_4X4 && max_tx_size >= TX_16X16) {
-    tx_size += aom_read(r, tx_probs[1], ACCT_STR);
+    tx_size += aom_read(r, tx_probs[TX_8X8], ACCT_STR);
     if (tx_size != TX_8X8 && max_tx_size >= TX_32X32)
-      tx_size += aom_read(r, tx_probs[2], ACCT_STR);
+      tx_size += aom_read(r, tx_probs[TX_16X16], ACCT_STR);
   }
 
   if (counts) ++get_tx_counts(max_tx_size, ctx, &counts->tx)[tx_size];

av1/encoder/bitstream.c

@@ -264,11 +264,11 @@ static void write_selected_tx_size(const AV1_COMMON *cm, const MACROBLOCKD *xd,
   const TX_SIZE max_tx_size = max_txsize_lookup[bsize];
   const aom_prob *const tx_probs =
       get_tx_probs2(max_tx_size, xd, &cm->fc->tx_probs);
-  aom_write(w, tx_size != TX_4X4, tx_probs[0]);
+  aom_write(w, tx_size != TX_4X4, tx_probs[TX_4X4]);
   if (tx_size != TX_4X4 && max_tx_size >= TX_16X16) {
-    aom_write(w, tx_size != TX_8X8, tx_probs[1]);
+    aom_write(w, tx_size != TX_8X8, tx_probs[TX_8X8]);
     if (tx_size != TX_8X8 && max_tx_size >= TX_32X32)
-      aom_write(w, tx_size != TX_16X16, tx_probs[2]);
+      aom_write(w, tx_size != TX_16X16, tx_probs[TX_16X16]);
   }
 }
 
@@ -1374,7 +1374,7 @@ static void update_coef_probs(AV1_COMP *cpi, aom_writer *w) {
   const TX_MODE tx_mode = cpi->common.tx_mode;
   const TX_SIZE max_tx_size = tx_mode_to_biggest_tx_size[tx_mode];
   TX_SIZE tx_size;
-  for (tx_size = 0; tx_size <= max_tx_size; ++tx_size) {
+  for (tx_size = TX_4X4; tx_size <= max_tx_size; ++tx_size) {
     av1_coeff_stats frame_branch_ct[PLANE_TYPES];
     av1_coeff_probs_model frame_coef_probs[PLANE_TYPES];
     if (cpi->td.counts->tx.tx_totals[tx_size] <= 20 ||

av1/encoder/rdopt.c

@@ -380,6 +380,11 @@ int64_t av1_highbd_block_error_c(const tran_low_t *coeff,
  * 16th coefficient in a 4x4 block or the 64th coefficient in a 8x8 block,
  * were non-zero). */
 static const int16_t band_counts[TX_SIZES][8] = {
+#if CONFIG_CB4X4
+  {
+      1, 2, 2, 3, 0, 0, 0,
+  },
+#endif
   { 1, 2, 3, 4, 3, 16 - 13, 0 },
   { 1, 2, 3, 4, 11, 64 - 21, 0 },
   { 1, 2, 3, 4, 11, 256 - 21, 0 },
@@ -744,7 +749,7 @@ static void choose_tx_size_from_rd(const AV1_COMP *const cpi, MACROBLOCK *x,
     last_rd = INT64_MAX;
     for (n = start_tx; n >= end_tx; --n) {
       int r_tx_size = 0;
-      for (m = 0; m <= n - (n == (int)max_tx_size); ++m) {
+      for (m = TX_4X4; m <= n - (n == (int)max_tx_size); ++m) {
         if (m == n)
           r_tx_size += av1_cost_zero(tx_probs[m]);
         else
@@ -3939,6 +3944,7 @@ void av1_rd_pick_intra_mode_sb(const AV1_COMP *cpi, MACROBLOCK *x,
   }
   max_uv_tx_size = get_uv_tx_size_impl(
       xd->mi[0]->mbmi.tx_size, bsize, pd[1].subsampling_x, pd[1].subsampling_y);
+
   rd_pick_intra_sbuv_mode(cpi, x, &rate_uv, &rate_uv_tokenonly, &dist_uv,
                           &uv_skip, AOMMAX(BLOCK_8X8, bsize), max_uv_tx_size);