diff --git a/stan/math/prim/core/init_threadpool_tbb.hpp b/stan/math/prim/core/init_threadpool_tbb.hpp index ca96dc3a8b2..99e4f361a56 100644 --- a/stan/math/prim/core/init_threadpool_tbb.hpp +++ b/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 /** diff --git a/test/unit/math/prim/core/get_num_threads_test.cpp b/test/unit/math/prim/core/get_num_threads_test.cpp index 5b86e308c89..c65e2d05a49 100644 --- a/test/unit/math/prim/core/get_num_threads_test.cpp +++ b/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")); +} 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 index 4cefd6b96c0..46bfd4b3cc2 100644 --- a/test/unit/math/prim/core/init_threadpool_tbb_late_test.cpp +++ b/test/unit/math/prim/core/init_threadpool_tbb_late_test.cpp @@ -7,26 +7,19 @@ #ifdef TBB_INTERFACE_NEW -#include - 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 @@ -34,6 +27,7 @@ TEST(intel_tbb_new_late_init, check_status) { #include 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; diff --git a/test/unit/math/prim/core/init_threadpool_tbb_test.cpp b/test/unit/math/prim/core/init_threadpool_tbb_test.cpp index b20adbf2e2e..b9522099cc8 100644 --- a/test/unit/math/prim/core/init_threadpool_tbb_test.cpp +++ b/test/unit/math/prim/core/init_threadpool_tbb_test.cpp @@ -5,19 +5,34 @@ #include 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 diff --git a/test/unit/math/prim/core/set_num_threads_test.cpp b/test/unit/math/prim/core/set_num_threads_test.cpp index f46eeec0760..e6d2c1758bc 100644 --- a/test/unit/math/prim/core/set_num_threads_test.cpp +++ b/test/unit/math/prim/core/set_num_threads_test.cpp @@ -1,14 +1,22 @@ #include #include +#include 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(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)); + } } diff --git a/test/unit/math/prim/functor/utils_threads.hpp b/test/unit/math/prim/functor/utils_threads.hpp index 77de14ffeb2..83187bfd439 100644 --- a/test/unit/math/prim/functor/utils_threads.hpp +++ b/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 +#include +#include #include +// 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]; diff --git a/test/unit/math/rev/functor/map_rect_concurrent_prim_test.cpp b/test/unit/math/rev/functor/map_rect_concurrent_prim_test.cpp index b7fe369ae98..fcf2f8378c6 100644 --- a/test/unit/math/rev/functor/map_rect_concurrent_prim_test.cpp +++ b/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 job_params_d; std::vector > x_r; diff --git a/test/unit/math/rev/functor/map_rect_concurrent_test.cpp b/test/unit/math/rev/functor/map_rect_concurrent_test.cpp index 65b0dec034a..4617f12d798 100644 --- a/test/unit/math/rev/functor/map_rect_concurrent_test.cpp +++ b/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 job_params_d; std::vector > x_r; diff --git a/test/unit/math/rev/functor/map_rect_concurrent_threads_test.cpp b/test/unit/math/rev/functor/map_rect_concurrent_threads_test.cpp index 40b20451e73..b733304050b 100644 --- a/test/unit/math/rev/functor/map_rect_concurrent_threads_test.cpp +++ b/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 job_params_d; std::vector > x_r;