1 | #pragma once |
2 | |
3 | #include <c10/util/in_place.h> |
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 | union { |
59 | char dummy_; |
60 | typename ExclusivelyOwnedTraits<T>::repr_type repr_; |
61 | }; |
62 | |
63 | public: |
64 | ExclusivelyOwned() : repr_(EOT::nullRepr()) {} |
65 | |
66 | explicit ExclusivelyOwned(T&& t) : repr_(EOT::moveToRepr(std::move(t))) {} |
67 | |
68 | template <class... Args> |
69 | explicit ExclusivelyOwned(in_place_t, Args&&... args) |
70 | : repr_(EOT::createInPlace(std::forward<Args>(args)...)) {} |
71 | |
72 | ExclusivelyOwned(const ExclusivelyOwned&) = delete; |
73 | |
74 | ExclusivelyOwned(ExclusivelyOwned&& rhs) noexcept |
75 | : repr_(std::move(rhs.repr_)) { |
76 | rhs.repr_ = EOT::nullRepr(); |
77 | } |
78 | |
79 | ExclusivelyOwned& operator=(const ExclusivelyOwned&) = delete; |
80 | |
81 | ExclusivelyOwned& operator=(ExclusivelyOwned&& rhs) noexcept { |
82 | EOT::destroyOwned(repr_); |
83 | repr_ = std::move(rhs.repr_); |
84 | rhs.repr_ = EOT::nullRepr(); |
85 | return *this; |
86 | } |
87 | |
88 | ExclusivelyOwned& operator=(T&& rhs) noexcept { |
89 | EOT::destroyOwned(repr_); |
90 | repr_ = EOT::moveToRepr(std::move(rhs)); |
91 | return *this; |
92 | } |
93 | |
94 | ~ExclusivelyOwned() { |
95 | EOT::destroyOwned(repr_); |
96 | // Don't bother to call the destructor of repr_, since we already |
97 | // did specialized destruction for the exclusively-owned case in |
98 | // destroyOwned! |
99 | } |
100 | |
101 | // We don't provide this because it would require us to be able to |
102 | // differentiate an owned-but-empty T from a lack of T. This is |
103 | // particularly problematic for Tensor, which wants to use an |
104 | // undefined Tensor as its null state. |
105 | explicit operator bool() const noexcept = delete; |
106 | |
107 | operator T() && { |
108 | return take(); |
109 | } |
110 | |
111 | // NOTE: the equivalent operation on MaybeOwned is a moving |
112 | // operator*. For ExclusivelyOwned, take() and operator*() may well |
113 | // have different return types, so they are different functions. |
114 | T take() && { |
115 | return EOT::take(repr_); |
116 | } |
117 | |
118 | typename EOT::const_pointer_type operator->() const { |
119 | return get(); |
120 | } |
121 | |
122 | typename EOT::const_pointer_type get() const { |
123 | return EOT::getImpl(repr_); |
124 | } |
125 | |
126 | typename EOT::pointer_type operator->() { |
127 | return get(); |
128 | } |
129 | |
130 | typename EOT::pointer_type get() { |
131 | return EOT::getImpl(repr_); |
132 | } |
133 | |
134 | std::remove_pointer_t<typename EOT::const_pointer_type>& operator*() const { |
135 | return *get(); |
136 | } |
137 | |
138 | std::remove_pointer_t<typename EOT::pointer_type>& operator*() { |
139 | return *get(); |
140 | } |
141 | }; |
142 | |
143 | } // namespace c10 |
144 | |