performance - 为长度为 8 的一维卷积核加载 AVX 向量的高效代码

Question

一维卷积运算的实现通常需要加载数据向量，这些数据向量在每次迭代时按顺序逐步遍历数据缓冲区，偏移一个元素。

例如，考虑输入数据的缓冲区 X[0]、X[1]、...、X[n-1]，其中 n 大于内核长度的两倍。如果卷积长度为 3，并且我们可以在每个向量中容纳 8 个元素，我们可能首先需要一个具有 X[0]、X[1]、.. ., X[7]，然后是 X[1]、X[2]、...、X [8] 和最后一个 X[2]、X[3]、...、X[9] .

考虑内核长度和向量长度为​​ 8 的情况。我们必须加载八个向量，它们可能按顺序如下所示:

{  0   1   2   3   4   5   6   7  }
{  1   2   3   4   5   6   7   8  }
{  2   3   4   5   6   7   8   9  }
{  3   4   5   6   7   8   9  10  }
{  4   5   6   7   8   9  10  11  }
{  5   6   7   8   9  10  11  12  }
{  6   7   8   9  10  11  12  13  }
{  7   8   9  10  11  12  13  14  }

通过垂直减少这个序列，我们可以产生运行平均值或总和。即，这些向量的总和将是其第一个位置的前 8 个元素的总和。

考虑到列中元素的顺序并不重要。每列中元素的任何排列仍将产生相同的结果。对于卷积，可以通过改变内核中使用的常量的顺序来解释这种排列。

有没有一种更快的方法来加载这些利用这一点的向量？将未对齐负载的简单序列视为基线:

// Any sort of sliding window function, i.e. running mean, running max, convolution, etc.
void sliding_window(const float* input, unsigned length)
{
    for (unsigned i = 0; i < length - 7; i += 8) {
        for (unsigned j = 0; i < 8; j++) {
            __m256 v = _mm256_loadu_ps(input[i + j]);
            // reduction operation on v (e.g. max or fmadd) goes here
        }
    }
    // handle tail here
}

Answer 1

首先，您应该注意，如果您的卷积是可分离的，那么这通常是值得做的。简单的例子:

res[i] = x[i]+x[i+1]+x[i+2]+x[i+3]+x[i+4]+x[i+5]+x[i+6]+x[i+7];

这可以通过分三步与 [1 1] * [1 0 1] * [1 0 0 0 1] 进行卷积来完成，例如如下所示:

void sliding_window(float* output, const float* input, size_t length)
{
    // Nomenclature
    // aX input at i+X
    // bX convolution with [1 1] starting at i+X
    // cX convolution with [1 1] * [1 0 1] starting at i+X
    // dX convolution with [1 1] * [1 0 1] * [1 0 0 0 1] starting at i+X

    __m256 a0 = _mm256_load_ps(input), a8 = _mm256_load_ps(input + 8);
    __m256 b0 = _mm256_add_ps(a0, _mm256_loadu_ps(input+1)), b8 = _mm256_add_ps(a8, _mm256_loadu_ps(input+9));
    __m256 b4 = _mm256_permute2f128_ps(b0, b8, 1+16*2);
    __m256 b2 = _mm256_shuffle_ps(b0, b4, 2+3*4+0*16+1*64);
    __m256 c0 = _mm256_add_ps(b0, b2);

    for (unsigned i = 0; i < length - 25; i += 8) {
        // Convolute input with [1 1]
        __m256 a16 = _mm256_load_ps( input + i + 16);
        __m256 a17 = _mm256_loadu_ps(input + i + 17);
        __m256 b16 = _mm256_add_ps(a16, a17);

        // Convolute first convolution with [1 0 1]
        __m256 b12 = _mm256_permute2f128_ps(b8, b16, 1+16*2);
        __m256 b10 = _mm256_shuffle_ps(b8, b12, 2+3*4+0*16+1*64);
        __m256 c8 = _mm256_add_ps(b8, b10);

        // Convolute second convolution with [1 0 0 0 1]
        __m256 c4 = _mm256_permute2f128_ps(c0, c8, 1+16*2);
        __m256 d0 = _mm256_add_ps(c0, c4);

        // Store result
        _mm256_store_ps(output + i, d0);

        // rename registers for next iteration:
        b8 = b16;
        c0 = c8;
    }
    // handle tail here ...
}

您当然可以用 maxps 替换 addps。 Godbolt-演示:https://godbolt.org/z/W9K9o943o

总的来说，这需要 1 次对齐 + 1 次未对齐加载、3 次洗牌、3 次添加和 1 次存储来存储 8 个元素(实际上仅使用 AVX1)。在每个周期只有 1 次 shuffle 的 Intel CPU 上，这实际上可能只比简单的 8 负载、7 加法实现稍快一些(我没有对此进行基准测试)。在 Zen3 上，我不确定加载未对齐数据的实际成本。

如果你有一个不平凡的内核，那么可能很难确定它是否是可分离的。