Automatic Differentiation
 
Loading...
Searching...
No Matches
init_threadpool_tbb.hpp
Go to the documentation of this file.
1#ifndef STAN_MATH_PRIM_CORE_INIT_THREADPOOL_TBB_HPP
2#define STAN_MATH_PRIM_CORE_INIT_THREADPOOL_TBB_HPP
3
5
6#include <boost/lexical_cast.hpp>
7
8#ifndef TBB_INTERFACE_NEW
9#include <tbb/tbb_stddef.h>
10
11#if TBB_VERSION_MAJOR >= 2020
12#define TBB_INTERFACE_NEW
13#endif
14#endif
15
16#ifdef TBB_INTERFACE_NEW
17#include <tbb/global_control.h>
18#include <tbb/task_arena.h>
19#else
20#include <tbb/task_scheduler_init.h>
21#endif
22
23#include <cstdlib>
24#include <thread>
25
26namespace stan {
27namespace math {
28namespace internal {
29
46inline int get_num_threads() {
47 int num_threads = 1;
48#ifdef STAN_THREADS
49 const char* env_stan_num_threads = std::getenv("STAN_NUM_THREADS");
50 if (env_stan_num_threads != nullptr) {
51 try {
52 const int env_num_threads
53 = boost::lexical_cast<int>(env_stan_num_threads);
54 if (env_num_threads > 0) {
55 num_threads = env_num_threads;
56 } else if (env_num_threads == -1) {
57 num_threads = std::thread::hardware_concurrency();
58 } else {
59 invalid_argument("get_num_threads(int)", "STAN_NUM_THREADS",
60 env_stan_num_threads,
61 "The STAN_NUM_THREADS environment variable is '",
62 "' but it must be positive or -1");
63 }
64 } catch (const boost::bad_lexical_cast&) {
65 invalid_argument("get_num_threads(int)", "STAN_NUM_THREADS",
66 env_stan_num_threads,
67 "The STAN_NUM_THREADS environment variable is '",
68 "' but it must be a positive number or -1");
69 }
70 }
71#endif
72 return num_threads;
73}
74
75} // namespace internal
76
77#ifdef TBB_INTERFACE_NEW
100inline tbb::task_arena& init_threadpool_tbb(int n_threads = 0) {
101 int tbb_max_threads = 1;
102#ifdef STAN_THREADS
103 if (n_threads == 0) {
104 tbb_max_threads = internal::get_num_threads();
105 } else if (n_threads > 0) {
106 tbb_max_threads = n_threads;
107 } else if (n_threads == -1) {
108 tbb_max_threads = std::thread::hardware_concurrency();
109 } else {
110 invalid_argument("init_threadpool_tbb(int)", "n_threads", n_threads,
111 "The number of threads is '",
112 "' but it must be positive or -1");
113 }
114#endif
115 static tbb::global_control tbb_gc(
116 tbb::global_control::max_allowed_parallelism, tbb_max_threads);
117 static tbb::task_arena tbb_arena(tbb_max_threads, 1);
118 tbb_arena.initialize();
119
120 return tbb_arena;
121}
122#else
145inline tbb::task_scheduler_init& init_threadpool_tbb(int n_threads = 0) {
146 int tbb_max_threads = 1;
147#ifdef STAN_THREADS
148 if (n_threads == 0) {
149 tbb_max_threads = internal::get_num_threads();
150 } else if (n_threads > 0) {
151 tbb_max_threads = n_threads;
152 } else if (n_threads == -1) {
153 tbb_max_threads = std::thread::hardware_concurrency();
154 } else {
155 invalid_argument("init_threadpool_tbb(int)", "n_threads", n_threads,
156 "The number of threads is '",
157 "' but it must be positive or -1");
158 }
159#endif
160 static tbb::task_scheduler_init tbb_scheduler(tbb_max_threads, 0);
161 return tbb_scheduler;
162}
163#endif
164
165} // namespace math
166} // namespace stan
167
168#endif
int get_num_threads()
Get number of threads to use.
tbb::task_arena & init_threadpool_tbb(int n_threads=0)
Initialize the Intel TBB threadpool and global scheduler through the tbb::task_arena object.
void invalid_argument(const char *function, const char *name, const T &y, const char *msg1, const char *msg2)
Throw an invalid_argument exception with a consistently formatted message.
The lgamma implementation in stan-math is based on either the reentrant safe lgamma_r implementation ...