encodeframe.c 230 KB
Newer Older
Jingning Han's avatar
Jingning Han committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/*
 *  Copyright (c) 2010 The WebM project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <limits.h>
#include <math.h>
#include <stdio.h>

Yaowu Xu's avatar
Yaowu Xu committed
15 16 17
#include "./av1_rtcd.h"
#include "./aom_dsp_rtcd.h"
#include "./aom_config.h"
Jingning Han's avatar
Jingning Han committed
18

Yaowu Xu's avatar
Yaowu Xu committed
19
#include "aom_dsp/aom_dsp_common.h"
20
#include "aom_ports/mem.h"
Yaowu Xu's avatar
Yaowu Xu committed
21
#include "aom_ports/aom_timer.h"
22
#include "aom_ports/system_state.h"
Jingning Han's avatar
Jingning Han committed
23

24 25 26 27 28 29 30 31 32 33 34
#include "av1/common/common.h"
#include "av1/common/entropy.h"
#include "av1/common/entropymode.h"
#include "av1/common/idct.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"
Jingning Han's avatar
Jingning Han committed
35

36 37 38
#include "av1/encoder/aq_complexity.h"
#include "av1/encoder/aq_cyclicrefresh.h"
#include "av1/encoder/aq_variance.h"
39
#if CONFIG_SUPERTX
40
#include "av1/encoder/cost.h"
41
#endif
42
#if CONFIG_GLOBAL_MOTION
43
#include "av1/encoder/global_motion.h"
44
#endif
45 46 47 48 49 50 51 52 53
#include "av1/encoder/encodeframe.h"
#include "av1/encoder/encodemb.h"
#include "av1/encoder/encodemv.h"
#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"
Jingning Han's avatar
Jingning Han committed
54

Yaowu Xu's avatar
Yaowu Xu committed
55
#if CONFIG_AOM_HIGHBITDEPTH
56
#define IF_HBD(...) __VA_ARGS__
57
#else
58
#define IF_HBD(...)
Yaowu Xu's avatar
Yaowu Xu committed
59
#endif  // CONFIG_AOM_HIGHBITDEPTH
60

Yaowu Xu's avatar
Yaowu Xu committed
61
static void encode_superblock(AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
62 63
                              int output_enabled, int mi_row, int mi_col,
                              BLOCK_SIZE bsize, PICK_MODE_CONTEXT *ctx);
Jingning Han's avatar
Jingning Han committed
64

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

Yaowu Xu's avatar
Yaowu Xu committed
68 69 70
static int check_intra_sb(AV1_COMP *cpi, const TileInfo *const tile, int mi_row,
                          int mi_col, BLOCK_SIZE bsize, PC_TREE *pc_tree);
static void predict_superblock(AV1_COMP *cpi, ThreadData *td,
71 72 73
#if CONFIG_EXT_INTER
                               int mi_row_ori, int mi_col_ori,
#endif  // CONFIG_EXT_INTER
74 75 76 77
                               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);
Yaowu Xu's avatar
Yaowu Xu committed
78
static void predict_sb_complex(AV1_COMP *cpi, ThreadData *td,
79 80
                               const TileInfo *const tile, int mi_row,
                               int mi_col, int mi_row_ori, int mi_col_ori,
81
                               int output_enabled, BLOCK_SIZE bsize,
82 83
                               BLOCK_SIZE top_bsize, uint8_t *dst_buf[3],
                               int dst_stride[3], PC_TREE *pc_tree);
Yaowu Xu's avatar
Yaowu Xu committed
84
static void update_state_sb_supertx(AV1_COMP *cpi, ThreadData *td,
85 86
                                    const TileInfo *const tile, int mi_row,
                                    int mi_col, BLOCK_SIZE bsize,
87
                                    int output_enabled, PC_TREE *pc_tree);
Yaowu Xu's avatar
Yaowu Xu committed
88
static void rd_supertx_sb(AV1_COMP *cpi, ThreadData *td,
89 90 91
                          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);
92 93
#endif  // CONFIG_SUPERTX

Jingning Han's avatar
Jingning Han committed
94 95 96 97
// 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.
Yaowu Xu's avatar
Yaowu Xu committed
98
static const uint8_t AV1_VAR_OFFS[MAX_SB_SIZE] = {
99 100 101 102 103
  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,
104
#if CONFIG_EXT_PARTITION
105 106 107 108 109
  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
110
#endif  // CONFIG_EXT_PARTITION
Jingning Han's avatar
Jingning Han committed
111 112
};

Yaowu Xu's avatar
Yaowu Xu committed
113 114
#if CONFIG_AOM_HIGHBITDEPTH
static const uint16_t AV1_HIGH_VAR_OFFS_8[MAX_SB_SIZE] = {
115 116 117 118 119
  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,
120
#if CONFIG_EXT_PARTITION
121 122 123 124 125
  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
126
#endif  // CONFIG_EXT_PARTITION
Jingning Han's avatar
Jingning Han committed
127 128
};

Yaowu Xu's avatar
Yaowu Xu committed
129
static const uint16_t AV1_HIGH_VAR_OFFS_10[MAX_SB_SIZE] = {
130 131 132 133 134 135 136 137
  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,
138
#if CONFIG_EXT_PARTITION
139 140 141 142 143 144 145 146
  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
147
#endif  // CONFIG_EXT_PARTITION
Jingning Han's avatar
Jingning Han committed
148 149
};

Yaowu Xu's avatar
Yaowu Xu committed
150
static const uint16_t AV1_HIGH_VAR_OFFS_12[MAX_SB_SIZE] = {
151 152 153 154 155 156 157 158 159 160
  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,
161
#if CONFIG_EXT_PARTITION
162 163 164 165 166 167 168 169 170 171
  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
172
#endif  // CONFIG_EXT_PARTITION
Jingning Han's avatar
Jingning Han committed
173
};
Yaowu Xu's avatar
Yaowu Xu committed
174
#endif  // CONFIG_AOM_HIGHBITDEPTH
Jingning Han's avatar
Jingning Han committed
175

Yaowu Xu's avatar
Yaowu Xu committed
176 177 178
unsigned int av1_get_sby_perpixel_variance(AV1_COMP *cpi,
                                           const struct buf_2d *ref,
                                           BLOCK_SIZE bs) {
Jingning Han's avatar
Jingning Han committed
179
  unsigned int sse;
180
  const unsigned int var =
Yaowu Xu's avatar
Yaowu Xu committed
181
      cpi->fn_ptr[bs].vf(ref->buf, ref->stride, AV1_VAR_OFFS, 0, &sse);
Jingning Han's avatar
Jingning Han committed
182 183 184
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}

Yaowu Xu's avatar
Yaowu Xu committed
185 186 187 188
#if CONFIG_AOM_HIGHBITDEPTH
unsigned int av1_high_get_sby_perpixel_variance(AV1_COMP *cpi,
                                                const struct buf_2d *ref,
                                                BLOCK_SIZE bs, int bd) {
Jingning Han's avatar
Jingning Han committed
189 190 191
  unsigned int var, sse;
  switch (bd) {
    case 10:
Yaowu Xu's avatar
Yaowu Xu committed
192 193 194
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_10), 0, &sse);
Jingning Han's avatar
Jingning Han committed
195 196
      break;
    case 12:
Yaowu Xu's avatar
Yaowu Xu committed
197 198 199
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_12), 0, &sse);
Jingning Han's avatar
Jingning Han committed
200 201 202
      break;
    case 8:
    default:
203 204
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
Yaowu Xu's avatar
Yaowu Xu committed
205
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8), 0, &sse);
Jingning Han's avatar
Jingning Han committed
206 207 208 209
      break;
  }
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}
Yaowu Xu's avatar
Yaowu Xu committed
210
#endif  // CONFIG_AOM_HIGHBITDEPTH
Jingning Han's avatar
Jingning Han committed
211

Yaowu Xu's avatar
Yaowu Xu committed
212
static unsigned int get_sby_perpixel_diff_variance(AV1_COMP *cpi,
Jingning Han's avatar
Jingning Han committed
213 214 215 216 217 218 219 220 221 222 223 224 225 226
                                                   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]);
}

Yaowu Xu's avatar
Yaowu Xu committed
227 228
static BLOCK_SIZE get_rd_var_based_fixed_partition(AV1_COMP *cpi, MACROBLOCK *x,
                                                   int mi_row, int mi_col) {
229 230
  unsigned int var = get_sby_perpixel_diff_variance(
      cpi, &x->plane[0].src, mi_row, mi_col, BLOCK_64X64);
Jingning Han's avatar
Jingning Han committed
231 232 233 234 235 236 237 238 239 240 241 242
  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.
Yaowu Xu's avatar
Yaowu Xu committed
243
static void set_mode_info_offsets(AV1_COMP *const cpi, MACROBLOCK *const x,
244
                                  MACROBLOCKD *const xd, int mi_row,
245
                                  int mi_col) {
Yaowu Xu's avatar
Yaowu Xu committed
246
  AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
247 248 249
  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;
250
  x->mbmi_ext = cpi->mbmi_ext_base + (mi_row * cm->mi_cols + mi_col);
Jingning Han's avatar
Jingning Han committed
251 252
}

Yaowu Xu's avatar
Yaowu Xu committed
253
static void set_offsets_without_segment_id(AV1_COMP *cpi,
254
                                           const TileInfo *const tile,
255 256
                                           MACROBLOCK *const x, int mi_row,
                                           int mi_col, BLOCK_SIZE bsize) {
Yaowu Xu's avatar
Yaowu Xu committed
257
  AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
258 259 260 261 262 263
  MACROBLOCKD *const xd = &x->e_mbd;
  const int mi_width = num_8x8_blocks_wide_lookup[bsize];
  const int mi_height = num_8x8_blocks_high_lookup[bsize];

  set_skip_context(xd, mi_row, mi_col);

264
  set_mode_info_offsets(cpi, x, xd, mi_row, mi_col);
Jingning Han's avatar
Jingning Han committed
265

266 267
#if CONFIG_VAR_TX
  xd->above_txfm_context = cm->above_txfm_context + mi_col;
268
  xd->left_txfm_context =
269
      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
270
  xd->max_tx_size = max_txsize_lookup[bsize];
271 272
#endif

Jingning Han's avatar
Jingning Han committed
273
  // Set up destination pointers.
Yaowu Xu's avatar
Yaowu Xu committed
274
  av1_setup_dst_planes(xd->plane, get_frame_new_buffer(cm), mi_row, mi_col);
Jingning Han's avatar
Jingning Han committed
275 276 277

  // Set up limit values for MV components.
  // Mv beyond the range do not produce new/different prediction block.
Yaowu Xu's avatar
Yaowu Xu committed
278 279 280 281
  x->mv_row_min = -(((mi_row + mi_height) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_col_min = -(((mi_col + mi_width) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_row_max = (cm->mi_rows - mi_row) * MI_SIZE + AOM_INTERP_EXTEND;
  x->mv_col_max = (cm->mi_cols - mi_col) * MI_SIZE + AOM_INTERP_EXTEND;
Jingning Han's avatar
Jingning Han committed
282 283 284

  // 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)));
285 286
  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, cm->mi_rows,
                 cm->mi_cols);
Jingning Han's avatar
Jingning Han committed
287 288

  // Set up source buffers.
Yaowu Xu's avatar
Yaowu Xu committed
289
  av1_setup_src_planes(x, cpi->Source, mi_row, mi_col);
Jingning Han's avatar
Jingning Han committed
290 291 292 293 294

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

Yaowu Xu's avatar
Yaowu Xu committed
295
  // required by av1_append_sub8x8_mvs_for_idx() and av1_find_best_ref_mvs()
296 297 298
  xd->tile = *tile;
}

Yaowu Xu's avatar
Yaowu Xu committed
299
static void set_offsets(AV1_COMP *cpi, const TileInfo *const tile,
300 301
                        MACROBLOCK *const x, int mi_row, int mi_col,
                        BLOCK_SIZE bsize) {
Yaowu Xu's avatar
Yaowu Xu committed
302
  AV1_COMMON *const cm = &cpi->common;
303 304 305 306 307 308 309 310
  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;

Jingning Han's avatar
Jingning Han committed
311 312
  // Setup segment ID.
  if (seg->enabled) {
Geza Lore's avatar
Geza Lore committed
313
    if (!cpi->vaq_refresh) {
314 315
      const uint8_t *const map =
          seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map;
Jingning Han's avatar
Jingning Han committed
316 317
      mbmi->segment_id = get_segment_id(cm, map, bsize, mi_row, mi_col);
    }
Yaowu Xu's avatar
Yaowu Xu committed
318
    av1_init_plane_quantizers(cpi, x, mbmi->segment_id);
Jingning Han's avatar
Jingning Han committed
319 320 321 322 323 324

    x->encode_breakout = cpi->segment_encode_breakout[mbmi->segment_id];
  } else {
    mbmi->segment_id = 0;
    x->encode_breakout = cpi->encode_breakout;
  }
325 326 327 328

#if CONFIG_SUPERTX
  mbmi->segment_id_supertx = MAX_SEGMENTS;
#endif  // CONFIG_SUPERTX
Jingning Han's avatar
Jingning Han committed
329 330
}

331
#if CONFIG_SUPERTX
Yaowu Xu's avatar
Yaowu Xu committed
332
static void set_offsets_supertx(AV1_COMP *cpi, ThreadData *td,
333 334
                                const TileInfo *const tile, int mi_row,
                                int mi_col, BLOCK_SIZE bsize) {
335
  MACROBLOCK *const x = &td->mb;
Yaowu Xu's avatar
Yaowu Xu committed
336
  AV1_COMMON *const cm = &cpi->common;
337 338 339 340 341 342 343 344
  MACROBLOCKD *const xd = &x->e_mbd;
  const int mi_width = num_8x8_blocks_wide_lookup[bsize];
  const int mi_height = num_8x8_blocks_high_lookup[bsize];

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

  // 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)));
345 346
  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, cm->mi_rows,
                 cm->mi_cols);
347 348
}

Yaowu Xu's avatar
Yaowu Xu committed
349
static void set_offsets_extend(AV1_COMP *cpi, ThreadData *td,
350 351
                               const TileInfo *const tile, int mi_row_pred,
                               int mi_col_pred, int mi_row_ori, int mi_col_ori,
352
                               BLOCK_SIZE bsize_pred) {
353 354 355 356
  // 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;
Yaowu Xu's avatar
Yaowu Xu committed
357
  AV1_COMMON *const cm = &cpi->common;
358 359 360 361 362 363 364 365
  MACROBLOCKD *const xd = &x->e_mbd;
  const int mi_width = num_8x8_blocks_wide_lookup[bsize_pred];
  const int mi_height = num_8x8_blocks_high_lookup[bsize_pred];

  set_mode_info_offsets(cpi, x, xd, mi_row_ori, mi_col_ori);

  // Set up limit values for MV components.
  // Mv beyond the range do not produce new/different prediction block.
Yaowu Xu's avatar
Yaowu Xu committed
366 367 368 369
  x->mv_row_min = -(((mi_row_pred + mi_height) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_col_min = -(((mi_col_pred + mi_width) * MI_SIZE) + AOM_INTERP_EXTEND);
  x->mv_row_max = (cm->mi_rows - mi_row_pred) * MI_SIZE + AOM_INTERP_EXTEND;
  x->mv_col_max = (cm->mi_cols - mi_col_pred) * MI_SIZE + AOM_INTERP_EXTEND;
370 371 372 373 374

  // Set up distance of MB to edge of frame in 1/8th pel units.
  assert(!(mi_col_pred & (mi_width - 1)) && !(mi_row_pred & (mi_height - 1)));
  set_mi_row_col(xd, tile, mi_row_pred, mi_height, mi_col_pred, mi_width,
                 cm->mi_rows, cm->mi_cols);
375 376
  xd->up_available = (mi_row_ori > tile->mi_row_start);
  xd->left_available = (mi_col_ori > tile->mi_col_start);
377 378 379 380

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

Yaowu Xu's avatar
Yaowu Xu committed
383
static void set_segment_id_supertx(const AV1_COMP *const cpi,
384 385
                                   MACROBLOCK *const x, const int mi_row,
                                   const int mi_col, const BLOCK_SIZE bsize) {
Yaowu Xu's avatar
Yaowu Xu committed
386
  const AV1_COMMON *cm = &cpi->common;
387 388
  const struct segmentation *seg = &cm->seg;
  const int miw =
Yaowu Xu's avatar
Yaowu Xu committed
389
      AOMMIN(num_8x8_blocks_wide_lookup[bsize], cm->mi_cols - mi_col);
390
  const int mih =
Yaowu Xu's avatar
Yaowu Xu committed
391
      AOMMIN(num_8x8_blocks_high_lookup[bsize], cm->mi_rows - mi_row);
392 393 394 395
  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;
396

397 398
  if (!seg->enabled) {
    seg_id_supertx = 0;
399
    x->encode_breakout = cpi->encode_breakout;
400 401
  } else {
    // Find the minimum segment_id
402 403 404
    for (r = 0; r < mih; r++)
      for (c = 0; c < miw; c++)
        seg_id_supertx =
Yaowu Xu's avatar
Yaowu Xu committed
405
            AOMMIN(mip[r * cm->mi_stride + c]->mbmi.segment_id, seg_id_supertx);
406 407 408
    assert(0 <= seg_id_supertx && seg_id_supertx < MAX_SEGMENTS);

    // Initialize plane quantisers
Yaowu Xu's avatar
Yaowu Xu committed
409
    av1_init_plane_quantizers(cpi, x, seg_id_supertx);
410
    x->encode_breakout = cpi->segment_encode_breakout[seg_id_supertx];
411
  }
412 413

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

Yaowu Xu's avatar
Yaowu Xu committed
420
static void set_block_size(AV1_COMP *const cpi, MACROBLOCK *const x,
421
                           MACROBLOCKD *const xd, int mi_row, int mi_col,
Jingning Han's avatar
Jingning Han committed
422 423
                           BLOCK_SIZE bsize) {
  if (cpi->common.mi_cols > mi_col && cpi->common.mi_rows > mi_row) {
424
    set_mode_info_offsets(cpi, x, xd, mi_row, mi_col);
Jingning Han's avatar
Jingning Han committed
425 426 427 428
    xd->mi[0]->mbmi.sb_type = bsize;
  }
}

Yaowu Xu's avatar
Yaowu Xu committed
429
static void set_vt_partitioning(AV1_COMP *cpi, MACROBLOCK *const x,
430 431 432
                                MACROBLOCKD *const xd, VAR_TREE *vt, int mi_row,
                                int mi_col, const int64_t *const threshold,
                                const BLOCK_SIZE *const bsize_min) {
Yaowu Xu's avatar
Yaowu Xu committed
433
  AV1_COMMON *const cm = &cpi->common;
434 435 436 437
  const int hbw = num_8x8_blocks_wide_lookup[vt->bsize] / 2;
  const int hbh = num_8x8_blocks_high_lookup[vt->bsize] / 2;
  const int has_cols = mi_col + hbw < cm->mi_cols;
  const int has_rows = mi_row + hbh < cm->mi_rows;
Jingning Han's avatar
Jingning Han committed
438

439
  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return;
Jingning Han's avatar
Jingning Han committed
440

441 442 443 444
  assert(vt->bsize >= BLOCK_8X8);

  assert(hbh == hbw);

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

450
  if (vt->force_split || (!has_cols && !has_rows)) goto split;
Jingning Han's avatar
Jingning Han committed
451 452 453 454

  // 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.
455
  if (vt->bsize == bsize_min[0]) {
456
    if (has_cols && has_rows && vt->variances.none.variance < threshold[0]) {
457 458 459 460 461 462 463 464 465 466 467
      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;
Jingning Han's avatar
Jingning Han committed
468
    }
469
  } else if (vt->bsize > bsize_min[0]) {
Jingning Han's avatar
Jingning Han committed
470 471
    // For key frame: take split for bsize above 32X32 or very high variance.
    if (cm->frame_type == KEY_FRAME &&
472
        (vt->bsize > BLOCK_32X32 ||
473
         vt->variances.none.variance > (threshold[0] << 4))) {
474
      goto split;
Jingning Han's avatar
Jingning Han committed
475 476
    }
    // If variance is low, take the bsize (no split).
477
    if (has_cols && has_rows && vt->variances.none.variance < threshold[0]) {
478 479
      set_block_size(cpi, x, xd, mi_row, mi_col, vt->bsize);
      return;
Jingning Han's avatar
Jingning Han committed
480 481 482
    }

    // Check vertical split.
483 484 485 486
    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] &&
Jingning Han's avatar
Jingning Han committed
487 488
          get_plane_block_size(subsize, &xd->plane[1]) < BLOCK_INVALID) {
        set_block_size(cpi, x, xd, mi_row, mi_col, subsize);
489 490
        set_block_size(cpi, x, xd, mi_row, mi_col + hbw, subsize);
        return;
Jingning Han's avatar
Jingning Han committed
491 492 493
      }
    }
    // Check horizontal split.
494 495 496 497
    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] &&
Jingning Han's avatar
Jingning Han committed
498 499
          get_plane_block_size(subsize, &xd->plane[1]) < BLOCK_INVALID) {
        set_block_size(cpi, x, xd, mi_row, mi_col, subsize);
500 501
        set_block_size(cpi, x, xd, mi_row + hbh, mi_col, subsize);
        return;
Jingning Han's avatar
Jingning Han committed
502 503
      }
    }
504
  }
Jingning Han's avatar
Jingning Han committed
505

506 507 508 509 510 511 512 513 514 515 516
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;
}
Jingning Han's avatar
Jingning Han committed
517 518 519 520 521 522
}

// 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.
Yaowu Xu's avatar
Yaowu Xu committed
523 524
static void set_vbp_thresholds(AV1_COMP *cpi, int64_t thresholds[], int q) {
  AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
525 526
  const int is_key_frame = (cm->frame_type == KEY_FRAME);
  const int threshold_multiplier = is_key_frame ? 20 : 1;
527 528
  const int64_t threshold_base =
      (int64_t)(threshold_multiplier * cpi->y_dequant[q][1]);
Jingning Han's avatar
Jingning Han committed
529
  if (is_key_frame) {
530
    thresholds[1] = threshold_base;
Jingning Han's avatar
Jingning Han committed
531
    thresholds[2] = threshold_base >> 2;
532 533
    thresholds[3] = threshold_base >> 2;
    thresholds[4] = threshold_base << 2;
Jingning Han's avatar
Jingning Han committed
534
  } else {
535
    thresholds[2] = threshold_base;
Jingning Han's avatar
Jingning Han committed
536
    if (cm->width <= 352 && cm->height <= 288) {
537 538
      thresholds[1] = threshold_base >> 2;
      thresholds[3] = threshold_base << 3;
Jingning Han's avatar
Jingning Han committed
539
    } else {
540 541
      thresholds[1] = threshold_base;
      thresholds[2] = (5 * threshold_base) >> 2;
Jingning Han's avatar
Jingning Han committed
542
      if (cm->width >= 1920 && cm->height >= 1080)
543 544
        thresholds[2] = (7 * threshold_base) >> 2;
      thresholds[3] = threshold_base << cpi->oxcf.speed;
Jingning Han's avatar
Jingning Han committed
545 546
    }
  }
547
  thresholds[0] = INT64_MIN;
Jingning Han's avatar
Jingning Han committed
548 549
}

Yaowu Xu's avatar
Yaowu Xu committed
550 551
void av1_set_variance_partition_thresholds(AV1_COMP *cpi, int q) {
  AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
552 553 554 555 556 557 558 559 560 561 562 563 564 565 566
  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
567 568 569
        cpi->vbp_threshold_sad = (cpi->y_dequant[q][1] << 1) > 1000
                                     ? (cpi->y_dequant[q][1] << 1)
                                     : 1000;
Jingning Han's avatar
Jingning Han committed
570 571 572 573 574 575 576
      cpi->vbp_bsize_min = BLOCK_16X16;
    }
    cpi->vbp_threshold_minmax = 15 + (q >> 3);
  }
}

// Compute the minmax over the 8x8 subblocks.
577 578
static int compute_minmax_8x8(const uint8_t *src, int src_stride,
                              const uint8_t *ref, int ref_stride,
Yaowu Xu's avatar
Yaowu Xu committed
579
#if CONFIG_AOM_HIGHBITDEPTH
580
                              int highbd,
Jingning Han's avatar
Jingning Han committed
581
#endif
582
                              int pixels_wide, int pixels_high) {
Jingning Han's avatar
Jingning Han committed
583 584 585 586 587
  int k;
  int minmax_max = 0;
  int minmax_min = 255;
  // Loop over the 4 8x8 subblocks.
  for (k = 0; k < 4; k++) {
588 589
    const int x8_idx = ((k & 1) << 3);
    const int y8_idx = ((k >> 1) << 3);
Jingning Han's avatar
Jingning Han committed
590 591 592
    int min = 0;
    int max = 0;
    if (x8_idx < pixels_wide && y8_idx < pixels_high) {
593 594
      const int src_offset = y8_idx * src_stride + x8_idx;
      const int ref_offset = y8_idx * ref_stride + x8_idx;
Yaowu Xu's avatar
Yaowu Xu committed
595
#if CONFIG_AOM_HIGHBITDEPTH
596
      if (highbd) {
Yaowu Xu's avatar
Yaowu Xu committed
597
        aom_highbd_minmax_8x8(src + src_offset, src_stride, ref + ref_offset,
598
                              ref_stride, &min, &max);
Jingning Han's avatar
Jingning Han committed
599
      } else {
Yaowu Xu's avatar
Yaowu Xu committed
600
        aom_minmax_8x8(src + src_offset, src_stride, ref + ref_offset,
601
                       ref_stride, &min, &max);
Jingning Han's avatar
Jingning Han committed
602 603
      }
#else
Yaowu Xu's avatar
Yaowu Xu committed
604
      aom_minmax_8x8(src + src_offset, src_stride, ref + ref_offset, ref_stride,
Jingning Han's avatar
Jingning Han committed
605 606
                     &min, &max);
#endif
607 608
      if ((max - min) > minmax_max) minmax_max = (max - min);
      if ((max - min) < minmax_min) minmax_min = (max - min);
Jingning Han's avatar
Jingning Han committed
609 610 611 612 613
    }
  }
  return (minmax_max - minmax_min);
}

Yaowu Xu's avatar
Yaowu Xu committed
614
#if CONFIG_AOM_HIGHBITDEPTH
615 616 617
static INLINE int avg_4x4(const uint8_t *const src, const int stride,
                          const int highbd) {
  if (highbd) {
Yaowu Xu's avatar
Yaowu Xu committed
618
    return aom_highbd_avg_4x4(src, stride);
619
  } else {
Yaowu Xu's avatar
Yaowu Xu committed
620
    return aom_avg_4x4(src, stride);
621 622
  }
}
Jingning Han's avatar
Jingning Han committed
623
#else
624
static INLINE int avg_4x4(const uint8_t *const src, const int stride) {
Yaowu Xu's avatar
Yaowu Xu committed
625
  return aom_avg_4x4(src, stride);
626
}
Jingning Han's avatar
Jingning Han committed
627
#endif
628

Yaowu Xu's avatar
Yaowu Xu committed
629
#if CONFIG_AOM_HIGHBITDEPTH
630 631 632
static INLINE int avg_8x8(const uint8_t *const src, const int stride,
                          const int highbd) {
  if (highbd) {
Yaowu Xu's avatar
Yaowu Xu committed
633
    return aom_highbd_avg_8x8(src, stride);
634
  } else {
Yaowu Xu's avatar
Yaowu Xu committed
635
    return aom_avg_8x8(src, stride);
Jingning Han's avatar
Jingning Han committed
636 637
  }
}
638 639
#else
static INLINE int avg_8x8(const uint8_t *const src, const int stride) {
Yaowu Xu's avatar
Yaowu Xu committed
640
  return aom_avg_8x8(src, stride);
641 642
}
#endif
Jingning Han's avatar
Jingning Han committed
643

644
static void init_variance_tree(VAR_TREE *const vt,
Yaowu Xu's avatar
Yaowu Xu committed
645
#if CONFIG_AOM_HIGHBITDEPTH
646
                               const int highbd,
Jingning Han's avatar
Jingning Han committed
647
#endif
648
                               BLOCK_SIZE bsize, BLOCK_SIZE leaf_size,
649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665
                               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;

Yaowu Xu's avatar
Yaowu Xu committed
666
#if CONFIG_AOM_HIGHBITDEPTH
667
  vt->highbd = highbd;
Yaowu Xu's avatar
Yaowu Xu committed
668
#endif  // CONFIG_AOM_HIGHBITDEPTH
669 670 671 672 673 674

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

    init_variance_tree(vt->split[0],
Yaowu Xu's avatar
Yaowu Xu committed
675
#if CONFIG_AOM_HIGHBITDEPTH
676
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
677 678 679
#endif  // CONFIG_AOM_HIGHBITDEPTH
                       subsize, leaf_size, AOMMIN(px, width),
                       AOMMIN(px, height), src, src_stride, ref, ref_stride);
680
    init_variance_tree(vt->split[1],
Yaowu Xu's avatar
Yaowu Xu committed
681
#if CONFIG_AOM_HIGHBITDEPTH
682
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
683 684
#endif  // CONFIG_AOM_HIGHBITDEPTH
                       subsize, leaf_size, width - px, AOMMIN(px, height),
685
                       src + px, src_stride, ref + px, ref_stride);
686
    init_variance_tree(vt->split[2],
Yaowu Xu's avatar
Yaowu Xu committed
687
#if CONFIG_AOM_HIGHBITDEPTH
688
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
689 690
#endif  // CONFIG_AOM_HIGHBITDEPTH
                       subsize, leaf_size, AOMMIN(px, width), height - px,
691 692
                       src + px * src_stride, src_stride, ref + px * ref_stride,
                       ref_stride);
693
    init_variance_tree(vt->split[3],
Yaowu Xu's avatar
Yaowu Xu committed
694
#if CONFIG_AOM_HIGHBITDEPTH
695
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
696
#endif  // CONFIG_AOM_HIGHBITDEPTH
697
                       subsize, leaf_size, width - px, height - px,
698 699 700 701 702 703 704
                       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.
705
static void fill_variance_tree(VAR_TREE *const vt, const BLOCK_SIZE leaf_size) {
706 707 708 709 710 711 712 713 714
  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 {
Jingning Han's avatar
Jingning Han committed
715 716
    unsigned int sse = 0;
    int sum = 0;
717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746
    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);

747
    if (vt->bsize <= BLOCK_16X16) fill_variance_node(vt);
748 749 750 751 752 753
  } 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;
754
    const unsigned int sse = sum * sum;
755 756 757 758 759
    assert(vt->bsize == BLOCK_4X4);
    fill_variance(sse, sum, 0, &vt->variances.none);
  }
}

760
static int check_split_key_frame(VAR_TREE *const vt, const int64_t threshold) {
761 762 763 764 765 766 767 768 769 770 771
  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;
}

Yaowu Xu's avatar
Yaowu Xu committed
772
static int check_split(AV1_COMP *const cpi, VAR_TREE *const vt,
773
                       const int segment_id, const int64_t *const thresholds) {
774 775
  if (vt->bsize == BLOCK_16X16) {
    vt->force_split = vt->variances.none.variance > thresholds[0];
776 777
    if (!vt->force_split && vt->variances.none.variance > thresholds[-1] &&
        !cyclic_refresh_segment_id_boosted(segment_id)) {
778 779 780
      // 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.
781 782
      int minmax =
          compute_minmax_8x8(vt->src, vt->src_stride, vt->ref, vt->ref_stride,
Yaowu Xu's avatar
Yaowu Xu committed
783
#if CONFIG_AOM_HIGHBITDEPTH
784
                             vt->highbd,
Jingning Han's avatar
Jingning Han committed
785
#endif
786
                             vt->width, vt->height);
787 788 789
      vt->force_split = minmax > cpi->vbp_threshold_minmax;
    }
  } else {
790 791 792 793 794 795 796 797
    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);
798 799 800

    if (vt->bsize == BLOCK_32X32 && !vt->force_split) {
      vt->force_split = vt->variances.none.variance > thresholds[0];
Jingning Han's avatar
Jingning Han committed
801 802
    }
  }
803 804

  return vt->force_split;
Jingning Han's avatar
Jingning Han committed
805 806 807
}

// This function chooses partitioning based on the variance between source and
808 809
// reconstructed last (or golden), where variance is computed for down-sampled
// inputs.
Yaowu Xu's avatar
Yaowu Xu committed
810
static void choose_partitioning(AV1_COMP *const cpi, ThreadData *const td,
811
                                const TileInfo *const tile, MACROBLOCK *const x,
812
                                const int mi_row, const int mi_col) {
Yaowu Xu's avatar
Yaowu Xu committed
813
  AV1_COMMON *const cm = &cpi->common;
814 815 816 817 818 819 820
  MACROBLOCKD *const xd = &x->e_mbd;
  VAR_TREE *const vt = td->var_root[cm->mib_size_log2 - MIN_MIB_SIZE_LOG2];
  int i;
  const uint8_t *src;
  const uint8_t *ref;
  int src_stride;
  int ref_stride;
821 822
  int pixels_wide = 8 * num_8x8_blocks_wide_lookup[cm->sb_size];
  int pixels_high = 8 * num_8x8_blocks_high_lookup[cm->sb_size];
823
  int64_t thresholds[5] = {
824 825
    cpi->vbp_thresholds[0], cpi->vbp_thresholds[1], cpi->vbp_thresholds[2],
    cpi->vbp_thresholds[3], cpi->vbp_thresholds[4],
826
  };
827 828
  BLOCK_SIZE bsize_min[5] = { BLOCK_16X16, BLOCK_16X16, BLOCK_16X16,
                              cpi->vbp_bsize_min, BLOCK_8X8 };
829 830 831
  const int start_level = cm->sb_size == BLOCK_64X64 ? 1 : 0;
  const int64_t *const thre = thresholds + start_level;
  const BLOCK_SIZE *const bmin = bsize_min + start_level;
Jingning Han's avatar
Jingning Han committed
832 833 834 835 836

  const int is_key_frame = (cm->frame_type == KEY_FRAME);
  const int low_res = (cm->width <= 352 && cm->height <= 288);

  int segment_id = CR_SEGMENT_ID_BASE;
837

Jingning Han's avatar
Jingning Han committed
838
  if (cpi->oxcf.aq_mode == CYCLIC_REFRESH_AQ && cm->seg.enabled) {
839 840
    const uint8_t *const map =
        cm->seg.update_map ? cpi->segmentation_map : cm->last_frame_seg_map;
841
    segment_id = get_segment_id(cm, map, cm->sb_size, mi_row, mi_col);
Jingning Han's avatar
Jingning Han committed
842 843

    if (cyclic_refresh_segment_id_boosted(segment_id)) {
Yaowu Xu's avatar
Yaowu Xu committed
844
      int q = av1_get_qindex(&cm->seg, segment_id, cm->base_qindex);
Jingning Han's avatar
Jingning Han committed
845 846 847 848
      set_vbp_thresholds(cpi, thresholds, q);
    }
  }

849
  set_offsets(cpi, tile, x, mi_row, mi_col, cm->sb_size);
Jingning Han's avatar
Jingning Han committed
850

851 852
  if (xd->mb_to_right_edge < 0) pixels_wide += (xd->mb_to_right_edge >> 3);
  if (xd->mb_to_bottom_edge < 0) pixels_high += (xd->mb_to_bottom_edge >> 3);
Jingning Han's avatar
Jingning Han committed
853

854 855
  src = x->plane[0].src.buf;
  src_stride = x->plane[0].src.stride;
Jingning Han's avatar
Jingning Han committed
856

Yunqing Wang's avatar
Yunqing Wang committed
857
  if (!is_key_frame) {
Jingning Han's avatar
Jingning Han committed
858 859 860
    MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
    unsigned int uv_sad;
    const YV12_BUFFER_CONFIG *yv12 = get_ref_frame_buffer(cpi, LAST_FRAME);
861
    const YV12_BUFFER_CONFIG *yv12_g = get_ref_frame_buffer(cpi, GOLDEN_FRAME);
Jingning Han's avatar
Jingning Han committed
862
    unsigned int y_sad, y_sad_g;
863

864 865 866
    const int hbs = cm->mib_size / 2;
    const int split_vert = mi_col + hbs >= cm->mi_cols;
    const int split_horz = mi_row + hbs >= cm->mi_rows;
867 868
    BLOCK_SIZE bsize;

869
    if (split_vert && split_horz)
870
      bsize = get_subsize(cm->sb_size, PARTITION_SPLIT);
871
    else if (split_vert)
872
      bsize = get_subsize(cm->sb_size, PARTITION_VERT);
873
    else if (split_horz)
874
      bsize = get_subsize(cm->sb_size, PARTITION_HORZ);
875
    else
876
      bsize = cm->sb_size;
Jingning Han's avatar
Jingning Han committed
877 878 879 880

    assert(yv12 != NULL);

    if (yv12_g && yv12_g != yv12) {
Yaowu Xu's avatar
Yaowu Xu committed
881 882
      av1_setup_pre_planes(xd, 0, yv12_g, mi_row, mi_col,
                           &cm->frame_refs[GOLDEN_FRAME - 1].sf);
883 884 885
      y_sad_g = cpi->fn_ptr[bsize].sdf(
          x->plane[0].src.buf, x->plane[0].src.stride, xd->plane[0].pre[0].buf,
          xd->plane[0].pre[0].stride);
Jingning Han's avatar
Jingning Han committed
886 887 888 889
    } else {
      y_sad_g = UINT_MAX;
    }

Yaowu Xu's avatar
Yaowu Xu committed
890 891