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>
 
#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 <charconv>
#include <cstdlib>
#include <string_view>
#include <thread>
 
namespace stan {
namespace math {
namespace internal {
 
inline int get_num_threads() {
#ifdef STAN_THREADS
  const char* env_stan_num_threads = std::getenv("STAN_NUM_THREADS");
  if (env_stan_num_threads == nullptr) {
    return 1;
  }
 
  const std::string_view value(env_stan_num_threads);
  int num_threads;
  const auto [end, error]
      = std::from_chars(value.begin(), value.end(), num_threads);
  if (error != std::errc() || end != value.end()
      || (num_threads < 1 && num_threads != -1)) {
    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");
  }
 
  return num_threads == -1 ? std::thread::hardware_concurrency() : num_threads;
#else
  return 1;
#endif
}
 
}  // 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