Rewrite functions related to first pass block stats

Change-Id: I28679f88e2911b06eef5cbc83ecb62b8c69e4c53

vp8/vp8_cx_iface.c

@@ -1291,6 +1291,7 @@ static vpx_codec_enc_cfg_map_t vp8e_usage_cfg_map[] =
         VPX_VBR,            /* rc_end_usage */
 #if VPX_ENCODER_ABI_VERSION > (1 + VPX_CODEC_ABI_VERSION)
         {0},                /* rc_twopass_stats_in */
+        {0},                /* rc_firstpass_mb_stats_in */
 #endif
         256,                /* rc_target_bandwidth */
         4,                  /* rc_min_quantizer */

vp9/encoder/vp9_encoder.c

@@ -200,13 +200,6 @@ static void dealloc_compressor_data(VP9_COMP *cpi) {
     cpi->source_diff_var = NULL;
   }
 
-#if CONFIG_FP_MB_STATS
-  if (cpi->use_fp_mb_stats) {
-    vpx_free(cpi->twopass.this_frame_mb_stats.mb_stats);
-    cpi->twopass.this_frame_mb_stats.mb_stats = NULL;
-  }
-#endif
-
   for (i = 0; i < MAX_LAG_BUFFERS; ++i) {
     vp9_free_frame_buffer(&cpi->svc.scaled_frames[i]);
   }
@@ -793,11 +786,11 @@ VP9_COMP *vp9_create_compressor(VP9EncoderConfig *oxcf) {
 #if CONFIG_FP_MB_STATS
   cpi->use_fp_mb_stats = 0;
   if (cpi->use_fp_mb_stats) {
-    // a place holder for the mb stats obtained from the first pass
-    CHECK_MEM_ERROR(cm, cpi->twopass.this_frame_mb_stats.mb_stats,
-                    vpx_calloc(cm->MBs * sizeof(FIRSTPASS_MB_STATS), 1));
+    // a place holder used to store the first pass mb stats in the first pass
+    CHECK_MEM_ERROR(cm, cpi->twopass.frame_mb_stats_buf,
+                    vpx_calloc(cm->MBs * sizeof(uint8_t), 1));
   } else {
-    cpi->twopass.this_frame_mb_stats.mb_stats = NULL;
+    cpi->twopass.frame_mb_stats_buf = NULL;
   }
 #endif
 
@@ -940,6 +933,21 @@ VP9_COMP *vp9_create_compressor(VP9EncoderConfig *oxcf) {
 
       vp9_init_second_pass_spatial_svc(cpi);
     } else {
+#if CONFIG_FP_MB_STATS
+      if (cpi->use_fp_mb_stats) {
+        const size_t psz = cpi->common.MBs * sizeof(uint8_t);
+        const int ps = (int)(oxcf->firstpass_mb_stats_in.sz / psz);
+
+        cpi->twopass.firstpass_mb_stats.mb_stats_start =
+            oxcf->firstpass_mb_stats_in.buf;
+        cpi->twopass.firstpass_mb_stats.mb_stats_in =
+            cpi->twopass.firstpass_mb_stats.mb_stats_start;
+        cpi->twopass.firstpass_mb_stats.mb_stats_end =
+            cpi->twopass.firstpass_mb_stats.mb_stats_start +
+            (ps - 1) * cpi->common.MBs * sizeof(uint8_t);
+      }
+#endif
+
       cpi->twopass.stats_in_start = oxcf->two_pass_stats_in.buf;
       cpi->twopass.stats_in = cpi->twopass.stats_in_start;
       cpi->twopass.stats_in_end = &cpi->twopass.stats_in[packets - 1];
@@ -1130,6 +1138,13 @@ void vp9_remove_compressor(VP9_COMP *cpi) {
     vpx_free(cpi->mbgraph_stats[i].mb_stats);
   }
 
+#if CONFIG_FP_MB_STATS
+  if (cpi->use_fp_mb_stats) {
+    vpx_free(cpi->twopass.frame_mb_stats_buf);
+    cpi->twopass.frame_mb_stats_buf = NULL;
+  }
+#endif
+
   vp9_remove_common(&cpi->common);
   vpx_free(cpi);
 

vp9/encoder/vp9_encoder.h

@@ -227,6 +227,10 @@ typedef struct VP9EncoderConfig {
   struct vpx_fixed_buf         two_pass_stats_in;
   struct vpx_codec_pkt_list  *output_pkt_list;
 
+#if CONFIG_FP_MB_STATS
+  struct vpx_fixed_buf         firstpass_mb_stats_in;
+#endif
+
   vp8e_tuning tuning;
 } VP9EncoderConfig;
 

vp9/encoder/vp9_firstpass.c

@@ -98,34 +98,6 @@ static const FIRSTPASS_STATS *read_frame_stats(const TWO_PASS *p, int offset) {
   return &p->stats_in[offset];
 }
 
-#if CONFIG_FP_MB_STATS
-static int input_mb_stats(FIRSTPASS_FRAME_MB_STATS *fp_frame_stats,
-                          const VP9_COMMON *const cm) {
-  FILE *fpfile;
-  int ret;
-
-  fpfile = fopen("firstpass_mb.stt", "r");
-  fseek(fpfile, cm->current_video_frame * cm->MBs * sizeof(FIRSTPASS_MB_STATS),
-        SEEK_SET);
-  ret = fread(fp_frame_stats->mb_stats, sizeof(FIRSTPASS_MB_STATS), cm->MBs,
-              fpfile);
-  fclose(fpfile);
-  if (ret < cm->MBs) {
-    return EOF;
-  }
-  return 1;
-}
-
-static void output_mb_stats(FIRSTPASS_FRAME_MB_STATS *fp_frame_stats,
-                          const VP9_COMMON *const cm) {
-  FILE *fpfile;
-
-  fpfile = fopen("firstpass_mb.stt", "a");
-  fwrite(fp_frame_stats->mb_stats, sizeof(FIRSTPASS_MB_STATS), cm->MBs, fpfile);
-  fclose(fpfile);
-}
-#endif
-
 static int input_stats(TWO_PASS *p, FIRSTPASS_STATS *fps) {
   if (p->stats_in >= p->stats_in_end)
     return EOF;
@@ -175,6 +147,27 @@ static void output_stats(FIRSTPASS_STATS *stats,
 #endif
 }
 
+#if CONFIG_FP_MB_STATS
+static int input_fpmb_stats(FIRSTPASS_MB_STATS *firstpass_mb_stats,
+                            VP9_COMMON *cm, uint8_t **this_frame_mb_stats) {
+  if (firstpass_mb_stats->mb_stats_in > firstpass_mb_stats->mb_stats_end)
+    return EOF;
+
+  *this_frame_mb_stats = firstpass_mb_stats->mb_stats_in;
+  firstpass_mb_stats->mb_stats_in += cm->MBs * sizeof(uint8_t);
+  return 1;
+}
+
+static void output_fpmb_stats(uint8_t *this_frame_mb_stats, VP9_COMMON *cm,
+                         struct vpx_codec_pkt_list *pktlist) {
+  struct vpx_codec_cx_pkt pkt;
+  pkt.kind = VPX_CODEC_FPMB_STATS_PKT;
+  pkt.data.firstpass_mb_stats.buf = this_frame_mb_stats;
+  pkt.data.firstpass_mb_stats.sz = cm->MBs * sizeof(uint8_t);
+  vpx_codec_pkt_list_add(pktlist, &pkt);
+}
+#endif
+
 static void zero_stats(FIRSTPASS_STATS *section) {
   section->frame      = 0.0;
   section->intra_error = 0.0;
@@ -473,7 +466,9 @@ void vp9_first_pass(VP9_COMP *cpi) {
   const YV12_BUFFER_CONFIG *first_ref_buf = lst_yv12;
 
 #if CONFIG_FP_MB_STATS
-  FIRSTPASS_FRAME_MB_STATS *this_frame_mb_stats = &twopass->this_frame_mb_stats;
+  if (cpi->use_fp_mb_stats) {
+    vp9_zero_array(cpi->twopass.frame_mb_stats_buf, cm->MBs);
+  }
 #endif
 
   vp9_clear_system_state();
@@ -614,12 +609,7 @@ void vp9_first_pass(VP9_COMP *cpi) {
 
 #if CONFIG_FP_MB_STATS
       if (cpi->use_fp_mb_stats) {
-        this_frame_mb_stats->mb_stats[mb_row * cm->mb_cols + mb_col].mode =
-            DC_PRED;
-        this_frame_mb_stats->mb_stats[mb_row * cm->mb_cols + mb_col].err =
-            this_error;
-        this_frame_mb_stats->mb_stats[mb_row * cm->mb_cols + mb_col].mv.as_int
-            = 0;
+        // TODO(pengchong): store some related block statistics here
       }
 #endif
 
@@ -750,12 +740,7 @@ void vp9_first_pass(VP9_COMP *cpi) {
 
 #if CONFIG_FP_MB_STATS
           if (cpi->use_fp_mb_stats) {
-            this_frame_mb_stats->mb_stats[mb_row * cm->mb_cols + mb_col].mode =
-                NEWMV;
-            this_frame_mb_stats->mb_stats[mb_row * cm->mb_cols + mb_col].err =
-                motion_error;
-            this_frame_mb_stats->mb_stats[mb_row * cm->mb_cols + mb_col].mv.
-                as_int = mv.as_int;
+            // TODO(pengchong): save some related block statistics here
           }
 #endif
 
@@ -866,7 +851,7 @@ void vp9_first_pass(VP9_COMP *cpi) {
 
 #if CONFIG_FP_MB_STATS
     if (cpi->use_fp_mb_stats) {
-      output_mb_stats(this_frame_mb_stats, cm);
+      output_fpmb_stats(twopass->frame_mb_stats_buf, cm, cpi->output_pkt_list);
     }
 #endif
   }
@@ -2250,7 +2235,8 @@ void vp9_rc_get_second_pass_params(VP9_COMP *cpi) {
 
 #if CONFIG_FP_MB_STATS
   if (cpi->use_fp_mb_stats) {
-    input_mb_stats(&twopass->this_frame_mb_stats, cm);
+    input_fpmb_stats(&twopass->firstpass_mb_stats, cm,
+                     &twopass->this_frame_mb_stats);
   }
 #endif
 }

vp9/encoder/vp9_firstpass.h

@@ -20,14 +20,10 @@ extern "C" {
 
 #if CONFIG_FP_MB_STATS
 typedef struct {
-  PREDICTION_MODE mode;
-  int err;
-  int_mv mv;
+  uint8_t *mb_stats_in;
+  uint8_t *mb_stats_start;
+  uint8_t *mb_stats_end;
 } FIRSTPASS_MB_STATS;
-
-typedef struct {
-  FIRSTPASS_MB_STATS *mb_stats;
-} FIRSTPASS_FRAME_MB_STATS;
 #endif
 
 typedef struct {
@@ -89,7 +85,9 @@ typedef struct {
   double gf_intra_err_min;
 
 #if CONFIG_FP_MB_STATS
-  FIRSTPASS_FRAME_MB_STATS this_frame_mb_stats;
+  uint8_t *frame_mb_stats_buf;
+  uint8_t *this_frame_mb_stats;
+  FIRSTPASS_MB_STATS firstpass_mb_stats;
 #endif
 
   // Projected total bits available for a key frame group of frames

vp9/vp9_cx_iface.c

@@ -274,6 +274,7 @@ static vpx_codec_err_t validate_config(vpx_codec_alg_priv_t *ctx,
         ERROR("rc_twopass_stats_in missing EOS stats packet");
     }
   }
+
   if (cfg->g_profile <= (unsigned int)PROFILE_1 &&
       extra_cfg->bit_depth > BITS_8)
     ERROR("High bit-depth not supported in profile < 2");
@@ -376,6 +377,10 @@ static vpx_codec_err_t set_encoder_config(
   oxcf->two_pass_stats_in      =  cfg->rc_twopass_stats_in;
   oxcf->output_pkt_list        =  extra_cfg->pkt_list;
 
+#if CONFIG_FP_MB_STATS
+  oxcf->firstpass_mb_stats_in  = cfg->rc_firstpass_mb_stats_in;
+#endif
+
   oxcf->arnr_max_frames = extra_cfg->arnr_max_frames;
   oxcf->arnr_strength   = extra_cfg->arnr_strength;
   oxcf->arnr_type       = extra_cfg->arnr_type;
@@ -666,6 +671,7 @@ static vpx_codec_err_t encoder_init(vpx_codec_ctx_t *ctx,
 
     priv->extra_cfg = extracfg_map[i].cfg;
     priv->extra_cfg.pkt_list = &priv->pkt_list.head;
+
      // Maximum buffer size approximated based on having multiple ARF.
     priv->cx_data_sz = priv->cfg.g_w * priv->cfg.g_h * 3 / 2 * 8;
 
@@ -1259,6 +1265,7 @@ static vpx_codec_enc_cfg_map_t encoder_usage_cfg_map[] = {
       VPX_VBR,            // rc_end_usage
 #if VPX_ENCODER_ABI_VERSION > (1 + VPX_CODEC_ABI_VERSION)
       {NULL, 0},          // rc_twopass_stats_in
+      {NULL, 0},          // rc_firstpass_mb_stats_in
 #endif
       256,                // rc_target_bandwidth
       0,                  // rc_min_quantizer

vpx/vpx_encoder.h

@@ -155,6 +155,7 @@ extern "C" {
   enum vpx_codec_cx_pkt_kind {
     VPX_CODEC_CX_FRAME_PKT,    /**< Compressed video frame */
     VPX_CODEC_STATS_PKT,       /**< Two-pass statistics for this frame */
+    VPX_CODEC_FPMB_STATS_PKT,  /**< first pass mb statistics for this frame */
     VPX_CODEC_PSNR_PKT,        /**< PSNR statistics for this frame */
 #ifdef CONFIG_SPATIAL_SVC
     VPX_CODEC_SPATIAL_SVC_LAYER_SIZES, /**< Sizes for each layer in this frame*/
@@ -188,6 +189,7 @@ extern "C" {
 
       } frame;  /**< data for compressed frame packet */
       struct vpx_fixed_buf twopass_stats;  /**< data for two-pass packet */
+      struct vpx_fixed_buf firstpass_mb_stats; /**< first pass mb packet */
       struct vpx_psnr_pkt {
         unsigned int samples[4];  /**< Number of samples, total/y/u/v */
         uint64_t     sse[4];      /**< sum squared error, total/y/u/v */
@@ -452,6 +454,12 @@ extern "C" {
      */
     struct vpx_fixed_buf   rc_twopass_stats_in;
 
+    /*!\brief first pass mb stats buffer.
+     *
+     * A buffer containing all of the first pass mb stats packets produced
+     * in the first pass, concatenated.
+     */
+    struct vpx_fixed_buf   rc_firstpass_mb_stats_in;
 
     /*!\brief Target data rate
      *

vpxenc.c

@@ -141,6 +141,10 @@ static const arg_def_t pass_arg         = ARG_DEF(NULL, "pass", 1,
                                                   "Pass to execute (1/2)");
 static const arg_def_t fpf_name         = ARG_DEF(NULL, "fpf", 1,
                                                   "First pass statistics file name");
+#if CONFIG_FP_MB_STATS
+static const arg_def_t fpmbf_name         = ARG_DEF(NULL, "fpmbf", 1,
+                                      "First pass block statistics file name");
+#endif
 static const arg_def_t limit = ARG_DEF(NULL, "limit", 1,
                                        "Stop encoding after n input frames");
 static const arg_def_t skip = ARG_DEF(NULL, "skip", 1,
@@ -572,6 +576,9 @@ struct stream_config {
   struct vpx_codec_enc_cfg  cfg;
   const char               *out_fn;
   const char               *stats_fn;
+#if CONFIG_FP_MB_STATS
+  const char               *fpmb_stats_fn;
+#endif
   stereo_format_t           stereo_fmt;
   int                       arg_ctrls[ARG_CTRL_CNT_MAX][2];
   int                       arg_ctrl_cnt;
@@ -597,6 +604,9 @@ struct stream_state {
   uint64_t                  cx_time;
   size_t                    nbytes;
   stats_io_t                stats;
+#if CONFIG_FP_MB_STATS
+  stats_io_t                fpmb_stats;
+#endif
   struct vpx_image         *img;
   vpx_codec_ctx_t           decoder;
   int                       mismatch_seen;
@@ -873,6 +883,10 @@ static int parse_stream_params(struct VpxEncoderConfig *global,
       config->out_fn = arg.val;
     } else if (arg_match(&arg, &fpf_name, argi)) {
       config->stats_fn = arg.val;
+#if CONFIG_FP_MB_STATS
+    } else if (arg_match(&arg, &fpmbf_name, argi)) {
+      config->fpmb_stats_fn = arg.val;
+#endif
     } else if (arg_match(&arg, &use_ivf, argi)) {
       config->write_webm = 0;
     } else if (arg_match(&arg, &threads, argi)) {
@@ -1029,6 +1043,17 @@ static void validate_stream_config(const struct stream_state *stream,
         fatal("Stream %d: duplicate stats file (from stream %d)",
               streami->index, stream->index);
     }
+
+#if CONFIG_FP_MB_STATS
+    /* Check for two streams sharing a mb stats file. */
+    if (streami != stream) {
+      const char *a = stream->config.fpmb_stats_fn;
+      const char *b = streami->config.fpmb_stats_fn;
+      if (a && b && !strcmp(a, b))
+        fatal("Stream %d: duplicate mb stats file (from stream %d)",
+              streami->index, stream->index);
+    }
+#endif
   }
 }
 
@@ -1200,11 +1225,27 @@ static void setup_pass(struct stream_state *stream,
       fatal("Failed to open statistics store");
   }
 
+#if CONFIG_FP_MB_STATS
+  if (stream->config.fpmb_stats_fn) {
+    if (!stats_open_file(&stream->fpmb_stats,
+                         stream->config.fpmb_stats_fn, pass))
+      fatal("Failed to open mb statistics store");
+  } else {
+    if (!stats_open_mem(&stream->fpmb_stats, pass))
+      fatal("Failed to open mb statistics store");
+  }
+#endif
+
   stream->config.cfg.g_pass = global->passes == 2
                               ? pass ? VPX_RC_LAST_PASS : VPX_RC_FIRST_PASS
                             : VPX_RC_ONE_PASS;
-  if (pass)
+  if (pass) {
     stream->config.cfg.rc_twopass_stats_in = stats_get(&stream->stats);
+#if CONFIG_FP_MB_STATS
+    stream->config.cfg.rc_firstpass_mb_stats_in =
+        stats_get(&stream->fpmb_stats);
+#endif
+  }
 
   stream->cx_time = 0;
   stream->nbytes = 0;
@@ -1388,6 +1429,14 @@ static void get_cx_data(struct stream_state *stream,
                     pkt->data.twopass_stats.sz);
         stream->nbytes += pkt->data.raw.sz;
         break;
+#if CONFIG_FP_MB_STATS
+      case VPX_CODEC_FPMB_STATS_PKT:
+        stats_write(&stream->fpmb_stats,
+                    pkt->data.firstpass_mb_stats.buf,
+                    pkt->data.firstpass_mb_stats.sz);
+        stream->nbytes += pkt->data.raw.sz;
+        break;
+#endif
       case VPX_CODEC_PSNR_PKT:
 
         if (global->show_psnr) {
@@ -1778,6 +1827,10 @@ int main(int argc, const char **argv_) {
 
     FOREACH_STREAM(stats_close(&stream->stats, global.passes - 1));
 
+#if CONFIG_FP_MB_STATS
+    FOREACH_STREAM(stats_close(&stream->fpmb_stats, global.passes - 1));
+#endif
+
     if (global.pass)
       break;
   }