diff options
Diffstat (limited to 'third_party/aom/av1/encoder/encodeframe.c')
| -rw-r--r-- | third_party/aom/av1/encoder/encodeframe.c | 7160 |
1 files changed, 7160 insertions, 0 deletions
diff --git a/third_party/aom/av1/encoder/encodeframe.c b/third_party/aom/av1/encoder/encodeframe.c new file mode 100644 index 0000000000..d254157e72 --- /dev/null +++ b/third_party/aom/av1/encoder/encodeframe.c @@ -0,0 +1,7160 @@ +/* + * Copyright (c) 2016, Alliance for Open Media. All rights reserved + * + * This source code is subject to the terms of the BSD 2 Clause License and + * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License + * was not distributed with this source code in the LICENSE file, you can + * obtain it at www.aomedia.org/license/software. If the Alliance for Open + * Media Patent License 1.0 was not distributed with this source code in the + * PATENTS file, you can obtain it at www.aomedia.org/license/patent. + */ + +#include <limits.h> +#include <math.h> +#include <stdio.h> + +#include "./av1_rtcd.h" +#include "./aom_dsp_rtcd.h" +#include "./aom_config.h" + +#include "aom_dsp/aom_dsp_common.h" +#include "aom_dsp/binary_codes_writer.h" +#include "aom_ports/mem.h" +#include "aom_ports/aom_timer.h" +#include "aom_ports/system_state.h" + +#include "av1/common/common.h" +#include "av1/common/entropy.h" +#include "av1/common/entropymode.h" +#include "av1/common/idct.h" +#include "av1/common/mv.h" +#include "av1/common/mvref_common.h" +#include "av1/common/pred_common.h" +#include "av1/common/quant_common.h" +#include "av1/common/reconintra.h" +#include "av1/common/reconinter.h" +#include "av1/common/seg_common.h" +#include "av1/common/tile_common.h" + +#include "av1/encoder/aq_complexity.h" +#include "av1/encoder/aq_cyclicrefresh.h" +#include "av1/encoder/aq_variance.h" +#if CONFIG_SUPERTX +#include "av1/encoder/cost.h" +#endif +#if CONFIG_GLOBAL_MOTION || CONFIG_WARPED_MOTION +#include "av1/common/warped_motion.h" +#endif // CONFIG_GLOBAL_MOTION || CONFIG_WARPED_MOTION +#if CONFIG_GLOBAL_MOTION +#include "av1/encoder/global_motion.h" +#endif // CONFIG_GLOBAL_MOTION +#include "av1/encoder/encodeframe.h" +#include "av1/encoder/encodemb.h" +#include "av1/encoder/encodemv.h" +#if CONFIG_LV_MAP +#include "av1/encoder/encodetxb.h" +#endif +#include "av1/encoder/ethread.h" +#include "av1/encoder/extend.h" +#include "av1/encoder/rd.h" +#include "av1/encoder/rdopt.h" +#include "av1/encoder/segmentation.h" +#include "av1/encoder/tokenize.h" +#if CONFIG_PVQ +#include "av1/common/pvq.h" +#include "av1/encoder/pvq_encoder.h" +#endif +#if CONFIG_HIGHBITDEPTH +#define IF_HBD(...) __VA_ARGS__ +#else +#define IF_HBD(...) +#endif // CONFIG_HIGHBITDEPTH + +static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td, + TOKENEXTRA **t, RUN_TYPE dry_run, int mi_row, + int mi_col, BLOCK_SIZE bsize, + PICK_MODE_CONTEXT *ctx, int *rate); + +#if CONFIG_SUPERTX +static int check_intra_b(PICK_MODE_CONTEXT *ctx); + +static int check_intra_sb(const AV1_COMP *cpi, const TileInfo *const tile, + int mi_row, int mi_col, BLOCK_SIZE bsize, + PC_TREE *pc_tree); +static void predict_superblock(const AV1_COMP *const cpi, ThreadData *td, +#if CONFIG_EXT_INTER + int mi_row_ori, int mi_col_ori, +#endif // CONFIG_EXT_INTER + int mi_row_pred, int mi_col_pred, + BLOCK_SIZE bsize_pred, int b_sub8x8, int block); +static int check_supertx_sb(BLOCK_SIZE bsize, TX_SIZE supertx_size, + PC_TREE *pc_tree); +static void predict_sb_complex(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, + int mi_col, int mi_row_ori, int mi_col_ori, + RUN_TYPE dry_run, BLOCK_SIZE bsize, + BLOCK_SIZE top_bsize, uint8_t *dst_buf[3], + int dst_stride[3], PC_TREE *pc_tree); +static void update_state_sb_supertx(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, + int mi_col, BLOCK_SIZE bsize, + RUN_TYPE dry_run, PC_TREE *pc_tree); +static void rd_supertx_sb(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, int mi_col, + BLOCK_SIZE bsize, int *tmp_rate, int64_t *tmp_dist, + TX_TYPE *best_tx, PC_TREE *pc_tree); +#endif // CONFIG_SUPERTX + +// This is used as a reference when computing the source variance for the +// purposes of activity masking. +// Eventually this should be replaced by custom no-reference routines, +// which will be faster. +static const uint8_t AV1_VAR_OFFS[MAX_SB_SIZE] = { + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, +#if CONFIG_EXT_PARTITION + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128 +#endif // CONFIG_EXT_PARTITION +}; + +#if CONFIG_HIGHBITDEPTH +static const uint16_t AV1_HIGH_VAR_OFFS_8[MAX_SB_SIZE] = { + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, +#if CONFIG_EXT_PARTITION + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, + 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128 +#endif // CONFIG_EXT_PARTITION +}; + +static const uint16_t AV1_HIGH_VAR_OFFS_10[MAX_SB_SIZE] = { + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, +#if CONFIG_EXT_PARTITION + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, + 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4, 128 * 4 +#endif // CONFIG_EXT_PARTITION +}; + +static const uint16_t AV1_HIGH_VAR_OFFS_12[MAX_SB_SIZE] = { + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, +#if CONFIG_EXT_PARTITION + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, 128 * 16, + 128 * 16 +#endif // CONFIG_EXT_PARTITION +}; +#endif // CONFIG_HIGHBITDEPTH + +unsigned int av1_get_sby_perpixel_variance(const AV1_COMP *cpi, + const struct buf_2d *ref, + BLOCK_SIZE bs) { + unsigned int sse; + const unsigned int var = + cpi->fn_ptr[bs].vf(ref->buf, ref->stride, AV1_VAR_OFFS, 0, &sse); + return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]); +} + +#if CONFIG_HIGHBITDEPTH +unsigned int av1_high_get_sby_perpixel_variance(const AV1_COMP *cpi, + const struct buf_2d *ref, + BLOCK_SIZE bs, int bd) { + unsigned int var, sse; + switch (bd) { + case 10: + var = + cpi->fn_ptr[bs].vf(ref->buf, ref->stride, + CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_10), 0, &sse); + break; + case 12: + var = + cpi->fn_ptr[bs].vf(ref->buf, ref->stride, + CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_12), 0, &sse); + break; + case 8: + default: + var = + cpi->fn_ptr[bs].vf(ref->buf, ref->stride, + CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8), 0, &sse); + break; + } + return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]); +} +#endif // CONFIG_HIGHBITDEPTH + +static unsigned int get_sby_perpixel_diff_variance(const AV1_COMP *const cpi, + const struct buf_2d *ref, + int mi_row, int mi_col, + BLOCK_SIZE bs) { + unsigned int sse, var; + uint8_t *last_y; + const YV12_BUFFER_CONFIG *last = get_ref_frame_buffer(cpi, LAST_FRAME); + + assert(last != NULL); + last_y = + &last->y_buffer[mi_row * MI_SIZE * last->y_stride + mi_col * MI_SIZE]; + var = cpi->fn_ptr[bs].vf(ref->buf, ref->stride, last_y, last->y_stride, &sse); + return ROUND_POWER_OF_TWO(var, num_pels_log2_lookup[bs]); +} + +static BLOCK_SIZE get_rd_var_based_fixed_partition(AV1_COMP *cpi, MACROBLOCK *x, + int mi_row, int mi_col) { + unsigned int var = get_sby_perpixel_diff_variance( + cpi, &x->plane[0].src, mi_row, mi_col, BLOCK_64X64); + if (var < 8) + return BLOCK_64X64; + else if (var < 128) + return BLOCK_32X32; + else if (var < 2048) + return BLOCK_16X16; + else + return BLOCK_8X8; +} + +// Lighter version of set_offsets that only sets the mode info +// pointers. +static void set_mode_info_offsets(const AV1_COMP *const cpi, + MACROBLOCK *const x, MACROBLOCKD *const xd, + int mi_row, int mi_col) { + const AV1_COMMON *const cm = &cpi->common; + const int idx_str = xd->mi_stride * mi_row + mi_col; + xd->mi = cm->mi_grid_visible + idx_str; + xd->mi[0] = cm->mi + idx_str; + x->mbmi_ext = cpi->mbmi_ext_base + (mi_row * cm->mi_cols + mi_col); +} + +static void set_offsets_without_segment_id(const AV1_COMP *const cpi, + const TileInfo *const tile, + MACROBLOCK *const x, int mi_row, + int mi_col, BLOCK_SIZE bsize) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + + set_skip_context(xd, mi_row, mi_col); + + set_mode_info_offsets(cpi, x, xd, mi_row, mi_col); + +#if CONFIG_VAR_TX + xd->above_txfm_context = cm->above_txfm_context + mi_col; + xd->left_txfm_context = + xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK); + xd->max_tx_size = max_txsize_lookup[bsize]; +#endif + + // Set up destination pointers. + av1_setup_dst_planes(xd->plane, bsize, get_frame_new_buffer(cm), mi_row, + mi_col); + + // Set up limit values for MV components. + // Mv beyond the range do not produce new/different prediction block. + x->mv_limits.row_min = + -(((mi_row + mi_height) * MI_SIZE) + AOM_INTERP_EXTEND); + x->mv_limits.col_min = -(((mi_col + mi_width) * MI_SIZE) + AOM_INTERP_EXTEND); + x->mv_limits.row_max = (cm->mi_rows - mi_row) * MI_SIZE + AOM_INTERP_EXTEND; + x->mv_limits.col_max = (cm->mi_cols - mi_col) * MI_SIZE + AOM_INTERP_EXTEND; + + set_plane_n4(xd, mi_width, mi_height); + + // Set up distance of MB to edge of frame in 1/8th pel units. + assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1))); + set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, +#if CONFIG_DEPENDENT_HORZTILES + cm->dependent_horz_tiles, +#endif // CONFIG_DEPENDENT_HORZTILES + cm->mi_rows, cm->mi_cols); + + // Set up source buffers. + av1_setup_src_planes(x, cpi->source, mi_row, mi_col); + + // R/D setup. + x->rddiv = cpi->rd.RDDIV; + x->rdmult = cpi->rd.RDMULT; + + // required by av1_append_sub8x8_mvs_for_idx() and av1_find_best_ref_mvs() + xd->tile = *tile; +} + +static void set_offsets(const AV1_COMP *const cpi, const TileInfo *const tile, + MACROBLOCK *const x, int mi_row, int mi_col, + BLOCK_SIZE bsize) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + MB_MODE_INFO *mbmi; + const struct segmentation *const seg = &cm->seg; + + set_offsets_without_segment_id(cpi, tile, x, mi_row, mi_col, bsize); + + mbmi = &xd->mi[0]->mbmi; + + // Setup segment ID. + if (seg->enabled) { + if (!cpi->vaq_refresh) { + const uint8_t *const map = + seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map; + mbmi->segment_id = get_segment_id(cm, map, bsize, mi_row, mi_col); + } + av1_init_plane_quantizers(cpi, x, mbmi->segment_id); + } else { + mbmi->segment_id = 0; + } + +#if CONFIG_SUPERTX + mbmi->segment_id_supertx = MAX_SEGMENTS; +#endif // CONFIG_SUPERTX +} + +#if CONFIG_SUPERTX +static void set_offsets_supertx(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, + int mi_col, BLOCK_SIZE bsize) { + MACROBLOCK *const x = &td->mb; + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; +#if CONFIG_DEPENDENT_HORZTILES + set_mode_info_offsets(cpi, x, xd, mi_row, mi_col, cm->dependent_horz_tiles); +#else + set_mode_info_offsets(cpi, x, xd, mi_row, mi_col); +#endif + + // Set up distance of MB to edge of frame in 1/8th pel units. + assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1))); + set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width, +#if CONFIG_DEPENDENT_HORZTILES + cm->dependent_horz_tiles, +#endif // CONFIG_DEPENDENT_HORZTILES + cm->mi_rows, cm->mi_cols); +} + +static void set_offsets_extend(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row_pred, + int mi_col_pred, int mi_row_ori, int mi_col_ori, + BLOCK_SIZE bsize_pred) { + // Used in supertx + // (mi_row_ori, mi_col_ori, bsize_ori): region for mv + // (mi_row_pred, mi_col_pred, bsize_pred): region to predict + MACROBLOCK *const x = &td->mb; + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + const int mi_width = mi_size_wide[bsize_pred]; + const int mi_height = mi_size_high[bsize_pred]; + +#if CONFIG_DEPENDENT_HORZTILES + set_mode_info_offsets(cpi, x, xd, mi_row_ori, mi_col_ori, + cm->dependent_horz_tiles); +#else + set_mode_info_offsets(cpi, x, xd, mi_row_ori, mi_col_ori); +#endif + + // Set up limit values for MV components. + // Mv beyond the range do not produce new/different prediction block. + x->mv_limits.row_min = + -(((mi_row_pred + mi_height) * MI_SIZE) + AOM_INTERP_EXTEND); + x->mv_limits.col_min = + -(((mi_col_pred + mi_width) * MI_SIZE) + AOM_INTERP_EXTEND); + x->mv_limits.row_max = + (cm->mi_rows - mi_row_pred) * MI_SIZE + AOM_INTERP_EXTEND; + x->mv_limits.col_max = + (cm->mi_cols - mi_col_pred) * MI_SIZE + AOM_INTERP_EXTEND; + +// Set up distance of MB to edge of frame in 1/8th pel units. +#if !CONFIG_CB4X4 + assert(!(mi_col_pred & (mi_width - mi_size_wide[BLOCK_8X8])) && + !(mi_row_pred & (mi_height - mi_size_high[BLOCK_8X8]))); +#endif + set_mi_row_col(xd, tile, mi_row_pred, mi_height, mi_col_pred, mi_width, +#if CONFIG_DEPENDENT_HORZTILES + cm->dependent_horz_tiles, +#endif // CONFIG_DEPENDENT_HORZTILES + cm->mi_rows, cm->mi_cols); + xd->up_available = (mi_row_ori > tile->mi_row_start); + xd->left_available = (mi_col_ori > tile->mi_col_start); + + // R/D setup. + x->rddiv = cpi->rd.RDDIV; + x->rdmult = cpi->rd.RDMULT; +} + +static void set_segment_id_supertx(const AV1_COMP *const cpi, + MACROBLOCK *const x, const int mi_row, + const int mi_col, const BLOCK_SIZE bsize) { + const AV1_COMMON *cm = &cpi->common; + const struct segmentation *seg = &cm->seg; + const int miw = AOMMIN(mi_size_wide[bsize], cm->mi_cols - mi_col); + const int mih = AOMMIN(mi_size_high[bsize], cm->mi_rows - mi_row); + const int mi_offset = mi_row * cm->mi_stride + mi_col; + MODE_INFO **const mip = cm->mi_grid_visible + mi_offset; + int r, c; + int seg_id_supertx = MAX_SEGMENTS; + + if (!seg->enabled) { + seg_id_supertx = 0; + } else { + // Find the minimum segment_id + for (r = 0; r < mih; r++) + for (c = 0; c < miw; c++) + seg_id_supertx = + AOMMIN(mip[r * cm->mi_stride + c]->mbmi.segment_id, seg_id_supertx); + assert(0 <= seg_id_supertx && seg_id_supertx < MAX_SEGMENTS); + + // Initialize plane quantisers + av1_init_plane_quantizers(cpi, x, seg_id_supertx); + } + + // Assign the the segment_id back to segment_id_supertx + for (r = 0; r < mih; r++) + for (c = 0; c < miw; c++) + mip[r * cm->mi_stride + c]->mbmi.segment_id_supertx = seg_id_supertx; +} +#endif // CONFIG_SUPERTX + +static void set_block_size(AV1_COMP *const cpi, MACROBLOCK *const x, + MACROBLOCKD *const xd, int mi_row, int mi_col, + BLOCK_SIZE bsize) { + if (cpi->common.mi_cols > mi_col && cpi->common.mi_rows > mi_row) { + const int mi_width = AOMMAX(mi_size_wide[bsize], mi_size_wide[BLOCK_8X8]); + const int mi_height = AOMMAX(mi_size_high[bsize], mi_size_high[BLOCK_8X8]); + for (int r = 0; r < mi_height; ++r) { + for (int c = 0; c < mi_width; ++c) { + set_mode_info_offsets(cpi, x, xd, mi_row + r, mi_col + c); + xd->mi[0]->mbmi.sb_type = bsize; + } + } + } +} + +static void set_vt_partitioning(AV1_COMP *cpi, MACROBLOCK *const x, + MACROBLOCKD *const xd, VAR_TREE *vt, int mi_row, + int mi_col, const int64_t *const threshold, + const BLOCK_SIZE *const bsize_min) { + AV1_COMMON *const cm = &cpi->common; + const int hbw = mi_size_wide[vt->bsize] / 2; + const int hbh = mi_size_high[vt->bsize] / 2; + const int has_cols = mi_col + hbw < cm->mi_cols; + const int has_rows = mi_row + hbh < cm->mi_rows; + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + + assert(vt->bsize >= BLOCK_8X8); + + assert(hbh == hbw); + + if (vt->bsize == BLOCK_8X8 && cm->frame_type != KEY_FRAME) { + set_block_size(cpi, x, xd, mi_row, mi_col, BLOCK_8X8); + return; + } + + if (vt->force_split || (!has_cols && !has_rows)) goto split; + + // For bsize=bsize_min (16x16/8x8 for 8x8/4x4 downsampling), select if + // variance is below threshold, otherwise split will be selected. + // No check for vert/horiz split as too few samples for variance. + if (vt->bsize == bsize_min[0]) { + if (has_cols && has_rows && vt->variances.none.variance < threshold[0]) { + set_block_size(cpi, x, xd, mi_row, mi_col, vt->bsize); + return; + } else { + BLOCK_SIZE subsize = get_subsize(vt->bsize, PARTITION_SPLIT); + set_block_size(cpi, x, xd, mi_row, mi_col, subsize); + if (vt->bsize > BLOCK_8X8) { + set_block_size(cpi, x, xd, mi_row, mi_col + hbw, subsize); + set_block_size(cpi, x, xd, mi_row + hbh, mi_col, subsize); + set_block_size(cpi, x, xd, mi_row + hbh, mi_col + hbw, subsize); + } + return; + } + } else if (vt->bsize > bsize_min[0]) { + // For key frame: take split for bsize above 32X32 or very high variance. + if (cm->frame_type == KEY_FRAME && + (vt->bsize > BLOCK_32X32 || + vt->variances.none.variance > (threshold[0] << 4))) { + goto split; + } + // If variance is low, take the bsize (no split). + if (has_cols && has_rows && vt->variances.none.variance < threshold[0]) { + set_block_size(cpi, x, xd, mi_row, mi_col, vt->bsize); + return; + } + + // Check vertical split. + if (has_rows) { + BLOCK_SIZE subsize = get_subsize(vt->bsize, PARTITION_VERT); + if (vt->variances.vert[0].variance < threshold[0] && + vt->variances.vert[1].variance < threshold[0] && + get_plane_block_size(subsize, &xd->plane[1]) < BLOCK_INVALID) { + set_block_size(cpi, x, xd, mi_row, mi_col, subsize); + set_block_size(cpi, x, xd, mi_row, mi_col + hbw, subsize); + return; + } + } + // Check horizontal split. + if (has_cols) { + BLOCK_SIZE subsize = get_subsize(vt->bsize, PARTITION_HORZ); + if (vt->variances.horz[0].variance < threshold[0] && + vt->variances.horz[1].variance < threshold[0] && + get_plane_block_size(subsize, &xd->plane[1]) < BLOCK_INVALID) { + set_block_size(cpi, x, xd, mi_row, mi_col, subsize); + set_block_size(cpi, x, xd, mi_row + hbh, mi_col, subsize); + return; + } + } + } + +split : { + set_vt_partitioning(cpi, x, xd, vt->split[0], mi_row, mi_col, threshold + 1, + bsize_min + 1); + set_vt_partitioning(cpi, x, xd, vt->split[1], mi_row, mi_col + hbw, + threshold + 1, bsize_min + 1); + set_vt_partitioning(cpi, x, xd, vt->split[2], mi_row + hbh, mi_col, + threshold + 1, bsize_min + 1); + set_vt_partitioning(cpi, x, xd, vt->split[3], mi_row + hbh, mi_col + hbw, + threshold + 1, bsize_min + 1); + return; +} +} + +// Set the variance split thresholds for following the block sizes: +// 0 - threshold_64x64, 1 - threshold_32x32, 2 - threshold_16x16, +// 3 - vbp_threshold_8x8. vbp_threshold_8x8 (to split to 4x4 partition) is +// currently only used on key frame. +static void set_vbp_thresholds(AV1_COMP *cpi, int64_t thresholds[], int q) { + AV1_COMMON *const cm = &cpi->common; + const int is_key_frame = (cm->frame_type == KEY_FRAME); + const int threshold_multiplier = is_key_frame ? 20 : 1; + const int64_t threshold_base = + (int64_t)(threshold_multiplier * cpi->y_dequant[q][1]); + if (is_key_frame) { + thresholds[1] = threshold_base; + thresholds[2] = threshold_base >> 2; + thresholds[3] = threshold_base >> 2; + thresholds[4] = threshold_base << 2; + } else { + thresholds[2] = threshold_base; + if (cm->width <= 352 && cm->height <= 288) { + thresholds[1] = threshold_base >> 2; + thresholds[3] = threshold_base << 3; + } else { + thresholds[1] = threshold_base; + thresholds[2] = (5 * threshold_base) >> 2; + if (cm->width >= 1920 && cm->height >= 1080) + thresholds[2] = (7 * threshold_base) >> 2; + thresholds[3] = threshold_base << cpi->oxcf.speed; + } + } + thresholds[0] = INT64_MIN; +} + +void av1_set_variance_partition_thresholds(AV1_COMP *cpi, int q) { + AV1_COMMON *const cm = &cpi->common; + SPEED_FEATURES *const sf = &cpi->sf; + const int is_key_frame = (cm->frame_type == KEY_FRAME); + if (sf->partition_search_type != VAR_BASED_PARTITION && + sf->partition_search_type != REFERENCE_PARTITION) { + return; + } else { + set_vbp_thresholds(cpi, cpi->vbp_thresholds, q); + // The thresholds below are not changed locally. + if (is_key_frame) { + cpi->vbp_threshold_sad = 0; + cpi->vbp_bsize_min = BLOCK_8X8; + } else { + if (cm->width <= 352 && cm->height <= 288) + cpi->vbp_threshold_sad = 100; + else + cpi->vbp_threshold_sad = (cpi->y_dequant[q][1] << 1) > 1000 + ? (cpi->y_dequant[q][1] << 1) + : 1000; + cpi->vbp_bsize_min = BLOCK_16X16; + } + cpi->vbp_threshold_minmax = 15 + (q >> 3); + } +} + +// Compute the minmax over the 8x8 subblocks. +static int compute_minmax_8x8(const uint8_t *src, int src_stride, + const uint8_t *ref, int ref_stride, +#if CONFIG_HIGHBITDEPTH + int highbd, +#endif + int pixels_wide, int pixels_high) { + int k; + int minmax_max = 0; + int minmax_min = 255; + // Loop over the 4 8x8 subblocks. + for (k = 0; k < 4; k++) { + const int x8_idx = ((k & 1) << 3); + const int y8_idx = ((k >> 1) << 3); + int min = 0; + int max = 0; + if (x8_idx < pixels_wide && y8_idx < pixels_high) { + const int src_offset = y8_idx * src_stride + x8_idx; + const int ref_offset = y8_idx * ref_stride + x8_idx; +#if CONFIG_HIGHBITDEPTH + if (highbd) { + aom_highbd_minmax_8x8(src + src_offset, src_stride, ref + ref_offset, + ref_stride, &min, &max); + } else { + aom_minmax_8x8(src + src_offset, src_stride, ref + ref_offset, + ref_stride, &min, &max); + } +#else + aom_minmax_8x8(src + src_offset, src_stride, ref + ref_offset, ref_stride, + &min, &max); +#endif + if ((max - min) > minmax_max) minmax_max = (max - min); + if ((max - min) < minmax_min) minmax_min = (max - min); + } + } + return (minmax_max - minmax_min); +} + +#if CONFIG_HIGHBITDEPTH +static INLINE int avg_4x4(const uint8_t *const src, const int stride, + const int highbd) { + if (highbd) { + return aom_highbd_avg_4x4(src, stride); + } else { + return aom_avg_4x4(src, stride); + } +} +#else +static INLINE int avg_4x4(const uint8_t *const src, const int stride) { + return aom_avg_4x4(src, stride); +} +#endif + +#if CONFIG_HIGHBITDEPTH +static INLINE int avg_8x8(const uint8_t *const src, const int stride, + const int highbd) { + if (highbd) { + return aom_highbd_avg_8x8(src, stride); + } else { + return aom_avg_8x8(src, stride); + } +} +#else +static INLINE int avg_8x8(const uint8_t *const src, const int stride) { + return aom_avg_8x8(src, stride); +} +#endif + +static void init_variance_tree(VAR_TREE *const vt, +#if CONFIG_HIGHBITDEPTH + const int highbd, +#endif + BLOCK_SIZE bsize, BLOCK_SIZE leaf_size, + const int width, const int height, + const uint8_t *const src, const int src_stride, + const uint8_t *const ref, const int ref_stride) { + assert(bsize >= leaf_size); + + vt->bsize = bsize; + + vt->force_split = 0; + + vt->src = src; + vt->src_stride = src_stride; + vt->ref = ref; + vt->ref_stride = ref_stride; + + vt->width = width; + vt->height = height; + +#if CONFIG_HIGHBITDEPTH + vt->highbd = highbd; +#endif // CONFIG_HIGHBITDEPTH + + if (bsize > leaf_size) { + const BLOCK_SIZE subsize = get_subsize(bsize, PARTITION_SPLIT); + const int px = block_size_wide[subsize]; + + init_variance_tree(vt->split[0], +#if CONFIG_HIGHBITDEPTH + highbd, +#endif // CONFIG_HIGHBITDEPTH + subsize, leaf_size, AOMMIN(px, width), + AOMMIN(px, height), src, src_stride, ref, ref_stride); + init_variance_tree(vt->split[1], +#if CONFIG_HIGHBITDEPTH + highbd, +#endif // CONFIG_HIGHBITDEPTH + subsize, leaf_size, width - px, AOMMIN(px, height), + src + px, src_stride, ref + px, ref_stride); + init_variance_tree(vt->split[2], +#if CONFIG_HIGHBITDEPTH + highbd, +#endif // CONFIG_HIGHBITDEPTH + subsize, leaf_size, AOMMIN(px, width), height - px, + src + px * src_stride, src_stride, ref + px * ref_stride, + ref_stride); + init_variance_tree(vt->split[3], +#if CONFIG_HIGHBITDEPTH + highbd, +#endif // CONFIG_HIGHBITDEPTH + subsize, leaf_size, width - px, height - px, + src + px * src_stride + px, src_stride, + ref + px * ref_stride + px, ref_stride); + } +} + +// Fill the variance tree based on averaging pixel values (sub-sampling), at +// the leaf node size. +static void fill_variance_tree(VAR_TREE *const vt, const BLOCK_SIZE leaf_size) { + if (vt->bsize > leaf_size) { + fill_variance_tree(vt->split[0], leaf_size); + fill_variance_tree(vt->split[1], leaf_size); + fill_variance_tree(vt->split[2], leaf_size); + fill_variance_tree(vt->split[3], leaf_size); + fill_variance_node(vt); + } else if (vt->width <= 0 || vt->height <= 0) { + fill_variance(0, 0, 0, &vt->variances.none); + } else { + unsigned int sse = 0; + int sum = 0; + int src_avg; + int ref_avg; + assert(leaf_size == BLOCK_4X4 || leaf_size == BLOCK_8X8); + if (leaf_size == BLOCK_4X4) { + src_avg = avg_4x4(vt->src, vt->src_stride IF_HBD(, vt->highbd)); + ref_avg = avg_4x4(vt->ref, vt->ref_stride IF_HBD(, vt->highbd)); + } else { + src_avg = avg_8x8(vt->src, vt->src_stride IF_HBD(, vt->highbd)); + ref_avg = avg_8x8(vt->ref, vt->ref_stride IF_HBD(, vt->highbd)); + } + sum = src_avg - ref_avg; + sse = sum * sum; + fill_variance(sse, sum, 0, &vt->variances.none); + } +} + +static void refine_variance_tree(VAR_TREE *const vt, const int64_t threshold) { + if (vt->bsize >= BLOCK_8X8) { + if (vt->bsize == BLOCK_16X16) { + if (vt->variances.none.variance <= threshold) + return; + else + vt->force_split = 0; + } + + refine_variance_tree(vt->split[0], threshold); + refine_variance_tree(vt->split[1], threshold); + refine_variance_tree(vt->split[2], threshold); + refine_variance_tree(vt->split[3], threshold); + + if (vt->bsize <= BLOCK_16X16) fill_variance_node(vt); + } else if (vt->width <= 0 || vt->height <= 0) { + fill_variance(0, 0, 0, &vt->variances.none); + } else { + const int src_avg = avg_4x4(vt->src, vt->src_stride IF_HBD(, vt->highbd)); + const int ref_avg = avg_4x4(vt->ref, vt->ref_stride IF_HBD(, vt->highbd)); + const int sum = src_avg - ref_avg; + const unsigned int sse = sum * sum; + assert(vt->bsize == BLOCK_4X4); + fill_variance(sse, sum, 0, &vt->variances.none); + } +} + +static int check_split_key_frame(VAR_TREE *const vt, const int64_t threshold) { + if (vt->bsize == BLOCK_32X32) { + vt->force_split = vt->variances.none.variance > threshold; + } else { + vt->force_split |= check_split_key_frame(vt->split[0], threshold); + vt->force_split |= check_split_key_frame(vt->split[1], threshold); + vt->force_split |= check_split_key_frame(vt->split[2], threshold); + vt->force_split |= check_split_key_frame(vt->split[3], threshold); + } + return vt->force_split; +} + +static int check_split(AV1_COMP *const cpi, VAR_TREE *const vt, + const int segment_id, const int64_t *const thresholds) { + if (vt->bsize == BLOCK_16X16) { + vt->force_split = vt->variances.none.variance > thresholds[0]; + if (!vt->force_split && vt->variances.none.variance > thresholds[-1] && + !cyclic_refresh_segment_id_boosted(segment_id)) { + // We have some nominal amount of 16x16 variance (based on average), + // compute the minmax over the 8x8 sub-blocks, and if above threshold, + // force split to 8x8 block for this 16x16 block. + int minmax = + compute_minmax_8x8(vt->src, vt->src_stride, vt->ref, vt->ref_stride, +#if CONFIG_HIGHBITDEPTH + vt->highbd, +#endif + vt->width, vt->height); + vt->force_split = minmax > cpi->vbp_threshold_minmax; + } + } else { + vt->force_split |= + check_split(cpi, vt->split[0], segment_id, thresholds + 1); + vt->force_split |= + check_split(cpi, vt->split[1], segment_id, thresholds + 1); + vt->force_split |= + check_split(cpi, vt->split[2], segment_id, thresholds + 1); + vt->force_split |= + check_split(cpi, vt->split[3], segment_id, thresholds + 1); + + if (vt->bsize == BLOCK_32X32 && !vt->force_split) { + vt->force_split = vt->variances.none.variance > thresholds[0]; + } + } + + return vt->force_split; +} + +// This function chooses partitioning based on the variance between source and +// reconstructed last (or golden), where variance is computed for down-sampled +// inputs. +static void choose_partitioning(AV1_COMP *const cpi, ThreadData *const td, + const TileInfo *const tile, MACROBLOCK *const x, + const int mi_row, const int mi_col) { + AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + VAR_TREE *const vt = td->var_root[cm->mib_size_log2 - MIN_MIB_SIZE_LOG2]; +#if CONFIG_DUAL_FILTER + int i; +#endif + const uint8_t *src; + const uint8_t *ref; + int src_stride; + int ref_stride; + int pixels_wide = MI_SIZE * mi_size_wide[cm->sb_size]; + int pixels_high = MI_SIZE * mi_size_high[cm->sb_size]; + int64_t thresholds[5] = { + cpi->vbp_thresholds[0], cpi->vbp_thresholds[1], cpi->vbp_thresholds[2], + cpi->vbp_thresholds[3], cpi->vbp_thresholds[4], + }; + BLOCK_SIZE bsize_min[5] = { BLOCK_16X16, BLOCK_16X16, BLOCK_16X16, + cpi->vbp_bsize_min, BLOCK_8X8 }; + 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; + + 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; + + if (cpi->oxcf.aq_mode == CYCLIC_REFRESH_AQ && cm->seg.enabled) { + const uint8_t *const map = + cm->seg.update_map ? cpi->segmentation_map : cm->last_frame_seg_map; + segment_id = get_segment_id(cm, map, cm->sb_size, mi_row, mi_col); + + if (cyclic_refresh_segment_id_boosted(segment_id)) { + int q = av1_get_qindex(&cm->seg, segment_id, cm->base_qindex); + set_vbp_thresholds(cpi, thresholds, q); + } + } + + set_offsets(cpi, tile, x, mi_row, mi_col, cm->sb_size); + + 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); + + src = x->plane[0].src.buf; + src_stride = x->plane[0].src.stride; + + if (!is_key_frame) { + MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi; + const YV12_BUFFER_CONFIG *yv12 = get_ref_frame_buffer(cpi, LAST_FRAME); + const YV12_BUFFER_CONFIG *yv12_g = get_ref_frame_buffer(cpi, GOLDEN_FRAME); + unsigned int y_sad, y_sad_g; + + 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; + BLOCK_SIZE bsize; + + if (split_vert && split_horz) + bsize = get_subsize(cm->sb_size, PARTITION_SPLIT); + else if (split_vert) + bsize = get_subsize(cm->sb_size, PARTITION_VERT); + else if (split_horz) + bsize = get_subsize(cm->sb_size, PARTITION_HORZ); + else + bsize = cm->sb_size; + + assert(yv12 != NULL); + + if (yv12_g && yv12_g != yv12) { + av1_setup_pre_planes(xd, 0, yv12_g, mi_row, mi_col, + &cm->frame_refs[GOLDEN_FRAME - 1].sf); + 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); + } else { + y_sad_g = UINT_MAX; + } + + av1_setup_pre_planes(xd, 0, yv12, mi_row, mi_col, + &cm->frame_refs[LAST_FRAME - 1].sf); + mbmi->ref_frame[0] = LAST_FRAME; + mbmi->ref_frame[1] = NONE_FRAME; + mbmi->sb_type = cm->sb_size; + mbmi->mv[0].as_int = 0; +#if CONFIG_DUAL_FILTER + for (i = 0; i < 4; ++i) mbmi->interp_filter[i] = BILINEAR; +#else + mbmi->interp_filter = BILINEAR; +#endif + + y_sad = av1_int_pro_motion_estimation(cpi, x, bsize, mi_row, mi_col); + + if (y_sad_g < y_sad) { + av1_setup_pre_planes(xd, 0, yv12_g, mi_row, mi_col, + &cm->frame_refs[GOLDEN_FRAME - 1].sf); + 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; + } + + av1_build_inter_predictors_sb(xd, mi_row, mi_col, NULL, cm->sb_size); + + ref = xd->plane[0].dst.buf; + ref_stride = xd->plane[0].dst.stride; + + // 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) { + if (!split_vert && !split_horz) { + set_block_size(cpi, x, xd, mi_row, mi_col, cm->sb_size); + return; + } + } + } else { + ref = AV1_VAR_OFFS; + ref_stride = 0; +#if CONFIG_HIGHBITDEPTH + if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) { + switch (xd->bd) { + case 10: ref = CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_10); break; + case 12: ref = CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_12); break; + case 8: + default: ref = CONVERT_TO_BYTEPTR(AV1_HIGH_VAR_OFFS_8); break; + } + } +#endif // CONFIG_HIGHBITDEPTH + } + + init_variance_tree( + vt, +#if CONFIG_HIGHBITDEPTH + xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH, +#endif // CONFIG_HIGHBITDEPTH + cm->sb_size, (is_key_frame || low_res) ? BLOCK_4X4 : BLOCK_8X8, + pixels_wide, pixels_high, src, src_stride, ref, ref_stride); + + // Fill in the entire tree of variances and compute splits. + if (is_key_frame) { + fill_variance_tree(vt, BLOCK_4X4); + check_split_key_frame(vt, thre[1]); + } else { + fill_variance_tree(vt, BLOCK_8X8); + check_split(cpi, vt, segment_id, thre); + if (low_res) { + refine_variance_tree(vt, thre[1] << 1); + } + } + + vt->force_split |= mi_col + cm->mib_size > cm->mi_cols || + mi_row + cm->mib_size > cm->mi_rows; + + // Now go through the entire structure, splitting every block size until + // we get to one that's got a variance lower than our threshold. + set_vt_partitioning(cpi, x, xd, vt, mi_row, mi_col, thre, bmin); +} + +#if CONFIG_DUAL_FILTER +static void reset_intmv_filter_type(const AV1_COMMON *const cm, MACROBLOCKD *xd, + MB_MODE_INFO *mbmi) { + int dir; + for (dir = 0; dir < 2; ++dir) { + if (!has_subpel_mv_component(xd->mi[0], xd, dir) && + (mbmi->ref_frame[1] == NONE_FRAME || + !has_subpel_mv_component(xd->mi[0], xd, dir + 2))) + mbmi->interp_filter[dir] = (cm->interp_filter == SWITCHABLE) + ? EIGHTTAP_REGULAR + : cm->interp_filter; + mbmi->interp_filter[dir + 2] = mbmi->interp_filter[dir]; + } +} + +static void update_filter_type_count(FRAME_COUNTS *counts, + const MACROBLOCKD *xd, + const MB_MODE_INFO *mbmi) { + int dir; + for (dir = 0; dir < 2; ++dir) { + if (has_subpel_mv_component(xd->mi[0], xd, dir) || + (mbmi->ref_frame[1] > INTRA_FRAME && + has_subpel_mv_component(xd->mi[0], xd, dir + 2))) { + const int ctx = av1_get_pred_context_switchable_interp(xd, dir); + ++counts->switchable_interp[ctx][mbmi->interp_filter[dir]]; + } + } +} +#endif +#if CONFIG_GLOBAL_MOTION +static void update_global_motion_used(PREDICTION_MODE mode, BLOCK_SIZE bsize, + const MB_MODE_INFO *mbmi, + RD_COUNTS *rdc) { + if (mode == ZEROMV +#if CONFIG_EXT_INTER + || mode == ZERO_ZEROMV +#endif + ) { + const int num_4x4s = + num_4x4_blocks_wide_lookup[bsize] * num_4x4_blocks_high_lookup[bsize]; + int ref; + for (ref = 0; ref < 1 + has_second_ref(mbmi); ++ref) { + rdc->global_motion_used[mbmi->ref_frame[ref]] += num_4x4s; + } + } +} +#endif // CONFIG_GLOBAL_MOTION + +static void reset_tx_size(MACROBLOCKD *xd, MB_MODE_INFO *mbmi, + const TX_MODE tx_mode) { + if (xd->lossless[mbmi->segment_id]) { + mbmi->tx_size = TX_4X4; + } else if (tx_mode != TX_MODE_SELECT) { + mbmi->tx_size = + tx_size_from_tx_mode(mbmi->sb_type, tx_mode, is_inter_block(mbmi)); + } +} + +#if CONFIG_REF_MV +static void set_ref_and_pred_mvs(MACROBLOCK *const x, int_mv *const mi_pred_mv, + int8_t rf_type) { + MACROBLOCKD *const xd = &x->e_mbd; + MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi; + + const int bw = xd->n8_w << MI_SIZE_LOG2; + const int bh = xd->n8_h << MI_SIZE_LOG2; + int ref_mv_idx = mbmi->ref_mv_idx; + MB_MODE_INFO_EXT *const mbmi_ext = x->mbmi_ext; + CANDIDATE_MV *const curr_ref_mv_stack = mbmi_ext->ref_mv_stack[rf_type]; + +#if CONFIG_EXT_INTER + if (has_second_ref(mbmi)) { + // Special case: NEAR_NEWMV and NEW_NEARMV modes use 1 + mbmi->ref_mv_idx + // (like NEARMV) instead + if (mbmi->mode == NEAR_NEWMV || mbmi->mode == NEW_NEARMV) ref_mv_idx += 1; + + if (compound_ref0_mode(mbmi->mode) == NEWMV) { + int_mv this_mv = curr_ref_mv_stack[ref_mv_idx].this_mv; + clamp_mv_ref(&this_mv.as_mv, bw, bh, xd); + mbmi_ext->ref_mvs[mbmi->ref_frame[0]][0] = this_mv; + mbmi->pred_mv[0] = this_mv; + mi_pred_mv[0] = this_mv; + } + if (compound_ref1_mode(mbmi->mode) == NEWMV) { + int_mv this_mv = curr_ref_mv_stack[ref_mv_idx].comp_mv; + clamp_mv_ref(&this_mv.as_mv, bw, bh, xd); + mbmi_ext->ref_mvs[mbmi->ref_frame[1]][0] = this_mv; + mbmi->pred_mv[1] = this_mv; + mi_pred_mv[1] = this_mv; + } + } else { +#endif // CONFIG_EXT_INTER + if (mbmi->mode == NEWMV) { + int i; + for (i = 0; i < 1 + has_second_ref(mbmi); ++i) { + int_mv this_mv = (i == 0) ? curr_ref_mv_stack[ref_mv_idx].this_mv + : curr_ref_mv_stack[ref_mv_idx].comp_mv; + clamp_mv_ref(&this_mv.as_mv, bw, bh, xd); + mbmi_ext->ref_mvs[mbmi->ref_frame[i]][0] = this_mv; + mbmi->pred_mv[i] = this_mv; + mi_pred_mv[i] = this_mv; + } + } +#if CONFIG_EXT_INTER + } +#endif // CONFIG_EXT_INTER +} +#endif // CONFIG_REF_MV + +static void update_state(const AV1_COMP *const cpi, ThreadData *td, + PICK_MODE_CONTEXT *ctx, int mi_row, int mi_col, + BLOCK_SIZE bsize, RUN_TYPE dry_run) { + int i, x_idx, y; + const AV1_COMMON *const cm = &cpi->common; + RD_COUNTS *const rdc = &td->rd_counts; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + struct macroblock_plane *const p = x->plane; + struct macroblockd_plane *const pd = xd->plane; + MODE_INFO *mi = &ctx->mic; + MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi; + MODE_INFO *mi_addr = xd->mi[0]; + const struct segmentation *const seg = &cm->seg; + const int bw = mi_size_wide[mi->mbmi.sb_type]; + const int bh = mi_size_high[mi->mbmi.sb_type]; + const int x_mis = AOMMIN(bw, cm->mi_cols - mi_col); + const int y_mis = AOMMIN(bh, cm->mi_rows - mi_row); + MV_REF *const frame_mvs = cm->cur_frame->mvs + mi_row * cm->mi_cols + mi_col; + int w, h; + + const int mis = cm->mi_stride; + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + const int unify_bsize = CONFIG_CB4X4; + +#if CONFIG_REF_MV + int8_t rf_type; +#endif + +#if !CONFIG_SUPERTX + assert(mi->mbmi.sb_type == bsize); +#endif + + *mi_addr = *mi; + *x->mbmi_ext = ctx->mbmi_ext; + +#if CONFIG_DUAL_FILTER + reset_intmv_filter_type(cm, xd, mbmi); +#endif + +#if CONFIG_REF_MV + rf_type = av1_ref_frame_type(mbmi->ref_frame); + if (x->mbmi_ext->ref_mv_count[rf_type] > 1 && + (mbmi->sb_type >= BLOCK_8X8 || unify_bsize)) { + set_ref_and_pred_mvs(x, mi->mbmi.pred_mv, rf_type); + } +#endif // CONFIG_REF_MV + + // If segmentation in use + if (seg->enabled) { + // For in frame complexity AQ copy the segment id from the segment map. + if (cpi->oxcf.aq_mode == COMPLEXITY_AQ) { + const uint8_t *const map = + seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map; + mi_addr->mbmi.segment_id = get_segment_id(cm, map, bsize, mi_row, mi_col); + reset_tx_size(xd, &mi_addr->mbmi, cm->tx_mode); + } + // Else for cyclic refresh mode update the segment map, set the segment id + // and then update the quantizer. + if (cpi->oxcf.aq_mode == CYCLIC_REFRESH_AQ) { + av1_cyclic_refresh_update_segment(cpi, &xd->mi[0]->mbmi, mi_row, mi_col, + bsize, ctx->rate, ctx->dist, x->skip); + reset_tx_size(xd, &mi_addr->mbmi, cm->tx_mode); + } + } + + for (i = 0; i < MAX_MB_PLANE; ++i) { + p[i].coeff = ctx->coeff[i]; + p[i].qcoeff = ctx->qcoeff[i]; + pd[i].dqcoeff = ctx->dqcoeff[i]; +#if CONFIG_PVQ + pd[i].pvq_ref_coeff = ctx->pvq_ref_coeff[i]; +#endif + p[i].eobs = ctx->eobs[i]; +#if CONFIG_LV_MAP + p[i].txb_entropy_ctx = ctx->txb_entropy_ctx[i]; +#endif // CONFIG_LV_MAP + } +#if CONFIG_PALETTE + for (i = 0; i < 2; ++i) pd[i].color_index_map = ctx->color_index_map[i]; +#endif // CONFIG_PALETTE + + // Restore the coding context of the MB to that that was in place + // when the mode was picked for it + for (y = 0; y < mi_height; y++) + for (x_idx = 0; x_idx < mi_width; x_idx++) + if ((xd->mb_to_right_edge >> (3 + MI_SIZE_LOG2)) + mi_width > x_idx && + (xd->mb_to_bottom_edge >> (3 + MI_SIZE_LOG2)) + mi_height > y) { + xd->mi[x_idx + y * mis] = mi_addr; + } + +#if CONFIG_DELTA_Q && !CONFIG_EXT_DELTA_Q + if (cpi->oxcf.aq_mode > NO_AQ && cpi->oxcf.aq_mode < DELTA_AQ) + av1_init_plane_quantizers(cpi, x, xd->mi[0]->mbmi.segment_id); +#else + if (cpi->oxcf.aq_mode) + av1_init_plane_quantizers(cpi, x, xd->mi[0]->mbmi.segment_id); +#endif + + if (is_inter_block(mbmi) && mbmi->sb_type < BLOCK_8X8 && !unify_bsize) { + mbmi->mv[0].as_int = mi->bmi[3].as_mv[0].as_int; + mbmi->mv[1].as_int = mi->bmi[3].as_mv[1].as_int; + } + + x->skip = ctx->skip; + +#if CONFIG_VAR_TX + for (i = 0; i < 1; ++i) + memcpy(x->blk_skip[i], ctx->blk_skip[i], + sizeof(uint8_t) * ctx->num_4x4_blk); +#endif + + if (dry_run) return; + +#if CONFIG_INTERNAL_STATS + { + unsigned int *const mode_chosen_counts = + (unsigned int *)cpi->mode_chosen_counts; // Cast const away. + if (frame_is_intra_only(cm)) { + static const int kf_mode_index[] = { + THR_DC /*DC_PRED*/, + THR_V_PRED /*V_PRED*/, + THR_H_PRED /*H_PRED*/, + THR_D45_PRED /*D45_PRED*/, + THR_D135_PRED /*D135_PRED*/, + THR_D117_PRED /*D117_PRED*/, + THR_D153_PRED /*D153_PRED*/, + THR_D207_PRED /*D207_PRED*/, + THR_D63_PRED /*D63_PRED*/, +#if CONFIG_ALT_INTRA + THR_SMOOTH, /*SMOOTH_PRED*/ +#endif // CONFIG_ALT_INTRA + THR_TM /*TM_PRED*/, + }; + ++mode_chosen_counts[kf_mode_index[mbmi->mode]]; + } else { + // Note how often each mode chosen as best + ++mode_chosen_counts[ctx->best_mode_index]; + } + } +#endif + if (!frame_is_intra_only(cm)) { + if (is_inter_block(mbmi)) { + av1_update_mv_count(td); +#if CONFIG_GLOBAL_MOTION + if (bsize >= BLOCK_8X8) { + // TODO(sarahparker): global motion stats need to be handled per-tile + // to be compatible with tile-based threading. + update_global_motion_used(mbmi->mode, bsize, mbmi, rdc); + } else { + const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize]; + const int num_4x4_h = num_4x4_blocks_high_lookup[bsize]; + int idx, idy; + for (idy = 0; idy < 2; idy += num_4x4_h) { + for (idx = 0; idx < 2; idx += num_4x4_w) { + const int j = idy * 2 + idx; + update_global_motion_used(mi->bmi[j].as_mode, bsize, mbmi, rdc); + } + } + } +#endif // CONFIG_GLOBAL_MOTION + if (cm->interp_filter == SWITCHABLE +#if CONFIG_WARPED_MOTION + && mbmi->motion_mode != WARPED_CAUSAL +#endif // CONFIG_WARPED_MOTION +#if CONFIG_GLOBAL_MOTION + && !is_nontrans_global_motion(xd) +#endif // CONFIG_GLOBAL_MOTION + ) { +#if CONFIG_DUAL_FILTER + update_filter_type_count(td->counts, xd, mbmi); +#else + const int switchable_ctx = av1_get_pred_context_switchable_interp(xd); + ++td->counts->switchable_interp[switchable_ctx][mbmi->interp_filter]; +#endif + } + } + + rdc->comp_pred_diff[SINGLE_REFERENCE] += ctx->single_pred_diff; + rdc->comp_pred_diff[COMPOUND_REFERENCE] += ctx->comp_pred_diff; + rdc->comp_pred_diff[REFERENCE_MODE_SELECT] += ctx->hybrid_pred_diff; + } + + for (h = 0; h < y_mis; ++h) { + MV_REF *const frame_mv = frame_mvs + h * cm->mi_cols; + for (w = 0; w < x_mis; ++w) { + MV_REF *const mv = frame_mv + w; + mv->ref_frame[0] = mi->mbmi.ref_frame[0]; + mv->ref_frame[1] = mi->mbmi.ref_frame[1]; + mv->mv[0].as_int = mi->mbmi.mv[0].as_int; + mv->mv[1].as_int = mi->mbmi.mv[1].as_int; + } + } +} + +#if CONFIG_SUPERTX +static void update_state_supertx(const AV1_COMP *const cpi, ThreadData *td, + PICK_MODE_CONTEXT *ctx, int mi_row, int mi_col, + BLOCK_SIZE bsize, RUN_TYPE dry_run) { + int y, x_idx; +#if CONFIG_VAR_TX + int i; +#endif + const AV1_COMMON *const cm = &cpi->common; + RD_COUNTS *const rdc = &td->rd_counts; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + MODE_INFO *mi = &ctx->mic; + MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi; + MODE_INFO *mi_addr = xd->mi[0]; + const struct segmentation *const seg = &cm->seg; + const int mis = cm->mi_stride; + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + const int x_mis = AOMMIN(mi_width, cm->mi_cols - mi_col); + const int y_mis = AOMMIN(mi_height, cm->mi_rows - mi_row); + const int unify_bsize = CONFIG_CB4X4; + MV_REF *const frame_mvs = cm->cur_frame->mvs + mi_row * cm->mi_cols + mi_col; + int w, h; + +#if CONFIG_REF_MV + int8_t rf_type; +#endif + + *mi_addr = *mi; + *x->mbmi_ext = ctx->mbmi_ext; + assert(is_inter_block(mbmi)); + assert(mbmi->tx_size == ctx->mic.mbmi.tx_size); + +#if CONFIG_DUAL_FILTER + reset_intmv_filter_type(cm, xd, mbmi); +#endif + +#if CONFIG_REF_MV + rf_type = av1_ref_frame_type(mbmi->ref_frame); + if (x->mbmi_ext->ref_mv_count[rf_type] > 1 && + (mbmi->sb_type >= BLOCK_8X8 || unify_bsize)) { + set_ref_and_pred_mvs(x, mi->mbmi.pred_mv, rf_type); + } +#endif // CONFIG_REF_MV + + // If segmentation in use + if (seg->enabled) { + if (cpi->vaq_refresh) { + const int energy = + bsize <= BLOCK_16X16 ? x->mb_energy : av1_block_energy(cpi, x, bsize); + mi_addr->mbmi.segment_id = av1_vaq_segment_id(energy); + } else if (cpi->oxcf.aq_mode == CYCLIC_REFRESH_AQ) { + // For cyclic refresh mode, now update the segment map + // and set the segment id. + av1_cyclic_refresh_update_segment(cpi, &xd->mi[0]->mbmi, mi_row, mi_col, + bsize, ctx->rate, ctx->dist, 1); + } else { + // Otherwise just set the segment id based on the current segment map + const uint8_t *const map = + seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map; + mi_addr->mbmi.segment_id = get_segment_id(cm, map, bsize, mi_row, mi_col); + } + mi_addr->mbmi.segment_id_supertx = MAX_SEGMENTS; + } + + // Restore the coding context of the MB to that that was in place + // when the mode was picked for it + for (y = 0; y < mi_height; y++) + for (x_idx = 0; x_idx < mi_width; x_idx++) + if ((xd->mb_to_right_edge >> (3 + MI_SIZE_LOG2)) + mi_width > x_idx && + (xd->mb_to_bottom_edge >> (3 + MI_SIZE_LOG2)) + mi_height > y) { + xd->mi[x_idx + y * mis] = mi_addr; + } + +#if !CONFIG_CB4X4 + if (is_inter_block(mbmi) && mbmi->sb_type < BLOCK_8X8) { + mbmi->mv[0].as_int = mi->bmi[3].as_mv[0].as_int; + mbmi->mv[1].as_int = mi->bmi[3].as_mv[1].as_int; + } +#endif + + x->skip = ctx->skip; + +#if CONFIG_VAR_TX + for (i = 0; i < 1; ++i) + memcpy(x->blk_skip[i], ctx->blk_skip[i], + sizeof(uint8_t) * ctx->num_4x4_blk); + + if (!is_inter_block(mbmi) || mbmi->skip) + mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size); +#endif // CONFIG_VAR_TX + +#if CONFIG_VAR_TX + { + const TX_SIZE mtx = mbmi->tx_size; + const int num_4x4_blocks_wide = tx_size_wide_unit[mtx] >> 1; + const int num_4x4_blocks_high = tx_size_high_unit[mtx] >> 1; + int idy, idx; + mbmi->inter_tx_size[0][0] = mtx; + for (idy = 0; idy < num_4x4_blocks_high; ++idy) + for (idx = 0; idx < num_4x4_blocks_wide; ++idx) + mbmi->inter_tx_size[idy][idx] = mtx; + } +#endif // CONFIG_VAR_TX + // Turn motion variation off for supertx + mbmi->motion_mode = SIMPLE_TRANSLATION; + + if (dry_run) return; + + if (!frame_is_intra_only(cm)) { + av1_update_mv_count(td); + +#if CONFIG_GLOBAL_MOTION + if (is_inter_block(mbmi)) { + if (bsize >= BLOCK_8X8) { + // TODO(sarahparker): global motion stats need to be handled per-tile + // to be compatible with tile-based threading. + update_global_motion_used(mbmi->mode, bsize, mbmi, rdc); + } else { + const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize]; + const int num_4x4_h = num_4x4_blocks_high_lookup[bsize]; + int idx, idy; + for (idy = 0; idy < 2; idy += num_4x4_h) { + for (idx = 0; idx < 2; idx += num_4x4_w) { + const int j = idy * 2 + idx; + update_global_motion_used(mi->bmi[j].as_mode, bsize, mbmi, rdc); + } + } + } + } +#endif // CONFIG_GLOBAL_MOTION + + if (cm->interp_filter == SWITCHABLE +#if CONFIG_GLOBAL_MOTION + && !is_nontrans_global_motion(xd) +#endif // CONFIG_GLOBAL_MOTION + ) { +#if CONFIG_DUAL_FILTER + update_filter_type_count(td->counts, xd, mbmi); +#else + const int pred_ctx = av1_get_pred_context_switchable_interp(xd); + ++td->counts->switchable_interp[pred_ctx][mbmi->interp_filter]; +#endif + } + + rdc->comp_pred_diff[SINGLE_REFERENCE] += ctx->single_pred_diff; + rdc->comp_pred_diff[COMPOUND_REFERENCE] += ctx->comp_pred_diff; + rdc->comp_pred_diff[REFERENCE_MODE_SELECT] += ctx->hybrid_pred_diff; + } + + for (h = 0; h < y_mis; ++h) { + MV_REF *const frame_mv = frame_mvs + h * cm->mi_cols; + for (w = 0; w < x_mis; ++w) { + MV_REF *const mv = frame_mv + w; + mv->ref_frame[0] = mi->mbmi.ref_frame[0]; + mv->ref_frame[1] = mi->mbmi.ref_frame[1]; + mv->mv[0].as_int = mi->mbmi.mv[0].as_int; + mv->mv[1].as_int = mi->mbmi.mv[1].as_int; + } + } +} + +static void update_state_sb_supertx(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, + int mi_col, BLOCK_SIZE bsize, + RUN_TYPE dry_run, PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + struct macroblock_plane *const p = x->plane; + struct macroblockd_plane *const pd = xd->plane; + int hbs = mi_size_wide[bsize] / 2; +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; +#endif + PARTITION_TYPE partition = pc_tree->partitioning; + BLOCK_SIZE subsize = get_subsize(bsize, partition); + int i; +#if CONFIG_EXT_PARTITION_TYPES + BLOCK_SIZE bsize2 = get_subsize(bsize, PARTITION_SPLIT); +#endif + PICK_MODE_CONTEXT *pmc = NULL; + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + + if (bsize == BLOCK_16X16 && cpi->vaq_refresh) + x->mb_energy = av1_block_energy(cpi, x, bsize); + + switch (partition) { + case PARTITION_NONE: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->none, mi_row, mi_col, subsize, + dry_run); + break; + case PARTITION_VERT: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->vertical[0], mi_row, mi_col, + subsize, dry_run); + if (mi_col + hbs < cm->mi_cols && (bsize > BLOCK_8X8 || unify_bsize)) { + set_offsets_supertx(cpi, td, tile, mi_row, mi_col + hbs, subsize); + update_state_supertx(cpi, td, &pc_tree->vertical[1], mi_row, + mi_col + hbs, subsize, dry_run); + } + pmc = &pc_tree->vertical_supertx; + break; + case PARTITION_HORZ: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->horizontal[0], mi_row, mi_col, + subsize, dry_run); + if (mi_row + hbs < cm->mi_rows && (bsize > BLOCK_8X8 || unify_bsize)) { + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->horizontal[1], mi_row + hbs, + mi_col, subsize, dry_run); + } + pmc = &pc_tree->horizontal_supertx; + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_supertx(cpi, td, pc_tree->leaf_split[0], mi_row, mi_col, + subsize, dry_run); + } else { + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_sb_supertx(cpi, td, tile, mi_row, mi_col, subsize, dry_run, + pc_tree->split[0]); + set_offsets_supertx(cpi, td, tile, mi_row, mi_col + hbs, subsize); + update_state_sb_supertx(cpi, td, tile, mi_row, mi_col + hbs, subsize, + dry_run, pc_tree->split[1]); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col, subsize); + update_state_sb_supertx(cpi, td, tile, mi_row + hbs, mi_col, subsize, + dry_run, pc_tree->split[2]); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col + hbs, subsize); + update_state_sb_supertx(cpi, td, tile, mi_row + hbs, mi_col + hbs, + subsize, dry_run, pc_tree->split[3]); + } + pmc = &pc_tree->split_supertx; + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, bsize2); + update_state_supertx(cpi, td, &pc_tree->horizontala[0], mi_row, mi_col, + bsize2, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row, mi_col + hbs, bsize2); + update_state_supertx(cpi, td, &pc_tree->horizontala[1], mi_row, + mi_col + hbs, bsize2, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->horizontala[2], mi_row + hbs, + mi_col, subsize, dry_run); + pmc = &pc_tree->horizontala_supertx; + break; + case PARTITION_HORZ_B: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->horizontalb[0], mi_row, mi_col, + subsize, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col, bsize2); + update_state_supertx(cpi, td, &pc_tree->horizontalb[1], mi_row + hbs, + mi_col, bsize2, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col + hbs, bsize2); + update_state_supertx(cpi, td, &pc_tree->horizontalb[2], mi_row + hbs, + mi_col + hbs, bsize2, dry_run); + pmc = &pc_tree->horizontalb_supertx; + break; + case PARTITION_VERT_A: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, bsize2); + update_state_supertx(cpi, td, &pc_tree->verticala[0], mi_row, mi_col, + bsize2, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col, bsize2); + update_state_supertx(cpi, td, &pc_tree->verticala[1], mi_row + hbs, + mi_col, bsize2, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row, mi_col + hbs, subsize); + update_state_supertx(cpi, td, &pc_tree->verticala[2], mi_row, + mi_col + hbs, subsize, dry_run); + pmc = &pc_tree->verticala_supertx; + break; + case PARTITION_VERT_B: + set_offsets_supertx(cpi, td, tile, mi_row, mi_col, subsize); + update_state_supertx(cpi, td, &pc_tree->verticalb[0], mi_row, mi_col, + subsize, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row, mi_col + hbs, bsize2); + update_state_supertx(cpi, td, &pc_tree->verticalb[1], mi_row, + mi_col + hbs, bsize2, dry_run); + set_offsets_supertx(cpi, td, tile, mi_row + hbs, mi_col + hbs, bsize2); + update_state_supertx(cpi, td, &pc_tree->verticalb[2], mi_row + hbs, + mi_col + hbs, bsize2, dry_run); + pmc = &pc_tree->verticalb_supertx; + break; +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0); + } + + for (i = 0; i < MAX_MB_PLANE; ++i) { + if (pmc != NULL) { + p[i].coeff = pmc->coeff[i]; + p[i].qcoeff = pmc->qcoeff[i]; + pd[i].dqcoeff = pmc->dqcoeff[i]; + p[i].eobs = pmc->eobs[i]; + } else { + // These should never be used + p[i].coeff = NULL; + p[i].qcoeff = NULL; + pd[i].dqcoeff = NULL; + p[i].eobs = NULL; + } + } +} + +static void update_supertx_param(ThreadData *td, PICK_MODE_CONTEXT *ctx, + int best_tx, TX_SIZE supertx_size) { + MACROBLOCK *const x = &td->mb; +#if CONFIG_VAR_TX + int i; + + for (i = 0; i < 1; ++i) + memcpy(ctx->blk_skip[i], x->blk_skip[i], + sizeof(uint8_t) * ctx->num_4x4_blk); + ctx->mic.mbmi.min_tx_size = get_min_tx_size(supertx_size); +#endif // CONFIG_VAR_TX + ctx->mic.mbmi.tx_size = supertx_size; + ctx->skip = x->skip; + ctx->mic.mbmi.tx_type = best_tx; +} + +static void update_supertx_param_sb(const AV1_COMP *const cpi, ThreadData *td, + int mi_row, int mi_col, BLOCK_SIZE bsize, + int best_tx, TX_SIZE supertx_size, + PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + const int hbs = mi_size_wide[bsize] / 2; + PARTITION_TYPE partition = pc_tree->partitioning; + BLOCK_SIZE subsize = get_subsize(bsize, partition); +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; +#endif +#if CONFIG_EXT_PARTITION_TYPES + int i; +#endif + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + + switch (partition) { + case PARTITION_NONE: + update_supertx_param(td, &pc_tree->none, best_tx, supertx_size); + break; + case PARTITION_VERT: + update_supertx_param(td, &pc_tree->vertical[0], best_tx, supertx_size); + if (mi_col + hbs < cm->mi_cols && (bsize > BLOCK_8X8 || unify_bsize)) + update_supertx_param(td, &pc_tree->vertical[1], best_tx, supertx_size); + break; + case PARTITION_HORZ: + update_supertx_param(td, &pc_tree->horizontal[0], best_tx, supertx_size); + if (mi_row + hbs < cm->mi_rows && (bsize > BLOCK_8X8 || unify_bsize)) + update_supertx_param(td, &pc_tree->horizontal[1], best_tx, + supertx_size); + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + update_supertx_param(td, pc_tree->leaf_split[0], best_tx, supertx_size); + } else { + update_supertx_param_sb(cpi, td, mi_row, mi_col, subsize, best_tx, + supertx_size, pc_tree->split[0]); + update_supertx_param_sb(cpi, td, mi_row, mi_col + hbs, subsize, best_tx, + supertx_size, pc_tree->split[1]); + update_supertx_param_sb(cpi, td, mi_row + hbs, mi_col, subsize, best_tx, + supertx_size, pc_tree->split[2]); + update_supertx_param_sb(cpi, td, mi_row + hbs, mi_col + hbs, subsize, + best_tx, supertx_size, pc_tree->split[3]); + } + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + for (i = 0; i < 3; i++) + update_supertx_param(td, &pc_tree->horizontala[i], best_tx, + supertx_size); + break; + case PARTITION_HORZ_B: + for (i = 0; i < 3; i++) + update_supertx_param(td, &pc_tree->horizontalb[i], best_tx, + supertx_size); + break; + case PARTITION_VERT_A: + for (i = 0; i < 3; i++) + update_supertx_param(td, &pc_tree->verticala[i], best_tx, supertx_size); + break; + case PARTITION_VERT_B: + for (i = 0; i < 3; i++) + update_supertx_param(td, &pc_tree->verticalb[i], best_tx, supertx_size); + break; +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0); + } +} +#endif // CONFIG_SUPERTX + +#if CONFIG_MOTION_VAR && CONFIG_NCOBMC +static void set_mode_info_b(const AV1_COMP *const cpi, + const TileInfo *const tile, ThreadData *td, + int mi_row, int mi_col, BLOCK_SIZE bsize, + PICK_MODE_CONTEXT *ctx) { + MACROBLOCK *const x = &td->mb; + set_offsets(cpi, tile, x, mi_row, mi_col, bsize); + update_state(cpi, td, ctx, mi_row, mi_col, bsize, 1); +} + +static void set_mode_info_sb(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, TOKENEXTRA **tp, + int mi_row, int mi_col, BLOCK_SIZE bsize, + PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + const int hbs = mi_size_wide[bsize] / 2; + const PARTITION_TYPE partition = pc_tree->partitioning; + BLOCK_SIZE subsize = get_subsize(bsize, partition); +#if CONFIG_EXT_PARTITION_TYPES + const BLOCK_SIZE bsize2 = get_subsize(bsize, PARTITION_SPLIT); +#endif +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; + assert(bsize >= BLOCK_8X8); +#endif + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + + switch (partition) { + case PARTITION_NONE: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, subsize, &pc_tree->none); + break; + case PARTITION_VERT: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, subsize, + &pc_tree->vertical[0]); + if (mi_col + hbs < cm->mi_cols && (bsize > BLOCK_8X8 || unify_bsize)) { + set_mode_info_b(cpi, tile, td, mi_row, mi_col + hbs, subsize, + &pc_tree->vertical[1]); + } + break; + case PARTITION_HORZ: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, subsize, + &pc_tree->horizontal[0]); + if (mi_row + hbs < cm->mi_rows && (bsize > BLOCK_8X8 || unify_bsize)) { + set_mode_info_b(cpi, tile, td, mi_row + hbs, mi_col, subsize, + &pc_tree->horizontal[1]); + } + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + set_mode_info_b(cpi, tile, td, mi_row, mi_col, subsize, + pc_tree->leaf_split[0]); + } else { + set_mode_info_sb(cpi, td, tile, tp, mi_row, mi_col, subsize, + pc_tree->split[0]); + set_mode_info_sb(cpi, td, tile, tp, mi_row, mi_col + hbs, subsize, + pc_tree->split[1]); + set_mode_info_sb(cpi, td, tile, tp, mi_row + hbs, mi_col, subsize, + pc_tree->split[2]); + set_mode_info_sb(cpi, td, tile, tp, mi_row + hbs, mi_col + hbs, subsize, + pc_tree->split[3]); + } + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, bsize2, + &pc_tree->horizontala[0]); + set_mode_info_b(cpi, tile, td, mi_row, mi_col + hbs, bsize2, + &pc_tree->horizontala[1]); + set_mode_info_b(cpi, tile, td, mi_row + hbs, mi_col, subsize, + &pc_tree->horizontala[2]); + break; + case PARTITION_HORZ_B: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, subsize, + &pc_tree->horizontalb[0]); + set_mode_info_b(cpi, tile, td, mi_row + hbs, mi_col, bsize2, + &pc_tree->horizontalb[1]); + set_mode_info_b(cpi, tile, td, mi_row + hbs, mi_col + hbs, bsize2, + &pc_tree->horizontalb[2]); + break; + case PARTITION_VERT_A: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, bsize2, + &pc_tree->verticala[0]); + set_mode_info_b(cpi, tile, td, mi_row + hbs, mi_col, bsize2, + &pc_tree->verticala[1]); + set_mode_info_b(cpi, tile, td, mi_row, mi_col + hbs, subsize, + &pc_tree->verticala[2]); + break; + case PARTITION_VERT_B: + set_mode_info_b(cpi, tile, td, mi_row, mi_col, subsize, + &pc_tree->verticalb[0]); + set_mode_info_b(cpi, tile, td, mi_row, mi_col + hbs, bsize2, + &pc_tree->verticalb[1]); + set_mode_info_b(cpi, tile, td, mi_row + hbs, mi_col + hbs, bsize2, + &pc_tree->verticalb[2]); + break; +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0 && "Invalid partition type."); break; + } +} +#endif + +void av1_setup_src_planes(MACROBLOCK *x, const YV12_BUFFER_CONFIG *src, + int mi_row, int mi_col) { + uint8_t *const buffers[3] = { src->y_buffer, src->u_buffer, src->v_buffer }; + const int widths[3] = { src->y_crop_width, src->uv_crop_width, + src->uv_crop_width }; + const int heights[3] = { src->y_crop_height, src->uv_crop_height, + src->uv_crop_height }; + const int strides[3] = { src->y_stride, src->uv_stride, src->uv_stride }; + int i; + + // Set current frame pointer. + x->e_mbd.cur_buf = src; + + for (i = 0; i < MAX_MB_PLANE; i++) + setup_pred_plane(&x->plane[i].src, x->e_mbd.mi[0]->mbmi.sb_type, buffers[i], + widths[i], heights[i], strides[i], mi_row, mi_col, NULL, + x->e_mbd.plane[i].subsampling_x, + x->e_mbd.plane[i].subsampling_y); +} + +static int set_segment_rdmult(const AV1_COMP *const cpi, MACROBLOCK *const x, + int8_t segment_id) { + int segment_qindex; + const AV1_COMMON *const cm = &cpi->common; + av1_init_plane_quantizers(cpi, x, segment_id); + aom_clear_system_state(); + segment_qindex = av1_get_qindex(&cm->seg, segment_id, cm->base_qindex); + return av1_compute_rd_mult(cpi, segment_qindex + cm->y_dc_delta_q); +} + +static void rd_pick_sb_modes(const AV1_COMP *const cpi, TileDataEnc *tile_data, + MACROBLOCK *const x, int mi_row, int mi_col, + RD_STATS *rd_cost, +#if CONFIG_SUPERTX + int *totalrate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_TYPE partition, +#endif + BLOCK_SIZE bsize, PICK_MODE_CONTEXT *ctx, + int64_t best_rd) { + const AV1_COMMON *const cm = &cpi->common; + TileInfo *const tile_info = &tile_data->tile_info; + MACROBLOCKD *const xd = &x->e_mbd; + MB_MODE_INFO *mbmi; + struct macroblock_plane *const p = x->plane; + struct macroblockd_plane *const pd = xd->plane; + const AQ_MODE aq_mode = cpi->oxcf.aq_mode; + int i, orig_rdmult; + const int unify_bsize = CONFIG_CB4X4; + + aom_clear_system_state(); + +#if CONFIG_PVQ + x->pvq_speed = 1; + x->pvq_coded = 0; +#endif +#if CONFIG_CFL + // Don't store luma during RDO (we will store the best mode later). + x->cfl_store_y = 0; +#endif + + set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize); + mbmi = &xd->mi[0]->mbmi; + mbmi->sb_type = bsize; +#if CONFIG_RD_DEBUG + mbmi->mi_row = mi_row; + mbmi->mi_col = mi_col; +#endif +#if CONFIG_SUPERTX + // We set tx_size here as skip blocks would otherwise not set it. + // tx_size needs to be set at this point as supertx_enable in + // write_modes_sb is computed based on this, and if the garbage in memory + // just happens to be the supertx_size, then the packer will code this + // block as a supertx block, even if rdopt did not pick it as such. + mbmi->tx_size = max_txsize_lookup[bsize]; +#endif +#if CONFIG_EXT_PARTITION_TYPES + mbmi->partition = partition; +#endif + + for (i = 0; i < MAX_MB_PLANE; ++i) { + p[i].coeff = ctx->coeff[i]; + p[i].qcoeff = ctx->qcoeff[i]; + pd[i].dqcoeff = ctx->dqcoeff[i]; +#if CONFIG_PVQ + pd[i].pvq_ref_coeff = ctx->pvq_ref_coeff[i]; +#endif + p[i].eobs = ctx->eobs[i]; +#if CONFIG_LV_MAP + p[i].txb_entropy_ctx = ctx->txb_entropy_ctx[i]; +#endif + } + +#if CONFIG_PALETTE + for (i = 0; i < 2; ++i) pd[i].color_index_map = ctx->color_index_map[i]; +#endif // CONFIG_PALETTE + + ctx->skippable = 0; + ctx->pred_pixel_ready = 0; + + // Set to zero to make sure we do not use the previous encoded frame stats + mbmi->skip = 0; + +#if CONFIG_CB4X4 + x->skip_chroma_rd = + !is_chroma_reference(mi_row, mi_col, bsize, xd->plane[1].subsampling_x, + xd->plane[1].subsampling_y); +#endif + +#if CONFIG_HIGHBITDEPTH + if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) { + x->source_variance = av1_high_get_sby_perpixel_variance( + cpi, &x->plane[0].src, bsize, xd->bd); + } else { + x->source_variance = + av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize); + } +#else + x->source_variance = + av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize); +#endif // CONFIG_HIGHBITDEPTH + + // Save rdmult before it might be changed, so it can be restored later. + orig_rdmult = x->rdmult; + + if (aq_mode == VARIANCE_AQ) { + if (cpi->vaq_refresh) { + const int energy = + bsize <= BLOCK_16X16 ? x->mb_energy : av1_block_energy(cpi, x, bsize); + mbmi->segment_id = av1_vaq_segment_id(energy); + // Re-initialise quantiser + av1_init_plane_quantizers(cpi, x, mbmi->segment_id); + } + x->rdmult = set_segment_rdmult(cpi, x, mbmi->segment_id); + } else if (aq_mode == COMPLEXITY_AQ) { + x->rdmult = set_segment_rdmult(cpi, x, mbmi->segment_id); + } else if (aq_mode == CYCLIC_REFRESH_AQ) { + // If segment is boosted, use rdmult for that segment. + if (cyclic_refresh_segment_id_boosted(mbmi->segment_id)) + x->rdmult = av1_cyclic_refresh_get_rdmult(cpi->cyclic_refresh); + } + + // Find best coding mode & reconstruct the MB so it is available + // as a predictor for MBs that follow in the SB + if (frame_is_intra_only(cm)) { + av1_rd_pick_intra_mode_sb(cpi, x, rd_cost, bsize, ctx, best_rd); +#if CONFIG_SUPERTX + *totalrate_nocoef = 0; +#endif // CONFIG_SUPERTX + } else { + if (bsize >= BLOCK_8X8 || unify_bsize) { + if (segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) { + av1_rd_pick_inter_mode_sb_seg_skip(cpi, tile_data, x, mi_row, mi_col, + rd_cost, bsize, ctx, best_rd); +#if CONFIG_SUPERTX + *totalrate_nocoef = rd_cost->rate; +#endif // CONFIG_SUPERTX + } else { + av1_rd_pick_inter_mode_sb(cpi, tile_data, x, mi_row, mi_col, rd_cost, +#if CONFIG_SUPERTX + totalrate_nocoef, +#endif // CONFIG_SUPERTX + bsize, ctx, best_rd); +#if CONFIG_SUPERTX + assert(*totalrate_nocoef >= 0); +#endif // CONFIG_SUPERTX + } + } else { + if (segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) { + // The decoder rejects sub8x8 partitions when SEG_LVL_SKIP is set. + rd_cost->rate = INT_MAX; + } else { + av1_rd_pick_inter_mode_sub8x8(cpi, tile_data, x, mi_row, mi_col, + rd_cost, +#if CONFIG_SUPERTX + totalrate_nocoef, +#endif // CONFIG_SUPERTX + bsize, ctx, best_rd); +#if CONFIG_SUPERTX + assert(*totalrate_nocoef >= 0); +#endif // CONFIG_SUPERTX + } + } + } + + // Examine the resulting rate and for AQ mode 2 make a segment choice. + if ((rd_cost->rate != INT_MAX) && (aq_mode == COMPLEXITY_AQ) && + (bsize >= BLOCK_16X16) && + (cm->frame_type == KEY_FRAME || cpi->refresh_alt_ref_frame || + (cpi->refresh_golden_frame && !cpi->rc.is_src_frame_alt_ref))) { + av1_caq_select_segment(cpi, x, bsize, mi_row, mi_col, rd_cost->rate); + } + + x->rdmult = orig_rdmult; + + // TODO(jingning) The rate-distortion optimization flow needs to be + // refactored to provide proper exit/return handle. + if (rd_cost->rate == INT_MAX) rd_cost->rdcost = INT64_MAX; + + ctx->rate = rd_cost->rate; + ctx->dist = rd_cost->dist; +} + +#if CONFIG_REF_MV +static void update_inter_mode_stats(FRAME_COUNTS *counts, PREDICTION_MODE mode, + int16_t mode_context) { + int16_t mode_ctx = mode_context & NEWMV_CTX_MASK; + if (mode == NEWMV) { + ++counts->newmv_mode[mode_ctx][0]; + return; + } else { + ++counts->newmv_mode[mode_ctx][1]; + + if (mode_context & (1 << ALL_ZERO_FLAG_OFFSET)) { + return; + } + + mode_ctx = (mode_context >> ZEROMV_OFFSET) & ZEROMV_CTX_MASK; + if (mode == ZEROMV) { + ++counts->zeromv_mode[mode_ctx][0]; + return; + } else { + ++counts->zeromv_mode[mode_ctx][1]; + mode_ctx = (mode_context >> REFMV_OFFSET) & REFMV_CTX_MASK; + + if (mode_context & (1 << SKIP_NEARESTMV_OFFSET)) mode_ctx = 6; + if (mode_context & (1 << SKIP_NEARMV_OFFSET)) mode_ctx = 7; + if (mode_context & (1 << SKIP_NEARESTMV_SUB8X8_OFFSET)) mode_ctx = 8; + + ++counts->refmv_mode[mode_ctx][mode != NEARESTMV]; + } + } +} +#endif + +static void update_stats(const AV1_COMMON *const cm, ThreadData *td, int mi_row, + int mi_col +#if CONFIG_SUPERTX + , + int supertx_enabled +#endif + ) { +#if CONFIG_DELTA_Q + MACROBLOCK *x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; +#else + const MACROBLOCK *x = &td->mb; + const MACROBLOCKD *const xd = &x->e_mbd; +#endif + const MODE_INFO *const mi = xd->mi[0]; + const MB_MODE_INFO *const mbmi = &mi->mbmi; + const MB_MODE_INFO_EXT *const mbmi_ext = x->mbmi_ext; + const BLOCK_SIZE bsize = mbmi->sb_type; + const int unify_bsize = CONFIG_CB4X4; + +#if CONFIG_DELTA_Q + // delta quant applies to both intra and inter + const int super_block_upper_left = ((mi_row & 7) == 0) && ((mi_col & 7) == 0); + + if (cm->delta_q_present_flag && (bsize != BLOCK_64X64 || !mbmi->skip) && + super_block_upper_left) { + const int dq = (mbmi->current_q_index - xd->prev_qindex) / cm->delta_q_res; + const int absdq = abs(dq); + int i; + for (i = 0; i < AOMMIN(absdq, DELTA_Q_SMALL); ++i) { + td->counts->delta_q[i][1]++; + } + if (absdq < DELTA_Q_SMALL) td->counts->delta_q[absdq][0]++; + xd->prev_qindex = mbmi->current_q_index; +#if CONFIG_EXT_DELTA_Q + if (cm->delta_lf_present_flag) { + const int dlf = + (mbmi->current_delta_lf_from_base - xd->prev_delta_lf_from_base) / + cm->delta_lf_res; + const int absdlf = abs(dlf); + for (i = 0; i < AOMMIN(absdlf, DELTA_LF_SMALL); ++i) { + td->counts->delta_lf[i][1]++; + } + if (absdlf < DELTA_LF_SMALL) td->counts->delta_lf[absdlf][0]++; + xd->prev_delta_lf_from_base = mbmi->current_delta_lf_from_base; + } +#endif + } +#else + (void)mi_row; + (void)mi_col; +#endif + if (!frame_is_intra_only(cm)) { + FRAME_COUNTS *const counts = td->counts; + const int inter_block = is_inter_block(mbmi); + const int seg_ref_active = + segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_REF_FRAME); + if (!seg_ref_active) { +#if CONFIG_SUPERTX + if (!supertx_enabled) +#endif + counts->intra_inter[av1_get_intra_inter_context(xd)][inter_block]++; + // If the segment reference feature is enabled we have only a single + // reference frame allowed for the segment so exclude it from + // the reference frame counts used to work out probabilities. + if (inter_block) { + const MV_REFERENCE_FRAME ref0 = mbmi->ref_frame[0]; +#if CONFIG_EXT_REFS + const MV_REFERENCE_FRAME ref1 = mbmi->ref_frame[1]; +#endif // CONFIG_EXT_REFS + + if (cm->reference_mode == REFERENCE_MODE_SELECT) { +#if !SUB8X8_COMP_REF + if (mbmi->sb_type >= BLOCK_8X8) + counts->comp_inter[av1_get_reference_mode_context(cm, xd)] + [has_second_ref(mbmi)]++; +#else + counts->comp_inter[av1_get_reference_mode_context(cm, xd)] + [has_second_ref(mbmi)]++; +#endif + } + + if (has_second_ref(mbmi)) { +#if CONFIG_EXT_REFS + const int bit = (ref0 == GOLDEN_FRAME || ref0 == LAST3_FRAME); + + counts->comp_ref[av1_get_pred_context_comp_ref_p(cm, xd)][0][bit]++; + if (!bit) { + counts->comp_ref[av1_get_pred_context_comp_ref_p1(cm, xd)][1] + [ref0 == LAST_FRAME]++; + } else { + counts->comp_ref[av1_get_pred_context_comp_ref_p2(cm, xd)][2] + [ref0 == GOLDEN_FRAME]++; + } + + counts->comp_bwdref[av1_get_pred_context_comp_bwdref_p(cm, xd)][0] + [ref1 == ALTREF_FRAME]++; +#else + counts->comp_ref[av1_get_pred_context_comp_ref_p(cm, xd)][0] + [ref0 == GOLDEN_FRAME]++; +#endif // CONFIG_EXT_REFS + } else { +#if CONFIG_EXT_REFS + const int bit = (ref0 == ALTREF_FRAME || ref0 == BWDREF_FRAME); + + counts->single_ref[av1_get_pred_context_single_ref_p1(xd)][0][bit]++; + if (bit) { + counts->single_ref[av1_get_pred_context_single_ref_p2(xd)][1] + [ref0 != BWDREF_FRAME]++; + } else { + const int bit1 = !(ref0 == LAST2_FRAME || ref0 == LAST_FRAME); + counts + ->single_ref[av1_get_pred_context_single_ref_p3(xd)][2][bit1]++; + if (!bit1) { + counts->single_ref[av1_get_pred_context_single_ref_p4(xd)][3] + [ref0 != LAST_FRAME]++; + } else { + counts->single_ref[av1_get_pred_context_single_ref_p5(xd)][4] + [ref0 != LAST3_FRAME]++; + } + } +#else + counts->single_ref[av1_get_pred_context_single_ref_p1(xd)][0] + [ref0 != LAST_FRAME]++; + if (ref0 != LAST_FRAME) { + counts->single_ref[av1_get_pred_context_single_ref_p2(xd)][1] + [ref0 != GOLDEN_FRAME]++; + } +#endif // CONFIG_EXT_REFS + } + +#if CONFIG_EXT_INTER + if (cm->reference_mode != COMPOUND_REFERENCE && +#if CONFIG_SUPERTX + !supertx_enabled && +#endif + is_interintra_allowed(mbmi)) { + const int bsize_group = size_group_lookup[bsize]; + if (mbmi->ref_frame[1] == INTRA_FRAME) { + counts->interintra[bsize_group][1]++; + counts->interintra_mode[bsize_group][mbmi->interintra_mode]++; + if (is_interintra_wedge_used(bsize)) + counts->wedge_interintra[bsize][mbmi->use_wedge_interintra]++; + } else { + counts->interintra[bsize_group][0]++; + } + } +#endif // CONFIG_EXT_INTER + +#if CONFIG_MOTION_VAR || CONFIG_WARPED_MOTION + const MOTION_MODE motion_allowed = motion_mode_allowed( +#if CONFIG_GLOBAL_MOTION && SEPARATE_GLOBAL_MOTION + 0, xd->global_motion, +#endif // CONFIG_GLOBAL_MOTION && SEPARATE_GLOBAL_MOTION + mi); +#if CONFIG_SUPERTX + if (!supertx_enabled) +#endif // CONFIG_SUPERTX +#if CONFIG_EXT_INTER + if (mbmi->ref_frame[1] != INTRA_FRAME) +#endif // CONFIG_EXT_INTER +#if CONFIG_MOTION_VAR && CONFIG_WARPED_MOTION + { + if (motion_allowed == WARPED_CAUSAL) + counts->motion_mode[mbmi->sb_type][mbmi->motion_mode]++; + else if (motion_allowed == OBMC_CAUSAL) + counts->obmc[mbmi->sb_type][mbmi->motion_mode == OBMC_CAUSAL]++; + } +#else + if (motion_allowed > SIMPLE_TRANSLATION) + counts->motion_mode[mbmi->sb_type][mbmi->motion_mode]++; +#endif // CONFIG_MOTION_VAR && CONFIG_WARPED_MOTION +#endif // CONFIG_MOTION_VAR || CONFIG_WARPED_MOTION + +#if CONFIG_EXT_INTER + if (cm->reference_mode != SINGLE_REFERENCE && + is_inter_compound_mode(mbmi->mode) +#if CONFIG_MOTION_VAR || CONFIG_WARPED_MOTION + && mbmi->motion_mode == SIMPLE_TRANSLATION +#endif // CONFIG_MOTION_VAR || CONFIG_WARPED_MOTION + ) { + counts->compound_interinter[bsize][mbmi->interinter_compound_type]++; + } +#endif // CONFIG_EXT_INTER + } + } + + if (inter_block && + !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) { + int16_t mode_ctx; +#if !CONFIG_REF_MV + mode_ctx = mbmi_ext->mode_context[mbmi->ref_frame[0]]; +#endif + if (bsize >= BLOCK_8X8 || unify_bsize) { + const PREDICTION_MODE mode = mbmi->mode; +#if CONFIG_REF_MV +#if CONFIG_EXT_INTER + if (has_second_ref(mbmi)) { + mode_ctx = mbmi_ext->compound_mode_context[mbmi->ref_frame[0]]; + ++counts->inter_compound_mode[mode_ctx][INTER_COMPOUND_OFFSET(mode)]; + } else { +#endif // CONFIG_EXT_INTER + mode_ctx = av1_mode_context_analyzer(mbmi_ext->mode_context, + mbmi->ref_frame, bsize, -1); + update_inter_mode_stats(counts, mode, mode_ctx); +#if CONFIG_EXT_INTER + } +#endif // CONFIG_EXT_INTER + +#if CONFIG_EXT_INTER + if (mbmi->mode == NEWMV || mbmi->mode == NEW_NEWMV) { +#else + if (mbmi->mode == NEWMV) { +#endif + uint8_t ref_frame_type = av1_ref_frame_type(mbmi->ref_frame); + int idx; + + for (idx = 0; idx < 2; ++idx) { + if (mbmi_ext->ref_mv_count[ref_frame_type] > idx + 1) { + uint8_t drl_ctx = + av1_drl_ctx(mbmi_ext->ref_mv_stack[ref_frame_type], idx); + ++counts->drl_mode[drl_ctx][mbmi->ref_mv_idx != idx]; + + if (mbmi->ref_mv_idx == idx) break; + } + } + } + +#if CONFIG_EXT_INTER + if (have_nearmv_in_inter_mode(mbmi->mode)) { +#else + if (mbmi->mode == NEARMV) { +#endif + uint8_t ref_frame_type = av1_ref_frame_type(mbmi->ref_frame); + int idx; + + for (idx = 1; idx < 3; ++idx) { + if (mbmi_ext->ref_mv_count[ref_frame_type] > idx + 1) { + uint8_t drl_ctx = + av1_drl_ctx(mbmi_ext->ref_mv_stack[ref_frame_type], idx); + ++counts->drl_mode[drl_ctx][mbmi->ref_mv_idx != idx - 1]; + + if (mbmi->ref_mv_idx == idx - 1) break; + } + } + } +#else +#if CONFIG_EXT_INTER + if (is_inter_compound_mode(mode)) + ++counts->inter_compound_mode[mode_ctx][INTER_COMPOUND_OFFSET(mode)]; + else +#endif // CONFIG_EXT_INTER + ++counts->inter_mode[mode_ctx][INTER_OFFSET(mode)]; +#endif + } else { + const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize]; + const int num_4x4_h = num_4x4_blocks_high_lookup[bsize]; + int idx, idy; + for (idy = 0; idy < 2; idy += num_4x4_h) { + for (idx = 0; idx < 2; idx += num_4x4_w) { + const int j = idy * 2 + idx; + const PREDICTION_MODE b_mode = mi->bmi[j].as_mode; +#if CONFIG_REF_MV +#if CONFIG_EXT_INTER + if (has_second_ref(mbmi)) { + mode_ctx = mbmi_ext->compound_mode_context[mbmi->ref_frame[0]]; + ++counts->inter_compound_mode[mode_ctx] + [INTER_COMPOUND_OFFSET(b_mode)]; + } else { +#endif // CONFIG_EXT_INTER + mode_ctx = av1_mode_context_analyzer(mbmi_ext->mode_context, + mbmi->ref_frame, bsize, j); + update_inter_mode_stats(counts, b_mode, mode_ctx); +#if CONFIG_EXT_INTER + } +#endif // CONFIG_EXT_INTER +#else +#if CONFIG_EXT_INTER + if (is_inter_compound_mode(b_mode)) + ++counts->inter_compound_mode[mode_ctx] + [INTER_COMPOUND_OFFSET(b_mode)]; + else +#endif // CONFIG_EXT_INTER + ++counts->inter_mode[mode_ctx][INTER_OFFSET(b_mode)]; +#endif + } + } + } + } + } +} + +typedef struct { + ENTROPY_CONTEXT a[2 * MAX_MIB_SIZE * MAX_MB_PLANE]; + ENTROPY_CONTEXT l[2 * MAX_MIB_SIZE * MAX_MB_PLANE]; + PARTITION_CONTEXT sa[MAX_MIB_SIZE]; + PARTITION_CONTEXT sl[MAX_MIB_SIZE]; +#if CONFIG_VAR_TX + TXFM_CONTEXT *p_ta; + TXFM_CONTEXT *p_tl; + TXFM_CONTEXT ta[MAX_MIB_SIZE]; + TXFM_CONTEXT tl[MAX_MIB_SIZE]; +#endif +} RD_SEARCH_MACROBLOCK_CONTEXT; + +static void restore_context(MACROBLOCK *x, + const RD_SEARCH_MACROBLOCK_CONTEXT *ctx, int mi_row, + int mi_col, +#if CONFIG_PVQ + od_rollback_buffer *rdo_buf, +#endif + BLOCK_SIZE bsize) { + MACROBLOCKD *xd = &x->e_mbd; + int p; + const int num_4x4_blocks_wide = + block_size_wide[bsize] >> tx_size_wide_log2[0]; + const int num_4x4_blocks_high = + block_size_high[bsize] >> tx_size_high_log2[0]; + int mi_width = mi_size_wide[bsize]; + int mi_height = mi_size_high[bsize]; + for (p = 0; p < MAX_MB_PLANE; p++) { + memcpy(xd->above_context[p] + ((mi_col * 2) >> xd->plane[p].subsampling_x), + ctx->a + num_4x4_blocks_wide * p, + (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_wide) >> + xd->plane[p].subsampling_x); + memcpy(xd->left_context[p] + + ((mi_row & MAX_MIB_MASK) * 2 >> xd->plane[p].subsampling_y), + ctx->l + num_4x4_blocks_high * p, + (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_high) >> + xd->plane[p].subsampling_y); + } + memcpy(xd->above_seg_context + mi_col, ctx->sa, + sizeof(*xd->above_seg_context) * mi_width); + memcpy(xd->left_seg_context + (mi_row & MAX_MIB_MASK), ctx->sl, + sizeof(xd->left_seg_context[0]) * mi_height); +#if CONFIG_VAR_TX + xd->above_txfm_context = ctx->p_ta; + xd->left_txfm_context = ctx->p_tl; + memcpy(xd->above_txfm_context, ctx->ta, + sizeof(*xd->above_txfm_context) * mi_width); + memcpy(xd->left_txfm_context, ctx->tl, + sizeof(*xd->left_txfm_context) * mi_height); +#endif +#if CONFIG_PVQ + od_encode_rollback(&x->daala_enc, rdo_buf); +#endif +} + +static void save_context(const MACROBLOCK *x, RD_SEARCH_MACROBLOCK_CONTEXT *ctx, + int mi_row, int mi_col, +#if CONFIG_PVQ + od_rollback_buffer *rdo_buf, +#endif + BLOCK_SIZE bsize) { + const MACROBLOCKD *xd = &x->e_mbd; + int p; + const int num_4x4_blocks_wide = + block_size_wide[bsize] >> tx_size_wide_log2[0]; + const int num_4x4_blocks_high = + block_size_high[bsize] >> tx_size_high_log2[0]; + int mi_width = mi_size_wide[bsize]; + int mi_height = mi_size_high[bsize]; + + // buffer the above/left context information of the block in search. + for (p = 0; p < MAX_MB_PLANE; ++p) { + memcpy(ctx->a + num_4x4_blocks_wide * p, + xd->above_context[p] + (mi_col * 2 >> xd->plane[p].subsampling_x), + (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_wide) >> + xd->plane[p].subsampling_x); + memcpy(ctx->l + num_4x4_blocks_high * p, + xd->left_context[p] + + ((mi_row & MAX_MIB_MASK) * 2 >> xd->plane[p].subsampling_y), + (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_high) >> + xd->plane[p].subsampling_y); + } + memcpy(ctx->sa, xd->above_seg_context + mi_col, + sizeof(*xd->above_seg_context) * mi_width); + memcpy(ctx->sl, xd->left_seg_context + (mi_row & MAX_MIB_MASK), + sizeof(xd->left_seg_context[0]) * mi_height); +#if CONFIG_VAR_TX + memcpy(ctx->ta, xd->above_txfm_context, + sizeof(*xd->above_txfm_context) * mi_width); + memcpy(ctx->tl, xd->left_txfm_context, + sizeof(*xd->left_txfm_context) * mi_height); + ctx->p_ta = xd->above_txfm_context; + ctx->p_tl = xd->left_txfm_context; +#endif +#if CONFIG_PVQ + od_encode_checkpoint(&x->daala_enc, rdo_buf); +#endif +} + +static void encode_b(const AV1_COMP *const cpi, const TileInfo *const tile, + ThreadData *td, TOKENEXTRA **tp, int mi_row, int mi_col, + RUN_TYPE dry_run, BLOCK_SIZE bsize, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_TYPE partition, +#endif + PICK_MODE_CONTEXT *ctx, int *rate) { + MACROBLOCK *const x = &td->mb; +#if (CONFIG_MOTION_VAR && CONFIG_NCOBMC) | CONFIG_EXT_DELTA_Q + MACROBLOCKD *xd = &x->e_mbd; + MB_MODE_INFO *mbmi; +#if CONFIG_MOTION_VAR && CONFIG_NCOBMC + int check_ncobmc; +#endif +#endif + + set_offsets(cpi, tile, x, mi_row, mi_col, bsize); +#if CONFIG_EXT_PARTITION_TYPES + x->e_mbd.mi[0]->mbmi.partition = partition; +#endif + update_state(cpi, td, ctx, mi_row, mi_col, bsize, dry_run); +#if CONFIG_MOTION_VAR && CONFIG_NCOBMC + mbmi = &xd->mi[0]->mbmi; + const MOTION_MODE motion_allowed = motion_mode_allowed( +#if CONFIG_GLOBAL_MOTION && SEPARATE_GLOBAL_MOTION + 0, xd->global_motion, +#endif // CONFIG_GLOBAL_MOTION && SEPARATE_GLOBAL_MOTION + xd->mi[0]); + check_ncobmc = is_inter_block(mbmi) && motion_allowed >= OBMC_CAUSAL; + if (!dry_run && check_ncobmc) { + av1_check_ncobmc_rd(cpi, x, mi_row, mi_col); + av1_setup_dst_planes(x->e_mbd.plane, bsize, + get_frame_new_buffer(&cpi->common), mi_row, mi_col); + } +#endif + encode_superblock(cpi, td, tp, dry_run, mi_row, mi_col, bsize, ctx, rate); + + if (!dry_run) { +#if CONFIG_EXT_DELTA_Q + mbmi = &xd->mi[0]->mbmi; + if (bsize == BLOCK_64X64 && mbmi->skip == 1 && is_inter_block(mbmi) && + cpi->common.delta_lf_present_flag) { + mbmi->current_delta_lf_from_base = xd->prev_delta_lf_from_base; + } +#endif +#if CONFIG_SUPERTX + update_stats(&cpi->common, td, mi_row, mi_col, 0); +#else + update_stats(&cpi->common, td, mi_row, mi_col); +#endif + } +} + +static void encode_sb(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, TOKENEXTRA **tp, int mi_row, + int mi_col, RUN_TYPE dry_run, BLOCK_SIZE bsize, + PC_TREE *pc_tree, int *rate) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + const int hbs = mi_size_wide[bsize] / 2; + const int is_partition_root = bsize >= BLOCK_8X8; + const int ctx = is_partition_root + ? partition_plane_context(xd, mi_row, mi_col, +#if CONFIG_UNPOISON_PARTITION_CTX + mi_row + hbs < cm->mi_rows, + mi_col + hbs < cm->mi_cols, +#endif + bsize) + : -1; + const PARTITION_TYPE partition = pc_tree->partitioning; + const BLOCK_SIZE subsize = get_subsize(bsize, partition); +#if CONFIG_EXT_PARTITION_TYPES + const BLOCK_SIZE bsize2 = get_subsize(bsize, PARTITION_SPLIT); +#endif + +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; + assert(bsize >= BLOCK_8X8); +#endif + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + + if (!dry_run && ctx >= 0) td->counts->partition[ctx][partition]++; + +#if CONFIG_SUPERTX + if (!frame_is_intra_only(cm) && bsize <= MAX_SUPERTX_BLOCK_SIZE && + partition != PARTITION_NONE && !xd->lossless[0]) { + int supertx_enabled; + TX_SIZE supertx_size = max_txsize_lookup[bsize]; + supertx_enabled = check_supertx_sb(bsize, supertx_size, pc_tree); + if (supertx_enabled) { + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + int x_idx, y_idx, i; + uint8_t *dst_buf[3]; + int dst_stride[3]; + set_skip_context(xd, mi_row, mi_col); + set_mode_info_offsets(cpi, x, xd, mi_row, mi_col); + update_state_sb_supertx(cpi, td, tile, mi_row, mi_col, bsize, dry_run, + pc_tree); + + av1_setup_dst_planes(xd->plane, bsize, get_frame_new_buffer(cm), mi_row, + mi_col); + for (i = 0; i < MAX_MB_PLANE; i++) { + dst_buf[i] = xd->plane[i].dst.buf; + dst_stride[i] = xd->plane[i].dst.stride; + } + predict_sb_complex(cpi, td, tile, mi_row, mi_col, mi_row, mi_col, dry_run, + bsize, bsize, dst_buf, dst_stride, pc_tree); + + set_offsets_without_segment_id(cpi, tile, x, mi_row, mi_col, bsize); + set_segment_id_supertx(cpi, x, mi_row, mi_col, bsize); + + if (!x->skip) { + int this_rate = 0; + av1_encode_sb_supertx((AV1_COMMON *)cm, x, bsize); + av1_tokenize_sb_supertx(cpi, td, tp, dry_run, bsize, rate); + if (rate) *rate += this_rate; + } else { + xd->mi[0]->mbmi.skip = 1; + if (!dry_run) td->counts->skip[av1_get_skip_context(xd)][1]++; + reset_skip_context(xd, bsize); + } + if (!dry_run) { + for (y_idx = 0; y_idx < mi_height; y_idx++) + for (x_idx = 0; x_idx < mi_width; x_idx++) { + if ((xd->mb_to_right_edge >> (3 + MI_SIZE_LOG2)) + mi_width > + x_idx && + (xd->mb_to_bottom_edge >> (3 + MI_SIZE_LOG2)) + mi_height > + y_idx) { + xd->mi[x_idx + y_idx * cm->mi_stride]->mbmi.skip = + xd->mi[0]->mbmi.skip; + } + } + td->counts->supertx[partition_supertx_context_lookup[partition]] + [supertx_size][1]++; + td->counts->supertx_size[supertx_size]++; +#if CONFIG_EXT_TX + if (get_ext_tx_types(supertx_size, bsize, 1, cm->reduced_tx_set_used) > + 1 && + !xd->mi[0]->mbmi.skip) { + const int eset = + get_ext_tx_set(supertx_size, bsize, 1, cm->reduced_tx_set_used); + if (eset > 0) { + ++td->counts + ->inter_ext_tx[eset][supertx_size][xd->mi[0]->mbmi.tx_type]; + } + } +#else + if (supertx_size < TX_32X32 && !xd->mi[0]->mbmi.skip) { + ++td->counts->inter_ext_tx[supertx_size][xd->mi[0]->mbmi.tx_type]; + } +#endif // CONFIG_EXT_TX + } +#if CONFIG_EXT_PARTITION_TYPES + update_ext_partition_context(xd, mi_row, mi_col, subsize, bsize, + partition); +#else + if (partition != PARTITION_SPLIT || bsize == BLOCK_8X8) + update_partition_context(xd, mi_row, mi_col, subsize, bsize); +#endif +#if CONFIG_VAR_TX + set_txfm_ctxs(supertx_size, mi_width, mi_height, xd->mi[0]->mbmi.skip, + xd); +#endif // CONFIG_VAR_TX + return; + } else { + if (!dry_run) { + td->counts->supertx[partition_supertx_context_lookup[partition]] + [supertx_size][0]++; + } + } + } +#endif // CONFIG_SUPERTX + + switch (partition) { + case PARTITION_NONE: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, subsize, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + &pc_tree->none, rate); + break; + case PARTITION_VERT: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, subsize, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + &pc_tree->vertical[0], rate); + if (mi_col + hbs < cm->mi_cols && (bsize > BLOCK_8X8 || unify_bsize)) { + encode_b(cpi, tile, td, tp, mi_row, mi_col + hbs, dry_run, subsize, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + &pc_tree->vertical[1], rate); + } + break; + case PARTITION_HORZ: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, subsize, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + &pc_tree->horizontal[0], rate); + if (mi_row + hbs < cm->mi_rows && (bsize > BLOCK_8X8 || unify_bsize)) { + encode_b(cpi, tile, td, tp, mi_row + hbs, mi_col, dry_run, subsize, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + &pc_tree->horizontal[1], rate); + } + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, subsize, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + pc_tree->leaf_split[0], rate); + } else { + encode_sb(cpi, td, tile, tp, mi_row, mi_col, dry_run, subsize, + pc_tree->split[0], rate); + encode_sb(cpi, td, tile, tp, mi_row, mi_col + hbs, dry_run, subsize, + pc_tree->split[1], rate); + encode_sb(cpi, td, tile, tp, mi_row + hbs, mi_col, dry_run, subsize, + pc_tree->split[2], rate); + encode_sb(cpi, td, tile, tp, mi_row + hbs, mi_col + hbs, dry_run, + subsize, pc_tree->split[3], rate); + } + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, bsize2, partition, + &pc_tree->horizontala[0], rate); + encode_b(cpi, tile, td, tp, mi_row, mi_col + hbs, dry_run, bsize2, + partition, &pc_tree->horizontala[1], rate); + encode_b(cpi, tile, td, tp, mi_row + hbs, mi_col, dry_run, subsize, + partition, &pc_tree->horizontala[2], rate); + break; + case PARTITION_HORZ_B: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, subsize, partition, + &pc_tree->horizontalb[0], rate); + encode_b(cpi, tile, td, tp, mi_row + hbs, mi_col, dry_run, bsize2, + partition, &pc_tree->horizontalb[1], rate); + encode_b(cpi, tile, td, tp, mi_row + hbs, mi_col + hbs, dry_run, bsize2, + partition, &pc_tree->horizontalb[2], rate); + break; + case PARTITION_VERT_A: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, bsize2, partition, + &pc_tree->verticala[0], rate); + encode_b(cpi, tile, td, tp, mi_row + hbs, mi_col, dry_run, bsize2, + partition, &pc_tree->verticala[1], rate); + encode_b(cpi, tile, td, tp, mi_row, mi_col + hbs, dry_run, subsize, + partition, &pc_tree->verticala[2], rate); + + break; + case PARTITION_VERT_B: + encode_b(cpi, tile, td, tp, mi_row, mi_col, dry_run, subsize, partition, + &pc_tree->verticalb[0], rate); + encode_b(cpi, tile, td, tp, mi_row, mi_col + hbs, dry_run, bsize2, + partition, &pc_tree->verticalb[1], rate); + encode_b(cpi, tile, td, tp, mi_row + hbs, mi_col + hbs, dry_run, bsize2, + partition, &pc_tree->verticalb[2], rate); + break; +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0 && "Invalid partition type."); break; + } + +#if CONFIG_EXT_PARTITION_TYPES + update_ext_partition_context(xd, mi_row, mi_col, subsize, bsize, partition); +#else + if (partition != PARTITION_SPLIT || bsize == BLOCK_8X8) + update_partition_context(xd, mi_row, mi_col, subsize, bsize); +#endif // CONFIG_EXT_PARTITION_TYPES +} + +// Check to see if the given partition size is allowed for a specified number +// of mi block rows and columns remaining in the image. +// If not then return the largest allowed partition size +static BLOCK_SIZE find_partition_size(BLOCK_SIZE bsize, int rows_left, + int cols_left, int *bh, int *bw) { + if (rows_left <= 0 || cols_left <= 0) { + return AOMMIN(bsize, BLOCK_8X8); + } else { + for (; bsize > 0; bsize -= 3) { + *bh = mi_size_high[bsize]; + *bw = mi_size_wide[bsize]; + if ((*bh <= rows_left) && (*bw <= cols_left)) { + break; + } + } + } + return bsize; +} + +static void set_partial_sb_partition(const AV1_COMMON *const cm, MODE_INFO *mi, + int bh_in, int bw_in, + int mi_rows_remaining, + int mi_cols_remaining, BLOCK_SIZE bsize, + MODE_INFO **mib) { + int bh = bh_in; + int r, c; + for (r = 0; r < cm->mib_size; r += bh) { + int bw = bw_in; + for (c = 0; c < cm->mib_size; c += bw) { + const int index = r * cm->mi_stride + c; + mib[index] = mi + index; + mib[index]->mbmi.sb_type = find_partition_size( + bsize, mi_rows_remaining - r, mi_cols_remaining - c, &bh, &bw); + } + } +} + +// This function attempts to set all mode info entries in a given superblock +// to the same block partition size. +// However, at the bottom and right borders of the image the requested size +// may not be allowed in which case this code attempts to choose the largest +// allowable partition. +static void set_fixed_partitioning(AV1_COMP *cpi, const TileInfo *const tile, + MODE_INFO **mib, int mi_row, int mi_col, + BLOCK_SIZE bsize) { + AV1_COMMON *const cm = &cpi->common; + const int mi_rows_remaining = tile->mi_row_end - mi_row; + const int mi_cols_remaining = tile->mi_col_end - mi_col; + int block_row, block_col; + MODE_INFO *const mi_upper_left = cm->mi + mi_row * cm->mi_stride + mi_col; + int bh = mi_size_high[bsize]; + int bw = mi_size_wide[bsize]; + + assert((mi_rows_remaining > 0) && (mi_cols_remaining > 0)); + + // Apply the requested partition size to the SB if it is all "in image" + if ((mi_cols_remaining >= cm->mib_size) && + (mi_rows_remaining >= cm->mib_size)) { + for (block_row = 0; block_row < cm->mib_size; block_row += bh) { + for (block_col = 0; block_col < cm->mib_size; block_col += bw) { + int index = block_row * cm->mi_stride + block_col; + mib[index] = mi_upper_left + index; + mib[index]->mbmi.sb_type = bsize; + } + } + } else { + // Else this is a partial SB. + set_partial_sb_partition(cm, mi_upper_left, bh, bw, mi_rows_remaining, + mi_cols_remaining, bsize, mib); + } +} + +static void rd_use_partition(AV1_COMP *cpi, ThreadData *td, + TileDataEnc *tile_data, MODE_INFO **mib, + TOKENEXTRA **tp, int mi_row, int mi_col, + BLOCK_SIZE bsize, int *rate, int64_t *dist, +#if CONFIG_SUPERTX + int *rate_nocoef, +#endif + int do_recon, PC_TREE *pc_tree) { + AV1_COMMON *const cm = &cpi->common; + TileInfo *const tile_info = &tile_data->tile_info; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + const int bs = mi_size_wide[bsize]; + const int hbs = bs / 2; + int i; + const int pl = (bsize >= BLOCK_8X8) + ? partition_plane_context(xd, mi_row, mi_col, +#if CONFIG_UNPOISON_PARTITION_CTX + mi_row + hbs < cm->mi_rows, + mi_col + hbs < cm->mi_cols, +#endif + bsize) + : 0; + const PARTITION_TYPE partition = + (bsize >= BLOCK_8X8) ? get_partition(cm, mi_row, mi_col, bsize) + : PARTITION_NONE; + const BLOCK_SIZE subsize = get_subsize(bsize, partition); + RD_SEARCH_MACROBLOCK_CONTEXT x_ctx; + RD_STATS last_part_rdc, none_rdc, chosen_rdc; + BLOCK_SIZE sub_subsize = BLOCK_4X4; + int splits_below = 0; + BLOCK_SIZE bs_type = mib[0]->mbmi.sb_type; + int do_partition_search = 1; + PICK_MODE_CONTEXT *ctx_none = &pc_tree->none; + const int unify_bsize = CONFIG_CB4X4; +#if CONFIG_SUPERTX + int last_part_rate_nocoef = INT_MAX; + int none_rate_nocoef = INT_MAX; + int chosen_rate_nocoef = INT_MAX; +#endif +#if CONFIG_PVQ + od_rollback_buffer pre_rdo_buf; +#endif + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + + assert(num_4x4_blocks_wide_lookup[bsize] == + num_4x4_blocks_high_lookup[bsize]); + + av1_invalid_rd_stats(&last_part_rdc); + av1_invalid_rd_stats(&none_rdc); + av1_invalid_rd_stats(&chosen_rdc); + + pc_tree->partitioning = partition; + +#if CONFIG_VAR_TX + xd->above_txfm_context = cm->above_txfm_context + mi_col; + xd->left_txfm_context = + xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK); +#endif +#if !CONFIG_PVQ + save_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + save_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + + if (bsize == BLOCK_16X16 && cpi->vaq_refresh) { + set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize); + x->mb_energy = av1_block_energy(cpi, x, bsize); + } + + if (do_partition_search && + cpi->sf.partition_search_type == SEARCH_PARTITION && + cpi->sf.adjust_partitioning_from_last_frame) { + // Check if any of the sub blocks are further split. + if (partition == PARTITION_SPLIT && subsize > BLOCK_8X8) { + sub_subsize = get_subsize(subsize, PARTITION_SPLIT); + splits_below = 1; + for (i = 0; i < 4; i++) { + int jj = i >> 1, ii = i & 0x01; + MODE_INFO *this_mi = mib[jj * hbs * cm->mi_stride + ii * hbs]; + if (this_mi && this_mi->mbmi.sb_type >= sub_subsize) { + splits_below = 0; + } + } + } + + // If partition is not none try none unless each of the 4 splits are split + // even further.. + if (partition != PARTITION_NONE && !splits_below && + mi_row + hbs < cm->mi_rows && mi_col + hbs < cm->mi_cols) { + pc_tree->partitioning = PARTITION_NONE; + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &none_rdc, +#if CONFIG_SUPERTX + &none_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_NONE, +#endif + bsize, ctx_none, INT64_MAX); + + if (none_rdc.rate < INT_MAX) { + none_rdc.rate += cpi->partition_cost[pl][PARTITION_NONE]; + none_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, none_rdc.rate, none_rdc.dist); +#if CONFIG_SUPERTX + none_rate_nocoef += cpi->partition_cost[pl][PARTITION_NONE]; +#endif + } + +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + mib[0]->mbmi.sb_type = bs_type; + pc_tree->partitioning = partition; + } + } + + switch (partition) { + case PARTITION_NONE: + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc, +#if CONFIG_SUPERTX + &last_part_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_NONE, +#endif + bsize, ctx_none, INT64_MAX); + break; + case PARTITION_HORZ: + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc, +#if CONFIG_SUPERTX + &last_part_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_HORZ, +#endif + subsize, &pc_tree->horizontal[0], INT64_MAX); + if (last_part_rdc.rate != INT_MAX && bsize >= BLOCK_8X8 && + mi_row + hbs < cm->mi_rows) { + RD_STATS tmp_rdc; +#if CONFIG_SUPERTX + int rt_nocoef = 0; +#endif + PICK_MODE_CONTEXT *ctx_h = &pc_tree->horizontal[0]; + av1_init_rd_stats(&tmp_rdc); + update_state(cpi, td, ctx_h, mi_row, mi_col, subsize, 1); + encode_superblock(cpi, td, tp, DRY_RUN_NORMAL, mi_row, mi_col, subsize, + ctx_h, NULL); + rd_pick_sb_modes(cpi, tile_data, x, mi_row + hbs, mi_col, &tmp_rdc, +#if CONFIG_SUPERTX + &rt_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_HORZ, +#endif + subsize, &pc_tree->horizontal[1], INT64_MAX); + if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) { + av1_invalid_rd_stats(&last_part_rdc); +#if CONFIG_SUPERTX + last_part_rate_nocoef = INT_MAX; +#endif + break; + } + last_part_rdc.rate += tmp_rdc.rate; + last_part_rdc.dist += tmp_rdc.dist; + last_part_rdc.rdcost += tmp_rdc.rdcost; +#if CONFIG_SUPERTX + last_part_rate_nocoef += rt_nocoef; +#endif + } + break; + case PARTITION_VERT: + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc, +#if CONFIG_SUPERTX + &last_part_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_VERT, +#endif + subsize, &pc_tree->vertical[0], INT64_MAX); + if (last_part_rdc.rate != INT_MAX && bsize >= BLOCK_8X8 && + mi_col + hbs < cm->mi_cols) { + RD_STATS tmp_rdc; +#if CONFIG_SUPERTX + int rt_nocoef = 0; +#endif + PICK_MODE_CONTEXT *ctx_v = &pc_tree->vertical[0]; + av1_init_rd_stats(&tmp_rdc); + update_state(cpi, td, ctx_v, mi_row, mi_col, subsize, 1); + encode_superblock(cpi, td, tp, DRY_RUN_NORMAL, mi_row, mi_col, subsize, + ctx_v, NULL); + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col + hbs, &tmp_rdc, +#if CONFIG_SUPERTX + &rt_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_VERT, +#endif + subsize, &pc_tree->vertical[bsize > BLOCK_8X8], + INT64_MAX); + if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) { + av1_invalid_rd_stats(&last_part_rdc); +#if CONFIG_SUPERTX + last_part_rate_nocoef = INT_MAX; +#endif + break; + } + last_part_rdc.rate += tmp_rdc.rate; + last_part_rdc.dist += tmp_rdc.dist; + last_part_rdc.rdcost += tmp_rdc.rdcost; +#if CONFIG_SUPERTX + last_part_rate_nocoef += rt_nocoef; +#endif + } + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc, +#if CONFIG_SUPERTX + &last_part_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_SPLIT, +#endif + subsize, pc_tree->leaf_split[0], INT64_MAX); + break; + } + last_part_rdc.rate = 0; + last_part_rdc.dist = 0; + last_part_rdc.rdcost = 0; +#if CONFIG_SUPERTX + last_part_rate_nocoef = 0; +#endif + for (i = 0; i < 4; i++) { + int x_idx = (i & 1) * hbs; + int y_idx = (i >> 1) * hbs; + int jj = i >> 1, ii = i & 0x01; + RD_STATS tmp_rdc; +#if CONFIG_SUPERTX + int rt_nocoef; +#endif + if ((mi_row + y_idx >= cm->mi_rows) || (mi_col + x_idx >= cm->mi_cols)) + continue; + + av1_init_rd_stats(&tmp_rdc); + rd_use_partition(cpi, td, tile_data, + mib + jj * hbs * cm->mi_stride + ii * hbs, tp, + mi_row + y_idx, mi_col + x_idx, subsize, &tmp_rdc.rate, + &tmp_rdc.dist, +#if CONFIG_SUPERTX + &rt_nocoef, +#endif + i != 3, pc_tree->split[i]); + if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) { + av1_invalid_rd_stats(&last_part_rdc); +#if CONFIG_SUPERTX + last_part_rate_nocoef = INT_MAX; +#endif + break; + } + last_part_rdc.rate += tmp_rdc.rate; + last_part_rdc.dist += tmp_rdc.dist; +#if CONFIG_SUPERTX + last_part_rate_nocoef += rt_nocoef; +#endif + } + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_VERT_A: + case PARTITION_VERT_B: + case PARTITION_HORZ_A: + case PARTITION_HORZ_B: assert(0 && "Cannot handle extended partiton types"); +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0); break; + } + + if (last_part_rdc.rate < INT_MAX) { + last_part_rdc.rate += cpi->partition_cost[pl][partition]; + last_part_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, last_part_rdc.rate, last_part_rdc.dist); +#if CONFIG_SUPERTX + last_part_rate_nocoef += cpi->partition_cost[pl][partition]; +#endif + } + + if (do_partition_search && cpi->sf.adjust_partitioning_from_last_frame && + cpi->sf.partition_search_type == SEARCH_PARTITION && + partition != PARTITION_SPLIT && bsize > BLOCK_8X8 && + (mi_row + bs < cm->mi_rows || mi_row + hbs == cm->mi_rows) && + (mi_col + bs < cm->mi_cols || mi_col + hbs == cm->mi_cols)) { + BLOCK_SIZE split_subsize = get_subsize(bsize, PARTITION_SPLIT); + chosen_rdc.rate = 0; + chosen_rdc.dist = 0; +#if CONFIG_SUPERTX + chosen_rate_nocoef = 0; +#endif +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + pc_tree->partitioning = PARTITION_SPLIT; + + // Split partition. + for (i = 0; i < 4; i++) { + int x_idx = (i & 1) * hbs; + int y_idx = (i >> 1) * hbs; + RD_STATS tmp_rdc; +#if CONFIG_SUPERTX + int rt_nocoef = 0; +#endif +#if CONFIG_PVQ + od_rollback_buffer buf; +#endif + if ((mi_row + y_idx >= cm->mi_rows) || (mi_col + x_idx >= cm->mi_cols)) + continue; + +#if !CONFIG_PVQ + save_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + save_context(x, &x_ctx, mi_row, mi_col, &buf, bsize); +#endif + pc_tree->split[i]->partitioning = PARTITION_NONE; + rd_pick_sb_modes(cpi, tile_data, x, mi_row + y_idx, mi_col + x_idx, + &tmp_rdc, +#if CONFIG_SUPERTX + &rt_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_SPLIT, +#endif + split_subsize, &pc_tree->split[i]->none, INT64_MAX); + +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &buf, bsize); +#endif + if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) { + av1_invalid_rd_stats(&chosen_rdc); +#if CONFIG_SUPERTX + chosen_rate_nocoef = INT_MAX; +#endif + break; + } + + chosen_rdc.rate += tmp_rdc.rate; + chosen_rdc.dist += tmp_rdc.dist; +#if CONFIG_SUPERTX + chosen_rate_nocoef += rt_nocoef; +#endif + + if (i != 3) + encode_sb(cpi, td, tile_info, tp, mi_row + y_idx, mi_col + x_idx, + OUTPUT_ENABLED, split_subsize, pc_tree->split[i], NULL); + + chosen_rdc.rate += cpi->partition_cost[pl][PARTITION_NONE]; +#if CONFIG_SUPERTX + chosen_rate_nocoef += cpi->partition_cost[pl][PARTITION_SPLIT]; +#endif + } + if (chosen_rdc.rate < INT_MAX) { + chosen_rdc.rate += cpi->partition_cost[pl][PARTITION_SPLIT]; + chosen_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, chosen_rdc.rate, chosen_rdc.dist); +#if CONFIG_SUPERTX + chosen_rate_nocoef += cpi->partition_cost[pl][PARTITION_NONE]; +#endif + } + } + + // If last_part is better set the partitioning to that. + if (last_part_rdc.rdcost < chosen_rdc.rdcost) { + mib[0]->mbmi.sb_type = bsize; + if (bsize >= BLOCK_8X8) pc_tree->partitioning = partition; + chosen_rdc = last_part_rdc; +#if CONFIG_SUPERTX + chosen_rate_nocoef = last_part_rate_nocoef; +#endif + } + // If none was better set the partitioning to that. + if (none_rdc.rdcost < chosen_rdc.rdcost) { + if (bsize >= BLOCK_8X8) pc_tree->partitioning = PARTITION_NONE; + chosen_rdc = none_rdc; +#if CONFIG_SUPERTX + chosen_rate_nocoef = none_rate_nocoef; +#endif + } + +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + + // We must have chosen a partitioning and encoding or we'll fail later on. + // No other opportunities for success. + if (bsize == cm->sb_size) + assert(chosen_rdc.rate < INT_MAX && chosen_rdc.dist < INT64_MAX); + + if (do_recon) { + if (bsize == cm->sb_size) { + // NOTE: To get estimate for rate due to the tokens, use: + // int rate_coeffs = 0; + // encode_sb(cpi, td, tile_info, tp, mi_row, mi_col, DRY_RUN_COSTCOEFFS, + // bsize, pc_tree, &rate_coeffs); + encode_sb(cpi, td, tile_info, tp, mi_row, mi_col, OUTPUT_ENABLED, bsize, + pc_tree, NULL); + } else { + encode_sb(cpi, td, tile_info, tp, mi_row, mi_col, DRY_RUN_NORMAL, bsize, + pc_tree, NULL); + } + } + + *rate = chosen_rdc.rate; + *dist = chosen_rdc.dist; +#if CONFIG_SUPERTX + *rate_nocoef = chosen_rate_nocoef; +#endif +} + +/* clang-format off */ +static const BLOCK_SIZE min_partition_size[BLOCK_SIZES] = { +#if CONFIG_CB4X4 + BLOCK_2X2, BLOCK_2X2, BLOCK_2X2, // 2x2, 2x4, 4x2 +#endif + BLOCK_4X4, // 4x4 + BLOCK_4X4, BLOCK_4X4, BLOCK_4X4, // 4x8, 8x4, 8x8 + BLOCK_4X4, BLOCK_4X4, BLOCK_8X8, // 8x16, 16x8, 16x16 + BLOCK_8X8, BLOCK_8X8, BLOCK_16X16, // 16x32, 32x16, 32x32 + BLOCK_16X16, BLOCK_16X16, BLOCK_16X16, // 32x64, 64x32, 64x64 +#if CONFIG_EXT_PARTITION + BLOCK_16X16, BLOCK_16X16, BLOCK_16X16 // 64x128, 128x64, 128x128 +#endif // CONFIG_EXT_PARTITION +}; + +static const BLOCK_SIZE max_partition_size[BLOCK_SIZES] = { +#if CONFIG_CB4X4 + BLOCK_4X4, BLOCK_4X4, BLOCK_4X4, // 2x2, 2x4, 4x2 +#endif + BLOCK_8X8, // 4x4 + BLOCK_16X16, BLOCK_16X16, BLOCK_16X16, // 4x8, 8x4, 8x8 + BLOCK_32X32, BLOCK_32X32, BLOCK_32X32, // 8x16, 16x8, 16x16 + BLOCK_64X64, BLOCK_64X64, BLOCK_64X64, // 16x32, 32x16, 32x32 + BLOCK_LARGEST, BLOCK_LARGEST, BLOCK_LARGEST, // 32x64, 64x32, 64x64 +#if CONFIG_EXT_PARTITION + BLOCK_LARGEST, BLOCK_LARGEST, BLOCK_LARGEST // 64x128, 128x64, 128x128 +#endif // CONFIG_EXT_PARTITION +}; + +// Next square block size less or equal than current block size. +static const BLOCK_SIZE next_square_size[BLOCK_SIZES] = { +#if CONFIG_CB4X4 + BLOCK_2X2, BLOCK_2X2, BLOCK_2X2, // 2x2, 2x4, 4x2 +#endif + BLOCK_4X4, // 4x4 + BLOCK_4X4, BLOCK_4X4, BLOCK_8X8, // 4x8, 8x4, 8x8 + BLOCK_8X8, BLOCK_8X8, BLOCK_16X16, // 8x16, 16x8, 16x16 + BLOCK_16X16, BLOCK_16X16, BLOCK_32X32, // 16x32, 32x16, 32x32 + BLOCK_32X32, BLOCK_32X32, BLOCK_64X64, // 32x64, 64x32, 64x64 +#if CONFIG_EXT_PARTITION + BLOCK_64X64, BLOCK_64X64, BLOCK_128X128 // 64x128, 128x64, 128x128 +#endif // CONFIG_EXT_PARTITION +}; +/* clang-format on */ + +// Look at all the mode_info entries for blocks that are part of this +// partition and find the min and max values for sb_type. +// At the moment this is designed to work on a superblock but could be +// adjusted to use a size parameter. +// +// The min and max are assumed to have been initialized prior to calling this +// function so repeat calls can accumulate a min and max of more than one +// superblock. +static void get_sb_partition_size_range(const AV1_COMMON *const cm, + MACROBLOCKD *xd, MODE_INFO **mib, + BLOCK_SIZE *min_block_size, + BLOCK_SIZE *max_block_size) { + int i, j; + int index = 0; + + // Check the sb_type for each block that belongs to this region. + for (i = 0; i < cm->mib_size; ++i) { + for (j = 0; j < cm->mib_size; ++j) { + MODE_INFO *mi = mib[index + j]; + BLOCK_SIZE sb_type = mi ? mi->mbmi.sb_type : BLOCK_4X4; + *min_block_size = AOMMIN(*min_block_size, sb_type); + *max_block_size = AOMMAX(*max_block_size, sb_type); + } + index += xd->mi_stride; + } +} + +// Look at neighboring blocks and set a min and max partition size based on +// what they chose. +static void rd_auto_partition_range(AV1_COMP *cpi, const TileInfo *const tile, + MACROBLOCKD *const xd, int mi_row, + int mi_col, BLOCK_SIZE *min_block_size, + BLOCK_SIZE *max_block_size) { + AV1_COMMON *const cm = &cpi->common; + MODE_INFO **mi = xd->mi; + const int left_in_image = xd->left_available && mi[-1]; + const int above_in_image = xd->up_available && mi[-xd->mi_stride]; + const int mi_rows_remaining = tile->mi_row_end - mi_row; + const int mi_cols_remaining = tile->mi_col_end - mi_col; + int bh, bw; + BLOCK_SIZE min_size = BLOCK_4X4; + BLOCK_SIZE max_size = BLOCK_LARGEST; + + // Trap case where we do not have a prediction. + if (left_in_image || above_in_image || cm->frame_type != KEY_FRAME) { + // Default "min to max" and "max to min" + min_size = BLOCK_LARGEST; + max_size = BLOCK_4X4; + + // NOTE: each call to get_sb_partition_size_range() uses the previous + // passed in values for min and max as a starting point. + // Find the min and max partition used in previous frame at this location + if (cm->frame_type != KEY_FRAME) { + MODE_INFO **prev_mi = + &cm->prev_mi_grid_visible[mi_row * xd->mi_stride + mi_col]; + get_sb_partition_size_range(cm, xd, prev_mi, &min_size, &max_size); + } + // Find the min and max partition sizes used in the left superblock + if (left_in_image) { + MODE_INFO **left_sb_mi = &mi[-cm->mib_size]; + get_sb_partition_size_range(cm, xd, left_sb_mi, &min_size, &max_size); + } + // Find the min and max partition sizes used in the above suprblock. + if (above_in_image) { + MODE_INFO **above_sb_mi = &mi[-xd->mi_stride * cm->mib_size]; + get_sb_partition_size_range(cm, xd, above_sb_mi, &min_size, &max_size); + } + + // Adjust observed min and max for "relaxed" auto partition case. + if (cpi->sf.auto_min_max_partition_size == RELAXED_NEIGHBORING_MIN_MAX) { + min_size = min_partition_size[min_size]; + max_size = max_partition_size[max_size]; + } + } + + // Check border cases where max and min from neighbors may not be legal. + max_size = find_partition_size(max_size, mi_rows_remaining, mi_cols_remaining, + &bh, &bw); + min_size = AOMMIN(min_size, max_size); + + // Test for blocks at the edge of the active image. + // This may be the actual edge of the image or where there are formatting + // bars. + if (av1_active_edge_sb(cpi, mi_row, mi_col)) { + min_size = BLOCK_4X4; + } else { + min_size = AOMMIN(cpi->sf.rd_auto_partition_min_limit, min_size); + } + + // When use_square_partition_only is true, make sure at least one square + // partition is allowed by selecting the next smaller square size as + // *min_block_size. + if (cpi->sf.use_square_partition_only) { + min_size = AOMMIN(min_size, next_square_size[max_size]); + } + + *min_block_size = AOMMIN(min_size, cm->sb_size); + *max_block_size = AOMMIN(max_size, cm->sb_size); +} + +// TODO(jingning) refactor functions setting partition search range +static void set_partition_range(const AV1_COMMON *const cm, + const MACROBLOCKD *const xd, int mi_row, + int mi_col, BLOCK_SIZE bsize, + BLOCK_SIZE *const min_bs, + BLOCK_SIZE *const max_bs) { + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + int idx, idy; + + const int idx_str = cm->mi_stride * mi_row + mi_col; + MODE_INFO **const prev_mi = &cm->prev_mi_grid_visible[idx_str]; + BLOCK_SIZE min_size = BLOCK_64X64; // default values + BLOCK_SIZE max_size = BLOCK_4X4; + + if (prev_mi) { + for (idy = 0; idy < mi_height; ++idy) { + for (idx = 0; idx < mi_width; ++idx) { + const MODE_INFO *const mi = prev_mi[idy * cm->mi_stride + idx]; + const BLOCK_SIZE bs = mi ? mi->mbmi.sb_type : bsize; + min_size = AOMMIN(min_size, bs); + max_size = AOMMAX(max_size, bs); + } + } + } + + if (xd->left_available) { + for (idy = 0; idy < mi_height; ++idy) { + const MODE_INFO *const mi = xd->mi[idy * cm->mi_stride - 1]; + const BLOCK_SIZE bs = mi ? mi->mbmi.sb_type : bsize; + min_size = AOMMIN(min_size, bs); + max_size = AOMMAX(max_size, bs); + } + } + + if (xd->up_available) { + for (idx = 0; idx < mi_width; ++idx) { + const MODE_INFO *const mi = xd->mi[idx - cm->mi_stride]; + const BLOCK_SIZE bs = mi ? mi->mbmi.sb_type : bsize; + min_size = AOMMIN(min_size, bs); + max_size = AOMMAX(max_size, bs); + } + } + + if (min_size == max_size) { + min_size = min_partition_size[min_size]; + max_size = max_partition_size[max_size]; + } + + *min_bs = AOMMIN(min_size, cm->sb_size); + *max_bs = AOMMIN(max_size, cm->sb_size); +} + +static INLINE void store_pred_mv(MACROBLOCK *x, PICK_MODE_CONTEXT *ctx) { + memcpy(ctx->pred_mv, x->pred_mv, sizeof(x->pred_mv)); +} + +static INLINE void load_pred_mv(MACROBLOCK *x, PICK_MODE_CONTEXT *ctx) { + memcpy(x->pred_mv, ctx->pred_mv, sizeof(x->pred_mv)); +} + +#if CONFIG_FP_MB_STATS +const int qindex_skip_threshold_lookup[BLOCK_SIZES] = { + 0, + 10, + 10, + 30, + 40, + 40, + 60, + 80, + 80, + 90, + 100, + 100, + 120, +#if CONFIG_EXT_PARTITION + // TODO(debargha): What are the correct numbers here? + 130, + 130, + 150 +#endif // CONFIG_EXT_PARTITION +}; +const int qindex_split_threshold_lookup[BLOCK_SIZES] = { + 0, + 3, + 3, + 7, + 15, + 15, + 30, + 40, + 40, + 60, + 80, + 80, + 120, +#if CONFIG_EXT_PARTITION + // TODO(debargha): What are the correct numbers here? + 160, + 160, + 240 +#endif // CONFIG_EXT_PARTITION +}; +const int complexity_16x16_blocks_threshold[BLOCK_SIZES] = { + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 4, + 4, + 6, +#if CONFIG_EXT_PARTITION + // TODO(debargha): What are the correct numbers here? + 8, + 8, + 10 +#endif // CONFIG_EXT_PARTITION +}; + +typedef enum { + MV_ZERO = 0, + MV_LEFT = 1, + MV_UP = 2, + MV_RIGHT = 3, + MV_DOWN = 4, + MV_INVALID +} MOTION_DIRECTION; + +static INLINE MOTION_DIRECTION get_motion_direction_fp(uint8_t fp_byte) { + if (fp_byte & FPMB_MOTION_ZERO_MASK) { + return MV_ZERO; + } else if (fp_byte & FPMB_MOTION_LEFT_MASK) { + return MV_LEFT; + } else if (fp_byte & FPMB_MOTION_RIGHT_MASK) { + return MV_RIGHT; + } else if (fp_byte & FPMB_MOTION_UP_MASK) { + return MV_UP; + } else { + return MV_DOWN; + } +} + +static INLINE int get_motion_inconsistency(MOTION_DIRECTION this_mv, + MOTION_DIRECTION that_mv) { + if (this_mv == that_mv) { + return 0; + } else { + return abs(this_mv - that_mv) == 2 ? 2 : 1; + } +} +#endif + +#if CONFIG_EXT_PARTITION_TYPES +static void rd_test_partition3( + const AV1_COMP *const cpi, ThreadData *td, TileDataEnc *tile_data, + TOKENEXTRA **tp, PC_TREE *pc_tree, RD_STATS *best_rdc, + PICK_MODE_CONTEXT ctxs[3], PICK_MODE_CONTEXT *ctx, int mi_row, int mi_col, + BLOCK_SIZE bsize, PARTITION_TYPE partition, +#if CONFIG_SUPERTX + int64_t best_rd, int *best_rate_nocoef, RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx, +#endif + int mi_row0, int mi_col0, BLOCK_SIZE subsize0, int mi_row1, int mi_col1, + BLOCK_SIZE subsize1, int mi_row2, int mi_col2, BLOCK_SIZE subsize2) { + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + RD_STATS this_rdc, sum_rdc; +#if CONFIG_SUPERTX + const AV1_COMMON *const cm = &cpi->common; + TileInfo *const tile_info = &tile_data->tile_info; + int this_rate_nocoef, sum_rate_nocoef; + int abort_flag; + const int supertx_allowed = !frame_is_intra_only(cm) && + bsize <= MAX_SUPERTX_BLOCK_SIZE && + !xd->lossless[0]; +#endif + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx); + + rd_pick_sb_modes(cpi, tile_data, x, mi_row0, mi_col0, &sum_rdc, +#if CONFIG_SUPERTX + &sum_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + subsize0, &ctxs[0], best_rdc->rdcost); +#if CONFIG_SUPERTX + abort_flag = sum_rdc.rdcost >= best_rd; +#endif + +#if CONFIG_SUPERTX + if (sum_rdc.rdcost < INT64_MAX) { +#else + if (sum_rdc.rdcost < best_rdc->rdcost) { +#endif + PICK_MODE_CONTEXT *ctx_0 = &ctxs[0]; + update_state(cpi, td, ctx_0, mi_row0, mi_col0, subsize0, 1); + encode_superblock(cpi, td, tp, DRY_RUN_NORMAL, mi_row0, mi_col0, subsize0, + ctx_0, NULL); + + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_0); + +#if CONFIG_SUPERTX + rd_pick_sb_modes(cpi, tile_data, x, mi_row1, mi_col1, &this_rdc, + &this_rate_nocoef, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + subsize1, &ctxs[1], INT64_MAX - sum_rdc.rdcost); +#else + rd_pick_sb_modes(cpi, tile_data, x, mi_row1, mi_col1, &this_rdc, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + subsize1, &ctxs[1], best_rdc->rdcost - sum_rdc.rdcost); +#endif // CONFIG_SUPERTX + + if (this_rdc.rate == INT_MAX) { + sum_rdc.rdcost = INT64_MAX; +#if CONFIG_SUPERTX + sum_rate_nocoef = INT_MAX; +#endif + } else { + sum_rdc.rate += this_rdc.rate; + sum_rdc.dist += this_rdc.dist; + sum_rdc.rdcost += this_rdc.rdcost; +#if CONFIG_SUPERTX + sum_rate_nocoef += this_rate_nocoef; +#endif + } + +#if CONFIG_SUPERTX + if (sum_rdc.rdcost < INT64_MAX) { +#else + if (sum_rdc.rdcost < best_rdc->rdcost) { +#endif + PICK_MODE_CONTEXT *ctx_1 = &ctxs[1]; + update_state(cpi, td, ctx_1, mi_row1, mi_col1, subsize1, 1); + encode_superblock(cpi, td, tp, DRY_RUN_NORMAL, mi_row1, mi_col1, subsize1, + ctx_1, NULL); + + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_1); + +#if CONFIG_SUPERTX + rd_pick_sb_modes(cpi, tile_data, x, mi_row2, mi_col2, &this_rdc, + &this_rate_nocoef, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + subsize2, &ctxs[2], INT64_MAX - sum_rdc.rdcost); +#else + rd_pick_sb_modes(cpi, tile_data, x, mi_row2, mi_col2, &this_rdc, +#if CONFIG_EXT_PARTITION_TYPES + partition, +#endif + subsize2, &ctxs[2], best_rdc->rdcost - sum_rdc.rdcost); +#endif // CONFIG_SUPERTX + + if (this_rdc.rate == INT_MAX) { + sum_rdc.rdcost = INT64_MAX; +#if CONFIG_SUPERTX + sum_rate_nocoef = INT_MAX; +#endif + } else { + sum_rdc.rate += this_rdc.rate; + sum_rdc.dist += this_rdc.dist; + sum_rdc.rdcost += this_rdc.rdcost; +#if CONFIG_SUPERTX + sum_rate_nocoef += this_rate_nocoef; +#endif + } + +#if CONFIG_SUPERTX + if (supertx_allowed && !abort_flag && sum_rdc.rdcost < INT64_MAX) { + TX_SIZE supertx_size = max_txsize_lookup[bsize]; + const PARTITION_TYPE best_partition = pc_tree->partitioning; + pc_tree->partitioning = partition; + sum_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup[partition]] + [supertx_size], + 0); + sum_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); + + if (!check_intra_sb(cpi, tile_info, mi_row, mi_col, bsize, pc_tree)) { + TX_TYPE best_tx = DCT_DCT; + RD_STATS tmp_rdc = { sum_rate_nocoef, 0, 0 }; + + restore_context(x, x_ctx, mi_row, mi_col, bsize); + + rd_supertx_sb(cpi, td, tile_info, mi_row, mi_col, bsize, + &tmp_rdc.rate, &tmp_rdc.dist, &best_tx, pc_tree); + + tmp_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup[partition]] + [supertx_size], + 1); + tmp_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, tmp_rdc.rate, tmp_rdc.dist); + if (tmp_rdc.rdcost < sum_rdc.rdcost) { + sum_rdc = tmp_rdc; + update_supertx_param_sb(cpi, td, mi_row, mi_col, bsize, best_tx, + supertx_size, pc_tree); + } + } + + pc_tree->partitioning = best_partition; + } +#endif // CONFIG_SUPERTX + + if (sum_rdc.rdcost < best_rdc->rdcost) { + int pl = partition_plane_context(xd, mi_row, mi_col, +#if CONFIG_UNPOISON_PARTITION_CTX + has_rows, has_cols, +#endif + bsize); + sum_rdc.rate += cpi->partition_cost[pl][partition]; + sum_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); +#if CONFIG_SUPERTX + sum_rate_nocoef += cpi->partition_cost[pl][partition]; +#endif + if (sum_rdc.rdcost < best_rdc->rdcost) { +#if CONFIG_SUPERTX + *best_rate_nocoef = sum_rate_nocoef; + assert(*best_rate_nocoef >= 0); +#endif + *best_rdc = sum_rdc; + pc_tree->partitioning = partition; + } + } + } + } +} +#endif // CONFIG_EXT_PARTITION_TYPES + +// TODO(jingning,jimbankoski,rbultje): properly skip partition types that are +// unlikely to be selected depending on previous rate-distortion optimization +// results, for encoding speed-up. +static void rd_pick_partition(const AV1_COMP *const cpi, ThreadData *td, + TileDataEnc *tile_data, TOKENEXTRA **tp, + int mi_row, int mi_col, BLOCK_SIZE bsize, + RD_STATS *rd_cost, +#if CONFIG_SUPERTX + int *rate_nocoef, +#endif + int64_t best_rd, PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + TileInfo *const tile_info = &tile_data->tile_info; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + const int mi_step = mi_size_wide[bsize] / 2; + RD_SEARCH_MACROBLOCK_CONTEXT x_ctx; + const TOKENEXTRA *const tp_orig = *tp; + PICK_MODE_CONTEXT *ctx_none = &pc_tree->none; +#if CONFIG_UNPOISON_PARTITION_CTX + const int hbs = mi_size_wide[bsize] / 2; + const int has_rows = mi_row + hbs < cm->mi_rows; + const int has_cols = mi_col + hbs < cm->mi_cols; +#else + int tmp_partition_cost[PARTITION_TYPES]; +#endif + BLOCK_SIZE subsize; + RD_STATS this_rdc, sum_rdc, best_rdc; + const int bsize_at_least_8x8 = (bsize >= BLOCK_8X8); + int do_square_split = bsize_at_least_8x8; +#if CONFIG_CB4X4 + const int unify_bsize = 1; + const int pl = bsize_at_least_8x8 + ? partition_plane_context(xd, mi_row, mi_col, +#if CONFIG_UNPOISON_PARTITION_CTX + has_rows, has_cols, +#endif + bsize) + : 0; +#else + const int unify_bsize = 0; + const int pl = partition_plane_context(xd, mi_row, mi_col, +#if CONFIG_UNPOISON_PARTITION_CTX + has_rows, has_cols, +#endif + bsize); +#endif // CONFIG_CB4X4 + const int *partition_cost = cpi->partition_cost[pl]; +#if CONFIG_SUPERTX + int this_rate_nocoef, sum_rate_nocoef = 0, best_rate_nocoef = INT_MAX; + int abort_flag; + const int supertx_allowed = !frame_is_intra_only(cm) && + bsize <= MAX_SUPERTX_BLOCK_SIZE && + !xd->lossless[0]; +#endif // CONFIG_SUPERTX + + int do_rectangular_split = 1; +#if CONFIG_EXT_PARTITION_TYPES + BLOCK_SIZE bsize2 = get_subsize(bsize, PARTITION_SPLIT); +#endif + + // Override skipping rectangular partition operations for edge blocks + const int force_horz_split = (mi_row + mi_step >= cm->mi_rows); + const int force_vert_split = (mi_col + mi_step >= cm->mi_cols); + const int xss = x->e_mbd.plane[1].subsampling_x; + const int yss = x->e_mbd.plane[1].subsampling_y; + + BLOCK_SIZE min_size = x->min_partition_size; + BLOCK_SIZE max_size = x->max_partition_size; + +#if CONFIG_FP_MB_STATS + unsigned int src_diff_var = UINT_MAX; + int none_complexity = 0; +#endif + + int partition_none_allowed = !force_horz_split && !force_vert_split; + int partition_horz_allowed = + !force_vert_split && yss <= xss && bsize_at_least_8x8; + int partition_vert_allowed = + !force_horz_split && xss <= yss && bsize_at_least_8x8; + +#if CONFIG_PVQ + od_rollback_buffer pre_rdo_buf; +#endif + + (void)*tp_orig; + +#if !CONFIG_UNPOISON_PARTITION_CTX + if (force_horz_split || force_vert_split) { + tmp_partition_cost[PARTITION_NONE] = INT_MAX; + + if (!force_vert_split) { // force_horz_split only + tmp_partition_cost[PARTITION_VERT] = INT_MAX; + tmp_partition_cost[PARTITION_HORZ] = + av1_cost_bit(cm->fc->partition_prob[pl][PARTITION_HORZ], 0); + tmp_partition_cost[PARTITION_SPLIT] = + av1_cost_bit(cm->fc->partition_prob[pl][PARTITION_HORZ], 1); + } else if (!force_horz_split) { // force_vert_split only + tmp_partition_cost[PARTITION_HORZ] = INT_MAX; + tmp_partition_cost[PARTITION_VERT] = + av1_cost_bit(cm->fc->partition_prob[pl][PARTITION_VERT], 0); + tmp_partition_cost[PARTITION_SPLIT] = + av1_cost_bit(cm->fc->partition_prob[pl][PARTITION_VERT], 1); + } else { // force_ horz_split && force_vert_split horz_split + tmp_partition_cost[PARTITION_HORZ] = INT_MAX; + tmp_partition_cost[PARTITION_VERT] = INT_MAX; + tmp_partition_cost[PARTITION_SPLIT] = 0; + } + + partition_cost = tmp_partition_cost; + } +#endif + +#if CONFIG_VAR_TX +#ifndef NDEBUG + // Nothing should rely on the default value of this array (which is just + // leftover from encoding the previous block. Setting it to magic number + // when debugging. + memset(x->blk_skip[0], 234, sizeof(x->blk_skip[0])); +#endif // NDEBUG +#endif // CONFIG_VAR_TX + + assert(mi_size_wide[bsize] == mi_size_high[bsize]); + + av1_init_rd_stats(&this_rdc); + av1_init_rd_stats(&sum_rdc); + av1_invalid_rd_stats(&best_rdc); + best_rdc.rdcost = best_rd; + + set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize); + + if (bsize == BLOCK_16X16 && cpi->vaq_refresh) + x->mb_energy = av1_block_energy(cpi, x, bsize); + + if (cpi->sf.cb_partition_search && bsize == BLOCK_16X16) { + const int cb_partition_search_ctrl = + ((pc_tree->index == 0 || pc_tree->index == 3) + + get_chessboard_index(cm->current_video_frame)) & + 0x1; + + if (cb_partition_search_ctrl && bsize > min_size && bsize < max_size) + set_partition_range(cm, xd, mi_row, mi_col, bsize, &min_size, &max_size); + } + + // Determine partition types in search according to the speed features. + // The threshold set here has to be of square block size. + if (cpi->sf.auto_min_max_partition_size) { + const int no_partition_allowed = (bsize <= max_size && bsize >= min_size); + // Note: Further partitioning is NOT allowed when bsize == min_size already. + const int partition_allowed = (bsize <= max_size && bsize > min_size); + partition_none_allowed &= no_partition_allowed; + partition_horz_allowed &= partition_allowed || force_horz_split; + partition_vert_allowed &= partition_allowed || force_vert_split; + do_square_split &= bsize > min_size; + } + if (cpi->sf.use_square_partition_only) { + partition_horz_allowed &= force_horz_split; + partition_vert_allowed &= force_vert_split; + } + +#if CONFIG_VAR_TX + xd->above_txfm_context = cm->above_txfm_context + mi_col; + xd->left_txfm_context = + xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK); +#endif +#if !CONFIG_PVQ + save_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + save_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + +#if CONFIG_FP_MB_STATS + if (cpi->use_fp_mb_stats) { + set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize); + src_diff_var = get_sby_perpixel_diff_variance(cpi, &x->plane[0].src, mi_row, + mi_col, bsize); + } +#endif + +#if CONFIG_FP_MB_STATS + // Decide whether we shall split directly and skip searching NONE by using + // the first pass block statistics + if (cpi->use_fp_mb_stats && bsize >= BLOCK_32X32 && do_square_split && + partition_none_allowed && src_diff_var > 4 && + cm->base_qindex < qindex_split_threshold_lookup[bsize]) { + int mb_row = mi_row >> 1; + int mb_col = mi_col >> 1; + int mb_row_end = + AOMMIN(mb_row + num_16x16_blocks_high_lookup[bsize], cm->mb_rows); + int mb_col_end = + AOMMIN(mb_col + num_16x16_blocks_wide_lookup[bsize], cm->mb_cols); + int r, c; + + // compute a complexity measure, basically measure inconsistency of motion + // vectors obtained from the first pass in the current block + for (r = mb_row; r < mb_row_end; r++) { + for (c = mb_col; c < mb_col_end; c++) { + const int mb_index = r * cm->mb_cols + c; + + MOTION_DIRECTION this_mv; + MOTION_DIRECTION right_mv; + MOTION_DIRECTION bottom_mv; + + this_mv = + get_motion_direction_fp(cpi->twopass.this_frame_mb_stats[mb_index]); + + // to its right + if (c != mb_col_end - 1) { + right_mv = get_motion_direction_fp( + cpi->twopass.this_frame_mb_stats[mb_index + 1]); + none_complexity += get_motion_inconsistency(this_mv, right_mv); + } + + // to its bottom + if (r != mb_row_end - 1) { + bottom_mv = get_motion_direction_fp( + cpi->twopass.this_frame_mb_stats[mb_index + cm->mb_cols]); + none_complexity += get_motion_inconsistency(this_mv, bottom_mv); + } + + // do not count its left and top neighbors to avoid double counting + } + } + + if (none_complexity > complexity_16x16_blocks_threshold[bsize]) { + partition_none_allowed = 0; + } + } +#endif + + // PARTITION_NONE + if (partition_none_allowed) { + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &this_rdc, +#if CONFIG_SUPERTX + &this_rate_nocoef, +#endif +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_NONE, +#endif + bsize, ctx_none, best_rdc.rdcost); + if (this_rdc.rate != INT_MAX) { + if (bsize_at_least_8x8) { + this_rdc.rate += partition_cost[PARTITION_NONE]; + this_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, this_rdc.rate, this_rdc.dist); +#if CONFIG_SUPERTX + this_rate_nocoef += partition_cost[PARTITION_NONE]; +#endif + } + + if (this_rdc.rdcost < best_rdc.rdcost) { + // Adjust dist breakout threshold according to the partition size. + const int64_t dist_breakout_thr = + cpi->sf.partition_search_breakout_dist_thr >> + ((2 * (MAX_SB_SIZE_LOG2 - 2)) - + (b_width_log2_lookup[bsize] + b_height_log2_lookup[bsize])); + const int rate_breakout_thr = + cpi->sf.partition_search_breakout_rate_thr * + num_pels_log2_lookup[bsize]; + + best_rdc = this_rdc; +#if CONFIG_SUPERTX + best_rate_nocoef = this_rate_nocoef; + assert(best_rate_nocoef >= 0); +#endif + if (bsize_at_least_8x8) pc_tree->partitioning = PARTITION_NONE; + + // If all y, u, v transform blocks in this partition are skippable, and + // the dist & rate are within the thresholds, the partition search is + // terminated for current branch of the partition search tree. + // The dist & rate thresholds are set to 0 at speed 0 to disable the + // early termination at that speed. + if (!x->e_mbd.lossless[xd->mi[0]->mbmi.segment_id] && + (ctx_none->skippable && best_rdc.dist < dist_breakout_thr && + best_rdc.rate < rate_breakout_thr)) { + do_square_split = 0; + do_rectangular_split = 0; + } + +#if CONFIG_FP_MB_STATS + // Check if every 16x16 first pass block statistics has zero + // motion and the corresponding first pass residue is small enough. + // If that is the case, check the difference variance between the + // current frame and the last frame. If the variance is small enough, + // stop further splitting in RD optimization + if (cpi->use_fp_mb_stats && do_square_split && + cm->base_qindex > qindex_skip_threshold_lookup[bsize]) { + int mb_row = mi_row >> 1; + int mb_col = mi_col >> 1; + int mb_row_end = + AOMMIN(mb_row + num_16x16_blocks_high_lookup[bsize], cm->mb_rows); + int mb_col_end = + AOMMIN(mb_col + num_16x16_blocks_wide_lookup[bsize], cm->mb_cols); + int r, c; + + int skip = 1; + for (r = mb_row; r < mb_row_end; r++) { + for (c = mb_col; c < mb_col_end; c++) { + const int mb_index = r * cm->mb_cols + c; + if (!(cpi->twopass.this_frame_mb_stats[mb_index] & + FPMB_MOTION_ZERO_MASK) || + !(cpi->twopass.this_frame_mb_stats[mb_index] & + FPMB_ERROR_SMALL_MASK)) { + skip = 0; + break; + } + } + if (skip == 0) { + break; + } + } + if (skip) { + if (src_diff_var == UINT_MAX) { + set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize); + src_diff_var = get_sby_perpixel_diff_variance( + cpi, &x->plane[0].src, mi_row, mi_col, bsize); + } + if (src_diff_var < 8) { + do_square_split = 0; + do_rectangular_split = 0; + } + } + } +#endif + } + } +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + } + + // store estimated motion vector + if (cpi->sf.adaptive_motion_search) store_pred_mv(x, ctx_none); + + // PARTITION_SPLIT + // TODO(jingning): use the motion vectors given by the above search as + // the starting point of motion search in the following partition type check. + if (do_square_split) { + int reached_last_index = 0; + subsize = get_subsize(bsize, PARTITION_SPLIT); + if (bsize == BLOCK_8X8 && !unify_bsize) { +#if CONFIG_DUAL_FILTER + if (cpi->sf.adaptive_pred_interp_filter && partition_none_allowed) + pc_tree->leaf_split[0]->pred_interp_filter = + ctx_none->mic.mbmi.interp_filter[0]; +#else + if (cpi->sf.adaptive_pred_interp_filter && partition_none_allowed) + pc_tree->leaf_split[0]->pred_interp_filter = + ctx_none->mic.mbmi.interp_filter; +#endif +#if CONFIG_SUPERTX + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &sum_rdc, + &sum_rate_nocoef, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_SPLIT, +#endif + subsize, pc_tree->leaf_split[0], INT64_MAX); +#else + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &sum_rdc, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_SPLIT, +#endif + subsize, pc_tree->leaf_split[0], best_rdc.rdcost); +#endif // CONFIG_SUPERTX + if (sum_rdc.rate == INT_MAX) { + sum_rdc.rdcost = INT64_MAX; +#if CONFIG_SUPERTX + sum_rate_nocoef = INT_MAX; +#endif + } +#if CONFIG_SUPERTX + if (supertx_allowed && sum_rdc.rdcost < INT64_MAX) { + TX_SIZE supertx_size = max_txsize_lookup[bsize]; + const PARTITION_TYPE best_partition = pc_tree->partitioning; + + pc_tree->partitioning = PARTITION_SPLIT; + + sum_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup + [PARTITION_SPLIT]][supertx_size], + 0); + sum_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); + + if (is_inter_mode(pc_tree->leaf_split[0]->mic.mbmi.mode)) { + TX_TYPE best_tx = DCT_DCT; + RD_STATS tmp_rdc; + av1_init_rd_stats(&tmp_rdc); + tmp_rdc.rate = sum_rate_nocoef; + + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + + rd_supertx_sb(cpi, td, tile_info, mi_row, mi_col, bsize, + &tmp_rdc.rate, &tmp_rdc.dist, &best_tx, pc_tree); + + tmp_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup + [PARTITION_SPLIT]][supertx_size], + 1); + tmp_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, tmp_rdc.rate, tmp_rdc.dist); + if (tmp_rdc.rdcost < sum_rdc.rdcost) { + sum_rdc = tmp_rdc; + update_supertx_param_sb(cpi, td, mi_row, mi_col, bsize, best_tx, + supertx_size, pc_tree); + } + } + + pc_tree->partitioning = best_partition; + } +#endif // CONFIG_SUPERTX + reached_last_index = 1; + } else { + int idx; +#if CONFIG_SUPERTX + for (idx = 0; idx < 4 && sum_rdc.rdcost < INT64_MAX; ++idx) { +#else + for (idx = 0; idx < 4 && sum_rdc.rdcost < best_rdc.rdcost; ++idx) { +#endif // CONFIG_SUPERTX + const int x_idx = (idx & 1) * mi_step; + const int y_idx = (idx >> 1) * mi_step; + + if (mi_row + y_idx >= cm->mi_rows || mi_col + x_idx >= cm->mi_cols) + continue; + + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_none); + + pc_tree->split[idx]->index = idx; +#if CONFIG_SUPERTX + rd_pick_partition(cpi, td, tile_data, tp, mi_row + y_idx, + mi_col + x_idx, subsize, &this_rdc, &this_rate_nocoef, + INT64_MAX - sum_rdc.rdcost, pc_tree->split[idx]); +#else + rd_pick_partition( + cpi, td, tile_data, tp, mi_row + y_idx, mi_col + x_idx, subsize, + &this_rdc, best_rdc.rdcost - sum_rdc.rdcost, pc_tree->split[idx]); +#endif // CONFIG_SUPERTX + + if (this_rdc.rate == INT_MAX) { + sum_rdc.rdcost = INT64_MAX; +#if CONFIG_SUPERTX + sum_rate_nocoef = INT_MAX; +#endif // CONFIG_SUPERTX + break; + } else { + sum_rdc.rate += this_rdc.rate; + sum_rdc.dist += this_rdc.dist; + sum_rdc.rdcost += this_rdc.rdcost; +#if CONFIG_SUPERTX + sum_rate_nocoef += this_rate_nocoef; +#endif // CONFIG_SUPERTX + } + } + reached_last_index = (idx == 4); +#if CONFIG_SUPERTX + if (supertx_allowed && sum_rdc.rdcost < INT64_MAX && reached_last_index) { + TX_SIZE supertx_size = max_txsize_lookup[bsize]; + const PARTITION_TYPE best_partition = pc_tree->partitioning; + + pc_tree->partitioning = PARTITION_SPLIT; + + sum_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup + [PARTITION_SPLIT]][supertx_size], + 0); + sum_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); + + if (!check_intra_sb(cpi, tile_info, mi_row, mi_col, bsize, pc_tree)) { + TX_TYPE best_tx = DCT_DCT; + RD_STATS tmp_rdc; + av1_init_rd_stats(&tmp_rdc); + tmp_rdc.rate = sum_rate_nocoef; + + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + + rd_supertx_sb(cpi, td, tile_info, mi_row, mi_col, bsize, + &tmp_rdc.rate, &tmp_rdc.dist, &best_tx, pc_tree); + + tmp_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup + [PARTITION_SPLIT]][supertx_size], + 1); + tmp_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, tmp_rdc.rate, tmp_rdc.dist); + if (tmp_rdc.rdcost < sum_rdc.rdcost) { + sum_rdc = tmp_rdc; + update_supertx_param_sb(cpi, td, mi_row, mi_col, bsize, best_tx, + supertx_size, pc_tree); + } + } + + pc_tree->partitioning = best_partition; + } +#endif // CONFIG_SUPERTX + } + + if (reached_last_index && sum_rdc.rdcost < best_rdc.rdcost) { + sum_rdc.rate += partition_cost[PARTITION_SPLIT]; + sum_rdc.rdcost = RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); +#if CONFIG_SUPERTX + sum_rate_nocoef += partition_cost[PARTITION_SPLIT]; +#endif // CONFIG_SUPERTX + + if (sum_rdc.rdcost < best_rdc.rdcost) { + best_rdc = sum_rdc; +#if CONFIG_SUPERTX + best_rate_nocoef = sum_rate_nocoef; + assert(best_rate_nocoef >= 0); +#endif // CONFIG_SUPERTX + pc_tree->partitioning = PARTITION_SPLIT; + } + } else if (cpi->sf.less_rectangular_check) { + // skip rectangular partition test when larger block size + // gives better rd cost + do_rectangular_split &= !partition_none_allowed; + } +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + } // if (do_split) + + // PARTITION_HORZ + if (partition_horz_allowed && + (do_rectangular_split || av1_active_h_edge(cpi, mi_row, mi_step))) { + subsize = get_subsize(bsize, PARTITION_HORZ); + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_none); +#if CONFIG_DUAL_FILTER + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->horizontal[0].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter[0]; +#else + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->horizontal[0].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter; +#endif + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &sum_rdc, +#if CONFIG_SUPERTX + &sum_rate_nocoef, +#endif // CONFIG_SUPERTX +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_HORZ, +#endif + subsize, &pc_tree->horizontal[0], best_rdc.rdcost); + +#if CONFIG_SUPERTX + abort_flag = + (sum_rdc.rdcost >= best_rd && (bsize > BLOCK_8X8 || unify_bsize)) || + (sum_rdc.rate == INT_MAX && bsize == BLOCK_8X8); + if (sum_rdc.rdcost < INT64_MAX && +#else + if (sum_rdc.rdcost < best_rdc.rdcost && +#endif // CONFIG_SUPERTX + !force_horz_split && (bsize > BLOCK_8X8 || unify_bsize)) { + PICK_MODE_CONTEXT *ctx_h = &pc_tree->horizontal[0]; + update_state(cpi, td, ctx_h, mi_row, mi_col, subsize, 1); + encode_superblock(cpi, td, tp, DRY_RUN_NORMAL, mi_row, mi_col, subsize, + ctx_h, NULL); + + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_h); + +#if CONFIG_DUAL_FILTER + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->horizontal[1].pred_interp_filter = + ctx_h->mic.mbmi.interp_filter[0]; +#else + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->horizontal[1].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter; +#endif +#if CONFIG_SUPERTX + rd_pick_sb_modes(cpi, tile_data, x, mi_row + mi_step, mi_col, &this_rdc, + &this_rate_nocoef, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_HORZ, +#endif + subsize, &pc_tree->horizontal[1], INT64_MAX); +#else + rd_pick_sb_modes(cpi, tile_data, x, mi_row + mi_step, mi_col, &this_rdc, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_HORZ, +#endif + subsize, &pc_tree->horizontal[1], + best_rdc.rdcost - sum_rdc.rdcost); +#endif // CONFIG_SUPERTX + if (this_rdc.rate == INT_MAX) { + sum_rdc.rdcost = INT64_MAX; +#if CONFIG_SUPERTX + sum_rate_nocoef = INT_MAX; +#endif // CONFIG_SUPERTX + } else { + sum_rdc.rate += this_rdc.rate; + sum_rdc.dist += this_rdc.dist; + sum_rdc.rdcost += this_rdc.rdcost; +#if CONFIG_SUPERTX + sum_rate_nocoef += this_rate_nocoef; +#endif // CONFIG_SUPERTX + } + } + +#if CONFIG_SUPERTX + if (supertx_allowed && sum_rdc.rdcost < INT64_MAX && !abort_flag) { + TX_SIZE supertx_size = max_txsize_lookup[bsize]; + const PARTITION_TYPE best_partition = pc_tree->partitioning; + + pc_tree->partitioning = PARTITION_HORZ; + + sum_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup[PARTITION_HORZ]] + [supertx_size], + 0); + sum_rdc.rdcost = RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); + + if (!check_intra_sb(cpi, tile_info, mi_row, mi_col, bsize, pc_tree)) { + TX_TYPE best_tx = DCT_DCT; + RD_STATS tmp_rdc; + av1_init_rd_stats(&tmp_rdc); + tmp_rdc.rate = sum_rate_nocoef; + + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + + rd_supertx_sb(cpi, td, tile_info, mi_row, mi_col, bsize, &tmp_rdc.rate, + &tmp_rdc.dist, &best_tx, pc_tree); + + tmp_rdc.rate += av1_cost_bit( + cm->fc + ->supertx_prob[partition_supertx_context_lookup[PARTITION_HORZ]] + [supertx_size], + 1); + tmp_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, tmp_rdc.rate, tmp_rdc.dist); + if (tmp_rdc.rdcost < sum_rdc.rdcost) { + sum_rdc = tmp_rdc; + update_supertx_param_sb(cpi, td, mi_row, mi_col, bsize, best_tx, + supertx_size, pc_tree); + } + } + + pc_tree->partitioning = best_partition; + } +#endif // CONFIG_SUPERTX + + if (sum_rdc.rdcost < best_rdc.rdcost) { + sum_rdc.rate += partition_cost[PARTITION_HORZ]; + sum_rdc.rdcost = RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); +#if CONFIG_SUPERTX + sum_rate_nocoef += partition_cost[PARTITION_HORZ]; +#endif // CONFIG_SUPERTX + if (sum_rdc.rdcost < best_rdc.rdcost) { + best_rdc = sum_rdc; +#if CONFIG_SUPERTX + best_rate_nocoef = sum_rate_nocoef; + assert(best_rate_nocoef >= 0); +#endif // CONFIG_SUPERTX + pc_tree->partitioning = PARTITION_HORZ; + } + } +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + } + + // PARTITION_VERT + if (partition_vert_allowed && + (do_rectangular_split || av1_active_v_edge(cpi, mi_col, mi_step))) { + subsize = get_subsize(bsize, PARTITION_VERT); + + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_none); + +#if CONFIG_DUAL_FILTER + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->vertical[0].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter[0]; +#else + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->vertical[0].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter; +#endif + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &sum_rdc, +#if CONFIG_SUPERTX + &sum_rate_nocoef, +#endif // CONFIG_SUPERTX +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_VERT, +#endif + subsize, &pc_tree->vertical[0], best_rdc.rdcost); +#if CONFIG_SUPERTX + abort_flag = + (sum_rdc.rdcost >= best_rd && (bsize > BLOCK_8X8 || unify_bsize)) || + (sum_rdc.rate == INT_MAX && bsize == BLOCK_8X8); + if (sum_rdc.rdcost < INT64_MAX && +#else + if (sum_rdc.rdcost < best_rdc.rdcost && +#endif // CONFIG_SUPERTX + !force_vert_split && (bsize > BLOCK_8X8 || unify_bsize)) { + update_state(cpi, td, &pc_tree->vertical[0], mi_row, mi_col, subsize, 1); + encode_superblock(cpi, td, tp, DRY_RUN_NORMAL, mi_row, mi_col, subsize, + &pc_tree->vertical[0], NULL); + + if (cpi->sf.adaptive_motion_search) load_pred_mv(x, ctx_none); + +#if CONFIG_DUAL_FILTER + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->vertical[1].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter[0]; +#else + if (cpi->sf.adaptive_pred_interp_filter && bsize == BLOCK_8X8 && + partition_none_allowed) + pc_tree->vertical[1].pred_interp_filter = + ctx_none->mic.mbmi.interp_filter; +#endif +#if CONFIG_SUPERTX + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col + mi_step, &this_rdc, + &this_rate_nocoef, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_VERT, +#endif + subsize, &pc_tree->vertical[1], + INT64_MAX - sum_rdc.rdcost); +#else + rd_pick_sb_modes(cpi, tile_data, x, mi_row, mi_col + mi_step, &this_rdc, +#if CONFIG_EXT_PARTITION_TYPES + PARTITION_VERT, +#endif + subsize, &pc_tree->vertical[1], + best_rdc.rdcost - sum_rdc.rdcost); +#endif // CONFIG_SUPERTX + if (this_rdc.rate == INT_MAX) { + sum_rdc.rdcost = INT64_MAX; +#if CONFIG_SUPERTX + sum_rate_nocoef = INT_MAX; +#endif // CONFIG_SUPERTX + } else { + sum_rdc.rate += this_rdc.rate; + sum_rdc.dist += this_rdc.dist; + sum_rdc.rdcost += this_rdc.rdcost; +#if CONFIG_SUPERTX + sum_rate_nocoef += this_rate_nocoef; +#endif // CONFIG_SUPERTX + } + } +#if CONFIG_SUPERTX + if (supertx_allowed && sum_rdc.rdcost < INT64_MAX && !abort_flag) { + TX_SIZE supertx_size = max_txsize_lookup[bsize]; + const PARTITION_TYPE best_partition = pc_tree->partitioning; + + pc_tree->partitioning = PARTITION_VERT; + + sum_rdc.rate += av1_cost_bit( + cm->fc->supertx_prob[partition_supertx_context_lookup[PARTITION_VERT]] + [supertx_size], + 0); + sum_rdc.rdcost = RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); + + if (!check_intra_sb(cpi, tile_info, mi_row, mi_col, bsize, pc_tree)) { + TX_TYPE best_tx = DCT_DCT; + RD_STATS tmp_rdc; + av1_init_rd_stats(&tmp_rdc); + tmp_rdc.rate = sum_rate_nocoef; + + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + + rd_supertx_sb(cpi, td, tile_info, mi_row, mi_col, bsize, &tmp_rdc.rate, + &tmp_rdc.dist, &best_tx, pc_tree); + + tmp_rdc.rate += av1_cost_bit( + cm->fc + ->supertx_prob[partition_supertx_context_lookup[PARTITION_VERT]] + [supertx_size], + 1); + tmp_rdc.rdcost = + RDCOST(x->rdmult, x->rddiv, tmp_rdc.rate, tmp_rdc.dist); + if (tmp_rdc.rdcost < sum_rdc.rdcost) { + sum_rdc = tmp_rdc; + update_supertx_param_sb(cpi, td, mi_row, mi_col, bsize, best_tx, + supertx_size, pc_tree); + } + } + + pc_tree->partitioning = best_partition; + } +#endif // CONFIG_SUPERTX + + if (sum_rdc.rdcost < best_rdc.rdcost) { + sum_rdc.rate += partition_cost[PARTITION_VERT]; + sum_rdc.rdcost = RDCOST(x->rdmult, x->rddiv, sum_rdc.rate, sum_rdc.dist); +#if CONFIG_SUPERTX + sum_rate_nocoef += partition_cost[PARTITION_VERT]; +#endif // CONFIG_SUPERTX + if (sum_rdc.rdcost < best_rdc.rdcost) { + best_rdc = sum_rdc; +#if CONFIG_SUPERTX + best_rate_nocoef = sum_rate_nocoef; + assert(best_rate_nocoef >= 0); +#endif // CONFIG_SUPERTX + pc_tree->partitioning = PARTITION_VERT; + } + } +#if !CONFIG_PVQ + restore_context(x, &x_ctx, mi_row, mi_col, bsize); +#else + restore_context(x, &x_ctx, mi_row, mi_col, &pre_rdo_buf, bsize); +#endif + } + +#if CONFIG_EXT_PARTITION_TYPES + // PARTITION_HORZ_A + if (partition_horz_allowed && do_rectangular_split && bsize > BLOCK_8X8 && + partition_none_allowed) { + subsize = get_subsize(bsize, PARTITION_HORZ_A); + rd_test_partition3(cpi, td, tile_data, tp, pc_tree, &best_rdc, + pc_tree->horizontala, ctx_none, mi_row, mi_col, bsize, + PARTITION_HORZ_A, +#if CONFIG_SUPERTX + best_rd, &best_rate_nocoef, &x_ctx, +#endif + mi_row, mi_col, bsize2, mi_row, mi_col + mi_step, bsize2, + mi_row + mi_step, mi_col, subsize); + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + } + // PARTITION_HORZ_B + if (partition_horz_allowed && do_rectangular_split && bsize > BLOCK_8X8 && + partition_none_allowed) { + subsize = get_subsize(bsize, PARTITION_HORZ_B); + rd_test_partition3(cpi, td, tile_data, tp, pc_tree, &best_rdc, + pc_tree->horizontalb, ctx_none, mi_row, mi_col, bsize, + PARTITION_HORZ_B, +#if CONFIG_SUPERTX + best_rd, &best_rate_nocoef, &x_ctx, +#endif + mi_row, mi_col, subsize, mi_row + mi_step, mi_col, + bsize2, mi_row + mi_step, mi_col + mi_step, bsize2); + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + } + // PARTITION_VERT_A + if (partition_vert_allowed && do_rectangular_split && bsize > BLOCK_8X8 && + partition_none_allowed) { + subsize = get_subsize(bsize, PARTITION_VERT_A); + rd_test_partition3(cpi, td, tile_data, tp, pc_tree, &best_rdc, + pc_tree->verticala, ctx_none, mi_row, mi_col, bsize, + PARTITION_VERT_A, +#if CONFIG_SUPERTX + best_rd, &best_rate_nocoef, &x_ctx, +#endif + mi_row, mi_col, bsize2, mi_row + mi_step, mi_col, bsize2, + mi_row, mi_col + mi_step, subsize); + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + } + // PARTITION_VERT_B + if (partition_vert_allowed && do_rectangular_split && bsize > BLOCK_8X8 && + partition_none_allowed) { + subsize = get_subsize(bsize, PARTITION_VERT_B); + rd_test_partition3(cpi, td, tile_data, tp, pc_tree, &best_rdc, + pc_tree->verticalb, ctx_none, mi_row, mi_col, bsize, + PARTITION_VERT_B, +#if CONFIG_SUPERTX + best_rd, &best_rate_nocoef, &x_ctx, +#endif + mi_row, mi_col, subsize, mi_row, mi_col + mi_step, + bsize2, mi_row + mi_step, mi_col + mi_step, bsize2); + restore_context(x, &x_ctx, mi_row, mi_col, bsize); + } +#endif // CONFIG_EXT_PARTITION_TYPES + + // TODO(jbb): This code added so that we avoid static analysis + // warning related to the fact that best_rd isn't used after this + // point. This code should be refactored so that the duplicate + // checks occur in some sub function and thus are used... + (void)best_rd; + *rd_cost = best_rdc; +#if CONFIG_SUPERTX + *rate_nocoef = best_rate_nocoef; +#endif // CONFIG_SUPERTX + +#if CONFIG_CFL + // Store the luma for the best mode + x->cfl_store_y = 1; +#endif + if (best_rdc.rate < INT_MAX && best_rdc.dist < INT64_MAX && + pc_tree->index != 3) { + if (bsize == cm->sb_size) { +#if CONFIG_MOTION_VAR && CONFIG_NCOBMC + set_mode_info_sb(cpi, td, tile_info, tp, mi_row, mi_col, bsize, pc_tree); +#endif + encode_sb(cpi, td, tile_info, tp, mi_row, mi_col, OUTPUT_ENABLED, bsize, + pc_tree, NULL); + } else { + encode_sb(cpi, td, tile_info, tp, mi_row, mi_col, DRY_RUN_NORMAL, bsize, + pc_tree, NULL); + } + } +#if CONFIG_CFL + x->cfl_store_y = 0; +#endif + + if (bsize == cm->sb_size) { +#if !CONFIG_PVQ && !CONFIG_LV_MAP + assert(tp_orig < *tp || (tp_orig == *tp && xd->mi[0]->mbmi.skip)); +#endif + assert(best_rdc.rate < INT_MAX); + assert(best_rdc.dist < INT64_MAX); + } else { + assert(tp_orig == *tp); + } +} + +static void encode_rd_sb_row(AV1_COMP *cpi, ThreadData *td, + TileDataEnc *tile_data, int mi_row, + TOKENEXTRA **tp) { + AV1_COMMON *const cm = &cpi->common; + const TileInfo *const tile_info = &tile_data->tile_info; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + SPEED_FEATURES *const sf = &cpi->sf; + int mi_col; +#if CONFIG_EXT_PARTITION + const int leaf_nodes = 256; +#else + const int leaf_nodes = 64; +#endif // CONFIG_EXT_PARTITION + + // Initialize the left context for the new SB row + av1_zero_left_context(xd); + +#if CONFIG_DELTA_Q + // Reset delta for every tile + if (cm->delta_q_present_flag) + if (mi_row == tile_info->mi_row_start) xd->prev_qindex = cm->base_qindex; +#if CONFIG_EXT_DELTA_Q + if (cm->delta_lf_present_flag) + if (mi_row == tile_info->mi_row_start) xd->prev_delta_lf_from_base = 0; +#endif +#endif + + // Code each SB in the row + for (mi_col = tile_info->mi_col_start; mi_col < tile_info->mi_col_end; + mi_col += cm->mib_size) { + const struct segmentation *const seg = &cm->seg; + int dummy_rate; + int64_t dummy_dist; + RD_STATS dummy_rdc; +#if CONFIG_SUPERTX + int dummy_rate_nocoef; +#endif // CONFIG_SUPERTX + int i; + int seg_skip = 0; + + const int idx_str = cm->mi_stride * mi_row + mi_col; + MODE_INFO **mi = cm->mi_grid_visible + idx_str; + PC_TREE *const pc_root = td->pc_root[cm->mib_size_log2 - MIN_MIB_SIZE_LOG2]; + + av1_update_boundary_info(cm, tile_info, mi_row, mi_col); + + if (sf->adaptive_pred_interp_filter) { + for (i = 0; i < leaf_nodes; ++i) + td->leaf_tree[i].pred_interp_filter = SWITCHABLE; + + for (i = 0; i < leaf_nodes; ++i) { + td->pc_tree[i].vertical[0].pred_interp_filter = SWITCHABLE; + td->pc_tree[i].vertical[1].pred_interp_filter = SWITCHABLE; + td->pc_tree[i].horizontal[0].pred_interp_filter = SWITCHABLE; + td->pc_tree[i].horizontal[1].pred_interp_filter = SWITCHABLE; + } + } + + av1_zero(x->pred_mv); + pc_root->index = 0; + + if (seg->enabled) { + const uint8_t *const map = + seg->update_map ? cpi->segmentation_map : cm->last_frame_seg_map; + int segment_id = get_segment_id(cm, map, cm->sb_size, mi_row, mi_col); + seg_skip = segfeature_active(seg, segment_id, SEG_LVL_SKIP); + } + +#if CONFIG_DELTA_Q + if (cm->delta_q_present_flag) { + // Test mode for delta quantization + int sb_row = mi_row >> 3; + int sb_col = mi_col >> 3; + int sb_stride = (cm->width + MAX_SB_SIZE - 1) >> MAX_SB_SIZE_LOG2; + int index = ((sb_row * sb_stride + sb_col + 8) & 31) - 16; + + // Ensure divisibility of delta_qindex by delta_q_res + int offset_qindex = (index < 0 ? -index - 8 : index - 8); + int qmask = ~(cm->delta_q_res - 1); + int current_qindex = clamp(cm->base_qindex + offset_qindex, + cm->delta_q_res, 256 - cm->delta_q_res); + + current_qindex = + ((current_qindex - cm->base_qindex + cm->delta_q_res / 2) & qmask) + + cm->base_qindex; + assert(current_qindex > 0); + + xd->delta_qindex = current_qindex - cm->base_qindex; + set_offsets(cpi, tile_info, x, mi_row, mi_col, BLOCK_64X64); + xd->mi[0]->mbmi.current_q_index = current_qindex; +#if !CONFIG_EXT_DELTA_Q + xd->mi[0]->mbmi.segment_id = 0; +#endif // CONFIG_EXT_DELTA_Q + av1_init_plane_quantizers(cpi, x, xd->mi[0]->mbmi.segment_id); +#if CONFIG_EXT_DELTA_Q + if (cpi->oxcf.deltaq_mode == DELTA_Q_LF) { + int j, k; + int lfmask = ~(cm->delta_lf_res - 1); + int current_delta_lf_from_base = offset_qindex / 2; + current_delta_lf_from_base = + ((current_delta_lf_from_base + cm->delta_lf_res / 2) & lfmask); + + // pre-set the delta lf for loop filter. Note that this value is set + // before mi is assigned for each block in current superblock + for (j = 0; j < AOMMIN(cm->mib_size, cm->mi_rows - mi_row); j++) { + for (k = 0; k < AOMMIN(cm->mib_size, cm->mi_cols - mi_col); k++) { + cm->mi[(mi_row + j) * cm->mi_stride + (mi_col + k)] + .mbmi.current_delta_lf_from_base = current_delta_lf_from_base; + } + } + } +#endif // CONFIG_EXT_DELTA_Q + } +#endif // CONFIG_DELTA_Q + + x->source_variance = UINT_MAX; + if (sf->partition_search_type == FIXED_PARTITION || seg_skip) { + BLOCK_SIZE bsize; + set_offsets(cpi, tile_info, x, mi_row, mi_col, cm->sb_size); + bsize = seg_skip ? cm->sb_size : sf->always_this_block_size; + set_fixed_partitioning(cpi, tile_info, mi, mi_row, mi_col, bsize); + rd_use_partition(cpi, td, tile_data, mi, tp, mi_row, mi_col, cm->sb_size, + &dummy_rate, &dummy_dist, +#if CONFIG_SUPERTX + &dummy_rate_nocoef, +#endif // CONFIG_SUPERTX + 1, pc_root); + } else if (cpi->partition_search_skippable_frame) { + BLOCK_SIZE bsize; + set_offsets(cpi, tile_info, x, mi_row, mi_col, cm->sb_size); + bsize = get_rd_var_based_fixed_partition(cpi, x, mi_row, mi_col); + set_fixed_partitioning(cpi, tile_info, mi, mi_row, mi_col, bsize); + rd_use_partition(cpi, td, tile_data, mi, tp, mi_row, mi_col, cm->sb_size, + &dummy_rate, &dummy_dist, +#if CONFIG_SUPERTX + &dummy_rate_nocoef, +#endif // CONFIG_SUPERTX + 1, pc_root); + } else if (sf->partition_search_type == VAR_BASED_PARTITION) { + choose_partitioning(cpi, td, tile_info, x, mi_row, mi_col); + rd_use_partition(cpi, td, tile_data, mi, tp, mi_row, mi_col, cm->sb_size, + &dummy_rate, &dummy_dist, +#if CONFIG_SUPERTX + &dummy_rate_nocoef, +#endif // CONFIG_SUPERTX + 1, pc_root); + } else { + // If required set upper and lower partition size limits + if (sf->auto_min_max_partition_size) { + set_offsets(cpi, tile_info, x, mi_row, mi_col, cm->sb_size); + rd_auto_partition_range(cpi, tile_info, xd, mi_row, mi_col, + &x->min_partition_size, &x->max_partition_size); + } + rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, cm->sb_size, + &dummy_rdc, +#if CONFIG_SUPERTX + &dummy_rate_nocoef, +#endif // CONFIG_SUPERTX + INT64_MAX, pc_root); + } + } +#if CONFIG_SUBFRAME_PROB_UPDATE + if (cm->do_subframe_update && + cm->refresh_frame_context == REFRESH_FRAME_CONTEXT_BACKWARD) { + const int mi_rows_per_update = + MI_SIZE * AOMMAX(cm->mi_rows / MI_SIZE / COEF_PROBS_BUFS, 1); + if ((mi_row + MI_SIZE) % mi_rows_per_update == 0 && + mi_row + MI_SIZE < cm->mi_rows && + cm->coef_probs_update_idx < COEF_PROBS_BUFS - 1) { + TX_SIZE t; + SUBFRAME_STATS *subframe_stats = &cpi->subframe_stats; + + for (t = 0; t < TX_SIZES; ++t) + av1_full_to_model_counts(cpi->td.counts->coef[t], + cpi->td.rd_counts.coef_counts[t]); + av1_partial_adapt_probs(cm, mi_row, mi_col); + ++cm->coef_probs_update_idx; + av1_copy(subframe_stats->coef_probs_buf[cm->coef_probs_update_idx], + cm->fc->coef_probs); + av1_copy(subframe_stats->coef_counts_buf[cm->coef_probs_update_idx], + cpi->td.rd_counts.coef_counts); + av1_copy(subframe_stats->eob_counts_buf[cm->coef_probs_update_idx], + cm->counts.eob_branch); + av1_fill_token_costs(x->token_costs, cm->fc->coef_probs); + } + } +#endif // CONFIG_SUBFRAME_PROB_UPDATE +} + +static void init_encode_frame_mb_context(AV1_COMP *cpi) { + MACROBLOCK *const x = &cpi->td.mb; + AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + + // Copy data over into macro block data structures. + av1_setup_src_planes(x, cpi->source, 0, 0); + + av1_setup_block_planes(xd, cm->subsampling_x, cm->subsampling_y); +} + +#if !CONFIG_REF_ADAPT +static int check_dual_ref_flags(AV1_COMP *cpi) { + const int ref_flags = cpi->ref_frame_flags; + + if (segfeature_active(&cpi->common.seg, 1, SEG_LVL_REF_FRAME)) { + return 0; + } else { + return (!!(ref_flags & AOM_GOLD_FLAG) + !!(ref_flags & AOM_LAST_FLAG) + +#if CONFIG_EXT_REFS + !!(ref_flags & AOM_LAST2_FLAG) + !!(ref_flags & AOM_LAST3_FLAG) + + !!(ref_flags & AOM_BWD_FLAG) + +#endif // CONFIG_EXT_REFS + !!(ref_flags & AOM_ALT_FLAG)) >= 2; + } +} +#endif // !CONFIG_REF_ADAPT + +#if !CONFIG_VAR_TX +static void reset_skip_tx_size(AV1_COMMON *cm, TX_SIZE max_tx_size) { + int mi_row, mi_col; + const int mis = cm->mi_stride; + MODE_INFO **mi_ptr = cm->mi_grid_visible; + + for (mi_row = 0; mi_row < cm->mi_rows; ++mi_row, mi_ptr += mis) { + for (mi_col = 0; mi_col < cm->mi_cols; ++mi_col) { + if (txsize_sqr_up_map[mi_ptr[mi_col]->mbmi.tx_size] > max_tx_size) + mi_ptr[mi_col]->mbmi.tx_size = max_tx_size; + } + } +} +#endif + +static MV_REFERENCE_FRAME get_frame_type(const AV1_COMP *cpi) { + if (frame_is_intra_only(&cpi->common)) return INTRA_FRAME; +#if CONFIG_EXT_REFS + // We will not update the golden frame with an internal overlay frame + else if ((cpi->rc.is_src_frame_alt_ref && cpi->refresh_golden_frame) || + cpi->rc.is_src_frame_ext_arf) +#else + else if (cpi->rc.is_src_frame_alt_ref && cpi->refresh_golden_frame) +#endif + return ALTREF_FRAME; + else if (cpi->refresh_golden_frame || cpi->refresh_alt_ref_frame) + return GOLDEN_FRAME; + else + // TODO(zoeliu): To investigate whether a frame_type other than + // INTRA/ALTREF/GOLDEN/LAST needs to be specified seperately. + return LAST_FRAME; +} + +static TX_MODE select_tx_mode(const AV1_COMP *cpi, MACROBLOCKD *const xd) { + int i, all_lossless = 1; + + if (cpi->common.seg.enabled) { + for (i = 0; i < MAX_SEGMENTS; ++i) { + if (!xd->lossless[i]) { + all_lossless = 0; + break; + } + } + } else { + all_lossless = xd->lossless[0]; + } + if (all_lossless) return ONLY_4X4; + if (cpi->sf.tx_size_search_method == USE_LARGESTALL) + return ALLOW_32X32 + CONFIG_TX64X64; + else if (cpi->sf.tx_size_search_method == USE_FULL_RD || + cpi->sf.tx_size_search_method == USE_TX_8X8) + return TX_MODE_SELECT; + else + return cpi->common.tx_mode; +} + +void av1_init_tile_data(AV1_COMP *cpi) { + AV1_COMMON *const cm = &cpi->common; + const int tile_cols = cm->tile_cols; + const int tile_rows = cm->tile_rows; + int tile_col, tile_row; + TOKENEXTRA *pre_tok = cpi->tile_tok[0][0]; + unsigned int tile_tok = 0; + + if (cpi->tile_data == NULL || cpi->allocated_tiles < tile_cols * tile_rows) { + if (cpi->tile_data != NULL) aom_free(cpi->tile_data); + CHECK_MEM_ERROR( + cm, cpi->tile_data, + aom_memalign(32, tile_cols * tile_rows * sizeof(*cpi->tile_data))); + cpi->allocated_tiles = tile_cols * tile_rows; + + for (tile_row = 0; tile_row < tile_rows; ++tile_row) + for (tile_col = 0; tile_col < tile_cols; ++tile_col) { + TileDataEnc *const tile_data = + &cpi->tile_data[tile_row * tile_cols + tile_col]; + int i, j; + for (i = 0; i < BLOCK_SIZES; ++i) { + for (j = 0; j < MAX_MODES; ++j) { + tile_data->thresh_freq_fact[i][j] = 32; + tile_data->mode_map[i][j] = j; + } + } +#if CONFIG_PVQ + // This will be dynamically increased as more pvq block is encoded. + tile_data->pvq_q.buf_len = 1000; + CHECK_MEM_ERROR( + cm, tile_data->pvq_q.buf, + aom_malloc(tile_data->pvq_q.buf_len * sizeof(PVQ_INFO))); + tile_data->pvq_q.curr_pos = 0; +#endif + } + } + + for (tile_row = 0; tile_row < tile_rows; ++tile_row) { + for (tile_col = 0; tile_col < tile_cols; ++tile_col) { + TileInfo *const tile_info = + &cpi->tile_data[tile_row * tile_cols + tile_col].tile_info; + av1_tile_init(tile_info, cm, tile_row, tile_col); + + cpi->tile_tok[tile_row][tile_col] = pre_tok + tile_tok; + pre_tok = cpi->tile_tok[tile_row][tile_col]; + tile_tok = allocated_tokens(*tile_info); +#if CONFIG_PVQ + cpi->tile_data[tile_row * tile_cols + tile_col].pvq_q.curr_pos = 0; +#endif + } + } +} + +void av1_encode_tile(AV1_COMP *cpi, ThreadData *td, int tile_row, + int tile_col) { + AV1_COMMON *const cm = &cpi->common; + TileDataEnc *const this_tile = + &cpi->tile_data[tile_row * cm->tile_cols + tile_col]; + const TileInfo *const tile_info = &this_tile->tile_info; + TOKENEXTRA *tok = cpi->tile_tok[tile_row][tile_col]; + int mi_row; + +#if CONFIG_DEPENDENT_HORZTILES +#if CONFIG_TILE_GROUPS + if ((!cm->dependent_horz_tiles) || (tile_row == 0) || + tile_info->tg_horz_boundary) { +#else + if ((!cm->dependent_horz_tiles) || (tile_row == 0)) { +#endif + av1_zero_above_context(cm, tile_info->mi_col_start, tile_info->mi_col_end); + } +#else + av1_zero_above_context(cm, tile_info->mi_col_start, tile_info->mi_col_end); +#endif + + // Set up pointers to per thread motion search counters. + this_tile->m_search_count = 0; // Count of motion search hits. + this_tile->ex_search_count = 0; // Exhaustive mesh search hits. + td->mb.m_search_count_ptr = &this_tile->m_search_count; + td->mb.ex_search_count_ptr = &this_tile->ex_search_count; + +#if CONFIG_PVQ + td->mb.pvq_q = &this_tile->pvq_q; + + // TODO(yushin) : activity masking info needs be signaled by a bitstream + td->mb.daala_enc.use_activity_masking = AV1_PVQ_ENABLE_ACTIVITY_MASKING; + + if (td->mb.daala_enc.use_activity_masking) + td->mb.daala_enc.qm = OD_HVS_QM; // Hard coded. Enc/dec required to sync. + else + td->mb.daala_enc.qm = OD_FLAT_QM; // Hard coded. Enc/dec required to sync. + + { + // FIXME: Multiple segments support + int segment_id = 0; + int rdmult = set_segment_rdmult(cpi, &td->mb, segment_id); + int qindex = av1_get_qindex(&cm->seg, segment_id, cm->base_qindex); +#if CONFIG_HIGHBITDEPTH + const int quantizer_shift = td->mb.e_mbd.bd - 8; +#else + const int quantizer_shift = 0; +#endif // CONFIG_HIGHBITDEPTH + int64_t q_ac = OD_MAXI( + 1, av1_ac_quant(qindex, 0, cpi->common.bit_depth) >> quantizer_shift); + int64_t q_dc = OD_MAXI( + 1, av1_dc_quant(qindex, 0, cpi->common.bit_depth) >> quantizer_shift); + /* td->mb.daala_enc.pvq_norm_lambda = OD_PVQ_LAMBDA; */ + td->mb.daala_enc.pvq_norm_lambda = + (double)rdmult * (64 / 16) / (q_ac * q_ac * (1 << RDDIV_BITS)); + td->mb.daala_enc.pvq_norm_lambda_dc = + (double)rdmult * (64 / 16) / (q_dc * q_dc * (1 << RDDIV_BITS)); + // printf("%f\n", td->mb.daala_enc.pvq_norm_lambda); + } + od_init_qm(td->mb.daala_enc.state.qm, td->mb.daala_enc.state.qm_inv, + td->mb.daala_enc.qm == OD_HVS_QM ? OD_QM8_Q4_HVS : OD_QM8_Q4_FLAT); + + if (td->mb.daala_enc.use_activity_masking) { + int pli; + int use_masking = td->mb.daala_enc.use_activity_masking; + int segment_id = 0; + int qindex = av1_get_qindex(&cm->seg, segment_id, cm->base_qindex); + + for (pli = 0; pli < MAX_MB_PLANE; pli++) { + int i; + int q; + + q = qindex; + if (q <= OD_DEFAULT_QMS[use_masking][0][pli].interp_q << OD_COEFF_SHIFT) { + od_interp_qm(&td->mb.daala_enc.state.pvq_qm_q4[pli][0], q, + &OD_DEFAULT_QMS[use_masking][0][pli], NULL); + } else { + i = 0; + while (OD_DEFAULT_QMS[use_masking][i + 1][pli].qm_q4 != NULL && + q > OD_DEFAULT_QMS[use_masking][i + 1][pli].interp_q + << OD_COEFF_SHIFT) { + i++; + } + od_interp_qm(&td->mb.daala_enc.state.pvq_qm_q4[pli][0], q, + &OD_DEFAULT_QMS[use_masking][i][pli], + &OD_DEFAULT_QMS[use_masking][i + 1][pli]); + } + } + } + +#if CONFIG_DAALA_EC + od_ec_enc_init(&td->mb.daala_enc.w.ec, 65025); +#else +#error "CONFIG_PVQ currently requires CONFIG_DAALA_EC." +#endif + +#if CONFIG_DAALA_EC + od_ec_enc_reset(&td->mb.daala_enc.w.ec); +#else +#error "CONFIG_PVQ currently requires CONFIG_DAALA_EC." +#endif +#endif // #if CONFIG_PVQ + +#if CONFIG_EC_ADAPT + this_tile->tctx = *cm->fc; + td->mb.e_mbd.tile_ctx = &this_tile->tctx; +#endif // #if CONFIG_EC_ADAPT + +#if CONFIG_CFL + MACROBLOCKD *const xd = &td->mb.e_mbd; + xd->cfl = &this_tile->cfl; + cfl_init(xd->cfl, cm, xd->plane[AOM_PLANE_U].subsampling_x, + xd->plane[AOM_PLANE_U].subsampling_y); +#endif + +#if CONFIG_PVQ + td->mb.daala_enc.state.adapt = &this_tile->tctx.pvq_context; +#endif // CONFIG_PVQ + + for (mi_row = tile_info->mi_row_start; mi_row < tile_info->mi_row_end; + mi_row += cm->mib_size) { + encode_rd_sb_row(cpi, td, this_tile, mi_row, &tok); + } + + cpi->tok_count[tile_row][tile_col] = + (unsigned int)(tok - cpi->tile_tok[tile_row][tile_col]); + assert(cpi->tok_count[tile_row][tile_col] <= allocated_tokens(*tile_info)); +#if CONFIG_PVQ +#if CONFIG_DAALA_EC + od_ec_enc_clear(&td->mb.daala_enc.w.ec); +#else +#error "CONFIG_PVQ currently requires CONFIG_DAALA_EC." +#endif + + td->mb.pvq_q->last_pos = td->mb.pvq_q->curr_pos; + // rewind current position so that bitstream can be written + // from the 1st pvq block + td->mb.pvq_q->curr_pos = 0; + + td->mb.pvq_q = NULL; +#endif +} + +static void encode_tiles(AV1_COMP *cpi) { + AV1_COMMON *const cm = &cpi->common; + int tile_col, tile_row; + + av1_init_tile_data(cpi); + + for (tile_row = 0; tile_row < cm->tile_rows; ++tile_row) + for (tile_col = 0; tile_col < cm->tile_cols; ++tile_col) + av1_encode_tile(cpi, &cpi->td, tile_row, tile_col); +} + +#if CONFIG_FP_MB_STATS +static int input_fpmb_stats(FIRSTPASS_MB_STATS *firstpass_mb_stats, + AV1_COMMON *cm, uint8_t **this_frame_mb_stats) { + uint8_t *mb_stats_in = firstpass_mb_stats->mb_stats_start + + cm->current_video_frame * cm->MBs * sizeof(uint8_t); + + if (mb_stats_in > firstpass_mb_stats->mb_stats_end) return EOF; + + *this_frame_mb_stats = mb_stats_in; + + return 1; +} +#endif + +#if CONFIG_GLOBAL_MOTION +#define GLOBAL_TRANS_TYPES_ENC 3 // highest motion model to search +static int gm_get_params_cost(WarpedMotionParams *gm, + WarpedMotionParams *ref_gm, int allow_hp) { + assert(gm->wmtype < GLOBAL_TRANS_TYPES); + int params_cost = 0; + int trans_bits, trans_prec_diff; + switch (gm->wmtype) { + case HOMOGRAPHY: + case HORTRAPEZOID: + case VERTRAPEZOID: + if (gm->wmtype != HORTRAPEZOID) + params_cost += aom_count_signed_primitive_refsubexpfin( + GM_ROW3HOMO_MAX + 1, SUBEXPFIN_K, + (ref_gm->wmmat[6] >> GM_ROW3HOMO_PREC_DIFF), + (gm->wmmat[6] >> GM_ROW3HOMO_PREC_DIFF)); + if (gm->wmtype != VERTRAPEZOID) + params_cost += aom_count_signed_primitive_refsubexpfin( + GM_ROW3HOMO_MAX + 1, SUBEXPFIN_K, + (ref_gm->wmmat[7] >> GM_ROW3HOMO_PREC_DIFF), + (gm->wmmat[7] >> GM_ROW3HOMO_PREC_DIFF)); + // Fallthrough intended + case AFFINE: + case ROTZOOM: + params_cost += aom_count_signed_primitive_refsubexpfin( + GM_ALPHA_MAX + 1, SUBEXPFIN_K, + (ref_gm->wmmat[2] >> GM_ALPHA_PREC_DIFF) - (1 << GM_ALPHA_PREC_BITS), + (gm->wmmat[2] >> GM_ALPHA_PREC_DIFF) - (1 << GM_ALPHA_PREC_BITS)); + if (gm->wmtype != VERTRAPEZOID) + params_cost += aom_count_signed_primitive_refsubexpfin( + GM_ALPHA_MAX + 1, SUBEXPFIN_K, + (ref_gm->wmmat[3] >> GM_ALPHA_PREC_DIFF), + (gm->wmmat[3] >> GM_ALPHA_PREC_DIFF)); + if (gm->wmtype >= AFFINE) { + if (gm->wmtype != HORTRAPEZOID) + params_cost += aom_count_signed_primitive_refsubexpfin( + GM_ALPHA_MAX + 1, SUBEXPFIN_K, + (ref_gm->wmmat[4] >> GM_ALPHA_PREC_DIFF), + (gm->wmmat[4] >> GM_ALPHA_PREC_DIFF)); + params_cost += aom_count_signed_primitive_refsubexpfin( + GM_ALPHA_MAX + 1, SUBEXPFIN_K, + (ref_gm->wmmat[5] >> GM_ALPHA_PREC_DIFF) - + (1 << GM_ALPHA_PREC_BITS), + (gm->wmmat[5] >> GM_ALPHA_PREC_DIFF) - (1 << GM_ALPHA_PREC_BITS)); + } + // Fallthrough intended + case TRANSLATION: + trans_bits = (gm->wmtype == TRANSLATION) + ? GM_ABS_TRANS_ONLY_BITS - !allow_hp + : GM_ABS_TRANS_BITS; + trans_prec_diff = (gm->wmtype == TRANSLATION) + ? GM_TRANS_ONLY_PREC_DIFF + !allow_hp + : GM_TRANS_PREC_DIFF; + params_cost += aom_count_signed_primitive_refsubexpfin( + (1 << trans_bits) + 1, SUBEXPFIN_K, + (ref_gm->wmmat[0] >> trans_prec_diff), + (gm->wmmat[0] >> trans_prec_diff)); + params_cost += aom_count_signed_primitive_refsubexpfin( + (1 << trans_bits) + 1, SUBEXPFIN_K, + (ref_gm->wmmat[1] >> trans_prec_diff), + (gm->wmmat[1] >> trans_prec_diff)); + // Fallthrough intended + case IDENTITY: break; + default: assert(0); + } + return (params_cost << AV1_PROB_COST_SHIFT); +} +#endif // CONFIG_GLOBAL_MOTION + +static void encode_frame_internal(AV1_COMP *cpi) { + ThreadData *const td = &cpi->td; + MACROBLOCK *const x = &td->mb; + AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + RD_COUNTS *const rdc = &cpi->td.rd_counts; + int i; +#if CONFIG_TEMPMV_SIGNALING || CONFIG_EXT_REFS + const int last_fb_buf_idx = get_ref_frame_buf_idx(cpi, LAST_FRAME); +#endif // CONFIG_TEMPMV_SIGNALING || CONFIG_EXT_REFS + +#if CONFIG_ADAPT_SCAN + av1_deliver_eob_threshold(cm, xd); +#endif + + x->min_partition_size = AOMMIN(x->min_partition_size, cm->sb_size); + x->max_partition_size = AOMMIN(x->max_partition_size, cm->sb_size); +#if CONFIG_REF_MV + cm->setup_mi(cm); +#endif + + xd->mi = cm->mi_grid_visible; + xd->mi[0] = cm->mi; + + av1_zero(*td->counts); + av1_zero(rdc->coef_counts); + av1_zero(rdc->comp_pred_diff); + +#if CONFIG_GLOBAL_MOTION + av1_zero(rdc->global_motion_used); + if (cpi->common.frame_type == INTER_FRAME && cpi->source && + !cpi->global_motion_search_done) { + YV12_BUFFER_CONFIG *ref_buf; + int frame; + double params_by_motion[RANSAC_NUM_MOTIONS * (MAX_PARAMDIM - 1)]; + const double *params_this_motion; + int inliers_by_motion[RANSAC_NUM_MOTIONS]; + WarpedMotionParams tmp_wm_params; + static const double kInfiniteErrAdv = 1e12; + static const double kIdentityParams[MAX_PARAMDIM - 1] = { + 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0 + }; + + for (frame = LAST_FRAME; frame <= ALTREF_FRAME; ++frame) { + ref_buf = get_ref_frame_buffer(cpi, frame); + if (ref_buf) { + TransformationType model; + aom_clear_system_state(); + for (model = ROTZOOM; model < GLOBAL_TRANS_TYPES_ENC; ++model) { + double best_erroradvantage = kInfiniteErrAdv; + + // Initially set all params to identity. + for (i = 0; i < RANSAC_NUM_MOTIONS; ++i) { + memcpy(params_by_motion + (MAX_PARAMDIM - 1) * i, kIdentityParams, + (MAX_PARAMDIM - 1) * sizeof(*params_by_motion)); + } + + compute_global_motion_feature_based( + model, cpi->source, ref_buf, +#if CONFIG_HIGHBITDEPTH + cpi->common.bit_depth, +#endif // CONFIG_HIGHBITDEPTH + inliers_by_motion, params_by_motion, RANSAC_NUM_MOTIONS); + + for (i = 0; i < RANSAC_NUM_MOTIONS; ++i) { + if (inliers_by_motion[i] == 0) continue; + + params_this_motion = params_by_motion + (MAX_PARAMDIM - 1) * i; + convert_model_to_params(params_this_motion, &tmp_wm_params); + + if (tmp_wm_params.wmtype != IDENTITY) { + const double erroradv_this_motion = refine_integerized_param( + &tmp_wm_params, tmp_wm_params.wmtype, +#if CONFIG_HIGHBITDEPTH + xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH, xd->bd, +#endif // CONFIG_HIGHBITDEPTH + ref_buf->y_buffer, ref_buf->y_width, ref_buf->y_height, + ref_buf->y_stride, cpi->source->y_buffer, + cpi->source->y_width, cpi->source->y_height, + cpi->source->y_stride, 3); + if (erroradv_this_motion < best_erroradvantage) { + best_erroradvantage = erroradv_this_motion; + // Save the wm_params modified by refine_integerized_param() + // rather than motion index to avoid rerunning refine() below. + memcpy(&(cm->global_motion[frame]), &tmp_wm_params, + sizeof(WarpedMotionParams)); + } + } + } + if (cm->global_motion[frame].wmtype <= AFFINE) + if (!get_shear_params(&cm->global_motion[frame])) + set_default_warp_params(&cm->global_motion[frame]); + + if (cm->global_motion[frame].wmtype == TRANSLATION) { + cm->global_motion[frame].wmmat[0] = + convert_to_trans_prec(cm->allow_high_precision_mv, + cm->global_motion[frame].wmmat[0]) * + GM_TRANS_ONLY_DECODE_FACTOR; + cm->global_motion[frame].wmmat[1] = + convert_to_trans_prec(cm->allow_high_precision_mv, + cm->global_motion[frame].wmmat[1]) * + GM_TRANS_ONLY_DECODE_FACTOR; + } + + // If the best error advantage found doesn't meet the threshold for + // this motion type, revert to IDENTITY. + if (!is_enough_erroradvantage( + best_erroradvantage, + gm_get_params_cost(&cm->global_motion[frame], + &cm->prev_frame->global_motion[frame], + cm->allow_high_precision_mv))) { + set_default_warp_params(&cm->global_motion[frame]); + } + + if (cm->global_motion[frame].wmtype != IDENTITY) break; + } + aom_clear_system_state(); + } + cpi->gmparams_cost[frame] = + gm_get_params_cost(&cm->global_motion[frame], + &cm->prev_frame->global_motion[frame], + cm->allow_high_precision_mv) + + cpi->gmtype_cost[cm->global_motion[frame].wmtype] - + cpi->gmtype_cost[IDENTITY]; + } + cpi->global_motion_search_done = 1; + } + memcpy(cm->cur_frame->global_motion, cm->global_motion, + TOTAL_REFS_PER_FRAME * sizeof(WarpedMotionParams)); +#endif // CONFIG_GLOBAL_MOTION + + for (i = 0; i < MAX_SEGMENTS; ++i) { + const int qindex = cm->seg.enabled + ? av1_get_qindex(&cm->seg, i, cm->base_qindex) + : cm->base_qindex; + xd->lossless[i] = qindex == 0 && cm->y_dc_delta_q == 0 && + cm->uv_dc_delta_q == 0 && cm->uv_ac_delta_q == 0; + xd->qindex[i] = qindex; + } + + if (!cm->seg.enabled && xd->lossless[0]) x->optimize = 0; + + cm->tx_mode = select_tx_mode(cpi, xd); + +#if CONFIG_DELTA_Q + // Fix delta q resolution for the moment + cm->delta_q_res = DEFAULT_DELTA_Q_RES; +// Set delta_q_present_flag before it is used for the first time +#if CONFIG_EXT_DELTA_Q + cm->delta_lf_res = DEFAULT_DELTA_LF_RES; + // update delta_q_present_flag and delta_lf_present_flag based on base_qindex + cm->delta_q_present_flag &= cm->base_qindex > 0; + cm->delta_lf_present_flag &= cm->base_qindex > 0; +#else + cm->delta_q_present_flag = + cpi->oxcf.aq_mode == DELTA_AQ && cm->base_qindex > 0; +#endif // CONFIG_EXT_DELTA_Q +#endif + + av1_frame_init_quantizer(cpi); + + av1_initialize_rd_consts(cpi); + av1_initialize_me_consts(cpi, x, cm->base_qindex); + init_encode_frame_mb_context(cpi); +#if CONFIG_TEMPMV_SIGNALING + if (last_fb_buf_idx != INVALID_IDX) { + cm->prev_frame = &cm->buffer_pool->frame_bufs[last_fb_buf_idx]; + cm->use_prev_frame_mvs &= !cm->error_resilient_mode && + cm->width == cm->prev_frame->buf.y_width && + cm->height == cm->prev_frame->buf.y_height && + !cm->intra_only && !cm->prev_frame->intra_only; + } +#else + cm->use_prev_frame_mvs = + !cm->error_resilient_mode && cm->width == cm->last_width && + cm->height == cm->last_height && !cm->intra_only && cm->last_show_frame; +#endif + +#if CONFIG_EXT_REFS + // NOTE(zoeliu): As cm->prev_frame can take neither a frame of + // show_exisiting_frame=1, nor can it take a frame not used as + // a reference, it is probable that by the time it is being + // referred to, the frame buffer it originally points to may + // already get expired and have been reassigned to the current + // newly coded frame. Hence, we need to check whether this is + // the case, and if yes, we have 2 choices: + // (1) Simply disable the use of previous frame mvs; or + // (2) Have cm->prev_frame point to one reference frame buffer, + // e.g. LAST_FRAME. + if (cm->use_prev_frame_mvs && !enc_is_ref_frame_buf(cpi, cm->prev_frame)) { + // Reassign the LAST_FRAME buffer to cm->prev_frame. + cm->prev_frame = &cm->buffer_pool->frame_bufs[last_fb_buf_idx]; + } +#endif // CONFIG_EXT_REFS + + // Special case: set prev_mi to NULL when the previous mode info + // context cannot be used. + cm->prev_mi = + cm->use_prev_frame_mvs ? cm->prev_mip + cm->mi_stride + 1 : NULL; + +#if CONFIG_VAR_TX + x->txb_split_count = 0; +#if CONFIG_REF_MV + av1_zero(x->blk_skip_drl); +#endif +#endif + + if (cpi->sf.partition_search_type == VAR_BASED_PARTITION && + cpi->td.var_root[0] == NULL) + av1_setup_var_tree(&cpi->common, &cpi->td); + + { + struct aom_usec_timer emr_timer; + aom_usec_timer_start(&emr_timer); + +#if CONFIG_FP_MB_STATS + if (cpi->use_fp_mb_stats) { + input_fpmb_stats(&cpi->twopass.firstpass_mb_stats, cm, + &cpi->twopass.this_frame_mb_stats); + } +#endif + + // If allowed, encoding tiles in parallel with one thread handling one tile. + // TODO(geza.lore): The multi-threaded encoder is not safe with more than + // 1 tile rows, as it uses the single above_context et al arrays from + // cpi->common + if (AOMMIN(cpi->oxcf.max_threads, cm->tile_cols) > 1 && cm->tile_rows == 1) + av1_encode_tiles_mt(cpi); + else + encode_tiles(cpi); + + aom_usec_timer_mark(&emr_timer); + cpi->time_encode_sb_row += aom_usec_timer_elapsed(&emr_timer); + } + +#if 0 + // Keep record of the total distortion this time around for future use + cpi->last_frame_distortion = cpi->frame_distortion; +#endif +} + +void av1_encode_frame(AV1_COMP *cpi) { + AV1_COMMON *const cm = &cpi->common; +#if CONFIG_EXT_TX + // Indicates whether or not to use a default reduced set for ext-tx + // rather than the potential full set of 16 transforms + cm->reduced_tx_set_used = 0; +#endif // CONFIG_EXT_TX + + // In the longer term the encoder should be generalized to match the + // decoder such that we allow compound where one of the 3 buffers has a + // different sign bias and that buffer is then the fixed ref. However, this + // requires further work in the rd loop. For now the only supported encoder + // side behavior is where the ALT ref buffer has opposite sign bias to + // the other two. + if (!frame_is_intra_only(cm)) { +#if CONFIG_LOWDELAY_COMPOUND // Normative in encoder + cpi->allow_comp_inter_inter = 1; +#if CONFIG_EXT_REFS + cm->comp_fwd_ref[0] = LAST_FRAME; + cm->comp_fwd_ref[1] = LAST2_FRAME; + cm->comp_fwd_ref[2] = LAST3_FRAME; + cm->comp_fwd_ref[3] = GOLDEN_FRAME; + cm->comp_bwd_ref[0] = BWDREF_FRAME; + cm->comp_bwd_ref[1] = ALTREF_FRAME; +#else + cm->comp_fixed_ref = ALTREF_FRAME; + cm->comp_var_ref[0] = LAST_FRAME; + cm->comp_var_ref[1] = GOLDEN_FRAME; +#endif // CONFIG_EXT_REFS +#else + if ((cm->ref_frame_sign_bias[ALTREF_FRAME] == + cm->ref_frame_sign_bias[GOLDEN_FRAME]) || + (cm->ref_frame_sign_bias[ALTREF_FRAME] == + cm->ref_frame_sign_bias[LAST_FRAME])) { + cpi->allow_comp_inter_inter = 0; + } else { + cpi->allow_comp_inter_inter = 1; + +#if CONFIG_EXT_REFS + cm->comp_fwd_ref[0] = LAST_FRAME; + cm->comp_fwd_ref[1] = LAST2_FRAME; + cm->comp_fwd_ref[2] = LAST3_FRAME; + cm->comp_fwd_ref[3] = GOLDEN_FRAME; + cm->comp_bwd_ref[0] = BWDREF_FRAME; + cm->comp_bwd_ref[1] = ALTREF_FRAME; +#else + cm->comp_fixed_ref = ALTREF_FRAME; + cm->comp_var_ref[0] = LAST_FRAME; + cm->comp_var_ref[1] = GOLDEN_FRAME; +#endif // CONFIG_EXT_REFS + } +#endif + } else { + cpi->allow_comp_inter_inter = 0; + } + + if (cpi->sf.frame_parameter_update) { + int i; + RD_OPT *const rd_opt = &cpi->rd; + FRAME_COUNTS *counts = cpi->td.counts; + RD_COUNTS *const rdc = &cpi->td.rd_counts; + + // This code does a single RD pass over the whole frame assuming + // either compound, single or hybrid prediction as per whatever has + // worked best for that type of frame in the past. + // It also predicts whether another coding mode would have worked + // better than this coding mode. If that is the case, it remembers + // that for subsequent frames. + // It does the same analysis for transform size selection also. + // + // TODO(zoeliu): To investigate whether a frame_type other than + // INTRA/ALTREF/GOLDEN/LAST needs to be specified seperately. + const MV_REFERENCE_FRAME frame_type = get_frame_type(cpi); + int64_t *const mode_thrs = rd_opt->prediction_type_threshes[frame_type]; + const int is_alt_ref = frame_type == ALTREF_FRAME; + +/* prediction (compound, single or hybrid) mode selection */ +#if CONFIG_REF_ADAPT + // NOTE(zoeliu): "is_alt_ref" is true only for OVERLAY/INTNL_OVERLAY frames + if (is_alt_ref || !cpi->allow_comp_inter_inter) + cm->reference_mode = SINGLE_REFERENCE; + else + cm->reference_mode = REFERENCE_MODE_SELECT; +#else + if (is_alt_ref || !cpi->allow_comp_inter_inter) + cm->reference_mode = SINGLE_REFERENCE; + else if (mode_thrs[COMPOUND_REFERENCE] > mode_thrs[SINGLE_REFERENCE] && + mode_thrs[COMPOUND_REFERENCE] > mode_thrs[REFERENCE_MODE_SELECT] && + check_dual_ref_flags(cpi) && cpi->static_mb_pct == 100) + cm->reference_mode = COMPOUND_REFERENCE; + else if (mode_thrs[SINGLE_REFERENCE] > mode_thrs[REFERENCE_MODE_SELECT]) + cm->reference_mode = SINGLE_REFERENCE; + else + cm->reference_mode = REFERENCE_MODE_SELECT; +#endif // CONFIG_REF_ADAPT + +#if CONFIG_DUAL_FILTER + cm->interp_filter = SWITCHABLE; +#endif + + encode_frame_internal(cpi); + + for (i = 0; i < REFERENCE_MODES; ++i) + mode_thrs[i] = (mode_thrs[i] + rdc->comp_pred_diff[i] / cm->MBs) / 2; + + if (cm->reference_mode == REFERENCE_MODE_SELECT) { + int single_count_zero = 0; + int comp_count_zero = 0; + + for (i = 0; i < COMP_INTER_CONTEXTS; i++) { + single_count_zero += counts->comp_inter[i][0]; + comp_count_zero += counts->comp_inter[i][1]; + } + + if (comp_count_zero == 0) { + cm->reference_mode = SINGLE_REFERENCE; + av1_zero(counts->comp_inter); +#if !CONFIG_REF_ADAPT + } else if (single_count_zero == 0) { + cm->reference_mode = COMPOUND_REFERENCE; + av1_zero(counts->comp_inter); +#endif // !CONFIG_REF_ADAPT + } + } + +#if CONFIG_VAR_TX + if (cm->tx_mode == TX_MODE_SELECT && cpi->td.mb.txb_split_count == 0) + cm->tx_mode = ALLOW_32X32 + CONFIG_TX64X64; +#else + if (cm->tx_mode == TX_MODE_SELECT) { +#if CONFIG_TX64X64 + int count4x4 = 0; + int count8x8_8x8p = 0, count8x8_lp = 0; + int count16x16_16x16p = 0, count16x16_lp = 0; + int count32x32_32x32p = 0, count32x32_lp = 0; + int count64x64_64x64p = 0; + for (i = 0; i < TX_SIZE_CONTEXTS; ++i) { + // counts->tx_size[max_depth][context_idx][this_depth_level] + count4x4 += counts->tx_size[0][i][0]; + count4x4 += counts->tx_size[1][i][0]; + count4x4 += counts->tx_size[2][i][0]; + count4x4 += counts->tx_size[3][i][0]; + + count8x8_8x8p += counts->tx_size[0][i][1]; + count8x8_lp += counts->tx_size[1][i][1]; + count8x8_lp += counts->tx_size[2][i][1]; + count8x8_lp += counts->tx_size[3][i][1]; + + count16x16_16x16p += counts->tx_size[1][i][2]; + count16x16_lp += counts->tx_size[2][i][2]; + count16x16_lp += counts->tx_size[3][i][2]; + + count32x32_32x32p += counts->tx_size[2][i][3]; + count32x32_lp += counts->tx_size[3][i][3]; + + count64x64_64x64p += counts->tx_size[3][i][4]; + } +#if CONFIG_EXT_TX && CONFIG_RECT_TX + count4x4 += counts->tx_size_implied[0][TX_4X4]; + count4x4 += counts->tx_size_implied[1][TX_4X4]; + count4x4 += counts->tx_size_implied[2][TX_4X4]; + count4x4 += counts->tx_size_implied[3][TX_4X4]; + count8x8_8x8p += counts->tx_size_implied[1][TX_8X8]; + count8x8_lp += counts->tx_size_implied[2][TX_8X8]; + count8x8_lp += counts->tx_size_implied[3][TX_8X8]; + count8x8_lp += counts->tx_size_implied[4][TX_8X8]; + count16x16_16x16p += counts->tx_size_implied[2][TX_16X16]; + count16x16_lp += counts->tx_size_implied[3][TX_16X16]; + count16x16_lp += counts->tx_size_implied[4][TX_16X16]; + count32x32_32x32p += counts->tx_size_implied[3][TX_32X32]; + count32x32_lp += counts->tx_size_implied[4][TX_32X32]; + count64x64_64x64p += counts->tx_size_implied[4][TX_64X64]; +#endif // CONFIG_EXT_TX && CONFIG_RECT_TX + if (count4x4 == 0 && count16x16_lp == 0 && count16x16_16x16p == 0 && + count32x32_lp == 0 && count32x32_32x32p == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_16X16] == 0 && + cm->counts.supertx_size[TX_32X32] == 0 && + cm->counts.supertx_size[TX_64X64] == 0 && +#endif + count64x64_64x64p == 0) { + cm->tx_mode = ALLOW_8X8; + reset_skip_tx_size(cm, TX_8X8); + } else if (count8x8_8x8p == 0 && count8x8_lp == 0 && + count16x16_16x16p == 0 && count16x16_lp == 0 && + count32x32_32x32p == 0 && count32x32_lp == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_8X8] == 0 && + cm->counts.supertx_size[TX_16X16] == 0 && + cm->counts.supertx_size[TX_32X32] == 0 && + cm->counts.supertx_size[TX_64X64] == 0 && +#endif + count64x64_64x64p == 0) { + cm->tx_mode = ONLY_4X4; + reset_skip_tx_size(cm, TX_4X4); + } else if (count4x4 == 0 && count8x8_lp == 0 && count16x16_lp == 0 && + count32x32_lp == 0) { + cm->tx_mode = ALLOW_64X64; + } else if (count4x4 == 0 && count8x8_lp == 0 && count16x16_lp == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_64X64] == 0 && +#endif + count64x64_64x64p == 0) { + cm->tx_mode = ALLOW_32X32; + reset_skip_tx_size(cm, TX_32X32); + } else if (count4x4 == 0 && count8x8_lp == 0 && count32x32_lp == 0 && + count32x32_32x32p == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_32X32] == 0 && + cm->counts.supertx_size[TX_64X64] == 0 && +#endif + count64x64_64x64p == 0) { + cm->tx_mode = ALLOW_16X16; + reset_skip_tx_size(cm, TX_16X16); + } + +#else // CONFIG_TX64X64 + + int count4x4 = 0; + int count8x8_lp = 0, count8x8_8x8p = 0; + int count16x16_16x16p = 0, count16x16_lp = 0; + int count32x32 = 0; + for (i = 0; i < TX_SIZE_CONTEXTS; ++i) { + // counts->tx_size[max_depth][context_idx][this_depth_level] + count4x4 += counts->tx_size[0][i][0]; + count4x4 += counts->tx_size[1][i][0]; + count4x4 += counts->tx_size[2][i][0]; + + count8x8_8x8p += counts->tx_size[0][i][1]; + count8x8_lp += counts->tx_size[1][i][1]; + count8x8_lp += counts->tx_size[2][i][1]; + + count16x16_16x16p += counts->tx_size[1][i][2]; + count16x16_lp += counts->tx_size[2][i][2]; + count32x32 += counts->tx_size[2][i][3]; + } +#if CONFIG_EXT_TX && CONFIG_RECT_TX + count4x4 += counts->tx_size_implied[0][TX_4X4]; + count4x4 += counts->tx_size_implied[1][TX_4X4]; + count4x4 += counts->tx_size_implied[2][TX_4X4]; + count4x4 += counts->tx_size_implied[3][TX_4X4]; + count8x8_8x8p += counts->tx_size_implied[1][TX_8X8]; + count8x8_lp += counts->tx_size_implied[2][TX_8X8]; + count8x8_lp += counts->tx_size_implied[3][TX_8X8]; + count16x16_lp += counts->tx_size_implied[3][TX_16X16]; + count16x16_16x16p += counts->tx_size_implied[2][TX_16X16]; + count32x32 += counts->tx_size_implied[3][TX_32X32]; +#endif // CONFIG_EXT_TX && CONFIG_RECT_TX + if (count4x4 == 0 && count16x16_lp == 0 && count16x16_16x16p == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_16X16] == 0 && + cm->counts.supertx_size[TX_32X32] == 0 && +#endif // CONFIG_SUPERTX + count32x32 == 0) { + cm->tx_mode = ALLOW_8X8; + reset_skip_tx_size(cm, TX_8X8); + } else if (count8x8_8x8p == 0 && count16x16_16x16p == 0 && + count8x8_lp == 0 && count16x16_lp == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_8X8] == 0 && + cm->counts.supertx_size[TX_16X16] == 0 && + cm->counts.supertx_size[TX_32X32] == 0 && +#endif // CONFIG_SUPERTX + count32x32 == 0) { + cm->tx_mode = ONLY_4X4; + reset_skip_tx_size(cm, TX_4X4); + } else if (count8x8_lp == 0 && count16x16_lp == 0 && count4x4 == 0) { + cm->tx_mode = ALLOW_32X32; + } else if (count32x32 == 0 && count8x8_lp == 0 && +#if CONFIG_SUPERTX + cm->counts.supertx_size[TX_32X32] == 0 && +#endif // CONFIG_SUPERTX + count4x4 == 0) { + cm->tx_mode = ALLOW_16X16; + reset_skip_tx_size(cm, TX_16X16); + } +#endif // CONFIG_TX64X64 + } +#endif + } else { + encode_frame_internal(cpi); + } +} + +static void sum_intra_stats(FRAME_COUNTS *counts, MACROBLOCKD *xd, + const MODE_INFO *mi, const MODE_INFO *above_mi, + const MODE_INFO *left_mi, const int intraonly, + const int mi_row, const int mi_col) { + const MB_MODE_INFO *const mbmi = &mi->mbmi; + const PREDICTION_MODE y_mode = mbmi->mode; + const PREDICTION_MODE uv_mode = mbmi->uv_mode; + const BLOCK_SIZE bsize = mbmi->sb_type; + const int unify_bsize = CONFIG_CB4X4; + + if (bsize < BLOCK_8X8 && !unify_bsize) { + int idx, idy; + const int num_4x4_w = num_4x4_blocks_wide_lookup[bsize]; + const int num_4x4_h = num_4x4_blocks_high_lookup[bsize]; + for (idy = 0; idy < 2; idy += num_4x4_h) + for (idx = 0; idx < 2; idx += num_4x4_w) { + const int bidx = idy * 2 + idx; + const PREDICTION_MODE bmode = mi->bmi[bidx].as_mode; + if (intraonly) { + const PREDICTION_MODE a = av1_above_block_mode(mi, above_mi, bidx); + const PREDICTION_MODE l = av1_left_block_mode(mi, left_mi, bidx); + ++counts->kf_y_mode[a][l][bmode]; + } else { + ++counts->y_mode[0][bmode]; + } + } + } else { + if (intraonly) { + const PREDICTION_MODE above = av1_above_block_mode(mi, above_mi, 0); + const PREDICTION_MODE left = av1_left_block_mode(mi, left_mi, 0); + ++counts->kf_y_mode[above][left][y_mode]; + } else { + ++counts->y_mode[size_group_lookup[bsize]][y_mode]; + } +#if CONFIG_FILTER_INTRA + if (mbmi->mode == DC_PRED +#if CONFIG_PALETTE + && mbmi->palette_mode_info.palette_size[0] == 0 +#endif // CONFIG_PALETTE + ) { + const int use_filter_intra_mode = + mbmi->filter_intra_mode_info.use_filter_intra_mode[0]; + ++counts->filter_intra[0][use_filter_intra_mode]; + } + if (mbmi->uv_mode == DC_PRED +#if CONFIG_PALETTE + && mbmi->palette_mode_info.palette_size[1] == 0 +#endif // CONFIG_PALETTE + ) { + const int use_filter_intra_mode = + mbmi->filter_intra_mode_info.use_filter_intra_mode[1]; + ++counts->filter_intra[1][use_filter_intra_mode]; + } +#endif // CONFIG_FILTER_INTRA +#if CONFIG_EXT_INTRA && CONFIG_INTRA_INTERP + if (av1_is_directional_mode(mbmi->mode, bsize)) { + const int intra_filter_ctx = av1_get_pred_context_intra_interp(xd); + const int p_angle = + mode_to_angle_map[mbmi->mode] + mbmi->angle_delta[0] * ANGLE_STEP; + if (av1_is_intra_filter_switchable(p_angle)) + ++counts->intra_filter[intra_filter_ctx][mbmi->intra_filter]; + } +#endif // CONFIG_INTRA_INTERP && CONFIG_INTRA_INTERP + } + +#if CONFIG_CB4X4 + if (!is_chroma_reference(mi_row, mi_col, bsize, xd->plane[1].subsampling_x, + xd->plane[1].subsampling_y)) + return; +#else + (void)mi_row; + (void)mi_col; + (void)xd; +#endif + ++counts->uv_mode[y_mode][uv_mode]; +} + +#if CONFIG_VAR_TX +static void update_txfm_count(MACROBLOCK *x, MACROBLOCKD *xd, + FRAME_COUNTS *counts, TX_SIZE tx_size, int depth, + int blk_row, int blk_col) { + MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi; + const int tx_row = blk_row >> 1; + const int tx_col = blk_col >> 1; + const int max_blocks_high = max_block_high(xd, mbmi->sb_type, 0); + const int max_blocks_wide = max_block_wide(xd, mbmi->sb_type, 0); + int ctx = txfm_partition_context(xd->above_txfm_context + tx_col, + xd->left_txfm_context + tx_row, + mbmi->sb_type, tx_size); + const TX_SIZE plane_tx_size = mbmi->inter_tx_size[tx_row][tx_col]; + + if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return; + + if (tx_size == plane_tx_size) { + ++counts->txfm_partition[ctx][0]; + mbmi->tx_size = tx_size; + txfm_partition_update(xd->above_txfm_context + tx_col, + xd->left_txfm_context + tx_row, tx_size, tx_size); + } else { + const TX_SIZE sub_txs = sub_tx_size_map[tx_size]; + const int bs = tx_size_wide_unit[sub_txs]; + int i; + + ++counts->txfm_partition[ctx][1]; + ++x->txb_split_count; + + if (tx_size == TX_8X8) { + mbmi->inter_tx_size[tx_row][tx_col] = TX_4X4; + mbmi->tx_size = TX_4X4; + txfm_partition_update(xd->above_txfm_context + tx_col, + xd->left_txfm_context + tx_row, TX_4X4, tx_size); + return; + } + + for (i = 0; i < 4; ++i) { + int offsetr = (i >> 1) * bs; + int offsetc = (i & 0x01) * bs; + update_txfm_count(x, xd, counts, sub_txs, depth + 1, blk_row + offsetr, + blk_col + offsetc); + } + } +} + +static void tx_partition_count_update(const AV1_COMMON *const cm, MACROBLOCK *x, + BLOCK_SIZE plane_bsize, int mi_row, + int mi_col, FRAME_COUNTS *td_counts) { + MACROBLOCKD *xd = &x->e_mbd; + const int mi_width = block_size_wide[plane_bsize] >> tx_size_wide_log2[0]; + const int mi_height = block_size_high[plane_bsize] >> tx_size_wide_log2[0]; + TX_SIZE max_tx_size = get_vartx_max_txsize(&xd->mi[0]->mbmi, plane_bsize); + const int bh = tx_size_high_unit[max_tx_size]; + const int bw = tx_size_wide_unit[max_tx_size]; + int idx, idy; + + xd->above_txfm_context = cm->above_txfm_context + mi_col; + xd->left_txfm_context = + xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK); + + for (idy = 0; idy < mi_height; idy += bh) + for (idx = 0; idx < mi_width; idx += bw) + update_txfm_count(x, xd, td_counts, max_tx_size, mi_width != mi_height, + idy, idx); +} + +static void set_txfm_context(MACROBLOCKD *xd, TX_SIZE tx_size, int blk_row, + int blk_col) { + MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi; + const int tx_row = blk_row >> 1; + const int tx_col = blk_col >> 1; + const int max_blocks_high = max_block_high(xd, mbmi->sb_type, 0); + const int max_blocks_wide = max_block_wide(xd, mbmi->sb_type, 0); + const TX_SIZE plane_tx_size = mbmi->inter_tx_size[tx_row][tx_col]; + + if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return; + + if (tx_size == plane_tx_size) { + mbmi->tx_size = tx_size; + txfm_partition_update(xd->above_txfm_context + tx_col, + xd->left_txfm_context + tx_row, tx_size, tx_size); + + } else { + const TX_SIZE sub_txs = sub_tx_size_map[tx_size]; + const int bsl = tx_size_wide_unit[sub_txs]; + int i; + + if (tx_size == TX_8X8) { + mbmi->inter_tx_size[tx_row][tx_col] = TX_4X4; + mbmi->tx_size = TX_4X4; + txfm_partition_update(xd->above_txfm_context + tx_col, + xd->left_txfm_context + tx_row, TX_4X4, tx_size); + return; + } + + assert(bsl > 0); + for (i = 0; i < 4; ++i) { + int offsetr = (i >> 1) * bsl; + int offsetc = (i & 0x01) * bsl; + set_txfm_context(xd, sub_txs, blk_row + offsetr, blk_col + offsetc); + } + } +} + +static void tx_partition_set_contexts(const AV1_COMMON *const cm, + MACROBLOCKD *xd, BLOCK_SIZE plane_bsize, + int mi_row, int mi_col) { + const int mi_width = block_size_wide[plane_bsize] >> tx_size_wide_log2[0]; + const int mi_height = block_size_high[plane_bsize] >> tx_size_high_log2[0]; + TX_SIZE max_tx_size = get_vartx_max_txsize(&xd->mi[0]->mbmi, plane_bsize); + const int bh = tx_size_high_unit[max_tx_size]; + const int bw = tx_size_wide_unit[max_tx_size]; + int idx, idy; + + xd->above_txfm_context = cm->above_txfm_context + mi_col; + xd->left_txfm_context = + xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK); + + for (idy = 0; idy < mi_height; idy += bh) + for (idx = 0; idx < mi_width; idx += bw) + set_txfm_context(xd, max_tx_size, idy, idx); +} +#endif + +void av1_update_tx_type_count(const AV1_COMMON *cm, MACROBLOCKD *xd, +#if CONFIG_TXK_SEL + int block, int plane, +#endif + BLOCK_SIZE bsize, TX_SIZE tx_size, + FRAME_COUNTS *counts) { + MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi; + int is_inter = is_inter_block(mbmi); +#if !CONFIG_TXK_SEL + TX_TYPE tx_type = mbmi->tx_type; +#else + // Only y plane's tx_type is updated + if (plane > 0) return; + TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size); +#endif +#if CONFIG_EXT_TX + if (get_ext_tx_types(tx_size, bsize, is_inter, cm->reduced_tx_set_used) > 1 && + cm->base_qindex > 0 && !mbmi->skip && + !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) { + const int eset = + get_ext_tx_set(tx_size, bsize, is_inter, cm->reduced_tx_set_used); + if (eset > 0) { + if (is_inter) { + ++counts->inter_ext_tx[eset][txsize_sqr_map[tx_size]][tx_type]; + } else { + ++counts->intra_ext_tx[eset][txsize_sqr_map[tx_size]][mbmi->mode] + [tx_type]; + } + } + } +#else + (void)bsize; + if (tx_size < TX_32X32 && + ((!cm->seg.enabled && cm->base_qindex > 0) || + (cm->seg.enabled && xd->qindex[mbmi->segment_id] > 0)) && + !mbmi->skip && + !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) { + if (is_inter) { + ++counts->inter_ext_tx[tx_size][tx_type]; + } else { + ++counts->intra_ext_tx[tx_size][intra_mode_to_tx_type_context[mbmi->mode]] + [tx_type]; + } + } +#endif // CONFIG_EXT_TX +} + +static void encode_superblock(const AV1_COMP *const cpi, ThreadData *td, + TOKENEXTRA **t, RUN_TYPE dry_run, int mi_row, + int mi_col, BLOCK_SIZE bsize, + PICK_MODE_CONTEXT *ctx, int *rate) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + MODE_INFO **mi_8x8 = xd->mi; + MODE_INFO *mi = mi_8x8[0]; + MB_MODE_INFO *mbmi = &mi->mbmi; + const int seg_skip = + segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP); + const int mis = cm->mi_stride; + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + const int is_inter = is_inter_block(mbmi); +#if CONFIG_CB4X4 + const BLOCK_SIZE block_size = bsize; +#else + const BLOCK_SIZE block_size = AOMMAX(bsize, BLOCK_8X8); +#endif + +#if CONFIG_PVQ + x->pvq_speed = 0; + x->pvq_coded = (dry_run == OUTPUT_ENABLED) ? 1 : 0; +#endif +#if CONFIG_CFL + x->cfl_store_y = (dry_run == OUTPUT_ENABLED) ? 1 : 0; +#endif + + if (!is_inter) { + int plane; + mbmi->skip = 1; + for (plane = 0; plane < MAX_MB_PLANE; ++plane) { + av1_encode_intra_block_plane((AV1_COMMON *)cm, x, block_size, plane, 1, + mi_row, mi_col); + } + if (!dry_run) { + sum_intra_stats(td->counts, xd, mi, xd->above_mi, xd->left_mi, + frame_is_intra_only(cm), mi_row, mi_col); + } +#if CONFIG_PALETTE + if (bsize >= BLOCK_8X8 && !dry_run) { + for (plane = 0; plane <= 1; ++plane) { + if (mbmi->palette_mode_info.palette_size[plane] > 0) { + mbmi->palette_mode_info.palette_first_color_idx[plane] = + xd->plane[plane].color_index_map[0]; + // TODO(huisu): this increases the use of token buffer. Needs stretch + // test to verify. + av1_tokenize_palette_sb(cpi, td, plane, t, dry_run, bsize, rate); + } + } + } +#endif // CONFIG_PALETTE +#if CONFIG_VAR_TX + mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size); +#endif +#if CONFIG_LV_MAP + av1_update_txb_context(cpi, td, dry_run, block_size, rate, mi_row, mi_col); +#else // CONFIG_LV_MAP + av1_tokenize_sb(cpi, td, t, dry_run, block_size, rate, mi_row, mi_col); +#endif // CONFIG_LV_MAP + } else { + int ref; + const int is_compound = has_second_ref(mbmi); + + set_ref_ptrs(cm, xd, mbmi->ref_frame[0], mbmi->ref_frame[1]); + for (ref = 0; ref < 1 + is_compound; ++ref) { + YV12_BUFFER_CONFIG *cfg = get_ref_frame_buffer(cpi, mbmi->ref_frame[ref]); +#if CONFIG_INTRABC + assert(IMPLIES(!is_intrabc_block(mbmi), cfg)); +#else + assert(cfg != NULL); +#endif // !CONFIG_INTRABC + av1_setup_pre_planes(xd, ref, cfg, mi_row, mi_col, + &xd->block_refs[ref]->sf); + } + if (!(cpi->sf.reuse_inter_pred_sby && ctx->pred_pixel_ready) || seg_skip) + av1_build_inter_predictors_sby(xd, mi_row, mi_col, NULL, block_size); + + av1_build_inter_predictors_sbuv(xd, mi_row, mi_col, NULL, block_size); +#if CONFIG_MOTION_VAR + if (mbmi->motion_mode == OBMC_CAUSAL) { +#if CONFIG_NCOBMC + if (dry_run == OUTPUT_ENABLED) + av1_build_ncobmc_inter_predictors_sb(cm, xd, mi_row, mi_col); + else +#endif + av1_build_obmc_inter_predictors_sb(cm, xd, mi_row, mi_col); + } +#endif // CONFIG_MOTION_VAR + + av1_encode_sb((AV1_COMMON *)cm, x, block_size, mi_row, mi_col); +#if CONFIG_VAR_TX + if (mbmi->skip) mbmi->min_tx_size = get_min_tx_size(mbmi->tx_size); + av1_tokenize_sb_vartx(cpi, td, t, dry_run, mi_row, mi_col, block_size, + rate); +#else +#if CONFIG_LV_MAP + av1_update_txb_context(cpi, td, dry_run, block_size, rate, mi_row, mi_col); +#else // CONFIG_LV_MAP + av1_tokenize_sb(cpi, td, t, dry_run, block_size, rate, mi_row, mi_col); +#endif // CONFIG_LV_MAP +#endif + } + + if (!dry_run) { +#if CONFIG_VAR_TX + TX_SIZE tx_size = + is_inter && !mbmi->skip ? mbmi->min_tx_size : mbmi->tx_size; +#else + TX_SIZE tx_size = mbmi->tx_size; +#endif + if (cm->tx_mode == TX_MODE_SELECT && !xd->lossless[mbmi->segment_id] && +#if CONFIG_CB4X4 && (CONFIG_VAR_TX || CONFIG_EXT_TX) && CONFIG_RECT_TX + mbmi->sb_type > BLOCK_4X4 && +#else + mbmi->sb_type >= BLOCK_8X8 && +#endif + !(is_inter && (mbmi->skip || seg_skip))) { +#if CONFIG_VAR_TX + if (is_inter) { + tx_partition_count_update(cm, x, bsize, mi_row, mi_col, td->counts); + } else { + const int tx_size_ctx = get_tx_size_context(xd); + const int tx_size_cat = is_inter ? inter_tx_size_cat_lookup[bsize] + : intra_tx_size_cat_lookup[bsize]; + const TX_SIZE coded_tx_size = txsize_sqr_up_map[tx_size]; + const int depth = tx_size_to_depth(coded_tx_size); + ++td->counts->tx_size[tx_size_cat][tx_size_ctx][depth]; + if (tx_size != max_txsize_lookup[bsize]) ++x->txb_split_count; + } +#else + const int tx_size_ctx = get_tx_size_context(xd); + const int tx_size_cat = is_inter ? inter_tx_size_cat_lookup[bsize] + : intra_tx_size_cat_lookup[bsize]; + const TX_SIZE coded_tx_size = txsize_sqr_up_map[tx_size]; + const int depth = tx_size_to_depth(coded_tx_size); + + ++td->counts->tx_size[tx_size_cat][tx_size_ctx][depth]; +#endif +#if CONFIG_EXT_TX && CONFIG_RECT_TX + assert(IMPLIES(is_rect_tx(tx_size), is_rect_tx_allowed(xd, mbmi))); +#endif // CONFIG_EXT_TX && CONFIG_RECT_TX + } else { + int i, j; + TX_SIZE intra_tx_size; + // The new intra coding scheme requires no change of transform size + if (is_inter) { + if (xd->lossless[mbmi->segment_id]) { + intra_tx_size = TX_4X4; + } else { + intra_tx_size = tx_size_from_tx_mode(bsize, cm->tx_mode, 1); + } + } else { +#if CONFIG_EXT_TX && CONFIG_RECT_TX + intra_tx_size = tx_size; +#else + intra_tx_size = (bsize >= BLOCK_8X8) ? tx_size : TX_4X4; +#endif // CONFIG_EXT_TX && CONFIG_RECT_TX + } +#if CONFIG_EXT_TX && CONFIG_RECT_TX + ++td->counts->tx_size_implied[max_txsize_lookup[bsize]] + [txsize_sqr_up_map[tx_size]]; +#endif // CONFIG_EXT_TX && CONFIG_RECT_TX + + for (j = 0; j < mi_height; j++) + for (i = 0; i < mi_width; i++) + if (mi_col + i < cm->mi_cols && mi_row + j < cm->mi_rows) + mi_8x8[mis * j + i]->mbmi.tx_size = intra_tx_size; + +#if CONFIG_VAR_TX + mbmi->min_tx_size = get_min_tx_size(intra_tx_size); + if (intra_tx_size != max_txsize_lookup[bsize]) ++x->txb_split_count; +#endif + } + + ++td->counts->tx_size_totals[txsize_sqr_map[tx_size]]; + ++td->counts + ->tx_size_totals[txsize_sqr_map[get_uv_tx_size(mbmi, &xd->plane[1])]]; +#if !CONFIG_TXK_SEL + av1_update_tx_type_count(cm, xd, bsize, tx_size, td->counts); +#endif + } + +#if CONFIG_VAR_TX + if (cm->tx_mode == TX_MODE_SELECT && +#if CONFIG_CB4X4 + mbmi->sb_type > BLOCK_4X4 && +#else + mbmi->sb_type >= BLOCK_8X8 && +#endif + is_inter && !(mbmi->skip || seg_skip)) { + if (dry_run) tx_partition_set_contexts(cm, xd, bsize, mi_row, mi_col); + } else { + TX_SIZE tx_size = mbmi->tx_size; + // The new intra coding scheme requires no change of transform size + if (is_inter) + tx_size = tx_size_from_tx_mode(bsize, cm->tx_mode, is_inter); + else + tx_size = (bsize > BLOCK_4X4) ? tx_size : TX_4X4; + mbmi->tx_size = tx_size; + set_txfm_ctxs(tx_size, xd->n8_w, xd->n8_h, (mbmi->skip || seg_skip), xd); + } +#endif // CONFIG_VAR_TX +} + +#if CONFIG_SUPERTX +static int check_intra_b(PICK_MODE_CONTEXT *ctx) { + if (!is_inter_mode((&ctx->mic)->mbmi.mode)) return 1; +#if CONFIG_EXT_INTER + if (ctx->mic.mbmi.ref_frame[1] == INTRA_FRAME) return 1; +#endif // CONFIG_EXT_INTER + return 0; +} + +static int check_intra_sb(const AV1_COMP *const cpi, const TileInfo *const tile, + int mi_row, int mi_col, BLOCK_SIZE bsize, + PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + const int hbs = mi_size_wide[bsize] / 2; + const PARTITION_TYPE partition = pc_tree->partitioning; + const BLOCK_SIZE subsize = get_subsize(bsize, partition); +#if CONFIG_EXT_PARTITION_TYPES + int i; +#endif +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; +#endif + +#if !CONFIG_CB4X4 + assert(bsize >= BLOCK_8X8); +#endif + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return 1; + + switch (partition) { + case PARTITION_NONE: return check_intra_b(&pc_tree->none); break; + case PARTITION_VERT: + if (check_intra_b(&pc_tree->vertical[0])) return 1; + if (mi_col + hbs < cm->mi_cols && (bsize > BLOCK_8X8 || unify_bsize)) { + if (check_intra_b(&pc_tree->vertical[1])) return 1; + } + break; + case PARTITION_HORZ: + if (check_intra_b(&pc_tree->horizontal[0])) return 1; + if (mi_row + hbs < cm->mi_rows && (bsize > BLOCK_8X8 || unify_bsize)) { + if (check_intra_b(&pc_tree->horizontal[1])) return 1; + } + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + if (check_intra_b(pc_tree->leaf_split[0])) return 1; + } else { + if (check_intra_sb(cpi, tile, mi_row, mi_col, subsize, + pc_tree->split[0])) + return 1; + if (check_intra_sb(cpi, tile, mi_row, mi_col + hbs, subsize, + pc_tree->split[1])) + return 1; + if (check_intra_sb(cpi, tile, mi_row + hbs, mi_col, subsize, + pc_tree->split[2])) + return 1; + if (check_intra_sb(cpi, tile, mi_row + hbs, mi_col + hbs, subsize, + pc_tree->split[3])) + return 1; + } + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + for (i = 0; i < 3; i++) { + if (check_intra_b(&pc_tree->horizontala[i])) return 1; + } + break; + case PARTITION_HORZ_B: + for (i = 0; i < 3; i++) { + if (check_intra_b(&pc_tree->horizontalb[i])) return 1; + } + break; + case PARTITION_VERT_A: + for (i = 0; i < 3; i++) { + if (check_intra_b(&pc_tree->verticala[i])) return 1; + } + break; + case PARTITION_VERT_B: + for (i = 0; i < 3; i++) { + if (check_intra_b(&pc_tree->verticalb[i])) return 1; + } + break; +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0); + } + return 0; +} + +static int check_supertx_b(TX_SIZE supertx_size, PICK_MODE_CONTEXT *ctx) { + return ctx->mic.mbmi.tx_size == supertx_size; +} + +static int check_supertx_sb(BLOCK_SIZE bsize, TX_SIZE supertx_size, + PC_TREE *pc_tree) { + PARTITION_TYPE partition; + BLOCK_SIZE subsize; +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; +#endif + + partition = pc_tree->partitioning; + subsize = get_subsize(bsize, partition); + switch (partition) { + case PARTITION_NONE: return check_supertx_b(supertx_size, &pc_tree->none); + case PARTITION_VERT: + return check_supertx_b(supertx_size, &pc_tree->vertical[0]); + case PARTITION_HORZ: + return check_supertx_b(supertx_size, &pc_tree->horizontal[0]); + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) + return check_supertx_b(supertx_size, pc_tree->leaf_split[0]); + else + return check_supertx_sb(subsize, supertx_size, pc_tree->split[0]); +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + return check_supertx_b(supertx_size, &pc_tree->horizontala[0]); + case PARTITION_HORZ_B: + return check_supertx_b(supertx_size, &pc_tree->horizontalb[0]); + case PARTITION_VERT_A: + return check_supertx_b(supertx_size, &pc_tree->verticala[0]); + case PARTITION_VERT_B: + return check_supertx_b(supertx_size, &pc_tree->verticalb[0]); +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0); return 0; + } +} + +static void predict_superblock(const AV1_COMP *const cpi, ThreadData *td, +#if CONFIG_EXT_INTER + int mi_row_ori, int mi_col_ori, +#endif // CONFIG_EXT_INTER + int mi_row_pred, int mi_col_pred, + BLOCK_SIZE bsize_pred, int b_sub8x8, int block) { + // Used in supertx + // (mi_row_ori, mi_col_ori): location for mv + // (mi_row_pred, mi_col_pred, bsize_pred): region to predict + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + MODE_INFO *mi_8x8 = xd->mi[0]; + MODE_INFO *mi = mi_8x8; + MB_MODE_INFO *mbmi = &mi->mbmi; + int ref; + const int is_compound = has_second_ref(mbmi); + + set_ref_ptrs(cm, xd, mbmi->ref_frame[0], mbmi->ref_frame[1]); + + for (ref = 0; ref < 1 + is_compound; ++ref) { + YV12_BUFFER_CONFIG *cfg = get_ref_frame_buffer(cpi, mbmi->ref_frame[ref]); + av1_setup_pre_planes(xd, ref, cfg, mi_row_pred, mi_col_pred, + &xd->block_refs[ref]->sf); + } + + if (!b_sub8x8) + av1_build_inter_predictors_sb_extend(xd, +#if CONFIG_EXT_INTER + mi_row_ori, mi_col_ori, +#endif // CONFIG_EXT_INTER + mi_row_pred, mi_col_pred, bsize_pred); + else + av1_build_inter_predictors_sb_sub8x8_extend(xd, +#if CONFIG_EXT_INTER + mi_row_ori, mi_col_ori, +#endif // CONFIG_EXT_INTER + mi_row_pred, mi_col_pred, + bsize_pred, block); +} + +static void predict_b_extend(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int block, + int mi_row_ori, int mi_col_ori, int mi_row_pred, + int mi_col_pred, int mi_row_top, int mi_col_top, + uint8_t *dst_buf[3], int dst_stride[3], + BLOCK_SIZE bsize_top, BLOCK_SIZE bsize_pred, + RUN_TYPE dry_run, int b_sub8x8, int bextend) { + // Used in supertx + // (mi_row_ori, mi_col_ori): location for mv + // (mi_row_pred, mi_col_pred, bsize_pred): region to predict + // (mi_row_top, mi_col_top, bsize_top): region of the top partition size + // block: sub location of sub8x8 blocks + // b_sub8x8: 1: ori is sub8x8; 0: ori is not sub8x8 + // bextend: 1: region to predict is an extension of ori; 0: not + + MACROBLOCK *const x = &td->mb; + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCKD *const xd = &x->e_mbd; + int r = (mi_row_pred - mi_row_top) * MI_SIZE; + int c = (mi_col_pred - mi_col_top) * MI_SIZE; + const int mi_width_top = mi_size_wide[bsize_top]; + const int mi_height_top = mi_size_high[bsize_top]; + + if (mi_row_pred < mi_row_top || mi_col_pred < mi_col_top || + mi_row_pred >= mi_row_top + mi_height_top || + mi_col_pred >= mi_col_top + mi_width_top || mi_row_pred >= cm->mi_rows || + mi_col_pred >= cm->mi_cols) + return; + + set_offsets_extend(cpi, td, tile, mi_row_pred, mi_col_pred, mi_row_ori, + mi_col_ori, bsize_pred); + xd->plane[0].dst.stride = dst_stride[0]; + xd->plane[1].dst.stride = dst_stride[1]; + xd->plane[2].dst.stride = dst_stride[2]; + xd->plane[0].dst.buf = dst_buf[0] + + (r >> xd->plane[0].subsampling_y) * dst_stride[0] + + (c >> xd->plane[0].subsampling_x); + xd->plane[1].dst.buf = dst_buf[1] + + (r >> xd->plane[1].subsampling_y) * dst_stride[1] + + (c >> xd->plane[1].subsampling_x); + xd->plane[2].dst.buf = dst_buf[2] + + (r >> xd->plane[2].subsampling_y) * dst_stride[2] + + (c >> xd->plane[2].subsampling_x); + + predict_superblock(cpi, td, +#if CONFIG_EXT_INTER + mi_row_ori, mi_col_ori, +#endif // CONFIG_EXT_INTER + mi_row_pred, mi_col_pred, bsize_pred, b_sub8x8, block); + + if (!dry_run && !bextend) + update_stats(&cpi->common, td, mi_row_pred, mi_col_pred, 1); +} + +static void extend_dir(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int block, BLOCK_SIZE bsize, + BLOCK_SIZE top_bsize, int mi_row, int mi_col, + int mi_row_top, int mi_col_top, RUN_TYPE dry_run, + uint8_t *dst_buf[3], int dst_stride[3], int dir) { + // dir: 0-lower, 1-upper, 2-left, 3-right + // 4-lowerleft, 5-upperleft, 6-lowerright, 7-upperright + MACROBLOCKD *xd = &td->mb.e_mbd; + const int mi_width = mi_size_wide[bsize]; + const int mi_height = mi_size_high[bsize]; + int xss = xd->plane[1].subsampling_x; + int yss = xd->plane[1].subsampling_y; +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; +#endif + int b_sub8x8 = (bsize < BLOCK_8X8) && !unify_bsize ? 1 : 0; + int wide_unit, high_unit; + int i, j; + int ext_offset = 0; + + BLOCK_SIZE extend_bsize; + int mi_row_pred, mi_col_pred; + + if (dir == 0 || dir == 1) { // lower and upper + extend_bsize = + (mi_width == mi_size_wide[BLOCK_8X8] || bsize < BLOCK_8X8 || xss < yss) + ? BLOCK_8X8 + : BLOCK_16X8; + +#if CONFIG_CB4X4 + if (bsize < BLOCK_8X8) { + extend_bsize = BLOCK_4X4; + ext_offset = mi_size_wide[BLOCK_8X8]; + } +#endif + wide_unit = mi_size_wide[extend_bsize]; + high_unit = mi_size_high[extend_bsize]; + + mi_row_pred = mi_row + ((dir == 0) ? mi_height : -(mi_height + ext_offset)); + mi_col_pred = mi_col; + + for (j = 0; j < mi_height + ext_offset; j += high_unit) + for (i = 0; i < mi_width + ext_offset; i += wide_unit) + predict_b_extend(cpi, td, tile, block, mi_row, mi_col, mi_row_pred + j, + mi_col_pred + i, mi_row_top, mi_col_top, dst_buf, + dst_stride, top_bsize, extend_bsize, dry_run, b_sub8x8, + 1); + } else if (dir == 2 || dir == 3) { // left and right + extend_bsize = + (mi_height == mi_size_high[BLOCK_8X8] || bsize < BLOCK_8X8 || yss < xss) + ? BLOCK_8X8 + : BLOCK_8X16; +#if CONFIG_CB4X4 + if (bsize < BLOCK_8X8) { + extend_bsize = BLOCK_4X4; + ext_offset = mi_size_wide[BLOCK_8X8]; + } +#endif + wide_unit = mi_size_wide[extend_bsize]; + high_unit = mi_size_high[extend_bsize]; + + mi_row_pred = mi_row; + mi_col_pred = mi_col + ((dir == 3) ? mi_width : -(mi_width + ext_offset)); + + for (j = 0; j < mi_height + ext_offset; j += high_unit) + for (i = 0; i < mi_width + ext_offset; i += wide_unit) + predict_b_extend(cpi, td, tile, block, mi_row, mi_col, mi_row_pred + j, + mi_col_pred + i, mi_row_top, mi_col_top, dst_buf, + dst_stride, top_bsize, extend_bsize, dry_run, b_sub8x8, + 1); + } else { + extend_bsize = BLOCK_8X8; +#if CONFIG_CB4X4 + if (bsize < BLOCK_8X8) { + extend_bsize = BLOCK_4X4; + ext_offset = mi_size_wide[BLOCK_8X8]; + } +#endif + wide_unit = mi_size_wide[extend_bsize]; + high_unit = mi_size_high[extend_bsize]; + + mi_row_pred = mi_row + ((dir == 4 || dir == 6) ? mi_height + : -(mi_height + ext_offset)); + mi_col_pred = + mi_col + ((dir == 6 || dir == 7) ? mi_width : -(mi_width + ext_offset)); + + for (j = 0; j < mi_height + ext_offset; j += high_unit) + for (i = 0; i < mi_width + ext_offset; i += wide_unit) + predict_b_extend(cpi, td, tile, block, mi_row, mi_col, mi_row_pred + j, + mi_col_pred + i, mi_row_top, mi_col_top, dst_buf, + dst_stride, top_bsize, extend_bsize, dry_run, b_sub8x8, + 1); + } +} + +static void extend_all(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int block, BLOCK_SIZE bsize, + BLOCK_SIZE top_bsize, int mi_row, int mi_col, + int mi_row_top, int mi_col_top, RUN_TYPE dry_run, + uint8_t *dst_buf[3], int dst_stride[3]) { + assert(block >= 0 && block < 4); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 0); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 1); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 2); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 3); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 4); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 5); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 6); + extend_dir(cpi, td, tile, block, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride, 7); +} + +// This function generates prediction for multiple blocks, between which +// discontinuity around boundary is reduced by smoothing masks. The basic +// smoothing mask is a soft step function along horz/vert direction. In more +// complicated case when a block is split into 4 subblocks, the basic mask is +// first applied to neighboring subblocks (2 pairs) in horizontal direction and +// then applied to the 2 masked prediction mentioned above in vertical direction +// If the block is split into more than one level, at every stage, masked +// prediction is stored in dst_buf[] passed from higher level. +static void predict_sb_complex(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, + int mi_col, int mi_row_top, int mi_col_top, + RUN_TYPE dry_run, BLOCK_SIZE bsize, + BLOCK_SIZE top_bsize, uint8_t *dst_buf[3], + int dst_stride[3], PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + const int hbs = mi_size_wide[bsize] / 2; + const int is_partition_root = bsize >= BLOCK_8X8; + const int ctx = is_partition_root + ? partition_plane_context(xd, mi_row, mi_col, +#if CONFIG_UNPOISON_PARTITION_CTX + mi_row + hbs < cm->mi_rows, + mi_col + hbs < cm->mi_cols, +#endif + bsize) + : -1; + const PARTITION_TYPE partition = pc_tree->partitioning; + const BLOCK_SIZE subsize = get_subsize(bsize, partition); +#if CONFIG_EXT_PARTITION_TYPES + const BLOCK_SIZE bsize2 = get_subsize(bsize, PARTITION_SPLIT); +#endif + + int i; + uint8_t *dst_buf1[3], *dst_buf2[3], *dst_buf3[3]; + DECLARE_ALIGNED(16, uint8_t, tmp_buf1[MAX_MB_PLANE * MAX_TX_SQUARE * 2]); + DECLARE_ALIGNED(16, uint8_t, tmp_buf2[MAX_MB_PLANE * MAX_TX_SQUARE * 2]); + DECLARE_ALIGNED(16, uint8_t, tmp_buf3[MAX_MB_PLANE * MAX_TX_SQUARE * 2]); + int dst_stride1[3] = { MAX_TX_SIZE, MAX_TX_SIZE, MAX_TX_SIZE }; + int dst_stride2[3] = { MAX_TX_SIZE, MAX_TX_SIZE, MAX_TX_SIZE }; + int dst_stride3[3] = { MAX_TX_SIZE, MAX_TX_SIZE, MAX_TX_SIZE }; +#if CONFIG_CB4X4 + const int unify_bsize = 1; +#else + const int unify_bsize = 0; + assert(bsize >= BLOCK_8X8); +#endif + + if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return; + +#if CONFIG_HIGHBITDEPTH + if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) { + int len = sizeof(uint16_t); + dst_buf1[0] = CONVERT_TO_BYTEPTR(tmp_buf1); + dst_buf1[1] = CONVERT_TO_BYTEPTR(tmp_buf1 + MAX_TX_SQUARE * len); + dst_buf1[2] = CONVERT_TO_BYTEPTR(tmp_buf1 + 2 * MAX_TX_SQUARE * len); + dst_buf2[0] = CONVERT_TO_BYTEPTR(tmp_buf2); + dst_buf2[1] = CONVERT_TO_BYTEPTR(tmp_buf2 + MAX_TX_SQUARE * len); + dst_buf2[2] = CONVERT_TO_BYTEPTR(tmp_buf2 + 2 * MAX_TX_SQUARE * len); + dst_buf3[0] = CONVERT_TO_BYTEPTR(tmp_buf3); + dst_buf3[1] = CONVERT_TO_BYTEPTR(tmp_buf3 + MAX_TX_SQUARE * len); + dst_buf3[2] = CONVERT_TO_BYTEPTR(tmp_buf3 + 2 * MAX_TX_SQUARE * len); + } else { +#endif // CONFIG_HIGHBITDEPTH + dst_buf1[0] = tmp_buf1; + dst_buf1[1] = tmp_buf1 + MAX_TX_SQUARE; + dst_buf1[2] = tmp_buf1 + 2 * MAX_TX_SQUARE; + dst_buf2[0] = tmp_buf2; + dst_buf2[1] = tmp_buf2 + MAX_TX_SQUARE; + dst_buf2[2] = tmp_buf2 + 2 * MAX_TX_SQUARE; + dst_buf3[0] = tmp_buf3; + dst_buf3[1] = tmp_buf3 + MAX_TX_SQUARE; + dst_buf3[2] = tmp_buf3 + 2 * MAX_TX_SQUARE; +#if CONFIG_HIGHBITDEPTH + } +#endif // CONFIG_HIGHBITDEPTH + + if (!dry_run && ctx >= 0 && bsize < top_bsize) { + // Explicitly cast away const. + FRAME_COUNTS *const frame_counts = (FRAME_COUNTS *)&cm->counts; + frame_counts->partition[ctx][partition]++; + } + + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + } + + switch (partition) { + case PARTITION_NONE: + assert(bsize < top_bsize); + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + bsize, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize, top_bsize, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, dst_buf, dst_stride); + break; + case PARTITION_HORZ: + if (bsize == BLOCK_8X8 && !unify_bsize) { + // Fisrt half + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + BLOCK_8X8, dry_run, 1, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + + // Second half + predict_b_extend(cpi, td, tile, 2, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf1, dst_stride1, + top_bsize, BLOCK_8X8, dry_run, 1, 1); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 2, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + + // Smooth + xd->plane[0].dst.buf = dst_buf[0]; + xd->plane[0].dst.stride = dst_stride[0]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[0], dst_stride[0], dst_buf1[0], dst_stride1[0], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_HORZ, + 0); + } else { + // First half + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride, 0); + + if (mi_row + hbs < cm->mi_rows) { + // Second half + predict_b_extend(cpi, td, tile, 0, mi_row + hbs, mi_col, mi_row + hbs, + mi_col, mi_row_top, mi_col_top, dst_buf1, + dst_stride1, top_bsize, subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row + hbs, + mi_col, mi_row_top, mi_col_top, dry_run, dst_buf1, + dst_stride1); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row + hbs, + mi_col, mi_row_top, mi_col_top, dry_run, dst_buf1, + dst_stride1, 1); + + // Smooth + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_HORZ, i); + } + } + } + break; + case PARTITION_VERT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + // First half + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + BLOCK_8X8, dry_run, 1, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + + // Second half + predict_b_extend(cpi, td, tile, 1, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf1, dst_stride1, + top_bsize, BLOCK_8X8, dry_run, 1, 1); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 1, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + + // Smooth + xd->plane[0].dst.buf = dst_buf[0]; + xd->plane[0].dst.stride = dst_stride[0]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[0], dst_stride[0], dst_buf1[0], dst_stride1[0], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_VERT, + 0); + } else { + // bsize: not important, not useful + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride, 3); + + if (mi_col + hbs < cm->mi_cols) { + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col + hbs, mi_row, + mi_col + hbs, mi_row_top, mi_col_top, dst_buf1, + dst_stride1, top_bsize, subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, + mi_col + hbs, mi_row_top, mi_col_top, dry_run, dst_buf1, + dst_stride1); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row, + mi_col + hbs, mi_row_top, mi_col_top, dry_run, dst_buf1, + dst_stride1, 2); + + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_VERT, i); + } + } + } + break; + case PARTITION_SPLIT: + if (bsize == BLOCK_8X8 && !unify_bsize) { + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + BLOCK_8X8, dry_run, 1, 0); + predict_b_extend(cpi, td, tile, 1, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf1, dst_stride1, + top_bsize, BLOCK_8X8, dry_run, 1, 1); + predict_b_extend(cpi, td, tile, 2, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf2, dst_stride2, + top_bsize, BLOCK_8X8, dry_run, 1, 1); + predict_b_extend(cpi, td, tile, 3, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf3, dst_stride3, + top_bsize, BLOCK_8X8, dry_run, 1, 1); + + if (bsize < top_bsize) { + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + extend_all(cpi, td, tile, 1, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + extend_all(cpi, td, tile, 2, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf2, dst_stride2); + extend_all(cpi, td, tile, 3, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf3, dst_stride3); + } + } else { + predict_sb_complex(cpi, td, tile, mi_row, mi_col, mi_row_top, + mi_col_top, dry_run, subsize, top_bsize, dst_buf, + dst_stride, pc_tree->split[0]); + if (mi_row < cm->mi_rows && mi_col + hbs < cm->mi_cols) + predict_sb_complex(cpi, td, tile, mi_row, mi_col + hbs, mi_row_top, + mi_col_top, dry_run, subsize, top_bsize, dst_buf1, + dst_stride1, pc_tree->split[1]); + if (mi_row + hbs < cm->mi_rows && mi_col < cm->mi_cols) + predict_sb_complex(cpi, td, tile, mi_row + hbs, mi_col, mi_row_top, + mi_col_top, dry_run, subsize, top_bsize, dst_buf2, + dst_stride2, pc_tree->split[2]); + if (mi_row + hbs < cm->mi_rows && mi_col + hbs < cm->mi_cols) + predict_sb_complex(cpi, td, tile, mi_row + hbs, mi_col + hbs, + mi_row_top, mi_col_top, dry_run, subsize, + top_bsize, dst_buf3, dst_stride3, + pc_tree->split[3]); + } + for (i = 0; i < MAX_MB_PLANE; i++) { +#if !CONFIG_CB4X4 + if (bsize == BLOCK_8X8 && i != 0) + continue; // Skip <4x4 chroma smoothing +#endif + if (mi_row < cm->mi_rows && mi_col + hbs < cm->mi_cols) { + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_VERT, i); + if (mi_row + hbs < cm->mi_rows) { + av1_build_masked_inter_predictor_complex( + xd, dst_buf2[i], dst_stride2[i], dst_buf3[i], dst_stride3[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_VERT, i); + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf2[i], dst_stride2[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_HORZ, i); + } + } else if (mi_row + hbs < cm->mi_rows && mi_col < cm->mi_cols) { + if (bsize == BLOCK_8X8 && i != 0) + continue; // Skip <4x4 chroma smoothing + + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf2[i], dst_stride2[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_HORZ, i); + } + } + break; +#if CONFIG_EXT_PARTITION_TYPES + case PARTITION_HORZ_A: + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col + hbs, mi_row, + mi_col + hbs, mi_row_top, mi_col_top, dst_buf1, + dst_stride1, top_bsize, bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row, mi_col + hbs, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + + predict_b_extend(cpi, td, tile, 0, mi_row + hbs, mi_col, mi_row + hbs, + mi_col, mi_row_top, mi_col_top, dst_buf2, dst_stride2, + top_bsize, subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row + hbs, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf2, dst_stride2); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row + hbs, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf2, dst_stride2, 1); + + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_VERT, + i); + } + for (i = 0; i < MAX_MB_PLANE; i++) { + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf2[i], dst_stride2[i], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_HORZ, + i); + } + + break; + case PARTITION_VERT_A: + + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + + predict_b_extend(cpi, td, tile, 0, mi_row + hbs, mi_col, mi_row + hbs, + mi_col, mi_row_top, mi_col_top, dst_buf1, dst_stride1, + top_bsize, bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row + hbs, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col + hbs, mi_row, + mi_col + hbs, mi_row_top, mi_col_top, dst_buf2, + dst_stride2, top_bsize, subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col + hbs, + mi_row_top, mi_col_top, dry_run, dst_buf2, dst_stride2); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col + hbs, + mi_row_top, mi_col_top, dry_run, dst_buf2, dst_stride2, 2); + + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_HORZ, + i); + } + for (i = 0; i < MAX_MB_PLANE; i++) { + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf2[i], dst_stride2[i], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_VERT, + i); + } + break; + case PARTITION_HORZ_B: + + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride, 0); + + predict_b_extend(cpi, td, tile, 0, mi_row + hbs, mi_col, mi_row + hbs, + mi_col, mi_row_top, mi_col_top, dst_buf1, dst_stride1, + top_bsize, bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row + hbs, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + + predict_b_extend(cpi, td, tile, 0, mi_row + hbs, mi_col + hbs, + mi_row + hbs, mi_col + hbs, mi_row_top, mi_col_top, + dst_buf2, dst_stride2, top_bsize, bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row + hbs, + mi_col + hbs, mi_row_top, mi_col_top, dry_run, dst_buf2, + dst_stride2); + + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf1[i]; + xd->plane[i].dst.stride = dst_stride1[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf1[i], dst_stride1[i], dst_buf2[i], dst_stride2[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_VERT, i); + } + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_HORZ, + i); + } + break; + case PARTITION_VERT_B: + + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col, mi_row, mi_col, + mi_row_top, mi_col_top, dst_buf, dst_stride, top_bsize, + subsize, dry_run, 0, 0); + if (bsize < top_bsize) + extend_all(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride); + else + extend_dir(cpi, td, tile, 0, subsize, top_bsize, mi_row, mi_col, + mi_row_top, mi_col_top, dry_run, dst_buf, dst_stride, 3); + + predict_b_extend(cpi, td, tile, 0, mi_row, mi_col + hbs, mi_row, + mi_col + hbs, mi_row_top, mi_col_top, dst_buf1, + dst_stride1, top_bsize, bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row, mi_col + hbs, + mi_row_top, mi_col_top, dry_run, dst_buf1, dst_stride1); + + predict_b_extend(cpi, td, tile, 0, mi_row + hbs, mi_col + hbs, + mi_row + hbs, mi_col + hbs, mi_row_top, mi_col_top, + dst_buf2, dst_stride2, top_bsize, bsize2, dry_run, 0, 0); + extend_all(cpi, td, tile, 0, bsize2, top_bsize, mi_row + hbs, + mi_col + hbs, mi_row_top, mi_col_top, dry_run, dst_buf2, + dst_stride2); + + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf1[i]; + xd->plane[i].dst.stride = dst_stride1[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf1[i], dst_stride1[i], dst_buf2[i], dst_stride2[i], + mi_row, mi_col, mi_row_top, mi_col_top, bsize, top_bsize, + PARTITION_HORZ, i); + } + for (i = 0; i < MAX_MB_PLANE; i++) { + xd->plane[i].dst.buf = dst_buf[i]; + xd->plane[i].dst.stride = dst_stride[i]; + av1_build_masked_inter_predictor_complex( + xd, dst_buf[i], dst_stride[i], dst_buf1[i], dst_stride1[i], mi_row, + mi_col, mi_row_top, mi_col_top, bsize, top_bsize, PARTITION_VERT, + i); + } + break; +#endif // CONFIG_EXT_PARTITION_TYPES + default: assert(0); + } + +#if CONFIG_EXT_PARTITION_TYPES + if (bsize < top_bsize) + update_ext_partition_context(xd, mi_row, mi_col, subsize, bsize, partition); +#else + if (bsize < top_bsize && (partition != PARTITION_SPLIT || bsize == BLOCK_8X8)) + update_partition_context(xd, mi_row, mi_col, subsize, bsize); +#endif // CONFIG_EXT_PARTITION_TYPES +} + +static void rd_supertx_sb(const AV1_COMP *const cpi, ThreadData *td, + const TileInfo *const tile, int mi_row, int mi_col, + BLOCK_SIZE bsize, int *tmp_rate, int64_t *tmp_dist, + TX_TYPE *best_tx, PC_TREE *pc_tree) { + const AV1_COMMON *const cm = &cpi->common; + MACROBLOCK *const x = &td->mb; + MACROBLOCKD *const xd = &x->e_mbd; + int plane, pnskip, skippable, skippable_uv, rate_uv, this_rate, + base_rate = *tmp_rate; + int64_t sse, pnsse, sse_uv, this_dist, dist_uv; + uint8_t *dst_buf[3]; + int dst_stride[3]; + TX_SIZE tx_size; + MB_MODE_INFO *mbmi; + TX_TYPE tx_type, best_tx_nostx; +#if CONFIG_EXT_TX + int ext_tx_set; +#endif // CONFIG_EXT_TX + int tmp_rate_tx = 0, skip_tx = 0; + int64_t tmp_dist_tx = 0, rd_tx, bestrd_tx = INT64_MAX; + + set_skip_context(xd, mi_row, mi_col); + set_mode_info_offsets(cpi, x, xd, mi_row, mi_col); + update_state_sb_supertx(cpi, td, tile, mi_row, mi_col, bsize, 1, pc_tree); + av1_setup_dst_planes(xd->plane, bsize, get_frame_new_buffer(cm), mi_row, + mi_col); + for (plane = 0; plane < MAX_MB_PLANE; plane++) { + dst_buf[plane] = xd->plane[plane].dst.buf; + dst_stride[plane] = xd->plane[plane].dst.stride; + } + predict_sb_complex(cpi, td, tile, mi_row, mi_col, mi_row, mi_col, 1, bsize, + bsize, dst_buf, dst_stride, pc_tree); + + set_offsets_without_segment_id(cpi, tile, x, mi_row, mi_col, bsize); + set_segment_id_supertx(cpi, x, mi_row, mi_col, bsize); + + mbmi = &xd->mi[0]->mbmi; + best_tx_nostx = mbmi->tx_type; + + *best_tx = DCT_DCT; + + // chroma + skippable_uv = 1; + rate_uv = 0; + dist_uv = 0; + sse_uv = 0; + for (plane = 1; plane < MAX_MB_PLANE; ++plane) { +#if CONFIG_VAR_TX + ENTROPY_CONTEXT ctxa[2 * MAX_MIB_SIZE]; + ENTROPY_CONTEXT ctxl[2 * MAX_MIB_SIZE]; + const struct macroblockd_plane *const pd = &xd->plane[plane]; + RD_STATS this_rd_stats; + av1_init_rd_stats(&this_rd_stats); + + tx_size = max_txsize_lookup[bsize]; + tx_size = + uv_txsize_lookup[bsize][tx_size][cm->subsampling_x][cm->subsampling_y]; + av1_get_entropy_contexts(bsize, tx_size, pd, ctxa, ctxl); + + av1_subtract_plane(x, bsize, plane); + av1_tx_block_rd_b(cpi, x, tx_size, 0, 0, plane, 0, + get_plane_block_size(bsize, pd), &ctxa[0], &ctxl[0], + &this_rd_stats); + + this_rate = this_rd_stats.rate; + this_dist = this_rd_stats.dist; + pnsse = this_rd_stats.sse; + pnskip = this_rd_stats.skip; +#else + tx_size = max_txsize_lookup[bsize]; + tx_size = + uv_txsize_lookup[bsize][tx_size][cm->subsampling_x][cm->subsampling_y]; + av1_subtract_plane(x, bsize, plane); + av1_txfm_rd_in_plane_supertx(x, cpi, &this_rate, &this_dist, &pnskip, + &pnsse, INT64_MAX, plane, bsize, tx_size, 0); +#endif // CONFIG_VAR_TX + + rate_uv += this_rate; + dist_uv += this_dist; + sse_uv += pnsse; + skippable_uv &= pnskip; + } + + // luma + tx_size = max_txsize_lookup[bsize]; + av1_subtract_plane(x, bsize, 0); +#if CONFIG_EXT_TX + ext_tx_set = get_ext_tx_set(tx_size, bsize, 1, cm->reduced_tx_set_used); +#endif // CONFIG_EXT_TX + for (tx_type = DCT_DCT; tx_type < TX_TYPES; ++tx_type) { +#if CONFIG_VAR_TX + ENTROPY_CONTEXT ctxa[2 * MAX_MIB_SIZE]; + ENTROPY_CONTEXT ctxl[2 * MAX_MIB_SIZE]; + const struct macroblockd_plane *const pd = &xd->plane[0]; + RD_STATS this_rd_stats; +#endif // CONFIG_VAR_TX + +#if CONFIG_EXT_TX + if (!ext_tx_used_inter[ext_tx_set][tx_type]) continue; +#else + if (tx_size >= TX_32X32 && tx_type != DCT_DCT) continue; +#endif // CONFIG_EXT_TX + mbmi->tx_type = tx_type; + +#if CONFIG_VAR_TX + av1_init_rd_stats(&this_rd_stats); + av1_get_entropy_contexts(bsize, tx_size, pd, ctxa, ctxl); + av1_tx_block_rd_b(cpi, x, tx_size, 0, 0, 0, 0, bsize, &ctxa[0], &ctxl[0], + &this_rd_stats); + + this_rate = this_rd_stats.rate; + this_dist = this_rd_stats.dist; + pnsse = this_rd_stats.sse; + pnskip = this_rd_stats.skip; +#else + av1_txfm_rd_in_plane_supertx(x, cpi, &this_rate, &this_dist, &pnskip, + &pnsse, INT64_MAX, 0, bsize, tx_size, 0); +#endif // CONFIG_VAR_TX + +#if CONFIG_EXT_TX + if (get_ext_tx_types(tx_size, bsize, 1, cm->reduced_tx_set_used) > 1 && + !xd->lossless[xd->mi[0]->mbmi.segment_id] && this_rate != INT_MAX) { + if (ext_tx_set > 0) + this_rate += + cpi->inter_tx_type_costs[ext_tx_set][mbmi->tx_size][mbmi->tx_type]; + } +#else + if (tx_size < TX_32X32 && !xd->lossless[xd->mi[0]->mbmi.segment_id] && + this_rate != INT_MAX) { + this_rate += cpi->inter_tx_type_costs[tx_size][mbmi->tx_type]; + } +#endif // CONFIG_EXT_TX + *tmp_rate = rate_uv + this_rate; + *tmp_dist = dist_uv + this_dist; + sse = sse_uv + pnsse; + skippable = skippable_uv && pnskip; + if (skippable) { + *tmp_rate = av1_cost_bit(av1_get_skip_prob(cm, xd), 1); + x->skip = 1; + } else { + if (RDCOST(x->rdmult, x->rddiv, *tmp_rate, *tmp_dist) < + RDCOST(x->rdmult, x->rddiv, 0, sse)) { + *tmp_rate += av1_cost_bit(av1_get_skip_prob(cm, xd), 0); + x->skip = 0; + } else { + *tmp_dist = sse; + *tmp_rate = av1_cost_bit(av1_get_skip_prob(cm, xd), 1); + x->skip = 1; + } + } + *tmp_rate += base_rate; + rd_tx = RDCOST(x->rdmult, x->rddiv, *tmp_rate, *tmp_dist); + if (rd_tx < bestrd_tx * 0.99 || tx_type == DCT_DCT) { + *best_tx = tx_type; + bestrd_tx = rd_tx; + tmp_rate_tx = *tmp_rate; + tmp_dist_tx = *tmp_dist; + skip_tx = x->skip; + } + } + *tmp_rate = tmp_rate_tx; + *tmp_dist = tmp_dist_tx; + x->skip = skip_tx; +#if CONFIG_VAR_TX + for (plane = 0; plane < 1; ++plane) + memset(x->blk_skip[plane], x->skip, + sizeof(uint8_t) * pc_tree->none.num_4x4_blk); +#endif // CONFIG_VAR_TX + xd->mi[0]->mbmi.tx_type = best_tx_nostx; +} +#endif // CONFIG_SUPERTX |