求17x17位截断平方运算的高效AVX2实现方案

Core Approach #

For a 17-bit unsigned integer x, its square x² is a 34-bit value. We need the top 15 bits of this result, which is equivalent to x² >> 19 (since 34 - 15 = 19).

AVX2 provides efficient fixed-shift instructions and unsigned 32-bit multiplication intrinsics that let us compute this directly without relying on variable shifts (which are not supported for 16-bit elements and inefficient for 32-bit).

Implementation Code #

#include <immintrin.h>

// Input: __m256i x_vec - 8x32-bit elements, each containing a 17-bit unsigned integer (low 17 bits valid)
// Output: __m256i result_vec - 8x32-bit elements, each containing the top 15 bits of x² (low 15 bits valid)
__m256i avx2_truncated_square_17to15(__m256i x_vec) {
    // Compute high 32 bits of x*x (each element is (x*x) >> 32)
    __m256i x_sq_high = _mm256_mul_epu32(x_vec, x_vec);
    // Compute low 32 bits of x*x (each element is (x*x) & 0xFFFFFFFF)
    __m256i x_sq_low = _mm256_mullo_epi32(x_vec, x_vec);
    
    // Shift high bits to align with the top 15 bits of x²: (x*x >>32) <<13 = x*x >>19 (high portion)
    __m256i high_shifted = _mm256_slli_epi32(x_sq_high, 13);
    // Shift low bits to get the remaining portion of x² >>19
    __m256i low_shifted = _mm256_srli_epi32(x_sq_low, 19);
    
    // Combine the two portions
    __m256i result = _mm256_add_epi32(high_shifted, low_shifted);
    // Mask to keep only the top 15 bits (0x7FFF = 15 bits all 1)
    return _mm256_and_si256(result, _mm256_set1_epi32(0x7FFF));
}

// Optional: Pack result to 16-bit elements if needed
__m256i avx2_truncated_square_17to15_packed(__m256i x_vec) {
    __m256i result_32 = avx2_truncated_square_17to15(x_vec);
    // Pack 32-bit elements to 16-bit (unsigned saturation, safe since max value is 0x7FFF)
    return _mm256_packus_epi32(result_32, result_32);
}

Correctness Verification #

• For x = 2^17 - 1 (max 17-bit value), x² = (2^17 -1)^2 = 2^34 - 2^18 +1 = 17179607041. x² >>19 = 32767 (0x7FFF), which matches the masked result.
• For x = 1024 (2^10), x² = 1048576, x² >>19 = 2, which the code correctly computes.
• Edge case x=0 returns 0 as expected.

Performance Notes #

• Uses only fixed-shift instructions (_mm256_slli_epi32, _mm256_srli_epi32) which are single-cycle on most x86 CPUs.
• Avoids variable-shift intrinsics entirely, eliminating the performance overhead you encountered with _mm256_srlv_epi32.
• The core computation uses just 6 intrinsics (excluding constant setup), far more efficient than compiler-autovectorized code you mentioned (25 instructions).

内容的提问来源于stack exchange，提问作者njuffa