Support rectangular tx_type coding in var-tx

Support the transform block kernel coding for rectangular
transform block size in var-tx. This integrates txk-sel with
var-tx.

Change-Id: I9a8edd84812168f56c79b78cc9af34f6304b1d54
diff --git a/av1/common/blockd.h b/av1/common/blockd.h
index 7a629fe..638548f 100644
--- a/av1/common/blockd.h
+++ b/av1/common/blockd.h
@@ -1095,6 +1095,7 @@
                                   int block, TX_SIZE tx_size) {
   const MODE_INFO *const mi = xd->mi[0];
   const MB_MODE_INFO *const mbmi = &mi->mbmi;
+
 #if CONFIG_INTRABC && (!CONFIG_EXT_TX || CONFIG_TXK_SEL)
   // TODO(aconverse@google.com): Handle INTRABC + EXT_TX + TXK_SEL
   if (is_intrabc_block(mbmi)) return DCT_DCT;
@@ -1158,7 +1159,7 @@
   (void)tx_size;
   TX_TYPE tx_type;
   if (plane_type != PLANE_TYPE_Y || xd->lossless[mbmi->segment_id] ||
-      mbmi->tx_size >= TX_32X32) {
+      txsize_sqr_map[tx_size] >= TX_32X32) {
     tx_type = DCT_DCT;
   } else {
     tx_type = mbmi->txk_type[block];
diff --git a/av1/common/onyxc_int.h b/av1/common/onyxc_int.h
index 0462a6e..f913328 100644
--- a/av1/common/onyxc_int.h
+++ b/av1/common/onyxc_int.h
@@ -946,7 +946,6 @@
 #endif
 }
 
-#if CONFIG_VAR_TX
 // Disable array-bounds checks as the TX_SIZE enum contains values larger than
 // TX_SIZES_ALL (TX_INVALID) which make extending the array as a workaround
 // infeasible. The assert is enough for static analysis and this or other tools
@@ -962,6 +961,7 @@
 #pragma GCC diagnostic warning "-Warray-bounds"
 #endif
 
+#if CONFIG_VAR_TX
 static INLINE void set_txfm_ctx(TXFM_CONTEXT *txfm_ctx, uint8_t txs, int len) {
   int i;
   for (i = 0; i < len; ++i) txfm_ctx[i] = txs;
diff --git a/av1/decoder/decodemv.c b/av1/decoder/decodemv.c
index 0b99b42..476f8f1 100644
--- a/av1/decoder/decodemv.c
+++ b/av1/decoder/decodemv.c
@@ -917,16 +917,18 @@
                       int supertx_enabled,
 #endif
 #if CONFIG_TXK_SEL
-                      int block, int plane,
+                      int block, int plane, TX_SIZE tx_size,
 #endif
                       aom_reader *r) {
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
   const int inter_block = is_inter_block(mbmi);
+#if !CONFIG_TXK_SEL
 #if CONFIG_VAR_TX
   const TX_SIZE tx_size = inter_block ? mbmi->min_tx_size : mbmi->tx_size;
 #else
   const TX_SIZE tx_size = mbmi->tx_size;
 #endif
+#endif  // !CONFIG_TXK_SEL
 #if CONFIG_EC_ADAPT
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
 #else
diff --git a/av1/decoder/decodemv.h b/av1/decoder/decodemv.h
index ceaee1d..9538e96 100644
--- a/av1/decoder/decodemv.h
+++ b/av1/decoder/decodemv.h
@@ -37,7 +37,7 @@
                       int supertx_enabled,
 #endif
 #if CONFIG_TXK_SEL
-                      int block, int plane,
+                      int block, int plane, TX_SIZE tx_size,
 #endif
                       aom_reader *r);
 
diff --git a/av1/decoder/decodetxb.c b/av1/decoder/decodetxb.c
index 98262cd..cae88fb 100644
--- a/av1/decoder/decodetxb.c
+++ b/av1/decoder/decodetxb.c
@@ -78,7 +78,7 @@
   }
 
 #if CONFIG_TXK_SEL
-  av1_read_tx_type(cm, xd, block, plane, r);
+  av1_read_tx_type(cm, xd, block, plane, get_min_tx_size(tx_size), r);
 #endif
   TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
   const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index b620283..c665d2f 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1742,16 +1742,18 @@
                        const int supertx_enabled,
 #endif
 #if CONFIG_TXK_SEL
-                       int block, int plane,
+                       int block, int plane, TX_SIZE tx_size,
 #endif
                        aom_writer *w) {
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
   const int is_inter = is_inter_block(mbmi);
+#if !CONFIG_TXK_SEL
 #if CONFIG_VAR_TX
   const TX_SIZE tx_size = is_inter ? mbmi->min_tx_size : mbmi->tx_size;
 #else
   const TX_SIZE tx_size = mbmi->tx_size;
 #endif  // CONFIG_VAR_TX
+#endif  // !CONFIG_TXK_SEL
 #if CONFIG_EC_ADAPT
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
 #else
diff --git a/av1/encoder/bitstream.h b/av1/encoder/bitstream.h
index c75d808..5a57047 100644
--- a/av1/encoder/bitstream.h
+++ b/av1/encoder/bitstream.h
@@ -42,7 +42,7 @@
                        const int supertx_enabled,
 #endif
 #if CONFIG_TXK_SEL
-                       int block, int plane,
+                       int block, int plane, TX_SIZE tx_size,
 #endif
                        aom_writer *w);
 
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 37704e5..e3c6036 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -5649,6 +5649,7 @@
                               FRAME_COUNTS *counts) {
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
   int is_inter = is_inter_block(mbmi);
+
 #if !CONFIG_TXK_SEL
   TX_TYPE tx_type = mbmi->tx_type;
 #else
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 0101b16..853516a 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -93,7 +93,7 @@
 
   if (eob == 0) return;
 #if CONFIG_TXK_SEL
-  av1_write_tx_type(cm, xd, block, plane, w);
+  av1_write_tx_type(cm, xd, block, plane, get_min_tx_size(tx_size), w);
 #endif
 
   nz_map = cm->fc->nz_map[txs_ctx][plane_type];
@@ -1599,8 +1599,8 @@
   }
 
 #if CONFIG_TXK_SEL
-  av1_update_tx_type_count(cm, xd, block, plane, mbmi->sb_type, tx_size,
-                           td->counts);
+  av1_update_tx_type_count(cm, xd, block, plane, mbmi->sb_type,
+                           get_min_tx_size(tx_size), td->counts);
 #endif
 
   for (c = 0; c < eob; ++c) {
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index d05936a..f896745 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -2059,6 +2059,10 @@
                      TX_TYPE tx_type) {
   if (plane > 0) return 0;
 
+#if CONFIG_VAR_TX
+  tx_size = get_min_tx_size(tx_size);
+#endif
+
   const MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
   const int is_inter = is_inter_block(mbmi);
 #if CONFIG_EXT_TX
@@ -4152,6 +4156,13 @@
   MACROBLOCKD *xd = &x->e_mbd;
   const struct macroblock_plane *const p = &x->plane[plane];
   struct macroblockd_plane *const pd = &xd->plane[plane];
+
+#if CONFIG_TXK_SEL
+  av1_search_txk_type(cpi, x, plane, block, blk_row, blk_col, plane_bsize,
+                      tx_size, a, l, 0, rd_stats);
+  return;
+#endif
+
   int64_t tmp;
   tran_low_t *const dqcoeff = BLOCK_OFFSET(pd->dqcoeff, block);
   PLANE_TYPE plane_type = get_plane_type(plane);
@@ -4626,10 +4637,12 @@
                                     [mbmi->tx_type];
     }
   }
-#else   // CONFIG_EXT_TX
+#else  // CONFIG_EXT_TX
+#if !CONFIG_TXK_SEL
   if (mbmi->min_tx_size < TX_32X32 && !xd->lossless[xd->mi[0]->mbmi.segment_id])
     rd_stats->rate +=
         cpi->inter_tx_type_costs[mbmi->min_tx_size][mbmi->tx_type];
+#endif
 #endif  // CONFIG_EXT_TX
 
   if (rd_stats->skip)