ExclusivelyOwned.h 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. #pragma once
  2. #include <utility>
  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. typename ExclusivelyOwnedTraits<T>::repr_type repr_;
  55. public:
  56. ExclusivelyOwned() : repr_(EOT::nullRepr()) {}
  57. explicit ExclusivelyOwned(T&& t) : repr_(EOT::moveToRepr(std::move(t))) {}
  58. template <class... Args>
  59. explicit ExclusivelyOwned(std::in_place_t, Args&&... args)
  60. : repr_(EOT::createInPlace(std::forward<Args>(args)...)) {}
  61. ExclusivelyOwned(const ExclusivelyOwned&) = delete;
  62. ExclusivelyOwned(ExclusivelyOwned&& rhs) noexcept
  63. : repr_(std::move(rhs.repr_)) {
  64. rhs.repr_ = EOT::nullRepr();
  65. }
  66. ExclusivelyOwned& operator=(const ExclusivelyOwned&) = delete;
  67. ExclusivelyOwned& operator=(ExclusivelyOwned&& rhs) noexcept {
  68. EOT::destroyOwned(repr_);
  69. repr_ = std::move(rhs.repr_);
  70. rhs.repr_ = EOT::nullRepr();
  71. return *this;
  72. }
  73. ExclusivelyOwned& operator=(T&& rhs) noexcept {
  74. EOT::destroyOwned(repr_);
  75. repr_ = EOT::moveToRepr(std::move(rhs));
  76. return *this;
  77. }
  78. ~ExclusivelyOwned() {
  79. EOT::destroyOwned(repr_);
  80. // Don't bother to call the destructor of repr_, since we already
  81. // did specialized destruction for the exclusively-owned case in
  82. // destroyOwned!
  83. }
  84. // We don't provide this because it would require us to be able to
  85. // differentiate an owned-but-empty T from a lack of T. This is
  86. // particularly problematic for Tensor, which wants to use an
  87. // undefined Tensor as its null state.
  88. explicit operator bool() const noexcept = delete;
  89. operator T() && {
  90. return take();
  91. }
  92. // NOTE: the equivalent operation on MaybeOwned is a moving
  93. // operator*. For ExclusivelyOwned, take() and operator*() may well
  94. // have different return types, so they are different functions.
  95. T take() && {
  96. return EOT::take(repr_);
  97. }
  98. typename EOT::const_pointer_type operator->() const {
  99. return get();
  100. }
  101. typename EOT::const_pointer_type get() const {
  102. return EOT::getImpl(repr_);
  103. }
  104. typename EOT::pointer_type operator->() {
  105. return get();
  106. }
  107. typename EOT::pointer_type get() {
  108. return EOT::getImpl(repr_);
  109. }
  110. std::remove_pointer_t<typename EOT::const_pointer_type>& operator*() const {
  111. return *get();
  112. }
  113. std::remove_pointer_t<typename EOT::pointer_type>& operator*() {
  114. return *get();
  115. }
  116. };
  117. } // namespace c10