tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

encodetxb.c (35842B)


      1 /*
      2 * Copyright (c) 2017, Alliance for Open Media. All rights reserved.
      3 *
      4 * This source code is subject to the terms of the BSD 2 Clause License and
      5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
      6 * was not distributed with this source code in the LICENSE file, you can
      7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
      8 * Media Patent License 1.0 was not distributed with this source code in the
      9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
     10 */
     11 
     12 #include "av1/encoder/encodetxb.h"
     13 
     14 #include <stdint.h>
     15 
     16 #include "aom_ports/mem.h"
     17 #include "av1/common/blockd.h"
     18 #include "av1/common/idct.h"
     19 #include "av1/common/pred_common.h"
     20 #include "av1/common/scan.h"
     21 #include "av1/encoder/bitstream.h"
     22 #include "av1/encoder/cost.h"
     23 #include "av1/encoder/encodeframe.h"
     24 #include "av1/encoder/hash.h"
     25 #include "av1/encoder/rdopt.h"
     26 #include "av1/encoder/tokenize.h"
     27 
     28 void av1_alloc_txb_buf(AV1_COMP *cpi) {
     29  AV1_COMMON *cm = &cpi->common;
     30  CoeffBufferPool *coeff_buf_pool = &cpi->coeff_buffer_pool;
     31  const int num_sb_rows =
     32      CEIL_POWER_OF_TWO(cm->mi_params.mi_rows, cm->seq_params->mib_size_log2);
     33  const int num_sb_cols =
     34      CEIL_POWER_OF_TWO(cm->mi_params.mi_cols, cm->seq_params->mib_size_log2);
     35  const int size = num_sb_rows * num_sb_cols;
     36  const int num_planes = av1_num_planes(cm);
     37  const int subsampling_x = cm->seq_params->subsampling_x;
     38  const int subsampling_y = cm->seq_params->subsampling_y;
     39  const int luma_max_sb_square =
     40      1 << num_pels_log2_lookup[cm->seq_params->sb_size];
     41  const int chroma_max_sb_square =
     42      luma_max_sb_square >> (subsampling_x + subsampling_y);
     43  const int total_max_sb_square =
     44      (luma_max_sb_square + (num_planes - 1) * chroma_max_sb_square);
     45  if ((size_t)size > SIZE_MAX / (size_t)total_max_sb_square) {
     46    aom_internal_error(cm->error, AOM_CODEC_ERROR,
     47                       "A multiplication would overflow size_t");
     48  }
     49  const size_t num_tcoeffs = (size_t)size * (size_t)total_max_sb_square;
     50  const int txb_unit_size = TX_SIZE_W_MIN * TX_SIZE_H_MIN;
     51 
     52  av1_free_txb_buf(cpi);
     53  // TODO(jingning): This should be further reduced.
     54  CHECK_MEM_ERROR(cm, cpi->coeff_buffer_base,
     55                  aom_malloc(sizeof(*cpi->coeff_buffer_base) * size));
     56  if (sizeof(*coeff_buf_pool->tcoeff) > SIZE_MAX / num_tcoeffs) {
     57    aom_internal_error(cm->error, AOM_CODEC_ERROR,
     58                       "A multiplication would overflow size_t");
     59  }
     60  CHECK_MEM_ERROR(
     61      cm, coeff_buf_pool->tcoeff,
     62      aom_memalign(32, sizeof(*coeff_buf_pool->tcoeff) * num_tcoeffs));
     63  if (sizeof(*coeff_buf_pool->eobs) > SIZE_MAX / num_tcoeffs) {
     64    aom_internal_error(cm->error, AOM_CODEC_ERROR,
     65                       "A multiplication would overflow size_t");
     66  }
     67  CHECK_MEM_ERROR(
     68      cm, coeff_buf_pool->eobs,
     69      aom_malloc(sizeof(*coeff_buf_pool->eobs) * num_tcoeffs / txb_unit_size));
     70  if (sizeof(*coeff_buf_pool->entropy_ctx) > SIZE_MAX / num_tcoeffs) {
     71    aom_internal_error(cm->error, AOM_CODEC_ERROR,
     72                       "A multiplication would overflow size_t");
     73  }
     74  CHECK_MEM_ERROR(cm, coeff_buf_pool->entropy_ctx,
     75                  aom_malloc(sizeof(*coeff_buf_pool->entropy_ctx) *
     76                             num_tcoeffs / txb_unit_size));
     77 
     78  tran_low_t *tcoeff_ptr = coeff_buf_pool->tcoeff;
     79  uint16_t *eob_ptr = coeff_buf_pool->eobs;
     80  uint8_t *entropy_ctx_ptr = coeff_buf_pool->entropy_ctx;
     81  for (int i = 0; i < size; i++) {
     82    for (int plane = 0; plane < num_planes; plane++) {
     83      const int max_sb_square =
     84          (plane == AOM_PLANE_Y) ? luma_max_sb_square : chroma_max_sb_square;
     85      cpi->coeff_buffer_base[i].tcoeff[plane] = tcoeff_ptr;
     86      cpi->coeff_buffer_base[i].eobs[plane] = eob_ptr;
     87      cpi->coeff_buffer_base[i].entropy_ctx[plane] = entropy_ctx_ptr;
     88      tcoeff_ptr += max_sb_square;
     89      eob_ptr += max_sb_square / txb_unit_size;
     90      entropy_ctx_ptr += max_sb_square / txb_unit_size;
     91    }
     92  }
     93 }
     94 
     95 void av1_free_txb_buf(AV1_COMP *cpi) {
     96  CoeffBufferPool *coeff_buf_pool = &cpi->coeff_buffer_pool;
     97  aom_free(cpi->coeff_buffer_base);
     98  cpi->coeff_buffer_base = NULL;
     99  aom_free(coeff_buf_pool->tcoeff);
    100  coeff_buf_pool->tcoeff = NULL;
    101  aom_free(coeff_buf_pool->eobs);
    102  coeff_buf_pool->eobs = NULL;
    103  aom_free(coeff_buf_pool->entropy_ctx);
    104  coeff_buf_pool->entropy_ctx = NULL;
    105 }
    106 
    107 static void write_golomb(aom_writer *w, int level) {
    108  int x = level + 1;
    109  int i = x;
    110  int length = 0;
    111 
    112  while (i) {
    113    i >>= 1;
    114    ++length;
    115  }
    116  assert(length > 0);
    117 
    118  for (i = 0; i < length - 1; ++i) aom_write_bit(w, 0);
    119 
    120  for (i = length - 1; i >= 0; --i) aom_write_bit(w, (x >> i) & 0x01);
    121 }
    122 
    123 static const int8_t eob_to_pos_small[33] = {
    124  0, 1, 2,                                        // 0-2
    125  3, 3,                                           // 3-4
    126  4, 4, 4, 4,                                     // 5-8
    127  5, 5, 5, 5, 5, 5, 5, 5,                         // 9-16
    128  6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6  // 17-32
    129 };
    130 
    131 static const int8_t eob_to_pos_large[17] = {
    132  6,                               // place holder
    133  7,                               // 33-64
    134  8,  8,                           // 65-128
    135  9,  9,  9,  9,                   // 129-256
    136  10, 10, 10, 10, 10, 10, 10, 10,  // 257-512
    137  11                               // 513-
    138 };
    139 
    140 int av1_get_eob_pos_token(const int eob, int *const extra) {
    141  int t;
    142 
    143  if (eob < 33) {
    144    t = eob_to_pos_small[eob];
    145  } else {
    146    const int e = AOMMIN((eob - 1) >> 5, 16);
    147    t = eob_to_pos_large[e];
    148  }
    149 
    150  *extra = eob - av1_eob_group_start[t];
    151 
    152  return t;
    153 }
    154 
    155 #if CONFIG_ENTROPY_STATS
    156 static void update_eob_context(int cdf_idx, int eob, TX_SIZE tx_size,
    157                               TX_CLASS tx_class, PLANE_TYPE plane,
    158                               FRAME_CONTEXT *ec_ctx, FRAME_COUNTS *counts,
    159                               uint8_t allow_update_cdf) {
    160 #else
    161 static void update_eob_context(int eob, TX_SIZE tx_size, TX_CLASS tx_class,
    162                               PLANE_TYPE plane, FRAME_CONTEXT *ec_ctx,
    163                               uint8_t allow_update_cdf) {
    164 #endif
    165  int eob_extra;
    166  const int eob_pt = av1_get_eob_pos_token(eob, &eob_extra);
    167  TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
    168 
    169  const int eob_multi_size = txsize_log2_minus4[tx_size];
    170  const int eob_multi_ctx = (tx_class == TX_CLASS_2D) ? 0 : 1;
    171 
    172  switch (eob_multi_size) {
    173    case 0:
    174 #if CONFIG_ENTROPY_STATS
    175      ++counts->eob_multi16[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    176 #endif
    177      if (allow_update_cdf)
    178        update_cdf(ec_ctx->eob_flag_cdf16[plane][eob_multi_ctx], eob_pt - 1, 5);
    179      break;
    180    case 1:
    181 #if CONFIG_ENTROPY_STATS
    182      ++counts->eob_multi32[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    183 #endif
    184      if (allow_update_cdf)
    185        update_cdf(ec_ctx->eob_flag_cdf32[plane][eob_multi_ctx], eob_pt - 1, 6);
    186      break;
    187    case 2:
    188 #if CONFIG_ENTROPY_STATS
    189      ++counts->eob_multi64[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    190 #endif
    191      if (allow_update_cdf)
    192        update_cdf(ec_ctx->eob_flag_cdf64[plane][eob_multi_ctx], eob_pt - 1, 7);
    193      break;
    194    case 3:
    195 #if CONFIG_ENTROPY_STATS
    196      ++counts->eob_multi128[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    197 #endif
    198      if (allow_update_cdf) {
    199        update_cdf(ec_ctx->eob_flag_cdf128[plane][eob_multi_ctx], eob_pt - 1,
    200                   8);
    201      }
    202      break;
    203    case 4:
    204 #if CONFIG_ENTROPY_STATS
    205      ++counts->eob_multi256[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    206 #endif
    207      if (allow_update_cdf) {
    208        update_cdf(ec_ctx->eob_flag_cdf256[plane][eob_multi_ctx], eob_pt - 1,
    209                   9);
    210      }
    211      break;
    212    case 5:
    213 #if CONFIG_ENTROPY_STATS
    214      ++counts->eob_multi512[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    215 #endif
    216      if (allow_update_cdf) {
    217        update_cdf(ec_ctx->eob_flag_cdf512[plane][eob_multi_ctx], eob_pt - 1,
    218                   10);
    219      }
    220      break;
    221    case 6:
    222    default:
    223 #if CONFIG_ENTROPY_STATS
    224      ++counts->eob_multi1024[cdf_idx][plane][eob_multi_ctx][eob_pt - 1];
    225 #endif
    226      if (allow_update_cdf) {
    227        update_cdf(ec_ctx->eob_flag_cdf1024[plane][eob_multi_ctx], eob_pt - 1,
    228                   11);
    229      }
    230      break;
    231  }
    232 
    233  if (av1_eob_offset_bits[eob_pt] > 0) {
    234    int eob_ctx = eob_pt - 3;
    235    int eob_shift = av1_eob_offset_bits[eob_pt] - 1;
    236    int bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
    237 #if CONFIG_ENTROPY_STATS
    238    counts->eob_extra[cdf_idx][txs_ctx][plane][eob_pt][bit]++;
    239 #endif  // CONFIG_ENTROPY_STATS
    240    if (allow_update_cdf)
    241      update_cdf(ec_ctx->eob_extra_cdf[txs_ctx][plane][eob_ctx], bit, 2);
    242  }
    243 }
    244 
    245 static inline int get_nz_map_ctx(const uint8_t *const levels,
    246                                 const int coeff_idx, const int bhl,
    247                                 const int width, const int scan_idx,
    248                                 const int is_eob, const TX_SIZE tx_size,
    249                                 const TX_CLASS tx_class) {
    250  if (is_eob) {
    251    if (scan_idx == 0) return 0;
    252    if (scan_idx <= (width << bhl) / 8) return 1;
    253    if (scan_idx <= (width << bhl) / 4) return 2;
    254    return 3;
    255  }
    256  const int stats =
    257      get_nz_mag(levels + get_padded_idx(coeff_idx, bhl), bhl, tx_class);
    258  return get_nz_map_ctx_from_stats(stats, coeff_idx, bhl, tx_size, tx_class);
    259 }
    260 
    261 void av1_txb_init_levels_c(const tran_low_t *const coeff, const int width,
    262                           const int height, uint8_t *const levels) {
    263  const int stride = height + TX_PAD_HOR;
    264  uint8_t *ls = levels;
    265 
    266  memset(levels + stride * width, 0,
    267         sizeof(*levels) * (TX_PAD_BOTTOM * stride + TX_PAD_END));
    268 
    269  for (int i = 0; i < width; i++) {
    270    for (int j = 0; j < height; j++) {
    271      *ls++ = (uint8_t)clamp(abs(coeff[i * height + j]), 0, INT8_MAX);
    272    }
    273    for (int j = 0; j < TX_PAD_HOR; j++) {
    274      *ls++ = 0;
    275    }
    276  }
    277 }
    278 
    279 void av1_get_nz_map_contexts_c(const uint8_t *const levels,
    280                               const int16_t *const scan, const uint16_t eob,
    281                               const TX_SIZE tx_size, const TX_CLASS tx_class,
    282                               int8_t *const coeff_contexts) {
    283  const int bhl = get_txb_bhl(tx_size);
    284  const int width = get_txb_wide(tx_size);
    285  for (int i = 0; i < eob; ++i) {
    286    const int pos = scan[i];
    287    coeff_contexts[pos] = get_nz_map_ctx(levels, pos, bhl, width, i,
    288                                         i == eob - 1, tx_size, tx_class);
    289  }
    290 }
    291 
    292 void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *const x,
    293                          aom_writer *w, int blk_row, int blk_col, int plane,
    294                          int block, TX_SIZE tx_size) {
    295  MACROBLOCKD *xd = &x->e_mbd;
    296  const CB_COEFF_BUFFER *cb_coef_buff = x->cb_coef_buff;
    297  const PLANE_TYPE plane_type = get_plane_type(plane);
    298  const int txb_offset = x->mbmi_ext_frame->cb_offset[plane_type] /
    299                         (TX_SIZE_W_MIN * TX_SIZE_H_MIN);
    300  const uint16_t *eob_txb = cb_coef_buff->eobs[plane] + txb_offset;
    301  const uint16_t eob = eob_txb[block];
    302  const uint8_t *entropy_ctx = cb_coef_buff->entropy_ctx[plane] + txb_offset;
    303  const int txb_skip_ctx = entropy_ctx[block] & TXB_SKIP_CTX_MASK;
    304  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
    305  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
    306  aom_write_symbol(w, eob == 0, ec_ctx->txb_skip_cdf[txs_ctx][txb_skip_ctx], 2);
    307  if (eob == 0) return;
    308 
    309  const TX_TYPE tx_type =
    310      av1_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,
    311                      cm->features.reduced_tx_set_used);
    312  // Only y plane's tx_type is transmitted
    313  if (plane == 0) {
    314    av1_write_tx_type(cm, xd, tx_type, tx_size, w);
    315  }
    316 
    317  int eob_extra;
    318  const int eob_pt = av1_get_eob_pos_token(eob, &eob_extra);
    319  const int eob_multi_size = txsize_log2_minus4[tx_size];
    320  const TX_CLASS tx_class = tx_type_to_class[tx_type];
    321  const int eob_multi_ctx = (tx_class == TX_CLASS_2D) ? 0 : 1;
    322  switch (eob_multi_size) {
    323    case 0:
    324      aom_write_symbol(w, eob_pt - 1,
    325                       ec_ctx->eob_flag_cdf16[plane_type][eob_multi_ctx], 5);
    326      break;
    327    case 1:
    328      aom_write_symbol(w, eob_pt - 1,
    329                       ec_ctx->eob_flag_cdf32[plane_type][eob_multi_ctx], 6);
    330      break;
    331    case 2:
    332      aom_write_symbol(w, eob_pt - 1,
    333                       ec_ctx->eob_flag_cdf64[plane_type][eob_multi_ctx], 7);
    334      break;
    335    case 3:
    336      aom_write_symbol(w, eob_pt - 1,
    337                       ec_ctx->eob_flag_cdf128[plane_type][eob_multi_ctx], 8);
    338      break;
    339    case 4:
    340      aom_write_symbol(w, eob_pt - 1,
    341                       ec_ctx->eob_flag_cdf256[plane_type][eob_multi_ctx], 9);
    342      break;
    343    case 5:
    344      aom_write_symbol(w, eob_pt - 1,
    345                       ec_ctx->eob_flag_cdf512[plane_type][eob_multi_ctx], 10);
    346      break;
    347    default:
    348      aom_write_symbol(w, eob_pt - 1,
    349                       ec_ctx->eob_flag_cdf1024[plane_type][eob_multi_ctx], 11);
    350      break;
    351  }
    352 
    353  const int eob_offset_bits = av1_eob_offset_bits[eob_pt];
    354  if (eob_offset_bits > 0) {
    355    const int eob_ctx = eob_pt - 3;
    356    int eob_shift = eob_offset_bits - 1;
    357    int bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
    358    aom_write_symbol(w, bit,
    359                     ec_ctx->eob_extra_cdf[txs_ctx][plane_type][eob_ctx], 2);
    360    for (int i = 1; i < eob_offset_bits; i++) {
    361      eob_shift = eob_offset_bits - 1 - i;
    362      bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
    363      aom_write_bit(w, bit);
    364    }
    365  }
    366 
    367  const int width = get_txb_wide(tx_size);
    368  const int height = get_txb_high(tx_size);
    369  uint8_t levels_buf[TX_PAD_2D];
    370  uint8_t *const levels = set_levels(levels_buf, height);
    371  const tran_low_t *tcoeff_txb =
    372      cb_coef_buff->tcoeff[plane] + x->mbmi_ext_frame->cb_offset[plane_type];
    373  const tran_low_t *tcoeff = tcoeff_txb + BLOCK_OFFSET(block);
    374  av1_txb_init_levels(tcoeff, width, height, levels);
    375  const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);
    376  const int16_t *const scan = scan_order->scan;
    377  DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
    378  av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_class, coeff_contexts);
    379 
    380  const int bhl = get_txb_bhl(tx_size);
    381  for (int c = eob - 1; c >= 0; --c) {
    382    const int pos = scan[c];
    383    const int coeff_ctx = coeff_contexts[pos];
    384    const tran_low_t v = tcoeff[pos];
    385    const tran_low_t level = abs(v);
    386 
    387    if (c == eob - 1) {
    388      aom_write_symbol(
    389          w, AOMMIN(level, 3) - 1,
    390          ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx], 3);
    391    } else {
    392      aom_write_symbol(w, AOMMIN(level, 3),
    393                       ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx],
    394                       4);
    395    }
    396    if (level > NUM_BASE_LEVELS) {
    397      // level is above 1.
    398      const int base_range = level - 1 - NUM_BASE_LEVELS;
    399      const int br_ctx = get_br_ctx(levels, pos, bhl, tx_class);
    400      aom_cdf_prob *cdf =
    401          ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)][plane_type][br_ctx];
    402      for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
    403        const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
    404        aom_write_symbol(w, k, cdf, BR_CDF_SIZE);
    405        if (k < BR_CDF_SIZE - 1) break;
    406      }
    407    }
    408  }
    409 
    410  // Loop to code all signs in the transform block,
    411  // starting with the sign of DC (if applicable)
    412  for (int c = 0; c < eob; ++c) {
    413    const tran_low_t v = tcoeff[scan[c]];
    414    const tran_low_t level = abs(v);
    415    const int sign = (v < 0) ? 1 : 0;
    416    if (level) {
    417      if (c == 0) {
    418        const int dc_sign_ctx =
    419            (entropy_ctx[block] >> DC_SIGN_CTX_SHIFT) & DC_SIGN_CTX_MASK;
    420        aom_write_symbol(w, sign, ec_ctx->dc_sign_cdf[plane_type][dc_sign_ctx],
    421                         2);
    422      } else {
    423        aom_write_bit(w, sign);
    424      }
    425      if (level > COEFF_BASE_RANGE + NUM_BASE_LEVELS)
    426        write_golomb(w, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
    427    }
    428  }
    429 }
    430 
    431 void av1_write_intra_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
    432                               aom_writer *w, BLOCK_SIZE bsize) {
    433  MACROBLOCKD *xd = &x->e_mbd;
    434  const int num_planes = av1_num_planes(cm);
    435  int block[MAX_MB_PLANE] = { 0 };
    436  int row, col;
    437  assert(bsize == get_plane_block_size(bsize, xd->plane[0].subsampling_x,
    438                                       xd->plane[0].subsampling_y));
    439  const int max_blocks_wide = max_block_wide(xd, bsize, 0);
    440  const int max_blocks_high = max_block_high(xd, bsize, 0);
    441  const BLOCK_SIZE max_unit_bsize = BLOCK_64X64;
    442  int mu_blocks_wide = mi_size_wide[max_unit_bsize];
    443  int mu_blocks_high = mi_size_high[max_unit_bsize];
    444  mu_blocks_wide = AOMMIN(max_blocks_wide, mu_blocks_wide);
    445  mu_blocks_high = AOMMIN(max_blocks_high, mu_blocks_high);
    446 
    447  for (row = 0; row < max_blocks_high; row += mu_blocks_high) {
    448    for (col = 0; col < max_blocks_wide; col += mu_blocks_wide) {
    449      for (int plane = 0; plane < num_planes; ++plane) {
    450        if (plane && !xd->is_chroma_ref) break;
    451        const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
    452        const int stepr = tx_size_high_unit[tx_size];
    453        const int stepc = tx_size_wide_unit[tx_size];
    454        const int step = stepr * stepc;
    455        const struct macroblockd_plane *const pd = &xd->plane[plane];
    456        const int unit_height = ROUND_POWER_OF_TWO(
    457            AOMMIN(mu_blocks_high + row, max_blocks_high), pd->subsampling_y);
    458        const int unit_width = ROUND_POWER_OF_TWO(
    459            AOMMIN(mu_blocks_wide + col, max_blocks_wide), pd->subsampling_x);
    460        for (int blk_row = row >> pd->subsampling_y; blk_row < unit_height;
    461             blk_row += stepr) {
    462          for (int blk_col = col >> pd->subsampling_x; blk_col < unit_width;
    463               blk_col += stepc) {
    464            av1_write_coeffs_txb(cm, x, w, blk_row, blk_col, plane,
    465                                 block[plane], tx_size);
    466            block[plane] += step;
    467          }
    468        }
    469      }
    470    }
    471  }
    472 }
    473 
    474 uint8_t av1_get_txb_entropy_context(const tran_low_t *qcoeff,
    475                                    const SCAN_ORDER *scan_order, int eob) {
    476  const int16_t *const scan = scan_order->scan;
    477  int cul_level = 0;
    478  int c;
    479 
    480  if (eob == 0) return 0;
    481  for (c = 0; c < eob; ++c) {
    482    cul_level += abs(qcoeff[scan[c]]);
    483    if (cul_level > COEFF_CONTEXT_MASK) break;
    484  }
    485 
    486  cul_level = AOMMIN(COEFF_CONTEXT_MASK, cul_level);
    487  set_dc_sign(&cul_level, qcoeff[0]);
    488 
    489  return (uint8_t)cul_level;
    490 }
    491 
    492 static void update_tx_type_count(const AV1_COMP *cpi, const AV1_COMMON *cm,
    493                                 MACROBLOCKD *xd, int blk_row, int blk_col,
    494                                 int plane, TX_SIZE tx_size,
    495                                 FRAME_COUNTS *counts,
    496                                 uint8_t allow_update_cdf) {
    497  MB_MODE_INFO *mbmi = xd->mi[0];
    498  int is_inter = is_inter_block(mbmi);
    499  const int reduced_tx_set_used = cm->features.reduced_tx_set_used;
    500  FRAME_CONTEXT *fc = xd->tile_ctx;
    501 #if !CONFIG_ENTROPY_STATS
    502  (void)counts;
    503 #endif  // !CONFIG_ENTROPY_STATS
    504 
    505  // Only y plane's tx_type is updated
    506  if (plane > 0) return;
    507  const TX_TYPE tx_type = av1_get_tx_type(xd, PLANE_TYPE_Y, blk_row, blk_col,
    508                                          tx_size, reduced_tx_set_used);
    509  if (is_inter) {
    510    if (cpi->oxcf.txfm_cfg.use_inter_dct_only) {
    511      assert(tx_type == DCT_DCT);
    512    }
    513  } else {
    514    if (cpi->oxcf.txfm_cfg.use_intra_dct_only) {
    515      assert(tx_type == DCT_DCT);
    516    } else if (cpi->oxcf.txfm_cfg.use_intra_default_tx_only) {
    517      const TX_TYPE default_type = get_default_tx_type(
    518          PLANE_TYPE_Y, xd, tx_size, cpi->use_screen_content_tools);
    519      (void)default_type;
    520      // TODO(kyslov): We don't always respect use_intra_default_tx_only flag in
    521      // NonRD and REALTIME case. Specifically we ignore it in hybrid inta mode
    522      // search, when picking up intra mode in nonRD inter mode search and in RD
    523      // REALTIME mode when we limit TX type usage.
    524      // We need to fix txfm cfg for these cases. Meanwhile relieving the
    525      // assert.
    526      assert(tx_type == default_type || cpi->sf.rt_sf.use_nonrd_pick_mode ||
    527             cpi->oxcf.mode == REALTIME);
    528    }
    529  }
    530 
    531  if (get_ext_tx_types(tx_size, is_inter, reduced_tx_set_used) > 1 &&
    532      cm->quant_params.base_qindex > 0 && !mbmi->skip_txfm &&
    533      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) {
    534    const int eset = get_ext_tx_set(tx_size, is_inter, reduced_tx_set_used);
    535    if (eset > 0) {
    536      const TxSetType tx_set_type =
    537          av1_get_ext_tx_set_type(tx_size, is_inter, reduced_tx_set_used);
    538      if (is_inter) {
    539        if (allow_update_cdf) {
    540          update_cdf(fc->inter_ext_tx_cdf[eset][txsize_sqr_map[tx_size]],
    541                     av1_ext_tx_ind[tx_set_type][tx_type],
    542                     av1_num_ext_tx_set[tx_set_type]);
    543        }
    544 #if CONFIG_ENTROPY_STATS
    545        ++counts->inter_ext_tx[eset][txsize_sqr_map[tx_size]]
    546                              [av1_ext_tx_ind[tx_set_type][tx_type]];
    547 #endif  // CONFIG_ENTROPY_STATS
    548      } else {
    549        PREDICTION_MODE intra_dir;
    550        if (mbmi->filter_intra_mode_info.use_filter_intra)
    551          intra_dir = fimode_to_intradir[mbmi->filter_intra_mode_info
    552                                             .filter_intra_mode];
    553        else
    554          intra_dir = mbmi->mode;
    555 #if CONFIG_ENTROPY_STATS
    556        ++counts->intra_ext_tx[eset][txsize_sqr_map[tx_size]][intra_dir]
    557                              [av1_ext_tx_ind[tx_set_type][tx_type]];
    558 #endif  // CONFIG_ENTROPY_STATS
    559        if (allow_update_cdf) {
    560          update_cdf(
    561              fc->intra_ext_tx_cdf[eset][txsize_sqr_map[tx_size]][intra_dir],
    562              av1_ext_tx_ind[tx_set_type][tx_type],
    563              av1_num_ext_tx_set[tx_set_type]);
    564        }
    565      }
    566    }
    567  }
    568 }
    569 
    570 void av1_update_and_record_txb_context(int plane, int block, int blk_row,
    571                                       int blk_col, BLOCK_SIZE plane_bsize,
    572                                       TX_SIZE tx_size, void *arg) {
    573  struct tokenize_b_args *const args = arg;
    574  const AV1_COMP *cpi = args->cpi;
    575  const AV1_COMMON *cm = &cpi->common;
    576  ThreadData *const td = args->td;
    577  MACROBLOCK *const x = &td->mb;
    578  MACROBLOCKD *const xd = &x->e_mbd;
    579  struct macroblock_plane *p = &x->plane[plane];
    580  struct macroblockd_plane *pd = &xd->plane[plane];
    581  const int eob = p->eobs[block];
    582  const int block_offset = BLOCK_OFFSET(block);
    583  tran_low_t *qcoeff = p->qcoeff + block_offset;
    584  const PLANE_TYPE plane_type = pd->plane_type;
    585  const TX_TYPE tx_type =
    586      av1_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,
    587                      cm->features.reduced_tx_set_used);
    588  const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);
    589  tran_low_t *tcoeff;
    590  assert(args->dry_run != DRY_RUN_COSTCOEFFS);
    591  if (args->dry_run == OUTPUT_ENABLED) {
    592    MB_MODE_INFO *mbmi = xd->mi[0];
    593    TXB_CTX txb_ctx;
    594    get_txb_ctx(plane_bsize, tx_size, plane,
    595                pd->above_entropy_context + blk_col,
    596                pd->left_entropy_context + blk_row, &txb_ctx);
    597    const int bhl = get_txb_bhl(tx_size);
    598    const int width = get_txb_wide(tx_size);
    599    const int height = get_txb_high(tx_size);
    600    const uint8_t allow_update_cdf = args->allow_update_cdf;
    601    const TX_SIZE txsize_ctx = get_txsize_entropy_ctx(tx_size);
    602    FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
    603 #if CONFIG_ENTROPY_STATS
    604    int cdf_idx = cm->coef_cdf_category;
    605    ++td->counts->txb_skip[cdf_idx][txsize_ctx][txb_ctx.txb_skip_ctx][eob == 0];
    606 #endif  // CONFIG_ENTROPY_STATS
    607    if (allow_update_cdf) {
    608      update_cdf(ec_ctx->txb_skip_cdf[txsize_ctx][txb_ctx.txb_skip_ctx],
    609                 eob == 0, 2);
    610    }
    611 
    612    CB_COEFF_BUFFER *cb_coef_buff = x->cb_coef_buff;
    613    const int txb_offset = x->mbmi_ext_frame->cb_offset[plane_type] /
    614                           (TX_SIZE_W_MIN * TX_SIZE_H_MIN);
    615    uint16_t *eob_txb = cb_coef_buff->eobs[plane] + txb_offset;
    616    uint8_t *const entropy_ctx = cb_coef_buff->entropy_ctx[plane] + txb_offset;
    617    entropy_ctx[block] = txb_ctx.txb_skip_ctx;
    618    eob_txb[block] = eob;
    619 
    620    if (eob == 0) {
    621      av1_set_entropy_contexts(xd, pd, plane, plane_bsize, tx_size, 0, blk_col,
    622                               blk_row);
    623      return;
    624    }
    625    const int segment_id = mbmi->segment_id;
    626    const int seg_eob = av1_get_tx_eob(&cpi->common.seg, segment_id, tx_size);
    627    tran_low_t *tcoeff_txb =
    628        cb_coef_buff->tcoeff[plane] + x->mbmi_ext_frame->cb_offset[plane_type];
    629    tcoeff = tcoeff_txb + block_offset;
    630    memcpy(tcoeff, qcoeff, sizeof(*tcoeff) * seg_eob);
    631 
    632    uint8_t levels_buf[TX_PAD_2D];
    633    uint8_t *const levels = set_levels(levels_buf, height);
    634    av1_txb_init_levels(tcoeff, width, height, levels);
    635    update_tx_type_count(cpi, cm, xd, blk_row, blk_col, plane, tx_size,
    636                         td->counts, allow_update_cdf);
    637 
    638    const TX_CLASS tx_class = tx_type_to_class[tx_type];
    639    const int16_t *const scan = scan_order->scan;
    640 
    641    // record tx type usage
    642    td->rd_counts.tx_type_used[tx_size][tx_type]++;
    643 
    644 #if CONFIG_ENTROPY_STATS
    645    update_eob_context(cdf_idx, eob, tx_size, tx_class, plane_type, ec_ctx,
    646                       td->counts, allow_update_cdf);
    647 #else
    648    update_eob_context(eob, tx_size, tx_class, plane_type, ec_ctx,
    649                       allow_update_cdf);
    650 #endif
    651 
    652    DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
    653    av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_class,
    654                            coeff_contexts);
    655 
    656    for (int c = eob - 1; c >= 0; --c) {
    657      const int pos = scan[c];
    658      const int coeff_ctx = coeff_contexts[pos];
    659      const tran_low_t v = qcoeff[pos];
    660      const tran_low_t level = abs(v);
    661      /* abs_sum_level is needed to decide the job scheduling order of
    662       * pack bitstream multi-threading. This data is not needed if
    663       * multi-threading is disabled. */
    664      if (cpi->mt_info.pack_bs_mt_enabled) td->abs_sum_level += level;
    665 
    666      if (allow_update_cdf) {
    667        if (c == eob - 1) {
    668          assert(coeff_ctx < 4);
    669          update_cdf(
    670              ec_ctx->coeff_base_eob_cdf[txsize_ctx][plane_type][coeff_ctx],
    671              AOMMIN(level, 3) - 1, 3);
    672        } else {
    673          update_cdf(ec_ctx->coeff_base_cdf[txsize_ctx][plane_type][coeff_ctx],
    674                     AOMMIN(level, 3), 4);
    675        }
    676      }
    677      if (c == eob - 1) {
    678        assert(coeff_ctx < 4);
    679 #if CONFIG_ENTROPY_STATS
    680        ++td->counts->coeff_base_eob_multi[cdf_idx][txsize_ctx][plane_type]
    681                                          [coeff_ctx][AOMMIN(level, 3) - 1];
    682      } else {
    683        ++td->counts->coeff_base_multi[cdf_idx][txsize_ctx][plane_type]
    684                                      [coeff_ctx][AOMMIN(level, 3)];
    685 #endif
    686      }
    687      if (level > NUM_BASE_LEVELS) {
    688        const int base_range = level - 1 - NUM_BASE_LEVELS;
    689        const int br_ctx = get_br_ctx(levels, pos, bhl, tx_class);
    690        for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
    691          const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
    692          if (allow_update_cdf) {
    693            update_cdf(ec_ctx->coeff_br_cdf[AOMMIN(txsize_ctx, TX_32X32)]
    694                                           [plane_type][br_ctx],
    695                       k, BR_CDF_SIZE);
    696          }
    697          for (int lps = 0; lps < BR_CDF_SIZE - 1; lps++) {
    698 #if CONFIG_ENTROPY_STATS
    699            ++td->counts->coeff_lps[AOMMIN(txsize_ctx, TX_32X32)][plane_type]
    700                                   [lps][br_ctx][lps == k];
    701 #endif  // CONFIG_ENTROPY_STATS
    702            if (lps == k) break;
    703          }
    704 #if CONFIG_ENTROPY_STATS
    705          ++td->counts->coeff_lps_multi[cdf_idx][AOMMIN(txsize_ctx, TX_32X32)]
    706                                       [plane_type][br_ctx][k];
    707 #endif
    708          if (k < BR_CDF_SIZE - 1) break;
    709        }
    710      }
    711    }
    712    // Update the context needed to code the DC sign (if applicable)
    713    if (tcoeff[0] != 0) {
    714      const int dc_sign = (tcoeff[0] < 0) ? 1 : 0;
    715      const int dc_sign_ctx = txb_ctx.dc_sign_ctx;
    716 #if CONFIG_ENTROPY_STATS
    717      ++td->counts->dc_sign[plane_type][dc_sign_ctx][dc_sign];
    718 #endif  // CONFIG_ENTROPY_STATS
    719      if (allow_update_cdf)
    720        update_cdf(ec_ctx->dc_sign_cdf[plane_type][dc_sign_ctx], dc_sign, 2);
    721      entropy_ctx[block] |= dc_sign_ctx << DC_SIGN_CTX_SHIFT;
    722    }
    723  } else {
    724    tcoeff = qcoeff;
    725  }
    726  const uint8_t cul_level =
    727      av1_get_txb_entropy_context(tcoeff, scan_order, eob);
    728  av1_set_entropy_contexts(xd, pd, plane, plane_bsize, tx_size, cul_level,
    729                           blk_col, blk_row);
    730 }
    731 
    732 void av1_record_txb_context(int plane, int block, int blk_row, int blk_col,
    733                            BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
    734                            void *arg) {
    735  struct tokenize_b_args *const args = arg;
    736  const AV1_COMP *cpi = args->cpi;
    737  const AV1_COMMON *cm = &cpi->common;
    738  ThreadData *const td = args->td;
    739  MACROBLOCK *const x = &td->mb;
    740  MACROBLOCKD *const xd = &x->e_mbd;
    741  struct macroblock_plane *p = &x->plane[plane];
    742  struct macroblockd_plane *pd = &xd->plane[plane];
    743  const int eob = p->eobs[block];
    744  const int block_offset = BLOCK_OFFSET(block);
    745  tran_low_t *qcoeff = p->qcoeff + block_offset;
    746  const PLANE_TYPE plane_type = pd->plane_type;
    747  const TX_TYPE tx_type =
    748      av1_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,
    749                      cm->features.reduced_tx_set_used);
    750  const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);
    751  tran_low_t *tcoeff;
    752  assert(args->dry_run != DRY_RUN_COSTCOEFFS);
    753  if (args->dry_run == OUTPUT_ENABLED) {
    754    MB_MODE_INFO *mbmi = xd->mi[0];
    755    TXB_CTX txb_ctx;
    756    get_txb_ctx(plane_bsize, tx_size, plane,
    757                pd->above_entropy_context + blk_col,
    758                pd->left_entropy_context + blk_row, &txb_ctx);
    759 #if CONFIG_ENTROPY_STATS
    760    const TX_SIZE txsize_ctx = get_txsize_entropy_ctx(tx_size);
    761    const int bhl = get_txb_bhl(tx_size);
    762    const int width = get_txb_wide(tx_size);
    763    const int height = get_txb_high(tx_size);
    764    int cdf_idx = cm->coef_cdf_category;
    765    ++td->counts->txb_skip[cdf_idx][txsize_ctx][txb_ctx.txb_skip_ctx][eob == 0];
    766 #endif  // CONFIG_ENTROPY_STATS
    767 
    768    CB_COEFF_BUFFER *cb_coef_buff = x->cb_coef_buff;
    769    const int txb_offset = x->mbmi_ext_frame->cb_offset[plane_type] /
    770                           (TX_SIZE_W_MIN * TX_SIZE_H_MIN);
    771    uint16_t *eob_txb = cb_coef_buff->eobs[plane] + txb_offset;
    772    uint8_t *const entropy_ctx = cb_coef_buff->entropy_ctx[plane] + txb_offset;
    773    entropy_ctx[block] = txb_ctx.txb_skip_ctx;
    774    eob_txb[block] = eob;
    775 
    776    if (eob == 0) {
    777      av1_set_entropy_contexts(xd, pd, plane, plane_bsize, tx_size, 0, blk_col,
    778                               blk_row);
    779      return;
    780    }
    781    const int segment_id = mbmi->segment_id;
    782    const int seg_eob = av1_get_tx_eob(&cpi->common.seg, segment_id, tx_size);
    783    tran_low_t *tcoeff_txb =
    784        cb_coef_buff->tcoeff[plane] + x->mbmi_ext_frame->cb_offset[plane_type];
    785    tcoeff = tcoeff_txb + block_offset;
    786    memcpy(tcoeff, qcoeff, sizeof(*tcoeff) * seg_eob);
    787 
    788 #if CONFIG_ENTROPY_STATS
    789    uint8_t levels_buf[TX_PAD_2D];
    790    uint8_t *const levels = set_levels(levels_buf, height);
    791    av1_txb_init_levels(tcoeff, width, height, levels);
    792    update_tx_type_count(cpi, cm, xd, blk_row, blk_col, plane, tx_size,
    793                         td->counts, 0 /*allow_update_cdf*/);
    794 
    795    const TX_CLASS tx_class = tx_type_to_class[tx_type];
    796    const bool do_coeff_scan = true;
    797 #else
    798    const bool do_coeff_scan = cpi->mt_info.pack_bs_mt_enabled;
    799 #endif
    800    const int16_t *const scan = scan_order->scan;
    801 
    802    // record tx type usage
    803    td->rd_counts.tx_type_used[tx_size][tx_type]++;
    804 
    805 #if CONFIG_ENTROPY_STATS
    806    FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
    807    update_eob_context(cdf_idx, eob, tx_size, tx_class, plane_type, ec_ctx,
    808                       td->counts, 0 /*allow_update_cdf*/);
    809 
    810    DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
    811    av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_class,
    812                            coeff_contexts);
    813 #endif
    814 
    815    for (int c = eob - 1; (c >= 0) && do_coeff_scan; --c) {
    816      const int pos = scan[c];
    817      const tran_low_t v = qcoeff[pos];
    818      const tran_low_t level = abs(v);
    819      /* abs_sum_level is needed to decide the job scheduling order of
    820       * pack bitstream multi-threading. This data is not needed if
    821       * multi-threading is disabled. */
    822      if (cpi->mt_info.pack_bs_mt_enabled) td->abs_sum_level += level;
    823 
    824 #if CONFIG_ENTROPY_STATS
    825      const int coeff_ctx = coeff_contexts[pos];
    826      if (c == eob - 1) {
    827        assert(coeff_ctx < 4);
    828        ++td->counts->coeff_base_eob_multi[cdf_idx][txsize_ctx][plane_type]
    829                                          [coeff_ctx][AOMMIN(level, 3) - 1];
    830      } else {
    831        ++td->counts->coeff_base_multi[cdf_idx][txsize_ctx][plane_type]
    832                                      [coeff_ctx][AOMMIN(level, 3)];
    833      }
    834      if (level > NUM_BASE_LEVELS) {
    835        const int base_range = level - 1 - NUM_BASE_LEVELS;
    836        const int br_ctx = get_br_ctx(levels, pos, bhl, tx_class);
    837        for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
    838          const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
    839          for (int lps = 0; lps < BR_CDF_SIZE - 1; lps++) {
    840            ++td->counts->coeff_lps[AOMMIN(txsize_ctx, TX_32X32)][plane_type]
    841                                   [lps][br_ctx][lps == k];
    842            if (lps == k) break;
    843          }
    844          ++td->counts->coeff_lps_multi[cdf_idx][AOMMIN(txsize_ctx, TX_32X32)]
    845                                       [plane_type][br_ctx][k];
    846          if (k < BR_CDF_SIZE - 1) break;
    847        }
    848      }
    849 #endif
    850    }
    851    // Update the context needed to code the DC sign (if applicable)
    852    if (tcoeff[0] != 0) {
    853      const int dc_sign_ctx = txb_ctx.dc_sign_ctx;
    854 #if CONFIG_ENTROPY_STATS
    855      const int dc_sign = (tcoeff[0] < 0) ? 1 : 0;
    856      ++td->counts->dc_sign[plane_type][dc_sign_ctx][dc_sign];
    857 #endif  // CONFIG_ENTROPY_STATS
    858      entropy_ctx[block] |= dc_sign_ctx << DC_SIGN_CTX_SHIFT;
    859    }
    860  } else {
    861    tcoeff = qcoeff;
    862  }
    863  const uint8_t cul_level =
    864      av1_get_txb_entropy_context(tcoeff, scan_order, eob);
    865  av1_set_entropy_contexts(xd, pd, plane, plane_bsize, tx_size, cul_level,
    866                           blk_col, blk_row);
    867 }
    868 
    869 void av1_update_intra_mb_txb_context(const AV1_COMP *cpi, ThreadData *td,
    870                                     RUN_TYPE dry_run, BLOCK_SIZE bsize,
    871                                     uint8_t allow_update_cdf) {
    872  const AV1_COMMON *const cm = &cpi->common;
    873  const int num_planes = av1_num_planes(cm);
    874  MACROBLOCK *const x = &td->mb;
    875  MACROBLOCKD *const xd = &x->e_mbd;
    876  MB_MODE_INFO *const mbmi = xd->mi[0];
    877  struct tokenize_b_args arg = { cpi, td, 0, allow_update_cdf, dry_run };
    878  if (mbmi->skip_txfm) {
    879    av1_reset_entropy_context(xd, bsize, num_planes);
    880    return;
    881  }
    882  const foreach_transformed_block_visitor visit =
    883      allow_update_cdf ? av1_update_and_record_txb_context
    884                       : av1_record_txb_context;
    885 
    886  for (int plane = 0; plane < num_planes; ++plane) {
    887    if (plane && !xd->is_chroma_ref) break;
    888    const struct macroblockd_plane *const pd = &xd->plane[plane];
    889    const int ss_x = pd->subsampling_x;
    890    const int ss_y = pd->subsampling_y;
    891    const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, ss_x, ss_y);
    892    av1_foreach_transformed_block_in_plane(xd, plane_bsize, plane, visit, &arg);
    893  }
    894 }
    895 
    896 CB_COEFF_BUFFER *av1_get_cb_coeff_buffer(const struct AV1_COMP *cpi, int mi_row,
    897                                         int mi_col) {
    898  const AV1_COMMON *const cm = &cpi->common;
    899  const int mib_size_log2 = cm->seq_params->mib_size_log2;
    900  const int stride =
    901      CEIL_POWER_OF_TWO(cm->mi_params.mi_cols, cm->seq_params->mib_size_log2);
    902  const int offset =
    903      (mi_row >> mib_size_log2) * stride + (mi_col >> mib_size_log2);
    904  return cpi->coeff_buffer_base + offset;
    905 }