[CMake 3/7] Isolate TBB construction tests from process-wide initialization

stan/math/prim/core/init_threadpool_tbb.hpp

@@ -68,6 +68,25 @@ inline int get_num_threads() {
 #endif
 }
 
+#ifdef TBB_INTERFACE_NEW
+/** Own the concurrency control before the arena, so the arena dies first. */
+struct threadpool_tbb_state {
+  tbb::global_control control;
+  tbb::task_arena arena;
+
+  explicit threadpool_tbb_state(int num_threads)
+      : control(tbb::global_control::max_allowed_parallelism, num_threads),
+        arena(num_threads, 1) {
+    arena.initialize();
+  }
+};
+
+/** Construct independently owned thread-pool state for an explicit count. */
+inline threadpool_tbb_state make_threadpool_tbb(int num_threads) {
+  return threadpool_tbb_state(num_threads);
+}
+#endif
+
 }  // namespace internal
 
 #ifdef TBB_INTERFACE_NEW
@@ -83,13 +102,12 @@ inline int get_num_threads() {
  * of get_num_threads. The TBB scheduler will be activated by calling
  * this function.
  *
- * The function returns a reference to the static
- * tbb::global_control instance.
+ * The first call fixes the concurrency configuration for the process.
  *
  * @param n_threads The maximum number of threads available to the tbb. If not
  *  set will search for the environment variable `STAN_NUM_THREADS`. A value of
  *  -1 will assume all detectable threads are available for the process.
- * @return reference to the static tbb::global_control
+ * @return reference to the static tbb::task_arena
  * @throws std::runtime_error if n_threads (defaults to zero) is not provided
  * and the value of STAN_NUM_THREADS environment variable is invalid.
  */
@@ -108,12 +126,10 @@ inline tbb::task_arena& init_threadpool_tbb(int n_threads = 0) {
                      "' 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();
+  static auto state = internal::make_threadpool_tbb(tbb_max_threads);
+  state.arena.initialize();
 
-  return tbb_arena;
+  return state.arena;
 }
 #else
 /**

test/unit/math/prim/core/get_num_threads_test.cpp

@@ -5,6 +5,7 @@
 
 #ifdef STAN_THREADS
 TEST(get_num_threads, correct_values_stan_threads) {
+  scoped_stan_num_threads restore_environment;
   set_n_threads("10");
   EXPECT_EQ(stan::math::internal::get_num_threads(), 10);
 
@@ -17,6 +18,7 @@ TEST(get_num_threads, correct_values_stan_threads) {
 }
 
 TEST(get_num_threads, incorrect_values) {
+  scoped_stan_num_threads restore_environment;
   set_n_threads("abc");
   EXPECT_THROW_MSG(stan::math::internal::get_num_threads(),
                    std::invalid_argument, "positive number or -1");
@@ -35,6 +37,7 @@ TEST(get_num_threads, incorrect_values) {
 }
 #else
 TEST(get_num_threads, correct_values_no_stan_threads) {
+  scoped_stan_num_threads restore_environment;
   set_n_threads("10");
   EXPECT_EQ(stan::math::internal::get_num_threads(), 1);
 
@@ -45,3 +48,24 @@ TEST(get_num_threads, correct_values_no_stan_threads) {
   EXPECT_EQ(stan::math::internal::get_num_threads(), 1);
 }
 #endif
+
+TEST(get_num_threads, restores_environment) {
+  scoped_stan_num_threads restore_environment;
+  set_n_threads("7");
+  {
+    scoped_stan_num_threads restore_previous;
+    set_n_threads("4");
+  }
+  EXPECT_STREQ("7", std::getenv("STAN_NUM_THREADS"));
+
+#ifdef _WIN32
+  _putenv_s("STAN_NUM_THREADS", "");
+#else
+  unsetenv("STAN_NUM_THREADS");
+#endif
+  {
+    scoped_stan_num_threads restore_absence;
+    set_n_threads("2");
+  }
+  EXPECT_EQ(nullptr, std::getenv("STAN_NUM_THREADS"));
+}

test/unit/math/prim/core/init_threadpool_tbb_late_test.cpp

@@ -7,33 +7,27 @@
 
 #ifdef TBB_INTERFACE_NEW
 
-#include <tbb/global_control.h>
-
 TEST(intel_tbb_new_late_init, check_status) {
-  const int num_threads = tbb::global_control::max_allowed_parallelism;
-
-  if (num_threads > 1) {
-    set_n_threads(num_threads - 1);
-    tbb::task_arena& tbb_arena = stan::math::init_threadpool_tbb();
-    tbb_arena.execute([&]() {
-      EXPECT_TRUE(tbb_arena.is_active());
-
-      // STAN_NUM_THREADS is not being honored if we have first
-      // initialized the TBB scheduler outside of init_threadpool_tbb
-#ifdef STAN_THREADS
-      EXPECT_EQ(num_threads, tbb::this_task_arena::max_concurrency());
-#else
-      EXPECT_EQ(num_threads, 1);
-#endif
+  const int num_threads = 4;
+  tbb::task_arena existing_arena(num_threads, 1);
+  existing_arena.execute([num_threads]() {
+    EXPECT_EQ(num_threads, tbb::this_task_arena::max_concurrency());
+    auto state = stan::math::internal::make_threadpool_tbb(num_threads - 1);
+    EXPECT_TRUE(state.arena.is_active());
+    EXPECT_EQ(num_threads - 1, state.arena.max_concurrency());
+    state.arena.execute([nested_threads = num_threads - 1]() {
+      EXPECT_EQ(nested_threads, tbb::this_task_arena::max_concurrency());
     });
-  }
+    EXPECT_EQ(num_threads, tbb::this_task_arena::max_concurrency());
+  });
 }
 
 #else
 
 #include <tbb/task_scheduler_init.h>
 
 TEST(intel_tbb_late_init, check_status) {
+  scoped_stan_num_threads restore_environment;
   const int num_threads = tbb::task_scheduler_init::default_num_threads();
   tbb::task_scheduler_init tbb_scheduler;
 

test/unit/math/prim/core/init_threadpool_tbb_test.cpp

@@ -5,19 +5,34 @@
 #include <test/unit/math/prim/functor/utils_threads.hpp>
 
 TEST(intel_tbb_new_init, check_status) {
+  scoped_stan_num_threads restore_environment;
   set_n_threads(-1);
-  auto& tbb_init = stan::math::init_threadpool_tbb();
-  EXPECT_TRUE(tbb_init.is_active());
-
-#ifdef STAN_THREADS
-  EXPECT_EQ(std::thread::hardware_concurrency(), tbb_init.max_concurrency());
-#else
-  EXPECT_EQ(1, tbb_init.max_concurrency());
+  auto& arena = stan::math::init_threadpool_tbb();
+  ASSERT_TRUE(arena.is_active());
+  const int original_threads = arena.max_concurrency();
+#ifndef STAN_THREADS
+  EXPECT_EQ(1, original_threads);
 #endif
 
-  auto& tbb_reinit = stan::math::init_threadpool_tbb();
-  EXPECT_TRUE(tbb_init.is_active());
+  // Other tests may have initialized the singleton before this test.
+  set_n_threads(original_threads + 1);
+  auto& reinit = stan::math::init_threadpool_tbb();
+  EXPECT_EQ(&arena, &reinit);
+  EXPECT_TRUE(reinit.is_active());
+  EXPECT_EQ(original_threads, reinit.max_concurrency());
+  auto& explicit_reinit = stan::math::init_threadpool_tbb(original_threads + 2);
+  EXPECT_EQ(&arena, &explicit_reinit);
+  EXPECT_EQ(original_threads, explicit_reinit.max_concurrency());
+}
 
-  tbb_init.terminate();
-  EXPECT_FALSE(tbb_init.is_active());
+#ifdef STAN_THREADS
+TEST(intel_tbb_new_init, validates_after_initialization) {
+  scoped_stan_num_threads restore_environment;
+  stan::math::init_threadpool_tbb(2);
+  EXPECT_THROW_MSG(stan::math::init_threadpool_tbb(-2), std::invalid_argument,
+                   "positive or -1");
+  set_n_threads("invalid");
+  EXPECT_THROW_MSG(stan::math::init_threadpool_tbb(), std::invalid_argument,
+                   "positive number or -1");
 }
+#endif

test/unit/math/prim/core/set_num_threads_test.cpp

@@ -1,14 +1,22 @@
 #include <stan/math/prim/core.hpp>
 #include <gtest/gtest.h>
+#include <algorithm>
 
 TEST(intel_tbb_new_init, set_threads) {
-  int num_threads = std::thread::hardware_concurrency();
-  auto& tbb_init = stan::math::init_threadpool_tbb(num_threads - 1);
-  EXPECT_TRUE(tbb_init.is_active());
-
-#ifdef STAN_THREADS
-  EXPECT_EQ(num_threads - 1, tbb_init.max_concurrency());
-#else
-  EXPECT_EQ(1, tbb_init.max_concurrency());
-#endif
+  using control = tbb::global_control;
+  const auto previous = control::active_value(control::max_allowed_parallelism);
+  for (int num_threads : {1, 2, 4}) {
+    {
+      auto state = stan::math::internal::make_threadpool_tbb(num_threads);
+      EXPECT_TRUE(state.arena.is_active());
+      EXPECT_EQ(num_threads, state.arena.max_concurrency());
+      EXPECT_EQ(std::min(previous, static_cast<std::size_t>(num_threads)),
+                control::active_value(control::max_allowed_parallelism));
+      state.arena.execute([num_threads]() {
+        EXPECT_EQ(num_threads, tbb::this_task_arena::max_concurrency());
+      });
+    }
+    EXPECT_EQ(previous,
+              control::active_value(control::max_allowed_parallelism));
+  }
 }

test/unit/math/prim/functor/utils_threads.hpp

@@ -1,9 +1,33 @@
 #ifndef TEST_UNIT_MATH_PRIM_FUNCTOR_UTILS_THREADS_HPP
 #define TEST_UNIT_MATH_PRIM_FUNCTOR_UTILS_THREADS_HPP
 
-#include <stdlib.h>
+#include <cstdlib>
+#include <cstdio>
 #include <string>
 
+// Restore the caller's environment, including an originally absent variable.
+class scoped_stan_num_threads {
+  const bool was_set_;
+  const std::string previous_;
+
+ public:
+  scoped_stan_num_threads()
+      : was_set_(std::getenv("STAN_NUM_THREADS") != nullptr),
+        previous_(was_set_ ? std::getenv("STAN_NUM_THREADS") : "") {}
+
+  ~scoped_stan_num_threads() {
+#ifdef _WIN32
+    _putenv_s("STAN_NUM_THREADS", previous_.c_str());
+#else
+    if (was_set_) {
+      setenv("STAN_NUM_THREADS", previous_.c_str(), 1);
+    } else {
+      unsetenv("STAN_NUM_THREADS");
+    }
+#endif
+  }
+};
+
 // utility to set number of threads to use
 inline void set_n_threads(int num_threads) {
   static char env_string[256];

test/unit/math/rev/functor/map_rect_concurrent_prim_test.cpp

@@ -19,6 +19,7 @@ STAN_REGISTER_MAP_RECT(0, hard_work)
 STAN_REGISTER_MAP_RECT(1, hard_work)
 
 struct map_rect_con_prim : public ::testing::Test {
+  scoped_stan_num_threads restore_environment;
   Eigen::VectorXd shared_params_d;
   std::vector<Eigen::VectorXd> job_params_d;
   std::vector<std::vector<double> > x_r;

test/unit/math/rev/functor/map_rect_concurrent_test.cpp

@@ -18,6 +18,7 @@
 STAN_REGISTER_MAP_RECT(0, hard_work)
 
 struct map_rect_con : public ::testing::Test {
+  scoped_stan_num_threads restore_environment;
   Eigen::VectorXd shared_params_d;
   std::vector<Eigen::VectorXd> job_params_d;
   std::vector<std::vector<double> > x_r;

test/unit/math/rev/functor/map_rect_concurrent_threads_test.cpp

@@ -46,6 +46,7 @@ inline void setup_job(int N, Eigen::VectorXd& shared_params_d,
 }
 
 struct map_rect_con_threads : public ::testing::Test {
+  scoped_stan_num_threads restore_environment;
   Eigen::VectorXd shared_params_d;
   std::vector<Eigen::VectorXd> job_params_d;
   std::vector<std::vector<double> > x_r;