intrabc: support var-tx

Support recursive tx block partition.

On the screen content testset, 0.2% gain for keyframe encoding.

Change-Id: I623e6fbb910fef9c91617e02edf420019f67d189

av1/decoder/decodemv.c

@@ -1133,7 +1133,36 @@ static void read_intrabc_info(AV1_COMMON *const cm, MACROBLOCKD *const xd,
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
   mbmi->use_intrabc = aom_read_symbol(r, ec_ctx->intrabc_cdf, 2, ACCT_STR);
   if (mbmi->use_intrabc) {
+#if CONFIG_VAR_TX
+    const BLOCK_SIZE bsize = mbmi->sb_type;
+    const int width = block_size_wide[bsize] >> tx_size_wide_log2[0];
+    const int height = block_size_high[bsize] >> tx_size_high_log2[0];
+    int idx, idy;
+    if ((cm->tx_mode == TX_MODE_SELECT && block_signals_txsize(bsize) &&
+         !xd->lossless[mbmi->segment_id] && !mbmi->skip)) {
+      const TX_SIZE max_tx_size = max_txsize_rect_lookup[bsize];
+      const int bh = tx_size_high_unit[max_tx_size];
+      const int bw = tx_size_wide_unit[max_tx_size];
+      int init_depth =
+          (height != width) ? RECT_VARTX_DEPTH_INIT : SQR_VARTX_DEPTH_INIT;
+      mbmi->min_tx_size = TX_SIZES_ALL;
+      for (idy = 0; idy < height; idy += bh) {
+        for (idx = 0; idx < width; idx += bw) {
+          read_tx_size_vartx(cm, xd, mbmi, xd->counts, max_tx_size, init_depth,
+                             idy, idx, r);
+        }
+      }
+    } else {
+      mbmi->tx_size = read_tx_size(cm, xd, 1, !mbmi->skip, r);
+      for (idy = 0; idy < height; ++idy)
+        for (idx = 0; idx < width; ++idx)
+          mbmi->inter_tx_size[idy >> 1][idx >> 1] = mbmi->tx_size;
+      mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size);
+      set_txfm_ctxs(mbmi->tx_size, xd->n8_w, xd->n8_h, mbmi->skip, xd);
+    }
+#else
     mbmi->tx_size = read_tx_size(cm, xd, 1, !mbmi->skip, r);
+#endif  // CONFIG_VAR_TX
     mbmi->mode = mbmi->uv_mode = UV_DC_PRED;
     mbmi->interp_filters = av1_broadcast_interp_filter(BILINEAR);
 
@@ -1149,19 +1178,8 @@ static void read_intrabc_info(AV1_COMMON *const cm, MACROBLOCKD *const xd,
 
     int_mv dv_ref = nearestmv.as_int == 0 ? nearmv : nearestmv;
     if (dv_ref.as_int == 0) av1_find_ref_dv(&dv_ref, mi_row, mi_col);
-    const BLOCK_SIZE bsize = mbmi->sb_type;
     xd->corrupted |=
         !assign_dv(cm, xd, &mbmi->mv[0], &dv_ref, mi_row, mi_col, bsize, r);
-#if CONFIG_VAR_TX
-    // TODO(aconverse@google.com): Evaluate allowing VAR TX on intrabc blocks
-    const int width = block_size_wide[bsize] >> tx_size_wide_log2[0];
-    const int height = block_size_high[bsize] >> tx_size_high_log2[0];
-    int idx, idy;
-    for (idy = 0; idy < height; ++idy)
-      for (idx = 0; idx < width; ++idx)
-        mbmi->inter_tx_size[idy >> 1][idx >> 1] = mbmi->tx_size;
-    mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size);
-#endif  // CONFIG_VAR_TX
 #if CONFIG_EXT_TX && !CONFIG_TXK_SEL
     av1_read_tx_type(cm, xd, r);
 #endif  // CONFIG_EXT_TX && !CONFIG_TXK_SEL
@@ -1232,6 +1250,14 @@ static void read_intra_frame_mode_info(AV1_COMMON *const cm,
   mbmi->ref_frame[1] = NONE_FRAME;
 
 #if CONFIG_INTRABC
+#if CONFIG_VAR_TX
+  if (cm->allow_screen_content_tools) {
+    xd->above_txfm_context =
+        cm->above_txfm_context + (mi_col << TX_UNIT_WIDE_LOG2);
+    xd->left_txfm_context = xd->left_txfm_context_buffer +
+                            ((mi_row & MAX_MIB_MASK) << TX_UNIT_HIGH_LOG2);
+  }
+#endif  // CONFIG_VAR_TX
   if (av1_allow_intrabc(bsize, cm)) {
     read_intrabc_info(cm, xd, mi_row, mi_col, r);
     if (is_intrabc_block(mbmi)) return;
@@ -1239,6 +1265,10 @@ static void read_intra_frame_mode_info(AV1_COMMON *const cm,
 #endif  // CONFIG_INTRABC
 
   mbmi->tx_size = read_tx_size(cm, xd, 0, 1, r);
+#if CONFIG_INTRABC && CONFIG_VAR_TX
+  if (cm->allow_screen_content_tools)
+    set_txfm_ctxs(mbmi->tx_size, xd->n8_w, xd->n8_h, mbmi->skip, xd);
+#endif  // CONFIG_INTRABC && CONFIG_VAR_TX
 
 #if CONFIG_CB4X4
   (void)i;

av1/encoder/bitstream.c

@@ -2036,7 +2036,31 @@ static void write_intrabc_info(AV1_COMMON *cm, MACROBLOCKD *xd,
   if (use_intrabc) {
     assert(mbmi->mode == DC_PRED);
     assert(mbmi->uv_mode == UV_DC_PRED);
-    if (enable_tx_size && !mbmi->skip) write_selected_tx_size(cm, xd, w);
+    if ((enable_tx_size && !mbmi->skip)) {
+#if CONFIG_VAR_TX
+      const BLOCK_SIZE bsize = mbmi->sb_type;
+      const TX_SIZE max_tx_size = get_vartx_max_txsize(mbmi, bsize, 0);
+      const int bh = tx_size_high_unit[max_tx_size];
+      const int bw = tx_size_wide_unit[max_tx_size];
+      const int width = block_size_wide[bsize] >> tx_size_wide_log2[0];
+      const int height = block_size_high[bsize] >> tx_size_wide_log2[0];
+      int init_depth =
+          (height != width) ? RECT_VARTX_DEPTH_INIT : SQR_VARTX_DEPTH_INIT;
+      int idx, idy;
+      for (idy = 0; idy < height; idy += bh) {
+        for (idx = 0; idx < width; idx += bw) {
+          write_tx_size_vartx(cm, xd, mbmi, max_tx_size, init_depth, idy, idx,
+                              w);
+        }
+      }
+#else
+      write_selected_tx_size(cm, xd, w);
+#endif
+    } else {
+#if CONFIG_VAR_TX
+      set_txfm_ctxs(mbmi->tx_size, xd->n8_w, xd->n8_h, mbmi->skip, xd);
+#endif  // CONFIG_VAR_TX
+    }
     int_mv dv_ref = mbmi_ext->ref_mvs[INTRA_FRAME][0];
     av1_encode_dv(w, &mbmi->mv[0].as_mv, &dv_ref.as_mv, &ec_ctx->ndvc);
 #if CONFIG_EXT_TX && !CONFIG_TXK_SEL
@@ -2124,6 +2148,10 @@ static void write_mb_modes_kf(AV1_COMMON *cm, MACROBLOCKD *xd,
 #endif  // CONFIG_INTRABC
 
   if (enable_tx_size) write_selected_tx_size(cm, xd, w);
+#if CONFIG_INTRABC && CONFIG_VAR_TX
+  if (cm->allow_screen_content_tools)
+    set_txfm_ctxs(mbmi->tx_size, xd->n8_w, xd->n8_h, mbmi->skip, xd);
+#endif  // CONFIG_INTRABC && CONFIG_VAR_TX
 
   if (bsize >= BLOCK_8X8 || unify_bsize) {
     write_intra_mode_kf(cm, ec_ctx, mi, above_mi, left_mi, 0, mbmi->mode, w);
@@ -2306,6 +2334,14 @@ static void write_mbmi_b(AV1_COMP *cpi, const TileInfo *const tile,
                  cm->mi_rows, cm->mi_cols);
 
   if (frame_is_intra_only(cm)) {
+#if CONFIG_INTRABC && CONFIG_VAR_TX
+    if (cm->allow_screen_content_tools) {
+      xd->above_txfm_context =
+          cm->above_txfm_context + (mi_col << TX_UNIT_WIDE_LOG2);
+      xd->left_txfm_context = xd->left_txfm_context_buffer +
+                              ((mi_row & MAX_MIB_MASK) << TX_UNIT_HIGH_LOG2);
+    }
+#endif  // CONFIG_INTRABC && CONFIG_VAR_TX
     write_mb_modes_kf(cm, xd,
 #if CONFIG_INTRABC
                       cpi->td.mb.mbmi_ext,

av1/encoder/encodeframe.c

@@ -4867,16 +4867,6 @@ static void tx_partition_count_update(const AV1_COMMON *const cm, MACROBLOCK *x,
   int init_depth =
       (mi_height != mi_width) ? RECT_VARTX_DEPTH_INIT : SQR_VARTX_DEPTH_INIT;
 
-#if CONFIG_INTRABC
-  // Intrabc doesn't support var-tx yet. So no need to update tx partition
-  // info., except for the split count (otherwise common->tx_mode may be
-  // modified, causing mismatch).
-  if (is_intrabc_block(&x->e_mbd.mi[0]->mbmi)) {
-    if (x->e_mbd.mi[0]->mbmi.tx_size != max_tx_size) ++x->txb_split_count;
-    return;
-  }
-#endif  // CONFIG_INTRABC
-
   xd->above_txfm_context =
       cm->above_txfm_context + (mi_col << TX_UNIT_WIDE_LOG2);
   xd->left_txfm_context = xd->left_txfm_context_buffer +

av1/encoder/rdopt.c

@@ -9744,24 +9744,27 @@ static int64_t rd_pick_intrabc_mode_sb(const AV1_COMP *cpi, MACROBLOCK *x,
     const int rate_mode = x->intrabc_cost[1];
     RD_STATS rd_stats, rd_stats_uv;
     av1_subtract_plane(x, bsize, 0);
+#if CONFIG_VAR_TX
+    if (cm->tx_mode == TX_MODE_SELECT && !xd->lossless[mbmi->segment_id]) {
+      select_tx_type_yrd(cpi, x, &rd_stats, bsize, INT64_MAX);
+    } else {
+      int idx, idy;
+      super_block_yrd(cpi, x, &rd_stats, bsize, INT64_MAX);
+      for (idy = 0; idy < xd->n8_h; ++idy)
+        for (idx = 0; idx < xd->n8_w; ++idx)
+          mbmi->inter_tx_size[idy][idx] = mbmi->tx_size;
+      memset(x->blk_skip[0], rd_stats.skip,
+             sizeof(uint8_t) * xd->n8_h * xd->n8_w * 4);
+    }
+#else
     super_block_yrd(cpi, x, &rd_stats, bsize, INT64_MAX);
+#endif  // CONFIG_VAR_TX
     super_block_uvrd(cpi, x, &rd_stats_uv, bsize, INT64_MAX);
     av1_merge_rd_stats(&rd_stats, &rd_stats_uv);
 #if CONFIG_RD_DEBUG
     mbmi->rd_stats = rd_stats;
 #endif
 
-#if CONFIG_VAR_TX
-    // TODO(aconverse@google.com): Evaluate allowing VAR TX on intrabc blocks
-    const int width = block_size_wide[bsize] >> tx_size_wide_log2[0];
-    const int height = block_size_high[bsize] >> tx_size_high_log2[0];
-    int idx, idy;
-    for (idy = 0; idy < height; ++idy)
-      for (idx = 0; idx < width; ++idx)
-        mbmi->inter_tx_size[idy >> 1][idx >> 1] = mbmi->tx_size;
-    mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size);
-#endif  // CONFIG_VAR_TX
-
     const aom_prob skip_prob = av1_get_skip_prob(cm, xd);
 
     RD_STATS rdc_noskip;