encodetxb.c 84.3 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
  int64_t nz_rd;
Dake He's avatar
Dake He committed
52
53
54
55
56
57
  int64_t rd_diff;
  int cost_diff;
  int64_t dist_diff;
  int new_eob;
} LevelDownStats;

58
void av1_alloc_txb_buf(AV1_COMP *cpi) {
59
#if 0
60
61
62
63
64
65
66
67
68
69
70
71
  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));
  }
72
#else
73
  AV1_COMMON *cm = &cpi->common;
74
75
  int size = ((cm->mi_rows >> cm->mib_size_log2) + 1) *
             ((cm->mi_cols >> cm->mib_size_log2) + 1);
76

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

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

95
96
void av1_set_coeff_buffer(const AV1_COMP *const cpi, MACROBLOCK *const x,
                          int mi_row, int mi_col) {
97
98
99
  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);
100
101
102
103
104
105
106
107
108
109
110
111
  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;
  }
}

112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
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
128
129
130
131
132
133
134
static INLINE tran_low_t get_lower_coeff(tran_low_t qc) {
  if (qc == 0) {
    return 0;
  }
  return qc > 0 ? qc - 1 : qc + 1;
}

135
136
137
138
139
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
140
  int sgn = qc < 0 ? -1 : 1;
141
#if CONFIG_NEW_QUANT
142
143
  int dqcoeff = av1_dequant_coeff_nuq(abs(qc), dqv, nq_dq, shift);
  return sgn * dqcoeff;
144
145
#endif  // CONFIG_NEW_QUANT

Dake He's avatar
Dake He committed
146
147
148
149
150
  return sgn * ((abs(qc) * dqv) >> shift);
}

static INLINE int64_t get_coeff_dist(tran_low_t tcoeff, tran_low_t dqcoeff,
                                     int shift) {
151
152
153
154
#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;
155
  const int64_t error = (diff * diff + depth_round) >> depth_shift;
156
157
  (void)shift;
#else
Dake He's avatar
Dake He committed
158
159
  const int64_t diff = (tcoeff - dqcoeff) * (1 << shift);
  const int64_t error = diff * diff;
160
#endif
Dake He's avatar
Dake He committed
161
162
163
  return error;
}

Jingning Han's avatar
Jingning Han committed
164
165
void av1_update_eob_context(int eob, int seg_eob, TX_SIZE tx_size,
                            TX_TYPE tx_type, PLANE_TYPE plane,
166
167
                            FRAME_CONTEXT *ec_ctx, FRAME_COUNTS *counts,
                            uint8_t allow_update_cdf) {
Linfeng Zhang's avatar
Linfeng Zhang committed
168
169
170
  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);
171
  TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
Dake He's avatar
Dake He committed
172

173
174
175
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
  (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;
  }
Jingning Han's avatar
Jingning Han committed
219

Angie Chiang's avatar
Angie Chiang committed
220
221
222
  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
223
    counts->eob_extra[txs_ctx][plane][eob_pt][bit]++;
224
225
    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
226
  }
Dake He's avatar
Dake He committed
227
228
229
}

static int get_eob_cost(int eob, int seg_eob,
230
                        const LV_MAP_EOB_COST *txb_eob_costs,
231
                        const LV_MAP_COEFF_COST *txb_costs, TX_TYPE tx_type) {
Linfeng Zhang's avatar
Linfeng Zhang committed
232
233
234
  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
235
  int eob_cost = 0;
236
237
238
239
  (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];

Dake He's avatar
Dake He committed
240
  if (k_eob_offset_bits[eob_pt] > 0) {
Angie Chiang's avatar
Angie Chiang committed
241
242
243
244
245
246
    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
247
248
249
250
251
252
      eob_cost += av1_cost_bit(128, bit);
    }
  }
  return eob_cost;
}

253
static int get_coeff_cost(const tran_low_t qc, const int scan_idx,
254
                          const int is_eob, const TxbInfo *const txb_info,
Linfeng Zhang's avatar
Linfeng Zhang committed
255
256
                          const LV_MAP_COEFF_COST *const txb_costs,
                          const int coeff_ctx);
Dake He's avatar
Dake He committed
257

Linfeng Zhang's avatar
Linfeng Zhang committed
258
static void get_dist_cost_stats(LevelDownStats *const stats, const int scan_idx,
259
                                const int is_eob,
Linfeng Zhang's avatar
Linfeng Zhang committed
260
                                const LV_MAP_COEFF_COST *const txb_costs,
261
262
                                const TxbInfo *const txb_info,
                                int has_nz_tail) {
Linfeng Zhang's avatar
Linfeng Zhang committed
263
  const int16_t *const scan = txb_info->scan_order->scan;
Dake He's avatar
Dake He committed
264
265
  const int coeff_idx = scan[scan_idx];
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
266
  const uint8_t *const levels = txb_info->levels;
Dake He's avatar
Dake He committed
267
268
  stats->new_eob = -1;
  stats->update = 0;
269
270
  stats->rd_low = 0;
  stats->rd = 0;
271
272
  // TODO(mfo): explore if there's a better way to prevent compiler init
  // warnings
273
274
275
276
  stats->nz_rd = 0;
  stats->dist_low = 0;
  stats->rate_low = 0;
  stats->low_qc = 0;
Dake He's avatar
Dake He committed
277
278
279

  const tran_low_t tqc = txb_info->tcoeff[coeff_idx];
  const int dqv = txb_info->dequant[coeff_idx != 0];
280
281
282
283
#if CONFIG_NEW_QUANT
  const tran_low_t *nq_dequant_val = txb_info->nq_dequant_vals[coeff_idx != 0];
#endif  // CONFIG_NEW_QUANT

284
285
286
287
288
  const int coeff_ctx =
      get_nz_map_ctx(levels, coeff_idx, txb_info->bwl, txb_info->height,
                     scan_idx, is_eob, txb_info->tx_size, txb_info->tx_type);
  const int qc_cost =
      get_coeff_cost(qc, scan_idx, is_eob, txb_info, txb_costs, coeff_ctx);
289
290
291
292
293
294
295
296
297
298
299
  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
300

301
302
    // distortion difference when coefficient is quantized to 0
    const tran_low_t dqc0 = qcoeff_to_dqcoeff(0,
303
#if CONFIG_NEW_QUANT
304
                                              nq_dequant_val,
305
#endif  // CONFIG_NEW_QUANT
306
                                              dqv, txb_info->shift);
Dake He's avatar
Dake He committed
307

308
309
310
    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
311
  }
312

Dake He's avatar
Dake He committed
313
314
315
  stats->rd = RDCOST(txb_info->rdmult, stats->rate, stats->dist);

  stats->low_qc = get_lower_coeff(qc);
Dake He's avatar
Dake He committed
316
317
318
319

  if (is_eob && stats->low_qc == 0) {
    stats->rd_low = stats->rd;  // disable selection of low_qc in this case.
  } else {
320
321
322
323
    if (stats->low_qc == 0) {
      stats->dist_low = 0;
    } else {
      stats->low_dqc = qcoeff_to_dqcoeff(stats->low_qc,
324
#if CONFIG_NEW_QUANT
325
                                         nq_dequant_val,
326
#endif  // CONFIG_NEW_QUANT
327
328
329
330
331
                                         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;
    }
332
    const int low_qc_cost = get_coeff_cost(stats->low_qc, scan_idx, is_eob,
333
                                           txb_info, txb_costs, coeff_ctx);
Dake He's avatar
Dake He committed
334
335
336
    stats->rate_low = low_qc_cost;
    stats->rd_low = RDCOST(txb_info->rdmult, stats->rate_low, stats->dist_low);
  }
337
  if ((has_nz_tail < 2) && ((scan_idx == txb_info->eob - 1) || !is_eob)) {
338
339
    (void)levels;
    const int coeff_ctx_temp =
340
341
        get_nz_map_ctx(levels, coeff_idx, txb_info->bwl, txb_info->height,
                       scan_idx, 1, txb_info->tx_size, txb_info->tx_type);
342
343
344
345
346
347
348
349
350
351
352
353
354
    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;
  }
Dake He's avatar
Dake He committed
355
356
}

357
358
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
359
  txb_info->qcoeff[coeff_idx] = qc;
Linfeng Zhang's avatar
Linfeng Zhang committed
360
  txb_info->levels[get_padded_idx(coeff_idx, txb_info->bwl)] =
361
      (uint8_t)clamp(abs(qc), 0, INT8_MAX);
362
363
364
365
366
}

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
367
  const int dqv = txb_info->dequant[coeff_idx != 0];
368
369
370
371
372
373
374
375
#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
376
377
}

378
static INLINE void av1_txb_init_levels(const tran_low_t *const coeff,
379
                                       const int width, const int height,
380
381
                                       uint8_t *const levels) {
  const int stride = width + TX_PAD_HOR;
382
  uint8_t *ls = levels;
383
384
385

  memset(levels - TX_PAD_TOP * stride, 0,
         sizeof(*levels) * TX_PAD_TOP * stride);
386
387
  memset(levels + stride * height, 0,
         sizeof(*levels) * (TX_PAD_BOTTOM * stride + TX_PAD_END));
388

389
390
  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
391
      *ls++ = (uint8_t)clamp(abs(coeff[i * width + j]), 0, INT8_MAX);
392
393
394
395
    }
    for (int j = 0; j < TX_PAD_HOR; j++) {
      *ls++ = 0;
    }
396
397
398
  }
}

399
400
401
402
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
403
404
405
406
  const int bwl = get_txb_bwl(tx_size);
  const int height = get_txb_high(tx_size);
  for (int i = 0; i < eob; ++i) {
    const int pos = scan[i];
407
408
    coeff_contexts[pos] = get_nz_map_ctx(levels, pos, bwl, height, i,
                                         i == eob - 1, tx_size, tx_type);
Linfeng Zhang's avatar
Linfeng Zhang committed
409
410
411
  }
}

412
void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
Luc Trudeau's avatar
Luc Trudeau committed
413
414
                          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
415
                          uint16_t eob, TXB_CTX *txb_ctx) {
416
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
417
  const PLANE_TYPE plane_type = get_plane_type(plane);
418
  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);
419
  const TX_TYPE tx_type =
Luc Trudeau's avatar
Luc Trudeau committed
420
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, tx_size);
Angie Chiang's avatar
Angie Chiang committed
421
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
Linfeng Zhang's avatar
Linfeng Zhang committed
422
  const int16_t *const scan = scan_order->scan;
423
  const int seg_eob = av1_get_max_eob(tx_size);
424
  int c;
425
426
427
  const int bwl = get_txb_bwl(tx_size);
  const int width = get_txb_wide(tx_size);
  const int height = get_txb_high(tx_size);
428
  int update_eob = -1;
429
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
430
431
  uint8_t levels_buf[TX_PAD_2D];
  uint8_t *const levels = set_levels(levels_buf, width);
432
  DECLARE_ALIGNED(16, uint8_t, level_counts[MAX_TX_SQUARE]);
Linfeng Zhang's avatar
Linfeng Zhang committed
433
  DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
434

435
436
  aom_write_bin(w, eob == 0,
                ec_ctx->txb_skip_cdf[txs_ctx][txb_ctx->txb_skip_ctx], 2);
437
#if CONFIG_TXK_SEL
438
439
440
  if (plane == 0 && eob == 0) {
    assert(tx_type == DCT_DCT);
  }
441
#endif
442
  if (eob == 0) return;
443

444
  av1_txb_init_levels(tcoeff, width, height, levels);
445

Angie Chiang's avatar
Angie Chiang committed
446
#if CONFIG_TXK_SEL
447
  av1_write_tx_type(cm, xd, blk_row, blk_col, plane, tx_size, w);
Angie Chiang's avatar
Angie Chiang committed
448
#endif
449

Linfeng Zhang's avatar
Linfeng Zhang committed
450
451
452
  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
453

454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
  (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;
  }

Dake He's avatar
Dake He committed
488
  if (k_eob_offset_bits[eob_pt] > 0) {
Angie Chiang's avatar
Angie Chiang committed
489
490
491
492
493
494
495
    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
496
497
498
      aom_write_bit(w, bit);
    }
  }
499

Linfeng Zhang's avatar
Linfeng Zhang committed
500
501
502
  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
503
    const int pos = scan[c];
Linfeng Zhang's avatar
Linfeng Zhang committed
504
505
    const int coeff_ctx = coeff_contexts[pos];
    const tran_low_t v = tcoeff[pos];
506

507
508
509
    if (c == eob - 1) {
      aom_write_symbol(
          w, AOMMIN(abs(v), 3) - 1,
Dake He's avatar
Dake He committed
510
          ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx], 3);
511
512
513
514
515
    } else {
      aom_write_symbol(w, AOMMIN(abs(v), 3),
                       ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx],
                       4);
    }
516
517
  }
  update_eob = eob - 1;
518

519
520
521
  // Loop to code all signs in the transform block,
  // starting with the sign of DC (if applicable)
  for (c = 0; c < eob; ++c) {
522
523
524
    const tran_low_t v = tcoeff[scan[c]];
    const tran_low_t level = abs(v);
    const int sign = (v < 0) ? 1 : 0;
525
    if (level == 0) continue;
526
527

    if (c == 0) {
528
529
      aom_write_bin(w, sign,
                    ec_ctx->dc_sign_cdf[plane_type][txb_ctx->dc_sign_ctx], 2);
530
531
532
    } else {
      aom_write_bit(w, sign);
    }
533
534
  }

Linfeng Zhang's avatar
Linfeng Zhang committed
535
536
  if (update_eob >= 0) {
    for (c = update_eob; c >= 0; --c) {
Linfeng Zhang's avatar
Linfeng Zhang committed
537
538
      const int pos = scan[c];
      const tran_low_t level = abs(tcoeff[pos]);
Linfeng Zhang's avatar
Linfeng Zhang committed
539
540
      int idx;
      int ctx;
541

Linfeng Zhang's avatar
Linfeng Zhang committed
542
543
      if (level <= NUM_BASE_LEVELS) continue;

544
      // level is above 1.
Linfeng Zhang's avatar
Linfeng Zhang committed
545
      const int base_range = level - 1 - NUM_BASE_LEVELS;
546
#if USE_CAUSAL_BR_CTX
Linfeng Zhang's avatar
Linfeng Zhang committed
547
      ctx = get_br_ctx(levels, pos, bwl, level_counts[pos], tx_type);
548
549

#else
Linfeng Zhang's avatar
Linfeng Zhang committed
550
      ctx = get_br_ctx(levels, pos, bwl, level_counts[pos]);
551
#endif
552
      for (idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
Linfeng Zhang's avatar
Linfeng Zhang committed
553
        const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
Dake He's avatar
Dake He committed
554
555
        aom_write_symbol(w, k,
#if 0
556
            ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_16X16)][plane_type][ctx],
Dake He's avatar
Dake He committed
557
558
559
560
561
#else
                         ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)]
                                             [plane_type][ctx],
#endif
                         BR_CDF_SIZE);
562
563
564
        if (k < BR_CDF_SIZE - 1) break;
      }
      if (base_range < COEFF_BASE_RANGE) continue;
Linfeng Zhang's avatar
Linfeng Zhang committed
565
      // use 0-th order Golomb code to handle the residual level.
Linfeng Zhang's avatar
Linfeng Zhang committed
566
567
      write_golomb(w,
                   abs(tcoeff[pos]) - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
Linfeng Zhang's avatar
Linfeng Zhang committed
568
    }
569
570
  }
}
571

572
573
574
575
576
577
typedef struct encode_txb_args {
  const AV1_COMMON *cm;
  MACROBLOCK *x;
  aom_writer *w;
} ENCODE_TXB_ARGS;

578
579
580
static void write_coeffs_txb_wrap(const AV1_COMMON *cm, MACROBLOCK *x,
                                  aom_writer *w, int plane, int block,
                                  int blk_row, int blk_col, TX_SIZE tx_size) {
581
582
583
584
585
  MACROBLOCKD *xd = &x->e_mbd;
  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
586
587
  av1_write_coeffs_txb(cm, xd, w, blk_row, blk_col, plane, tx_size, tcoeff, eob,
                       &txb_ctx);
588
589
}

590
591
592
void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x, int mi_row,
                         int mi_col, aom_writer *w, int plane,
                         BLOCK_SIZE bsize) {
Angie Chiang's avatar
Angie Chiang committed
593
  MACROBLOCKD *xd = &x->e_mbd;
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
  const struct macroblockd_plane *const pd = &xd->plane[plane];
  const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
  const int stepr = tx_size_high_unit[tx_size];
  const int stepc = tx_size_wide_unit[tx_size];
  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
  int row, col;
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);

  if (!is_chroma_reference(mi_row, mi_col, bsize, pd->subsampling_x,
                           pd->subsampling_y))
    return;

  int blk_row, blk_col;
  int block = 0;
  const int step = stepr * stepc;
  const BLOCK_SIZE max_unit_bsize = get_plane_block_size(BLOCK_64X64, pd);
  int mu_blocks_wide = block_size_wide[max_unit_bsize] >> tx_size_wide_log2[0];
  int mu_blocks_high = block_size_high[max_unit_bsize] >> tx_size_high_log2[0];
  mu_blocks_wide = AOMMIN(max_blocks_wide, mu_blocks_wide);
  mu_blocks_high = AOMMIN(max_blocks_high, mu_blocks_high);

  for (row = 0; row < max_blocks_high; row += mu_blocks_high) {
    const int unit_height = AOMMIN(mu_blocks_high + row, max_blocks_high);
    for (col = 0; col < max_blocks_wide; col += mu_blocks_wide) {
      const int unit_width = AOMMIN(mu_blocks_wide + col, max_blocks_wide);

      for (blk_row = row; blk_row < unit_height; blk_row += stepr) {
        for (blk_col = col; blk_col < unit_width; blk_col += stepc) {
623
624
          write_coeffs_txb_wrap(cm, x, w, plane, block, blk_row, blk_col,
                                tx_size);
625
626
627
628
629
          block += step;
        }
      }
    }
  }
Angie Chiang's avatar
Angie Chiang committed
630
631
}

632
static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
Angie Chiang's avatar
Angie Chiang committed
633
                              const int *coeff_lps) {
634
635
  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
636
  (void)ctx;
637
638
  if (abs_qc >= min_level) {
    if (abs_qc >= max_level)
639
      return coeff_lps[COEFF_BASE_RANGE];  // COEFF_BASE_RANGE * cost0;
640
    else
641
      return coeff_lps[(abs_qc - min_level)];  //  * cost0 + cost1;
642
643
644
645
646
  } else {
    return 0;
  }
}

647
648
649
650
651
652
653
// 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);
654
  const PLANE_TYPE plane_type = get_plane_type(plane);
655
  const TX_TYPE tx_type =
Luc Trudeau's avatar
Luc Trudeau committed
656
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, tx_size);
657
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
658
659
660
661
  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;
662
  const int txb_skip_ctx = txb_ctx->txb_skip_ctx;
663
664
665
  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
666
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
Linfeng Zhang's avatar
Linfeng Zhang committed
667
  const int16_t *const scan = scan_order->scan;
668
669
  uint8_t levels_buf[TX_PAD_2D];
  uint8_t *const levels = set_levels(levels_buf, width);
670
  DECLARE_ALIGNED(16, uint8_t, level_counts[MAX_TX_SQUARE]);
Linfeng Zhang's avatar
Linfeng Zhang committed
671
  DECLARE_ALIGNED(16, int8_t, coeff_contexts[MAX_TX_SQUARE]);
672
673
  const LV_MAP_COEFF_COST *const coeff_costs =
      &x->coeff_costs[txs_ctx][plane_type];
674

675
676
677
  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];
678
679
  // eob must be greater than 0 here.
  assert(eob > 0);
680
  cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
681

682
  av1_txb_init_levels(qcoeff, width, height, levels);
683

Angie Chiang's avatar
Angie Chiang committed
684
#if CONFIG_TXK_SEL
685
  cost += av1_tx_type_cost(cm, x, xd, mbmi->sb_type, plane, tx_size, tx_type);
Angie Chiang's avatar
Angie Chiang committed
686
#endif
687

688
  const int seg_eob = av1_get_max_eob(tx_size);
689
  int eob_cost = get_eob_cost(eob, seg_eob, eob_costs, coeff_costs, tx_type);
Dake He's avatar
Dake He committed
690
  cost += eob_cost;
Linfeng Zhang's avatar
Linfeng Zhang committed
691
692
693

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

Dake He's avatar
Dake He committed
694
  for (c = eob - 1; c >= 0; --c) {
Linfeng Zhang's avatar
Linfeng Zhang committed
695
696
697
698
    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
699
    const int coeff_ctx = coeff_contexts[pos];
700
    if (c == eob - 1) {
Dake He's avatar
Dake He committed
701
      cost += coeff_costs->base_eob_cost[coeff_ctx][AOMMIN(level, 3) - 1];
702
703
704
    } else {
      cost += coeff_costs->base_cost[coeff_ctx][AOMMIN(level, 3)];
    }
Dake He's avatar
Dake He committed
705
706
707
708
709
710
711
712
713
714
715
716
717

    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);
      }
      if (level > NUM_BASE_LEVELS) {
        int ctx;
718
#if USE_CAUSAL_BR_CTX
Linfeng Zhang's avatar
Linfeng Zhang committed
719
        ctx = get_br_ctx(levels, pos, bwl, level_counts[pos], tx_type);
720
#else
Linfeng Zhang's avatar
Linfeng Zhang committed
721
        ctx = get_br_ctx(levels, pos, bwl, level_counts[pos]);
722
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
723
        const int base_range = level - 1 - NUM_BASE_LEVELS;
Dake He's avatar
Dake He committed
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
        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);
        }
      }
    }
  }
749
750
  return cost;
}
Angie Chiang's avatar
Angie Chiang committed
751

752
753
754
755
756
757
758
759
760
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;
}

761
static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
762
                                    const int (*dc_sign_cost)[2],
763
764
765
766
                                    int dc_sign_ctx) {
  const int sign = (qc < 0) ? 1 : 0;
  // sign bit cost
  if (coeff_idx == 0) {
767
    return dc_sign_cost[dc_sign_ctx][sign];
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
  } 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;
  }
}

790
void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
Angie Chiang's avatar
Angie Chiang committed
791
  // gen_nz_count_arr
Linfeng Zhang's avatar
Linfeng Zhang committed
792
  const int16_t *const scan = txb_info->scan_order->scan;
Angie Chiang's avatar
Angie Chiang committed
793
794
  const int bwl = txb_info->bwl;
  const int height = txb_info->height;
Linfeng Zhang's avatar
Linfeng Zhang committed
795
  const tran_low_t *const qcoeff = txb_info->qcoeff;
796
  const uint8_t *const levels = txb_info->levels;
797
798
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
799
800
801
802
  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
803

Linfeng Zhang's avatar
Linfeng Zhang committed
804
    txb_cache->nz_count_arr[coeff_idx] =
Linfeng Zhang's avatar
Linfeng Zhang committed
805
        get_nz_count(levels + get_padded_idx(coeff_idx, bwl), bwl,
Linfeng Zhang's avatar
Linfeng Zhang committed
806
                     tx_type_to_class[txb_info->tx_type]);
Dake He's avatar
Dake He committed
807

808
809
810
    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
811
812
813
814

    // 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
815
816
    int count[NUM_BASE_LEVELS];
    get_base_count_mag(base_mag, count, qcoeff, bwl, height, row, col);
Angie Chiang's avatar
Angie Chiang committed
817
818

    for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
Angie Chiang's avatar
Angie Chiang committed
819
820
      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
821
      const int level = i + 1;
Angie Chiang's avatar
Angie Chiang committed
822
      txb_cache->base_ctx_arr[i][coeff_idx] =
823
          base_ctx_table[row != 0][col != 0][base_mag[0] > level][count[i]];
Angie Chiang's avatar
Angie Chiang committed
824
825
826
827
828
829
    }

    // 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
830
831
    *br_count = get_br_count_mag(br_mag, qcoeff, bwl, height, row, col,
                                 NUM_BASE_LEVELS);
Angie Chiang's avatar
Angie Chiang committed
832
833
834
    txb_cache->br_ctx_arr[coeff_idx] =
        get_br_ctx_from_count_mag(row, col, *br_count, br_mag[0]);
  }
835
836
}

837
838
839
static INLINE const int *get_level_prob(int level, int coeff_idx,
                                        const TxbCache *txb_cache,
                                        const LV_MAP_COEFF_COST *txb_costs) {
840
841
842
  if (level < 1 + NUM_BASE_LEVELS) {
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
    return &txb_costs->base_cost[ctx][level];
843
844
  } else if (level >= 1 + NUM_BASE_LEVELS &&
             level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
845
    const int ctx = txb_cache->br_ctx_arr[coeff_idx];
846
    return txb_costs->lps_cost[ctx];
847
  } else if (level >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
848
    // printf("get_level_prob does not support golomb\n");
849
850
851
852
853
854
855
856
    assert(0);
    return 0;
  } else {
    assert(0);
    return 0;
  }
}

Angie Chiang's avatar
Angie Chiang committed
857
858
859
860
861
862
863
static INLINE void update_mag_arr(int *mag_arr, int abs_qc) {
  if (mag_arr[0] == abs_qc) {
    mag_arr[1] -= 1;
    assert(mag_arr[1] >= 0);
  }
}

Angie Chiang's avatar
Angie Chiang committed
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
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
906
907
908
909
910
911
912
913
static INLINE int get_mag_from_mag_arr(const int *mag_arr) {
  int mag;
  if (mag_arr[1] > 0) {
    mag = mag_arr[0];
  } else if (mag_arr[0] > 0) {
    mag = mag_arr[0] - 1;
  } else {
    // no neighbor
    assert(mag_arr[0] == 0 && mag_arr[1] == 0);
    mag = 0;
  }
  return mag;
}

static int neighbor_level_down_update(int *new_count, int *new_mag, int count,
                                      const int *mag, int coeff_idx,
                                      tran_low_t abs_nb_coeff, int nb_coeff_idx,
                                      int level, const TxbInfo *txb_info) {
  *new_count = count;
  *new_mag = get_mag_from_mag_arr(mag);

  int update = 0;
  // check if br_count changes
  if (abs_nb_coeff == level) {
    update = 1;
    *new_count -= 1;
    assert(*new_count >= 0);
  }
  const int row = coeff_idx >> txb_info->bwl;
  const int col = coeff_idx - (row << txb_info->bwl);
  const int nb_row = nb_coeff_idx >> txb_info->bwl;
  const int nb_col = nb_coeff_idx - (nb_row << txb_info->bwl);

  // check if mag changes
  if (nb_row >= row && nb_col >= col) {
    if (abs_nb_coeff == mag[0]) {
      assert(mag[1] > 0);
      if (mag[1] == 1) {
        // the nb is the only qc with max mag
        *new_mag -= 1;
        assert(*new_mag >= 0);
        update = 1;
      }
    }
  }
  return update;
}

static int try_neighbor_level_down_br(int coeff_idx, int nb_coeff_idx,
                                      const TxbCache *txb_cache,
914
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
                                      const TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);
  const int level = NUM_BASE_LEVELS + 1;
  if (abs_qc < level) return 0;

  const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
  const tran_low_t abs_nb_coeff = abs(nb_coeff);
  const int count = txb_cache->br_count_arr[coeff_idx];
  const int *mag = txb_cache->br_mag_arr[coeff_idx];
  int new_count;
  int new_mag;
  const int update =
      neighbor_level_down_update(&new_count, &new_mag, count, mag, coeff_idx,
                                 abs_nb_coeff, nb_coeff_idx, level, txb_info);
  if (update) {
    const int row = coeff_idx >> txb_info->bwl;
    const int col = coeff_idx - (row << txb_info->bwl);
933
    const int ctx = txb_cache->br_ctx_arr[coeff_idx];
934
    const int org_cost = get_br_cost(abs_qc, ctx, txb_costs->lps_cost[ctx]);
Angie Chiang's avatar
Angie Chiang committed
935
936

    const int new_ctx = get_br_ctx_from_count_mag(row, col, new_count, new_mag);
937
938
    const int new_cost =
        get_br_cost(abs_qc, new_ctx, txb_costs->lps_cost[new_ctx]);
Angie Chiang's avatar
Angie Chiang committed
939
940
941
942
943
944
945
946
947
    const int cost_diff = -org_cost + new_cost;
    return cost_diff;
  } else {
    return 0;
  }
}

static int try_neighbor_level_down_base(int coeff_idx, int nb_coeff_idx,
                                        const TxbCache *txb_cache,
948
                                        const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
949
                                        const TxbInfo *txb_info) {
950
951
952
953
954
955
956
  // TODO(olah): not implemented yet
  (void)coeff_idx;
  (void)nb_coeff_idx;
  (void)txb_cache;
  (void)txb_costs;
  (void)txb_info;
  return 0;
Angie Chiang's avatar
Angie Chiang committed
957
958
959
960
}

static int try_neighbor_level_down_nz(int coeff_idx, int nb_coeff_idx,
                                      const TxbCache *txb_cache,
961
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
962
963
964
965
966
967
968
                                      TxbInfo *txb_info) {
  // assume eob doesn't change
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);
  const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
  const tran_low_t abs_nb_coeff = abs(nb_coeff);
  if (abs_nb_coeff != 1) return 0;
Linfeng Zhang's avatar
Linfeng Zhang committed
969
  const int16_t *const iscan = txb_info->scan_order->iscan;
Angie Chiang's avatar
Angie Chiang committed
970
971
972
973
974
  const int scan_idx = iscan[coeff_idx];
  if (scan_idx == txb_info->seg_eob) return 0;
  const int nb_scan_idx = iscan[nb_coeff_idx];
  if (nb_scan_idx < scan_idx) {
    const int count = txb_cache->nz_count_arr[coeff_idx];
Linfeng Zhang's avatar
Linfeng Zhang committed
975
    (void)count;
Angie Chiang's avatar
Angie Chiang committed
976
    assert(count > 0);
977
    update_qcoeff(nb_coeff_idx, get_lower_coeff(nb_coeff), txb_info);
Linfeng Zhang's avatar
Linfeng Zhang committed
978
    const int new_ctx = get_nz_map_ctx_from_stats(
979
980
        0, coeff_idx, txb_info->bwl, txb_info->tx_size,
        tx_type_to_class[txb_info->tx_type]);
981
    update_qcoeff(nb_coeff_idx, nb_coeff, txb_info);
982
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
983
984
    const int org_cost = txb_costs->base_cost[ctx][AOMMIN(abs_qc, 3)];
    const int new_cost = txb_costs->base_cost[new_ctx][AOMMIN(abs_qc, 3)];
Angie Chiang's avatar
Angie Chiang committed
985
986
987
988
989
990
991
992
993
    const int cost_diff = new_cost - org_cost;
    return cost_diff;
  } else {
    return 0;
  }
}

static int try_self_level_down(tran_low_t *low_coeff, int coeff_idx,
                               const TxbCache *txb_cache,
994
995
                               const LV_MAP_COEFF_COST *txb_costs,
                               TxbInfo *txb_info) {
996
997
998
999
1000
1001
1002
1003
1004
1005
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  if (qc == 0) {
    *low_coeff = 0;
    return 0;
  }
  const tran_low_t abs_qc = abs(qc);
  *low_coeff = get_lower_coeff(qc);
  int cost_diff;
  if (*low_coeff == 0) {
    const int scan_idx = txb_info->scan_order->iscan[coeff_idx];
1006
1007
1008
1009
    const int *level_cost =
        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_costs);
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
Jingning Han's avatar
Jingning Han committed
1010

Dake He's avatar
Dake He committed
1011
    if (scan_idx < txb_info->eob - 1) {
1012
1013
1014
1015
      // When level-0, we code the binary of abs_qc > level
      // but when level-k k > 0 we code the binary of abs_qc == level
      // That's why wee need this special treatment for level-0 map
      // TODO(angiebird): make leve-0 consistent to other levels
1016
      cost_diff = -level_cost[1] + low_level_cost[0] - low_level_cost[1];
1017
    } else {
1018
      cost_diff = -level_cost[1];
1019
1020
1021
    }

    const int sign_cost = get_sign_bit_cost(
1022
        qc, coeff_idx, txb_costs->dc_sign_cost, txb_info->txb_ctx->dc_sign_ctx);
1023
    cost_diff -= sign_cost;
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
  } else if (abs_qc <= NUM_BASE_LEVELS) {
    const int *level_cost =
        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_costs);
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
    cost_diff = -level_cost[1] + low_level_cost[1] - low_level_cost[0];
  } else if (abs_qc == NUM_BASE_LEVELS + 1) {
    const int *level_cost =
        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_costs);
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
    cost_diff = -level_cost[0] + low_level_cost[1] - low_level_cost[0];
1036
  } else if (abs_qc < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
1037
1038
1039
1040
1041
    const int *level_cost =
        get_level_prob(abs_qc, coeff_idx, txb_cache, txb_costs);
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);

1042
1043
    cost_diff = -level_cost[abs_qc - 1 - NUM_BASE_LEVELS] +
                low_level_cost[abs(*low_coeff) - 1 - NUM_BASE_LEVELS];
1044
  } else if (abs_qc == 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
1045
1046
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
1047
1048
    cost_diff = -get_golomb_cost(abs_qc) - low_level_cost[COEFF_BASE_RANGE] +
                low_level_cost[COEFF_BASE_RANGE - 1];
1049
1050
1051
1052
1053
1054
1055
1056
  } else {
    assert(abs_qc > 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE);
    const tran_low_t abs_low_coeff = abs(*low_coeff);
    cost_diff = -get_golomb_cost(abs_qc) + get_golomb_cost(abs_low_coeff);
  }
  return cost_diff;
}

Angie Chiang's avatar
Angie Chiang committed
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
#define COST_MAP_SIZE 5
#define COST_MAP_OFFSET 2

static INLINE int check_nz_neighbor(tran_low_t qc) { return abs(qc) == 1; }

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

static INLINE int check_br_neighbor(tran_low_t qc) {
  return abs(qc) > BR_MAG_OFFSET;
}

1070
#define FAST_OPTIMIZE_TXB 1
1071
1072
1073

#if FAST_OPTIMIZE_TXB
#define ALNB_REF_OFFSET_NUM 2
1074
static const int alnb_ref_offset[ALNB_REF_OFFSET_NUM][2] = {
1075
1076
1077
  { -1, 0 }, { 0, -1 },
};
#define NB_REF_OFFSET_NUM 4
1078
static const int nb_ref_offset[NB_REF_OFFSET_NUM][2] = {
1079
1080
1081
1082
  { -1, 0 }, { 0, -1 }, { 1, 0 }, { 0, 1 },
};
#endif  // FAST_OPTIMIZE_TXB

Angie Chiang's avatar
Angie Chiang committed
1083
1084
// TODO(angiebird): add static to this function once it's called
int try_level_down(int coeff_idx, const TxbCache *txb_cache,
1085
                   const LV_MAP_COEFF_COST *txb_costs, TxbInfo *txb_info,
1086
1087
1088
1089
                   int (*cost_map)[COST_MAP_SIZE], int fast_mode) {
#if !FAST_OPTIMIZE_TXB
  (void)fast_mode;
#endif
Angie Chiang's avatar
Angie Chiang committed
1090
1091
1092
1093
1094