1 #pragma once 2 3 #include <utility> 4 5 namespace c10 { 6 7 // See example implementation in TensorBase.h and TensorBody.h. 8 // Synopsis: 9 // 10 // repr_type -- type to use to store an owned T in ExclusivelyOwned. 11 // 12 // pointer_type -- pointer-esque type to return from 13 // ExclusivelyOwned's get() and operator*() methods. 14 // 15 // const_pointer_type -- similar to pointer_type, used for the const methods. 16 // 17 // static repr_type nullRepr() -- return a null instance of repr_type. 18 // 19 // template <class... Args> 20 // static repr_type createInPlace(Args&&... args) -- used by the in-place 21 // ExclusivelyOwned constructor. 22 // 23 // static repr_type moveToRepr(T&& x) -- move the given x into an 24 // instance of repr_type. used by the ExclusivelyOwned(T&&) 25 // constructor. 26 // 27 // static void destroyOwned(repr_type x) -- free memory for a 28 // known-exclusively-owned instance of x. Replaces calling repr_type's 29 // destructor. Being able to implement this more efficiently than 30 // repr_type's destructor is the main reason to use ExclusivelyOwned 31 // for a type. 32 // 33 // static T take(repr_type&) -- move out of the given repr_type into an owned T. 34 // 35 // static pointer_type getImpl(const repr_type&) -- return a pointer 36 // to the given repr_type. May take repr_type by value if that is more 37 // efficient. 38 template <typename T> 39 struct ExclusivelyOwnedTraits; 40 41 /// ExclusivelyOwned is a smart-pointer-like wrapper around an 42 /// exclusively-owned instance of some type T that normally has 43 /// mandatory reference counting (currently just Tensor). If you have 44 /// an isolated piece of code that knows that it has sole ownership of 45 /// an object of one of these types (i.e., because you created it 46 /// directly or using a factory function) and that object will not 47 /// escape from that isolated piece of code, then moving the object 48 /// into an ExclusivelyOwned will avoid an atomic reference count 49 /// decrement at destruction time. 50 /// 51 /// If you directly create the Tensor in the first 52 /// place, you can use the in_place constructor of ExclusivelyOwned to 53 /// additionally avoid doing any stores to initialize the refcount & 54 /// weakcount. 55 template <typename T> 56 class ExclusivelyOwned { 57 using EOT = ExclusivelyOwnedTraits<T>; 58 typename ExclusivelyOwnedTraits<T>::repr_type repr_; 59 60 public: ExclusivelyOwned()61 ExclusivelyOwned() : repr_(EOT::nullRepr()) {} 62 ExclusivelyOwned(T && t)63 explicit ExclusivelyOwned(T&& t) : repr_(EOT::moveToRepr(std::move(t))) {} 64 65 template <class... Args> ExclusivelyOwned(std::in_place_t,Args &&...args)66 explicit ExclusivelyOwned(std::in_place_t, Args&&... args) 67 : repr_(EOT::createInPlace(std::forward<Args>(args)...)) {} 68 69 ExclusivelyOwned(const ExclusivelyOwned&) = delete; 70 ExclusivelyOwned(ExclusivelyOwned && rhs)71 ExclusivelyOwned(ExclusivelyOwned&& rhs) noexcept 72 : repr_(std::move(rhs.repr_)) { 73 rhs.repr_ = EOT::nullRepr(); 74 } 75 76 ExclusivelyOwned& operator=(const ExclusivelyOwned&) = delete; 77 78 ExclusivelyOwned& operator=(ExclusivelyOwned&& rhs) noexcept { 79 EOT::destroyOwned(repr_); 80 repr_ = std::move(rhs.repr_); 81 rhs.repr_ = EOT::nullRepr(); 82 return *this; 83 } 84 85 ExclusivelyOwned& operator=(T&& rhs) noexcept { 86 EOT::destroyOwned(repr_); 87 repr_ = EOT::moveToRepr(std::move(rhs)); 88 return *this; 89 } 90 ~ExclusivelyOwned()91 ~ExclusivelyOwned() { 92 EOT::destroyOwned(repr_); 93 // Don't bother to call the destructor of repr_, since we already 94 // did specialized destruction for the exclusively-owned case in 95 // destroyOwned! 96 } 97 98 // We don't provide this because it would require us to be able to 99 // differentiate an owned-but-empty T from a lack of T. This is 100 // particularly problematic for Tensor, which wants to use an 101 // undefined Tensor as its null state. 102 explicit operator bool() const noexcept = delete; 103 T()104 operator T() && { 105 return take(); 106 } 107 108 // NOTE: the equivalent operation on MaybeOwned is a moving 109 // operator*. For ExclusivelyOwned, take() and operator*() may well 110 // have different return types, so they are different functions. take()111 T take() && { 112 return EOT::take(repr_); 113 } 114 115 typename EOT::const_pointer_type operator->() const { 116 return get(); 117 } 118 get()119 typename EOT::const_pointer_type get() const { 120 return EOT::getImpl(repr_); 121 } 122 123 typename EOT::pointer_type operator->() { 124 return get(); 125 } 126 get()127 typename EOT::pointer_type get() { 128 return EOT::getImpl(repr_); 129 } 130 131 std::remove_pointer_t<typename EOT::const_pointer_type>& operator*() const { 132 return *get(); 133 } 134 135 std::remove_pointer_t<typename EOT::pointer_type>& operator*() { 136 return *get(); 137 } 138 }; 139 140 } // namespace c10 141