/*
 * Copyright (c) 2020, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
 * Media Patent License 1.0 was not distributed with this source code in the
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
 */

#include "aom_ports/system_state.h"

#include "av1/common/blockd.h"
#include "av1/common/enums.h"
#include "av1/common/reconintra.h"
#include "av1/encoder/aq_complexity.h"
#include "av1/encoder/aq_variance.h"
#include "av1/encoder/block.h"
#include "av1/encoder/context_tree.h"
#include "av1/encoder/encoder.h"
#include "av1/encoder/encodeframe.h"
#include "av1/encoder/encodemv.h"
#include "av1/encoder/partition_search.h"
#include "av1/encoder/partition_search_utils.h"
#include "av1/encoder/partition_strategy.h"
#include "av1/encoder/reconinter_enc.h"
#include "av1/encoder/tokenize.h"
#include "av1/encoder/var_based_part.h"

static void pick_sb_modes(AV1_COMP *const cpi, TileDataEnc *tile_data,
                          MACROBLOCK *const x, int mi_row, int mi_col,
                          RD_STATS *rd_cost, PARTITION_TYPE partition,
                          BLOCK_SIZE bsize, PICK_MODE_CONTEXT *ctx,
                          RD_STATS best_rd, int pick_mode_type) {
  AV1_COMMON *const cm = &cpi->common;
  const int num_planes = av1_num_planes(cm);
  TileInfo *const tile_info = &tile_data->tile_info;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *mbmi;
  MB_MODE_INFO *ctx_mbmi = &ctx->mic;
  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;

#if CONFIG_COLLECT_COMPONENT_TIMING
  start_timing(cpi, rd_pick_sb_modes_time);
#endif

  if (best_rd.rdcost < 0) {
    ctx->rd_stats.rdcost = INT64_MAX;
    ctx->rd_stats.skip = 0;
    av1_invalid_rd_stats(rd_cost);
    return;
  }

  aom_clear_system_state();

  av1_enc_set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize,
                      &ctx->chroma_ref_info);

  mbmi = xd->mi[0];

  if (ctx->rd_mode_is_ready) {
    assert(ctx_mbmi->sb_type == bsize);
    assert(ctx_mbmi->partition == partition);
    *mbmi = *ctx_mbmi;
    rd_cost->rate = ctx->rd_stats.rate;
    rd_cost->dist = ctx->rd_stats.dist;
    rd_cost->rdcost = ctx->rd_stats.rdcost;
  } else {
    mbmi->sb_type = bsize;
    mbmi->partition = partition;
#if CONFIG_DSPL_RESIDUAL
    mbmi->dspl_type = DSPL_NONE;
#endif  // CONFIG_DSPL_RESIDUAL
  }

  mbmi->chroma_ref_info = ctx->chroma_ref_info;

#if CONFIG_RD_DEBUG
  mbmi->mi_row = mi_row;
  mbmi->mi_col = mi_col;
#endif

  for (i = 0; i < num_planes; ++i) {
    p[i].coeff = ctx->coeff[i];
    p[i].qcoeff = ctx->qcoeff[i];
    pd[i].dqcoeff = ctx->dqcoeff[i];
    p[i].eobs = ctx->eobs[i];
    p[i].txb_entropy_ctx = ctx->txb_entropy_ctx[i];
  }

  for (i = 0; i < 2; ++i) pd[i].color_index_map = ctx->color_index_map[i];

  if (!ctx->rd_mode_is_ready) {
    ctx->skippable = 0;

    // Set to zero to make sure we do not use the previous encoded frame stats
    mbmi->skip = 0;

    // Reset skip mode flag.
    mbmi->skip_mode = 0;
  }

  x->skip_chroma_rd = !mbmi->chroma_ref_info.is_chroma_ref;

  if (ctx->rd_mode_is_ready) {
    x->skip = ctx->rd_stats.skip;
    *x->mbmi_ext = ctx->mbmi_ext;
    return;
  }

  if (is_cur_buf_hbd(xd)) {
    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);
  }

  // If the threshold for disabling wedge search is zero, it means the feature
  // should not be used. Use a value that will always succeed in the check.
  if (cpi->sf.disable_wedge_search_edge_thresh == 0) {
    x->edge_strength = UINT16_MAX;
    x->edge_strength_x = UINT16_MAX;
    x->edge_strength_y = UINT16_MAX;
  } else {
    EdgeInfo ei =
        av1_get_edge_info(&x->plane[0].src, bsize, is_cur_buf_hbd(xd), xd->bd);
    x->edge_strength = ei.magnitude;
    x->edge_strength_x = ei.x;
    x->edge_strength_y = ei.y;
  }

  // Default initialization of the threshold for R-D optimization of
  // coefficients for mode decision
  x->coeff_opt_dist_threshold =
      get_rd_opt_coeff_thresh(cpi->coeff_opt_dist_threshold, 0, 0);

  // Save rdmult before it might be changed, so it can be restored later.
  const int orig_rdmult = x->rdmult;
  av1_setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, aq_mode, mbmi);
  // Set error per bit for current rdmult
  set_error_per_bit(x, x->rdmult);
  av1_rd_cost_update(x->rdmult, &best_rd);

#if CONFIG_DERIVED_INTRA_MODE
  mbmi->use_derived_intra_mode[0] = 0;
  mbmi->use_derived_intra_mode[1] = 0;
#endif  // CONFIG_DERIVED_INTRA_MODE
#if CONFIG_DERIVED_MV
  mbmi->derived_mv_allowed = mbmi->use_derived_mv = 0;
#endif  // CONFIG_DERIVED_MV

  // 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)) {
#if CONFIG_COLLECT_COMPONENT_TIMING
    start_timing(cpi, av1_rd_pick_intra_mode_sb_time);
#endif
    av1_rd_pick_intra_mode_sb(cpi, x, rd_cost, bsize, ctx, best_rd.rdcost);
#if CONFIG_COLLECT_COMPONENT_TIMING
    end_timing(cpi, av1_rd_pick_intra_mode_sb_time);
#endif
  } else {
#if CONFIG_COLLECT_COMPONENT_TIMING
    start_timing(cpi, av1_rd_pick_inter_mode_sb_time);
#endif
    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.rdcost);
    } else {
      // TODO(kyslov): do the same for pick_intra_mode and
      //               pick_inter_mode_sb_seg_skip
      switch (pick_mode_type) {
#if !CONFIG_REALTIME_ONLY
        case PICK_MODE_RD:
          av1_rd_pick_inter_mode_sb(cpi, tile_data, x, rd_cost, bsize, ctx,
                                    best_rd.rdcost);
          break;
#endif
        case PICK_MODE_NONRD:
          av1_nonrd_pick_inter_mode_sb(cpi, tile_data, x, rd_cost, bsize, ctx,
                                       best_rd.rdcost);
          break;
        case PICK_MODE_FAST_NONRD:
          av1_fast_nonrd_pick_inter_mode_sb(cpi, tile_data, x, rd_cost, bsize,
                                            ctx, best_rd.rdcost);
          break;
        default: assert(0 && "Unknown pick mode type.");
      }
    }
#if CONFIG_COLLECT_COMPONENT_TIMING
    end_timing(cpi, av1_rd_pick_inter_mode_sb_time);
#endif
  }

  // 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->current_frame.frame_type == KEY_FRAME ||
       cpi->refresh_alt_ref_frame || cpi->refresh_alt2_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->rd_stats.rate = rd_cost->rate;
  ctx->rd_stats.dist = rd_cost->dist;
  ctx->rd_stats.rdcost = rd_cost->rdcost;

#if CONFIG_COLLECT_COMPONENT_TIMING
  end_timing(cpi, rd_pick_sb_modes_time);
#endif
}

void av1_nonrd_use_partition(AV1_COMP *cpi, ThreadData *td,
                             TileDataEnc *tile_data, MB_MODE_INFO **mib,
                             TOKENEXTRA **tp, int mi_row, int mi_col,
                             BLOCK_SIZE bsize, PC_TREE *pc_tree,
                             PARTITION_TREE *ptree) {
  AV1_COMMON *const cm = &cpi->common;
  TileInfo *const tile_info = &tile_data->tile_info;
  const SPEED_FEATURES *const sf = &cpi->sf;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  const int ss_x = xd->plane[1].subsampling_x;
  const int ss_y = xd->plane[1].subsampling_y;
  // Only square blocks from 8x8 to 128x128 are supported
  assert(bsize >= BLOCK_8X8 && bsize <= BLOCK_128X128);
  const int bs = mi_size_wide[bsize];
  const int hbs = bs / 2;
  const PARTITION_TYPE partition =
      (bsize >= BLOCK_8X8) ? get_partition(cm, mi_row, mi_col, bsize)
                           : PARTITION_NONE;
  const BLOCK_SIZE subsize = get_partition_subsize(bsize, partition);
  RD_STATS dummy_cost;
  av1_invalid_rd_stats(&dummy_cost);
  RD_STATS invalid_rd;
  av1_invalid_rd_stats(&invalid_rd);

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

  assert(mi_size_wide[bsize] == mi_size_high[bsize]);

  pc_tree->partitioning = partition;

  assert(ptree);
  ptree->partition = partition;
  ptree->bsize = bsize;
  ptree->mi_row = mi_row;
  ptree->mi_col = mi_col;
  PARTITION_TREE *parent = ptree->parent;
  set_chroma_ref_info(mi_row, mi_col, ptree->index, bsize,
                      &ptree->chroma_ref_info,
                      parent ? &parent->chroma_ref_info : NULL,
                      parent ? parent->bsize : BLOCK_INVALID,
                      parent ? parent->partition : PARTITION_NONE, ss_x, ss_y);

  // int num_splittable_sub_blocks = 0;
  switch (partition) {
    case PARTITION_SPLIT:
      ptree->sub_tree[0] = av1_alloc_ptree_node(ptree, 0);
      ptree->sub_tree[1] = av1_alloc_ptree_node(ptree, 1);
      ptree->sub_tree[2] = av1_alloc_ptree_node(ptree, 2);
      ptree->sub_tree[3] = av1_alloc_ptree_node(ptree, 3);
      break;
#if CONFIG_EXT_RECUR_PARTITIONS
    case PARTITION_HORZ:
    case PARTITION_VERT:
      ptree->sub_tree[0] = av1_alloc_ptree_node(ptree, 0);
      ptree->sub_tree[1] = av1_alloc_ptree_node(ptree, 1);
      break;
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    default: break;
  }

  xd->above_txfm_context = cm->above_txfm_context[tile_info->tile_row] + mi_col;
  xd->left_txfm_context =
      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);

  switch (partition) {
    case PARTITION_NONE:
      pc_tree->none =
          av1_alloc_pmc(cm, mi_row, mi_col, bsize, pc_tree, PARTITION_NONE, 0,
                        ss_x, ss_y, &td->shared_coeff_buf);
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &dummy_cost,
                    PARTITION_NONE, bsize, pc_tree->none, invalid_rd,
                    sf->use_fast_nonrd_pick_mode ? PICK_MODE_FAST_NONRD
                                                 : PICK_MODE_NONRD);
      av1_encode_b(cpi, tile_data, td, tp, mi_row, mi_col, 0, bsize, partition,
                   pc_tree->none, NULL);
      break;
    case PARTITION_VERT:
#if CONFIG_EXT_RECUR_PARTITIONS
      ptree->sub_tree[0]->partition = PARTITION_NONE;
      ptree->sub_tree[1]->partition = PARTITION_NONE;

      pc_tree->vertical[0] = av1_alloc_pc_tree_node(
          mi_row, mi_col, subsize, pc_tree, PARTITION_VERT, 0, 0, ss_x, ss_y);
      pc_tree->vertical[1] =
          av1_alloc_pc_tree_node(mi_row, mi_col + hbs, subsize, pc_tree,
                                 PARTITION_VERT, 1, 1, ss_x, ss_y);

      av1_nonrd_use_partition(cpi, td, tile_data, mib, tp, mi_row, mi_col,
                              subsize, pc_tree->vertical[0],
                              ptree->sub_tree[0]);
#else
      for (int i = 0; i < 2; ++i) {
        pc_tree->vertical[i] =
            av1_alloc_pmc(cm, mi_row, mi_col + hbs * i, subsize, pc_tree,
                          PARTITION_VERT, i, ss_x, ss_y, &td->shared_coeff_buf);
      }
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &dummy_cost,
                    PARTITION_VERT, subsize, pc_tree->vertical[0], invalid_rd,
                    sf->use_fast_nonrd_pick_mode ? PICK_MODE_FAST_NONRD
                                                 : PICK_MODE_NONRD);
      av1_encode_b(cpi, tile_data, td, tp, mi_row, mi_col, 0, subsize,
                   PARTITION_VERT, pc_tree->vertical[0], NULL);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      if (mi_col + hbs < cm->mi_cols && bsize > BLOCK_8X8) {
#if CONFIG_EXT_RECUR_PARTITIONS
        av1_nonrd_use_partition(cpi, td, tile_data, mib + hbs, tp, mi_row,
                                mi_col + hbs, subsize, pc_tree->vertical[1],
                                ptree->sub_tree[1]);
#else
        pick_sb_modes(cpi, tile_data, x, mi_row, mi_col + hbs, &dummy_cost,
                      PARTITION_VERT, subsize, pc_tree->vertical[1], invalid_rd,
                      sf->use_fast_nonrd_pick_mode ? PICK_MODE_FAST_NONRD
                                                   : PICK_MODE_NONRD);
        av1_encode_b(cpi, tile_data, td, tp, mi_row, mi_col + hbs, 0, subsize,
                     PARTITION_VERT, pc_tree->vertical[1], NULL);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      }
      break;
    case PARTITION_HORZ:
#if CONFIG_EXT_RECUR_PARTITIONS
      ptree->sub_tree[0]->partition = PARTITION_NONE;
      ptree->sub_tree[1]->partition = PARTITION_NONE;

      pc_tree->horizontal[0] = av1_alloc_pc_tree_node(
          mi_row, mi_col, subsize, pc_tree, PARTITION_HORZ, 0, 0, ss_x, ss_y);
      pc_tree->horizontal[1] =
          av1_alloc_pc_tree_node(mi_row + hbs, mi_col, subsize, pc_tree,
                                 PARTITION_HORZ, 1, 1, ss_x, ss_y);

      av1_nonrd_use_partition(cpi, td, tile_data, mib, tp, mi_row, mi_col,
                              subsize, pc_tree->horizontal[0],
                              ptree->sub_tree[0]);
#else
      for (int i = 0; i < 2; ++i) {
        pc_tree->horizontal[i] =
            av1_alloc_pmc(cm, mi_row + hbs * i, mi_col, subsize, pc_tree,
                          PARTITION_HORZ, i, ss_x, ss_y, &td->shared_coeff_buf);
      }
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &dummy_cost,
                    PARTITION_HORZ, subsize, pc_tree->horizontal[0], invalid_rd,
                    sf->use_fast_nonrd_pick_mode ? PICK_MODE_FAST_NONRD
                                                 : PICK_MODE_NONRD);
      av1_encode_b(cpi, tile_data, td, tp, mi_row, mi_col, 0, subsize,
                   PARTITION_HORZ, pc_tree->horizontal[0], NULL);
#endif  // CONFIG_EXT_RECUR_PARTITIONS

      if (mi_row + hbs < cm->mi_rows && bsize > BLOCK_8X8) {
#if CONFIG_EXT_RECUR_PARTITIONS
        av1_nonrd_use_partition(cpi, td, tile_data, mib + hbs * cm->mi_stride,
                                tp, mi_row + hbs, mi_col, subsize,
                                pc_tree->horizontal[1], ptree->sub_tree[1]);
#else
        pick_sb_modes(cpi, tile_data, x, mi_row + hbs, mi_col, &dummy_cost,
                      PARTITION_HORZ, subsize, pc_tree->horizontal[1],
                      invalid_rd,
                      sf->use_fast_nonrd_pick_mode ? PICK_MODE_FAST_NONRD
                                                   : PICK_MODE_NONRD);
        av1_encode_b(cpi, tile_data, td, tp, mi_row + hbs, mi_col, 0, subsize,
                     PARTITION_HORZ, pc_tree->horizontal[1], NULL);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      }
      break;
    case PARTITION_SPLIT:
      for (int i = 0; i < 4; i++) {
        int x_idx = (i & 1) * hbs;
        int y_idx = (i >> 1) * hbs;
        int jj = i >> 1, ii = i & 0x01;

        if ((mi_row + y_idx >= cm->mi_rows) || (mi_col + x_idx >= cm->mi_cols))
          continue;

        pc_tree->split[i] = av1_alloc_pc_tree_node(
            mi_row + y_idx, mi_col + x_idx, subsize, pc_tree, PARTITION_SPLIT,
            i, i == 3, ss_x, ss_y);
        av1_nonrd_use_partition(cpi, td, tile_data,
                                mib + jj * hbs * cm->mi_stride + ii * hbs, tp,
                                mi_row + y_idx, mi_col + x_idx, subsize,
                                pc_tree->split[i], ptree->sub_tree[i]);
      }
      break;
#if CONFIG_EXT_RECUR_PARTITIONS
    case PARTITION_HORZ_3:
    case PARTITION_VERT_3:
#else
    case PARTITION_VERT_A:
    case PARTITION_VERT_B:
    case PARTITION_HORZ_A:
    case PARTITION_HORZ_B:
    case PARTITION_HORZ_4:
    case PARTITION_VERT_4:
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      assert(0 && "Cannot handle extended partition types");
    default: assert(0); break;
  }

  ptree->is_settled = 1;
  if (partition != PARTITION_SPLIT || bsize == BLOCK_8X8)
    update_partition_context(xd, mi_row, mi_col, subsize, bsize);
}

#if !CONFIG_REALTIME_ONLY
static void set_partial_sb_partition(const AV1_COMMON *const cm,
                                     MB_MODE_INFO *mi, int bh_in, int bw_in,
                                     int mi_rows_remaining,
                                     int mi_cols_remaining, BLOCK_SIZE bsize,
                                     MB_MODE_INFO **mib) {
  int bh = bh_in;
  int r, c;
  for (r = 0; r < cm->seq_params.mib_size; r += bh) {
    int bw = bw_in;
    for (c = 0; c < cm->seq_params.mib_size; c += bw) {
      const int index = r * cm->mi_stride + c;
      mib[index] = mi + index;
      mib[index]->sb_type = find_partition_size(
          bsize, mi_rows_remaining - r, mi_cols_remaining - c, &bh, &bw);
    }
  }
}

void av1_set_fixed_partitioning(AV1_COMP *cpi, const TileInfo *const tile,
                                MB_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;
  MB_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->seq_params.mib_size) &&
      (mi_rows_remaining >= cm->seq_params.mib_size)) {
    for (block_row = 0; block_row < cm->seq_params.mib_size; block_row += bh) {
      for (block_col = 0; block_col < cm->seq_params.mib_size;
           block_col += bw) {
        int index = block_row * cm->mi_stride + block_col;
        mib[index] = mi_upper_left + index;
        mib[index]->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);
  }
}

void av1_rd_use_partition(AV1_COMP *cpi, ThreadData *td, TileDataEnc *tile_data,
                          MB_MODE_INFO **mib, TOKENEXTRA **tp, int mi_row,
                          int mi_col, BLOCK_SIZE bsize, int *rate,
                          int64_t *dist, int do_recon, PC_TREE *pc_tree) {
  AV1_COMMON *const cm = &cpi->common;
  const int num_planes = av1_num_planes(cm);
  TileInfo *const tile_info = &tile_data->tile_info;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  const int ss_x = xd->plane[1].subsampling_x;
  const int ss_y = xd->plane[1].subsampling_y;
  const int bw = mi_size_wide[bsize];
  const int bh = mi_size_high[bsize];
  const int hbw = bw / 2;
  const int hbh = bh / 2;
  const int pl = (bsize >= BLOCK_8X8)
                     ? partition_plane_context(xd, mi_row, mi_col, bsize)
                     : 0;
  const PARTITION_TYPE partition =
      (bsize >= BLOCK_8X8) ? get_partition(cm, mi_row, mi_col, bsize)
                           : PARTITION_NONE;
  const BLOCK_SIZE subsize = get_partition_subsize(bsize, partition);
  RD_SEARCH_MACROBLOCK_CONTEXT x_ctx;
  RD_STATS last_part_rdc, none_rdc, chosen_rdc, invalid_rdc;
  BLOCK_SIZE sub_subsize = BLOCK_4X4;
  int splits_below = 0;
  BLOCK_SIZE bs_type = mib[0]->sb_type;
  int do_partition_search = 1;

  if (pc_tree->none == NULL) {
    pc_tree->none =
        av1_alloc_pmc(cm, mi_row, mi_col, bsize, pc_tree, PARTITION_NONE, 0,
                      ss_x, ss_y, &td->shared_coeff_buf);
  }
  PICK_MODE_CONTEXT *ctx_none = pc_tree->none;

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

#if !CONFIG_EXT_RECUR_PARTITIONS
  assert(mi_size_wide[bsize] == mi_size_high[bsize]);
#endif

  av1_invalid_rd_stats(&last_part_rdc);
  av1_invalid_rd_stats(&none_rdc);
  av1_invalid_rd_stats(&chosen_rdc);
  av1_invalid_rd_stats(&invalid_rdc);

  pc_tree->partitioning = partition;

  xd->above_txfm_context = cm->above_txfm_context[tile_info->tile_row] + mi_col;
  xd->left_txfm_context =
      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
  av1_save_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);

  if (bsize == BLOCK_16X16 && cpi->vaq_refresh) {
    av1_enc_set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize,
                        &pc_tree->chroma_ref_info);
    x->mb_energy = av1_log_block_var(cpi, x, bsize);
  }

  // Save rdmult before it might be changed, so it can be restored later.
  const int orig_rdmult = x->rdmult;
  av1_setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, NO_AQ, NULL);

  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_partition_subsize(subsize, PARTITION_SPLIT);
      splits_below = 1;
      for (int i = 0; i < 4; i++) {
        int jj = i >> 1, ii = i & 0x01;
        MB_MODE_INFO *this_mi = mib[jj * hbh * cm->mi_stride + ii * hbw];
        if (this_mi && this_mi->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 + hbh < cm->mi_rows && mi_col + hbw < cm->mi_cols) {
      pc_tree->partitioning = PARTITION_NONE;
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &none_rdc,
                    PARTITION_NONE, bsize, ctx_none, invalid_rdc, PICK_MODE_RD);

      if (none_rdc.rate < INT_MAX) {
        none_rdc.rate += x->partition_cost[pl][PARTITION_NONE];
        none_rdc.rdcost = RDCOST(x->rdmult, none_rdc.rate, none_rdc.dist);
      }

      av1_restore_context(cm, x, &x_ctx, mi_row, mi_col, bsize, num_planes);
      mib[0]->sb_type = bs_type;
      pc_tree->partitioning = partition;
    }
  }

  for (int i = 0; i < 4; ++i) {
    pc_tree->split[i] = av1_alloc_pc_tree_node(
        mi_row + (i >> 1) * hbh, mi_col + (i & 1) * hbw, subsize, pc_tree,
        PARTITION_SPLIT, i, i == 3, ss_x, ss_y);
  }
  switch (partition) {
    case PARTITION_NONE:
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc,
                    PARTITION_NONE, bsize, ctx_none, invalid_rdc, PICK_MODE_RD);
      break;
    case PARTITION_HORZ:
#if CONFIG_EXT_RECUR_PARTITIONS
      pc_tree->horizontal[0] = av1_alloc_pc_tree_node(
          mi_row, mi_col, subsize, pc_tree, PARTITION_HORZ, 0, 0, ss_x, ss_y);
      pc_tree->horizontal[1] =
          av1_alloc_pc_tree_node(mi_row + hbh, mi_col, subsize, pc_tree,
                                 PARTITION_HORZ, 1, 1, ss_x, ss_y);
      av1_rd_use_partition(cpi, td, tile_data, mib, tp, mi_row, mi_col, subsize,
                           &last_part_rdc.rate, &last_part_rdc.dist, 1,
                           pc_tree->horizontal[0]);
#else
      for (int i = 0; i < 2; ++i) {
        pc_tree->horizontal[i] =
            av1_alloc_pmc(cm, mi_row + hbh * i, mi_col, subsize, pc_tree,
                          PARTITION_HORZ, i, ss_x, ss_y, &td->shared_coeff_buf);
      }
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc,
                    PARTITION_HORZ, subsize, pc_tree->horizontal[0],
                    invalid_rdc, PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      if (last_part_rdc.rate != INT_MAX && bsize >= BLOCK_8X8 &&
          mi_row + hbh < cm->mi_rows) {
        RD_STATS tmp_rdc;
        av1_init_rd_stats(&tmp_rdc);
#if CONFIG_EXT_RECUR_PARTITIONS
        av1_rd_use_partition(cpi, td, tile_data, mib + hbh * cm->mi_stride, tp,
                             mi_row + hbh, mi_col, subsize, &tmp_rdc.rate,
                             &tmp_rdc.dist, 0, pc_tree->horizontal[1]);
#else
        const PICK_MODE_CONTEXT *const ctx_h = pc_tree->horizontal[0];
        av1_update_state(cpi, td, ctx_h, mi_row, mi_col, subsize, 1);
        av1_encode_superblock(cpi, tile_data, td, tp, DRY_RUN_NORMAL, subsize,
                              NULL);
        pick_sb_modes(cpi, tile_data, x, mi_row + hbh, mi_col, &tmp_rdc,
                      PARTITION_HORZ, subsize, pc_tree->horizontal[1],
                      invalid_rdc, PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
        if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) {
          av1_invalid_rd_stats(&last_part_rdc);
          break;
        }
        last_part_rdc.rate += tmp_rdc.rate;
        last_part_rdc.dist += tmp_rdc.dist;
        last_part_rdc.rdcost += tmp_rdc.rdcost;
      }
      break;
    case PARTITION_VERT:
#if CONFIG_EXT_RECUR_PARTITIONS
      pc_tree->vertical[0] = av1_alloc_pc_tree_node(
          mi_row, mi_col, subsize, pc_tree, PARTITION_VERT, 0, 0, ss_x, ss_y);
      pc_tree->vertical[1] =
          av1_alloc_pc_tree_node(mi_row, mi_col + hbw, subsize, pc_tree,
                                 PARTITION_VERT, 1, 1, ss_x, ss_y);
      av1_rd_use_partition(cpi, td, tile_data, mib, tp, mi_row, mi_col, subsize,
                           &last_part_rdc.rate, &last_part_rdc.dist, 1,
                           pc_tree->vertical[0]);
#else
      for (int i = 0; i < 2; ++i) {
        pc_tree->vertical[i] =
            av1_alloc_pmc(cm, mi_row, mi_col + hbw * i, subsize, pc_tree,
                          PARTITION_VERT, i, ss_x, ss_y, &td->shared_coeff_buf);
      }
      pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &last_part_rdc,
                    PARTITION_VERT, subsize, pc_tree->vertical[0], invalid_rdc,
                    PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      if (last_part_rdc.rate != INT_MAX && bsize >= BLOCK_8X8 &&
          mi_col + hbw < cm->mi_cols) {
        RD_STATS tmp_rdc;
        av1_init_rd_stats(&tmp_rdc);

#if CONFIG_EXT_RECUR_PARTITIONS
        av1_rd_use_partition(cpi, td, tile_data, mib + hbw, tp, mi_row,
                             mi_col + hbw, subsize, &tmp_rdc.rate,
                             &tmp_rdc.dist, 0, pc_tree->vertical[1]);
#else
        const PICK_MODE_CONTEXT *const ctx_v = pc_tree->vertical[0];
        av1_update_state(cpi, td, ctx_v, mi_row, mi_col, subsize, 1);
        av1_encode_superblock(cpi, tile_data, td, tp, DRY_RUN_NORMAL, subsize,
                              NULL);
        pick_sb_modes(cpi, tile_data, x, mi_row, mi_col + hbw, &tmp_rdc,
                      PARTITION_VERT, subsize,
                      pc_tree->vertical[bsize > BLOCK_8X8], invalid_rdc,
                      PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
        if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) {
          av1_invalid_rd_stats(&last_part_rdc);
          break;
        }
        last_part_rdc.rate += tmp_rdc.rate;
        last_part_rdc.dist += tmp_rdc.dist;
        last_part_rdc.rdcost += tmp_rdc.rdcost;
      }
      break;
    case PARTITION_SPLIT:
      last_part_rdc.rate = 0;
      last_part_rdc.dist = 0;
      last_part_rdc.rdcost = 0;
      for (int i = 0; i < 4; i++) {
        int x_idx = (i & 1) * hbw;
        int y_idx = (i >> 1) * hbh;
        int jj = i >> 1, ii = i & 0x01;
        RD_STATS tmp_rdc;
        if ((mi_row + y_idx >= cm->mi_rows) || (mi_col + x_idx >= cm->mi_cols))
          continue;

        av1_init_rd_stats(&tmp_rdc);
        av1_rd_use_partition(
            cpi, td, tile_data, mib + jj * hbh * cm->mi_stride + ii * hbw, tp,
            mi_row + y_idx, mi_col + x_idx, subsize, &tmp_rdc.rate,
            &tmp_rdc.dist, i != 3, pc_tree->split[i]);
        if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) {
          av1_invalid_rd_stats(&last_part_rdc);
          break;
        }
        last_part_rdc.rate += tmp_rdc.rate;
        last_part_rdc.dist += tmp_rdc.dist;
      }
      break;
#if CONFIG_EXT_RECUR_PARTITIONS
    case PARTITION_HORZ_3:
    case PARTITION_VERT_3:
#else
    case PARTITION_VERT_A:
    case PARTITION_VERT_B:
    case PARTITION_HORZ_A:
    case PARTITION_HORZ_B:
    case PARTITION_HORZ_4:
    case PARTITION_VERT_4:
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      assert(0 && "Cannot handle extended partition types");
    default: assert(0); break;
  }

  if (last_part_rdc.rate < INT_MAX) {
    last_part_rdc.rate += x->partition_cost[pl][partition];
    last_part_rdc.rdcost =
        RDCOST(x->rdmult, last_part_rdc.rate, last_part_rdc.dist);
  }

  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 + bh < cm->mi_rows || mi_row + hbh == cm->mi_rows) &&
      (mi_col + bw < cm->mi_cols || mi_col + hbw == cm->mi_cols)) {
    BLOCK_SIZE split_subsize = get_partition_subsize(bsize, PARTITION_SPLIT);
    chosen_rdc.rate = 0;
    chosen_rdc.dist = 0;

    av1_restore_context(cm, x, &x_ctx, mi_row, mi_col, bsize, num_planes);
    pc_tree->partitioning = PARTITION_SPLIT;

    // Split partition.
    for (int i = 0; i < 4; i++) {
      int x_idx = (i & 1) * hbw;
      int y_idx = (i >> 1) * hbh;
      RD_STATS tmp_rdc;

      if ((mi_row + y_idx >= cm->mi_rows) || (mi_col + x_idx >= cm->mi_cols))
        continue;

      av1_save_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);
      pc_tree->split[i]->partitioning = PARTITION_NONE;
      if (pc_tree->split[i]->none != NULL)
        pc_tree->split[i]->none =
            av1_alloc_pmc(cm, mi_row + y_idx, mi_col + x_idx, split_subsize,
                          pc_tree->split[i], PARTITION_NONE, 0, ss_x, ss_y,
                          &td->shared_coeff_buf);
      pick_sb_modes(cpi, tile_data, x, mi_row + y_idx, mi_col + x_idx, &tmp_rdc,
                    PARTITION_SPLIT, split_subsize, pc_tree->split[i]->none,
                    invalid_rdc, PICK_MODE_RD);

      av1_restore_context(cm, x, &x_ctx, mi_row, mi_col, bsize, num_planes);
      if (tmp_rdc.rate == INT_MAX || tmp_rdc.dist == INT64_MAX) {
        av1_invalid_rd_stats(&chosen_rdc);
        break;
      }

      chosen_rdc.rate += tmp_rdc.rate;
      chosen_rdc.dist += tmp_rdc.dist;

      if (i != 3)
        av1_encode_sb(cpi, td, tile_data, tp, mi_row + y_idx, mi_col + x_idx,
                      DRY_RUN_NORMAL, split_subsize, pc_tree->split[i], NULL,
                      NULL);

      chosen_rdc.rate += x->partition_cost[pl][PARTITION_NONE];
    }
    if (chosen_rdc.rate < INT_MAX) {
#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
      chosen_rdc.rate += x->partition_cost[pl][PARTITION_SPLIT];
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
      chosen_rdc.rdcost = RDCOST(x->rdmult, chosen_rdc.rate, chosen_rdc.dist);
    }
  }

  // If last_part is better set the partitioning to that.
  if (last_part_rdc.rdcost < chosen_rdc.rdcost) {
    mib[0]->sb_type = bsize;
    if (bsize >= BLOCK_8X8) pc_tree->partitioning = partition;
    chosen_rdc = last_part_rdc;
  }
  // 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;
  }

  av1_restore_context(cm, x, &x_ctx, mi_row, mi_col, bsize, num_planes);

  // We must have chosen a partitioning and encoding or we'll fail later on.
  // No other opportunities for success.
  if (bsize == cm->seq_params.sb_size)
    assert(chosen_rdc.rate < INT_MAX && chosen_rdc.dist < INT64_MAX);

  if (do_recon) {
    if (bsize == cm->seq_params.sb_size) {
      // NOTE: To get estimate for rate due to the tokens, use:
      // int rate_coeffs = 0;
      // encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, DRY_RUN_COSTCOEFFS,
      //           bsize, pc_tree, &rate_coeffs);
      x->cb_offset = 0;
      av1_reset_ptree_in_sbi(xd->sbi);
      av1_encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, OUTPUT_ENABLED,
                    bsize, pc_tree, xd->sbi->ptree_root, NULL);
    } else {
      av1_encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, DRY_RUN_NORMAL,
                    bsize, pc_tree, NULL, NULL);
    }
  }

  *rate = chosen_rdc.rate;
  *dist = chosen_rdc.dist;
  x->rdmult = orig_rdmult;
}
#if !CONFIG_EXT_RECUR_PARTITIONS
// Try searching for an encoding for the given subblock. Returns zero if the
// rdcost is already too high (to tell the caller not to bother searching for
// encodings of further subblocks)
static int rd_try_subblock(AV1_COMP *const cpi, ThreadData *td,
                           TileDataEnc *tile_data, TOKENEXTRA **tp, int is_last,
                           int mi_row, int mi_col, BLOCK_SIZE subsize,
                           RD_STATS best_rdcost, RD_STATS *sum_rdc,
                           PARTITION_TYPE partition,
                           PICK_MODE_CONTEXT *this_ctx) {
  MACROBLOCK *const x = &td->mb;
  const int orig_mult = x->rdmult;
  av1_setup_block_rdmult(cpi, x, mi_row, mi_col, subsize, NO_AQ, NULL);

  av1_rd_cost_update(x->rdmult, &best_rdcost);

  RD_STATS rdcost_remaining;
  av1_rd_stats_subtraction(x->rdmult, &best_rdcost, sum_rdc, &rdcost_remaining);
  RD_STATS this_rdc;
  pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &this_rdc, partition,
                subsize, this_ctx, rdcost_remaining, PICK_MODE_RD);

  if (this_rdc.rate == INT_MAX) {
    sum_rdc->rdcost = INT64_MAX;
  } else {
    sum_rdc->rate += this_rdc.rate;
    sum_rdc->dist += this_rdc.dist;
    av1_rd_cost_update(x->rdmult, sum_rdc);
  }

  if (sum_rdc->rdcost >= best_rdcost.rdcost) {
    x->rdmult = orig_mult;
    return 0;
  }

  if (!is_last) {
    av1_update_state(cpi, td, this_ctx, mi_row, mi_col, subsize, 1);
    av1_encode_superblock(cpi, tile_data, td, tp, DRY_RUN_NORMAL, subsize,
                          NULL);
  }

  x->rdmult = orig_mult;
  return 1;
}
#endif  // !CONFIG_EXT_RECUR_PARTITIONS

#if CONFIG_EXT_RECUR_PARTITIONS
typedef struct {
  SIMPLE_MOTION_DATA_TREE *sms_tree;
  PC_TREE *pc_tree;
  PICK_MODE_CONTEXT *ctx;
  int mi_row;
  int mi_col;
  BLOCK_SIZE bsize;
  PARTITION_TYPE partition;
  int is_last_subblock;
  int is_splittable;
  int max_sq_part;
  int min_sq_part;
} SUBBLOCK_RDO_DATA;

// Try searching for an encoding for the given subblock. Returns zero if the
// rdcost is already too high (to tell the caller not to bother searching for
// encodings of further subblocks)
static int rd_try_subblock_new(AV1_COMP *const cpi, ThreadData *td,
                               TileDataEnc *tile_data, TOKENEXTRA **tp,
                               SUBBLOCK_RDO_DATA *rdo_data,
                               RD_STATS best_rdcost, RD_STATS *sum_rdc,
                               SB_MULTI_PASS_MODE multi_pass_mode) {
  MACROBLOCK *const x = &td->mb;
  const int orig_mult = x->rdmult;
  const int mi_row = rdo_data->mi_row;
  const int mi_col = rdo_data->mi_col;
  const BLOCK_SIZE bsize = rdo_data->bsize;

  av1_setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, NO_AQ, NULL);

  av1_rd_cost_update(x->rdmult, &best_rdcost);

  RD_STATS rdcost_remaining;
  av1_rd_stats_subtraction(x->rdmult, &best_rdcost, sum_rdc, &rdcost_remaining);
  RD_STATS this_rdc;

  if (rdo_data->is_splittable) {
    if (!av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, bsize,
                               rdo_data->max_sq_part, rdo_data->min_sq_part,
                               &this_rdc, rdcost_remaining, rdo_data->pc_tree,
                               rdo_data->sms_tree, NULL, multi_pass_mode))
      return 0;
  } else {
    const BLOCK_SIZE sb_size = cpi->common.seq_params.sb_size;
    SimpleMotionData *sms_data =
        av1_get_sms_data_entry(x->sms_bufs, mi_row, mi_col, bsize, sb_size);
    av1_set_best_mode_cache(x, sms_data->mode_cache);

    pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &this_rdc,
                  rdo_data->partition, bsize, rdo_data->ctx, rdcost_remaining,
                  PICK_MODE_RD);

    x->inter_mode_cache = NULL;
    if (this_rdc.rate != INT_MAX) {
      av1_add_mode_search_context_to_cache(sms_data, rdo_data->ctx);
    }
  }

  if (this_rdc.rate == INT_MAX) {
    sum_rdc->rdcost = INT64_MAX;
  } else {
    sum_rdc->rate += this_rdc.rate;
    sum_rdc->dist += this_rdc.dist;
    av1_rd_cost_update(x->rdmult, sum_rdc);
  }

  if (sum_rdc->rdcost >= best_rdcost.rdcost) {
    x->rdmult = orig_mult;
    return 0;
  }

  if (!rdo_data->is_last_subblock && !rdo_data->is_splittable) {
    av1_update_state(cpi, td, rdo_data->ctx, mi_row, mi_col, bsize, 1);
    av1_encode_superblock(cpi, tile_data, td, tp, DRY_RUN_NORMAL, bsize, NULL);
  }

  x->rdmult = orig_mult;
  return 1;
}
#else   // !CONFIG_EXT_RECUR_PARTITIONS
static bool rd_test_partition3(AV1_COMP *const cpi, ThreadData *td,
                               TileDataEnc *tile_data, TOKENEXTRA **tp,
                               PC_TREE *pc_tree, RD_STATS *best_rdc,
                               PICK_MODE_CONTEXT *ctxs[3], int mi_row,
                               int mi_col, BLOCK_SIZE bsize,
                               PARTITION_TYPE partition, 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) {
  const MACROBLOCK *const x = &td->mb;
  const MACROBLOCKD *const xd = &x->e_mbd;
  const int pl = partition_plane_context(xd, mi_row, mi_col, bsize);
  RD_STATS sum_rdc;
  av1_init_rd_stats(&sum_rdc);
  sum_rdc.rate = x->partition_cost[pl][partition];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);
  if (!rd_try_subblock(cpi, td, tile_data, tp, 0, mi_row0, mi_col0, subsize0,
                       *best_rdc, &sum_rdc, partition, ctxs[0]))
    return false;

  if (!rd_try_subblock(cpi, td, tile_data, tp, 0, mi_row1, mi_col1, subsize1,
                       *best_rdc, &sum_rdc, partition, ctxs[1]))
    return false;

  if (!rd_try_subblock(cpi, td, tile_data, tp, 1, mi_row2, mi_col2, subsize2,
                       *best_rdc, &sum_rdc, partition, ctxs[2]))
    return false;

  av1_rd_cost_update(x->rdmult, &sum_rdc);
  if (sum_rdc.rdcost >= best_rdc->rdcost) return false;
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, sum_rdc.dist);
  if (sum_rdc.rdcost >= best_rdc->rdcost) return false;

  *best_rdc = sum_rdc;
  pc_tree->partitioning = partition;
  return true;
}
#endif  // CONFIG_EXT_RECUR_PARTITIONS

static INLINE void prune_partitions_after_none(
    PartitionSearchState *search_state, AV1_COMP *const cpi, MACROBLOCK *x,
    SIMPLE_MOTION_DATA_TREE *sms_tree, const PICK_MODE_CONTEXT *ctx_none,
    const RD_STATS *this_rdc, unsigned int *pb_source_variance) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const BLOCK_SIZE bsize = blk_params->bsize;

#if CONFIG_EXT_RECUR_PARTITIONS
  (void)sms_tree;
#else   // !CONFIG_EXT_RECUR_PARTITIONS
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  if (!frame_is_intra_only(cm) &&
#if CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      search_state->do_rectangular_split &&
#else   // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      (search_state->do_square_split || search_state->do_rectangular_split) &&
#endif  // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      !x->e_mbd.lossless[xd->mi[0]->segment_id] && ctx_none->skippable) {
    const int use_ml_based_breakout =
#if CONFIG_EXT_RECUR_PARTITIONS
        is_square_block(bsize) &&
#endif  // CONFIG_EXT_RECUR_PARTITIONS
        bsize <= cpi->sf.use_square_partition_only_threshold &&
        bsize > BLOCK_4X4 && xd->bd == 8;
    if (use_ml_based_breakout) {
      if (av1_ml_predict_breakout(cpi, bsize, x, this_rdc,
                                  *pb_source_variance)) {
#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
        search_state->do_square_split = 0;
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
        search_state->do_rectangular_split = 0;
      }
    }

    // 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)) -
         (mi_size_wide_log2[bsize] + mi_size_high_log2[bsize]));
    const int rate_breakout_thr = cpi->sf.partition_search_breakout_rate_thr *
                                  num_pels_log2_lookup[bsize];

    // 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 (this_rdc->dist < dist_breakout_thr &&
        this_rdc->rate < rate_breakout_thr) {
#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
      search_state->do_square_split = 0;
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
      search_state->do_rectangular_split = 0;
    }
  }

#if !CONFIG_EXT_RECUR_PARTITIONS
  if (cpi->sf.simple_motion_search_early_term_none && cm->show_frame &&
      !frame_is_intra_only(cm) && bsize >= BLOCK_16X16 &&
      mi_row + blk_params->mi_step_h < cm->mi_rows &&
      mi_col + blk_params->mi_step_w < cm->mi_cols &&
      this_rdc->rdcost < INT64_MAX && this_rdc->rdcost >= 0 &&
      this_rdc->rate < INT_MAX && this_rdc->rate >= 0 &&
      (search_state->do_square_split || search_state->do_rectangular_split)) {
    av1_simple_motion_search_early_term_none(
        cpi, x, sms_tree, mi_row, mi_col, bsize, this_rdc,
        &search_state->terminate_partition_search);
  }
#endif  // !CONFIG_EXT_RECUR_PARTITIONS
}

static INLINE void search_partition_none(
    PartitionSearchState *search_state, AV1_COMP *const cpi, ThreadData *td,
    TileDataEnc *tile_data, RD_STATS *best_rdc, PC_TREE *pc_tree,
    SIMPLE_MOTION_DATA_TREE *sms_tree, RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx,
    unsigned int *pb_source_variance, int64_t *none_rd, int64_t *part_none_rd) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

#if CONFIG_EXT_RECUR_PARTITIONS
  (void)part_none_rd;
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  RD_STATS this_rdc;

  pc_tree->none =
      av1_alloc_pmc(cm, mi_row, mi_col, bsize, pc_tree, PARTITION_NONE, 0,
                    blk_params->ss_x, blk_params->ss_y, &td->shared_coeff_buf);
  PICK_MODE_CONTEXT *ctx_none = pc_tree->none;

  if (blk_params->is_le_min_sq_part && blk_params->has_rows &&
      blk_params->has_cols)
    search_state->partition_none_allowed = 1;
  if (search_state->terminate_partition_search ||
      !search_state->partition_none_allowed || blk_params->is_gt_max_sq_part) {
    return;
  }

  int pt_cost = 0;
  if (search_state->is_block_splittable) {
    pt_cost = search_state->partition_cost[PARTITION_NONE] < INT_MAX
                  ? search_state->partition_cost[PARTITION_NONE]
                  : 0;
  }
  RD_STATS partition_rdcost;
  av1_init_rd_stats(&partition_rdcost);
  partition_rdcost.rate = pt_cost;
  av1_rd_cost_update(x->rdmult, &partition_rdcost);
  RD_STATS best_remain_rdcost;
  av1_rd_stats_subtraction(x->rdmult, best_rdc, &partition_rdcost,
                           &best_remain_rdcost);
#if CONFIG_COLLECT_PARTITION_STATS
  if (best_remain_rdcost >= 0) {
    partition_attempts[PARTITION_NONE] += 1;
    aom_usec_timer_start(&partition_timer);
    partition_timer_on = 1;
  }
#endif
#if CONFIG_EXT_RECUR_PARTITIONS
  SimpleMotionData *sms_data = av1_get_sms_data_entry(
      x->sms_bufs, mi_row, mi_col, bsize, cm->seq_params.sb_size);
  av1_set_best_mode_cache(x, sms_data->mode_cache);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &this_rdc, PARTITION_NONE,
                bsize, ctx_none, best_remain_rdcost, PICK_MODE_RD);
#if CONFIG_EXT_RECUR_PARTITIONS
  x->inter_mode_cache = NULL;
  if (this_rdc.rate != INT_MAX) {
    av1_add_mode_search_context_to_cache(sms_data, ctx_none);
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  av1_rd_cost_update(x->rdmult, &this_rdc);
#if CONFIG_COLLECT_PARTITION_STATS
  if (partition_timer_on) {
    aom_usec_timer_mark(&partition_timer);
    int64_t time = aom_usec_timer_elapsed(&partition_timer);
    partition_times[PARTITION_NONE] += time;
    partition_timer_on = 0;
  }
#endif
  *pb_source_variance = x->source_variance;
  if (none_rd) *none_rd = this_rdc.rdcost;
  search_state->none_rd = this_rdc.rdcost;

  // Whether the search produced a valid result.
  if (this_rdc.rate != INT_MAX) {
    // Record picked ref frame to prune ref frames for other partition types.
    if (cpi->sf.prune_ref_frame_for_rect_partitions) {
      const int ref_type = av1_ref_frame_type(ctx_none->mic.ref_frame);
      av1_update_picked_ref_frames_mask(
          x, ref_type, bsize, cm->seq_params.mib_size, mi_row, mi_col);
    }

    // Calculate the total cost and update the best partition.
    if (search_state->is_block_splittable) {
      this_rdc.rate += pt_cost;
      this_rdc.rdcost = RDCOST(x->rdmult, this_rdc.rate, this_rdc.dist);
    }

#if !CONFIG_EXT_RECUR_PARTITIONS
    *part_none_rd = this_rdc.rdcost;
#endif  // !CONFIG_EXT_RECUR_PARTITIONS
    if (this_rdc.rdcost < best_rdc->rdcost) {
      // Update if the current partition is the best
      *best_rdc = this_rdc;
      search_state->found_best_partition = true;
      pc_tree->partitioning = PARTITION_NONE;

      prune_partitions_after_none(search_state, cpi, x, sms_tree, ctx_none,
                                  &this_rdc, pb_source_variance);
    }
  }

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
static INLINE void search_partition_split(
    PartitionSearchState *search_state, AV1_COMP *const cpi, ThreadData *td,
    TileDataEnc *tile_data, TOKENEXTRA **tp, RD_STATS *best_rdc,
    int64_t *part_split_rd, PC_TREE *pc_tree, SIMPLE_MOTION_DATA_TREE *sms_tree,
    RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx, SB_MULTI_PASS_MODE multi_pass_mode) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE max_sq_part = blk_params->max_sq_part;
  const BLOCK_SIZE min_sq_part = blk_params->min_sq_part;
  const BLOCK_SIZE bsize = blk_params->bsize;
  const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);

#if CONFIG_EXT_RECUR_PARTITIONS && KEEP_PARTITION_SPLIT
  const int search_split = !search_state->terminate_partition_search &&
                           search_state->do_square_split &&
                           is_partition_valid(bsize, PARTITION_SPLIT);
#elif !CONFIG_EXT_RECUR_PARTITIONS
  const int search_split = (!search_state->terminate_partition_search &&
                            search_state->do_square_split) ||
                           blk_params->is_gt_max_sq_part;
#endif  // CONFIG_EXT_RECUR_PARTITIONS
#if !CONFIG_EXT_RECUR_PARTITIONS || KEEP_PARTITION_SPLIT
  if (!search_split) {
    return;
  }
#endif  // !CONFIF_EXT_RECUR_PARTITIONS || KEEP_PARTITION_SPLIT

  for (int i = 0; i < 4; ++i) {
    const int x_idx = (i & 1) * blk_params->mi_step_w;
    const int y_idx = (i >> 1) * blk_params->mi_step_h;
    pc_tree->split[i] = av1_alloc_pc_tree_node(
        mi_row + y_idx, mi_col + x_idx, subsize, pc_tree, PARTITION_SPLIT, i,
        i == 3, blk_params->ss_x, blk_params->ss_y);
  }

  RD_STATS sum_rdc;
  av1_init_rd_stats(&sum_rdc);
  sum_rdc.rate = search_state->partition_cost[PARTITION_SPLIT];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);

  int idx;
  for (idx = 0; idx < 4 && sum_rdc.rdcost < best_rdc->rdcost; ++idx) {
    const int x_idx = (idx & 1) * blk_params->mi_step_w;
    const int y_idx = (idx >> 1) * blk_params->mi_step_h;

    if (mi_row + y_idx >= cm->mi_rows || mi_col + x_idx >= cm->mi_cols)
      continue;

    RD_STATS this_rdc;
    int64_t *p_split_rd = &search_state->split_rd[idx];

    RD_STATS best_remain_rdcost;
    av1_rd_stats_subtraction(x->rdmult, best_rdc, &sum_rdc,
                             &best_remain_rdcost);

    int curr_quad_tree_idx = 0;
    if (frame_is_intra_only(cm) && bsize <= BLOCK_64X64) {
      curr_quad_tree_idx = x->quad_tree_idx;
      x->quad_tree_idx = 4 * curr_quad_tree_idx + idx + 1;
    }
    if (!av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row + y_idx,
                               mi_col + x_idx, subsize, max_sq_part,
                               min_sq_part, &this_rdc, best_remain_rdcost,
                               pc_tree->split[idx],
#if CONFIG_EXT_RECUR_PARTITIONS
                               sms_tree ? sms_tree->split[idx] : NULL,
                               p_split_rd, multi_pass_mode)) {
#else
                               sms_tree->split[idx], p_split_rd,
                               multi_pass_mode)) {
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      av1_invalid_rd_stats(&sum_rdc);
      break;
    }
    if (frame_is_intra_only(cm) && bsize <= BLOCK_64X64) {
      x->quad_tree_idx = curr_quad_tree_idx;
    }

    sum_rdc.rate += this_rdc.rate;
    sum_rdc.dist += this_rdc.dist;
    av1_rd_cost_update(x->rdmult, &sum_rdc);
#if !CONFIG_EXT_RECUR_PARTITIONS
    if (idx <= 1 && (bsize <= BLOCK_8X8 ||
                     pc_tree->split[idx]->partitioning == PARTITION_NONE)) {
      const MB_MODE_INFO *const mbmi = &pc_tree->split[idx]->none->mic;
      const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
      // Neither palette mode nor cfl predicted
      if (pmi->palette_size[0] == 0 && pmi->palette_size[1] == 0) {
        if (mbmi->uv_mode != UV_CFL_PRED)
          search_state->split_ctx_is_ready[idx] = 1;
      }
    }
#endif  // !CONFIG_EXT_RECUR_PARTITIONS
  }
#if CONFIG_COLLECT_PARTITION_STATS
  if (partition_timer_on) {
    aom_usec_timer_mark(&partition_timer);
    int64_t time = aom_usec_timer_elapsed(&partition_timer);
    partition_times[PARTITION_SPLIT] += time;
    partition_timer_on = 0;
  }
#endif
  const int reached_last_index = (idx == 4);

  *part_split_rd = sum_rdc.rdcost;
  if (reached_last_index && sum_rdc.rdcost < best_rdc->rdcost) {
    sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, sum_rdc.dist);
    if (sum_rdc.rdcost < best_rdc->rdcost) {
      *best_rdc = sum_rdc;
      search_state->found_best_partition = true;
      pc_tree->partitioning = PARTITION_SPLIT;
    }
  } else if (cpi->sf.less_rectangular_check_level > 0) {
    // Skip rectangular partition test when partition type none gives better
    // rd than partition type split.
    if (cpi->sf.less_rectangular_check_level == 2 || idx <= 2) {
      const int partition_none_valid = search_state->none_rd > 0;
      const int partition_none_better = search_state->none_rd < sum_rdc.rdcost;
      search_state->do_rectangular_split &=
          !(partition_none_valid && partition_none_better);
    }
  }

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)

#if CONFIG_EXT_RECUR_PARTITIONS
static INLINE int is_bsize_pruning_cand(BLOCK_SIZE bsize) {
  if (bsize == BLOCK_INVALID) {
    return 0;
  }

  const int avg_bsize = (block_size_wide[bsize] + block_size_high[bsize]) / 2;
  return avg_bsize <= 32;
}
#endif  // CONFIG_EXT_RECUR_PARTITIONS

static INLINE void search_partition_horz(PartitionSearchState *search_state,
                                         AV1_COMP *const cpi, ThreadData *td,
                                         TileDataEnc *tile_data,
                                         TOKENEXTRA **tp, RD_STATS *best_rdc,
                                         PC_TREE *pc_tree,
                                         RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx,
                                         SB_MULTI_PASS_MODE multi_pass_mode) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
#if CONFIG_EXT_RECUR_PARTITIONS
  const TileInfo *const tile_info = &tile_data->tile_info;
  const BLOCK_SIZE max_sq_part = blk_params->max_sq_part;
  const BLOCK_SIZE min_sq_part = blk_params->min_sq_part;
#else
  (void)multi_pass_mode;
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ);

  RD_STATS sum_rdc, this_rdc;
  av1_init_rd_stats(&sum_rdc);
  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state->partition_rect_allowed[HORZ]));

  if (search_state->terminate_partition_search ||
      !search_state->partition_rect_allowed[HORZ] ||
      search_state->prune_rect_part[HORZ] ||
#if CONFIG_EXT_RECUR_PARTITIONS
      !is_partition_valid(bsize, PARTITION_HORZ) ||
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      (!search_state->do_rectangular_split &&
       !av1_active_h_edge(cpi, mi_row, blk_params->mi_step_h)) ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  const int part_h_rate = search_state->partition_cost[PARTITION_HORZ];
  if (part_h_rate == INT_MAX ||
      RDCOST(x->rdmult, part_h_rate, 0) >= best_rdc->rdcost) {
    return;
  }
#if CONFIG_EXT_RECUR_PARTITIONS
  pc_tree->horizontal[0] =
      av1_alloc_pc_tree_node(mi_row, mi_col, subsize, pc_tree, PARTITION_HORZ,
                             0, 0, blk_params->ss_x, blk_params->ss_y);
  pc_tree->horizontal[1] = av1_alloc_pc_tree_node(
      mi_row + blk_params->mi_step_h, mi_col, subsize, pc_tree, PARTITION_HORZ,
      1, 1, blk_params->ss_x, blk_params->ss_y);

  if (ENABLE_FAST_RECUR_PARTITION && !frame_is_intra_only(cm) &&
      !x->must_find_valid_partition && is_bsize_pruning_cand(bsize)) {
    SMSPartitionStats part_data;
    const SimpleMotionData *up =
        av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, subsize);
    const SimpleMotionData *down = av1_get_sms_data(
        cpi, tile_info, x, mi_row + blk_params->mi_step_h, mi_col, subsize);
    part_data.sms_data[0] = up;
    part_data.sms_data[1] = down;
    part_data.num_sub_parts = 2;
    part_data.part_rate = part_h_rate;

    if (best_rdc->rdcost < INT64_MAX &&
        (mi_row + 2 * blk_params->mi_step_h <= cm->mi_rows) &&
        (mi_col + 2 * blk_params->mi_step_w <= cm->mi_cols) &&
        av1_prune_new_part(&search_state->none_data, &part_data, x->rdmult,
                           bsize)) {
      const BLOCK_SIZE subsubsize =
          get_partition_subsize(subsize, PARTITION_VERT);
      if (subsubsize == BLOCK_INVALID) {
        return;
      }

      // Do one more check to deal with recursion
      SMSPartitionStats subpart_data;
      const SimpleMotionData *upleft =
          av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, subsubsize);
      const SimpleMotionData *upright =
          av1_get_sms_data(cpi, tile_info, x, mi_row,
                           mi_col + blk_params->mi_step_w, subsubsize);
      const SimpleMotionData *downleft =
          av1_get_sms_data(cpi, tile_info, x, mi_row + blk_params->mi_step_h,
                           mi_col, subsubsize);
      const SimpleMotionData *downright =
          av1_get_sms_data(cpi, tile_info, x, mi_row + blk_params->mi_step_h,
                           mi_col + blk_params->mi_step_w, subsubsize);
      subpart_data.sms_data[0] = upleft;
      subpart_data.sms_data[1] = upright;
      subpart_data.sms_data[2] = downleft;
      subpart_data.sms_data[3] = downright;
      subpart_data.num_sub_parts = 4;
      subpart_data.part_rate = 0;
      if (av1_prune_new_part(&search_state->none_data, &subpart_data, x->rdmult,
                             bsize)) {
        return;
      }
    }
  }
#else
  for (int i = 0; i < 2; ++i) {
    pc_tree->horizontal[i] =
        av1_alloc_pmc(cm, mi_row + blk_params->mi_step_h * i, mi_col, subsize,
                      pc_tree, PARTITION_HORZ, i, blk_params->ss_x,
                      blk_params->ss_y, &td->shared_coeff_buf);
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  // Search first part
  sum_rdc.rate = search_state->partition_cost[PARTITION_HORZ];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);
  RD_STATS best_remain_rdcost;
  av1_rd_stats_subtraction(x->rdmult, best_rdc, &sum_rdc, &best_remain_rdcost);
#if CONFIG_EXT_RECUR_PARTITIONS
  av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, subsize,
                        max_sq_part, min_sq_part, &this_rdc, best_remain_rdcost,
                        pc_tree->horizontal[0], NULL, NULL, multi_pass_mode);
#else
  pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &this_rdc, PARTITION_HORZ,
                subsize, pc_tree->horizontal[0], best_remain_rdcost,
                PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  av1_rd_cost_update(x->rdmult, &this_rdc);

  if (this_rdc.rate == INT_MAX) {
    sum_rdc.rdcost = INT64_MAX;
  } else {
    sum_rdc.rate += this_rdc.rate;
    sum_rdc.dist += this_rdc.dist;
    av1_rd_cost_update(x->rdmult, &sum_rdc);
  }
  search_state->rect_part_rd[HORZ][0] = this_rdc.rdcost;

  if (sum_rdc.rdcost < best_rdc->rdcost && blk_params->has_rows) {
#if !CONFIG_EXT_RECUR_PARTITIONS
    const PICK_MODE_CONTEXT *const ctx_h = pc_tree->horizontal[0];
    const MB_MODE_INFO *const mbmi = &pc_tree->horizontal[0]->mic;
    const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
    // Neither palette mode nor cfl predicted
    if (pmi->palette_size[0] == 0 && pmi->palette_size[1] == 0) {
      if (mbmi->uv_mode != UV_CFL_PRED)
        search_state->rect_ctx_is_ready[HORZ] = 1;
    }
    av1_update_state(cpi, td, ctx_h, mi_row, mi_col, subsize, 1);
    av1_encode_superblock(cpi, tile_data, td, tp, DRY_RUN_NORMAL, subsize,
                          NULL);
#endif  // !CONFIG_EXT_RECUR_PARTITIONS

    av1_rd_stats_subtraction(x->rdmult, best_rdc, &sum_rdc,
                             &best_remain_rdcost);

    // Search second part
#if CONFIG_EXT_RECUR_PARTITIONS
    av1_rd_pick_partition(
        cpi, td, tile_data, tp, mi_row + blk_params->mi_step_h, mi_col, subsize,
        max_sq_part, min_sq_part, &this_rdc, best_remain_rdcost,
        pc_tree->horizontal[1], NULL, NULL, multi_pass_mode);
#else
    pick_sb_modes(cpi, tile_data, x, mi_row + blk_params->mi_step_h, mi_col,
                  &this_rdc, PARTITION_HORZ, subsize, pc_tree->horizontal[1],
                  best_remain_rdcost, PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    av1_rd_cost_update(x->rdmult, &this_rdc);
    search_state->rect_part_rd[HORZ][1] = this_rdc.rdcost;

    if (this_rdc.rate == INT_MAX) {
      sum_rdc.rdcost = INT64_MAX;
    } else {
      sum_rdc.rate += this_rdc.rate;
      sum_rdc.dist += this_rdc.dist;
      av1_rd_cost_update(x->rdmult, &sum_rdc);
    }
  }

  if (sum_rdc.rdcost < best_rdc->rdcost) {
    sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, sum_rdc.dist);
    if (sum_rdc.rdcost < best_rdc->rdcost) {
      *best_rdc = sum_rdc;
      search_state->found_best_partition = true;
      pc_tree->partitioning = PARTITION_HORZ;
    }
  }

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void search_partition_vert(PartitionSearchState *search_state,
                                         AV1_COMP *const cpi, ThreadData *td,
                                         TileDataEnc *tile_data,
                                         TOKENEXTRA **tp, RD_STATS *best_rdc,
                                         PC_TREE *pc_tree,
                                         RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx,
                                         SB_MULTI_PASS_MODE multi_pass_mode) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
#if CONFIG_EXT_RECUR_PARTITIONS
  const TileInfo *const tile_info = &tile_data->tile_info;
  const BLOCK_SIZE max_sq_part = blk_params->max_sq_part;
  const BLOCK_SIZE min_sq_part = blk_params->min_sq_part;
#else
  (void)multi_pass_mode;
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_VERT);

  RD_STATS sum_rdc, this_rdc;
  av1_init_rd_stats(&sum_rdc);

  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state->partition_rect_allowed[VERT]));

  if (search_state->terminate_partition_search ||
      !search_state->partition_rect_allowed[VERT] ||
      search_state->prune_rect_part[VERT] ||
#if CONFIG_EXT_RECUR_PARTITIONS
      !is_partition_valid(bsize, PARTITION_VERT) ||
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      (!search_state->do_rectangular_split &&
       !av1_active_v_edge(cpi, mi_col, blk_params->mi_step_w)) ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  const int part_v_rate = search_state->partition_cost[PARTITION_VERT];
  if (part_v_rate == INT_MAX ||
      RDCOST(x->rdmult, part_v_rate, 0) >= best_rdc->rdcost) {
    return;
  }
  av1_init_rd_stats(&sum_rdc);
#if CONFIG_EXT_RECUR_PARTITIONS
  pc_tree->vertical[0] =
      av1_alloc_pc_tree_node(mi_row, mi_col, subsize, pc_tree, PARTITION_VERT,
                             0, 0, blk_params->ss_x, blk_params->ss_y);
  pc_tree->vertical[1] = av1_alloc_pc_tree_node(
      mi_row, mi_col + blk_params->mi_step_w, subsize, pc_tree, PARTITION_VERT,
      1, 1, blk_params->ss_x, blk_params->ss_y);

  if (ENABLE_FAST_RECUR_PARTITION && !frame_is_intra_only(cm) &&
      !x->must_find_valid_partition && is_bsize_pruning_cand(bsize)) {
    const SimpleMotionData *left =
        av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, subsize);
    const SimpleMotionData *right = av1_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col + blk_params->mi_step_w, subsize);

    SMSPartitionStats part_data;
    part_data.sms_data[0] = left;
    part_data.sms_data[1] = right;
    part_data.num_sub_parts = 2;
    part_data.part_rate = part_v_rate;

    if (best_rdc->rdcost < INT64_MAX &&
        (mi_row + 2 * blk_params->mi_step_h <= cm->mi_rows) &&
        (mi_col + 2 * blk_params->mi_step_w <= cm->mi_cols) &&
        av1_prune_new_part(&search_state->none_data, &part_data, x->rdmult,
                           bsize)) {
      const BLOCK_SIZE subsubsize =
          get_partition_subsize(subsize, PARTITION_HORZ);
      if (subsubsize == BLOCK_INVALID) {
        return;
      }

      // Do one more check to deal with recursion

      SMSPartitionStats subpart_data;
      const SimpleMotionData *upleft =
          av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, subsubsize);
      const SimpleMotionData *upright =
          av1_get_sms_data(cpi, tile_info, x, mi_row,
                           mi_col + blk_params->mi_step_w, subsubsize);
      const SimpleMotionData *downleft =
          av1_get_sms_data(cpi, tile_info, x, mi_row + blk_params->mi_step_h,
                           mi_col, subsubsize);
      const SimpleMotionData *downright =
          av1_get_sms_data(cpi, tile_info, x, mi_row + blk_params->mi_step_h,
                           mi_col + blk_params->mi_step_w, subsubsize);
      subpart_data.sms_data[0] = upleft;
      subpart_data.sms_data[1] = upright;
      subpart_data.sms_data[2] = downleft;
      subpart_data.sms_data[3] = downright;
      subpart_data.num_sub_parts = 4;
      subpart_data.part_rate = 0;
      if (av1_prune_new_part(&search_state->none_data, &subpart_data, x->rdmult,
                             bsize)) {
        return;
      }
    }
  }
#else
  for (int i = 0; i < 2; ++i) {
    pc_tree->vertical[i] =
        av1_alloc_pmc(cm, mi_row, mi_col + blk_params->mi_step_w * i, subsize,
                      pc_tree, PARTITION_VERT, i, blk_params->ss_x,
                      blk_params->ss_y, &td->shared_coeff_buf);
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  sum_rdc.rate = search_state->partition_cost[PARTITION_VERT];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);
  RD_STATS best_remain_rdcost;
  av1_rd_stats_subtraction(x->rdmult, best_rdc, &sum_rdc, &best_remain_rdcost);
#if CONFIG_EXT_RECUR_PARTITIONS
  av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row, mi_col, subsize,
                        max_sq_part, min_sq_part, &this_rdc, best_remain_rdcost,
                        pc_tree->vertical[0], NULL, NULL, multi_pass_mode);
#else
  pick_sb_modes(cpi, tile_data, x, mi_row, mi_col, &this_rdc, PARTITION_VERT,
                subsize, pc_tree->vertical[0], best_remain_rdcost,
                PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  av1_rd_cost_update(x->rdmult, &this_rdc);

  if (this_rdc.rate == INT_MAX) {
    sum_rdc.rdcost = INT64_MAX;
  } else {
    sum_rdc.rate += this_rdc.rate;
    sum_rdc.dist += this_rdc.dist;
    av1_rd_cost_update(x->rdmult, &sum_rdc);
  }
  search_state->rect_part_rd[VERT][0] = this_rdc.rdcost;
  if (sum_rdc.rdcost < best_rdc->rdcost && blk_params->has_cols) {
#if !CONFIG_EXT_RECUR_PARTITIONS
    const MB_MODE_INFO *const mbmi = &pc_tree->vertical[0]->mic;
    const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
    // Neither palette mode nor cfl predicted
    if (pmi->palette_size[0] == 0 && pmi->palette_size[1] == 0) {
      if (mbmi->uv_mode != UV_CFL_PRED)
        search_state->rect_ctx_is_ready[VERT] = 1;
    }
    av1_update_state(cpi, td, pc_tree->vertical[0], mi_row, mi_col, subsize, 1);
    av1_encode_superblock(cpi, tile_data, td, tp, DRY_RUN_NORMAL, subsize,
                          NULL);
#endif  // !CONFIG_EXT_RECUR_PARTITIONS

    av1_rd_stats_subtraction(x->rdmult, best_rdc, &sum_rdc,
                             &best_remain_rdcost);
#if CONFIG_EXT_RECUR_PARTITIONS
    av1_rd_pick_partition(cpi, td, tile_data, tp, mi_row,
                          mi_col + blk_params->mi_step_w, subsize, max_sq_part,
                          min_sq_part, &this_rdc, best_remain_rdcost,
                          pc_tree->vertical[1], NULL, NULL, multi_pass_mode);
#else
    pick_sb_modes(cpi, tile_data, x, mi_row, mi_col + blk_params->mi_step_w,
                  &this_rdc, PARTITION_VERT, subsize, pc_tree->vertical[1],
                  best_remain_rdcost, PICK_MODE_RD);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    av1_rd_cost_update(x->rdmult, &this_rdc);
    search_state->rect_part_rd[VERT][1] = this_rdc.rdcost;

    if (this_rdc.rate == INT_MAX) {
      sum_rdc.rdcost = INT64_MAX;
    } else {
      sum_rdc.rate += this_rdc.rate;
      sum_rdc.dist += this_rdc.dist;
      av1_rd_cost_update(x->rdmult, &sum_rdc);
    }
  }

  av1_rd_cost_update(x->rdmult, &sum_rdc);
  if (sum_rdc.rdcost < best_rdc->rdcost) {
    *best_rdc = sum_rdc;
    search_state->found_best_partition = true;
    pc_tree->partitioning = PARTITION_VERT;
  }

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

// TODO(any): Streamline all the ab_partition searches
#if !CONFIG_EXT_RECUR_PARTITIONS
static INLINE void prune_ab_partitions(AV1_COMP *cpi, PC_TREE *pc_tree,
                                       PartitionSearchState *search_state,
                                       const MACROBLOCK *x,
                                       const RD_STATS *best_rdc,
                                       unsigned int pb_source_variance,
                                       int ext_partition_allowed) {
  const BLOCK_SIZE bsize = search_state->part_blk_params.bsize;
  int horzab_partition_allowed =
      ext_partition_allowed & cpi->oxcf.enable_ab_partitions;
  int vertab_partition_allowed =
      ext_partition_allowed & cpi->oxcf.enable_ab_partitions;

#if CONFIG_DIST_8X8
  if (x->using_dist_8x8) {
    if (block_size_high[bsize] <= 8 || block_size_wide[bsize] <= 8) {
      horzab_partition_allowed = 0;
      vertab_partition_allowed = 0;
    }
  }
#endif

  if (cpi->sf.prune_ext_partition_types_search_level) {
    if (cpi->sf.prune_ext_partition_types_search_level == 1) {
      // TODO(debargha,huisu@google.com): may need to tune the threshold for
      // pb_source_variance.
      horzab_partition_allowed &= (pc_tree->partitioning == PARTITION_HORZ ||
                                   (pc_tree->partitioning == PARTITION_NONE &&
                                    pb_source_variance < 32) ||
                                   pc_tree->partitioning == PARTITION_SPLIT);
      vertab_partition_allowed &= (pc_tree->partitioning == PARTITION_VERT ||
                                   (pc_tree->partitioning == PARTITION_NONE &&
                                    pb_source_variance < 32) ||
                                   pc_tree->partitioning == PARTITION_SPLIT);
    } else {
      horzab_partition_allowed &= (pc_tree->partitioning == PARTITION_HORZ ||
                                   pc_tree->partitioning == PARTITION_SPLIT);
      vertab_partition_allowed &= (pc_tree->partitioning == PARTITION_VERT ||
                                   pc_tree->partitioning == PARTITION_SPLIT);
    }

    clip_partition_search_state_rd(search_state);
  }
  search_state->partition_ab_allowed[HORZ][PART_A] = horzab_partition_allowed;
  search_state->partition_ab_allowed[HORZ][PART_B] = horzab_partition_allowed;
  if (cpi->sf.prune_ext_partition_types_search_level) {
    const int64_t horz_a_rd = search_state->rect_part_rd[HORZ][1] +
                              search_state->split_rd[0] +
                              search_state->split_rd[1];
    const int64_t horz_b_rd = search_state->rect_part_rd[HORZ][0] +
                              search_state->split_rd[2] +
                              search_state->split_rd[3];
    switch (cpi->sf.prune_ext_partition_types_search_level) {
      case 1:
        search_state->partition_ab_allowed[HORZ][PART_A] &=
            (horz_a_rd / 16 * 14 < best_rdc->rdcost);
        search_state->partition_ab_allowed[HORZ][PART_B] &=
            (horz_b_rd / 16 * 14 < best_rdc->rdcost);
        break;
      case 2:
      default:
        search_state->partition_ab_allowed[HORZ][PART_A] &=
            (horz_a_rd / 16 * 15 < best_rdc->rdcost);
        search_state->partition_ab_allowed[HORZ][PART_B] &=
            (horz_b_rd / 16 * 15 < best_rdc->rdcost);
        break;
    }
  }

  search_state->partition_ab_allowed[VERT][PART_A] = vertab_partition_allowed;
  search_state->partition_ab_allowed[VERT][PART_B] = vertab_partition_allowed;
  if (cpi->sf.prune_ext_partition_types_search_level) {
    const int64_t vert_a_rd = search_state->rect_part_rd[VERT][1] +
                              search_state->split_rd[0] +
                              search_state->split_rd[2];
    const int64_t vert_b_rd = search_state->rect_part_rd[VERT][0] +
                              search_state->split_rd[1] +
                              search_state->split_rd[3];
    switch (cpi->sf.prune_ext_partition_types_search_level) {
      case 1:
        search_state->partition_ab_allowed[VERT][PART_A] &=
            (vert_a_rd / 16 * 14 < best_rdc->rdcost);
        search_state->partition_ab_allowed[VERT][PART_B] &=
            (vert_b_rd / 16 * 14 < best_rdc->rdcost);
        break;
      case 2:
      default:
        search_state->partition_ab_allowed[VERT][PART_A] &=
            (vert_a_rd / 16 * 15 < best_rdc->rdcost);
        search_state->partition_ab_allowed[VERT][PART_B] &=
            (vert_b_rd / 16 * 15 < best_rdc->rdcost);
        break;
    }
  }

  if (cpi->sf.ml_prune_ab_partition && ext_partition_allowed &&
#if CONFIG_EXT_RECUR_PARTITIONS
      is_square_block(bsize) &&
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      search_state->partition_rect_allowed[HORZ] &&
      search_state->partition_rect_allowed[VERT]) {
    // TODO(huisu@google.com): x->source_variance may not be the current
    // block's variance. The correct one to use is pb_source_variance. Need to
    // re-train the model to fix it.
    av1_ml_prune_ab_partition(
        bsize, pc_tree->partitioning, get_unsigned_bits(x->source_variance),
        best_rdc->rdcost, search_state->rect_part_rd[HORZ],
        search_state->rect_part_rd[VERT], search_state->split_rd,
        &search_state->partition_ab_allowed[HORZ][PART_A],
        &search_state->partition_ab_allowed[HORZ][PART_B],
        &search_state->partition_ab_allowed[VERT][PART_A],
        &search_state->partition_ab_allowed[VERT][PART_B]);
  }

  search_state->partition_ab_allowed[HORZ][PART_A] &=
      cpi->oxcf.enable_ab_partitions;
  search_state->partition_ab_allowed[HORZ][PART_B] &=
      cpi->oxcf.enable_ab_partitions;
  search_state->partition_ab_allowed[VERT][PART_A] &=
      cpi->oxcf.enable_ab_partitions;
  search_state->partition_ab_allowed[VERT][PART_B] &=
      cpi->oxcf.enable_ab_partitions;
}

static INLINE void search_partition_horza(PartitionSearchState *search_state,
                                          AV1_COMP *const cpi, ThreadData *td,
                                          TileDataEnc *tile_data,
                                          TOKENEXTRA **tp, RD_STATS *best_rdc,
                                          PC_TREE *pc_tree,
                                          RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE rec_subsize = get_partition_subsize(bsize, PARTITION_HORZ_A);
  const BLOCK_SIZE sqr_subsize = get_partition_subsize(bsize, PARTITION_SPLIT);

  if (search_state->terminate_partition_search ||
      !search_state->partition_rect_allowed[HORZ] ||
      !search_state->partition_ab_allowed[HORZ][PART_A] ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  pc_tree->horizontala[0] = av1_alloc_pmc(
      cm, mi_row, mi_col, sqr_subsize, pc_tree, PARTITION_HORZ_A, 0,
      blk_params->ss_x, blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->horizontala[1] =
      av1_alloc_pmc(cm, mi_row, mi_col + blk_params->mi_step_w, sqr_subsize,
                    pc_tree, PARTITION_HORZ_A, 1, blk_params->ss_x,
                    blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->horizontala[2] =
      av1_alloc_pmc(cm, mi_row + blk_params->mi_step_h, mi_col, rec_subsize,
                    pc_tree, PARTITION_HORZ_A, 2, blk_params->ss_x,
                    blk_params->ss_y, &td->shared_coeff_buf);

  pc_tree->horizontala[0]->rd_mode_is_ready = 0;
  pc_tree->horizontala[1]->rd_mode_is_ready = 0;
  pc_tree->horizontala[2]->rd_mode_is_ready = 0;
  if (search_state->split_ctx_is_ready[0]) {
    av1_copy_tree_context(pc_tree->horizontala[0], pc_tree->split[0]->none);
    pc_tree->horizontala[0]->mic.partition = PARTITION_HORZ_A;
    pc_tree->horizontala[0]->rd_mode_is_ready = 1;
    if (search_state->split_ctx_is_ready[1]) {
      av1_copy_tree_context(pc_tree->horizontala[1], pc_tree->split[1]->none);
      pc_tree->horizontala[1]->mic.partition = PARTITION_HORZ_A;
      pc_tree->horizontala[1]->rd_mode_is_ready = 1;
    }
  }

  search_state->found_best_partition |= rd_test_partition3(
      cpi, td, tile_data, tp, pc_tree, best_rdc, pc_tree->horizontala, mi_row,
      mi_col, bsize, PARTITION_HORZ_A, mi_row, mi_col, sqr_subsize, mi_row,
      mi_col + blk_params->mi_step_w, sqr_subsize,
      mi_row + blk_params->mi_step_h, mi_col, rec_subsize);

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void search_partition_horzb(PartitionSearchState *search_state,
                                          AV1_COMP *const cpi, ThreadData *td,
                                          TileDataEnc *tile_data,
                                          TOKENEXTRA **tp, RD_STATS *best_rdc,
                                          PC_TREE *pc_tree,
                                          RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE rec_subsize = get_partition_subsize(bsize, PARTITION_HORZ_B);
  const BLOCK_SIZE sqr_subsize = get_partition_subsize(bsize, PARTITION_SPLIT);

  if (search_state->terminate_partition_search ||
      !search_state->partition_rect_allowed[HORZ] ||
      !search_state->partition_ab_allowed[HORZ][PART_B] ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  pc_tree->horizontalb[0] = av1_alloc_pmc(
      cm, mi_row, mi_col, rec_subsize, pc_tree, PARTITION_HORZ_B, 0,
      blk_params->ss_x, blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->horizontalb[1] =
      av1_alloc_pmc(cm, mi_row + blk_params->mi_step_h, mi_col, sqr_subsize,
                    pc_tree, PARTITION_HORZ_B, 1, blk_params->ss_x,
                    blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->horizontalb[2] = av1_alloc_pmc(
      cm, mi_row + blk_params->mi_step_h, mi_col + blk_params->mi_step_w,
      sqr_subsize, pc_tree, PARTITION_HORZ_B, 2, blk_params->ss_x,
      blk_params->ss_y, &td->shared_coeff_buf);

  pc_tree->horizontalb[0]->rd_mode_is_ready = 0;
  pc_tree->horizontalb[1]->rd_mode_is_ready = 0;
  pc_tree->horizontalb[2]->rd_mode_is_ready = 0;
  if (search_state->rect_ctx_is_ready[HORZ]) {
    av1_copy_tree_context(pc_tree->horizontalb[0], pc_tree->horizontal[0]);
    pc_tree->horizontalb[0]->mic.partition = PARTITION_HORZ_B;
    pc_tree->horizontalb[0]->rd_mode_is_ready = 1;
  }

  search_state->found_best_partition |= rd_test_partition3(
      cpi, td, tile_data, tp, pc_tree, best_rdc, pc_tree->horizontalb, mi_row,
      mi_col, bsize, PARTITION_HORZ_B, mi_row, mi_col, rec_subsize,
      mi_row + blk_params->mi_step_h, mi_col, sqr_subsize,
      mi_row + blk_params->mi_step_h, mi_col + blk_params->mi_step_w,
      sqr_subsize);

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void search_partition_verta(PartitionSearchState *search_state,
                                          AV1_COMP *const cpi, ThreadData *td,
                                          TileDataEnc *tile_data,
                                          TOKENEXTRA **tp, RD_STATS *best_rdc,
                                          PC_TREE *pc_tree,
                                          RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE rec_subsize = get_partition_subsize(bsize, PARTITION_VERT_A);
  const BLOCK_SIZE sqr_subsize = get_partition_subsize(bsize, PARTITION_SPLIT);

  if (search_state->terminate_partition_search ||
      !search_state->partition_rect_allowed[VERT] ||
      !search_state->partition_ab_allowed[VERT][PART_A] ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  pc_tree->verticala[0] = av1_alloc_pmc(
      cm, mi_row, mi_col, sqr_subsize, pc_tree, PARTITION_VERT_A, 0,
      blk_params->ss_x, blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->verticala[1] =
      av1_alloc_pmc(cm, mi_row + blk_params->mi_step_h, mi_col, sqr_subsize,
                    pc_tree, PARTITION_VERT_A, 1, blk_params->ss_x,
                    blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->verticala[2] =
      av1_alloc_pmc(cm, mi_row, mi_col + blk_params->mi_step_w, rec_subsize,
                    pc_tree, PARTITION_VERT_A, 2, blk_params->ss_x,
                    blk_params->ss_y, &td->shared_coeff_buf);

  pc_tree->verticala[0]->rd_mode_is_ready = 0;
  pc_tree->verticala[1]->rd_mode_is_ready = 0;
  pc_tree->verticala[2]->rd_mode_is_ready = 0;
  if (search_state->split_ctx_is_ready[0]) {
    av1_copy_tree_context(pc_tree->verticala[0], pc_tree->split[0]->none);
    pc_tree->verticala[0]->mic.partition = PARTITION_VERT_A;
    pc_tree->verticala[0]->rd_mode_is_ready = 1;
  }

  search_state->found_best_partition |= rd_test_partition3(
      cpi, td, tile_data, tp, pc_tree, best_rdc, pc_tree->verticala, mi_row,
      mi_col, bsize, PARTITION_VERT_A, mi_row, mi_col, sqr_subsize,
      mi_row + blk_params->mi_step_h, mi_col, sqr_subsize, mi_row,
      mi_col + blk_params->mi_step_w, rec_subsize);

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void search_partition_vertb(PartitionSearchState *search_state,
                                          AV1_COMP *const cpi, ThreadData *td,
                                          TileDataEnc *tile_data,
                                          TOKENEXTRA **tp, RD_STATS *best_rdc,
                                          PC_TREE *pc_tree,
                                          RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE rec_subsize = get_partition_subsize(bsize, PARTITION_VERT_B);
  const BLOCK_SIZE sqr_subsize = get_partition_subsize(bsize, PARTITION_SPLIT);

  if (search_state->terminate_partition_search ||
      !search_state->partition_rect_allowed[VERT] ||
      !search_state->partition_ab_allowed[VERT][PART_B] ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  pc_tree->verticalb[0] = av1_alloc_pmc(
      cm, mi_row, mi_col, rec_subsize, pc_tree, PARTITION_VERT_B, 0,
      blk_params->ss_x, blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->verticalb[1] =
      av1_alloc_pmc(cm, mi_row, mi_col + blk_params->mi_step_w, sqr_subsize,
                    pc_tree, PARTITION_VERT_B, 1, blk_params->ss_x,
                    blk_params->ss_y, &td->shared_coeff_buf);
  pc_tree->verticalb[2] = av1_alloc_pmc(
      cm, mi_row + blk_params->mi_step_h, mi_col + blk_params->mi_step_w,
      sqr_subsize, pc_tree, PARTITION_VERT_B, 2, blk_params->ss_x,
      blk_params->ss_y, &td->shared_coeff_buf);

  pc_tree->verticalb[0]->rd_mode_is_ready = 0;
  pc_tree->verticalb[1]->rd_mode_is_ready = 0;
  pc_tree->verticalb[2]->rd_mode_is_ready = 0;
  if (search_state->rect_ctx_is_ready[VERT]) {
    av1_copy_tree_context(pc_tree->verticalb[0], pc_tree->vertical[0]);
    pc_tree->verticalb[0]->mic.partition = PARTITION_VERT_B;
    pc_tree->verticalb[0]->rd_mode_is_ready = 1;
  }

  search_state->found_best_partition |= rd_test_partition3(
      cpi, td, tile_data, tp, pc_tree, best_rdc, pc_tree->verticalb, mi_row,
      mi_col, bsize, PARTITION_VERT_B, mi_row, mi_col, rec_subsize, mi_row,
      mi_col + blk_params->mi_step_w, sqr_subsize,
      mi_row + blk_params->mi_step_h, mi_col + blk_params->mi_step_w,
      sqr_subsize);

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void prune_partition_4(AV1_COMP *cpi, PC_TREE *pc_tree,
                                     PartitionSearchState *search_state,
                                     MACROBLOCK *x, const RD_STATS *best_rdc,
                                     unsigned int pb_source_variance,
                                     int ext_partition_allowed) {
  const PartitionBlkParams *blk_params = &search_state->part_blk_params;

  const BLOCK_SIZE bsize = blk_params->bsize;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const int ss_x = blk_params->ss_x, ss_y = blk_params->ss_y;

  // partition4_allowed is 1 if we can use a PARTITION_HORZ_4 or
  // PARTITION_VERT_4 for this block. This is almost the same as
  // ext_partition_allowed, except that we don't allow 128x32 or 32x128
  // blocks, so we require that bsize is not BLOCK_128X128.
  const int partition4_allowed = cpi->oxcf.enable_1to4_partitions &&
                                 ext_partition_allowed &&
                                 bsize != BLOCK_128X128;

  const int is_chroma_size_valid_horz4 = check_is_chroma_size_valid(
      PARTITION_HORZ_4, bsize, mi_row, mi_col, ss_x, ss_y, pc_tree);

  const int is_chroma_size_valid_vert4 = check_is_chroma_size_valid(
      PARTITION_VERT_4, bsize, mi_row, mi_col, ss_x, ss_y, pc_tree);

  search_state->partition_4_allowed[HORZ] =
      partition4_allowed && search_state->partition_rect_allowed[HORZ] &&
      is_chroma_size_valid_horz4;
  search_state->partition_4_allowed[VERT] =
      partition4_allowed && search_state->partition_rect_allowed[VERT] &&
      is_chroma_size_valid_vert4;
  if (cpi->sf.prune_ext_partition_types_search_level == 2) {
    search_state->partition_4_allowed[HORZ] &=
        (pc_tree->partitioning == PARTITION_HORZ ||
         pc_tree->partitioning == PARTITION_HORZ_A ||
         pc_tree->partitioning == PARTITION_HORZ_B ||
         pc_tree->partitioning == PARTITION_SPLIT ||
         pc_tree->partitioning == PARTITION_NONE);
    search_state->partition_4_allowed[VERT] &=
        (pc_tree->partitioning == PARTITION_VERT ||
         pc_tree->partitioning == PARTITION_VERT_A ||
         pc_tree->partitioning == PARTITION_VERT_B ||
         pc_tree->partitioning == PARTITION_SPLIT ||
         pc_tree->partitioning == PARTITION_NONE);
  }
  if (cpi->sf.ml_prune_4_partition && partition4_allowed &&
      search_state->partition_rect_allowed[HORZ] &&
      search_state->partition_rect_allowed[VERT]) {
    av1_ml_prune_4_partition(
        cpi, x, bsize, pc_tree->partitioning, best_rdc->rdcost,
        search_state->rect_part_rd[HORZ], search_state->rect_part_rd[VERT],
        search_state->split_rd, &search_state->partition_4_allowed[HORZ],
        &search_state->partition_4_allowed[VERT], pb_source_variance, mi_row,
        mi_col);
  }

#if CONFIG_DIST_8X8
  if (x->using_dist_8x8) {
    if (block_size_high[bsize] <= 16 || block_size_wide[bsize] <= 16) {
      search_state->partition_4_allowed[HORZ] = 0;
      search_state->partition_4_allowed[VERT] = 0;
    }
  }
#endif

  if (blk_params->width < (blk_params->min_partition_size_1d << 2)) {
    search_state->partition_4_allowed[HORZ] = 0;
    search_state->partition_4_allowed[VERT] = 0;
  }
}

static INLINE void search_partition_horz4(PartitionSearchState *search_state,
                                          AV1_COMP *const cpi, ThreadData *td,
                                          TileDataEnc *tile_data,
                                          TOKENEXTRA **tp, RD_STATS *best_rdc,
                                          PC_TREE *pc_tree,
                                          RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state->partition_4_allowed[HORZ]));

  if (search_state->terminate_partition_search ||
      !search_state->partition_4_allowed[HORZ] || !blk_params->has_rows ||
      (!search_state->do_rectangular_split &&
       !av1_active_h_edge(cpi, mi_row, blk_params->mi_step_h)) ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  RD_STATS sum_rdc;
  av1_init_rd_stats(&sum_rdc);
  BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ_4);
  const int quarter_step = mi_size_high[bsize] / 4;

  sum_rdc.rate = search_state->partition_cost[PARTITION_HORZ_4];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);

  for (int i = 0; i < 4; ++i) {
    pc_tree->horizontal4[i] =
        av1_alloc_pmc(cm, mi_row + i * quarter_step, mi_col, subsize, pc_tree,
                      PARTITION_HORZ_4, i, blk_params->ss_x, blk_params->ss_y,
                      &td->shared_coeff_buf);
  }

  for (int i = 0; i < 4; ++i) {
    const int this_mi_row = mi_row + i * quarter_step;

    if (i > 0 && this_mi_row >= cm->mi_rows) break;

    PICK_MODE_CONTEXT *ctx_this = pc_tree->horizontal4[i];

    ctx_this->rd_mode_is_ready = 0;
    if (!rd_try_subblock(cpi, td, tile_data, tp, (i == 3), this_mi_row, mi_col,
                         subsize, *best_rdc, &sum_rdc, PARTITION_HORZ_4,
                         ctx_this)) {
      av1_invalid_rd_stats(&sum_rdc);
      break;
    }
  }

  av1_rd_cost_update(x->rdmult, &sum_rdc);
  if (sum_rdc.rdcost < best_rdc->rdcost) {
    *best_rdc = sum_rdc;
    search_state->found_best_partition = true;
    pc_tree->partitioning = PARTITION_HORZ_4;
  }

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void search_partition_vert4(PartitionSearchState *search_state,
                                          AV1_COMP *const cpi, ThreadData *td,
                                          TileDataEnc *tile_data,
                                          TOKENEXTRA **tp, RD_STATS *best_rdc,
                                          PC_TREE *pc_tree,
                                          RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx) {
  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state->partition_4_allowed[VERT]));
  if (search_state->terminate_partition_search ||
      !search_state->partition_4_allowed[VERT] || !blk_params->has_cols ||
      (!search_state->do_rectangular_split &&
       !av1_active_v_edge(cpi, mi_row, blk_params->mi_step_h)) ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_VERT_4);
  RD_STATS sum_rdc;
  av1_init_rd_stats(&sum_rdc);

  const int quarter_step = mi_size_wide[bsize] / 4;

  sum_rdc.rate = search_state->partition_cost[PARTITION_VERT_4];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);

  for (int i = 0; i < 4; ++i) {
    pc_tree->vertical4[i] =
        av1_alloc_pmc(cm, mi_row, mi_col + i * quarter_step, subsize, pc_tree,
                      PARTITION_VERT_4, i, blk_params->ss_x, blk_params->ss_y,
                      &td->shared_coeff_buf);
  }

  for (int i = 0; i < 4; ++i) {
    const int this_mi_col = mi_col + i * quarter_step;

    if (i > 0 && this_mi_col >= cm->mi_cols) break;

    PICK_MODE_CONTEXT *ctx_this = pc_tree->vertical4[i];

    ctx_this->rd_mode_is_ready = 0;
    if (!rd_try_subblock(cpi, td, tile_data, tp, (i == 3), mi_row, this_mi_col,
                         subsize, *best_rdc, &sum_rdc, PARTITION_VERT_4,
                         ctx_this)) {
      av1_invalid_rd_stats(&sum_rdc);
      break;
    }
  }

  av1_rd_cost_update(x->rdmult, &sum_rdc);
  if (sum_rdc.rdcost < best_rdc->rdcost) {
    *best_rdc = sum_rdc;
    search_state->found_best_partition = true;
    pc_tree->partitioning = PARTITION_VERT_4;
  }
  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}
#else
static INLINE void prune_partition_3(AV1_COMP *cpi, PC_TREE *pc_tree,
                                     PartitionSearchState *search_state,
                                     MACROBLOCK *x, const RD_STATS *best_rdc,
                                     unsigned int pb_source_variance,
                                     int ext_partition_allowed) {
  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const BLOCK_SIZE bsize = blk_params->bsize;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;

  // partition3_allowed is 1 if we can use a PARTITION_HORZ_3 or
  // PARTITION_VERT_3 for this block. This is almost the same as
  // ext_partition_allowed, except that we don't allow 128x32 or 32x128
  // blocks, so we require that bsize is not BLOCK_128X128.
  int partition3_allowed = cpi->oxcf.enable_1to3_partitions &&
                           ext_partition_allowed && bsize != BLOCK_128X128;

  const int is_chroma_size_valid_horz3 =
      check_is_chroma_size_valid(PARTITION_HORZ_3, bsize, mi_row, mi_col,
                                 blk_params->ss_x, blk_params->ss_y, pc_tree);

  const int is_chroma_size_valid_vert3 =
      check_is_chroma_size_valid(PARTITION_VERT_3, bsize, mi_row, mi_col,
                                 blk_params->ss_x, blk_params->ss_y, pc_tree);

  search_state->partition_3_allowed[HORZ] =
      partition3_allowed && search_state->partition_rect_allowed[HORZ] &&
      is_chroma_size_valid_horz3;
  search_state->partition_3_allowed[VERT] =
      partition3_allowed && search_state->partition_rect_allowed[VERT] &&
      is_chroma_size_valid_vert3;
  if (cpi->sf.prune_ext_partition_types_search_level == 2) {
    search_state->partition_3_allowed[HORZ] &=
        (pc_tree->partitioning == PARTITION_HORZ ||
         pc_tree->partitioning == PARTITION_SPLIT ||
         pc_tree->partitioning == PARTITION_NONE);
    search_state->partition_3_allowed[VERT] &=
        (pc_tree->partitioning == PARTITION_VERT ||
         pc_tree->partitioning == PARTITION_SPLIT ||
         pc_tree->partitioning == PARTITION_NONE);
  }

  partition3_allowed &= (search_state->partition_3_allowed[HORZ] ||
                         search_state->partition_3_allowed[VERT]);

  // TODO(urvang): Rename speed feature, and change behavior / make it work.
  // currently, it's a hack (still dividing into 4 subparts to get score).
  if (cpi->sf.ml_prune_4_partition && partition3_allowed &&
      is_square_block(bsize) && search_state->partition_rect_allowed[HORZ] &&
      search_state->partition_rect_allowed[VERT]) {
    clip_partition_search_state_rd(search_state);
    av1_ml_prune_4_partition(
        cpi, x, bsize, pc_tree->partitioning, best_rdc->rdcost,
        search_state->rect_part_rd[HORZ], search_state->rect_part_rd[VERT],
        search_state->split_rd, &search_state->partition_3_allowed[HORZ],
        &search_state->partition_3_allowed[VERT], pb_source_variance, mi_row,
        mi_col);
  }

#if CONFIG_DIST_8X8
  if (x->using_dist_8x8) {
    if (block_size_high[bsize] <= 16 || block_size_wide[bsize] <= 16) {
      search_state->partition_3_allowed[HORZ] = 0;
      search_state->partition_3_allowed[VERT] = 0;
    }
  }
#endif

  if (blk_params->width < (blk_params->min_partition_size_1d << 2)) {
    search_state->partition_3_allowed[HORZ] = 0;
    search_state->partition_3_allowed[VERT] = 0;
  }
}

static INLINE void search_partition_horz_3(PartitionSearchState *search_state,
                                           AV1_COMP *const cpi, ThreadData *td,
                                           TileDataEnc *tile_data,
                                           TOKENEXTRA **tp, RD_STATS *best_rdc,
                                           PC_TREE *pc_tree,
                                           RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx,
                                           SB_MULTI_PASS_MODE multi_pass_mode) {
  const AV1_COMMON *const cm = &cpi->common;
  TileInfo *const tile_info = &tile_data->tile_info;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE sml_subsize = get_partition_subsize(bsize, PARTITION_HORZ_3);
  const BLOCK_SIZE big_subsize = get_partition_subsize(bsize, PARTITION_HORZ);

  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state->partition_3_allowed[HORZ]));

  if (search_state->terminate_partition_search ||
      !search_state->partition_3_allowed[HORZ] || !blk_params->has_rows ||
      !is_partition_valid(bsize, PARTITION_HORZ_3) ||
      !(search_state->do_rectangular_split ||
        av1_active_h_edge(cpi, mi_row, blk_params->mi_step_h)) ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  const int part_h3_rate = search_state->partition_cost[PARTITION_HORZ_3];
  if (part_h3_rate == INT_MAX ||
      RDCOST(x->rdmult, part_h3_rate, 0) >= best_rdc->rdcost) {
    return;
  }
  RD_STATS sum_rdc;
  av1_init_rd_stats(&sum_rdc);
  const int quarter_step = mi_size_high[bsize] / 4;

  sum_rdc.rate = search_state->partition_cost[PARTITION_HORZ_3];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);

  const int step_multipliers[3] = { 0, 1, 2 };
  const BLOCK_SIZE subblock_sizes[3] = { sml_subsize, big_subsize,
                                         sml_subsize };

  pc_tree->horizontal3[0] = av1_alloc_pc_tree_node(
      mi_row, mi_col, subblock_sizes[0], pc_tree, PARTITION_HORZ_3, 0, 0,
      blk_params->ss_x, blk_params->ss_y);
  pc_tree->horizontal3[1] = av1_alloc_pc_tree_node(
      mi_row + quarter_step, mi_col, subblock_sizes[1], pc_tree,
      PARTITION_HORZ_3, 1, 0, blk_params->ss_x, blk_params->ss_y);
  pc_tree->horizontal3[2] = av1_alloc_pc_tree_node(
      mi_row + quarter_step * 3, mi_col, subblock_sizes[2], pc_tree,
      PARTITION_HORZ_3, 2, 1, blk_params->ss_x, blk_params->ss_y);

  if (ENABLE_FAST_RECUR_PARTITION && !frame_is_intra_only(cm) &&
      !x->must_find_valid_partition && is_bsize_pruning_cand(bsize)) {
    const SimpleMotionData *up =
        av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, subblock_sizes[0]);
    const SimpleMotionData *middle = av1_get_sms_data(
        cpi, tile_info, x, mi_row + quarter_step, mi_col, subblock_sizes[1]);
    const SimpleMotionData *down =
        av1_get_sms_data(cpi, tile_info, x, mi_row + 3 * quarter_step, mi_col,
                         subblock_sizes[2]);

    SMSPartitionStats part_data;
    part_data.sms_data[0] = up;
    part_data.sms_data[1] = middle;
    part_data.sms_data[2] = down;
    part_data.num_sub_parts = 3;
    part_data.part_rate = part_h3_rate;

    if (best_rdc->rdcost < INT64_MAX &&
        (mi_row + 2 * blk_params->mi_step_h <= cm->mi_rows) &&
        (mi_col + 2 * blk_params->mi_step_w <= cm->mi_cols) &&
        av1_prune_new_part(&search_state->none_data, &part_data, x->rdmult,
                           bsize)) {
      const BLOCK_SIZE midsize = subblock_sizes[1];
      const BLOCK_SIZE subsubsize =
          get_partition_subsize(midsize, PARTITION_VERT);
      if (subsubsize == BLOCK_INVALID) {
        return;
      }

      // Do one more check to deal with recursion
      SMSPartitionStats subpart_data;
      const SimpleMotionData *midleft =
          av1_get_sms_data(cpi, tile_info, x, mi_row + quarter_step,
                           mi_col + 2 * quarter_step, subsubsize);
      const SimpleMotionData *midright =
          av1_get_sms_data(cpi, tile_info, x, mi_row + quarter_step,
                           mi_col + 2 * quarter_step, subsubsize);
      subpart_data.sms_data[0] = up;
      subpart_data.sms_data[1] = midleft;
      subpart_data.sms_data[2] = midright;
      subpart_data.sms_data[3] = down;
      subpart_data.num_sub_parts = 4;
      subpart_data.part_rate = 0;
      if (av1_prune_new_part(&search_state->none_data, &subpart_data, x->rdmult,
                             bsize)) {
        return;
      }
    }
  }

  int this_mi_row = mi_row;
  for (int i = 0; i < 3; ++i) {
    this_mi_row += quarter_step * step_multipliers[i];

    if (i > 0 && this_mi_row >= cm->mi_rows) break;

    SUBBLOCK_RDO_DATA rdo_data = { NULL,
                                   pc_tree->horizontal3[i],
                                   NULL,
                                   this_mi_row,
                                   mi_col,
                                   subblock_sizes[i],
                                   PARTITION_HORZ_3,
                                   i == 2,
                                   1,
                                   blk_params->max_sq_part,
                                   blk_params->min_sq_part };

    if (!rd_try_subblock_new(cpi, td, tile_data, tp, &rdo_data, *best_rdc,
                             &sum_rdc, multi_pass_mode)) {
      av1_invalid_rd_stats(&sum_rdc);
      break;
    }
  }

  av1_rd_cost_update(x->rdmult, &sum_rdc);
  if (sum_rdc.rdcost < best_rdc->rdcost) {
    *best_rdc = sum_rdc;
    search_state->found_best_partition = true;
    pc_tree->partitioning = PARTITION_HORZ_3;
  }

  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

static INLINE void search_partition_vert_3(PartitionSearchState *search_state,
                                           AV1_COMP *const cpi, ThreadData *td,
                                           TileDataEnc *tile_data,
                                           TOKENEXTRA **tp, RD_STATS *best_rdc,
                                           PC_TREE *pc_tree,
                                           RD_SEARCH_MACROBLOCK_CONTEXT *x_ctx,
                                           SB_MULTI_PASS_MODE multi_pass_mode) {
  const AV1_COMMON *const cm = &cpi->common;
  TileInfo *const tile_info = &tile_data->tile_info;
  MACROBLOCK *const x = &td->mb;
  const int num_planes = av1_num_planes(cm);

  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const BLOCK_SIZE bsize = blk_params->bsize;

  const BLOCK_SIZE sml_subsize = get_partition_subsize(bsize, PARTITION_VERT_3);
  const BLOCK_SIZE big_subsize = get_partition_subsize(bsize, PARTITION_VERT);

  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state->partition_3_allowed[VERT]));

  if (search_state->terminate_partition_search ||
      !search_state->partition_3_allowed[VERT] || !blk_params->has_cols ||
      !is_partition_valid(bsize, PARTITION_VERT_3) ||
      !(search_state->do_rectangular_split ||
        av1_active_v_edge(cpi, mi_row, blk_params->mi_step_h)) ||
      blk_params->is_gt_max_sq_part) {
    return;
  }

  const int part_v3_rate = search_state->partition_cost[PARTITION_VERT_3];
  if (part_v3_rate == INT_MAX ||
      RDCOST(x->rdmult, part_v3_rate, 0) >= best_rdc->rdcost) {
    return;
  }

  RD_STATS sum_rdc;
  av1_init_rd_stats(&sum_rdc);
  const int quarter_step = mi_size_wide[bsize] / 4;

  sum_rdc.rate = search_state->partition_cost[PARTITION_VERT_3];
  sum_rdc.rdcost = RDCOST(x->rdmult, sum_rdc.rate, 0);

  const int step_multipliers[3] = { 0, 1, 2 };
  const BLOCK_SIZE subblock_sizes[3] = { sml_subsize, big_subsize,
                                         sml_subsize };

  pc_tree->vertical3[0] = av1_alloc_pc_tree_node(
      mi_row, mi_col, subblock_sizes[0], pc_tree, PARTITION_VERT_3, 0, 0,
      blk_params->ss_x, blk_params->ss_y);
  pc_tree->vertical3[1] = av1_alloc_pc_tree_node(
      mi_row, mi_col + quarter_step, subblock_sizes[1], pc_tree,
      PARTITION_VERT_3, 1, 0, blk_params->ss_x, blk_params->ss_y);
  pc_tree->vertical3[2] = av1_alloc_pc_tree_node(
      mi_row, mi_col + quarter_step * 3, subblock_sizes[2], pc_tree,
      PARTITION_VERT_3, 2, 1, blk_params->ss_x, blk_params->ss_y);

  if (ENABLE_FAST_RECUR_PARTITION && !frame_is_intra_only(cm) &&
      !x->must_find_valid_partition && is_bsize_pruning_cand(bsize)) {
    const SimpleMotionData *left =
        av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, subblock_sizes[0]);
    const SimpleMotionData *middle = av1_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col + quarter_step, subblock_sizes[1]);
    const SimpleMotionData *right =
        av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col + 3 * quarter_step,
                         subblock_sizes[2]);

    SMSPartitionStats part_data;
    part_data.sms_data[0] = left;
    part_data.sms_data[1] = middle;
    part_data.sms_data[2] = right;
    part_data.num_sub_parts = 3;
    part_data.part_rate = part_v3_rate;

    if (best_rdc->rdcost < INT64_MAX &&
        (mi_row + 2 * blk_params->mi_step_h <= cm->mi_rows) &&
        (mi_col + 2 * blk_params->mi_step_w <= cm->mi_cols) &&
        av1_prune_new_part(&search_state->none_data, &part_data, x->rdmult,
                           bsize)) {
      const BLOCK_SIZE midsize = subblock_sizes[1];
      const BLOCK_SIZE subsubsize =
          get_partition_subsize(midsize, PARTITION_HORZ);
      if (subsubsize == BLOCK_INVALID) {
        return;
      }

      // Do one more check to deal with recursion
      SMSPartitionStats subpart_data;
      const SimpleMotionData *leftmid = av1_get_sms_data(
          cpi, tile_info, x, mi_row, mi_col + quarter_step, subsubsize);
      const SimpleMotionData *rightmid =
          av1_get_sms_data(cpi, tile_info, x, mi_row + 2 * quarter_step,
                           mi_col + quarter_step, subsubsize);
      subpart_data.sms_data[0] = left;
      subpart_data.sms_data[1] = leftmid;
      subpart_data.sms_data[2] = rightmid;
      subpart_data.sms_data[3] = right;
      subpart_data.num_sub_parts = 4;
      subpart_data.part_rate = 0;
      if (av1_prune_new_part(&search_state->none_data, &subpart_data, x->rdmult,
                             bsize)) {
        return;
      }
    }
  }

  int this_mi_col = mi_col;
  for (int i = 0; i < 3; ++i) {
    this_mi_col += quarter_step * step_multipliers[i];

    if (i > 0 && this_mi_col >= cm->mi_cols) break;

    SUBBLOCK_RDO_DATA rdo_data = { NULL,
                                   pc_tree->vertical3[i],
                                   NULL,
                                   mi_row,
                                   this_mi_col,
                                   subblock_sizes[i],
                                   PARTITION_VERT_3,
                                   i == 2,
                                   1,
                                   blk_params->max_sq_part,
                                   blk_params->min_sq_part };

    if (!rd_try_subblock_new(cpi, td, tile_data, tp, &rdo_data, *best_rdc,
                             &sum_rdc, multi_pass_mode)) {
      av1_invalid_rd_stats(&sum_rdc);
      break;
    }
  }

  av1_rd_cost_update(x->rdmult, &sum_rdc);
  if (sum_rdc.rdcost < best_rdc->rdcost) {
    *best_rdc = sum_rdc;
    search_state->found_best_partition = true;
    pc_tree->partitioning = PARTITION_VERT_3;
  }
  av1_restore_context(cm, x, x_ctx, mi_row, mi_col, bsize, num_planes);
}

#define PRUNE_WITH_PREV_PARTITION(cur_partition) \
  (prev_partition != PARTITION_INVALID && prev_partition != (cur_partition))
#endif  // CONFIG_EXT_RECUR_PARTITIONS

bool av1_rd_pick_partition(AV1_COMP *const cpi, ThreadData *td,
                           TileDataEnc *tile_data, TOKENEXTRA **tp, int mi_row,
                           int mi_col, BLOCK_SIZE bsize, BLOCK_SIZE max_sq_part,
                           BLOCK_SIZE min_sq_part, RD_STATS *rd_cost,
                           RD_STATS best_rdc, PC_TREE *pc_tree,
                           SIMPLE_MOTION_DATA_TREE *sms_tree, int64_t *none_rd,
                           SB_MULTI_PASS_MODE multi_pass_mode) {
  const AV1_COMMON *const cm = &cpi->common;
  const int num_planes = av1_num_planes(cm);
  TileInfo *const tile_info = &tile_data->tile_info;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  RD_SEARCH_MACROBLOCK_CONTEXT x_ctx;
  const TOKENEXTRA *const tp_orig = *tp;
  PartitionSearchState search_state;
  av1_init_partition_search_state(&search_state, x, cpi, pc_tree, mi_row,
                                  mi_col, bsize, max_sq_part, min_sq_part);
  const PartitionBlkParams *blk_params = &search_state.part_blk_params;
#if CONFIG_EXT_RECUR_PARTITIONS
  const PARTITION_TYPE prev_partition =
      av1_get_prev_partition(cpi, x, mi_row, mi_col, bsize);

  if (sms_tree != NULL)
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    sms_tree->partitioning = PARTITION_NONE;

  if (best_rdc.rdcost < 0) {
    av1_invalid_rd_stats(rd_cost);
    return search_state.found_best_partition;
  }

#if CONFIG_EXT_RECUR_PARTITIONS
  // Check whether there is a counterpart pc_tree node with the same size
  // and the same neighboring context at the same location but from a different
  // partition path. If yes directly copy the RDO decision made for the
  // counterpart.
  PC_TREE *counterpart_block = av1_look_for_counterpart_block(pc_tree);
  if (counterpart_block) {
    if (counterpart_block->rd_cost.rate != INT_MAX) {
      av1_copy_pc_tree_recursive(cm, pc_tree, counterpart_block,
                                 blk_params->ss_x, blk_params->ss_y,
                                 &td->shared_coeff_buf, num_planes);
      *rd_cost = pc_tree->rd_cost;
      assert(bsize != cm->seq_params.sb_size);
      if (bsize == cm->seq_params.sb_size) exit(0);

      if (!pc_tree->is_last_subblock) {
        av1_encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, DRY_RUN_NORMAL,
                      bsize, pc_tree, NULL, NULL);
      }
      return true;
    } else {
      av1_invalid_rd_stats(rd_cost);
      return false;
    }
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  if (frame_is_intra_only(cm) && bsize == BLOCK_64X64) {
    x->quad_tree_idx = 0;
    x->cnn_output_valid = 0;
  }

  if (bsize == cm->seq_params.sb_size) x->must_find_valid_partition = 0;

  if (none_rd) *none_rd = 0;

  (void)*tp_orig;

#if CONFIG_COLLECT_PARTITION_STATS
  int partition_decisions[EXT_PARTITION_TYPES] = { 0 };
  int partition_attempts[EXT_PARTITION_TYPES] = { 0 };
  int64_t partition_times[EXT_PARTITION_TYPES] = { 0 };
  struct aom_usec_timer partition_timer = { 0 };
  int partition_timer_on = 0;
#if CONFIG_COLLECT_PARTITION_STATS == 2
  PartitionStats *part_stats = &cpi->partition_stats;
#endif
#endif

#ifndef NDEBUG
  // Nothing should rely on the default value of this array (which is just
  // leftover from encoding the previous block. Setting it to fixed pattern
  // when debugging.
  // bit 0, 1, 2 are blk_skip of each plane
  // bit 4, 5, 6 are initialization checking of each plane
  memset(x->blk_skip, 0x77, sizeof(x->blk_skip));
#endif  // NDEBUG

#if !CONFIG_EXT_RECUR_PARTITIONS
  assert(block_size_wide[bsize] == block_size_high[bsize]);
#endif  // !CONFIG_EXT_RECUR_PARTITIONS

  av1_enc_set_offsets(cpi, tile_info, x, mi_row, mi_col, bsize,
                      &pc_tree->chroma_ref_info);

  // Save rdmult before it might be changed, so it can be restored later.
  const int orig_rdmult = x->rdmult;
  av1_setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, NO_AQ, NULL);

  av1_rd_cost_update(x->rdmult, &best_rdc);

  if (bsize == BLOCK_16X16 && cpi->vaq_refresh)
    x->mb_energy = av1_log_block_var(cpi, x, bsize);

  xd->above_txfm_context = cm->above_txfm_context[tile_info->tile_row] + mi_col;
  xd->left_txfm_context =
      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);
  av1_save_context(x, &x_ctx, mi_row, mi_col, bsize, num_planes);

#if !CONFIG_EXT_RECUR_PARTITIONS
  const int try_intra_cnn_split =
      frame_is_intra_only(cm) && cpi->sf.intra_cnn_split &&
      cm->seq_params.sb_size >= BLOCK_64X64 && bsize <= BLOCK_64X64 &&
      bsize >= BLOCK_8X8 && mi_row + mi_size_high[bsize] <= cm->mi_rows &&
      mi_col + mi_size_wide[bsize] <= cm->mi_cols;

  if (try_intra_cnn_split) {
    av1_intra_mode_cnn_partition(&cpi->common, x, bsize, x->quad_tree_idx,
                                 &search_state.partition_none_allowed,
                                 &search_state.partition_rect_allowed[HORZ],
                                 &search_state.partition_rect_allowed[VERT],
                                 &search_state.do_rectangular_split,
                                 &search_state.do_square_split);
  }

  // Use simple_motion_search to prune partitions. This must be done prior to
  // PARTITION_SPLIT to propagate the initial mvs to a smaller blocksize.
  const int try_split_only =
      cpi->sf.simple_motion_search_split && search_state.do_square_split &&
      bsize >= BLOCK_8X8 && mi_row + mi_size_high[bsize] <= cm->mi_rows &&
      mi_col + mi_size_wide[bsize] <= cm->mi_cols && !frame_is_intra_only(cm) &&
      !av1_superres_scaled(cm);

#if CONFIG_EXT_RECUR_PARTITIONS
  if (try_split_only && sms_tree) {
#else
  if (try_split_only) {
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    av1_simple_motion_search_based_split(
        cpi, x, sms_tree, mi_row, mi_col, bsize,
        &search_state.partition_none_allowed,
        &search_state.partition_rect_allowed[HORZ],
        &search_state.partition_rect_allowed[VERT],
        &search_state.do_rectangular_split, &search_state.do_square_split);
  }

  const int try_prune_rect =
      cpi->sf.simple_motion_search_prune_rect && !frame_is_intra_only(cm) &&
      search_state.do_rectangular_split &&
      (search_state.do_square_split || search_state.partition_none_allowed ||
       (search_state.prune_rect_part[HORZ] &&
        search_state.prune_rect_part[VERT])) &&
      (search_state.partition_rect_allowed[HORZ] ||
       search_state.partition_rect_allowed[VERT]) &&
      bsize >= BLOCK_8X8;

#if CONFIG_EXT_RECUR_PARTITIONS
  if (try_prune_rect && sms_tree) {
#else
  if (try_prune_rect) {
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    av1_simple_motion_search_prune_part(
        cpi, x, sms_tree, mi_row, mi_col, bsize,
        &search_state.partition_none_allowed,
        &search_state.partition_rect_allowed[HORZ],
        &search_state.partition_rect_allowed[VERT],
        &search_state.do_square_split, &search_state.do_rectangular_split,
        &search_state.prune_rect_part[HORZ],
        &search_state.prune_rect_part[VERT]);
  }
#endif  // !CONFIG_EXT_RECUR_PARTITIONS

  // Max and min square partition levels are defined as the partition nodes that
  // the recursive function rd_pick_partition() can reach. To implement this:
  // only PARTITION_NONE is allowed if the current node equals min_sq_part,
  // only PARTITION_SPLIT is allowed if the current node exceeds max_sq_part.
  assert(block_size_wide[min_sq_part] == block_size_high[min_sq_part]);
  assert(block_size_wide[max_sq_part] == block_size_high[max_sq_part]);
  assert(min_sq_part <= max_sq_part);
  assert(blk_params->min_partition_size_1d <=
         blk_params->max_partition_size_1d);
  if (blk_params->is_gt_max_sq_part) {
    // If current block size is larger than max, only allow split.
    search_state.partition_none_allowed = 0;
#if CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
    search_state.partition_rect_allowed[HORZ] = 1;
    search_state.partition_rect_allowed[VERT] = 1;
#else   // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
    search_state.partition_rect_allowed[HORZ] = 0;
    search_state.partition_rect_allowed[VERT] = 0;
    search_state.do_square_split = 1;
#endif  // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
  } else if (blk_params->is_le_min_sq_part) {
    // If current block size is less or equal to min, only allow none if valid
    // block large enough; only allow split otherwise.
    search_state.partition_rect_allowed[HORZ] = 0;
    search_state.partition_rect_allowed[VERT] = 0;
    // only disable square split when current block is not at the picture
    // boundary. Otherwise, inherit the square split flag from previous logic
#if CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
    search_state.partition_none_allowed = 1;
#else   // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
    if (blk_params->has_rows && blk_params->has_cols)
      search_state.do_square_split = 0;
    search_state.partition_none_allowed = !search_state.do_square_split;
#endif  // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
  }

BEGIN_PARTITION_SEARCH:
  if (x->must_find_valid_partition) {
    init_partition_allowed(&search_state, cpi, pc_tree);
#if CONFIG_EXT_RECUR_PARTITIONS
    if (!is_square_block(bsize)) {
      if (!search_state.partition_rect_allowed[HORZ] &&
          !search_state.partition_rect_allowed[VERT] &&
          !search_state.partition_none_allowed) {
        if (block_size_wide[bsize] > block_size_high[bsize])
          search_state.partition_rect_allowed[VERT] = 1;
        else
          search_state.partition_rect_allowed[HORZ] = 1;
      }
    }
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  }

  // Partition block source pixel variance.
  unsigned int pb_source_variance = UINT_MAX;

#if CONFIG_DIST_8X8
  if (x->using_dist_8x8) {
    if (block_size_high[bsize] <= 8) partition_rect_allowed[HORZ] = 0;
    if (block_size_wide[bsize] <= 8) partition_rect_allowed[VERT] = 0;
    if (block_size_high[bsize] <= 8 || block_size_wide[bsize] <= 8)
      do_square_split = 0;
  }
#endif

  // PARTITION_NONE
#if CONFIG_EXT_RECUR_PARTITIONS
  if (ENABLE_FAST_RECUR_PARTITION && !frame_is_intra_only(cm)) {
    const SimpleMotionData *whole =
        av1_get_sms_data(cpi, tile_info, x, mi_row, mi_col, bsize);
    search_state.none_data.sms_data[0] = whole;
    search_state.none_data.num_sub_parts = 1;
    search_state.none_data.part_rate =
        search_state.partition_cost[PARTITION_NONE];
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  int64_t part_none_rd = INT64_MAX;

#if CONFIG_EXT_RECUR_PARTITIONS
  if (IMPLIES(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG),
              !PRUNE_WITH_PREV_PARTITION(PARTITION_NONE))) {
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    search_partition_none(&search_state, cpi, td, tile_data, &best_rdc, pc_tree,
                          sms_tree, &x_ctx, &pb_source_variance, none_rd,
                          &part_none_rd);
#if CONFIG_EXT_RECUR_PARTITIONS
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  // PARTITION_SPLIT
#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
  int64_t part_split_rd = INT64_MAX;
  search_partition_split(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         &part_split_rd, pc_tree, sms_tree, &x_ctx,
                         multi_pass_mode);
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)

#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
  if (cpi->sf.ml_early_term_after_part_split_level &&
      !frame_is_intra_only(cm) && !search_state.terminate_partition_search &&
#if CONFIG_EXT_RECUR_PARTITIONS
      sms_tree &&
#endif  // CONFIG_EXT_RECUR_PARTITIONS
      search_state.do_rectangular_split &&
      (search_state.partition_rect_allowed[HORZ] ||
       search_state.partition_rect_allowed[VERT])) {
    av1_ml_early_term_after_split(cpi, x, sms_tree, bsize, best_rdc.rdcost,
                                  part_none_rd, part_split_rd,
                                  search_state.split_rd, mi_row, mi_col,
                                  &search_state.terminate_partition_search);
  }
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)

  if (!cpi->sf.ml_early_term_after_part_split_level &&
      cpi->sf.ml_prune_rect_partition && !frame_is_intra_only(cm) &&
      (search_state.partition_rect_allowed[HORZ] ||
       search_state.partition_rect_allowed[VERT]) &&
      !(search_state.prune_rect_part[HORZ] ||
        search_state.prune_rect_part[VERT]) &&
      !search_state.terminate_partition_search) {
    av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes,
                         &pc_tree->chroma_ref_info);
    av1_ml_prune_rect_partition(cpi, x, bsize, best_rdc.rdcost,
                                search_state.none_rd, search_state.split_rd,
                                &search_state.prune_rect_part[HORZ],
                                &search_state.prune_rect_part[VERT]);
  }

  // PARTITION_HORZ
#if CONFIG_EXT_RECUR_PARTITIONS
  if (IMPLIES(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG),
              !PRUNE_WITH_PREV_PARTITION(PARTITION_HORZ))) {
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    search_partition_horz(&search_state, cpi, td, tile_data, tp, &best_rdc,
                          pc_tree, &x_ctx, multi_pass_mode);
#if CONFIG_EXT_RECUR_PARTITIONS
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  // PARTITION_VERT
#if CONFIG_EXT_RECUR_PARTITIONS
  if (IMPLIES(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG),
              !PRUNE_WITH_PREV_PARTITION(PARTITION_VERT))) {
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    search_partition_vert(&search_state, cpi, td, tile_data, tp, &best_rdc,
                          pc_tree, &x_ctx, multi_pass_mode);
#if CONFIG_EXT_RECUR_PARTITIONS
  }
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  if (pb_source_variance == UINT_MAX) {
    av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes,
                         &pc_tree->chroma_ref_info);
    if (is_cur_buf_hbd(xd)) {
      pb_source_variance = av1_high_get_sby_perpixel_variance(
          cpi, &x->plane[0].src, bsize, xd->bd);
    } else {
      pb_source_variance =
          av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize);
    }
  }

  assert(IMPLIES(!cpi->oxcf.enable_rect_partitions,
                 !search_state.do_rectangular_split));

  const int ext_partition_allowed = search_state.do_rectangular_split &&
                                    bsize > BLOCK_8X8 && blk_params->has_rows &&
                                    blk_params->has_cols;

#if !CONFIG_EXT_RECUR_PARTITIONS
  // The standard AB partitions are allowed whenever ext-partition-types are
  // allowed
  prune_ab_partitions(cpi, pc_tree, &search_state, x, &best_rdc,
                      pb_source_variance, ext_partition_allowed);

  // PARTITION_HORZ_A
  search_partition_horza(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         pc_tree, &x_ctx);

  // PARTITION_HORZ_B
  search_partition_horzb(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         pc_tree, &x_ctx);

  // PARTITION_VERT_A
  search_partition_verta(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         pc_tree, &x_ctx);

  // PARTITION_VERT_B
  search_partition_vertb(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         pc_tree, &x_ctx);
#endif  // !CONFIG_EXT_RECUR_PARTITIONS

#if CONFIG_EXT_RECUR_PARTITIONS
  prune_partition_3(cpi, pc_tree, &search_state, x, &best_rdc,
                    pb_source_variance, ext_partition_allowed);

  // PARTITION_HORZ_3
  if (IMPLIES(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG),
              !PRUNE_WITH_PREV_PARTITION(PARTITION_HORZ_3))) {
    search_partition_horz_3(&search_state, cpi, td, tile_data, tp, &best_rdc,
                            pc_tree, &x_ctx, multi_pass_mode);
  }

  // PARTITION_VERT_3
  if (IMPLIES(should_reuse_mode(x, REUSE_PARTITION_MODE_FLAG),
              !PRUNE_WITH_PREV_PARTITION(PARTITION_VERT_3))) {
    search_partition_vert_3(&search_state, cpi, td, tile_data, tp, &best_rdc,
                            pc_tree, &x_ctx, multi_pass_mode);
  }
#else
  prune_partition_4(cpi, pc_tree, &search_state, x, &best_rdc,
                    pb_source_variance, ext_partition_allowed);

  // PARTITION_HORZ_4
  search_partition_horz4(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         pc_tree, &x_ctx);

  // PARTITION_VERT_4
  search_partition_vert4(&search_state, cpi, td, tile_data, tp, &best_rdc,
                         pc_tree, &x_ctx);
#endif  // CONFIG_EXT_RECUR_PARTITIONS

#if CONFIG_EXT_RECUR_PARTITIONS
  if (sms_tree)
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    sms_tree->partitioning = pc_tree->partitioning;
  if (bsize == cm->seq_params.sb_size && !search_state.found_best_partition) {
    // Did not find a valid partition, go back and search again, with less
    // constraint on which partition types to search.
    x->must_find_valid_partition = 1;
#if CONFIG_COLLECT_PARTITION_STATS == 2
    part_stats->partition_redo += 1;
#endif
    goto BEGIN_PARTITION_SEARCH;
  }

  *rd_cost = best_rdc;
  pc_tree->rd_cost = best_rdc;
  if (!search_state.found_best_partition) {
    av1_invalid_rd_stats(&pc_tree->rd_cost);
  } else {
#if CONFIG_EXT_RECUR_PARTITIONS
    av1_cache_best_partition(x->sms_bufs, mi_row, mi_col, bsize,
                             cm->seq_params.sb_size, pc_tree->partitioning);
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  }

#if CONFIG_COLLECT_PARTITION_STATS
  if (best_rdc.rate < INT_MAX && best_rdc.dist < INT64_MAX) {
    partition_decisions[pc_tree->partitioning] += 1;
  }
#endif

#if CONFIG_COLLECT_PARTITION_STATS == 1
  // If CONFIG_COLLECT_PARTITION_STATS is 1, then print out the stats for each
  // prediction block
  FILE *f = fopen("data.csv", "a");
  fprintf(f, "%d,%d,%d,", bsize, cm->show_frame, frame_is_intra_only(cm));
  for (int idx = 0; idx < EXT_PARTITION_TYPES; idx++) {
    fprintf(f, "%d,", partition_decisions[idx]);
  }
  for (int idx = 0; idx < EXT_PARTITION_TYPES; idx++) {
    fprintf(f, "%d,", partition_attempts[idx]);
  }
  for (int idx = 0; idx < EXT_PARTITION_TYPES; idx++) {
    fprintf(f, "%ld,", partition_times[idx]);
  }
  fprintf(f, "\n");
  fclose(f);
#endif

#if CONFIG_COLLECT_PARTITION_STATS == 2
  // If CONFIG_COLLECTION_PARTITION_STATS is 2, then we print out the stats for
  // the whole clip. So we need to pass the information upstream to the encoder
  const int bsize_idx = av1_get_bsize_idx_for_part_stats(bsize);
  int *agg_attempts = part_stats->partition_attempts[bsize_idx];
  int *agg_decisions = part_stats->partition_decisions[bsize_idx];
  int64_t *agg_times = part_stats->partition_times[bsize_idx];
  for (int idx = 0; idx < EXT_PARTITION_TYPES; idx++) {
    agg_attempts[idx] += partition_attempts[idx];
    agg_decisions[idx] += partition_decisions[idx];
    agg_times[idx] += partition_times[idx];
  }
#endif

#if CONFIG_EXT_RECUR_PARTITIONS
  if (sms_tree)
#endif  // CONFIG_EXT_RECUR_PARTITIONS
    sms_tree->partitioning = pc_tree->partitioning;
  int pc_tree_dealloc = 0;
  if (search_state.found_best_partition) {
#if CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
    assert(pc_tree->partitioning != PARTITION_SPLIT);
#endif  // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
    if (bsize == cm->seq_params.sb_size) {
      const int emit_output = multi_pass_mode != SB_DRY_PASS;
      const RUN_TYPE run_type = emit_output ? OUTPUT_ENABLED : DRY_RUN_NORMAL;

      x->cb_offset = 0;

      av1_reset_ptree_in_sbi(xd->sbi);
      av1_encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, run_type, bsize,
                    pc_tree, xd->sbi->ptree_root, NULL);
      av1_free_pc_tree_recursive(pc_tree, num_planes, 0, 0);
      pc_tree_dealloc = 1;
    } else if (!pc_tree->is_last_subblock) {
      av1_encode_sb(cpi, td, tile_data, tp, mi_row, mi_col, DRY_RUN_NORMAL,
                    bsize, pc_tree, NULL, NULL);
    }
  }

  int keep_tree = 0;
#if CONFIG_EXT_RECUR_PARTITIONS
  keep_tree = should_reuse_mode(x, REUSE_INTER_MODE_IN_INTERFRAME_FLAG |
                                       REUSE_INTRA_MODE_IN_INTERFRAME_FLAG);
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  if (!pc_tree_dealloc && !keep_tree) {
    av1_free_pc_tree_recursive(pc_tree, num_planes, 1, 1);
  }

  if (bsize == cm->seq_params.sb_size) {
    assert(best_rdc.rate < INT_MAX);
    assert(best_rdc.dist < INT64_MAX);
  } else {
    assert(tp_orig == *tp);
  }

  x->rdmult = orig_rdmult;
  return search_state.found_best_partition;
}
#endif  // !CONFIG_REALTIME_ONLY