1// ======================================================================== //
2// Copyright 2009-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#pragma once
18
19#include "platform.h"
20
21#if !defined(_WIN32)
22 #include <pthread.h>
23 #include <sched.h>
24 #if defined(__APPLE__)
25 #include <mach/thread_policy.h>
26 #endif
27#endif
28
29#include <vector>
30#include <mutex>
31
32namespace oidn {
33
34 // --------------------------------------------------------------------------
35 // ThreadLocal
36 // --------------------------------------------------------------------------
37
38 // Wrapper which makes any variable thread-local
39 template<typename T>
40 class ThreadLocal : public Verbose
41 {
42 private:
43 #if defined(_WIN32)
44 DWORD key;
45 #else
46 pthread_key_t key;
47 #endif
48
49 std::vector<T*> instances;
50 std::mutex mutex;
51
52 public:
53 ThreadLocal(int verbose = 0)
54 : Verbose(verbose)
55 {
56 #if defined(_WIN32)
57 key = TlsAlloc();
58 if (key == TLS_OUT_OF_INDEXES)
59 OIDN_FATAL("TlsAlloc failed");
60 #else
61 if (pthread_key_create(&key, nullptr) != 0)
62 OIDN_FATAL("pthread_key_create failed");
63 #endif
64 }
65
66 ~ThreadLocal()
67 {
68 std::lock_guard<std::mutex> lock(mutex);
69 for (T* ptr : instances)
70 delete ptr;
71
72 #if defined(_WIN32)
73 if (!TlsFree(key))
74 OIDN_WARNING("TlsFree failed");
75 #else
76 if (pthread_key_delete(key) != 0)
77 OIDN_WARNING("pthread_key_delete failed");
78 #endif
79 }
80
81 T& get()
82 {
83 #if defined(_WIN32)
84 T* ptr = (T*)TlsGetValue(key);
85 #else
86 T* ptr = (T*)pthread_getspecific(key);
87 #endif
88
89 if (ptr)
90 return *ptr;
91
92 ptr = new T;
93 std::lock_guard<std::mutex> lock(mutex);
94 instances.push_back(ptr);
95
96 #if defined(_WIN32)
97 if (!TlsSetValue(key, ptr))
98 OIDN_FATAL("TlsSetValue failed");
99 #else
100 if (pthread_setspecific(key, ptr) != 0)
101 OIDN_FATAL("pthread_setspecific failed");
102 #endif
103
104 return *ptr;
105 }
106 };
107
108#if defined(_WIN32)
109
110 // --------------------------------------------------------------------------
111 // ThreadAffinity - Windows
112 // --------------------------------------------------------------------------
113
114 class ThreadAffinity : public Verbose
115 {
116 private:
117 typedef BOOL (WINAPI *GetLogicalProcessorInformationExFunc)(LOGICAL_PROCESSOR_RELATIONSHIP,
118 PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX,
119 PDWORD);
120
121 typedef BOOL (WINAPI *SetThreadGroupAffinityFunc)(HANDLE,
122 CONST GROUP_AFFINITY*,
123 PGROUP_AFFINITY);
124
125 GetLogicalProcessorInformationExFunc pGetLogicalProcessorInformationEx = nullptr;
126 SetThreadGroupAffinityFunc pSetThreadGroupAffinity = nullptr;
127
128 std::vector<GROUP_AFFINITY> affinities; // thread affinities
129 std::vector<GROUP_AFFINITY> oldAffinities; // original thread affinities
130
131 public:
132 ThreadAffinity(int numThreadsPerCore = INT_MAX, int verbose = 0);
133
134 int getNumThreads() const
135 {
136 return (int)affinities.size();
137 }
138
139 // Sets the affinity (0..numThreads-1) of the thread after saving the current affinity
140 void set(int threadIndex);
141
142 // Restores the affinity of the thread
143 void restore(int threadIndex);
144 };
145
146#elif defined(__linux__)
147
148 // --------------------------------------------------------------------------
149 // ThreadAffinity - Linux
150 // --------------------------------------------------------------------------
151
152 class ThreadAffinity : public Verbose
153 {
154 private:
155 std::vector<cpu_set_t> affinities; // thread affinities
156 std::vector<cpu_set_t> oldAffinities; // original thread affinities
157
158 public:
159 ThreadAffinity(int numThreadsPerCore = INT_MAX, int verbose = 0);
160
161 int getNumThreads() const
162 {
163 return (int)affinities.size();
164 }
165
166 // Sets the affinity (0..numThreads-1) of the thread after saving the current affinity
167 void set(int threadIndex);
168
169 // Restores the affinity of the thread
170 void restore(int threadIndex);
171 };
172
173#elif defined(__APPLE__)
174
175 // --------------------------------------------------------------------------
176 // ThreadAffinity - macOS
177 // --------------------------------------------------------------------------
178
179 class ThreadAffinity : public Verbose
180 {
181 private:
182 std::vector<thread_affinity_policy> affinities; // thread affinities
183 std::vector<thread_affinity_policy> oldAffinities; // original thread affinities
184
185 public:
186 ThreadAffinity(int numThreadsPerCore = INT_MAX, int verbose = 0);
187
188 int getNumThreads() const
189 {
190 return (int)affinities.size();
191 }
192
193 // Sets the affinity (0..numThreads-1) of the thread after saving the current affinity
194 void set(int threadIndex);
195
196 // Restores the affinity of the thread
197 void restore(int threadIndex);
198 };
199
200#endif
201
202} // namespace oidn
203