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_DAALA_EC || CONFIG_ANS)
Yaowu Xu's avatar
Yaowu Xu committed
268
const struct av1_token av1_coef_encodings[ENTROPY_TOKENS] = {
269 270
  { 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
271
};
272
#endif  // !(CONFIG_DAALA_EC || CONFIG_ANS)
Jingning Han's avatar
Jingning Han committed
273

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

299 300
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
301
                                  TX_SIZE tx_size, void *arg) {
302
  struct tokenize_b_args *const args = arg;
Jingning Han's avatar
Jingning Han committed
303 304 305 306 307
  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];
308
  (void)plane_bsize;
309 310
  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
                   blk_row);
Jingning Han's avatar
Jingning Han committed
311 312
}

313
#if CONFIG_DAALA_EC || CONFIG_ANS
314 315 316
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)],
317 318
                             int eob_val, int first_val, int32_t extra,
                             uint8_t token) {
319 320 321 322
  (*t)->token = token;
  (*t)->extra = extra;
  (*t)->tail_cdf = tail_cdf;
  (*t)->head_cdf = head_cdf;
323
  (*t)->eob_val = eob_val;
324
  (*t)->first_val = first_val;
325 326 327
  (*t)++;
}

328
#else
329 330 331
static INLINE void add_token(TOKENEXTRA **t, const aom_prob *context_tree,
                             int32_t extra, uint8_t token,
                             uint8_t skip_eob_node, unsigned int *counts) {
Jingning Han's avatar
Jingning Han committed
332 333 334 335 336 337 338
  (*t)->token = token;
  (*t)->extra = extra;
  (*t)->context_tree = context_tree;
  (*t)->skip_eob_node = skip_eob_node;
  (*t)++;
  ++counts[token];
}
339
#endif  // CONFIG_DAALA_EC || CONFIG_ANS
Angie Chiang's avatar
Angie Chiang committed
340
#endif  // !CONFIG_PVQ || CONFIG_VAR_TX
Jingning Han's avatar
Jingning Han committed
341

Fangwen Fu's avatar
Fangwen Fu committed
342
#if CONFIG_PALETTE
343 344 345 346 347 348 349
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;
350
  const uint8_t *const color_map = xd->plane[plane].color_index_map;
351
  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
352
  const int n = pmi->palette_size[plane];
353
  int i, j;
354
  int this_rate = 0;
355
  uint8_t color_order[PALETTE_MAX_SIZE];
clang-format's avatar
clang-format committed
356 357 358 359
  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;
360 361 362 363
  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
364

Fangwen Fu's avatar
Fangwen Fu committed
365 366 367 368 369 370
#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
371 372
  for (i = 0; i < rows; ++i) {
    for (j = (i == 0 ? 1 : 0); j < cols; ++j) {
Fangwen Fu's avatar
Fangwen Fu committed
373
#endif  // CONFIG_PALETTE_THROUGHPUT
374
      int color_new_idx;
375
      const int color_ctx = av1_get_palette_color_index_context(
376
          color_map, plane_block_width, i, j, n, color_order, &color_new_idx);
377
      assert(color_new_idx >= 0 && color_new_idx < n);
378
      if (dry_run == DRY_RUN_COSTCOEFFS)
379 380
        this_rate += cpi->palette_y_color_cost[n - PALETTE_MIN_SIZE][color_ctx]
                                              [color_new_idx];
381
      (*t)->token = color_new_idx;
382
      (*t)->context_tree = probs[n - PALETTE_MIN_SIZE][color_ctx];
hui su's avatar
hui su committed
383 384 385 386
      (*t)->skip_eob_node = 0;
      ++(*t);
    }
  }
387
  if (rate) *rate += this_rate;
hui su's avatar
hui su committed
388
}
389
#endif  // CONFIG_PALETTE
hui su's avatar
hui su committed
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 425 426 427 428 429 430 431
#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

432
static void tokenize_b(int plane, int block, int blk_row, int blk_col,
433
                       BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
434
#if !CONFIG_PVQ
435
  struct tokenize_b_args *const args = arg;
436
  const AV1_COMP *cpi = args->cpi;
Angie Chiang's avatar
Angie Chiang committed
437
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
438 439 440 441
  ThreadData *const td = args->td;
  MACROBLOCK *const x = &td->mb;
  MACROBLOCKD *const xd = &x->e_mbd;
  TOKENEXTRA **tp = args->tp;
442
  uint8_t token_cache[MAX_TX_SQUARE];
Jingning Han's avatar
Jingning Han committed
443 444 445 446 447
  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;
448
  TOKENEXTRA *t = *tp; /* store tokens starting here */
449
  const int eob = p->eobs[block];
Jingning Han's avatar
Jingning Han committed
450 451
  const PLANE_TYPE type = pd->plane_type;
  const tran_low_t *qcoeff = BLOCK_OFFSET(p->qcoeff, block);
452
#if CONFIG_SUPERTX
Yaowu Xu's avatar
Yaowu Xu committed
453
  const int segment_id = AOMMIN(mbmi->segment_id, mbmi->segment_id_supertx);
454
#else
Jingning Han's avatar
Jingning Han committed
455
  const int segment_id = mbmi->segment_id;
456
#endif  // CONFIG_SUEPRTX
Jingning Han's avatar
Jingning Han committed
457
  const int16_t *scan, *nb;
458
  const TX_TYPE tx_type = get_tx_type(type, xd, block, tx_size);
459
  const SCAN_ORDER *const scan_order =
Angie Chiang's avatar
Angie Chiang committed
460
      get_scan(cm, tx_size, tx_type, is_inter_block(mbmi));
Jingning Han's avatar
Jingning Han committed
461
  const int ref = is_inter_block(mbmi);
clang-format's avatar
clang-format committed
462
  unsigned int(*const counts)[COEFF_CONTEXTS][ENTROPY_TOKENS] =
463
      td->rd_counts.coef_counts[txsize_sqr_map[tx_size]][type][ref];
464
#if !(CONFIG_DAALA_EC || CONFIG_ANS)
clang-format's avatar
clang-format committed
465
  aom_prob(*const coef_probs)[COEFF_CONTEXTS][UNCONSTRAINED_NODES] =
466
      cpi->common.fc->coef_probs[txsize_sqr_map[tx_size]][type][ref];
467
#endif  // !(CONFIG_DAALA_EC || CONFIG_ANS)
468
#if CONFIG_EC_ADAPT
469
  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
470
#elif CONFIG_DAALA_EC || CONFIG_ANS
471 472
  FRAME_CONTEXT *ec_ctx = cpi->common.fc;
#endif
473
#if CONFIG_DAALA_EC || CONFIG_ANS
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 485
#else
  int skip_eob = 0;
486
#endif  // CONFIG_DAALA_EC || CONFIG_ANS
487
  const int seg_eob = get_tx_eob(&cpi->common.seg, segment_id, tx_size);
clang-format's avatar
clang-format committed
488
  unsigned int(*const eob_branch)[COEFF_CONTEXTS] =
489
      td->counts->eob_branch[txsize_sqr_map[tx_size]][type][ref];
Jingning Han's avatar
Jingning Han committed
490 491 492
  const uint8_t *const band = get_band_translate(tx_size);
  int16_t token;
  EXTRABIT extra;
493
  (void)plane_bsize;
494 495
  pt = get_entropy_context(tx_size, pd->above_context + blk_col,
                           pd->left_context + blk_row);
496 497
  scan = scan_order->scan;
  nb = scan_order->neighbors;
Jingning Han's avatar
Jingning Han committed
498 499
  c = 0;

500
#if CONFIG_DAALA_EC || CONFIG_ANS
501
  if (eob == 0)
502 503
    add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt], 1,
              1, 0, BLOCK_Z_TOKEN);
504

505
  ++blockz_count[pt][eob != 0];
506

Jingning Han's avatar
Jingning Han committed
507
  while (c < eob) {
508
    int v = qcoeff[scan[c]];
509
    first_val = (c == 0);
Jingning Han's avatar
Jingning Han committed
510

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

520
      av1_get_token_extra(v, &token, &extra);
Jingning Han's avatar
Jingning Han committed
521

522
      add_token(&t, &coef_tail_cdfs[band[c]][pt], &coef_head_cdfs[band[c]][pt],
523
                eob_val, first_val, extra, (uint8_t)token);
524

525 526 527 528 529
      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;
      }
530 531 532

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

    av1_get_token_extra(v, &token, &extra);

543 544
    add_token(&t, coef_probs[band[c]][pt], extra, (uint8_t)token,
              (uint8_t)skip_eob, counts[band[c]][pt]);
545 546 547 548 549 550 551

    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) {
552 553
    add_token(&t, coef_probs[band[c]][pt], 0, EOB_TOKEN, 0,
              counts[band[c]][pt]);
Jingning Han's avatar
Jingning Han committed
554 555
    ++eob_branch[band[c]][pt];
  }
556
#endif  // CONFIG_DAALA_EC || CONFIG_ANS
Jingning Han's avatar
Jingning Han committed
557

558 559 560 561 562
#if CONFIG_COEF_INTERLEAVE
  t->token = EOSB_TOKEN;
  t++;
#endif

Jingning Han's avatar
Jingning Han committed
563 564
  *tp = t;

Angie Chiang's avatar
Angie Chiang committed
565 566 567 568 569
#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.
570 571
  av1_update_scan_count_facade((AV1_COMMON *)cm, td->counts, tx_size, tx_type,
                               qcoeff, c);
Angie Chiang's avatar
Angie Chiang committed
572 573
#endif

574
  av1_set_contexts(xd, pd, plane, tx_size, c > 0, blk_col, blk_row);
575 576 577
#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
578
}
Fangwen Fu's avatar
Fangwen Fu committed
579

Jingning Han's avatar
Jingning Han committed
580 581 582 583
struct is_skippable_args {
  uint16_t *eobs;
  int *skippable;
};
584
static void is_skippable(int plane, int block, int blk_row, int blk_col,
585
                         BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *argv) {
Jingning Han's avatar
Jingning Han committed
586 587 588 589
  struct is_skippable_args *args = argv;
  (void)plane;
  (void)plane_bsize;
  (void)tx_size;
590 591
  (void)blk_row;
  (void)blk_col;
Jingning Han's avatar
Jingning Han committed
592 593 594 595
  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
596 597
//              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
598
  int result = 1;
599
  struct is_skippable_args args = { x->plane[plane].eobs, &result };
Yaowu Xu's avatar
Yaowu Xu committed
600 601
  av1_foreach_transformed_block_in_plane(&x->e_mbd, bsize, plane, is_skippable,
                                         &args);
Jingning Han's avatar
Jingning Han committed
602 603 604
  return result;
}

605
#if CONFIG_VAR_TX
606
void tokenize_vartx(ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,
607 608
                    TX_SIZE tx_size, BLOCK_SIZE plane_bsize, int blk_row,
                    int blk_col, int block, int plane, void *arg) {
609 610 611 612
  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];
613
  const BLOCK_SIZE bsize = txsize_to_bsize[tx_size];
614 615
  const int tx_row = blk_row >> (1 - pd->subsampling_y);
  const int tx_col = blk_col >> (1 - pd->subsampling_x);
616 617
  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
618
  TX_SIZE plane_tx_size;
619

620
  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
621

622 623 624
  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
625

626
  if (tx_size == plane_tx_size) {
627
    plane_bsize = get_plane_block_size(mbmi->sb_type, pd);
628 629 630 631 632 633 634 635 636 637 638 639
#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
640 641
    if (!dry_run)
      tokenize_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
642
    else if (dry_run == DRY_RUN_NORMAL)
643 644
      set_entropy_context_b(plane, block, blk_row, blk_col, plane_bsize,
                            tx_size, arg);
645 646
    else if (dry_run == DRY_RUN_COSTCOEFFS)
      cost_coeffs_b(plane, block, blk_row, blk_col, plane_bsize, tx_size, arg);
647
#endif
648
  } else {
649
    // Half the block size in transform block unit.
650 651
    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
    const int bsl = tx_size_wide_unit[sub_txs];
652 653 654 655 656
    int i;

    assert(bsl > 0);

    for (i = 0; i < 4; ++i) {
657 658 659 660
      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];
661

662
      if (offsetr >= max_blocks_high || offsetc >= max_blocks_wide) continue;
663

664
      tokenize_vartx(td, t, dry_run, sub_txs, plane_bsize, offsetr, offsetc,
665 666
                     block, plane, arg);
      block += step;
667 668 669 670
    }
  }
}

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

  if (mbmi->skip) {
691
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
692
    av1_reset_skip_context(xd, mi_row, mi_col, bsize);
693
#if !CONFIG_LV_MAP
694
    if (dry_run) *t = t_backup;
695
#endif
696 697 698
    return;
  }

699 700
  if (!dry_run) td->counts->skip[ctx][0] += skip_inc;
#if !CONFIG_LV_MAP
701 702
  else
    *t = t_backup;
703
#endif
704 705

  for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
706
#if CONFIG_CB4X4
707 708 709
    if (!is_chroma_reference(mi_row, mi_col, bsize,
                             xd->plane[plane].subsampling_x,
                             xd->plane[plane].subsampling_y)) {
710
#if !CONFIG_PVQ || !CONFIG_LV_MAP
711 712 713 714 715 716 717 718
      if (!dry_run) {
        (*t)->token = EOSB_TOKEN;
        (*t)++;
      }
#endif
      continue;
    }
#endif
719
    const struct macroblockd_plane *const pd = &xd->plane[plane];
720 721 722 723
#if CONFIG_CB4X4 && !CONFIG_CHROMA_2X2
    const BLOCK_SIZE plane_bsize =
        AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
#else
724
    const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
725
#endif
726 727
    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];
728
    const TX_SIZE max_tx_size = get_vartx_max_txsize(mbmi, plane_bsize);
729
    const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
730 731
    int bw = block_size_wide[txb_size] >> tx_size_wide_log2[0];
    int bh = block_size_high[txb_size] >> tx_size_wide_log2[0];
732 733
    int idx, idy;
    int block = 0;
734
    int step = tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];
735
    for (idy = 0; idy < mi_height; idy += bh) {
736
      for (idx = 0; idx < mi_width; idx += bw) {
737 738
        tokenize_vartx(td, t, dry_run, max_tx_size, plane_bsize, idy, idx,
                       block, plane, &arg);
739 740 741
        block += step;
      }
    }
742

743
#if !CONFIG_LV_MAP
744 745 746 747
    if (!dry_run) {
      (*t)->token = EOSB_TOKEN;
      (*t)++;
    }
748
#endif
749
  }
750
  if (rate) *rate += arg.this_rate;
751
}
752
#endif  // CONFIG_VAR_TX
753

754
void av1_tokenize_sb(const AV1_COMP *cpi, ThreadData *td, TOKENEXTRA **t,
755 756
                     RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
                     const int mi_row, const int mi_col) {
757
  const AV1_COMMON *const cm = &cpi->common;
Jingning Han's avatar
Jingning Han committed
758 759 760
  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
761
  const int ctx = av1_get_skip_context(xd);
762 763
  const int skip_inc =
      !segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);
764
  struct tokenize_b_args arg = { cpi, td, t, 0 };
Jingning Han's avatar
Jingning Han committed
765
  if (mbmi->skip) {
766
    if (!dry_run) td->counts->skip[ctx][1] += skip_inc;
767
    av1_reset_skip_context(xd, mi_row, mi_col, bsize);
Jingning Han's avatar
Jingning Han committed
768 769 770 771
    return;
  }

  if (!dry_run) {
772 773 774 775
#if CONFIG_COEF_INTERLEAVE
    td->counts->skip[ctx][0] += skip_inc;
    av1_foreach_transformed_block_interleave(xd, bsize, tokenize_b, &arg);
#else
776 777
    int plane;

Jingning Han's avatar
Jingning Han committed
778
    td->counts->skip[ctx][0] += skip_inc;
779
    for (plane = 0; plane < MAX_MB_PLANE; ++plane) {
780
#if CONFIG_CB4X4
781 782 783
      if (!is_chroma_reference(mi_row, mi_col, bsize,
                               xd->plane[plane].subsampling_x,
                               xd->plane[plane].subsampling_y)) {
784
#if !CONFIG_PVQ
785 786
        (*t)->token = EOSB_TOKEN;
        (*t)++;
787
#endif
788 789 790 791 792 793
        continue;
      }
#else
      (void)mi_row;
      (void)mi_col;
#endif
Yaowu Xu's avatar
Yaowu Xu committed
794 795
      av1_foreach_transformed_block_in_plane(xd, bsize, plane, tokenize_b,
                                             &arg);
796
#if !CONFIG_PVQ
797