使用 AVX 一次性完成 4 个水平双精度求和

Posted 2023-02-16

【中文标题】使用 AVX 一次性完成 4 个水平双精度求和

【英文标题】：4 horizontal double-precision sums in one go with AVX

【发布时间】：2012-05-31 12:03:13

【问题描述】：

问题可以描述如下。

输入

__m256d a, b, c, d

输出

__m256d s = a[0]+a[1]+a[2]+a[3], b[0]+b[1]+b[2]+b[3], 
             c[0]+c[1]+c[2]+c[3], d[0]+d[1]+d[2]+d[3]

到目前为止我所做的工作

这似乎很容易：两个 VHADD 之间有一些改组，但实际上结合 AVX 的所有排列不能生成实现该目标所需的排列。让我解释一下：

VHADD x, a, b => x = a[0]+a[1], b[0]+b[1], a[2]+a[3], b[2]+b[3]
VHADD y, c, d => y = c[0]+c[1], d[0]+d[1], c[2]+c[3], d[2]+d[3]

我是否能够以相同的方式排列 x 和 y 以获得

x1 = a[0]+a[1], a[2]+a[3], c[0]+c[1], c[2]+c[3]
y1 = b[0]+b[1], b[2]+b[3], d[0]+d[1], d[2]+d[3]

然后

VHADD s, x1, y1 => s1 = a[0]+a[1]+a[2]+a[3], b[0]+b[1]+b[2]+b[3], 
                         c[0]+c[1]+c[2]+c[3], d[0]+d[1]+d[2]+d[3]

这是我想要的结果。

所以我只需要找到如何执行

x,y => x[0], x[2], y[0], y[2], x[1], x[3], y[1], y[3]

不幸的是，我得出的结论是，使用 VSHUFPD、VBLENDPD、VPERMILPD、VPERM2F128、VUNPCKHPD、VUNPCKLPD 的任何组合都证明是不可能的。问题的关键在于，在__m256d的实例u中，u[1]和u[2]是不可能交换的。

问题

这真的是死胡同吗？还是我错过了排列指令？

【问题讨论】：

您可以通过在输入中交换 b 和 c 来简化任务。这样，您只需交换内部的两个双打。无论如何，这可以通过使用vperm2f128、两个备用寄存器、vpermilps 交换第一个/第二个和第三个/第四个以及最后一个vunpack[hl]pd 组合来完成粗心的源代码阅读器的组合来实现。

另一个想法（可能更快）是输入矩阵的转置（沿对角线镜像）并进行垂直求和。即转化为a0, b0, c0, d0等，简单加起来。

我也想过交换 b 和 c。由于这段代码可能会被我以外的其他人重用，我希望我能找到一个没有这种扭曲的解决方案！至于您的第二个建议，我无法重新组织代码以首先生成转置矩阵，并且我看不到执行转置的有效方法，所以我会说死路一条。

【参考方案1】：

VHADD 指令后面是常规的VADD。下面的代码应该给你你想要的：

// a[0]+a[1], b[0]+b[1], a[2]+a[3], b[2]+b[3]
__m256d sumab = _mm256_hadd_pd(a, b);
// c[0]+c[1], d[0]+d[1], c[2]+c[3], d[2]+d[3]
__m256d sumcd = _mm256_hadd_pd(c, d);

// a[0]+a[1], b[0]+b[1], c[2]+c[3], d[2]+d[3]
__m256d blend = _mm256_blend_pd(sumab, sumcd, 0b1100);
// a[2]+a[3], b[2]+b[3], c[0]+c[1], d[0]+d[1]
__m256d perm = _mm256_permute2f128_pd(sumab, sumcd, 0x21);

__m256d sum =  _mm256_add_pd(perm, blend);

这给出了 5 条指令的结果。我希望我得到了正确的常量。

您提出的排列当然可以完成，但需要多条指令。抱歉，我没有回答你的那部分问题。

编辑：我无法抗拒，这是完整的排列。 （再一次，尽我最大的努力使常量正确。）您可以看到交换u[1] 和u[2] 是可能的，只需要一些工作。在第一代中跨越 128 位的障碍是很困难的。 AVX。我还想说VADD 比VHADD 更可取，因为VADD 具有两倍的吞吐量，即使它执行相同数量的添加。

// x[0],x[1],x[2],x[3]
__m256d x;

// x[1],x[0],x[3],x[2]
__m256d xswap = _mm256_permute_pd(x, 0b0101);

// x[3],x[2],x[1],x[0]
__m256d xflip128 = _mm256_permute2f128_pd(xswap, xswap, 0x01);

// x[0],x[2],x[1],x[3] -- not imposssible to swap x[1] and x[2]
__m256d xblend = _mm256_blend_pd(x, xflip128, 0b0110);

// repeat the same for y
// y[0],y[2],y[1],y[3]
__m256d yblend;

// x[0],x[2],y[0],y[2]
__m256d x02y02 = _mm256_permute2f128_pd(xblend, yblend, 0x20);

// x[1],x[3],y[1],y[3]
__m256d x13y13 = _mm256_permute2f128_pd(xblend, yblend, 0x31);

【讨论】：

聪明！谢谢！不，我向您保证，我提出的排列无法完成，但您的方案使其成为多余。并且通过 5 条指令，您给出了最有效的解决方案。

@ljbou：在当前一代的 AVX 中，最多可以在 4 条指令中执行 __m256d 的任意置换：一条 VPERM2F128，两条 VSHUFPD，和一条VBLENDPD。传入的AVX2 in Haswell 功能更强大，允许在一条指令中进行任意排列（我认为是VPERMPD）。

这真是个天才，在水平添加上挣扎了几个小时。您是否知道任何以尽可能少的指令/周期实现这些常用操作的库？

【参考方案2】：

我不知道有任何指令可以让您进行这种排列。 AVX 指令的操作通常使得寄存器的高 128 位和低 128 位有些独立；将两半的值混合的能力并不多。我能想到的最佳实现将基于对this question 的回答：

__m128d horizontal_add_pd(__m256d x1, __m256d x2)

    // calculate 4 two-element horizontal sums:
    // lower 64 bits contain x1[0] + x1[1]
    // next 64 bits contain x2[0] + x1[1]
    // next 64 bits contain x1[2] + x1[3]
    // next 64 bits contain x2[2] + x2[3]
    __m256d sum = _mm256_hadd_pd(x1, x2);
    // extract upper 128 bits of result
    __m128d sum_high = _mm256_extractf128_pd(sum1, 1);
    // add upper 128 bits of sum to its lower 128 bits
    __m128d result = _mm_add_pd(sum_high, (__m128d) sum);
    // lower 64 bits of result contain the sum of x1[0], x1[1], x1[2], x1[3]
    // upper 64 bits of result contain the sum of x2[0], x2[1], x2[2], x2[3]
    return result;


__m256d a, b, c, d;
__m128d res1 = horizontal_add_pd(a, b);
__m128d res2 = horizontal_add_pd(c, d);
// At this point:
//     res1 contains a's horizontal sum in bits 0-63
//     res1 contains b's horizontal sum in bits 64-127
//     res2 contains c's horizontal sum in bits 0-63
//     res2 contains d's horizontal sum in bits 64-127
// cast res1 to a __m256d, then insert res2 into the upper 128 bits of the result
__m256d sum = _mm256_insertf128_pd(_mm256_castpd128_pd256(res1), res2, 1);
// At this point:
//     sum contains a's horizontal sum in bits 0-63
//     sum contains b's horizontal sum in bits 64-127
//     sum contains c's horizontal sum in bits 128-191
//     sum contains d's horizontal sum in bits 192-255

这应该是你想要的。以上应该在 7 条指令中是可行的（强制转换不应该真正做任何事情；它只是给编译器的一个注释，以改变它处理 res1 中的值的方式），假设短 horizontal_add_pd() 函数可以是由您的编译器内联，并且您有足够的可用寄存器。

【讨论】：

正如 drhirsch 指出的那样，您可以将其视为相当于 4 个垂直总和加上一些开销，您的解决方案相当于 3/4，而 Norbert 的解决方案下降到 1/4。

在所有支持它的 Intel 和 AMD CPU 上都已解码为 3 uop（2x shuffle 馈送垂直添加）。不值得在两个输入相同的情况下使用，除非您正在优化代码大小而不是速度。提取高半部分并添加到低半部分，从 256 下降到 128，然后从 128 下降到 64（例如 vunpcklpd）。

请参阅Get sum of values stored in __m256d with SSE/AVX 了解编译为高效 asm 的方法。 （但这些都没有回答这个问题，这不是关于一个寄存器的 hsum，而是关于 sum+transpose，是_mm256_hadd_pd 的少数用例之一。）