UniqueVoidPtr.h 4.5 KB

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