init_threadpool_tbb.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_CORE_INIT_THREADPOOL_TBB_HPP
#define STAN_MATH_PRIM_CORE_INIT_THREADPOOL_TBB_HPP
 
#include <stan/math/prim/err/invalid_argument.hpp>
 
#include <boost/lexical_cast.hpp>
 
#ifndef TBB_INTERFACE_NEW
#include <tbb/tbb_stddef.h>
 
#if TBB_VERSION_MAJOR >= 2020
#define TBB_INTERFACE_NEW
#endif
#endif
 
#ifdef TBB_INTERFACE_NEW
#include <tbb/global_control.h>
#include <tbb/task_arena.h>
#else
#include <tbb/task_scheduler_init.h>
#endif
 
#include <cstdlib>
#include <thread>
 
namespace stan {
namespace math {
namespace internal {
 
inline int get_num_threads() {
  int num_threads = 1;
#ifdef STAN_THREADS
  const char* env_stan_num_threads = std::getenv("STAN_NUM_THREADS");
  if (env_stan_num_threads != nullptr) {
    try {
      const int env_num_threads
          = boost::lexical_cast<int>(env_stan_num_threads);
      if (env_num_threads > 0) {
        num_threads = env_num_threads;
      } else if (env_num_threads == -1) {
        num_threads = std::thread::hardware_concurrency();
      } else {
        invalid_argument("get_num_threads(int)", "STAN_NUM_THREADS",
                         env_stan_num_threads,
                         "The STAN_NUM_THREADS environment variable is '",
                         "' but it must be positive or -1");
      }
    } catch (const boost::bad_lexical_cast&) {
      invalid_argument("get_num_threads(int)", "STAN_NUM_THREADS",
                       env_stan_num_threads,
                       "The STAN_NUM_THREADS environment variable is '",
                       "' but it must be a positive number or -1");
    }
  }
#endif
  return num_threads;
}
 
}  // namespace internal
 
#ifdef TBB_INTERFACE_NEW
inline tbb::task_arena& init_threadpool_tbb(int n_threads = 0) {
  int tbb_max_threads = 1;
#ifdef STAN_THREADS
  if (n_threads == 0) {
    tbb_max_threads = internal::get_num_threads();
  } else if (n_threads > 0) {
    tbb_max_threads = n_threads;
  } else if (n_threads == -1) {
    tbb_max_threads = std::thread::hardware_concurrency();
  } else {
    invalid_argument("init_threadpool_tbb(int)", "n_threads", n_threads,
                     "The number of threads is '",
                     "' but it must be positive or -1");
  }
#endif
  static tbb::global_control tbb_gc(
      tbb::global_control::max_allowed_parallelism, tbb_max_threads);
  static tbb::task_arena tbb_arena(tbb_max_threads, 1);
  tbb_arena.initialize();
 
  return tbb_arena;
}
#else
inline tbb::task_scheduler_init& init_threadpool_tbb(int n_threads = 0) {
  int tbb_max_threads = 1;
#ifdef STAN_THREADS
  if (n_threads == 0) {
    tbb_max_threads = internal::get_num_threads();
  } else if (n_threads > 0) {
    tbb_max_threads = n_threads;
  } else if (n_threads == -1) {
    tbb_max_threads = std::thread::hardware_concurrency();
  } else {
    invalid_argument("init_threadpool_tbb(int)", "n_threads", n_threads,
                     "The number of threads is '",
                     "' but it must be positive or -1");
  }
#endif
  static tbb::task_scheduler_init tbb_scheduler(tbb_max_threads, 0);
  return tbb_scheduler;
}
#endif
 
}  // namespace math
}  // namespace stan
 
#endif