tokenize.c 36.6 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
};
Yaowu Xu's avatar
Yaowu Xu committed
129
const int av1_cat6_high_cost[64] = {
130
131
132
133
134
135
136
  88,    2251,  2727,  4890,  3148,  5311,  5787,  7950,  3666,  5829,  6305,
  8468,  6726,  8889,  9365,  11528, 3666,  5829,  6305,  8468,  6726,  8889,
  9365,  11528, 7244,  9407,  9883,  12046, 10304, 12467, 12943, 15106, 3666,
  5829,  6305,  8468,  6726,  8889,  9365,  11528, 7244,  9407,  9883,  12046,
  10304, 12467, 12943, 15106, 7244,  9407,  9883,  12046, 10304, 12467, 12943,
  15106, 10822, 12985, 13461, 15624, 13882, 16045, 16521, 18684
};
Jingning Han's avatar
Jingning Han committed
137

Yaowu Xu's avatar
Yaowu Xu committed
138
139
#if CONFIG_AOM_HIGHBITDEPTH
const int av1_cat6_high10_high_cost[256] = {
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
  94,    2257,  2733,  4896,  3154,  5317,  5793,  7956,  3672,  5835,  6311,
  8474,  6732,  8895,  9371,  11534, 3672,  5835,  6311,  8474,  6732,  8895,
  9371,  11534, 7250,  9413,  9889,  12052, 10310, 12473, 12949, 15112, 3672,
  5835,  6311,  8474,  6732,  8895,  9371,  11534, 7250,  9413,  9889,  12052,
  10310, 12473, 12949, 15112, 7250,  9413,  9889,  12052, 10310, 12473, 12949,
  15112, 10828, 12991, 13467, 15630, 13888, 16051, 16527, 18690, 4187,  6350,
  6826,  8989,  7247,  9410,  9886,  12049, 7765,  9928,  10404, 12567, 10825,
  12988, 13464, 15627, 7765,  9928,  10404, 12567, 10825, 12988, 13464, 15627,
  11343, 13506, 13982, 16145, 14403, 16566, 17042, 19205, 7765,  9928,  10404,
  12567, 10825, 12988, 13464, 15627, 11343, 13506, 13982, 16145, 14403, 16566,
  17042, 19205, 11343, 13506, 13982, 16145, 14403, 16566, 17042, 19205, 14921,
  17084, 17560, 19723, 17981, 20144, 20620, 22783, 4187,  6350,  6826,  8989,
  7247,  9410,  9886,  12049, 7765,  9928,  10404, 12567, 10825, 12988, 13464,
  15627, 7765,  9928,  10404, 12567, 10825, 12988, 13464, 15627, 11343, 13506,
  13982, 16145, 14403, 16566, 17042, 19205, 7765,  9928,  10404, 12567, 10825,
  12988, 13464, 15627, 11343, 13506, 13982, 16145, 14403, 16566, 17042, 19205,
  11343, 13506, 13982, 16145, 14403, 16566, 17042, 19205, 14921, 17084, 17560,
  19723, 17981, 20144, 20620, 22783, 8280,  10443, 10919, 13082, 11340, 13503,
  13979, 16142, 11858, 14021, 14497, 16660, 14918, 17081, 17557, 19720, 11858,
  14021, 14497, 16660, 14918, 17081, 17557, 19720, 15436, 17599, 18075, 20238,
  18496, 20659, 21135, 23298, 11858, 14021, 14497, 16660, 14918, 17081, 17557,
  19720, 15436, 17599, 18075, 20238, 18496, 20659, 21135, 23298, 15436, 17599,
  18075, 20238, 18496, 20659, 21135, 23298, 19014, 21177, 21653, 23816, 22074,
  24237, 24713, 26876
};
Yaowu Xu's avatar
Yaowu Xu committed
165
const int av1_cat6_high12_high_cost[1024] = {
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
  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,
  15118, 10834, 12997, 13473, 15636, 13894, 16057, 16533, 18696, 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,  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, 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, 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, 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, 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, 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, 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, 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
261
262
#endif

Yaowu Xu's avatar
Yaowu Xu committed
263
const av1_extra_bit av1_extra_bits[ENTROPY_TOKENS] = {
Alex Converse's avatar
Alex Converse committed
264
265
266
267
268
269
270
271
272
273
274
275
  { 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
  { av1_cat6_prob, 14, CAT6_MIN_VAL, 0 },         // CATEGORY6_TOKEN
  { 0, 0, 0, zero_cost }                          // EOB_TOKEN
Jingning Han's avatar
Jingning Han committed
276
277
};

Yaowu Xu's avatar
Yaowu Xu committed
278
279
#if CONFIG_AOM_HIGHBITDEPTH
const av1_extra_bit av1_extra_bits_high10[ENTROPY_TOKENS] = {
Alex Converse's avatar
Alex Converse committed
280
281
282
283
284
285
286
287
288
289
290
291
  { 0, 0, 0, zero_cost },                                // ZERO
  { 0, 0, 1, sign_cost },                                // ONE
  { 0, 0, 2, sign_cost },                                // TWO
  { 0, 0, 3, sign_cost },                                // THREE
  { 0, 0, 4, sign_cost },                                // FOUR
  { av1_cat1_prob_high10, 1, CAT1_MIN_VAL, cat1_cost },  // CAT1
  { av1_cat2_prob_high10, 2, CAT2_MIN_VAL, cat2_cost },  // CAT2
  { av1_cat3_prob_high10, 3, CAT3_MIN_VAL, cat3_cost },  // CAT3
  { av1_cat4_prob_high10, 4, CAT4_MIN_VAL, cat4_cost },  // CAT4
  { av1_cat5_prob_high10, 5, CAT5_MIN_VAL, cat5_cost },  // CAT5
  { av1_cat6_prob_high10, 16, CAT6_MIN_VAL, 0 },         // CAT6
  { 0, 0, 0, zero_cost }                                 // EOB
Jingning Han's avatar
Jingning Han committed
292
};
Yaowu Xu's avatar
Yaowu Xu committed
293
const av1_extra_bit av1_extra_bits_high12[ENTROPY_TOKENS] = {
Alex Converse's avatar
Alex Converse committed
294
295
296
297
298
299
300
301
302
303
304
305
  { 0, 0, 0, zero_cost },                                // ZERO
  { 0, 0, 1, sign_cost },                                // ONE
  { 0, 0, 2, sign_cost },                                // TWO
  { 0, 0, 3, sign_cost },                                // THREE
  { 0, 0, 4, sign_cost },                                // FOUR
  { av1_cat1_prob_high12, 1, CAT1_MIN_VAL, cat1_cost },  // CAT1
  { av1_cat2_prob_high12, 2, CAT2_MIN_VAL, cat2_cost },  // CAT2
  { av1_cat3_prob_high12, 3, CAT3_MIN_VAL, cat3_cost },  // CAT3
  { av1_cat4_prob_high12, 4, CAT4_MIN_VAL, cat4_cost },  // CAT4
  { av1_cat5_prob_high12, 5, CAT5_MIN_VAL, cat5_cost },  // CAT5
  { av1_cat6_prob_high12, 18, CAT6_MIN_VAL, 0 },         // CAT6
  { 0, 0, 0, zero_cost }                                 // EOB
Jingning Han's avatar
Jingning Han committed
306
307
308
};
#endif

309
#if !CONFIG_RANS
Yaowu Xu's avatar
Yaowu Xu committed
310
const struct av1_token av1_coef_encodings[ENTROPY_TOKENS] = {
311
312
  { 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
313
};
314
#endif  // !CONFIG_RANS
Jingning Han's avatar
Jingning Han committed
315
316

struct tokenize_b_args {
317
  const AV1_COMP *cpi;
Jingning Han's avatar
Jingning Han committed
318
319
  ThreadData *td;
  TOKENEXTRA **tp;
320
  int this_rate;
Jingning Han's avatar
Jingning Han committed
321
322
};

323
#if !CONFIG_PVQ
324
325
326
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;
327
  const AV1_COMMON *cm = &args->cpi->common;
328
329
330
331
332
333
334
335
336
  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);
  const TX_TYPE tx_type = get_tx_type(type, xd, block, tx_size);
Angie Chiang's avatar
Angie Chiang committed
337
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, ref);
338
339
  int pt = get_entropy_context(tx_size, pd->above_context + blk_col,
                               pd->left_context + blk_row);
340
  int rate = av1_cost_coeffs(cm, x, plane, block, pt, tx_size, scan_order->scan,
341
                             scan_order->neighbors, 0);
342
  args->this_rate += rate;
343
344
  (void)plane_bsize;
  av1_set_contexts(xd, pd, tx_size, p->eobs[block] > 0, blk_col, blk_row);
345
346
}

347
348
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
349
                                  TX_SIZE tx_size, void *arg) {
350
  struct tokenize_b_args *const args = arg;
Jingning Han's avatar
Jingning Han committed
351
352
353
354
355
  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];
356
357
  (void)plane_bsize;
  av1_set_contexts(xd, pd, tx_size, p->eobs[block] > 0, blk_col, blk_row);
Jingning Han's avatar
Jingning Han committed
358
359
}

Yaowu Xu's avatar
Yaowu Xu committed
360
static INLINE void add_token(TOKENEXTRA **t, const aom_prob *context_tree,
361
#if CONFIG_EC_MULTISYMBOL
362
                             aom_cdf_prob (*token_cdf)[ENTROPY_TOKENS],
363
#endif  // CONFIG_EC_MULTISYMBOL
Jingning Han's avatar
Jingning Han committed
364
                             int32_t extra, uint8_t token,
365
                             uint8_t skip_eob_node, unsigned int *counts) {
Jingning Han's avatar
Jingning Han committed
366
367
368
  (*t)->token = token;
  (*t)->extra = extra;
  (*t)->context_tree = context_tree;
369
#if CONFIG_EC_MULTISYMBOL
370
  (*t)->token_cdf = token_cdf;
371
#endif  // CONFIG_EC_MULTISYMBOL
Jingning Han's avatar
Jingning Han committed
372
373
374
375
376
377
378
  (*t)->skip_eob_node = skip_eob_node;
  (*t)++;
  ++counts[token];
}

static INLINE int get_tx_eob(const struct segmentation *seg, int segment_id,
                             TX_SIZE tx_size) {
379
  const int eob_max = tx_size_2d[tx_size];
Jingning Han's avatar
Jingning Han committed
380
381
382
  return segfeature_active(seg, segment_id, SEG_LVL_SKIP) ? 0 : eob_max;
}

383
#if CONFIG_PALETTE
384
385
386
387
388
389
390
391
392
393
394
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;
  const uint8_t *const color_map = xd->plane[plane != 0].color_index_map;
  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
  const int n = pmi->palette_size[plane != 0];
  int i, j;
395
  int this_rate = 0;
396
  uint8_t color_order[PALETTE_MAX_SIZE];
397
  const int rows = (4 * num_4x4_blocks_high_lookup[bsize]) >>
398
                   (xd->plane[plane != 0].subsampling_y);
399
  const int cols = (4 * num_4x4_blocks_wide_lookup[bsize]) >>
400
                   (xd->plane[plane != 0].subsampling_x);
clang-format's avatar
clang-format committed
401
  const aom_prob(*const probs)[PALETTE_COLOR_CONTEXTS][PALETTE_COLORS - 1] =
Yaowu Xu's avatar
Yaowu Xu committed
402
403
      plane == 0 ? av1_default_palette_y_color_prob
                 : av1_default_palette_uv_color_prob;
hui su's avatar
hui su committed
404
405
406

  for (i = 0; i < rows; ++i) {
    for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
407
408
409
410
      int color_new_idx;
      const int color_ctx = av1_get_palette_color_context(
          color_map, cols, i, j, n, color_order, &color_new_idx);
      assert(color_new_idx >= 0 && color_new_idx < n);
411
      if (dry_run == DRY_RUN_COSTCOEFFS)
412
413
        this_rate += cpi->palette_y_color_cost[n - 2][color_ctx][color_new_idx];
      (*t)->token = color_new_idx;
414
      (*t)->context_tree = probs[n - 2][color_ctx];
hui su's avatar
hui su committed
415
416
417
418
      (*t)->skip_eob_node = 0;
      ++(*t);
    }
  }
419
  if (rate) *rate += this_rate;
hui su's avatar
hui su committed
420
}
421
#endif  // CONFIG_PALETTE
hui su's avatar
hui su committed
422

423
static void tokenize_b(int plane, int block, int blk_row, int blk_col,
424
425
                       BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
  struct tokenize_b_args *const args = arg;
426
  const AV1_COMP *cpi = args->cpi;
Angie Chiang's avatar
Angie Chiang committed
427
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
428
429
430
431
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  TOKENEXTRA **tp = args->tp;
432
  uint8_t token_cache[MAX_TX_SQUARE];
Jingning Han's avatar
Jingning Han committed
433
434
435
436
437
  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;
438
  TOKENEXTRA *t = *tp; /* store tokens starting here */
Jingning Han's avatar
Jingning Han committed
439
440
441
  int eob = p->eobs[block];
  const PLANE_TYPE type = pd->plane_type;
  const tran_low_t *qcoeff = BLOCK_OFFSET(p->qcoeff, block);
442
#if CONFIG_SUPERTX
Yaowu Xu's avatar
Yaowu Xu committed
443
  const int segment_id = AOMMIN(mbmi->segment_id, mbmi->segment_id_supertx);
444
#else
Jingning Han's avatar
Jingning Han committed
445
  const int segment_id = mbmi->segment_id;
446
#endif  // CONFIG_SUEPRTX
Jingning Han's avatar
Jingning Han committed
447
  const int16_t *scan, *nb;
hui su's avatar
hui su committed
448
  const TX_TYPE tx_type = get_tx_type(type, xd, block, tx_size);
449
  const SCAN_ORDER *const scan_order =
Angie Chiang's avatar
Angie Chiang committed
450
      get_scan(cm, tx_size, tx_type, is_inter_block(mbmi));
Jingning Han's avatar
Jingning Han committed
451
  const int ref = is_inter_block(mbmi);
clang-format's avatar
clang-format committed
452
  unsigned int(*const counts)[COEFF_CONTEXTS][ENTROPY_TOKENS] =
453
      td->rd_counts.coef_counts[txsize_sqr_map[tx_size]][type][ref];
hui su's avatar
hui su committed
454
#if CONFIG_ENTROPY
455
  const aom_prob(*coef_probs)[COEFF_CONTEXTS][UNCONSTRAINED_NODES] =
hui su's avatar
hui su committed
456
      cpi->subframe_stats.coef_probs_buf[cpi->common.coef_probs_update_idx]
457
                                        [txsize_sqr_map[tx_size]][type][ref];
hui su's avatar
hui su committed
458
#else
clang-format's avatar
clang-format committed
459
  aom_prob(*const coef_probs)[COEFF_CONTEXTS][UNCONSTRAINED_NODES] =
460
      cpi->common.fc->coef_probs[txsize_sqr_map[tx_size]][type][ref];
hui su's avatar
hui su committed
461
#endif  // CONFIG_ENTROPY
462
#if CONFIG_EC_MULTISYMBOL
Alex Converse's avatar
Alex Converse committed
463
464
  aom_cdf_prob(*const coef_cdfs)[COEFF_CONTEXTS][ENTROPY_TOKENS] =
      cpi->common.fc->coef_cdfs[tx_size][type][ref];
465
#endif
clang-format's avatar
clang-format committed
466
  unsigned int(*const eob_branch)[COEFF_CONTEXTS] =
467
      td->counts->eob_branch[txsize_sqr_map[tx_size]][type][ref];
Jingning Han's avatar
Jingning Han committed
468
469
  const uint8_t *const band = get_band_translate(tx_size);
  const int seg_eob = get_tx_eob(&cpi->common.seg, segment_id, tx_size);
470
  int skip_eob = 0;
Jingning Han's avatar
Jingning Han committed
471
472
  int16_t token;
  EXTRABIT extra;
473
  (void)plane_bsize;
474
475
  pt = get_entropy_context(tx_size, pd->above_context + blk_col,
                           pd->left_context + blk_row);
476
477
  scan = scan_order->scan;
  nb = scan_order->neighbors;
Jingning Han's avatar
Jingning Han committed
478
479
480
  c = 0;

  while (c < eob) {
481
482
    const int v = qcoeff[scan[c]];
    eob_branch[band[c]][pt] += !skip_eob;
Jingning Han's avatar
Jingning Han committed
483

Yaowu Xu's avatar
Yaowu Xu committed
484
    av1_get_token_extra(v, &token, &extra);
Jingning Han's avatar
Jingning Han committed
485

486
    add_token(&t, coef_probs[band[c]][pt],
487
#if CONFIG_EC_MULTISYMBOL
488
              &coef_cdfs[band[c]][pt],
489
#endif
490
              extra, (uint8_t)token, (uint8_t)skip_eob, counts[band[c]][pt]);
Jingning Han's avatar
Jingning Han committed
491

Yaowu Xu's avatar
Yaowu Xu committed
492
    token_cache[scan[c]] = av1_pt_energy_class[token];
Jingning Han's avatar
Jingning Han committed
493
494
    ++c;
    pt = get_coef_context(nb, token_cache, c);
495
    skip_eob = (token == ZERO_TOKEN);
Jingning Han's avatar
Jingning Han committed
496
497
  }
  if (c < seg_eob) {
Alex Converse's avatar
Alex Converse committed
498
    add_token(&t, coef_probs[band[c]][pt],
499
#if CONFIG_EC_MULTISYMBOL
Alex Converse's avatar
Alex Converse committed
500
501
502
              NULL,
#endif
              0, EOB_TOKEN, 0, counts[band[c]][pt]);
Jingning Han's avatar
Jingning Han committed
503
504
505
506
507
    ++eob_branch[band[c]][pt];
  }

  *tp = t;

Angie Chiang's avatar
Angie Chiang committed
508
509
510
511
512
513
514
515
#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.
  av1_update_scan_count_facade((AV1_COMMON *)cm, tx_size, tx_type, qcoeff, c);
#endif

516
  av1_set_contexts(xd, pd, tx_size, c > 0, blk_col, blk_row);
Jingning Han's avatar
Jingning Han committed
517
}
518
#endif
Jingning Han's avatar
Jingning Han committed
519
520
521
522
523

struct is_skippable_args {
  uint16_t *eobs;
  int *skippable;
};
524
static void is_skippable(int plane, int block, int blk_row, int blk_col,
525
                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *argv) {
Jingning Han's avatar
Jingning Han committed
526
527
528
529
  struct is_skippable_args *args = argv;
  (void)plane;
  (void)plane_bsize;
  (void)tx_size;
530
531
  (void)blk_row;
  (void)blk_col;
Jingning Han's avatar
Jingning Han committed
532
533
534
535
  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
536
537
//              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
538
  int result = 1;
539
  struct is_skippable_args args = { x->plane[plane].eobs, &result };
Yaowu Xu's avatar
Yaowu Xu committed
540
541
  av1_foreach_transformed_block_in_plane(&x->e_mbd, bsize, plane, is_skippable,
                                         &args);
Jingning Han's avatar
Jingning Han committed
542
543
544
  return result;
}

545
static void has_high_freq_coeff(int plane, int block, int blk_row, int blk_col,
Jingning Han's avatar
Jingning Han committed
546
547
548
549
                                BLOCK_SIZE plane_bsize, TX_SIZE tx_size,
                                void *argv) {
  struct is_skippable_args *args = argv;
  int eobs = (tx_size == TX_4X4) ? 3 : 10;
550
551
552
553
  (void)plane;
  (void)plane_bsize;
  (void)blk_row;
  (void)blk_col;
Jingning Han's avatar
Jingning Han committed
554
555
556
557

  *(args->skippable) |= (args->eobs[block] > eobs);
}

Yaowu Xu's avatar
Yaowu Xu committed
558
int av1_has_high_freq_in_plane(MACROBLOCK *x, BLOCK_SIZE bsize, int plane) {
Jingning Han's avatar
Jingning Han committed
559
  int result = 0;
560
  struct is_skippable_args args = { x->plane[plane].eobs, &result };
Yaowu Xu's avatar
Yaowu Xu committed
561
562
  av1_foreach_transformed_block_in_plane(&x->e_mbd, bsize, plane,
                                         has_high_freq_coeff, &args);
Jingning Han's avatar
Jingning Han committed
563
564
  return result;
}
565
566
567
568
569
#if CONFIG_PVQ
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) {
    q->buf_len = 2 * q->buf_len + 1;
570
571
    CHECK_MEM_ERROR(cm, q->buf,
                    aom_realloc(q->buf, q->buf_len * sizeof(PVQ_INFO)));
572
  }
573
  // memcpy(q->buf + q->curr_pos, pvq, sizeof(PVQ_INFO));
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
  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;
Jingning Han's avatar
Jingning Han committed
595

596
597
  assert(block < MAX_PVQ_BLOCKS_IN_SB);
  pvq_info = &x->pvq[block][plane];
598
  add_pvq_block((AV1_COMMON * const)cm, x, pvq_info);
599
600
}
#endif
601
#if CONFIG_VAR_TX
602
void tokenize_vartx(ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,
603
604
                    TX_SIZE tx_size, BLOCK_SIZE plane_bsize, int blk_row,
                    int blk_col, int block, int plane, void *arg) {
605
606
607
608
  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];
609
  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
610
611
  const int tx_row = blk_row >> (1 - pd->subsampling_y);
  const int tx_col = blk_col >> (1 - pd->subsampling_x);
612
613
  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
614
  TX_SIZE plane_tx_size;
615

616
617
  assert(tx_size < TX_SIZES);

618
  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
619

620
621
622
  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
623

624
  if (tx_size == plane_tx_size) {
625
    plane_bsize = get_plane_block_size(mbmi->sb_type, pd);
626
627
    if (!dry_run)
      tokenize_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
628
    else if (dry_run == DRY_RUN_NORMAL)
629
630
      set_entropy_context_b(plane, block, blk_row, blk_col, plane_bsize,
                            tx_size, arg);
631
632
    else if (dry_run == DRY_RUN_COSTCOEFFS)
      cost_coeffs_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
633
  } else {
634
    // Half the block size in transform block unit.
635
636
    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
    const int bsl = tx_size_wide_unit[sub_txs];
637
638
639
640
641
    int i;

    assert(bsl > 0);

    for (i = 0; i < 4; ++i) {
642
643
644
645
      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];
646

647
      if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;
648

649
      tokenize_vartx(td, t, dry_run, sub_txs, plane_bsize, offsetr, offsetc,
650
651
                     block, plane, arg);
      block += step;
652
653
654
655
    }
  }
}

656
void av1_tokenize_sb_vartx(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
657
658
                           RUN_TYPE dry_run, int mi_row, int mi_col,
                           BLOCK_SIZE bsize, int *rate) {
659
  const AV1_COMMON *const cm = &cpi->common;
660
661
662
663
  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
664
  const int ctx = av1_get_skip_context(xd);
665
666
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
667
  struct tokenize_b_args arg = { cpi, td, t, 0 };
668
  int plane;
669
  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return;
670
671

  if (mbmi->skip) {
672
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
673
    reset_skip_context(xd, bsize);
674
    if (dry_run) *t = t_backup;
675
676
677
678
679
680
681
682
683
684
685
    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);
686
687
    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];
688
    const TX_SIZE max_tx_size = max_txsize_lookup[plane_bsize];
689
    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
690
691
    int bw = block_size_wide[txb_size] >> tx_size_wide_log2[0];
    int bh = block_size_high[txb_size] >> tx_size_wide_log2[0];
692
693
    int idx, idy;
    int block = 0;
694
    int step = tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];
695
    for (idy = 0; idy < mi_height; idy += bh) {
696
      for (idx = 0; idx < mi_width; idx += bw) {
697
698
        tokenize_vartx(td, t, dry_run, max_tx_size, plane_bsize, idy, idx,
                       block, plane, &arg);
699
700
701
        block += step;
      }
    }
702
703
704
705
706

    if (!dry_run) {
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
707
  }
708
  if (rate) *rate += arg.this_rate;
709
}
710
#endif  // CONFIG_VAR_TX
711

712
void av1_tokenize_sb(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
713
                     RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate) {
714
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
715
716
717
  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
718
  const int ctx = av1_get_skip_context(xd);
719
720
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
721
  struct tokenize_b_args arg = { cpi, td, t, 0 };
Jingning Han's avatar
Jingning Han committed
722
  if (mbmi->skip) {
723
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
Jingning Han's avatar
Jingning Han committed
724
725
726
727
    reset_skip_context(xd, bsize);
    return;
  }

728
#if !CONFIG_PVQ
Jingning Han's avatar
Jingning Han committed
729
  if (!dry_run) {
730
731
    int plane;

Jingning Han's avatar
Jingning Han committed
732
    td->counts->skip[ctx][0] += skip_inc;
733
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
Yaowu Xu's avatar
Yaowu Xu committed
734
735
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_b,
                                             &arg);
736
737
738
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
739
  } else if (dry_run == DRY_RUN_NORMAL) {
Yaowu Xu's avatar
Yaowu Xu committed
740
    av1_foreach_transformed_block(xd, bsize, set_entropy_context_b, &arg);
741
742
  } else if (dry_run == DRY_RUN_COSTCOEFFS) {
    av1_foreach_transformed_block(xd, bsize, cost_coeffs_b, &arg);
Jingning Han's avatar
Jingning Han committed
743
  }
744
745
746
747
748
749
750
751
752
753
754
#else
  if (!dry_run) {
    int plane;

    td->counts->skip[ctx][0] += skip_inc;

    for (plane = 0; plane < MAX_MB_PLANE; ++plane)
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_pvq,
                                             &arg);
  }
#endif
755
  if (rate) *rate += arg.this_rate;
Jingning Han's avatar
Jingning Han committed
756
}
757
758

#if CONFIG_SUPERTX
759
760
761
762
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;
763
764
765
  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
766
  const int ctx = av1_get_skip_context(xd);
767
768
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id_supertx, SEG_LVL_SKIP);
769
  struct tokenize_b_args arg = { cpi, td, t, 0 };
770
  if (mbmi->skip) {
771
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
772
    reset_skip_context(xd, bsize);
773
    if (dry_run) *t = t_backup;
774
775
776
777
778
779
780
781
    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
782
783
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_b,
                                             &arg);
784
785
786
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
787
  } else if (dry_run == DRY_RUN_NORMAL) {
Yaowu Xu's avatar
Yaowu Xu committed
788
    av1_foreach_transformed_block(xd, bsize, set_entropy_context_b, &arg);
789
    *t = t_backup;
790
791
  } else if (dry_run == DRY_RUN_COSTCOEFFS) {
    av1_foreach_transformed_block(xd, bsize, cost_coeffs_b, &arg);
792
  }
793
  if (rate) *rate += arg.this_rate;
794
795
}
#endif  // CONFIG_SUPERTX