1 /*
2 * Copyright (c) 2022, Alliance for Open Media. All rights reserved.
3 *
4 * This source code is subject to the terms of the BSD 2 Clause License and
5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6 * was not distributed with this source code in the LICENSE file, you can
7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8 * Media Patent License 1.0 was not distributed with this source code in the
9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10 */
11
12 #include <immintrin.h>
13 #include "config/aom_dsp_rtcd.h"
14 #include "aom/aom_integer.h"
15 #include "aom_dsp/x86/quantize_x86.h"
16
load_b_values_avx2(const int16_t * zbin_ptr,__m256i * zbin,const int16_t * round_ptr,__m256i * round,const int16_t * quant_ptr,__m256i * quant,const int16_t * dequant_ptr,__m256i * dequant,const int16_t * shift_ptr,__m256i * shift,int log_scale)17 static inline void load_b_values_avx2(const int16_t *zbin_ptr, __m256i *zbin,
18 const int16_t *round_ptr, __m256i *round,
19 const int16_t *quant_ptr, __m256i *quant,
20 const int16_t *dequant_ptr,
21 __m256i *dequant,
22 const int16_t *shift_ptr, __m256i *shift,
23 int log_scale) {
24 *zbin = _mm256_castsi128_si256(_mm_load_si128((const __m128i *)zbin_ptr));
25 *zbin = _mm256_permute4x64_epi64(*zbin, 0x54);
26 if (log_scale > 0) {
27 const __m256i rnd = _mm256_set1_epi16((int16_t)(1 << (log_scale - 1)));
28 *zbin = _mm256_add_epi16(*zbin, rnd);
29 *zbin = _mm256_srai_epi16(*zbin, log_scale);
30 }
31 // Subtracting 1 here eliminates a _mm256_cmpeq_epi16() instruction when
32 // calculating the zbin mask. (See quantize_b_logscale{0,1,2}_16)
33 *zbin = _mm256_sub_epi16(*zbin, _mm256_set1_epi16(1));
34
35 *round = _mm256_castsi128_si256(_mm_load_si128((const __m128i *)round_ptr));
36 *round = _mm256_permute4x64_epi64(*round, 0x54);
37 if (log_scale > 0) {
38 const __m256i rnd = _mm256_set1_epi16((int16_t)(1 << (log_scale - 1)));
39 *round = _mm256_add_epi16(*round, rnd);
40 *round = _mm256_srai_epi16(*round, log_scale);
41 }
42
43 *quant = _mm256_castsi128_si256(_mm_load_si128((const __m128i *)quant_ptr));
44 *quant = _mm256_permute4x64_epi64(*quant, 0x54);
45 *dequant =
46 _mm256_castsi128_si256(_mm_load_si128((const __m128i *)dequant_ptr));
47 *dequant = _mm256_permute4x64_epi64(*dequant, 0x54);
48 *shift = _mm256_castsi128_si256(_mm_load_si128((const __m128i *)shift_ptr));
49 *shift = _mm256_permute4x64_epi64(*shift, 0x54);
50 }
51
load_coefficients_avx2(const tran_low_t * coeff_ptr)52 static inline __m256i load_coefficients_avx2(const tran_low_t *coeff_ptr) {
53 const __m256i coeff1 = _mm256_load_si256((__m256i *)coeff_ptr);
54 const __m256i coeff2 = _mm256_load_si256((__m256i *)(coeff_ptr + 8));
55 return _mm256_packs_epi32(coeff1, coeff2);
56 }
57
store_coefficients_avx2(__m256i coeff_vals,tran_low_t * coeff_ptr)58 static inline void store_coefficients_avx2(__m256i coeff_vals,
59 tran_low_t *coeff_ptr) {
60 __m256i coeff_sign = _mm256_srai_epi16(coeff_vals, 15);
61 __m256i coeff_vals_lo = _mm256_unpacklo_epi16(coeff_vals, coeff_sign);
62 __m256i coeff_vals_hi = _mm256_unpackhi_epi16(coeff_vals, coeff_sign);
63 _mm256_store_si256((__m256i *)coeff_ptr, coeff_vals_lo);
64 _mm256_store_si256((__m256i *)(coeff_ptr + 8), coeff_vals_hi);
65 }
66
quantize_b_logscale0_16(const tran_low_t * coeff_ptr,tran_low_t * qcoeff_ptr,tran_low_t * dqcoeff_ptr,__m256i * v_quant,__m256i * v_dequant,__m256i * v_round,__m256i * v_zbin,__m256i * v_quant_shift)67 static AOM_FORCE_INLINE __m256i quantize_b_logscale0_16(
68 const tran_low_t *coeff_ptr, tran_low_t *qcoeff_ptr,
69 tran_low_t *dqcoeff_ptr, __m256i *v_quant, __m256i *v_dequant,
70 __m256i *v_round, __m256i *v_zbin, __m256i *v_quant_shift) {
71 const __m256i v_coeff = load_coefficients_avx2(coeff_ptr);
72 const __m256i v_abs_coeff = _mm256_abs_epi16(v_coeff);
73 const __m256i v_zbin_mask = _mm256_cmpgt_epi16(v_abs_coeff, *v_zbin);
74
75 if (_mm256_movemask_epi8(v_zbin_mask) == 0) {
76 _mm256_store_si256((__m256i *)qcoeff_ptr, _mm256_setzero_si256());
77 _mm256_store_si256((__m256i *)dqcoeff_ptr, _mm256_setzero_si256());
78 _mm256_store_si256((__m256i *)(qcoeff_ptr + 8), _mm256_setzero_si256());
79 _mm256_store_si256((__m256i *)(dqcoeff_ptr + 8), _mm256_setzero_si256());
80 return _mm256_setzero_si256();
81 }
82
83 // tmp = v_zbin_mask ? (int64_t)abs_coeff + log_scaled_round : 0
84 const __m256i v_tmp_rnd =
85 _mm256_and_si256(_mm256_adds_epi16(v_abs_coeff, *v_round), v_zbin_mask);
86 // tmp32 = (int)(((((tmp * quant_ptr[rc != 0]) >> 16) + tmp) *
87 // quant_shift_ptr[rc != 0]) >>
88 // (16 - log_scale + AOM_QM_BITS));
89 const __m256i v_tmp32_a = _mm256_mulhi_epi16(v_tmp_rnd, *v_quant);
90 const __m256i v_tmp32_b = _mm256_add_epi16(v_tmp32_a, v_tmp_rnd);
91 const __m256i v_tmp32 = _mm256_mulhi_epi16(v_tmp32_b, *v_quant_shift);
92 const __m256i v_nz_mask = _mm256_cmpgt_epi16(v_tmp32, _mm256_setzero_si256());
93 const __m256i v_qcoeff = _mm256_sign_epi16(v_tmp32, v_coeff);
94 const __m256i v_dqcoeff = _mm256_mullo_epi16(v_qcoeff, *v_dequant);
95 store_coefficients_avx2(v_qcoeff, qcoeff_ptr);
96 store_coefficients_avx2(v_dqcoeff, dqcoeff_ptr);
97 return v_nz_mask;
98 }
99
get_max_lane_eob(const int16_t * iscan,__m256i v_eobmax,__m256i v_mask)100 static inline __m256i get_max_lane_eob(const int16_t *iscan, __m256i v_eobmax,
101 __m256i v_mask) {
102 const __m256i v_iscan = _mm256_loadu_si256((const __m256i *)iscan);
103 const __m256i v_iscan_perm = _mm256_permute4x64_epi64(v_iscan, 0xD8);
104 const __m256i v_iscan_plus1 = _mm256_sub_epi16(v_iscan_perm, v_mask);
105 const __m256i v_nz_iscan = _mm256_and_si256(v_iscan_plus1, v_mask);
106 return _mm256_max_epi16(v_eobmax, v_nz_iscan);
107 }
108
accumulate_eob256(__m256i eob256)109 static inline int16_t accumulate_eob256(__m256i eob256) {
110 const __m128i eob_lo = _mm256_castsi256_si128(eob256);
111 const __m128i eob_hi = _mm256_extractf128_si256(eob256, 1);
112 __m128i eob = _mm_max_epi16(eob_lo, eob_hi);
113 __m128i eob_shuffled = _mm_shuffle_epi32(eob, 0xe);
114 eob = _mm_max_epi16(eob, eob_shuffled);
115 eob_shuffled = _mm_shufflelo_epi16(eob, 0xe);
116 eob = _mm_max_epi16(eob, eob_shuffled);
117 eob_shuffled = _mm_shufflelo_epi16(eob, 0x1);
118 eob = _mm_max_epi16(eob, eob_shuffled);
119 return _mm_extract_epi16(eob, 1);
120 }
121
aom_quantize_b_avx2(const tran_low_t * coeff_ptr,intptr_t n_coeffs,const int16_t * zbin_ptr,const int16_t * round_ptr,const int16_t * quant_ptr,const int16_t * quant_shift_ptr,tran_low_t * qcoeff_ptr,tran_low_t * dqcoeff_ptr,const int16_t * dequant_ptr,uint16_t * eob_ptr,const int16_t * scan,const int16_t * iscan)122 void aom_quantize_b_avx2(const tran_low_t *coeff_ptr, intptr_t n_coeffs,
123 const int16_t *zbin_ptr, const int16_t *round_ptr,
124 const int16_t *quant_ptr,
125 const int16_t *quant_shift_ptr, tran_low_t *qcoeff_ptr,
126 tran_low_t *dqcoeff_ptr, const int16_t *dequant_ptr,
127 uint16_t *eob_ptr, const int16_t *scan,
128 const int16_t *iscan) {
129 (void)scan;
130 __m256i v_zbin, v_round, v_quant, v_dequant, v_quant_shift;
131 __m256i v_eobmax = _mm256_setzero_si256();
132
133 load_b_values_avx2(zbin_ptr, &v_zbin, round_ptr, &v_round, quant_ptr,
134 &v_quant, dequant_ptr, &v_dequant, quant_shift_ptr,
135 &v_quant_shift, 0);
136
137 // Do DC and first 15 AC.
138 __m256i v_nz_mask =
139 quantize_b_logscale0_16(coeff_ptr, qcoeff_ptr, dqcoeff_ptr, &v_quant,
140 &v_dequant, &v_round, &v_zbin, &v_quant_shift);
141
142 v_eobmax = get_max_lane_eob(iscan, v_eobmax, v_nz_mask);
143
144 v_round = _mm256_unpackhi_epi64(v_round, v_round);
145 v_quant = _mm256_unpackhi_epi64(v_quant, v_quant);
146 v_dequant = _mm256_unpackhi_epi64(v_dequant, v_dequant);
147 v_quant_shift = _mm256_unpackhi_epi64(v_quant_shift, v_quant_shift);
148 v_zbin = _mm256_unpackhi_epi64(v_zbin, v_zbin);
149
150 for (intptr_t count = n_coeffs - 16; count > 0; count -= 16) {
151 coeff_ptr += 16;
152 qcoeff_ptr += 16;
153 dqcoeff_ptr += 16;
154 iscan += 16;
155 v_nz_mask =
156 quantize_b_logscale0_16(coeff_ptr, qcoeff_ptr, dqcoeff_ptr, &v_quant,
157 &v_dequant, &v_round, &v_zbin, &v_quant_shift);
158
159 v_eobmax = get_max_lane_eob(iscan, v_eobmax, v_nz_mask);
160 }
161
162 *eob_ptr = accumulate_eob256(v_eobmax);
163 }
164
quantize_b_logscale_16(const tran_low_t * coeff_ptr,tran_low_t * qcoeff_ptr,tran_low_t * dqcoeff_ptr,__m256i * v_quant,__m256i * v_dequant,__m256i * v_round,__m256i * v_zbin,__m256i * v_quant_shift,int log_scale)165 static AOM_FORCE_INLINE __m256i quantize_b_logscale_16(
166 const tran_low_t *coeff_ptr, tran_low_t *qcoeff_ptr,
167 tran_low_t *dqcoeff_ptr, __m256i *v_quant, __m256i *v_dequant,
168 __m256i *v_round, __m256i *v_zbin, __m256i *v_quant_shift, int log_scale) {
169 const __m256i v_coeff = load_coefficients_avx2(coeff_ptr);
170 const __m256i v_abs_coeff = _mm256_abs_epi16(v_coeff);
171 const __m256i v_zbin_mask = _mm256_cmpgt_epi16(v_abs_coeff, *v_zbin);
172
173 if (_mm256_movemask_epi8(v_zbin_mask) == 0) {
174 _mm256_store_si256((__m256i *)qcoeff_ptr, _mm256_setzero_si256());
175 _mm256_store_si256((__m256i *)dqcoeff_ptr, _mm256_setzero_si256());
176 _mm256_store_si256((__m256i *)(qcoeff_ptr + 8), _mm256_setzero_si256());
177 _mm256_store_si256((__m256i *)(dqcoeff_ptr + 8), _mm256_setzero_si256());
178 return _mm256_setzero_si256();
179 }
180
181 // tmp = v_zbin_mask ? (int64_t)abs_coeff + log_scaled_round : 0
182 const __m256i v_tmp_rnd =
183 _mm256_and_si256(_mm256_adds_epi16(v_abs_coeff, *v_round), v_zbin_mask);
184 // tmp32 = (int)(((((tmp * quant_ptr[rc != 0]) >> 16) + tmp) *
185 // quant_shift_ptr[rc != 0]) >>
186 // (16 - log_scale + AOM_QM_BITS));
187 const __m256i v_tmp32_a = _mm256_mulhi_epi16(v_tmp_rnd, *v_quant);
188 const __m256i v_tmp32_b = _mm256_add_epi16(v_tmp32_a, v_tmp_rnd);
189 const __m256i v_tmp32_hi = _mm256_slli_epi16(
190 _mm256_mulhi_epi16(v_tmp32_b, *v_quant_shift), log_scale);
191 const __m256i v_tmp32_lo = _mm256_srli_epi16(
192 _mm256_mullo_epi16(v_tmp32_b, *v_quant_shift), 16 - log_scale);
193 const __m256i v_tmp32 = _mm256_or_si256(v_tmp32_hi, v_tmp32_lo);
194 const __m256i v_dqcoeff_hi = _mm256_slli_epi16(
195 _mm256_mulhi_epi16(v_tmp32, *v_dequant), 16 - log_scale);
196 const __m256i v_dqcoeff_lo =
197 _mm256_srli_epi16(_mm256_mullo_epi16(v_tmp32, *v_dequant), log_scale);
198 const __m256i v_dqcoeff =
199 _mm256_sign_epi16(_mm256_or_si256(v_dqcoeff_hi, v_dqcoeff_lo), v_coeff);
200 const __m256i v_qcoeff = _mm256_sign_epi16(v_tmp32, v_coeff);
201 const __m256i v_nz_mask = _mm256_cmpgt_epi16(v_tmp32, _mm256_setzero_si256());
202 store_coefficients_avx2(v_qcoeff, qcoeff_ptr);
203 store_coefficients_avx2(v_dqcoeff, dqcoeff_ptr);
204 return v_nz_mask;
205 }
206
quantize_b_no_qmatrix_avx2(const tran_low_t * coeff_ptr,intptr_t n_coeffs,const int16_t * zbin_ptr,const int16_t * round_ptr,const int16_t * quant_ptr,const int16_t * quant_shift_ptr,tran_low_t * qcoeff_ptr,tran_low_t * dqcoeff_ptr,const int16_t * dequant_ptr,uint16_t * eob_ptr,const int16_t * iscan,int log_scale)207 static AOM_FORCE_INLINE void quantize_b_no_qmatrix_avx2(
208 const tran_low_t *coeff_ptr, intptr_t n_coeffs, const int16_t *zbin_ptr,
209 const int16_t *round_ptr, const int16_t *quant_ptr,
210 const int16_t *quant_shift_ptr, tran_low_t *qcoeff_ptr,
211 tran_low_t *dqcoeff_ptr, const int16_t *dequant_ptr, uint16_t *eob_ptr,
212 const int16_t *iscan, int log_scale) {
213 __m256i v_zbin, v_round, v_quant, v_dequant, v_quant_shift;
214 __m256i v_eobmax = _mm256_setzero_si256();
215
216 load_b_values_avx2(zbin_ptr, &v_zbin, round_ptr, &v_round, quant_ptr,
217 &v_quant, dequant_ptr, &v_dequant, quant_shift_ptr,
218 &v_quant_shift, log_scale);
219
220 // Do DC and first 15 AC.
221 __m256i v_nz_mask = quantize_b_logscale_16(
222 coeff_ptr, qcoeff_ptr, dqcoeff_ptr, &v_quant, &v_dequant, &v_round,
223 &v_zbin, &v_quant_shift, log_scale);
224
225 v_eobmax = get_max_lane_eob(iscan, v_eobmax, v_nz_mask);
226
227 v_round = _mm256_unpackhi_epi64(v_round, v_round);
228 v_quant = _mm256_unpackhi_epi64(v_quant, v_quant);
229 v_dequant = _mm256_unpackhi_epi64(v_dequant, v_dequant);
230 v_quant_shift = _mm256_unpackhi_epi64(v_quant_shift, v_quant_shift);
231 v_zbin = _mm256_unpackhi_epi64(v_zbin, v_zbin);
232
233 for (intptr_t count = n_coeffs - 16; count > 0; count -= 16) {
234 coeff_ptr += 16;
235 qcoeff_ptr += 16;
236 dqcoeff_ptr += 16;
237 iscan += 16;
238 v_nz_mask = quantize_b_logscale_16(coeff_ptr, qcoeff_ptr, dqcoeff_ptr,
239 &v_quant, &v_dequant, &v_round, &v_zbin,
240 &v_quant_shift, log_scale);
241
242 v_eobmax = get_max_lane_eob(iscan, v_eobmax, v_nz_mask);
243 }
244
245 *eob_ptr = accumulate_eob256(v_eobmax);
246 }
247
aom_quantize_b_32x32_avx2(const tran_low_t * coeff_ptr,intptr_t n_coeffs,const int16_t * zbin_ptr,const int16_t * round_ptr,const int16_t * quant_ptr,const int16_t * quant_shift_ptr,tran_low_t * qcoeff_ptr,tran_low_t * dqcoeff_ptr,const int16_t * dequant_ptr,uint16_t * eob_ptr,const int16_t * scan,const int16_t * iscan)248 void aom_quantize_b_32x32_avx2(const tran_low_t *coeff_ptr, intptr_t n_coeffs,
249 const int16_t *zbin_ptr,
250 const int16_t *round_ptr,
251 const int16_t *quant_ptr,
252 const int16_t *quant_shift_ptr,
253 tran_low_t *qcoeff_ptr, tran_low_t *dqcoeff_ptr,
254 const int16_t *dequant_ptr, uint16_t *eob_ptr,
255 const int16_t *scan, const int16_t *iscan) {
256 (void)scan;
257 quantize_b_no_qmatrix_avx2(coeff_ptr, n_coeffs, zbin_ptr, round_ptr,
258 quant_ptr, quant_shift_ptr, qcoeff_ptr,
259 dqcoeff_ptr, dequant_ptr, eob_ptr, iscan, 1);
260 }
261
aom_quantize_b_64x64_avx2(const tran_low_t * coeff_ptr,intptr_t n_coeffs,const int16_t * zbin_ptr,const int16_t * round_ptr,const int16_t * quant_ptr,const int16_t * quant_shift_ptr,tran_low_t * qcoeff_ptr,tran_low_t * dqcoeff_ptr,const int16_t * dequant_ptr,uint16_t * eob_ptr,const int16_t * scan,const int16_t * iscan)262 void aom_quantize_b_64x64_avx2(const tran_low_t *coeff_ptr, intptr_t n_coeffs,
263 const int16_t *zbin_ptr,
264 const int16_t *round_ptr,
265 const int16_t *quant_ptr,
266 const int16_t *quant_shift_ptr,
267 tran_low_t *qcoeff_ptr, tran_low_t *dqcoeff_ptr,
268 const int16_t *dequant_ptr, uint16_t *eob_ptr,
269 const int16_t *scan, const int16_t *iscan) {
270 (void)scan;
271 quantize_b_no_qmatrix_avx2(coeff_ptr, n_coeffs, zbin_ptr, round_ptr,
272 quant_ptr, quant_shift_ptr, qcoeff_ptr,
273 dqcoeff_ptr, dequant_ptr, eob_ptr, iscan, 2);
274 }
275