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

Question

我分析了我的代码，代码中最昂贵的部分是帖子中包含的循环。我想使用 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 的 float 组，output 是 float 的大输出数组。您的代码在 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];

这个版本的代码(连同输入/输出数据的声明)应该出现在问题本身中，而无需其他人清理/简化困惑。

---

现在很明显这段代码适用one-dimensional convolution filter ，这是信号处理中很常见的事情。例如，它可以在 Python 中使用 numpy.convolve 进行计算。功能。内核的长度非常小，为 10，所以 Fast Fourier Transform与蛮力方法相比，不会提供任何好处。鉴于这个问题是众所周知的，你可以阅读很多关于向量化小核卷积的文章。我会关注the article by hgomersall .

首先，让我们摆脱反向索引。显然，我们可以在运行主算法之前反转内核。之后，我们必须计算所谓的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 位 float 打包到一个 256 位 AVX 寄存器中。我们将通过 i 向量化外循环，这意味着:

• 每次迭代 i 的循环将提前 8。
• 外循环内的每个值都变成一个 8 元素包，这样包的第 k 个元素在标量版本的外循环的第 (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 CPU 上的结果可能会有所不同。

现在如果你真的想展开一些东西，由 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 毫秒(比第一个矢量化版本稍微好一点)。

所有元素的简单复制(从字节到 float 的转换)对我来说需要 40 毫秒，这意味着这段代码还没有内存限制，还有一些改进空间。