c++ - 如何编写可移植的 simd 代码以实现复杂的乘法归约

Question

我想编写快速 simd 代码来计算复杂数组的乘法归约。在标准 C 中，这是:

#include <complex.h>
complex float f(complex float x[], int n ) {
   complex float p = 1.0;
   for (int i = 0; i < n; i++)
      p *= x[i];
   return p;
}

n 最多为 50。

Gcc 不能自动矢量化复数乘法，但是，我很乐意假设 gcc 编译器，如果我知道我想以 sse3 为目标，我可以关注 How to enable sse3 autovectorization in gcc并写:

typedef float v4sf __attribute__ ((vector_size (16)));
typedef union {
  v4sf v;
  float e[4];
} float4
typedef struct {
  float4 x;
  float4 y;
} complex4;
static complex4 complex4_mul(complex4 a, complex4 b) {
  return (complex4){a.x.v*b.x.v -a.y.v*b.y.v, a.y.v*b.x.v + a.x.v*b.y.v};
}
complex4 f4(complex4 x[], int n) {
  v4sf one = {1,1,1,1};
  complex4 p = {one,one};
  for (int i = 0; i < n; i++) p = complex4_mul(p, x[i]);
  return p;
}

这确实可以使用 gcc 生成快速矢量化汇编代码。虽然您仍然需要将输入填充为 4 的倍数。您得到的程序集是:

.L3:
    vmovaps xmm0, XMMWORD PTR 16[rsi]
    add     rsi, 32
    vmulps  xmm1, xmm0, xmm2
    vmulps  xmm0, xmm0, xmm3
    vfmsubps        xmm1, xmm3, XMMWORD PTR -32[rsi], xmm1
    vmovaps xmm3, xmm1
    vfmaddps        xmm2, xmm2, XMMWORD PTR -32[rsi], xmm0
    cmp     rdx, rsi
    jne     .L3

但是，它是为精确的 simd 指令集而设计的，对于需要更改代码的 avx2 或 avx512 来说不是最佳的。

How can you write C or C++ code for which gcc will produce optimal code when compiled for any of sse, avx2 or avx512? That is, do you always have to write separate functions by hand for each different width of SIMD register?

Are there any open source libraries that make this easier?

Answer 1

这是一个使用 Eigen library 的示例:

#include <Eigen/Core>
std::complex<float> f(const std::complex<float> *x, int n)
{
    return Eigen::VectorXcf::Map(x, n).prod();
}

如果您使用 clang 或 g++ 编译它并启用 sse 或 avx(和 -O2)，您应该会得到相当不错的机器代码。它也适用于其他一些架构，如 Altivec 或 NEON。如果您知道 x 的第一个条目已对齐，您可以使用 MapAligned而不是 Map .

如果您在编译时碰巧知道 vector 的大小，您会得到更好的代码:

template<int n>
std::complex<float> f(const std::complex<float> *x)
{
    return Eigen::Matrix<std::complex<float>, n, 1> >::MapAligned(x).prod();
}

注意:上面的函数直接对应函数f的OP。 但是，正如@PeterCordes 指出的那样，存储交错的复数通常是不好的，因为这将需要大量的洗牌来进行乘法运算。相反，应该以一种可以一次直接加载一个数据包的方式存储实部和虚部。

编辑/附录:要实现像复数乘法这样的数组结构，您实际上可以编写如下代码:

typedef Eigen::Array<float, 8, 1> v8sf; // Eigen::Array allows element-wise standard operations
typedef std::complex<v8sf> complex8;
complex8 prod(const complex8& a, const complex8& b)
{
    return a*b;
}

或更通用(使用 C++11):

template<int size, typename Scalar = float> using complexX = std::complex<Eigen::Array<Scalar, size, 1> >;

template<int size>
complexX<size> prod(const complexX<size>& a, const complexX<size>& b)
{
    return a*b;
}

当使用 -mavx -O2 编译时，这会编译成这样的东西(使用 g++-5.4):

    vmovaps 32(%rsi), %ymm1
    movq    %rdi, %rax
    vmovaps (%rsi), %ymm0
    vmovaps 32(%rdi), %ymm3
    vmovaps (%rdi), %ymm4
    vmulps  %ymm0, %ymm3, %ymm2
    vmulps  %ymm4, %ymm1, %ymm5
    vmulps  %ymm4, %ymm0, %ymm0
    vmulps  %ymm3, %ymm1, %ymm1
    vaddps  %ymm5, %ymm2, %ymm2
    vsubps  %ymm1, %ymm0, %ymm0
    vmovaps %ymm2, 32(%rdi)
    vmovaps %ymm0, (%rdi)
    vzeroupper
    ret

由于对我来说不是很明显的原因，这实际上隐藏在由实际方法调用的方法中，它只是在一些内存中移动——我不知道为什么 Eigen/gcc 不假设参数已经正确对齐。如果我用 clang 3.8.0 编译相同(和相同的参数)，它被编译为:

    vmovaps (%rsi), %ymm0
    vmovaps %ymm0, (%rdi)
    vmovaps 32(%rsi), %ymm0
    vmovaps %ymm0, 32(%rdi)
    vmovaps (%rdi), %ymm1
    vmovaps (%rdx), %ymm2
    vmovaps 32(%rdx), %ymm3
    vmulps  %ymm2, %ymm1, %ymm4
    vmulps  %ymm3, %ymm0, %ymm5
    vsubps  %ymm5, %ymm4, %ymm4
    vmulps  %ymm3, %ymm1, %ymm1
    vmulps  %ymm0, %ymm2, %ymm0
    vaddps  %ymm1, %ymm0, %ymm0
    vmovaps %ymm0, 32(%rdi)
    vmovaps %ymm4, (%rdi)
    movq    %rdi, %rax
    vzeroupper
    retq

同样，一开始的内存运动很奇怪，但至少那是矢量化的。但是，对于 gcc 和 clang，当在循环中调用时，它会被优化掉:

complex8 f8(complex8 x[], int n) {
    if(n==0)
        return complex8(v8sf::Ones(),v8sf::Zero()); // I guess you want p = 1 + 0*i at the beginning?

    complex8 p = x[0];
    for (int i = 1; i < n; i++) p = prod(p, x[i]);
    return p;
}

这里的不同之处在于，clang 会将外部循环展开为每个循环 2 次乘法。另一方面，gcc 在使用 -mfma 编译时将使用 fused-multiply-add 指令。 .

f8函数当然也可以推广到任意维度:

template<int size>
complexX<size> fX(complexX<size> x[], int n) {
    using S= typename complexX<size>::value_type;
    if(n==0)
        return complexX<size>(S::Ones(),S::Zero());

    complexX<size> p = x[0];
    for (int i = 1; i < n; i++) p *=x[i];
    return p;
}

为了减少 complexX<N>到单个 std::complex可以使用以下函数:

// only works for powers of two
template<int size> EIGEN_ALWAYS_INLINE
std::complex<float> redux(const complexX<size>& var) {
    complexX<size/2> a(var.real().template head<size/2>(), var.imag().template head<size/2>());
    complexX<size/2> b(var.real().template tail<size/2>(), var.imag().template tail<size/2>());
    return redux(a*b);
}
template<> EIGEN_ALWAYS_INLINE
std::complex<float> redux(const complexX<1>& var) {
    return std::complex<float>(var.real()[0], var.imag()[0]);
}

但是，根据我使用的是 clang 还是 g++，我得到的汇编程序输出完全不同。总体而言，g++ 倾向于无法内联加载输入参数，并且 clang 无法使用 FMA 操作(YMMV ...) 本质上，无论如何您都需要检查生成的汇编代码。更重要的是，您应该对代码进行基准测试(不确定该例程对您的整体问题有多大影响)。

另外，我想指出 Eigen 实际上是一个线性代数库。利用它来生成纯可移植 SIMD 代码并不是真正的设计目的。