1#pragma once
2
3#include <c10/util/in_place.h>
4
5namespace 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.
38template <typename T>
39struct 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.
55template <typename T>
56class 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