encoder.c

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

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

#include "./aom_config.h"

#include "av1/common/alloccommon.h"
#if CONFIG_CLPF
#include "aom/aom_image.h"
#include "av1/common/clpf.h"
#include "av1/encoder/clpf_rdo.h"
#endif
#if CONFIG_DERING
#include "av1/common/dering.h"
#endif  // CONFIG_DERING
#include "av1/common/filter.h"
#include "av1/common/idct.h"
#include "av1/common/reconinter.h"
#include "av1/common/reconintra.h"
#include "av1/common/tile_common.h"

#include "av1/encoder/aq_complexity.h"
#include "av1/encoder/aq_cyclicrefresh.h"
#include "av1/encoder/aq_variance.h"
#include "av1/encoder/bitstream.h"
#if CONFIG_ANS
#include "aom_dsp/buf_ans.h"
#endif
#include "av1/encoder/context_tree.h"
#include "av1/encoder/encodeframe.h"
#include "av1/encoder/encodemv.h"
#include "av1/encoder/encoder.h"
#include "av1/encoder/ethread.h"
#include "av1/encoder/firstpass.h"
#include "av1/encoder/mbgraph.h"
#include "av1/encoder/picklpf.h"
#if CONFIG_LOOP_RESTORATION
#include "av1/encoder/pickrst.h"
#endif  // CONFIG_LOOP_RESTORATION
#include "av1/encoder/ratectrl.h"
#include "av1/encoder/rd.h"
#include "av1/encoder/resize.h"
#include "av1/encoder/segmentation.h"
#include "av1/encoder/speed_features.h"
#include "av1/encoder/temporal_filter.h"

#include "./av1_rtcd.h"
#include "./aom_dsp_rtcd.h"
#include "./aom_scale_rtcd.h"
#include "aom_dsp/psnr.h"
#if CONFIG_INTERNAL_STATS
#include "aom_dsp/ssim.h"
#endif
#include "aom_dsp/aom_dsp_common.h"
#include "aom_dsp/aom_filter.h"
#include "aom_ports/aom_timer.h"
#include "aom_ports/mem.h"
#include "aom_ports/system_state.h"
#include "aom_scale/aom_scale.h"
#if CONFIG_BITSTREAM_DEBUG
#include "aom_util/debug_util.h"
#endif  // CONFIG_BITSTREAM_DEBUG

#define AM_SEGMENT_ID_INACTIVE 7
#define AM_SEGMENT_ID_ACTIVE 0

#define SHARP_FILTER_QTHRESH 0 /* Q threshold for 8-tap sharp filter */

#define ALTREF_HIGH_PRECISION_MV 1     // Whether to use high precision mv
                                       //  for altref computation.
#define HIGH_PRECISION_MV_QTHRESH 200  // Q threshold for high precision
                                       // mv. Choose a very high value for
                                       // now so that HIGH_PRECISION is always
                                       // chosen.
// #define OUTPUT_YUV_REC
#ifdef OUTPUT_YUV_DENOISED
FILE *yuv_denoised_file = NULL;
#endif
#ifdef OUTPUT_YUV_SKINMAP
FILE *yuv_skinmap_file = NULL;
#endif
#ifdef OUTPUT_YUV_REC
FILE *yuv_rec_file;
#define FILE_NAME_LEN 100
#endif

#if 0
FILE *framepsnr;
FILE *kf_list;
FILE *keyfile;
#endif

#if CONFIG_INTERNAL_STATS
typedef enum { Y, U, V, ALL } STAT_TYPE;
#endif  // CONFIG_INTERNAL_STATS

static INLINE void Scale2Ratio(AOM_SCALING mode, int *hr, int *hs) {
  switch (mode) {
    case NORMAL:
      *hr = 1;
      *hs = 1;
      break;
    case FOURFIVE:
      *hr = 4;
      *hs = 5;
      break;
    case THREEFIVE:
      *hr = 3;
      *hs = 5;
      break;
    case ONETWO:
      *hr = 1;
      *hs = 2;
      break;
    default:
      *hr = 1;
      *hs = 1;
      assert(0);
      break;
  }
}

// Mark all inactive blocks as active. Other segmentation features may be set
// so memset cannot be used, instead only inactive blocks should be reset.
static void suppress_active_map(AV1_COMP *cpi) {
  unsigned char *const seg_map = cpi->segmentation_map;
  int i;
  if (cpi->active_map.enabled || cpi->active_map.update)
    for (i = 0; i < cpi->common.mi_rows * cpi->common.mi_cols; ++i)
      if (seg_map[i] == AM_SEGMENT_ID_INACTIVE)
        seg_map[i] = AM_SEGMENT_ID_ACTIVE;
}

static void apply_active_map(AV1_COMP *cpi) {
  struct segmentation *const seg = &cpi->common.seg;
  unsigned char *const seg_map = cpi->segmentation_map;
  const unsigned char *const active_map = cpi->active_map.map;
  int i;

  assert(AM_SEGMENT_ID_ACTIVE == CR_SEGMENT_ID_BASE);

  if (frame_is_intra_only(&cpi->common)) {
    cpi->active_map.enabled = 0;
    cpi->active_map.update = 1;
  }

  if (cpi->active_map.update) {
    if (cpi->active_map.enabled) {
      for (i = 0; i < cpi->common.mi_rows * cpi->common.mi_cols; ++i)
        if (seg_map[i] == AM_SEGMENT_ID_ACTIVE) seg_map[i] = active_map[i];
      av1_enable_segmentation(seg);
      av1_enable_segfeature(seg, AM_SEGMENT_ID_INACTIVE, SEG_LVL_SKIP);
      av1_enable_segfeature(seg, AM_SEGMENT_ID_INACTIVE, SEG_LVL_ALT_LF);
      // Setting the data to -MAX_LOOP_FILTER will result in the computed loop
      // filter level being zero regardless of the value of seg->abs_delta.
      av1_set_segdata(seg, AM_SEGMENT_ID_INACTIVE, SEG_LVL_ALT_LF,
                      -MAX_LOOP_FILTER);
    } else {
      av1_disable_segfeature(seg, AM_SEGMENT_ID_INACTIVE, SEG_LVL_SKIP);
      av1_disable_segfeature(seg, AM_SEGMENT_ID_INACTIVE, SEG_LVL_ALT_LF);
      if (seg->enabled) {
        seg->update_data = 1;
        seg->update_map = 1;
      }
    }
    cpi->active_map.update = 0;
  }
}

int av1_set_active_map(AV1_COMP *cpi, unsigned char *new_map_16x16, int rows,
                       int cols) {
  if (rows == cpi->common.mb_rows && cols == cpi->common.mb_cols) {
    unsigned char *const active_map_8x8 = cpi->active_map.map;
    const int mi_rows = cpi->common.mi_rows;
    const int mi_cols = cpi->common.mi_cols;
    cpi->active_map.update = 1;
    if (new_map_16x16) {
      int r, c;
      for (r = 0; r < mi_rows; ++r) {
        for (c = 0; c < mi_cols; ++c) {
          active_map_8x8[r * mi_cols + c] =
              new_map_16x16[(r >> 1) * cols + (c >> 1)]
                  ? AM_SEGMENT_ID_ACTIVE
                  : AM_SEGMENT_ID_INACTIVE;
        }
      }
      cpi->active_map.enabled = 1;
    } else {
      cpi->active_map.enabled = 0;
    }
    return 0;
  } else {
    return -1;
  }
}

int av1_get_active_map(AV1_COMP *cpi, unsigned char *new_map_16x16, int rows,
                       int cols) {
  if (rows == cpi->common.mb_rows && cols == cpi->common.mb_cols &&
      new_map_16x16) {
    unsigned char *const seg_map_8x8 = cpi->segmentation_map;
    const int mi_rows = cpi->common.mi_rows;
    const int mi_cols = cpi->common.mi_cols;
    memset(new_map_16x16, !cpi->active_map.enabled, rows * cols);
    if (cpi->active_map.enabled) {
      int r, c;
      for (r = 0; r < mi_rows; ++r) {
        for (c = 0; c < mi_cols; ++c) {
          // Cyclic refresh segments are considered active despite not having
          // AM_SEGMENT_ID_ACTIVE
          new_map_16x16[(r >> 1) * cols + (c >> 1)] |=
              seg_map_8x8[r * mi_cols + c] != AM_SEGMENT_ID_INACTIVE;
        }
      }
    }
    return 0;
  } else {
    return -1;
  }
}

void av1_set_high_precision_mv(AV1_COMP *cpi, int allow_high_precision_mv) {
  MACROBLOCK *const mb = &cpi->td.mb;
  cpi->common.allow_high_precision_mv = allow_high_precision_mv;

#if CONFIG_REF_MV
  if (cpi->common.allow_high_precision_mv) {
    int i;
    for (i = 0; i < NMV_CONTEXTS; ++i) {
      mb->mv_cost_stack[i] = mb->nmvcost_hp[i];
      mb->mvsadcost = mb->nmvsadcost_hp;
    }
  } else {
    int i;
    for (i = 0; i < NMV_CONTEXTS; ++i) {
      mb->mv_cost_stack[i] = mb->nmvcost[i];
      mb->mvsadcost = mb->nmvsadcost;
    }
  }
#else
  if (cpi->common.allow_high_precision_mv) {
    mb->mvcost = mb->nmvcost_hp;
    mb->mvsadcost = mb->nmvcost_hp;
  } else {
    mb->mvcost = mb->nmvcost;
    mb->mvsadcost = mb->nmvcost;
  }
#endif
}

static BLOCK_SIZE select_sb_size(const AV1_COMP *const cpi) {
#if CONFIG_EXT_PARTITION
  if (cpi->oxcf.superblock_size == AOM_SUPERBLOCK_SIZE_64X64)
    return BLOCK_64X64;

  if (cpi->oxcf.superblock_size == AOM_SUPERBLOCK_SIZE_128X128)
    return BLOCK_128X128;

  assert(cpi->oxcf.superblock_size == AOM_SUPERBLOCK_SIZE_DYNAMIC);

  assert(IMPLIES(cpi->common.tile_cols > 1,
                 cpi->common.tile_width % MAX_MIB_SIZE == 0));
  assert(IMPLIES(cpi->common.tile_rows > 1,
                 cpi->common.tile_height % MAX_MIB_SIZE == 0));

  // TODO(any): Possibly could improve this with a heuristic.
  return BLOCK_128X128;
#else
  (void)cpi;
  return BLOCK_64X64;
#endif  //  CONFIG_EXT_PARTITION
}

static void setup_frame(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;
  // Set up entropy context depending on frame type. The decoder mandates
  // the use of the default context, index 0, for keyframes and inter
  // frames where the error_resilient_mode or intra_only flag is set. For
  // other inter-frames the encoder currently uses only two contexts;
  // context 1 for ALTREF frames and context 0 for the others.
  if (frame_is_intra_only(cm) || cm->error_resilient_mode) {
    av1_setup_past_independence(cm);
  } else {
#if CONFIG_EXT_REFS
    const GF_GROUP *gf_group = &cpi->twopass.gf_group;
    if (gf_group->rf_level[gf_group->index] == GF_ARF_LOW)
      cm->frame_context_idx = EXT_ARF_FRAME;
    else if (cpi->refresh_alt_ref_frame)
      cm->frame_context_idx = ARF_FRAME;
#else
    if (cpi->refresh_alt_ref_frame) cm->frame_context_idx = ARF_FRAME;
#endif  // CONFIG_EXT_REFS
    else if (cpi->rc.is_src_frame_alt_ref)
      cm->frame_context_idx = OVERLAY_FRAME;
    else if (cpi->refresh_golden_frame)
      cm->frame_context_idx = GLD_FRAME;
#if CONFIG_EXT_REFS
    else if (cpi->refresh_bwd_ref_frame)
      cm->frame_context_idx = BRF_FRAME;
#endif  // CONFIG_EXT_REFS
    else
      cm->frame_context_idx = REGULAR_FRAME;
  }

  if (cm->frame_type == KEY_FRAME) {
    cpi->refresh_golden_frame = 1;
    cpi->refresh_alt_ref_frame = 1;
    av1_zero(cpi->interp_filter_selected);
  } else {
    *cm->fc = cm->frame_contexts[cm->frame_context_idx];
    av1_zero(cpi->interp_filter_selected[0]);
  }

  cpi->vaq_refresh = 0;

  set_sb_size(cm, select_sb_size(cpi));
}

static void av1_enc_setup_mi(AV1_COMMON *cm) {
  int i;
  cm->mi = cm->mip + cm->mi_stride + 1;
  memset(cm->mip, 0, cm->mi_stride * (cm->mi_rows + 1) * sizeof(*cm->mip));
  cm->prev_mi = cm->prev_mip + cm->mi_stride + 1;
  // Clear top border row
  memset(cm->prev_mip, 0, sizeof(*cm->prev_mip) * cm->mi_stride);
  // Clear left border column
  for (i = 1; i < cm->mi_rows + 1; ++i)
    memset(&cm->prev_mip[i * cm->mi_stride], 0, sizeof(*cm->prev_mip));

  cm->mi_grid_visible = cm->mi_grid_base + cm->mi_stride + 1;
  cm->prev_mi_grid_visible = cm->prev_mi_grid_base + cm->mi_stride + 1;

  memset(cm->mi_grid_base, 0,
         cm->mi_stride * (cm->mi_rows + 1) * sizeof(*cm->mi_grid_base));
}

static int av1_enc_alloc_mi(AV1_COMMON *cm, int mi_size) {
  cm->mip = aom_calloc(mi_size, sizeof(*cm->mip));
  if (!cm->mip) return 1;
  cm->prev_mip = aom_calloc(mi_size, sizeof(*cm->prev_mip));
  if (!cm->prev_mip) return 1;
  cm->mi_alloc_size = mi_size;

  cm->mi_grid_base = (MODE_INFO **)aom_calloc(mi_size, sizeof(MODE_INFO *));
  if (!cm->mi_grid_base) return 1;
  cm->prev_mi_grid_base =
      (MODE_INFO **)aom_calloc(mi_size, sizeof(MODE_INFO *));
  if (!cm->prev_mi_grid_base) return 1;

  return 0;
}

static void av1_enc_free_mi(AV1_COMMON *cm) {
  aom_free(cm->mip);
  cm->mip = NULL;
  aom_free(cm->prev_mip);
  cm->prev_mip = NULL;
  aom_free(cm->mi_grid_base);
  cm->mi_grid_base = NULL;
  aom_free(cm->prev_mi_grid_base);
  cm->prev_mi_grid_base = NULL;
}

static void av1_swap_mi_and_prev_mi(AV1_COMMON *cm) {
  // Current mip will be the prev_mip for the next frame.
  MODE_INFO **temp_base = cm->prev_mi_grid_base;
  MODE_INFO *temp = cm->prev_mip;
  cm->prev_mip = cm->mip;
  cm->mip = temp;

  // Update the upper left visible macroblock ptrs.
  cm->mi = cm->mip + cm->mi_stride + 1;
  cm->prev_mi = cm->prev_mip + cm->mi_stride + 1;

  cm->prev_mi_grid_base = cm->mi_grid_base;
  cm->mi_grid_base = temp_base;
  cm->mi_grid_visible = cm->mi_grid_base + cm->mi_stride + 1;
  cm->prev_mi_grid_visible = cm->prev_mi_grid_base + cm->mi_stride + 1;
}

void av1_initialize_enc(void) {
  static volatile int init_done = 0;

  if (!init_done) {
    av1_rtcd();
    aom_dsp_rtcd();
    aom_scale_rtcd();
    av1_init_intra_predictors();
    av1_init_me_luts();
    av1_rc_init_minq_luts();
    av1_entropy_mv_init();
    av1_encode_token_init();
#if CONFIG_EXT_INTER
    av1_init_wedge_masks();
#endif
    init_done = 1;
  }
}

static void dealloc_compressor_data(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;
  int i;

  aom_free(cpi->mbmi_ext_base);
  cpi->mbmi_ext_base = NULL;

#if CONFIG_PVQ
  if (cpi->oxcf.pass != 1) {
    const int tile_cols = 1 << cm->log2_tile_cols;
    const int tile_rows = 1 << cm->log2_tile_rows;
    int tile_col, tile_row;

    for (tile_row = 0; tile_row < tile_rows; ++tile_row)
      for (tile_col = 0; tile_col < tile_cols; ++tile_col) {
        TileDataEnc *tile_data =
            &cpi->tile_data[tile_row * tile_cols + tile_col];
        aom_free(tile_data->pvq_q.buf);
      }
  }
#endif
  aom_free(cpi->tile_data);
  cpi->tile_data = NULL;

  // Delete sementation map
  aom_free(cpi->segmentation_map);
  cpi->segmentation_map = NULL;

  av1_cyclic_refresh_free(cpi->cyclic_refresh);
  cpi->cyclic_refresh = NULL;

  aom_free(cpi->active_map.map);
  cpi->active_map.map = NULL;

  // Free up-sampled reference buffers.
  for (i = 0; i < (REF_FRAMES + 1); i++)
    aom_free_frame_buffer(&cpi->upsampled_ref_bufs[i].buf);

  av1_free_ref_frame_buffers(cm->buffer_pool);
  av1_free_context_buffers(cm);

  aom_free_frame_buffer(&cpi->last_frame_uf);
#if CONFIG_LOOP_RESTORATION
  av1_free_restoration_buffers(cm);
  aom_free_frame_buffer(&cpi->last_frame_db);
  aom_free_frame_buffer(&cpi->trial_frame_rst);
  aom_free(cpi->extra_rstbuf);
  av1_free_restoration_struct(&cpi->rst_search);
#endif  // CONFIG_LOOP_RESTORATION
  aom_free_frame_buffer(&cpi->scaled_source);
  aom_free_frame_buffer(&cpi->scaled_last_source);
  aom_free_frame_buffer(&cpi->alt_ref_buffer);
  av1_lookahead_destroy(cpi->lookahead);

  aom_free(cpi->tile_tok[0][0]);
  cpi->tile_tok[0][0] = 0;

  av1_free_pc_tree(&cpi->td);
  av1_free_var_tree(&cpi->td);

#if CONFIG_PALETTE
  if (cpi->common.allow_screen_content_tools)
    aom_free(cpi->td.mb.palette_buffer);
#endif  // CONFIG_PALETTE

  if (cpi->source_diff_var != NULL) {
    aom_free(cpi->source_diff_var);
    cpi->source_diff_var = NULL;
  }
#if CONFIG_ANS
  aom_buf_ans_free(&cpi->buf_ans);
#endif  // CONFIG_ANS
}

static void save_coding_context(AV1_COMP *cpi) {
  CODING_CONTEXT *const cc = &cpi->coding_context;
  AV1_COMMON *cm = &cpi->common;
#if CONFIG_REF_MV
  int i;
#endif

// Stores a snapshot of key state variables which can subsequently be
// restored with a call to av1_restore_coding_context. These functions are
// intended for use in a re-code loop in av1_compress_frame where the
// quantizer value is adjusted between loop iterations.
#if CONFIG_REF_MV
  for (i = 0; i < NMV_CONTEXTS; ++i) {
    av1_copy(cc->nmv_vec_cost[i], cpi->td.mb.nmv_vec_cost[i]);
    av1_copy(cc->nmv_costs, cpi->nmv_costs);
    av1_copy(cc->nmv_costs_hp, cpi->nmv_costs_hp);
  }
#else
  av1_copy(cc->nmvjointcost, cpi->td.mb.nmvjointcost);
#endif

  av1_copy(cc->nmvcosts, cpi->nmvcosts);
  av1_copy(cc->nmvcosts_hp, cpi->nmvcosts_hp);

  av1_copy(cc->last_ref_lf_deltas, cm->lf.last_ref_deltas);
  av1_copy(cc->last_mode_lf_deltas, cm->lf.last_mode_deltas);

  cc->fc = *cm->fc;
}

static void restore_coding_context(AV1_COMP *cpi) {
  CODING_CONTEXT *const cc = &cpi->coding_context;
  AV1_COMMON *cm = &cpi->common;
#if CONFIG_REF_MV
  int i;
#endif

// Restore key state variables to the snapshot state stored in the
// previous call to av1_save_coding_context.
#if CONFIG_REF_MV
  for (i = 0; i < NMV_CONTEXTS; ++i) {
    av1_copy(cpi->td.mb.nmv_vec_cost[i], cc->nmv_vec_cost[i]);
    av1_copy(cpi->nmv_costs, cc->nmv_costs);
    av1_copy(cpi->nmv_costs_hp, cc->nmv_costs_hp);
  }
#else
  av1_copy(cpi->td.mb.nmvjointcost, cc->nmvjointcost);
#endif

  av1_copy(cpi->nmvcosts, cc->nmvcosts);
  av1_copy(cpi->nmvcosts_hp, cc->nmvcosts_hp);

  av1_copy(cm->lf.last_ref_deltas, cc->last_ref_lf_deltas);
  av1_copy(cm->lf.last_mode_deltas, cc->last_mode_lf_deltas);

  *cm->fc = cc->fc;
}

static void configure_static_seg_features(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;
  const RATE_CONTROL *const rc = &cpi->rc;
  struct segmentation *const seg = &cm->seg;

  int high_q = (int)(rc->avg_q > 48.0);
  int qi_delta;

  // Disable and clear down for KF
  if (cm->frame_type == KEY_FRAME) {
    // Clear down the global segmentation map
    memset(cpi->segmentation_map, 0, cm->mi_rows * cm->mi_cols);
    seg->update_map = 0;
    seg->update_data = 0;
    cpi->static_mb_pct = 0;

    // Disable segmentation
    av1_disable_segmentation(seg);

    // Clear down the segment features.
    av1_clearall_segfeatures(seg);
  } else if (cpi->refresh_alt_ref_frame) {
    // If this is an alt ref frame
    // Clear down the global segmentation map
    memset(cpi->segmentation_map, 0, cm->mi_rows * cm->mi_cols);
    seg->update_map = 0;
    seg->update_data = 0;
    cpi->static_mb_pct = 0;

    // Disable segmentation and individual segment features by default
    av1_disable_segmentation(seg);
    av1_clearall_segfeatures(seg);

    // Scan frames from current to arf frame.
    // This function re-enables segmentation if appropriate.
    av1_update_mbgraph_stats(cpi);

    // If segmentation was enabled set those features needed for the
    // arf itself.
    if (seg->enabled) {
      seg->update_map = 1;
      seg->update_data = 1;

      qi_delta =
          av1_compute_qdelta(rc, rc->avg_q, rc->avg_q * 0.875, cm->bit_depth);
      av1_set_segdata(seg, 1, SEG_LVL_ALT_Q, qi_delta - 2);
      av1_set_segdata(seg, 1, SEG_LVL_ALT_LF, -2);

      av1_enable_segfeature(seg, 1, SEG_LVL_ALT_Q);
      av1_enable_segfeature(seg, 1, SEG_LVL_ALT_LF);

      // Where relevant assume segment data is delta data
      seg->abs_delta = SEGMENT_DELTADATA;
    }
  } else if (seg->enabled) {
    // All other frames if segmentation has been enabled

    // First normal frame in a valid gf or alt ref group
    if (rc->frames_since_golden == 0) {
      // Set up segment features for normal frames in an arf group
      if (rc->source_alt_ref_active) {
        seg->update_map = 0;
        seg->update_data = 1;
        seg->abs_delta = SEGMENT_DELTADATA;

        qi_delta =
            av1_compute_qdelta(rc, rc->avg_q, rc->avg_q * 1.125, cm->bit_depth);
        av1_set_segdata(seg, 1, SEG_LVL_ALT_Q, qi_delta + 2);
        av1_enable_segfeature(seg, 1, SEG_LVL_ALT_Q);

        av1_set_segdata(seg, 1, SEG_LVL_ALT_LF, -2);
        av1_enable_segfeature(seg, 1, SEG_LVL_ALT_LF);

        // Segment coding disabled for compred testing
        if (high_q || (cpi->static_mb_pct == 100)) {
          av1_set_segdata(seg, 1, SEG_LVL_REF_FRAME, ALTREF_FRAME);
          av1_enable_segfeature(seg, 1, SEG_LVL_REF_FRAME);
          av1_enable_segfeature(seg, 1, SEG_LVL_SKIP);
        }
      } else {
        // Disable segmentation and clear down features if alt ref
        // is not active for this group

        av1_disable_segmentation(seg);

        memset(cpi->segmentation_map, 0, cm->mi_rows * cm->mi_cols);

        seg->update_map = 0;
        seg->update_data = 0;

        av1_clearall_segfeatures(seg);
      }
    } else if (rc->is_src_frame_alt_ref) {
      // Special case where we are coding over the top of a previous
      // alt ref frame.
      // Segment coding disabled for compred testing

      // Enable ref frame features for segment 0 as well
      av1_enable_segfeature(seg, 0, SEG_LVL_REF_FRAME);
      av1_enable_segfeature(seg, 1, SEG_LVL_REF_FRAME);

      // All mbs should use ALTREF_FRAME
      av1_clear_segdata(seg, 0, SEG_LVL_REF_FRAME);
      av1_set_segdata(seg, 0, SEG_LVL_REF_FRAME, ALTREF_FRAME);
      av1_clear_segdata(seg, 1, SEG_LVL_REF_FRAME);
      av1_set_segdata(seg, 1, SEG_LVL_REF_FRAME, ALTREF_FRAME);

      // Skip all MBs if high Q (0,0 mv and skip coeffs)
      if (high_q) {
        av1_enable_segfeature(seg, 0, SEG_LVL_SKIP);
        av1_enable_segfeature(seg, 1, SEG_LVL_SKIP);
      }
      // Enable data update
      seg->update_data = 1;
    } else {
      // All other frames.

      // No updates.. leave things as they are.
      seg->update_map = 0;
      seg->update_data = 0;
    }
  }
}

static void update_reference_segmentation_map(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;
  MODE_INFO **mi_8x8_ptr = cm->mi_grid_visible;
  uint8_t *cache_ptr = cm->last_frame_seg_map;
  int row, col;

  for (row = 0; row < cm->mi_rows; row++) {
    MODE_INFO **mi_8x8 = mi_8x8_ptr;
    uint8_t *cache = cache_ptr;
    for (col = 0; col < cm->mi_cols; col++, mi_8x8++, cache++)
      cache[0] = mi_8x8[0]->mbmi.segment_id;
    mi_8x8_ptr += cm->mi_stride;
    cache_ptr += cm->mi_cols;
  }
}

static void alloc_raw_frame_buffers(AV1_COMP *cpi) {
  AV1_COMMON *cm = &cpi->common;
  const AV1EncoderConfig *oxcf = &cpi->oxcf;

  if (!cpi->lookahead)
    cpi->lookahead = av1_lookahead_init(oxcf->width, oxcf->height,
                                        cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                                        cm->use_highbitdepth,
#endif
                                        oxcf->lag_in_frames);
  if (!cpi->lookahead)
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate lag buffers");

  // TODO(agrange) Check if ARF is enabled and skip allocation if not.
  if (aom_realloc_frame_buffer(&cpi->alt_ref_buffer, oxcf->width, oxcf->height,
                               cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                               cm->use_highbitdepth,
#endif
                               AOM_BORDER_IN_PIXELS, cm->byte_alignment, NULL,
                               NULL, NULL))
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate altref buffer");
}

static void alloc_util_frame_buffers(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;
  if (aom_realloc_frame_buffer(&cpi->last_frame_uf, cm->width, cm->height,
                               cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                               cm->use_highbitdepth,
#endif
                               AOM_BORDER_IN_PIXELS, cm->byte_alignment, NULL,
                               NULL, NULL))
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate last frame buffer");

#if CONFIG_LOOP_RESTORATION
  if (aom_realloc_frame_buffer(&cpi->last_frame_db, cm->width, cm->height,
                               cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                               cm->use_highbitdepth,
#endif
                               AOM_BORDER_IN_PIXELS, cm->byte_alignment, NULL,
                               NULL, NULL))
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate last frame deblocked buffer");
  if (aom_realloc_frame_buffer(&cpi->trial_frame_rst, cm->width, cm->height,
                               cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                               cm->use_highbitdepth,
#endif
                               AOM_BORDER_IN_PIXELS, cm->byte_alignment, NULL,
                               NULL, NULL))
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate trial restored frame buffer");
  cpi->extra_rstbuf = (uint8_t *)aom_realloc(
      cpi->extra_rstbuf, RESTORATION_TILEPELS_MAX * sizeof(int32_t));
  if (!cpi->extra_rstbuf)
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate extra rstbuf for restoration");
  av1_alloc_restoration_struct(&cpi->rst_search, cm->width, cm->height);
#endif  // CONFIG_LOOP_RESTORATION

  if (aom_realloc_frame_buffer(&cpi->scaled_source, cm->width, cm->height,
                               cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                               cm->use_highbitdepth,
#endif
                               AOM_BORDER_IN_PIXELS, cm->byte_alignment, NULL,
                               NULL, NULL))
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate scaled source buffer");

  if (aom_realloc_frame_buffer(&cpi->scaled_last_source, cm->width, cm->height,
                               cm->subsampling_x, cm->subsampling_y,
#if CONFIG_AOM_HIGHBITDEPTH
                               cm->use_highbitdepth,
#endif
                               AOM_BORDER_IN_PIXELS, cm->byte_alignment, NULL,
                               NULL, NULL))
    aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                       "Failed to allocate scaled last source buffer");
}

static int alloc_context_buffers_ext(AV1_COMP *cpi) {
  AV1_COMMON *cm = &cpi->common;
  int mi_size = cm->mi_cols * cm->mi_rows;

  cpi->mbmi_ext_base = aom_calloc(mi_size, sizeof(*cpi->mbmi_ext_base));
  if (!cpi->mbmi_ext_base) return 1;

  return 0;
}

void av1_alloc_compressor_data(AV1_COMP *cpi) {
  AV1_COMMON *cm = &cpi->common;

  av1_alloc_context_buffers(cm, cm->width, cm->height);

  alloc_context_buffers_ext(cpi);

  aom_free(cpi->tile_tok[0][0]);

  {
    unsigned int tokens = get_token_alloc(cm->mb_rows, cm->mb_cols);
    CHECK_MEM_ERROR(cm, cpi->tile_tok[0][0],
                    aom_calloc(tokens, sizeof(*cpi->tile_tok[0][0])));
#if CONFIG_ANS
    aom_buf_ans_alloc(&cpi->buf_ans, &cm->error, tokens);
#endif  // CONFIG_ANS
  }

  av1_setup_pc_tree(&cpi->common, &cpi->td);
}

void av1_new_framerate(AV1_COMP *cpi, double framerate) {
  cpi->framerate = framerate < 0.1 ? 30 : framerate;
  av1_rc_update_framerate(cpi);
}

static void set_tile_info(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;

#if CONFIG_EXT_TILE
#if CONFIG_EXT_PARTITION
  if (cpi->oxcf.superblock_size != AOM_SUPERBLOCK_SIZE_64X64) {
    cm->tile_width = clamp(cpi->oxcf.tile_columns, 1, 32);
    cm->tile_height = clamp(cpi->oxcf.tile_rows, 1, 32);
    cm->tile_width <<= MAX_MIB_SIZE_LOG2;
    cm->tile_height <<= MAX_MIB_SIZE_LOG2;
  } else {
    cm->tile_width = clamp(cpi->oxcf.tile_columns, 1, 64);
    cm->tile_height = clamp(cpi->oxcf.tile_rows, 1, 64);
    cm->tile_width <<= MAX_MIB_SIZE_LOG2 - 1;
    cm->tile_height <<= MAX_MIB_SIZE_LOG2 - 1;
  }
#else
  cm->tile_width = clamp(cpi->oxcf.tile_columns, 1, 64);
  cm->tile_height = clamp(cpi->oxcf.tile_rows, 1, 64);
  cm->tile_width <<= MAX_MIB_SIZE_LOG2;
  cm->tile_height <<= MAX_MIB_SIZE_LOG2;
#endif  // CONFIG_EXT_PARTITION

  cm->tile_width = AOMMIN(cm->tile_width, cm->mi_cols);
  cm->tile_height = AOMMIN(cm->tile_height, cm->mi_rows);

  assert(cm->tile_width >> MAX_MIB_SIZE <= 32);
  assert(cm->tile_height >> MAX_MIB_SIZE <= 32);

  // Get the number of tiles
  cm->tile_cols = 1;
  while (cm->tile_cols * cm->tile_width < cm->mi_cols) ++cm->tile_cols;

  cm->tile_rows = 1;
  while (cm->tile_rows * cm->tile_height < cm->mi_rows) ++cm->tile_rows;
#else
  int min_log2_tile_cols, max_log2_tile_cols;
  av1_get_tile_n_bits(cm->mi_cols, &min_log2_tile_cols, &max_log2_tile_cols);

  cm->log2_tile_cols =
      clamp(cpi->oxcf.tile_columns, min_log2_tile_cols, max_log2_tile_cols);
  cm->log2_tile_rows = cpi->oxcf.tile_rows;

  cm->tile_cols = 1 << cm->log2_tile_cols;
  cm->tile_rows = 1 << cm->log2_tile_rows;

  cm->tile_width = ALIGN_POWER_OF_TWO(cm->mi_cols, MAX_MIB_SIZE_LOG2);
  cm->tile_width >>= cm->log2_tile_cols;
  cm->tile_height = ALIGN_POWER_OF_TWO(cm->mi_rows, MAX_MIB_SIZE_LOG2);
  cm->tile_height >>= cm->log2_tile_rows;

  // round to integer multiples of max superblock size
  cm->tile_width = ALIGN_POWER_OF_TWO(cm->tile_width, MAX_MIB_SIZE_LOG2);
  cm->tile_height = ALIGN_POWER_OF_TWO(cm->tile_height, MAX_MIB_SIZE_LOG2);
#endif  // CONFIG_EXT_TILE
}

static void update_frame_size(AV1_COMP *cpi) {
  AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;

  av1_set_mb_mi(cm, cm->width, cm->height);
  av1_init_context_buffers(cm);
  av1_init_macroblockd(cm, xd,
#if CONFIG_PVQ
                       NULL,
#endif
                       NULL);
  memset(cpi->mbmi_ext_base, 0,
         cm->mi_rows * cm->mi_cols * sizeof(*cpi->mbmi_ext_base));

  set_tile_info(cpi);
}

static void init_buffer_indices(AV1_COMP *cpi) {
#if CONFIG_EXT_REFS
  int fb_idx;
  for (fb_idx = 0; fb_idx < LAST_REF_FRAMES; ++fb_idx)
    cpi->lst_fb_idxes[fb_idx] = fb_idx;
  cpi->gld_fb_idx = LAST_REF_FRAMES;
  cpi->bwd_fb_idx = LAST_REF_FRAMES + 1;
  cpi->alt_fb_idx = LAST_REF_FRAMES + 2;
  for (fb_idx = 0; fb_idx < MAX_EXT_ARFS + 1; ++fb_idx)
    cpi->arf_map[fb_idx] = LAST_REF_FRAMES + 2 + fb_idx;
#else
  cpi->lst_fb_idx = 0;
  cpi->gld_fb_idx = 1;
  cpi->alt_fb_idx = 2;
#endif  // CONFIG_EXT_REFS
}

static void init_config(struct AV1_COMP *cpi, AV1EncoderConfig *oxcf) {
  AV1_COMMON *const cm = &cpi->common;

  cpi->oxcf = *oxcf;
  cpi->framerate = oxcf->init_framerate;

  cm->profile = oxcf->profile;
  cm->bit_depth = oxcf->bit_depth;
#if CONFIG_AOM_HIGHBITDEPTH
  cm->use_highbitdepth = oxcf->use_highbitdepth;
#endif
  cm->color_space = oxcf->color_space;
  cm->color_range = oxcf->color_range;

  cm->width = oxcf->width;
  cm->height = oxcf->height;
  av1_alloc_compressor_data(cpi);

  // Single thread case: use counts in common.
  cpi->td.counts = &cm->counts;

  // change includes all joint functionality
  av1_change_config(cpi, oxcf);

  cpi->static_mb_pct = 0;
  cpi->ref_frame_flags = 0;

  init_buffer_indices(cpi);
}

static void set_rc_buffer_sizes(RATE_CONTROL *rc,
                                const AV1EncoderConfig *oxcf) {
  const int64_t bandwidth = oxcf->target_bandwidth;
  const int64_t starting = oxcf->starting_buffer_level_ms;
  const int64_t optimal = oxcf->optimal_buffer_level_ms;
  const int64_t maximum = oxcf->maximum_buffer_size_ms;

  rc->starting_buffer_level = starting * bandwidth / 1000;
  rc->optimal_buffer_level =
      (optimal == 0) ? bandwidth / 8 : optimal * bandwidth / 1000;
  rc->maximum_buffer_size =
      (maximum == 0) ? bandwidth / 8 : maximum * bandwidth / 1000;
}

#if CONFIG_AOM_HIGHBITDEPTH
#define HIGHBD_BFP(BT, SDF, SDAF, VF, SVF, SVAF, SDX3F, SDX8F, SDX4DF) \
  cpi->fn_ptr[BT].sdf = SDF;                                           \
  cpi->fn_ptr[BT].sdaf = SDAF;                                         \
  cpi->fn_ptr[BT].vf = VF;                                             \
  cpi->fn_ptr[BT].svf = SVF;                                           \
  cpi->fn_ptr[BT].svaf = SVAF;                                         \
  cpi->fn_ptr[BT].sdx3f = SDX3F;                                       \
  cpi->fn_ptr[BT].sdx8f = SDX8F;                                       \
  cpi->fn_ptr[BT].sdx4df = SDX4DF;

#define MAKE_BFP_SAD_WRAPPER(fnname)                                           \
  static unsigned int fnname##_bits8(const uint8_t *src_ptr,                   \
                                     int source_stride,                        \
                                     const uint8_t *ref_ptr, int ref_stride) { \
    return fnname(src_ptr, source_stride, ref_ptr, ref_stride);                \
  }                                                                            \
  static unsigned int fnname##_bits10(                                         \
      const uint8_t *src_ptr, int source_stride, const uint8_t *ref_ptr,       \
      int ref_stride) {                                                        \
    return fnname(src_ptr, source_stride, ref_ptr, ref_stride) >> 2;           \
  }                                                                            \
  static unsigned int fnname##_bits12(                                         \
      const uint8_t *src_ptr, int source_stride, const uint8_t *ref_ptr,       \
      int ref_stride) {                                                        \
    return fnname(src_ptr, source_stride, ref_ptr, ref_stride) >> 4;           \
  }

#define MAKE_BFP_SADAVG_WRAPPER(fnname)                                        \
  static unsigned int fnname##_bits8(                                          \
      const uint8_t *src_ptr, int source_stride, const uint8_t *ref_ptr,       \
      int ref_stride, const uint8_t *second_pred) {                            \
    return fnname(src_ptr, source_stride, ref_ptr, ref_stride, second_pred);   \
  }                                                                            \
  static unsigned int fnname##_bits10(                                         \
      const uint8_t *src_ptr, int source_stride, const uint8_t *ref_ptr,       \
      int ref_stride, const uint8_t *second_pred) {                            \
    return fnname(src_ptr, source_stride, ref_ptr, ref_stride, second_pred) >> \
           2;                                                                  \
  }                                                                            \
  static unsigned int fnname##_bits12(                                         \
      const uint8_t *src_ptr, int source_stride, const uint8_t *ref_ptr,       \
      int ref_stride, const uint8_t *second_pred) {                            \
    return fnname(src_ptr, source_stride, ref_ptr, ref_stride, second_pred) >> \
           4;                                                                  \
  }

#define MAKE_BFP_SAD3_WRAPPER(fnname)                                    \
  static void fnname##_bits8(const uint8_t *src_ptr, int source_stride,  \
                             const uint8_t *ref_ptr, int ref_stride,     \
                             unsigned int *sad_array) {                  \
    fnname(src_ptr, source_stride, ref_ptr, ref_stride, sad_array);      \
  }                                                                      \
  static void fnname##_bits10(const uint8_t *src_ptr, int source_stride, \
                              const uint8_t *ref_ptr, int ref_stride,    \
                              unsigned int *sad_array) {                 \
    int i;                                                               \
    fnname(src_ptr, source_stride, ref_ptr, ref_stride, sad_array);      \
    for (i = 0; i < 3; i++) sad_array[i] >>= 2;                          \
  }                                                                      \