[x86]: Improve av1_quantize_fp_avx2(). 1.13x to 1.49x faster than the previous version depending on the last nonzero coeff position. Bug: b/235228922 Change-Id: I0cf3fed0d74da3cca13faea799baa92f57ff03aa
diff --git a/av1/encoder/x86/av1_quantize_avx2.c b/av1/encoder/x86/av1_quantize_avx2.c index 591edd7..d6b0e17 100644 --- a/av1/encoder/x86/av1_quantize_avx2.c +++ b/av1/encoder/x86/av1_quantize_avx2.c
@@ -95,6 +95,21 @@ } \ } while (0) +static INLINE __m256i load_coefficients_avx2(const tran_low_t *coeff_ptr) { + const __m256i coeff1 = _mm256_load_si256((__m256i *)coeff_ptr); + const __m256i coeff2 = _mm256_load_si256((__m256i *)(coeff_ptr + 8)); + return _mm256_packs_epi32(coeff1, coeff2); +} + +static INLINE void store_coefficients_avx2(__m256i coeff_vals, + tran_low_t *coeff_ptr) { + __m256i coeff_sign = _mm256_srai_epi16(coeff_vals, 15); + __m256i coeff_vals_lo = _mm256_unpacklo_epi16(coeff_vals, coeff_sign); + __m256i coeff_vals_hi = _mm256_unpackhi_epi16(coeff_vals, coeff_sign); + _mm256_store_si256((__m256i *)coeff_ptr, coeff_vals_lo); + _mm256_store_si256((__m256i *)(coeff_ptr + 8), coeff_vals_hi); +} + static INLINE uint16_t quant_gather_eob(__m256i eob) { const __m128i eob_lo = _mm256_castsi256_si128(eob); const __m128i eob_hi = _mm256_extractf128_si256(eob, 1); @@ -104,34 +119,6 @@ return INT16_MAX - _mm_extract_epi16(eob_s, 0); } -static INLINE void quantize(const __m256i *thr, const __m256i *qp, __m256i *c, - const int16_t *iscan_ptr, tran_low_t *qcoeff, - tran_low_t *dqcoeff, __m256i *eob) { - const __m256i abs_coeff = _mm256_abs_epi16(*c); - __m256i mask = _mm256_cmpgt_epi16(abs_coeff, *thr); - mask = _mm256_or_si256(mask, _mm256_cmpeq_epi16(abs_coeff, *thr)); - const int nzflag = _mm256_movemask_epi8(mask); - - if (nzflag) { - __m256i q = _mm256_adds_epi16(abs_coeff, qp[0]); - q = _mm256_mulhi_epi16(q, qp[1]); - q = _mm256_sign_epi16(q, *c); - const __m256i dq = _mm256_mullo_epi16(q, qp[2]); - - store_two_quan(q, qcoeff, dq, dqcoeff); - const __m256i zero = _mm256_setzero_si256(); - const __m256i iscan = _mm256_loadu_si256((const __m256i *)iscan_ptr); - const __m256i zero_coeff = _mm256_cmpeq_epi16(dq, zero); - const __m256i nzero_coeff = _mm256_cmpeq_epi16(zero_coeff, zero); - __m256i cur_eob = _mm256_sub_epi16(iscan, nzero_coeff); - cur_eob = _mm256_and_si256(cur_eob, nzero_coeff); - *eob = _mm256_max_epi16(*eob, cur_eob); - } else { - write_zero(qcoeff); - write_zero(dqcoeff); - } -} - static INLINE __m256i scan_eob_256(const __m256i *iscan_ptr, __m256i *coeff256) { const __m256i iscan = _mm256_loadu_si256(iscan_ptr); @@ -235,6 +222,42 @@ *eob_ptr = accumulate_eob(eob); } +static AOM_FORCE_INLINE __m256i get_max_lane_eob(const int16_t *iscan, + __m256i v_eobmax, + __m256i v_mask) { + const __m256i v_iscan = _mm256_loadu_si256((const __m256i *)iscan); + const __m256i v_iscan_perm = _mm256_permute4x64_epi64(v_iscan, 0xD8); + const __m256i v_iscan_plus1 = _mm256_sub_epi16(v_iscan_perm, v_mask); + const __m256i v_nz_iscan = _mm256_and_si256(v_iscan_plus1, v_mask); + return _mm256_max_epi16(v_eobmax, v_nz_iscan); +} + +static AOM_FORCE_INLINE void quantize_fp_16( + const __m256i *thr, const __m256i *qp, const tran_low_t *coeff_ptr, + const int16_t *iscan_ptr, tran_low_t *qcoeff_ptr, tran_low_t *dqcoeff_ptr, + __m256i *eob) { + const __m256i coeff = load_coefficients_avx2(coeff_ptr); + const __m256i abs_coeff = _mm256_abs_epi16(coeff); + const __m256i mask = _mm256_cmpgt_epi16(abs_coeff, *thr); + const int nzflag = _mm256_movemask_epi8(mask); + + if (nzflag) { + const __m256i tmp_rnd = _mm256_adds_epi16(abs_coeff, qp[0]); + const __m256i abs_q = _mm256_mulhi_epi16(tmp_rnd, qp[1]); + const __m256i q = _mm256_sign_epi16(abs_q, coeff); + const __m256i dq = _mm256_mullo_epi16(q, qp[2]); + const __m256i nz_mask = _mm256_cmpgt_epi16(abs_q, _mm256_setzero_si256()); + + store_coefficients_avx2(q, qcoeff_ptr); + store_coefficients_avx2(dq, dqcoeff_ptr); + + *eob = get_max_lane_eob(iscan_ptr, *eob, nz_mask); + } else { + write_zero(qcoeff_ptr); + write_zero(dqcoeff_ptr); + } +} + void av1_quantize_fp_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, @@ -245,17 +268,18 @@ (void)scan_ptr; (void)zbin_ptr; (void)quant_shift_ptr; - const unsigned int step = 16; - __m256i qp[3]; - __m256i coeff, thr; const int log_scale = 0; + const int step = 16; + __m256i qp[3], thr; + __m256i eob = _mm256_setzero_si256(); init_qp(round_ptr, quant_ptr, dequant_ptr, log_scale, &thr, qp); - read_coeff(coeff_ptr, &coeff); + // Subtracting 1 here eliminates a _mm256_cmpeq_epi16() instruction when + // calculating the zbin mask. + thr = _mm256_sub_epi16(thr, _mm256_set1_epi16(1)); - __m256i eob = _mm256_setzero_si256(); - quantize(&thr, qp, &coeff, iscan_ptr, qcoeff_ptr, dqcoeff_ptr, &eob); + quantize_fp_16(&thr, qp, coeff_ptr, iscan_ptr, qcoeff_ptr, dqcoeff_ptr, &eob); coeff_ptr += step; qcoeff_ptr += step; @@ -263,11 +287,14 @@ iscan_ptr += step; n_coeffs -= step; - update_qp(log_scale, &thr, qp); + qp[0] = _mm256_permute2x128_si256(qp[0], qp[0], 0x11); + qp[1] = _mm256_permute2x128_si256(qp[1], qp[1], 0x11); + qp[2] = _mm256_permute2x128_si256(qp[2], qp[2], 0x11); + thr = _mm256_permute2x128_si256(thr, thr, 0x11); while (n_coeffs > 0) { - read_coeff(coeff_ptr, &coeff); - quantize(&thr, qp, &coeff, iscan_ptr, qcoeff_ptr, dqcoeff_ptr, &eob); + quantize_fp_16(&thr, qp, coeff_ptr, iscan_ptr, qcoeff_ptr, dqcoeff_ptr, + &eob); coeff_ptr += step; qcoeff_ptr += step;