Integrate TXK_SEL with EXT_INTRA

Record the best_txk_type in calc_rd_given_intra_angle()

Change-Id: I174709096d019d31567517b5aec6c2858ee68bc3

av1/encoder/encodetxb.c

diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 7edda58..c8da510 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -2471,9 +2471,8 @@
     }
 
     const int is_inter = is_inter_block(mbmi);
-    const TxSetType tx_set_type =
-        get_ext_tx_set_type(get_min_tx_size(tx_size), mbmi->sb_type, is_inter,
-                            cm->reduced_tx_set_used);
+    const TxSetType tx_set_type = get_ext_tx_set_type(
+        tx_size, mbmi->sb_type, is_inter, cm->reduced_tx_set_used);
     if (!av1_ext_tx_used[tx_set_type][tx_type]) continue;
 
     RD_STATS this_rd_stats;

av1/encoder/rdopt.c

diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index d34ca5f..eca60ed 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -3103,13 +3103,18 @@
     const AV1_COMP *const cpi, MACROBLOCK *x, BLOCK_SIZE bsize, int mode_cost,
     int64_t best_rd_in, int8_t angle_delta, int max_angle_delta, int *rate,
     RD_STATS *rd_stats, int *best_angle_delta, TX_SIZE *best_tx_size,
-    TX_TYPE *best_tx_type, int64_t *best_rd, int64_t *best_model_rd) {
+    TX_TYPE *best_tx_type, int64_t *best_rd, int64_t *best_model_rd,
+    TX_TYPE *best_txk_type) {
   int this_rate;
   RD_STATS tokenonly_rd_stats;
   int64_t this_rd, this_model_rd;
   MB_MODE_INFO *mbmi = &x->e_mbd.mi[0]->mbmi;
   assert(!is_inter_block(mbmi));
 
+#if !CONFIG_TXK_SEL
+  (void)best_txk_type;
+#endif
+
   mbmi->angle_delta[0] = angle_delta;
   this_model_rd = intra_model_yrd(cpi, x, bsize, mode_cost);
   if (*best_model_rd != INT64_MAX &&
@@ -3130,6 +3135,11 @@
   this_rd = RDCOST(x->rdmult, this_rate, tokenonly_rd_stats.dist);
 
   if (this_rd < *best_rd) {
+#if CONFIG_TXK_SEL
+    memcpy(best_txk_type, mbmi->txk_type,
+           sizeof(*best_txk_type) *
+               (MAX_SB_SQUARE / (TX_SIZE_W_MIN * TX_SIZE_H_MIN)));
+#endif
     *best_rd = this_rd;
     *best_angle_delta = mbmi->angle_delta[0];
     *best_tx_size = mbmi->tx_size;
@@ -3158,6 +3168,11 @@
   int64_t this_rd, best_rd_in, rd_cost[2 * (MAX_ANGLE_DELTA + 2)];
   TX_SIZE best_tx_size = mic->mbmi.tx_size;
   TX_TYPE best_tx_type = mbmi->tx_type;
+#if CONFIG_TXK_SEL
+  TX_TYPE best_txk_type[MAX_SB_SQUARE / (TX_SIZE_W_MIN * TX_SIZE_H_MIN)];
+#else
+  TX_TYPE *best_txk_type = NULL;
+#endif
 
   for (i = 0; i < 2 * (MAX_ANGLE_DELTA + 2); ++i) rd_cost[i] = INT64_MAX;
 
@@ -3169,7 +3184,7 @@
       this_rd = calc_rd_given_intra_angle(
           cpi, x, bsize, mode_cost, best_rd_in, (1 - 2 * i) * angle_delta,
           MAX_ANGLE_DELTA, rate, rd_stats, &best_angle_delta, &best_tx_size,
-          &best_tx_type, &best_rd, best_model_rd);
+          &best_tx_type, &best_rd, best_model_rd, best_txk_type);
       rd_cost[2 * angle_delta + i] = this_rd;
       if (first_try && this_rd == INT64_MAX) return best_rd;
       first_try = 0;
@@ -3193,7 +3208,7 @@
         calc_rd_given_intra_angle(
             cpi, x, bsize, mode_cost, best_rd, (1 - 2 * i) * angle_delta,
             MAX_ANGLE_DELTA, rate, rd_stats, &best_angle_delta, &best_tx_size,
-            &best_tx_type, &best_rd, best_model_rd);
+            &best_tx_type, &best_rd, best_model_rd, best_txk_type);
       }
     }
   }
@@ -3201,6 +3216,11 @@
   mbmi->tx_size = best_tx_size;
   mbmi->angle_delta[0] = best_angle_delta;
   mbmi->tx_type = best_tx_type;
+#if CONFIG_TXK_SEL
+  memcpy(mbmi->txk_type, best_txk_type,
+         sizeof(*best_txk_type) *
+             (MAX_SB_SQUARE / (TX_SIZE_W_MIN * TX_SIZE_H_MIN)));
+#endif
   return best_rd;
 }