ExclusivelyOwned.h 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. #pragma once
  2. #include <c10/util/in_place.h>
  3. namespace c10 {
  4. // See example implementation in TensorBase.h and TensorBody.h.
  5. // Synopsis:
  6. //
  7. // repr_type -- type to use to store an owned T in ExclusivelyOwned.
  8. //
  9. // pointer_type -- pointer-esque type to return from
  10. // ExclusivelyOwned's get() and operator*() methods.
  11. //
  12. // const_pointer_type -- similar to pointer_type, used for the const methods.
  13. //
  14. // static repr_type nullRepr() -- return a null instance of repr_type.
  15. //
  16. // template <class... Args>
  17. // static repr_type createInPlace(Args&&... args) -- used by the in-place
  18. // ExclusivelyOwned constructor.
  19. //
  20. // static repr_type moveToRepr(T&& x) -- move the given x into an
  21. // instance of repr_type. used by the ExclusivelyOwned(T&&)
  22. // constructor.
  23. //
  24. // static void destroyOwned(repr_type x) -- free memory for a
  25. // known-exclusively-owned instance of x. Replaces calling repr_type's
  26. // destructor. Being able to implement this more efficiently than
  27. // repr_type's destructor is the main reason to use ExclusivelyOwned
  28. // for a type.
  29. //
  30. // static T take(repr_type&) -- move out of the given repr_type into an owned T.
  31. //
  32. // static pointer_type getImpl(const repr_type&) -- return a pointer
  33. // to the given repr_type. May take repr_type by value if that is more
  34. // efficient.
  35. template <typename T>
  36. struct ExclusivelyOwnedTraits;
  37. /// ExclusivelyOwned is a smart-pointer-like wrapper around an
  38. /// exclusively-owned instance of some type T that normally has
  39. /// mandatory reference counting (currently just Tensor). If you have
  40. /// an isolated piece of code that knows that it has sole ownership of
  41. /// an object of one of these types (i.e., because you created it
  42. /// directly or using a factory function) and that object will not
  43. /// escape from that isolated piece of code, then moving the object
  44. /// into an ExclusivelyOwned will avoid an atomic reference count
  45. /// decrement at destruction time.
  46. ///
  47. /// If you directly create the Tensor in the first
  48. /// place, you can use the in_place constructor of ExclusivelyOwned to
  49. /// additionally avoid doing any stores to initialize the refcount &
  50. /// weakcount.
  51. template <typename T>
  52. class ExclusivelyOwned {
  53. using EOT = ExclusivelyOwnedTraits<T>;
  54. union {
  55. char dummy_;
  56. typename ExclusivelyOwnedTraits<T>::repr_type repr_;
  57. };
  58. public:
  59. ExclusivelyOwned() : repr_(EOT::nullRepr()) {}
  60. explicit ExclusivelyOwned(T&& t) : repr_(EOT::moveToRepr(std::move(t))) {}
  61. template <class... Args>
  62. explicit ExclusivelyOwned(in_place_t, Args&&... args)
  63. : repr_(EOT::createInPlace(std::forward<Args>(args)...)) {}
  64. ExclusivelyOwned(const ExclusivelyOwned&) = delete;
  65. ExclusivelyOwned(ExclusivelyOwned&& rhs) noexcept
  66. : repr_(std::move(rhs.repr_)) {
  67. rhs.repr_ = EOT::nullRepr();
  68. }
  69. ExclusivelyOwned& operator=(const ExclusivelyOwned&) = delete;
  70. ExclusivelyOwned& operator=(ExclusivelyOwned&& rhs) noexcept {
  71. EOT::destroyOwned(repr_);
  72. repr_ = std::move(rhs.repr_);
  73. rhs.repr_ = EOT::nullRepr();
  74. return *this;
  75. }
  76. ExclusivelyOwned& operator=(T&& rhs) noexcept {
  77. EOT::destroyOwned(repr_);
  78. repr_ = EOT::moveToRepr(std::move(rhs));
  79. return *this;
  80. }
  81. ~ExclusivelyOwned() {
  82. EOT::destroyOwned(repr_);
  83. // Don't bother to call the destructor of repr_, since we already
  84. // did specialized destruction for the exclusively-owned case in
  85. // destroyOwned!
  86. }
  87. // We don't provide this because it would require us to be able to
  88. // differentiate an owned-but-empty T from a lack of T. This is
  89. // particularly problematic for Tensor, which wants to use an
  90. // undefined Tensor as its null state.
  91. explicit operator bool() const noexcept = delete;
  92. operator T() && {
  93. return take();
  94. }
  95. // NOTE: the equivalent operation on MaybeOwned is a moving
  96. // operator*. For ExclusivelyOwned, take() and operator*() may well
  97. // have different return types, so they are different functions.
  98. T take() && {
  99. return EOT::take(repr_);
  100. }
  101. typename EOT::const_pointer_type operator->() const {
  102. return get();
  103. }
  104. typename EOT::const_pointer_type get() const {
  105. return EOT::getImpl(repr_);
  106. }
  107. typename EOT::pointer_type operator->() {
  108. return get();
  109. }
  110. typename EOT::pointer_type get() {
  111. return EOT::getImpl(repr_);
  112. }
  113. std::remove_pointer_t<typename EOT::const_pointer_type>& operator*() const {
  114. return *get();
  115. }
  116. std::remove_pointer_t<typename EOT::pointer_type>& operator*() {
  117. return *get();
  118. }
  119. };
  120. } // namespace c10