tokenize.c 36 KB
Newer Older
Jingning Han's avatar
Jingning Han committed
1
/*
Yaowu Xu's avatar
Yaowu Xu committed
2
 * Copyright (c) 2016, Alliance for Open Media. All rights reserved
Jingning Han's avatar
Jingning Han committed
3
 *
Yaowu Xu's avatar
Yaowu Xu committed
4
5
6
7
8
9
 * 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.
Jingning Han's avatar
Jingning Han committed
10
11
12
13
14
15
16
 */

#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <string.h>

Yaowu Xu's avatar
Yaowu Xu committed
17
#include "aom_mem/aom_mem.h"
Jingning Han's avatar
Jingning Han committed
18

19
20
21
22
#include "av1/common/entropy.h"
#include "av1/common/pred_common.h"
#include "av1/common/scan.h"
#include "av1/common/seg_common.h"
Jingning Han's avatar
Jingning Han committed
23

24
25
#include "av1/encoder/cost.h"
#include "av1/encoder/encoder.h"
26
#include "av1/encoder/rdopt.h"
27
#include "av1/encoder/tokenize.h"
Jingning Han's avatar
Jingning Han committed
28
29

static const TOKENVALUE dct_cat_lt_10_value_tokens[] = {
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
  { 9, 63 }, { 9, 61 }, { 9, 59 }, { 9, 57 }, { 9, 55 }, { 9, 53 }, { 9, 51 },
  { 9, 49 }, { 9, 47 }, { 9, 45 }, { 9, 43 }, { 9, 41 }, { 9, 39 }, { 9, 37 },
  { 9, 35 }, { 9, 33 }, { 9, 31 }, { 9, 29 }, { 9, 27 }, { 9, 25 }, { 9, 23 },
  { 9, 21 }, { 9, 19 }, { 9, 17 }, { 9, 15 }, { 9, 13 }, { 9, 11 }, { 9, 9 },
  { 9, 7 },  { 9, 5 },  { 9, 3 },  { 9, 1 },  { 8, 31 }, { 8, 29 }, { 8, 27 },
  { 8, 25 }, { 8, 23 }, { 8, 21 }, { 8, 19 }, { 8, 17 }, { 8, 15 }, { 8, 13 },
  { 8, 11 }, { 8, 9 },  { 8, 7 },  { 8, 5 },  { 8, 3 },  { 8, 1 },  { 7, 15 },
  { 7, 13 }, { 7, 11 }, { 7, 9 },  { 7, 7 },  { 7, 5 },  { 7, 3 },  { 7, 1 },
  { 6, 7 },  { 6, 5 },  { 6, 3 },  { 6, 1 },  { 5, 3 },  { 5, 1 },  { 4, 1 },
  { 3, 1 },  { 2, 1 },  { 1, 1 },  { 0, 0 },  { 1, 0 },  { 2, 0 },  { 3, 0 },
  { 4, 0 },  { 5, 0 },  { 5, 2 },  { 6, 0 },  { 6, 2 },  { 6, 4 },  { 6, 6 },
  { 7, 0 },  { 7, 2 },  { 7, 4 },  { 7, 6 },  { 7, 8 },  { 7, 10 }, { 7, 12 },
  { 7, 14 }, { 8, 0 },  { 8, 2 },  { 8, 4 },  { 8, 6 },  { 8, 8 },  { 8, 10 },
  { 8, 12 }, { 8, 14 }, { 8, 16 }, { 8, 18 }, { 8, 20 }, { 8, 22 }, { 8, 24 },
  { 8, 26 }, { 8, 28 }, { 8, 30 }, { 9, 0 },  { 9, 2 },  { 9, 4 },  { 9, 6 },
  { 9, 8 },  { 9, 10 }, { 9, 12 }, { 9, 14 }, { 9, 16 }, { 9, 18 }, { 9, 20 },
  { 9, 22 }, { 9, 24 }, { 9, 26 }, { 9, 28 }, { 9, 30 }, { 9, 32 }, { 9, 34 },
  { 9, 36 }, { 9, 38 }, { 9, 40 }, { 9, 42 }, { 9, 44 }, { 9, 46 }, { 9, 48 },
  { 9, 50 }, { 9, 52 }, { 9, 54 }, { 9, 56 }, { 9, 58 }, { 9, 60 }, { 9, 62 }
Jingning Han's avatar
Jingning Han committed
49
};
Yaowu Xu's avatar
Yaowu Xu committed
50
const TOKENVALUE *av1_dct_cat_lt_10_value_tokens =
51
52
53
    dct_cat_lt_10_value_tokens +
    (sizeof(dct_cat_lt_10_value_tokens) / sizeof(*dct_cat_lt_10_value_tokens)) /
        2;
Julia Robson's avatar
Julia Robson committed
54
55
56
// The corresponding costs of the extrabits for the tokens in the above table
// are stored in the table below. The values are obtained from looking up the
// entry for the specified extrabits in the table corresponding to the token
Yaowu Xu's avatar
Yaowu Xu committed
57
// (as defined in cost element av1_extra_bits)
Julia Robson's avatar
Julia Robson committed
58
59
// e.g. {9, 63} maps to cat5_cost[63 >> 1], {1, 1} maps to sign_cost[1 >> 1]
static const int dct_cat_lt_10_value_cost[] = {
60
61
62
63
64
65
66
67
68
69
70
  3773, 3750, 3704, 3681, 3623, 3600, 3554, 3531, 3432, 3409, 3363, 3340, 3282,
  3259, 3213, 3190, 3136, 3113, 3067, 3044, 2986, 2963, 2917, 2894, 2795, 2772,
  2726, 2703, 2645, 2622, 2576, 2553, 3197, 3116, 3058, 2977, 2881, 2800, 2742,
  2661, 2615, 2534, 2476, 2395, 2299, 2218, 2160, 2079, 2566, 2427, 2334, 2195,
  2023, 1884, 1791, 1652, 1893, 1696, 1453, 1256, 1229, 864,  512,  512,  512,
  512,  0,    512,  512,  512,  512,  864,  1229, 1256, 1453, 1696, 1893, 1652,
  1791, 1884, 2023, 2195, 2334, 2427, 2566, 2079, 2160, 2218, 2299, 2395, 2476,
  2534, 2615, 2661, 2742, 2800, 2881, 2977, 3058, 3116, 3197, 2553, 2576, 2622,
  2645, 2703, 2726, 2772, 2795, 2894, 2917, 2963, 2986, 3044, 3067, 3113, 3136,
  3190, 3213, 3259, 3282, 3340, 3363, 3409, 3432, 3531, 3554, 3600, 3623, 3681,
  3704, 3750, 3773,
Julia Robson's avatar
Julia Robson committed
71
};
Yaowu Xu's avatar
Yaowu Xu committed
72
const int *av1_dct_cat_lt_10_value_cost =
73
74
    dct_cat_lt_10_value_cost +
    (sizeof(dct_cat_lt_10_value_cost) / sizeof(*dct_cat_lt_10_value_cost)) / 2;
Jingning Han's avatar
Jingning Han committed
75
76

// Array indices are identical to previously-existing CONTEXT_NODE indices
77
/* clang-format off */
Yaowu Xu's avatar
Yaowu Xu committed
78
const aom_tree_index av1_coef_tree[TREE_SIZE(ENTROPY_TOKENS)] = {
Jingning Han's avatar
Jingning Han committed
79
80
81
82
83
84
85
86
87
88
89
90
  -EOB_TOKEN, 2,                       // 0  = EOB
  -ZERO_TOKEN, 4,                      // 1  = ZERO
  -ONE_TOKEN, 6,                       // 2  = ONE
  8, 12,                               // 3  = LOW_VAL
  -TWO_TOKEN, 10,                      // 4  = TWO
  -THREE_TOKEN, -FOUR_TOKEN,           // 5  = THREE
  14, 16,                              // 6  = HIGH_LOW
  -CATEGORY1_TOKEN, -CATEGORY2_TOKEN,  // 7  = CAT_ONE
  18, 20,                              // 8  = CAT_THREEFOUR
  -CATEGORY3_TOKEN, -CATEGORY4_TOKEN,  // 9  = CAT_THREE
  -CATEGORY5_TOKEN, -CATEGORY6_TOKEN   // 10 = CAT_FIVE
};
91
92
93
94
95
96
97
98
99
100
101
/* clang-format on */

static const int16_t zero_cost[] = { 0 };
static const int16_t sign_cost[1] = { 512 };
static const int16_t cat1_cost[1 << 1] = { 864, 1229 };
static const int16_t cat2_cost[1 << 2] = { 1256, 1453, 1696, 1893 };
static const int16_t cat3_cost[1 << 3] = { 1652, 1791, 1884, 2023,
                                           2195, 2334, 2427, 2566 };
static const int16_t cat4_cost[1 << 4] = { 2079, 2160, 2218, 2299, 2395, 2476,
                                           2534, 2615, 2661, 2742, 2800, 2881,
                                           2977, 3058, 3116, 3197 };
102
static const int16_t cat5_cost[1 << 5] = {
103
104
105
106
  2553, 2576, 2622, 2645, 2703, 2726, 2772, 2795, 2894, 2917, 2963,
  2986, 3044, 3067, 3113, 3136, 3190, 3213, 3259, 3282, 3340, 3363,
  3409, 3432, 3531, 3554, 3600, 3623, 3681, 3704, 3750, 3773
};
Yaowu Xu's avatar
Yaowu Xu committed
107
const int16_t av1_cat6_low_cost[256] = {
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
  3378, 3390, 3401, 3413, 3435, 3447, 3458, 3470, 3517, 3529, 3540, 3552, 3574,
  3586, 3597, 3609, 3671, 3683, 3694, 3706, 3728, 3740, 3751, 3763, 3810, 3822,
  3833, 3845, 3867, 3879, 3890, 3902, 3973, 3985, 3996, 4008, 4030, 4042, 4053,
  4065, 4112, 4124, 4135, 4147, 4169, 4181, 4192, 4204, 4266, 4278, 4289, 4301,
  4323, 4335, 4346, 4358, 4405, 4417, 4428, 4440, 4462, 4474, 4485, 4497, 4253,
  4265, 4276, 4288, 4310, 4322, 4333, 4345, 4392, 4404, 4415, 4427, 4449, 4461,
  4472, 4484, 4546, 4558, 4569, 4581, 4603, 4615, 4626, 4638, 4685, 4697, 4708,
  4720, 4742, 4754, 4765, 4777, 4848, 4860, 4871, 4883, 4905, 4917, 4928, 4940,
  4987, 4999, 5010, 5022, 5044, 5056, 5067, 5079, 5141, 5153, 5164, 5176, 5198,
  5210, 5221, 5233, 5280, 5292, 5303, 5315, 5337, 5349, 5360, 5372, 4988, 5000,
  5011, 5023, 5045, 5057, 5068, 5080, 5127, 5139, 5150, 5162, 5184, 5196, 5207,
  5219, 5281, 5293, 5304, 5316, 5338, 5350, 5361, 5373, 5420, 5432, 5443, 5455,
  5477, 5489, 5500, 5512, 5583, 5595, 5606, 5618, 5640, 5652, 5663, 5675, 5722,
  5734, 5745, 5757, 5779, 5791, 5802, 5814, 5876, 5888, 5899, 5911, 5933, 5945,
  5956, 5968, 6015, 6027, 6038, 6050, 6072, 6084, 6095, 6107, 5863, 5875, 5886,
  5898, 5920, 5932, 5943, 5955, 6002, 6014, 6025, 6037, 6059, 6071, 6082, 6094,
  6156, 6168, 6179, 6191, 6213, 6225, 6236, 6248, 6295, 6307, 6318, 6330, 6352,
  6364, 6375, 6387, 6458, 6470, 6481, 6493, 6515, 6527, 6538, 6550, 6597, 6609,
  6620, 6632, 6654, 6666, 6677, 6689, 6751, 6763, 6774, 6786, 6808, 6820, 6831,
  6843, 6890, 6902, 6913, 6925, 6947, 6959, 6970, 6982
};
129
const int av1_cat6_high_cost[CAT6_HIGH_COST_ENTRIES] = {
130
131
132
133
134
  100,   2263,  2739,  4902,  3160,  5323,  5799,  7962,  3678,  5841,  6317,
  8480,  6738,  8901,  9377,  11540, 3678,  5841,  6317,  8480,  6738,  8901,
  9377,  11540, 7256,  9419,  9895,  12058, 10316, 12479, 12955, 15118, 3678,
  5841,  6317,  8480,  6738,  8901,  9377,  11540, 7256,  9419,  9895,  12058,
  10316, 12479, 12955, 15118, 7256,  9419,  9895,  12058, 10316, 12479, 12955,
135
136
137
138
139
140
141
142
  15118, 10834, 12997, 13473, 15636, 13894, 16057, 16533, 18696,
#if CONFIG_AOM_HIGHBITDEPTH
  4193,  6356,  6832,  8995,  7253,  9416,  9892,  12055, 7771,  9934,  10410,
  12573, 10831, 12994, 13470, 15633, 7771,  9934,  10410, 12573, 10831, 12994,
  13470, 15633, 11349, 13512, 13988, 16151, 14409, 16572, 17048, 19211, 7771,
  9934,  10410, 12573, 10831, 12994, 13470, 15633, 11349, 13512, 13988, 16151,
  14409, 16572, 17048, 19211, 11349, 13512, 13988, 16151, 14409, 16572, 17048,
  19211, 14927, 17090, 17566, 19729, 17987, 20150, 20626, 22789, 4193,  6356,
143
144
145
146
147
  6832,  8995,  7253,  9416,  9892,  12055, 7771,  9934,  10410, 12573, 10831,
  12994, 13470, 15633, 7771,  9934,  10410, 12573, 10831, 12994, 13470, 15633,
  11349, 13512, 13988, 16151, 14409, 16572, 17048, 19211, 7771,  9934,  10410,
  12573, 10831, 12994, 13470, 15633, 11349, 13512, 13988, 16151, 14409, 16572,
  17048, 19211, 11349, 13512, 13988, 16151, 14409, 16572, 17048, 19211, 14927,
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
  17090, 17566, 19729, 17987, 20150, 20626, 22789, 8286,  10449, 10925, 13088,
  11346, 13509, 13985, 16148, 11864, 14027, 14503, 16666, 14924, 17087, 17563,
  19726, 11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726, 15442, 17605,
  18081, 20244, 18502, 20665, 21141, 23304, 11864, 14027, 14503, 16666, 14924,
  17087, 17563, 19726, 15442, 17605, 18081, 20244, 18502, 20665, 21141, 23304,
  15442, 17605, 18081, 20244, 18502, 20665, 21141, 23304, 19020, 21183, 21659,
  23822, 22080, 24243, 24719, 26882, 4193,  6356,  6832,  8995,  7253,  9416,
  9892,  12055, 7771,  9934,  10410, 12573, 10831, 12994, 13470, 15633, 7771,
  9934,  10410, 12573, 10831, 12994, 13470, 15633, 11349, 13512, 13988, 16151,
  14409, 16572, 17048, 19211, 7771,  9934,  10410, 12573, 10831, 12994, 13470,
  15633, 11349, 13512, 13988, 16151, 14409, 16572, 17048, 19211, 11349, 13512,
  13988, 16151, 14409, 16572, 17048, 19211, 14927, 17090, 17566, 19729, 17987,
  20150, 20626, 22789, 8286,  10449, 10925, 13088, 11346, 13509, 13985, 16148,
  11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726, 11864, 14027, 14503,
  16666, 14924, 17087, 17563, 19726, 15442, 17605, 18081, 20244, 18502, 20665,
  21141, 23304, 11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726, 15442,
  17605, 18081, 20244, 18502, 20665, 21141, 23304, 15442, 17605, 18081, 20244,
  18502, 20665, 21141, 23304, 19020, 21183, 21659, 23822, 22080, 24243, 24719,
  26882, 8286,  10449, 10925, 13088, 11346, 13509, 13985, 16148, 11864, 14027,
167
168
169
170
  14503, 16666, 14924, 17087, 17563, 19726, 11864, 14027, 14503, 16666, 14924,
  17087, 17563, 19726, 15442, 17605, 18081, 20244, 18502, 20665, 21141, 23304,
  11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726, 15442, 17605, 18081,
  20244, 18502, 20665, 21141, 23304, 15442, 17605, 18081, 20244, 18502, 20665,
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
  21141, 23304, 19020, 21183, 21659, 23822, 22080, 24243, 24719, 26882, 12379,
  14542, 15018, 17181, 15439, 17602, 18078, 20241, 15957, 18120, 18596, 20759,
  19017, 21180, 21656, 23819, 15957, 18120, 18596, 20759, 19017, 21180, 21656,
  23819, 19535, 21698, 22174, 24337, 22595, 24758, 25234, 27397, 15957, 18120,
  18596, 20759, 19017, 21180, 21656, 23819, 19535, 21698, 22174, 24337, 22595,
  24758, 25234, 27397, 19535, 21698, 22174, 24337, 22595, 24758, 25234, 27397,
  23113, 25276, 25752, 27915, 26173, 28336, 28812, 30975, 4193,  6356,  6832,
  8995,  7253,  9416,  9892,  12055, 7771,  9934,  10410, 12573, 10831, 12994,
  13470, 15633, 7771,  9934,  10410, 12573, 10831, 12994, 13470, 15633, 11349,
  13512, 13988, 16151, 14409, 16572, 17048, 19211, 7771,  9934,  10410, 12573,
  10831, 12994, 13470, 15633, 11349, 13512, 13988, 16151, 14409, 16572, 17048,
  19211, 11349, 13512, 13988, 16151, 14409, 16572, 17048, 19211, 14927, 17090,
  17566, 19729, 17987, 20150, 20626, 22789, 8286,  10449, 10925, 13088, 11346,
  13509, 13985, 16148, 11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726,
  11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726, 15442, 17605, 18081,
  20244, 18502, 20665, 21141, 23304, 11864, 14027, 14503, 16666, 14924, 17087,
  17563, 19726, 15442, 17605, 18081, 20244, 18502, 20665, 21141, 23304, 15442,
  17605, 18081, 20244, 18502, 20665, 21141, 23304, 19020, 21183, 21659, 23822,
  22080, 24243, 24719, 26882, 8286,  10449, 10925, 13088, 11346, 13509, 13985,
190
191
192
193
194
  16148, 11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726, 11864, 14027,
  14503, 16666, 14924, 17087, 17563, 19726, 15442, 17605, 18081, 20244, 18502,
  20665, 21141, 23304, 11864, 14027, 14503, 16666, 14924, 17087, 17563, 19726,
  15442, 17605, 18081, 20244, 18502, 20665, 21141, 23304, 15442, 17605, 18081,
  20244, 18502, 20665, 21141, 23304, 19020, 21183, 21659, 23822, 22080, 24243,
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
  24719, 26882, 12379, 14542, 15018, 17181, 15439, 17602, 18078, 20241, 15957,
  18120, 18596, 20759, 19017, 21180, 21656, 23819, 15957, 18120, 18596, 20759,
  19017, 21180, 21656, 23819, 19535, 21698, 22174, 24337, 22595, 24758, 25234,
  27397, 15957, 18120, 18596, 20759, 19017, 21180, 21656, 23819, 19535, 21698,
  22174, 24337, 22595, 24758, 25234, 27397, 19535, 21698, 22174, 24337, 22595,
  24758, 25234, 27397, 23113, 25276, 25752, 27915, 26173, 28336, 28812, 30975,
  8286,  10449, 10925, 13088, 11346, 13509, 13985, 16148, 11864, 14027, 14503,
  16666, 14924, 17087, 17563, 19726, 11864, 14027, 14503, 16666, 14924, 17087,
  17563, 19726, 15442, 17605, 18081, 20244, 18502, 20665, 21141, 23304, 11864,
  14027, 14503, 16666, 14924, 17087, 17563, 19726, 15442, 17605, 18081, 20244,
  18502, 20665, 21141, 23304, 15442, 17605, 18081, 20244, 18502, 20665, 21141,
  23304, 19020, 21183, 21659, 23822, 22080, 24243, 24719, 26882, 12379, 14542,
  15018, 17181, 15439, 17602, 18078, 20241, 15957, 18120, 18596, 20759, 19017,
  21180, 21656, 23819, 15957, 18120, 18596, 20759, 19017, 21180, 21656, 23819,
  19535, 21698, 22174, 24337, 22595, 24758, 25234, 27397, 15957, 18120, 18596,
  20759, 19017, 21180, 21656, 23819, 19535, 21698, 22174, 24337, 22595, 24758,
  25234, 27397, 19535, 21698, 22174, 24337, 22595, 24758, 25234, 27397, 23113,
  25276, 25752, 27915, 26173, 28336, 28812, 30975, 12379, 14542, 15018, 17181,
213
214
215
216
217
  15439, 17602, 18078, 20241, 15957, 18120, 18596, 20759, 19017, 21180, 21656,
  23819, 15957, 18120, 18596, 20759, 19017, 21180, 21656, 23819, 19535, 21698,
  22174, 24337, 22595, 24758, 25234, 27397, 15957, 18120, 18596, 20759, 19017,
  21180, 21656, 23819, 19535, 21698, 22174, 24337, 22595, 24758, 25234, 27397,
  19535, 21698, 22174, 24337, 22595, 24758, 25234, 27397, 23113, 25276, 25752,
218
219
220
221
222
223
224
  27915, 26173, 28336, 28812, 30975, 16472, 18635, 19111, 21274, 19532, 21695,
  22171, 24334, 20050, 22213, 22689, 24852, 23110, 25273, 25749, 27912, 20050,
  22213, 22689, 24852, 23110, 25273, 25749, 27912, 23628, 25791, 26267, 28430,
  26688, 28851, 29327, 31490, 20050, 22213, 22689, 24852, 23110, 25273, 25749,
  27912, 23628, 25791, 26267, 28430, 26688, 28851, 29327, 31490, 23628, 25791,
  26267, 28430, 26688, 28851, 29327, 31490, 27206, 29369, 29845, 32008, 30266,
  32429, 32905, 35068
Jingning Han's avatar
Jingning Han committed
225
#endif
226
227
228
229
230
};

const uint8_t av1_cat6_skipped_bits_discount[8] = {
  0, 3, 6, 9, 12, 18, 24, 30
};
Jingning Han's avatar
Jingning Han committed
231

Yaowu Xu's avatar
Yaowu Xu committed
232
const av1_extra_bit av1_extra_bits[ENTROPY_TOKENS] = {
Alex Converse's avatar
Alex Converse committed
233
234
235
236
237
238
239
240
241
242
  { 0, 0, 0, zero_cost },                         // ZERO_TOKEN
  { 0, 0, 1, sign_cost },                         // ONE_TOKEN
  { 0, 0, 2, sign_cost },                         // TWO_TOKEN
  { 0, 0, 3, sign_cost },                         // THREE_TOKEN
  { 0, 0, 4, sign_cost },                         // FOUR_TOKEN
  { av1_cat1_prob, 1, CAT1_MIN_VAL, cat1_cost },  // CATEGORY1_TOKEN
  { av1_cat2_prob, 2, CAT2_MIN_VAL, cat2_cost },  // CATEGORY2_TOKEN
  { av1_cat3_prob, 3, CAT3_MIN_VAL, cat3_cost },  // CATEGORY3_TOKEN
  { av1_cat4_prob, 4, CAT4_MIN_VAL, cat4_cost },  // CATEGORY4_TOKEN
  { av1_cat5_prob, 5, CAT5_MIN_VAL, cat5_cost },  // CATEGORY5_TOKEN
243
  { av1_cat6_prob, 18, CAT6_MIN_VAL, 0 },         // CATEGORY6_TOKEN
Alex Converse's avatar
Alex Converse committed
244
  { 0, 0, 0, zero_cost }                          // EOB_TOKEN
Jingning Han's avatar
Jingning Han committed
245
246
};

247
#if !CONFIG_EC_MULTISYMBOL
Yaowu Xu's avatar
Yaowu Xu committed
248
const struct av1_token av1_coef_encodings[ENTROPY_TOKENS] = {
249
250
  { 2, 2 },  { 6, 3 },   { 28, 5 },  { 58, 6 },  { 59, 6 },  { 60, 6 },
  { 61, 6 }, { 124, 7 }, { 125, 7 }, { 126, 7 }, { 127, 7 }, { 0, 1 }
Jingning Han's avatar
Jingning Han committed
251
};
252
#endif  // !CONFIG_EC_MULTISYMBOL
Jingning Han's avatar
Jingning Han committed
253
254

struct tokenize_b_args {
255
  const AV1_COMP *cpi;
Jingning Han's avatar
Jingning Han committed
256
257
  ThreadData *td;
  TOKENEXTRA **tp;
258
  int this_rate;
Jingning Han's avatar
Jingning Han committed
259
260
};

261
#if !CONFIG_PVQ || CONFIG_VAR_TX
262
263
264
static void cost_coeffs_b(int plane, int block, int blk_row, int blk_col,
                          BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
  struct tokenize_b_args *const args = arg;
265
  const AV1_COMMON *cm = &args->cpi->common;
266
267
268
269
270
271
272
273
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
  struct macroblock_plane *p = &x->plane[plane];
  struct macroblockd_plane *pd = &xd->plane[plane];
  const PLANE_TYPE type = pd->plane_type;
  const int ref = is_inter_block(mbmi);
274
275
  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
  const TX_TYPE tx_type = get_tx_type(type, xd, block_raster_idx, tx_size);
Angie Chiang's avatar
Angie Chiang committed
276
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, ref);
277
278
279
  int rate = av1_cost_coeffs(cm, x, plane, block, tx_size, scan_order,
                             pd->above_context + blk_col,
                             pd->left_context + blk_row, 0);
280
  args->this_rate += rate;
281
  (void)plane_bsize;
282
283
  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
                   blk_row);
284
285
}

286
287
static void set_entropy_context_b(int plane, int block, int blk_row,
                                  int blk_col, BLOCK_SIZE plane_bsize,
Jingning Han's avatar
Jingning Han committed
288
                                  TX_SIZE tx_size, void *arg) {
289
  struct tokenize_b_args *const args = arg;
Jingning Han's avatar
Jingning Han committed
290
291
292
293
294
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  struct macroblock_plane *p = &x->plane[plane];
  struct macroblockd_plane *pd = &xd->plane[plane];
295
  (void)plane_bsize;
296
297
  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
                   blk_row);
Jingning Han's avatar
Jingning Han committed
298
299
}

300
#if CONFIG_NEW_TOKENSET
301
302
303
static INLINE void add_token(TOKENEXTRA **t,
                             aom_cdf_prob (*tail_cdf)[CDF_SIZE(ENTROPY_TOKENS)],
                             aom_cdf_prob (*head_cdf)[CDF_SIZE(ENTROPY_TOKENS)],
304
305
                             int eob_val, int first_val, int32_t extra,
                             uint8_t token) {
306
307
308
309
  (*t)->token = token;
  (*t)->extra = extra;
  (*t)->tail_cdf = tail_cdf;
  (*t)->head_cdf = head_cdf;
310
  (*t)->eob_val = eob_val;
311
  (*t)->first_val = first_val;
312
313
314
  (*t)++;
}

Angie Chiang's avatar
Angie Chiang committed
315
#else  // CONFIG_NEW_TOKENSET
316
317
static INLINE void add_token(
    TOKENEXTRA **t, const aom_prob *context_tree,
318
#if CONFIG_EC_MULTISYMBOL
319
    aom_cdf_prob (*token_cdf)[CDF_SIZE(ENTROPY_TOKENS)],
320
#endif  // CONFIG_EC_MULTISYMBOL
321
    int32_t extra, uint8_t token, uint8_t skip_eob_node, unsigned int *counts) {
Jingning Han's avatar
Jingning Han committed
322
323
324
  (*t)->token = token;
  (*t)->extra = extra;
  (*t)->context_tree = context_tree;
325
326
327
#if CONFIG_EC_MULTISYMBOL
  (*t)->token_cdf = token_cdf;
#endif  // CONFIG_EC_MULTISYMBOL
Jingning Han's avatar
Jingning Han committed
328
329
330
331
  (*t)->skip_eob_node = skip_eob_node;
  (*t)++;
  ++counts[token];
}
Angie Chiang's avatar
Angie Chiang committed
332
333
#endif  // CONFIG_NEW_TOKENSET
#endif  // !CONFIG_PVQ || CONFIG_VAR_TX
Jingning Han's avatar
Jingning Han committed
334

Fangwen Fu's avatar
Fangwen Fu committed
335
#if CONFIG_PALETTE
336
337
338
339
340
341
342
void av1_tokenize_palette_sb(const AV1_COMP *cpi,
                             const struct ThreadData *const td, int plane,
                             TOKENEXTRA **t, RUN_TYPE dry_run, BLOCK_SIZE bsize,
                             int *rate) {
  const MACROBLOCK *const x = &td->mb;
  const MACROBLOCKD *const xd = &x->e_mbd;
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
343
  const uint8_t *const color_map = xd->plane[plane].color_index_map;
344
  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
345
  const int n = pmi->palette_size[plane];
346
  int i, j;
347
  int this_rate = 0;
348
  uint8_t color_order[PALETTE_MAX_SIZE];
clang-format's avatar
clang-format committed
349
350
351
352
  const aom_prob(
      *const probs)[PALETTE_COLOR_INDEX_CONTEXTS][PALETTE_COLORS - 1] =
      plane == 0 ? av1_default_palette_y_color_index_prob
                 : av1_default_palette_uv_color_index_prob;
353
354
355
356
  int plane_block_width, rows, cols;
  av1_get_block_dimensions(bsize, plane, xd, &plane_block_width, NULL, &rows,
                           &cols);
  assert(plane == 0 || plane == 1);
hui su's avatar
hui su committed
357

Fangwen Fu's avatar
Fangwen Fu committed
358
359
360
361
362
363
#if CONFIG_PALETTE_THROUGHPUT
  int k;
  for (k = 1; k < rows + cols - 1; ++k) {
    for (j = AOMMIN(k, cols - 1); j >= AOMMAX(0, k - rows + 1); --j) {
      i = k - j;
#else
hui su's avatar
hui su committed
364
365
  for (i = 0; i < rows; ++i) {
    for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
Fangwen Fu's avatar
Fangwen Fu committed
366
#endif  // CONFIG_PALETTE_THROUGHPUT
367
      int color_new_idx;
368
      const int color_ctx = av1_get_palette_color_index_context(
369
          color_map, plane_block_width, i, j, n, color_order, &color_new_idx);
370
      assert(color_new_idx >= 0 && color_new_idx < n);
371
      if (dry_run == DRY_RUN_COSTCOEFFS)
372
373
        this_rate += cpi->palette_y_color_cost[n - PALETTE_MIN_SIZE][color_ctx]
                                              [color_new_idx];
374
      (*t)->token = color_new_idx;
375
      (*t)->context_tree = probs[n - PALETTE_MIN_SIZE][color_ctx];
hui su's avatar
hui su committed
376
377
378
379
      (*t)->skip_eob_node = 0;
      ++(*t);
    }
  }
380
  if (rate) *rate += this_rate;
hui su's avatar
hui su committed
381
}
382
#endif  // CONFIG_PALETTE
hui su's avatar
hui su committed
383

384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
#if CONFIG_PVQ
static void add_pvq_block(AV1_COMMON *const cm, MACROBLOCK *const x,
                          PVQ_INFO *pvq) {
  PVQ_QUEUE *q = x->pvq_q;
  if (q->curr_pos >= q->buf_len) {
    int new_buf_len = 2 * q->buf_len + 1;
    PVQ_INFO *new_buf;
    CHECK_MEM_ERROR(cm, new_buf, aom_malloc(new_buf_len * sizeof(PVQ_INFO)));
    memcpy(new_buf, q->buf, q->buf_len * sizeof(PVQ_INFO));
    aom_free(q->buf);
    q->buf = new_buf;
    q->buf_len = new_buf_len;
  }
  OD_COPY(q->buf + q->curr_pos, pvq, 1);
  ++q->curr_pos;
}

// NOTE: This does not actually generate tokens, instead we store the encoding
// decisions made for PVQ in a queue that we will read from when
// actually writing the bitstream in write_modes_b
static void tokenize_pvq(int plane, int block, int blk_row, int blk_col,
                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
  struct tokenize_b_args *const args = arg;
  const AV1_COMP *cpi = args->cpi;
  const AV1_COMMON *const cm = &cpi->common;
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  PVQ_INFO *pvq_info;

  (void)block;
  (void)blk_row;
  (void)blk_col;
  (void)plane_bsize;
  (void)tx_size;

  assert(block < MAX_PVQ_BLOCKS_IN_SB);
  pvq_info = &x->pvq[block][plane];
  add_pvq_block((AV1_COMMON * const)cm, x, pvq_info);
}
#endif  // CONFIG_PVQ

425
static void tokenize_b(int plane, int block, int blk_row, int blk_col,
426
                       BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
427
#if !CONFIG_PVQ
428
  struct tokenize_b_args *const args = arg;
429
  const AV1_COMP *cpi = args->cpi;
Angie Chiang's avatar
Angie Chiang committed
430
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
431
432
433
434
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  TOKENEXTRA **tp = args->tp;
435
  uint8_t token_cache[MAX_TX_SQUARE];
Jingning Han's avatar
Jingning Han committed
436
437
438
439
440
  struct macroblock_plane *p = &x->plane[plane];
  struct macroblockd_plane *pd = &xd->plane[plane];
  MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
  int pt; /* near block/prev token context index */
  int c;
441
  TOKENEXTRA *t = *tp; /* store tokens starting here */
442
  const int eob = p->eobs[block];
Jingning Han's avatar
Jingning Han committed
443
444
  const PLANE_TYPE type = pd->plane_type;
  const tran_low_t *qcoeff = BLOCK_OFFSET(p->qcoeff, block);
445
#if CONFIG_SUPERTX
Yaowu Xu's avatar
Yaowu Xu committed
446
  const int segment_id = AOMMIN(mbmi->segment_id, mbmi->segment_id_supertx);
447
#else
Jingning Han's avatar
Jingning Han committed
448
  const int segment_id = mbmi->segment_id;
449
#endif  // CONFIG_SUEPRTX
Jingning Han's avatar
Jingning Han committed
450
  const int16_t *scan, *nb;
451
452
  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
  const TX_TYPE tx_type = get_tx_type(type, xd, block_raster_idx, tx_size);
453
  const SCAN_ORDER *const scan_order =
Angie Chiang's avatar
Angie Chiang committed
454
      get_scan(cm, tx_size, tx_type, is_inter_block(mbmi));
Jingning Han's avatar
Jingning Han committed
455
  const int ref = is_inter_block(mbmi);
clang-format's avatar
clang-format committed
456
  unsigned int(*const counts)[COEFF_CONTEXTS][ENTROPY_TOKENS] =
457
      td->rd_counts.coef_counts[txsize_sqr_map[tx_size]][type][ref];
458
#if !CONFIG_NEW_TOKENSET
hui su's avatar
hui su committed
459
#if CONFIG_SUBFRAME_PROB_UPDATE
460
  const aom_prob(*coef_probs)[COEFF_CONTEXTS][UNCONSTRAINED_NODES] =
hui su's avatar
hui su committed
461
      cpi->subframe_stats.coef_probs_buf[cpi->common.coef_probs_update_idx]
462
                                        [txsize_sqr_map[tx_size]][type][ref];
hui su's avatar
hui su committed
463
#else
clang-format's avatar
clang-format committed
464
  aom_prob(*const coef_probs)[COEFF_CONTEXTS][UNCONSTRAINED_NODES] =
465
      cpi->common.fc->coef_probs[txsize_sqr_map[tx_size]][type][ref];
hui su's avatar
hui su committed
466
#endif  // CONFIG_SUBFRAME_PROB_UPDATE
467
#endif  // !CONFIG_NEW_TOKENSET
468
#if CONFIG_EC_ADAPT
469
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
470
#elif CONFIG_EC_MULTISYMBOL
471
472
  FRAME_CONTEXT *ec_ctx = cpi->common.fc;
#endif
473
#if CONFIG_NEW_TOKENSET
474
  aom_cdf_prob(
475
      *const coef_head_cdfs)[COEFF_CONTEXTS][CDF_SIZE(ENTROPY_TOKENS)] =
476
      ec_ctx->coef_head_cdfs[txsize_sqr_map[tx_size]][type][ref];
477
  aom_cdf_prob(
478
      *const coef_tail_cdfs)[COEFF_CONTEXTS][CDF_SIZE(ENTROPY_TOKENS)] =
479
      ec_ctx->coef_tail_cdfs[txsize_sqr_map[tx_size]][type][ref];
480
481
  unsigned int(*const blockz_count)[2] =
      td->counts->blockz_count[txsize_sqr_map[tx_size]][type][ref];
482
  int eob_val;
483
  int first_val = 1;
484
#else
485
#if CONFIG_EC_MULTISYMBOL
486
  aom_cdf_prob(*const coef_cdfs)[COEFF_CONTEXTS][CDF_SIZE(ENTROPY_TOKENS)] =
487
      ec_ctx->coef_cdfs[txsize_sqr_map[tx_size]][type][ref];
488
#endif
489
490
  int skip_eob = 0;
#endif
491
  const int seg_eob = get_tx_eob(&cpi->common.seg, segment_id, tx_size);
clang-format's avatar
clang-format committed
492
  unsigned int(*const eob_branch)[COEFF_CONTEXTS] =
493
      td->counts->eob_branch[txsize_sqr_map[tx_size]][type][ref];
Jingning Han's avatar
Jingning Han committed
494
495
496
  const uint8_t *const band = get_band_translate(tx_size);
  int16_t token;
  EXTRABIT extra;
497
  (void)plane_bsize;
498
499
  pt = get_entropy_context(tx_size, pd->above_context + blk_col,
                           pd->left_context + blk_row);
500
501
  scan = scan_order->scan;
  nb = scan_order->neighbors;
Jingning Han's avatar
Jingning Han committed
502
503
  c = 0;

504
#if CONFIG_NEW_TOKENSET
505
  if (eob == 0)
506
507
    add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt], 1,
              1, 0, BLOCK_Z_TOKEN);
508

509
  ++blockz_count[pt][eob != 0];
510

Jingning Han's avatar
Jingning Han committed
511
  while (c < eob) {
512
    int v = qcoeff[scan[c]];
513
    first_val = (c == 0);
Jingning Han's avatar
Jingning Han committed
514

515
516
    if (!v) {
      add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt],
517
                0, first_val, 0, ZERO_TOKEN);
518
519
520
      ++counts[band[c]][pt][ZERO_TOKEN];
      token_cache[scan[c]] = 0;
    } else {
521
522
      eob_val =
          (c + 1 == eob) ? (c + 1 == seg_eob ? LAST_EOB : EARLY_EOB) : NO_EOB;
Jingning Han's avatar
Jingning Han committed
523

524
      av1_get_token_extra(v, &token, &extra);
Jingning Han's avatar
Jingning Han committed
525

526
      add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt],
527
                eob_val, first_val, extra, (uint8_t)token);
528

529
530
531
532
533
      if (eob_val != LAST_EOB) {
        ++counts[band[c]][pt][token];
        ++eob_branch[band[c]][pt];
        counts[band[c]][pt][EOB_TOKEN] += eob_val != NO_EOB;
      }
534
535
536

      token_cache[scan[c]] = av1_pt_energy_class[token];
    }
Jingning Han's avatar
Jingning Han committed
537
    ++c;
538
    pt = get_coef_context(nb, token_cache, AOMMIN(c, eob - 1));
Jingning Han's avatar
Jingning Han committed
539
  }
540
541
542
543
544
545
546
#else
  while (c < eob) {
    const int v = qcoeff[scan[c]];
    eob_branch[band[c]][pt] += !skip_eob;

    av1_get_token_extra(v, &token, &extra);

547
548
549
550
551
    add_token(&t, coef_probs[band[c]][pt],
#if CONFIG_EC_MULTISYMBOL
              &coef_cdfs[band[c]][pt],
#endif
              extra, (uint8_t)token, (uint8_t)skip_eob, counts[band[c]][pt]);
552
553
554
555
556
557
558

    token_cache[scan[c]] = av1_pt_energy_class[token];
    ++c;
    pt = get_coef_context(nb, token_cache, c);
    skip_eob = (token == ZERO_TOKEN);
  }
  if (c < seg_eob) {
559
560
561
562
563
    add_token(&t, coef_probs[band[c]][pt],
#if CONFIG_EC_MULTISYMBOL
              NULL,
#endif
              0, EOB_TOKEN, 0, counts[band[c]][pt]);
Jingning Han's avatar
Jingning Han committed
564
565
    ++eob_branch[band[c]][pt];
  }
566
#endif  // CONFIG_NEW_TOKENSET
Jingning Han's avatar
Jingning Han committed
567

568
569
570
571
572
#if CONFIG_COEF_INTERLEAVE
  t->token = EOSB_TOKEN;
  t++;
#endif

Jingning Han's avatar
Jingning Han committed
573
574
  *tp = t;

Angie Chiang's avatar
Angie Chiang committed
575
576
577
578
579
#if CONFIG_ADAPT_SCAN
  // Since dqcoeff is not available here, we pass qcoeff into
  // av1_update_scan_count_facade(). The update behavior should be the same
  // because av1_update_scan_count_facade() only cares if coefficients are zero
  // or not.
580
581
  av1_update_scan_count_facade((AV1_COMMON *)cm, td->counts, tx_size, tx_type,
                               qcoeff, c);
Angie Chiang's avatar
Angie Chiang committed
582
583
#endif

584
  av1_set_contexts(xd, pd, plane, tx_size, c > 0, blk_col, blk_row);
585
586
587
#else   // !CONFIG_PVQ
  tokenize_pvq(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
#endif  // !CONFIG_PVQ
Jingning Han's avatar
Jingning Han committed
588
}
Fangwen Fu's avatar
Fangwen Fu committed
589

Jingning Han's avatar
Jingning Han committed
590
591
592
593
struct is_skippable_args {
  uint16_t *eobs;
  int *skippable;
};
594
static void is_skippable(int plane, int block, int blk_row, int blk_col,
595
                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *argv) {
Jingning Han's avatar
Jingning Han committed
596
597
598
599
  struct is_skippable_args *args = argv;
  (void)plane;
  (void)plane_bsize;
  (void)tx_size;
600
601
  (void)blk_row;
  (void)blk_col;
Jingning Han's avatar
Jingning Han committed
602
603
604
605
  args->skippable[0] &= (!args->eobs[block]);
}

// TODO(yaowu): rewrite and optimize this function to remove the usage of
Yaowu Xu's avatar
Yaowu Xu committed
606
607
//              av1_foreach_transform_block() and simplify is_skippable().
int av1_is_skippable_in_plane(MACROBLOCK *x, BLOCK_SIZE bsize, int plane) {
Jingning Han's avatar
Jingning Han committed
608
  int result = 1;
609
  struct is_skippable_args args = { x->plane[plane].eobs, &result };
Yaowu Xu's avatar
Yaowu Xu committed
610
611
  av1_foreach_transformed_block_in_plane(&x->e_mbd, bsize, plane, is_skippable,
                                         &args);
Jingning Han's avatar
Jingning Han committed
612
613
614
  return result;
}

615
#if CONFIG_VAR_TX
616
void tokenize_vartx(ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,
617
618
                    TX_SIZE tx_size, BLOCK_SIZE plane_bsize, int blk_row,
                    int blk_col, int block, int plane, void *arg) {
619
620
621
622
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  const struct macroblockd_plane *const pd = &xd->plane[plane];
623
  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
624
625
  const int tx_row = blk_row >> (1 - pd->subsampling_y);
  const int tx_col = blk_col >> (1 - pd->subsampling_x);
626
627
  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
Debargha Mukherjee's avatar
Debargha Mukherjee committed
628
  TX_SIZE plane_tx_size;
629

630
  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
631

632
633
634
  plane_tx_size =
      plane ? uv_txsize_lookup[bsize][mbmi->inter_tx_size[tx_row][tx_col]][0][0]
            : mbmi->inter_tx_size[tx_row][tx_col];
Debargha Mukherjee's avatar
Debargha Mukherjee committed
635

636
  if (tx_size == plane_tx_size) {
637
    plane_bsize = get_plane_block_size(mbmi->sb_type, pd);
638
639
    if (!dry_run)
      tokenize_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
640
    else if (dry_run == DRY_RUN_NORMAL)
641
642
      set_entropy_context_b(plane, block, blk_row, blk_col, plane_bsize,
                            tx_size, arg);
643
644
    else if (dry_run == DRY_RUN_COSTCOEFFS)
      cost_coeffs_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
645
  } else {
646
    // Half the block size in transform block unit.
647
648
    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
    const int bsl = tx_size_wide_unit[sub_txs];
649
650
651
652
653
    int i;

    assert(bsl > 0);

    for (i = 0; i < 4; ++i) {
654
655
656
657
      const int offsetr = blk_row + ((i >> 1) * bsl);
      const int offsetc = blk_col + ((i & 0x01) * bsl);

      int step = tx_size_wide_unit[sub_txs] * tx_size_high_unit[sub_txs];
658

659
      if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;
660

661
      tokenize_vartx(td, t, dry_run, sub_txs, plane_bsize, offsetr, offsetc,
662
663
                     block, plane, arg);
      block += step;
664
665
666
667
    }
  }
}

668
void av1_tokenize_sb_vartx(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
669
670
                           RUN_TYPE dry_run, int mi_row, int mi_col,
                           BLOCK_SIZE bsize, int *rate) {
671
  const AV1_COMMON *const cm = &cpi->common;
672
673
674
675
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  TOKENEXTRA *t_backup = *t;
Yaowu Xu's avatar
Yaowu Xu committed
676
  const int ctx = av1_get_skip_context(xd);
677
678
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
679
  struct tokenize_b_args arg = { cpi, td, t, 0 };
680
  int plane;
681
  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return;
682
683

  if (mbmi->skip) {
684
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
685
    reset_skip_context(xd, bsize);
686
    if (dry_run) *t = t_backup;
687
688
689
690
691
692
693
694
695
696
697
    return;
  }

  if (!dry_run)
    td->counts->skip[ctx][0] += skip_inc;
  else
    *t = t_backup;

  for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
    const struct macroblockd_plane *const pd = &xd->plane[plane];
    const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
698
699
    const int mi_width = block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
    const int mi_height = block_size_high[plane_bsize] >> tx_size_wide_log2[0];
700
    const TX_SIZE max_tx_size = max_txsize_rect_lookup[plane_bsize];
701
    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
702
703
    int bw = block_size_wide[txb_size] >> tx_size_wide_log2[0];
    int bh = block_size_high[txb_size] >> tx_size_wide_log2[0];
704
705
    int idx, idy;
    int block = 0;
706
    int step = tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];
707
    for (idy = 0; idy < mi_height; idy += bh) {
708
      for (idx = 0; idx < mi_width; idx += bw) {
709
710
        tokenize_vartx(td, t, dry_run, max_tx_size, plane_bsize, idy, idx,
                       block, plane, &arg);
711
712
713
        block += step;
      }
    }
714
715
716
717
718

    if (!dry_run) {
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
719
  }
720
  if (rate) *rate += arg.this_rate;
721
}
722
#endif  // CONFIG_VAR_TX
723

724
void av1_tokenize_sb(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
725
726
                     RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
                     const int mi_row, const int mi_col) {
727
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
728
729
730
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
Yaowu Xu's avatar
Yaowu Xu committed
731
  const int ctx = av1_get_skip_context(xd);
732
733
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
734
  struct tokenize_b_args arg = { cpi, td, t, 0 };
Jingning Han's avatar
Jingning Han committed
735
  if (mbmi->skip) {
736
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
Jingning Han's avatar
Jingning Han committed
737
738
739
740
741
    reset_skip_context(xd, bsize);
    return;
  }

  if (!dry_run) {
742
743
744
745
#if CONFIG_COEF_INTERLEAVE
    td->counts->skip[ctx][0] += skip_inc;
    av1_foreach_transformed_block_interleave(xd, bsize, tokenize_b, &arg);
#else
746
747
    int plane;

Jingning Han's avatar
Jingning Han committed
748
    td->counts->skip[ctx][0] += skip_inc;
749
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
750
751
#if CONFIG_CB4X4
      if (bsize < BLOCK_8X8 && plane && !is_chroma_reference(mi_row, mi_col)) {
752
#if !CONFIG_PVQ
753
754
        (*t)->token = EOSB_TOKEN;
        (*t)++;
755
#endif
756
757
758
759
760
761
        continue;
      }
#else
      (void)mi_row;
      (void)mi_col;
#endif
Yaowu Xu's avatar
Yaowu Xu committed
762
763
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_b,
                                             &arg);
764
#if !CONFIG_PVQ
765
766
      (*t)->token = EOSB_TOKEN;
      (*t)++;
767
#endif  // !CONFIG_PVQ
768
    }
769
#endif
770
771
772
  }
#if !CONFIG_PVQ
  else if (dry_run == DRY_RUN_NORMAL) {
773
    int plane;
774
775
776
777
778
779
780
781
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
#if CONFIG_CB4X4
      if (bsize < BLOCK_8X8 && plane && !is_chroma_reference(mi_row, mi_col))
        continue;
#else
      (void)mi_row;
      (void)mi_col;
#endif
782
783
      av1_foreach_transformed_block_in_plane(xd, bsize, plane,
                                             set_entropy_context_b, &arg);
784
    }
785
  } else if (dry_run == DRY_RUN_COSTCOEFFS) {
786
    int plane;
787
788
789
790
791
792
793
794
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
#if CONFIG_CB4X4
      if (bsize < BLOCK_8X8 && plane && !is_chroma_reference(mi_row, mi_col))
        continue;
#else
      (void)mi_row;
      (void)mi_col;
#endif
795
796
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, cost_coeffs_b,
                                             &arg);
797
    }
Jingning Han's avatar
Jingning Han committed
798
  }
799
#endif  // !CONFIG_PVQ
800

801
  if (rate) *rate += arg.this_rate;
Jingning Han's avatar
Jingning Han committed
802
}
803
804

#if CONFIG_SUPERTX
805
806
807
808
void av1_tokenize_sb_supertx(const AV1_COMP *cpi, ThreadData *td,
                             TOKENEXTRA **t, RUN_TYPE dry_run, BLOCK_SIZE bsize,
                             int *rate) {
  const AV1_COMMON *const cm = &cpi->common;
809
810
811
  MACROBLOCKD *const xd = &td->mb.e_mbd;
  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  TOKENEXTRA *t_backup = *t;
Yaowu Xu's avatar
Yaowu Xu committed
812
  const int ctx = av1_get_skip_context(xd);
813
814
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id_supertx, SEG_LVL_SKIP);
815
  struct tokenize_b_args arg = { cpi, td, t, 0 };
816
  if (mbmi->skip) {
817
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
818
    reset_skip_context(xd, bsize);
819
    if (dry_run) *t = t_backup;
820
821
822
823
824
825
826
827
    return;
  }

  if (!dry_run) {
    int plane;
    td->counts->skip[ctx][0] += skip_inc;

    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
Yaowu Xu's avatar
Yaowu Xu committed
828
829
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_b,
                                             &arg);
830
831
832
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
833
  } else if (dry_run == DRY_RUN_NORMAL) {
834
835
836
837
    int plane;
    for (plane = 0; plane < MAX_MB_PLANE; ++plane)
      av1_foreach_transformed_block_in_plane(xd, bsize, plane,
                                             set_entropy_context_b, &arg);
838
    *t = t_backup;
839
  } else if (dry_run == DRY_RUN_COSTCOEFFS) {
840
841
842
843
    int plane;
    for (plane = 0; plane < MAX_MB_PLANE; ++plane)
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, cost_coeffs_b,
                                             &arg);
844
  }
845
  if (rate) *rate += arg.this_rate;
846
847
}
#endif  // CONFIG_SUPERTX