Prepare for sb level multi-pass encoding

This starts the process of adding multipass encoding at the superblock
level.

At the first stage of the process, we add the ability to encode
the same block twice while maintaining the same output to ensure that
there is no state leakage from pass to pass.

BUG=aomedia:2573

Change-Id: Ib9b7eaf7d47a658d44a1bbe8b435b671548fedd0
diff --git a/aom/aomcx.h b/aom/aomcx.h
index a8202fd..8007ed2 100644
--- a/aom/aomcx.h
+++ b/aom/aomcx.h
@@ -1189,6 +1189,14 @@
    * 0 : off, 1 : enable EXT_TILE_DEBUG
    */
   AV1E_ENABLE_EXT_TILE_DEBUG = 154,
+
+  /*!\brief Codec control function to enable the superblock multipass unit test
+   * in AV1 to ensure that the encoder does not leak state between different
+   * passes. Please note that this is only used in sb_multipass unit test.
+   *
+   * 0 : off, 1 : on
+   */
+  AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST = 155,
 };
 
 /*!\brief aom 1-D scaling mode
@@ -1675,6 +1683,9 @@
 AOM_CTRL_USE_TYPE(AV1E_SET_SVC_REF_FRAME_CONFIG, aom_svc_ref_frame_config_t *)
 #define AOME_CTRL_AV1E_SET_SVC_REF_FRAME_CONFIG
 
+AOM_CTRL_USE_TYPE(AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST, unsigned int)
+#define AOM_CTRL_AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST
+
 /*!\endcond */
 /*! @} - end defgroup aom_encoder */
 #ifdef __cplusplus
diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c
index 16eeec6..6e9a6b5 100644
--- a/av1/av1_cx_iface.c
+++ b/av1/av1_cx_iface.c
@@ -153,6 +153,7 @@
   COST_UPDATE_TYPE mode_cost_upd_freq;
   COST_UPDATE_TYPE mv_cost_upd_freq;
   unsigned int ext_tile_debug;
+  unsigned int sb_multipass_unit_test;
 };
 
 static struct av1_extracfg default_extra_cfg = {
@@ -279,6 +280,7 @@
   COST_UPD_SB,  // mode_cost_upd_freq
   COST_UPD_SB,  // mv_cost_upd_freq
   0,            // ext_tile_debug
+  0,            // sb_multipass_unit_test
 };
 
 struct aom_codec_alg_priv {
@@ -424,6 +426,7 @@
         "or kf_max_dist instead.");
 
   RANGE_CHECK_HI(extra_cfg, motion_vector_unit_test, 2);
+  RANGE_CHECK_HI(extra_cfg, sb_multipass_unit_test, 1);
   RANGE_CHECK_HI(extra_cfg, ext_tile_debug, 1);
   RANGE_CHECK_HI(extra_cfg, enable_auto_alt_ref, 1);
   RANGE_CHECK_HI(extra_cfg, enable_auto_bwd_ref, 2);
@@ -975,6 +978,7 @@
 
   oxcf->frame_periodic_boost = extra_cfg->frame_periodic_boost;
   oxcf->motion_vector_unit_test = extra_cfg->motion_vector_unit_test;
+  oxcf->sb_multipass_unit_test = extra_cfg->sb_multipass_unit_test;
   oxcf->ext_tile_debug = extra_cfg->ext_tile_debug;
 
   oxcf->chroma_subsampling_x = extra_cfg->chroma_subsampling_x;
@@ -1770,6 +1774,15 @@
   extra_cfg.min_cr = CAST(AV1E_SET_MIN_CR, args);
   return update_extra_cfg(ctx, &extra_cfg);
 }
+
+static aom_codec_err_t ctrl_enable_sb_multipass_unit_test(
+    aom_codec_alg_priv_t *ctx, va_list args) {
+  struct av1_extracfg extra_cfg = ctx->extra_cfg;
+  extra_cfg.sb_multipass_unit_test =
+      CAST(AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST, args);
+  return update_extra_cfg(ctx, &extra_cfg);
+}
+
 static aom_codec_err_t create_frame_stats_buffer(
     FIRSTPASS_STATS **frame_stats_buffer, STATS_BUFFER_CTX *stats_buf_context,
     int num_lap_buffers) {
@@ -2682,6 +2695,7 @@
   { AV1E_SET_SVC_LAYER_ID, ctrl_set_layer_id },
   { AV1E_SET_SVC_PARAMS, ctrl_set_svc_params },
   { AV1E_SET_SVC_REF_FRAME_CONFIG, ctrl_set_svc_ref_frame_config },
+  { AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST, ctrl_enable_sb_multipass_unit_test },
 
   // Getters
   { AOME_GET_LAST_QUANTIZER, ctrl_get_quantizer },
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 34894e0..e4ac06e 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -146,8 +146,45 @@
   128 * 16, 128 * 16
 };
 
+typedef struct {
+  ENTROPY_CONTEXT a[MAX_MIB_SIZE * MAX_MB_PLANE];
+  ENTROPY_CONTEXT l[MAX_MIB_SIZE * MAX_MB_PLANE];
+  PARTITION_CONTEXT sa[MAX_MIB_SIZE];
+  PARTITION_CONTEXT sl[MAX_MIB_SIZE];
+  TXFM_CONTEXT *p_ta;
+  TXFM_CONTEXT *p_tl;
+  TXFM_CONTEXT ta[MAX_MIB_SIZE];
+  TXFM_CONTEXT tl[MAX_MIB_SIZE];
+} RD_SEARCH_MACROBLOCK_CONTEXT;
+
 enum { PICK_MODE_RD = 0, PICK_MODE_NONRD };
 
+enum {
+  SB_SINGLE_PASS,  // Single pass encoding: all ctxs get updated normally
+  SB_DRY_PASS,     // First pass of multi-pass: does not update the ctxs
+  SB_WET_PASS      // Second pass of multi-pass: finalize and update the ctx
+} UENUM1BYTE(SB_MULTI_PASS_MODE);
+
+// This struct is used to store the statistics used by sb-level multi-pass
+// encoding. Currently, this is only used to make a copy of the state before we
+// perform the first pass
+typedef struct SB_FIRST_PASS_STATS {
+  RD_SEARCH_MACROBLOCK_CONTEXT x_ctx;
+  RD_COUNTS rd_count;
+
+  int split_count;
+  FRAME_COUNTS fc;
+  InterModeRdModel inter_mode_rd_models[BLOCK_SIZES_ALL];
+  int thresh_freq_fact[BLOCK_SIZES_ALL][MAX_MODES];
+  int m_search_count;
+  int ex_search_count;
+  int current_qindex;
+
+#if CONFIG_INTERNAL_STATS
+  unsigned int mode_chosen_counts[MAX_MODES];
+#endif  // CONFIG_INTERNAL_STATS
+} SB_FIRST_PASS_STATS;
+
 unsigned int av1_get_sby_perpixel_variance(const AV1_COMP *cpi,
                                            const struct buf_2d *ref,
                                            BLOCK_SIZE bs) {
@@ -1470,17 +1507,6 @@
   }
 }
 
-typedef struct {
-  ENTROPY_CONTEXT a[MAX_MIB_SIZE * MAX_MB_PLANE];
-  ENTROPY_CONTEXT l[MAX_MIB_SIZE * MAX_MB_PLANE];
-  PARTITION_CONTEXT sa[MAX_MIB_SIZE];
-  PARTITION_CONTEXT sl[MAX_MIB_SIZE];
-  TXFM_CONTEXT *p_ta;
-  TXFM_CONTEXT *p_tl;
-  TXFM_CONTEXT ta[MAX_MIB_SIZE];
-  TXFM_CONTEXT tl[MAX_MIB_SIZE];
-} RD_SEARCH_MACROBLOCK_CONTEXT;
-
 static AOM_INLINE void restore_context(MACROBLOCK *x,
                                        const RD_SEARCH_MACROBLOCK_CONTEXT *ctx,
                                        int mi_row, int mi_col, BLOCK_SIZE bsize,
@@ -2585,7 +2611,8 @@
                               int mi_row, int mi_col, BLOCK_SIZE bsize,
                               BLOCK_SIZE max_sq_part, BLOCK_SIZE min_sq_part,
                               RD_STATS *rd_cost, RD_STATS best_rdc,
-                              PC_TREE *pc_tree, int64_t *none_rd) {
+                              PC_TREE *pc_tree, int64_t *none_rd,
+                              SB_MULTI_PASS_MODE multi_pass_mode) {
   const AV1_COMMON *const cm = &cpi->common;
   const int num_planes = av1_num_planes(cm);
   TileInfo *const tile_info = &tile_data->tile_info;
@@ -2978,7 +3005,7 @@
       if (!rd_pick_partition(cpi, td, tile_data, tp, mi_row + y_idx,
                              mi_col + x_idx, subsize, max_sq_part, min_sq_part,
                              &this_rdc, best_remain_rdcost, pc_tree->split[idx],
-                             p_split_rd)) {
+                             p_split_rd, multi_pass_mode)) {
         av1_invalid_rd_stats(&sum_rdc);
         break;
       }
@@ -3720,8 +3747,11 @@
 
   if (found_best_partition && pc_tree->index != 3) {
     if (bsize == cm->seq_params.sb_size) {
+      const int emit_output = multi_pass_mode != SB_DRY_PASS;
+      const RUN_TYPE run_type = emit_output ? OUTPUT_ENABLED : DRY_RUN_NORMAL;
+
       x->cb_offset = 0;
-      encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, OUTPUT_ENABLED, bsize,
+      encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, run_type, bsize,
                 pc_tree, NULL);
     } else {
       encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, DRY_RUN_NORMAL, bsize,
@@ -4288,6 +4318,157 @@
                       pc_root);
 }
 
+// Memset the mbmis at the current superblock to 0
+static INLINE void reset_mbmi(AV1_COMMON *cm, int mi_row, int mi_col) {
+  const BLOCK_SIZE sb_size = cm->seq_params.sb_size;
+  // size of sb in unit of mi_grid (BLOCK_4X4)
+  const int sb_size_mi = mi_size_wide[sb_size];
+  const int mi_alloc_size_1d = mi_size_wide[cm->mi_alloc_bsize];
+  // size of sb in unit of allocated mi size
+  const int sb_size_alloc_mi = mi_size_wide[sb_size] / mi_alloc_size_1d;
+  assert(cm->mi_alloc_stride % sb_size_alloc_mi == 0 &&
+         "mi is not allocated as a multiple of sb!");
+  assert(cm->mi_stride % sb_size_mi == 0 &&
+         "mi_grid is not allocated as a multiple of sb!");
+
+  const int mi_rows = mi_size_high[sb_size];
+  for (int cur_mi_row = 0; cur_mi_row < mi_rows; cur_mi_row++) {
+    assert(get_mi_grid_idx(cm, 0, mi_col + mi_alloc_size_1d) < cm->mi_stride);
+    const int mi_grid_idx = get_mi_grid_idx(cm, mi_row + cur_mi_row, mi_col);
+    const int alloc_mi_idx = get_alloc_mi_idx(cm, mi_row + cur_mi_row, mi_col);
+    memset(&cm->mi_grid_base[mi_grid_idx], 0,
+           sb_size_mi * sizeof(*cm->mi_grid_base));
+    memset(&cm->tx_type_map[mi_grid_idx], 0,
+           sb_size_mi * sizeof(*cm->tx_type_map));
+    if (cur_mi_row % mi_alloc_size_1d == 0) {
+      memset(&cm->mi[alloc_mi_idx], 0, sb_size_alloc_mi * sizeof(*cm->mi));
+    }
+  }
+}
+
+static INLINE void backup_sb_state(SB_FIRST_PASS_STATS *sb_fp_stats,
+                                   const AV1_COMP *cpi, ThreadData *td,
+                                   const TileDataEnc *tile_data, int mi_row,
+                                   int mi_col) {
+  MACROBLOCK *x = &td->mb;
+  MACROBLOCKD *xd = &x->e_mbd;
+  const TileInfo *tile_info = &tile_data->tile_info;
+
+  const AV1_COMMON *cm = &cpi->common;
+  const int num_planes = av1_num_planes(cm);
+  const BLOCK_SIZE sb_size = cm->seq_params.sb_size;
+
+  xd->above_txfm_context = cm->above_txfm_context[tile_info->tile_row] + mi_col;
+  xd->left_txfm_context =
+      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
+  save_context(x, &sb_fp_stats->x_ctx, mi_row, mi_col, sb_size, num_planes);
+
+  sb_fp_stats->rd_count = cpi->td.rd_counts;
+  sb_fp_stats->split_count = cpi->td.mb.txb_split_count;
+
+  sb_fp_stats->fc = *td->counts;
+
+  memcpy(sb_fp_stats->inter_mode_rd_models, tile_data->inter_mode_rd_models,
+         sizeof(sb_fp_stats->inter_mode_rd_models));
+
+  memcpy(sb_fp_stats->thresh_freq_fact, x->thresh_freq_fact,
+         sizeof(sb_fp_stats->thresh_freq_fact));
+
+  sb_fp_stats->m_search_count = *x->m_search_count_ptr;
+  sb_fp_stats->ex_search_count = *x->ex_search_count_ptr;
+
+  const int alloc_mi_idx = get_alloc_mi_idx(cm, mi_row, mi_col);
+  sb_fp_stats->current_qindex = cm->mi[alloc_mi_idx].current_qindex;
+
+#if CONFIG_INTERNAL_STATS
+  memcpy(sb_fp_stats->mode_chosen_counts, cpi->mode_chosen_counts,
+         sizeof(sb_fp_stats->mode_chosen_counts));
+#endif  // CONFIG_INTERNAL_STATS
+}
+
+static INLINE void restore_sb_state(const SB_FIRST_PASS_STATS *sb_fp_stats,
+                                    AV1_COMP *cpi, ThreadData *td,
+                                    TileDataEnc *tile_data, int mi_row,
+                                    int mi_col) {
+  MACROBLOCK *x = &td->mb;
+
+  const AV1_COMMON *cm = &cpi->common;
+  const int num_planes = av1_num_planes(cm);
+  const BLOCK_SIZE sb_size = cm->seq_params.sb_size;
+
+  restore_context(x, &sb_fp_stats->x_ctx, mi_row, mi_col, sb_size, num_planes);
+
+  cpi->td.rd_counts = sb_fp_stats->rd_count;
+  cpi->td.mb.txb_split_count = sb_fp_stats->split_count;
+
+  *td->counts = sb_fp_stats->fc;
+
+  memcpy(tile_data->inter_mode_rd_models, sb_fp_stats->inter_mode_rd_models,
+         sizeof(sb_fp_stats->inter_mode_rd_models));
+  memcpy(x->thresh_freq_fact, sb_fp_stats->thresh_freq_fact,
+         sizeof(sb_fp_stats->thresh_freq_fact));
+
+  *x->m_search_count_ptr = sb_fp_stats->m_search_count;
+  *x->ex_search_count_ptr = sb_fp_stats->ex_search_count;
+
+  const int alloc_mi_idx = get_alloc_mi_idx(cm, mi_row, mi_col);
+  cm->mi[alloc_mi_idx].current_qindex = sb_fp_stats->current_qindex;
+
+#if CONFIG_INTERNAL_STATS
+  memcpy(cpi->mode_chosen_counts, sb_fp_stats->mode_chosen_counts,
+         sizeof(sb_fp_stats->mode_chosen_counts));
+#endif  // CONFIG_INTERNAL_STATS
+}
+
+// This function initializes the stats for encode_rd_sb.
+static INLINE void init_encode_rd_sb(AV1_COMP *cpi, ThreadData *td,
+                                     const TileDataEnc *tile_data,
+                                     PC_TREE *pc_root, RD_STATS *rd_cost,
+                                     int mi_row, int mi_col,
+                                     int gather_tpl_data) {
+  const AV1_COMMON *cm = &cpi->common;
+  const TileInfo *tile_info = &tile_data->tile_info;
+  MACROBLOCK *x = &td->mb;
+
+  const SPEED_FEATURES *sf = &cpi->sf;
+  const int use_simple_motion_search =
+      (sf->part_sf.simple_motion_search_split ||
+       sf->part_sf.simple_motion_search_prune_rect ||
+       sf->part_sf.simple_motion_search_early_term_none ||
+       sf->part_sf.ml_early_term_after_part_split_level) &&
+      !frame_is_intra_only(cm);
+  if (use_simple_motion_search) {
+    init_simple_motion_search_mvs(pc_root);
+  }
+
+#if !CONFIG_REALTIME_ONLY
+  x->sb_energy_level = 0;
+  x->cnn_output_valid = 0;
+  if (gather_tpl_data) {
+    if (cm->delta_q_info.delta_q_present_flag) {
+      const int num_planes = av1_num_planes(cm);
+      const BLOCK_SIZE sb_size = cm->seq_params.sb_size;
+      setup_delta_q(cpi, td, x, tile_info, mi_row, mi_col, num_planes);
+      av1_tpl_rdmult_setup_sb(cpi, x, sb_size, mi_row, mi_col);
+    }
+    if (cpi->oxcf.enable_tpl_model) {
+      adjust_rdmult_tpl_model(cpi, x, mi_row, mi_col);
+    }
+  }
+#else
+  (void)tile_info;
+  (void)mi_row;
+  (void)mi_col;
+  (void)gather_tpl_data;
+#endif
+
+  // Reset hash state for transform/mode rd hash information
+  reset_hash_records(x, cpi->sf.tx_sf.use_inter_txb_hash);
+  av1_zero(x->picked_ref_frames_mask);
+  av1_zero(x->pred_mv);
+  av1_invalid_rd_stats(rd_cost);
+}
+
 static AOM_INLINE void encode_rd_sb(AV1_COMP *cpi, ThreadData *td,
                                     TileDataEnc *tile_data,
                                     PC_TREE *const pc_root, TOKENEXTRA **tp,
@@ -4303,31 +4484,11 @@
   int64_t dummy_dist;
   RD_STATS dummy_rdc;
 
-  if ((sf->part_sf.simple_motion_search_split ||
-       sf->part_sf.simple_motion_search_prune_rect ||
-       sf->part_sf.simple_motion_search_early_term_none ||
-       sf->part_sf.ml_early_term_after_part_split_level) &&
-      !frame_is_intra_only(cm)) {
-    init_simple_motion_search_mvs(pc_root);
-  }
-#if !CONFIG_REALTIME_ONLY
-  const int num_planes = av1_num_planes(cm);
-  x->sb_energy_level = 0;
-  td->mb.cnn_output_valid = 0;
-  if (cm->delta_q_info.delta_q_present_flag) {
-    setup_delta_q(cpi, td, x, tile_info, mi_row, mi_col, num_planes);
-    av1_tpl_rdmult_setup_sb(cpi, x, sb_size, mi_row, mi_col);
-  }
-  if (cpi->oxcf.enable_tpl_model)
-    adjust_rdmult_tpl_model(cpi, x, mi_row, mi_col);
-#else
+#if CONFIG_REALTIME_ONLY
   (void)seg_skip;
-#endif
-  // Reset hash state for transform/mode rd hash information
-  reset_hash_records(x, cpi->sf.tx_sf.use_inter_txb_hash);
-  av1_zero(x->picked_ref_frames_mask);
-  av1_zero(x->pred_mv);
-  av1_invalid_rd_stats(&dummy_rdc);
+#endif  // CONFIG_REALTIME_ONLY
+
+  init_encode_rd_sb(cpi, td, tile_data, pc_root, &dummy_rdc, mi_row, mi_col, 1);
 
   if (sf->part_sf.partition_search_type == VAR_BASED_PARTITION) {
     set_offsets_without_segment_id(cpi, tile_info, x, mi_row, mi_col, sb_size);
@@ -4351,11 +4512,10 @@
     rd_use_partition(cpi, td, tile_data, mi, tp, mi_row, mi_col, sb_size,
                      &dummy_rate, &dummy_dist, 1, pc_root);
   } else {
-    x->valid_cost_b = 0;
     // No stats for overlay frames. Exclude key frame.
     x->valid_cost_b =
-        get_tpl_stats_b(cpi, cm->seq_params.sb_size, mi_row, mi_col,
-                        x->intra_cost_b, x->inter_cost_b, &x->cost_stride);
+        get_tpl_stats_b(cpi, sb_size, mi_row, mi_col, x->intra_cost_b,
+                        x->inter_cost_b, &x->cost_stride);
 
     reset_partition(pc_root, sb_size);
 
@@ -4375,9 +4535,33 @@
 
     min_sq_size = AOMMIN(min_sq_size, max_sq_size);
 
-    rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, sb_size,
-                      max_sq_size, min_sq_size, &dummy_rdc, dummy_rdc, pc_root,
-                      NULL);
+    const int num_passes = cpi->oxcf.sb_multipass_unit_test ? 2 : 1;
+
+    if (num_passes == 1) {
+      rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, sb_size,
+                        max_sq_size, min_sq_size, &dummy_rdc, dummy_rdc,
+                        pc_root, NULL, SB_SINGLE_PASS);
+    } else {
+      // First pass
+      SB_FIRST_PASS_STATS sb_fp_stats;
+      backup_sb_state(&sb_fp_stats, cpi, td, tile_data, mi_row, mi_col);
+      rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, sb_size,
+                        max_sq_size, min_sq_size, &dummy_rdc, dummy_rdc,
+                        pc_root, NULL, SB_DRY_PASS);
+
+      // Second pass
+      init_encode_rd_sb(cpi, td, tile_data, pc_root, &dummy_rdc, mi_row, mi_col,
+                        0);
+      reset_mbmi(&cpi->common, mi_row, mi_col);
+      reset_partition(pc_root, sb_size);
+
+      restore_sb_state(&sb_fp_stats, cpi, td, tile_data, mi_row, mi_col);
+
+      rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, sb_size,
+                        max_sq_size, min_sq_size, &dummy_rdc, dummy_rdc,
+                        pc_root, NULL, SB_WET_PASS);
+    }
+
 #if CONFIG_COLLECT_COMPONENT_TIMING
     end_timing(cpi, rd_pick_partition_time);
 #endif
diff --git a/av1/encoder/encoder.h b/av1/encoder/encoder.h
index 997c673..cfe326e 100644
--- a/av1/encoder/encoder.h
+++ b/av1/encoder/encoder.h
@@ -371,6 +371,7 @@
   unsigned int full_still_picture_hdr;
   int enable_dual_filter;
   unsigned int motion_vector_unit_test;
+  unsigned int sb_multipass_unit_test;
   unsigned int ext_tile_debug;
   int enable_rect_partitions;
   int enable_ab_partitions;
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 8ea5891..8d91210 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -63,6 +63,10 @@
          "Invalid sb_size for intra_cnn!");
   const int bsize_idx = convert_bsize_to_idx(bsize);
 
+  if (bsize == BLOCK_128X128) {
+    return;
+  }
+
   // Precompute the CNN part and cache the result in MACROBLOCK
   if (bsize == BLOCK_64X64 && !x->cnn_output_valid) {
     aom_clear_system_state();
diff --git a/test/sb_multipass_test.cc b/test/sb_multipass_test.cc
new file mode 100644
index 0000000..bf7f66a
--- /dev/null
+++ b/test/sb_multipass_test.cc
@@ -0,0 +1,152 @@
+/*
+ * Copyright (c) 2020, Alliance for Open Media. All rights reserved
+ *
+ * This source code is subject to the terms of the BSD 2 Clause License and
+ * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
+ * was not distributed with this source code in the LICENSE file, you can
+ * obtain it at www.aomedia.org/license/software. If the Alliance for Open
+ * Media Patent License 1.0 was not distributed with this source code in the
+ * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
+ */
+
+#include <initializer_list>
+#include <string>
+#include <vector>
+#include "third_party/googletest/src/googletest/include/gtest/gtest.h"
+#include "test/codec_factory.h"
+#include "test/encode_test_driver.h"
+#include "test/md5_helper.h"
+#include "test/util.h"
+#include "test/yuv_video_source.h"
+
+namespace {
+class AV1SBMultipassTest : public ::libaom_test::CodecTestWith2Params<int, int>,
+                           public ::libaom_test::EncoderTest {
+ protected:
+  AV1SBMultipassTest()
+      : EncoderTest(GET_PARAM(0)), set_cpu_used_(GET_PARAM(1)),
+        row_mt_(GET_PARAM(2)) {
+    init_flags_ = AOM_CODEC_USE_PSNR;
+    aom_codec_dec_cfg_t cfg = aom_codec_dec_cfg_t();
+    cfg.w = 1280;
+    cfg.h = 720;
+    cfg.allow_lowbitdepth = 1;
+    decoder_ = codec_->CreateDecoder(cfg, 0);
+    if (decoder_->IsAV1()) {
+      decoder_->Control(AV1_SET_DECODE_TILE_ROW, -1);
+      decoder_->Control(AV1_SET_DECODE_TILE_COL, -1);
+    }
+
+    size_enc_.clear();
+    md5_dec_.clear();
+    md5_enc_.clear();
+  }
+  virtual ~AV1SBMultipassTest() { delete decoder_; }
+
+  virtual void SetUp() {
+    InitializeConfig();
+    SetMode(::libaom_test::kTwoPassGood);
+
+    cfg_.g_lag_in_frames = 5;
+    cfg_.rc_end_usage = AOM_VBR;
+    cfg_.rc_2pass_vbr_minsection_pct = 5;
+    cfg_.rc_2pass_vbr_maxsection_pct = 2000;
+
+    cfg_.rc_max_quantizer = 56;
+    cfg_.rc_min_quantizer = 0;
+  }
+
+  virtual void PreEncodeFrameHook(::libaom_test::VideoSource *video,
+                                  ::libaom_test::Encoder *encoder) {
+    if (video->frame() == 0) {
+      SetTileSize(encoder);
+      encoder->Control(AOME_SET_CPUUSED, set_cpu_used_);
+      encoder->Control(AV1E_ENABLE_SB_MULTIPASS_UNIT_TEST, use_multipass_);
+      encoder->Control(AV1E_SET_ROW_MT, row_mt_);
+
+      encoder->Control(AOME_SET_ENABLEAUTOALTREF, 1);
+      encoder->Control(AOME_SET_ARNR_MAXFRAMES, 7);
+      encoder->Control(AOME_SET_ARNR_STRENGTH, 5);
+    }
+  }
+
+  virtual void SetTileSize(libaom_test::Encoder *encoder) {
+    encoder->Control(AV1E_SET_TILE_COLUMNS, 1);
+    encoder->Control(AV1E_SET_TILE_ROWS, 1);
+  }
+
+  virtual void FramePktHook(const aom_codec_cx_pkt_t *pkt) {
+    size_enc_.push_back(pkt->data.frame.sz);
+
+    ::libaom_test::MD5 md5_enc;
+    md5_enc.Add(reinterpret_cast<uint8_t *>(pkt->data.frame.buf),
+                pkt->data.frame.sz);
+    md5_enc_.push_back(md5_enc.Get());
+
+    const aom_codec_err_t res = decoder_->DecodeFrame(
+        reinterpret_cast<uint8_t *>(pkt->data.frame.buf), pkt->data.frame.sz);
+    if (res != AOM_CODEC_OK) {
+      abort_ = true;
+      ASSERT_EQ(AOM_CODEC_OK, res);
+    }
+    const aom_image_t *img = decoder_->GetDxData().Next();
+
+    if (img) {
+      ::libaom_test::MD5 md5_res;
+      md5_res.Add(img);
+      md5_dec_.push_back(md5_res.Get());
+    }
+  }
+
+  void DoTest() {
+    ::libaom_test::YUVVideoSource video(
+        "niklas_640_480_30.yuv", AOM_IMG_FMT_I420, 640, 480, 30, 1, 0, 6);
+    cfg_.rc_target_bitrate = 1000;
+
+    // Encode while coding each sb once
+    use_multipass_ = false;
+    ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+    std::vector<size_t> single_pass_size_enc;
+    std::vector<std::string> single_pass_md5_enc;
+    std::vector<std::string> single_pass_md5_dec;
+    single_pass_size_enc = size_enc_;
+    single_pass_md5_enc = md5_enc_;
+    single_pass_md5_dec = md5_dec_;
+    size_enc_.clear();
+    md5_enc_.clear();
+    md5_dec_.clear();
+
+    // Encode while coding each sb twice
+    use_multipass_ = true;
+    ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+    std::vector<size_t> multi_pass_size_enc;
+    std::vector<std::string> multi_pass_md5_enc;
+    std::vector<std::string> multi_pass_md5_dec;
+    multi_pass_size_enc = size_enc_;
+    multi_pass_md5_enc = md5_enc_;
+    multi_pass_md5_dec = md5_dec_;
+    size_enc_.clear();
+    md5_enc_.clear();
+    md5_dec_.clear();
+
+    // Check that the vectors are equal.
+    ASSERT_EQ(single_pass_size_enc, multi_pass_size_enc);
+    ASSERT_EQ(single_pass_md5_enc, multi_pass_md5_enc);
+    ASSERT_EQ(single_pass_md5_dec, multi_pass_md5_dec);
+  }
+
+  bool use_multipass_;
+  int set_cpu_used_;
+  bool row_mt_;
+  ::libaom_test::Decoder *decoder_;
+  std::vector<size_t> size_enc_;
+  std::vector<std::string> md5_enc_;
+  std::vector<std::string> md5_dec_;
+};
+
+TEST_P(AV1SBMultipassTest, TwoPassMatchTest) { DoTest(); }
+
+AV1_INSTANTIATE_TEST_CASE(AV1SBMultipassTest, ::testing::Range(0, 6),
+                          ::testing::Range(0, 2));
+
+}  // namespace
diff --git a/test/test.cmake b/test/test.cmake
index ea7c179..05f0864 100644
--- a/test/test.cmake
+++ b/test/test.cmake
@@ -127,6 +127,7 @@
                 "${AOM_ROOT}/test/ec_test.cc"
                 "${AOM_ROOT}/test/ethread_test.cc"
                 "${AOM_ROOT}/test/film_grain_table_test.cc"
+                "${AOM_ROOT}/test/sb_multipass_test.cc"
                 "${AOM_ROOT}/test/segment_binarization_sync.cc"
                 "${AOM_ROOT}/test/superframe_test.cc"
                 "${AOM_ROOT}/test/tile_independence_test.cc"