tokenize.c 37.5 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
27
28
#if CONFIG_LV_MAP
#include "av1/encoder/encodetxb.c"
#endif
29
#include "av1/encoder/rdopt.h"
30
#include "av1/encoder/tokenize.h"
Jingning Han's avatar
Jingning Han committed
31
32

static const TOKENVALUE dct_cat_lt_10_value_tokens[] = {
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
  { 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
52
};
Yaowu Xu's avatar
Yaowu Xu committed
53
const TOKENVALUE *av1_dct_cat_lt_10_value_tokens =
54
55
56
    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
57
58
59
// 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
60
// (as defined in cost element av1_extra_bits)
Julia Robson's avatar
Julia Robson committed
61
62
// 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[] = {
63
64
65
66
67
68
69
70
71
72
73
  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
74
};
Yaowu Xu's avatar
Yaowu Xu committed
75
const int *av1_dct_cat_lt_10_value_cost =
76
77
    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
78
79

// Array indices are identical to previously-existing CONTEXT_NODE indices
80
/* clang-format off */
Yaowu Xu's avatar
Yaowu Xu committed
81
const aom_tree_index av1_coef_tree[TREE_SIZE(ENTROPY_TOKENS)] = {
Jingning Han's avatar
Jingning Han committed
82
83
84
85
86
87
88
89
90
91
92
93
  -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
};
94
95
96
97
98
99
100
101
102
103
104
/* 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 };
105
static const int16_t cat5_cost[1 << 5] = {
106
107
108
109
  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
110
const int16_t av1_cat6_low_cost[256] = {
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
  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
};
132
const int av1_cat6_high_cost[CAT6_HIGH_COST_ENTRIES] = {
133
134
135
136
137
  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,
138
  15118, 10834, 12997, 13473, 15636, 13894, 16057, 16533, 18696,
139
#if CONFIG_HIGHBITDEPTH
140
141
142
143
144
145
  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,
146
147
148
149
150
  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,
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
  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,
170
171
172
173
  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,
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
  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,
193
194
195
196
197
  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,
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
  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,
216
217
218
219
220
  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,
221
222
223
224
225
226
227
  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
228
#endif
229
230
231
232
233
};

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

235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
#if CONFIG_NEW_MULTISYMBOL
const av1_extra_bit av1_extra_bits[ENTROPY_TOKENS] = {
  { 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_cdf, 1, CAT1_MIN_VAL, cat1_cost },  // CATEGORY1_TOKEN
  { av1_cat2_cdf, 2, CAT2_MIN_VAL, cat2_cost },  // CATEGORY2_TOKEN
  { av1_cat3_cdf, 3, CAT3_MIN_VAL, cat3_cost },  // CATEGORY3_TOKEN
  { av1_cat4_cdf, 4, CAT4_MIN_VAL, cat4_cost },  // CATEGORY4_TOKEN
  { av1_cat5_cdf, 5, CAT5_MIN_VAL, cat5_cost },  // CATEGORY5_TOKEN
  { av1_cat6_cdf, 18, CAT6_MIN_VAL, 0 },         // CATEGORY6_TOKEN
  { 0, 0, 0, zero_cost }                         // EOB_TOKEN
};
#else
Yaowu Xu's avatar
Yaowu Xu committed
251
const av1_extra_bit av1_extra_bits[ENTROPY_TOKENS] = {
Alex Converse's avatar
Alex Converse committed
252
253
254
255
256
257
258
259
260
261
  { 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
262
  { av1_cat6_prob, 18, CAT6_MIN_VAL, 0 },         // CATEGORY6_TOKEN
Alex Converse's avatar
Alex Converse committed
263
  { 0, 0, 0, zero_cost }                          // EOB_TOKEN
Jingning Han's avatar
Jingning Han committed
264
};
265
#endif
Jingning Han's avatar
Jingning Han committed
266

267
#if !CONFIG_PVQ || CONFIG_VAR_TX
268
269
270
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;
271
272
  const AV1_COMP *const cpi = args->cpi;
  const AV1_COMMON *cm = &cpi->common;
273
274
275
276
277
278
279
  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;
280
281
  const TX_TYPE tx_type =
      av1_get_tx_type(type, xd, blk_row, blk_col, block, tx_size);
Angie Chiang's avatar
Angie Chiang committed
282
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
Jingning Han's avatar
Jingning Han committed
283
284
285
  const int rate = av1_cost_coeffs(
      cpi, x, plane, blk_row, blk_col, block, tx_size, scan_order,
      pd->above_context + blk_col, pd->left_context + blk_row, 0);
286
  args->this_rate += rate;
287
  (void)plane_bsize;
288
289
  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
                   blk_row);
290
291
}

292
293
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
294
                                  TX_SIZE tx_size, void *arg) {
295
  struct tokenize_b_args *const args = arg;
Jingning Han's avatar
Jingning Han committed
296
297
298
299
300
  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];
301
  (void)plane_bsize;
302
303
  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
                   blk_row);
Jingning Han's avatar
Jingning Han committed
304
305
}

306
307
308
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)],
309
310
                             int eob_val, int first_val, int32_t extra,
                             uint8_t token) {
311
312
313
314
  (*t)->token = token;
  (*t)->extra = extra;
  (*t)->tail_cdf = tail_cdf;
  (*t)->head_cdf = head_cdf;
315
  (*t)->eob_val = eob_val;
316
  (*t)->first_val = first_val;
317
  (*t)++;
318
319
320
321
322
323
324
325
326
327
328

  if (token == BLOCK_Z_TOKEN) {
    update_cdf(*head_cdf, 0, HEAD_TOKENS + 1);
  } else {
    if (eob_val != LAST_EOB) {
      const int symb = 2 * AOMMIN(token, TWO_TOKEN) - eob_val + first_val;
      update_cdf(*head_cdf, symb, HEAD_TOKENS + first_val);
    }
    if (token > ONE_TOKEN)
      update_cdf(*tail_cdf, token - TWO_TOKEN, TAIL_TOKENS);
  }
329
}
Angie Chiang's avatar
Angie Chiang committed
330
#endif  // !CONFIG_PVQ || CONFIG_VAR_TX
Jingning Han's avatar
Jingning Han committed
331

332
static int cost_and_tokenize_map(Av1ColorMapParam *param, TOKENEXTRA **t,
333
334
                                 int calc_rate) {
  const uint8_t *const color_map = param->color_map;
335
336
  MapCdf map_cdf = param->map_cdf;
  ColorCost color_cost = param->color_cost;
337
338
339
340
  const int plane_block_width = param->plane_width;
  const int rows = param->rows;
  const int cols = param->cols;
  const int n = param->n_colors;
hui su's avatar
hui su committed
341

342
343
  int this_rate = 0;
  uint8_t color_order[PALETTE_MAX_SIZE];
Fangwen Fu's avatar
Fangwen Fu committed
344
#if CONFIG_PALETTE_THROUGHPUT
345
346
347
  for (int k = 1; k < rows + cols - 1; ++k) {
    for (int j = AOMMIN(k, cols - 1); j >= AOMMAX(0, k - rows + 1); --j) {
      int i = k - j;
Fangwen Fu's avatar
Fangwen Fu committed
348
#else
349
350
  for (int i = 0; i < rows; ++i) {
    for (int j = (i == 0 ? 1 : 0); j < cols; ++j) {
Fangwen Fu's avatar
Fangwen Fu committed
351
#endif  // CONFIG_PALETTE_THROUGHPUT
352
      int color_new_idx;
353
      const int color_ctx = av1_get_palette_color_index_context(
354
          color_map, plane_block_width, i, j, n, color_order, &color_new_idx);
355
      assert(color_new_idx >= 0 && color_new_idx < n);
356
      if (calc_rate) {
Sarah Parker's avatar
Sarah Parker committed
357
358
359
360
        this_rate +=
            (*color_cost)[n - PALETTE_MIN_SIZE][color_ctx][color_new_idx];
      } else {
        (*t)->token = color_new_idx;
361
        (*t)->color_map_cdf = map_cdf[n - PALETTE_MIN_SIZE][color_ctx];
Sarah Parker's avatar
Sarah Parker committed
362
        ++(*t);
363
      }
hui su's avatar
hui su committed
364
365
    }
  }
Sarah Parker's avatar
Sarah Parker committed
366
367
368
369
  if (calc_rate) return this_rate;
  return 0;
}

370
static void get_palette_params(const MACROBLOCK *const x, int plane,
371
                               BLOCK_SIZE bsize, Av1ColorMapParam *params) {
372
373
374
375
376
377
378
379
380
381
382
  const MACROBLOCKD *const xd = &x->e_mbd;
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
  params->color_map = xd->plane[plane].color_index_map;
  params->map_cdf = plane ? xd->tile_ctx->palette_uv_color_index_cdf
                          : xd->tile_ctx->palette_y_color_index_cdf;
  params->color_cost =
      plane ? &x->palette_uv_color_cost : &x->palette_y_color_cost;
  params->n_colors = pmi->palette_size[plane];
  av1_get_block_dimensions(bsize, plane, xd, &params->plane_width, NULL,
                           &params->rows, &params->cols);
Sarah Parker's avatar
Sarah Parker committed
383
384
}

Sarah Parker's avatar
Sarah Parker committed
385
386
387
#if CONFIG_MRC_TX && SIGNAL_ANY_MRC_MASK
static void get_mrc_params(const MACROBLOCK *const x, int block,
                           TX_SIZE tx_size, Av1ColorMapParam *params) {
388
  memset(params, 0, sizeof(*params));
Sarah Parker's avatar
Sarah Parker committed
389
390
391
392
393
394
395
396
397
398
399
400
  const MACROBLOCKD *const xd = &x->e_mbd;
  const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
  const int is_inter = is_inter_block(mbmi);
  params->color_map = BLOCK_OFFSET(xd->mrc_mask, block);
  params->map_cdf = is_inter ? xd->tile_ctx->mrc_mask_inter_cdf
                             : xd->tile_ctx->mrc_mask_intra_cdf;
  params->color_cost =
      is_inter ? &x->mrc_mask_inter_cost : &x->mrc_mask_intra_cost;
  params->n_colors = 2;
  params->plane_width = tx_size_wide[tx_size];
  params->rows = tx_size_high[tx_size];
  params->cols = tx_size_wide[tx_size];
401
}
Sarah Parker's avatar
Sarah Parker committed
402
#endif  // CONFIG_MRC_TX && SIGNAL_ANY_MRC_MASK
403
404

static void get_color_map_params(const MACROBLOCK *const x, int plane,
Sarah Parker's avatar
Sarah Parker committed
405
406
                                 int block, BLOCK_SIZE bsize, TX_SIZE tx_size,
                                 COLOR_MAP_TYPE type,
407
                                 Av1ColorMapParam *params) {
Sarah Parker's avatar
Sarah Parker committed
408
409
  (void)block;
  (void)tx_size;
410
411
412
  memset(params, 0, sizeof(*params));
  switch (type) {
    case PALETTE_MAP: get_palette_params(x, plane, bsize, params); break;
Sarah Parker's avatar
Sarah Parker committed
413
414
415
#if CONFIG_MRC_TX && SIGNAL_ANY_MRC_MASK
    case MRC_MAP: get_mrc_params(x, block, tx_size, params); break;
#endif  // CONFIG_MRC_TX && SIGNAL_ANY_MRC_MASK
416
417
418
419
    default: assert(0 && "Invalid color map type"); return;
  }
}

Sarah Parker's avatar
Sarah Parker committed
420
421
int av1_cost_color_map(const MACROBLOCK *const x, int plane, int block,
                       BLOCK_SIZE bsize, TX_SIZE tx_size, COLOR_MAP_TYPE type) {
Sarah Parker's avatar
Sarah Parker committed
422
  assert(plane == 0 || plane == 1);
423
  Av1ColorMapParam color_map_params;
Sarah Parker's avatar
Sarah Parker committed
424
425
  get_color_map_params(x, plane, block, bsize, tx_size, type,
                       &color_map_params);
426
427
428
  return cost_and_tokenize_map(&color_map_params, NULL, 1);
}

Sarah Parker's avatar
Sarah Parker committed
429
430
void av1_tokenize_color_map(const MACROBLOCK *const x, int plane, int block,
                            TOKENEXTRA **t, BLOCK_SIZE bsize, TX_SIZE tx_size,
431
432
                            COLOR_MAP_TYPE type) {
  assert(plane == 0 || plane == 1);
Sarah Parker's avatar
Sarah Parker committed
433
434
435
436
437
438
439
440
#if CONFIG_MRC_TX
  if (type == MRC_MAP) {
    const int is_inter = is_inter_block(&x->e_mbd.mi[0]->mbmi);
    if ((is_inter && !SIGNAL_MRC_MASK_INTER) ||
        (!is_inter && !SIGNAL_MRC_MASK_INTRA))
      return;
  }
#endif  // CONFIG_MRC_TX
441
  Av1ColorMapParam color_map_params;
Sarah Parker's avatar
Sarah Parker committed
442
443
  get_color_map_params(x, plane, block, bsize, tx_size, type,
                       &color_map_params);
Sarah Parker's avatar
Sarah Parker committed
444
  // The first color index does not use context or entropy.
445
446
  (*t)->token = color_map_params.color_map[0];
  (*t)->color_map_cdf = NULL;
Sarah Parker's avatar
Sarah Parker committed
447
  ++(*t);
448
  cost_and_tokenize_map(&color_map_params, t, 0);
hui su's avatar
hui su committed
449
450
}

451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
#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];
clang-format's avatar
clang-format committed
488
  add_pvq_block((AV1_COMMON * const) cm, x, pvq_info);
489
490
491
}
#endif  // CONFIG_PVQ

492
static void tokenize_b(int plane, int block, int blk_row, int blk_col,
493
                       BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
494
#if !CONFIG_PVQ
495
  struct tokenize_b_args *const args = arg;
496
  const AV1_COMP *cpi = args->cpi;
Angie Chiang's avatar
Angie Chiang committed
497
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
498
499
500
501
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  TOKENEXTRA **tp = args->tp;
502
  uint8_t token_cache[MAX_TX_SQUARE];
Jingning Han's avatar
Jingning Han committed
503
504
505
506
507
  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;
508
  TOKENEXTRA *t = *tp; /* store tokens starting here */
509
  const int eob = p->eobs[block];
Jingning Han's avatar
Jingning Han committed
510
511
512
513
  const PLANE_TYPE type = pd->plane_type;
  const tran_low_t *qcoeff = BLOCK_OFFSET(p->qcoeff, block);
  const int segment_id = mbmi->segment_id;
  const int16_t *scan, *nb;
514
515
  const TX_TYPE tx_type =
      av1_get_tx_type(type, xd, blk_row, blk_col, block, tx_size);
Angie Chiang's avatar
Angie Chiang committed
516
  const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
Jingning Han's avatar
Jingning Han committed
517
  const int ref = is_inter_block(mbmi);
518
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
519
  aom_cdf_prob(
520
      *const coef_head_cdfs)[COEFF_CONTEXTS][CDF_SIZE(ENTROPY_TOKENS)] =
521
      ec_ctx->coef_head_cdfs[txsize_sqr_map[tx_size]][type][ref];
522
  aom_cdf_prob(
523
      *const coef_tail_cdfs)[COEFF_CONTEXTS][CDF_SIZE(ENTROPY_TOKENS)] =
524
      ec_ctx->coef_tail_cdfs[txsize_sqr_map[tx_size]][type][ref];
525
  int eob_val;
526
  int first_val = 1;
hui su's avatar
hui su committed
527
  const int seg_eob = av1_get_tx_eob(&cpi->common.seg, segment_id, tx_size);
Jingning Han's avatar
Jingning Han committed
528
529
530
  const uint8_t *const band = get_band_translate(tx_size);
  int16_t token;
  EXTRABIT extra;
531
  (void)plane_bsize;
532
533
  pt = get_entropy_context(tx_size, pd->above_context + blk_col,
                           pd->left_context + blk_row);
534
535
  scan = scan_order->scan;
  nb = scan_order->neighbors;
Jingning Han's avatar
Jingning Han committed
536
537
  c = 0;

Sarah Parker's avatar
Sarah Parker committed
538
539
540
541
542
#if CONFIG_MRC_TX && SIGNAL_ANY_MRC_MASK
  if (tx_type == MRC_DCT)
    av1_tokenize_color_map(x, plane, block, &t, plane_bsize, tx_size, MRC_MAP);
#endif  // CONFIG_MRC_TX && SIGNAL_ANY_MRC_MASK

543
  if (eob == 0)
544
545
    add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt], 1,
              1, 0, BLOCK_Z_TOKEN);
546

Jingning Han's avatar
Jingning Han committed
547
  while (c < eob) {
548
    int v = qcoeff[scan[c]];
549
    first_val = (c == 0);
Jingning Han's avatar
Jingning Han committed
550

551
552
    if (!v) {
      add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt],
553
                0, first_val, 0, ZERO_TOKEN);
554
555
      token_cache[scan[c]] = 0;
    } else {
556
557
      eob_val =
          (c + 1 == eob) ? (c + 1 == seg_eob ? LAST_EOB : EARLY_EOB) : NO_EOB;
558
559
      av1_get_token_extra(v, &token, &extra);
      add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt],
560
                eob_val, first_val, extra, (uint8_t)token);
561
562
      token_cache[scan[c]] = av1_pt_energy_class[token];
    }
Jingning Han's avatar
Jingning Han committed
563
    ++c;
564
    pt = get_coef_context(nb, token_cache, AOMMIN(c, eob - 1));
Jingning Han's avatar
Jingning Han committed
565
566
  }

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

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

Angie Chiang's avatar
Angie Chiang committed
574
575
576
577
578
#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.
579
580
  av1_update_scan_count_facade((AV1_COMMON *)cm, td->counts, tx_size, tx_type,
                               qcoeff, c);
Angie Chiang's avatar
Angie Chiang committed
581
582
#endif

583
  av1_set_contexts(xd, pd, plane, tx_size, c > 0, blk_col, blk_row);
584
585
586
#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
587
}
Fangwen Fu's avatar
Fangwen Fu committed
588

Jingning Han's avatar
Jingning Han committed
589
590
591
592
struct is_skippable_args {
  uint16_t *eobs;
  int *skippable;
};
593
static void is_skippable(int plane, int block, int blk_row, int blk_col,
594
                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *argv) {
Jingning Han's avatar
Jingning Han committed
595
596
597
598
  struct is_skippable_args *args = argv;
  (void)plane;
  (void)plane_bsize;
  (void)tx_size;
599
600
  (void)blk_row;
  (void)blk_col;
Jingning Han's avatar
Jingning Han committed
601
602
603
604
  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
605
606
//              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
607
  int result = 1;
608
  struct is_skippable_args args = { x->plane[plane].eobs, &result };
Yaowu Xu's avatar
Yaowu Xu committed
609
610
  av1_foreach_transformed_block_in_plane(&x->e_mbd, bsize, plane, is_skippable,
                                         &args);
Jingning Han's avatar
Jingning Han committed
611
612
613
  return result;
}

614
#if CONFIG_VAR_TX
615
void tokenize_vartx(ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,
616
617
                    TX_SIZE tx_size, BLOCK_SIZE plane_bsize, int blk_row,
                    int blk_col, int block, int plane, void *arg) {
618
619
620
621
  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];
622
  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
623
624
  const int tx_row = blk_row >> (1 - pd->subsampling_y);
  const int tx_col = blk_col >> (1 - pd->subsampling_x);
625
626
  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
627
  TX_SIZE plane_tx_size;
628

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

631
632
633
  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
634

635
  if (tx_size == plane_tx_size) {
636
    plane_bsize = get_plane_block_size(mbmi->sb_type, pd);
637
638
639
640
641
642
643
644
645
646
647
648
#if CONFIG_LV_MAP
    if (!dry_run) {
      av1_update_and_record_txb_context(plane, block, blk_row, blk_col,
                                        plane_bsize, tx_size, arg);
    } else if (dry_run == DRY_RUN_NORMAL) {
      av1_update_txb_context_b(plane, block, blk_row, blk_col, plane_bsize,
                               tx_size, arg);
    } else {
      printf("DRY_RUN_COSTCOEFFS is not supported yet\n");
      assert(0);
    }
#else
649
650
    if (!dry_run)
      tokenize_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
651
    else if (dry_run == DRY_RUN_NORMAL)
652
653
      set_entropy_context_b(plane, block, blk_row, blk_col, plane_bsize,
                            tx_size, arg);
654
655
    else if (dry_run == DRY_RUN_COSTCOEFFS)
      cost_coeffs_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
656
#endif
657
  } else {
Yue Chen's avatar
Yue Chen committed
658
659
660
661
#if CONFIG_RECT_TX_EXT
    int is_qttx = plane_tx_size == quarter_txsize_lookup[plane_bsize];
    const TX_SIZE sub_txs = is_qttx ? plane_tx_size : sub_tx_size_map[tx_size];
#else
662
    // Half the block size in transform block unit.
663
    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
Yue Chen's avatar
Yue Chen committed
664
#endif
665
    const int bsl = tx_size_wide_unit[sub_txs];
666
667
668
669
670
    int i;

    assert(bsl > 0);

    for (i = 0; i < 4; ++i) {
Yue Chen's avatar
Yue Chen committed
671
672
673
674
675
676
677
678
679
#if CONFIG_RECT_TX_EXT
      int is_wide_tx = tx_size_wide_unit[sub_txs] > tx_size_high_unit[sub_txs];
      const int offsetr =
          is_qttx ? (is_wide_tx ? i * tx_size_high_unit[sub_txs] : 0)
                  : blk_row + ((i >> 1) * bsl);
      const int offsetc =
          is_qttx ? (is_wide_tx ? 0 : i * tx_size_wide_unit[sub_txs])
                  : blk_col + ((i & 0x01) * bsl);
#else
680
681
      const int offsetr = blk_row + ((i >> 1) * bsl);
      const int offsetc = blk_col + ((i & 0x01) * bsl);
Yue Chen's avatar
Yue Chen committed
682
#endif
683
684

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

686
      if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;
687

688
      tokenize_vartx(td, t, dry_run, sub_txs, plane_bsize, offsetr, offsetc,
689
690
                     block, plane, arg);
      block += step;
691
692
693
694
    }
  }
}

695
void av1_tokenize_sb_vartx(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
696
697
                           RUN_TYPE dry_run, int mi_row, int mi_col,
                           BLOCK_SIZE bsize, int *rate) {
698
  const AV1_COMMON *const cm = &cpi->common;
699
700
701
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
702
703
704
#if CONFIG_LV_MAP
  (void)t;
#else
705
  TOKENEXTRA *t_backup = *t;
706
#endif
707
  struct tokenize_b_args arg = { cpi, td, t, 0 };
708
  int plane;
709
  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return;
710
711

  if (mbmi->skip) {
712
    av1_reset_skip_context(xd, mi_row, mi_col, bsize);
713
#if !CONFIG_LV_MAP
714
    if (dry_run) *t = t_backup;
715
#endif
716
717
718
    return;
  }

719
#if !CONFIG_LV_MAP
720
721
  else
    *t = t_backup;
722
#endif
723
724

  for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
725
#if CONFIG_CB4X4
726
727
728
    if (!is_chroma_reference(mi_row, mi_col, bsize,
                             xd->plane[plane].subsampling_x,
                             xd->plane[plane].subsampling_y)) {
729
#if !CONFIG_PVQ && !CONFIG_LV_MAP
730
731
732
733
734
735
736
737
      if (!dry_run) {
        (*t)->token = EOSB_TOKEN;
        (*t)++;
      }
#endif
      continue;
    }
#endif
738
    const struct macroblockd_plane *const pd = &xd->plane[plane];
739
#if CONFIG_CHROMA_SUB8X8
740
741
742
    const BLOCK_SIZE plane_bsize =
        AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
#else
743
    const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
744
#endif
745
746
    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];
747
748
    const TX_SIZE max_tx_size = get_vartx_max_txsize(
        mbmi, plane_bsize, pd->subsampling_x || pd->subsampling_y);
749
    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
750
751
    int bw = block_size_wide[txb_size] >> tx_size_wide_log2[0];
    int bh = block_size_high[txb_size] >> tx_size_wide_log2[0];
752
753
    int idx, idy;
    int block = 0;
754
    int step = tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776

    const BLOCK_SIZE max_unit_bsize = get_plane_block_size(BLOCK_64X64, pd);
    int mu_blocks_wide =
        block_size_wide[max_unit_bsize] >> tx_size_wide_log2[0];
    int mu_blocks_high =
        block_size_high[max_unit_bsize] >> tx_size_high_log2[0];

    mu_blocks_wide = AOMMIN(mi_width, mu_blocks_wide);
    mu_blocks_high = AOMMIN(mi_height, mu_blocks_high);

    for (idy = 0; idy < mi_height; idy += mu_blocks_high) {
      for (idx = 0; idx < mi_width; idx += mu_blocks_wide) {
        int blk_row, blk_col;
        const int unit_height = AOMMIN(mu_blocks_high + idy, mi_height);
        const int unit_width = AOMMIN(mu_blocks_wide + idx, mi_width);
        for (blk_row = idy; blk_row < unit_height; blk_row += bh) {
          for (blk_col = idx; blk_col < unit_width; blk_col += bw) {
            tokenize_vartx(td, t, dry_run, max_tx_size, plane_bsize, blk_row,
                           blk_col, block, plane, &arg);
            block += step;
          }
        }
777
778
      }
    }
779
#if !CONFIG_LV_MAP
780
781
782
783
    if (!dry_run) {
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
784
#endif
785
  }
786
  if (rate) *rate += arg.this_rate;
787
}
788
#endif  // CONFIG_VAR_TX
789

790
void av1_tokenize_sb(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
791
792
                     RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
                     const int mi_row, const int mi_col) {
Jingning Han's avatar
Jingning Han committed
793
794
795
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
796
  struct tokenize_b_args arg = { cpi, td, t, 0 };
Jingning Han's avatar
Jingning Han committed
797
  if (mbmi->skip) {
798
    av1_reset_skip_context(xd, mi_row, mi_col, bsize);
Jingning Han's avatar
Jingning Han committed
799
800
801
802
    return;
  }

  if (!dry_run) {
803
804
805
#if CONFIG_COEF_INTERLEAVE
    av1_foreach_transformed_block_interleave(xd, bsize, tokenize_b, &arg);
#else
806
807
808
    int plane;

    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
809
#if CONFIG_CB4X4
810
811
812
      if (!is_chroma_reference(mi_row, mi_col, bsize,
                               xd->plane[plane].subsampling_x,
                               xd->plane[plane].subsampling_y)) {
813
#if !CONFIG_PVQ
814
815
        (*t)->token = EOSB_TOKEN;
        (*t)++;
816
#endif
817
818
819
820
821
822
        continue;
      }
#else
      (void)mi_row;
      (void)mi_col;
#endif
Yaowu Xu's avatar
Yaowu Xu committed
823
824
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_b,
                                             &arg);
825
#if !CONFIG_PVQ
826
827
      (*t)->token = EOSB_TOKEN;
      (*t)++;
828
#endif  // !CONFIG_PVQ
829
    }
830
#endif
831
832
833
  }
#if !CONFIG_PVQ
  else if (dry_run == DRY_RUN_NORMAL) {
834
    int plane;
835
836
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
#if CONFIG_CB4X4
837
838
839
      if (!is_chroma_reference(mi_row, mi_col, bsize,
                               xd->plane[plane].subsampling_x,
                               xd->plane[plane].subsampling_y))
840
841
842
843
844
        continue;
#else
      (void)mi_row;
      (void)mi_col;
#endif
845
846
      av1_foreach_transformed_block_in_plane(xd, bsize, plane,
                                             set_entropy_context_b, &arg);
847
    }
848
  } else if (dry_run == DRY_RUN_COSTCOEFFS) {
849
    int plane;
850
851
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
#if CONFIG_CB4X4
852
853
854
      if (!is_chroma_reference(mi_row, mi_col, bsize,
                               xd->plane[plane].subsampling_x,
                               xd->plane[plane].subsampling_y))
855
856
857
858
859
        continue;
#else
      (void)mi_row;
      (void)mi_col;
#endif
860
861
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, cost_coeffs_b,
                                             &arg);
862
    }
Jingning Han's avatar
Jingning Han committed
863
  }
864
#endif  // !CONFIG_PVQ
865

866
  if (rate) *rate += arg.this_rate;
Jingning Han's avatar
Jingning Han committed
867
}