Add txb_coeff_cost_map into TOKEN_STATS
This is to facilitate debugging process in var_tx experiment
Change-Id: Ibd5ea7f6054c598b8e686abb4e8158ef28c67aab
diff --git a/aom_dsp/bitwriter.h b/aom_dsp/bitwriter.h
index b437669..f96bf90 100644
--- a/aom_dsp/bitwriter.h
+++ b/aom_dsp/bitwriter.h
@@ -28,6 +28,7 @@
#include "aom_dsp/prob.h"
#if CONFIG_RD_DEBUG
+#include "av1/common/blockd.h"
#include "av1/encoder/cost.h"
#endif
@@ -43,7 +44,28 @@
typedef struct aom_dk_writer aom_writer;
#endif
-typedef struct TOKEN_STATS { int64_t cost; } TOKEN_STATS;
+typedef struct TOKEN_STATS {
+ int cost;
+#if CONFIG_VAR_TX
+#if CONFIG_RD_DEBUG
+ int txb_coeff_cost_map[TXB_COEFF_COST_MAP_SIZE][TXB_COEFF_COST_MAP_SIZE];
+#endif
+#endif
+} TOKEN_STATS;
+
+static INLINE void init_token_stats(TOKEN_STATS *token_stats) {
+#if CONFIG_VAR_TX
+#if CONFIG_RD_DEBUG
+ int r, c;
+ for (r = 0; r < TXB_COEFF_COST_MAP_SIZE; ++r) {
+ for (c = 0; c < TXB_COEFF_COST_MAP_SIZE; ++c) {
+ token_stats->txb_coeff_cost_map[r][c] = 0;
+ }
+ }
+#endif
+#endif
+ token_stats->cost = 0;
+}
static INLINE void aom_start_encode(aom_writer *bc, uint8_t *buffer) {
#if CONFIG_ANS
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 3b66fa7..9111cb7 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -847,7 +847,13 @@
: mbmi->inter_tx_size[tx_row][tx_col];
if (tx_size == plane_tx_size) {
- pack_mb_tokens(w, tp, tok_end, bit_depth, tx_size, token_stats);
+ TOKEN_STATS tmp_token_stats;
+ init_token_stats(&tmp_token_stats);
+ pack_mb_tokens(w, tp, tok_end, bit_depth, tx_size, &tmp_token_stats);
+#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 int bsl = block_size_wide[bsize] >> (tx_size_wide_log2[0] + 1);
int i;
@@ -1728,6 +1734,33 @@
printf("&& mi->bmi[0].as_mode == %d\n", mi->bmi[0].as_mode);
}
}
+static int rd_token_stats_mismatch(RD_STATS *rd_stats, TOKEN_STATS *token_stats,
+ int plane) {
+ if (rd_stats->txb_coeff_cost[plane] != token_stats->cost) {
+ int r, c;
+ printf("\nplane %d rd_stats->txb_coeff_cost %d token_stats->cost %d\n",
+ plane, rd_stats->txb_coeff_cost[plane], token_stats->cost);
+#if CONFIG_VAR_TX
+ printf("rd txb_coeff_cost_map\n");
+ for (r = 0; r < TXB_COEFF_COST_MAP_SIZE; ++r) {
+ for (c = 0; c < TXB_COEFF_COST_MAP_SIZE; ++c) {
+ printf("%d ", rd_stats->txb_coeff_cost_map[plane][r][c]);
+ }
+ printf("\n");
+ }
+
+ printf("pack txb_coeff_cost_map\n");
+ for (r = 0; r < TXB_COEFF_COST_MAP_SIZE; ++r) {
+ for (c = 0; c < TXB_COEFF_COST_MAP_SIZE; ++c) {
+ printf("%d ", token_stats->txb_coeff_cost_map[r][c]);
+ }
+ printf("\n");
+ }
+#endif
+ return 1;
+ }
+ return 0;
+}
#endif
#if CONFIG_PVQ
@@ -1756,9 +1789,6 @@
MODE_INFO *m;
int plane;
int bh, bw;
-#if CONFIG_RD_DEBUG
- int64_t txb_coeff_cost[MAX_MB_PLANE] = { 0 };
-#endif
#if CONFIG_RANS
(void)tok;
(void)tok_end;
@@ -1873,7 +1903,7 @@
#endif
TOKEN_STATS token_stats;
- token_stats.cost = 0;
+ init_token_stats(&token_stats);
#if CONFIG_EXT_TX && CONFIG_RECT_TX
@@ -1896,6 +1926,12 @@
block += step;
}
}
+#if CONFIG_RD_DEBUG
+ if (rd_token_stats_mismatch(&m->mbmi.rd_stats, &token_stats, plane)) {
+ dump_mode_info(m);
+ assert(0);
+ }
+#endif
} else {
TX_SIZE tx = plane ? get_uv_tx_size(&m->mbmi, &xd->plane[plane])
: m->mbmi.tx_size;
@@ -1914,25 +1950,10 @@
pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx, &token_stats);
#endif // CONFIG_VAR_TX
-#if CONFIG_RD_DEBUG
- txb_coeff_cost[plane] += token_stats.cost;
-#else
- (void)token_stats;
-#endif
-
assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
(*tok)++;
}
}
-
-#if CONFIG_RD_DEBUG
- for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
- if (m->mbmi.rd_stats.txb_coeff_cost[plane] != txb_coeff_cost[plane]) {
- dump_mode_info(m);
- assert(0);
- }
- }
-#endif // CONFIG_RD_DEBUG
#else
// PVQ writes its tokens (i.e. symbols) here.
if (!m->mbmi.skip) {
