SIMD AVX2 向量计算

• _mm256_fmadd_ps: 能够在单个操作中执行乘法和加法，从而提高浮点计算的精度和性能。
• _mm256_sub_ps : Intel Advanced Vector Extensions (AVX) 指令集中用于从两个 AVX 寄存器中逐元素进行单精度浮点数减法的内联函数。这个函数允许同时对 8 个单精度浮点数进行减法操作，适合于执行高效率的向量化并行减法运算。
• _mm_add_ps: 是 Intel 提供的 SSE（SIMD 流扩展）指令集，用于累加存储在 SSE 寄存器中的四个单精度浮点数。它专门处理 128 位 __m128 数据类型，这使得它对于优化浮点数组的操作非常有效。
• _mm_cvtss_f32: 将 SSE 寄存器 （__m128） 的最低单精度浮点元素转换为标量浮点值。
• _mm256_load_ps: Intel Advanced Vector Extensions (AVX) 指令集中用于从内存中加载 256 位数据到 AVX 寄存器的内联函数。这个函数专门用于加载 8 个连续的单精度浮点数（floats），是处理高性能计算任务时常用的指令之一。

注意：使用avx需要保证输入的数据内存对齐，申请数组时需要注意，不然会发生段错误。

样例代码：

#include <iostream>
#include <immintrin.h>
#include <cstdlib>
#include <ctime>
#include <chrono>   // For high_resolution_clock

// Helper function to reduce an __m256 to a single float value
float _mm256_reduce_add_ps(__m256 x) {
    const __m128 x128 = _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x));
    const __m128 x64 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128));
    const __m128 x32 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55));
    return _mm_cvtss_f32(x32);
}

// Compute the squared L2 distance between two vectors
float VectorL2SquaredDistance(int dim, float *ax, float *bx, bool use_avx) {
    float distance = 0.0f;

    if (dim % 8 == 0 && use_avx == true) {
        uint16_t niters = dim / 8;
        __m256 sum = _mm256_setzero_ps();

        for (uint16_t j = 0; j < niters; j++) {
            __m256 a_vec = _mm256_load_ps(ax + 8 * j);
            __m256 b_vec = _mm256_load_ps(bx + 8 * j);
            __m256 tmp_vec = _mm256_sub_ps(a_vec, b_vec);
            sum = _mm256_fmadd_ps(tmp_vec, tmp_vec, sum);
        }

        distance = _mm256_reduce_add_ps(sum);
        // std::cout << "使用 AVX" << std::endl;
    } else {
        for (int i = 0; i < dim; i++) {
            float diff = ax[i] - bx[i];
            distance += diff * diff;
        }
        // std::cout << "未使用 AVX" << std::endl;
    }
    return distance;
}

int main() {
    // cmd: g++ main.cpp -mavx2 -mfma && ./a.out
    const int dim = 960; // Vector dimension
    float* ax;
    float* bx;

    // Allocate aligned memory
    if (posix_memalign((void**)&ax, 32, dim * sizeof(float))) {
        std::cerr << "Failed to allocate memory for ax" << std::endl;
        return 1;
    }

    if (posix_memalign((void**)&bx, 32, dim * sizeof(float))) {
        std::cerr << "Failed to allocate memory for bx" << std::endl;
        free(ax);  // Free previously allocated memory
        return 1;
    }

    size_t test_cnt = 100000;

    {
        auto start = std::chrono::high_resolution_clock::now();

        float sum = 0;
        for (int i = 0; i < test_cnt; i++) {
            // Initialize random values for the vectors
            std::srand(0);
            for (int i = 0; i < dim; i++) {
                // ax[i] = static_cast<float>(std::rand()) / RAND_MAX;
                // bx[i] = static_cast<float>(std::rand()) / RAND_MAX;
                ax[i] = static_cast<float>(i);
                bx[i] = static_cast<float>(i+1);
            }

            float distance = VectorL2SquaredDistance(dim, ax, bx, false);
            sum += distance;
        }
        std::cout << "not use AVX: Squared L2 distance: " << sum << std::endl;

        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> elapsed = end - start;
        std::cout << "  Elapsed time: " << elapsed.count() << " ms\n";
    }
    {
        auto start = std::chrono::high_resolution_clock::now();

        float sum = 0;
        for (int i = 0; i < test_cnt; i++) {
            // Initialize random values for the vectors
            std::srand(0);
            for (int i = 0; i < dim; i++) {
                ax[i] = static_cast<float>(i);
                bx[i] = static_cast<float>(i+1);
            }

            float distance = VectorL2SquaredDistance(dim, ax, bx, true);
            sum += distance;
        }
        std::cout << "use AVX:Squared L2 distance: " << sum << std::endl;

        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> elapsed = end - start;
        std::cout << "  Elapsed time: " << elapsed.count() << " ms\n";
    }

    // Clean up
    std::free(ax);
    std::free(bx);

    return 0;
}

测试结果如下，开O3优化时差异可以达到近两倍:

$ g++ main.cpp -mavx2 -mfma && ./a.out
not use AVX: Squared L2 distance: 9.6e+07
  Elapsed time: 550.872 ms
use AVX:Squared L2 distance: 9.6e+07
  Elapsed time: 379.223 ms
$ g++ main.cpp -O3 -mfma -std=c++17 -march=native -mavx2 && ./a.out
not use AVX: Squared L2 distance: 9.6e+07
  Elapsed time: 196.283 ms
use AVX:Squared L2 distance: 9.6e+07
  Elapsed time: 111.351 ms