thread_local.h

Go to the documentation of this file.

 
#ifndef DMLC_THREAD_LOCAL_H_
#define DMLC_THREAD_LOCAL_H_
 
#include <mutex>
#include <memory>
#include <vector>
#include "./base.h"
 
namespace dmlc {
 
// macro hanlding for threadlocal variables
#ifdef __GNUC__
  #define MX_THREAD_LOCAL __thread
#elif __STDC_VERSION__ >= 201112L
  #define  MX_THREAD_LOCAL _Thread_local
#elif defined(_MSC_VER)
  #define MX_THREAD_LOCAL __declspec(thread)
#endif
 
#if DMLC_CXX11_THREAD_LOCAL == 0
#pragma message("Warning: CXX11 thread_local is not formally supported")
#endif
 
template<typename T>
class ThreadLocalStore {
 public:
  static T* Get() {
#if DMLC_CXX11_THREAD_LOCAL && DMLC_MODERN_THREAD_LOCAL == 1
    static thread_local T inst;
    return &inst;
#else
    static MX_THREAD_LOCAL T* ptr = nullptr;
    if (ptr == nullptr) {
      ptr = new T();
      // Syntactic work-around for the nvcc of the initial cuda v10.1 release,
      // which fails to compile 'Singleton()->' below. Fixed in v10.1 update 1.
      (*Singleton()).RegisterDelete(ptr);
    }
    return ptr;
#endif
  }
 
 private:
  ThreadLocalStore() {}
  ~ThreadLocalStore() {
    for (size_t i = 0; i < data_.size(); ++i) {
      delete data_[i];
    }
  }
  static ThreadLocalStore<T> *Singleton() {
    static ThreadLocalStore<T> inst;
    return &inst;
  }
  void RegisterDelete(T *str) {
    std::unique_lock<std::mutex> lock(mutex_);
    data_.push_back(str);
    lock.unlock();
  }
  std::mutex mutex_;
  std::vector<T*> data_;
};
 
}  // namespace dmlc
 
#endif  // DMLC_THREAD_LOCAL_H_