xref: /aosp_15_r20/external/pytorch/c10/util/ExclusivelyOwned.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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