encodetxb.c 99.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
/*
 * Copyright (c) 2017, 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.
 */

12
#include "aom_ports/mem.h"
13
#include "av1/common/scan.h"
Angie Chiang's avatar
Angie Chiang committed
14
#include "av1/common/blockd.h"
15
#include "av1/common/idct.h"
Angie Chiang's avatar
Angie Chiang committed
16
#include "av1/common/pred_common.h"
17 18
#include "av1/encoder/bitstream.h"
#include "av1/encoder/encodeframe.h"
19
#include "av1/encoder/cost.h"
20
#include "av1/encoder/encodetxb.h"
21
#include "av1/encoder/hash.h"
Angie Chiang's avatar
Angie Chiang committed
22
#include "av1/encoder/rdopt.h"
Angie Chiang's avatar
Angie Chiang committed
23
#include "av1/encoder/tokenize.h"
24

Angie Chiang's avatar
Angie Chiang committed
25 26
#define TEST_OPTIMIZE_TXB 0

27 28 29 30 31 32 33 34 35 36 37 38 39
static int hbt_hash_needs_init = 1;
static CRC_CALCULATOR crc_calculator;
static CRC_CALCULATOR crc_calculator2;
static const int HBT_HASH_EOB = 16;  // also the length in opt_qcoeff

typedef struct OptTxbQcoeff {
  uint32_t hbt_hash_match;
  double hits;
  tran_low_t opt_qcoeff[16];
} OptTxbQcoeff;

OptTxbQcoeff hbt_hash_table[65536][16];

Dake He's avatar
Dake He committed
40 41 42 43 44 45 46 47 48 49 50
typedef struct LevelDownStats {
  int update;
  tran_low_t low_qc;
  tran_low_t low_dqc;
  int64_t dist0;
  int rate;
  int rate_low;
  int64_t dist;
  int64_t dist_low;
  int64_t rd;
  int64_t rd_low;
Dake He's avatar
Dake He committed
51 52 53
#if CONFIG_LV_MAP_MULTI
  int64_t nz_rd;
#else
Dake He's avatar
Dake He committed
54
  int nz_rate;  // for eob
Dake He's avatar
Dake He committed
55
#endif
Dake He's avatar
Dake He committed
56 57 58 59 60 61
  int64_t rd_diff;
  int cost_diff;
  int64_t dist_diff;
  int new_eob;
} LevelDownStats;

62
void av1_alloc_txb_buf(AV1_COMP *cpi) {
63
#if 0
64 65 66 67 68 69 70 71 72 73 74 75
  AV1_COMMON *cm = &cpi->common;
  int mi_block_size = 1 << MI_SIZE_LOG2;
  // TODO(angiebird): Make sure cm->subsampling_x/y is set correctly, and then
  // use precise buffer size according to cm->subsampling_x/y
  int pixel_stride = mi_block_size * cm->mi_cols;
  int pixel_height = mi_block_size * cm->mi_rows;
  int i;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    CHECK_MEM_ERROR(
        cm, cpi->tcoeff_buf[i],
        aom_malloc(sizeof(*cpi->tcoeff_buf[i]) * pixel_stride * pixel_height));
  }
76
#else
77
  AV1_COMMON *cm = &cpi->common;
78 79
  int size = ((cm->mi_rows >> cm->mib_size_log2) + 1) *
             ((cm->mi_cols >> cm->mib_size_log2) + 1);
80

81
  av1_free_txb_buf(cpi);
82 83 84
  // TODO(jingning): This should be further reduced.
  CHECK_MEM_ERROR(cm, cpi->coeff_buffer_base,
                  aom_malloc(sizeof(*cpi->coeff_buffer_base) * size));
85
#endif
86 87 88
}

void av1_free_txb_buf(AV1_COMP *cpi) {
89
#if 0
90 91 92 93
  int i;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    aom_free(cpi->tcoeff_buf[i]);
  }
94
#else
95
  aom_free(cpi->coeff_buffer_base);
96
#endif
97 98
}

99 100
void av1_set_coeff_buffer(const AV1_COMP *const cpi, MACROBLOCK *const x,
                          int mi_row, int mi_col) {
101 102 103
  int mib_size_log2 = cpi->common.mib_size_log2;
  int stride = (cpi->common.mi_cols >> mib_size_log2) + 1;
  int offset = (mi_row >> mib_size_log2) * stride + (mi_col >> mib_size_log2);
104 105 106 107 108 109 110 111 112 113 114 115
  CB_COEFF_BUFFER *coeff_buf = &cpi->coeff_buffer_base[offset];
  const int txb_offset = x->cb_offset / (TX_SIZE_W_MIN * TX_SIZE_H_MIN);
  for (int plane = 0; plane < MAX_MB_PLANE; ++plane) {
    x->mbmi_ext->tcoeff[plane] = coeff_buf->tcoeff[plane] + x->cb_offset;
    x->mbmi_ext->eobs[plane] = coeff_buf->eobs[plane] + txb_offset;
    x->mbmi_ext->txb_skip_ctx[plane] =
        coeff_buf->txb_skip_ctx[plane] + txb_offset;
    x->mbmi_ext->dc_sign_ctx[plane] =
        coeff_buf->dc_sign_ctx[plane] + txb_offset;
  }
}

116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131
static void write_golomb(aom_writer *w, int level) {
  int x = level + 1;
  int i = x;
  int length = 0;

  while (i) {
    i >>= 1;
    ++length;
  }
  assert(length > 0);

  for (i = 0; i < length - 1; ++i) aom_write_bit(w, 0);

  for (i = length - 1; i >= 0; --i) aom_write_bit(w, (x >> i) & 0x01);
}

Dake He's avatar
Dake He committed
132 133 134 135 136 137 138
static INLINE tran_low_t get_lower_coeff(tran_low_t qc) {
  if (qc == 0) {
    return 0;
  }
  return qc > 0 ? qc - 1 : qc + 1;
}

139 140 141 142 143
static INLINE tran_low_t qcoeff_to_dqcoeff(tran_low_t qc,
#if CONFIG_NEW_QUANT
                                           const tran_low_t *nq_dq,
#endif  // CONFIG_NEW_QUANT
                                           int dqv, int shift) {
Dake He's avatar
Dake He committed
144
  int sgn = qc < 0 ? -1 : 1;
145
#if CONFIG_NEW_QUANT
146 147
  int dqcoeff = av1_dequant_coeff_nuq(abs(qc), dqv, nq_dq, shift);
  return sgn * dqcoeff;
148 149
#endif  // CONFIG_NEW_QUANT

Dake He's avatar
Dake He committed
150 151 152 153 154
  return sgn * ((abs(qc) * dqv) >> shift);
}

static INLINE int64_t get_coeff_dist(tran_low_t tcoeff, tran_low_t dqcoeff,
                                     int shift) {
155 156 157 158
#if CONFIG_DAALA_TX
  int depth_shift = (TX_COEFF_DEPTH - 11) * 2;
  int depth_round = depth_shift > 1 ? (1 << (depth_shift - 1)) : 0;
  const int64_t diff = tcoeff - dqcoeff;
159
  const int64_t error = (diff * diff + depth_round) >> depth_shift;
160 161
  (void)shift;
#else
Dake He's avatar
Dake He committed
162 163
  const int64_t diff = (tcoeff - dqcoeff) * (1 << shift);
  const int64_t error = diff * diff;
164
#endif
Dake He's avatar
Dake He committed
165 166 167
  return error;
}

Jingning Han's avatar
Jingning Han committed
168 169
void av1_update_eob_context(int eob, int seg_eob, TX_SIZE tx_size,
                            TX_TYPE tx_type, PLANE_TYPE plane,
170 171
                            FRAME_CONTEXT *ec_ctx, FRAME_COUNTS *counts,
                            uint8_t allow_update_cdf) {
Linfeng Zhang's avatar
Linfeng Zhang committed
172 173 174
  int eob_extra, dummy;
  const int eob_pt = get_eob_pos_token(eob, &eob_extra);
  const int max_eob_pt = get_eob_pos_token(seg_eob, &dummy);
175
  TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
Dake He's avatar
Dake He committed
176

177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224
#if CONFIG_LV_MAP_MULTI
  (void)max_eob_pt;
  const int eob_multi_size = txsize_log2_minus4[tx_size];
  const int eob_multi_ctx = (tx_type_to_class[tx_type] == TX_CLASS_2D) ? 0 : 1;

  switch (eob_multi_size) {
    case 0:
      ++counts->eob_multi16[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf16[plane][eob_multi_ctx], eob_pt - 1, 5);
      break;
    case 1:
      ++counts->eob_multi32[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf32[plane][eob_multi_ctx], eob_pt - 1, 6);
      break;
    case 2:
      ++counts->eob_multi64[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf64[plane][eob_multi_ctx], eob_pt - 1, 7);
      break;
    case 3:
      ++counts->eob_multi128[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf128[plane][eob_multi_ctx], eob_pt - 1,
                   8);
      break;
    case 4:
      ++counts->eob_multi256[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf256[plane][eob_multi_ctx], eob_pt - 1,
                   9);
      break;
    case 5:
      ++counts->eob_multi512[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf512[plane][eob_multi_ctx], eob_pt - 1,
                   10);
      break;
    case 6:
    default:
      ++counts->eob_multi1024[plane][eob_multi_ctx][eob_pt - 1];
      if (allow_update_cdf)
        update_cdf(ec_ctx->eob_flag_cdf1024[plane][eob_multi_ctx], eob_pt - 1,
                   11);
      break;
  }
#else
Dake He's avatar
Dake He committed
225
  for (int i = 1; i < max_eob_pt; i++) {
Jingning Han's avatar
Jingning Han committed
226 227
    int eob_pos_ctx = av1_get_eob_pos_ctx(tx_type, i);
    counts->eob_flag[txs_ctx][plane][eob_pos_ctx][eob_pt == i]++;
228 229 230
    if (allow_update_cdf)
      update_cdf(ec_ctx->eob_flag_cdf[txs_ctx][plane][eob_pos_ctx], eob_pt == i,
                 2);
Dake He's avatar
Dake He committed
231 232 233 234
    if (eob_pt == i) {
      break;
    }
  }
235
#endif
Jingning Han's avatar
Jingning Han committed
236

Angie Chiang's avatar
Angie Chiang committed
237 238 239
  if (k_eob_offset_bits[eob_pt] > 0) {
    int eob_shift = k_eob_offset_bits[eob_pt] - 1;
    int bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
Jingning Han's avatar
Jingning Han committed
240
    counts->eob_extra[txs_ctx][plane][eob_pt][bit]++;
241 242
    if (allow_update_cdf)
      update_cdf(ec_ctx->eob_extra_cdf[txs_ctx][plane][eob_pt], bit, 2);
Angie Chiang's avatar
Angie Chiang committed
243
  }
Dake He's avatar
Dake He committed
244 245 246
}

static int get_eob_cost(int eob, int seg_eob,
247 248 249 250 251 252 253
#if CONFIG_LV_MAP_MULTI
                        const LV_MAP_EOB_COST *txb_eob_costs,
                        const LV_MAP_COEFF_COST *txb_costs,
#else
                        const LV_MAP_COEFF_COST *txb_costs,
#endif
                        TX_TYPE tx_type) {
Linfeng Zhang's avatar
Linfeng Zhang committed
254 255 256
  int eob_extra, dummy;
  const int eob_pt = get_eob_pos_token(eob, &eob_extra);
  const int max_eob_pt = get_eob_pos_token(seg_eob, &dummy);
Dake He's avatar
Dake He committed
257
  int eob_cost = 0;
258 259 260 261 262
#if CONFIG_LV_MAP_MULTI
  (void)max_eob_pt;
  const int eob_multi_ctx = (tx_type_to_class[tx_type] == TX_CLASS_2D) ? 0 : 1;
  eob_cost = txb_eob_costs->eob_cost[eob_multi_ctx][eob_pt - 1];
#else
Dake He's avatar
Dake He committed
263
  for (int i = 1; i < max_eob_pt; i++) {
Jingning Han's avatar
Jingning Han committed
264
    int eob_pos_ctx = av1_get_eob_pos_ctx(tx_type, i);
Dake He's avatar
Dake He committed
265 266 267 268 269
    eob_cost += txb_costs->eob_cost[eob_pos_ctx][eob_pt == i];
    if (eob_pt == i) {
      break;
    }
  }
270 271
#endif

Dake He's avatar
Dake He committed
272
  if (k_eob_offset_bits[eob_pt] > 0) {
Angie Chiang's avatar
Angie Chiang committed
273 274 275 276 277 278
    int eob_shift = k_eob_offset_bits[eob_pt] - 1;
    int bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
    eob_cost += txb_costs->eob_extra_cost[eob_pt][bit];
    for (int i = 1; i < k_eob_offset_bits[eob_pt]; i++) {
      eob_shift = k_eob_offset_bits[eob_pt] - 1 - i;
      bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
Dake He's avatar
Dake He committed
279 280 281 282 283 284
      eob_cost += av1_cost_bit(128, bit);
    }
  }
  return eob_cost;
}

285
static int get_coeff_cost(const tran_low_t qc, const int scan_idx,
286 287 288
#if CONFIG_LV_MAP_MULTI
                          const int is_eob,
#endif
289
                          const TxbInfo *const txb_info,
Linfeng Zhang's avatar
Linfeng Zhang committed
290 291
                          const LV_MAP_COEFF_COST *const txb_costs,
                          const int coeff_ctx);
Dake He's avatar
Dake He committed
292

Linfeng Zhang's avatar
Linfeng Zhang committed
293
static void get_dist_cost_stats(LevelDownStats *const stats, const int scan_idx,
294 295 296
#if CONFIG_LV_MAP_MULTI
                                const int is_eob,
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
297
                                const LV_MAP_COEFF_COST *const txb_costs,
298 299
                                const TxbInfo *const txb_info,
                                int has_nz_tail) {
Linfeng Zhang's avatar
Linfeng Zhang committed
300
  const int16_t *const scan = txb_info->scan_order->scan;
Dake He's avatar
Dake He committed
301 302
  const int coeff_idx = scan[scan_idx];
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
303
  const uint8_t *const levels = txb_info->levels;
Dake He's avatar
Dake He committed
304 305
  stats->new_eob = -1;
  stats->update = 0;
306 307
  stats->rd_low = 0;
  stats->rd = 0;
308 309 310 311 312 313 314 315 316 317
// TODO(mfo): explore if there's a better way to prevent compiler init
// warnings
#if CONFIG_LV_MAP_MULTI
  stats->nz_rd = 0;
#else
  stats->nz_rate = 0;
#endif
  stats->dist_low = 0;
  stats->rate_low = 0;
  stats->low_qc = 0;
Dake He's avatar
Dake He committed
318 319 320

  const tran_low_t tqc = txb_info->tcoeff[coeff_idx];
  const int dqv = txb_info->dequant[coeff_idx != 0];
321 322 323 324
#if CONFIG_NEW_QUANT
  const tran_low_t *nq_dequant_val = txb_info->nq_dequant_vals[coeff_idx != 0];
#endif  // CONFIG_NEW_QUANT

Linfeng Zhang's avatar
Linfeng Zhang committed
325 326 327 328 329
  const int coeff_ctx = get_nz_map_ctx(levels, coeff_idx, txb_info->bwl,
#if CONFIG_LV_MAP_MULTI
                                       txb_info->height, scan_idx, is_eob,
#endif
                                       txb_info->tx_size, txb_info->tx_type);
330
  const int qc_cost = get_coeff_cost(qc, scan_idx,
331
#if CONFIG_LV_MAP_MULTI
332
                                     is_eob,
333
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
334
                                     txb_info, txb_costs, coeff_ctx);
335 336 337 338 339 340 341 342 343 344 345
  if (qc == 0) {
    stats->dist = 0;
    stats->rate = qc_cost;
    return;
  } else {
    const tran_low_t dqc = qcoeff_to_dqcoeff(qc,
#if CONFIG_NEW_QUANT
                                             nq_dequant_val,
#endif  // CONFIG_NEW_QUANT
                                             dqv, txb_info->shift);
    const int64_t dqc_dist = get_coeff_dist(tqc, dqc, txb_info->shift);
Dake He's avatar
Dake He committed
346

347 348
    // distortion difference when coefficient is quantized to 0
    const tran_low_t dqc0 = qcoeff_to_dqcoeff(0,
349
#if CONFIG_NEW_QUANT
350
                                              nq_dequant_val,
351
#endif  // CONFIG_NEW_QUANT
352
                                              dqv, txb_info->shift);
Dake He's avatar
Dake He committed
353

354 355 356
    stats->dist0 = get_coeff_dist(tqc, dqc0, txb_info->shift);
    stats->dist = dqc_dist - stats->dist0;
    stats->rate = qc_cost;
Dake He's avatar
Dake He committed
357
  }
358

Dake He's avatar
Dake He committed
359 360 361
  stats->rd = RDCOST(txb_info->rdmult, stats->rate, stats->dist);

  stats->low_qc = get_lower_coeff(qc);
Dake He's avatar
Dake He committed
362 363 364 365 366 367 368 369

#if CONFIG_LV_MAP_MULTI
  if (is_eob && stats->low_qc == 0) {
    stats->rd_low = stats->rd;  // disable selection of low_qc in this case.
  } else {
#else
  {
#endif
370 371 372 373
    if (stats->low_qc == 0) {
      stats->dist_low = 0;
    } else {
      stats->low_dqc = qcoeff_to_dqcoeff(stats->low_qc,
374
#if CONFIG_NEW_QUANT
375
                                         nq_dequant_val,
376
#endif  // CONFIG_NEW_QUANT
377 378 379 380 381
                                         dqv, txb_info->shift);
      const int64_t low_dqc_dist =
          get_coeff_dist(tqc, stats->low_dqc, txb_info->shift);
      stats->dist_low = low_dqc_dist - stats->dist0;
    }
382
    const int low_qc_cost = get_coeff_cost(stats->low_qc, scan_idx,
383
#if CONFIG_LV_MAP_MULTI
384
                                           is_eob,
385
#endif
386
                                           txb_info, txb_costs, coeff_ctx);
Dake He's avatar
Dake He committed
387 388 389
    stats->rate_low = low_qc_cost;
    stats->rd_low = RDCOST(txb_info->rdmult, stats->rate_low, stats->dist_low);
  }
390
#if CONFIG_LV_MAP_MULTI
391 392 393 394 395
  if ((has_nz_tail < 2) && ((scan_idx == txb_info->eob - 1) || !is_eob))
#else
  if ((has_nz_tail < 2) && (scan_idx == txb_info->eob - 1))
#endif
  {
Linfeng Zhang's avatar
Linfeng Zhang committed
396
#if CONFIG_LV_MAP_MULTI
397 398 399 400 401
    (void)levels;
    const int coeff_ctx_temp =
        get_nz_map_ctx(levels, coeff_idx, txb_info->bwl,
#if CONFIG_LV_MAP_MULTI
                       txb_info->height, scan_idx, 1,
Dake He's avatar
Dake He committed
402
#endif
403 404 405 406 407 408 409 410 411 412 413 414 415
                       txb_info->tx_size, txb_info->tx_type);
    const int qc_eob_cost =
        get_coeff_cost(qc, scan_idx, 1, txb_info, txb_costs, coeff_ctx_temp);
    int64_t rd_eob = RDCOST(txb_info->rdmult, qc_eob_cost, stats->dist);
    if (stats->low_qc != 0) {
      const int low_qc_eob_cost = get_coeff_cost(
          stats->low_qc, scan_idx, 1, txb_info, txb_costs, coeff_ctx_temp);
      int64_t rd_eob_low =
          RDCOST(txb_info->rdmult, low_qc_eob_cost, stats->dist_low);
      rd_eob = (rd_eob > rd_eob_low) ? rd_eob_low : rd_eob;
    }

    stats->nz_rd = AOMMIN(stats->rd_low, stats->rd) - rd_eob;
416
#else
417 418
    const int is_nz = (stats->rd_low < stats->rd && stats->low_qc == 0) ? 0 : 1;
    stats->nz_rate = txb_costs->nz_map_cost[coeff_ctx][is_nz];
419
#endif
420
  }
Dake He's avatar
Dake He committed
421 422
}

423 424
static INLINE void update_qcoeff(const int coeff_idx, const tran_low_t qc,
                                 const TxbInfo *const txb_info) {
Dake He's avatar
Dake He committed
425
  txb_info->qcoeff[coeff_idx] = qc;
Linfeng Zhang's avatar
Linfeng Zhang committed
426
  txb_info->levels[get_padded_idx(coeff_idx, txb_info->bwl)] =
427
      (uint8_t)clamp(abs(qc), 0, INT8_MAX);
428 429 430 431 432
}

static INLINE void update_coeff(const int coeff_idx, const tran_low_t qc,
                                const TxbInfo *const txb_info) {
  update_qcoeff(coeff_idx, qc, txb_info);
Dake He's avatar
Dake He committed
433
  const int dqv = txb_info->dequant[coeff_idx != 0];
434 435 436 437 438 439 440 441
#if CONFIG_NEW_QUANT
  const tran_low_t *nq_dequant_val = txb_info->nq_dequant_vals[coeff_idx != 0];
#endif  // CONFIG_NEW_QUANT
  txb_info->dqcoeff[coeff_idx] = qcoeff_to_dqcoeff(qc,
#if CONFIG_NEW_QUANT
                                                   nq_dequant_val,
#endif  // CONFIG_NEW_QUANT
                                                   dqv, txb_info->shift);
Dake He's avatar
Dake He committed
442 443
}

444
static INLINE void av1_txb_init_levels(const tran_low_t *const coeff,
445
                                       const int width, const int height,
446 447
                                       uint8_t *const levels) {
  const int stride = width + TX_PAD_HOR;
448
  uint8_t *ls = levels;
449 450 451

  memset(levels - TX_PAD_TOP * stride, 0,
         sizeof(*levels) * TX_PAD_TOP * stride);
452 453
  memset(levels + stride * height, 0,
         sizeof(*levels) * (TX_PAD_BOTTOM * stride + TX_PAD_END));
454

455 456
  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
457
      *ls++ = (uint8_t)clamp(abs(coeff[i * width + j]), 0, INT8_MAX);
458 459 460 461
    }
    for (int j = 0; j < TX_PAD_HOR; j++) {
      *ls++ = 0;
    }
462 463 464
  }
}

465 466 467 468
void av1_get_nz_map_contexts_c(const uint8_t *const levels,
                               const int16_t *const scan, const uint16_t eob,
                               const TX_SIZE tx_size, const TX_TYPE tx_type,
                               int8_t *const coeff_contexts) {
Linfeng Zhang's avatar
Linfeng Zhang committed
469 470 471 472 473 474 475 476 477 478 479 480 481 482
  const int bwl = get_txb_bwl(tx_size);
#if CONFIG_LV_MAP_MULTI
  const int height = get_txb_high(tx_size);
#endif
  for (int i = 0; i < eob; ++i) {
    const int pos = scan[i];
    coeff_contexts[pos] = get_nz_map_ctx(levels, pos, bwl,
#if CONFIG_LV_MAP_MULTI
                                         height, i, i == eob - 1,
#endif
                                         tx_size, tx_type);
  }
}

483
void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
Luc Trudeau's avatar
Luc Trudeau committed
484 485
                          aom_writer *w, int blk_row, int blk_col, int plane,
                          TX_SIZE tx_size, const tran_low_t *tcoeff,
Jingning Han's avatar
Jingning Han committed
486
                          uint16_t eob, TXB_CTX *txb_ctx) {
487
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
488
  const PLANE_TYPE plane_type = get_plane_type(plane);
489
  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
490
  const TX_TYPE tx_type =
Luc Trudeau's avatar
Luc Trudeau committed
491
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, tx_size);
Angie Chiang's avatar
Angie Chiang committed
492
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
Linfeng Zhang's avatar
Linfeng Zhang committed
493
  const int16_t *const scan = scan_order->scan;
494
  const int seg_eob = av1_get_max_eob(tx_size);
495
  int c;
496 497 498
  const int bwl = get_txb_bwl(tx_size);
  const int width = get_txb_wide(tx_size);
  const int height = get_txb_high(tx_size);
499
  int update_eob = -1;
500
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
501 502
  uint8_t levels_buf[TX_PAD_2D];
  uint8_t *const levels = set_levels(levels_buf, width);
503
  DECLARE_ALIGNED(16, uint8_t, level_counts[MAX_TX_SQUARE]);
Linfeng Zhang's avatar
Linfeng Zhang committed
504
  DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
505

506 507
  aom_write_bin(w, eob == 0,
                ec_ctx->txb_skip_cdf[txs_ctx][txb_ctx->txb_skip_ctx], 2);
508
#if CONFIG_TXK_SEL
509 510 511
  if (plane == 0 && eob == 0) {
    assert(tx_type == DCT_DCT);
  }
512
#endif
513
  if (eob == 0) return;
514

515
  av1_txb_init_levels(tcoeff, width, height, levels);
516

Angie Chiang's avatar
Angie Chiang committed
517
#if CONFIG_TXK_SEL
518
  av1_write_tx_type(cm, xd, blk_row, blk_col, plane, tx_size, w);
Angie Chiang's avatar
Angie Chiang committed
519
#endif
520

Linfeng Zhang's avatar
Linfeng Zhang committed
521 522 523
  int eob_extra, dummy;
  const int eob_pt = get_eob_pos_token(eob, &eob_extra);
  const int max_eob_pt = get_eob_pos_token(seg_eob, &dummy);
Dake He's avatar
Dake He committed
524

525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560
#if CONFIG_LV_MAP_MULTI
  (void)max_eob_pt;
  const int eob_multi_size = txsize_log2_minus4[tx_size];
  const int eob_multi_ctx = (tx_type_to_class[tx_type] == TX_CLASS_2D) ? 0 : 1;
  switch (eob_multi_size) {
    case 0:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf16[plane_type][eob_multi_ctx], 5);
      break;
    case 1:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf32[plane_type][eob_multi_ctx], 6);
      break;
    case 2:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf64[plane_type][eob_multi_ctx], 7);
      break;
    case 3:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf128[plane_type][eob_multi_ctx], 8);
      break;
    case 4:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf256[plane_type][eob_multi_ctx], 9);
      break;
    case 5:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf512[plane_type][eob_multi_ctx], 10);
      break;
    default:
      aom_write_symbol(w, eob_pt - 1,
                       ec_ctx->eob_flag_cdf1024[plane_type][eob_multi_ctx], 11);
      break;
  }

#else
Dake He's avatar
Dake He committed
561
  for (int i = 1; i < max_eob_pt; i++) {
Jingning Han's avatar
Jingning Han committed
562
    int eob_pos_ctx = av1_get_eob_pos_ctx(tx_type, i);
Dake He's avatar
Dake He committed
563 564 565 566 567 568 569 570 571

    aom_write_bin(w, eob_pt == i,
                  ec_ctx->eob_flag_cdf[txs_ctx][plane_type][eob_pos_ctx], 2);
    // aom_write_symbol(w, eob_pt == i,
    // ec_ctx->eob_flag_cdf[AOMMIN(txs_ctx,3)][plane_type][eob_pos_ctx], 2);
    if (eob_pt == i) {
      break;
    }
  }
572
#endif
Dake He's avatar
Dake He committed
573 574

  if (k_eob_offset_bits[eob_pt] > 0) {
Angie Chiang's avatar
Angie Chiang committed
575 576 577 578 579 580 581
    int eob_shift = k_eob_offset_bits[eob_pt] - 1;
    int bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
    aom_write_bin(w, bit, ec_ctx->eob_extra_cdf[txs_ctx][plane_type][eob_pt],
                  2);
    for (int i = 1; i < k_eob_offset_bits[eob_pt]; i++) {
      eob_shift = k_eob_offset_bits[eob_pt] - 1 - i;
      bit = (eob_extra & (1 << eob_shift)) ? 1 : 0;
Dake He's avatar
Dake He committed
582 583 584
      aom_write_bit(w, bit);
    }
  }
585

Linfeng Zhang's avatar
Linfeng Zhang committed
586 587 588
  av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_type, coeff_contexts);

  for (c = eob - 1; c >= 0; --c) {
Linfeng Zhang's avatar
Linfeng Zhang committed
589
    const int pos = scan[c];
Linfeng Zhang's avatar
Linfeng Zhang committed
590 591
    const int coeff_ctx = coeff_contexts[pos];
    const tran_low_t v = tcoeff[pos];
592

593
#if CONFIG_LV_MAP_MULTI
594 595 596
    if (c == eob - 1) {
      aom_write_symbol(
          w, AOMMIN(abs(v), 3) - 1,
Dake He's avatar
Dake He committed
597
          ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx], 3);
598 599 600 601 602
    } else {
      aom_write_symbol(w, AOMMIN(abs(v), 3),
                       ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx],
                       4);
    }
Dake He's avatar
Dake He committed
603

604
#else
Linfeng Zhang's avatar
Linfeng Zhang committed
605
    const int is_nz = (v != 0);
606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621

    if (c < eob - 1) {
      aom_write_bin(w, is_nz,
                    ec_ctx->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], 2);
    }
    if (is_nz) {
      const int level = abs(v);
      int k;
      for (k = 0; k < NUM_BASE_LEVELS; ++k) {
        int is_k = (level > (k + 1));
        int ctx = coeff_ctx;
        aom_write_bin(w, is_k,
                      ec_ctx->coeff_base_cdf[txs_ctx][plane_type][k][ctx], 2);
        if (is_k == 0) break;
      }
    }
622
#endif
623 624
  }
  update_eob = eob - 1;
625

626 627 628
  // Loop to code all signs in the transform block,
  // starting with the sign of DC (if applicable)
  for (c = 0; c < eob; ++c) {
629 630 631
    const tran_low_t v = tcoeff[scan[c]];
    const tran_low_t level = abs(v);
    const int sign = (v < 0) ? 1 : 0;
632
    if (level == 0) continue;
633 634

    if (c == 0) {
635 636
      aom_write_bin(w, sign,
                    ec_ctx->dc_sign_cdf[plane_type][txb_ctx->dc_sign_ctx], 2);
637 638 639
    } else {
      aom_write_bit(w, sign);
    }
640 641
  }

Linfeng Zhang's avatar
Linfeng Zhang committed
642
  if (update_eob >= 0) {
643
#if !CONFIG_LV_MAP_MULTI
Linfeng Zhang's avatar
Linfeng Zhang committed
644
    av1_get_br_level_counts(levels, width, height, level_counts);
645
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
646
    for (c = update_eob; c >= 0; --c) {
Linfeng Zhang's avatar
Linfeng Zhang committed
647 648
      const int pos = scan[c];
      const tran_low_t level = abs(tcoeff[pos]);
Linfeng Zhang's avatar
Linfeng Zhang committed
649 650
      int idx;
      int ctx;
651

Linfeng Zhang's avatar
Linfeng Zhang committed
652 653
      if (level <= NUM_BASE_LEVELS) continue;

654 655
// level is above 1.
#if !CONFIG_LV_MAP_MULTI
Linfeng Zhang's avatar
Linfeng Zhang committed
656
      ctx = get_br_ctx(levels, pos, bwl, level_counts[pos]);
657
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
658
      const int base_range = level - 1 - NUM_BASE_LEVELS;
659
#if CONFIG_LV_MAP_MULTI
660
#if USE_CAUSAL_BR_CTX
Linfeng Zhang's avatar
Linfeng Zhang committed
661
      ctx = get_br_ctx(levels, pos, bwl, level_counts[pos], tx_type);
662 663

#else
Linfeng Zhang's avatar
Linfeng Zhang committed
664
      ctx = get_br_ctx(levels, pos, bwl, level_counts[pos]);
665
#endif
666
      for (idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
Linfeng Zhang's avatar
Linfeng Zhang committed
667
        const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
Dake He's avatar
Dake He committed
668 669
        aom_write_symbol(w, k,
#if 0
670
            ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_16X16)][plane_type][ctx],
Dake He's avatar
Dake He committed
671 672 673 674 675
#else
                         ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)]
                                             [plane_type][ctx],
#endif
                         BR_CDF_SIZE);
676 677 678 679
        if (k < BR_CDF_SIZE - 1) break;
      }
      if (base_range < COEFF_BASE_RANGE) continue;
#else
Linfeng Zhang's avatar
Linfeng Zhang committed
680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702
      int br_set_idx = 0;
      int br_base = 0;
      int br_offset = 0;

      if (base_range >= COEFF_BASE_RANGE)
        br_set_idx = BASE_RANGE_SETS;
      else
        br_set_idx = coeff_to_br_index[base_range];

      for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
        aom_write_bin(w, idx == br_set_idx,
                      ec_ctx->coeff_br_cdf[txs_ctx][plane_type][idx][ctx], 2);
        if (idx == br_set_idx) {
          br_base = br_index_to_coeff[br_set_idx];
          br_offset = base_range - br_base;
          int extra_bits = (1 << br_extra_bits[idx]) - 1;
          for (int tok = 0; tok < extra_bits; ++tok) {
            if (tok == br_offset) {
              aom_write_bin(w, 1,
                            ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2);
              break;
            }
            aom_write_bin(w, 0, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
703
                          2);
704
          }
Linfeng Zhang's avatar
Linfeng Zhang committed
705 706
          //        aom_write_literal(w, br_offset, br_extra_bits[idx]);
          break;
707 708 709
        }
      }

Linfeng Zhang's avatar
Linfeng Zhang committed
710
      if (br_set_idx < BASE_RANGE_SETS) continue;
711
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
712
      // use 0-th order Golomb code to handle the residual level.
Linfeng Zhang's avatar
Linfeng Zhang committed
713 714
      write_golomb(w,
                   abs(tcoeff[pos]) - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
Linfeng Zhang's avatar
Linfeng Zhang committed
715
    }
716 717
  }
}
718

719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737
typedef struct encode_txb_args {
  const AV1_COMMON *cm;
  MACROBLOCK *x;
  aom_writer *w;
} ENCODE_TXB_ARGS;

static void av1_write_coeffs_txb_wrap(int plane, int block, int blk_row,
                                      int blk_col, BLOCK_SIZE plane_bsize,
                                      TX_SIZE tx_size, void *arg) {
  (void)plane_bsize;
  ENCODE_TXB_ARGS *enc_args = (ENCODE_TXB_ARGS *)arg;
  const AV1_COMMON *cm = enc_args->cm;
  MACROBLOCK *x = enc_args->x;
  MACROBLOCKD *xd = &x->e_mbd;
  aom_writer *w = enc_args->w;
  tran_low_t *tcoeff = BLOCK_OFFSET(x->mbmi_ext->tcoeff[plane], block);
  uint16_t eob = x->mbmi_ext->eobs[plane][block];
  TXB_CTX txb_ctx = { x->mbmi_ext->txb_skip_ctx[plane][block],
                      x->mbmi_ext->dc_sign_ctx[plane][block] };
Luc Trudeau's avatar
Luc Trudeau committed
738 739
  av1_write_coeffs_txb(cm, xd, w, blk_row, blk_col, plane, tx_size, tcoeff, eob,
                       &txb_ctx);
740 741
}

Angie Chiang's avatar
Angie Chiang committed
742
void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
743 744
                         aom_writer *w, int plane, BLOCK_SIZE bsize) {
  ENCODE_TXB_ARGS enc_args = { cm, x, w };
Angie Chiang's avatar
Angie Chiang committed
745
  MACROBLOCKD *xd = &x->e_mbd;
746 747
  av1_foreach_transformed_block_in_plane(xd, bsize, plane,
                                         av1_write_coeffs_txb_wrap, &enc_args);
Angie Chiang's avatar
Angie Chiang committed
748 749
}

750
static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
Angie Chiang's avatar
Angie Chiang committed
751
                              const int *coeff_lps) {
752 753
  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
754
  (void)ctx;
755 756
  if (abs_qc >= min_level) {
    if (abs_qc >= max_level)
757
      return coeff_lps[COEFF_BASE_RANGE];  // COEFF_BASE_RANGE * cost0;
758
    else
759
      return coeff_lps[(abs_qc - min_level)];  //  * cost0 + cost1;
760 761 762 763 764
  } else {
    return 0;
  }
}

765
#if !CONFIG_LV_MAP_MULTI
766
static INLINE int get_base_cost(tran_low_t abs_qc, int ctx,
767
                                const int coeff_base[2], int base_idx) {
768
  const int level = base_idx + 1;
769
  (void)ctx;
770 771 772
  if (abs_qc < level)
    return 0;
  else
773
    return coeff_base[abs_qc == level];
774
}
775
#endif
776

777 778 779 780 781 782 783
// Note: don't call this function when eob is 0.
int av1_cost_coeffs_txb(const AV1_COMMON *const cm, const MACROBLOCK *x,
                        const int plane, const int blk_row, const int blk_col,
                        const int block, const TX_SIZE tx_size,
                        const TXB_CTX *const txb_ctx) {
  const MACROBLOCKD *const xd = &x->e_mbd;
  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
784
  const PLANE_TYPE plane_type = get_plane_type(plane);
785
  const TX_TYPE tx_type =
Luc Trudeau's avatar
Luc Trudeau committed
786
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, tx_size);
787
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
788 789 790 791
  const struct macroblock_plane *p = &x->plane[plane];
  const int eob = p->eobs[block];
  const tran_low_t *const qcoeff = BLOCK_OFFSET(p->qcoeff, block);
  int c, cost;
792
  const int txb_skip_ctx = txb_ctx->txb_skip_ctx;
793 794 795
  const int bwl = get_txb_bwl(tx_size);
  const int width = get_txb_wide(tx_size);
  const int height = get_txb_high(tx_size);
Angie Chiang's avatar
Angie Chiang committed
796
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
Linfeng Zhang's avatar
Linfeng Zhang committed
797
  const int16_t *const scan = scan_order->scan;
798 799
  uint8_t levels_buf[TX_PAD_2D];
  uint8_t *const levels = set_levels(levels_buf, width);
800
  DECLARE_ALIGNED(16, uint8_t, level_counts[MAX_TX_SQUARE]);
Linfeng Zhang's avatar
Linfeng Zhang committed
801
  DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
802 803
  const LV_MAP_COEFF_COST *const coeff_costs =
      &x->coeff_costs[txs_ctx][plane_type];
804

805 806 807 808 809
#if CONFIG_LV_MAP_MULTI
  const int eob_multi_size = txsize_log2_minus4[tx_size];
  const LV_MAP_EOB_COST *const eob_costs =
      &x->eob_costs[eob_multi_size][plane_type];
#endif
810 811
  // eob must be greater than 0 here.
  assert(eob > 0);
812
  cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
813

814
  av1_txb_init_levels(qcoeff, width, height, levels);
815

Angie Chiang's avatar
Angie Chiang committed
816
#if CONFIG_TXK_SEL
817
  cost += av1_tx_type_cost(cm, x, xd, mbmi->sb_type, plane, tx_size, tx_type);
Angie Chiang's avatar
Angie Chiang committed
818
#endif
819

820
  const int seg_eob = av1_get_max_eob(tx_size);
821 822 823 824 825 826 827
  int eob_cost = get_eob_cost(eob, seg_eob,
#if CONFIG_LV_MAP_MULTI
                              eob_costs, coeff_costs,
#else
                              coeff_costs,
#endif
                              tx_type);
828
#if !CONFIG_LV_MAP_MULTI
Linfeng Zhang's avatar
Linfeng Zhang committed
829
  av1_get_br_level_counts(levels, width, height, level_counts);
830
#endif
Dake He's avatar
Dake He committed
831
  cost += eob_cost;
Linfeng Zhang's avatar
Linfeng Zhang committed
832 833 834

  av1_get_nz_map_contexts(levels, scan, eob, tx_size, tx_type, coeff_contexts);

Dake He's avatar
Dake He committed
835
  for (c = eob - 1; c >= 0; --c) {
Linfeng Zhang's avatar
Linfeng Zhang committed
836 837 838 839
    const int pos = scan[c];
    const tran_low_t v = qcoeff[pos];
    const int is_nz = (v != 0);
    const int level = abs(v);
Linfeng Zhang's avatar
Linfeng Zhang committed
840
    const int coeff_ctx = coeff_contexts[pos];
841
#if CONFIG_LV_MAP_MULTI
842
    if (c == eob - 1) {
Dake He's avatar
Dake He committed
843
      cost += coeff_costs->base_eob_cost[coeff_ctx][AOMMIN(level, 3) - 1];
844 845 846
    } else {
      cost += coeff_costs->base_cost[coeff_ctx][AOMMIN(level, 3)];
    }
847
#else   // CONFIG_LV_MAP_MULTI
Dake He's avatar
Dake He committed
848 849 850
    if (c < eob - 1) {
      cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
    }
851
#endif  // CONFIG_LV_MAP_MULTI
Dake He's avatar
Dake He committed
852 853 854 855 856 857 858 859 860 861 862

    if (is_nz) {
      int sign = (v < 0) ? 1 : 0;

      // sign bit cost
      if (c == 0) {
        int dc_sign_ctx = txb_ctx->dc_sign_ctx;
        cost += coeff_costs->dc_sign_cost[dc_sign_ctx][sign];
      } else {
        cost += av1_cost_bit(128, sign);
      }
863
#if !CONFIG_LV_MAP_MULTI
864 865 866 867 868 869 870 871
      int k;
      for (k = 0; k < NUM_BASE_LEVELS; ++k) {
        int is_k = (level > (k + 1));
        int ctx = coeff_ctx;
        // get_base_ctx_from_b(c, k, b0, b1, b2);
        cost += coeff_costs->base_cost[k][ctx][is_k];
        if (is_k == 0) break;
      }
872
#endif  // CONFIG_LV_MAP_MULTI
Dake He's avatar
Dake He committed
873 874
      if (level > NUM_BASE_LEVELS) {
        int ctx;
875
#if CONFIG_LV_MAP_MULTI && USE_CAUSAL_BR_CTX
Linfeng Zhang's avatar
Linfeng Zhang committed
876
        ctx = get_br_ctx(levels, pos, bwl, level_counts[pos], tx_type);
877
#else
Linfeng Zhang's avatar
Linfeng Zhang committed
878
        ctx = get_br_ctx(levels, pos, bwl, level_counts[pos]);
879
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
880
        const int base_range = level - 1 - NUM_BASE_LEVELS;
Dake He's avatar
Dake He committed
881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905
        if (base_range < COEFF_BASE_RANGE) {
          cost += coeff_costs->lps_cost[ctx][base_range];
        } else {
          cost += coeff_costs->lps_cost[ctx][COEFF_BASE_RANGE];
        }

        if (level >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
          // residual cost
          int r = level - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
          int ri = r;
          int length = 0;

          while (ri) {
            ri >>= 1;
            ++length;
          }

          for (ri = 0; ri < length - 1; ++ri) cost += av1_cost_bit(128, 0);

          for (ri = length - 1; ri >= 0; --ri)
            cost += av1_cost_bit(128, (r >> ri) & 0x01);
        }
      }
    }
  }
906 907
  return cost;
}
Angie Chiang's avatar
Angie Chiang committed
908

909 910 911 912 913 914 915 916 917
static INLINE int has_base(tran_low_t qc, int base_idx) {
  const int level = base_idx + 1;
  return abs(qc) >= level;
}

static INLINE int has_br(tran_low_t qc) {
  return abs(qc) >= 1 + NUM_BASE_LEVELS;
}

918
static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
919
                                    const int (*dc_sign_cost)[2],
920 921 922 923
                                    int dc_sign_ctx) {
  const int sign = (qc < 0) ? 1 : 0;
  // sign bit cost
  if (coeff_idx == 0) {
924
    return dc_sign_cost[dc_sign_ctx][sign];
925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946
  } else {
    return av1_cost_bit(128, sign);
  }
}
static INLINE int get_golomb_cost(int abs_qc) {
  if (abs_qc >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
    // residual cost
    int r = abs_qc - COEFF_BASE_RANGE - NUM_BASE_LEVELS;
    int ri = r;
    int length = 0;

    while (ri) {
      ri >>= 1;
      ++length;
    }

    return av1_cost_literal(2 * length - 1);
  } else {
    return 0;
  }
}

947
void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
Angie Chiang's avatar
Angie Chiang committed
948
  // gen_nz_count_arr
Linfeng Zhang's avatar
Linfeng Zhang committed
949
  const int16_t *const scan = txb_info->scan_order->scan;
Angie Chiang's avatar
Angie Chiang committed
950 951
  const int bwl = txb_info->bwl;
  const int height = txb_info->height;
Linfeng Zhang's avatar
Linfeng Zhang committed
952
  const tran_low_t *const qcoeff = txb_info->qcoeff;
953
  const uint8_t *const levels = txb_info->levels;
954 955
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
956 957 958 959
  for (int c = 0; c < txb_info->eob; ++c) {
    const int coeff_idx = scan[c];  // raster order
    const int row = coeff_idx >> bwl;
    const int col = coeff_idx - (row << bwl);
Jingning Han's avatar
Jingning Han committed
960

Linfeng Zhang's avatar
Linfeng Zhang committed
961
    txb_cache->nz_count_arr[coeff_idx] =
Linfeng Zhang's avatar
Linfeng Zhang committed
962
        get_nz_count(levels + get_padded_idx(coeff_idx, bwl), bwl,
Linfeng Zhang's avatar
Linfeng Zhang committed
963
                     tx_type_to_class[txb_info->tx_type]);
Dake He's avatar
Dake He committed
964

965 966 967
    txb_cache->nz_ctx_arr[coeff_idx] =
        get_nz_map_ctx_from_stats(0, coeff_idx, bwl, txb_info->tx_size,
                                  tx_type_to_class[txb_info->tx_type]);
Angie Chiang's avatar
Angie Chiang committed
968 969 970 971

    // gen_base_count_mag_arr
    if (!has_base(qcoeff[coeff_idx], 0)) continue;
    int *base_mag = txb_cache->base_mag_arr[coeff_idx];
Angie Chiang's avatar
Angie Chiang committed
972 973
    int count[NUM_BASE_LEVELS];
    get_base_count_mag(base_mag, count, qcoeff, bwl, height, row, col);
Angie Chiang's avatar
Angie Chiang committed
974 975

    for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
Angie Chiang's avatar
Angie Chiang committed
976 977
      if (!has_base(qcoeff[coeff_idx], i)) break;
      txb_cache->base_count_arr[i][coeff_idx] = count[i];
Angie Chiang's avatar
Angie Chiang committed
978
      const int level = i + 1;
Angie Chiang's avatar
Angie Chiang committed
979
      txb_cache->base_ctx_arr[i][coeff_idx] =
980
          base_ctx_table[row != 0][col != 0][base_mag[0] > level][count[i]];
Angie Chiang's avatar
Angie Chiang committed
981 982 983 984 985 986
    }

    // gen_br_count_mag_arr
    if (!has_br(qcoeff[coeff_idx])) continue;
    int *br_count = txb_cache->br_count_arr + coeff_idx;
    int *br_mag = txb_cache->br_mag_arr[coeff_idx];
Angie Chiang's avatar
Angie Chiang committed
987 988
    *br_count = get_br_count_mag(br_mag, qcoeff, bwl, height, row, col,
                                 NUM_BASE_LEVELS);
Angie Chiang's avatar
Angie Chiang committed
989 990 991
    txb_cache->br_ctx_arr[coeff_idx] =
        get_br_ctx_from_count_mag(row, col, *br_count, br_mag[0]);
  }
992 993
}

994 995 996
static INLINE const int *get_level_prob(int level, int coeff_idx,
                                        const TxbCache *txb_cache,
                                        const LV_MAP_COEFF_COST *txb_costs) {
997 998 999 1000 1001
#if CONFIG_LV_MAP_MULTI
  if (level < 1 + NUM_BASE_LEVELS