使用AVX实现的C++整数向量点积未达预期提速效果

已实现的代码 #

普通循环版本 #

int dotProductNoAVX(const std::vector<int>& vec1, const std::vector<int>& vec2)
{
    // 用普通循环计算点积
    int dotProduct = 0;
    for (std::size_t i = 0; i < vec1.size(); ++i)
    {
        dotProduct += vec1[i] * vec2[i];
    }

    return dotProduct;
}

AVX256实现版本 #

// 用AVX intrinsics计算向量点积
int dotProductAVX(const std::vector<int>& vec1, const std::vector<int>& vec2)
{
    // 确保处理的元素数量是8的倍数（AVX寄存器大小）
    std::size_t size = vec1.size();
    std::size_t avxSize = size / 8 * 8;

    // 初始化AVX求和寄存器
    __m256i sum = _mm256_setzero_si256();

    // AVX并行计算点积
    for (std::size_t i = 0; i < avxSize; i += 8)
    {
        __m256i vec1_part = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vec1[i]));
        __m256i vec2_part = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vec2[i]));
        sum = _mm256_add_epi32(sum, _mm256_mullo_epi32(vec1_part, vec2_part));
    }

    // 将AVX寄存器结果存入数组，再累加所有元素
    int result[8];
    _mm256_storeu_si256(reinterpret_cast<__m256i*>(result), sum);
    int dotProduct = result[0] + result[1] + result[2] + result[3] + result[4] + result[5] + result[6] + result[7];

    // 处理剩余的非8倍数元素
    for (std::size_t i = avxSize; i < size; ++i)
    {
        dotProduct += vec1[i] * vec2[i];
    }

    return dotProduct;
}

测试情况 #

测试环境为VS2017、i7-12700K（支持AVX），测试使用的main函数及辅助代码如下：

std::vector<int> generateRandomVector(std::size_t size)
{
    std::vector<int> randomVector;
    randomVector.reserve(size);

    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<int> distribution(1, 100);

    for (std::size_t i = 0; i < size; ++i)
    {
        randomVector.push_back(distribution(gen));
    }

    return randomVector;
}

int main()
{
    // 指定向量大小
    std::size_t vectorSize = 100;
    int iterations = 1000000;
    // 生成两个随机向量
    std::vector<int> vector1 = generateRandomVector(vectorSize);
    std::vector<int> vector2 = generateRandomVector(vectorSize);

    // 测试AVX版本耗时
    int resultAVX;
    auto startAVX = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < iterations; i++)
        resultAVX = dotProductAVX(vector1, vector2);
    auto endAVX = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> durationAVX = endAVX - startAVX;

    // 测试普通版本耗时
    int resultNoAVX;
    auto startNoAVX = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < iterations; i++)
        resultNoAVX = dotProductNoAVX(vector1, vector2);
    auto endNoAVX = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> durationNoAVX = endNoAVX - startNoAVX;

    // 输出结果和耗时
    std::cout << "Dot product without AVX: " << resultNoAVX << std::endl;
    std::cout << "Time taken without AVX: " << durationNoAVX.count() << " seconds" << std::endl;

    std::cout << "Dot product with AVX: " << resultAVX << std::endl;
    std::cout << "Time taken with AVX: " << durationAVX.count() << " seconds" << std::endl;

    return 0;
}

测试结果显示两个版本速度大致相同，仅AVX版本偶尔快10-20%。性能分析发现，AVX版本中将寄存器结果存入数组再逐个累加的步骤耗时较长，需要优化。

优化方案 #

1. 寄存器内完成水平求和，避免内存操作 #

原来的代码需要把AVX寄存器中的8个int结果存入内存数组，再逐一累加，这会带来不必要的内存读写开销。可以使用AVX的水平加法指令_mm256_hadd_epi32在寄存器内完成求和：

优化后的AVX函数示例：

int dotProductAVXOptimized(const std::vector<int>& vec1, const std::vector<int>& vec2)
{
    std::size_t size = vec1.size();
    std::size_t avxSize = size / 8 * 8;

    __m256i sum = _mm256_setzero_si256();

    for (std::size_t i = 0; i < avxSize; i += 8)
    {
        __m256i vec1_part = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vec1[i]));
        __m256i vec2_part = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&vec2[i]));
        sum = _mm256_add_epi32(sum, _mm256_mullo_epi32(vec1_part, vec2_part));
    }

    // 寄存器内水平求和8个int元素
    __m256i hsum1 = _mm256_hadd_epi32(sum, sum);
    __m256i hsum2 = _mm256_hadd_epi32(hsum1, hsum1);
    // 提取最终求和结果
    int dotProduct = _mm_cvtsi128_si32(_mm256_castsi256_si128(hsum2)) + 
                     _mm_extract_epi32(_mm256_extractf128_si256(hsum2, 1), 0);

    // 处理剩余元素
    for (std::size_t i = avxSize; i < size; ++i)
    {
        dotProduct += vec1[i] * vec2[i];
    }

    return dotProduct;
}

2. 检查编译器自动向量化 #

VS2017在/O2优化级别下，普通循环版本可能已经被编译器自动转换为AVX指令，导致手动AVX版本优势不明显。可以在普通循环前添加#pragma loop(no_vector)禁用自动向量化，测试手动AVX的真实性能。

3. 增大测试向量尺寸 #

当前测试的向量大小为100，AVX仅处理96个元素（12次循环），初始化和收尾操作的开销占比过高，导致并行优势无法体现。换成更大的向量尺寸（比如10000或更大），AVX的并行计算优势会更显著。

4. 内存对齐优化 #

当前使用_mm256_loadu_si256进行非对齐加载，如果能保证向量内存是32字节对齐的（AVX寄存器大小），可以改用_mm256_load_si256，减少对齐检查的开销。可以通过自定义vector分配器，使用std::aligned_alloc分配对齐内存来实现。

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