1 | #pragma once |
2 | |
3 | #include <c10/macros/Macros.h> |
4 | |
5 | /** |
6 | * Android versions with libgnustl incorrectly handle thread_local C++ |
7 | * qualifier with composite types. NDK up to r17 version is affected. |
8 | * |
9 | * (A fix landed on Jun 4 2018: |
10 | * https://android-review.googlesource.com/c/toolchain/gcc/+/683601) |
11 | * |
12 | * In such cases, use c10::ThreadLocal<T> wrapper |
13 | * which is `pthread_*` based with smart pointer semantics. |
14 | * |
15 | * In addition, convenient macro C10_DEFINE_TLS_static is available. |
16 | * To define static TLS variable of type std::string, do the following |
17 | * ``` |
18 | * C10_DEFINE_TLS_static(std::string, str_tls_); |
19 | * /////// |
20 | * { |
21 | * *str_tls_ = "abc"; |
22 | * assert(str_tls_->length(), 3); |
23 | * } |
24 | * ``` |
25 | * |
26 | * (see c10/test/util/ThreadLocal_test.cpp for more examples) |
27 | */ |
28 | #if !defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
29 | |
30 | #if defined(C10_ANDROID) && defined(__GLIBCXX__) && __GLIBCXX__ < 20180604 |
31 | #define C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE |
32 | #endif // defined(C10_ANDROID) && defined(__GLIBCXX__) && __GLIBCXX__ < 20180604 |
33 | |
34 | #endif // !defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
35 | |
36 | #if defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
37 | #include <c10/util/Exception.h> |
38 | #include <errno.h> |
39 | #include <pthread.h> |
40 | #include <memory> |
41 | namespace c10 { |
42 | |
43 | /** |
44 | * @brief Temporary thread_local C++ qualifier replacement for Android |
45 | * based on `pthread_*`. |
46 | * To be used with composite types that provide default ctor. |
47 | */ |
48 | template <typename Type> |
49 | class ThreadLocal { |
50 | public: |
51 | ThreadLocal() { |
52 | pthread_key_create( |
53 | &key_, [](void* buf) { delete static_cast<Type*>(buf); }); |
54 | } |
55 | |
56 | ~ThreadLocal() { |
57 | if (void* current = pthread_getspecific(key_)) { |
58 | delete static_cast<Type*>(current); |
59 | } |
60 | |
61 | pthread_key_delete(key_); |
62 | } |
63 | |
64 | ThreadLocal(const ThreadLocal&) = delete; |
65 | ThreadLocal& operator=(const ThreadLocal&) = delete; |
66 | |
67 | Type& get() { |
68 | if (void* current = pthread_getspecific(key_)) { |
69 | return *static_cast<Type*>(current); |
70 | } |
71 | |
72 | std::unique_ptr<Type> ptr = std::make_unique<Type>(); |
73 | if (0 == pthread_setspecific(key_, ptr.get())) { |
74 | return *ptr.release(); |
75 | } |
76 | |
77 | int err = errno; |
78 | TORCH_INTERNAL_ASSERT(false, "pthread_setspecific() failed, errno = " , err); |
79 | } |
80 | |
81 | Type& operator*() { |
82 | return get(); |
83 | } |
84 | |
85 | Type* operator->() { |
86 | return &get(); |
87 | } |
88 | |
89 | private: |
90 | pthread_key_t key_; |
91 | }; |
92 | |
93 | } // namespace c10 |
94 | |
95 | #define C10_DEFINE_TLS_static(Type, Name) static ::c10::ThreadLocal<Type> Name |
96 | |
97 | #define C10_DECLARE_TLS_class_static(Class, Type, Name) \ |
98 | static ::c10::ThreadLocal<Type> Name |
99 | |
100 | #define C10_DEFINE_TLS_class_static(Class, Type, Name) \ |
101 | ::c10::ThreadLocal<Type> Class::Name |
102 | |
103 | #else // defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
104 | |
105 | namespace c10 { |
106 | |
107 | /** |
108 | * @brief Default thread_local implementation for non-Android cases. |
109 | * To be used with composite types that provide default ctor. |
110 | */ |
111 | template <typename Type> |
112 | class ThreadLocal { |
113 | public: |
114 | using Accessor = Type* (*)(); |
115 | explicit ThreadLocal(Accessor accessor) : accessor_(accessor) {} |
116 | |
117 | ThreadLocal(const ThreadLocal&) = delete; |
118 | ThreadLocal& operator=(const ThreadLocal&) = delete; |
119 | |
120 | Type& get() { |
121 | return *accessor_(); |
122 | } |
123 | |
124 | Type& operator*() { |
125 | return get(); |
126 | } |
127 | |
128 | Type* operator->() { |
129 | return &get(); |
130 | } |
131 | |
132 | private: |
133 | Accessor accessor_; |
134 | }; |
135 | |
136 | } // namespace c10 |
137 | |
138 | #define C10_DEFINE_TLS_static(Type, Name) \ |
139 | static ::c10::ThreadLocal<Type> Name([]() { \ |
140 | static thread_local Type var; \ |
141 | return &var; \ |
142 | }) |
143 | |
144 | #define C10_DECLARE_TLS_class_static(Class, Type, Name) \ |
145 | static ::c10::ThreadLocal<Type> Name |
146 | |
147 | #define C10_DEFINE_TLS_class_static(Class, Type, Name) \ |
148 | ::c10::ThreadLocal<Type> Class::Name([]() { \ |
149 | static thread_local Type var; \ |
150 | return &var; \ |
151 | }) |
152 | |
153 | #endif // defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
154 | |