encodetxb.c 86.6 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"
Angie Chiang's avatar
Angie Chiang committed
21
#include "av1/encoder/rdopt.h"
22
#include "av1/encoder/subexp.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

Dake He's avatar
Dake He committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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;
  int nz_rate;  // for eob
  int64_t rd_diff;
  int cost_diff;
  int64_t dist_diff;
  int new_eob;
} LevelDownStats;

45
void av1_alloc_txb_buf(AV1_COMP *cpi) {
46
#if 0
47
48
49
50
51
52
53
54
55
56
57
58
  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));
  }
59
#else
60
61
62
63
  AV1_COMMON *cm = &cpi->common;
  int size = ((cm->mi_rows >> MAX_MIB_SIZE_LOG2) + 1) *
             ((cm->mi_cols >> MAX_MIB_SIZE_LOG2) + 1);

64
  av1_free_txb_buf(cpi);
65
66
67
  // TODO(jingning): This should be further reduced.
  CHECK_MEM_ERROR(cm, cpi->coeff_buffer_base,
                  aom_malloc(sizeof(*cpi->coeff_buffer_base) * size));
68
#endif
69
70
71
}

void av1_free_txb_buf(AV1_COMP *cpi) {
72
#if 0
73
74
75
76
  int i;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    aom_free(cpi->tcoeff_buf[i]);
  }
77
#else
78
  aom_free(cpi->coeff_buffer_base);
79
#endif
80
81
}

82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
void av1_set_coeff_buffer(const AV1_COMP *const cpi, MACROBLOCK *const x,
                          int mi_row, int mi_col) {
  int stride = (cpi->common.mi_cols >> MAX_MIB_SIZE_LOG2) + 1;
  int offset =
      (mi_row >> MAX_MIB_SIZE_LOG2) * stride + (mi_col >> MAX_MIB_SIZE_LOG2);
  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;
  }
}

99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
static INLINE tran_low_t get_lower_coeff(tran_low_t qc) {
  if (qc == 0) {
    return 0;
  }
  return qc > 0 ? qc - 1 : qc + 1;
}

static INLINE tran_low_t qcoeff_to_dqcoeff(tran_low_t qc, int dqv, int shift) {
  int sgn = qc < 0 ? -1 : 1;
  return sgn * ((abs(qc) * dqv) >> shift);
}

static INLINE int64_t get_coeff_dist(tran_low_t tcoeff, tran_low_t dqcoeff,
                                     int shift) {
129
130
131
132
133
134
135
#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;
  const int64_t error = diff * diff + depth_round >> depth_shift;
  (void)shift;
#else
Dake He's avatar
Dake He committed
136
137
  const int64_t diff = (tcoeff - dqcoeff) * (1 << shift);
  const int64_t error = diff * diff;
138
#endif
Dake He's avatar
Dake He committed
139
140
141
  return error;
}

Jingning Han's avatar
Jingning Han committed
142
143
void av1_update_eob_context(int eob, int seg_eob, TX_SIZE tx_size,
                            TX_TYPE tx_type, PLANE_TYPE plane,
144
145
                            FRAME_CONTEXT *ec_ctx, FRAME_COUNTS *counts,
                            uint8_t allow_update_cdf) {
Dake He's avatar
Dake He committed
146
147
148
149
  int16_t eob_extra;
  int16_t eob_pt = get_eob_pos_token(eob, &eob_extra);
  int16_t dummy;
  int16_t max_eob_pt = get_eob_pos_token(seg_eob, &dummy);
Jingning Han's avatar
Jingning Han committed
150
  TX_SIZE txs_ctx = get_txsize_context(tx_size);
Dake He's avatar
Dake He committed
151
152

  for (int i = 1; i < max_eob_pt; i++) {
Jingning Han's avatar
Jingning Han committed
153
154
    int eob_pos_ctx = av1_get_eob_pos_ctx(tx_type, i);
    counts->eob_flag[txs_ctx][plane][eob_pos_ctx][eob_pt == i]++;
155
156
157
    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
158
159
160
161
    if (eob_pt == i) {
      break;
    }
  }
Jingning Han's avatar
Jingning Han committed
162

Angie Chiang's avatar
Angie Chiang committed
163
164
165
  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
166
    counts->eob_extra[txs_ctx][plane][eob_pt][bit]++;
167
168
    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
169
  }
Dake He's avatar
Dake He committed
170
171
172
}

static int get_eob_cost(int eob, int seg_eob,
Jingning Han's avatar
Jingning Han committed
173
                        const LV_MAP_COEFF_COST *txb_costs, TX_TYPE tx_type) {
Dake He's avatar
Dake He committed
174
175
176
177
178
179
180
181
  int16_t eob_extra;
  int16_t eob_pt = get_eob_pos_token(eob, &eob_extra);
  int16_t dummy;
  int16_t max_eob_pt = get_eob_pos_token(seg_eob, &dummy);
  int eob_cost = 0;

  // printf("Enc: [%d, %d], (%d, %d) ", seg_eob, eob, eob_pt, eob_extra);
  for (int i = 1; i < max_eob_pt; i++) {
Jingning Han's avatar
Jingning Han committed
182
    int eob_pos_ctx = av1_get_eob_pos_ctx(tx_type, i);
Dake He's avatar
Dake He committed
183
184
185
186
187
188
    eob_cost += txb_costs->eob_cost[eob_pos_ctx][eob_pt == i];
    if (eob_pt == i) {
      break;
    }
  }
  if (k_eob_offset_bits[eob_pt] > 0) {
Angie Chiang's avatar
Angie Chiang committed
189
190
191
192
193
194
    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
195
196
197
198
199
200
      eob_cost += av1_cost_bit(128, bit);
    }
  }
  return eob_cost;
}

201
static int get_coeff_cost(const tran_low_t qc, const int scan_idx,
202
203
204
#if CONFIG_LV_MAP_MULTI
                          const int is_eob,
#endif
205
206
                          const TxbInfo *const txb_info,
                          const LV_MAP_COEFF_COST *const txb_costs);
Dake He's avatar
Dake He committed
207

Linfeng Zhang's avatar
Linfeng Zhang committed
208
static void get_dist_cost_stats(LevelDownStats *stats, int scan_idx,
209
210
211
#if CONFIG_LV_MAP_MULTI
                                const int is_eob,
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
212
213
                                const LV_MAP_COEFF_COST *txb_costs,
                                TxbInfo *txb_info) {
Dake He's avatar
Dake He committed
214
215
216
  const int16_t *scan = txb_info->scan_order->scan;
  const int coeff_idx = scan[scan_idx];
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
217
  const uint8_t *const levels = txb_info->levels;
Dake He's avatar
Dake He committed
218
219
220
221
222
223
224
225
  stats->new_eob = -1;
  stats->update = 0;

  const tran_low_t tqc = txb_info->tcoeff[coeff_idx];
  const int dqv = txb_info->dequant[coeff_idx != 0];

  const tran_low_t dqc = qcoeff_to_dqcoeff(qc, dqv, txb_info->shift);
  const int64_t dqc_dist = get_coeff_dist(tqc, dqc, txb_info->shift);
226
227
228
#if CONFIG_LV_MAP_MULTI
  const int qc_cost = get_coeff_cost(qc, scan_idx, is_eob, txb_info, txb_costs);
#else
Dake He's avatar
Dake He committed
229
  const int qc_cost = get_coeff_cost(qc, scan_idx, txb_info, txb_costs);
230
#endif
Dake He's avatar
Dake He committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246

  // distortion difference when coefficient is quantized to 0
  const tran_low_t dqc0 = qcoeff_to_dqcoeff(0, dqv, txb_info->shift);
  stats->dist0 = get_coeff_dist(tqc, dqc0, txb_info->shift);
  stats->dist = dqc_dist - stats->dist0;
  stats->rate = qc_cost;

  if (qc == 0) {
    return;
  }
  stats->rd = RDCOST(txb_info->rdmult, stats->rate, stats->dist);

  stats->low_qc = get_lower_coeff(qc);
  stats->low_dqc = qcoeff_to_dqcoeff(stats->low_qc, dqv, txb_info->shift);
  const int64_t low_dqc_dist =
      get_coeff_dist(tqc, stats->low_dqc, txb_info->shift);
247
248
249
250
#if CONFIG_LV_MAP_MULTI
  const int low_qc_cost =
      get_coeff_cost(stats->low_qc, scan_idx, is_eob, txb_info, txb_costs);
#else
Dake He's avatar
Dake He committed
251
252
  const int low_qc_cost =
      get_coeff_cost(stats->low_qc, scan_idx, txb_info, txb_costs);
253
#endif
Dake He's avatar
Dake He committed
254
255
256
257
  stats->dist_low = low_dqc_dist - stats->dist0;
  stats->rate_low = low_qc_cost;
  stats->rd_low = RDCOST(txb_info->rdmult, stats->rate_low, stats->dist_low);

258
#if CONFIG_LV_MAP_MULTI
259
  (void)levels;
260
  if ((stats->rd_low < stats->rd) && (stats->low_qc == 0)) {
261
    stats->nz_rate = low_qc_cost;
262
  } else {
263
264
265
266
267
268
269
270
271
    if (stats->rd_low < stats->rd) {
      const int low_qc_eob_cost =
          get_coeff_cost(stats->low_qc, scan_idx, 1, txb_info, txb_costs);
      stats->nz_rate = low_qc_cost - low_qc_eob_cost;
    } else {
      const int qc_eob_cost =
          get_coeff_cost(qc, scan_idx, 1, txb_info, txb_costs);
      stats->nz_rate = qc_cost - qc_eob_cost;
    }
272
273
  }
#else
274
  int coeff_ctx = get_nz_map_ctx(levels, scan_idx, scan, txb_info->bwl,
275
                                 txb_info->height, txb_info->tx_type);
Dake He's avatar
Dake He committed
276
277
278
279
280
  if ((stats->rd_low < stats->rd) && (stats->low_qc == 0)) {
    stats->nz_rate = txb_costs->nz_map_cost[coeff_ctx][0];
  } else {
    stats->nz_rate = txb_costs->nz_map_cost[coeff_ctx][1];
  }
281
#endif
Dake He's avatar
Dake He committed
282
283
}

284
285
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
286
  txb_info->qcoeff[coeff_idx] = qc;
287
  txb_info->levels[get_paded_idx(coeff_idx, txb_info->bwl)] =
288
      (uint8_t)clamp(abs(qc), 0, INT8_MAX);
289
290
291
292
293
}

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
294
295
296
297
  const int dqv = txb_info->dequant[coeff_idx != 0];
  txb_info->dqcoeff[coeff_idx] = qcoeff_to_dqcoeff(qc, dqv, txb_info->shift);
}

298
static INLINE void av1_txb_init_levels(const tran_low_t *const coeff,
299
                                       const int width, const int height,
300
301
                                       uint8_t *const levels) {
  const int stride = width + TX_PAD_HOR;
302
  uint8_t *ls = levels;
303
304
305

  memset(levels - TX_PAD_TOP * stride, 0,
         sizeof(*levels) * TX_PAD_TOP * stride);
306
307
  memset(levels + stride * height, 0,
         sizeof(*levels) * (TX_PAD_BOTTOM * stride + TX_PAD_END));
308

309
310
  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
311
      *ls++ = (uint8_t)clamp(abs(coeff[i * width + j]), 0, INT8_MAX);
312
313
314
315
    }
    for (int j = 0; j < TX_PAD_HOR; j++) {
      *ls++ = 0;
    }
316
317
318
  }
}

319
void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
Jingning Han's avatar
Jingning Han committed
320
321
322
                          aom_writer *w, int blk_row, int blk_col, int block,
                          int plane, TX_SIZE tx_size, const tran_low_t *tcoeff,
                          uint16_t eob, TXB_CTX *txb_ctx) {
323
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
324
  const PLANE_TYPE plane_type = get_plane_type(plane);
325
  const TX_SIZE txs_ctx = get_txsize_context(tx_size);
326
327
  const TX_TYPE tx_type =
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, block, tx_size);
Angie Chiang's avatar
Angie Chiang committed
328
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
329
  const int16_t *scan = scan_order->scan;
330
  const int seg_eob = tx_size_2d[tx_size];
331
  int c;
332
  const int bwl = tx_size_wide_log2[tx_size];
333
  const int width = tx_size_wide[tx_size];
334
  const int height = tx_size_high[tx_size];
335
  int update_eob = -1;
336
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
337
338
  uint8_t levels_buf[TX_PAD_2D];
  uint8_t *const levels = set_levels(levels_buf, width);
339
  DECLARE_ALIGNED(16, uint8_t, level_counts[MAX_TX_SQUARE]);
340

Jingning Han's avatar
Jingning Han committed
341
342
  (void)blk_row;
  (void)blk_col;
343
344
  aom_write_bin(w, eob == 0,
                ec_ctx->txb_skip_cdf[txs_ctx][txb_ctx->txb_skip_ctx], 2);
345
  if (eob == 0) return;
346

347
  av1_txb_init_levels(tcoeff, width, height, levels);
348

Angie Chiang's avatar
Angie Chiang committed
349
#if CONFIG_TXK_SEL
350
351
  av1_write_tx_type(cm, xd, blk_row, blk_col, block, plane,
                    get_min_tx_size(tx_size), w);
Angie Chiang's avatar
Angie Chiang committed
352
#endif
353

Dake He's avatar
Dake He committed
354
355
356
357
358
359
360
  int16_t eob_extra;
  int16_t eob_pt = get_eob_pos_token(eob, &eob_extra);
  int16_t dummy;
  int16_t max_eob_pt = get_eob_pos_token(seg_eob, &dummy);

  // printf("Enc: [%d, %d], (%d, %d) ", seg_eob, eob, eob_pt, eob_extra);
  for (int i = 1; i < max_eob_pt; i++) {
Jingning Han's avatar
Jingning Han committed
361
    int eob_pos_ctx = av1_get_eob_pos_ctx(tx_type, i);
Dake He's avatar
Dake He committed
362
363
364
365
366
367
368
369
370
371
372

    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;
    }
  }

  if (k_eob_offset_bits[eob_pt] > 0) {
Angie Chiang's avatar
Angie Chiang committed
373
374
375
376
377
378
379
    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
380
381
382
383
      aom_write_bit(w, bit);
      // printf("%d ", bit);
    }
  }
384
385
386
387
388
389

#if USE_CAUSAL_BASE_CTX
  int coeff_ctx = 0;
  for (int i = 0; i < eob; ++i) {
    c = eob - 1 - i;

390
391
392
393
#if CONFIG_LV_MAP_MULTI
    coeff_ctx =
        get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, c == eob - 1);
    tran_low_t v = tcoeff[scan[c]];
394
395
396
397
398
399
400
401
402
403
404
405
406
407
#if USE_BASE_EOB_ALPHABET
    if (c == eob - 1) {
      aom_write_symbol(
          w, AOMMIN(abs(v), 3) - 1,
          ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type]
                                    [coeff_ctx - SIG_COEF_CONTEXTS +
                                     SIG_COEF_CONTEXTS_EOB],
          3);
    } else {
      aom_write_symbol(w, AOMMIN(abs(v), 3),
                       ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx],
                       4);
    }
#else
408
409
    aom_write_symbol(w, AOMMIN(abs(v), 3),
                     ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx], 4);
410
#endif
411
412
#else
    coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
    tran_low_t v = tcoeff[scan[c]];
    int is_nz = (v != 0);

    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;
      }
    }
431
#endif
432
433
434
  }
  update_eob = eob - 1;
#else
Dake He's avatar
Dake He committed
435
436
  for (int i = 1; i < eob; ++i) {
    c = eob - 1 - i;
437
    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
Dake He's avatar
Dake He committed
438
439
440
441
442
443
444
445
446

    tran_low_t v = tcoeff[scan[c]];
    int is_nz = (v != 0);

    aom_write_bin(w, is_nz, ec_ctx->nz_map_cdf[txs_ctx][plane_type][coeff_ctx],
                  2);
  }

  for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
447
    av1_get_base_level_counts(levels, i, width, height, level_counts);
448
    for (c = eob - 1; c >= 0; --c) {
449
      const tran_low_t level = abs(tcoeff[scan[c]]);
450
451
452
453
      int ctx;

      if (level <= i) continue;

454
      ctx = get_base_ctx(levels, scan[c], bwl, i, level_counts[scan[c]]);
455
456

      if (level == i + 1) {
457
458
        aom_write_bin(w, 1, ec_ctx->coeff_base_cdf[txs_ctx][plane_type][i][ctx],
                      2);
459
460
        continue;
      }
461

462
463
      aom_write_bin(w, 0, ec_ctx->coeff_base_cdf[txs_ctx][plane_type][i][ctx],
                    2);
464
465
466
      update_eob = AOMMAX(update_eob, c);
    }
  }
467
#endif
468

469
470
471
  // Loop to code all signs in the transform block,
  // starting with the sign of DC (if applicable)
  for (c = 0; c < eob; ++c) {
472
473
474
    const tran_low_t v = tcoeff[scan[c]];
    const tran_low_t level = abs(v);
    const int sign = (v < 0) ? 1 : 0;
475
    if (level == 0) continue;
476
477

    if (c == 0) {
478
#if LV_MAP_PROB
479
480
      aom_write_bin(w, sign,
                    ec_ctx->dc_sign_cdf[plane_type][txb_ctx->dc_sign_ctx], 2);
481
482
483
#else
      aom_write(w, sign, ec_ctx->dc_sign[plane_type][txb_ctx->dc_sign_ctx]);
#endif
484
485
486
    } else {
      aom_write_bit(w, sign);
    }
487
488
  }

Linfeng Zhang's avatar
Linfeng Zhang committed
489
490
491
492
493
494
  if (update_eob >= 0) {
    av1_get_br_level_counts(levels, width, height, level_counts);
    for (c = update_eob; c >= 0; --c) {
      const tran_low_t level = abs(tcoeff[scan[c]]);
      int idx;
      int ctx;
495

Linfeng Zhang's avatar
Linfeng Zhang committed
496
497
498
499
500
501
      if (level <= NUM_BASE_LEVELS) continue;

      // level is above 1.
      ctx = get_br_ctx(levels, scan[c], bwl, level_counts[scan[c]]);

      int base_range = level - 1 - NUM_BASE_LEVELS;
502
503
504
#if CONFIG_LV_MAP_MULTI
      for (idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {
        int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);
505
506
        aom_write_symbol(w, k, ec_ctx->coeff_br_cdf[txs_ctx][plane_type][ctx],
                         BR_CDF_SIZE);
507
508
509
510
        if (k < BR_CDF_SIZE - 1) break;
      }
      if (base_range < COEFF_BASE_RANGE) continue;
#else
Linfeng Zhang's avatar
Linfeng Zhang committed
511
512
513
514
515
516
517
518
519
520
      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) {
521
        // printf("br: %d %d %d %d\n", txs_ctx, plane_type, idx, ctx);
Linfeng Zhang's avatar
Linfeng Zhang committed
522
523
524
525
526
527
528
529
530
531
532
533
534
        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],
535
                          2);
536
          }
Linfeng Zhang's avatar
Linfeng Zhang committed
537
538
          //        aom_write_literal(w, br_offset, br_extra_bits[idx]);
          break;
539
540
541
        }
      }

Linfeng Zhang's avatar
Linfeng Zhang committed
542
      if (br_set_idx < BASE_RANGE_SETS) continue;
543
#endif
Linfeng Zhang's avatar
Linfeng Zhang committed
544
545
546
547
      // use 0-th order Golomb code to handle the residual level.
      write_golomb(
          w, abs(tcoeff[scan[c]]) - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
    }
548
549
  }
}
550

Angie Chiang's avatar
Angie Chiang committed
551
552
553
554
555
556
void av1_write_coeffs_mb(const AV1_COMMON *const cm, MACROBLOCK *x,
                         aom_writer *w, int plane) {
  MACROBLOCKD *xd = &x->e_mbd;
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
  BLOCK_SIZE bsize = mbmi->sb_type;
  struct macroblockd_plane *pd = &xd->plane[plane];
557
558
  const BLOCK_SIZE plane_bsize =
      AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
559
560
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
561
  const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
Angie Chiang's avatar
Angie Chiang committed
562
563
564
565
566
  const int bkw = tx_size_wide_unit[tx_size];
  const int bkh = tx_size_high_unit[tx_size];
  const int step = tx_size_wide_unit[tx_size] * tx_size_high_unit[tx_size];
  int row, col;
  int block = 0;
567
568
  for (row = 0; row < max_blocks_high; row += bkh) {
    for (col = 0; col < max_blocks_wide; col += bkw) {
Angie Chiang's avatar
Angie Chiang committed
569
570
571
572
      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] };
Jingning Han's avatar
Jingning Han committed
573
574
      av1_write_coeffs_txb(cm, xd, w, row, col, block, plane, tx_size, tcoeff,
                           eob, &txb_ctx);
Angie Chiang's avatar
Angie Chiang committed
575
576
577
578
579
      block += step;
    }
  }
}

580
#if !USE_CAUSAL_BASE_CTX
581
582
static INLINE void get_base_ctx_set(const uint8_t *const levels,
                                    const int c,  // raster order
583
                                    const int bwl,
584
585
586
                                    int ctx_set[NUM_BASE_LEVELS]) {
  const int row = c >> bwl;
  const int col = c - (row << bwl);
587
588
  const int width = 1 << bwl;
  const int stride = width + TX_PAD_HOR;
589
590
  int mag_count[NUM_BASE_LEVELS] = { 0 };
  int nb_mag[NUM_BASE_LEVELS][3] = { { 0 } };
591
592
593
594
  int idx;
  int i;

  for (idx = 0; idx < BASE_CONTEXT_POSITION_NUM; ++idx) {
595
596
597
    const int ref_row = row + base_ref_offset[idx][0];
    const int ref_col = col + base_ref_offset[idx][1];
    const int pos = ref_row * stride + ref_col;
598
    const uint8_t abs_coeff = levels[pos];
599
600
601

    for (i = 0; i < NUM_BASE_LEVELS; ++i) {
      ctx_set[i] += abs_coeff > i;
602
603
604
605
606
607
      if (base_ref_offset[idx][0] == 0 && base_ref_offset[idx][1] == 1)
        nb_mag[i][0] = abs_coeff;
      if (base_ref_offset[idx][0] == 1 && base_ref_offset[idx][1] == 0)
        nb_mag[i][1] = abs_coeff;
      if (base_ref_offset[idx][0] == 1 && base_ref_offset[idx][1] == 1)
        nb_mag[i][2] = abs_coeff;
608
609
610
611
    }
  }

  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
612
613
614
    for (idx = 0; idx < 3; ++idx) mag_count[i] += nb_mag[i][idx] > i + 1;
    ctx_set[i] = get_base_ctx_from_count_mag(row, col, ctx_set[i],
                                             AOMMIN(2, mag_count[i]));
615
616
617
  }
  return;
}
618
#endif
619

620
static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
Angie Chiang's avatar
Angie Chiang committed
621
                              const int *coeff_lps) {
622
623
  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
624
  (void)ctx;
625
626
  if (abs_qc >= min_level) {
    if (abs_qc >= max_level)
627
      return coeff_lps[COEFF_BASE_RANGE];  // COEFF_BASE_RANGE * cost0;
628
    else
629
      return coeff_lps[(abs_qc - min_level)];  //  * cost0 + cost1;
630
631
632
633
634
635
  } else {
    return 0;
  }
}

static INLINE int get_base_cost(tran_low_t abs_qc, int ctx,
636
                                const int coeff_base[2], int base_idx) {
637
  const int level = base_idx + 1;
638
  (void)ctx;
639
640
641
  if (abs_qc < level)
    return 0;
  else
642
    return coeff_base[abs_qc == level];
643
644
}

645
int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
Jingning Han's avatar
Jingning Han committed
646
647
                        int blk_row, int blk_col, int block, TX_SIZE tx_size,
                        TXB_CTX *txb_ctx) {
648
  MACROBLOCKD *const xd = &x->e_mbd;
649
  TX_SIZE txs_ctx = get_txsize_context(tx_size);
650
  const PLANE_TYPE plane_type = get_plane_type(plane);
651
652
  const TX_TYPE tx_type =
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, block, tx_size);
653
654
655
656
657
658
659
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
  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;
  int txb_skip_ctx = txb_ctx->txb_skip_ctx;

660
  const int bwl = tx_size_wide_log2[tx_size];
661
  const int width = tx_size_wide[tx_size];
662
663
  const int height = tx_size_high[tx_size];

Angie Chiang's avatar
Angie Chiang committed
664
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
665
  const int16_t *scan = scan_order->scan;
666
667
  uint8_t levels_buf[TX_PAD_2D];
  uint8_t *const levels = set_levels(levels_buf, width);
668
  DECLARE_ALIGNED(16, uint8_t, level_counts[MAX_TX_SQUARE]);
669

670
671
  LV_MAP_COEFF_COST *coeff_costs = &x->coeff_costs[txs_ctx][plane_type];

672
673
674
  cost = 0;

  if (eob == 0) {
675
    cost = coeff_costs->txb_skip_cost[txb_skip_ctx][1];
676
677
    return cost;
  }
678
  cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
679

680
  av1_txb_init_levels(qcoeff, width, height, levels);
681

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

Dake He's avatar
Dake He committed
686
  const int seg_eob = tx_size_2d[tx_size];
Jingning Han's avatar
Jingning Han committed
687
  int eob_cost = get_eob_cost(eob, seg_eob, coeff_costs, tx_type);
Dake He's avatar
Dake He committed
688

Linfeng Zhang's avatar
Linfeng Zhang committed
689
  av1_get_br_level_counts(levels, width, height, level_counts);
Dake He's avatar
Dake He committed
690
  cost += eob_cost;
691
692
693
#if USE_CAUSAL_BASE_CTX
  int coeff_ctx = 0;
#endif
Dake He's avatar
Dake He committed
694
695
696
697
  for (c = eob - 1; c >= 0; --c) {
    tran_low_t v = qcoeff[scan[c]];
    int is_nz = (v != 0);
    int level = abs(v);
698
699
700
#if CONFIG_LV_MAP_MULTI
    coeff_ctx =
        get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, c == eob - 1);
701
702
703
704
705
706
707
708
709
#if USE_BASE_EOB_ALPHABET
    if (c == eob - 1) {
      cost += coeff_costs
                  ->base_eob_cost[coeff_ctx - SIG_COEF_CONTEXTS +
                                  SIG_COEF_CONTEXTS_EOB][AOMMIN(level, 3) - 1];
    } else {
      cost += coeff_costs->base_cost[coeff_ctx][AOMMIN(level, 3)];
    }
#else
710
    cost += coeff_costs->base_cost[coeff_ctx][AOMMIN(level, 3)];
711
#endif
712
#else  // CONFIG_LV_MAP_MULTI
713
#if USE_CAUSAL_BASE_CTX
714
    coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
715
#endif
Dake He's avatar
Dake He committed
716
717

    if (c < eob - 1) {
718
#if !USE_CAUSAL_BASE_CTX
719
      int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
720
#endif
Dake He's avatar
Dake He committed
721
722
      cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
    }
723
#endif  // CONFIG_LV_MAP_MULTI
Dake He's avatar
Dake He committed
724
725

    if (is_nz) {
726
#if !USE_CAUSAL_BASE_CTX
Dake He's avatar
Dake He committed
727
      int ctx_ls[NUM_BASE_LEVELS] = { 0 };
728
#endif
Dake He's avatar
Dake He committed
729
730
731
732
733
734
735
736
737
      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);
      }
738
#if !CONFIG_LV_MAP_MULTI
739
740
741
742
743
744
745
746
747
748
#if USE_CAUSAL_BASE_CTX
      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;
      }
#else
749
      get_base_ctx_set(levels, scan[c], bwl, ctx_ls);
Dake He's avatar
Dake He committed
750
751
752
753
754
755
756
757
758
759
760

      int i;
      for (i = 0; i < NUM_BASE_LEVELS; ++i) {
        if (level <= i) continue;

        if (level == i + 1) {
          cost += coeff_costs->base_cost[i][ctx_ls[i]][1];
          continue;
        }
        cost += coeff_costs->base_cost[i][ctx_ls[i]][0];
      }
761
#endif
762
#endif  // CONFIG_LV_MAP_MULTI
Dake He's avatar
Dake He committed
763
764
      if (level > NUM_BASE_LEVELS) {
        int ctx;
Linfeng Zhang's avatar
Linfeng Zhang committed
765
        ctx = get_br_ctx(levels, scan[c], bwl, level_counts[scan[c]]);
Dake He's avatar
Dake He committed
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
        int base_range = level - 1 - NUM_BASE_LEVELS;
        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);
        }
      }
    }
  }
792
793
  return cost;
}
Angie Chiang's avatar
Angie Chiang committed
794

795
796
797
798
799
800
801
802
803
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;
}

804
static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
805
                                    const int (*dc_sign_cost)[2],
806
807
808
809
                                    int dc_sign_ctx) {
  const int sign = (qc < 0) ? 1 : 0;
  // sign bit cost
  if (coeff_idx == 0) {
810
    return dc_sign_cost[dc_sign_ctx][sign];
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
  } 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;
  }
}

833
void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
Angie Chiang's avatar
Angie Chiang committed
834
  // gen_nz_count_arr
835
  const int16_t *scan = txb_info->scan_order->scan;
Angie Chiang's avatar
Angie Chiang committed
836
837
  const int bwl = txb_info->bwl;
  const int height = txb_info->height;
Linfeng Zhang's avatar
Linfeng Zhang committed
838
  const tran_low_t *const qcoeff = txb_info->qcoeff;
839
  const uint8_t *const levels = txb_info->levels;
840
841
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
842
843
844
845
  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
846

847
848
    txb_cache->nz_count_arr[coeff_idx] =
        get_nz_count(levels, bwl, row, col, get_tx_class(txb_info->tx_type));
Dake He's avatar
Dake He committed
849

Angie Chiang's avatar
Angie Chiang committed
850
    const int nz_count = txb_cache->nz_count_arr[coeff_idx];
851
852
853
854
855
856
    txb_cache->nz_ctx_arr[coeff_idx] =
        get_nz_map_ctx_from_count(nz_count, coeff_idx, bwl, height,
#if USE_CAUSAL_BASE_CTX
                                  0,
#endif
                                  txb_info->tx_type);
Angie Chiang's avatar
Angie Chiang committed
857
858
859
860

    // 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
861
862
    int count[NUM_BASE_LEVELS];
    get_base_count_mag(base_mag, count, qcoeff, bwl, height, row, col);
Angie Chiang's avatar
Angie Chiang committed
863
864

    for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
Angie Chiang's avatar
Angie Chiang committed
865
866
      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
867
      const int level = i + 1;
Angie Chiang's avatar
Angie Chiang committed
868
      txb_cache->base_ctx_arr[i][coeff_idx] =
869
          base_ctx_table[row != 0][col != 0][base_mag[0] > level][count[i]];
Angie Chiang's avatar
Angie Chiang committed
870
871
872
873
874
875
    }

    // 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
876
877
    *br_count = get_br_count_mag(br_mag, qcoeff, bwl, height, row, col,
                                 NUM_BASE_LEVELS);
Angie Chiang's avatar
Angie Chiang committed
878
879
880
    txb_cache->br_ctx_arr[coeff_idx] =
        get_br_ctx_from_count_mag(row, col, *br_count, br_mag[0]);
  }
881
882
}

883
884
885
static INLINE const int *get_level_prob(int level, int coeff_idx,
                                        const TxbCache *txb_cache,
                                        const LV_MAP_COEFF_COST *txb_costs) {
886
887
888
889
890
#if CONFIG_LV_MAP_MULTI
  if (level < 1 + NUM_BASE_LEVELS) {
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
    return &txb_costs->base_cost[ctx][level];
#else
891
  if (level == 0) {
892
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
893
    return txb_costs->nz_map_cost[ctx];
894
895
  } else if (level >= 1 && level < 1 + NUM_BASE_LEVELS) {
    const int idx = level - 1;
896
    const int ctx = txb_cache->base_ctx_arr[idx][coeff_idx];
897
    return txb_costs->base_cost[idx][ctx];
898
#endif
899
900
  } else if (level >= 1 + NUM_BASE_LEVELS &&
             level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
901
    const int ctx = txb_cache->br_ctx_arr[coeff_idx];
902
    return txb_costs->lps_cost[ctx];
903
  } else if (level >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
904
    // printf("get_level_prob does not support golomb\n");
905
906
907
908
909
910
911
912
    assert(0);
    return 0;
  } else {
    assert(0);
    return 0;
  }
}

Angie Chiang's avatar
Angie Chiang committed
913
914
915
916
917
918
919
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
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
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,
970
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
                                      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);
989
    const int ctx = txb_cache->br_ctx_arr[coeff_idx];
990
    const int org_cost = get_br_cost(abs_qc, ctx, txb_costs->lps_cost[ctx]);
Angie Chiang's avatar
Angie Chiang committed
991
992

    const int new_ctx = get_br_ctx_from_count_mag(row, col, new_count, new_mag);
993
994
    const int new_cost =
        get_br_cost(abs_qc, new_ctx, txb_costs->lps_cost[new_ctx]);
Angie Chiang's avatar
Angie Chiang committed
995
996
997
998
999
1000
1001
1002
1003
    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,
1004
                                        const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
1005
                                        const TxbInfo *txb_info) {
1006
1007
1008
1009
1010
1011
1012
1013
1014
#if CONFIG_LV_MAP_MULTI
  // TODO(olah): not implemented yet
  (void)coeff_idx;
  (void)nb_coeff_idx;
  (void)txb_cache;
  (void)txb_costs;
  (void)txb_info;
  return 0;
#else
Angie Chiang's avatar
Angie Chiang committed
1015
1016
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);
1017
1018
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037

  int cost_diff = 0;
  for (int base_idx = 0; base_idx < NUM_BASE_LEVELS; ++base_idx) {
    const int level = base_idx + 1;
    if (abs_qc < level) continue;

    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->base_count_arr[base_idx][coeff_idx];
    const int *mag = txb_cache->base_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);
1038
      const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx];
1039
1040
      const int org_cost = get_base_cost(
          abs_qc, ctx, txb_costs->base_cost[base_idx][ctx], base_idx);
Angie Chiang's avatar
Angie Chiang committed
1041
1042

      const int new_ctx =
1043
          base_ctx_table[row != 0][col != 0][new_mag > level][new_count];
1044
1045
      const int new_cost = get_base_cost(
          abs_qc, new_ctx, txb_costs->base_cost[base_idx][new_ctx], base_idx);
Angie Chiang's avatar
Angie Chiang committed
1046
1047
1048
1049
      cost_diff += -org_cost + new_cost;
    }
  }
  return cost_diff;
1050
#endif
Angie Chiang's avatar
Angie Chiang committed
1051
1052
1053
1054
}

static int try_neighbor_level_down_nz(int coeff_idx, int nb_coeff_idx,
                                      const TxbCache *txb_cache,
1055
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
                                      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;
  const int16_t *iscan = txb_info->scan_order->iscan;
  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];
    assert(count > 0);
1070
    update_qcoeff(nb_coeff_idx, get_lower_coeff(nb_coeff), txb_info);
1071
1072
1073
1074
1075
1076
    const int new_ctx = get_nz_map_ctx_from_count(
        count - 1, coeff_idx, txb_info->bwl, txb_info->height,
#if USE_CAUSAL_BASE_CTX
        0,
#endif
        txb_info->tx_type);
1077
    update_qcoeff(nb_coeff_idx, nb_coeff, txb_info);
1078
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
1079
1080
1081
1082
#if CONFIG_LV_MAP_MULTI
    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)];
#else
Angie Chiang's avatar
Angie Chiang committed
1083
    const int is_nz = abs_qc > 0;
1084
1085
    const int org_cost = txb_costs->nz_map_cost[ctx][is_nz];
    const int new_cost = txb_costs->nz_map_cost[new_ctx][is_nz];
1086
#endif
Angie Chiang's avatar
Angie Chiang committed
1087
1088
1089
1090
1091
1092
1093
1094
1095
    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,
1096
1097
                               const LV_MAP_COEFF_COST *txb_costs,
                               TxbInfo *txb_info) {
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
  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];
1108
1109
1110
1111
    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
1112

Dake He's avatar
Dake He committed
1113
    if (scan_idx < txb_info->eob - 1) {
1114
1115
1116
1117
      // 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
1118
      cost_diff = -level_cost[1] + low_level_cost[0] - low_level_cost[1];
1119
    } else {
1120
      cost_diff = -level_cost[1];
1121
1122
1123
    }

    const int sign_cost = get_sign_bit_cost(
1124
        qc, coeff_idx, txb_costs->dc_sign_cost, txb_info->txb_ctx->dc_sign_ctx);
1125
    cost_diff -= sign_cost;
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
  } 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