Fix the new TBB interface and use tbb::task_arena

Signed-off-by: Hamada S. Badr <hamada.s.badr@gmail.com>
stan-dev · Dec 15, 2020 · 861ce93 · 861ce93
1 parent 343a3b7
commit 861ce93
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 13 deletions.
diff --git a/make/compiler_flags b/make/compiler_flags
@@ -247,7 +247,7 @@ else
 CXXFLAGS_TBB ?= -I $(TBB)/include
 endif
 
-LDFLAGS_TBB ?= -Wl,-L,"$(TBB_LIB)" -Wl,-rpath,"$(TBB_LIB) -ltbb"
+LDFLAGS_TBB ?= -Wl,-L,"$(TBB_LIB)" -Wl,-rpath,"$(TBB_LIB)" -ltbb
 LDLIBS_TBB ?=
 
 else

diff --git a/stan/math/prim/core/init_threadpool_tbb.hpp b/stan/math/prim/core/init_threadpool_tbb.hpp
@@ -67,7 +67,7 @@ inline int get_num_threads() {
 #ifdef TBB_INTERFACE_NEW
 /**
  * Initialize the Intel TBB threadpool and global scheduler through
- * the tbb::global_control object. In case an instance of the
+ * the tbb::task_arena object. In case an instance of the
  * tbb::task_scheduler_object has been instantiated prior to calling
  * this function, then any subsequent initialization is ignored by the
  * Intel TBB.
@@ -84,13 +84,14 @@ inline int get_num_threads() {
  * @throws std::runtime_error if the value of STAN_NUM_THREADS env. variable
  * is invalid
  */
-inline tbb::global_control& init_threadpool_tbb() {
+inline tbb::task_arena& init_threadpool_tbb() {
   int tbb_max_threads = internal::get_num_threads();
 
-  static tbb::global_control tbb_scheduler(
-      tbb::global_control::max_allowed_parallelism, tbb_max_threads);
+  static tbb::task_arena tbb_arena(tbb_max_threads, 1,
+                                   tbb::task_arena::priority::normal);
+  tbb_arena.initialize();
 
-  return tbb_scheduler;
+  return tbb_arena;
 }
 #else
 /**

diff --git a/test/unit/math/prim/core/init_threadpool_tbb_late_test.cpp b/test/unit/math/prim/core/init_threadpool_tbb_late_test.cpp
@@ -10,12 +10,12 @@
 #include <tbb/task_arena.h>
 
 TEST(intel_tbb_new_late_init, check_status) {
-  const int num_threads = tbb::global_control::default_num_threads();
-  tbb::global_control tbb_scheduler;
+  const int num_threads = tbb::global_control::max_allowed_parallelism;
+  tbb::task_arena tbb_arena;
 
   if (num_threads > 1) {
     set_n_threads(num_threads - 1);
-    tbb::global_control& tbb_init = stan::math::init_threadpool_tbb();
+    tbb::task_arena& tbb_init = stan::math::init_threadpool_tbb();
     EXPECT_TRUE(tbb_init.is_active());
 
     // STAN_NUM_THREADS is not being honored if we have first

diff --git a/test/unit/math/prim/core/init_threadpool_tbb_test.cpp b/test/unit/math/prim/core/init_threadpool_tbb_test.cpp
@@ -11,13 +11,12 @@
 
 TEST(intel_tbb_new_init, check_status) {
   set_n_threads(-1);
-  tbb::global_control& tbb_init = stan::math::init_threadpool_tbb();
+  tbb::task_arena& tbb_init = stan::math::init_threadpool_tbb();
   EXPECT_TRUE(tbb_init.is_active());
 
-  EXPECT_EQ(std::thread::hardware_concurrency(),
-            tbb::this_task_arena::max_concurrency());
+  EXPECT_EQ(std::thread::hardware_concurrency(), tbb_init.max_concurrency());
 
-  tbb::global_control& tbb_reinit = stan::math::init_threadpool_tbb();
+  tbb::task_arena& tbb_reinit = stan::math::init_threadpool_tbb();
   EXPECT_TRUE(tbb_init.is_active());
 
   tbb_init.terminate();