Treelite
parallel_for.h
Go to the documentation of this file.
1 
7 #ifndef TREELITE_THREADING_UTILS_PARALLEL_FOR_H_
8 #define TREELITE_THREADING_UTILS_PARALLEL_FOR_H_
9 
10 #include <treelite/omp.h>
11 #include <treelite/logging.h>
12 #include <treelite/omp_exception.h>
13 #include <type_traits>
14 #include <algorithm>
15 #include <exception>
16 #include <mutex>
17 #include <cstddef>
18 #include <cstdint>
19 
20 namespace treelite {
21 namespace threading_utils {
22 
23 inline int OmpGetThreadLimit() {
24  int limit = omp_get_thread_limit();
25  TREELITE_CHECK_GE(limit, 1) << "Invalid thread limit for OpenMP.";
26  return limit;
27 }
28 
29 inline int MaxNumThread() {
30  return std::min(std::min(omp_get_num_procs(), omp_get_max_threads()), OmpGetThreadLimit());
31 }
32 
36 struct ThreadConfig {
37  std::uint32_t nthread;
38 };
39 
46 inline ThreadConfig ConfigureThreadConfig(int nthread) {
47  if (nthread <= 0) {
48  nthread = MaxNumThread();
49  TREELITE_CHECK_GE(nthread, 1) << "Invalid number of threads configured in OpenMP";
50  } else {
51  TREELITE_CHECK_LE(nthread, MaxNumThread())
52  << "nthread cannot exceed " << MaxNumThread() << " (configured by OpenMP).";
53  }
54  return ThreadConfig{static_cast<std::uint32_t>(nthread)};
55 }
56 
57 // OpenMP schedule
59  enum {
60  kAuto,
61  kDynamic,
62  kStatic,
63  kGuided,
64  } sched;
65  std::size_t chunk{0};
66 
67  ParallelSchedule static Auto() { return ParallelSchedule{kAuto}; }
68  ParallelSchedule static Dynamic(std::size_t n = 0) { return ParallelSchedule{kDynamic, n}; }
69  ParallelSchedule static Static(std::size_t n = 0) { return ParallelSchedule{kStatic, n}; }
70  ParallelSchedule static Guided() { return ParallelSchedule{kGuided}; }
71 };
72 
73 template <typename IndexType, typename FuncType>
74 inline void ParallelFor(IndexType begin, IndexType end, const ThreadConfig& thread_config,
75  ParallelSchedule sched, FuncType func) {
76  if (begin == end) {
77  return;
78  }
79 
80 #if defined(_MSC_VER)
81  // msvc doesn't support unsigned integer as openmp index.
82  using OmpInd = std::conditional_t<std::is_signed<IndexType>::value, IndexType, std::int64_t>;
83 #else
84  using OmpInd = IndexType;
85 #endif
86 
87  OMPException exc;
88  switch (sched.sched) {
89  case ParallelSchedule::kAuto: {
90 #pragma omp parallel for num_threads(thread_config.nthread)
91  for (OmpInd i = begin; i < end; ++i) {
92  exc.Run(func, static_cast<IndexType>(i), omp_get_thread_num());
93  }
94  break;
95  }
96  case ParallelSchedule::kDynamic: {
97  if (sched.chunk == 0) {
98 #pragma omp parallel for num_threads(thread_config.nthread) schedule(dynamic)
99  for (OmpInd i = begin; i < end; ++i) {
100  exc.Run(func, static_cast<IndexType>(i), omp_get_thread_num());
101  }
102  } else {
103 #pragma omp parallel for num_threads(thread_config.nthread) schedule(dynamic, sched.chunk)
104  for (OmpInd i = begin; i < end; ++i) {
105  exc.Run(func, static_cast<IndexType>(i), omp_get_thread_num());
106  }
107  }
108  break;
109  }
110  case ParallelSchedule::kStatic: {
111  if (sched.chunk == 0) {
112 #pragma omp parallel for num_threads(thread_config.nthread) schedule(static)
113  for (OmpInd i = begin; i < end; ++i) {
114  exc.Run(func, static_cast<IndexType>(i), omp_get_thread_num());
115  }
116  } else {
117 #pragma omp parallel for num_threads(thread_config.nthread) schedule(static, sched.chunk)
118  for (OmpInd i = begin; i < end; ++i) {
119  exc.Run(func, static_cast<IndexType>(i), omp_get_thread_num());
120  }
121  }
122  break;
123  }
124  case ParallelSchedule::kGuided: {
125 #pragma omp parallel for num_threads(thread_config.nthread) schedule(guided)
126  for (OmpInd i = begin; i < end; ++i) {
127  exc.Run(func, static_cast<IndexType>(i), omp_get_thread_num());
128  }
129  break;
130  }
131  }
132  exc.Rethrow();
133 }
134 
135 } // namespace threading_utils
136 } // namespace treelite
137 
138 #endif // TREELITE_THREADING_UTILS_PARALLEL_FOR_H_
Represent thread configuration, to be used with parallel loops.
Definition: parallel_for.h:36
logging facility for Treelite
ThreadConfig ConfigureThreadConfig(int nthread)
Create therad configuration.
Definition: parallel_for.h:46
Utility to propagate exceptions throws inside an OpenMP block.
void Run(Function f, Parameters... params)
Parallel OMP blocks should be placed within Run to save exception.
Definition: omp_exception.h:31
OMP Exception class catches, saves and rethrows exception from OMP blocks.
Definition: omp_exception.h:19
compatiblity wrapper for systems that don&#39;t support OpenMP
void Rethrow()
should be called from the main thread to rethrow the exception
Definition: omp_exception.h:50