Lazy.h 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <atomic>
  4. #include <utility>
  5. namespace c10 {
  6. /**
  7. * Thread-safe lazy value with opportunistic concurrency: on concurrent first
  8. * access, the factory may be called by multiple threads, but only one result is
  9. * stored and its reference returned to all the callers.
  10. *
  11. * Value is heap-allocated; this optimizes for the case in which the value is
  12. * never actually computed.
  13. */
  14. template <class T>
  15. class OptimisticLazy {
  16. public:
  17. OptimisticLazy() = default;
  18. OptimisticLazy(const OptimisticLazy& other) {
  19. if (T* value = other.value_.load(std::memory_order_acquire)) {
  20. value_ = new T(*value);
  21. }
  22. }
  23. OptimisticLazy(OptimisticLazy&& other) noexcept
  24. : value_(other.value_.exchange(nullptr, std::memory_order_acq_rel)) {}
  25. ~OptimisticLazy() {
  26. reset();
  27. }
  28. template <class Factory>
  29. T& ensure(const Factory& factory) {
  30. if (T* value = value_.load(std::memory_order_acquire)) {
  31. return *value;
  32. }
  33. T* value = new T(factory());
  34. T* old = nullptr;
  35. if (!value_.compare_exchange_strong(
  36. old, value, std::memory_order_release, std::memory_order_acquire)) {
  37. delete value;
  38. value = old;
  39. }
  40. return *value;
  41. }
  42. // The following methods are not thread-safe: they should not be called
  43. // concurrently with any other method.
  44. OptimisticLazy& operator=(const OptimisticLazy& other) {
  45. *this = OptimisticLazy{other};
  46. return *this;
  47. }
  48. OptimisticLazy& operator=(OptimisticLazy&& other) noexcept {
  49. if (this != &other) {
  50. reset();
  51. value_.store(
  52. other.value_.exchange(nullptr, std::memory_order_acquire),
  53. std::memory_order_release);
  54. }
  55. return *this;
  56. }
  57. void reset() {
  58. if (T* old = value_.load(std::memory_order_relaxed)) {
  59. value_.store(nullptr, std::memory_order_relaxed);
  60. delete old;
  61. }
  62. }
  63. private:
  64. std::atomic<T*> value_{nullptr};
  65. };
  66. /**
  67. * Interface for a value that is computed on first access.
  68. */
  69. template <class T>
  70. class LazyValue {
  71. public:
  72. virtual ~LazyValue() = default;
  73. virtual const T& get() const = 0;
  74. };
  75. /**
  76. * Convenience thread-safe LazyValue implementation with opportunistic
  77. * concurrency.
  78. */
  79. template <class T>
  80. class OptimisticLazyValue : public LazyValue<T> {
  81. public:
  82. const T& get() const override {
  83. return value_.ensure([this] { return compute(); });
  84. }
  85. private:
  86. virtual T compute() const = 0;
  87. mutable OptimisticLazy<T> value_;
  88. };
  89. /**
  90. * Convenience immutable (thus thread-safe) LazyValue implementation for cases
  91. * in which the value is not actually lazy.
  92. */
  93. template <class T>
  94. class PrecomputedLazyValue : public LazyValue<T> {
  95. public:
  96. PrecomputedLazyValue(T value) : value_(std::move(value)) {}
  97. const T& get() const override {
  98. return value_;
  99. }
  100. private:
  101. T value_;
  102. };
  103. } // namespace c10
  104. #else
  105. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  106. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)