1 | /*! |
2 | * Copyright (c) 2015 by Contributors |
3 | * \file thread_local.h |
4 | * \brief Portable thread local storage. |
5 | */ |
6 | #ifndef DMLC_THREAD_LOCAL_H_ |
7 | #define DMLC_THREAD_LOCAL_H_ |
8 | |
9 | #include <mutex> |
10 | #include <memory> |
11 | #include <vector> |
12 | #include "./base.h" |
13 | |
14 | namespace dmlc { |
15 | |
16 | // macro hanlding for threadlocal variables |
17 | #ifdef __GNUC__ |
18 | #define MX_THREAD_LOCAL __thread |
19 | #elif __STDC_VERSION__ >= 201112L |
20 | #define MX_THREAD_LOCAL _Thread_local |
21 | #elif defined(_MSC_VER) |
22 | #define MX_THREAD_LOCAL __declspec(thread) |
23 | #endif |
24 | |
25 | #if DMLC_CXX11_THREAD_LOCAL == 0 |
26 | #pragma message("Warning: CXX11 thread_local is not formally supported") |
27 | #endif |
28 | |
29 | /*! |
30 | * \brief A threadlocal store to store threadlocal variables. |
31 | * Will return a thread local singleton of type T |
32 | * \tparam T the type we like to store |
33 | */ |
34 | template<typename T> |
35 | class ThreadLocalStore { |
36 | public: |
37 | /*! \return get a thread local singleton */ |
38 | static T* Get() { |
39 | #if DMLC_CXX11_THREAD_LOCAL && DMLC_MODERN_THREAD_LOCAL == 1 |
40 | static thread_local T inst; |
41 | return &inst; |
42 | #else |
43 | static MX_THREAD_LOCAL T* ptr = nullptr; |
44 | if (ptr == nullptr) { |
45 | ptr = new T(); |
46 | // Syntactic work-around for the nvcc of the initial cuda v10.1 release, |
47 | // which fails to compile 'Singleton()->' below. Fixed in v10.1 update 1. |
48 | (*Singleton()).RegisterDelete(ptr); |
49 | } |
50 | return ptr; |
51 | #endif |
52 | } |
53 | |
54 | private: |
55 | /*! \brief constructor */ |
56 | ThreadLocalStore() {} |
57 | /*! \brief destructor */ |
58 | ~ThreadLocalStore() { |
59 | for (size_t i = 0; i < data_.size(); ++i) { |
60 | delete data_[i]; |
61 | } |
62 | } |
63 | /*! \return singleton of the store */ |
64 | static ThreadLocalStore<T> *Singleton() { |
65 | static ThreadLocalStore<T> inst; |
66 | return &inst; |
67 | } |
68 | /*! |
69 | * \brief register str for internal deletion |
70 | * \param str the string pointer |
71 | */ |
72 | void RegisterDelete(T *str) { |
73 | std::unique_lock<std::mutex> lock(mutex_); |
74 | data_.push_back(str); |
75 | lock.unlock(); |
76 | } |
77 | /*! \brief internal mutex */ |
78 | std::mutex mutex_; |
79 | /*!\brief internal data */ |
80 | std::vector<T*> data_; |
81 | }; |
82 | |
83 | } // namespace dmlc |
84 | |
85 | #endif // DMLC_THREAD_LOCAL_H_ |
86 | |