如何使用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 -- 但 Amem == Z_curr 应该（在这种情况下）没问题）。

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 的倍数的用例，您可以将其保留）。

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

How to use AVX instructions to optimize ReLU written in C

c

vectorization

conditional-operator