我想使用 AVX 提高这段代码的性能

Question

我分析了我的代码，代码中最昂贵的部分是 post 中包含的循环。我想使用 AVX 提高这个循环的性能。我尝试过手动展开循环，虽然这确实提高了性能，但改进并不令人满意。

int N = 100000000;
int8_t* data = new int8_t[N];
for(int i = 0; i< N; i++) { data[i] = 1 ;}
std::array<float, 10> f  = {1,2,3,4,5,6,7,8,9,10};
std::vector<float> output(N, 0);
int k = 0;
for (int i = k; i < N; i = i + 2) {
    for (int j = 0; j < 10; j++, k = j + 1) {
        output[i] += f[j] * data[i - k];
        output[i + 1] += f[j] * data[i - k + 1];
    }
}

我可以就如何处理这个问题提供一些指导吗？

Answer 1

我假设 data 是一个大的有符号字节输入数组，f 是一个小的长度为 10 的浮点数数组，而 output 是大的输出数组花车。您的代码在 i 的前 10 次迭代中超出范围，因此我将从 i 开始，而不是从 10 开始。这是原始代码的干净版本：

int s = 10;
for (int i = s; i < N; i += 2) {
    for (int j = 0; j < 10; j++) {
        output[i]   += f[j] * data[i-j-1];
        output[i+1] += f[j] * data[i-j];
    }
}

事实证明，i 处理两次迭代不会改变任何东西，因此我们进一步简化为：

for (int i = s; i < N; i++)
    for (int j = 0; j < 10; j++)
        output[i] += f[j] * data[i-j-1];

这个版本的代码（连同 input/output 数据的声明）应该出现在问题本身中，而其他人不必 clean/simplify 搞砸。

现在很明显这个代码适用one-dimensional convolution filter, which is a very common thing in signal processing. For instance, it can by computed in Python using numpy.convolve function. The kernel has very small length 10, so Fast Fourier Transform won't provide any benefits compared to bruteforce approach. Given that the problem is well-known, you can read a lot of articles on vectorizing small-kernel convolution. I will follow the article by hgomersall。

首先，让我们去掉反向索引。显然，我们可以在运行主算法之前反转内核。之后，我们必须计算 so-called cross-correlation 而不是卷积。简而言之，我们沿着输入数组移动内核数组，并为每个可能的偏移量计算它们之间的点积。

std::reverse(f.data(), f.data() + 10);
for (int i = s; i < N; i++) {
    int b = i-10;
    float res = 0.0;
    for (int j = 0; j < 10; j++)
        res += f[j] * data[b+j];
    output[i] = res;
}

为了对其进行矢量化，让我们一次计算 8 个连续的点积。回想一下，我们可以将八个 32 位浮点数打包到一个 256 位 AVX 寄存器中。我们将通过 i 向量化外循环，这意味着：

每次迭代 i 的循环将提前 8。
外循环内的每个值都变成一个 8 元素包，这样包的 k-th 元素在标量版本的外循环的第 (i+k) 次迭代中保留该值.

这是结果代码：

//reverse the kernel
__m256 revKernel[10];
for (size_t i = 0; i < 10; i++)
    revKernel[i] = _mm256_set1_ps(f[9-i]); //every component will have same value
//note: you have to compute the last 16 values separately!
for (size_t i = s; i + 16 <= N; i += 8) {
    int b = i-10;
    __m256 res = _mm256_setzero_ps();
    for (size_t j = 0; j < 10; j++) {
        //load: data[b+j], data[b+j+1], data[b+j+2], ..., data[b+j+15]
        __m128i bytes = _mm_loadu_si128((__m128i*)&data[b+j]);
        //convert first 8 bytes of loaded 16-byte pack into 8 floats
        __m256 floats = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(bytes));
        //compute res = res + floats * revKernel[j] elementwise
        res = _mm256_fmadd_ps(revKernel[j], floats, res);
    }
    //store 8 values packed in res into: output[i], output[i+1], ..., output[i+7]
    _mm256_storeu_ps(&output[i], res);
}

对于 1 亿个元素，这段代码在我的机器上花费了大约 120 毫秒，而原始标量实现花费了 850 毫秒。当心：我有 Ryzen 1600 CPU，因此在 Intel CPUs 上的结果可能有些不同。

现在如果你真的想展开一些东西，由 10 个内核元素组成的内循环是完美的地方。这是如何完成的：

__m256 revKernel[10];
for (size_t i = 0; i < 10; i++)
    revKernel[i] = _mm256_set1_ps(f[9-i]);
for (size_t i = s; i + 16 <= N; i += 8) {
    size_t b = i-10;
    __m256 res = _mm256_setzero_ps();
    #define DOIT(j) {\
        __m128i bytes = _mm_loadu_si128((__m128i*)&data[b+j]); \
        __m256 floats = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(bytes)); \
        res = _mm256_fmadd_ps(revKernel[j], floats, res); \
    }
    DOIT(0);
    DOIT(1);
    DOIT(2);
    DOIT(3);
    DOIT(4);
    DOIT(5);
    DOIT(6);
    DOIT(7);
    DOIT(8);
    DOIT(9);
    _mm256_storeu_ps(&output[i], res);
}

在我的机器上需要 110 毫秒（比第一个矢量化版本稍微好一点）。

所有元素的简单复制（从字节到浮点数的转换）对我来说需要 40 毫秒，这意味着这段代码还不是 memory-bound，还有一些改进空间。

我想使用 AVX 提高这段代码的性能

I would like to improve the performance of this code using AVX

c++

optimization

avx

compiler-optimization

avx2