optimization - 计算大点积的最快方法是什么？

Question

考虑这个片段：

double dot(double* a, double* b, int n) {
  double sum = 0;
  for (int i = 0; i < n; ++i) sum += a[i] * b[i];
  return sum;
}

如何使用内在函数或汇编程序加速它？

笔记：

• 您可以假设最新的架构，包括 AVX 扩展。
• n是几百。
• dot 本身将在紧密循环中使用

Answer 1

这是一个简单的 SSE 实现：

#include "pmmintrin.h"

__m128d vsum = _mm_set1_pd(0.0);
double sum = 0.0;
int k;

// process 2 elements per iteration
for (k = 0; k < n - 1; k += 2)
{
    __m128d va = _mm_loadu_pd(&a[k]);
    __m128d vb = _mm_loadu_pd(&b[k]);
    __m128d vs = _mm_mul_pd(va, vb);
    vsum = _mm_add_pd(vsum, vs);
}

// horizontal sum of 2 partial dot products
vsum = _mm_hadd_pd(vsum, vsum);
_mm_store_sd(&sum, vsum);

// clean up any remaining elements
for ( ; k < n; ++k)
{
    sum += a[k] * b[k];
}

请注意，如果您可以保证 a 和 b 是 16 字节对齐的，那么您可以使用_mm_load_pd而不是_mm_loadu_pd可能有助于提高性能，尤其是在较旧的（Nehalem 之前）CPU 上。

另请注意，对于这样的循环，相对于加载数量而言，算术指令非常少，那么性能可能会受到内存带宽的限制，并且在实践中可能无法实现矢量化的预期加速。

---

如果您想以 AVX 为目标 CPU，那么这是从上述 SSE 实现到每次迭代处理 4 个元素而不是 2 个的相当简单的转换：

#include "immintrin.h"

__m256d vsum = _mm256_set1_pd(0.0);
double sum = 0.0;
int k;

// process 4 elements per iteration
for (k = 0; k < n - 3; k += 4)
{
    __m256d va = _mm256_loadu_pd(&a[k]);
    __m256d vb = _mm256_loadu_pd(&b[k]);
    __m256d vs = _mm256_mul_pd(va, vb);
    vsum = _mm256_add_pd(vsum, vs);
}

// horizontal sum of 4 partial dot products
vsum = _mm256_hadd_pd(_mm256_permute2f128_pd(vsum, vsum, 0x20), _mm256_permute2f128_pd(vsum, vsum, 0x31));
vsum = _mm256_hadd_pd(_mm256_permute2f128_pd(vsum, vsum, 0x20), _mm256_permute2f128_pd(vsum, vsum, 0x31));
_mm256_store_sd(&sum, vsum);

// clean up any remaining elements
for ( ; k < n; ++k)
{
    sum += a[k] * b[k];
}