如何使用 AVX 指令优化用 C 编写的 ReLU

Question

我正在尝试优化以下简化的 ReLU 模拟代码。该代码使用三元运算，这可能会妨碍编译器的自动矢量化。如何矢量化这段代码？

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mkl.h>

void relu(double* &Amem, double* Z_curr, int bo)
{
    for (int i=0; i<bo; ++i) {
        Amem[i] = Z_curr[i] > 0 ? Z_curr[i] : Z_curr[i]*0.1;
    }
}

int main()
{
    int i, j;
    int batch_size = 16384;
    int output_dim = 21;
    // double* Amem = new double[batch_size*output_dim];
    // double* Z_curr = new double[batch_size*output_dim];
    double* Amem = (double *)mkl_malloc(batch_size*output_dim*sizeof( double ), 64 );
    double* Z_curr = (double *)mkl_malloc(batch_size*output_dim*sizeof( double ), 64 );
    memset(Amem, 0, sizeof(double)*batch_size*output_dim);
    for (i=0; i<batch_size*output_dim; ++i) {
        Z_curr[i] = -1+2*((double)rand())/RAND_MAX;
    }
    relu(Amem, Z_curr, batch_size*output_dim);
}

要编译它，如果您有 MKL，则使用以下内容，否则使用普通 g++ -O3。

g++ -O3 ex.cxx -L${MKLROOT}/lib/intel64 -lmkl_intel_ilp64 -lmkl_intel_thread -lmkl_core -liomp5

到目前为止，我已尝试将 -march=skylake-avx512 添加为编译器选项，但它并没有像我使用选项 -fopt-info-vec-all 进行编译时发现的那样对循环进行矢量化：

ex.cxx:9:16: missed: couldn't vectorize loop
ex.cxx:9:16: missed: not vectorized: control flow in loop.
ex.cxx:6:6: note: vectorized 0 loops in function.
ex.cxx:9:16: missed: couldn't vectorize loop
ex.cxx:9:16: missed: not vectorized: control flow in loop.

这是我目前花费的时间：

time ./a.out
real    0m0.034s
user    0m0.026s
sys     0m0.009s

Answer 1

通过引用传递指针通常没有任何好处（除非您想修改指针本身）。此外，您可以使用（非标准）__restrict 关键字帮助您的编译器，告诉它输入和输出之间不会发生别名（当然，这可能会给出错误的结果，例如，Amem == Z_curr+1 -- 但是@ 987654326@ 应该（在这种情况下）没问题）。

void relu(double* __restrict Amem, double* Z_curr, int bo)

单独使用它，clang 实际上能够使用vcmpltpd 和屏蔽移动（出于某些原因，仅使用 256 位寄存器）对循环进行矢量化。

如果您将表达式简化为 std::max(Z_curr[i], 0.1*Z_curr[i])，即使 gcc 也很容易对其进行矢量化：https://godbolt.org/z/eTv4PnMWb

一般来说，我建议使用不同的编译器和不同的编译选项编译代码的关键例程（有时尝试-ffast-math 可以向您展示简化表达式的方法）并查看生成的代码。为了可移植性，您可以将生成的代码重新翻译成内在函数（或者保持原样，如果您关心的每个编译器都能提供足够好的结果）。

为了完整起见，这里是一个可能的使用内在函数的手动矢量化实现：

void relu_avx512(double* __restrict Amem, double* Z_curr, int bo)
{
    int i;
    for (i=0; i<=bo-8; i+=8)
    {
        __m512d z = _mm512_loadu_pd(Z_curr+i);
        __mmask8 positive = _mm512_cmplt_pd_mask (_mm512_setzero_pd(), z);
        __m512d res = _mm512_mask_mul_pd(z, positive, z, _mm512_set1_pd(0.9));
        _mm512_storeu_pd(Amem+i, res);
    }
    // remaining elements in scalar loop
    for (; i<bo; ++i) {
        Amem[i] = 0.0 < Z_curr[i] ? Z_curr[i] : Z_curr[i]*0.1;;
    }
}

Godbolt：https://godbolt.org/z/s6br5KEEc（如果你在 clang 上使用 -O2 或 -O3 编译它，它会大量展开清理循环，即使它不能超过 7 次迭代。理论上，你可以这样做其余元素带有掩码或重叠存储（或者您的用例可能保证大小为 8 的倍数，您可以将其保留）。