Support level map coeffs in var-tx encoding
Use recursive transform block coding order.
Change-Id: I6d4671fe669e8a1897e034973de181078272cbfd
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 0cf7c5e..a012cb0 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1129,7 +1129,81 @@
}
#endif // !CONFIG_PVG
-#if CONFIG_VAR_TX && !CONFIG_COEF_INTERLEAVE && !CONFIG_LV_MAP
+#if CONFIG_VAR_TX && !CONFIG_COEF_INTERLEAVE
+#if CONFIG_LV_MAP
+static void pack_txb_tokens(aom_writer *w,
+#if CONFIG_LV_MAP
+ AV1_COMMON *cm,
+#endif // CONFIG_LV_MAP
+ const TOKENEXTRA **tp,
+ const TOKENEXTRA *const tok_end,
+#if CONFIG_PVQ || CONFIG_LV_MAP
+ MACROBLOCK *const x,
+#endif
+ MACROBLOCKD *xd, MB_MODE_INFO *mbmi, int plane,
+ BLOCK_SIZE plane_bsize, aom_bit_depth_t bit_depth,
+ int block, int blk_row, int blk_col,
+ TX_SIZE tx_size, TOKEN_STATS *token_stats) {
+ const struct macroblockd_plane *const pd = &xd->plane[plane];
+ const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
+ const int tx_row = blk_row >> (1 - pd->subsampling_y);
+ const int tx_col = blk_col >> (1 - pd->subsampling_x);
+ TX_SIZE plane_tx_size;
+ const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
+ const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
+
+ if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
+
+ plane_tx_size =
+ plane ? uv_txsize_lookup[bsize][mbmi->inter_tx_size[tx_row][tx_col]][0][0]
+ : mbmi->inter_tx_size[tx_row][tx_col];
+
+ if (tx_size == plane_tx_size) {
+ TOKEN_STATS tmp_token_stats;
+ init_token_stats(&tmp_token_stats);
+
+#if !CONFIG_PVQ
+ tran_low_t *tcoeff = BLOCK_OFFSET(x->mbmi_ext->tcoeff[plane], block);
+ uint16_t eob = x->mbmi_ext->eobs[plane][block];
+ TXB_CTX txb_ctx = { x->mbmi_ext->txb_skip_ctx[plane][block],
+ x->mbmi_ext->dc_sign_ctx[plane][block] };
+ av1_write_coeffs_txb(cm, xd, w, block, plane, tcoeff, eob, &txb_ctx);
+#else
+ pack_pvq_tokens(w, x, xd, plane, bsize, tx_size);
+#endif
+#if CONFIG_RD_DEBUG
+ token_stats->txb_coeff_cost_map[blk_row][blk_col] = tmp_token_stats.cost;
+ token_stats->cost += tmp_token_stats.cost;
+#endif
+ } else {
+ const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
+ const int bsl = tx_size_wide_unit[sub_txs];
+ int i;
+
+ assert(bsl > 0);
+
+ for (i = 0; i < 4; ++i) {
+ const int offsetr = blk_row + (i >> 1) * bsl;
+ const int offsetc = blk_col + (i & 0x01) * bsl;
+ const int step = tx_size_wide_unit[sub_txs] * tx_size_high_unit[sub_txs];
+
+ if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;
+
+ pack_txb_tokens(w,
+#if CONFIG_LV_MAP
+ cm,
+#endif
+ tp, tok_end,
+#if CONFIG_PVQ || CONFIG_LV_MAP
+ x,
+#endif
+ xd, mbmi, plane, plane_bsize, bit_depth, block, offsetr,
+ offsetc, sub_txs, token_stats);
+ block += step;
+ }
+ }
+}
+#else // CONFIG_LV_MAP
static void pack_txb_tokens(aom_writer *w, const TOKENEXTRA **tp,
const TOKENEXTRA *const tok_end,
#if CONFIG_PVQ
@@ -1189,7 +1263,8 @@
}
}
}
-#endif
+#endif // CONFIG_LV_MAP
+#endif // CONFIG_VAR_TX
static void write_segment_id(aom_writer *w, const struct segmentation *seg,
struct segmentation_probs *segp, int segment_id) {
@@ -2604,22 +2679,16 @@
const int bkh = tx_size_high_unit[max_tx_size];
for (row = 0; row < num_4x4_h; row += bkh) {
for (col = 0; col < num_4x4_w; col += bkw) {
+ pack_txb_tokens(w,
#if CONFIG_LV_MAP
- tran_low_t *tcoeff =
- BLOCK_OFFSET(x->mbmi_ext->tcoeff[plane], block);
- uint16_t eob = x->mbmi_ext->eobs[plane][block];
- TXB_CTX txb_ctx = { x->mbmi_ext->txb_skip_ctx[plane][block],
- x->mbmi_ext->dc_sign_ctx[plane][block] };
- av1_write_coeffs_txb(cm, xd, w, block, plane, tcoeff, eob,
- &txb_ctx);
-#else
- pack_txb_tokens(w, tok, tok_end,
-#if CONFIG_PVQ
+ cm,
+#endif
+ tok, tok_end,
+#if CONFIG_PVQ || CONFIG_LV_MAP
x,
#endif
xd, mbmi, plane, plane_bsize, cm->bit_depth, block,
row, col, max_tx_size, &token_stats);
-#endif // CONFIG_LV_MAP
block += step;
}
}
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 788427a..06d04a1 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -388,9 +388,9 @@
return cul_level;
}
-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) {
+void av1_update_txb_context_b(int plane, int block, int blk_row, int blk_col,
+ BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+ void *arg) {
struct tokenize_b_args *const args = arg;
const AV1_COMP *cpi = args->cpi;
const AV1_COMMON *cm = &cpi->common;
@@ -412,9 +412,9 @@
av1_set_contexts(xd, pd, plane, tx_size, cul_level, 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) {
+void av1_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) {
struct tokenize_b_args *const args = arg;
const AV1_COMP *cpi = args->cpi;
const AV1_COMMON *cm = &cpi->common;
@@ -575,10 +575,10 @@
if (!dry_run) {
td->counts->skip[ctx][0] += skip_inc;
av1_foreach_transformed_block(xd, bsize, mi_row, mi_col,
- update_and_record_txb_context, &arg);
+ av1_update_and_record_txb_context, &arg);
} else if (dry_run == DRY_RUN_NORMAL) {
- av1_foreach_transformed_block(xd, bsize, mi_row, mi_col, update_txb_context,
- &arg);
+ av1_foreach_transformed_block(xd, bsize, mi_row, mi_col,
+ av1_update_txb_context_b, &arg);
} else {
printf("DRY_RUN_COSTCOEFFS is not supported yet\n");
assert(0);
diff --git a/av1/encoder/encodetxb.h b/av1/encoder/encodetxb.h
index 552d47b..e1aca2c 100644
--- a/av1/encoder/encodetxb.h
+++ b/av1/encoder/encodetxb.h
@@ -39,6 +39,14 @@
const int mi_row, const int mi_col);
void av1_write_txb_probs(AV1_COMP *cpi, aom_writer *w);
+void av1_update_txb_context_b(int plane, int block, int blk_row, int blk_col,
+ BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
+ void *arg);
+
+void av1_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);
+
#if CONFIG_TXK_SEL
int64_t av1_search_txk_type(const AV1_COMP *cpi, MACROBLOCK *x, int plane,
int block, int blk_row, int blk_col,
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 75bd737..d925287 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -23,6 +23,9 @@
#include "av1/encoder/cost.h"
#include "av1/encoder/encoder.h"
+#if CONFIG_LV_MAP
+#include "av1/encoder/encodetxb.c"
+#endif
#include "av1/encoder/rdopt.h"
#include "av1/encoder/tokenize.h"
@@ -622,6 +625,18 @@
if (tx_size == plane_tx_size) {
plane_bsize = get_plane_block_size(mbmi->sb_type, pd);
+#if CONFIG_LV_MAP
+ if (!dry_run) {
+ av1_update_and_record_txb_context(plane, block, blk_row, blk_col,
+ plane_bsize, tx_size, arg);
+ } else if (dry_run == DRY_RUN_NORMAL) {
+ av1_update_txb_context_b(plane, block, blk_row, blk_col, plane_bsize,
+ tx_size, arg);
+ } else {
+ printf("DRY_RUN_COSTCOEFFS is not supported yet\n");
+ assert(0);
+ }
+#else
if (!dry_run)
tokenize_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
else if (dry_run == DRY_RUN_NORMAL)
@@ -629,6 +644,7 @@
tx_size, arg);
else if (dry_run == DRY_RUN_COSTCOEFFS)
cost_coeffs_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
+#endif
} else {
// Half the block size in transform block unit.
const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
@@ -659,7 +675,11 @@
MACROBLOCK *const x = &td->mb;
MACROBLOCKD *const xd = &x->e_mbd;
MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
+#if CONFIG_LV_MAP
+ (void)t;
+#else
TOKENEXTRA *t_backup = *t;
+#endif
const int ctx = av1_get_skip_context(xd);
const int skip_inc =
!segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
@@ -670,21 +690,24 @@
if (mbmi->skip) {
if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
reset_skip_context(xd, bsize);
+#if !CONFIG_LV_MAP
if (dry_run) *t = t_backup;
+#endif
return;
}
- if (!dry_run)
- td->counts->skip[ctx][0] += skip_inc;
+ if (!dry_run) td->counts->skip[ctx][0] += skip_inc;
+#if !CONFIG_LV_MAP
else
*t = t_backup;
+#endif
for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
#if CONFIG_CB4X4
if (!is_chroma_reference(mi_row, mi_col, bsize,
xd->plane[plane].subsampling_x,
xd->plane[plane].subsampling_y)) {
-#if !CONFIG_PVQ
+#if !CONFIG_PVQ || !CONFIG_LV_MAP
if (!dry_run) {
(*t)->token = EOSB_TOKEN;
(*t)++;
@@ -717,10 +740,12 @@
}
}
+#if !CONFIG_LV_MAP
if (!dry_run) {
(*t)->token = EOSB_TOKEN;
(*t)++;
}
+#endif
}
if (rate) *rate += arg.this_rate;
}
