ThreadLocal.h 4.2 KB

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