使用自定义运算符的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 并行 for reduction(...) 来处理多线程情况，并使用自定义约减器，如下所示，但是如何在这里使用 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倍的加速。

如果我要自己实现SIMD而不使用OpenMP，我会在一个向量中维护4个均值，在另一个向量中维护4个计数（假设使用AVX指令），并使用operator+(double rhs)的向量化版本不断添加4元素双精度向量。完成后，我将使用operator+=中的数学方法添加4对均值和计数。我该如何指示OpenMP执行此操作？

- chippies

---

你的自定义reducer有递归，我猜想编译器无法解决...那么mean += data[i] / N怎么样？ - GHL

明白了，但是您希望如何使用SIMD加速计算呢？因为要计算每个运行总和（即循环的每次迭代），您需要前一个结果的结果。我认为，为了使SIMD起作用，您需要让计算独立进行？ - GHL

@chippies，你验证过并行版本与串行版本得到的结果相同吗？ - Z boson

是的，我有。结果在小数点后12位相同，这很好。 - chippies

@Zboson 这是我的理解：在每个线程开始时，会创建一个私有的 mean 副本，其中 mean.mean=0 和 mean.count=0。该线程循环遍历 N 数组元素的子集，执行 mean += data[i]，生成具有一些值的 mean.mean 和处理的元素数量的 mean.count 的 mean。该私有的 mean 副本进入约简，其中使用两个 RunningMean 对象的 += 运算符。即使在副本进入约简之前只有一个 data 条目进入了 mean 的私有副本，数学仍然可以正确地工作。 - chippies

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

- Z boson

---

这是正确的数学方法来合并两个均值，但对于 mean(n+1) 的特殊情况，即添加一个数据点，有一个特殊的方程式，我将其标记为 X。mean(n+1) = mean(n) + (X - mean(n)) / (n+1)。我想在您的 for 循环中使用此方程式，即代替 tmp += data[i];，因为它比仅将所有数字相加更具数值稳定性。 - chippies

@chippies，mean(n+1) = mean(n) + (X - mean(n)) / (n+1) 的问题在于它不容易适应SIMD。我认为它不会自动向量化。你是否查看了汇编代码以查看其是否向量化？我认为我可以使用内部函数或矢量类手动向量化它。也许你应该重新发布你的问题，并加上SSE/AVX标签，而不是OpenMP的内容。你已经解决了OpenMP部分的MIMD问题。 - Z boson

@chippies，你的问题让我想起了使用OpenMP和SIMD进行前缀和（因为它有依赖性）的经验。我做过这个，虽然不太容易但最终还是成功了。https://dev59.com/52Ik5IYBdhLWcg3wMLmI。将你的公式贴上AVX标签，看看是否有人能弄清楚它。 - Z boson

查看反汇编代码显示，GCC（MinGW x86-64）实际上已经在使用XMM寄存器和指令vaddsd和vdivsd。由于这些是标量指令，而不是打包/SIMD指令，我猜测这并不是真正的向量化，而是来自GCC的一些奇怪优化。无论如何，当我尝试使用OpenMP SIMD时，在某些情况下会出现相同的指令，因此OpenMP无法将我的代码向量化。 - chippies

谢谢你提供的链接。我本来希望OpenMP也能有一种方式来将自定义约简操作向量化，可惜没有。感谢你的帮助。 - chippies

Answer 2

KISS的答案是：只需在循环外计算平均值。并行化以下代码：

double sum = 0.0;
for(size_t i = 0; i < N; i++) sum += data[i];
double mean = sum/N;

这个总和很容易并行化，但你不会看到任何 SIMD 优化的效果：它纯粹是内存绑定，CPU 只会等待来自内存的数据。如果 N 和 1024 一样小，那么甚至并行化都没有什么意义，同步开销会消耗所有收益。

- cmaster - reinstate monica

---

OP已经说明他想按照自己的方式做，因为“它在数值上更加稳定”。如果N只有1024，我同意使用MIMD没有意义，但在这种情况下，SIMD仍然是有意义的（我展示了如何做到这一点）。但是可以假设OP在每次迭代之间进行其他计算，在这种情况下可能不会受到内存限制。似乎人们并没有认真对待OP的问题。当他们看到依赖关系时，他们放弃了，或者只是说“以显而易见的方式计算平均值”（这就是我在第一次回答中所做的，如果您查看我的编辑）。 - Z boson

@Zboson 我没有看到关于数值稳定性的评论。然而，从精度角度来看，求和和除法的方法稍微好一些，因为它避免了由于除法而产生的不必要的舍入误差。加法的舍入误差是无法避免的。的确，求和方法更容易达到无穷大。但这需要非常大的数字，几乎不用担心。如果求和是在整数中完成的，那么这将是一个问题，但使用 double 累加器就不是了。 - cmaster - reinstate monica

这可能是真的。如果是这样，那么显而易见的解决方案就是走的路线。但OP还说：“在我的Core i7 2600k上使用8个线程可以让我加速3.2倍”，这告诉我他正在处理超过1024个元素和/或在每次迭代之间做其他事情。无论如何，他不可能通过OpenMP获得任何加速，因为只有1024个元素。实际上，由于开销，它会更糟。即使有更多的元素，它也会受到内存限制，因此他必须要做更多的事情。 - Z boson

@Zboson，你说得对，我测试的数据量更多——810241024个元素。这项工作的一些背景是，方差也有类似的方程式。由于所有这些方程都遵循相似的模式，在实现更复杂的情况之前，我想在简单情况下（即均值）进行练习。使用我在问题中链接的PDF中的方程式会对高阶统计的准确性产生更大的影响。要了解它有多糟糕，请看这里。 - chippies

@chippies，是的，我想这只是你想要完成更复杂计算的第一步。我希望我的解决方案能够帮助你进行下一步。如果SIMD代码比你预期的要慢，尝试将其展开四次。加法的延迟为3，因此必须至少展开三次，但你可能会被除法的吞吐量所限制。 - Z boson