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