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>
41namespace 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 */
48template <typename Type>
49class 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
105namespace 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 */
111template <typename Type>
112class 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