1 | #pragma once |
2 | #include <memory> |
3 | |
4 | #include <c10/macros/Macros.h> |
5 | |
6 | namespace c10 { |
7 | |
8 | using DeleterFnPtr = void (*)(void*); |
9 | |
10 | namespace detail { |
11 | |
12 | // Does not delete anything |
13 | C10_API void deleteNothing(void*); |
14 | |
15 | // A detail::UniqueVoidPtr is an owning smart pointer like unique_ptr, but |
16 | // with three major differences: |
17 | // |
18 | // 1) It is specialized to void |
19 | // |
20 | // 2) It is specialized for a function pointer deleter |
21 | // void(void* ctx); i.e., the deleter doesn't take a |
22 | // reference to the data, just to a context pointer |
23 | // (erased as void*). In fact, internally, this pointer |
24 | // is implemented as having an owning reference to |
25 | // context, and a non-owning reference to data; this is why |
26 | // you release_context(), not release() (the conventional |
27 | // API for release() wouldn't give you enough information |
28 | // to properly dispose of the object later.) |
29 | // |
30 | // 3) The deleter is guaranteed to be called when the unique |
31 | // pointer is destructed and the context is non-null; this is different |
32 | // from std::unique_ptr where the deleter is not called if the |
33 | // data pointer is null. |
34 | // |
35 | // Some of the methods have slightly different types than std::unique_ptr |
36 | // to reflect this. |
37 | // |
38 | class UniqueVoidPtr { |
39 | private: |
40 | // Lifetime tied to ctx_ |
41 | void* data_; |
42 | std::unique_ptr<void, DeleterFnPtr> ctx_; |
43 | |
44 | public: |
45 | UniqueVoidPtr() : data_(nullptr), ctx_(nullptr, &deleteNothing) {} |
46 | explicit UniqueVoidPtr(void* data) |
47 | : data_(data), ctx_(nullptr, &deleteNothing) {} |
48 | UniqueVoidPtr(void* data, void* ctx, DeleterFnPtr ctx_deleter) |
49 | : data_(data), ctx_(ctx, ctx_deleter ? ctx_deleter : &deleteNothing) {} |
50 | void* operator->() const { |
51 | return data_; |
52 | } |
53 | void clear() { |
54 | ctx_ = nullptr; |
55 | data_ = nullptr; |
56 | } |
57 | void* get() const { |
58 | return data_; |
59 | } |
60 | void* get_context() const { |
61 | return ctx_.get(); |
62 | } |
63 | void* release_context() { |
64 | return ctx_.release(); |
65 | } |
66 | std::unique_ptr<void, DeleterFnPtr>&& move_context() { |
67 | return std::move(ctx_); |
68 | } |
69 | C10_NODISCARD bool compare_exchange_deleter( |
70 | DeleterFnPtr expected_deleter, |
71 | DeleterFnPtr new_deleter) { |
72 | if (get_deleter() != expected_deleter) |
73 | return false; |
74 | ctx_ = std::unique_ptr<void, DeleterFnPtr>(ctx_.release(), new_deleter); |
75 | return true; |
76 | } |
77 | |
78 | template <typename T> |
79 | T* cast_context(DeleterFnPtr expected_deleter) const { |
80 | if (get_deleter() != expected_deleter) |
81 | return nullptr; |
82 | return static_cast<T*>(get_context()); |
83 | } |
84 | operator bool() const { |
85 | return data_ || ctx_; |
86 | } |
87 | DeleterFnPtr get_deleter() const { |
88 | return ctx_.get_deleter(); |
89 | } |
90 | }; |
91 | |
92 | // Note [How UniqueVoidPtr is implemented] |
93 | // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ |
94 | // UniqueVoidPtr solves a common problem for allocators of tensor data, which |
95 | // is that the data pointer (e.g., float*) which you are interested in, is not |
96 | // the same as the context pointer (e.g., DLManagedTensor) which you need |
97 | // to actually deallocate the data. Under a conventional deleter design, you |
98 | // have to store extra context in the deleter itself so that you can actually |
99 | // delete the right thing. Implementing this with standard C++ is somewhat |
100 | // error-prone: if you use a std::unique_ptr to manage tensors, the deleter will |
101 | // not be called if the data pointer is nullptr, which can cause a leak if the |
102 | // context pointer is non-null (and the deleter is responsible for freeing both |
103 | // the data pointer and the context pointer). |
104 | // |
105 | // So, in our reimplementation of unique_ptr, which just store the context |
106 | // directly in the unique pointer, and attach the deleter to the context |
107 | // pointer itself. In simple cases, the context pointer is just the pointer |
108 | // itself. |
109 | |
110 | inline bool operator==(const UniqueVoidPtr& sp, std::nullptr_t) noexcept { |
111 | return !sp; |
112 | } |
113 | inline bool operator==(std::nullptr_t, const UniqueVoidPtr& sp) noexcept { |
114 | return !sp; |
115 | } |
116 | inline bool operator!=(const UniqueVoidPtr& sp, std::nullptr_t) noexcept { |
117 | return sp; |
118 | } |
119 | inline bool operator!=(std::nullptr_t, const UniqueVoidPtr& sp) noexcept { |
120 | return sp; |
121 | } |
122 | |
123 | } // namespace detail |
124 | } // namespace c10 |
125 | |