Add coef_interleave experiment
This commit adds an experiment to interleave the coding of transform
coefficients from YUV planes. The experiment can be enabled at config
time by --enable-coef-interleave.
Change-Id: Ifd92f9c367304bca9732f13fa026eb8996363677
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 2c64880..67322e7 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -1883,6 +1883,82 @@
if (supertx_enabled) return;
#endif // CONFIG_SUPERTX
+#if CONFIG_COEF_INTERLEAVE
+ if (!m->mbmi.skip) {
+ const struct macroblockd_plane *const pd_y = &xd->plane[0];
+ const struct macroblockd_plane *const pd_c = &xd->plane[1];
+ const TX_SIZE tx_log2_y = m->mbmi.tx_size;
+ const TX_SIZE tx_log2_c = get_uv_tx_size(&m->mbmi, pd_c);
+ const int tx_sz_y = (1 << tx_log2_y);
+ const int tx_sz_c = (1 << tx_log2_c);
+
+ const BLOCK_SIZE plane_bsize_y =
+ get_plane_block_size(AOMMAX(m->mbmi.sb_type, 3), pd_y);
+ const BLOCK_SIZE plane_bsize_c =
+ get_plane_block_size(AOMMAX(m->mbmi.sb_type, 3), pd_c);
+
+ const int num_4x4_w_y = num_4x4_blocks_wide_lookup[plane_bsize_y];
+ const int num_4x4_w_c = num_4x4_blocks_wide_lookup[plane_bsize_c];
+ const int num_4x4_h_y = num_4x4_blocks_high_lookup[plane_bsize_y];
+ const int num_4x4_h_c = num_4x4_blocks_high_lookup[plane_bsize_c];
+
+ const int max_4x4_w_y = get_max_4x4_size(num_4x4_w_y, xd->mb_to_right_edge,
+ pd_y->subsampling_x);
+ const int max_4x4_h_y = get_max_4x4_size(num_4x4_h_y, xd->mb_to_bottom_edge,
+ pd_y->subsampling_y);
+ const int max_4x4_w_c = get_max_4x4_size(num_4x4_w_c, xd->mb_to_right_edge,
+ pd_c->subsampling_x);
+ const int max_4x4_h_c = get_max_4x4_size(num_4x4_h_c, xd->mb_to_bottom_edge,
+ pd_c->subsampling_y);
+
+ // The max_4x4_w/h may be smaller than tx_sz under some corner cases,
+ // i.e. when the SB is splitted by tile boundaries.
+ const int tu_num_w_y = (max_4x4_w_y + tx_sz_y - 1) / tx_sz_y;
+ const int tu_num_h_y = (max_4x4_h_y + tx_sz_y - 1) / tx_sz_y;
+ const int tu_num_w_c = (max_4x4_w_c + tx_sz_c - 1) / tx_sz_c;
+ const int tu_num_h_c = (max_4x4_h_c + tx_sz_c - 1) / tx_sz_c;
+ const int tu_num_y = tu_num_w_y * tu_num_h_y;
+ const int tu_num_c = tu_num_w_c * tu_num_h_c;
+
+ int tu_idx_y = 0, tu_idx_c = 0;
+ TOKEN_STATS token_stats;
+ init_token_stats(&token_stats);
+
+ assert(*tok < tok_end);
+
+ while (tu_idx_y < tu_num_y) {
+ pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx_log2_y, &token_stats);
+ assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
+ (*tok)++;
+ tu_idx_y++;
+
+ if (tu_idx_c < tu_num_c) {
+ pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx_log2_c, &token_stats);
+ assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
+ (*tok)++;
+
+ pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx_log2_c, &token_stats);
+ assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
+ (*tok)++;
+
+ tu_idx_c++;
+ }
+ }
+
+ // In 422 case, it's possilbe that Chroma has more TUs than Luma
+ while (tu_idx_c < tu_num_c) {
+ pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx_log2_c, &token_stats);
+ assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
+ (*tok)++;
+
+ pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx_log2_c, &token_stats);
+ assert(*tok < tok_end && (*tok)->token == EOSB_TOKEN);
+ (*tok)++;
+
+ tu_idx_c++;
+ }
+ }
+#else // CONFIG_COEF_INTERLEAVE
if (!m->mbmi.skip) {
assert(*tok < tok_end);
for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
@@ -1965,6 +2041,7 @@
(*tok)++;
}
}
+#endif // CONFIG_COEF_INTERLEAVE
#else
// PVQ writes its tokens (i.e. symbols) here.
if (!m->mbmi.skip) {
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 0012559..0e6b815 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -503,6 +503,11 @@
++eob_branch[band[c]][pt];
}
+#if CONFIG_COEF_INTERLEAVE
+ t->token = EOSB_TOKEN;
+ t++;
+#endif
+
*tp = t;
#if CONFIG_ADAPT_SCAN
@@ -725,6 +730,10 @@
#if !CONFIG_PVQ
if (!dry_run) {
+#if CONFIG_COEF_INTERLEAVE
+ td->counts->skip[ctx][0] += skip_inc;
+ av1_foreach_transformed_block_interleave(xd, bsize, tokenize_b, &arg);
+#else
int plane;
td->counts->skip[ctx][0] += skip_inc;
@@ -734,6 +743,7 @@
(*t)->token = EOSB_TOKEN;
(*t)++;
}
+#endif
} else if (dry_run == DRY_RUN_NORMAL) {
av1_foreach_transformed_block(xd, bsize, set_entropy_context_b, &arg);
} else if (dry_run == DRY_RUN_COSTCOEFFS) {
