OmniSciDB  72c90bc290
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
threading_tbb.h
Go to the documentation of this file.
1 
2 #include <tbb/blocked_range.h>
3 #include <tbb/parallel_for.h>
4 #include <tbb/parallel_reduce.h>
5 #include <tbb/task_arena.h>
6 #include <tbb/task_group.h>
7 #include <functional>
8 
9 namespace threading_tbb {
10 
11 using tbb::blocked_range;
12 using tbb::task_arena;
13 using tbb::task_group;
14 namespace this_task_arena {
15 using namespace tbb::this_task_arena;
16 }
17 extern tbb::task_arena g_tbb_arena;
18 
19 template <typename... X>
20 void parallel_for(X&&... x) {
21  g_tbb_arena.execute([&] {
22  this_task_arena::isolate([&] { tbb::parallel_for(std::forward<X>(x)...); });
23  });
24 }
25 
26 template <typename... X>
27 auto parallel_reduce(X&&... x) -> decltype(tbb::parallel_reduce(std::forward<X>(x)...)) {
28  return g_tbb_arena.execute([&] {
29  return this_task_arena::isolate(
30  [&] { return tbb::parallel_reduce(std::forward<X>(x)...); });
31  });
32 }
33 
34 template <typename T>
35 struct tbb_packaged_task : tbb::task_group {
37  tbb_packaged_task() : value_(T()) {}
38 };
39 
40 template <>
41 struct tbb_packaged_task<void> : tbb::task_group {};
42 
43 template <typename T>
44 struct future {
45  std::unique_ptr<tbb_packaged_task<T>> task_;
46  future() = default;
47  future(future&&) = default;
48  future(std::unique_ptr<tbb_packaged_task<T>>&& p) : task_(std::move(p)) {}
49  void wait() {
50  g_tbb_arena.execute([this] { task_->wait(); });
51  }
52  T& get() {
53  wait();
54  return task_->value_;
55  }
56 };
57 
58 template <>
59 struct future<void> {
60  std::unique_ptr<tbb_packaged_task<void>> task_;
61  future() = default;
62  future(future&&) = default;
63  future(std::unique_ptr<tbb_packaged_task<void>>&& p) : task_(std::move(p)) {}
64  void wait() {
65  g_tbb_arena.execute([this] { task_->wait(); });
66  }
67  void get() { wait(); }
68 };
69 
70 template <typename Fn,
71  typename... Args,
72  typename Result = std::result_of_t<Fn && (Args && ...)>>
73 future<Result> async(Fn&& fn, Args&&... args) {
74  auto f = std::bind(std::forward<Fn>(fn), std::forward<Args>(args)...);
75  auto ptask = std::make_unique<tbb_packaged_task<Result>>();
76 #if TBB_INTERFACE_VERSION >= 12040
77  g_tbb_arena.enqueue(ptask->defer(f));
78 #else
79  g_tbb_arena.execute([&] { ptask->run(f); });
80 #endif
81  return future<Result>(std::move(ptask));
82 }
83 
84 } // namespace threading_tbb
future(std::unique_ptr< tbb_packaged_task< void >> &&p)
Definition: threading_tbb.h:63
future(std::unique_ptr< tbb_packaged_task< T >> &&p)
Definition: threading_tbb.h:48
tbb::task_arena g_tbb_arena
auto parallel_reduce(X &&...x) -> decltype(tbb::parallel_reduce(std::forward< X >(x)...))
Definition: threading_tbb.h:27
Value parallel_reduce(const blocked_range< Int > &range, const Value &identity, const RealBody &real_body, const Reduction &reduction, const Partitioner &p=Partitioner())
Parallel iteration with reduction.
torch::Tensor f(torch::Tensor x, torch::Tensor W_target, torch::Tensor b_target)
void parallel_for(const blocked_range< Int > &range, const Body &body, const Partitioner &p=Partitioner())
std::unique_ptr< tbb_packaged_task< T > > task_
Definition: threading_tbb.h:45
std::unique_ptr< tbb_packaged_task< void > > task_
Definition: threading_tbb.h:60
future< Result > async(Fn &&fn, Args &&...args)
Definition: threading_tbb.h:73
void parallel_for(X &&...x)
Definition: threading_tbb.h:20