c++ - Intel cpu 上的 SIMD 前缀和

Question

我需要实现一个前缀和算法，并且需要它尽可能快。
例如:

[3, 1,  7,  0,  4,  1,  6,  3]

应该给:

[3, 4, 11, 11, 15, 16, 22, 25]

有没有办法使用 SSE SIMD CPU 指令来做到这一点？

我的第一个想法是递归地对每一对进行并行求和，直到所有总和都被计算如下!

//in parallel do 
for (int i = 0; i < z.length; i++) {
    z[i] = x[i << 1] + x[(i << 1) + 1];
}

为了让算法更清晰一点，z并不是最终的输出，而是用来计算输出的。

int[] w = computePrefixSum(z);
for (int i = 1; i < ouput.length; i++) {
    ouput[i] = (i % 2 == 0) ? (x[i] + ouput[i - 1]) :  w[(i - 1) >> 1];
}

Answer 1

我所知道的最快的并行前缀求和算法是并行运行两次总和，并在第二次传递中使用 SSE。

在第一遍中，您并行计算部分总和并存储每个部分总和的总和。在第二遍中，您将前一个部分总和的总和添加到下一个部分总和。您可以使用多个线程(例如使用 OpenMP)并行运行两个 channel 。第二遍您也可以使用 SIMD，因为每个部分总和都会添加一个常数值。

假设数组的 n 个元素、m 个核心和 w 的 SIMD 宽度，则时间成本应该是

n/m + n/(m*w) = (n/m)*(1+1/w)

由于第一次不使用SIMD，所以时间成本总是大于n/m

例如，对于 SIMD_width 为 4 的四个内核(带 SSE 的四个 32 位 float )，成本将为 5n/16。或者比时间成本为 n 的顺序代码快 3.2 倍。使用超线程，速度会更快。

在特殊情况下，也可以在第一遍使用 SIMD。那么时间成本就是

2*n/(m*w)

我发布了一般情况下的代码，该代码使用 OpenMP 作为 SSE 代码的线程和内在函数，并在以下链接中讨论了有关特殊情况的详细信息 parallel-prefix-cumulative-sum-with-sse

编辑: 我设法为第一遍找到了一个 SIMD 版本，它的速度大约是顺序代码的两倍。现在我的四核 Ivy 桥系统总共提升了大约 7 个。

编辑: 对于较大的数组，一个问题是在第一次通过后，大多数值已从缓存中清除。我想出了一个解决方案，它在一个 block 内并行运行，但串行运行每个 block 。 chunk_size 是一个应该调整的值。例如，我将其设置为 1MB = 256K float 。现在，当值仍在二级缓存中时，第二遍就完成了。这样做对大型数组有很大的改进。

这是 SSE 的代码。 AVX 代码的速度差不多，所以我没有在这里发布。进行前缀和的函数是scan_omp_SSEp2_SSEp1_chunk。向它传递一个 float 组 a 并用累积和填充数组 s。

__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_shuffle_ps(_mm_setzero_ps(), x, 0x40)); 
    return x;
}

float pass1_SSE(float *a, float *s, const int n) {
    __m128 offset = _mm_setzero_ps();
    #pragma omp for schedule(static) nowait
    for (int i = 0; i < n / 4; i++) {
        __m128 x = _mm_load_ps(&a[4 * i]);
        __m128 out = scan_SSE(x);
        out = _mm_add_ps(out, offset);
        _mm_store_ps(&s[4 * i], out);
        offset = _mm_shuffle_ps(out, out, _MM_SHUFFLE(3, 3, 3, 3));
    }
    float tmp[4];
    _mm_store_ps(tmp, offset);
    return tmp[3];
}

void pass2_SSE(float *s, __m128 offset, const int n) {
    #pragma omp for schedule(static)
    for (int i = 0; i<n/4; i++) {
        __m128 tmp1 = _mm_load_ps(&s[4 * i]);
        tmp1 = _mm_add_ps(tmp1, offset);
        _mm_store_ps(&s[4 * i], tmp1);
    }
}

void scan_omp_SSEp2_SSEp1_chunk(float a[], float s[], int n) {
    float *suma;
    const int chunk_size = 1<<18;
    const int nchunks = n%chunk_size == 0 ? n / chunk_size : n / chunk_size + 1;
    //printf("nchunks %d\n", nchunks);
    #pragma omp parallel
    {
        const int ithread = omp_get_thread_num();
        const int nthreads = omp_get_num_threads();

        #pragma omp single
        {
            suma = new float[nthreads + 1];
            suma[0] = 0;
        }

        float offset2 = 0.0f;
        for (int c = 0; c < nchunks; c++) {
            const int start = c*chunk_size;
            const int chunk = (c + 1)*chunk_size < n ? chunk_size : n - c*chunk_size;
            suma[ithread + 1] = pass1_SSE(&a[start], &s[start], chunk);
            #pragma omp barrier
            #pragma omp single
            {
                float tmp = 0;
                for (int i = 0; i < (nthreads + 1); i++) {
                    tmp += suma[i];
                    suma[i] = tmp;
                }
            }
            __m128 offset = _mm_set1_ps(suma[ithread]+offset2);
            pass2_SSE(&s[start], offset, chunk);
            #pragma omp barrier
            offset2 = s[start + chunk-1];
        }
    }
    delete[] suma;
}