kernel function by using RcppParallel with RcppArmadillo

library(kernlab)
library(Rcpp)
library(RcppArmadillo)
library(RcppParallel)
sourceCpp(code = '
// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
using namespace Rcpp;
using namespace arma;

// [[Rcpp::export]]
NumericMatrix kernelMatrix_Arma(NumericMatrix Xr, NumericMatrix Centerr, double sigma) {
  uword n = Xr.nrow(), b = Centerr.nrow(), row_index, col_index;
  mat X(Xr.begin(), n, Xr.ncol(), false),
      Center(Centerr.begin(), b, Centerr.ncol(), false),
      KerX(X*Center.t());
  colvec X_sq = sum(square(X), 1) / 2;
  rowvec Center_sq = (sum(square(Center), 1)).t() / 2;
  KerX.each_row() -= Center_sq;
  KerX.each_col() -= X_sq;
  KerX *= 1 / (sigma * sigma);
  KerX = exp(KerX);
  return wrap(KerX);
}')

sourceCpp(code = '
// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
#include <omp.h>
// [[Rcpp::plugins(openmp)]]
using namespace Rcpp;
using namespace arma;

// [[Rcpp::export]]
NumericMatrix kernelMatrix_openmp(NumericMatrix Xr, NumericMatrix Centerr, double sigma) {
  omp_set_num_threads(omp_get_max_threads());
  uword n = Xr.nrow(), b = Centerr.nrow(), row_index, col_index;
  mat X(Xr.begin(), n, Xr.ncol(), false);
  mat Center(Centerr.begin(), b, Centerr.ncol(), false);
  mat KerX(n, b);
  #pragma omp parallel private(row_index, col_index)
  for (row_index = 0; row_index < n; row_index++)
  {
    #pragma omp for nowait
    for (col_index = 0; col_index < b; col_index++)
    {
      KerX(row_index, col_index) = exp(sum(square(X.row(row_index)
        - Center.row(col_index))) / (-2.0 * sigma * sigma));
    }
  }
  return wrap(KerX);
}')

sourceCpp(code = '
// [[Rcpp::depends(RcppArmadillo, RcppParallel)]]
#define ARMA_DONT_USE_CXX11
#include <RcppArmadillo.h>
#include <RcppParallel.h>
using namespace Rcpp;
using namespace arma;
using namespace RcppParallel;

struct KernelCompute: public Worker {
  mat& X;
  mat& Center;
  double sigma;
  mat& output;
  KernelCompute(mat& X, mat& Center, double sigma, mat& output) : X(X), Center(Center), sigma(sigma), output(output) {}
  void operator()(std::size_t begin, std::size_t end) {
    for (uword row_index = begin; row_index < end; row_index++)
    {
      for (uword col_index = 0; col_index < Center.n_rows; col_index++)
        output(row_index, col_index) = exp(sum(square(X.row(row_index)
          - Center.row(col_index))) / (-2.0 * sigma * sigma));
    }
  }
};

// [[Rcpp::export]]
NumericMatrix kernelMatrix_tbb(NumericMatrix Xr, NumericMatrix Centerr, double sigma) {
  uword n = Xr.nrow(), b = Centerr.nrow();
  mat X(Xr.begin(), n, Xr.ncol(), false),
      Center(Centerr.begin(), b, Centerr.ncol(), false), KerX(n, b);
  KernelCompute kernelCompute(X, Center, sigma, KerX);
  parallelFor(0, X.n_rows, kernelCompute);
  return wrap(KerX);
}')

N = 3000
p = 100
b = 500
X = matrix(rnorm(N*p), ncol = p)
center = X[sample(1:N, b),]
sigma = 50
kernel_X = kernelMatrix(rbfdot(sigma=1/(2*sigma^2)), X, center)
kernel_X_arma = kernelMatrix_Arma(X, center, sigma)
kernel_X_openmp = kernelMatrix_openmp(X, center, sigma)
kernel_X_tbb = kernelMatrix_tbb(X, center, sigma)
## test
all.equal(kernel_X@.Data, kernel_X_arma)
all.equal(kernel_X@.Data, kernel_X_openmp)
all.equal(kernel_X@.Data, kernel_X_tbb)
# TRUE

library(rbenchmark)
benchmark(
  arma = kernelMatrix_Arma(X, center, sigma),
  openmp = kernelMatrix_openmp(X, center, sigma),
  tbb = kernelMatrix_tbb(X, center, sigma),
  kernlab = kernelMatrix(rbfdot(sigma=1/(2*sigma^2)), X, center),
  columns=c("test", "replications","elapsed", "relative"),
  replications=20, order="relative")
## p = 100
#      test replications elapsed relative
# 3     tbb           20    1.51    1.000
# 2  openmp           20    1.75    1.159
# 1    arma           20    1.97    1.305
# 4 kernlab           20    3.25    2.152

## p = 300
#      test replications elapsed relative
# 1    arma           20    2.20    1.000
# 4 kernlab           20    3.62    1.645
# 3     tbb           20   27.52   12.509
# 2  openmp           20   28.48   12.945