1/*
2 Copyright (c) 2005-2019 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_tests_harness_concurrency_H
18#define tbb_tests_harness_concurrency_H
19
20#if _WIN32||_WIN64
21#include "tbb/machine/windows_api.h"
22#elif __linux__
23#include <unistd.h>
24#include <sys/sysinfo.h>
25#include <string.h>
26#include <sched.h>
27#elif __FreeBSD__
28#include <unistd.h>
29#include <errno.h>
30#include <string.h>
31#include <sys/param.h> // Required by <sys/cpuset.h>
32#include <sys/cpuset.h>
33#endif
34
35#include <limits.h>
36
37namespace Harness {
38 static int maxProcs = 0;
39 static int GetMaxProcs() {
40 if ( !maxProcs ) {
41#if _WIN32||_WIN64
42 SYSTEM_INFO si;
43 GetNativeSystemInfo(&si);
44 maxProcs = si.dwNumberOfProcessors;
45#elif __linux__
46 maxProcs = get_nprocs();
47#else /* __FreeBSD__ */
48 maxProcs = sysconf(_SC_NPROCESSORS_ONLN);
49#endif
50 }
51 return maxProcs;
52 }
53
54 int LimitNumberOfThreads(int max_threads) {
55 ASSERT( max_threads >= 1 , "The limited number of threads should be positive." );
56 maxProcs = GetMaxProcs();
57 if ( maxProcs < max_threads )
58 // Suppose that process mask is not set so the number of available threads equals maxProcs
59 return maxProcs;
60
61#if _WIN32||_WIN64
62 ASSERT( max_threads <= 64 , "LimitNumberOfThreads doesn't support max_threads to be more than 64 on Windows." );
63 DWORD_PTR mask = 1;
64 for ( int i = 1; i < max_threads; ++i )
65 mask |= mask << 1;
66 bool err = !SetProcessAffinityMask( GetCurrentProcess(), mask );
67#else /* !WIN */
68#if __linux__
69 typedef cpu_set_t mask_t;
70#if __TBB_MAIN_THREAD_AFFINITY_BROKEN
71#define setaffinity(mask) sched_setaffinity(0 /*get the mask of the calling thread*/, sizeof(mask_t), &mask)
72#else
73#define setaffinity(mask) sched_setaffinity(getpid(), sizeof(mask_t), &mask)
74#endif
75#else /* __FreeBSD__ */
76 typedef cpuset_t mask_t;
77#if __TBB_MAIN_THREAD_AFFINITY_BROKEN
78#define setaffinity(mask) cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_TID, -1, sizeof(mask_t), &mask)
79#else
80#define setaffinity(mask) cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, -1, sizeof(mask_t), &mask)
81#endif
82#endif /* __FreeBSD__ */
83 mask_t newMask;
84 CPU_ZERO(&newMask);
85
86 int maskSize = (int)sizeof(mask_t) * CHAR_BIT;
87 ASSERT_WARNING( maskSize >= maxProcs, "The mask size doesn't seem to be big enough to call setaffinity. The call may return an error." );
88
89 ASSERT( max_threads <= (int)sizeof(mask_t) * CHAR_BIT , "The mask size is not enough to set the requested number of threads." );
90 for ( int i = 0; i < max_threads; ++i )
91 CPU_SET( i, &newMask );
92 int err = setaffinity( newMask );
93#endif /* !WIN */
94 ASSERT( !err, "Setting process affinity failed" );
95
96 return max_threads;
97 }
98
99} // namespace Harness
100
101#endif /* tbb_tests_harness_concurrency_H */
102