Add av1_cost_coeffs_txb() for lv_map experiment
Change-Id: I44842387207b19f8e0c3894d3f4e8d0646a4cafd
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 62db8cb..4128d50 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -10,6 +10,7 @@
*/
#include "av1/common/scan.h"
+#include "av1/encoder/cost.h"
#include "av1/encoder/encodetxb.h"
static void write_golomb(aom_writer *w, int level) {
@@ -131,3 +132,171 @@
write_golomb(w, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
}
}
+
+static INLINE void get_base_ctx_set(const tran_low_t *tcoeffs,
+ int c, // raster order
+ const int bwl,
+ int ctx_set[NUM_BASE_LEVELS]) {
+ const int row = c >> bwl;
+ const int col = c - (row << bwl);
+ const int stride = 1 << bwl;
+ int mag[NUM_BASE_LEVELS] = { 0 };
+ int idx;
+ tran_low_t abs_coeff;
+ int i;
+
+ for (idx = 0; idx < BASE_CONTEXT_POSITION_NUM; ++idx) {
+ int ref_row = row + base_ref_offset[idx][0];
+ int ref_col = col + base_ref_offset[idx][1];
+ int pos = (ref_row << bwl) + ref_col;
+
+ if (ref_row < 0 || ref_col < 0 || ref_row >= stride || ref_col >= stride)
+ continue;
+
+ abs_coeff = abs(tcoeffs[pos]);
+
+ for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+ ctx_set[i] += abs_coeff > i;
+ if (base_ref_offset[idx][0] >= 0 && base_ref_offset[idx][1] >= 0)
+ mag[i] |= abs_coeff > (i + 1);
+ }
+ }
+
+ for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+ ctx_set[i] = (ctx_set[i] + 1) >> 1;
+
+ if (row == 0 && col == 0)
+ ctx_set[i] = (ctx_set[i] << 1) + mag[i];
+ else if (row == 0)
+ ctx_set[i] = 8 + (ctx_set[i] << 1) + mag[i];
+ else if (col == 0)
+ ctx_set[i] = 18 + (ctx_set[i] << 1) + mag[i];
+ else
+ ctx_set[i] = 28 + (ctx_set[i] << 1) + mag[i];
+ }
+ return;
+}
+
+int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
+ int block, TXB_CTX *txb_ctx, int *cul_level) {
+ MACROBLOCKD *const xd = &x->e_mbd;
+ const TX_SIZE tx_size = get_tx_size(plane, xd);
+ const PLANE_TYPE plane_type = get_plane_type(plane);
+ const TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
+ MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
+ const struct macroblock_plane *p = &x->plane[plane];
+ const int eob = p->eobs[block];
+ const tran_low_t *const qcoeff = BLOCK_OFFSET(p->qcoeff, block);
+ int c, cost;
+ const int seg_eob = AOMMIN(eob, (16 << (tx_size << 1)) - 1);
+ int txb_skip_ctx = txb_ctx->txb_skip_ctx;
+ aom_prob *nz_map = xd->fc->nz_map[tx_size][plane_type];
+
+ const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
+ *cul_level = 0;
+ // txb_mask is only initialized for once here. After that, it will be set when
+ // coding zero map and then reset when coding level 1 info.
+ uint8_t txb_mask[32 * 32] = { 0 };
+ aom_prob(*coeff_base)[COEFF_BASE_CONTEXTS] =
+ xd->fc->coeff_base[tx_size][plane_type];
+
+ const SCAN_ORDER *const scan_order =
+ get_scan(cm, tx_size, tx_type, is_inter_block(mbmi));
+ const int16_t *scan = scan_order->scan;
+
+ cost = 0;
+
+ if (eob == 0) {
+ cost = av1_cost_bit(xd->fc->txb_skip[tx_size][txb_skip_ctx], 1);
+ return cost;
+ }
+
+ cost = av1_cost_bit(xd->fc->txb_skip[tx_size][txb_skip_ctx], 0);
+
+ for (c = 0; c < eob; ++c) {
+ tran_low_t v = qcoeff[scan[c]];
+ int is_nz = (v != 0);
+ int level = abs(v);
+
+ if (c < seg_eob) {
+ int coeff_ctx = get_nz_map_ctx(qcoeff, txb_mask, scan[c], bwl);
+ cost += av1_cost_bit(nz_map[coeff_ctx], is_nz);
+ }
+
+ if (is_nz) {
+ int ctx_ls[NUM_BASE_LEVELS] = { 0 };
+ int sign = (v < 0) ? 1 : 0;
+
+ // sign bit cost
+ if (c == 0) {
+ int dc_sign_ctx = txb_ctx->dc_sign_ctx;
+
+ cost += av1_cost_bit(xd->fc->dc_sign[plane_type][dc_sign_ctx], sign);
+ } else {
+ cost += av1_cost_bit(128, sign);
+ }
+
+ get_base_ctx_set(qcoeff, scan[c], bwl, ctx_ls);
+
+ int i;
+ for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+ if (level <= i) continue;
+
+ if (level == i + 1) {
+ cost += av1_cost_bit(coeff_base[i][ctx_ls[i]], 1);
+ *cul_level += level;
+ continue;
+ }
+ cost += av1_cost_bit(coeff_base[i][ctx_ls[i]], 0);
+ }
+
+ if (level > NUM_BASE_LEVELS) {
+ int idx;
+ int ctx;
+ *cul_level += level;
+
+ ctx = get_level_ctx(qcoeff, scan[c], bwl);
+
+ for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
+ if (level == (idx + 1 + NUM_BASE_LEVELS)) {
+ cost +=
+ av1_cost_bit(xd->fc->coeff_lps[tx_size][plane_type][ctx], 1);
+ break;
+ }
+ cost += av1_cost_bit(xd->fc->coeff_lps[tx_size][plane_type][ctx], 0);
+ }
+
+ if (idx >= COEFF_BASE_RANGE) {
+ // residual cost
+ int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
+ int ri = r;
+ int length = 0;
+
+ while (ri) {
+ ri >>= 1;
+ ++length;
+ }
+
+ for (ri = 0; ri < length - 1; ++ri) cost += av1_cost_bit(128, 0);
+
+ for (ri = length - 1; ri >= 0; --ri)
+ cost += av1_cost_bit(128, (r >> ri) & 0x01);
+ }
+ }
+
+ if (c < seg_eob) {
+ int eob_ctx = get_eob_ctx(qcoeff, scan[c], bwl);
+ cost += av1_cost_bit(xd->fc->eob_flag[tx_size][plane_type][eob_ctx],
+ c == (eob - 1));
+ }
+ }
+
+ txb_mask[scan[c]] = 1;
+ }
+
+ *cul_level = AOMMIN(63, *cul_level);
+ // DC value
+ set_dc_sign(cul_level, qcoeff[0]);
+
+ return cost;
+}
diff --git a/av1/encoder/encodetxb.h b/av1/encoder/encodetxb.h
index b6675f1..b3fe8e6 100644
--- a/av1/encoder/encodetxb.h
+++ b/av1/encoder/encodetxb.h
@@ -16,6 +16,7 @@
#include "av1/common/blockd.h"
#include "av1/common/onyxc_int.h"
#include "av1/common/txb_common.h"
+#include "av1/encoder/block.h"
#include "aom_dsp/bitwriter.h"
#ifdef __cplusplus
extern "C" {
