|
#pragma once |
|
|
|
#include <c10/macros/Macros.h> |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
#if !defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
|
|
|
#if defined(C10_ANDROID) && defined(__GLIBCXX__) && __GLIBCXX__ < 20180604 |
|
#define C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE |
|
#endif |
|
|
|
#endif |
|
|
|
#if defined(C10_PREFER_CUSTOM_THREAD_LOCAL_STORAGE) |
|
#include <c10/util/Exception.h> |
|
#include <errno.h> |
|
#include <pthread.h> |
|
#include <memory> |
|
namespace c10 { |
|
|
|
|
|
|
|
|
|
|
|
|
|
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_; |
|
}; |
|
|
|
} |
|
|
|
#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 |
|
|
|
namespace c10 { |
|
|
|
|
|
|
|
|
|
|
|
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_; |
|
}; |
|
|
|
} |
|
|
|
#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 |
|
|