encodetxb.c 71.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
/*
 * 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.
 */

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

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

26
void av1_alloc_txb_buf(AV1_COMP *cpi) {
27
#if 0
28
29
30
31
32
33
34
35
36
37
38
39
  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));
  }
40
#else
41
42
43
44
45
46
47
  AV1_COMMON *cm = &cpi->common;
  int size = ((cm->mi_rows >> MAX_MIB_SIZE_LOG2) + 1) *
             ((cm->mi_cols >> MAX_MIB_SIZE_LOG2) + 1);

  // TODO(jingning): This should be further reduced.
  CHECK_MEM_ERROR(cm, cpi->coeff_buffer_base,
                  aom_malloc(sizeof(*cpi->coeff_buffer_base) * size));
48
#endif
49
50
51
}

void av1_free_txb_buf(AV1_COMP *cpi) {
52
#if 0
53
54
55
56
  int i;
  for (i = 0; i < MAX_MB_PLANE; ++i) {
    aom_free(cpi->tcoeff_buf[i]);
  }
57
#else
58
  aom_free(cpi->coeff_buffer_base);
59
#endif
60
61
}

62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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;
  }
}

79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
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);
}

void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
Jingning Han's avatar
Jingning Han committed
96
97
98
                          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) {
99
#if !LV_MAP_PROB
100
101
  aom_prob *nz_map;
  aom_prob *eob_flag;
102
#endif
103
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
104
  const PLANE_TYPE plane_type = get_plane_type(plane);
105
  const TX_SIZE txs_ctx = get_txsize_context(tx_size);
106
107
  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
108
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
109
110
111
112
  const int16_t *scan = scan_order->scan;
  int c;
  int is_nz;
  const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
113
  const int height = tx_size_high[tx_size];
114
  const int seg_eob = tx_size_2d[tx_size];
115
  uint16_t update_eob = 0;
116
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
117

Jingning Han's avatar
Jingning Han committed
118
119
120
  (void)blk_row;
  (void)blk_col;

121
122
#if LV_MAP_PROB
  aom_write_symbol(w, eob == 0,
123
                   ec_ctx->txb_skip_cdf[txs_ctx][txb_ctx->txb_skip_ctx], 2);
124
#else
125
  aom_write(w, eob == 0, ec_ctx->txb_skip[txs_ctx][txb_ctx->txb_skip_ctx]);
126
#endif
127
128

  if (eob == 0) return;
Angie Chiang's avatar
Angie Chiang committed
129
#if CONFIG_TXK_SEL
130
131
  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
132
#endif
133

134
#if !LV_MAP_PROB
135
136
  nz_map = ec_ctx->nz_map[txs_ctx][plane_type];
  eob_flag = ec_ctx->eob_flag[txs_ctx][plane_type];
137
#endif
138
139

  for (c = 0; c < eob; ++c) {
Angie Chiang's avatar
Angie Chiang committed
140
    int coeff_ctx = get_nz_map_ctx(tcoeff, scan[c], bwl, height, tx_type);
141
    int eob_ctx = get_eob_ctx(tcoeff, scan[c], txs_ctx);
142
143
144
145
146
147

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

    if (c == seg_eob - 1) break;

148
149
#if LV_MAP_PROB
    aom_write_symbol(w, is_nz,
150
                     ec_ctx->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], 2);
151
#else
152
    aom_write(w, is_nz, nz_map[coeff_ctx]);
153
#endif
154
155

    if (is_nz) {
156
157
#if LV_MAP_PROB
      aom_write_symbol(w, c == (eob - 1),
158
                       ec_ctx->eob_flag_cdf[txs_ctx][plane_type][eob_ctx], 2);
159
#else
160
      aom_write(w, c == (eob - 1), eob_flag[eob_ctx]);
161
#endif
162
163
164
165
166
    }
  }

  int i;
  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
167
#if !LV_MAP_PROB
168
    aom_prob *coeff_base = ec_ctx->coeff_base[txs_ctx][plane_type][i];
169
#endif
170
171
172
173
174
175
176
177
178
    update_eob = 0;
    for (c = eob - 1; c >= 0; --c) {
      tran_low_t v = tcoeff[scan[c]];
      tran_low_t level = abs(v);
      int sign = (v < 0) ? 1 : 0;
      int ctx;

      if (level <= i) continue;

179
      ctx = get_base_ctx(tcoeff, scan[c], bwl, height, i + 1);
180
181

      if (level == i + 1) {
182
183
#if LV_MAP_PROB
        aom_write_symbol(
184
            w, 1, ec_ctx->coeff_base_cdf[txs_ctx][plane_type][i][ctx], 2);
185
#else
186
        aom_write(w, 1, coeff_base[ctx]);
187
#endif
188
        if (c == 0) {
189
190
#if LV_MAP_PROB
          aom_write_symbol(
191
              w, sign, ec_ctx->dc_sign_cdf[plane_type][txb_ctx->dc_sign_ctx],
192
193
              2);
#else
194
          aom_write(w, sign, ec_ctx->dc_sign[plane_type][txb_ctx->dc_sign_ctx]);
195
#endif
196
197
198
199
200
        } else {
          aom_write_bit(w, sign);
        }
        continue;
      }
201
202
203

#if LV_MAP_PROB
      aom_write_symbol(w, 0,
204
                       ec_ctx->coeff_base_cdf[txs_ctx][plane_type][i][ctx], 2);
205
#else
206
      aom_write(w, 0, coeff_base[ctx]);
207
#endif
208
209
210
211
212
213
214
215
216
217
218
219
220
221
      update_eob = AOMMAX(update_eob, c);
    }
  }

  for (c = update_eob; c >= 0; --c) {
    tran_low_t v = tcoeff[scan[c]];
    tran_low_t level = abs(v);
    int sign = (v < 0) ? 1 : 0;
    int idx;
    int ctx;

    if (level <= NUM_BASE_LEVELS) continue;

    if (c == 0) {
222
223
#if LV_MAP_PROB
      aom_write_symbol(
224
          w, sign, ec_ctx->dc_sign_cdf[plane_type][txb_ctx->dc_sign_ctx], 2);
225
#else
226
      aom_write(w, sign, ec_ctx->dc_sign[plane_type][txb_ctx->dc_sign_ctx]);
227
#endif
228
229
230
231
232
    } else {
      aom_write_bit(w, sign);
    }

    // level is above 1.
233
    ctx = get_br_ctx(tcoeff, scan[c], bwl, height);
234
235
236
237
238
239
240
241
242
243
244
245
246
247

#if BR_NODE
    int base_range = level - 1 - NUM_BASE_LEVELS;
    int br_set_idx = 0;
    int br_base = 0;
    int br_offset = 0;

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

    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
      aom_write_symbol(w, idx == br_set_idx,
248
                       ec_ctx->coeff_br_cdf[txs_ctx][plane_type][idx][ctx], 2);
249
250
251
252
253
254
255
      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_symbol(
256
                w, 1, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2);
257
258
259
            break;
          }
          aom_write_symbol(w, 0,
260
                           ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2);
261
262
263
264
265
266
267
268
        }
        //        aom_write_literal(w, br_offset, br_extra_bits[idx]);
        break;
      }
    }

    if (br_set_idx < BASE_RANGE_SETS) continue;
#else  // BR_NODE
269
270
    for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
      if (level == (idx + 1 + NUM_BASE_LEVELS)) {
271
#if LV_MAP_PROB
272
        aom_write_symbol(w, 1, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
273
274
                         2);
#else
275
        aom_write(w, 1, ec_ctx->coeff_lps[txs_ctx][plane_type][ctx]);
276
#endif
277
278
        break;
      }
279
#if LV_MAP_PROB
280
      aom_write_symbol(w, 0, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
281
282
                       2);
#else
283
      aom_write(w, 0, ec_ctx->coeff_lps[txs_ctx][plane_type][ctx]);
284
#endif
285
286
    }
    if (idx < COEFF_BASE_RANGE) continue;
287
#endif  // BR_NODE
288
289
290
291
292

    // use 0-th order Golomb code to handle the residual level.
    write_golomb(w, level - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
  }
}
293

Angie Chiang's avatar
Angie Chiang committed
294
295
296
297
298
299
300
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];

301
#if CONFIG_CHROMA_SUB8X8
302
303
  const BLOCK_SIZE plane_bsize =
      AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
304
305
#elif CONFIG_CB4X4
  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
Angie Chiang's avatar
Angie Chiang committed
306
307
308
309
#else
  const BLOCK_SIZE plane_bsize =
      get_plane_block_size(AOMMAX(bsize, BLOCK_8X8), pd);
#endif
310
311
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
312
  const TX_SIZE tx_size = av1_get_tx_size(plane, xd);
Angie Chiang's avatar
Angie Chiang committed
313
314
315
316
317
  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;
318
319
  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
320
321
322
323
      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
324
325
      av1_write_coeffs_txb(cm, xd, w, row, col, block, plane, tx_size, tcoeff,
                           eob, &txb_ctx);
Angie Chiang's avatar
Angie Chiang committed
326
327
328
329
330
      block += step;
    }
  }
}

331
332
static INLINE void get_base_ctx_set(const tran_low_t *tcoeffs,
                                    int c,  // raster order
333
                                    const int bwl, const int height,
334
335
336
337
338
339
340
341
342
343
344
345
346
347
                                    int ctx_set[NUM_BASE_LEVELS]) {
  const int row = c >> bwl;
  const int col = c - (row << bwl);
  const int stride = 1 << bwl;
  int mag[NUM_BASE_LEVELS] = { 0 };
  int idx;
  tran_low_t abs_coeff;
  int i;

  for (idx = 0; idx < BASE_CONTEXT_POSITION_NUM; ++idx) {
    int ref_row = row + base_ref_offset[idx][0];
    int ref_col = col + base_ref_offset[idx][1];
    int pos = (ref_row << bwl) + ref_col;

348
    if (ref_row < 0 || ref_col < 0 || ref_row >= height || ref_col >= stride)
349
350
351
352
353
354
355
356
357
358
359
360
      continue;

    abs_coeff = abs(tcoeffs[pos]);

    for (i = 0; i < NUM_BASE_LEVELS; ++i) {
      ctx_set[i] += abs_coeff > i;
      if (base_ref_offset[idx][0] >= 0 && base_ref_offset[idx][1] >= 0)
        mag[i] |= abs_coeff > (i + 1);
    }
  }

  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
361
    ctx_set[i] = get_base_ctx_from_count_mag(row, col, ctx_set[i], mag[i]);
362
363
364
365
  }
  return;
}

366
static INLINE int get_br_cost(tran_low_t abs_qc, int ctx,
367
                              const int coeff_lps[COEFF_BASE_RANGE + 1]) {
368
369
  const tran_low_t min_level = 1 + NUM_BASE_LEVELS;
  const tran_low_t max_level = 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE;
370
  (void)ctx;
371
372
  if (abs_qc >= min_level) {
    if (abs_qc >= max_level)
373
      return coeff_lps[COEFF_BASE_RANGE];  // COEFF_BASE_RANGE * cost0;
374
    else
375
      return coeff_lps[(abs_qc - min_level)];  //  * cost0 + cost1;
376
377
378
379
380
381
  } else {
    return 0;
  }
}

static INLINE int get_base_cost(tran_low_t abs_qc, int ctx,
382
                                const int coeff_base[2], int base_idx) {
383
  const int level = base_idx + 1;
384
  (void)ctx;
385
386
387
  if (abs_qc < level)
    return 0;
  else
388
    return coeff_base[abs_qc == level];
389
390
}

391
int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
Jingning Han's avatar
Jingning Han committed
392
393
                        int blk_row, int blk_col, int block, TX_SIZE tx_size,
                        TXB_CTX *txb_ctx) {
394
  MACROBLOCKD *const xd = &x->e_mbd;
395
  TX_SIZE txs_ctx = get_txsize_context(tx_size);
396
  const PLANE_TYPE plane_type = get_plane_type(plane);
397
398
  const TX_TYPE tx_type =
      av1_get_tx_type(plane_type, xd, blk_row, blk_col, block, tx_size);
399
400
401
402
403
  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;
404
  const int seg_eob = AOMMIN(eob, tx_size_2d[tx_size] - 1);
405
406
407
  int txb_skip_ctx = txb_ctx->txb_skip_ctx;

  const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
408
409
  const int height = tx_size_high[tx_size];

Angie Chiang's avatar
Angie Chiang committed
410
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
411
412
  const int16_t *scan = scan_order->scan;

413
414
  LV_MAP_COEFF_COST *coeff_costs = &x->coeff_costs[txs_ctx][plane_type];

415
416
417
  cost = 0;

  if (eob == 0) {
418
    cost = coeff_costs->txb_skip_cost[txb_skip_ctx][1];
419
420
    return cost;
  }
421
  cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
422

Angie Chiang's avatar
Angie Chiang committed
423
#if CONFIG_TXK_SEL
424
  cost += av1_tx_type_cost(cm, x, xd, mbmi->sb_type, plane, tx_size, tx_type);
Angie Chiang's avatar
Angie Chiang committed
425
#endif
426

427
428
429
430
431
432
  for (c = 0; c < eob; ++c) {
    tran_low_t v = qcoeff[scan[c]];
    int is_nz = (v != 0);
    int level = abs(v);

    if (c < seg_eob) {
Angie Chiang's avatar
Angie Chiang committed
433
      int coeff_ctx = get_nz_map_ctx(qcoeff, scan[c], bwl, height, tx_type);
434
      cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
435
436
437
438
439
440
441
442
443
    }

    if (is_nz) {
      int ctx_ls[NUM_BASE_LEVELS] = { 0 };
      int sign = (v < 0) ? 1 : 0;

      // sign bit cost
      if (c == 0) {
        int dc_sign_ctx = txb_ctx->dc_sign_ctx;
444
        cost += coeff_costs->dc_sign_cost[dc_sign_ctx][sign];
445
446
447
448
      } else {
        cost += av1_cost_bit(128, sign);
      }

449
      get_base_ctx_set(qcoeff, scan[c], bwl, height, ctx_ls);
450
451
452
453
454
455

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

        if (level == i + 1) {
456
          cost += coeff_costs->base_cost[i][ctx_ls[i]][1];
457
458
          continue;
        }
459
        cost += coeff_costs->base_cost[i][ctx_ls[i]][0];
460
461
462
463
464
      }

      if (level > NUM_BASE_LEVELS) {
        int idx;
        int ctx;
465
        ctx = get_br_ctx(qcoeff, scan[c], bwl, height);
466
467
#if BR_NODE
        int base_range = level - 1 - NUM_BASE_LEVELS;
468
469
470
471
472
        if (base_range < COEFF_BASE_RANGE) {
          cost += coeff_costs->lps_cost[ctx][base_range];
          continue;
        } else {
          cost += coeff_costs->lps_cost[ctx][COEFF_BASE_RANGE];
473
        }
474

475
476
        idx = COEFF_BASE_RANGE;
#else
477
478
        for (idx = 0; idx < COEFF_BASE_RANGE; ++idx) {
          if (level == (idx + 1 + NUM_BASE_LEVELS)) {
479
            cost += coeff_costs->lps_cost[ctx][1];
480
481
            break;
          }
482
          cost += coeff_costs->lps_cost[ctx][0];
483
        }
484
#endif
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
        if (idx >= 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);
        }
      }

      if (c < seg_eob) {
504
        int eob_ctx = get_eob_ctx(qcoeff, scan[c], txs_ctx);
505
        cost += coeff_costs->eob_cost[eob_ctx][c == (eob - 1)];
506
507
508
509
510
511
      }
    }
  }

  return cost;
}
Angie Chiang's avatar
Angie Chiang committed
512

513
514
515
516
517
518
519
520
521
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;
}

522
static INLINE int get_sign_bit_cost(tran_low_t qc, int coeff_idx,
523
                                    const int (*dc_sign_cost)[2],
524
525
526
527
                                    int dc_sign_ctx) {
  const int sign = (qc < 0) ? 1 : 0;
  // sign bit cost
  if (coeff_idx == 0) {
528
    return dc_sign_cost[dc_sign_ctx][sign];
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
  } 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;
  }
}

551
void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
Angie Chiang's avatar
Angie Chiang committed
552
  // gen_nz_count_arr
553
  const int16_t *scan = txb_info->scan_order->scan;
Angie Chiang's avatar
Angie Chiang committed
554
555
556
  const int bwl = txb_info->bwl;
  const int height = txb_info->height;
  tran_low_t *qcoeff = txb_info->qcoeff;
557
558
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
559
560
561
562
563
  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);
    txb_cache->nz_count_arr[coeff_idx] =
Angie Chiang's avatar
Angie Chiang committed
564
        get_nz_count(qcoeff, bwl, height, row, col);
Angie Chiang's avatar
Angie Chiang committed
565
    const int nz_count = txb_cache->nz_count_arr[coeff_idx];
Angie Chiang's avatar
Angie Chiang committed
566
567
    txb_cache->nz_ctx_arr[coeff_idx] =
        get_nz_map_ctx_from_count(nz_count, coeff_idx, bwl, txb_info->tx_type);
Angie Chiang's avatar
Angie Chiang committed
568
569
570
571

    // 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
572
573
    int count[NUM_BASE_LEVELS];
    get_base_count_mag(base_mag, count, qcoeff, bwl, height, row, col);
Angie Chiang's avatar
Angie Chiang committed
574
575

    for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
Angie Chiang's avatar
Angie Chiang committed
576
577
      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
578
      const int level = i + 1;
Angie Chiang's avatar
Angie Chiang committed
579
      txb_cache->base_ctx_arr[i][coeff_idx] =
580
          base_ctx_table[row != 0][col != 0][base_mag[0] > level][count[i]];
Angie Chiang's avatar
Angie Chiang committed
581
582
583
584
585
586
    }

    // 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
587
588
    *br_count = get_br_count_mag(br_mag, qcoeff, bwl, height, row, col,
                                 NUM_BASE_LEVELS);
Angie Chiang's avatar
Angie Chiang committed
589
590
591
    txb_cache->br_ctx_arr[coeff_idx] =
        get_br_ctx_from_count_mag(row, col, *br_count, br_mag[0]);
  }
592
593
}

594
595
596
static INLINE const int *get_level_prob(int level, int coeff_idx,
                                        const TxbCache *txb_cache,
                                        const LV_MAP_COEFF_COST *txb_costs) {
597
  if (level == 0) {
598
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
599
    return txb_costs->nz_map_cost[ctx];
600
601
  } else if (level >= 1 && level < 1 + NUM_BASE_LEVELS) {
    const int idx = level - 1;
602
    const int ctx = txb_cache->base_ctx_arr[idx][coeff_idx];
603
    return txb_costs->base_cost[idx][ctx];
604
605
  } else if (level >= 1 + NUM_BASE_LEVELS &&
             level < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
606
    const int ctx = txb_cache->br_ctx_arr[coeff_idx];
607
    return txb_costs->lps_cost[ctx];
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
  } else if (level >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
    printf("get_level_prob does not support golomb\n");
    assert(0);
    return 0;
  } else {
    assert(0);
    return 0;
  }
}

static INLINE tran_low_t get_lower_coeff(tran_low_t qc) {
  if (qc == 0) {
    return 0;
  }
  return qc > 0 ? qc - 1 : qc + 1;
}

Angie Chiang's avatar
Angie Chiang committed
625
626
627
628
629
630
631
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
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
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,
682
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
                                      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);
701
    const int ctx = txb_cache->br_ctx_arr[coeff_idx];
702
    const int org_cost = get_br_cost(abs_qc, ctx, txb_costs->lps_cost[ctx]);
Angie Chiang's avatar
Angie Chiang committed
703
704

    const int new_ctx = get_br_ctx_from_count_mag(row, col, new_count, new_mag);
705
706
    const int new_cost =
        get_br_cost(abs_qc, new_ctx, txb_costs->lps_cost[new_ctx]);
Angie Chiang's avatar
Angie Chiang committed
707
708
709
710
711
712
713
714
715
    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,
716
                                        const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
717
718
719
                                        const TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const tran_low_t abs_qc = abs(qc);
720
721
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740

  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);
741
      const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx];
742
743
      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
744
745

      const int new_ctx =
746
          base_ctx_table[row != 0][col != 0][new_mag > level][new_count];
747
748
      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
749
750
751
752
753
754
755
756
      cost_diff += -org_cost + new_cost;
    }
  }
  return cost_diff;
}

static int try_neighbor_level_down_nz(int coeff_idx, int nb_coeff_idx,
                                      const TxbCache *txb_cache,
757
                                      const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
                                      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);
    txb_info->qcoeff[nb_coeff_idx] = get_lower_coeff(nb_coeff);
Angie Chiang's avatar
Angie Chiang committed
773
774
    const int new_ctx = get_nz_map_ctx_from_count(
        count - 1, coeff_idx, txb_info->bwl, txb_info->tx_type);
Angie Chiang's avatar
Angie Chiang committed
775
    txb_info->qcoeff[nb_coeff_idx] = nb_coeff;
776
    const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
Angie Chiang's avatar
Angie Chiang committed
777
    const int is_nz = abs_qc > 0;
778
779
    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];
Angie Chiang's avatar
Angie Chiang committed
780
781
782
783
784
785
786
787
788
    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,
789
790
                               const LV_MAP_COEFF_COST *txb_costs,
                               TxbInfo *txb_info) {
791
792
793
794
795
796
797
798
799
800
  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];
801
802
803
804
    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);
805
806
807
808
809
    if (scan_idx < txb_info->seg_eob) {
      // 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
810
      cost_diff = -level_cost[1] + low_level_cost[0] - low_level_cost[1];
811
    } else {
812
      cost_diff = -level_cost[1];
813
814
815
816
    }

    if (scan_idx < txb_info->seg_eob) {
      const int eob_ctx =
817
          get_eob_ctx(txb_info->qcoeff, coeff_idx, txb_info->txs_ctx);
818
819
      cost_diff -=
          txb_costs->eob_cost[eob_ctx][scan_idx == (txb_info->eob - 1)];
820
821
822
    }

    const int sign_cost = get_sign_bit_cost(
823
        qc, coeff_idx, txb_costs->dc_sign_cost, txb_info->txb_ctx->dc_sign_ctx);
824
    cost_diff -= sign_cost;
825
826
827
828
829
830
831
832
833
834
835
836
  } 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];
837
  } else if (abs_qc < 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
838
839
840
841
842
    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);

843
844
    cost_diff = -level_cost[abs_qc - 1 - NUM_BASE_LEVELS] +
                low_level_cost[abs(*low_coeff) - 1 - NUM_BASE_LEVELS];
845
  } else if (abs_qc == 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {
846
847
    const int *low_level_cost =
        get_level_prob(abs(*low_coeff), coeff_idx, txb_cache, txb_costs);
848
849
    cost_diff = -get_golomb_cost(abs_qc) - low_level_cost[COEFF_BASE_RANGE] +
                low_level_cost[COEFF_BASE_RANGE - 1];
850
851
852
853
854
855
856
857
  } 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
858
859
860
861
862
863
864
865
866
867
868
869
870
#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;
}

871
872
873
874
875
876
877
878
879
880
881
882
883
#define FAST_OPTIMIZE_TXB 1

#if FAST_OPTIMIZE_TXB
#define ALNB_REF_OFFSET_NUM 2
static int alnb_ref_offset[ALNB_REF_OFFSET_NUM][2] = {
  { -1, 0 }, { 0, -1 },
};
#define NB_REF_OFFSET_NUM 4
static int nb_ref_offset[NB_REF_OFFSET_NUM][2] = {
  { -1, 0 }, { 0, -1 }, { 1, 0 }, { 0, 1 },
};
#endif  // FAST_OPTIMIZE_TXB

Angie Chiang's avatar
Angie Chiang committed
884
885
// TODO(angiebird): add static to this function once it's called
int try_level_down(int coeff_idx, const TxbCache *txb_cache,
886
                   const LV_MAP_COEFF_COST *txb_costs, TxbInfo *txb_info,
Angie Chiang's avatar
Angie Chiang committed
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
                   int (*cost_map)[COST_MAP_SIZE]) {
  if (cost_map) {
    for (int i = 0; i < COST_MAP_SIZE; ++i) av1_zero(cost_map[i]);
  }

  tran_low_t qc = txb_info->qcoeff[coeff_idx];
  tran_low_t low_coeff;
  if (qc == 0) return 0;
  int accu_cost_diff = 0;

  const int16_t *iscan = txb_info->scan_order->iscan;
  const int eob = txb_info->eob;
  const int scan_idx = iscan[coeff_idx];
  if (scan_idx < eob) {
    const int cost_diff = try_self_level_down(&low_coeff, coeff_idx, txb_cache,
902
                                              txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
903
904
905
906
907
908
909
910
    if (cost_map)
      cost_map[0 + COST_MAP_OFFSET][0 + COST_MAP_OFFSET] = cost_diff;
    accu_cost_diff += cost_diff;
  }

  const int row = coeff_idx >> txb_info->bwl;
  const int col = coeff_idx - (row << txb_info->bwl);
  if (check_nz_neighbor(qc)) {
911
912
913
914
915
916
917
918
919
920
#if FAST_OPTIMIZE_TXB
    int(*ref_offset)[2] = alnb_ref_offset;
    const int ref_num = ALNB_REF_OFFSET_NUM;
#else
    int(*ref_offset)[2] = sig_ref_offset;
    const int ref_num = SIG_REF_OFFSET_NUM;
#endif
    for (int i = 0; i < ref_num; ++i) {
      const int nb_row = row - ref_offset[i][0];
      const int nb_col = col - ref_offset[i][1];
Angie Chiang's avatar
Angie Chiang committed
921
      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
922

923
924
      if (nb_row < 0 || nb_col < 0 || nb_row >= txb_info->height ||
          nb_col >= txb_info->stride)
925
926
        continue;

Angie Chiang's avatar
Angie Chiang committed
927
      const int nb_scan_idx = iscan[nb_coeff_idx];
928
      if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
929
        const int cost_diff = try_neighbor_level_down_nz(
930
            nb_coeff_idx, coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
931
932
933
934
935
936
937
938
939
        if (cost_map)
          cost_map[nb_row - row + COST_MAP_OFFSET]
                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
        accu_cost_diff += cost_diff;
      }
    }
  }

  if (check_base_neighbor(qc)) {
940
941
942
943
944
945
946
947
948
949
#if FAST_OPTIMIZE_TXB
    int(*ref_offset)[2] = nb_ref_offset;
    const int ref_num = NB_REF_OFFSET_NUM;
#else
    int(*ref_offset)[2] = base_ref_offset;
    const int ref_num = BASE_CONTEXT_POSITION_NUM;
#endif
    for (int i = 0; i < ref_num; ++i) {
      const int nb_row = row - ref_offset[i][0];
      const int nb_col = col - ref_offset[i][1];
Angie Chiang's avatar
Angie Chiang committed
950
      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
951

952
953
      if (nb_row < 0 || nb_col < 0 || nb_row >= txb_info->height ||
          nb_col >= txb_info->stride)
954
955
        continue;

Angie Chiang's avatar
Angie Chiang committed
956
      const int nb_scan_idx = iscan[nb_coeff_idx];
957
      if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
958
        const int cost_diff = try_neighbor_level_down_base(
959
            nb_coeff_idx, coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
960
961
962
963
964
965
966
967
968
        if (cost_map)
          cost_map[nb_row - row + COST_MAP_OFFSET]
                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
        accu_cost_diff += cost_diff;
      }
    }
  }

  if (check_br_neighbor(qc)) {
969
970
971
972
973
974
975
976
977
978
#if FAST_OPTIMIZE_TXB
    int(*ref_offset)[2] = nb_ref_offset;
    const int ref_num = NB_REF_OFFSET_NUM;
#else
    int(*ref_offset)[2] = br_ref_offset;
    const int ref_num = BR_CONTEXT_POSITION_NUM;
#endif
    for (int i = 0; i < ref_num; ++i) {
      const int nb_row = row - ref_offset[i][0];
      const int nb_col = col - ref_offset[i][1];
Angie Chiang's avatar
Angie Chiang committed
979
      const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
980

981
982
      if (nb_row < 0 || nb_col < 0 || nb_row >= txb_info->height ||
          nb_col >= txb_info->stride)
983
984
        continue;

Angie Chiang's avatar
Angie Chiang committed
985
      const int nb_scan_idx = iscan[nb_coeff_idx];
986
      if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
987
        const int cost_diff = try_neighbor_level_down_br(
988
            nb_coeff_idx, coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
989
990
991
992
993
994
995
996
997
998
999
        if (cost_map)
          cost_map[nb_row - row + COST_MAP_OFFSET]
                  [nb_col - col + COST_MAP_OFFSET] += cost_diff;
        accu_cost_diff += cost_diff;
      }
    }
  }

  return accu_cost_diff;
}

Angie Chiang's avatar
Angie Chiang committed
1000
static int get_low_coeff_cost(int coeff_idx, const TxbCache *txb_cache,
1001
                              const LV_MAP_COEFF_COST *txb_costs,
Angie Chiang's avatar
Angie Chiang committed
1002
1003
1004
1005
1006
1007
1008
                              const TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const int abs_qc = abs(qc);
  assert(abs_qc <= 1);
  int cost = 0;
  const int scan_idx = txb_info->scan_order->iscan[coeff_idx];
  if (scan_idx < txb_info->seg_eob) {
1009
1010
    const int *level_cost = get_level_prob(0, coeff_idx, txb_cache, txb_costs);
    cost += level_cost[qc != 0];
Angie Chiang's avatar
Angie Chiang committed
1011
1012
1013
1014
  }

  if (qc != 0) {
    const int base_idx = 0;
1015
    const int ctx = txb_cache->base_ctx_arr[base_idx][coeff_idx];
1016
1017
    cost += get_base_cost(abs_qc, ctx, txb_costs->base_cost[base_idx][ctx],
                          base_idx);
Angie Chiang's avatar
Angie Chiang committed
1018
1019
    if (scan_idx < txb_info->seg_eob) {
      const int eob_ctx =
1020
          get_eob_ctx(txb_info->qcoeff, coeff_idx, txb_info->txs_ctx);
1021
      cost += txb_costs->eob_cost[eob_ctx][scan_idx == (txb_info->eob - 1)];
Angie Chiang's avatar
Angie Chiang committed
1022
    }
1023
    cost += get_sign_bit_cost(qc, coeff_idx, txb_costs->dc_sign_cost,
Angie Chiang's avatar
Angie Chiang committed
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
                              txb_info->txb_ctx->dc_sign_ctx);
  }
  return cost;
}

static INLINE void set_eob(TxbInfo *txb_info, int eob) {
  txb_info->eob = eob;
  txb_info->seg_eob = AOMMIN(eob, tx_size_2d[txb_info->tx_size] - 1);
}

// TODO(angiebird): add static to this function once it's called
int try_change_eob(int *new_eob, int coeff_idx, const TxbCache *txb_cache,
1036
                   const LV_MAP_COEFF_COST *txb_costs, TxbInfo *txb_info) {
Angie Chiang's avatar
Angie Chiang committed
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
  assert(txb_info->eob > 0);
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const int abs_qc = abs(qc);
  if (abs_qc != 1) {
    *new_eob = -1;
    return 0;
  }
  const int16_t *iscan = txb_info->scan_order->iscan;
  const int16_t *scan = txb_info->scan_order->scan;
  const int scan_idx = iscan[coeff_idx];
  *new_eob = 0;
  int cost_diff = 0;
1049
  cost_diff -= get_low_coeff_cost(coeff_idx, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
1050
  // int coeff_cost =
1051
  //     get_coeff_cost(qc, scan_idx, txb_info, txb_probs);
Angie Chiang's avatar
Angie Chiang committed
1052
1053
1054
  // if (-cost_diff != coeff_cost) {
  //   printf("-cost_diff %d coeff_cost %d\n", -cost_diff, coeff_cost);
  //   get_low_coeff_cost(coeff_idx, txb_cache, txb_probs, txb_info);
1055
  //   get_coeff_cost(qc, scan_idx, txb_info, txb_probs);
Angie Chiang's avatar
Angie Chiang committed
1056
1057
1058
1059
1060
1061
1062
  // }
  for (int si = scan_idx - 1; si >= 0; --si) {
    const int ci = scan[si];
    if (txb_info->qcoeff[ci] != 0) {
      *new_eob = si + 1;
      break;
    } else {
1063
      cost_diff -= get_low_coeff_cost(ci, txb_cache, txb_costs, txb_info);
Angie Chiang's avatar
Angie Chiang committed
1064
1065
1066
1067
1068
    }
  }

  const int org_eob = txb_info->eob;
  set_eob(txb_info, *new_eob);
1069
  cost_diff += try_level_down(coeff_idx, txb_cache, txb_costs, txb_info, NULL);
Angie Chiang's avatar
Angie Chiang committed
1070
1071
1072
1073
1074
1075
  set_eob(txb_info, org_eob);

  if (*new_eob > 0) {
    // Note that get_eob_ctx does NOT actually account for qcoeff, so we don't
    // need to lower down the qcoeff here
    const int eob_ctx =
1076
        get_eob_ctx(txb_info->qcoeff, scan[*new_eob - 1], txb_info->txs_ctx);
1077
1078
    cost_diff -= txb_costs->eob_cost[eob_ctx][0];
    cost_diff += txb_costs->eob_cost[eob_ctx][1];
Angie Chiang's avatar
Angie Chiang committed
1079
1080
  } else {
    const int txb_skip_ctx = txb_info->txb_ctx->txb_skip_ctx;
1081
1082
    cost_diff -= txb_costs->txb_skip_cost[txb_skip_ctx][0];
    cost_diff += txb_costs->txb_skip_cost[txb_skip_ctx][1];
Angie Chiang's avatar
Angie Chiang committed
1083
1084
1085
  }
  return cost_diff;
}
Angie Chiang's avatar
Angie Chiang committed
1086
1087
1088
1089
1090
1091

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

Angie Chiang's avatar
Angie Chiang committed
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
// TODO(angiebird): add static to this function it's called
void update_level_down(int coeff_idx, TxbCache *txb_cache, TxbInfo *txb_info) {
  const tran_low_t qc = txb_info->qcoeff[coeff_idx];
  const int abs_qc = abs(qc);
  if (qc == 0) return;
  const tran_low_t low_coeff = get_lower_coeff(qc);
  txb_info->qcoeff[coeff_idx] = low_coeff;
  const int dqv = txb_info->dequant[coeff_idx != 0];
  txb_info->dqcoeff[coeff_idx] =
      qcoeff_to_dqcoeff(low_coeff, dqv, txb_info->shift);

  const int row = coeff_idx >> txb_info->bwl;
  const int col = coeff_idx - (row << txb_info->bwl);
  const int eob = txb_info->eob;
  const int16_t *iscan = txb_info->scan_order->iscan;
  for (int i = 0; i < SIG_REF_OFFSET_NUM; ++i) {
    const int nb_row = row - sig_ref_offset[i][0];
    const int nb_col = col - sig_ref_offset[i][1];
1110
1111
1112
1113
1114

    if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
          nb_col < txb_info->stride))
      continue;

Angie Chiang's avatar
Angie Chiang committed
1115
1116
    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
    const int nb_scan_idx = iscan[nb_coeff_idx];
1117
    if (nb_scan_idx < eob) {
Angie Chiang's avatar
Angie Chiang committed
1118
1119
1120
1121
1122
1123
1124
1125
      const int scan_idx = iscan[coeff_idx];
      if (scan_idx < nb_scan_idx) {
        const int level = 1;
        if (abs_qc == level) {
          txb_cache->nz_count_arr[nb_coeff_idx] -= 1;
          assert(txb_cache->nz_count_arr[nb_coeff_idx] >= 0);
        }
        const int count = txb_cache->nz_count_arr[nb_coeff_idx];
Angie Chiang's avatar
Angie Chiang committed
1126
1127
        txb_cache->nz_ctx_arr[nb_coeff_idx] = get_nz_map_ctx_from_count(
            count, nb_coeff_idx, txb_info->bwl, txb_info->tx_type);
1128
        // int ref_ctx = get_nz_map_ctx(txb_info->qcoeff, nb_coeff_idx,
Angie Chiang's avatar
Angie Chiang committed
1129
        // txb_info->bwl, tx_type);
1130
        // if (ref_ctx != txb_cache->nz_ctx_arr[nb_coeff_idx])
Angie Chiang's avatar
Angie Chiang committed
1131
        //   printf("nz ctx %d ref_ctx %d\n",
1132
        //   txb_cache->nz_ctx_arr[nb_coeff_idx], ref_ctx);
Angie Chiang's avatar
Angie Chiang committed
1133
1134
1135
1136
      }
    }
  }

1137
1138
  const BASE_CTX_TABLE *base_ctx_table =
      txb_info->coeff_ctx_table->base_ctx_table;
Angie Chiang's avatar
Angie Chiang committed
1139
1140
1141
1142
  for (int i = 0; i < BASE_CONTEXT_POSITION_NUM; ++i) {
    const int nb_row = row - base_ref_offset[i][0];
    const int nb_col = col - base_ref_offset[i][1];
    const int nb_coeff_idx = nb_row * txb_info->stride + nb_col;
1143
1144
1145
1146
1147

    if (!(nb_row >= 0 && nb_col >= 0 && nb_row < txb_info->height &&
          nb_col < txb_info->stride))
      continue;

Angie Chiang's avatar
Angie Chiang committed
1148
1149
1150
    const tran_low_t nb_coeff = txb_info->qcoeff[nb_coeff_idx];
    if