使用 Rcpp 和 OpenMP 在 R 中多线程和 SIMD 矢量化 Mandelbrot

Multithreaded & SIMD vectorized Mandelbrot in R using Rcpp & OpenMP

作为 OpenMP & Rcpp 性能测试,我想检查使用最直接和简单的 Rcpp+[=12= 计算 R 中的 Mandelbrot 集的速度有多快] 执行。目前我所做的是:

#include <Rcpp.h>
#include <omp.h>
// [[Rcpp::plugins(openmp)]]

using namespace Rcpp;

// [[Rcpp::export]]
Rcpp::NumericMatrix mandelRcpp(const double x_min, const double x_max, const double y_min, const double y_max,
                         const int res_x, const int res_y, const int nb_iter) {
  Rcpp::NumericMatrix ret(res_x, res_y);
  double x_step = (x_max - x_min) / res_x;
  double y_step = (y_max - y_min) / res_y;
  int r,c;
#pragma omp parallel for default(shared) private(c) schedule(dynamic,1)
  for (r = 0; r < res_y; r++) {
    for (c = 0; c < res_x; c++) {
      double zx = 0.0, zy = 0.0, new_zx;
      double cx = x_min + c*x_step, cy = y_min + r*y_step;
      int n = 0;
      for (n=0;  (zx*zx + zy*zy < 4.0 ) && ( n < nb_iter ); n++ ) {
        new_zx = zx*zx - zy*zy + cx;
        zy = 2.0*zx*zy + cy;
        zx = new_zx;
      }
      ret(c,r) = n;
    }
  }
  return ret;
}

然后在 R 中:

library(Rcpp)
sourceCpp("mandelRcpp.cpp")
xlims=c(-0.74877,-0.74872);
ylims=c(0.065053,0.065103);
x_res=y_res=1080L; nb_iter=10000L;
system.time(m <- mandelRcpp(xlims[[1]], xlims[[2]], ylims[[1]], ylims[[2]], x_res, y_res, nb_iter)) 
# 0.92s
rainbow=c(rgb(0.47,0.11,0.53),rgb(0.27,0.18,0.73),rgb(0.25,0.39,0.81),rgb(0.30,0.57,0.75),rgb(0.39,0.67,0.60),rgb(0.51,0.73,0.44),rgb(0.67,0.74,0.32),rgb(0.81,0.71,0.26),rgb(0.89,0.60,0.22),rgb(0.89,0.39,0.18),rgb(0.86,0.13,0.13))
    cols=c(colorRampPalette(rainbow)(100),rev(colorRampPalette(rainbow)(100)),"black") # palette
par(mar=c(0, 0, 0, 0))
system.time(image(m^(1/7), col=cols, asp=diff(ylims)/diff(xlims), axes=F, useRaster=T)) 
# 0.5s

我不确定除了 OpenMP 多线程之外是否还有其他明显的速度改进我可以利用,例如通过 simd 矢量化? (在 openmp #pragma 中使用 simd 选项似乎没有做任何事情)

PS 起初我的代码崩溃了,但后来我发现通过将 ret[r,c] = n; 替换为 ret(r,c) = n; 解决了这个问题 按照下面的答案中的建议使用 Armadillo 类 可以使事情稍微快一些,尽管时间几乎相同。还翻转了 xy,所以当用 image() 绘制时它会以正确的方向出现。使用 8 线程速度约为比矢量化普通 R Mandelbrot 版本快 350 倍here and also about 7.3 times faster than the (non-multithreaded) Python/Numba version here (similar to PyCUDA or PyOpenCL speeds), so quite happy with that... Rasterizing/display now seems the bottleneck in R....

不要OpenMPRcpp*Vector*Matrix 对象,因为它们掩盖了 SEXP 单线程的函数/内存分配。 OpenMP 是 multi-threaded approach

这就是代码崩溃的原因。

绕过此限制的一种方法是使用非R 数据结构来存储结果。以下其中一项就足够了:arma::matEigen::MatrixXdstd::vector<T>... 因为我喜欢犰狳,所以我会将 res 矩阵更改为 arma::matRcpp::NumericMatrix。因此,以下将并行执行您的代码:

#include <RcppArmadillo.h> // Note the changed include and new attribute
// [[Rcpp::depends(RcppArmadillo)]]

// Avoid including header if openmp not on system
#ifdef _OPENMP
#include <omp.h>
#endif
// [[Rcpp::plugins(openmp)]]

// Note the changed return type
// [[Rcpp::export]]
arma::mat mandelRcpp(const double x_min, const double x_max,
                     const double y_min, const double y_max,
                     const int res_x, const int res_y, const int nb_iter) {
  arma::mat ret(res_x, res_y); // note change
  double x_step = (x_max - x_min) / res_x;
  double y_step = (y_max - y_min) / res_y;
  unsigned r,c;

  #pragma omp parallel for shared(res)
  for (r = 0; r < res_y; r++) {
    for (c = 0; c < res_x; c++) {
      double zx = 0.0, zy = 0.0, new_zx;
      double cx = x_min + c*x_step, cy = y_min + r*y_step;
      unsigned n = 0;
      for (;  (zx*zx + zy*zy < 4.0 ) && ( n < nb_iter ); n++ ) {
        new_zx = zx*zx - zy*zy + cx;
        zy = 2.0*zx*zy + cy;
        zx = new_zx;
      }

      if(n == nb_iter) {
        n = 0;
      }

      ret(r, c) = n;
    }
  }

  return ret;
}

使用测试代码(注意 yx 未定义,因此我假设 y = ylimsx = xlims)我们有:

xlims = ylims = c(-2.0, 2.0)

x_res = y_res = 400L
nb_iter = 256L

system.time(m <-
              mandelRcpp(xlims[[1]], xlims[[2]],
                         ylims[[1]], ylims[[2]], 
                         x_res, y_res, nb_iter))

rainbow = c(
  rgb(0.47, 0.11, 0.53),
  rgb(0.27, 0.18, 0.73),
  rgb(0.25, 0.39, 0.81),
  rgb(0.30, 0.57, 0.75),
  rgb(0.39, 0.67, 0.60),
  rgb(0.51, 0.73, 0.44),
  rgb(0.67, 0.74, 0.32),
  rgb(0.81, 0.71, 0.26),
  rgb(0.89, 0.60, 0.22),
  rgb(0.89, 0.39, 0.18),
  rgb(0.86, 0.13, 0.13)
)

cols = c(colorRampPalette(rainbow)(100),
         rev(colorRampPalette(rainbow)(100)),
         "black") # palette
par(mar = c(0, 0, 0, 0))

image(m,
      col = cols,
      asp = diff(range(ylims)) / diff(range(xlims)),
      axes = F)

对于:

我继续使用 GCC 和 Clang 的矢量扩展对 OP 的代码进行矢量化。在展示我是如何做到这一点之前,让我展示一下使用以下硬件的性能:

Skylake (SKL) at 3.1 GHz with 4 cores
Knights Landing (KNL) at 1.5 GHz with 68 cores
ARMv8 Cortex-A57 arch64 (Nvidia Jetson TX1) 4 cores at ? GHz

nb_iter = 1000000
                        GCC             Clang
SKL_scalar              6m5,422s
SKL_SSE41               3m18,058s
SKL_AVX2                1m37,843s       1m39,943s
SKL_scalar_omp          0m52,237s
SKL_SSE41_omp           0m29,624s       0m31,356s
SKL_AVX2_omp            0m14,156s       0m16,783s

ARM_scalar              15m28.285s
ARM_vector              9m26.384s
ARM_scalar_omp          3m54.242s
ARM_vector_omp          2m21.780s

KNL_scalar              19m34.121s
KNL_SSE41               11m30.280s
KNL_AVX2                5m0.005s        6m39.568s
KNL_AVX512              2m40.934s       6m20.061s
KNL_scalar_omp          0m9.108s
KNL_SSE41_omp           0m6.666s        0m6.992s
KNL_AVX2_omp            0m2.973s        0m3.988s
KNL_AVX512_omp          0m1.761s        0m3.335s

KNL 相对于 SKL 的理论加速是

(68 cores/4 cores)*(1.5 GHz/3.1 Ghz)*
(8 doubles per lane/4 doubles per lane) = 16.45

我详细介绍了 GCC 和 Clang 的矢量扩展功能 here。要对此处的 OP 代码进行矢量化,我们需要定义三个额外的矢量运算。

1.广播

对于向量 v 和标量 s GCC 不能做到 v = s 但 Clang 可以。但是我找到了一个适用于 GCC 和 Clang 的不错的解决方案。例如

vsi v = s - (vsi){};

2。 any() 函数 like in OpenCL or like in R.

我想出的最好的办法是通用函数

static bool any(vli const & x) {
  for(int i=0; i<VLI_SIZE; i++) if(x[i]) return true;
  return false;
}

Clang 实际上生成相对 efficient code for this using the ptest instruction (but not for AVX512) 但 GCC 不会。

3。压缩

计算以 64 位双精度数完成,但结果以 32 位整数形式写出。因此,使用 64 位整数完成两次计算,然后将这两次计算压缩为一个 32 位整数向量。我想出了一个通用的解决方案,Clang 做得很好

static vsi compress(vli const & lo, vli const & hi) {
  vsi lo2 = (vsi)lo, hi2 = (vsi)hi, z;
  for(int i=0; i<VLI_SIZE; i++) z[i+0*VLI_SIZE] = lo2[2*i];
  for(int i=0; i<VLI_SIZE; i++) z[i+1*VLI_SIZE] = hi2[2*i];
  return z;
}

以下解决方案有效 better for GCC but is no better for Clang。但由于这个功能并不重要,所以我只使用通用版本。

static vsi compress(vli const & low, vli const & high) {
#if defined(__clang__)
  return __builtin_shufflevector((vsi)low, (vsi)high, MASK);
#else
  return __builtin_shuffle((vsi)low, (vsi)high, (vsi){MASK});
#endif
}

这些定义不依赖任何特定于 x86 的内容,并且代码(定义如下)针对 ARM 处理器以及 GCC 和 Clang 进行编译。


既然这里定义了这些就是代码

#include <string.h>
#include <inttypes.h>
#include <Rcpp.h>

using namespace Rcpp;

#ifdef _OPENMP
#include <omp.h>
#endif
// [[Rcpp::plugins(openmp)]]
// [[Rcpp::plugins(cpp14)]]

#if defined ( __AVX512F__ ) || defined ( __AVX512__ )
static const int SIMD_SIZE = 64;
#elif defined ( __AVX2__ )
static const int SIMD_SIZE = 32;
#else
static const int SIMD_SIZE = 16;
#endif

static const int VSI_SIZE = SIMD_SIZE/sizeof(int32_t);
static const int VLI_SIZE = SIMD_SIZE/sizeof(int64_t);
static const int VDF_SIZE = SIMD_SIZE/sizeof(double);

#if defined(__clang__)
typedef int32_t vsi __attribute__ ((ext_vector_type(VSI_SIZE)));
typedef int64_t vli __attribute__ ((ext_vector_type(VLI_SIZE)));
typedef double  vdf __attribute__ ((ext_vector_type(VDF_SIZE)));
#else
typedef int32_t vsi __attribute__ ((vector_size (SIMD_SIZE)));
typedef int64_t vli __attribute__ ((vector_size (SIMD_SIZE)));
typedef double  vdf __attribute__ ((vector_size (SIMD_SIZE)));
#endif

static bool any(vli const & x) {
  for(int i=0; i<VLI_SIZE; i++) if(x[i]) return true;
  return false;
}

static vsi compress(vli const & lo, vli const & hi) {
  vsi lo2 = (vsi)lo, hi2 = (vsi)hi, z;
  for(int i=0; i<VLI_SIZE; i++) z[i+0*VLI_SIZE] = lo2[2*i];
  for(int i=0; i<VLI_SIZE; i++) z[i+1*VLI_SIZE] = hi2[2*i];
  return z;
}

// [[Rcpp::export]]
IntegerVector frac(double x_min, double x_max, double y_min,  double y_max, int res_x, int res_y, int nb_iter) {
  IntegerVector out(res_x*res_y);
  vdf x_minv = x_min - (vdf){}, y_minv = y_min - (vdf){};
  vdf x_stepv = (x_max - x_min)/res_x - (vdf){}, y_stepv = (y_max - y_min)/res_y - (vdf){};
  double a[VDF_SIZE] __attribute__ ((aligned(SIMD_SIZE)));
  for(int i=0; i<VDF_SIZE; i++) a[i] = 1.0*i;
  vdf vi0 = *(vdf*)a;

  #pragma omp parallel for schedule(dynamic) collapse(2)
  for (int r = 0; r < res_y; r++) {
    for (int c = 0; c < res_x/(VSI_SIZE); c++) {
      vli nv[2] = {0 - (vli){}, 0 - (vli){}};
      for(int j=0; j<2; j++) {
        vdf c2 = 1.0*VDF_SIZE*(2*c+j) + vi0;
        vdf zx = 0.0 - (vdf){}, zy = 0.0 - (vdf){}, new_zx;
        vdf cx = x_minv + c2*x_stepv, cy = y_minv + r*y_stepv;
        vli t = -1 - (vli){};
        for (int n = 0; any(t = zx*zx + zy*zy < 4.0) && n < nb_iter; n++, nv[j] -= t) {
          new_zx = zx*zx - zy*zy + cx;
          zy = 2.0*zx*zy + cy;
          zx = new_zx;
        }
      }
      vsi sp = compress(nv[0], nv[1]);
      memcpy(&out[r*res_x + VSI_SIZE*c], (int*)&sp, SIMD_SIZE);
    }
  }
  return out;
}

R代码和OP的代码几乎一样

library(Rcpp)
sourceCpp("frac.cpp", verbose=TRUE, rebuild=TRUE)                                                                                                                                                         
xlims=c(-0.74877,-0.74872);
ylims=c(0.065053,0.065103);
x_res=y_res=1080L; nb_iter=100000L;

t = system.time(m <- frac(xlims[[1]], xlims[[2]], ylims[[1]], ylims[[2]], x_res, y_res, nb_iter))
print(t)
m2 = matrix(m, ncol = x_res)

rainbow = c(
  rgb(0.47, 0.11, 0.53),
  rgb(0.27, 0.18, 0.73),
  rgb(0.25, 0.39, 0.81),
  rgb(0.30, 0.57, 0.75),
  rgb(0.39, 0.67, 0.60),
  rgb(0.51, 0.73, 0.44),
  rgb(0.67, 0.74, 0.32),
  rgb(0.81, 0.71, 0.26),
  rgb(0.89, 0.60, 0.22),
  rgb(0.89, 0.39, 0.18),
  rgb(0.86, 0.13, 0.13)
)

cols = c(colorRampPalette(rainbow)(100),
         rev(colorRampPalette(rainbow)(100)),"black") # palette                                                                                                                  
par(mar = c(0, 0, 0, 0))
image(m2^(1/7), col=cols, asp=diff(ylims)/diff(xlims), axes=F, useRaster=T)

要针对 GCC 或 Clang 进行编译,请将文件 ~/.R/Makevars 更改为

CXXFLAGS= -Wall -std=c++14 -O3 -march=native -ffp-contract=fast -fopenmp
#uncomment the following two lines for clang    
#CXX=clang-5.0
#LDFLAGS= -lomp

如果您在让 OpenMP 为 Clang 工作时遇到问题,请参阅 this


该代码生成大致相同的图像。