encodeframe.c

/*
 * Copyright (c) 2016, 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 <limits.h>
#include <math.h>
#include <stdio.h>

#include "./av1_rtcd.h"
#include "./aom_dsp_rtcd.h"
#include "./aom_config.h"

#include "aom_dsp/aom_dsp_common.h"
#include "aom_dsp/binary_codes_writer.h"
#include "aom_ports/mem.h"
#include "aom_ports/aom_timer.h"
#include "aom_ports/system_state.h"

#include "av1/common/common.h"
#include "av1/common/entropy.h"
#include "av1/common/entropymode.h"
#include "av1/common/idct.h"
#include "av1/common/mv.h"
#include "av1/common/mvref_common.h"
#include "av1/common/pred_common.h"
#include "av1/common/quant_common.h"
#include "av1/common/reconintra.h"
#include "av1/common/reconinter.h"
#include "av1/common/seg_common.h"
#include "av1/common/tile_common.h"

#include "av1/encoder/aq_complexity.h"
#include "av1/encoder/aq_cyclicrefresh.h"
#include "av1/encoder/aq_variance.h"
#if CONFIG_SUPERTX
#include "av1/encoder/cost.h"
#endif
#if CONFIG_GLOBAL_MOTION || CONFIG_WARPED_MOTION
#include "av1/common/warped_motion.h"
#endif  // CONFIG_GLOBAL_MOTION || CONFIG_WARPED_MOTION
#if CONFIG_GLOBAL_MOTION
#include "av1/encoder/global_motion.h"
#endif  // CONFIG_GLOBAL_MOTION
#include "av1/encoder/encodeframe.h"
#include "av1/encoder/encodemb.h"
#include "av1/encoder/encodemv.h"
#if CONFIG_LV_MAP
#include "av1/encoder/encodetxb.h"
#endif
#include "av1/encoder/ethread.h"
#include "av1/encoder/extend.h"
#include "av1/encoder/rd.h"
#include "av1/encoder/rdopt.h"
#include "av1/encoder/segmentation.h"
#include "av1/encoder/tokenize.h"
#if CONFIG_PVQ
#include "av1/common/pvq.h"
#include "av1/encoder/pvq_encoder.h"
#endif
#if CONFIG_HIGHBITDEPTH
#define IF_HBD(...) __VA_ARGS__
#else
#define IF_HBD(...)
#endif  // CONFIG_HIGHBITDEPTH

static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td,
                              TOKENEXTRA **t, RUN_TYPE dry_run, int mi_row,
                              int mi_col, BLOCK_SIZE bsize,
                              PICK_MODE_CONTEXT *ctx, int *rate);

#if CONFIG_SUPERTX
static int check_intra_b(PICK_MODE_CONTEXT *ctx);

static int check_intra_sb(const AV1_COMP *cpi, const TileInfo *const tile,
                          int mi_row, int mi_col, BLOCK_SIZE bsize,
                          PC_TREE *pc_tree);
static void predict_superblock(const AV1_COMP *const cpi, ThreadData *td,
#if CONFIG_EXT_INTER
                               int mi_row_ori, int mi_col_ori,
#endif  // CONFIG_EXT_INTER
                               int mi_row_pred, int mi_col_pred,
                               BLOCK_SIZE bsize_pred, int b_sub8x8, int block);
static int check_supertx_sb(BLOCK_SIZE bsize, TX_SIZE supertx_size,
                            PC_TREE *pc_tree);
static void predict_sb_complex(const AV1_COMP *const cpi, ThreadData *td,
                               const TileInfo *const tile, int mi_row,
                               int mi_col, int mi_row_ori, int mi_col_ori,
                               RUN_TYPE dry_run, BLOCK_SIZE bsize,
                               BLOCK_SIZE top_bsize, uint8_t *dst_buf[3],
                               int dst_stride[3], PC_TREE *pc_tree);
static void update_state_sb_supertx(const AV1_COMP *const cpi, ThreadData *td,
                                    const TileInfo *const tile, int mi_row,
                                    int mi_col, BLOCK_SIZE bsize,
                                    RUN_TYPE dry_run, PC_TREE *pc_tree);
static void rd_supertx_sb(const AV1_COMP *const cpi, ThreadData *td,
                          const TileInfo *const tile, int mi_row, int mi_col,
                          BLOCK_SIZE bsize, int *tmp_rate, int64_t *tmp_dist,
                          TX_TYPE *best_tx, PC_TREE *pc_tree);
#endif  // CONFIG_SUPERTX

// This is used as a reference when computing the source variance for the
//  purposes of activity masking.
// Eventually this should be replaced by custom no-reference routines,
//  which will be faster.
static const uint8_t AV1_VAR_OFFS[MAX_SB_SIZE] = {
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
#if CONFIG_EXT_PARTITION
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128
#endif  // CONFIG_EXT_PARTITION
};

#if CONFIG_HIGHBITDEPTH
static const uint16_t AV1_HIGH_VAR_OFFS_8[MAX_SB_SIZE] = {
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
#if CONFIG_EXT_PARTITION
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
  128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128
#endif  // CONFIG_EXT_PARTITION
};

static const uint16_t AV1_HIGH_VAR_OFFS_10[MAX_SB_SIZE] = {
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
#if CONFIG_EXT_PARTITION
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4,
  128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4
#endif  // CONFIG_EXT_PARTITION
};

static const uint16_t AV1_HIGH_VAR_OFFS_12[MAX_SB_SIZE] = {
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16,
#if CONFIG_EXT_PARTITION
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16,
  128 * 16
#endif  // CONFIG_EXT_PARTITION
};
#endif  // CONFIG_HIGHBITDEPTH

unsigned int av1_get_sby_perpixel_variance(const AV1_COMP *cpi,
                                           const struct buf_2d *ref,
                                           BLOCK_SIZE bs) {
  unsigned int sse;
  const unsigned int var =
      cpi->fn_ptr[bs].vf(ref->buf, ref->stride, AV1_VAR_OFFS, 0, &sse);
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}

#if CONFIG_HIGHBITDEPTH
unsigned int av1_high_get_sby_perpixel_variance(const AV1_COMP *cpi,
                                                const struct buf_2d *ref,
                                                BLOCK_SIZE bs, int bd) {
  unsigned int var, sse;
  switch (bd) {
    case 10:
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_10), 0, &sse);
      break;
    case 12:
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_12), 0, &sse);
      break;
    case 8:
    default:
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8), 0, &sse);
      break;
  }
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}
#endif  // CONFIG_HIGHBITDEPTH

static unsigned int get_sby_perpixel_diff_variance(const AV1_COMP *const cpi,
                                                   const struct buf_2d *ref,
                                                   int mi_row, int mi_col,
                                                   BLOCK_SIZE bs) {
  unsigned int sse, var;
  uint8_t *last_y;
  const YV12_BUFFER_CONFIG *last = get_ref_frame_buffer(cpi, LAST_FRAME);

  assert(last != NULL);
  last_y =
      &last->y_buffer[mi_row * MI_SIZE * last->y_stride + mi_col * MI_SIZE];
  var = cpi->fn_ptr[bs].vf(ref->buf, ref->stride, last_y, last->y_stride, &sse);
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}

static BLOCK_SIZE get_rd_var_based_fixed_partition(AV1_COMP *cpi, MACROBLOCK *x,
                                                   int mi_row, int mi_col) {
  unsigned int var = get_sby_perpixel_diff_variance(
      cpi, &x->plane[0].src, mi_row, mi_col, BLOCK_64X64);
  if (var < 8)
    return BLOCK_64X64;
  else if (var < 128)
    return BLOCK_32X32;
  else if (var < 2048)
    return BLOCK_16X16;
  else
    return BLOCK_8X8;
}

// Lighter version of set_offsets that only sets the mode info
// pointers.
static void set_mode_info_offsets(const AV1_COMP *const cpi,
                                  MACROBLOCK *const x, MACROBLOCKD *const xd,
                                  int mi_row, int mi_col) {
  const AV1_COMMON *const cm = &cpi->common;
  const int idx_str = xd->mi_stride * mi_row + mi_col;
  xd->mi = cm->mi_grid_visible + idx_str;
  xd->mi[0] = cm->mi + idx_str;
  x->mbmi_ext = cpi->mbmi_ext_base + (mi_row * cm->mi_cols + mi_col);
}

static void set_offsets_without_segment_id(const AV1_COMP *const cpi,
                                           const TileInfo *const tile,
                                           MACROBLOCK *const x, int mi_row,
                                           int mi_col, BLOCK_SIZE bsize) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
  const int mi_width = mi_size_wide[bsize];
  const int mi_height = mi_size_high[bsize];

  set_skip_context(xd, mi_row, mi_col);

  set_mode_info_offsets(cpi, x, xd, mi_row, mi_col);

#if CONFIG_VAR_TX
  xd->above_txfm_context = cm->above_txfm_context + mi_col;
  xd->left_txfm_context =
      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
  xd->max_tx_size = max_txsize_lookup[bsize];
#endif

  // Set up destination pointers.
  av1_setup_dst_planes(xd->plane, get_frame_new_buffer(cm), mi_row, mi_col);

  // Set up limit values for MV components.
  // Mv beyond the range do not produce new/different prediction block.
  x->mv_limits.row_min =
      -(((mi_row + mi_height) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_limits.col_min = -(((mi_col + mi_width) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_limits.row_max = (cm->mi_rows - mi_row) * MI_SIZE + AOM_INTERP_EXTEND;
  x->mv_limits.col_max = (cm->mi_cols - mi_col) * MI_SIZE + AOM_INTERP_EXTEND;

  set_plane_n4(xd, mi_width, mi_height);

  // Set up distance of MB to edge of frame in 1/8th pel units.
  assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));
#if CONFIG_DEPENDENT_HORZTILES
  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, cm->mi_rows,
                 cm->mi_cols, cm->dependent_horz_tiles);
#else
  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, cm->mi_rows,
                 cm->mi_cols);
#endif

  // Set up source buffers.
  av1_setup_src_planes(x, cpi->source, mi_row, mi_col);

  // R/D setup.
  x->rddiv = cpi->rd.RDDIV;
  x->rdmult = cpi->rd.RDMULT;

  // required by av1_append_sub8x8_mvs_for_idx() and av1_find_best_ref_mvs()
  xd->tile = *tile;
}

static void set_offsets(const AV1_COMP *const cpi, const TileInfo *const tile,
                        MACROBLOCK *const x, int mi_row, int mi_col,
                        BLOCK_SIZE bsize) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *mbmi;
  const struct segmentation *const seg = &cm->seg;

  set_offsets_without_segment_id(cpi, tile, x, mi_row, mi_col, bsize);

  mbmi = &xd->mi[0]->mbmi;

  // Setup segment ID.
  if (seg->enabled) {
    if (!cpi->vaq_refresh) {
      const uint8_t *const map =
          seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map;
      mbmi->segment_id = get_segment_id(cm, map, bsize, mi_row, mi_col);
    }
    av1_init_plane_quantizers(cpi, x, mbmi->segment_id);
  } else {
    mbmi->segment_id = 0;
  }

#if CONFIG_SUPERTX
  mbmi->segment_id_supertx = MAX_SEGMENTS;
#endif  // CONFIG_SUPERTX
}

#if CONFIG_SUPERTX
static void set_offsets_supertx(const AV1_COMP *const cpi, ThreadData *td,
                                const TileInfo *const tile, int mi_row,
                                int mi_col, BLOCK_SIZE bsize) {
  MACROBLOCK *const x = &td->mb;
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
  const int mi_width = mi_size_wide[bsize];
  const int mi_height = mi_size_high[bsize];
#if CONFIG_DEPENDENT_HORZTILES
  set_mode_info_offsets(cpi, x, xd, mi_row, mi_col, cm->dependent_horz_tiles);
#else
  set_mode_info_offsets(cpi, x, xd, mi_row, mi_col);
#endif

  // Set up distance of MB to edge of frame in 1/8th pel units.
  assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));
  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, cm->mi_rows,
                 cm->mi_cols);
}

static void set_offsets_extend(const AV1_COMP *const cpi, ThreadData *td,
                               const TileInfo *const tile, int mi_row_pred,
                               int mi_col_pred, int mi_row_ori, int mi_col_ori,
                               BLOCK_SIZE bsize_pred) {
  // Used in supertx
  // (mi_row_ori, mi_col_ori, bsize_ori): region for mv
  // (mi_row_pred, mi_col_pred, bsize_pred): region to predict
  MACROBLOCK *const x = &td->mb;
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
  const int mi_width = mi_size_wide[bsize_pred];
  const int mi_height = mi_size_high[bsize_pred];

#if CONFIG_DEPENDENT_HORZTILES
  set_mode_info_offsets(cpi, x, xd, mi_row_ori, mi_col_ori,
                        cm->dependent_horz_tiles);
#else
  set_mode_info_offsets(cpi, x, xd, mi_row_ori, mi_col_ori);
#endif

  // Set up limit values for MV components.
  // Mv beyond the range do not produce new/different prediction block.
  x->mv_limits.row_min =
      -(((mi_row_pred + mi_height) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_limits.col_min =
      -(((mi_col_pred + mi_width) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_limits.row_max =
      (cm->mi_rows - mi_row_pred) * MI_SIZE + AOM_INTERP_EXTEND;
  x->mv_limits.col_max =
      (cm->mi_cols - mi_col_pred) * MI_SIZE + AOM_INTERP_EXTEND;

// Set up distance of MB to edge of frame in 1/8th pel units.
#if !CONFIG_CB4X4
  assert(!(mi_col_pred & (mi_width - mi_size_wide[BLOCK_8X8])) &&
         !(mi_row_pred & (mi_height - mi_size_high[BLOCK_8X8])));
#endif
  set_mi_row_col(xd, tile, mi_row_pred, mi_height, mi_col_pred, mi_width,
                 cm->mi_rows, cm->mi_cols);
  xd->up_available = (mi_row_ori > tile->mi_row_start);
  xd->left_available = (mi_col_ori > tile->mi_col_start);

  // R/D setup.
  x->rddiv = cpi->rd.RDDIV;
  x->rdmult = cpi->rd.RDMULT;
}

static void set_segment_id_supertx(const AV1_COMP *const cpi,
                                   MACROBLOCK *const x, const int mi_row,
                                   const int mi_col, const BLOCK_SIZE bsize) {
  const AV1_COMMON *cm = &cpi->common;
  const struct segmentation *seg = &cm->seg;
  const int miw = AOMMIN(mi_size_wide[bsize], cm->mi_cols - mi_col);
  const int mih = AOMMIN(mi_size_high[bsize], cm->mi_rows - mi_row);
  const int mi_offset = mi_row * cm->mi_stride + mi_col;
  MODE_INFO **const mip = cm->mi_grid_visible + mi_offset;
  int r, c;
  int seg_id_supertx = MAX_SEGMENTS;

  if (!seg->enabled) {
    seg_id_supertx = 0;
  } else {
    // Find the minimum segment_id
    for (r = 0; r < mih; r++)
      for (c = 0; c < miw; c++)
        seg_id_supertx =
            AOMMIN(mip[r * cm->mi_stride + c]->mbmi.segment_id, seg_id_supertx);
    assert(0 <= seg_id_supertx && seg_id_supertx < MAX_SEGMENTS);

    // Initialize plane quantisers
    av1_init_plane_quantizers(cpi, x, seg_id_supertx);
  }

  // Assign the the segment_id back to segment_id_supertx
  for (r = 0; r < mih; r++)
    for (c = 0; c < miw; c++)
      mip[r * cm->mi_stride + c]->mbmi.segment_id_supertx = seg_id_supertx;
}
#endif  // CONFIG_SUPERTX

static void set_block_size(AV1_COMP *const cpi, MACROBLOCK *const x,
                           MACROBLOCKD *const xd, int mi_row, int mi_col,
                           BLOCK_SIZE bsize) {
  if (cpi->common.mi_cols > mi_col && cpi->common.mi_rows > mi_row) {
    const int mi_width = AOMMAX(mi_size_wide[bsize], mi_size_wide[BLOCK_8X8]);
    const int mi_height = AOMMAX(mi_size_high[bsize], mi_size_high[BLOCK_8X8]);
    for (int r = 0; r < mi_height; ++r) {
      for (int c = 0; c < mi_width; ++c) {
        set_mode_info_offsets(cpi, x, xd, mi_row + r, mi_col + c);
        xd->mi[0]->mbmi.sb_type = bsize;
      }
    }
  }
}

static void set_vt_partitioning(AV1_COMP *cpi, MACROBLOCK *const x,
                                MACROBLOCKD *const xd, VAR_TREE *vt, int mi_row,
                                int mi_col, const int64_t *const threshold,
                                const BLOCK_SIZE *const bsize_min) {
  AV1_COMMON *const cm = &cpi->common;
  const int hbw = mi_size_wide[vt->bsize] / 2;
  const int hbh = mi_size_high[vt->bsize] / 2;
  const int has_cols = mi_col + hbw < cm->mi_cols;
  const int has_rows = mi_row + hbh < cm->mi_rows;

  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return;

  assert(vt->bsize >= BLOCK_8X8);

  assert(hbh == hbw);

  if (vt->bsize == BLOCK_8X8 && cm->frame_type != KEY_FRAME) {
    set_block_size(cpi, x, xd, mi_row, mi_col, BLOCK_8X8);
    return;
  }

  if (vt->force_split || (!has_cols && !has_rows)) goto split;

  // For bsize=bsize_min (16x16/8x8 for 8x8/4x4 downsampling), select if
  // variance is below threshold, otherwise split will be selected.
  // No check for vert/horiz split as too few samples for variance.
  if (vt->bsize == bsize_min[0]) {
    if (has_cols && has_rows && vt->variances.none.variance < threshold[0]) {
      set_block_size(cpi, x, xd, mi_row, mi_col, vt->bsize);
      return;
    } else {
      BLOCK_SIZE subsize = get_subsize(vt->bsize, PARTITION_SPLIT);
      set_block_size(cpi, x, xd, mi_row, mi_col, subsize);
      if (vt->bsize > BLOCK_8X8) {
        set_block_size(cpi, x, xd, mi_row, mi_col + hbw, subsize);
        set_block_size(cpi, x, xd, mi_row + hbh, mi_col, subsize);
        set_block_size(cpi, x, xd, mi_row + hbh, mi_col + hbw, subsize);
      }
      return;
    }
  } else if (vt->bsize > bsize_min[0]) {
    // For key frame: take split for bsize above 32X32 or very high variance.
    if (cm->frame_type == KEY_FRAME &&
        (vt->bsize > BLOCK_32X32 ||
         vt->variances.none.variance > (threshold[0] << 4))) {
      goto split;
    }
    // If variance is low, take the bsize (no split).
    if (has_cols && has_rows && vt->variances.none.variance < threshold[0]) {
      set_block_size(cpi, x, xd, mi_row, mi_col, vt->bsize);
      return;
    }

    // Check vertical split.
    if (has_rows) {
      BLOCK_SIZE subsize = get_subsize(vt->bsize, PARTITION_VERT);
      if (vt->variances.vert[0].variance < threshold[0] &&
          vt->variances.vert[1].variance < threshold[0] &&
          get_plane_block_size(subsize, &xd->plane[1]) < BLOCK_INVALID) {
        set_block_size(cpi, x, xd, mi_row, mi_col, subsize);
        set_block_size(cpi, x, xd, mi_row, mi_col + hbw, subsize);
        return;
      }
    }
    // Check horizontal split.
    if (has_cols) {
      BLOCK_SIZE subsize = get_subsize(vt->bsize, PARTITION_HORZ);
      if (vt->variances.horz[0].variance < threshold[0] &&
          vt->variances.horz[1].variance < threshold[0] &&
          get_plane_block_size(subsize, &xd->plane[1]) < BLOCK_INVALID) {
        set_block_size(cpi, x, xd, mi_row, mi_col, subsize);
        set_block_size(cpi, x, xd, mi_row + hbh, mi_col, subsize);
        return;
      }
    }
  }

split : {
  set_vt_partitioning(cpi, x, xd, vt->split[0], mi_row, mi_col, threshold + 1,
                      bsize_min + 1);
  set_vt_partitioning(cpi, x, xd, vt->split[1], mi_row, mi_col + hbw,
                      threshold + 1, bsize_min + 1);
  set_vt_partitioning(cpi, x, xd, vt->split[2], mi_row + hbh, mi_col,
                      threshold + 1, bsize_min + 1);
  set_vt_partitioning(cpi, x, xd, vt->split[3], mi_row + hbh, mi_col + hbw,
                      threshold + 1, bsize_min + 1);
  return;
}
}

// Set the variance split thresholds for following the block sizes:
// 0 - threshold_64x64, 1 - threshold_32x32, 2 - threshold_16x16,
// 3 - vbp_threshold_8x8. vbp_threshold_8x8 (to split to 4x4 partition) is
// currently only used on key frame.
static void set_vbp_thresholds(AV1_COMP *cpi, int64_t thresholds[], int q) {
  AV1_COMMON *const cm = &cpi->common;
  const int is_key_frame = (cm->frame_type == KEY_FRAME);
  const int threshold_multiplier = is_key_frame ? 20 : 1;
  const int64_t threshold_base =
      (int64_t)(threshold_multiplier * cpi->y_dequant[q][1]);
  if (is_key_frame) {
    thresholds[1] = threshold_base;
    thresholds[2] = threshold_base >> 2;
    thresholds[3] = threshold_base >> 2;
    thresholds[4] = threshold_base << 2;
  } else {
    thresholds[2] = threshold_base;
    if (cm->width <= 352 && cm->height <= 288) {
      thresholds[1] = threshold_base >> 2;
      thresholds[3] = threshold_base << 3;
    } else {
      thresholds[1] = threshold_base;
      thresholds[2] = (5 * threshold_base) >> 2;
      if (cm->width >= 1920 && cm->height >= 1080)
        thresholds[2] = (7 * threshold_base) >> 2;
      thresholds[3] = threshold_base << cpi->oxcf.speed;
    }
  }
  thresholds[0] = INT64_MIN;
}

void av1_set_variance_partition_thresholds(AV1_COMP *cpi, int q) {
  AV1_COMMON *const cm = &cpi->common;
  SPEED_FEATURES *const sf = &cpi->sf;
  const int is_key_frame = (cm->frame_type == KEY_FRAME);
  if (sf->partition_search_type != VAR_BASED_PARTITION &&
      sf->partition_search_type != REFERENCE_PARTITION) {
    return;
  } else {
    set_vbp_thresholds(cpi, cpi->vbp_thresholds, q);
    // The thresholds below are not changed locally.
    if (is_key_frame) {
      cpi->vbp_threshold_sad = 0;
      cpi->vbp_bsize_min = BLOCK_8X8;
    } else {
      if (cm->width <= 352 && cm->height <= 288)
        cpi->vbp_threshold_sad = 100;
      else
        cpi->vbp_threshold_sad = (cpi->y_dequant[q][1] << 1) > 1000
                                     ? (cpi->y_dequant[q][1] << 1)
                                     : 1000;
      cpi->vbp_bsize_min = BLOCK_16X16;
    }
    cpi->vbp_threshold_minmax = 15 + (q >> 3);
  }
}

// Compute the minmax over the 8x8 subblocks.
static int compute_minmax_8x8(const uint8_t *src, int src_stride,
                              const uint8_t *ref, int ref_stride,
#if CONFIG_HIGHBITDEPTH
                              int highbd,
#endif
                              int pixels_wide, int pixels_high) {
  int k;
  int minmax_max = 0;
  int minmax_min = 255;
  // Loop over the 4 8x8 subblocks.
  for (k = 0; k < 4; k++) {
    const int x8_idx = ((k & 1) << 3);
    const int y8_idx = ((k >> 1) << 3);
    int min = 0;
    int max = 0;
    if (x8_idx < pixels_wide && y8_idx < pixels_high) {
      const int src_offset = y8_idx * src_stride + x8_idx;
      const int ref_offset = y8_idx * ref_stride + x8_idx;
#if CONFIG_HIGHBITDEPTH
      if (highbd) {
        aom_highbd_minmax_8x8(src + src_offset, src_stride, ref + ref_offset,
                              ref_stride, &min, &max);
      } else {
        aom_minmax_8x8(src + src_offset, src_stride, ref + ref_offset,
                       ref_stride, &min, &max);
      }
#else
      aom_minmax_8x8(src + src_offset, src_stride, ref + ref_offset, ref_stride,
                     &min, &max);
#endif
      if ((max - min) > minmax_max) minmax_max = (max - min);
      if ((max - min) < minmax_min) minmax_min = (max - min);
    }
  }
  return (minmax_max - minmax_min);
}

#if CONFIG_HIGHBITDEPTH
static INLINE int avg_4x4(const uint8_t *const src, const int stride,
                          const int highbd) {
  if (highbd) {
    return aom_highbd_avg_4x4(src, stride);
  } else {
    return aom_avg_4x4(src, stride);
  }
}
#else
static INLINE int avg_4x4(const uint8_t *const src, const int stride) {
  return aom_avg_4x4(src, stride);
}
#endif

#if CONFIG_HIGHBITDEPTH
static INLINE int avg_8x8(const uint8_t *const src, const int stride,
                          const int highbd) {
  if (highbd) {
    return aom_highbd_avg_8x8(src, stride);
  } else {
    return aom_avg_8x8(src, stride);
  }
}
#else
static INLINE int avg_8x8(const uint8_t *const src, const int stride) {
  return aom_avg_8x8(src, stride);
}
#endif

static void init_variance_tree(VAR_TREE *const vt,
#if CONFIG_HIGHBITDEPTH
                               const int highbd,
#endif
                               BLOCK_SIZE bsize, BLOCK_SIZE leaf_size,
                               const int width, const int height,
                               const uint8_t *const src, const int src_stride,
                               const uint8_t *const ref, const int ref_stride) {
  assert(bsize >= leaf_size);

  vt->bsize = bsize;

  vt->force_split = 0;

  vt->src = src;
  vt->src_stride = src_stride;
  vt->ref = ref;
  vt->ref_stride = ref_stride;

  vt->width = width;
  vt->height = height;

#if CONFIG_HIGHBITDEPTH
  vt->highbd = highbd;
#endif  // CONFIG_HIGHBITDEPTH

  if (bsize > leaf_size) {
    const BLOCK_SIZE subsize = get_subsize(bsize, PARTITION_SPLIT);
    const int px = block_size_wide[subsize];

    init_variance_tree(vt->split[0],
#if CONFIG_HIGHBITDEPTH
                       highbd,
#endif  // CONFIG_HIGHBITDEPTH
                       subsize, leaf_size, AOMMIN(px, width),
                       AOMMIN(px, height), src, src_stride, ref, ref_stride);
    init_variance_tree(vt->split[1],
#if CONFIG_HIGHBITDEPTH
                       highbd,
#endif  // CONFIG_HIGHBITDEPTH
                       subsize, leaf_size, width - px, AOMMIN(px, height),
                       src + px, src_stride, ref + px, ref_stride);
    init_variance_tree(vt->split[2],
#if CONFIG_HIGHBITDEPTH
                       highbd,
#endif  // CONFIG_HIGHBITDEPTH
                       subsize, leaf_size, AOMMIN(px, width), height - px,
                       src + px * src_stride, src_stride, ref + px * ref_stride,
                       ref_stride);
    init_variance_tree(vt->split[3],
#if CONFIG_HIGHBITDEPTH
                       highbd,
#endif  // CONFIG_HIGHBITDEPTH
                       subsize, leaf_size, width - px, height - px,
                       src + px * src_stride + px, src_stride,
                       ref + px * ref_stride + px, ref_stride);
  }
}

// Fill the variance tree based on averaging pixel values (sub-sampling), at
// the leaf node size.
static void fill_variance_tree(VAR_TREE *const vt, const BLOCK_SIZE leaf_size) {
  if (vt->bsize > leaf_size) {
    fill_variance_tree(vt->split[0], leaf_size);
    fill_variance_tree(vt->split[1], leaf_size);
    fill_variance_tree(vt->split[2], leaf_size);
    fill_variance_tree(vt->split[3], leaf_size);
    fill_variance_node(vt);
  } else if (vt->width <= 0 || vt->height <= 0) {
    fill_variance(0, 0, 0, &vt->variances.none);
  } else {
    unsigned int sse = 0;
    int sum = 0;
    int src_avg;
    int ref_avg;
    assert(leaf_size == BLOCK_4X4 || leaf_size == BLOCK_8X8);
    if (leaf_size == BLOCK_4X4) {
      src_avg = avg_4x4(vt->src, vt->src_stride IF_HBD(, vt->highbd));
      ref_avg = avg_4x4(vt->ref, vt->ref_stride IF_HBD(, vt->highbd));
    } else {
      src_avg = avg_8x8(vt->src, vt->src_stride IF_HBD(, vt->highbd));
      ref_avg = avg_8x8(vt->ref, vt->ref_stride IF_HBD(, vt->highbd));
    }
    sum = src_avg - ref_avg;
    sse = sum * sum;
    fill_variance(sse, sum, 0, &vt->variances.none);
  }
}

static void refine_variance_tree(VAR_TREE *const vt, const int64_t threshold) {
  if (vt->bsize >= BLOCK_8X8) {
    if (vt->bsize == BLOCK_16X16) {
      if (vt->variances.none.variance <= threshold)
        return;
      else
        vt->force_split = 0;
    }

    refine_variance_tree(vt->split[0], threshold);
    refine_variance_tree(vt->split[1], threshold);
    refine_variance_tree(vt->split[2], threshold);
    refine_variance_tree(vt->split[3], threshold);

    if (vt->bsize <= BLOCK_16X16) fill_variance_node(vt);
  } else if (vt->width <= 0 || vt->height <= 0) {
    fill_variance(0, 0, 0, &vt->variances.none);
  } else {
    const int src_avg = avg_4x4(vt->src, vt->src_stride IF_HBD(, vt->highbd));
    const int ref_avg = avg_4x4(vt->ref, vt->ref_stride IF_HBD(, vt->highbd));
    const int sum = src_avg - ref_avg;
    const unsigned int sse = sum * sum;
    assert(vt->bsize == BLOCK_4X4);
    fill_variance(sse, sum, 0, &vt->variances.none);
  }
}

static int check_split_key_frame(VAR_TREE *const vt, const int64_t threshold) {
  if (vt->bsize == BLOCK_32X32) {
    vt->force_split = vt->variances.none.variance > threshold;
  } else {
    vt->force_split |= check_split_key_frame(vt->split[0], threshold);
    vt->force_split |= check_split_key_frame(vt->split[1], threshold);
    vt->force_split |= check_split_key_frame(vt->split[2], threshold);
    vt->force_split |= check_split_key_frame(vt->split[3], threshold);
  }
  return vt->force_split;
}

static int check_split(AV1_COMP *const cpi, VAR_TREE *const vt,
                       const int segment_id, const int64_t *const thresholds) {
  if (vt->bsize == BLOCK_16X16) {
    vt->force_split = vt->variances.none.variance > thresholds[0];
    if (!vt->force_split && vt->variances.none.variance > thresholds[-1] &&
        !cyclic_refresh_segment_id_boosted(segment_id)) {
      // We have some nominal amount of 16x16 variance (based on average),
      // compute the minmax over the 8x8 sub-blocks, and if above threshold,
      // force split to 8x8 block for this 16x16 block.
      int minmax =
          compute_minmax_8x8(vt->src, vt->src_stride, vt->ref, vt->ref_stride,
#if CONFIG_HIGHBITDEPTH
                             vt->highbd,
#endif
                             vt->width, vt->height);
      vt->force_split = minmax > cpi->vbp_threshold_minmax;
    }
  } else {
    vt->force_split |=
        check_split(cpi, vt->split[0], segment_id, thresholds + 1);
    vt->force_split |=
        check_split(cpi, vt->split[1], segment_id, thresholds + 1);
    vt->force_split |=
        check_split(cpi, vt->split[2], segment_id, thresholds + 1);
    vt->force_split |=
        check_split(cpi, vt->split[3], segment_id, thresholds + 1);

    if (vt->bsize == BLOCK_32X32 && !vt->force_split) {
      vt->force_split = vt->variances.none.variance > thresholds[0];
    }
  }

  return vt->force_split;
}

// This function chooses partitioning based on the variance between source and
// reconstructed last (or golden), where variance is computed for down-sampled
// inputs.
static void choose_partitioning(AV1_COMP *const cpi, ThreadData *const <