我想使用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的小浮点数组，而output是大的浮点输出数组。 您的代码去出界的第10次迭代的由i ，所以我将开始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];

此版本的代码（以及输入/输出数据的声明）应该已经存在于问题本身中，而其他代码不必清理/简化混乱情况。

---

现在很明显，该代码应用了一维卷积滤波器 ，这在信号处理中非常普遍。 例如，它可以使用numpy.convolve函数在Python中进行计算。 内核的长度非常小10，因此与bruteforce方法相比， 快速傅里叶变换不会提供任何好处。 鉴于这个问题是众所周知的，您可以阅读许多关于向量化小内核卷积的文章。 我将关注hgomersall的文章 。

首先，让我们摆脱反向索引。 显然，我们可以在运行主算法之前反转内核。 之后，我们必须计算所谓的互相关而不是卷积。 简而言之，我们沿着输入数组移动内核数组，并针对每个可能的偏移量计算它们之间的点积。

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对外循环进行矢量化处理，这意味着：

• 我的循环将在每次迭代中前进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毫秒（比第一个矢量化版本好一点）。

对我来说，所有元素的简单副本（从字节到浮点的转换）花了40毫秒，这意味着此代码尚未受内存限制，还有一些改进的余地。