与 SSE 的并行前缀（累积）和

Question

我正在寻找一些关于如何使用 SSE 进行并行前缀求和的建议。我有兴趣在整数、浮点数或双精度数组上执行此操作。

我想出了两个解决方案。一个特例和一个一般情况。在这两种情况下，解决方案都在与 OpenMP 并行的两遍中在阵列上运行。对于特殊情况，我在两次通行证上都使用 SSE。对于一般情况，我只在第二遍使用它。

我的主要问题是在一般情况下如何在第一次通过时使用 SSE？ 以下链接 simd-prefix-sum-on-intel-cpu 显示了字节的改进，但 32 位数据类型没有。

特殊情况被称为特殊的原因是它要求数组采用特殊格式。例如，假设数组aof 浮点数只有 16 个元素。那么如果数组是这样重新排列的（结构数组到数组结构）：

a[0] a[1] ...a[15] -> a[0] a[4] a[8] a[12] a[1] a[5] a[9] a[13]...a[3] a[7] a[11] a[15]

SSE 垂直总和可用于两个通道。但是，这只有在数组已经采用特殊格式并且输出可以以特殊格式使用时才有效。否则必须对输入和输出进行昂贵的重新排列，这将使其比一般情况慢得多。

也许我应该为前缀和考虑一种不同的算法（例如二叉树）？

一般情况的代码：

void prefix_sum_omp_sse(double a[], double s[], int n) {
    double *suma;
    #pragma omp parallel
    {
        const int ithread = omp_get_thread_num();
        const int nthreads = omp_get_num_threads();
        #pragma omp single
        {
            suma = new double[nthreads + 1];
            suma[0] = 0;
        }
        double sum = 0;
        #pragma omp for schedule(static) nowait //first parallel pass
        for (int i = 0; i<n; i++) {
            sum += a[i];
            s[i] = sum;
        }
        suma[ithread + 1] = sum;
        #pragma omp barrier
        #pragma omp single
        {
            double tmp = 0;
            for (int i = 0; i<(nthreads + 1); i++) {
                tmp += suma[i];
                suma[i] = tmp;
            }
        }
        __m128d offset = _mm_set1_pd(suma[ithread]);
        #pragma omp for schedule(static) //second parallel pass with SSE as well
        for (int i = 0; i<n/4; i++) {       
            __m128d tmp1 = _mm_load_pd(&s[4*i]);
            tmp1 = _mm_add_pd(tmp1, offset);    
            __m128d tmp2 = _mm_load_pd(&s[4*i+2]);
            tmp2 = _mm_add_pd(tmp2, offset);
            _mm_store_pd(&s[4*i], tmp1);
            _mm_store_pd(&s[4*i+2], tmp2);
        }
    }
    delete[] suma;
}

Answer 1

这是我第一次回答自己的问题，但似乎很合适。基于 hirschhornsalz 16 字节前缀总和的答案simd-prefix-sum-on-intel-cpu 我想出了一个解决方案，在第一遍使用 SIMD 处理 4、8 和 16 个 32 位字。

一般理论如下。对于n 字的顺序扫描，它需要n 加法（n-1 扫描 n 个字，并从前一组扫描的字中再进行一次加法）。然而，使用 SIMD n 字可以在 log₂(n) 个加法和相等数量的移位加上一个额外的加法和广播中扫描以从先前的 SIMD 扫描中携带。因此，对于n 的某些值，SIMD 方法将获胜。

让我们看一下 SSE、AVX 和 AVX-512 的 32 位字：

4 32-bit words (SSE):      2 shifts, 3 adds, 1 broadcast       sequential: 4 adds
8 32-bit words (AVX):      3 shifts, 4 adds, 1 broadcast       sequential: 8 adds
16 32 bit-words (AVX-512): 4 shifts, 5 adds, 1 broadcast       sequential: 16 adds

基于此，在 AVX-512 之前，SIMD 似乎对扫描 32 位字没有用处。这还假设轮班和广播只能在 1 条指令中完成。 SSE 是这样，但not for AVX and maybe not even for AVX2。

无论如何，我将一些工作和测试的代码放在一起，这些代码使用 SSE 进行前缀求和。

inline __m128 scan_SSE(__m128 x) {
    x = _mm_add_ps(x, _mm_castsi128_ps(_mm_slli_si128(_mm_castps_si128(x), 4))); 
    x = _mm_add_ps(x, _mm_castsi128_ps(_mm_slli_si128(_mm_castps_si128(x), 8)));
    return x;
}

void prefix_sum_SSE(float *a, float *s, const int n) {
__m128 offset = _mm_setzero_ps();
for (int i = 0; i < n; i+=4) {
    __m128 x = _mm_load_ps(&a[i]);
    __m128 out = scan_SSE(x);
    out = _mm_add_ps(out, offset);
    _mm_store_ps(&s[i], out);
    offset = _mm_shuffle_ps(out, out, _MM_SHUFFLE(3, 3, 3, 3)); 
}

注意scan_SSE 函数有两个加法（_mm_add_ps）和两个移位（_mm_slli_si128）。强制转换仅用于使编译器满意并且不会转换为指令。然后在prefix_sum_SSE 数组的主循环内，使用另一个添加和一个随机播放。这总共是 6 次操作，而顺序总和只有 4 次加法。

这是 AVX 的有效解决方案：

inline __m256 scan_AVX(__m256 x) {
    __m256 t0, t1;
    //shift1_AVX + add
    t0 = _mm256_permute_ps(x, _MM_SHUFFLE(2, 1, 0, 3));
    t1 = _mm256_permute2f128_ps(t0, t0, 41);
    x = _mm256_add_ps(x, _mm256_blend_ps(t0, t1, 0x11));
    //shift2_AVX + add
    t0 = _mm256_permute_ps(x, _MM_SHUFFLE(1, 0, 3, 2));
    t1 = _mm256_permute2f128_ps(t0, t0, 41);
    x = _mm256_add_ps(x, _mm256_blend_ps(t0, t1, 0x33));
    //shift3_AVX + add
    x = _mm256_add_ps(x,_mm256_permute2f128_ps(x, x, 41));
    return x;
}

void prefix_sum_AVX(float *a, float *s, const int n) {
    __m256 offset = _mm256_setzero_ps();
    for (int i = 0; i < n; i += 8) {
        __m256 x = _mm256_loadu_ps(&a[i]);
        __m256 out = scan_AVX(x);
        out = _mm256_add_ps(out, offset);
        _mm256_storeu_ps(&s[i], out);
        //broadcast last element
        __m256 t0 = _mm256_permute2f128_ps(out, out, 0x11);
        offset = _mm256_permute_ps(t0, 0xff);
    }   
}

三个班次需要 7 个内在函数。广播需要 2 个内在函数。所以加上 4 个加法，就是 13 个内在函数。对于 AVX2，移位只需要 5 个内在函数，因此总共需要 11 个内在函数。顺序求和只需要 8 次加法。因此，很可能 AVX 和 AVX2 都不适用于第一遍。

编辑：

所以我最终对此进行了基准测试，结果出乎意料。 SSE 和 AVX 代码的速度都大约是以下顺序代码的两倍：

void scan(float a[], float s[], int n) {
    float sum = 0;
    for (int i = 0; i<n; i++) {
        sum += a[i];
        s[i] = sum;
    }
}

我猜这是由于指令级并行。

这样就回答了我自己的问题。在一般情况下，我成功地将 SIMD 用于 pass1。当我在我的 4 核 ivy 桥接系统上将它与 OpenMP 结合使用时，对于 512k 浮点数，总速度提升大约为 7。