Add av1_write_coeffs_mb

This function go through each transform block in the
prediction block and call av1_write_coeffs_txb to
pack coefficients into the bitstream

Change-Id: I6dedebef6cf8957f9173241a7de60e9936bc0be8

av1/encoder/bitstream.c

@@ -43,6 +43,9 @@
 #if CONFIG_ANS
 #include "aom_dsp/buf_ans.h"
 #endif  // CONFIG_ANS
+#if CONFIG_LV_MAP
+#include "av1/encoder/encodetxb.h"
+#endif  // CONFIG_LV_MAP
 #include "av1/encoder/bitstream.h"
 #include "av1/encoder/cost.h"
 #include "av1/encoder/encodemv.h"
@@ -881,7 +884,8 @@ static void pack_mb_tokens(aom_writer *w, const TOKENEXTRA **tp,
 
   *tp = p;
 }
-#else
+#else  //  CONFIG_NEW_TOKENSET
+#if !CONFIG_LV_MAP
 static void pack_mb_tokens(aom_writer *w, const TOKENEXTRA **tp,
                            const TOKENEXTRA *const stop,
                            aom_bit_depth_t bit_depth, const TX_SIZE tx_size,
@@ -977,7 +981,8 @@ static void pack_mb_tokens(aom_writer *w, const TOKENEXTRA **tp,
 
   *tp = p;
 }
-#endif
+#endif  // !CONFIG_LV_MAP
+#endif  // CONFIG_NEW_TOKENSET
 #else   // !CONFIG_PVQ
 static PVQ_INFO *get_pvq_block(PVQ_QUEUE *pvq_q) {
   PVQ_INFO *pvq;
@@ -2142,7 +2147,7 @@ static void write_tokens_b(AV1_COMP *cpi, const TileInfo *const tile,
   MB_MODE_INFO *const mbmi = &m->mbmi;
   int plane;
   int bh, bw;
-#if CONFIG_PVQ
+#if CONFIG_PVQ || CONFIG_LV_MAP
   MACROBLOCK *const x = &cpi->td.mb;
   (void)tok;
   (void)tok_end;
@@ -2329,7 +2334,13 @@ static void write_tokens_b(AV1_COMP *cpi, const TileInfo *const tile,
       TOKEN_STATS token_stats;
 #if !CONFIG_PVQ
       init_token_stats(&token_stats);
+#if CONFIG_LV_MAP
+      (void)tx;
+      av1_write_coeffs_mb(cm, x, w, plane);
+#else   // CONFIG_LV_MAP
       pack_mb_tokens(w, tok, tok_end, cm->bit_depth, tx, &token_stats);
+#endif  // CONFIG_LV_MAP
+
 #else
       (void)token_stats;
       pack_pvq_tokens(w, x, xd, plane, mbmi->sb_type, tx);

av1/encoder/encodetxb.c

@@ -13,7 +13,6 @@
 #include "av1/common/blockd.h"
 #include "av1/common/pred_common.h"
 #include "av1/encoder/cost.h"
-#include "av1/encoder/encoder.h"
 #include "av1/encoder/encodetxb.h"
 #include "av1/encoder/tokenize.h"
 
@@ -188,6 +187,39 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
   }
 }
 
+void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
+                         aom_writer *w, int plane) {
+  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];
+
+#if CONFIG_CB4X4
+  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
+#else
+  const BLOCK_SIZE plane_bsize =
+      get_plane_block_size(AOMMAX(bsize, BLOCK_8X8), pd);
+#endif
+  const int num_4x4_w = block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
+  const int num_4x4_h = block_size_high[plane_bsize] >> tx_size_high_log2[0];
+  TX_SIZE tx_size = 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 < num_4x4_h; row += bkh) {
+    for (col = 0; col < num_4x4_w; 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, block, plane, tcoeff, eob, &txb_ctx);
+      block += step;
+    }
+  }
+}
+
 static INLINE void get_base_ctx_set(const tran_low_t *tcoeffs,
                                     int c,  // raster order
                                     const int bwl,
@@ -539,7 +571,10 @@ void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,
     td->counts->skip[ctx][0] += skip_inc;
     av1_foreach_transformed_block(xd, bsize, update_and_record_txb_context,
                                   &arg);
-  } else {
+  } else if (dry_run == DRY_RUN_NORMAL) {
     av1_foreach_transformed_block(xd, bsize, update_txb_context, &arg);
+  } else {
+    printf("DRY_RUN_COSTCOEFFS is not supported yet\n");
+    assert(0);
   }
 }

av1/encoder/encodetxb.h

@@ -17,6 +17,7 @@
 #include "av1/common/onyxc_int.h"
 #include "av1/common/txb_common.h"
 #include "av1/encoder/block.h"
+#include "av1/encoder/encoder.h"
 #include "aom_dsp/bitwriter.h"
 #ifdef __cplusplus
 extern "C" {
@@ -28,7 +29,8 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
                           aom_writer *w, int block, int plane,
                           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);
 void av1_update_txb_context(const AV1_COMP *cpi, ThreadData *td,
                             RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
                             const int mi_row, const int mi_col);