Add get_coeff_cost() and get_txb_cost()

Change-Id: I085f2bc706fde41afbee5ff48b56acc095f804c2

av1/common/txb_common.h

@@ -130,17 +130,18 @@ static INLINE int get_level_ctx(const tran_low_t *tcoeffs,
   return 8 + ctx;
 }
 
-static int sig_ref_offset[11][2] = {
+#define SIG_REF_OFFSET_NUM 11
+static int sig_ref_offset[SIG_REF_OFFSET_NUM][2] = {
   { -2, -1 }, { -2, 0 }, { -2, 1 }, { -1, -2 }, { -1, -1 }, { -1, 0 },
   { -1, 1 },  { 0, -2 }, { 0, -1 }, { 1, -2 },  { 1, -1 },
 };
 
 static INLINE int get_nz_map_ctx(const tran_low_t *tcoeffs,
                                  const uint8_t *txb_mask,
-                                 const int c,  // raster order
+                                 const int coeff_idx,  // raster order
                                  const int bwl) {
-  const int row = c >> bwl;
-  const int col = c - (row << bwl);
+  const int row = coeff_idx >> bwl;
+  const int col = coeff_idx - (row << bwl);
   int ctx = 0;
   int idx;
   int stride = 1 << bwl;
@@ -166,7 +167,7 @@ static INLINE int get_nz_map_ctx(const tran_low_t *tcoeffs,
     return 5 + ctx;
   }
 
-  for (idx = 0; idx < 11; ++idx) {
+  for (idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
     int ref_row = row + sig_ref_offset[idx][0];
     int ref_col = col + sig_ref_offset[idx][1];
     int pos;
@@ -200,14 +201,93 @@ static INLINE int get_nz_map_ctx(const tran_low_t *tcoeffs,
   return 14 + ctx;
 }
 
+static INLINE int get_nz_count(const tran_low_t *tcoeffs, int stride, int row,
+                               int col, const int16_t *iscan) {
+  int count = 0;
+  const int pos = row * stride + col;
+  for (int idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
+    const int ref_row = row + sig_ref_offset[idx][0];
+    const int ref_col = col + sig_ref_offset[idx][1];
+    if (ref_row < 0 || ref_col < 0 || ref_row >= stride || ref_col >= stride)
+      continue;
+    const int nb_pos = ref_row * stride + ref_col;
+    if (iscan[nb_pos] < iscan[pos]) count += (tcoeffs[nb_pos] != 0);
+  }
+  return count;
+}
+
+// TODO(angiebird): optimize this function by generate a table that maps from
+// count to ctx
+static INLINE int get_nz_map_ctx_from_count(int count,
+                                            const tran_low_t *tcoeffs,
+                                            int coeff_idx,  // raster order
+                                            int bwl, const int16_t *iscan) {
+  const int row = coeff_idx >> bwl;
+  const int col = coeff_idx - (row << bwl);
+  int ctx = 0;
+
+  if (row == 0 && col == 0) return 0;
+
+  if (row == 0 && col == 1) return 1 + (tcoeffs[0] != 0);
+
+  if (row == 1 && col == 0) return 3 + (tcoeffs[0] != 0);
+
+  if (row == 1 && col == 1) {
+    int pos;
+    ctx = (tcoeffs[0] != 0);
+
+    if (iscan[1] < iscan[coeff_idx]) ctx += (tcoeffs[1] != 0);
+    pos = 1 << bwl;
+    if (iscan[pos] < iscan[coeff_idx]) ctx += (tcoeffs[pos] != 0);
+
+    ctx = (ctx + 1) >> 1;
+
+    assert(5 + ctx <= 7);
+
+    return 5 + ctx;
+  }
+
+  if (row == 0) {
+    ctx = (count + 1) >> 1;
+
+    assert(ctx < 3);
+    return 8 + ctx;
+  }
+
+  if (col == 0) {
+    ctx = (count + 1) >> 1;
+
+    assert(ctx < 3);
+    return 11 + ctx;
+  }
+
+  ctx = count >> 1;
+
+  assert(14 + ctx < 20);
+
+  return 14 + ctx;
+}
+
+// TODO(angiebird): merge this function with get_nz_map_ctx() after proper
+// testing
+static INLINE int get_nz_map_ctx2(const tran_low_t *tcoeffs,
+                                  const int coeff_idx,  // raster order
+                                  const int bwl, const int16_t *iscan) {
+  int stride = 1 << bwl;
+  const int row = coeff_idx >> bwl;
+  const int col = coeff_idx - (row << bwl);
+  int count = get_nz_count(tcoeffs, stride, row, col, iscan);
+  return get_nz_map_ctx_from_count(count, tcoeffs, coeff_idx, bwl, iscan);
+}
+
 static INLINE int get_eob_ctx(const tran_low_t *tcoeffs,
-                              const int c,  // raster order
+                              const int coeff_idx,  // raster order
                               const int bwl) {
   (void)tcoeffs;
-  if (bwl == 2) return av1_coeff_band_4x4[c];
-  if (bwl == 3) return av1_coeff_band_8x8[c];
-  if (bwl == 4) return av1_coeff_band_16x16[c];
-  if (bwl == 5) return av1_coeff_band_32x32[c];
+  if (bwl == 2) return av1_coeff_band_4x4[coeff_idx];
+  if (bwl == 3) return av1_coeff_band_8x8[coeff_idx];
+  if (bwl == 4) return av1_coeff_band_16x16[coeff_idx];
+  if (bwl == 5) return av1_coeff_band_32x32[coeff_idx];
 
   assert(0);
   return 0;

av1/encoder/encodetxb.c

@@ -251,6 +251,32 @@ static INLINE void get_base_ctx_set(const tran_low_t *tcoeffs,
   return;
 }
 
+static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
+                              const aom_prob *coeff_lps) {
+  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
+  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
+  if (abs_qc >= min_level) {
+    const int cost0 = av1_cost_bit(coeff_lps[ctx], 0);
+    const int cost1 = av1_cost_bit(coeff_lps[ctx], 1);
+    if (abs_qc >= max_level)
+      return COEFF_BASE_RANGE * cost0;
+    else
+      return (abs_qc - min_level) * cost0 + cost1;
+  } else {
+    return 0;
+  }
+}
+
+static INLINE int get_base_cost(tran_low_t abs_qc, int ctx,
+                                aom_prob (*coeff_base)[COEFF_BASE_CONTEXTS],
+                                int base_idx) {
+  const int level = base_idx + 1;
+  if (abs_qc < level)
+    return 0;
+  else
+    return av1_cost_bit(coeff_base[base_idx][ctx], abs_qc == level);
+}
+
 int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,
                         int block, TXB_CTX *txb_ctx) {
   const AV1_COMMON *const cm = &cpi->common;
@@ -373,6 +399,96 @@ int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,
   return cost;
 }
 
+static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
+                                    const aom_prob *dc_sign_prob,
+                                    int dc_sign_ctx) {
+  const int sign = (qc < 0) ? 1 : 0;
+  // sign bit cost
+  if (coeff_idx == 0) {
+    return av1_cost_bit(dc_sign_prob[dc_sign_ctx], sign);
+  } else {
+    return av1_cost_bit(128, sign);
+  }
+}
+static INLINE int get_golomb_cost(int abs_qc) {
+  if (abs_qc >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
+    // residual cost
+    int r = abs_qc - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
+    int ri = r;
+    int length = 0;
+
+    while (ri) {
+      ri >>= 1;
+      ++length;
+    }
+
+    return av1_cost_literal(2 * length - 1);
+  } else {
+    return 0;
+  }
+}
+
+static int get_coeff_cost(tran_low_t qc, int scan_idx, TxbInfo *txb_info,
+                          TxbProbs *txb_probs) {
+  const TXB_CTX *txb_ctx = txb_info->txb_ctx;
+  const int is_nz = (qc != 0);
+  const tran_low_t abs_qc = abs(qc);
+  int cost = 0;
+  const int16_t *scan = txb_info->scan_order->scan;
+  const int16_t *iscan = txb_info->scan_order->iscan;
+
+  if (scan_idx < txb_info->seg_eob) {
+    int coeff_ctx =
+        get_nz_map_ctx2(txb_info->qcoeff, scan[scan_idx], txb_info->bwl, iscan);
+    cost += av1_cost_bit(txb_probs->nz_map[coeff_ctx], is_nz);
+  }
+
+  if (is_nz) {
+    cost += get_sign_bit_cost(qc, scan_idx, txb_probs->dc_sign_prob,
+                              txb_ctx->dc_sign_ctx);
+
+    int ctx_ls[NUM_BASE_LEVELS] = { 0 };
+    get_base_ctx_set(txb_info->qcoeff, scan[scan_idx], txb_info->bwl, ctx_ls);
+
+    int i;
+    for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+      cost += get_base_cost(abs_qc, ctx_ls[i], txb_probs->coeff_base, i);
+    }
+
+    if (abs_qc > NUM_BASE_LEVELS) {
+      int ctx = get_level_ctx(txb_info->qcoeff, scan[scan_idx], txb_info->bwl);
+      cost += get_br_cost(abs_qc, ctx, txb_probs->coeff_lps);
+      cost += get_golomb_cost(abs_qc);
+    }
+
+    if (scan_idx < txb_info->seg_eob) {
+      int eob_ctx =
+          get_eob_ctx(txb_info->qcoeff, scan[scan_idx], txb_info->bwl);
+      cost += av1_cost_bit(txb_probs->eob_flag[eob_ctx],
+                           scan_idx == (txb_info->eob - 1));
+    }
+  }
+  return cost;
+}
+
+// TODO(angiebird): make this static once it's called
+int get_txb_cost(TxbInfo *txb_info, TxbProbs *txb_probs) {
+  int cost = 0;
+  int txb_skip_ctx = txb_info->txb_ctx->txb_skip_ctx;
+  const int16_t *scan = txb_info->scan_order->scan;
+  if (txb_info->eob == 0) {
+    cost = av1_cost_bit(txb_probs->txb_skip[txb_skip_ctx], 1);
+    return cost;
+  }
+  cost = av1_cost_bit(txb_probs->txb_skip[txb_skip_ctx], 0);
+  for (int c = 0; c < txb_info->eob; ++c) {
+    tran_low_t qc = txb_info->qcoeff[scan[c]];
+    int coeff_cost = get_coeff_cost(qc, c, txb_info, txb_probs);
+    cost += coeff_cost;
+  }
+  return cost;
+}
+
 int av1_get_txb_entropy_context(const tran_low_t *qcoeff,
                                 const SCAN_ORDER *scan_order, int eob) {
   const int16_t *scan = scan_order->scan;

av1/encoder/encodetxb.h

@@ -22,6 +22,47 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+
+typedef struct TxbInfo {
+  tran_low_t *qcoeff;
+  tran_low_t *dqcoeff;
+  const tran_low_t *tcoeff;
+  const int16_t *dequant;
+  int shift;
+  TX_SIZE tx_size;
+  int bwl;
+  int stride;
+  int eob;
+  int seg_eob;
+  const SCAN_ORDER *scan_order;
+  TXB_CTX *txb_ctx;
+  int64_t rdmult;
+  int64_t rddiv;
+} TxbInfo;
+
+typedef struct TxbCache {
+  int nz_count_arr[MAX_TX_SQUARE];
+  int nz_ctx_arr[MAX_TX_SQUARE][2];
+  int base_count_arr[NUM_BASE_LEVELS][MAX_TX_SQUARE];
+  int base_mag_arr[MAX_TX_SQUARE]
+                  [2];  // [0]: max magnitude [1]: num of max magnitude
+  int base_ctx_arr[NUM_BASE_LEVELS][MAX_TX_SQUARE][2];  // [1]: not used
+
+  int br_count_arr[MAX_TX_SQUARE];
+  int br_mag_arr[MAX_TX_SQUARE]
+                [2];  // [0]: max magnitude [1]: num of max magnitude
+  int br_ctx_arr[MAX_TX_SQUARE][2];  // [1]: not used
+} TxbCache;
+
+typedef struct TxbProbs {
+  const aom_prob *dc_sign_prob;
+  const aom_prob *nz_map;
+  aom_prob (*coeff_base)[COEFF_BASE_CONTEXTS];
+  const aom_prob *coeff_lps;
+  const aom_prob *eob_flag;
+  const aom_prob *txb_skip;
+} TxbProbs;
+
 void av1_alloc_txb_buf(AV1_COMP *cpi);
 void av1_free_txb_buf(AV1_COMP *cpi);
 int av1_cost_coeffs_txb(const AV1_COMP *const cpi, MACROBLOCK *x, int plane,