math - 快速反范数函数

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 缓存的小循环有用。使用单个 float 也会更快，因为它的触发器数量是两倍，而 sqrt 是少数几个实际上比 float 慢的函数之一。

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));