Add frame level buffer to store txfm coeffs

Doing tokenize in the last step of RD loop and then doing packing
tokens in bitstream packing phase is hard for debugging.

Therefore, we create a frame-level buffer to store the txfm coeffs
from the reconstruction in RD loop and then in bitstream packing
phase, we can code the txfm coeffs directly.

Change-Id: I999470eef6e038317a91585df2bdfc20aca3573e

av1/encoder/block.h

@@ -67,6 +67,9 @@ typedef unsigned int av1_coeff_cost[PLANE_TYPES][REF_TYPES][COEF_BANDS][2]
 typedef struct {
   int_mv ref_mvs[MODE_CTX_REF_FRAMES][MAX_MV_REF_CANDIDATES];
   int16_t mode_context[MODE_CTX_REF_FRAMES];
+#if CONFIG_LV_MAP
+  tran_low_t *tcoeff[MAX_MB_PLANE];
+#endif
 #if CONFIG_REF_MV
   uint8_t ref_mv_count[MODE_CTX_REF_FRAMES];
   CANDIDATE_MV ref_mv_stack[MODE_CTX_REF_FRAMES][MAX_REF_MV_STACK_SIZE];

av1/encoder/encoder.c

@@ -39,6 +39,9 @@
 #include "av1/encoder/encodeframe.h"
 #include "av1/encoder/encodemv.h"
 #include "av1/encoder/encoder.h"
+#if CONFIG_LV_MAP
+#include "av1/encoder/encodetxb.h"
+#endif
 #include "av1/encoder/ethread.h"
 #include "av1/encoder/firstpass.h"
 #include "av1/encoder/mbgraph.h"
@@ -451,6 +454,9 @@ static void dealloc_compressor_data(AV1_COMP *cpi) {
     aom_free_frame_buffer(&cpi->upsampled_ref_bufs[i].buf);
 
   av1_free_ref_frame_buffers(cm->buffer_pool);
+#if CONFIG_LV_MAP
+  av1_free_txb_buf(cpi);
+#endif
   av1_free_context_buffers(cm);
 
   aom_free_frame_buffer(&cpi->last_frame_uf);
@@ -789,6 +795,10 @@ void av1_alloc_compressor_data(AV1_COMP *cpi) {
 
   av1_alloc_context_buffers(cm, cm->width, cm->height);
 
+#if CONFIG_LV_MAP
+  av1_alloc_txb_buf(cpi);
+#endif
+
   alloc_context_buffers_ext(cpi);
 
   aom_free(cpi->tile_tok[0][0]);
@@ -922,6 +932,9 @@ static void update_frame_size(AV1_COMP *cpi) {
                        NULL);
   memset(cpi->mbmi_ext_base, 0,
          cm->mi_rows * cm->mi_cols * sizeof(*cpi->mbmi_ext_base));
+#if CONFIG_LV_MAP
+  av1_reset_txb_buf(cpi);
+#endif
 
   set_tile_info(cpi);
 }

av1/encoder/encoder.h

@@ -667,6 +667,9 @@ typedef struct AV1_COMP {
 #if CONFIG_REFERENCE_BUFFER
   SequenceHeader seq_params;
 #endif
+#if CONFIG_LV_MAP
+  tran_low_t *tcoeff_buf[MAX_MB_PLANE];
+#endif
 } AV1_COMP;
 
 void av1_initialize_enc(void);

av1/encoder/encodetxb.c

@@ -11,8 +11,60 @@
 
 #include "av1/common/scan.h"
 #include "av1/encoder/cost.h"
+#include "av1/encoder/encoder.h"
 #include "av1/encoder/encodetxb.h"
 
+void av1_alloc_txb_buf(AV1_COMP *cpi) {
+  AV1_COMMON *cm = &cpi->common;
+  int mi_block_size = 1 << MI_SIZE_LOG2;
+  // TODO(angiebird): Make sure cm->subsampling_x/y is set correctly, and then
+  // use precise buffer size according to cm->subsampling_x/y
+  int pixel_stride = mi_block_size * cm->mi_cols;
+  int pixel_height = mi_block_size * cm->mi_rows;
+  int i;
+  for (i = 0; i < MAX_MB_PLANE; ++i) {
+    CHECK_MEM_ERROR(
+        cm, cpi->tcoeff_buf[i],
+        aom_malloc(sizeof(*cpi->tcoeff_buf[i]) * pixel_stride * pixel_height));
+  }
+}
+
+void av1_reset_txb_buf(AV1_COMP *cpi) {
+  AV1_COMMON *cm = &cpi->common;
+  int mi_block_size = 1 << MI_SIZE_LOG2;
+  int plane;
+  for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
+    // TODO(angiebird): Set the subsampling_x/y to cm->subsampling_x/y
+    int subsampling_x = 0;
+    int subsampling_y = 0;
+    int pixel_stride = (mi_block_size * cm->mi_cols) >> subsampling_x;
+    int pixel_height = (mi_block_size * cm->mi_rows) >> subsampling_y;
+
+    // TODO(angiebird): Check if we really need this initialization
+    memset(cpi->tcoeff_buf[plane], 0,
+           pixel_stride * pixel_height * sizeof(*cpi->tcoeff_buf[plane]));
+
+    int mi_row, mi_col;
+    for (mi_row = 0; mi_row < cm->mi_rows; ++mi_row) {
+      for (mi_col = 0; mi_col < cm->mi_cols; ++mi_col) {
+        MB_MODE_INFO_EXT *mbmi_ext =
+            cpi->mbmi_ext_base + mi_row * cm->mi_cols + mi_col;
+        int pixel_row = (mi_block_size * mi_row) >> subsampling_y;
+        int pixel_col = (mi_block_size * mi_col) >> subsampling_x;
+        mbmi_ext->tcoeff[plane] =
+            cpi->tcoeff_buf[plane] + pixel_row * pixel_stride + pixel_col;
+      }
+    }
+  }
+}
+
+void av1_free_txb_buf(AV1_COMP *cpi) {
+  int i;
+  for (i = 0; i < MAX_MB_PLANE; ++i) {
+    aom_free(cpi->tcoeff_buf[i]);
+  }
+}
+
 static void write_golomb(aom_writer *w, int level) {
   int x = level + 1;
   int i = x;

av1/encoder/encodetxb.h

@@ -21,7 +21,9 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
-
+void av1_alloc_txb_buf(AV1_COMP *cpi);
+void av1_reset_txb_buf(AV1_COMP *cpi);
+void av1_free_txb_buf(AV1_COMP *cpi);
 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,