quantize_avx2.c (13084B)
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 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 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 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 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 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 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 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 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 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 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 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 }