Merge pull request #2445 from stan-dev/feature/issue-2444-nested-paralellism

Feature/issue 2444 nested paralellism

Author: rok-cesnovar

stan/math/fwd/functor/reduce_sum.hpp

@@ -4,10 +4,6 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
 
-#include <tbb/task_arena.h>
-#include <tbb/parallel_reduce.h>
-#include <tbb/blocked_range.h>
-
 #include <algorithm>
 #include <tuple>
 #include <vector>

stan/math/rev/core/init_chainablestack.hpp

@@ -20,8 +20,9 @@ namespace math {
  * hook ensures that each worker thread has an initialized AD tape
  * ready for use.
  *
- * Refer to https://software.intel.com/en-us/node/506314 for details
- * on the observer concept.
+ * Refer to
+ * https://software.intel.com/content/www/us/en/develop/documentation/tbb-documentation/top/intel-threading-building-blocks-developer-reference/task-scheduler/taskschedulerobserver.html
+ * for details on the observer concept.
  */
 class ad_tape_observer final : public tbb::task_scheduler_observer {
   using stack_ptr = std::unique_ptr<ChainableStack>;

stan/math/rev/functor/map_rect_concurrent.hpp

@@ -45,10 +45,18 @@ map_rect_concurrent(
   };
 
 #ifdef STAN_THREADS
-  tbb::parallel_for(tbb::blocked_range<std::size_t>(0, num_jobs),
-                    [&](const tbb::blocked_range<size_t>& r) {
-                      execute_chunk(r.begin(), r.end());
-                    });
+  // we must use task isolation as described here:
+  // https://software.intel.com/content/www/us/en/develop/documentation/tbb-documentation/top/intel-threading-building-blocks-developer-guide/task-isolation.html
+  // this is to ensure that the thread local AD tape ressource is
+  // not being modified from a different task which may happen
+  // whenever this function is being used itself in a parallel
+  // context (like running multiple chains for Stan)
+  tbb::this_task_arena::isolate([&] {
+    tbb::parallel_for(tbb::blocked_range<std::size_t>(0, num_jobs),
+                      [&](const tbb::blocked_range<size_t>& r) {
+                        execute_chunk(r.begin(), r.end());
+                      });
+  });
 #else
   execute_chunk(0, num_jobs);
 #endif

stan/math/rev/functor/reduce_sum.hpp

@@ -4,6 +4,7 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/functor.hpp>
 #include <stan/math/rev/core.hpp>
+
 #include <tbb/task_arena.h>
 #include <tbb/parallel_reduce.h>
 #include <tbb/blocked_range.h>
@@ -249,15 +250,23 @@ struct reduce_sum_impl<ReduceFunction, require_var_t<ReturnType>, ReturnType,
                              std::forward<Vec>(vmapped),
                              std::forward<Args>(args)...);
 
-    if (auto_partitioning) {
-      tbb::parallel_reduce(
-          tbb::blocked_range<std::size_t>(0, num_terms, grainsize), worker);
-    } else {
-      tbb::simple_partitioner partitioner;
-      tbb::parallel_deterministic_reduce(
-          tbb::blocked_range<std::size_t>(0, num_terms, grainsize), worker,
-          partitioner);
-    }
+    // we must use task isolation as described here:
+    // https://software.intel.com/content/www/us/en/develop/documentation/tbb-documentation/top/intel-threading-building-blocks-developer-guide/task-isolation.html
+    // this is to ensure that the thread local AD tape ressource is
+    // not being modified from a different task which may happen
+    // whenever this function is being used itself in a parallel
+    // context (like running multiple chains for Stan)
+    tbb::this_task_arena::isolate([&] {
+      if (auto_partitioning) {
+        tbb::parallel_reduce(
+            tbb::blocked_range<std::size_t>(0, num_terms, grainsize), worker);
+      } else {
+        tbb::simple_partitioner partitioner;
+        tbb::parallel_deterministic_reduce(
+            tbb::blocked_range<std::size_t>(0, num_terms, grainsize), worker,
+            partitioner);
+      }
+    });
 
     for (size_t i = 0; i < num_vars_shared_terms; ++i) {
       partials[num_vars_sliced_terms + i] = worker.args_adjoints_.coeff(i);