Reusing existing memory for BLAS operations in R - r

Question

I have an operation inside tight loop in R that I need to optimize. It's updating the weights inside an IRLS algorithm by calculating the Schur product of a vector and a matrix. That is, it multiplies each element in the matrix by the corresponding row value in the vector, producing a result of the same dimensions as the matrix. In overly simplified schematic form, it looks like this:
reweight = function(iter, w, Q) {
for (i in 1:iter) {
wT = w * Q
}
}
In normal R code, a new matrix of dim() [rows,cols] is created on each iteration:
cols = 1000
rows = 1000000
w = runif(rows)
Q = matrix(1.0, rows, cols)
Rprofmem()
reweight(5, w, Q)
Rprofmem(NULL)
nate#ubuntu:~/R$ less Rprofmem.out
8000000040 :"reweight"
8000000040 :"reweight"
8000000040 :"reweight"
8000000040 :"reweight"
8000000040 :"reweight"
And if the matrix is large (multiple GB), the cost of the memory allocation exceeds the time spent on the numeric operation:
nate#ubuntu:~/R$ perf record -p `pgrep R` sleep 5 && perf report
49.93% R [kernel.kallsyms] [k] clear_page_c_e
47.67% R libR.so [.] real_binary
0.57% R [kernel.kallsyms] [k] get_page_from_freelist
0.35% R [kernel.kallsyms] [k] clear_huge_page
0.34% R libR.so [.] RunGenCollect
0.20% R [kernel.kallsyms] [k] clear_page
It also consumes a lot of memory:
USER PID VSZ RSS COMMAND
nate 17099 22.5GB 22.5GB /usr/local/lib/R/bin/exec/R --vanilla
If the matrix is smaller (several MB) but the number of iterations is larger, the memory usage is more reasonable, but at the cost of the garbage collector using more time than the numeric calculations:
cols = 100
rows = 10000
w = runif(rows)
Q = matrix(1.0, rows, cols)
reweight(1000, w, Q)
(note that this is a new process starting from scratch)
61.51% R libR.so [.] RunGenCollect
26.40% R libR.so [.] real_binary
7.94% R libR.so [.] SortNodes
2.79% R [kernel.kallsyms] [k] clear_page_c_e
USER PID VSZ RSS COMMAND
nate 17099 191MB 72MB /usr/local/lib/R/bin/exec/R --vanilla
If I write my own function with Rcpp that does the work in place, I can get the memory allocation that I want:
library(Rcpp)
cppFunction('
void weightMatrix(NumericVector w,
NumericMatrix Q,
NumericMatrix wQ) {
size_t numRows = Q.rows();
for (size_t row = 0; row < numRows; row++) {
wQ(row,_) = w(row) * Q(row,_);
}
return;
}
')
reweightCPP = function(iter, w, Q) {
# Initialize workspace to non-NA
wQ = matrix(1.0, nrow(Q), ncol(Q))
for (i in 1:iter) {
weightMatrix(w, Q, wQ)
}
}
cols = 100
rows = 10000
w = runif(rows)
Q = matrix(1.0, rows, cols)
wQ = matrix(NA, rows, cols)
Rprofmem()
reweightCPP(5, w, Q)
Rprofmem(NULL)
nate#ubuntu:~/R$ less Rprofmem.out
8000040 :"matrix" "reweightCPP"
2544 :"<Anonymous>" "weightMatrix" "reweightCPP"
2544 :"<Anonymous>" "weightMatrix" "reweightCPP"
2544 :"<Anonymous>" "weightMatrix" "reweightCPP"
2544 :"<Anonymous>" "weightMatrix" "reweightCPP"
2544 :"<Anonymous>" "weightMatrix" "reweightCPP"
(What's the 2544 bytes of allocation for? It seems to be an Rcpp constant. Is there any way I can avoid it?)
Performance is still suboptimal due to the Rcpp sugar:
76.53% R sourceCpp_82335.so [.] _Z12weightMatrixN4Rcpp6VectorILi14ENS_15PreserveStorageEEENS_6MatrixILi14ES1_EES4_
10.46% R libR.so [.] Rf_getAttrib
9.53% R libR.so [.] getAttrib0
2.06% R libR.so [.] Rf_isMatrix
0.42% R libR.so [.] INTEGER
But I can mostly fix that by resorting to lower level C++:
cppFunction('
void weightMatrix(NumericVector w_,
NumericMatrix Q_,
NumericMatrix wQ_) {
size_t numCols = Q_.ncol();
size_t numRows = Q_.nrow();
double * __restrict__ w = &w_[0];
double * __restrict__ Q = &Q_[0];
double * __restrict__ wQ = &wQ_[0];
for (size_t row = 0; row < numRows; row++) {
size_t colOffset = 0;
for (size_t col = 0; col < numCols; col++) {
wQ[colOffset + row] = w[row] * Q[colOffset + row];
colOffset += numRows;
}
}
return;
}
')
99.18% R sourceCpp_59392.so [.] sourceCpp_48203_weightMatrix
0.06% R libR.so [.] PutRNGstate
0.06% R libR.so [.] do_begin
0.06% R libR.so [.] Rf_eval
That said, I still haven't figured out to get the compiler to reliably generate efficient assembly without resorting to using SIMD intrinsics to force the use of VMULPD. Even with the ugly '__restrict__' attributes, in the form shown here it seems compelled to invert my loop order and do a lot of unnecessary work. But presumably I'll find the magic cross-compiler syntax eventually, or more likely, call out to a Fortran BLAS function.
Which brings me to my questions:
Is there any way that I can get the performance I want without going to all this trouble? Failing that, is there any way that I can at least hide it behind the scenes so that the end user in R can use "wQ = w * Q" and have it magically reuse wQ instead of allocating and throwing away another giant matrix?
The BLAS wrappers in R seem to do a fairly good job for cases where the answer can be written into one of the operands (Q = w * Q), but I haven't found any way to do this when I need a "3rd party" workspace. Is there maybe some reasonable way to define a method for %=% that will convert "wQ = w * Q" to "op_mult(w, Q, wQ)"?
To preempt the question as to whether it matters: yes, I've measured, and it matters. The use case is an ensemble of cross-validated logistic regressions inside a loop handling large arrays of longitudinal data (http://cran.r-project.org/web/packages/ltmle/ltmle.pdf). It will be called millions (if not billions) of times per analysis. A good optimization of this function would help to get the runtime from "impossible" down to "days". A great optimization (or rather the combination of several such optimizations) might get it down to "hours" or even "minutes".
Edit: In the comments, Henrik correctly points out that the example loop has been simplified to the point that it simply repeats the same calculation multiple times. I hoped this would focus the issue, but perhaps it confuses it. In the real version, there will be more steps in the loop such that the 'w' in the 'w * Q' is different each iteration. Below is a poorly tested draft version of the actual functions. This one is a "semi-optimized" logistic regression in straight R based on O'Leary's QR Newton IRLS described by Bryan Lewis.
logistic_irls_qrnewton = function(A, y, maxIter=25, targetSSE=1e-16) {
# warn user below on first weight less than threshold
tinyWeightsFound = FALSE
tiny = sqrt(.Machine$double.eps)
# decompose A to QR (only once, Choleski done in loop)
QR = qr(A) # A[rows=samples, cols=covariates]
Q = qr.Q(QR) # Q[rows, cols] (same dimensions as A)
R = qr.R(QR) # R[cols, cols] (upper right triangular)
# copying now prevents copying each time y is used as argument
y = y + 0; # y[rows]
# first pass is outside loop since initial values are constant
iter = 1
t = (y - 0.5) * 4.0 # t[rows] = (y - m) * initial weight
C = chol(crossprod(Q, Q)) # C[rows, rows]
t = crossprod(Q,t)
s = forwardsolve(t(C), t) # s[cols]
s = backsolve(C, s))
t = Q %*% s
sse = crossprod(s) # sum of squared errors
print(as.vector(sse))
converged = ifelse(sse < targetSSE, 1, 0)
while (converged == 0 && iter < maxIter) {
iter = iter + 1
# only t is required as an input
dim(t) = NULL # matrix to vector to counteract crossprod
e = exp(t)
m = e / (e + 1) # mu = exp(eta) / (1 + exp(eta))
d = m / (e + 1) # mu.eta = exp(eta) / (1 + exp(eta))^2
w = d * d / (m - m^2) # W = (1 / variance) = 1 / (mu * (1 - mu))
if(tinyWeightsFound == FALSE && min(w) < tiny) {
print("Tiny weights found")
tinyWeightsFound = TRUE
}
t = crossprod(Q, w * (((y - m) / d) + t))
C = chol(crossprod(Q, w * Q))
n = forwardsolve(t(C), t)
n = backsolve(C, n)
t = Q %*% n
sse = crossprod(n - s) # divergence from previous
s = n # save divergence for difference from next
print(as.vector(sse))
if (sse < targetSSE) converged = iter
}
if (converged == 0) {
print(paste("Failed to converge after", iter, "iterations"))
print(paste("Final SSE was", sse))
} else {
print(paste("Convergence after iteration", iter))
}
coefficients = backsolve(R, crossprod(Q,t))
dim(coefficients) = NULL # return as a vector
coefficients
}