encodeframe.c 234 KB
Newer Older
Jingning Han's avatar
Jingning Han committed
1
/*
Yaowu Xu's avatar
Yaowu Xu committed
2
 * Copyright (c) 2016, Alliance for Open Media. All rights reserved
Jingning Han's avatar
Jingning Han committed
3
 *
Yaowu Xu's avatar
Yaowu Xu committed
4
5
6
7
8
9
 * 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.
Jingning Han's avatar
Jingning Han committed
10
11
12
13
14
15
 */

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

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

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

25
26
27
28
29
30
31
32
33
34
35
#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
36

37
38
39
#include "av1/encoder/aq_complexity.h"
#include "av1/encoder/aq_cyclicrefresh.h"
#include "av1/encoder/aq_variance.h"
40
#if CONFIG_SUPERTX
41
#include "av1/encoder/cost.h"
42
#endif
43
#if CONFIG_GLOBAL_MOTION
44
#include "av1/common/warped_motion.h"
45
#include "av1/encoder/global_motion.h"
46
#endif
47
48
49
50
51
52
53
54
55
#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
56

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

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

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

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

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

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

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

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

Yaowu Xu's avatar
Yaowu Xu committed
178
179
180
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
181
  unsigned int sse;
182
  const unsigned int var =
Yaowu Xu's avatar
Yaowu Xu committed
183
      cpi->fn_ptr[bs].vf(ref->buf, ref->stride, AV1_VAR_OFFS, 0, &sse);
Jingning Han's avatar
Jingning Han committed
184
185
186
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}

Yaowu Xu's avatar
Yaowu Xu committed
187
188
189
190
#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
191
192
193
  unsigned int var, sse;
  switch (bd) {
    case 10:
Yaowu Xu's avatar
Yaowu Xu committed
194
195
196
      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
197
198
      break;
    case 12:
Yaowu Xu's avatar
Yaowu Xu committed
199
200
201
      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
202
203
204
      break;
    case 8:
    default:
205
206
      var =
          cpi->fn_ptr[bs].vf(ref->buf, ref->stride,
Yaowu Xu's avatar
Yaowu Xu committed
207
                             CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8), 0, &sse);
Jingning Han's avatar
Jingning Han committed
208
209
210
211
      break;
  }
  return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]);
}
Yaowu Xu's avatar
Yaowu Xu committed
212
#endif  // CONFIG_AOM_HIGHBITDEPTH
Jingning Han's avatar
Jingning Han committed
213

Yaowu Xu's avatar
Yaowu Xu committed
214
static unsigned int get_sby_perpixel_diff_variance(AV1_COMP *cpi,
Jingning Han's avatar
Jingning Han committed
215
216
217
218
219
220
221
222
223
224
225
226
227
228
                                                   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
229
230
static BLOCK_SIZE get_rd_var_based_fixed_partition(AV1_COMP *cpi, MACROBLOCK *x,
                                                   int mi_row, int mi_col) {
231
232
  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
233
234
235
236
237
238
239
240
241
242
243
244
  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
245
static void set_mode_info_offsets(AV1_COMP *const cpi, MACROBLOCK *const x,
246
                                  MACROBLOCKD *const xd, int mi_row,
247
                                  int mi_col) {
Yaowu Xu's avatar
Yaowu Xu committed
248
  AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
249
250
251
  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;
252
  x->mbmi_ext = cpi->mbmi_ext_base + (mi_row * cm->mi_cols + mi_col);
Jingning Han's avatar
Jingning Han committed
253
254
}

Yaowu Xu's avatar
Yaowu Xu committed
255
static void set_offsets_without_segment_id(AV1_COMP *cpi,
256
                                           const TileInfo *const tile,
257
258
                                           MACROBLOCK *const x, int mi_row,
                                           int mi_col, BLOCK_SIZE bsize) {
Yaowu Xu's avatar
Yaowu Xu committed
259
  AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
260
261
262
263
264
265
  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);

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

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

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

  // 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
280
281
282
283
  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
284
285
286

  // 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)));
287
288
  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
289
290

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

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

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

Yaowu Xu's avatar
Yaowu Xu committed
301
static void set_offsets(AV1_COMP *cpi, const TileInfo *const tile,
302
303
                        MACROBLOCK *const x, int mi_row, int mi_col,
                        BLOCK_SIZE bsize) {
Yaowu Xu's avatar
Yaowu Xu committed
304
  AV1_COMMON *const cm = &cpi->common;
305
306
307
308
309
310
311
312
  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
313
314
  // Setup segment ID.
  if (seg->enabled) {
Geza Lore's avatar
Geza Lore committed
315
    if (!cpi->vaq_refresh) {
316
317
      const uint8_t *const map =
          seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map;
Jingning Han's avatar
Jingning Han committed
318
319
      mbmi->segment_id = get_segment_id(cm, map, bsize, mi_row, mi_col);
    }
Yaowu Xu's avatar
Yaowu Xu committed
320
    av1_init_plane_quantizers(cpi, x, mbmi->segment_id);
Jingning Han's avatar
Jingning Han committed
321
322
323
324
325
326

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

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

333
#if CONFIG_SUPERTX
Yaowu Xu's avatar
Yaowu Xu committed
334
static void set_offsets_supertx(AV1_COMP *cpi, ThreadData *td,
335
336
                                const TileInfo *const tile, int mi_row,
                                int mi_col, BLOCK_SIZE bsize) {
337
  MACROBLOCK *const x = &td->mb;
Yaowu Xu's avatar
Yaowu Xu committed
338
  AV1_COMMON *const cm = &cpi->common;
339
340
341
342
343
344
345
346
  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)));
347
348
  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, cm->mi_rows,
                 cm->mi_cols);
349
350
}

Yaowu Xu's avatar
Yaowu Xu committed
351
static void set_offsets_extend(AV1_COMP *cpi, ThreadData *td,
352
353
                               const TileInfo *const tile, int mi_row_pred,
                               int mi_col_pred, int mi_row_ori, int mi_col_ori,
354
                               BLOCK_SIZE bsize_pred) {
355
356
357
358
  // 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
359
  AV1_COMMON *const cm = &cpi->common;
360
361
362
363
364
365
366
367
  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
368
369
370
371
  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;
372
373
374
375
376

  // 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);
377
378
  xd->up_available = (mi_row_ori > tile->mi_row_start);
  xd->left_available = (mi_col_ori > tile->mi_col_start);
379
380
381
382

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

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

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

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

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

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

Yaowu Xu's avatar
Yaowu Xu committed
431
static void set_vt_partitioning(AV1_COMP *cpi, MACROBLOCK *const x,
432
433
434
                                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
435
  AV1_COMMON *const cm = &cpi->common;
436
437
438
439
  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
440

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

443
444
445
446
  assert(vt->bsize >= BLOCK_8X8);

  assert(hbh == hbw);

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

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

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

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

508
509
510
511
512
513
514
515
516
517
518
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
519
520
521
522
523
524
}

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

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

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

Yaowu Xu's avatar
Yaowu Xu committed
616
#if CONFIG_AOM_HIGHBITDEPTH
617
618
619
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
620
    return aom_highbd_avg_4x4(src, stride);
621
  } else {
Yaowu Xu's avatar
Yaowu Xu committed
622
    return aom_avg_4x4(src, stride);
623
624
  }
}
Jingning Han's avatar
Jingning Han committed
625
#else
626
static INLINE int avg_4x4(const uint8_t *const src, const int stride) {
Yaowu Xu's avatar
Yaowu Xu committed
627
  return aom_avg_4x4(src, stride);
628
}
Jingning Han's avatar
Jingning Han committed
629
#endif
630

Yaowu Xu's avatar
Yaowu Xu committed
631
#if CONFIG_AOM_HIGHBITDEPTH
632
633
634
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
635
    return aom_highbd_avg_8x8(src, stride);
636
  } else {
Yaowu Xu's avatar
Yaowu Xu committed
637
    return aom_avg_8x8(src, stride);
Jingning Han's avatar
Jingning Han committed
638
639
  }
}
640
641
#else
static INLINE int avg_8x8(const uint8_t *const src, const int stride) {
Yaowu Xu's avatar
Yaowu Xu committed
642
  return aom_avg_8x8(src, stride);
643
644
}
#endif
Jingning Han's avatar
Jingning Han committed
645

646
static void init_variance_tree(VAR_TREE *const vt,
Yaowu Xu's avatar
Yaowu Xu committed
647
#if CONFIG_AOM_HIGHBITDEPTH
648
                               const int highbd,
Jingning Han's avatar
Jingning Han committed
649
#endif
650
                               BLOCK_SIZE bsize, BLOCK_SIZE leaf_size,
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
                               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
668
#if CONFIG_AOM_HIGHBITDEPTH
669
  vt->highbd = highbd;
Yaowu Xu's avatar
Yaowu Xu committed
670
#endif  // CONFIG_AOM_HIGHBITDEPTH
671
672
673
674
675
676

  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
677
#if CONFIG_AOM_HIGHBITDEPTH
678
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
679
680
681
#endif  // CONFIG_AOM_HIGHBITDEPTH
                       subsize, leaf_size, AOMMIN(px, width),
                       AOMMIN(px, height), src, src_stride, ref, ref_stride);
682
    init_variance_tree(vt->split[1],
Yaowu Xu's avatar
Yaowu Xu committed
683
#if CONFIG_AOM_HIGHBITDEPTH
684
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
685
686
#endif  // CONFIG_AOM_HIGHBITDEPTH
                       subsize, leaf_size, width - px, AOMMIN(px, height),
687
                       src + px, src_stride, ref + px, ref_stride);
688
    init_variance_tree(vt->split[2],
Yaowu Xu's avatar
Yaowu Xu committed
689
#if CONFIG_AOM_HIGHBITDEPTH
690
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
691
692
#endif  // CONFIG_AOM_HIGHBITDEPTH
                       subsize, leaf_size, AOMMIN(px, width), height - px,
693
694
                       src + px * src_stride, src_stride, ref + px * ref_stride,
                       ref_stride);
695
    init_variance_tree(vt->split[3],
Yaowu Xu's avatar
Yaowu Xu committed
696
#if CONFIG_AOM_HIGHBITDEPTH
697
                       highbd,
Yaowu Xu's avatar
Yaowu Xu committed
698
#endif  // CONFIG_AOM_HIGHBITDEPTH
699
                       subsize, leaf_size, width - px, height - px,
700
701
702
703
704
705
706
                       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.
707
static void fill_variance_tree(VAR_TREE *const vt, const BLOCK_SIZE leaf_size) {
708
709
710
711
712
713
714
715
716
  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
717
718
    unsigned int sse = 0;
    int sum = 0;
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
747
748
    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);

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

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

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

  return vt->force_split;
Jingning Han's avatar
Jingning Han committed
807
808
809
}

// This function chooses partitioning based on the variance between source and
810
811
// reconstructed last (or golden), where variance is computed for down-sampled
// inputs.
Yaowu Xu's avatar
Yaowu Xu committed
812
static void choose_partitioning(AV1_COMP *const cpi, ThreadData *const td,
813
                                const TileInfo *const tile, MACROBLOCK *const x,
814
                                const int mi_row, const int mi_col) {
Yaowu Xu's avatar
Yaowu Xu committed
815
  AV1_COMMON *const cm = &cpi->common;
816
817
818
819
820
821
822
  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;
823
824
  int pixels_wide = 8 * num_8x8_blocks_wide_lookup[cm->sb_size];
  int pixels_high = 8 * num_8x8_blocks_high_lookup[cm->sb_size];
825
  int64_t thresholds[5] = {
826
827
    cpi->vbp_thresholds[0], cpi->vbp_thresholds[1], cpi->vbp_thresholds[2],
    cpi->vbp_thresholds[3], cpi->vbp_thresholds[4],
828
  };
829
830
  BLOCK_SIZE bsize_min[5] = { BLOCK_16X16, BLOCK_16X16, BLOCK_16X16,
                              cpi->vbp_bsize_min, BLOCK_8X8 };
831
832
833
  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
834
835
836
837
838

  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;
839

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

    if (cyclic_refresh_segment_id_boosted(segment_id)) {
Yaowu Xu's avatar
Yaowu Xu committed
846
      int q = av1_get_qindex(&cm->seg, segment_id, cm->base_qindex);
847
      assert(q == xd->qindex[segment_id]);
Jingning Han's avatar
Jingning Han committed
848
849
850
851
      set_vbp_thresholds(cpi, thresholds, q);
    }
  }

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

854
855
  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
856

857
858
  src = x->plane[0].src.buf;
  src_stride = x->plane[0].src.stride;
Jingning Han's avatar
Jingning Han committed
859

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

867
868
869
    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;
870
871
    BLOCK_SIZE bsize;

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

    assert(yv12 != NULL);

    if (yv12_g && yv12_g != yv12) {
Yaowu Xu's avatar
Yaowu Xu committed
884
885
      av1_setup_pre_planes(xd, 0, yv12_g, mi_row, mi_col,
                           &cm->frame_refs[GOLDEN_FRAME - 1].sf);
886
887
888
      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
889
890
891
892
    } else {
      y_sad_g = UINT_MAX;
    }

Yaowu Xu's avatar
Yaowu Xu committed
893
894
    av1_setup_pre_planes(xd, 0, yv12, mi_row, mi_col,
                         &cm->frame_refs[LAST_FRAME - 1].sf);
Jingning Han's avatar
Jingning Han committed
895
896
    mbmi->ref_frame[0] = LAST_FRAME;
    mbmi->ref_frame[1] = NONE;
897
    mbmi->sb_type = cm->sb_size;
Jingning Han's avatar
Jingning Han committed
898
    mbmi->mv[0].as_int = 0;
899
#if CONFIG_DUAL_FILTER
900
    for (i = 0; i < 4; ++i) mbmi->interp_filter[i] = BILINEAR;
901
#else
Jingning Han's avatar
Jingning Han committed
902
    mbmi->interp_filter = BILINEAR;
903
#endif
Jingning Han's avatar
Jingning Han committed
904

Yaowu Xu's avatar
Yaowu Xu committed
905
    y_sad = av1_int_pro_motion_estimation(cpi, x, bsize, mi_row, mi_col);
906

Jingning Han's avatar
Jingning Han committed
907
    if (y_sad_g < y_sad) {
Yaowu Xu's avatar
Yaowu Xu committed
908
909
      av1_setup_pre_planes(xd, 0, yv12_g, mi_row, mi_col,
                           &cm->frame_refs[GOLDEN_FRAME - 1].sf);
Jingning Han's avatar
Jingning Han committed
910
911
912
913
914
915
916
      mbmi->ref_frame[0] = GOLDEN_FRAME;
      mbmi->mv[0].as_int = 0;
      y_sad = y_sad_g;
    } else {
      x->pred_mv[LAST_FRAME] = mbmi->mv[0].as_mv;
    }

Yaowu Xu's avatar
Yaowu Xu committed
917
    av1_build_inter_predictors_sb(xd, mi_row, mi_col, cm->sb_size);
Jingning Han's avatar
Jingning Han committed
918

919
    for (i = 1; i < MAX_MB_PLANE; ++i) {
920
      struct macroblock_plane *p = &x->plane[i];
Jingning Han's avatar
Jingning Han committed
921
922
923
924
925
926
      struct macroblockd_plane *pd = &xd->plane[i];
      const BLOCK_SIZE bs = get_plane_block_size(bsize, pd);

      if (bs == BLOCK_INVALID)
        uv_sad = UINT_MAX;
      else
927
928
        uv_sad = cpi->fn_ptr[bs].sdf(p->src.buf, p->src.stride, pd->dst.buf,
                                     pd->dst.stride);
Jingning Han's avatar
Jingning Han committed
929
930
931
932

      x->color_sensitivity[i - 1] = uv_sad > (y_sad >> 2);
    }

933
934
    ref = xd->plane[0].dst.buf;
    ref_stride = xd->plane[0].dst.stride;
Jingning Han's avatar
Jingning Han committed
935

936
937
938
    // If the y_sad is very small, take the largest partition and exit.
    // Don't check on boosted segment for now, as largest is suppressed there.
    if (segment_id == CR_SEGMENT_ID_BASE && y_sad < cpi->vbp_threshold_sad) {
939
      if (!split_vert && !split_horz) {
940
        set_block_size(cpi, x, xd, mi_row, mi_col,<