Fix inefficient coding of filter_intra modes

Include a few RDO fixes that properly account for the extra flag
cost at the right place. Also entropy coding of the flag is made
contexted on tx_size and based on cdf framework.

Change-Id: I52f98ace7e253ccc08917a4a7b517515d1d58a98
diff --git a/av1/common/entropy.c b/av1/common/entropy.c
index df5a08b..d640463 100644
--- a/av1/common/entropy.c
+++ b/av1/common/entropy.c
@@ -1736,6 +1736,10 @@
   AVERAGE_TILE_CDFS(palette_uv_size_cdf)
   AVERAGE_TILE_CDFS(palette_y_color_index_cdf)
   AVERAGE_TILE_CDFS(palette_uv_color_index_cdf)
+#if CONFIG_FILTER_INTRA
+  AVERAGE_TILE_CDFS(filter_intra_cdfs)
+  AVERAGE_TILE_CDFS(filter_intra_mode_cdf)
+#endif
 #if CONFIG_MRC_TX
   AVERAGE_TILE_CDFS(mrc_mask_intra_cdf)
 #endif  // CONFIG_MRC_TX
diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index e942777..f079711 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -1672,11 +1672,16 @@
 #endif  // CONFIG_JNT_COMP
 
 #if CONFIG_FILTER_INTRA
-static const aom_prob default_filter_intra_probs[2] = { 103, 231 };
-
 static const aom_cdf_prob default_filter_intra_mode_cdf[2][CDF_SIZE(
-    FILTER_INTRA_MODES)] = { { AOM_CDF6(24348, 27690, 30315, 30864, 31701) },
-                             { AOM_CDF6(31977, 32117, 32300, 32484, 32719) } };
+    FILTER_INTRA_MODES)] = { { AOM_CDF6(22207, 23158, 24144, 24278, 30434) },
+                             { AOM_CDF6(32768, 32768, 32768, 32768, 32768) } };
+
+static const aom_cdf_prob default_filter_intra_cdfs[TX_SIZES_ALL][CDF_SIZE(2)] =
+    { { AOM_CDF2(10985) }, { AOM_CDF2(10985) }, { AOM_CDF2(16645) },
+      { AOM_CDF2(27378) }, { AOM_CDF2(10985) }, { AOM_CDF2(10985) },
+      { AOM_CDF2(15723) }, { AOM_CDF2(12373) }, { AOM_CDF2(27199) },
+      { AOM_CDF2(24217) }, { AOM_CDF2(32767) }, { AOM_CDF2(32767) },
+      { AOM_CDF2(32767) }, { AOM_CDF2(32767) } };
 #endif  // CONFIG_FILTER_INTRA
 
 // FIXME(someone) need real defaults here
@@ -2253,6 +2258,62 @@
 // It is possible to re-train this model and bring back the 0.14% loss in CIF
 // set key frame coding. This reduction in context model does not change the
 // key frame coding stats for mid and high resolution sets.
+#if CONFIG_FILTER_INTRA
+const aom_cdf_prob
+    default_kf_y_mode_cdf[KF_MODE_CONTEXTS][KF_MODE_CONTEXTS][CDF_SIZE(
+        INTRA_MODES)] = {
+      { { AOM_CDF13(13234, 14775, 17115, 18040, 18783, 19420, 20510, 22129,
+                    23183, 28738, 30120, 32138) },
+        { AOM_CDF13(8983, 14623, 16290, 17124, 17864, 18817, 19593, 20876,
+                    22359, 27820, 29791, 31566) },
+        { AOM_CDF13(7091, 8084, 17897, 18490, 19057, 19428, 20811, 22624, 23265,
+                    28288, 29341, 31870) },
+        { AOM_CDF13(11191, 12808, 14120, 16182, 16785, 17440, 18159, 20280,
+                    22697, 28431, 30235, 32276) },
+        { AOM_CDF13(8208, 9510, 11986, 12851, 15212, 16786, 19400, 22224, 23146,
+                    28889, 30200, 32375) } },
+      { { AOM_CDF13(6308, 15986, 17454, 18110, 18739, 19867, 20479, 21575,
+                    22972, 28087, 30042, 31489) },
+        { AOM_CDF13(3549, 21993, 22593, 22968, 23262, 24052, 24280, 24856,
+                    26026, 29057, 30818, 31543) },
+        { AOM_CDF13(4371, 9956, 16063, 16680, 17207, 17870, 18692, 20142, 21261,
+                    26613, 28301, 30433) },
+        { AOM_CDF13(6445, 12764, 13699, 15338, 15922, 16891, 17304, 18868,
+                    22816, 28105, 30472, 31907) },
+        { AOM_CDF13(4300, 11014, 12466, 13258, 15028, 17584, 19170, 21448,
+                    22945, 28207, 30041, 31659) } },
+      { { AOM_CDF13(9111, 10159, 16955, 17625, 18268, 18703, 20078, 22004,
+                    22761, 28166, 29334, 31990) },
+        { AOM_CDF13(7107, 11104, 15591, 16340, 17066, 17802, 18721, 20303,
+                    21481, 26882, 28699, 30978) },
+        { AOM_CDF13(4546, 4935, 22442, 22717, 22960, 23087, 24171, 25671, 25939,
+                    29333, 29866, 32023) },
+        { AOM_CDF13(8332, 9555, 12646, 14689, 15340, 15873, 16872, 19939, 21942,
+                    27812, 29508, 31923) },
+        { AOM_CDF13(6413, 7233, 13108, 13895, 15332, 16187, 19121, 22694, 23365,
+                    28639, 29686, 32187) } },
+      { { AOM_CDF13(9584, 11586, 12990, 15322, 15927, 16732, 17406, 19225,
+                    22484, 28555, 30321, 32279) },
+        { AOM_CDF13(5907, 11662, 12625, 14955, 15491, 16403, 16865, 18074,
+                    23261, 28508, 30584, 32057) },
+        { AOM_CDF13(5759, 7323, 12581, 14779, 15363, 15946, 16851, 19330, 21902,
+                    27860, 29214, 31747) },
+        { AOM_CDF13(7166, 8714, 9430, 14479, 14672, 14953, 15184, 17239, 24798,
+                    29350, 31021, 32371) },
+        { AOM_CDF13(6318, 8140, 9595, 12354, 13754, 15324, 16681, 19701, 22723,
+                    28616, 30226, 32279) } },
+      { { AOM_CDF13(8669, 9875, 12300, 13093, 15518, 17458, 19843, 22083, 22927,
+                    28780, 30271, 32364) },
+        { AOM_CDF13(6600, 10422, 12153, 12937, 15218, 18211, 19914, 21744,
+                    22975, 28393, 30393, 31970) },
+        { AOM_CDF13(5512, 6207, 14265, 14897, 16246, 17175, 19865, 22553, 23178,
+                    28445, 29511, 31980) },
+        { AOM_CDF13(8195, 9407, 10830, 13261, 14443, 15761, 16922, 20311, 22151,
+                    28230, 30109, 32220) },
+        { AOM_CDF13(5612, 6462, 8166, 8737, 14316, 17802, 21788, 25554, 26080,
+                    30083, 30983, 32457) } }
+    };
+#else
 const aom_cdf_prob
     default_kf_y_mode_cdf[KF_MODE_CONTEXTS][KF_MODE_CONTEXTS][CDF_SIZE(
         INTRA_MODES)] = {
@@ -2317,6 +2378,7 @@
                       29653, 30954, 32215) },
       },
     };
+#endif
 #else
 const aom_cdf_prob
     default_kf_y_mode_cdf[INTRA_MODES][INTRA_MODES][CDF_SIZE(INTRA_MODES)] = {
@@ -3182,7 +3244,7 @@
   av1_copy(fc->seg.pred_cdf, default_segment_pred_cdf);
 #endif
 #if CONFIG_FILTER_INTRA
-  av1_copy(fc->filter_intra_probs, default_filter_intra_probs);
+  av1_copy(fc->filter_intra_cdfs, default_filter_intra_cdfs);
   av1_copy(fc->filter_intra_mode_cdf, default_filter_intra_mode_cdf);
 #endif  // CONFIG_FILTER_INTRA
 #if CONFIG_LOOP_RESTORATION
@@ -3444,12 +3506,6 @@
     fc->delta_lf_prob[i] =
         mode_mv_merge_probs(pre_fc->delta_lf_prob[i], counts->delta_lf[i]);
 #endif  // CONFIG_EXT_DELTA_Q
-#if CONFIG_FILTER_INTRA
-  for (i = 0; i < PLANE_TYPES; ++i) {
-    fc->filter_intra_probs[i] = av1_mode_mv_merge_probs(
-        pre_fc->filter_intra_probs[i], counts->filter_intra[i]);
-  }
-#endif  // CONFIG_FILTER_INTRA
 }
 
 static void set_default_lf_deltas(struct loopfilter *lf) {
diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index 508f46a..1925d24 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -291,7 +291,7 @@
   int initialized;
   struct segmentation_probs seg;
 #if CONFIG_FILTER_INTRA
-  aom_prob filter_intra_probs[PLANE_TYPES];
+  aom_cdf_prob filter_intra_cdfs[TX_SIZES_ALL][CDF_SIZE(2)];
   aom_cdf_prob filter_intra_mode_cdf[PLANE_TYPES][CDF_SIZE(FILTER_INTRA_MODES)];
 #endif  // CONFIG_FILTER_INTRA
 #if CONFIG_LOOP_RESTORATION
@@ -361,7 +361,11 @@
 // Note: This structure should only contain 'unsigned int' fields, or
 // aggregates built solely from 'unsigned int' fields/elements
 #if CONFIG_ENTROPY_STATS
+#if CONFIG_KF_CTX
+  unsigned int kf_y_mode[KF_MODE_CONTEXTS][KF_MODE_CONTEXTS][INTRA_MODES];
+#else
   unsigned int kf_y_mode[INTRA_MODES][INTRA_MODES][INTRA_MODES];
+#endif
   unsigned int angle_delta[DIRECTIONAL_MODES][2 * MAX_ANGLE_DELTA + 1];
   unsigned int y_mode[BLOCK_SIZE_GROUPS][INTRA_MODES];
   unsigned int uv_mode[INTRA_MODES][UV_INTRA_MODES];
@@ -472,8 +476,8 @@
 #endif  // CONFIG_ENTROPY_STATS
   struct seg_counts seg;
 #if CONFIG_FILTER_INTRA
-  unsigned int filter_intra[PLANE_TYPES][2];
   unsigned int filter_intra_mode[PLANE_TYPES][FILTER_INTRA_MODES];
+  unsigned int filter_intra_tx[TX_SIZES_ALL][2];
 #endif  // CONFIG_FILTER_INTRA
 #if CONFIG_LPF_SB
   unsigned int lpf_reuse[LPF_REUSE_CONTEXT][2];
diff --git a/av1/decoder/decodemv.c b/av1/decoder/decodemv.c
index 4aebae4..7752e49 100644
--- a/av1/decoder/decodemv.c
+++ b/av1/decoder/decodemv.c
@@ -912,27 +912,21 @@
 }
 
 #if CONFIG_FILTER_INTRA
-static void read_filter_intra_mode_info(AV1_COMMON *const cm,
-                                        MACROBLOCKD *const xd, aom_reader *r) {
+static void read_filter_intra_mode_info(MACROBLOCKD *const xd, aom_reader *r) {
   MODE_INFO *const mi = xd->mi[0];
   MB_MODE_INFO *const mbmi = &mi->mbmi;
-  FRAME_COUNTS *counts = xd->counts;
   FILTER_INTRA_MODE_INFO *filter_intra_mode_info =
       &mbmi->filter_intra_mode_info;
 
   if (mbmi->mode == DC_PRED && mbmi->palette_mode_info.palette_size[0] == 0 &&
       av1_filter_intra_allowed_txsize(mbmi->tx_size)) {
-    filter_intra_mode_info->use_filter_intra_mode[0] =
-        aom_read(r, cm->fc->filter_intra_probs[0], ACCT_STR);
+    filter_intra_mode_info->use_filter_intra_mode[0] = aom_read_symbol(
+        r, xd->tile_ctx->filter_intra_cdfs[mbmi->tx_size], 2, ACCT_STR);
     if (filter_intra_mode_info->use_filter_intra_mode[0]) {
       filter_intra_mode_info->filter_intra_mode[0] =
           aom_read_symbol(r, xd->tile_ctx->filter_intra_mode_cdf[0],
                           FILTER_INTRA_MODES, ACCT_STR);
     }
-    if (counts) {
-      ++counts
-            ->filter_intra[0][filter_intra_mode_info->use_filter_intra_mode[0]];
-    }
   }
 }
 #endif  // CONFIG_FILTER_INTRA
@@ -1272,7 +1266,7 @@
 #if CONFIG_FILTER_INTRA
   mbmi->filter_intra_mode_info.use_filter_intra_mode[0] = 0;
   mbmi->filter_intra_mode_info.use_filter_intra_mode[1] = 0;
-  read_filter_intra_mode_info(cm, xd, r);
+  read_filter_intra_mode_info(xd, r);
 #endif  // CONFIG_FILTER_INTRA
 
 #if !CONFIG_TXK_SEL
@@ -1750,7 +1744,7 @@
 #if CONFIG_FILTER_INTRA
   mbmi->filter_intra_mode_info.use_filter_intra_mode[0] = 0;
   mbmi->filter_intra_mode_info.use_filter_intra_mode[1] = 0;
-  read_filter_intra_mode_info(cm, xd, r);
+  read_filter_intra_mode_info(xd, r);
 #endif  // CONFIG_FILTER_INTRA
 }
 
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 23b60c5..58f9422 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -959,14 +959,13 @@
 }
 
 #if CONFIG_FILTER_INTRA
-static void write_filter_intra_mode_info(const AV1_COMMON *const cm,
-                                         const MACROBLOCKD *xd,
+static void write_filter_intra_mode_info(const MACROBLOCKD *xd,
                                          const MB_MODE_INFO *const mbmi,
                                          aom_writer *w) {
   if (mbmi->mode == DC_PRED && mbmi->palette_mode_info.palette_size[0] == 0 &&
       av1_filter_intra_allowed_txsize(mbmi->tx_size)) {
-    aom_write(w, mbmi->filter_intra_mode_info.use_filter_intra_mode[0],
-              cm->fc->filter_intra_probs[0]);
+    aom_write_symbol(w, mbmi->filter_intra_mode_info.use_filter_intra_mode[0],
+                     xd->tile_ctx->filter_intra_cdfs[mbmi->tx_size], 2);
     if (mbmi->filter_intra_mode_info.use_filter_intra_mode[0]) {
       const FILTER_INTRA_MODE mode =
           mbmi->filter_intra_mode_info.filter_intra_mode[0];
@@ -1492,7 +1491,7 @@
     if (av1_allow_palette(cm->allow_screen_content_tools, bsize))
       write_palette_mode_info(cm, xd, mi, w);
 #if CONFIG_FILTER_INTRA
-    write_filter_intra_mode_info(cm, xd, mbmi, w);
+    write_filter_intra_mode_info(xd, mbmi, w);
 #endif  // CONFIG_FILTER_INTRA
   } else {
     int16_t mode_ctx;
@@ -1809,7 +1808,7 @@
   if (av1_allow_palette(cm->allow_screen_content_tools, bsize))
     write_palette_mode_info(cm, xd, mi, w);
 #if CONFIG_FILTER_INTRA
-  write_filter_intra_mode_info(cm, xd, mbmi, w);
+  write_filter_intra_mode_info(xd, mbmi, w);
 #endif  // CONFIG_FILTER_INTRA
 
 #if !CONFIG_TXK_SEL
diff --git a/av1/encoder/block.h b/av1/encoder/block.h
index bd79d1f..cc640a4 100644
--- a/av1/encoder/block.h
+++ b/av1/encoder/block.h
@@ -230,6 +230,7 @@
   int intra_uv_mode_cost[INTRA_MODES][UV_INTRA_MODES];
   int y_mode_costs[INTRA_MODES][INTRA_MODES][INTRA_MODES];
 #if CONFIG_FILTER_INTRA
+  int filter_intra_cost[TX_SIZES_ALL][2];
   int filter_intra_mode_cost[PLANE_TYPES][FILTER_INTRA_MODES];
 #endif
   int switchable_interp_costs[SWITCHABLE_FILTER_CONTEXTS][SWITCHABLE_FILTERS];
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index e03e22e..51c4371 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -4498,7 +4498,13 @@
 #if CONFIG_ENTROPY_STATS
     const PREDICTION_MODE above = av1_above_block_mode(mi, above_mi, 0);
     const PREDICTION_MODE left = av1_left_block_mode(mi, left_mi, 0);
+#if CONFIG_KF_CTX
+    int above_ctx = intra_mode_context[above];
+    int left_ctx = intra_mode_context[left];
+    ++counts->kf_y_mode[above_ctx][left_ctx][y_mode];
+#else
     ++counts->kf_y_mode[above][left][y_mode];
+#endif
 #endif  // CONFIG_ENTROPY_STATS
     if (allow_update_cdf)
       update_cdf(get_y_mode_cdf(fc, mi, above_mi, left_mi, 0), y_mode,
@@ -4516,15 +4522,18 @@
       av1_filter_intra_allowed_txsize(mbmi->tx_size)) {
     const int use_filter_intra_mode =
         mbmi->filter_intra_mode_info.use_filter_intra_mode[0];
-    ++counts->filter_intra[0][use_filter_intra_mode];
 #if CONFIG_ENTROPY_STATS
     ++counts->filter_intra_mode[0][mbmi->filter_intra_mode_info
                                        .filter_intra_mode[0]];
+    ++counts->filter_intra_tx[mbmi->tx_size][use_filter_intra_mode];
 #endif  // CONFIG_ENTROPY_STATS
-    if (allow_update_cdf)
+    if (allow_update_cdf) {
       update_cdf(fc->filter_intra_mode_cdf[0],
                  mbmi->filter_intra_mode_info.filter_intra_mode[0],
                  FILTER_INTRA_MODES);
+      update_cdf(fc->filter_intra_cdfs[mbmi->tx_size], use_filter_intra_mode,
+                 2);
+    }
   }
 #endif  // CONFIG_FILTER_INTRA
 #if CONFIG_EXT_INTRA && CONFIG_EXT_INTRA_MOD
diff --git a/av1/encoder/encoder.c b/av1/encoder/encoder.c
index b941dc8..da21286 100644
--- a/av1/encoder/encoder.c
+++ b/av1/encoder/encoder.c
@@ -918,6 +918,7 @@
 
 static void set_tile_info(AV1_COMP *cpi) {
   AV1_COMMON *const cm = &cpi->common;
+  (void)cm;
 #if CONFIG_DEPENDENT_HORZTILES
   int tile_row, tile_col, num_tiles_in_tg;
   int tg_row_start, tg_col_start;
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index 6191136..1ef57f6 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -143,6 +143,9 @@
 #if CONFIG_FILTER_INTRA
   av1_cost_tokens_from_cdf(x->filter_intra_mode_cost[0],
                            fc->filter_intra_mode_cdf[0], NULL);
+  for (i = 0; i < TX_SIZES_ALL; ++i)
+    av1_cost_tokens_from_cdf(x->filter_intra_cost[i], fc->filter_intra_cdfs[i],
+                             NULL);
 #endif
 
   for (i = 0; i < SWITCHABLE_FILTER_CONTEXTS; ++i)
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index 08f052d..02408b6 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -2878,13 +2878,13 @@
   }
 #endif  // CONFIG_EXT_INTRA
 #if CONFIG_FILTER_INTRA
-  if (mbmi->mode == DC_PRED) {
-    const aom_prob prob = cpi->common.fc->filter_intra_probs[0];
+  if (mbmi->mode == DC_PRED && av1_filter_intra_allowed_txsize(mbmi->tx_size)) {
     if (mbmi->filter_intra_mode_info.use_filter_intra_mode[0]) {
       const int mode = mbmi->filter_intra_mode_info.filter_intra_mode[0];
-      mode_cost += av1_cost_bit(prob, 1) + x->filter_intra_mode_cost[0][mode];
+      mode_cost += x->filter_intra_cost[mbmi->tx_size][1] +
+                   x->filter_intra_mode_cost[0][mode];
     } else {
-      mode_cost += av1_cost_bit(prob, 0);
+      mode_cost += x->filter_intra_cost[mbmi->tx_size][0];
     }
   }
 #endif  // CONFIG_FILTER_INTRA
@@ -3153,7 +3153,7 @@
     super_block_yrd(cpi, x, &tokenonly_rd_stats, bsize, *best_rd);
     if (tokenonly_rd_stats.rate == INT_MAX) continue;
     this_rate = tokenonly_rd_stats.rate +
-                av1_cost_bit(cpi->common.fc->filter_intra_probs[0], 1) +
+                x->filter_intra_cost[mbmi->tx_size][1] +
                 x->filter_intra_mode_cost[0][mode] + mode_cost;
     this_rd = RDCOST(x->rdmult, this_rate, tokenonly_rd_stats.dist);
 
@@ -3567,7 +3567,7 @@
     }
 #if CONFIG_FILTER_INTRA
     if (mbmi->mode == DC_PRED && av1_filter_intra_allowed_txsize(mbmi->tx_size))
-      this_rate += av1_cost_bit(cpi->common.fc->filter_intra_probs[0], 0);
+      this_rate += x->filter_intra_cost[mbmi->tx_size][0];
 #endif  // CONFIG_FILTER_INTRA
 #if CONFIG_EXT_INTRA
     if (is_directional_mode) {
@@ -9760,11 +9760,10 @@
         int64_t best_rd_tmp = INT64_MAX;
         if (rate_y != INT_MAX &&
             av1_filter_intra_allowed_txsize(best_tx_size)) {
-          best_rd_tmp = RDCOST(
-              x->rdmult,
-              rate_y + av1_cost_bit(cpi->common.fc->filter_intra_probs[0], 0) +
-                  intra_mode_cost[mbmi->mode],
-              distortion_y);
+          best_rd_tmp = RDCOST(x->rdmult,
+                               rate_y + x->filter_intra_cost[mbmi->tx_size][0] +
+                                   intra_mode_cost[mbmi->mode],
+                               distortion_y);
         }
 
         mbmi->filter_intra_mode_info.use_filter_intra_mode[0] = 1;
@@ -9777,11 +9776,10 @@
           super_block_yrd(cpi, x, &rd_stats_y_fi, bsize, best_rd);
           if (rd_stats_y_fi.rate == INT_MAX) continue;
 
-          this_rate_tmp =
-              rd_stats_y_fi.rate +
-              av1_cost_bit(cpi->common.fc->filter_intra_probs[0], 1) +
-              x->filter_intra_mode_cost[0][fi_mode] +
-              intra_mode_cost[mbmi->mode];
+          this_rate_tmp = rd_stats_y_fi.rate +
+                          x->filter_intra_cost[mbmi->tx_size][1] +
+                          x->filter_intra_mode_cost[0][fi_mode] +
+                          intra_mode_cost[mbmi->mode];
           this_rd_tmp = RDCOST(x->rdmult, this_rate_tmp, rd_stats_y_fi.dist);
 
           if (this_rd_tmp < best_rd_tmp) {
@@ -9877,10 +9875,11 @@
       }
 #endif  // CONFIG_EXT_INTRA
 #if CONFIG_FILTER_INTRA
-      if (mbmi->mode == DC_PRED) {
+      if (mbmi->mode == DC_PRED &&
+          av1_filter_intra_allowed_txsize(mbmi->tx_size)) {
         rate2 +=
-            av1_cost_bit(cm->fc->filter_intra_probs[0],
-                         mbmi->filter_intra_mode_info.use_filter_intra_mode[0]);
+            x->filter_intra_cost[mbmi->tx_size][mbmi->filter_intra_mode_info
+                                                    .use_filter_intra_mode[0]];
         if (mbmi->filter_intra_mode_info.use_filter_intra_mode[0]) {
           rate2 += x->filter_intra_mode_cost[0][mbmi->filter_intra_mode_info
                                                     .filter_intra_mode[0]];
diff --git a/tools/aom_entropy_optimizer.c b/tools/aom_entropy_optimizer.c
index 98d0b20..f2e13b2 100644
--- a/tools/aom_entropy_optimizer.c
+++ b/tools/aom_entropy_optimizer.c
@@ -317,7 +317,16 @@
 
   int cts_each_dim[10];
 
-  /* Intra mode (keyframe luma) */
+/* Intra mode (keyframe luma) */
+#if CONFIG_KF_CTX
+  cts_each_dim[0] = KF_MODE_CONTEXTS;
+  cts_each_dim[1] = KF_MODE_CONTEXTS;
+  cts_each_dim[2] = INTRA_MODES;
+  optimize_cdf_table(&fc.kf_y_mode[0][0][0], probsfile, 3, cts_each_dim,
+                     "const aom_cdf_prob\n"
+                     "default_kf_y_mode_cdf[KF_MODE_CONTEXTS][KF_MODE_CONTEXTS]"
+                     "[CDF_SIZE(INTRA_MODES)]");
+#else
   cts_each_dim[0] = INTRA_MODES;
   cts_each_dim[1] = INTRA_MODES;
   cts_each_dim[2] = INTRA_MODES;
@@ -325,6 +334,7 @@
       &fc.kf_y_mode[0][0][0], probsfile, 3, cts_each_dim,
       "const aom_cdf_prob\n"
       "default_kf_y_mode_cdf[INTRA_MODES][INTRA_MODES][CDF_SIZE(INTRA_MODES)]");
+#endif
 
   cts_each_dim[0] = DIRECTIONAL_MODES;
   cts_each_dim[1] = 2 * MAX_ANGLE_DELTA + 1;
@@ -676,19 +686,17 @@
 /* filter_intra experiment */
 #if CONFIG_FILTER_INTRA
   cts_each_dim[0] = PLANE_TYPES;
-  cts_each_dim[1] = 2;
-  optimize_entropy_table(&fc.filter_intra[0][0], probsfile, 2, cts_each_dim,
-                         NULL, 1,
-                         "static const aom_prob default_filter_intra_probs[2]");
-  optimize_cdf_table(
-      &fc.filter_intra[0][0], probsfile, 2, cts_each_dim,
-      "static const aom_cdf_prob default_filter_intra_cdf[2][CDF_SIZE(2)]");
-  cts_each_dim[0] = PLANE_TYPES;
   cts_each_dim[1] = FILTER_INTRA_MODES;
   optimize_cdf_table(
       &fc.filter_intra_mode[0][0], probsfile, 2, cts_each_dim,
       "static const aom_cdf_prob "
       "default_filter_intra_mode_cdf[2][CDF_SIZE(FILTER_INTRA_MODES)]");
+
+  cts_each_dim[0] = TX_SIZES_ALL;
+  cts_each_dim[1] = 2;
+  optimize_cdf_table(&fc.filter_intra_tx[0][0], probsfile, 2, cts_each_dim,
+                     "static const aom_cdf_prob "
+                     "default_filter_intra_cdfs[TX_SIZES_ALL][CDF_SIZE(2)]");
 #endif
 
 #if CONFIG_LV_MAP