如何用AVX512/AVX2计算8位整数向量L2距离及累加寄存器值

核心问题：累加AVX512寄存器中的16个int32值 #

要高效累加__m512i寄存器内的16个32位整数，最优方案是使用AVX512F指令集的水平归约指令_mm512_reduce_add_epi32——它会直接将寄存器内所有int32元素求和，返回一个标量int32值，代码简洁且性能最优。

如果需要兼容不支持reduce指令的老编译器，也可以手动拆分寄存器逐步累加，但不推荐（代码繁琐且性能较差）。

完整修正与优化后的代码 #

以下代码修复了原版本的语法错误、逻辑问题，并补充了余数处理、类型规范等细节：

#include <immintrin.h>
#include <cstddef>

static float L2SqrSQ8(const int8_t* pVect1v, const int8_t* pVect2v, size_t qty) {
    const int8_t* x = pVect1v;
    const int8_t* y = pVect2v;
    __m512i sum = _mm512_setzero_si512();
    size_t i;

    // 批量处理16个int8元素（对应转换为16个int32）
    for (i = 0; i + 16 <= qty; i += 16) {
        __m128i xx = _mm_loadu_si128(reinterpret_cast<const __m128i*>(x + i));
        __m128i yy = _mm_loadu_si128(reinterpret_cast<const __m128i*>(y + i));
        
        __m512i xx_ext = _mm512_cvtepi8_epi32(xx);
        __m512i yy_ext = _mm512_cvtepi8_epi32(yy);
        
        __m512i sub = _mm512_sub_epi32(xx_ext, yy_ext);
        __m512i square = _mm512_mullo_epi32(sub, sub);
        sum = _mm512_add_epi32(sum, square);
    }

    // 累加寄存器内所有int32值
    int32_t total = _mm512_reduce_add_epi32(sum);

    // 处理剩余不足16个的元素（标量兜底）
    for (; i < qty; ++i) {
        int32_t diff = static_cast<int32_t>(x[i]) - static_cast<int32_t>(y[i]);
        total += diff * diff;
    }

    // 转换为float返回
    return static_cast<float>(total);
}

int main(int argc, char** argv) {
    int8_t a0[32] = {1, 2, 3, -4, 5, 6, 7, 8, 9, -127, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, -24, 25, 26, 27, 28, 29, 30, 31, 32};
    int8_t a1[32] = {2, 3, 4, 5, 6, 7, 8, 9, 10, 127, 12, -13, 14, 15, 16, 17, 18, 19, -20, 21, 22, 23, 24, -25, 26, 27, 28, 29, 30, 31, 32, 33};
    float result = L2SqrSQ8(a0, a1, 32);
    return 0;
}

代码评审与优化建议 #

1. 移除硬编码逻辑：原代码强制qty=32忽略传入参数，修复后支持任意长度的向量计算
2. 修复语法错误：补全_mm512_mullo_epi32后的分号，统一循环变量i的类型为size_t避免重复定义
3. 内存对齐优化：如果向量内存是16字节对齐的，改用_mm_load_si128替代_mm_loadu_si128可提升加载性能
4. 类型规范：使用reinterpret_cast替代C风格强制转换，符合C++语法规范
5. 余数处理：新增标量循环处理不足16个的剩余元素，避免越界计算
6. 编译选项：编译时需启用AVX512F支持（如GCC使用-mavx512f选项）

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