快速矢量化 rsqrt 和 SSE/AVX 倒数取决于精度

Question

假设有必要计算压缩浮点数据的倒数或倒数平方根。两者都可以通过以下方式轻松完成：

__m128 recip_float4_ieee(__m128 x) { return _mm_div_ps(_mm_set1_ps(1.0f), x); }
__m128 rsqrt_float4_ieee(__m128 x) { return _mm_div_ps(_mm_set1_ps(1.0f), _mm_sqrt_ps(x)); }

这工作得很好但很慢：根据the guide，他们在 Sandy Bridge 上需要 14 和 28 个周期（吞吐量）。相应的 AVX 版本在 Haswell 上花费的时间几乎相同。

另一方面，可以使用以下版本：

__m128 recip_float4_half(__m128 x) { return _mm_rcp_ps(x); }
__m128 rsqrt_float4_half(__m128 x) { return _mm_rsqrt_ps(x); }

它们只需要一两个时间周期（吞吐量），从而显着提升性能。但是，它们非常近似：它们产生的结果相对误差小于 1.5 * 2^-12。假设单精度浮点数的机器 epsilon 是 2^-24，我们可以说这个近似值大约有 一半 的精度。

似乎可以添加 Newton-Raphson 迭代以产生具有单一精度的结果（可能不如 IEEE 标准要求的精确，通过），请参阅GCC、ICC，讨论在LLVM。理论上，同样的方法可以用于双精度值，产生 half 或 single 或 double 精度。

我对这种方法在浮点和双精度数据类型以及所有（半精度、单精度、双精度）精度上的工作实现很感兴趣。不需要处理特殊情况（除以零、sqrt(-1)、inf/nan 等）。此外，我不清楚这些例程中的哪一个会比简单的 IEEE 兼容解决方案更快，哪个会更慢。

以下是对答案的一些小限制，请：

1. 在代码示例中使用内部函数。汇编依赖于编译器，所以用处不大。
2. 对函数使用类似的命名约定。
3. 实施例程，将包含密集封装的浮点/双精度值的单个 SSE/AVX 寄存器作为输入。如果有相当大的性能提升，您还可以发布将多个寄存器作为输入的例程（两个 reg 可能是可行的）。
4. 如果两个 SSE/AVX 版本完全等同于将 _mm 更改为 _mm256 （反之亦然），则不要发布它们。

欢迎任何性能评估、测量和讨论。

总结

以下是单精度浮点数的版本，带有一次 NR 迭代：

__m128 recip_float4_single(__m128 x) {
  __m128 res = _mm_rcp_ps(x);
  __m128 muls = _mm_mul_ps(x, _mm_mul_ps(res, res));
  return res =  _mm_sub_ps(_mm_add_ps(res, res), muls);
}
__m128 rsqrt_float4_single(__m128 x) {
  __m128 three = _mm_set1_ps(3.0f), half = _mm_set1_ps(0.5f);
  __m128 res = _mm_rsqrt_ps(x); 
  __m128 muls = _mm_mul_ps(_mm_mul_ps(x, res), res); 
  return res = _mm_mul_ps(_mm_mul_ps(half, res), _mm_sub_ps(three, muls)); 
}

Answer given by Peter Cordes 解释了如何创建其他版本，并包含全面的理论性能分析。

您可以在此处找到所有已实施的带有基准的解决方案： recip_rsqrt_benchmark。

在 Ivy Bridge 上获得的吞吐量结果如下所示。仅对单寄存器 SSE 实现进行了基准测试。花费的时间以每次调用的周期数给出。第一个数字是半精度（无 NR），第二个是单精度（1 次 NR 迭代），第三个是 2 次 NR 迭代。

1. recip float 需要 1、4 个周期，而 7 个周期。
2. rsqrt float 需要 1、6 个周期，而 14 个周期。
3. recip double 需要 3、6、9 个周期，而 14 个周期。
4. double 上的 rsqrt 需要 3、8、13 个周期，而 28 个周期。

警告：我必须创造性地对原始结果进行四舍五入...

Answer 1

在实践中有很多算法的例子。例如：

Newton Raphson with SSE2 - can someone explain me these 3 lines 有一个解释one of Intel's examples 使用的迭代的答案。

对于假设 Haswell 的性能分析（与以前的设计不同，它可以在两个执行端口上 FP mul），我将在此处重现代码（每行一个操作）。请参阅 http://agner.org/optimize/ 获取指令吞吐量和延迟表，以及有关如何了解更多背景的文档。

// execute (aka dispatch) on cycle 1, results ready on cycle 6
nr = _mm_rsqrt_ps( x );

// both of these execute on cycle 6, results ready on cycle 11
xnr = _mm_mul_ps( x, nr );         // dep on nr
half_nr = _mm_mul_ps( half, nr );  // dep on nr

// can execute on cycle 11, result ready on cycle 16
muls = _mm_mul_ps( xnr , nr );     // dep on xnr

// can execute on cycle 16, result ready on cycle 19
three_minus_muls = _mm_sub_ps( three, muls );  // dep on muls

// can execute on cycle 19, result ready on cycle 24
result = _mm_mul_ps( half_nr, three_minus_muls ); // dep on three_minus_muls
// result is an approximation of 1/sqrt(x), with ~22 to 23 bits of precision in the mantissa.

如果它不是依赖链的一部分，那么这里有很多空间可以让其他计算重叠。但是，如果您的代码每次迭代的数据都依赖于前一次的数据，那么使用sqrtps 的 11 周期延迟可能会更好。或者即使每个循环迭代足够长，乱序执行也无法通过重叠独立迭代来隐藏它。

要获得sqrt(x) 而不是1/sqrt(x)，请乘以x（如果x 可以为零，则进行修正，例如通过使用a 的结果屏蔽(_mm_andn_ps) CMPPS 对 0.0)。最佳方法是将half_nr 替换为half_xnr = _mm_mul_ps( half, xnr );。这不会延长 dep 链，因为half_xnr 可以从第 11 周期开始，但直到结束（第 19 周期）才需要。与可用的 FMA 相同：总延迟不会增加。

如果 11 位精度就足够了（没有牛顿迭代），Intel's optimization manual (section 11.12.3) 建议使用 rcpps(rsqrt(x))，这比乘以原始 x 更差，至少使用 AVX。它可能会使用 128 位 SSE 保存 movdqa 指令，但 256b rcpps 比 256b mul 或 fma 慢。 （并且它允许您在最后一步使用 FMA 而不是 mul 免费将 sqrt 结果添加到某些东西中）。

---

此循环的 SSE 版本（不考虑任何移动指令）为 6 uop。这意味着它在 Haswell 上的吞吐量应该是每 3 个周期一个（假设两个执行端口可以处理 FP mul，并且 rsqrt 位于 FP add/sub 的相反端口上）。在 SnB/IvB（可能还有 Nehalem）上，它应该有 每 5 个周期一个的吞吐量，因为 mulps 和 rsqrtps 竞争端口 0。subps 在端口 1 上，这不是瓶颈。

对于 Haswell，我们可以使用 FMA 将减法与 mul 结合起来。但是，分隔符/sqrt 单元的宽度不是 256b，因此与其他所有内容不同的是，ymm regs 上的 divps/sqrtps/rsqrtps/rcpps 需要额外的微指令和额外的周期。

// vrsqrtps ymm has higher latency
// execute on cycle 1, results ready on cycle 8
nr = _mm256_rsqrt_ps( x );

// both of can execute on cycle 8, results ready on cycle 13
xnr = _mm256_mul_ps( x, nr );         // dep on nr
half_nr = _mm256_mul_ps( half, nr );  // dep on nr

// can execute on cycle 13, result ready on cycle 18
three_minus_muls = _mm256_fnmadd_ps( xnr, nr, three );  // -(xnr*nr) + 3

// can execute on cycle 18, result ready on cycle 23
result = _mm256_mul_ps( half_nr, three_minus_muls ); // dep on three_minus_muls

我们使用 FMA 节省了 3 个周期。这可以通过使用 2-cyle-slower 256b rsqrt 来抵消，净增益减少了 1 个周期的延迟（对于两倍的宽度来说非常好）。 SnB/IvB AVX 是 128b 的 24c 延迟，256b 的 26c 延迟。

256b FMA 版本总共使用 7 个微指令。 （VRSQRTPS 是 3 uop，2 用于 p0，1 用于 p1/5。）256b mulps 和 fma 都是单 uop 指令，都可以在端口 0 或端口 1 上运行。（p0 仅在 pre-Haswell 上） .所以吞吐量应该是：每 3c 一个，如果 OOO 引擎将 uops 调度到最佳执行端口。 （即，来自 rsqrt 的 shuffle uop 总是转到 p5，而不是 p1，它会占用 mul/fma 带宽。）至于与其他计算重叠（不仅仅是自身的独立执行），它又是非常轻量级的。只有 7 个微指令和 23 个循环的 dep 链为其他事情留下了很大的空间，而这些微指令位于重新排序缓冲区中。

如果这是一个巨大的 dep 链中的一个步骤，没有其他任何事情发生（甚至没有独立的下一次迭代），那么 256b vsqrtps 是 19 个周期延迟，吞吐量为每 14 个周期一个。 （哈斯韦尔）。如果您仍然确实需要倒数，那么 256b vdivps 也有 18-21c 延迟，每 14c 吞吐量一个。所以对于正常的 sqrt，指令具有较低的延迟。对于 recip sqrt，近似值的迭代会减少延迟。 （如果它与自身重叠，则吞吐量会更高。如果与没有除法单元的其他东西重叠，sqrtps 不是问题。）

sqrtps 与rsqrt + 牛顿迭代相比，如果它是循环体的一部分，并且有足够的其他非除法和非平方工作正在进行，除法单元未饱和，则它可能是吞吐量的胜利。

如果您需要sqrt(x)，而不是1/sqrt(x)，则尤其如此。。例如在带有 AVX2 的 Haswell 上，在适合 L3 缓存的浮点数组上进行复制 + arcsinh 循环，实现为 fastlog(v + sqrt(v*v + 1))，以与实际 VSQRTPS 或 VRSQRTPS + Newton-Raphson 迭代大致相同的吞吐量运行。 （即使使用very fast approximation for log()，所以总循环体大约是 9 个 FMA/add/mul/convert 操作和 2 个布尔值，加上 VSQRTPS ymm。仅使用 fastlog(v2_plus_1 * rsqrt(v2_plus_1) + v2_plus_1) 可以加快速度，因此它不会成为瓶颈内存带宽，但它可能会成为延迟的瓶颈（因此乱序执行无法利用独立迭代的所有并行性）。

其他精度

对于半精度，没有指令可以对半浮点数进行计算。您需要在加载/存储时使用转换说明进行即时转换。

对于双精度，没有rsqrtpd。据推测，如果您不需要完全精度，则应该首先使用浮点数。因此，您可以转换为浮点数并返回，然后执行完全相同的算法，但使用 pd 而不是 ps 内在函数。或者，您可以将数据保持浮动一段时间。例如将两个双精度的 ymm 寄存器转换成一个单精度的 ymm 寄存器。

不幸的是，没有一条指令需要两个双精度的 ymm 寄存器并输出一个单精度的 ymm。你必须去 ymm->xmm 两次，然后 _mm256_insertf128_ps 一个 xmm 到另一个的高 128。但是您可以将 256b ymm 向量输入到相同的算法中。

如果您要在之后立即转换回双精度，则分别对单精度的两个 128b 寄存器进行 rsqrt + Newton-Raphson 迭代可能是有意义的。插入/提取的额外 2 微指令和 256b rsqrt 的额外 2 微指令开始加起来，更不用说 vinsertf128 / vextractf128 的 3 周期延迟。计算将重叠，两个结果将相隔几个周期准备好。 （或者相隔 1 个周期，如果您有一个特殊版本的代码可以同时对 2 个输入进行交错操作）。

请记住，单精度的指数范围比双精度要小，因此转换可能会溢出到 +Inf 或下溢到 0.0。如果您不确定，请务必使用普通的_mm_sqrt_pd。

---

对于 AVX512F，有_mm512_rsqrt14_pd( __m512d a)。有了AVX512ER (KNL but not SKX or Cannonlake)，就有了_mm512_rsqrt28_pd。 _ps 版本当然也存在。在更多情况下，14 位尾数精度可能足以在没有牛顿迭代的情况下使用。

牛顿迭代仍然不会像常规 sqrt 那样为您提供正确舍入的单精度浮点数。