ThreadLocal.h 3.8 KB

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