c++ - 使用自定义运算符减少 OpenMP SIMD

Question

我有以下循环，我想使用 #pragma omp simd 加速:

#define N 1024
double* data = new double[N];
// Generate data, not important how.

double mean = 0.0
for (size_t i = 0; i < N; i++) {
    mean += (data[i] - mean) / (i+1);
}

如我所料，直接将 #pragma omp simd 放在循环之前没有任何影响(我正在检查运行时间)。我可以使用 #pragma omp parallel for reduction(...) 和自定义 reducer 轻松解决多线程情况，如下所示，但是我如何在这里使用 OpenMP SIMD？

我正在使用以下类来实现 + 和 += 运算符，以将 double 添加到运行平均值以及组合两个运行装置:

class RunningMean {
    private:
        double mean;
        size_t count;

    public:
        RunningMean(): mean(0), count(0) {}
        RunningMean(double m, size_t c): mean(m), count(c) {}

        RunningMean operator+(RunningMean& rhs) {
            size_t c = this->count + rhs.count;
            double m = (this->mean*this->count + rhs.mean*rhs.count) / c;
            return RunningMean(m, c);
        }

        RunningMean operator+(double rhs) {
            size_t c = this->count + 1;
            double m = this->mean + (rhs - this->mean) / c;
            return RunningMean(m, c);
        }

        RunningMean& operator+=(const RunningMean& rhs) {
            this->mean = this->mean*this->count + rhs.mean*rhs.count;
            this->count += rhs.count;
            this->mean /= this->count;
            return *this;
        }

        RunningMean& operator+=(double rhs) {
            this->count++;
            this->mean += (rhs - this->mean) / this->count;
            return *this;
        }

        double getMean() { return mean; }
        size_t getCount() { return count; }
};

这方面的数学来自 http://prod.sandia.gov/techlib/access-control.cgi/2008/086212.pdf .对于多线程、非 SIMD 并行缩减，我执行以下操作:

#pragma omp declare reduction (runningmean : RunningMean : omp_out += omp_in)
RunningMean mean;
#pragma omp parallel for reduction(runningmean:mean)
for (size_t i = 0; i < N; i++)
    mean += data[i];

这让我在使用 8 个线程的 Core i7 2600k 上获得了 3.2 倍的加速。

如果我要在没有 OpenMP 的情况下自己实现 SIMD，我将只在一个 vector 中保留 4 个均值，在另一个 vector 中保留 4 个计数(假设使用 AVX 指令)并继续使用向量化添加 4 元素 double vector operator+(double rhs) 的版本。完成后，我将使用 operator+= 中的数学添加生成的 4 对均值和计数。我如何指示 OpenMP 执行此操作？

Answer 1

问题是

mean += (data[i] - mean) / (i+1);

不容易服从 SIMD。但是，通过仔细研究数学，可以毫不费力地将其矢量化。

关键的论坛是

mean(n+m) = (n*mean(n) + m*mean(m))/(n+m)

显示了如何将 n 个数的均值和 m 个数的均值相加。这可以在您的运算符定义 RunningMean operator+(RunningMean& rhs) 中看到。这解释了为什么您的并行代码有效。如果我们对您的 C++ 代码进行反卷积，我认为这会更清楚:

double mean = 0.0;
int count = 0;
#pragma omp parallel
{
    double mean_private = 0.0;
    int count_private = 0;
    #pragma omp for nowait
    for(size_t i=0; i<N; i++) {
        count_private ++;
        mean_private += (data[i] - mean_private)/count_private;
    }
    #pragma omp critical
    {
        mean = (count_private*mean_private + count*mean);
        count += count_private;
        mean /= count;
    }
}

但我们可以对 SIMD 使用相同的想法(并将它们组合在一起)。但是让我们先做只有 SIMD 的部分。使用 AVX，我们可以同时处理四种并行方式。每个平行均值将处理以下数据元素:

mean 1 data elements: 0,  4,  8, 12,...
mean 2 data elements: 1,  5,  9, 13,...
mean 3 data elements: 2,  6, 10, 14,...
mean 4 data elements: 3,  7, 11, 15,...

首先，我们遍历了所有元素，然后将四个平行的和相加并除以四(因为每个和都遍历了 N/4 个元素)。

下面是代码

double mean = 0.0;
__m256d mean4 = _mm256_set1_pd(0.0);
__m256d count4 = _mm256_set1_pd(0.0);
for(size_t i=0; i<N/4; i++) {
    count4 = _mm256_add_pd(count4,_mm256_set1_pd(1.0));
    __m256d t1 = _mm256_loadu_pd(&data[4*i]);
    __m256d t2 = _mm256_div_pd(_mm256_sub_pd(t1, mean4), count4);
    mean4 = _mm256_add_pd(t2, mean4);   
}
__m256d t1 = _mm256_hadd_pd(mean4,mean4);
__m128d t2 = _mm256_extractf128_pd(t1,1);
__m128d t3 = _mm_add_sd(_mm256_castpd256_pd128(t1),t2);
mean = _mm_cvtsd_f64(t3)/4;
int count = 0;
double mean2 = 0;
for(size_t i=4*(N/4); i<N; i++) {
    count++;
    mean2 += (data[i] - mean2)/count;
}
mean = (4*(N/4)*mean + count*mean2)/N;

最后，我们可以将它与 OpenMP 结合起来，像这样充分利用 SIMD 和 MIMD 的优势

double mean = 0.0;
int count = 0;
#pragma omp parallel 
{
    double mean_private = 0.0;
    int count_private = 0;
    __m256d mean4 = _mm256_set1_pd(0.0);
    __m256d count4 = _mm256_set1_pd(0.0);
    #pragma omp for nowait
    for(size_t i=0; i<N/4; i++) {
        count_private++;
        count4 = _mm256_add_pd(count4,_mm256_set1_pd(1.0));
        __m256d t1 = _mm256_loadu_pd(&data[4*i]);
        __m256d t2 = _mm256_div_pd(_mm256_sub_pd(t1, mean4), count4);
        mean4 = _mm256_add_pd(t2, mean4);
    }
    __m256d t1 = _mm256_hadd_pd(mean4,mean4);
    __m128d t2 = _mm256_extractf128_pd(t1,1);
    __m128d t3 = _mm_add_sd(_mm256_castpd256_pd128(t1),t2);
    mean_private = _mm_cvtsd_f64(t3)/4;

    #pragma omp critical
    {
        mean = (count_private*mean_private + count*mean);
        count += count_private;
        mean /= count;
    }   
}
int count2 = 0;
double mean2 = 0;
for(size_t i=4*(N/4); i<N; i++) {
    count2++;
    mean2 += (data[i] - mean2)/count2;
}
mean = (4*(N/4)*mean + count2*mean2)/N;

这是一个工作示例(使用 -O3 -mavx -fopenmp 编译)

#include <stdio.h>
#include <stdlib.h>
#include <x86intrin.h>

double mean_simd(double *data, const int N) {
    double mean = 0.0;
    __m256d mean4 = _mm256_set1_pd(0.0);
    __m256d count4 = _mm256_set1_pd(0.0);
    for(size_t i=0; i<N/4; i++) {
        count4 = _mm256_add_pd(count4,_mm256_set1_pd(1.0));
        __m256d t1 = _mm256_loadu_pd(&data[4*i]);
        __m256d t2 = _mm256_div_pd(_mm256_sub_pd(t1, mean4), count4);
        mean4 = _mm256_add_pd(t2, mean4);           
    }
    __m256d t1 = _mm256_hadd_pd(mean4,mean4);
    __m128d t2 = _mm256_extractf128_pd(t1,1);
    __m128d t3 = _mm_add_sd(_mm256_castpd256_pd128(t1),t2);
    mean = _mm_cvtsd_f64(t3)/4;
    int count = 0;
    double mean2 = 0;
    for(size_t i=4*(N/4); i<N; i++) {
        count++;
        mean2 += (data[i] - mean2)/count;
    }
    mean = (4*(N/4)*mean + count*mean2)/N;
    return mean;
}

double mean_simd_omp(double *data, const int N) {
    double mean = 0.0;
    int count = 0;
    #pragma omp parallel 
    {
        double mean_private = 0.0;
        int count_private = 0;
        __m256d mean4 = _mm256_set1_pd(0.0);
        __m256d count4 = _mm256_set1_pd(0.0);
        #pragma omp for nowait
        for(size_t i=0; i<N/4; i++) {
            count_private++;
            count4 = _mm256_add_pd(count4,_mm256_set1_pd(1.0));
            __m256d t1 = _mm256_loadu_pd(&data[4*i]);
            __m256d t2 = _mm256_div_pd(_mm256_sub_pd(t1, mean4), count4);
            mean4 = _mm256_add_pd(t2, mean4);
        }
        __m256d t1 = _mm256_hadd_pd(mean4,mean4);
        __m128d t2 = _mm256_extractf128_pd(t1,1);
        __m128d t3 = _mm_add_sd(_mm256_castpd256_pd128(t1),t2);
        mean_private = _mm_cvtsd_f64(t3)/4;

        #pragma omp critical
        {
            mean = (count_private*mean_private + count*mean);
            count += count_private;
            mean /= count;
        }   
    }
    int count2 = 0;
    double mean2 = 0;
    for(size_t i=4*(N/4); i<N; i++) {
        count2++;
        mean2 += (data[i] - mean2)/count2;
    }
    mean = (4*(N/4)*mean + count2*mean2)/N;
    return mean;
}


int main() {
    const int N = 1001;
    double data[N];

    for(int i=0; i<N; i++) data[i] = 1.0*rand()/RAND_MAX;
    float sum = 0; for(int i=0; i<N; i++) sum+= data[i]; sum/=N;
    printf("mean %f\n", sum);
    printf("mean_simd %f\n", mean_simd(data, N);
    printf("mean_simd_omp %f\n", mean_simd_omp(data, N));
}