仅包含 3 个元素的两个向量的 AVX 优化加法

AVX-optimized addition of two vectors containing only 3 elements

我有一些这样的代码:

void add_v3_v3(float r[3], const float a[3])
{
  r[0] += a[0];
  r[1] += a[1];
  r[2] += a[2];
}  

我想将其转换为 AVX 代码,但据我所知,AVX 仅在向量具有四个元素(x、y、z、w)时才有效,而这里我只有三个元素(x、 y, z).有什么想法吗?

将它们矢量化到 AVX 中的最简单方法是简单地考虑数组,例如

// this does not vectorise
void add_v3_v3(float r[3], const float a[3])
{
  r[0] += a[0];
  r[1] += a[1];
  r[2] += a[2];
}  

// ... but this will
void add_many_v3_v3(float r[], const float a[], int count)
{
  for(int i = 0; i < count; ++i)
    add_v3_v3(r + i*3, a + i*3);
}  

https://godbolt.org/z/6vhnWrh3a

它生成一个内部循环,使用 'ps' 变体使用 256 位 YMM 寄存器一次处理 8 个浮点数:

.L7:
        vmovups ymm3, YMMWORD PTR [rax+32]
        vmovups ymm4, YMMWORD PTR [rax+64]
        vmovups ymm5, YMMWORD PTR [rax]
        vaddps  ymm1, ymm3, YMMWORD PTR [rdx+32]
        vaddps  ymm0, ymm4, YMMWORD PTR [rdx+64]
        vaddps  ymm2, ymm5, YMMWORD PTR [rdx]
        add     rax, 96
        vmovups YMMWORD PTR [rax-64], ymm1
        vmovups YMMWORD PTR [rax-96], ymm2
        vmovups YMMWORD PTR [rax-32], ymm0
        add     rdx, 96
        cmp     rax, r8
        jne     .L7

或使用 ZMM 寄存器一次 16 个,如果您启用了 AVX512f,例如

.L7:
        vmovups zmm3, ZMMWORD PTR [rax+64]
        vmovups zmm5, ZMMWORD PTR [rax]
        add     rax, 192
        add     rdx, 192
        vmovups zmm4, ZMMWORD PTR [rax-64]
        vaddps  zmm1, zmm3, ZMMWORD PTR [rdx-128]
        vaddps  zmm0, zmm4, ZMMWORD PTR [rdx-64]
        vaddps  zmm2, zmm5, ZMMWORD PTR [rdx-192]
        vmovups ZMMWORD PTR [rax-128], zmm1
        vmovups ZMMWORD PTR [rax-192], zmm2
        vmovups ZMMWORD PTR [rax-64], zmm0
        cmp     r8, rax
        jne     .L7

最糟糕的事情是:

typedef __m128 vec3;

这样做只会浪费第 4 个浮点数,这意味着您只会获得 3 倍的性能提升(而不是上面的 8 倍或 16 倍)。好吧,不完全正确...编译器 可能 能够将这些 vec3 添加中的 2 个融合到一个 __m256 op 中,使您增加 6 倍,但它不会'不如上面的好。

然而,这种方法通常失败的地方是点积和叉积。

float dot_v3_v3(float a[3], const float b[3])
{
  return a[0] * b[0] +
         a[1] * b[1] +
         a[2] * b[2];

}

一般 SIMD 指令喜欢在每个通道上执行相同的操作。点积和叉积涉及跨通道操作,因此要么最终忽略 SIMD,要么产生一些略微次优的 SIMD 使用,并进行大量数据混洗。

避免这种情况的唯一实用方法是将代码转换为数组格式 (SOA) 结构。例如

struct float16 {
   float f[16];
   float operator [](int i) const { return f[i]; }
   float& operator [](int i) { return f[i]; }
};

struct vec3x16 {
   float16 x;
   float16 y;
   float16 z;
};

float16 dot_v3_v3_x16(vec3x16 a, vec3x16 b) {
   float16 r;
   for(int i = 0; i < 16; ++i)
   {
     r[i] = a.x[i] * b.x[i] +
            a.y[i] * b.y[i] +
            a.z[i] * b.z[i];
   }
   return r;
}

产生:https://godbolt.org/z/qhE5xW91s

dot_v3_v3_x16(vec3x16, vec3x16):
        push    rbp
        mov     rax, rdi
        mov     rbp, rsp
        vmovups zmm1, ZMMWORD PTR [rbp+272]
        vmulps  zmm0, zmm1, ZMMWORD PTR [rbp+80]
        vmovups zmm2, ZMMWORD PTR [rbp+208]
        vfmadd231ps     zmm0, zmm2, ZMMWORD PTR [rbp+16]
        vmovups zmm3, ZMMWORD PTR [rbp+336]
        vfmadd231ps     zmm0, zmm3, ZMMWORD PTR [rbp+144]
        vmovups ZMMWORD PTR [rdi], zmm0
        vzeroupper
        pop     rbp
        ret

内联将删除其中的大部分内容,为您提供大约 3 条指令的 16 个点积(+ 一些 movs 用于不在寄存器中的数据)。

然而,在结构上,在实践中到处使用 SOA 代码 比仅仅坚持标准的 vec3 类型更难。

所以要么完全忽略 AVX,而只是尝试从数组的角度思考(这会给你相当快的结果,代码库更容易维护和扩展)

或者,完全以数组格式的结构工作,这将为您提供非常快的代码,但预计新功能的开发时间会显着增加。 (另请注意,某些 cmath 函数可能无法向量化为 AVX2 或 AVX512)