Rework base range entropy coding in level map system

Replace the truncated geometric distribution model with the grouped
leaves structure for more efficient probability modeling.
Each group has its own Geometric distribution

This give us 0.2% gain on lowres

Change-Id: If5c73dd429bd5183a8aa81042f8f56937b1d8a6a

aom_dsp/prob.h

@@ -48,6 +48,8 @@ typedef uint16_t aom_cdf_prob;
 
 #define LV_MAP_PROB 1
 
+#define BR_NODE 1
+
 #define aom_prob_half ((aom_prob)128)
 
 typedef int8_t aom_tree_index;

av1/common/entropy.h

@@ -78,7 +78,8 @@ extern "C" {
 #define LEVEL_CONTEXTS (BR_TMP_OFFSET * BR_REF_CAT)
 
 #define NUM_BASE_LEVELS 2
-#define COEFF_BASE_RANGE (15 - NUM_BASE_LEVELS)
+#define COEFF_BASE_RANGE (16 - NUM_BASE_LEVELS)
+#define BASE_RANGE_SETS 3
 
 #define COEFF_CONTEXT_BITS 6
 #define COEFF_CONTEXT_MASK ((1 << COEFF_CONTEXT_BITS) - 1)

av1/common/entropymode.h

@@ -175,6 +175,9 @@ typedef struct frame_contexts {
   aom_prob coeff_base[TX_SIZES][PLANE_TYPES][NUM_BASE_LEVELS]
                      [COEFF_BASE_CONTEXTS];
   aom_prob coeff_lps[TX_SIZES][PLANE_TYPES][LEVEL_CONTEXTS];
+#if BR_NODE
+  aom_prob coeff_br[TX_SIZES][PLANE_TYPES][BASE_RANGE_SETS][LEVEL_CONTEXTS];
+#endif
 
 #if LV_MAP_PROB
   aom_cdf_prob txb_skip_cdf[TX_SIZES][TXB_SKIP_CONTEXTS][CDF_SIZE(2)];
@@ -187,6 +190,10 @@ typedef struct frame_contexts {
                              [COEFF_BASE_CONTEXTS][CDF_SIZE(2)];
   aom_cdf_prob coeff_lps_cdf[TX_SIZES][PLANE_TYPES][LEVEL_CONTEXTS]
                             [CDF_SIZE(2)];
+#if BR_NODE
+  aom_cdf_prob coeff_br_cdf[TX_SIZES][PLANE_TYPES][BASE_RANGE_SETS]
+                           [LEVEL_CONTEXTS][CDF_SIZE(2)];
+#endif
 #endif  // LV_MAP_PROB
 #endif
 
@@ -435,6 +442,8 @@ typedef struct FRAME_COUNTS {
   unsigned int coeff_base[TX_SIZES][PLANE_TYPES][NUM_BASE_LEVELS]
                          [COEFF_BASE_CONTEXTS][2];
   unsigned int coeff_lps[TX_SIZES][PLANE_TYPES][LEVEL_CONTEXTS][2];
+  unsigned int coeff_br[TX_SIZES][PLANE_TYPES][BASE_RANGE_SETS][LEVEL_CONTEXTS]
+                       [2];
 #endif  // CONFIG_LV_MAP
 
   av1_blockz_count_model blockz_count[TX_SIZES][PLANE_TYPES];

av1/common/txb_common.c

@@ -160,10 +160,20 @@ void av1_init_txb_probs(FRAME_CONTEXT *fc) {
         fc->coeff_lps_cdf[tx_size][plane][ctx][1] = AOM_ICDF(32768);
         fc->coeff_lps_cdf[tx_size][plane][ctx][2] = 0;
       }
+#if BR_NODE
+      for (int br = 0; br < BASE_RANGE_SETS; ++br) {
+        for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx) {
+          fc->coeff_br_cdf[tx_size][plane][br][ctx][0] = AOM_ICDF(
+              128 * (aom_cdf_prob)fc->coeff_br[tx_size][plane][br][ctx]);
+          fc->coeff_br_cdf[tx_size][plane][br][ctx][1] = AOM_ICDF(32768);
+          fc->coeff_br_cdf[tx_size][plane][br][ctx][2] = 0;
+        }
+      }
+#endif  // BR_NODE
     }
   }
 }
-#endif
+#endif  // LV_MAP_PROB
 
 void av1_adapt_txb_probs(AV1_COMMON *cm, unsigned int count_sat,
                          unsigned int update_factor) {
@@ -211,11 +221,23 @@ void av1_adapt_txb_probs(AV1_COMMON *cm, unsigned int count_sat,
   }
 
   for (tx_size = 0; tx_size < TX_SIZES; ++tx_size) {
-    for (plane = 0; plane < PLANE_TYPES; ++plane)
-      for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx)
+    for (plane = 0; plane < PLANE_TYPES; ++plane) {
+      for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx) {
         fc->coeff_lps[tx_size][plane][ctx] = merge_probs(
             pre_fc->coeff_lps[tx_size][plane][ctx],
             counts->coeff_lps[tx_size][plane][ctx], count_sat, update_factor);
+      }
+#if BR_NODE
+      for (int br = 0; br < BASE_RANGE_SETS; ++br) {
+        for (ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx) {
+          fc->coeff_br[tx_size][plane][br][ctx] =
+              merge_probs(pre_fc->coeff_br[tx_size][plane][br][ctx],
+                          counts->coeff_br[tx_size][plane][br][ctx], count_sat,
+                          update_factor);
+        }
+      }
+#endif  // BR_NODE
+    }
   }
 }
 

av1/common/txb_common.h

@@ -173,10 +173,22 @@ static int br_ref_offset[BR_CONTEXT_POSITION_NUM][2] = {
   /* clang-format on*/
 };
 
-static int br_level_map[9] = {
+static const int br_level_map[9] = {
   0, 0, 1, 1, 2, 2, 3, 3, 3,
 };
 
+static const int coeff_to_br_index[COEFF_BASE_RANGE] = {
+  0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2,
+};
+
+static const int br_index_to_coeff[BASE_RANGE_SETS] = {
+  0, 2, 6,
+};
+
+static const int br_extra_bits[BASE_RANGE_SETS] = {
+  1, 2, 3,
+};
+
 #define BR_MAG_OFFSET 1
 // TODO(angiebird): optimize this function by using a table to map from
 // count/mag to ctx
@@ -214,7 +226,7 @@ static INLINE int get_br_ctx_from_count_mag(int row, int col, int count,
     offset = 0;
   else if (mag <= 3)
     offset = 1;
-  else if (mag <= 6)
+  else if (mag <= 5)
     offset = 2;
   else
     offset = 3;

av1/decoder/decodetxb.c

@@ -216,6 +216,52 @@ uint8_t av1_read_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
 
     if (ec_ctx->coeff_lps[txs_ctx][plane_type][ctx] == 0) exit(0);
 
+#if BR_NODE
+    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+#if LV_MAP_PROB
+      if (aom_read_symbol(r,
+                          cm->fc->coeff_br_cdf[txs_ctx][plane_type][idx][ctx],
+                          2, ACCT_STR))
+#else   // LV_MAP_PROB
+      if (aom_read(r, cm->fc->coeff_br[txs_ctx][plane_type][idx][ctx],
+                   ACCT_STR))
+#endif  // LV_MAP_PROB
+      {
+        int extra_bits = (1 << br_extra_bits[idx]) - 1;
+        //        int br_offset = aom_read_literal(r, extra_bits, ACCT_STR);
+        int br_offset = 0;
+        int tok;
+        if (counts) ++counts->coeff_br[txs_ctx][plane_type][ctx][idx][1];
+        for (tok = 0; tok < extra_bits; ++tok) {
+#if LV_MAP_PROB
+          if (aom_read_symbol(r,
+                              cm->fc->coeff_lps_cdf[txs_ctx][plane_type][ctx],
+                              2, ACCT_STR))
+#else
+          if (aom_read(r, cm->fc->coeff_lps[txs_ctx][plane_type][ctx],
+                       ACCT_STR))
+#endif
+          {
+            br_offset = tok;
+            if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][1];
+            break;
+          }
+          if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][0];
+        }
+        if (tok == extra_bits) br_offset = extra_bits;
+
+        int br_base = br_index_to_coeff[idx];
+
+        *v = NUM_BASE_LEVELS + 1 + br_base + br_offset;
+        cul_level += *v;
+        if (sign) *v = -(*v);
+        break;
+      }
+      if (counts) ++counts->coeff_br[txs_ctx][plane_type][ctx][idx][0];
+    }
+
+    if (idx < BASE_RANGE_SETS) continue;
+#else
     for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
 #if LV_MAP_PROB
       if (aom_read_symbol(r, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2,
@@ -234,6 +280,7 @@ uint8_t av1_read_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
       if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][0];
     }
     if (idx < COEFF_BASE_RANGE) continue;
+#endif
 
     // decode 0-th order Golomb code
     *v = read_golomb(xd, r) + COEFF_BASE_RANGE + 1 + NUM_BASE_LEVELS;

av1/encoder/block.h

@@ -75,6 +75,9 @@ typedef struct {
   int dc_sign_cost[DC_SIGN_CONTEXTS][2];
   int base_cost[NUM_BASE_LEVELS][COEFF_BASE_CONTEXTS][2];
   int lps_cost[LEVEL_CONTEXTS][2];
+#if BR_NODE
+  int br_cost[BASE_RANGE_SETS][LEVEL_CONTEXTS][2];
+#endif
 } LV_MAP_COEFF_COST;
 
 typedef struct {

av1/encoder/encodetxb.c

@@ -231,6 +231,41 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
 
     // level is above 1.
     ctx = get_br_ctx(tcoeff, scan[c], bwl, height);
+
+#if BR_NODE
+    int base_range = level - 1 - NUM_BASE_LEVELS;
+    int br_set_idx = 0;
+    int br_base = 0;
+    int br_offset = 0;
+
+    if (base_range >= COEFF_BASE_RANGE)
+      br_set_idx = BASE_RANGE_SETS;
+    else
+      br_set_idx = coeff_to_br_index[base_range];
+
+    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+      aom_write_symbol(w, idx == br_set_idx,
+                       cm->fc->coeff_br_cdf[txs_ctx][plane_type][idx][ctx], 2);
+      if (idx == br_set_idx) {
+        br_base = br_index_to_coeff[br_set_idx];
+        br_offset = base_range - br_base;
+        int extra_bits = (1 << br_extra_bits[idx]) - 1;
+        for (int tok = 0; tok < extra_bits; ++tok) {
+          if (tok == br_offset) {
+            aom_write_symbol(
+                w, 1, cm->fc->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2);
+            break;
+          }
+          aom_write_symbol(w, 0,
+                           cm->fc->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2);
+        }
+        //        aom_write_literal(w, br_offset, br_extra_bits[idx]);
+        break;
+      }
+    }
+
+    if (br_set_idx < BASE_RANGE_SETS) continue;
+#else  // BR_NODE
     for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
       if (level == (idx + 1 + NUM_BASE_LEVELS)) {
 #if LV_MAP_PROB
@@ -249,6 +284,7 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
 #endif
     }
     if (idx < COEFF_BASE_RANGE) continue;
+#endif  // BR_NODE
 
     // use 0-th order Golomb code to handle the residual level.
     write_golomb(w, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
@@ -428,9 +464,35 @@ int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
       if (level > NUM_BASE_LEVELS) {
         int idx;
         int ctx;
-
         ctx = get_br_ctx(qcoeff, scan[c], bwl, height);
+#if BR_NODE
+        int base_range = level - 1 - NUM_BASE_LEVELS;
+        int br_set_idx = base_range < COEFF_BASE_RANGE
+                             ? coeff_to_br_index[base_range]
+                             : BASE_RANGE_SETS;
+
+        for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+          if (br_set_idx == idx) {
+            int br_base = br_index_to_coeff[br_set_idx];
+            int br_offset = base_range - br_base;
+            int extra_bits = (1 << br_extra_bits[idx]) - 1;
+            cost += coeff_costs->br_cost[idx][ctx][1];
+            for (int tok = 0; tok < extra_bits; ++tok) {
+              if (tok == br_offset) {
+                cost += coeff_costs->lps_cost[ctx][1];
+                break;
+              }
+              cost += coeff_costs->lps_cost[ctx][0];
+            }
+            //            cost += extra_bits * av1_cost_bit(128, 1);
+            break;
+          }
+          cost += coeff_costs->br_cost[idx][ctx][0];
+        }
+        if (idx < BASE_RANGE_SETS) continue;
 
+        idx = COEFF_BASE_RANGE;
+#else
         for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
           if (level == (idx + 1 + NUM_BASE_LEVELS)) {
             cost += coeff_costs->lps_cost[ctx][1];
@@ -438,7 +500,7 @@ int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
           }
           cost += coeff_costs->lps_cost[ctx][0];
         }
-
+#endif
         if (idx >= COEFF_BASE_RANGE) {
           // residual cost
           int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
@@ -1681,6 +1743,45 @@ void av1_update_and_record_txb_context(int plane, int block, int blk_row,
 
     // level is above 1.
     ctx = get_br_ctx(tcoeff, scan[c], bwl, height);
+
+#if BR_NODE
+    int base_range = level - 1 - NUM_BASE_LEVELS;
+    int br_set_idx = base_range < COEFF_BASE_RANGE
+                         ? coeff_to_br_index[base_range]
+                         : BASE_RANGE_SETS;
+
+    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+      if (idx == br_set_idx) {
+        int br_base = br_index_to_coeff[br_set_idx];
+        int br_offset = base_range - br_base;
+        ++td->counts->coeff_br[txsize_ctx][plane_type][idx][ctx][1];
+#if LV_MAP_PROB
+        update_cdf(ec_ctx->coeff_br_cdf[txsize_ctx][plane_type][idx][ctx], 1,
+                   2);
+#endif
+        int extra_bits = (1 << br_extra_bits[idx]) - 1;
+        for (int tok = 0; tok < extra_bits; ++tok) {
+          if (br_offset == tok) {
+            ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][1];
+#if LV_MAP_PROB
+            update_cdf(ec_ctx->coeff_lps_cdf[txsize_ctx][plane_type][ctx], 1,
+                       2);
+#endif
+            break;
+          }
+          ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][0];
+#if LV_MAP_PROB
+          update_cdf(ec_ctx->coeff_lps_cdf[txsize_ctx][plane_type][ctx], 0, 2);
+#endif
+        }
+        break;
+      }
+      ++td->counts->coeff_br[txsize_ctx][plane_type][idx][ctx][0];
+#if LV_MAP_PROB
+      update_cdf(ec_ctx->coeff_br_cdf[txsize_ctx][plane_type][idx][ctx], 0, 2);
+#endif
+    }
+#else  // BR_NODE
     for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
       if (level == (idx + 1 + NUM_BASE_LEVELS)) {
         ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][1];
@@ -1695,7 +1796,7 @@ void av1_update_and_record_txb_context(int plane, int block, int blk_row,
 #endif
     }
     if (idx < COEFF_BASE_RANGE) continue;
-
+#endif  // BR_NODE
     // use 0-th order Golomb code to handle the residual level.
   }
 

av1/encoder/encodetxb.h

@@ -64,6 +64,9 @@ typedef struct TxbProbs {
   const aom_prob *coeff_lps;
   const aom_prob *eob_flag;
   const aom_prob *txb_skip;
+#if BR_NODE
+  const aom_prob *coeff_br;
+#endif
 } TxbProbs;
 
 void av1_alloc_txb_buf(AV1_COMP *cpi);

av1/encoder/rd.c

@@ -500,10 +500,18 @@ void av1_fill_coeff_costs(MACROBLOCK *x, FRAME_CONTEXT *fc) {
               pcost->base_cost[layer][ctx],
               fc->coeff_base_cdf[tx_size][plane][layer][ctx], NULL);
 
+#if BR_NODE
+      for (int br = 0; br < BASE_RANGE_SETS; ++br)
+        for (int ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx)
+          av1_cost_tokens_from_cdf(pcost->br_cost[br][ctx],
+                                   fc->coeff_br_cdf[tx_size][plane][br][ctx],
+                                   NULL);
+#endif  // BR_NODE
+
       for (int ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx)
         av1_cost_tokens_from_cdf(pcost->lps_cost[ctx],
                                  fc->coeff_lps_cdf[tx_size][plane][ctx], NULL);
-#else
+#else   // LV_MAP_PROB
       for (int ctx = 0; ctx < TXB_SKIP_CONTEXTS; ++ctx)
         get_rate_cost(fc->txb_skip[tx_size][ctx], pcost->txb_skip_cost[ctx]);
 
@@ -523,11 +531,11 @@ void av1_fill_coeff_costs(MACROBLOCK *x, FRAME_CONTEXT *fc) {
 
       for (int ctx = 0; ctx < LEVEL_CONTEXTS; ++ctx)
         get_rate_cost(fc->coeff_lps[tx_size][plane][ctx], pcost->lps_cost[ctx]);
-#endif
+#endif  // LV_MAP_PROB
     }
   }
 }
-#endif
+#endif  // CONFIG_LV_MAP
 
 void av1_fill_token_costs_from_cdf(av1_coeff_cost *cost,
                                    coeff_cdf_model (*cdf)[PLANE_TYPES]) {

tools/aom_entropy_optimizer.c

@@ -862,7 +862,24 @@ int main(int argc, const char **argv) {
                      "static const aom_cdf_prob "
                      "default_coeff_lps_cdf[TX_SIZES][PLANE_TYPES][LEVEL_"
                      "CONTEXTS][CDF_SIZE(2)]");
-#endif
+
+#if BR_NODE
+  cts_each_dim[0] = TX_SIZES;
+  cts_each_dim[1] = PLANE_TYPES;
+  cts_each_dim[2] = BASE_RANGE_SETS;
+  cts_each_dim[3] = LEVEL_CONTEXTS;
+  cts_each_dim[4] = 2;
+  optimize_entropy_table(&fc.coeff_br[0][0][0][0][0], probsfile, 5,
+                         cts_each_dim, NULL, 1,
+                         "static const aom_prob "
+                         "default_coeff_br[TX_SIZES][PLANE_TYPES][BASE_RANGE_"
+                         "SETS][LEVEL_CONTEXTS]");
+  optimize_cdf_table(&fc.coeff_br[0][0][0][0][0], probsfile, 5, cts_each_dim,
+                     "static const aom_cdf_prob "
+                     "default_coeff_br_cdf[TX_SIZES][PLANE_TYPES][BASE_RANGE_"
+                     "SETS][LEVEL_CONTEXTS][CDF_SIZE(2)]");
+#endif  // BR_NODE
+#endif  // CONFIG_LV_MAP
 
   fclose(statsfile);
   fclose(logfile);