File size: 3,883 Bytes
7e50900
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#pragma once

#include <c10/macros/Macros.h>

/**
 * Android versions with libgnustl incorrectly handle thread_local C++
 * qualifier with composite types. NDK up to r17 version is affected.
 *
 * (A fix landed on Jun 4 2018:
 * https://android-review.googlesource.com/c/toolchain/gcc/+/683601)
 *
 * In such cases, use c10::ThreadLocal<T> wrapper
 * which is `pthread_*` based with smart pointer semantics.
 *
 * In addition, convenient macro C10_DEFINE_TLS_static is available.
 * To define static TLS variable of type std::string, do the following
 * ```
 *  C10_DEFINE_TLS_static(std::string, str_tls_);
 *  ///////
 *  {
 *    *str_tls_ = "abc";
 *    assert(str_tls_->length(), 3);
 *  }
 * ```
 *
 * (see c10/test/util/ThreadLocal_test.cpp for more examples)
 */
#if !defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE)

#if defined(C10_ANDROID) && defined(__GLIBCXX__) && __GLIBCXX__ < 20180604
#define C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE
#endif // defined(C10_ANDROID) && defined(__GLIBCXX__) && __GLIBCXX__ < 20180604

#endif // !defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE)

#if defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE)
#include <c10/util/Exception.h>
#include <errno.h>
#include <pthread.h>
#include <memory>
namespace c10 {

/**
 * @brief Temporary thread_local C++ qualifier replacement for Android
 * based on `pthread_*`.
 * To be used with composite types that provide default ctor.
 */
template <typename Type>
class ThreadLocal {
 public:
  ThreadLocal() {
    pthread_key_create(
        &key_, [](void* buf) { delete static_cast<Type*>(buf); });
  }

  ~ThreadLocal() {
    if (void* current = pthread_getspecific(key_)) {
      delete static_cast<Type*>(current);
    }

    pthread_key_delete(key_);
  }

  ThreadLocal(const ThreadLocal&) = delete;
  ThreadLocal& operator=(const ThreadLocal&) = delete;

  Type& get() {
    if (void* current = pthread_getspecific(key_)) {
      return *static_cast<Type*>(current);
    }

    std::unique_ptr<Type> ptr = std::make_unique<Type>();
    if (0 == pthread_setspecific(key_, ptr.get())) {
      return *ptr.release();
    }

    int err = errno;
    TORCH_INTERNAL_ASSERT(false, "pthread_setspecific() failed, errno = ", err);
  }

  Type& operator*() {
    return get();
  }

  Type* operator->() {
    return &get();
  }

 private:
  pthread_key_t key_;
};

} // namespace c10

#define C10_DEFINE_TLS_static(Type, Name) static ::c10::ThreadLocal<Type> Name

#define C10_DECLARE_TLS_class_static(Class, Type, Name) \
  static ::c10::ThreadLocal<Type> Name

#define C10_DEFINE_TLS_class_static(Class, Type, Name) \
  ::c10::ThreadLocal<Type> Class::Name

#else // defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE)

namespace c10 {

/**
 * @brief Default thread_local implementation for non-Android cases.
 * To be used with composite types that provide default ctor.
 */
template <typename Type>
class ThreadLocal {
 public:
  using Accessor = Type* (*)();
  explicit ThreadLocal(Accessor accessor) : accessor_(accessor) {}

  ThreadLocal(const ThreadLocal&) = delete;
  ThreadLocal& operator=(const ThreadLocal&) = delete;

  Type& get() {
    return *accessor_();
  }

  Type& operator*() {
    return get();
  }

  Type* operator->() {
    return &get();
  }

 private:
  Accessor accessor_;
};

} // namespace c10

#define C10_DEFINE_TLS_static(Type, Name)     \
  static ::c10::ThreadLocal<Type> Name([]() { \
    static thread_local Type var;             \
    return &var;                              \
  })

#define C10_DECLARE_TLS_class_static(Class, Type, Name) \
  static ::c10::ThreadLocal<Type> Name

#define C10_DEFINE_TLS_class_static(Class, Type, Name) \
  ::c10::ThreadLocal<Type> Class::Name([]() {          \
    static thread_local Type var;                      \
    return &var;                                       \
  })

#endif // defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE)