Fix av1_write_coeffs_mb when ext_partition is on
Use av1_foreach_transformed_block_in_plane.
Change-Id: I22f92b601f6a90e92a93dbff8492df86c717ab6f
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index b597ec7..1c04554 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -2097,7 +2097,7 @@
#endif // CONFIG_RD_DEBUG
} else {
#if CONFIG_LV_MAP
- av1_write_coeffs_mb(cm, x, w, plane);
+ av1_write_coeffs_mb(cm, x, w, plane, bsize);
#else
const TX_SIZE tx = av1_get_tx_size(plane, xd);
const int bkw = tx_size_wide_unit[tx];
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index c5a4a38..3f3da19 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -533,33 +533,35 @@
#endif
}
-void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
- aom_writer *w, int plane) {
+typedef struct encode_txb_args {
+ const AV1_COMMON *cm;
+ MACROBLOCK *x;
+ aom_writer *w;
+} ENCODE_TXB_ARGS;
+
+static void av1_write_coeffs_txb_wrap(int plane, int block, int blk_row,
+ int blk_col, BLOCK_SIZE plane_bsize,
+ TX_SIZE tx_size, void *arg) {
+ (void)plane_bsize;
+ ENCODE_TXB_ARGS *enc_args = (ENCODE_TXB_ARGS *)arg;
+ const AV1_COMMON *cm = enc_args->cm;
+ MACROBLOCK *x = enc_args->x;
MACROBLOCKD *xd = &x->e_mbd;
- MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
- BLOCK_SIZE bsize = mbmi->sb_type;
- struct macroblockd_plane *pd = &xd->plane[plane];
- const BLOCK_SIZE plane_bsize =
- AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
- const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
- const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
- const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
- const int bkw = tx_size_wide_unit[tx_size];
- const int bkh = tx_size_high_unit[tx_size];
- const int step = tx_size_wide_unit[tx_size] * tx_size_high_unit[tx_size];
- int row, col;
- int block = 0;
- for (row = 0; row < max_blocks_high; row += bkh) {
- for (col = 0; col < max_blocks_wide; col += bkw) {
- 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, row, col, block, plane, tx_size, tcoeff,
- eob, &txb_ctx);
- block += step;
- }
- }
+ aom_writer *w = enc_args->w;
+ 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, blk_row, blk_col, block, plane, tx_size,
+ tcoeff, eob, &txb_ctx);
+}
+
+void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
+ aom_writer *w, int plane, BLOCK_SIZE bsize) {
+ ENCODE_TXB_ARGS enc_args = { cm, x, w };
+ MACROBLOCKD *xd = &x->e_mbd;
+ av1_foreach_transformed_block_in_plane(xd, bsize, plane,
+ av1_write_coeffs_txb_wrap, &enc_args);
}
static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
diff --git a/av1/encoder/encodetxb.h b/av1/encoder/encodetxb.h
index f7edd35..2db05d6 100644
--- a/av1/encoder/encodetxb.h
+++ b/av1/encoder/encodetxb.h
@@ -78,7 +78,7 @@
int plane, TX_SIZE tx_size, const tran_low_t *tcoeff,
uint16_t eob, TXB_CTX *txb_ctx);
void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
- aom_writer *w, int plane);
+ aom_writer *w, int plane, BLOCK_SIZE bsize);
int av1_get_txb_entropy_context(const tran_low_t *qcoeff,
const SCAN_ORDER *scan_order, int eob);
void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,
