performance - 快速矢量化 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 .理论上，相同的方法可用于 double 值，产生半精度或单精度或 double 。

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

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

在代码示例中使用内在函数。汇编依赖于编译器，所以不太有用。

对函数使用类似的命名约定。

实现将包含密集浮点/ double 值的单个 SSE/AVX 寄存器作为输入的例程。如果有相当大的性能提升，您还可以发布将多个寄存器作为输入的例程(两个寄存器可能是可行的)。

如果 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 迭代。

float 的recip需要 1、4 周期与 7 循环。

rsqrt on float 需要 1, 6 周期与 14 循环。

双倍奖励 3、6、9 周期与 14 循环。

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 可以为零，则进行修正，例如通过使用 CMPPS 对 0.0 的结果屏蔽( _mm_andn_ps ))。最佳方式是替换 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 免费将 sqrt 结果添加到某些内容中，而不是在最后一步添加 mul)。

这个循环的 SSE 版本，不考虑任何移动指令，是 6 uop。这意味着它在 Haswell 上的吞吐量应该是 每 3 个周期一个 (假设两个执行端口可以处理 FP mul，而 rsqrt 位于 FP add/sub 的相反端口上)。在 SnB/IvB(可能还有 Nehalem)上，它的吞吐量应该是 每 5 个周期一个 ，因为 mulps 和 rsqrtps 竞争端口 0。 subps 在端口 1 上，这不是瓶颈。
对于 Haswell，我们可以使用 FMA 将减法与乘法结合起来。但是，分隔符/sqrt 单元不是 256b 宽，因此与其他所有内容不同，ymm regs 上的 divps/sqrtps/rsqrtps/rcpps 需要额外的 uops 和额外的周期。

// 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 uop。 ( VRSQRTPS 为 3 uop，p0 为 2，p1/5 为 1。)256b mulps 和 fma 都是单 uop 指令，都可以在端口 0 或端口 1 上运行。(p0 仅在 pre-Haswell 上) .所以吞吐量应该是: 每 3c 一个 , 如果 OOO 引擎将 uops 分派(dispatch)到最佳执行端口。 (即来自 rsqrt 的 shuffle uop 总是转到 p5，永远不会到达 p1，在那里它会占用 mul/fma 带宽。)至于与其他计算重叠(不仅仅是自身的独立执行)，它再次非常轻量级。只有 7 个 uops 和 23 个周期的 dep 链为其他事情留下了很大的空间，而这些 uops 位于重新排序缓冲区中。
如果这是一个巨大的 dep 链中的一个步骤，没有其他任何事情(甚至不是独立的下一次迭代)，那么 256b vsqrtps是 19 个周期的延迟，每 14 个周期的吞吐量为 1。 (哈斯韦尔)。如果你仍然真的需要倒数，那么 256b vdivps还具有 18-21c 延迟，每 14c 吞吐量有一个延迟。所以对于普通的 sqrt，指令的延迟较低。对于 recip sqrt，近似的迭代是更少的延迟。 (如果它与自身重叠，则吞吐量更大。如果与其他不是除法单元的东西重叠， sqrtps 不是问题。) sqrtps与 rsqrt 相比可能是吞吐量的胜利+ Newton 迭代，如果它是循环体的一部分，并且有足够多的其他非除法和非 sqrt 工作在进行，则除法单元未饱和。
如果您需要 sqrt(x) 尤其如此，不是 1/sqrt(x) .例如在带有 AVX2 的 Haswell 上，在适合 L3 缓存的浮点数组上执行 copy+arcsinh 循环，实现为 fastlog(v + sqrt(v*v + 1)) ，与真正的 VSQRTPS 或 VRSQRTPS + Newton-Raphson 迭代以大致相同的吞吐量运行。 (即使使用 very fast approximation for log() ，所以总循环体大约有 9 个 FMA/add/mul/convert 操作和 2 个 bool 值，加上 VSQRTPS ymm。仅使用 fastlog(v2_plus_1 * rsqrt(v2_plus_1) + v2_plus_1) 有一个加速，所以它不是瓶颈内存带宽，但它可能是延迟的瓶颈(因此乱序执行不能利​​用独立迭代的所有并行性)。
其他精度
对于 半精度 ，没有对半浮点数进行计算的指令。您打算在加载/存储时使用转换说明即时转换。
对于 double ，没有 rsqrtpd .据推测，这种想法是，如果您不需要完全精度，则应该首先使用 float。所以你可以转换为 float 和返回，然后执行完全相同的算法，但使用 pd而不是 ps内在因素。或者，您可以将数据保持为 float 状态一段时间。例如将两个 ymm double 寄存器转换为一个 ymm 单打寄存器。
不幸的是，没有一条指令需要两个 ymm 的 double 寄存器并输出 ymm 的单精度。你必须去 ymm->xmm 两次，然后 _mm256_insertf128_ps一个 xmm 到另一个的高 128。但是你可以将那个 256b ymm 向量提供给相同的算法。
如果您要立即转换回 double ，则分别对单曲的两个 128b 寄存器进行 rsqrt + Newton-Raphson 迭代可能是有意义的。插入/提取的额外 2 uops 和 256b rsqrt 的额外 2 uops 开始加起来，更不用说 vinsertf128 的 3 周期延迟了。/ vextractf128 .计算将重叠，并且两个结果将在几个周期之外准备好。 (或者间隔 1 个周期，如果您有一个特殊版本的代码用于一次对 2 个输入进行交错操作)。
请记住，单精度的指数范围比 double 小，因此转换可能会溢出到 +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 那样为您提供正确舍入的单精度浮点数。