ExclusivelyOwned.h 4.6 KB

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