Line data    Source code

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