1 /*
2     Copyright (c) 2020-2021 Intel Corporation
3 
4     Licensed under the Apache License, Version 2.0 (the "License");
5     you may not use this file except in compliance with the License.
6     You may obtain a copy of the License at
7 
8         http://www.apache.org/licenses/LICENSE-2.0
9 
10     Unless required by applicable law or agreed to in writing, software
11     distributed under the License is distributed on an "AS IS" BASIS,
12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13     See the License for the specific language governing permissions and
14     limitations under the License.
15 */
16 
17 #ifndef __TBB_test_common_utils_concurrency_limit_H
18 #define __TBB_test_common_utils_concurrency_limit_H
19 
20 #include "config.h"
21 #include "utils_assert.h"
22 #include "utils_report.h"
23 #include "oneapi/tbb/task_arena.h"
24 #include "oneapi/tbb/task_scheduler_observer.h"
25 #include "oneapi/tbb/enumerable_thread_specific.h"
26 
27 #include <cstddef>
28 #include <vector>
29 #include <algorithm>
30 
31 #if _WIN32 || _WIN64
32 #include <windows.h>
33 #elif __unix__
34 #include <unistd.h>
35 #if __linux__
36 #include <sys/sysinfo.h>
37 #endif
38 #include <string.h>
39 #include <sched.h>
40 #elif __FreeBSD__
41 #include <unistd.h>
42 #include <errno.h>
43 #include <string.h>
44 #include <sys/param.h>
45 #include <sys/cpuset.h>
46 #endif
47 #include <thread>
48 
49 namespace utils {
50 using thread_num_type = std::size_t;
51 
52 inline thread_num_type get_platform_max_threads() {
53     static thread_num_type platform_max_threads = tbb::this_task_arena::max_concurrency();
54     return platform_max_threads;
55 }
56 
57 inline std::vector<thread_num_type> concurrency_range(thread_num_type max_threads) {
58     std::vector<thread_num_type> threads_range;
59     thread_num_type step = 1;
60     for(thread_num_type thread_num = 1; thread_num <= max_threads; thread_num += step++)
61         threads_range.push_back(thread_num);
62     if(threads_range.back() != max_threads)
63         threads_range.push_back(max_threads);
64     // rotate in order to make threads_range non-monotonic
65     std::rotate(threads_range.begin(), threads_range.begin() + threads_range.size()/2, threads_range.end());
66     return threads_range;
67 }
68 
69 inline std::vector<thread_num_type> concurrency_range() {
70     static std::vector<thread_num_type> threads_range = concurrency_range(get_platform_max_threads());
71     return threads_range;
72 }
73 
74 #if !__TBB_TEST_SKIP_AFFINITY
75 
76 static int maxProcs = 0;
77 
78 static int get_max_procs() {
79     if (!maxProcs) {
80 #if _WIN32||_WIN64
81         DWORD_PTR pam, sam, m = 1;
82         GetProcessAffinityMask( GetCurrentProcess(), &pam, &sam );
83         int nproc = 0;
84         for ( std::size_t i = 0; i < sizeof(DWORD_PTR) * CHAR_BIT; ++i, m <<= 1 ) {
85             if ( pam & m )
86                 ++nproc;
87         }
88         maxProcs = nproc;
89 #elif __linux__
90         cpu_set_t mask;
91         int result = 0;
92         sched_getaffinity(0, sizeof(cpu_set_t), &mask);
93         int nproc = sysconf(_SC_NPROCESSORS_ONLN);
94         for (int i = 0; i < nproc; ++i) {
95             if (CPU_ISSET(i, &mask)) ++result;
96         }
97         maxProcs = result;
98 #else // FreeBSD
99         maxProcs = sysconf(_SC_NPROCESSORS_ONLN);
100 #endif
101     }
102     return maxProcs;
103 }
104 
105 int get_start_affinity_process() {
106 #if __linux__
107     cpu_set_t mask;
108     sched_getaffinity(0, sizeof(cpu_set_t), &mask);
109 
110     int result = -1;
111 
112     int nproc = sysconf(_SC_NPROCESSORS_ONLN);
113     for (int i = 0; i < nproc; ++i) {
114         if (CPU_ISSET(i, &mask)) {
115             result = i;
116             break;
117         }
118     }
119     ASSERT(result != -1, nullptr);
120     return result;
121 #else
122     // TODO: add affinity support for Windows and FreeBSD
123     return 0;
124 #endif
125 }
126 
127 int limit_number_of_threads( int max_threads ) {
128     ASSERT(max_threads >= 1,"The limited number of threads should be positive");
129     maxProcs = get_max_procs();
130     if (maxProcs < max_threads) {
131         // Suppose that process mask is not set so the number of available threads equals maxProcs
132         return maxProcs;
133     }
134 #if _WIN32 || _WIN64
135     ASSERT(max_threads <= 64, "LimitNumberOfThreads doesn't support max_threads to be more than 64 on Windows");
136     DWORD_PTR mask = 1;
137     for (int i = 1; i < max_threads; ++i) {
138         mask |= mask << 1;
139     }
140     bool err = !SetProcessAffinityMask(GetCurrentProcess(), mask);
141 #else
142 #if __linux__
143     using mask_t = cpu_set_t;
144 #define setaffinity(mask) sched_setaffinity(getpid(), sizeof(mask_t), &mask)
145 #else /*__FreeBSD*/
146     using mask_t = cpuset_t;
147 #define setaffinity(mask) cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, -1, sizeof(mask_t), &mask)
148 #endif
149 
150     mask_t new_mask;
151     CPU_ZERO(&new_mask);
152 
153     int mask_size = int(sizeof(mask_t) * CHAR_BIT);
154     if ( mask_size < maxProcs ) {
155         REPORT("The mask size doesn't seem to be big enough to call setaffinity. The call may return an error.");
156     }
157 
158     ASSERT(max_threads <= int(sizeof(mask_t) * CHAR_BIT), "The mask size is not enough to set the requested number of threads.");
159     int st = get_start_affinity_process();
160     for (int i = st; i < st + max_threads; ++i) {
161         CPU_SET(i, &new_mask);
162     }
163     int err = setaffinity(new_mask);
164 #endif
165     ASSERT(!err, "Setting process affinity failed");
166     return max_threads;
167 }
168 
169 #endif // __TBB_TEST_SKIP_AFFINITY
170 
171 // TODO: consider using cpuset_setaffinity/sched_getaffinity on FreeBSD to enable the functionality
172 #define OS_AFFINITY_SYSCALL_PRESENT (__linux__ && !__ANDROID__)
173 
174 #if OS_AFFINITY_SYSCALL_PRESENT
175 void get_thread_affinity_mask(std::size_t& ncpus, std::vector<int>& free_indexes) {
176     cpu_set_t* mask = nullptr;
177     ncpus = sizeof(cpu_set_t) * CHAR_BIT;
178     do {
179         mask = CPU_ALLOC(ncpus);
180         if (!mask) break;
181         const size_t size = CPU_ALLOC_SIZE(ncpus);
182         CPU_ZERO_S(size, mask);
183         const int err = sched_getaffinity(0, size, mask);
184         if (!err) break;
185 
186         CPU_FREE(mask);
187         mask = NULL;
188         if (errno != EINVAL) break;
189         ncpus <<= 1;
190     } while (ncpus < 16 * 1024 /* some reasonable limit */ );
191     ASSERT(mask, "Failed to obtain process affinity mask.");
192 
193     const size_t size = CPU_ALLOC_SIZE(ncpus);
194     const int num_cpus = CPU_COUNT_S(size, mask);
195     for (int i = 0; i < num_cpus; ++i) {
196         if (CPU_ISSET_S(i, size, mask)) {
197             free_indexes.push_back(i);
198         }
199     }
200 
201     CPU_FREE(mask);
202 }
203 
204 void pin_thread_imp(std::size_t ncpus, std::vector<int>& free_indexes, std::atomic<int>& curr_idx) {
205     const size_t size = CPU_ALLOC_SIZE(ncpus);
206 
207     ASSERT(free_indexes.size() > 0, nullptr);
208     int mapped_idx = free_indexes[curr_idx++ % free_indexes.size()];
209 
210     cpu_set_t *target_mask = CPU_ALLOC(ncpus);
211     ASSERT(target_mask, nullptr);
212     CPU_ZERO_S(size, target_mask);
213     CPU_SET_S(mapped_idx, size, target_mask);
214     const int err = sched_setaffinity(0, size, target_mask);
215     ASSERT(err == 0, "Failed to set thread affinity");
216 
217     CPU_FREE(target_mask);
218 }
219 #endif
220 
221 class thread_pinner {
222 public:
223     thread_pinner() {
224         tbb::detail::suppress_unused_warning(thread_index);
225 #if OS_AFFINITY_SYSCALL_PRESENT
226         get_thread_affinity_mask(ncpus, free_indexes);
227 #endif
228     }
229 
230     void pin_thread() {
231 #if OS_AFFINITY_SYSCALL_PRESENT
232         pin_thread_imp(ncpus, free_indexes, thread_index);
233 #endif
234     }
235 
236 private:
237 #if OS_AFFINITY_SYSCALL_PRESENT
238     std::size_t ncpus;
239     std::vector<int> free_indexes{};
240 #endif
241     std::atomic<int> thread_index{};
242 };
243 
244 class pinning_observer : public tbb::task_scheduler_observer {
245     thread_pinner pinner;
246     tbb::enumerable_thread_specific<bool> register_threads;
247 public:
248     pinning_observer(tbb::task_arena& arena) : tbb::task_scheduler_observer(arena), pinner() {
249         observe(true);
250     }
251 
252     void on_scheduler_entry( bool ) override {
253         bool& is_pinned = register_threads.local();
254         if (is_pinned) return;
255 
256         pinner.pin_thread();
257 
258         is_pinned = true;
259     }
260 
261     ~pinning_observer() { }
262 };
263 
264 #if __unix__
265 #include <sched.h>
266 #endif
267 
268 bool can_change_thread_priority() {
269 #if __unix__
270     pthread_t this_thread = pthread_self();
271     sched_param old_params;
272     int old_policy;
273     int err = pthread_getschedparam(this_thread, &old_policy, &old_params);
274     ASSERT(err == 0, nullptr);
275 
276     sched_param params;
277     params.sched_priority = sched_get_priority_max(SCHED_FIFO);
278     ASSERT(params.sched_priority != -1, nullptr);
279     err = pthread_setschedparam(this_thread, SCHED_FIFO, &params);
280     if (err == 0) {
281         err = pthread_setschedparam(this_thread, old_policy, &old_params);
282         ASSERT(err == 0, nullptr);
283     }
284     return err == 0;
285 #endif
286     return false;
287 }
288 
289 void increase_thread_priority() {
290 #if __unix__
291     pthread_t this_thread = pthread_self();
292     sched_param params;
293     params.sched_priority = sched_get_priority_max(SCHED_FIFO);
294     ASSERT(params.sched_priority != -1, nullptr);
295     int err = pthread_setschedparam(this_thread, SCHED_FIFO, &params);
296     ASSERT(err == 0, "Can not change thread priority.");
297 #endif
298 }
299 
300 void decrease_thread_priority() {
301 #if __unix__
302     pthread_t this_thread = pthread_self();
303     sched_param params;
304     params.sched_priority = sched_get_priority_min(SCHED_FIFO);
305     ASSERT(params.sched_priority != -1, nullptr);
306     int err = pthread_setschedparam(this_thread, SCHED_FIFO, &params);
307     ASSERT(err == 0, "Can not change thread priority.");
308 #endif
309 }
310 
311 } // namespace utils
312 
313 #endif // __TBB_test_common_utils_concurrency_limit_H
314