1#pragma once
2#include <memory>
3
4#include <c10/macros/Macros.h>
5
6namespace c10 {
7
8using DeleterFnPtr = void (*)(void*);
9
10namespace detail {
11
12// Does not delete anything
13C10_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//
38class 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
110inline bool operator==(const UniqueVoidPtr& sp, std::nullptr_t) noexcept {
111 return !sp;
112}
113inline bool operator==(std::nullptr_t, const UniqueVoidPtr& sp) noexcept {
114 return !sp;
115}
116inline bool operator!=(const UniqueVoidPtr& sp, std::nullptr_t) noexcept {
117 return sp;
118}
119inline bool operator!=(std::nullptr_t, const UniqueVoidPtr& sp) noexcept {
120 return sp;
121}
122
123} // namespace detail
124} // namespace c10
125