Draft of av1_update_txb_context

This draft version only pass compiling check, it's not working yet.

The following goal is to use new coding system when doing bitstream
packing but keep old coding system in RD loop.

Change-Id: I224a1581d1cc5c67d73e71558fb77d9faf9c2470

av1/common/blockd.c

@@ -156,6 +156,18 @@ void av1_foreach_transformed_block_in_plane(
   }
 }
 
+#if CONFIG_LV_MAP
+void av1_foreach_transformed_block(const MACROBLOCKD *const xd,
+                                   BLOCK_SIZE bsize,
+                                   foreach_transformed_block_visitor visit,
+                                   void *arg) {
+  int plane;
+
+  for (plane = 0; plane < MAX_MB_PLANE; ++plane)
+    av1_foreach_transformed_block_in_plane(xd, bsize, plane, visit, arg);
+}
+#endif
+
 #if CONFIG_DAALA_DIST
 void av1_foreach_8x8_transformed_block_in_plane(
     const MACROBLOCKD *const xd, BLOCK_SIZE bsize, int plane,

av1/common/blockd.h

@@ -976,6 +976,13 @@ void av1_foreach_transformed_block_in_plane(
     const MACROBLOCKD *const xd, BLOCK_SIZE bsize, int plane,
     foreach_transformed_block_visitor visit, void *arg);
 
+#if CONFIG_LV_MAP
+void av1_foreach_transformed_block(const MACROBLOCKD *const xd,
+                                   BLOCK_SIZE bsize,
+                                   foreach_transformed_block_visitor visit,
+                                   void *arg);
+#endif
+
 #if CONFIG_DAALA_DIST
 void av1_foreach_8x8_transformed_block_in_plane(
     const MACROBLOCKD *const xd, BLOCK_SIZE bsize, int plane,

av1/common/common_data.h

@@ -639,6 +639,12 @@ static const TX_SIZE txsize_vert_map[TX_SIZES_ALL] = {
   TX_8X8,    // TX_32X8
 };
 
+#if CONFIG_CB4X4
+#define TX_SIZE_W_MIN 2
+#else
+#define TX_SIZE_W_MIN 4
+#endif
+
 // Transform block width in pixels
 static const int tx_size_wide[TX_SIZES_ALL] = {
 #if CONFIG_CB4X4
@@ -651,6 +657,12 @@ static const int tx_size_wide[TX_SIZES_ALL] = {
   4,  8, 8,  16, 16, 32, 4, 16, 8, 32
 };
 
+#if CONFIG_CB4X4
+#define TX_SIZE_H_MIN 2
+#else
+#define TX_SIZE_H_MIN 4
+#endif
+
 // Transform block height in pixels
 static const int tx_size_high[TX_SIZES_ALL] = {
 #if CONFIG_CB4X4

av1/common/txb_common.h

@@ -307,4 +307,70 @@ static INLINE void set_dc_sign(int *cul_level, tran_low_t v) {
   else if (v > 0)
     *cul_level += 2 << COEFF_CONTEXT_BITS;
 }
+
+static INLINE int get_dc_sign_ctx(int dc_sign) {
+  int dc_sign_ctx = 0;
+  if (dc_sign < 0)
+    dc_sign_ctx = 1;
+  else if (dc_sign > 0)
+    dc_sign_ctx = 2;
+
+  return dc_sign_ctx;
+}
+
+static INLINE int get_txb_skip_context(BLOCK_SIZE bsize, TX_SIZE tx_size,
+                                       int plane, const ENTROPY_CONTEXT *a,
+                                       const ENTROPY_CONTEXT *l, int *dc_sign) {
+  const int tx_size_in_blocks = 1 << tx_size;
+  int ctx_offset = (plane == 0) ? 0 : 7;
+  int k;
+
+  if (bsize > txsize_to_bsize[tx_size]) ctx_offset += 3;
+
+  *dc_sign = 0;
+  for (k = 0; k < tx_size_in_blocks; ++k) {
+    int sign = a[k] >> COEFF_CONTEXT_BITS;
+    if (sign == 1)
+      --*dc_sign;
+    else if (sign == 2)
+      ++*dc_sign;
+    else if (sign != 0)
+      exit(0);
+
+    sign = l[k] >> 6;
+    if (sign == 1)
+      --*dc_sign;
+    else if (sign == 2)
+      ++*dc_sign;
+    else if (sign != 0)
+      exit(0);
+  }
+
+  if (plane == 0) {
+    int top = 0;
+    int left = 0;
+    for (k = 0; k < tx_size_in_blocks; ++k) {
+      top = AOMMAX(top, (a[k] & COEFF_CONTEXT_MASK));
+      left = AOMMAX(left, (l[k] & COEFF_CONTEXT_MASK));
+    }
+    top = AOMMIN(top, 255);
+    left = AOMMIN(left, 255);
+
+    if (bsize == txsize_to_bsize[tx_size])
+      return 0;
+    else if (top == 0 && left == 0)
+      return 1;
+    else if (top == 0 || left == 0)
+      return 2 + (AOMMAX(top, left) > 3);
+    else if (AOMMAX(top, left) <= 3)
+      return 4;
+    else if (AOMMIN(top, left) <= 3)
+      return 5;
+    else
+      return 6;
+  } else {
+    int ctx_base = get_entropy_context(tx_size, a, l);
+    return ctx_offset + ctx_base;
+  }
+}
 #endif  // AV1_COMMON_TXB_COMMON_H_

av1/encoder/block.h

@@ -69,6 +69,11 @@ typedef struct {
   int16_t mode_context[MODE_CTX_REF_FRAMES];
 #if CONFIG_LV_MAP
   tran_low_t *tcoeff[MAX_MB_PLANE];
+  uint16_t eobs[MAX_MB_PLANE][MAX_SB_SQUARE / (TX_SIZE_W_MIN * TX_SIZE_H_MIN)];
+  uint8_t txb_skip_ctx[MAX_MB_PLANE]
+                      [MAX_SB_SQUARE / (TX_SIZE_W_MIN * TX_SIZE_H_MIN)];
+  int dc_sign_ctx[MAX_MB_PLANE]
+                 [MAX_SB_SQUARE / (TX_SIZE_W_MIN * TX_SIZE_H_MIN)];
 #endif
 #if CONFIG_REF_MV
   uint8_t ref_mv_count[MODE_CTX_REF_FRAMES];

av1/encoder/encodeframe.c

@@ -49,6 +49,9 @@
 #include "av1/encoder/encodeframe.h"
 #include "av1/encoder/encodemb.h"
 #include "av1/encoder/encodemv.h"
+#if CONFIG_LV_MAP
+#include "av1/encoder/encodetxb.h"
+#endif
 #include "av1/encoder/ethread.h"
 #include "av1/encoder/extend.h"
 #include "av1/encoder/rd.h"
@@ -5690,7 +5693,11 @@ static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td,
 #if CONFIG_VAR_TX
     mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size);
 #endif
+#if CONFIG_LV_MAP
+    av1_update_txb_context(cpi, td, dry_run, block_size, rate, mi_row, mi_col);
+#else   // CONFIG_LV_MAP
     av1_tokenize_sb(cpi, td, t, dry_run, block_size, rate, mi_row, mi_col);
+#endif  // CONFIG_LV_MAP
   } else {
     int ref;
     const int is_compound = has_second_ref(mbmi);
@@ -5748,7 +5755,11 @@ static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td,
     av1_tokenize_sb_vartx(cpi, td, t, dry_run, mi_row, mi_col, block_size,
                           rate);
 #else
+#if CONFIG_LV_MAP
+    av1_update_txb_context(cpi, td, dry_run, block_size, rate, mi_row, mi_col);
+#else   // CONFIG_LV_MAP
     av1_tokenize_sb(cpi, td, t, dry_run, block_size, rate, mi_row, mi_col);
+#endif  // CONFIG_LV_MAP
 #endif
   }
 

av1/encoder/encodetxb.c

@@ -10,9 +10,12 @@
  */
 
 #include "av1/common/scan.h"
+#include "av1/common/blockd.h"
+#include "av1/common/pred_common.h"
 #include "av1/encoder/cost.h"
 #include "av1/encoder/encoder.h"
 #include "av1/encoder/encodetxb.h"
+#include "av1/encoder/tokenize.h"
 
 void av1_alloc_txb_buf(AV1_COMP *cpi) {
   AV1_COMMON *cm = &cpi->common;
@@ -352,3 +355,191 @@ int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
 
   return cost;
 }
+
+typedef struct TxbParams {
+  const AV1_COMP *cpi;
+  ThreadData *td;
+  int rate;
+} TxbParams;
+
+static void update_txb_context(int plane, int block, int blk_row, int blk_col,
+                               BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+                               void *arg) {
+  TxbParams *const args = arg;
+  ThreadData *const td = args->td;
+  MACROBLOCK *const x = &td->mb;
+  MACROBLOCKD *const xd = &x->e_mbd;
+  struct macroblock_plane *p = &x->plane[plane];
+  struct macroblockd_plane *pd = &xd->plane[plane];
+  (void)plane_bsize;
+  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
+                   blk_row);
+}
+
+static void update_and_record_txb_context(int plane, int block, int blk_row,
+                                          int blk_col, BLOCK_SIZE plane_bsize,
+                                          TX_SIZE tx_size, void *arg) {
+  TxbParams *const args = arg;
+  const AV1_COMP *cpi = args->cpi;
+  const AV1_COMMON *cm = &cpi->common;
+  ThreadData *const td = args->td;
+  MACROBLOCK *const x = &td->mb;
+  MACROBLOCKD *const xd = &x->e_mbd;
+  struct macroblock_plane *p = &x->plane[plane];
+  struct macroblockd_plane *pd = &xd->plane[plane];
+  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
+  int eob = p->eobs[block], update_eob = 0;
+  const PLANE_TYPE plane_type = pd->plane_type;
+  const tran_low_t *qcoeff = BLOCK_OFFSET(p->qcoeff, block);
+  tran_low_t *tcoeff = BLOCK_OFFSET(x->mbmi_ext->tcoeff[plane], block);
+  const int segment_id = mbmi->segment_id;
+  const int16_t *scan, *nb;
+  const 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, is_inter_block(mbmi));
+  const int ref = is_inter_block(mbmi);
+  unsigned int(*const counts)[COEFF_CONTEXTS][ENTROPY_TOKENS] =
+      td->rd_counts.coef_counts[tx_size][plane_type][ref];
+  const uint8_t *const band = get_band_translate(tx_size);
+  const int seg_eob = get_tx_eob(&cpi->common.seg, segment_id, tx_size);
+  int c, i;
+  int dc_sign;
+  int txb_skip_ctx = get_txb_skip_context(plane_bsize, tx_size, plane,
+                                          pd->above_context + blk_col,
+                                          pd->left_context + blk_row, &dc_sign);
+  const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
+  int cul_level = 0;
+  unsigned int(*nz_map_count)[SIG_COEF_CONTEXTS][2];
+  uint8_t txb_mask[32 * 32] = { 0 };
+
+  nz_map_count = &td->counts->nz_map[tx_size][plane_type];
+
+  scan = scan_order->scan;
+  nb = scan_order->neighbors;
+
+  memcpy(tcoeff, qcoeff, sizeof(*tcoeff) * seg_eob);
+
+  (void)nb;
+  (void)counts;
+  (void)band;
+
+  ++td->counts->txb_skip[tx_size][txb_skip_ctx][eob == 0];
+  x->mbmi_ext->txb_skip_ctx[plane][block] = txb_skip_ctx;
+
+  x->mbmi_ext->eobs[plane][block] = eob;
+
+  if (eob == 0) {
+    av1_set_contexts(xd, pd, plane, tx_size, 0, blk_col, blk_row);
+    return;
+  }
+
+  // update_tx_type_count(cm, mbmi, td, plane, block);
+
+  for (c = 0; c < eob; ++c) {
+    tran_low_t v = qcoeff[scan[c]];
+    int is_nz = (v != 0);
+    int coeff_ctx = get_nz_map_ctx(tcoeff, txb_mask, scan[c], bwl);
+    int eob_ctx = get_eob_ctx(tcoeff, scan[c], bwl);
+
+    if (c == seg_eob - 1) break;
+
+    ++(*nz_map_count)[coeff_ctx][is_nz];
+
+    if (is_nz) {
+      ++td->counts->eob_flag[tx_size][plane_type][eob_ctx][c == (eob - 1)];
+    }
+    txb_mask[scan[c]] = 1;
+  }
+
+  // Reverse process order to handle coefficient level and sign.
+  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+    update_eob = 0;
+    for (c = eob - 1; c >= 0; --c) {
+      tran_low_t v = qcoeff[scan[c]];
+      tran_low_t level = abs(v);
+      int ctx;
+
+      if (level <= i) continue;
+
+      ctx = get_base_ctx(tcoeff, scan[c], bwl, i + 1);
+
+      if (level == i + 1) {
+        ++td->counts->coeff_base[tx_size][plane_type][i][ctx][1];
+        if (c == 0) {
+          int dc_sign_ctx = get_dc_sign_ctx(dc_sign);
+
+          ++td->counts->dc_sign[plane_type][dc_sign_ctx][v < 0];
+          x->mbmi_ext->dc_sign_ctx[plane][block] = dc_sign_ctx;
+        }
+        cul_level += level;
+        continue;
+      }
+      ++td->counts->coeff_base[tx_size][plane_type][i][ctx][0];
+      update_eob = AOMMAX(update_eob, c);
+    }
+  }
+
+  for (c = update_eob; c >= 0; --c) {
+    tran_low_t v = qcoeff[scan[c]];
+    tran_low_t level = abs(v);
+    int idx;
+    int ctx;
+
+    if (level <= NUM_BASE_LEVELS) continue;
+
+    cul_level += level;
+    if (c == 0) {
+      int dc_sign_ctx = get_dc_sign_ctx(dc_sign);
+
+      ++td->counts->dc_sign[plane_type][dc_sign_ctx][v < 0];
+      x->mbmi_ext->dc_sign_ctx[plane][block] = dc_sign_ctx;
+    }
+
+    // level is above 1.
+    ctx = get_level_ctx(tcoeff, scan[c], bwl);
+    for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
+      if (level == (idx + 1 + NUM_BASE_LEVELS)) {
+        ++td->counts->coeff_lps[tx_size][plane_type][ctx][1];
+        break;
+      }
+      ++td->counts->coeff_lps[tx_size][plane_type][ctx][0];
+    }
+    if (idx < COEFF_BASE_RANGE) continue;
+
+    // use 0-th order Golomb code to handle the residual level.
+  }
+  cul_level = AOMMIN(63, cul_level);
+
+  // DC value
+  set_dc_sign(&cul_level, tcoeff[0]);
+  av1_set_contexts(xd, pd, plane, tx_size, cul_level, blk_col, blk_row);
+}
+
+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) {
+  const AV1_COMMON *const cm = &cpi->common;
+  MACROBLOCK *const x = &td->mb;
+  MACROBLOCKD *const xd = &x->e_mbd;
+  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
+  const int ctx = av1_get_skip_context(xd);
+  const int skip_inc =
+      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
+  struct TxbParams arg = { cpi, td, 0 };
+  (void)rate;
+  (void)mi_row;
+  (void)mi_col;
+  if (mbmi->skip) {
+    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
+    reset_skip_context(xd, bsize);
+    return;
+  }
+
+  if (!dry_run) {
+    td->counts->skip[ctx][0] += skip_inc;
+    av1_foreach_transformed_block(xd, bsize, update_and_record_txb_context,
+                                  &arg);
+  } else {
+    av1_foreach_transformed_block(xd, bsize, update_txb_context, &arg);
+  }
+}

av1/encoder/encodetxb.h

@@ -28,6 +28,10 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
                           aom_writer *w, int block, int plane,
                           const tran_low_t *tcoeff, uint16_t eob,
                           TXB_CTX *txb_ctx);
+
+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);
 #ifdef __cplusplus
 }
 #endif

av1/encoder/tokenize.c

@@ -313,7 +313,7 @@ static INLINE void add_token(TOKENEXTRA **t,
   (*t)++;
 }
 
-#else
+#else  // CONFIG_NEW_TOKENSET
 static INLINE void add_token(
     TOKENEXTRA **t, const aom_prob *context_tree,
 #if CONFIG_EC_MULTISYMBOL
@@ -330,14 +330,8 @@ static INLINE void add_token(
   (*t)++;
   ++counts[token];
 }
-#endif
-
-static INLINE int get_tx_eob(const struct segmentation *seg, int segment_id,
-                             TX_SIZE tx_size) {
-  const int eob_max = tx_size_2d[tx_size];
-  return segfeature_active(seg, segment_id, SEG_LVL_SKIP) ? 0 : eob_max;
-}
-#endif  // !CONFIG_PVQ
+#endif  // CONFIG_NEW_TOKENSET
+#endif  // !CONFIG_PVQ || CONFIG_VAR_TX
 
 #if CONFIG_PALETTE
 void av1_tokenize_palette_sb(const AV1_COMP *cpi,

av1/encoder/tokenize.h

@@ -136,6 +136,14 @@ static INLINE int av1_get_token_cost(int v, int16_t *token, int cat6_bits) {
   return av1_dct_cat_lt_10_value_cost[v];
 }
 
+#if !CONFIG_PVQ || CONFIG_VAR_TX
+static INLINE int get_tx_eob(const struct segmentation *seg, int segment_id,
+                             TX_SIZE tx_size) {
+  const int eob_max = tx_size_2d[tx_size];
+  return segfeature_active(seg, segment_id, SEG_LVL_SKIP) ? 0 : eob_max;
+}
+#endif
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif