Add forward probability update for lv_map

The feature is implemented in the following two functions.
av1_write_txb_probs
av1_read_txb_probs

Change-Id: I0b646e17ec54d7a10a77a6853439217091455af1

av1/decoder/decodeframe.c

@@ -2438,7 +2438,7 @@ static void setup_bool_decoder(const uint8_t *data, const uint8_t *data_end,
                        "Failed to allocate bool decoder %d", 1);
 }
 
-#if !CONFIG_PVQ && !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
+#if !CONFIG_PVQ && !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET) && !CONFIG_LV_MAP
 static void read_coef_probs_common(av1_coeff_probs_model *coef_probs,
                                    aom_reader *r) {
   int i, j, k, l, m;
@@ -4521,11 +4521,16 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
   if (cm->tx_mode == TX_MODE_SELECT) read_tx_size_probs(fc, &r);
 #endif
 
+#if CONFIG_LV_MAP
+  av1_read_txb_probs(fc, cm->tx_mode, &r);
+#else  // CONFIG_LV_MAP
 #if !CONFIG_PVQ
 #if !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
   read_coef_probs(fc, cm->tx_mode, &r);
-#endif
+#endif  // !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
 #endif  // !CONFIG_PVQ
+#endif  // CONFIG_LV_MAP
+
 #if CONFIG_VAR_TX
   for (k = 0; k < TXFM_PARTITION_CONTEXTS; ++k)
     av1_diff_update_prob(&r, &fc->txfm_partition_prob[k], ACCT_STR);

av1/decoder/decodetxb.c

@@ -13,6 +13,7 @@
 #include "av1/common/idct.h"
 #include "av1/common/txb_common.h"
 #include "av1/decoder/decodetxb.h"
+#include "av1/decoder/dsubexp.h"
 
 #define ACCT_STR __func__
 
@@ -235,3 +236,43 @@ uint8_t av1_read_coeffs_txb_facade(AV1_COMMON *cm, MACROBLOCKD *xd,
   av1_set_contexts(xd, pd, plane, tx_size, cul_level, col, row);
   return cul_level;
 }
+
+static void read_txb_probs(FRAME_CONTEXT *fc, const TX_SIZE tx_size,
+                           aom_reader *r) {
+  int plane, ctx, level;
+
+  if (aom_read_bit(r, ACCT_STR) == 0) return;
+
+  for (ctx = 0; ctx < TXB_SKIP_CONTEXTS; ++ctx)
+    av1_diff_update_prob(r, &fc->txb_skip[tx_size][ctx], ACCT_STR);
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane)
+    for (ctx = 0; ctx < SIG_COEF_CONTEXTS; ++ctx)
+      av1_diff_update_prob(r, &fc->nz_map[tx_size][plane][ctx], ACCT_STR);
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane)
+    for (ctx = 0; ctx < EOB_COEF_CONTEXTS; ++ctx)
+      av1_diff_update_prob(r, &fc->eob_flag[tx_size][plane][ctx], ACCT_STR);
+
+  for (level = 0; level < NUM_BASE_LEVELS; ++level)
+    for (plane = 0; plane < PLANE_TYPES; ++plane)
+      for (ctx = 0; ctx < COEFF_BASE_CONTEXTS; ++ctx)
+        av1_diff_update_prob(r, &fc->coeff_base[tx_size][plane][level][ctx],
+                             ACCT_STR);
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane)
+    for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx)
+      av1_diff_update_prob(r, &fc->coeff_lps[tx_size][plane][ctx], ACCT_STR);
+}
+
+void av1_read_txb_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;
+  int ctx, plane;
+  for (plane = 0; plane < PLANE_TYPES; ++plane)
+    for (ctx = 0; ctx < DC_SIGN_CONTEXTS; ++ctx)
+      av1_diff_update_prob(r, &fc->dc_sign[plane][ctx], ACCT_STR);
+
+  for (tx_size = TX_4X4; tx_size <= max_tx_size; ++tx_size)
+    read_txb_probs(fc, tx_size, r);
+}

av1/decoder/decodetxb.h

@@ -27,4 +27,5 @@ uint8_t av1_read_coeffs_txb_facade(AV1_COMMON *cm, MACROBLOCKD *xd,
                                    aom_reader *r, int row, int col, int block,
                                    int plane, tran_low_t *tcoeffs,
                                    int16_t *max_scan_line, int *eob);
+void av1_read_txb_probs(FRAME_CONTEXT *fc, TX_MODE tx_mode, aom_reader *r);
 #endif  //  DECODETXB_H_

av1/encoder/bitstream.c

@@ -2849,6 +2849,7 @@ static void write_modes(AV1_COMP *const cpi, const TileInfo *const tile,
 #endif
 }
 
+#if !CONFIG_LV_MAP
 #if !CONFIG_PVQ && !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
 static void build_tree_distribution(AV1_COMP *cpi, TX_SIZE tx_size,
                                     av1_coeff_stats *coef_branch_ct,
@@ -3325,8 +3326,9 @@ static void update_coef_probs(AV1_COMP *cpi, aom_writer *w) {
   }
 #endif  // CONFIG_SUBFRAME_PROB_UPDATE
 }
-#endif
+#endif  // !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
 #endif  // !CONFIG_EC_ADAPT
+#endif  // !CONFIG_LV_MAP
 
 #if CONFIG_LOOP_RESTORATION
 static void encode_restoration_mode(AV1_COMMON *cm,
@@ -4601,11 +4603,16 @@ static uint32_t write_compressed_header(AV1_COMP *cpi, uint8_t *data) {
 #if !CONFIG_EC_ADAPT
   update_txfm_probs(cm, header_bc, counts);
 #endif
+#if CONFIG_LV_MAP
+  av1_write_txb_probs(cpi, header_bc);
+#else
 #if !CONFIG_PVQ
 #if !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
   update_coef_probs(cpi, header_bc);
 #endif  // !(CONFIG_EC_ADAPT && CONFIG_NEW_TOKENSET)
 #endif  // CONFIG_PVQ
+#endif  // CONFIG_LV_MAP
+
 #if CONFIG_VAR_TX
   update_txfm_partition_probs(cm, header_bc, counts, probwt);
 #endif

av1/encoder/encodetxb.c

@@ -14,6 +14,7 @@
 #include "av1/common/pred_common.h"
 #include "av1/encoder/cost.h"
 #include "av1/encoder/encodetxb.h"
+#include "av1/encoder/subexp.h"
 #include "av1/encoder/tokenize.h"
 
 void av1_alloc_txb_buf(AV1_COMP *cpi) {
@@ -565,3 +566,133 @@ void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,
     assert(0);
   }
 }
+
+static void find_new_prob(unsigned int *branch_cnt, aom_prob *oldp,
+                          int *savings, int *update, aom_writer *const bc) {
+  const aom_prob upd = DIFF_UPDATE_PROB;
+  int u = 0;
+  aom_prob newp = get_binary_prob(branch_cnt[0], branch_cnt[1]);
+  int s = av1_prob_diff_update_savings_search(branch_cnt, *oldp, &newp, upd, 1);
+
+  if (s > 0 && newp != *oldp) u = 1;
+
+  if (u)
+    *savings += s - (int)(av1_cost_zero(upd));  // TODO(jingning): 1?
+  else
+    *savings -= (int)(av1_cost_zero(upd));
+
+  if (update) {
+    ++update[u];
+    return;
+  }
+
+  aom_write(bc, u, upd);
+  if (u) {
+    /* send/use new probability */
+    av1_write_prob_diff_update(bc, newp, *oldp);
+    *oldp = newp;
+  }
+}
+
+static void write_txb_probs(aom_writer *const bc, AV1_COMP *cpi,
+                            TX_SIZE tx_size) {
+  FRAME_CONTEXT *fc = cpi->common.fc;
+  FRAME_COUNTS *counts = cpi->td.counts;
+  int savings = 0;
+  int update[2] = { 0, 0 };
+  int plane, ctx, level;
+
+  for (ctx = 0; ctx < TXB_SKIP_CONTEXTS; ++ctx) {
+    find_new_prob(counts->txb_skip[tx_size][ctx], &fc->txb_skip[tx_size][ctx],
+                  &savings, update, bc);
+  }
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane) {
+    for (ctx = 0; ctx < SIG_COEF_CONTEXTS; ++ctx) {
+      find_new_prob(counts->nz_map[tx_size][plane][ctx],
+                    &fc->nz_map[tx_size][plane][ctx], &savings, update, bc);
+    }
+  }
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane) {
+    for (ctx = 0; ctx < EOB_COEF_CONTEXTS; ++ctx) {
+      find_new_prob(counts->eob_flag[tx_size][plane][ctx],
+                    &fc->eob_flag[tx_size][plane][ctx], &savings, update, bc);
+    }
+  }
+
+  for (level = 0; level < NUM_BASE_LEVELS; ++level) {
+    for (plane = 0; plane < PLANE_TYPES; ++plane) {
+      for (ctx = 0; ctx < COEFF_BASE_CONTEXTS; ++ctx) {
+        find_new_prob(counts->coeff_base[tx_size][plane][level][ctx],
+                      &fc->coeff_base[tx_size][plane][level][ctx], &savings,
+                      update, bc);
+      }
+    }
+  }
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane) {
+    for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx) {
+      find_new_prob(counts->coeff_lps[tx_size][plane][ctx],
+                    &fc->coeff_lps[tx_size][plane][ctx], &savings, update, bc);
+    }
+  }
+
+  // Decide if to update the model for this tx_size
+  if (update[1] == 0 || savings < 0) {
+    aom_write_bit(bc, 0);
+    return;
+  }
+  aom_write_bit(bc, 1);
+
+  for (ctx = 0; ctx < TXB_SKIP_CONTEXTS; ++ctx) {
+    find_new_prob(counts->txb_skip[tx_size][ctx], &fc->txb_skip[tx_size][ctx],
+                  &savings, NULL, bc);
+  }
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane) {
+    for (ctx = 0; ctx < SIG_COEF_CONTEXTS; ++ctx) {
+      find_new_prob(counts->nz_map[tx_size][plane][ctx],
+                    &fc->nz_map[tx_size][plane][ctx], &savings, NULL, bc);
+    }
+  }
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane) {
+    for (ctx = 0; ctx < EOB_COEF_CONTEXTS; ++ctx) {
+      find_new_prob(counts->eob_flag[tx_size][plane][ctx],
+                    &fc->eob_flag[tx_size][plane][ctx], &savings, NULL, bc);
+    }
+  }
+
+  for (level = 0; level < NUM_BASE_LEVELS; ++level) {
+    for (plane = 0; plane < PLANE_TYPES; ++plane) {
+      for (ctx = 0; ctx < COEFF_BASE_CONTEXTS; ++ctx) {
+        find_new_prob(counts->coeff_base[tx_size][plane][level][ctx],
+                      &fc->coeff_base[tx_size][plane][level][ctx], &savings,
+                      NULL, bc);
+      }
+    }
+  }
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane) {
+    for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx) {
+      find_new_prob(counts->coeff_lps[tx_size][plane][ctx],
+                    &fc->coeff_lps[tx_size][plane][ctx], &savings, NULL, bc);
+    }
+  }
+}
+
+void av1_write_txb_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;
+  int ctx, plane;
+
+  for (plane = 0; plane < PLANE_TYPES; ++plane)
+    for (ctx = 0; ctx < DC_SIGN_CONTEXTS; ++ctx)
+      av1_cond_prob_diff_update(w, &cpi->common.fc->dc_sign[plane][ctx],
+                                cpi->td.counts->dc_sign[plane][ctx], 1);
+
+  for (tx_size = TX_4X4; tx_size <= max_tx_size; ++tx_size)
+    write_txb_probs(w, cpi, tx_size);
+}

av1/encoder/encodetxb.h

@@ -33,6 +33,7 @@ void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
 void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,
                             RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
                             const int mi_row, const int mi_col);
+void av1_write_txb_probs(AV1_COMP *cpi, aom_writer *w);
 #ifdef __cplusplus
 }
 #endif