快速反范数函数

Question

我有一个程序几乎把他所有的时间都花在计算循环上，比如

for(int j = 0; j < BIGNUMBER; j++)
    for(int i = 0; i < SMALLNUMBER; i++)
        result += var[i] / sqrt((A[i].x-B[j].x)*(A[i].x-B[j].x)+(A[i].y-B[j].y)*(A[i].y-B[j].y)+(A[i].z-B[j].z)*(A[i].z-B[j].z));

其中1.0/sqrt(...) 计算两个向量A[i] = {A[i].x, A[i].y, A[i].z} 和B[j] = {B[j].x, B[j].y, B[j].z} 之差的范数的倒数，这也是循环中开销最大的部分。

有没有办法优化循环，即使有一些精度损失？

---

更新：

这里是非向量化循环步骤的汇编代码，每条指令的延迟都比较差。您清楚地看到平方根的倒数是瓶颈：

movsd   A(%rip), %xmm0      1
movsd   A+8(%rip), %xmm2    1
subsd   B(%rip), %xmm0      3
subsd   B+8(%rip), %xmm2    3
movsd   A+16(%rip), %xmm1   1
mulsd   %xmm0, %xmm0        5
subsd   B+16(%rip), %xmm1   3
mulsd   %xmm2, %xmm2        5
mulsd   %xmm1, %xmm1        5
addsd   %xmm2, %xmm0        3
addsd   %xmm1, %xmm0        3
movsd   .LC0(%rip), %xmm1   1
unpcklpd    %xmm0, %xmm0    1
cvtpd2ps    %xmm0, %xmm0    4
unpcklps    %xmm0, %xmm0    3
cvtps2pd    %xmm0, %xmm0    2
sqrtsd  %xmm0, %xmm0        58
divsd   %xmm0, %xmm1        32
mulsd   var(%rip), %xmm1    5
addsd   result(%rip), %xmm1 3
cvttsd2si   %xmm1, %eax     3
movsd   %xmm1, result(%rip) 1

（顺便说一句，我不明白它为什么会这样做unpcklpd cvtpd2ps unpcklps cvtps2pd。）

Answer 1

如果您可以以 AoSoA 形式 (xxyyzzxxyyzzxxyyzz...) 排列向量，则可以使用 SSE 或 AVX (xxxxyyyyzzzz...) 非常有效地做到这一点。在下面的代码中，我假设 SSE2 具有 vec_size=2 但很容易将其更改为 AVX。但是您的代码可能是内存限制而不是计算限制，因此这仅对适合 L1 缓存的小循环有用。使用单浮点数也会更快，因为它的翻牌次数是翻倍的两倍，而 sqrt 是少数几个实际上比浮点数慢的函数之一。

resultv = _mm_setzero_pd(0);
for(int j = 0; j < BIGNUMBER; j+=vec_size) {
    bx = _mm_load_pd(&B[3*j+0*vec_size]);
    by = _mm_load_pd(&B[3*j+1*vec_size]);
    bz = _mm_load_pd(&B[3*j+2*vec_size]);
    for(int i = 0; i < SMALLNUMBER; i+=vec_size) {
        ax = _mm_load_pd(&A[3*i+0*vec_size]);
        ay = _mm_load_pd(&A[3*i+1*vec_size]);
        az = _mm_load_pd(&A[3*i+2*vec_size]);
        dx = _mm_sub_pd(ax,bx);
        dy = _mm_sub_pd(ay,by);
        dz = _mm_sub_pd(az,bz);
        mag2 = _mm_add_pd(_mm_add_pd(_mm_mul_pd(dx,dx),_mm_mul_pd(dy,dy)), _mm_mul_pd(dz,dz));
        varv = _mm_load_pd(&var[i]);        
        resultv = _mm_add_pd(_mm_div_pd(varv, _mm_sqrt_pd(mag2)), resultv);
        //resultv = _mm_add_pd(_mm_mul_pd(varv, _mm_rsqrt_pd(mag2)), resultv);
    }
}
result = _mm_cvtsd_f64(_mm_hadd_pd(resultv,resultv));

Answer 2

很容易假设sqrt 是“瓶颈”，但如果您这样做this，您可以确定。 应该是这样的

  (A[i].x-B[j].x)*(A[i].x-B[j].x)
+ (A[i].y-B[j].y)*(A[i].y-B[j].y)
+ (A[i].z-B[j].z)*(A[i].z-B[j].z)

反而是瓶颈。

如果您依赖编译器来优化所有这些索引操作，它可能会，但我想确定。作为第一个剪辑，我会在内部循环中写这个：

// where Foo is the class of A[i] and B[j]
Foo& a = A[i];
Foo& b = B[j];
double dx = a.x - b.x, dy = a.y - b.y, dz = a.z - b.z;
result += var[i] / sqrt( dx*dx + dy*dy + dz*dz );

这对编译器来说更容易优化。 然后，如果索引仍然是一个瓶颈，我会将Foo& b = B[j] 提升出内部循环，并单步执行指针而不是编写A[i]。 如果样本显示内部 for 语句本身（测试终止和递增）需要相当长的时间，我会展开一些。

Answer 3

您在评论中说the bottleneck is still the computation of the inverse square root. 幸运的是，这在图形中出现了很多，并且有一个非常奇特的算法可以做到这一点。有一篇关于它的 wikipedia 文章，quake 中的一个实现和一个 SO 问题 3。

Answer 4

这是 SIMD (SSE) 指令的良好工作场所。如果您的编译器支持自动矢量化，请启用此选项（调整数据结构布局以获得实际增益）。如果它支持内在函数，您可以使用它们。

编辑
上面的汇编代码没有利用矢量化。它使用标量 SSE 指令。 您可能会尝试帮助编译器 - 制作（X，Y，Z，Dummy）浮点结构，而不是双精度。 SSE 向量指令可以同时处理 4 个浮点数（或 2 个双精度数）。
（我认为一些数学库已经包含向量范数的 SIMD 函数）

附：您可以在问题中添加 SSE 标签

Answer 5

通过通过(1.0f/sqrt(float(...))) 强制进行单精度计算并使用#pragma GCC optimize ("Ofast") 作为函数，我能够得到rsqrtss 指令，整个函数的加速速度很好（大约快2 倍）。它实际上破坏了自动矢量化（可能是因为单精度和双精度的混合）。

汇编代码：

movsd   56(%rax), %xmm0     
addq    $120, %rax
movsd   -72(%rax), %xmm2    
subsd   %xmm5, %xmm0
movsd   -56(%rax), %xmm1    
subsd   %xmm4, %xmm2
subsd   %xmm6, %xmm1
mulsd   %xmm0, %xmm0
mulsd   %xmm2, %xmm2
mulsd   %xmm1, %xmm1
addsd   %xmm2, %xmm0
addsd   %xmm1, %xmm0
unpcklpd    %xmm0, %xmm0
cvtpd2ps    %xmm0, %xmm0
rsqrtss %xmm0, %xmm1
mulss   %xmm1, %xmm0
mulss   %xmm1, %xmm0
mulss   %xmm7, %xmm1
addss   %xmm8, %xmm0
mulss   %xmm1, %xmm0
mulss   -40(%rax), %xmm0
cmpq    %rdx, %rax
unpcklps    %xmm0, %xmm0
cvtps2pd    %xmm0, %xmm0
addsd   %xmm0, %xmm3

但我不明白最后的附加乘法。