Support throttling and resetting thread use via OpenMP functions

ChangeLog

@@ -1,3 +1,19 @@
+2023-10-29  Dirk Eddelbuettel  <[email protected]>
+
+	* src/RcppArmadillo.cpp (armadillo_get_number_of_omp_threads)
+	(armadillo_set_number_of_omp_threads): New helper functions
+        * src/RcppExports.cpp: Regenerated
+        * R/RcppExports.R: Idem
+	* man/armadillo_get_number_of_omp_threads.Rd: Documentation
+
+	* R/init.R (.onLoad): Store initial thread count
+	* R/init.R (armadillo_throttle_cores, armadillo_reset_cores):
+	Tread throtte and reset helper functions
+	* man/armadillo_throttle_cores.Rd: Documentation
+
+	* man/fastLm.Rd: Illustration of use of throttle and reset function
+	* NAMESPACE: Export new functions
+
 2023-10-14  Dirk Eddelbuettel  <[email protected]>
 
 	* DESCRIPTION (Version, Date): RcppArmadillo 0.12.6.5.0

DESCRIPTION

@@ -26,3 +26,4 @@ Imports: Rcpp (>= 0.11.0), stats, utils, methods
 Suggests: tinytest, Matrix (>= 1.3.0), pkgKitten, reticulate, slam
 URL: https://github.com/RcppCore/RcppArmadillo, https://dirk.eddelbuettel.com/code/rcpp.armadillo.html
 BugReports: https://github.com/RcppCore/RcppArmadillo/issues
+RoxygenNote: 6.0.1

NAMESPACE

@@ -10,12 +10,16 @@ export("fastLmPure",
        "RcppArmadillo.package.skeleton",
        "armadillo_version",
        "armadillo_set_seed",
-       "armadillo_set_seed_random")
+       "armadillo_set_seed_random",
+
+       "armadillo_throttle_cores",
+       "armadillo_reset_cores",
+       "armadillo_set_number_of_omp_threads",
+       "armadillo_set_number_of_omp_threads"
+       )
 S3method("fastLm", "default")
 S3method("fastLm", "formula")
 S3method("predict", "fastLm")
 S3method("print", "fastLm")
 S3method("summary", "fastLm")
 S3method("print", "summary.fastLm")
-
-

R/RcppExports.R

@@ -53,6 +53,21 @@ armadillo_set_seed <- function(val) {
     invisible(.Call(`_RcppArmadillo_armadillo_set_seed`, val))
 }
 
+#' Report (or Set) Maximum Number of OpenMP Threads
+#'
+#' @param n Number of threads to be set
+#' @return For the getter, and on a system with OpenMP, the maximum
+#' number of threads that OpenMP may be using and on systems without it,
+#' one. The setter does not return a value.
+armadillo_get_number_of_omp_threads <- function() {
+    .Call(`_RcppArmadillo_armadillo_get_number_of_omp_threads`)
+}
+
+#' @rdname armadillo_get_number_of_omp_threads
+armadillo_set_number_of_omp_threads <- function(n) {
+    invisible(.Call(`_RcppArmadillo_armadillo_set_number_of_omp_threads`, n))
+}
+
 fastLm_impl <- function(X, y) {
     .Call(`_RcppArmadillo_fastLm_impl`, X, y)
 }

R/init.R

@@ -0,0 +1,40 @@
+## init.R: Startup
+##
+## Copyright (C)  2023  Dirk Eddelbuettel
+##
+## This file is part of RcppArmadillo.
+##
+## RcppArmadillo is free software: you can redistribute it and/or modify it
+## under the terms of the GNU General Public License as published by
+## the Free Software Foundation, either version 2 of the License, or
+## (at your option) any later version.
+##
+## RcppArmadillo is distributed in the hope that it will be useful, but
+## WITHOUT ANY WARRANTY; without even the implied warranty of
+## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+## GNU General Public License for more details.
+##
+## You should have received a copy of the GNU General Public License
+## along with RcppArmadillo.  If not, see <http://www.gnu.org/licenses/>.
+
+.pkgenv <- new.env(parent=emptyenv())
+
+.onLoad <- function(libname, pkgname) {
+    .pkgenv[["omp_threads"]] <- armadillo_get_number_of_omp_threads()
+}
+
+##' Throttle (or Reset) (Rcpp)Armadillo to Two Cores
+##'
+##' Helper functions to throttle use of cores by RcppArmadillo-internal
+##' code on systems with OpenMP. On package load, the initial value is
+##' saved and used to reset the value.
+##' @param n Integer value of desired cores, default is two
+armadillo_throttle_cores <- function(n = 2) {
+    armadillo_set_number_of_omp_threads(n)
+}
+
+##' @rdname armadillo_throttle_cores
+armadillo_reset_cores <- function() {
+    n <- .pkgenv[["omp_threads"]]
+    armadillo_set_number_of_omp_threads(n)
+}

man/fastLm.Rd

@@ -69,6 +69,11 @@ fastLm(X, \dots)
   Romain Francois, Dirk Eddelbuettel, Douglas Bates and Binxiang Ni.
 }
 \examples{
+  \dontshow{
+      ## as an illustration, the example is computationally inexpensive
+      ## and does not require this per se
+      armadillo_throttle_cores(2)
+  }
   data(trees, package="datasets")
 
   ## bare-bones direct interface
@@ -89,5 +94,7 @@ fastLm(X, \dots)
   dd$y <- mm \%*\% seq_len(ncol(mm)) + rnorm(nrow(mm), sd = 0.1)
   summary(lm(y ~ f1 * f2, dd))     # detects rank deficiency
   summary(fastLm(y ~ f1 * f2, dd)) # some huge coefficients
+
+  \dontshow{armadillo_reset_cores()}
 }
 \keyword{regression}

src/RcppArmadillo.cpp

@@ -1,7 +1,7 @@
 
 // RcppArmadillo.cpp: Rcpp/Armadillo glue
 //
-// Copyright (C)  2010 - 2020  Dirk Eddelbuettel, Romain Francois and Douglas Bates
+// Copyright (C)  2010 - 2023  Dirk Eddelbuettel, Romain Francois and Douglas Bates
 //
 // This file is part of RcppArmadillo.
 //
@@ -90,3 +90,26 @@ void armadillo_set_seed(unsigned int val) {
     //Rcpp::Rcout << "Setting value " << val << std::endl;
     arma::arma_rng::set_seed(val);  			    // set the seed to given value
 }
+
+//' Report (or Set) Maximum Number of OpenMP Threads
+//'
+//' @param n Number of threads to be set
+//' @return For the getter, and on a system with OpenMP, the maximum
+//' number of threads that OpenMP may be using and on systems without it,
+//' one. The setter does not return a value.
+// [[Rcpp::export]]
+int armadillo_get_number_of_omp_threads() {
+#ifdef _OPENMP
+    return omp_get_max_threads();
+#else
+    return 1;
+#endif
+}
+
+//' @rdname armadillo_get_number_of_omp_threads
+// [[Rcpp::export]]
+void armadillo_set_number_of_omp_threads(int n) {
+#ifdef _OPENMP
+    omp_set_num_threads(n);
+#endif
+}

src/RcppExports.cpp

@@ -41,6 +41,26 @@ BEGIN_RCPP
     return R_NilValue;
 END_RCPP
 }
+// armadillo_get_number_of_omp_threads
+int armadillo_get_number_of_omp_threads();
+RcppExport SEXP _RcppArmadillo_armadillo_get_number_of_omp_threads() {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    rcpp_result_gen = Rcpp::wrap(armadillo_get_number_of_omp_threads());
+    return rcpp_result_gen;
+END_RCPP
+}
+// armadillo_set_number_of_omp_threads
+void armadillo_set_number_of_omp_threads(int n);
+RcppExport SEXP _RcppArmadillo_armadillo_set_number_of_omp_threads(SEXP nSEXP) {
+BEGIN_RCPP
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< int >::type n(nSEXP);
+    armadillo_set_number_of_omp_threads(n);
+    return R_NilValue;
+END_RCPP
+}
 // fastLm_impl
 Rcpp::List fastLm_impl(const arma::mat& X, const arma::colvec& y);
 RcppExport SEXP _RcppArmadillo_fastLm_impl(SEXP XSEXP, SEXP ySEXP) {
@@ -58,6 +78,8 @@ static const R_CallMethodDef CallEntries[] = {
     {"_RcppArmadillo_armadillo_version", (DL_FUNC) &_RcppArmadillo_armadillo_version, 1},
     {"_RcppArmadillo_armadillo_set_seed_random", (DL_FUNC) &_RcppArmadillo_armadillo_set_seed_random, 0},
     {"_RcppArmadillo_armadillo_set_seed", (DL_FUNC) &_RcppArmadillo_armadillo_set_seed, 1},
+    {"_RcppArmadillo_armadillo_get_number_of_omp_threads", (DL_FUNC) &_RcppArmadillo_armadillo_get_number_of_omp_threads, 0},
+    {"_RcppArmadillo_armadillo_set_number_of_omp_threads", (DL_FUNC) &_RcppArmadillo_armadillo_set_number_of_omp_threads, 1},
     {"_RcppArmadillo_fastLm_impl", (DL_FUNC) &_RcppArmadillo_fastLm_impl, 2},
     {NULL, NULL, 0}
 };