在 AVX 中进行线性平均的最短方法

Posted 2023-02-16

【中文标题】在 AVX 中进行线性平均的最短方法

【英文标题】：Shortest way to do Linear Averaging in AVX

【发布时间】：2021-02-10 14:14:14

【问题描述】：

我有一个向量 Inp = A0, 0, A2, 0, A4, 0, A6, 0, ...;

我的兴趣是 Out = A0, mean(A0, A2), A2, mean(A2, A4), A4, mean(A4, A6), A6, ...;

Inp 和 Out 是浮点变量。 输入长度总是偶数。因此，Out(end) = inp(end -1);

编辑：

我的实现，

float *vecAPtr;
__m256 vecA;
__m256 vecB;
__m256 res1;
__m256 res2;
__m256 *AVXInp = (__m256*)Inp;
for(i = 0;i<inpLength;i = i+8)

  if( (inpLength - i) <= 8)
      //Normal C code
      Out[0] = Inp[0];
      Out[1] = (Inp[0]+Inp[2])/2;
      soon... (A simple for loop will take care of this part, not shown here)
   
else

vecA = _mm256_permutevar8x32_ps(*AVXInp, _mm256_set_epi32(1, 1, 6, 1, 4, 1, 2, 1) );
//vecA = [0 A2 0 A4 0 A6 0 0];

vecB = _mm256_permute_ps( *AVXInp, 0b10010001);
//vecB = [0 A0 0 A2 0 A4 0 A6];

vecAPtr = (float *)&vecA;
vecAPtr[7] = *( (float *)(AVXInp+1));
//vecA = [0 A2 0 A4 0 A6 0 A8];

res1 = _mm256_add_ps(vecA, vecB);
res2 = _mm256_mul_ps(res1, _mm256_set1_ps(0.5));
*AVXInp = _mm256_add_ps( *AVXInp, res2);

还有比这个更好的版本吗？

【问题讨论】：

A0、A2、...的类型是什么？

@PaulR 它们是浮点变量..

输入中的那些0元素是对空间和带宽的浪费；我想你有他们的理由？仅作记录，vecAPtr[7] = *( (float *)(AVXInp+1)); 可能很糟糕。最好只做另一个 32 字节的未对齐加载。虽然如果你幸运的话，编译器会发现它可以用vbroadcastss + vblendps 来完成。如果你不走运，它会编译成更糟糕的洗牌。

最好进行加载以获取所需的所有元素并跳过零，因为引入零可能会便宜得多！

【参考方案1】：

我将首先使用_mm256_moveldup_ps 复制偶数索引元素，然后加载从Inp+1 开始的向量并将其添加到其中。然后将奇数元素按0.5 缩放（偶数元素按1.0）：

void foo(float* out, float const* inp, size_t size)
    __m256 const scale = _mm256_setr_ps(1.f,.5f,1.f,.5f,1.f,.5f,1.f,.5f);

    for(size_t i=0; i<size-9; i+=8)
        __m256 even = _mm256_moveldup_ps(_mm256_loadu_ps(inp+i));
                                                // [a0, a0,      a2, a2,    ...]
        __m256 odd = _mm256_loadu_ps(inp+i+1);  // [ 0,    a2,    0,    a4, ...]
        __m256 sum = _mm256_add_ps(even,odd);   // [a0, a0+a2,   a2, a2+a4, ...]
        __m256 res = _mm256_mul_ps(scale, sum); // [a0,(a0+a2)/2,a2,(a2+a4)/2, ...]
        _mm256_storeu_ps(out+i, res);
    

这不处理最后的元素。此外，它假定您可以返回 a0+0，其中预期 a0（只有在 a0=-0 时才会有所不同）并且 a0+a2 不会溢出。

假设编译器融合了加载和_mm256_moveldup_ps，生成的vmovsldup甚至不需要shuffle操作，所以每次迭代只需要两次加载、两次算术和一次存储（因此应该有1个周期的吞吐量，忽略循环开销）。