使用Rcpp提取子矩阵为何较慢？求更快实现方案

问题描述 #

在开发R包时，使用以下Rcpp代码从矩阵中提取子矩阵，发现性能远不如原生R的提取方式：

原始Rcpp实现 #

NumericMatrix extractColumns(NumericMatrix A, IntegerVector indices) {
  int n = indices.size();
  int m = A.nrow();
  NumericMatrix B(m, n);
  for(int i = 0; i < n; i++) {
    B(_, i) = A(_, indices(i));
  }
  return B;
}

基准测试代码 #

#include <Rcpp.h>
using namespace Rcpp;

// [[Rcpp::export]]
NumericMatrix extractColumns(NumericMatrix A, IntegerVector indices) {
  int n = indices.size();
  int m = A.nrow();
  NumericMatrix B(m, n);
  for(int i = 0; i < n; i++) {
    B(_, i) = A(_, indices(i));
  }
  return B;
}

/*** R
library(microbenchmark)
A <- matrix(rgamma(100000L, 5, 1), nrow = 2L, ncol = 50000L)
indices <- 2000L:40000L 
microbenchmark(
  R    = A[, indices],
  Rcpp = extractColumns(A, indices - 1L),
  times = 5L
)
*/

RcppEigen尝试（性能依旧不佳） #

Eigen::MatrixXd extractColumns2(Eigen::MatrixXd A, Rcpp::IntegerVector indices) {
  int m = A.rows();
  int n = indices.size();
  Eigen::MatrixXd B(m, n);
  for(int i = 0; i < n; i++) {
    B.col(i) = A.col(indices(i));
  }
  return B;
}

注：提取的索引通常为非连续。

性能较差的原因 #

1. 逐列循环的额外开销
   原代码中每次循环执行列赋值操作（B(_, i) = A(_, indices(i))或B.col(i) = A.col(indices(i))），这类操作会触发Rcpp/Eigen内部的函数调用、边界检查和单次列拷贝。当提取的列数较多时，这些零散的操作开销会被放大，远不如批量内存操作高效。

2. 原生R的底层优化
   原生R的A[, indices]是通过高度优化的底层C代码实现的，它会一次性计算所有目标元素的内存地址，然后执行批量复制，避免了循环中多次函数调用的开销，即使索引非连续也能高效处理。

3. Eigen版本的额外拷贝
   原Eigen代码中参数Eigen::MatrixXd A是值传递，会导致传入矩阵时发生一次完整拷贝，改成const Eigen::MatrixXd& A可以减少这部分开销，但这不是主要性能瓶颈。

更快的Rcpp实现方式 #

方案1：直接操作内存的批量复制 #

通过直接访问矩阵的原始内存指针，结合std::copy进行批量元素复制，避免Rcpp列赋值的语法糖开销：

#include <Rcpp.h>
#include <algorithm> // 用于std::copy
using namespace Rcpp;

// [[Rcpp::export]]
NumericMatrix extractColumnsFast(NumericMatrix A, IntegerVector indices) {
  int m = A.nrow();
  int n = indices.size();
  NumericMatrix B(m, n);
  
  const double* pA = A.begin();
  double* pB = B.begin();
  int col_stride = m; // R矩阵是列优先存储，列间元素间隔为行数
  
  for (int i = 0; i < n; ++i) {
    int src_col = indices[i];
    // 计算当前列的起始地址，复制整列元素
    const double* src_start = pA + src_col * col_stride;
    std::copy(src_start, src_start + m, pB + i * col_stride);
  }
  
  return B;
}

方案2：使用RcppArmadillo的批量列提取 #

Armadillo库原生支持通过索引向量批量提取列，底层已做优化，代码简洁且性能接近原生R：

#include <RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]

// [[Rcpp::export]]
arma::mat extractColumnsArma(const arma::mat& A, const arma::uvec& indices) {
  return A.cols(indices);
}

基准测试验证 #

使用以下代码对比各方法性能：

library(microbenchmark)
A <- matrix(rgamma(100000L, 5, 1), nrow = 2L, ncol = 50000L)
indices <- 2000L:40000L 
microbenchmark(
  R         = A[, indices],
  Rcpp_old  = extractColumns(A, indices - 1L),
  Rcpp_fast = extractColumnsFast(A, indices - 1L),
  Rcpp_arma = extractColumnsArma(A, as.integer(indices - 1L)),
  times = 10L
)

测试结果会显示，优化后的Rcpp_fast和Rcpp_arma性能与原生R相当，甚至在部分场景下更优。

内容的提问来源于stack exchange，提问作者Stéphane Laurent